Feature Request: Hierarchical Difference-in-Differences with Random Effects

Feature Request: Hierarchical Difference-in-Differences with Random Effects

Summary

Add support for Hierarchical/Multilevel Difference-in-Differences models where individuals are nested within groups (e.g., students in schools, patients in clinics, workers in firms). This extends the existing DifferenceInDifferences class to handle clustered data structures with group-level random effects, using formula-based specification via the Formulaic package.

Motivation

Many real-world DiD applications involve nested data structures:

• Healthcare: Patients nested within clinics/hospitals
• Education: Students nested within schools/classrooms
• Labor economics: Workers nested within firms
• Policy evaluation: Individuals nested within geographic regions
• Marketing: Customers nested within stores/markets during promotional campaigns
• Sales: Sales reps nested within territories/regions for training interventions
• Retail: Transactions nested within retail locations for pricing experiments
• E-commerce: Users nested within marketing cohorts/campaigns for A/B tests
• SaaS: Accounts nested within customer success managers for onboarding experiments
• Advertising: Impressions nested within ad campaigns/platforms for campaign effectiveness

Currently, CausalPy's DifferenceInDifferences class does not explicitly model group-level heterogeneity through random effects. While users can include group indicators as fixed effects, this approach:

1. Ignores within-group correlation (leading to underestimated standard errors)
2. Doesn't pool information across groups (less efficient estimation)
3. Cannot estimate group-level treatment effect heterogeneity
4. Scales poorly with many groups (one parameter per group)

Detailed Example: Marketing Campaign Evaluation

This will be the basis of a new docs page (.ipynb notebook).

Scenario

A national retail chain with 200 stores wants to evaluate the effectiveness of a new customer loyalty program. The program is rolled out to 100 randomly selected stores (treatment) while the other 100 continue with the standard rewards program (control). All treated stores receive the program simultaneously (standard 2x2 DiD design, not staggered). The company tracks individual customer purchases over 6 months before and 6 months after the program launch.

Data Structure:

• Unit of observation: Individual customer transactions (500 customers per store on average = 100,000 total customers)
• Grouping: Customers nested within stores
• Time periods: Pre-intervention (months 1-6) and post-intervention (months 7-12)
• Outcome: Monthly purchase amount per customer
• Covariates: Customer age, tenure with brand, pre-intervention purchase history, store characteristics (size, region, urban/suburban)

Key Challenges:

1. Within-store correlation: Customers shopping at the same store are more similar to each other than to customers at other stores due to:

   • Shared local economic conditions
   • Similar demographic composition of the neighborhood
   • Store-specific factors (layout, staff quality, product mix)

2. Store-level heterogeneity: Different stores have different baseline sales levels and may respond differently to the loyalty program due to:

   • Varying competitive landscapes
   • Different customer demographics
   • Manager quality and implementation fidelity

3. Large number of groups: With 200 stores, including store fixed effects would add 199 parameters, reducing efficiency and making group-level inferences difficult

Why Regular DiD Falls Short

Standard DiD Analysis (Current CausalPy Approach):

# Treats all customer observations as independent
result_standard = cp.DifferenceInDifferences(
    df,
    formula="purchase_amount ~ 1 + post*treated + customer_age + customer_tenure",
    time_variable_name="month",
    group_variable_name="treated",
    model=cp.pymc_models.LinearRegression()
)

Problems with this approach:

1. Underestimated standard errors: Ignores that customers within the same store are correlated, leading to artificially narrow confidence intervals and inflated statistical significance
2. No store-level insights: Cannot answer "Which stores benefited most from the program?"
3. Inefficient estimation: Doesn't borrow strength across stores through partial pooling
4. Overly optimistic inference: Type I error rates exceed nominal levels due to clustering

What the Demo Notebook Will Show

The example notebook will demonstrate:

1. Data generation with realistic within-store correlation structure (ICC ≈ 0.15-0.25, typical for retail)

2. Side-by-side comparison of three approaches:

   • Naive DiD (ignoring clustering) → overconfident, narrow CIs
   • DiD with store fixed effects (200 parameters) → computationally expensive, no pooling
   • Hierarchical DiD (this implementation) → proper uncertainty, group-level insights

3. Empirical demonstration showing:

   • How naive DiD underestimates standard errors by 30-50% (typical with moderate ICC)
   • Coverage properties: 94% HDI from hierarchical model vs. <80% coverage from naive model
   • Efficiency gains from partial pooling (especially for stores with fewer customers)

4. Business insights enabled by hierarchical model:

   • Forest plot of store-specific effects: identify top/bottom performers
   • ICC calculation: "20% of variance in purchase amounts is between stores"
   • Heterogeneity analysis: correlate store random effects with store characteristics (size, region)
   • Targeted recommendations: "Stores in urban areas with high foot traffic showed 2x larger treatment effects"

This example will serve as a template for practitioners analyzing cluster-randomized or observational DiD studies in marketing, sales, healthcare, education, and other fields where clustering is inherent to the research design.

Future Extensions for Documentation: Sensitivity analysis examining ICC impact, parameterization comparisons, and sample size trade-offs can be added in follow-up notebooks or a dedicated sensitivity analysis section.

Proposed Solution

Implement a HierarchicalDifferenceInDifferences class that extends the standard DiD model with group-level random effects, specified directly in the formula using lme4-style syntax.

Note: This will be a Bayesian-only implementation using PyMC. No OLS/frequentist support will be provided for hierarchical DiD, as:

• Random effects estimation requires hierarchical priors (natural in Bayesian framework)
• Full uncertainty quantification for group-level effects
• Better handling of small group sizes through partial pooling
• Consistent with CausalPy's Bayesian modeling capabilities

Model Structure

For individual $i$ in group $j$ at time $t$:

$$Y_{ijt} = \alpha + \alpha_j + \beta \cdot \text{post}_t + \gamma \cdot \text{treated}_j + \delta \cdot (\text{post}t \times \text{treated}j) + \mathbf{X}{ijt}\boldsymbol{\beta} + \epsilon{ijt}$$

Where:

• $Y_{ijt}$ is the outcome for individual $i$ in group $j$ at time $t$
• $\alpha$ is the global intercept (fixed effect)
• $\alpha_j \sim N(0, \sigma_{\alpha}^2)$ are group-level random intercepts (deviation from global mean)
• $\beta$, $\gamma$ are fixed effect coefficients for time and treatment indicators
• $\delta$ is the Average Treatment Effect on the Treated (ATT) (coefficient on interaction term)
• $\mathbf{X}_{ijt}$ is a vector of individual-level covariates (e.g., age, gender, prior behavior)
• $\boldsymbol{\beta}$ is a vector of fixed effect coefficients for the covariates
• $\sigma_{\alpha}$ captures between-group variability (standard deviation of group effects)
• $\epsilon_{ijt} \sim N(0, \sigma^2)$ is individual-level error

Initial Implementation Scope (v1.0):

The first version will support:

• ✅ Random intercepts: (1|group) - group-specific baseline differences
• ✅ Random slopes: (x|group) or (1 + x|group) - group-specific effects of covariates
• ✅ Single grouping variable: One level of nesting

Unsupported in v1.0 (will raise clear exceptions with guidance):

• ❌ Multiple grouping variables: (1|group1) + (1|group2) (crossed random effects)
• ❌ Nested grouping: (1|school/classroom) (deeply nested structures)
• ❌ Complex interactions in random effects: (x:z|group)

Future Extensions (subsequent releases):

• Nested random effects for multi-level hierarchies
• Crossed random effects for non-nested structures
• Correlated vs. uncorrelated random effects specifications
• Hierarchical staggered DiD: Combine hierarchical modeling with staggered treatment adoption (extension of both HierarchicalDifferenceInDifferences and StaggeredDifferenceInDifferences)
• Sensitivity analysis documentation: Additional notebook or section exploring ICC impact on inference, parameterization comparisons, and sample size trade-offs

Data Validation and Requirements

The HierarchicalDifferenceInDifferences class will perform comprehensive data validation on initialization and raise informative exceptions when assumptions are violated.

Core Requirements:

1. Balanced Panel Structure

   • All units must be observed in all time periods (pre and post)
   • Exception: BadIndexException with message: "Unbalanced panel detected. Unit {id} has {n} observations but expected {expected}. Hierarchical DiD requires balanced panels where all units are observed in all time periods."
   • Validation: Check that data.groupby(unit_variable_name).size() is constant

2. No Missing Outcome Values

   • The outcome variable must not have any missing values
   • Exception: DataException with message: "Outcome variable contains {n_missing} missing values. Please remove or impute missing values before fitting."
   • Validation: data[outcome].isna().sum() == 0

3. Group Membership Integrity

   • Each unit must belong to exactly one group throughout the panel
   • No units can switch groups between time periods
   • Exception: DataException with message: "Unit {id} belongs to multiple groups: {groups}. Each unit must belong to a single group in hierarchical DiD."
   • Validation: Check that data.groupby(unit_variable_name)[group_variable].nunique() == 1 for all units

4. Treatment Assignment at Group Level

   • Treatment status must be constant within groups (all units in a group have same treatment)
   • No individual-level variation in treatment within groups
   • Exception: DataException with message: "Group {id} has mixed treatment status. Treatment must be assigned at the group level (all units within a group must have the same treatment status)."
   • Validation: data.groupby(group_variable)['treated'].nunique() == 1 for all groups

5. Sufficient Observations Per Group

   • Each group must have a minimum number of observations (default: ≥ 2)
   • Warning if any group has < 5 observations (low information for random effects)
   • Exception: DataException with message: "Group {id} has only {n} observations. Minimum of {min_obs} required per group."
   • Validation: data.groupby(group_variable).size() >= min_obs_per_group

6. Clear Pre/Post Period Structure

   • The time variable must clearly separate into pre and post periods
   • Exception: DataException with message: "Cannot identify pre/post periods. Ensure 'post_treatment' or equivalent indicator is correctly specified."
   • Validation: Check that time/treatment indicator creates non-empty pre and post periods

7. No Staggered Treatment Adoption

   • All treated groups must receive treatment at the same time (standard 2x2 DiD)
   • No variation in treatment timing across groups
   • Exception: DataException with message: "Staggered treatment adoption detected. Groups receive treatment at different times. Hierarchical DiD v1.0 only supports simultaneous treatment adoption (standard 2x2 design). For staggered adoption, use StaggeredDifferenceInDifferences or wait for future hierarchical staggered DiD support."
   • Validation: Check that all treated groups have the same treatment start time
   • Implementation: For each treated group, identify when treatment begins; ensure all groups have the same start time

8. Valid Random Effects Formula

   • Random effects syntax must be supported in v1.0 (single grouping variable, no nesting/crossing)
   • Exception: NotImplementedError (already implemented in _validate_random_effects())
   • See "Initial Implementation Scope" above for supported syntax

9. Appropriate Data Types

   • Grouping variables must be categorical or coercible to categorical
   • Time variables must be numeric or datetime
   • Outcome must be numeric
   • Exception: DataException with message: "Variable {var} has invalid type {type}. Expected {expected_type}."

10. Sufficient Groups for Random Effects

    • Minimum of 5 groups recommended (3 absolute minimum)
    • Warning if < 10 groups: "Only {n} groups detected. Random effects estimates may be unstable with few groups. Consider at least 10 groups for reliable hierarchical modeling."
    • Exception if < 3 groups: DataException with message: "Only {n} groups detected. Hierarchical models require at least 3 groups for random effects estimation. With fewer groups, use standard DiD with fixed effects."

11. No Perfect Multicollinearity

    • Check design matrix rank after formula parsing
    • Exception: FormulaException with message: "Design matrix is rank deficient. Check for perfect multicollinearity in covariates."

Implementation will follow patterns from existing experiment classes, using DataException, FormulaException, and BadIndexException with informative error messages.

API Design: Formula-Based Specification

Key Design Decision: Use Formulaic Instead of Separate Arguments

Rather than a separate random_effects argument, we integrate random effects into the formula using lme4/R-style syntax via the Formulaic package.

Why Formulaic?

• ✅ Supports mixed-effects formula syntax: (1|group), (1+x|group), etc.
• ✅ Familiar to R users (lme4 compatibility)
• ✅ Drop-in replacement for Patsy with minimal API changes
• ✅ Active maintenance, well-documented, type-hinted
• ✅ Lightweight (~500KB), pure Python
• ✅ Designed specifically for mixed-effects models

Example Usage

import causalpy as cp
import pandas as pd
import numpy as np

# Individual-level panel data with group membership
df = pd.DataFrame({
    'patient_id': range(1000),
    'clinic_id': np.repeat(range(50), 20),  # 50 clinics, 20 patients each
    'time': np.tile([0, 1], 500),
    'outcome': np.random.randn(1000),
    'post_treatment': np.tile([0, 1], 500),
    'treated': np.repeat([0]*500 + [1]*500, 1),
    'age': np.random.randint(20, 80, 1000),
})

# Random intercepts (clinic-level baseline differences)
# Non-centered parameterization is used by default for better sampling
result = cp.HierarchicalDifferenceInDifferences(
    df,
    formula="outcome ~ 1 + post_treatment*treated + age + (1|clinic_id)",
    #                                                      └── Random intercept
    time_variable_name="time",
    unit_variable_name="patient_id",
    model=cp.pymc_models.HierarchicalLinearRegression(
        sample_kwargs={"draws": 2000, "tune": 1000, "progressbar": False},
        non_centered=True  # Default, improves MCMC sampling efficiency
    )
)

result.summary()
result.plot()

# Access group-level effects
result.plot_group_effects()  # Forest plot of clinic random effects
result.icc  # Intraclass correlation coefficient

Advanced Formula Syntax

✅ Supported in v1.0:

# Random intercepts only (most common case)
formula = "outcome ~ post_treatment*treated + age + (1|clinic_id)"

# Random slopes (clinic-specific time trends)
formula = "outcome ~ post_treatment*treated + age + (1 + post_treatment|clinic_id)"
#                                                     └── Random slope for time

# Uncorrelated random intercepts and slopes
formula = "outcome ~ post_treatment*treated + (1|clinic_id) + (0 + post_treatment|clinic_id)"
#         Fixed effects ──────┘             Random intercept ─┘  Random slope (uncorrelated) ──┘

❌ Not supported in v1.0 (will raise NotImplementedError with guidance):

# Nested random effects (patients in clinics in regions)
formula = "outcome ~ post_treatment*treated + (1|region/clinic_id)"
# → Future release: Nested hierarchies

# Crossed random effects (patients see multiple doctors)
formula = "outcome ~ post_treatment*treated + (1|patient_id) + (1|doctor_id)"
# → Future release: Multiple grouping factors

# Group-specific treatment effects with interactions
formula = "outcome ~ post_treatment + treated + age + (post_treatment:treated|clinic_id)"
# → Future release: Complex interaction terms in random effects

Technical Implementation

1. Add Formulaic as Dependency

Update pyproject.toml:

[project]
dependencies = [
    "formulaic>=1.0.0",  # Mixed-effects formula support
    # ... existing dependencies
]

Migration Strategy:

• Start with optional import: gracefully degrade if not available
• Eventually make required for hierarchical experiments
• Consider migrating all experiments from Patsy to Formulaic (Patsy is less actively maintained)

2. Formula Parsing Logic

from typing import Tuple, Optional
import numpy as np
import pandas as pd

def parse_formula(
    formula: str, data: pd.DataFrame
) -> Tuple[np.ndarray, np.ndarray, Optional[dict], dict]:
    """
    Parse formula and detect random effects.
    
    Returns
    -------
    y : np.ndarray
        Outcome variable
    X_fixed : np.ndarray
        Fixed effects design matrix
    random_effects : dict or None
        Dictionary with random effects structure if present, else None
    formula_info : dict
        Dictionary containing:
        - 'fixed_effect_names': list of coefficient names for fixed effects
        - 'outcome_name': name of outcome variable
        - 'has_random_effects': bool indicating if formula has random effects
    """
    # Check if formula contains random effects
    has_random_effects = "|" in formula and "(" in formula
    
    formula_info = {
        'outcome_name': formula.split("~")[0].strip(),
        'has_random_effects': has_random_effects,
        'fixed_effect_names': [],
    }
    
    if has_random_effects:
        # Use Formulaic for mixed-effects formulas
        try:
            from formulaic import model_matrix
        except ImportError:
            raise ImportError(
                "Random effects syntax requires 'formulaic' package. "
                "Install with: pip install formulaic"
            )
        
        # Parse fixed effects portion (everything before the first random effect)
        # Remove random effects terms from formula for fixed effects matrix
        import re
        fixed_formula = re.sub(r'\+?\s*\([^)]+\|[^)]+\)', '', formula)
        
        # Parse mixed-effects formula
        y = model_matrix(formula.split("~")[0].strip(), data)
        X_fixed = model_matrix(fixed_formula.split("~")[1].strip(), data)
        
        # Store coefficient names
        formula_info['fixed_effect_names'] = list(X_fixed.columns)
        
        # Extract random effects structure
        random_effects = _extract_random_effects(formula, data)
        
        return np.asarray(y), np.asarray(X_fixed), random_effects, formula_info
    else:
        # Use Patsy for standard formulas (backward compatibility)
        from patsy import dmatrices
        y, X = dmatrices(formula, data)
        formula_info['fixed_effect_names'] = X.design_info.column_names
        return np.asarray(y), np.asarray(X), None, formula_info


def _extract_random_effects(formula: str, data: pd.DataFrame) -> dict:
    """
    Extract random effects structure from formula.
    
    Returns dict with:
    - 'groups': list of grouping variables
    - 'terms': list of terms for each group
    - 'type': 'intercept', 'slope', or 'both'
    """
    # Parse random effects terms (this is simplified - actual implementation 
    # would use Formulaic's parsed structure)
    import re
    random_terms = re.findall(r'\(([^)]+)\|([^)]+)\)', formula)
    
    random_structure = {
        'groups': [],
        'terms': [],
        'nested': False,
        'crossed': len(random_terms) > 1
    }
    
    for terms, group in random_terms:
        random_structure['groups'].append(group.strip())
        random_structure['terms'].append(terms.strip())
    
    return random_structure

3. New PyMC Model Class: HierarchicalLinearRegression

class HierarchicalLinearRegression(PyMCModel):
    """
    Hierarchical linear regression with group-level random effects.
    
    Supports random intercepts, random slopes, or both, as specified
    via formula syntax parsed by Formulaic.
    
    Model structure (centered parameterization):
        # Fixed effects
        beta ~ Normal(0, 10)
        
        # Random effects (group-level)
        sigma_group ~ HalfNormal(1)
        alpha_group ~ Normal(0, sigma_group)  [if random intercepts]
        beta_group ~ Normal(0, sigma_beta)    [if random slopes]
        
        # Linear predictor
        mu = alpha + X @ beta + alpha_group[group_idx] + ...
        
        # Likelihood
        sigma ~ HalfNormal(1)
        y ~ Normal(mu, sigma)
    
    Model structure (non-centered parameterization):
        # More efficient sampling for hierarchical models
        sigma_group ~ HalfNormal(1)
        alpha_group_raw ~ Normal(0, 1)
        alpha_group = alpha_group_raw * sigma_group
    
    Parameters
    ----------
    random_effects : dict
        Random effects structure from formula parsing containing:
        - 'groups': list of grouping variable names
        - 'terms': list of terms for each group
        - 'n_groups': int, number of groups
    sample_kwargs : dict
        Kwargs passed to pm.sample()
    priors : dict, optional
        Custom priors for parameters
    non_centered : bool, default=True
        Use non-centered parameterization for random effects.
        Recommended for most cases as it improves MCMC sampling efficiency.
    """
    
    def __init__(
        self, 
        random_effects: dict,
        sample_kwargs: dict | None = None,
        priors: dict | None = None,
        non_centered: bool = True,
    ):
        super().__init__(sample_kwargs=sample_kwargs, priors=priors)
        self.random_effects = random_effects
        self.non_centered = non_centered
        
        # Validate random effects structure
        self._validate_random_effects()
    
    def fit(
        self,
        X: xr.DataArray,
        y: xr.DataArray,
        group_idx: np.ndarray,
        coords: dict,
    ) -> az.InferenceData:
        """
        Fit the hierarchical model.
        
        Parameters
        ----------
        X : xr.DataArray
            Fixed effects design matrix with dims ["obs_ind", "coeffs"]
        y : xr.DataArray
            Outcome variable with dims ["obs_ind", "treated_units"]
        group_idx : np.ndarray
            Integer array mapping observations to group indices
        coords : dict
            Coordinate dictionary including 'groups' dimension
            
        Returns
        -------
        az.InferenceData
            Posterior samples and diagnostics
        """
        self.build_model(X, y, group_idx, coords)
        self.idata = pm.sample(**self.sample_kwargs)
        
        # Add posterior predictive samples
        self.idata.extend(pm.sample_posterior_predictive(self.idata))
        
        return self.idata
    
    default_priors = {
        "beta": Prior("Normal", mu=0, sigma=10, dims=["treated_units", "coeffs"]),
        "sigma_group": Prior("HalfNormal", sigma=1),
        "y_hat": Prior(
            "Normal",
            sigma=Prior("HalfNormal", sigma=1, dims=["treated_units"]),
            dims=["obs_ind", "treated_units"],
        ),
    }
    
    def _validate_random_effects(self) -> None:
        """Raise NotImplementedError for crossed, nested, or interaction random effects."""
        # Check: only single grouping variable, no '/' or ':' in groups/terms
        ...
    
    def build_model(self, X, y, group_idx, coords):
        """
        Build hierarchical model: fixed effects + random intercepts/slopes.
        
        Key structure:
        - alpha (global intercept), beta (fixed effects)
        - sigma_group ~ HalfNormal(1)
        - If non_centered: alpha_group = alpha_group_raw * sigma_group
        - Else: alpha_group ~ Normal(0, sigma_group)
        - mu = alpha + X @ beta + alpha_group[group_idx]
        - y ~ Normal(mu, sigma)
        """
        ...

4. New Experiment Class: HierarchicalDifferenceInDifferences

class HierarchicalDifferenceInDifferences(BaseExperiment):
    """
    Hierarchical Difference-in-Differences with group-level random effects.
    
    For panel data where individuals are nested within groups. Random effects
    are specified directly in the formula using lme4-style syntax via Formulaic.
    
    **Bayesian-only implementation**: This experiment class only supports PyMC
    models. OLS/frequentist estimation is not available for hierarchical DiD
    as random effects require hierarchical priors and partial pooling.
    
    Parameters
    ----------
    data : pd.DataFrame
        Individual-level panel data with group membership
    formula : str
        Patsy/Formulaic formula including random effects specification.
        Examples:
        - "y ~ 1 + post*treated + x + (1|group)" - random intercepts
        - "y ~ 1 + post*treated + (1 + post|group)" - random slopes
        - "y ~ 1 + post*treated + (1|region/clinic)" - nested random effects
    time_variable_name : str
        Column name for time periods
    unit_variable_name : str
        Column name for individual identifiers
    model : PyMCModel, optional
        A PyMC model. Will automatically use HierarchicalLinearRegression
        if random effects detected in formula and model=None.
    
    Attributes
    ----------
    random_effects : dict
        Parsed random effects structure from formula
    icc : float
        Intraclass correlation coefficient (proportion of variance between groups)
    group_effects : pd.DataFrame
        Posterior estimates of group-level random effects
    
    Examples
    --------
    >>> import causalpy as cp
    >>> df = pd.DataFrame({
    ...     'id': range(100),
    ...     'group': np.repeat(range(10), 10),
    ...     'time': np.tile([0, 1], 50),
    ...     'y': np.random.randn(100),
    ...     'post': np.tile([0, 1], 50),
    ...     'treated': np.repeat([0, 1], 50),
    ... })
    >>> result = cp.HierarchicalDifferenceInDifferences(
    ...     df,
    ...     formula="y ~ 1 + post*treated + (1|group)",
    ...     time_variable_name="time",
    ...     unit_variable_name="id",
    ... )
    >>> result.summary()
    >>> result.plot()
    >>> result.plot_group_effects()  # Forest plot of group random effects
    
    References
    ----------
    .. Bertrand, M., Duflo, E., & Mullainathan, S. (2004). How much should we 
       trust differences-in-differences estimates?
    .. Gelman, A., & Hill, J. (2006). Data analysis using regression and 
       multilevel/hierarchical models.
    """
    
    supports_ols = False  # Requires Bayesian inference for random effects
    supports_bayes = True
    
    def __init__(
        self,
        data: pd.DataFrame,
        formula: str,
        time_variable_name: str,
        unit_variable_name: str,
        model: PyMCModel | None = None,
        **kwargs: dict,
    ) -> None:
        # Store parameters
        self.expt_type = "Hierarchical Difference in Differences"
        self.formula = formula
        self.time_variable_name = time_variable_name
        self.unit_variable_name = unit_variable_name
        
        # Parse formula and extract design matrices
        y, X_fixed, self.random_effects, self._formula_info = parse_formula(formula, data)
        
        # Extract outcome variable name from formula
        self.outcome_variable_name = formula.split("~")[0].strip()
        
        # Store coefficient labels for printing and prediction
        self.labels = self._formula_info['fixed_effect_names']
        
        # Auto-select model if none provided
        if model is None and self.random_effects is not None:
            model = HierarchicalLinearRegression(
                random_effects=self.random_effects,
                sample_kwargs=kwargs.get('sample_kwargs', {})
            )
        
        super().__init__(model=model)
        
        # Validate and process data
        self.data = data.copy()
        self.data.index.name = "obs_ind"
        self.input_validation()
        
        # Create group index mapping
        self._create_group_indices()
        
        # Convert to xarray DataArrays (required for PyMC model fitting)
        self.X = xr.DataArray(
            X_fixed,
            dims=["obs_ind", "coeffs"],
            coords={
                "obs_ind": np.arange(X_fixed.shape[0]),
                "coeffs": self.labels,
            },
        )
        self.y = xr.DataArray(
            y,
            dims=["obs_ind", "treated_units"],
            coords={
                "obs_ind": np.arange(y.shape[0]),
                "treated_units": ["unit_0"],
            },
        )
        
        # Fit model
        self._fit_model()
        
        # Compute derived quantities
        self._compute_icc()
        self._extract_group_effects()
    
    # Key methods (implementations follow existing CausalPy patterns):
    
    def _create_group_indices(self) -> None:
        """Create integer indices for grouping variables."""
        ...
    
    def _fit_model(self) -> None:
        """Fit model with COORDS including 'groups' dimension."""
        ...
    
    def _compute_icc(self) -> None:
        """Compute ICC = σ²_group / (σ²_group + σ²)."""
        ...
    
    def _extract_group_effects(self) -> None:
        """Extract posterior summaries of group random effects."""
        ...
    
    def plot_group_effects(self, **kwargs) -> tuple[plt.Figure, plt.Axes]:
        """Forest plot of group-level random effects."""
        ...
    
    def summary(self, round_to: int | None = 2) -> None:
        """Print summary including ICC and variance components."""
        ...
    
    def _bayesian_plot(self, round_to: int | None = None, **kwargs) -> tuple[plt.Figure, plt.Axes]:
        """Standard DiD plot with treated/control trajectories and causal impact."""
        ...
    
    def get_plot_data_bayesian(self, **kwargs) -> pd.DataFrame:
        """Return data with predictions, HDI bounds, and group effects."""
        ...
    
    def effect_summary(self, *, direction="increase", alpha=0.05, **kwargs) -> EffectSummary:
        """Generate decision-ready summary with ATT, HDI, P(effect), ICC."""
        ...

Output and Diagnostics

Standard Output

• Fixed effects: ATT ($\delta$) with posterior distribution
• Random effects: Group-specific deviations ($\alpha_j$)
• Variance components:
  • Between-group variance ($\sigma_{\alpha}^2$)
  • Within-group variance ($\sigma^2$)
• ICC: Intraclass correlation coefficient - proportion of variance between vs. within groups

Visualization Methods

Plot Methods

Method	Purpose
result.plot()	Standard DiD plot: treated/control trajectories, counterfactual, causal impact arrow, 94% HDI bands
result.plot_group_effects()	Forest plot: group random effects (α_j) with 94% HDI, sorted by magnitude, ICC in title
result.plot_variance_components()	Posterior distributions: between-group variance, within-group variance, ICC distribution

For MCMC diagnostics, use ArviZ directly: az.plot_trace(result.model.idata), az.plot_ppc(result.model.idata)

Migration Path and Backward Compatibility

Phase 1: Optional Formulaic Support

• Add Formulaic as optional dependency
• Detect random effects syntax and route to appropriate parser
• Keep Patsy for all existing experiments
• No breaking changes

Phase 2: Hierarchical Experiments

• Implement HierarchicalDifferenceInDifferences using Formulaic
• Start with basic random effects support (single grouping variable)
• Thoroughly test against known results
• Implement comprehensive data validation (see "Data Validation and Requirements" section)
• Create API documentation with clear docstrings
• Raise clear exceptions for unsupported formula syntax with guidance

Phase 3: Example Documentation Notebook

• Create comprehensive .ipynb notebook in docs/source/notebooks/
• Follow naming convention: hdid_pymc.ipynb (hierarchical DiD with PyMC)
• Implement the "Detailed Example: Marketing Campaign Evaluation" scenario (see above)
• Notebook structure:
  1. Introduction:
     • Motivation for hierarchical DiD, when to use it
     • Key Assumptions:
       • Balanced panel (all units observed in all time periods)
       • No missing outcome data
       • Units nested within groups (no group switching)
       • Treatment assigned at group level (not individual level)
       • All treated groups receive treatment simultaneously (no staggered adoption)
       • Sufficient number of groups (≥10 recommended, ≥3 minimum)
       • Parallel trends assumption (as in standard DiD)
       • No anticipation effects
       • Stable unit treatment value assumption (SUTVA) within groups
     • When to use hierarchical DiD vs. standard DiD vs. staggered DiD
  2. Data Generation: Synthetic retail data with realistic within-store correlation (ICC ≈ 0.20)
  3. Exploratory Analysis: Visualize clustering structure, compute empirical ICC
  4. Three Approaches Comparison:
     • Naive DiD (ignoring clustering)
     • DiD with store fixed effects
     • Hierarchical DiD (this implementation)
  5. Results Comparison: Side-by-side parameter estimates, standard errors, coverage
  6. Visualizations:
     • Standard DiD plot with group trajectories (result.plot())
     • Forest plot of store random effects (result.plot_group_effects())
     • Variance components posterior distributions (result.plot_variance_components())
     • ICC interpretation and visualization
  7. Hierarchical Model Diagnostics:
     • Convergence diagnostics using ArviZ (az.plot_trace(), R-hat, ESS)
     • Posterior predictive checks using ArviZ (az.plot_ppc())
     • Model fit assessment and interpretation
  8. Business Insights:
     • Identify top/bottom performing stores from forest plot
     • Correlate store effects with characteristics (urban/suburban, size, region)
     • Practical recommendations for intervention rollout
  9. Recommendations: When to use hierarchical DiD, practical guidance, common pitfalls
• Include citations to key references (Bertrand et al. 2004, Gelman & Hill 2006)
• Add glossary links for technical terms (ICC, partial pooling, random effects)
• Ensure notebook runs in < 2 minutes (use lightweight MCMC settings for demo)

Testing Strategy

Test categories (all using pytest):

1. Random effects parsing: (1|group), (1+x|group) syntax
2. Parameterization: centered vs non-centered (default)
3. Unsupported syntax exceptions: crossed (1|g1)+(1|g2), nested (1|a/b), interactions (x:z|g)
4. Data validation exceptions: staggered treatment, unbalanced panels, missing outcomes, etc.
5. ICC computation: verify against known variance ratios
6. Backward compatibility: graceful ImportError if Formulaic not installed
7. Integration tests: full workflow with generate_hierarchical_did_data() helper

Related Work

Literature

1. Bertrand, Duflo, & Mullainathan (2004): "How Much Should We Trust Differences-In-Differences Estimates?" - Highlights clustering issues in DiD
2. Raudenbush & Bryk (2002): Hierarchical Linear Models - Foundation for multilevel modeling
3. Gelman & Hill (2006): Data Analysis Using Regression and Multilevel/Hierarchical Models - Bayesian approach to hierarchical models
4. Abadie et al. (2023): "When Should You Adjust Standard Errors for Clustering?" - Recent guidance on clustering

Existing Implementations

• R: lme4::lmer() for frequentist mixed models - inspiration for formula syntax
• R: brms for Bayesian hierarchical models (wraps Stan)
• Python: bambi for Bayesian GLMMs (uses Formulaic + PyMC)
• Python: statsmodels.MixedLM (frequentist only, no DiD-specific features)
• Stata: xtmixed, mixed commands

Gap: No Python package currently offers Bayesian hierarchical DiD with:

• Intuitive formula-based specification (lme4-style)
• Causal inference focus (not just generic mixed models)
• Visualization and diagnostics tailored for DiD
• Integration with a broader causal inference ecosystem

Benefits

1. Statistical: Proper accounting of within-group correlation, more efficient estimates, valid standard errors
2. Substantive: Estimate group-level treatment effect heterogeneity (which clinics/schools/stores respond better?)
3. Bayesian Advantages:
   • Natural handling of hierarchical structure through priors
   • Full posterior distributions for all parameters (fixed effects, random effects, variance components)
   • Partial pooling improves estimates for small groups
   • Direct probability statements about treatment effects
4. Practical: Familiar syntax for R users, extensible to complex random effects structures
5. Business Value: Essential for marketing/sales analytics where interventions are deployed at cluster level (stores, territories, markets) but outcomes measured at individual level (customers, reps, users)
6. Pedagogical: Demonstrates when/why hierarchical modeling matters in causal inference
7. Future-proof: Positions CausalPy for modern formula parsing (Formulaic vs legacy Patsy)

Implementation Checklist

Core Implementation (v1.0)

• Add Formulaic as optional dependency to pyproject.toml
• Implement formula parsing logic (place in causalpy/utils.py or new causalpy/formula_utils.py)
  • parse_formula() function returning (y, X, random_effects, formula_info)
  • _extract_random_effects() helper function
  • Return formula_info dict with fixed_effect_names for coefficient labels
  • Parse random intercepts: (1|group)
  • Parse random slopes: (x|group), (1 + x|group)
  • Validate single grouping variable only
  • Raise exceptions for unsupported syntax (crossed, nested, interactions)
• Create HierarchicalLinearRegression PyMC model (in pymc_models.py)
  • Inherit from PyMCModel
  • Implement fit(X, y, group_idx, coords) method
  • Implement build_model() method
  • Random intercepts with centered parameterization
  • Random intercepts with non-centered parameterization (default)
  • Random slopes support
  • non_centered parameter to toggle parameterization
  • Validation method for supported random effects structures (_validate_random_effects())
  • Define default_priors class attribute
• Create HierarchicalDifferenceInDifferences experiment class
  • Set expt_type = "Hierarchical Difference in Differences"
  • Extract outcome_variable_name from formula
  • Store coefficient labels for printing
  • Formula parsing integration (with _formula_info)
  • Group index creation and mapping (_create_group_indices())
  • Convert X, y to xarray DataArrays
  • Model fitting with COORDS including "groups" dimension (_fit_model())
  • Treatment effect extraction (coefficient on interaction term)
  • Implement _bayesian_plot() abstract method
  • Implement get_plot_data_bayesian() abstract method
  • Implement effect_summary() method with EffectSummary return
• Implement comprehensive data validation (input_validation() method)
  • Balanced panel check (all units observed in all periods)
  • No missing outcome values
  • Group membership integrity (units don't switch groups)
  • Treatment assignment at group level (no within-group variation)
  • Sufficient observations per group (min 2, warn if < 5)
  • Clear pre/post period structure
  • No staggered treatment adoption (all treated groups receive treatment at same time)
  • Valid random effects formula syntax (leverage existing _validate_random_effects())
  • Appropriate data types (categorical groups, numeric outcomes)
  • Sufficient groups for random effects (min 3, warn if < 10)
  • No perfect multicollinearity in design matrix
• Implement diagnostic methods
  • ICC calculation
  • Group effects extraction
  • plot_group_effects() - forest plot of random effects
  • plot_variance_components() - posterior distributions of variance components
  • Enhanced summary() with variance components
• Write comprehensive tests
  • Unit tests for formula parsing
  • Unit tests for validation (unsupported syntax raises exceptions)
  • Data validation tests (all 11 requirements from "Data Validation and Requirements")
    • Test unbalanced panel detection
    • Test missing outcome detection
    • Test group membership integrity violations
    • Test mixed treatment within groups
    • Test insufficient observations per group
    • Test clear pre/post period structure
    • Test staggered treatment adoption detection (should raise exception)
    • Test insufficient number of groups
    • Test data type validation
    • Test multicollinearity detection
  • Integration tests with known hierarchical DiD examples
  • Test centered vs non-centered parameterization
  • Test random intercepts only
  • Test random slopes
  • Performance testing with various group sizes (10, 50, 100, 500 groups)

Documentation & Examples

• Create hdid_pymc.ipynb notebook implementing "Marketing Campaign Evaluation" example (see Phase 3 above)
• Add to documentation with:
  • Theoretical background on hierarchical models
  • When to use vs standard DiD (decision tree)
  • Formula syntax guide with examples
  • Comparison to lme4/R syntax (cheat sheet)
  • Centered vs non-centered parameterization explanation
  • Troubleshooting guide for convergence issues
• Add inline documentation in code
  • Clear docstrings with examples
  • Exception messages with actionable guidance

Quality Assurance

• Code review by team
• Pre-commit checks pass
• Test coverage > 90% for new code
• Documentation builds without errors
• Example notebook runs successfully

Design Decisions & Rationale

1. Formulaic maturity: Formulaic v1.0+ is stable and production-ready, actively used by Bambi for Bayesian GLMMs
2. Complex random effects ✅ DECIDED: Start with basic support (single grouping variable, random intercepts/slopes). Raise clear exceptions for unsupported syntax (crossed, nested, complex interactions) with guidance on workarounds. This allows us to ship a robust v1.0 and extend incrementally.
3. Prior specification: Use weakly informative priors by default (HalfNormal(1) for variance components) to be minimally informative while aiding convergence. Users can override via priors argument.
4. Computational ✅ DECIDED: Support both centered and non-centered parameterization, with non-centered as default (non_centered=True). This improves MCMC sampling efficiency for hierarchical models, especially with small group sizes or weak between-group variation.

Open Questions (for future consideration)

1. API evolution: If hierarchical DiD proves successful, should other experiments migrate to Formulaic? (Can be addressed in separate issues)
2. Automatic parameterization selection: Should we auto-detect when to use non-centered vs centered based on data characteristics? (Can be deferred to v1.1+)
3. Performance optimization: For very large numbers of groups (>1000), are there computational strategies we should implement? (Address when real use cases emerge)

Example Use Cases

Healthcare: Clinic-Level Intervention

# 50 clinics, 20 patients each, half receive quality improvement intervention
result = cp.HierarchicalDifferenceInDifferences(
    patient_df,
    formula="hba1c ~ 1 + post_intervention*treated_clinic + patient_age + (1|clinic_id)",
    time_variable_name="quarter",
    unit_variable_name="patient_id",
)

# Which clinics had the strongest response?
result.plot_group_effects()

Education: School-Level Policy

# 100 schools, 30 students each, new curriculum in treated schools
result = cp.HierarchicalDifferenceInDifferences(
    student_df,
    formula="test_score ~ 1 + post*treated_school + ses + (1 + post|school_id)",
    #                                                         └── School-specific trends
    time_variable_name="year",
    unit_variable_name="student_id",
)

Marketing: Store-Level Promotional Campaign

# 200 retail stores, 500 customers per store, half receive new loyalty program
# Different stores have different baseline sales and customer demographics
result = cp.HierarchicalDifferenceInDifferences(
    transaction_df,
    formula="purchase_amount ~ 1 + post_launch*treated_store + customer_age + customer_tenure + (1|store_id)",
    time_variable_name="week",
    unit_variable_name="customer_id",
)

# Quantify store-level heterogeneity in campaign effectiveness
print(f"ICC: {result.icc:.3f}")  # What % of variance is between stores?
result.plot_group_effects()  # Which stores had strongest lift?

# Are certain store characteristics associated with better performance?
result.group_effects.merge(store_characteristics, on='store_id')

Sales: Territory-Level Training Intervention

# 50 sales territories, 10 reps per territory, half receive new sales training
# Reps nested within territories have correlated performance
result = cp.HierarchicalDifferenceInDifferences(
    sales_df,
    formula="monthly_revenue ~ 1 + post_training*treated_territory + rep_experience + (1 + post_training|territory_id)",
    #                                                                                    └── Territory-specific training effects
    time_variable_name="month",
    unit_variable_name="rep_id",
)

# Identify high-performing territories for best practice sharing
high_performers = result.group_effects[result.group_effects['mean'] > 0].sort_values('mean', ascending=False)

E-commerce: Market-Level Pricing Experiment

# 30 geographic markets, 1000 users per market, randomized pricing in half
# User behavior within markets is correlated due to local economic conditions
result = cp.HierarchicalDifferenceInDifferences(
    user_df,
    formula="conversion_rate ~ 1 + post_price_change*treated_market + user_ltv + (1|market_id)",
    time_variable_name="date",
    unit_variable_name="user_id",
)

# Account for market-level correlation when estimating price elasticity
print(f"Treatment Effect (ATT): {result.causal_impact.mean():.3f}")
print(f"Between-market variance explains {result.icc*100:.1f}% of total variance")

References

Clustering in Difference-in-Differences

• Bertrand, M., Duflo, E., & Mullainathan, S. (2004). How much should we trust differences-in-differences estimates? The Quarterly Journal of Economics, 119(1), 249-275.
• Abadie, A., Athey, S., Imbens, G. W., & Wooldridge, J. M. (2023). When should you adjust standard errors for clustering? The Quarterly Journal of Economics, 138(1), 1-35.

Hierarchical/Multilevel Modeling

• Gelman, A., & Hill, J. (2006). Data analysis using regression and multilevel/hierarchical models. Cambridge University Press.
• Raudenbush, S. W., & Bryk, A. S. (2002). Hierarchical linear models: Applications and data analysis methods (2nd ed.). Sage.
• Betancourt, M., & Girolami, M. (2015). Hamiltonian Monte Carlo for hierarchical models. Current Trends in Bayesian Methodology with Applications, 79, 30.

Treatment Effect Heterogeneity

• Athey, S., & Imbens, G. W. (2017). The econometrics of randomized experiments. In Handbook of Economic Field Experiments (Vol. 1, pp. 73-140). North-Holland.
• Meager, R. (2019). Understanding the average impact of microcredit expansions: A Bayesian hierarchical analysis of seven randomized experiments. American Economic Journal: Applied Economics, 11(1), 57-91.

Software and Implementation

• Formulaic Documentation: https://matthewwardrop.github.io/formulaic/
• Bambi (Bayesian GLMMs with Formulaic): https://bambinos.github.io/bambi/
• PyMC Documentation on Hierarchical Models: https://www.pymc.io/projects/docs/en/stable/learn/core_notebooks/multilevel_modeling.html

---

Related Issues: None currently