Draft (new feature) : Model to estimate when a intervention had effect

causalpy/pymc_models.py

@@ -497,3 +497,147 @@
                 )
             )
         return self.idata
+
+
+class InterventionTimeEstimator(PyMCModel):
+    r"""
+    Custom PyMC model to estimate the time an intervetnion took place.
+
+    defines the PyMC model :
+
+    .. math::
+        \alpha &\sim \mathrm{Normal}(0, 1) \\
+        \beta &\sim \mathrm{Normal}(0, 1) \\
+        s(t) &= \gamma_{i(t)} \quad \textrm{with} \quad \gamma_{k \in [0, ..., n_{seasons}-1]} \sim \mathrm{Normal}(0, 1)\\
+        base_{\mu}(t) &= \alpha + \beta \cdot t + s_t\\
+        \\
+        \tau &\sim \mathrm{Uniform}(0, 1) \\
+        w(t) &= sigmoid(t-\tau) \\
+        \\
+        level &\sim \mathrm{Normal}(0, 1) \\
+        trend &\sim \mathrm{Normal}(0, 1) \\
+        A &\sim \mathrm{Normal}(0, 1) \\
+        \lambda &\sim \mathrm{HalfNormal}(0, 1) \\
+        impulse(t) &= A \cdot exp(-\lambda \cdot |t-\tau|) \\
+        intervention(t) &= level + trend \cdot (t-\tau) + impulse_t\\
+        \\
+        \sigma &\sim \mathrm{Normal}(0, 1) \\
+        \mu(t) &= base_{\mu}(t) + w(t) \cdot intervention(t) \\
+        \\
+        y(t) &\sim \mathrm{Normal}(\mu (t), \sigma)
+
+    Example
+        --------
+        >>> import causalpy as cp
+        >>> import numpy as np
+        >>> from causalpy.pymc_models import InterventionTimeEstimator
+        >>> df = cp.load_data("its")
+        >>> y = df["y"].values
+        >>> t = df["t"].values
+        >>> coords = {"seasons": range(12)} # The data is monthly
+        >>> estimator = InterventionTimeEstimator()
+        >>> # We are trying to capture an impulse in the number of death per month due to Covid.
+        >>> estimator.fit(
+        ...     t,
+        ...     y,
+        ...     coords,
+        ...     priors={"impulse":[]}
+        ... )
+        Inference data...
+    """
+
+    def build_model(self, t, y, coords, time_range, grain_season, priors):
+        """
+        Defines the PyMC model
+
+        :param t: An array of values representing the time over which y is spread
+        :param y: An array of values representing our outcome y
+        :param coords: An optional dictionary with the coordinate names for our instruments.
+            In particular, used to determine the number of seasons.
+        :param time_range: An optional tuple providing a specific time_range where the
+            intervention effect should have taken place.
+        :param priors: An optional dictionary of priors for the parameters of the
+            different distributions.
+            :code:`priors = {"alpha":[0, 5], "beta":[0,2], "level":[5, 5], "impulse":[1, 2 ,3]}`
+        """
+
+        with self:
+            self.add_coords(coords)
+
+            if time_range is None:
+                time_range = (t.min(), t.max())
+
+            # --- Priors ---
+            switchpoint = pm.Uniform(
+                "switchpoint", lower=time_range[0], upper=time_range[1]
+            )
+            alpha = pm.Normal(name="alpha", mu=0, sigma=50)
+            beta = pm.Normal(name="beta", mu=0, sigma=50)
+            seasons = 0
+            if "seasons" in coords and len(coords["seasons"]) > 0:
+                season_idx = np.arange(len(y)) // grain_season % len(coords["seasons"])
+                seasons_effect = pm.Normal(
+                    "seasons_effect", mu=0, sigma=50, dims="seasons"
+                )
+                seasons = seasons_effect[season_idx]
+
+            # --- Intervention effect ---
+            level = trend = impulse = 0
+
+            if "level" in priors:
+                mu, sigma = (
+                    (0, 50)
+                    if len(priors["level"]) != 2
+                    else (priors["level"][0], priors["level"][1])
+                )
+                level = pm.Normal(
+                    "level",
+                    mu=mu,
+                    sigma=sigma,
+                )
+            if "trend" in priors:
+                mu, sigma = (
+                    (0, 50)
+                    if len(priors["trend"]) != 2
+                    else (priors["trend"][0], priors["trend"][1])
+                )
+                trend = pm.Normal("trend", mu=mu, sigma=sigma)
+            if "impulse" in priors:
+                mu, sigma1, sigma2 = (
+                    (0, 50, 50)
+                    if len(priors["impulse"]) != 3
+                    else (
+                        priors["impulse"][0],
+                        priors["impulse"][1],
+                        priors["impulse"][2],
+                    )
+                )
+                impulse_amplitude = pm.Normal("impulse_amplitude", mu=mu, sigma=sigma1)
+                decay_rate = pm.HalfNormal("decay_rate", sigma=sigma2)
+                impulse = impulse_amplitude * pm.math.exp(
+                    -decay_rate * abs(t - switchpoint)
+                )
+
+            # --- Parameterization ---
+            weight = pm.math.sigmoid(t - switchpoint)
+            # Compute and store the modelled time series
+            mu_ts = pm.Deterministic(name="mu_ts", var=alpha + beta * t + seasons)
+            # Compute and store the modelled intervention effect
+            mu_in = pm.Deterministic(
+                name="mu_in", var=level + trend * (t - switchpoint) + impulse
+            )
+            # Compute and store the the sum of the intervention and the time series
+            mu = pm.Deterministic("mu", mu_ts + weight * mu_in)
+            sigma = pm.HalfNormal("sigma", 1)
+
+            # --- Likelihood ---
+            pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y)
+
+    def fit(self, t, y, coords, time_range=None, grain_season=1, priors={}, n=1000):
+        """
+        Draw samples from posterior distribution
+        """
+        self.build_model(t, y, coords, time_range, grain_season, priors)
+        with self:
+            self.idata = pm.sample(n, progressbar=False, **self.sample_kwargs)
+        return self.idata