Add PiecewiseITS experiment for known interruption dates (#614)

* initial commit - MVP

* Expand Piecewise ITS notebook with new examples and explanations

Added detailed explanations comparing Piecewise ITS to Regression Discontinuity and Regression Kink designs. Introduced new real-world scenarios for level and slope changes, multiple interventions, and level-only models. Enhanced example code and output to illustrate these cases, improving clarity and practical guidance for users.

* Revise and clarify Piecewise ITS notebook explanations

Improved clarity and conciseness throughout the Piecewise Interrupted Time Series (ITS) notebook. Rewrote several sections for better readability, combined and streamlined example scenarios, and clarified distinctions between level and slope changes, as well as the relationship to regression discontinuity and regression kink designs.

* Refactor PiecewiseITS to use patsy step/ramp stateful transforms

Refactors the PiecewiseITS experiment to use flexible patsy formulas with new stateful step() and ramp() transforms for specifying level and slope changes at interventions. Adds the causalpy.transforms module with robust, datetime-aware step/ramp transforms, updates tests to cover new formula interface and transform behavior, and improves documentation and error handling. This enables more flexible modeling of multiple interventions and supports both numeric and datetime time columns.

* Expand PiecewiseITS notebook with formula API details

Added a new section describing the formula-based API for PiecewiseITS, including explanations of the custom step() and ramp() transforms, usage examples, and clarification on how the counterfactual is computed. This improves documentation clarity and helps users understand flexible model specification.

* Add effect_summary compatibility to PiecewiseITS

Implemented creation of post_impact, datapost, and post_pred attributes in PiecewiseITS for compatibility with effect_summary() from BaseExperiment. Added tests to verify effect_summary works for both OLS and PyMC models and that the new attributes are correctly created.

* Clarify usage of step and ramp transforms in docs

* Enhance piecewise ITS notebook with math and intro code

Added mathematical definitions for step and ramp functions using LaTeX for clarity, and moved import/setup code to the top of the notebook for better organization. Improved explanations of function arguments and removed duplicate import cell.

* Expand and clarify Piecewise ITS notebook introduction

The introductory markdown in the piecewise_its_pymc.ipynb notebook has been significantly expanded and reorganized. The new content provides clearer explanations of when to use Piecewise ITS, the distinction between level and slope changes, the mathematical model, and its relationship to regression discontinuity and regression kink designs. Redundant sections were removed and a more structured, didactic flow was introduced.

* Clarify level and slope change concepts in notebook

Expanded explanations of level and slope changes in piecewise ITS, referencing a new illustrative figure. Added a code cell to display the figure, and clarified the description of multiple interventions for improved instructional clarity.

* Add model formula table to piecewise ITS notebook

Inserted a markdown cell with a table summarizing model formulas for single and two intervention cases, covering level, slope, and combined effects. This provides clearer guidance on specifying models for each panel in the notebook.

* Revise and restructure piecewise ITS notebook intro

Condenses and reorganizes introductory explanations for piecewise interrupted time series (ITS), splitting out key concepts, model details, and comparisons to related methods into clearer, more focused sections. Adds collapsible dropdowns and card formatting for scenario examples, and improves clarity and flow for users learning the model and its API.

* Add extensive coverage tests for PiecewiseITS

Adds a comprehensive suite of tests for the PiecewiseITS class, including class and instance attribute checks, formula parsing, plotting, PyMC integration, counterfactuals, data generation, and error handling. Also updates the interrogate badge to reflect increased coverage.

* Add references and citations for piecewise ITS notebook

Added detailed references and in-text citations to the piecewise_its_pymc.ipynb notebook to support methodological explanations. Updated the references.bib file with key literature on segmented regression and interrupted time series analysis. Improved clarity on model parameterization and corrected the references section to use the Sphinx bibliography directive.

* run pre-commit checks

* Fail fast on PiecewiseITS threshold parsing.

Enforce a single step/ramp time variable and canonicalize interruption thresholds at initialization so invalid mixed-variable or malformed thresholds fail early with clear errors.

Co-authored-by: Cursor <cursoragent@cursor.com>

* Implement effect_summary for PiecewiseITS.

Add PiecewiseITS effect_summary support for both OLS and PyMC paths using shared reporting helpers so PiecewiseITS no longer falls back to BaseExperiment's NotImplementedError.

Co-authored-by: Cursor <cursoragent@cursor.com>

* Increase PiecewiseITS patch coverage for effect summary.

Add a focused test for unsupported period handling in PiecewiseITS effect_summary and simplify OLS array conversion logic to keep the new effect summary path fully exercised.

Co-authored-by: Cursor <cursoragent@cursor.com>

* Increase patch coverage for transforms and PiecewiseITS validation

Add tests for: datetime Series (non-Index) paths, chunked memorize_chunk
origin tracking, pd.Timestamp threshold parsing, step/ramp variable not
in data, and non-numeric/non-datetime time column validation.

Made-with: Cursor

* Add _default_model_class and minor plot fix for BaseExperiment compatibility

After rebase onto main, PiecewiseITS needs _default_model_class = LinearRegression
to comply with the enforced BaseExperiment abstract methods (#694). Also simplify
ax[1 + 1] to ax[2].

Made-with: Cursor

---------

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: drbenvincent

.pre-commit-config.yaml

@@ -24,7 +24,7 @@ repos:
       - id: check-toml
       - id: check-json
       - id: check-added-large-files
-        exclude: &exclude_pattern '(iv_weak_instruments|its_lift_test)\.ipynb'
+        exclude: &exclude_pattern '(iv_weak_instruments|its_lift_test|piecewise_its_pymc)\.ipynb'
         args: ["--maxkb=1500"]
       - id: check-merge-conflict
       - id: check-case-conflict

causalpy/__init__.py

@@ -16,13 +16,15 @@
 import causalpy.skl_models as skl_models
 import causalpy.variable_selection_priors as variable_selection_priors
 from causalpy.skl_models import create_causalpy_compatible_class
+from causalpy.transforms import ramp, step
 from causalpy.version import __version__
 
 from .data import load_data
 from .experiments.diff_in_diff import DifferenceInDifferences
 from .experiments.instrumental_variable import InstrumentalVariable
 from .experiments.interrupted_time_series import InterruptedTimeSeries
 from .experiments.inverse_propensity_weighting import InversePropensityWeighting
+from .experiments.piecewise_its import PiecewiseITS
 from .experiments.prepostnegd import PrePostNEGD
 from .experiments.regression_discontinuity import RegressionDiscontinuity
 from .experiments.regression_kink import RegressionKink
@@ -32,19 +34,22 @@
 
 __all__ = [
     "__version__",
-    "DifferenceInDifferences",
     "create_causalpy_compatible_class",
+    "DifferenceInDifferences",
     "extract_lift_for_mmm",
     "InstrumentalVariable",
     "InterruptedTimeSeries",
     "InversePropensityWeighting",
     "load_data",
+    "PiecewiseITS",
     "PrePostNEGD",
     "pymc_models",
+    "ramp",
     "RegressionDiscontinuity",
     "RegressionKink",
     "skl_models",
     "StaggeredDifferenceInDifferences",
+    "step",
     "SyntheticControl",
     "variable_selection_priors",
 ]

causalpy/data/simulate_data.py

@@ -645,3 +645,180 @@ def generate_staggered_did_data(
 
     df = pd.DataFrame(rows)
     return df
+
+
+def generate_piecewise_its_data(
+    N: int = 100,
+    interruption_times: list[int] | None = None,
+    baseline_intercept: float = 10.0,
+    baseline_slope: float = 0.1,
+    level_changes: list[float] | None = None,
+    slope_changes: list[float] | None = None,
+    noise_sigma: float = 1.0,
+    seed: int | None = None,
+) -> tuple[pd.DataFrame, dict]:
+    """
+    Generate piecewise Interrupted Time Series data with known ground truth parameters.
+
+    This function creates synthetic data for testing and demonstrating piecewise ITS
+    / segmented regression models. The data follows the model:
+
+    y_t = β₀ + β₁t + Σₖ(level_k · I_k(t) + slope_k · R_k(t)) + ε_t
+
+    Where:
+    - I_k(t) = 1 if t >= T_k else 0 (step function for level change)
+    - R_k(t) = max(0, t - T_k) (ramp function for slope change)
+
+    Parameters
+    ----------
+    N : int, default=100
+        Number of time points in the series.
+    interruption_times : list[int], optional
+        List of time indices where interruptions occur. Defaults to [50].
+    baseline_intercept : float, default=10.0
+        The intercept (β₀) of the baseline trend.
+    baseline_slope : float, default=0.1
+        The slope (β₁) of the baseline trend.
+    level_changes : list[float], optional
+        List of level changes at each interruption. Length must match
+        interruption_times. If None, defaults to [5.0] for single interruption.
+    slope_changes : list[float], optional
+        List of slope changes at each interruption. Length must match
+        interruption_times. If None, defaults to [0.0] (no slope change).
+    noise_sigma : float, default=1.0
+        Standard deviation of the Gaussian noise.
+    seed : int, optional
+        Random seed for reproducibility.
+
+    Returns
+    -------
+    df : pd.DataFrame
+        DataFrame with columns:
+        - 't': time index (0 to N-1)
+        - 'y': observed outcome with noise
+        - 'y_true': outcome without noise (ground truth)
+        - 'counterfactual': baseline trend without intervention effects
+        - 'effect': true causal effect at each time point
+    params : dict
+        Dictionary containing the true parameters:
+        - 'baseline_intercept': β₀
+        - 'baseline_slope': β₁
+        - 'level_changes': list of level changes
+        - 'slope_changes': list of slope changes
+        - 'interruption_times': list of interruption times
+        - 'noise_sigma': noise standard deviation
+
+    Examples
+    --------
+    >>> from causalpy.data.simulate_data import generate_piecewise_its_data
+    >>> # Single interruption with level and slope change
+    >>> df, params = generate_piecewise_its_data(
+    ...     N=100,
+    ...     interruption_times=[50],
+    ...     level_changes=[5.0],
+    ...     slope_changes=[0.2],
+    ...     seed=42,
+    ... )
+    >>> df.shape
+    (100, 5)
+
+    >>> # Multiple interruptions
+    >>> df, params = generate_piecewise_its_data(
+    ...     N=150,
+    ...     interruption_times=[50, 100],
+    ...     level_changes=[3.0, -2.0],
+    ...     slope_changes=[0.1, -0.15],
+    ...     seed=42,
+    ... )
+    >>> len(params["interruption_times"])
+    2
+
+    >>> # Level change only (no slope change)
+    >>> df, params = generate_piecewise_its_data(
+    ...     N=100,
+    ...     interruption_times=[50],
+    ...     level_changes=[5.0],
+    ...     slope_changes=[0.0],
+    ...     seed=42,
+    ... )
+    """
+    # Set defaults
+    if interruption_times is None:
+        interruption_times = [50]
+
+    n_interruptions = len(interruption_times)
+
+    if level_changes is None:
+        level_changes = [5.0] * n_interruptions
+
+    if slope_changes is None:
+        slope_changes = [0.0] * n_interruptions
+
+    # Validate inputs
+    if len(level_changes) != n_interruptions:
+        raise ValueError(
+            f"level_changes length ({len(level_changes)}) must match "
+            f"interruption_times length ({n_interruptions})"
+        )
+
+    if len(slope_changes) != n_interruptions:
+        raise ValueError(
+            f"slope_changes length ({len(slope_changes)}) must match "
+            f"interruption_times length ({n_interruptions})"
+        )
+
+    for t_k in interruption_times:
+        if t_k < 0 or t_k >= N:
+            raise ValueError(
+                f"Interruption time {t_k} is outside valid range [0, {N - 1}]"
+            )
+
+    # Set random seed
+    if seed is not None:
+        np.random.seed(seed)
+
+    # Generate time index
+    t = np.arange(N)
+
+    # Compute baseline (counterfactual)
+    counterfactual = baseline_intercept + baseline_slope * t
+
+    # Compute intervention effects
+    effect = np.zeros(N)
+    for k, t_k in enumerate(interruption_times):
+        # Step function: I_k(t) = 1 if t >= t_k
+        step = (t >= t_k).astype(float)
+        # Ramp function: R_k(t) = max(0, t - t_k)
+        ramp = np.maximum(0, t - t_k).astype(float)
+
+        effect += level_changes[k] * step + slope_changes[k] * ramp
+
+    # Compute true outcome (without noise)
+    y_true = counterfactual + effect
+
+    # Add noise
+    noise = np.random.normal(0, noise_sigma, N)
+    y = y_true + noise
+
+    # Create DataFrame
+    df = pd.DataFrame(
+        {
+            "t": t,
+            "y": y,
+            "y_true": y_true,
+            "counterfactual": counterfactual,
+            "effect": effect,
+        }
+    )
+
+    # Store parameters
+    params = {
+        "baseline_intercept": baseline_intercept,
+        "baseline_slope": baseline_slope,
+        "level_changes": level_changes,
+        "slope_changes": slope_changes,
+        "interruption_times": interruption_times,
+        "noise_sigma": noise_sigma,
+    }
+
+    return df, params

causalpy/experiments/__init__.py

@@ -17,6 +17,7 @@
 from .instrumental_variable import InstrumentalVariable
 from .interrupted_time_series import InterruptedTimeSeries
 from .inverse_propensity_weighting import InversePropensityWeighting
+from .piecewise_its import PiecewiseITS
 from .prepostnegd import PrePostNEGD
 from .regression_discontinuity import RegressionDiscontinuity
 from .regression_kink import RegressionKink
@@ -26,11 +27,12 @@
 __all__ = [
     "DifferenceInDifferences",
     "InstrumentalVariable",
+    "InterruptedTimeSeries",
     "InversePropensityWeighting",
+    "PiecewiseITS",
     "PrePostNEGD",
     "RegressionDiscontinuity",
     "RegressionKink",
     "StaggeredDifferenceInDifferences",
     "SyntheticControl",
-    "InterruptedTimeSeries",
 ]