Skip to content

statspai.matrix_completion

matrix_completion

Matrix Completion for Causal Panel Data.

Estimates treatment effects in panel settings by completing the counterfactual matrix of untreated potential outcomes using nuclear-norm regularisation (low-rank matrix completion).

References

Athey, S., Bayati, M., Doudchenko, N., Imbens, G., & Khosravi, K. (2021). Matrix Completion Methods for Causal Panel Data Models. JASA, 116(536), 1716-1730. [@athey2021matrix]

MCPanel

Matrix Completion for Causal Panels.

Parameters:

Name Type Description Default
data DataFrame
required
y str
required
unit str
required
time str
required
treat str
required
lambda_reg float
None
max_rank int
None
max_iter int
1000
tol float
1e-05
n_bootstrap int
200
alpha float
0.05
random_state int
42

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> from statspai.matrix_completion.mc_panel import MCPanel
>>> rng = np.random.default_rng(0)
>>> rows = []
>>> for u in range(8):
...     fe = rng.normal()
...     for t in range(6):
...         treated = int(u >= 6 and t >= 4)
...         y = fe + 0.3 * t + rng.normal(0, 0.2) + (1.5 if treated else 0.0)
...         rows.append((u, t, y, treated))
>>> panel = pd.DataFrame(rows, columns=['unit', 'period', 'y', 'treated'])
>>> est = MCPanel(data=panel, y='y', unit='unit', time='period',
...               treat='treated', n_bootstrap=50)
>>> res = est.fit()
>>> bool(res.estimate > 0)
True
References

[@athey2021matrix]

fit

fit() -> CausalResult

Run matrix completion and return treatment effect estimates.