trop 0.1.0__tar.gz → 0.1.2__tar.gz

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-0.1.2/LICENSE ADDED
@@ -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.
trop-0.1.2/PKG-INFO ADDED
@@ -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
+ ```
trop-0.1.2/README.md ADDED
@@ -0,0 +1,18 @@
1
+ # TROP: Triply Robust Panel Estimator
2
+
3
+ `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.
4
+
5
+
6
+ Reference:
7
+
8
+ > Susan Athey, Guido Imbens, Zhaonan Qu, Davide Viviano (2025).
9
+ > *Triply Robust Panel Estimators*.
10
+ > arXiv:2508.21536.
11
+
12
+ ---
13
+
14
+ ## Installation
15
+
16
+ ```
17
+ pip install trop
18
+ ```
@@ -0,0 +1,74 @@
1
+ [build-system]
2
+ requires = ["setuptools>=77", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "trop"
7
+ version = "0.1.2"
8
+ description = "Triply Robust Panel (TROP) estimator: weighted TWFE with optional low-rank adjustment."
9
+ readme = "README.md"
10
+ requires-python = ">=3.9"
11
+ license = "MIT"
12
+ license-files = ["LICENSE"]
13
+
14
+ authors = [
15
+ { name = "Susan Athey" },
16
+ { name = "Guido Imbens" },
17
+ { name = "Zhaonan Qu" },
18
+ { name = "Davide Viviano" }
19
+ ]
20
+
21
+ keywords = [
22
+ "causal-inference",
23
+ "panel-data",
24
+ "factor-models",
25
+ "difference-in-differences",
26
+ "synthetic-control",
27
+ "synthetic-controls",
28
+ "trop",
29
+ "twfe"
30
+ ]
31
+
32
+ classifiers = [
33
+ "Development Status :: 3 - Alpha",
34
+ "Intended Audience :: Science/Research",
35
+ "Operating System :: OS Independent",
36
+ "Programming Language :: Python :: 3",
37
+ "Programming Language :: Python :: 3 :: Only",
38
+ "Programming Language :: Python :: 3.9",
39
+ "Programming Language :: Python :: 3.10",
40
+ "Programming Language :: Python :: 3.11",
41
+ "Programming Language :: Python :: 3.12",
42
+ "Topic :: Scientific/Engineering",
43
+ "Topic :: Scientific/Engineering :: Mathematics",
44
+ ]
45
+
46
+ dependencies = [
47
+ "numpy>=1.23",
48
+ "cvxpy>=1.4",
49
+ "joblib>=1.2",
50
+ "osqp>=0.6.5",
51
+ "scs>=3.2.4",
52
+ ]
53
+
54
+ [project.urls]
55
+ Homepage = "https://github.com/zhaonanq/TROP"
56
+ Repository = "https://github.com/zhaonanq/TROP"
57
+ Issues = "https://github.com/zhaonanq/TROP/issues"
58
+
59
+ [project.optional-dependencies]
60
+ dev = [
61
+ "pytest>=7.0",
62
+ "ruff>=0.5.0",
63
+ "black>=24.0.0",
64
+ ]
65
+
66
+ [tool.setuptools]
67
+ package-dir = {"" = "src"}
68
+
69
+ [tool.setuptools.packages.find]
70
+ where = ["src"]
71
+
72
+ [tool.pytest.ini_options]
73
+ testpaths = ["tests"]
74
+ addopts = "-q"
trop-0.1.2/setup.cfg ADDED
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,9 @@
1
+ from .estimator import TROP_TWFE_average
2
+ from .cv import TROP_cv_single, TROP_cv_cycle, TROP_cv_joint
3
+
4
+ __all__ = [
5
+ "TROP_TWFE_average",
6
+ "TROP_cv_single",
7
+ "TROP_cv_cycle",
8
+ "TROP_cv_joint",
9
+ ]
@@ -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
@@ -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,12 @@
1
+ LICENSE
2
+ README.md
3
+ pyproject.toml
4
+ src/trop/__init__.py
5
+ src/trop/cv.py
6
+ src/trop/estimator.py
7
+ src/trop.egg-info/PKG-INFO
8
+ src/trop.egg-info/SOURCES.txt
9
+ src/trop.egg-info/dependency_links.txt
10
+ src/trop.egg-info/requires.txt
11
+ src/trop.egg-info/top_level.txt
12
+ tests/test_smoke_trop.py
@@ -0,0 +1,10 @@
1
+ numpy>=1.23
2
+ cvxpy>=1.4
3
+ joblib>=1.2
4
+ osqp>=0.6.5
5
+ scs>=3.2.4
6
+
7
+ [dev]
8
+ pytest>=7.0
9
+ ruff>=0.5.0
10
+ black>=24.0.0
@@ -0,0 +1 @@
1
+ trop
@@ -0,0 +1,46 @@
1
+ import numpy as np
2
+ import pytest
3
+
4
+ from trop import TROP_TWFE_average
5
+
6
+
7
+ def test_trop_twfe_average_smoke():
8
+ # Small synthetic panel
9
+ N, T = 8, 20
10
+ treated_periods = 5
11
+ treated_units = [0, 1]
12
+
13
+ rng = np.random.default_rng(0)
14
+ Y = rng.normal(size=(N, T))
15
+
16
+ # Treatment: treated_units get treated in last treated_periods
17
+ W = np.zeros((N, T))
18
+ W[treated_units, -treated_periods:] = 1.0
19
+
20
+ tau = TROP_TWFE_average(
21
+ Y=Y,
22
+ W=W,
23
+ treated_units=treated_units,
24
+ lambda_unit=0.1,
25
+ lambda_time=0.1,
26
+ lambda_nn=np.inf, # simplest path: no nuclear-norm component
27
+ treated_periods=treated_periods,
28
+ )
29
+
30
+ assert np.isfinite(tau), "tau should be a finite scalar"
31
+ assert isinstance(tau, float)
32
+
33
+
34
+ def test_invalid_shapes_raises():
35
+ Y = np.zeros((5, 10))
36
+ W = np.zeros((5, 9)) # mismatched
37
+ with pytest.raises(ValueError):
38
+ TROP_TWFE_average(
39
+ Y=Y,
40
+ W=W,
41
+ treated_units=[0],
42
+ lambda_unit=0.1,
43
+ lambda_time=0.1,
44
+ lambda_nn=np.inf,
45
+ treated_periods=2,
46
+ )
trop-0.1.0/PKG-INFO 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/README.md DELETED
@@ -1,3 +0,0 @@
1
- # Triply Robust Panel Estimators (TROP)
2
-
3
- This package will soon contain the replication files and implementation of the TROP estimator.
trop-0.1.0/pyproject.toml DELETED
@@ -1,21 +0,0 @@
1
- [project]
2
- name = "trop"
3
- version = "0.1.0"
4
- description = "Package that implements the Triply Robust Panel Estimator (TROP)"
5
- authors = [
6
- {name = "meganndare, zhaonanq"}
7
- ]
8
- license = {text = "MIT"}
9
- readme = "README.md"
10
- requires-python = ">=3.9"
11
- dependencies = [
12
- "numpy (==1.26.2)",
13
- "pandas (==2.1.3)",
14
- "cvxpy (==1.4.1)",
15
- "joblib (==1.3.2)"
16
- ]
17
-
18
-
19
- [build-system]
20
- requires = ["poetry-core>=2.0.0,<3.0.0"]
21
- build-backend = "poetry.core.masonry.api"
@@ -1 +0,0 @@
1
- # placeholder to test package publishing