ETWFERegression.build_model#

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

Define the ETWFE PyMC model.

Parameters:

  • X (DataArray) – Covariates with dims ["obs_ind", "coeffs"]. May have zero columns, in which case no X data node and no beta variable are created.

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

  • coords (dict[str, Any]) – Must supply obs_ind, treated_units, units, periods, cohorts and ev; plus coeffs when X has columns.

  • unit_idx (ndarray) – Integer position arrays of length n_obs, indexing into the units, periods, cohorts and ev coordinates respectively. All must be non-negative.

  • time_idx (ndarray) – Integer position arrays of length n_obs, indexing into the units, periods, cohorts and ev coordinates respectively. All must be non-negative.

  • cohort_idx (ndarray) – Integer position arrays of length n_obs, indexing into the units, periods, cohorts and ev coordinates respectively. All must be non-negative.

  • ev_idx (ndarray) – Integer position arrays of length n_obs, indexing into the units, periods, cohorts and ev coordinates respectively. All must be non-negative.

  • effect_indicator (ndarray) – Float array of length n_obs, 1.0 where the (cohort, event time) effect applies to that observation and 0.0 otherwise.

  • att_weights (ndarray) – Weight matrix of shape (n_cohorts, n_ev) summing to one over treated cells, used to aggregate tau into att.

  • dbar_unit (ndarray | None) – Mundlak treatment means of length n_obs. Required when conditioning="mundlak". Should be centred by the caller.

  • dbar_time (ndarray | None) – Mundlak treatment means of length n_obs. Required when conditioning="mundlak". Should be centred by the caller.

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

Return type:

None