trop 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- trop/__init__.py +9 -0
- trop/cv.py +427 -0
- trop/estimator.py +157 -0
- trop-0.1.2.dist-info/METADATA +53 -0
- trop-0.1.2.dist-info/RECORD +8 -0
- {trop-0.1.0.dist-info → trop-0.1.2.dist-info}/WHEEL +2 -1
- trop-0.1.2.dist-info/licenses/LICENSE +21 -0
- trop-0.1.2.dist-info/top_level.txt +1 -0
- trop/placeholder.py +0 -1
- trop-0.1.0.dist-info/METADATA +0 -24
- trop-0.1.0.dist-info/RECORD +0 -4
trop/__init__.py
ADDED
trop/cv.py
ADDED
|
@@ -0,0 +1,427 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from dataclasses import dataclass
|
|
4
|
+
from typing import Iterable, Optional, Sequence, Tuple, Union, List
|
|
5
|
+
|
|
6
|
+
import numpy as np
|
|
7
|
+
from joblib import Parallel, delayed
|
|
8
|
+
|
|
9
|
+
from .estimator import TROP_TWFE_average
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
ArrayLike = Union[np.ndarray, Sequence[Sequence[float]]]
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def _validate_panel(Y: np.ndarray, treated_periods: int, n_treated_units: int) -> None:
|
|
16
|
+
"""
|
|
17
|
+
Validate panel dimensions and basic placebo-treatment parameters.
|
|
18
|
+
|
|
19
|
+
Parameters
|
|
20
|
+
----------
|
|
21
|
+
Y:
|
|
22
|
+
Outcome panel of shape (N, T).
|
|
23
|
+
treated_periods:
|
|
24
|
+
Number of treated (post) periods assumed to be the final columns of the panel.
|
|
25
|
+
Must satisfy 1 <= treated_periods < T.
|
|
26
|
+
n_treated_units:
|
|
27
|
+
Number of treated units to sample without replacement from {0, ..., N-1}.
|
|
28
|
+
Must satisfy 1 <= n_treated_units < N.
|
|
29
|
+
|
|
30
|
+
Raises
|
|
31
|
+
------
|
|
32
|
+
ValueError
|
|
33
|
+
If Y is not 2D, or if treated_periods / n_treated_units are out of range.
|
|
34
|
+
"""
|
|
35
|
+
if Y.ndim != 2:
|
|
36
|
+
raise ValueError("Y must be a 2D array of shape (N, T).")
|
|
37
|
+
N, T = Y.shape
|
|
38
|
+
if treated_periods <= 0 or treated_periods >= T:
|
|
39
|
+
raise ValueError(f"treated_periods must be in [1, T-1]. Got treated_periods={treated_periods}, T={T}.")
|
|
40
|
+
if n_treated_units <= 0 or n_treated_units >= N:
|
|
41
|
+
raise ValueError(f"n_treated_units must be in [1, N-1]. Got n_treated_units={n_treated_units}, N={N}.")
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def _as_list(grid: Iterable[float]) -> List[float]:
|
|
45
|
+
"""
|
|
46
|
+
Convert a lambda grid iterable into a non-empty list of floats.
|
|
47
|
+
|
|
48
|
+
Parameters
|
|
49
|
+
----------
|
|
50
|
+
grid:
|
|
51
|
+
Iterable of candidate lambda values.
|
|
52
|
+
|
|
53
|
+
Returns
|
|
54
|
+
-------
|
|
55
|
+
List[float]
|
|
56
|
+
The grid converted to a list of floats.
|
|
57
|
+
|
|
58
|
+
Raises
|
|
59
|
+
------
|
|
60
|
+
ValueError
|
|
61
|
+
If the grid is empty.
|
|
62
|
+
"""
|
|
63
|
+
grid_list = list(grid)
|
|
64
|
+
if len(grid_list) == 0:
|
|
65
|
+
raise ValueError("lambda_grid must be non-empty.")
|
|
66
|
+
grid_list = [float(x) for x in grid_list]
|
|
67
|
+
return grid_list
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
def _simulate_ate(
|
|
72
|
+
seed: int,
|
|
73
|
+
Y: np.ndarray,
|
|
74
|
+
n_treated_units: int,
|
|
75
|
+
treated_periods: int,
|
|
76
|
+
lambda_unit: float,
|
|
77
|
+
lambda_time: float,
|
|
78
|
+
lambda_nn: float,
|
|
79
|
+
solver: Optional[str] = None,
|
|
80
|
+
verbose: bool = False,
|
|
81
|
+
) -> float:
|
|
82
|
+
"""
|
|
83
|
+
Simulate a single placebo ATE by randomly selecting treated units.
|
|
84
|
+
"""
|
|
85
|
+
rng = np.random.default_rng(seed)
|
|
86
|
+
N, _ = Y.shape
|
|
87
|
+
treated_units = rng.choice(N, size=n_treated_units, replace=False)
|
|
88
|
+
|
|
89
|
+
W = np.zeros_like(Y, dtype=float)
|
|
90
|
+
W[treated_units, -treated_periods:] = 1.0
|
|
91
|
+
|
|
92
|
+
return TROP_TWFE_average(
|
|
93
|
+
Y=Y,
|
|
94
|
+
W=W,
|
|
95
|
+
treated_units=treated_units,
|
|
96
|
+
lambda_unit=lambda_unit,
|
|
97
|
+
lambda_time=lambda_time,
|
|
98
|
+
lambda_nn=lambda_nn,
|
|
99
|
+
treated_periods=treated_periods,
|
|
100
|
+
solver=solver,
|
|
101
|
+
verbose=verbose,
|
|
102
|
+
)
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
def TROP_cv_single(
|
|
106
|
+
Y_control: ArrayLike,
|
|
107
|
+
n_treated_units: int,
|
|
108
|
+
treated_periods: int,
|
|
109
|
+
fixed_lambdas: Tuple[float, float] = (0.0, 0.0),
|
|
110
|
+
lambda_grid: Optional[Iterable[float]] = None,
|
|
111
|
+
lambda_cv: str = "unit",
|
|
112
|
+
*,
|
|
113
|
+
n_trials: int = 200,
|
|
114
|
+
n_jobs: int = -1,
|
|
115
|
+
prefer: str = "threads",
|
|
116
|
+
random_seed: int = 0,
|
|
117
|
+
solver: Optional[str] = None,
|
|
118
|
+
verbose: bool = False,
|
|
119
|
+
) -> float:
|
|
120
|
+
"""
|
|
121
|
+
Cross-validate one lambda parameter while keeping the other two fixed.
|
|
122
|
+
|
|
123
|
+
Parameters
|
|
124
|
+
----------
|
|
125
|
+
Y_control:
|
|
126
|
+
Control-only panel (N x T) used for placebo CV.
|
|
127
|
+
n_treated_units:
|
|
128
|
+
Number of placebo treated units to sample each trial.
|
|
129
|
+
treated_periods:
|
|
130
|
+
Number of placebo treated (post) periods (assumed final columns).
|
|
131
|
+
fixed_lambdas:
|
|
132
|
+
Tuple of two lambdas to hold fixed; interpretation depends on `lambda_cv`:
|
|
133
|
+
- lambda_cv='unit': fixed_lambdas=(lambda_time, lambda_nn)
|
|
134
|
+
- lambda_cv='time': fixed_lambdas=(lambda_unit, lambda_nn)
|
|
135
|
+
- lambda_cv='nn' : fixed_lambdas=(lambda_unit, lambda_time)
|
|
136
|
+
lambda_grid:
|
|
137
|
+
Grid of candidate values for the lambda being tuned.
|
|
138
|
+
If None, uses np.arange(0, 2, 0.2).
|
|
139
|
+
lambda_cv:
|
|
140
|
+
Which lambda to tune: {'unit','time','nn'}.
|
|
141
|
+
n_trials:
|
|
142
|
+
Number of placebo trials per lambda.
|
|
143
|
+
n_jobs:
|
|
144
|
+
joblib parallelism. -1 uses all available cores.
|
|
145
|
+
prefer:
|
|
146
|
+
joblib backend preference. Use 'threads' by default for solver stability.
|
|
147
|
+
random_seed:
|
|
148
|
+
Seed for generating trial seeds (deterministic CV).
|
|
149
|
+
solver, verbose:
|
|
150
|
+
Passed through to TROP_TWFE_average.
|
|
151
|
+
|
|
152
|
+
Returns
|
|
153
|
+
-------
|
|
154
|
+
float
|
|
155
|
+
Lambda value that minimizes RMSE of placebo ATEs.
|
|
156
|
+
"""
|
|
157
|
+
Y = np.asarray(Y_control, dtype=float)
|
|
158
|
+
_validate_panel(Y, treated_periods, n_treated_units)
|
|
159
|
+
|
|
160
|
+
if lambda_cv not in {"unit", "time", "nn"}:
|
|
161
|
+
raise ValueError("lambda_cv must be one of {'unit','time','nn'}.")
|
|
162
|
+
|
|
163
|
+
if lambda_grid is None:
|
|
164
|
+
lambda_grid_list = _as_list(np.arange(0.0, 2.0, 0.2))
|
|
165
|
+
else:
|
|
166
|
+
lambda_grid_list = _as_list(lambda_grid)
|
|
167
|
+
|
|
168
|
+
if n_trials <= 0:
|
|
169
|
+
raise ValueError("n_trials must be positive.")
|
|
170
|
+
if n_jobs == 0 or n_jobs < -1:
|
|
171
|
+
raise ValueError("n_jobs must be -1 or a positive integer.")
|
|
172
|
+
|
|
173
|
+
base_rng = np.random.default_rng(random_seed)
|
|
174
|
+
seeds = base_rng.integers(0, 2**32 - 1, size=n_trials, dtype=np.uint32)
|
|
175
|
+
|
|
176
|
+
scores: List[float] = []
|
|
177
|
+
|
|
178
|
+
for lamb in lambda_grid_list:
|
|
179
|
+
if lamb < 0:
|
|
180
|
+
raise ValueError("Lambda values must be nonnegative.")
|
|
181
|
+
|
|
182
|
+
if lambda_cv == "unit":
|
|
183
|
+
lambda_unit, lambda_time, lambda_nn = lamb, float(fixed_lambdas[0]), float(fixed_lambdas[1])
|
|
184
|
+
elif lambda_cv == "time":
|
|
185
|
+
lambda_unit, lambda_time, lambda_nn = float(fixed_lambdas[0]), lamb, float(fixed_lambdas[1])
|
|
186
|
+
else: # 'nn'
|
|
187
|
+
lambda_unit, lambda_time, lambda_nn = float(fixed_lambdas[0]), float(fixed_lambdas[1]), lamb
|
|
188
|
+
|
|
189
|
+
ates = Parallel(n_jobs=n_jobs, prefer=prefer)(
|
|
190
|
+
delayed(_simulate_ate)(
|
|
191
|
+
int(seed),
|
|
192
|
+
Y,
|
|
193
|
+
n_treated_units,
|
|
194
|
+
treated_periods,
|
|
195
|
+
lambda_unit,
|
|
196
|
+
lambda_time,
|
|
197
|
+
lambda_nn,
|
|
198
|
+
solver,
|
|
199
|
+
verbose,
|
|
200
|
+
)
|
|
201
|
+
for seed in seeds
|
|
202
|
+
)
|
|
203
|
+
|
|
204
|
+
ates_arr = np.asarray(ates, dtype=float)
|
|
205
|
+
ates_arr = ates_arr[np.isfinite(ates_arr)]
|
|
206
|
+
|
|
207
|
+
if ates_arr.size == 0:
|
|
208
|
+
raise RuntimeError(
|
|
209
|
+
f"All placebo trials failed or returned non-finite ATEs for lambda={lamb} "
|
|
210
|
+
f"(lambda_cv='{lambda_cv}'). Consider changing solver/settings."
|
|
211
|
+
)
|
|
212
|
+
|
|
213
|
+
scores.append(float(np.sqrt(np.mean(ates_arr**2))))
|
|
214
|
+
|
|
215
|
+
best_idx = int(np.argmin(scores))
|
|
216
|
+
return float(lambda_grid_list[best_idx])
|
|
217
|
+
|
|
218
|
+
|
|
219
|
+
def TROP_cv_cycle(
|
|
220
|
+
Y_control: ArrayLike,
|
|
221
|
+
n_treated_units: int,
|
|
222
|
+
treated_periods: int,
|
|
223
|
+
unit_grid: Sequence[float],
|
|
224
|
+
time_grid: Sequence[float],
|
|
225
|
+
nn_grid: Sequence[float],
|
|
226
|
+
lambdas_init: Optional[Tuple[float, float, float]] = None,
|
|
227
|
+
*,
|
|
228
|
+
max_iter: int = 50,
|
|
229
|
+
n_trials: int = 200,
|
|
230
|
+
n_jobs: int = -1,
|
|
231
|
+
prefer: str = "threads",
|
|
232
|
+
random_seed: int = 0,
|
|
233
|
+
solver: Optional[str] = None,
|
|
234
|
+
verbose: bool = False,
|
|
235
|
+
) -> Tuple[float, float, float]:
|
|
236
|
+
"""
|
|
237
|
+
Coordinate-descent style cross-validation for (lambda_unit, lambda_time, lambda_nn).
|
|
238
|
+
|
|
239
|
+
This routine alternates between optimizing lambda_unit, lambda_time, and lambda_nn
|
|
240
|
+
(via `TROP_cv_single`) while holding the other two fixed, until it reaches a fixed
|
|
241
|
+
point (no change in the selected lambdas) or until `max_iter` iterations are reached.
|
|
242
|
+
|
|
243
|
+
Parameters
|
|
244
|
+
----------
|
|
245
|
+
Y_control:
|
|
246
|
+
Control-only panel (N x T) used for placebo CV.
|
|
247
|
+
n_treated_units:
|
|
248
|
+
Number of placebo treated units to sample each trial.
|
|
249
|
+
treated_periods:
|
|
250
|
+
Number of placebo treated (post) periods (assumed final columns).
|
|
251
|
+
unit_grid:
|
|
252
|
+
Grid of candidate values for lambda_unit (unit-distance decay).
|
|
253
|
+
time_grid:
|
|
254
|
+
Grid of candidate values for lambda_time (time-distance decay).
|
|
255
|
+
nn_grid:
|
|
256
|
+
Grid of candidate values for lambda_nn (nuclear-norm penalty).
|
|
257
|
+
lambdas_init:
|
|
258
|
+
Optional initial values (lambda_unit, lambda_time, lambda_nn). If None, initializes
|
|
259
|
+
each lambda to the mean of its corresponding grid.
|
|
260
|
+
max_iter:
|
|
261
|
+
Maximum number of coordinate-descent iterations.
|
|
262
|
+
n_trials:
|
|
263
|
+
Number of placebo trials per grid point in each coordinate update.
|
|
264
|
+
n_jobs:
|
|
265
|
+
joblib parallelism. -1 uses all available cores.
|
|
266
|
+
prefer:
|
|
267
|
+
joblib backend preference. Use 'threads' by default for solver stability.
|
|
268
|
+
random_seed:
|
|
269
|
+
Seed for generating trial seeds (deterministic CV).
|
|
270
|
+
solver, verbose:
|
|
271
|
+
Passed through to TROP_TWFE_average.
|
|
272
|
+
|
|
273
|
+
Returns
|
|
274
|
+
-------
|
|
275
|
+
Tuple[float, float, float]
|
|
276
|
+
(lambda_unit, lambda_time, lambda_nn) at the converged fixed point.
|
|
277
|
+
|
|
278
|
+
Raises
|
|
279
|
+
------
|
|
280
|
+
RuntimeError
|
|
281
|
+
If the procedure does not converge to a fixed point within `max_iter`.
|
|
282
|
+
"""
|
|
283
|
+
|
|
284
|
+
Y = np.asarray(Y_control, dtype=float)
|
|
285
|
+
_validate_panel(Y, treated_periods, n_treated_units)
|
|
286
|
+
|
|
287
|
+
unit_grid_list = _as_list(unit_grid)
|
|
288
|
+
time_grid_list = _as_list(time_grid)
|
|
289
|
+
nn_grid_list = _as_list(nn_grid)
|
|
290
|
+
|
|
291
|
+
if lambdas_init is None:
|
|
292
|
+
lambda_unit = float(np.mean(unit_grid_list))
|
|
293
|
+
lambda_time = float(np.mean(time_grid_list))
|
|
294
|
+
lambda_nn = float(np.mean(nn_grid_list))
|
|
295
|
+
else:
|
|
296
|
+
lambda_unit, lambda_time, lambda_nn = map(float, lambdas_init)
|
|
297
|
+
|
|
298
|
+
for _ in range(max_iter):
|
|
299
|
+
old = (lambda_unit, lambda_time, lambda_nn)
|
|
300
|
+
|
|
301
|
+
lambda_unit = TROP_cv_single(
|
|
302
|
+
Y, n_treated_units, treated_periods,
|
|
303
|
+
fixed_lambdas=(lambda_time, lambda_nn),
|
|
304
|
+
lambda_grid=unit_grid_list,
|
|
305
|
+
lambda_cv="unit",
|
|
306
|
+
n_trials=n_trials, n_jobs=n_jobs, prefer=prefer,
|
|
307
|
+
random_seed=random_seed, solver=solver, verbose=verbose
|
|
308
|
+
)
|
|
309
|
+
|
|
310
|
+
lambda_time = TROP_cv_single(
|
|
311
|
+
Y, n_treated_units, treated_periods,
|
|
312
|
+
fixed_lambdas=(lambda_unit, lambda_nn),
|
|
313
|
+
lambda_grid=time_grid_list,
|
|
314
|
+
lambda_cv="time",
|
|
315
|
+
n_trials=n_trials, n_jobs=n_jobs, prefer=prefer,
|
|
316
|
+
random_seed=random_seed, solver=solver, verbose=verbose
|
|
317
|
+
)
|
|
318
|
+
|
|
319
|
+
lambda_nn = TROP_cv_single(
|
|
320
|
+
Y, n_treated_units, treated_periods,
|
|
321
|
+
fixed_lambdas=(lambda_unit, lambda_time),
|
|
322
|
+
lambda_grid=nn_grid_list,
|
|
323
|
+
lambda_cv="nn",
|
|
324
|
+
n_trials=n_trials, n_jobs=n_jobs, prefer=prefer,
|
|
325
|
+
random_seed=random_seed, solver=solver, verbose=verbose
|
|
326
|
+
)
|
|
327
|
+
|
|
328
|
+
new = (lambda_unit, lambda_time, lambda_nn)
|
|
329
|
+
if new == old:
|
|
330
|
+
return new
|
|
331
|
+
|
|
332
|
+
raise RuntimeError("TROP_cv_cycle did not converge (no fixed point) within max_iter.")
|
|
333
|
+
|
|
334
|
+
|
|
335
|
+
def TROP_cv_joint(
|
|
336
|
+
Y_control: ArrayLike,
|
|
337
|
+
n_treated_units: int,
|
|
338
|
+
treated_periods: int,
|
|
339
|
+
unit_grid: Sequence[float],
|
|
340
|
+
time_grid: Sequence[float],
|
|
341
|
+
nn_grid: Sequence[float],
|
|
342
|
+
*,
|
|
343
|
+
n_trials: int = 200,
|
|
344
|
+
n_jobs: int = -1,
|
|
345
|
+
prefer: str = "threads",
|
|
346
|
+
random_seed: int = 0,
|
|
347
|
+
solver: Optional[str] = None,
|
|
348
|
+
verbose: bool = False,
|
|
349
|
+
) -> Tuple[float, float, float]:
|
|
350
|
+
"""
|
|
351
|
+
Joint grid search over (lambda_unit, lambda_time, lambda_nn).
|
|
352
|
+
|
|
353
|
+
Parameters
|
|
354
|
+
----------
|
|
355
|
+
Y_control:
|
|
356
|
+
Control-only panel (N x T) used for placebo CV.
|
|
357
|
+
n_treated_units:
|
|
358
|
+
Number of placebo treated units to sample each trial.
|
|
359
|
+
treated_periods:
|
|
360
|
+
Number of placebo treated (post) periods (assumed final columns).
|
|
361
|
+
unit_grid:
|
|
362
|
+
Grid of candidate values for lambda_unit (unit-distance decay).
|
|
363
|
+
time_grid:
|
|
364
|
+
Grid of candidate values for lambda_time (time-distance decay).
|
|
365
|
+
nn_grid:
|
|
366
|
+
Grid of candidate values for lambda_nn (nuclear-norm penalty).
|
|
367
|
+
n_trials:
|
|
368
|
+
Number of placebo trials per (lambda_unit, lambda_time, lambda_nn) triple.
|
|
369
|
+
n_jobs:
|
|
370
|
+
joblib parallelism. -1 uses all available cores.
|
|
371
|
+
prefer:
|
|
372
|
+
joblib backend preference. Use 'threads' by default for solver stability.
|
|
373
|
+
random_seed:
|
|
374
|
+
Seed for generating trial seeds (deterministic CV).
|
|
375
|
+
solver, verbose:
|
|
376
|
+
Passed through to TROP_TWFE_average.
|
|
377
|
+
|
|
378
|
+
Returns
|
|
379
|
+
-------
|
|
380
|
+
Tuple[float, float, float]
|
|
381
|
+
(lambda_unit, lambda_time, lambda_nn) triple that minimizes the RMSE of placebo ATEs.
|
|
382
|
+
"""
|
|
383
|
+
|
|
384
|
+
Y = np.asarray(Y_control, dtype=float)
|
|
385
|
+
_validate_panel(Y, treated_periods, n_treated_units)
|
|
386
|
+
|
|
387
|
+
unit_grid_list = _as_list(unit_grid)
|
|
388
|
+
time_grid_list = _as_list(time_grid)
|
|
389
|
+
nn_grid_list = _as_list(nn_grid)
|
|
390
|
+
|
|
391
|
+
base_rng = np.random.default_rng(random_seed)
|
|
392
|
+
seeds = base_rng.integers(0, 2**32 - 1, size=n_trials, dtype=np.uint32)
|
|
393
|
+
|
|
394
|
+
best_params: Optional[Tuple[float, float, float]] = None
|
|
395
|
+
best_score: float = float("inf")
|
|
396
|
+
|
|
397
|
+
for lambda_unit in unit_grid_list:
|
|
398
|
+
for lambda_time in time_grid_list:
|
|
399
|
+
for lambda_nn in nn_grid_list:
|
|
400
|
+
ates = Parallel(n_jobs=n_jobs, prefer=prefer)(
|
|
401
|
+
delayed(_simulate_ate)(
|
|
402
|
+
int(seed),
|
|
403
|
+
Y,
|
|
404
|
+
n_treated_units,
|
|
405
|
+
treated_periods,
|
|
406
|
+
float(lambda_unit),
|
|
407
|
+
float(lambda_time),
|
|
408
|
+
float(lambda_nn),
|
|
409
|
+
solver,
|
|
410
|
+
verbose,
|
|
411
|
+
)
|
|
412
|
+
for seed in seeds
|
|
413
|
+
)
|
|
414
|
+
|
|
415
|
+
ates_arr = np.asarray(ates, dtype=float)
|
|
416
|
+
ates_arr = ates_arr[np.isfinite(ates_arr)]
|
|
417
|
+
if ates_arr.size == 0:
|
|
418
|
+
continue # skip invalid setting
|
|
419
|
+
|
|
420
|
+
score = float(np.sqrt(np.mean(ates_arr**2)))
|
|
421
|
+
if score < best_score:
|
|
422
|
+
best_score = score
|
|
423
|
+
best_params = (float(lambda_unit), float(lambda_time), float(lambda_nn))
|
|
424
|
+
|
|
425
|
+
if best_params is None:
|
|
426
|
+
raise RuntimeError("All parameter combinations failed during joint CV. Check solver/settings.")
|
|
427
|
+
return best_params
|
trop/estimator.py
ADDED
|
@@ -0,0 +1,157 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import math
|
|
4
|
+
from typing import Iterable, Optional, Sequence, Union
|
|
5
|
+
|
|
6
|
+
import numpy as np
|
|
7
|
+
import cvxpy as cp
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
ArrayLike = Union[np.ndarray, Sequence[Sequence[float]]]
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def TROP_TWFE_average(
|
|
14
|
+
Y: ArrayLike,
|
|
15
|
+
W: ArrayLike,
|
|
16
|
+
treated_units: Sequence[int],
|
|
17
|
+
lambda_unit: float,
|
|
18
|
+
lambda_time: float,
|
|
19
|
+
lambda_nn: float,
|
|
20
|
+
treated_periods: int = 10,
|
|
21
|
+
solver: Optional[str] = None,
|
|
22
|
+
verbose: bool = False,
|
|
23
|
+
) -> float:
|
|
24
|
+
"""
|
|
25
|
+
Triply Robust Panel (TROP) estimator in a TWFE framework with:
|
|
26
|
+
- distance-based unit weights (lambda_unit)
|
|
27
|
+
- distance-based time weights (lambda_time)
|
|
28
|
+
- optional low-rank regression adjustment (nuclear norm penalty, lambda_nn)
|
|
29
|
+
|
|
30
|
+
Parameters
|
|
31
|
+
----------
|
|
32
|
+
Y:
|
|
33
|
+
Outcome matrix of shape (N, T).
|
|
34
|
+
W:
|
|
35
|
+
Treatment indicator matrix of shape (N, T). Typically binary {0,1}.
|
|
36
|
+
Convention in the provided codebase: treated units are treated in the final
|
|
37
|
+
`treated_periods` columns, but the function will run for any W.
|
|
38
|
+
treated_units:
|
|
39
|
+
Indices of treated units (row indices into Y/W).
|
|
40
|
+
lambda_unit:
|
|
41
|
+
Nonnegative tuning parameter controlling unit-weight decay.
|
|
42
|
+
lambda_time:
|
|
43
|
+
Nonnegative tuning parameter controlling time-weight decay.
|
|
44
|
+
lambda_nn:
|
|
45
|
+
Nuclear norm penalty weight for the low-rank component L.
|
|
46
|
+
Use np.inf to disable low-rank adjustment (i.e., no L term).
|
|
47
|
+
treated_periods:
|
|
48
|
+
Number of treated (post) periods at the end of the panel used to define:
|
|
49
|
+
- the time-distance center
|
|
50
|
+
- the pre-period mask for unit distance computation
|
|
51
|
+
solver:
|
|
52
|
+
Optional CVXPY solver name. If None, chooses a reasonable default:
|
|
53
|
+
- if lambda_nn is finite: uses SCS (supports nuclear norm / SDP forms)
|
|
54
|
+
- if lambda_nn is infinite: uses OSQP (fast for pure quadratic problems)
|
|
55
|
+
verbose:
|
|
56
|
+
Passed to CVXPY solve().
|
|
57
|
+
|
|
58
|
+
Returns
|
|
59
|
+
-------
|
|
60
|
+
float
|
|
61
|
+
Estimated average treatment effect tau.
|
|
62
|
+
|
|
63
|
+
"""
|
|
64
|
+
Y = np.asarray(Y, dtype=float)
|
|
65
|
+
W = np.asarray(W, dtype=float)
|
|
66
|
+
|
|
67
|
+
if Y.ndim != 2 or W.ndim != 2:
|
|
68
|
+
raise ValueError(f"Y and W must be 2D arrays. Got Y.ndim={Y.ndim}, W.ndim={W.ndim}.")
|
|
69
|
+
if Y.shape != W.shape:
|
|
70
|
+
raise ValueError(f"Y and W must have the same shape. Got Y={Y.shape}, W={W.shape}.")
|
|
71
|
+
|
|
72
|
+
N, T = Y.shape
|
|
73
|
+
|
|
74
|
+
if not isinstance(treated_periods, int) or treated_periods <= 0:
|
|
75
|
+
raise ValueError("treated_periods must be a positive integer.")
|
|
76
|
+
if treated_periods >= T:
|
|
77
|
+
raise ValueError(f"treated_periods must be < T. Got treated_periods={treated_periods}, T={T}.")
|
|
78
|
+
|
|
79
|
+
treated_units_arr = np.asarray(treated_units, dtype=int)
|
|
80
|
+
if treated_units_arr.size == 0:
|
|
81
|
+
raise ValueError("treated_units must contain at least one unit index.")
|
|
82
|
+
if np.any(treated_units_arr < 0) or np.any(treated_units_arr >= N):
|
|
83
|
+
raise ValueError(f"treated_units contains out-of-range indices for N={N}: {treated_units_arr}")
|
|
84
|
+
|
|
85
|
+
if lambda_unit < 0 or lambda_time < 0:
|
|
86
|
+
raise ValueError("lambda_unit and lambda_time should be nonnegative.")
|
|
87
|
+
|
|
88
|
+
# ---------------------------------------------------------------------
|
|
89
|
+
# Distance-based time weights
|
|
90
|
+
# ---------------------------------------------------------------------
|
|
91
|
+
# Distance to the center of the treated block near the end of the panel.
|
|
92
|
+
# dist_time = abs(arange(T) - (T - treated_periods/2))
|
|
93
|
+
center = T - treated_periods / 2.0
|
|
94
|
+
dist_time = np.abs(np.arange(T, dtype=float) - center)
|
|
95
|
+
|
|
96
|
+
# ---------------------------------------------------------------------
|
|
97
|
+
# Distance-based unit weights
|
|
98
|
+
# ---------------------------------------------------------------------
|
|
99
|
+
average_treated = np.mean(Y[treated_units_arr, :], axis=0)
|
|
100
|
+
|
|
101
|
+
# Pre-period mask: 1 in pre, 0 in treated/post
|
|
102
|
+
mask = np.ones((N, T), dtype=float)
|
|
103
|
+
mask[:, -treated_periods:] = 0.0
|
|
104
|
+
|
|
105
|
+
# RMS distance to average treated trajectory over pre-periods
|
|
106
|
+
# dist_unit[i] = sqrt( sum_pre (avg_tr - Y_i)^2 / (#pre) )
|
|
107
|
+
A = np.sum(((average_treated - Y) ** 2) * mask, axis=1)
|
|
108
|
+
B = np.sum(mask, axis=1)
|
|
109
|
+
|
|
110
|
+
if np.any(B == 0):
|
|
111
|
+
raise ValueError(
|
|
112
|
+
"Pre-period mask has zero pre-periods for at least one unit."
|
|
113
|
+
)
|
|
114
|
+
|
|
115
|
+
dist_unit = np.sqrt(A / B)
|
|
116
|
+
|
|
117
|
+
# Convert distances to weights
|
|
118
|
+
delta_unit = np.exp(-lambda_unit * dist_unit) # shape (N,)
|
|
119
|
+
delta_time = np.exp(-lambda_time * dist_time) # shape (T,)
|
|
120
|
+
delta = np.outer(delta_unit, delta_time) # shape (N, T)
|
|
121
|
+
|
|
122
|
+
# ---------------------------------------------------------------------
|
|
123
|
+
# CVXPY problem: weighted TWFE
|
|
124
|
+
# ---------------------------------------------------------------------
|
|
125
|
+
unit_effects = cp.Variable((1, N))
|
|
126
|
+
time_effects = cp.Variable((1, T))
|
|
127
|
+
mu = cp.Variable() # intercept
|
|
128
|
+
tau = cp.Variable() # treatment effect
|
|
129
|
+
|
|
130
|
+
# Broadcast TWFE components to N x T
|
|
131
|
+
unit_factor = cp.kron(np.ones((T, 1)), unit_effects).T
|
|
132
|
+
time_factor = cp.kron(np.ones((N, 1)), time_effects)
|
|
133
|
+
|
|
134
|
+
is_low_rank = not math.isinf(float(lambda_nn))
|
|
135
|
+
|
|
136
|
+
if is_low_rank:
|
|
137
|
+
L = cp.Variable((N, T))
|
|
138
|
+
residual = Y - mu - unit_factor - time_factor - L - W * tau
|
|
139
|
+
loss = cp.sum_squares(cp.multiply(residual, delta)) + float(lambda_nn) * cp.norm(L, "nuc")
|
|
140
|
+
default_solver = "SCS" # robust choice for nuclear norm problems
|
|
141
|
+
else:
|
|
142
|
+
residual = Y - mu - unit_factor - time_factor - W * tau
|
|
143
|
+
loss = cp.sum_squares(cp.multiply(residual, delta))
|
|
144
|
+
default_solver = "OSQP" # fast for pure quadratic objective
|
|
145
|
+
|
|
146
|
+
prob = cp.Problem(cp.Minimize(loss))
|
|
147
|
+
|
|
148
|
+
chosen_solver = solver or default_solver
|
|
149
|
+
prob.solve(solver=chosen_solver, verbose=verbose)
|
|
150
|
+
|
|
151
|
+
if tau.value is None or not np.isfinite(tau.value):
|
|
152
|
+
raise RuntimeError(
|
|
153
|
+
"Optimization did not return a valid tau. "
|
|
154
|
+
f"Solver={chosen_solver}, status={prob.status}."
|
|
155
|
+
)
|
|
156
|
+
|
|
157
|
+
return float(tau.value)
|
|
@@ -0,0 +1,53 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: trop
|
|
3
|
+
Version: 0.1.2
|
|
4
|
+
Summary: Triply Robust Panel (TROP) estimator: weighted TWFE with optional low-rank adjustment.
|
|
5
|
+
Author: Susan Athey, Guido Imbens, Zhaonan Qu, Davide Viviano
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/zhaonanq/TROP
|
|
8
|
+
Project-URL: Repository, https://github.com/zhaonanq/TROP
|
|
9
|
+
Project-URL: Issues, https://github.com/zhaonanq/TROP/issues
|
|
10
|
+
Keywords: causal-inference,panel-data,factor-models,difference-in-differences,synthetic-control,synthetic-controls,trop,twfe
|
|
11
|
+
Classifier: Development Status :: 3 - Alpha
|
|
12
|
+
Classifier: Intended Audience :: Science/Research
|
|
13
|
+
Classifier: Operating System :: OS Independent
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
20
|
+
Classifier: Topic :: Scientific/Engineering
|
|
21
|
+
Classifier: Topic :: Scientific/Engineering :: Mathematics
|
|
22
|
+
Requires-Python: >=3.9
|
|
23
|
+
Description-Content-Type: text/markdown
|
|
24
|
+
License-File: LICENSE
|
|
25
|
+
Requires-Dist: numpy>=1.23
|
|
26
|
+
Requires-Dist: cvxpy>=1.4
|
|
27
|
+
Requires-Dist: joblib>=1.2
|
|
28
|
+
Requires-Dist: osqp>=0.6.5
|
|
29
|
+
Requires-Dist: scs>=3.2.4
|
|
30
|
+
Provides-Extra: dev
|
|
31
|
+
Requires-Dist: pytest>=7.0; extra == "dev"
|
|
32
|
+
Requires-Dist: ruff>=0.5.0; extra == "dev"
|
|
33
|
+
Requires-Dist: black>=24.0.0; extra == "dev"
|
|
34
|
+
Dynamic: license-file
|
|
35
|
+
|
|
36
|
+
# TROP: Triply Robust Panel Estimator
|
|
37
|
+
|
|
38
|
+
`trop` is a Python package implementing the **Triply Robust Panel (TROP)** estimator for average treatment effects (ATEs) in panel data. The core estimator is expressed as a weighted two-way fixed effects (TWFE) objective, with an optional low-rank regression adjustment via a nuclear-norm penalty.
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
Reference:
|
|
42
|
+
|
|
43
|
+
> Susan Athey, Guido Imbens, Zhaonan Qu, Davide Viviano (2025).
|
|
44
|
+
> *Triply Robust Panel Estimators*.
|
|
45
|
+
> arXiv:2508.21536.
|
|
46
|
+
|
|
47
|
+
---
|
|
48
|
+
|
|
49
|
+
## Installation
|
|
50
|
+
|
|
51
|
+
```
|
|
52
|
+
pip install trop
|
|
53
|
+
```
|
|
@@ -0,0 +1,8 @@
|
|
|
1
|
+
trop/__init__.py,sha256=B94vrDZevg2l6ijN4lut7wo0MYTEtBTTZhqmMQtq7Qg,205
|
|
2
|
+
trop/cv.py,sha256=7tumpAaAiiK_F8nUUAcAwnecNAfc2XghftgJsNWVsAQ,14017
|
|
3
|
+
trop/estimator.py,sha256=FWMO39GbL6k3Vz5g1V7SpR6t5wP3N81V5gGSFIe65Xw,6001
|
|
4
|
+
trop-0.1.2.dist-info/licenses/LICENSE,sha256=VqjvjioQz04uLYBj4ye0x-_Ss77-WTIuEWWCW_awEz8,1065
|
|
5
|
+
trop-0.1.2.dist-info/METADATA,sha256=YSyeONhxn4JO_NZBKT1ubVa5J96BhlWGtmRi-GJ0rLk,1997
|
|
6
|
+
trop-0.1.2.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
|
|
7
|
+
trop-0.1.2.dist-info/top_level.txt,sha256=jaqQZFm3D5B4vPBAKZtXfEAYnpl9FKsNHqlM49kcwTI,5
|
|
8
|
+
trop-0.1.2.dist-info/RECORD,,
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 zhaonanq
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
trop
|
trop/placeholder.py
DELETED
|
@@ -1 +0,0 @@
|
|
|
1
|
-
# placeholder to test package publishing
|
trop-0.1.0.dist-info/METADATA
DELETED
|
@@ -1,24 +0,0 @@
|
|
|
1
|
-
Metadata-Version: 2.4
|
|
2
|
-
Name: trop
|
|
3
|
-
Version: 0.1.0
|
|
4
|
-
Summary: Package that implements the Triply Robust Panel Estimator (TROP)
|
|
5
|
-
License: MIT
|
|
6
|
-
Author: meganndare, zhaonanq
|
|
7
|
-
Requires-Python: >=3.9
|
|
8
|
-
Classifier: License :: OSI Approved :: MIT License
|
|
9
|
-
Classifier: Programming Language :: Python :: 3
|
|
10
|
-
Classifier: Programming Language :: Python :: 3.9
|
|
11
|
-
Classifier: Programming Language :: Python :: 3.10
|
|
12
|
-
Classifier: Programming Language :: Python :: 3.11
|
|
13
|
-
Classifier: Programming Language :: Python :: 3.12
|
|
14
|
-
Classifier: Programming Language :: Python :: 3.13
|
|
15
|
-
Classifier: Programming Language :: Python :: 3.14
|
|
16
|
-
Requires-Dist: cvxpy (==1.4.1)
|
|
17
|
-
Requires-Dist: joblib (==1.3.2)
|
|
18
|
-
Requires-Dist: numpy (==1.26.2)
|
|
19
|
-
Requires-Dist: pandas (==2.1.3)
|
|
20
|
-
Description-Content-Type: text/markdown
|
|
21
|
-
|
|
22
|
-
# Triply Robust Panel Estimators (TROP)
|
|
23
|
-
|
|
24
|
-
This package will soon contain the replication files and implementation of the TROP estimator.
|
trop-0.1.0.dist-info/RECORD
DELETED
|
@@ -1,4 +0,0 @@
|
|
|
1
|
-
trop/placeholder.py,sha256=TIbWca9SwI6pN7kUOp932p_piNfZKYh7Zfnrm-Uqyxg,40
|
|
2
|
-
trop-0.1.0.dist-info/METADATA,sha256=Kk7ZnGxm-OplNoUeGONkbG_LFl0tsjOCOMPehQPihGA,895
|
|
3
|
-
trop-0.1.0.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
|
|
4
|
-
trop-0.1.0.dist-info/RECORD,,
|