ETWFERegression#

class causalpy.pymc_models.ETWFERegression[source]#

Extended two-way fixed effects (ETWFE) regression for staggered adoption.

This is the Wooldridge (2021) / Mundlak (1978) structural phrasing of the staggered difference-in-differences estimand: a saturated regression with one treatment effect per (adoption cohort, event time) cell.

\[\begin{split}\mu_{it} &= \alpha_i + \beta_t + E_{it}\,\tau_{g(i),\,k(i,t)} \;(+\; \gamma_u \bar{D}_i + \gamma_t \bar{D}_t) \;(+\; X_{it}\beta) \\ \bar{\tau}_k &\sim \mathrm{Normal}(0, \cdot) \\ \sigma_{dev} &\sim \mathrm{HalfNormal}(\cdot) \\ d_{gk} &\sim \mathrm{Normal}(0, 1) \\ \tau_{gk} &= \bar{\tau}_k + \sigma_{dev}\, d_{gk} \\ \mathrm{ATT} &= \sum_{g,k} w_{gk}\, \tau_{gk} \\ y_{it} &\sim \mathrm{Normal}(\mu_{it}, \sigma)\end{split}\]

The aggregated ATT is a pymc.Deterministic() inside the model, so it arrives with its own posterior rather than being reconstructed from differenced posterior predictive draws after the fact. That is the point of this estimator.

The cohort effects \(\tau_{gk}\) are partially pooled towards a common event-time profile \(\bar{\tau}_k\) via a non-centred parameterisation.

Notes

The constructor is inherited unchanged from PyMCModel, taking sample_kwargs (forwarded to pymc.sample()) and priors, a dictionary of pymc_extras.prior.Prior objects. Recognised prior keys are alpha_dummy, beta_t_dummy, mu_a, sd_a, a_z, sd_bt, beta_t_mundlak, g_u, g_t, tau_bar, sd_dev, dev, beta and y_hat.

All panel structure is supplied to fit(), not to __init__. This keeps the constructor signature identical to PyMCModel, so PyMCModel._clone() (and hence causalpy.checks.prior_sensitivity.PriorSensitivity) works unchanged.

Conditioning variants. conditioning="dummy" gives every unit a free intercept and gives beta_t a fixed-scale ZeroSumNormal. conditioning="mundlak" replaces the free intercepts with a non-centred hierarchical intercept (mu_a, sd_a, a_z), learns the scale of beta_t (sd_bt), and adds the Mundlak treatment means dbar_unit and dbar_time with coefficients g_u and g_t.

``g_t`` must not be interpreted. dbar_time is a deterministic function of \(t\) alone, so it lies exactly in the span of the time effects beta_t. g_t is therefore identified only by its prior; its posterior carries no information about the data. This is a property of the Mundlak device, not a bug. The ATT is unaffected: tau is identified off within-cell variation, which is orthogonal to any function of \(t\) alone. g_u is better behaved (the unit intercepts are only partially pooled, so shrinkage identifies it), but it is a nuisance parameter and shares a funnel with mu_a/sd_a.

Pass centred dbar_unit and dbar_time. Subtracting their means orthogonalises g_u against mu_a and improves geometry. It does not change the estimand.

Examples

>>> import numpy as np
>>> import xarray as xr
>>> from causalpy.pymc_models import ETWFERegression
>>> n_units, n_periods = 6, 4
>>> n = n_units * n_periods
>>> unit_idx = np.repeat(np.arange(n_units), n_periods)
>>> time_idx = np.tile(np.arange(n_periods), n_units)
>>> # units 0-2 adopt at t=2, units 3-5 never adopt: a single cohort, so
>>> # every row maps to cohort column 0
>>> cohort_idx = np.zeros(n, dtype=int)
>>> effect_indicator = ((unit_idx < 3) & (time_idx >= 2)).astype(float)
>>> ev_idx = np.where(effect_indicator > 0, time_idx - 2, 0)
>>> att_weights = np.array([[0.5, 0.5]])
>>> rng = np.random.default_rng(42)
>>> y = xr.DataArray(
...     rng.normal(size=(n, 1)),
...     dims=["obs_ind", "treated_units"],
...     coords={"obs_ind": np.arange(n), "treated_units": ["unit_0"]},
... )
>>> X = xr.DataArray(
...     np.empty((n, 0)),
...     dims=["obs_ind", "coeffs"],
...     coords={"obs_ind": np.arange(n), "coeffs": []},
... )
>>> coords = {
...     "obs_ind": np.arange(n),
...     "treated_units": ["unit_0"],
...     "units": np.arange(n_units),
...     "periods": np.arange(n_periods),
...     "cohorts": [2],
...     "ev": [0, 1],
... }
>>> model = ETWFERegression(sample_kwargs={"progressbar": False})
>>> model.fit(
...     X,
...     y,
...     coords,
...     unit_idx=unit_idx,
...     time_idx=time_idx,
...     cohort_idx=cohort_idx,
...     ev_idx=ev_idx,
...     effect_indicator=effect_indicator,
...     att_weights=att_weights,
...     conditioning="dummy",
... )
Inference data...

Methods

`ETWFERegression.add_coord`(name[, values, length])	Register a dimension coordinate with the model.
`ETWFERegression.add_coords`(coords, *[, lengths])	Vectorized version of `Model.add_coord`.
`ETWFERegression.add_named_variable`(var[, dims])	Add a random graph variable to the named variables of the model.
`ETWFERegression.build_model`(X, y, coords, *, ...)	Define the ETWFE PyMC model.
`ETWFERegression.check_start_vals`(start, **kwargs)	Check that the logp is defined and finite at the starting point.
`ETWFERegression.compile_d2logp`([vars, ...])	Compiled log probability density hessian function.
`ETWFERegression.compile_dlogp`([vars, jacobian])	Compiled log probability density gradient function.
`ETWFERegression.compile_fn`(outs, *[, ...])	Compiles a PyTensor function.
`ETWFERegression.compile_logp`([vars, ...])	Compiled log probability density function.
`ETWFERegression.copy`()	Clone the model.
`ETWFERegression.create_value_var`(rv_var, *, ...)	Create a `TensorVariable` that will be used as the random variable's "value" in log-likelihood graphs.
`ETWFERegression.d2logp`([vars, jacobian, ...])	Hessian of the models log-probability w.r.t.
`ETWFERegression.debug`([point, fn, verbose])	Debug model function at point.
`ETWFERegression.dlogp`([vars, jacobian])	Gradient of the models log-probability w.r.t.
`ETWFERegression.eval_rv_shapes`()	Evaluate shapes of untransformed AND transformed free variables.
`ETWFERegression.fit`(X, y, coords, *, ...[, ...])	Draw posterior, prior predictive and posterior predictive samples.
`ETWFERegression.get_context`([error_if_none, ...])
`ETWFERegression.initial_point`([random_seed])	Compute the initial point of the model.
`ETWFERegression.logp`([vars, jacobian, sum])	Elemwise log-probability of the model.
`ETWFERegression.logp_dlogp_function`([...])	Compile a PyTensor function that computes logp and gradient.
`ETWFERegression.make_obs_var`(rv_var, data, ...)	Create a TensorVariable for an observed random variable.
`ETWFERegression.name_for`(name)	Check if name has prefix and adds if needed.
`ETWFERegression.name_of`(name)	Check if name has prefix and deletes if needed.
`ETWFERegression.point_logps`([point, round_vals])	Compute the log probability of point for all random variables in the model.
`ETWFERegression.predict`([X, coords, ...])	Return the in-sample posterior predictive computed during `fit()`.
`ETWFERegression.print_coefficients`(labels[, ...])	Print a summary of the ETWFE parameters.
`ETWFERegression.priors_from_data`(X, y)	Build scale-adaptive priors from the outcome's location and spread.
`ETWFERegression.profile`(outs, *[, n, point, ...])	Compile and profile a PyTensor function which returns `outs` and takes values of model vars as a dict as an argument.
`ETWFERegression.register_data_var`(data[, dims])	Register a data variable with the model.
`ETWFERegression.register_rv`(rv_var, name, *)	Register an (un)observed random variable with the model.
`ETWFERegression.replace_rvs_by_values`(...)	Clone and replace random variables in graphs with their value variables.
`ETWFERegression.score`(X, y[, coords])	Score the Bayesian \(R^2\) given inputs `X` and outputs `y`.
`ETWFERegression.set_data`(name, values[, coords])	Change the values of a data variable in the model.
`ETWFERegression.set_dim`(name, new_length[, ...])	Update a mutable dimension.
`ETWFERegression.set_initval`(rv_var, initval)	Set an initial value (strategy) for a random variable.
`ETWFERegression.shape_from_dims`(dims)
`ETWFERegression.to_graphviz`(*[, var_names, ...])	Produce a graphviz Digraph from a PyMC model.

Attributes

`basic_RVs`	List of random variables the model is defined in terms of.
`continuous_value_vars`	All the continuous value variables in the model.
`coords`	Coordinate values for model dimensions.
`datalogp`	PyTensor scalar of log-probability of the observed variables and potential terms.
`default_priors`
`dim_lengths`	The symbolic lengths of dimensions in the model.
`discrete_value_vars`	All the discrete value variables in the model.
`isroot`
`observedlogp`	PyTensor scalar of log-probability of the observed variables.
`parent`
`potentiallogp`	PyTensor scalar of log-probability of the Potential terms.
`prefix`
`root`
`unobserved_RVs`	List of all random variables, including deterministic ones.
`unobserved_value_vars`	List of all random variables (including untransformed projections), as well as deterministics used as inputs and outputs of the model's log-likelihood graph.
`value_vars`	List of unobserved random variables used as inputs to the model's log-likelihood (which excludes deterministics).
`varlogp`	PyTensor scalar of log-probability of the unobserved random variables (excluding deterministic).
`varlogp_nojac`	PyTensor scalar of log-probability of the unobserved random variables (excluding deterministic) without jacobian term.

__init__(sample_kwargs=None, priors=None)#

Parameters:

sample_kwargs (dict[str, Any] | None) – Dictionary of kwargs that get unpacked and passed to the pymc.sample() function. Defaults to an empty dictionary if None.
priors (dict[str, Any] | None) – Dictionary of priors for the model. Defaults to None, in which case default priors are used.

Return type:

None

classmethod __new__(*args, **kwargs)#