ETWFERegression.fit#

ETWFERegression.fit(X, y, coords, *, unit_idx, time_idx, cohort_idx, ev_idx, effect_indicator, att_weights, dbar_unit=None, dbar_time=None, conditioning='mundlak')[source]#

Draw posterior, prior predictive and posterior predictive samples.

The signature is widened relative to PyMCModel.fit() to carry the panel structure; InstrumentalVariableRegression.fit() is the precedent for this in the codebase.

Parameters:

  • X (DataArray) – Covariate matrix with dims ["obs_ind", "coeffs"]. May have zero columns, in which case no beta is created.

  • y (DataArray) – Outcome with dims ["obs_ind", "treated_units"].

  • coords (dict[str, Any]) – Coordinate metadata; must supply units, periods, cohorts and ev alongside the usual obs_ind / treated_units.

  • unit_idx (ndarray) – Integer position arrays of length n_obs.

  • time_idx (ndarray) – Integer position arrays of length n_obs.

  • cohort_idx (ndarray) – Integer position arrays of length n_obs.

  • ev_idx (ndarray) – Integer position arrays of length n_obs.

  • effect_indicator (ndarray) – 0/1 float array of length n_obs.

  • att_weights (ndarray) – (n_cohorts, n_ev) weights summing to one.

  • dbar_unit (ndarray | None) – Centred Mundlak treatment means; required for conditioning="mundlak".

  • dbar_time (ndarray | None) – Centred Mundlak treatment means; required for conditioning="mundlak".

  • conditioning (Literal['dummy', 'mundlak']) – Conditioning strategy for the two-way effects.

Returns:

Inference data with posterior, prior_predictive and posterior_predictive groups. The posterior predictive group contains both y_hat and mu.

Return type:

arviz.InferenceData