trop 0.1.0__tar.gz → 0.1.2__tar.gz

trop-0.1.2/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 zhaonanq
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

trop-0.1.2/PKG-INFO

@@ -0,0 +1,53 @@
+Metadata-Version: 2.4
+Name: trop
+Version: 0.1.2
+Summary: Triply Robust Panel (TROP) estimator: weighted TWFE with optional low-rank adjustment.
+Author: Susan Athey, Guido Imbens, Zhaonan Qu, Davide Viviano
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/zhaonanq/TROP
+Project-URL: Repository, https://github.com/zhaonanq/TROP
+Project-URL: Issues, https://github.com/zhaonanq/TROP/issues
+Keywords: causal-inference,panel-data,factor-models,difference-in-differences,synthetic-control,synthetic-controls,trop,twfe
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.23
+Requires-Dist: cvxpy>=1.4
+Requires-Dist: joblib>=1.2
+Requires-Dist: osqp>=0.6.5
+Requires-Dist: scs>=3.2.4
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: ruff>=0.5.0; extra == "dev"
+Requires-Dist: black>=24.0.0; extra == "dev"
+Dynamic: license-file
+
+# TROP: Triply Robust Panel Estimator
+
+`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.
+
+
+Reference:
+
+> Susan Athey, Guido Imbens, Zhaonan Qu, Davide Viviano (2025).  
+> *Triply Robust Panel Estimators*.  
+> arXiv:2508.21536.
+
+---
+
+## Installation
+
+```
+pip install trop
+```

trop-0.1.2/README.md

@@ -0,0 +1,18 @@
+# TROP: Triply Robust Panel Estimator
+
+`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.
+
+
+Reference:
+
+> Susan Athey, Guido Imbens, Zhaonan Qu, Davide Viviano (2025).  
+> *Triply Robust Panel Estimators*.  
+> arXiv:2508.21536.
+
+---
+
+## Installation
+
+```
+pip install trop
+```

trop-0.1.2/pyproject.toml

@@ -0,0 +1,74 @@
+[build-system]
+requires = ["setuptools>=77", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "trop"
+version = "0.1.2"
+description = "Triply Robust Panel (TROP) estimator: weighted TWFE with optional low-rank adjustment."
+readme = "README.md"
+requires-python = ">=3.9"
+license = "MIT"
+license-files = ["LICENSE"]
+
+authors = [
+  { name = "Susan Athey" },
+  { name = "Guido Imbens" },
+  { name = "Zhaonan Qu" },
+  { name = "Davide Viviano" }
+]
+
+keywords = [
+  "causal-inference",
+  "panel-data",
+  "factor-models",
+  "difference-in-differences",
+  "synthetic-control",
+  "synthetic-controls",
+  "trop",
+  "twfe"
+]
+
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Intended Audience :: Science/Research",
+  "Operating System :: OS Independent",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3 :: Only",
+  "Programming Language :: Python :: 3.9",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Topic :: Scientific/Engineering",
+  "Topic :: Scientific/Engineering :: Mathematics",
+]
+
+dependencies = [
+  "numpy>=1.23",
+  "cvxpy>=1.4",
+  "joblib>=1.2",
+  "osqp>=0.6.5",
+  "scs>=3.2.4",
+]
+
+[project.urls]
+Homepage = "https://github.com/zhaonanq/TROP"
+Repository = "https://github.com/zhaonanq/TROP"
+Issues = "https://github.com/zhaonanq/TROP/issues"
+
+[project.optional-dependencies]
+dev = [
+  "pytest>=7.0",
+  "ruff>=0.5.0",
+  "black>=24.0.0",
+]
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-q"

trop-0.1.2/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

trop-0.1.2/src/trop/__init__.py

@@ -0,0 +1,9 @@
+from .estimator import TROP_TWFE_average
+from .cv import TROP_cv_single, TROP_cv_cycle, TROP_cv_joint
+
+__all__ = [
+    "TROP_TWFE_average",
+    "TROP_cv_single",
+    "TROP_cv_cycle",
+    "TROP_cv_joint",
+]

trop-0.1.2/src/trop/cv.py

@@ -0,0 +1,427 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable, Optional, Sequence, Tuple, Union, List
+
+import numpy as np
+from joblib import Parallel, delayed
+
+from .estimator import TROP_TWFE_average
+
+
+ArrayLike = Union[np.ndarray, Sequence[Sequence[float]]]
+
+
+def _validate_panel(Y: np.ndarray, treated_periods: int, n_treated_units: int) -> None:
+    """
+    Validate panel dimensions and basic placebo-treatment parameters.
+
+    Parameters
+    ----------
+    Y:
+        Outcome panel of shape (N, T).
+    treated_periods:
+        Number of treated (post) periods assumed to be the final columns of the panel.
+        Must satisfy 1 <= treated_periods < T.
+    n_treated_units:
+        Number of treated units to sample without replacement from {0, ..., N-1}.
+        Must satisfy 1 <= n_treated_units < N.
+
+    Raises
+    ------
+    ValueError
+        If Y is not 2D, or if treated_periods / n_treated_units are out of range.
+    """
+    if Y.ndim != 2:
+        raise ValueError("Y must be a 2D array of shape (N, T).")
+    N, T = Y.shape
+    if treated_periods <= 0 or treated_periods >= T:
+        raise ValueError(f"treated_periods must be in [1, T-1]. Got treated_periods={treated_periods}, T={T}.")
+    if n_treated_units <= 0 or n_treated_units >= N:
+        raise ValueError(f"n_treated_units must be in [1, N-1]. Got n_treated_units={n_treated_units}, N={N}.")
+
+
+def _as_list(grid: Iterable[float]) -> List[float]:
+    """
+    Convert a lambda grid iterable into a non-empty list of floats.
+
+    Parameters
+    ----------
+    grid:
+        Iterable of candidate lambda values.
+
+    Returns
+    -------
+    List[float]
+        The grid converted to a list of floats.
+
+    Raises
+    ------
+    ValueError
+        If the grid is empty.
+    """
+    grid_list = list(grid)
+    if len(grid_list) == 0:
+        raise ValueError("lambda_grid must be non-empty.")
+    grid_list = [float(x) for x in grid_list]
+    return grid_list
+
+
+
+def _simulate_ate(
+    seed: int,
+    Y: np.ndarray,
+    n_treated_units: int,
+    treated_periods: int,
+    lambda_unit: float,
+    lambda_time: float,
+    lambda_nn: float,
+    solver: Optional[str] = None,
+    verbose: bool = False,
+) -> float:
+    """
+    Simulate a single placebo ATE by randomly selecting treated units.
+    """
+    rng = np.random.default_rng(seed)
+    N, _ = Y.shape
+    treated_units = rng.choice(N, size=n_treated_units, replace=False)
+
+    W = np.zeros_like(Y, dtype=float)
+    W[treated_units, -treated_periods:] = 1.0
+
+    return TROP_TWFE_average(
+        Y=Y,
+        W=W,
+        treated_units=treated_units,
+        lambda_unit=lambda_unit,
+        lambda_time=lambda_time,
+        lambda_nn=lambda_nn,
+        treated_periods=treated_periods,
+        solver=solver,
+        verbose=verbose,
+    )
+
+
+def TROP_cv_single(
+    Y_control: ArrayLike,
+    n_treated_units: int,
+    treated_periods: int,
+    fixed_lambdas: Tuple[float, float] = (0.0, 0.0),
+    lambda_grid: Optional[Iterable[float]] = None,
+    lambda_cv: str = "unit",
+    *,
+    n_trials: int = 200,
+    n_jobs: int = -1,
+    prefer: str = "threads",
+    random_seed: int = 0,
+    solver: Optional[str] = None,
+    verbose: bool = False,
+) -> float:
+    """
+    Cross-validate one lambda parameter while keeping the other two fixed.
+
+    Parameters
+    ----------
+    Y_control:
+        Control-only panel (N x T) used for placebo CV.
+    n_treated_units:
+        Number of placebo treated units to sample each trial.
+    treated_periods:
+        Number of placebo treated (post) periods (assumed final columns).
+    fixed_lambdas:
+        Tuple of two lambdas to hold fixed; interpretation depends on `lambda_cv`:
+        - lambda_cv='unit': fixed_lambdas=(lambda_time, lambda_nn)
+        - lambda_cv='time': fixed_lambdas=(lambda_unit, lambda_nn)
+        - lambda_cv='nn'  : fixed_lambdas=(lambda_unit, lambda_time)
+    lambda_grid:
+        Grid of candidate values for the lambda being tuned.
+        If None, uses np.arange(0, 2, 0.2).
+    lambda_cv:
+        Which lambda to tune: {'unit','time','nn'}.
+    n_trials:
+        Number of placebo trials per lambda.
+    n_jobs:
+        joblib parallelism. -1 uses all available cores.
+    prefer:
+        joblib backend preference. Use 'threads' by default for solver stability.
+    random_seed:
+        Seed for generating trial seeds (deterministic CV).
+    solver, verbose:
+        Passed through to TROP_TWFE_average.
+
+    Returns
+    -------
+    float
+        Lambda value that minimizes RMSE of placebo ATEs.
+    """
+    Y = np.asarray(Y_control, dtype=float)
+    _validate_panel(Y, treated_periods, n_treated_units)
+
+    if lambda_cv not in {"unit", "time", "nn"}:
+        raise ValueError("lambda_cv must be one of {'unit','time','nn'}.")
+
+    if lambda_grid is None:
+        lambda_grid_list = _as_list(np.arange(0.0, 2.0, 0.2))
+    else:
+        lambda_grid_list = _as_list(lambda_grid)
+
+    if n_trials <= 0:
+        raise ValueError("n_trials must be positive.")
+    if n_jobs == 0 or n_jobs < -1:
+        raise ValueError("n_jobs must be -1 or a positive integer.")
+
+    base_rng = np.random.default_rng(random_seed)
+    seeds = base_rng.integers(0, 2**32 - 1, size=n_trials, dtype=np.uint32)
+
+    scores: List[float] = []
+
+    for lamb in lambda_grid_list:
+        if lamb < 0:
+            raise ValueError("Lambda values must be nonnegative.")
+
+        if lambda_cv == "unit":
+            lambda_unit, lambda_time, lambda_nn = lamb, float(fixed_lambdas[0]), float(fixed_lambdas[1])
+        elif lambda_cv == "time":
+            lambda_unit, lambda_time, lambda_nn = float(fixed_lambdas[0]), lamb, float(fixed_lambdas[1])
+        else:  # 'nn'
+            lambda_unit, lambda_time, lambda_nn = float(fixed_lambdas[0]), float(fixed_lambdas[1]), lamb
+
+        ates = Parallel(n_jobs=n_jobs, prefer=prefer)(
+            delayed(_simulate_ate)(
+                int(seed),
+                Y,
+                n_treated_units,
+                treated_periods,
+                lambda_unit,
+                lambda_time,
+                lambda_nn,
+                solver,
+                verbose,
+            )
+            for seed in seeds
+        )
+
+        ates_arr = np.asarray(ates, dtype=float)
+        ates_arr = ates_arr[np.isfinite(ates_arr)]
+
+        if ates_arr.size == 0:
+            raise RuntimeError(
+                f"All placebo trials failed or returned non-finite ATEs for lambda={lamb} "
+                f"(lambda_cv='{lambda_cv}'). Consider changing solver/settings."
+            )
+
+        scores.append(float(np.sqrt(np.mean(ates_arr**2))))
+
+    best_idx = int(np.argmin(scores))
+    return float(lambda_grid_list[best_idx])
+
+
+def TROP_cv_cycle(
+    Y_control: ArrayLike,
+    n_treated_units: int,
+    treated_periods: int,
+    unit_grid: Sequence[float],
+    time_grid: Sequence[float],
+    nn_grid: Sequence[float],
+    lambdas_init: Optional[Tuple[float, float, float]] = None,
+    *,
+    max_iter: int = 50,
+    n_trials: int = 200,
+    n_jobs: int = -1,
+    prefer: str = "threads",
+    random_seed: int = 0,
+    solver: Optional[str] = None,
+    verbose: bool = False,
+) -> Tuple[float, float, float]:
+    """
+    Coordinate-descent style cross-validation for (lambda_unit, lambda_time, lambda_nn).
+
+    This routine alternates between optimizing lambda_unit, lambda_time, and lambda_nn
+    (via `TROP_cv_single`) while holding the other two fixed, until it reaches a fixed
+    point (no change in the selected lambdas) or until `max_iter` iterations are reached.
+
+    Parameters
+    ----------
+    Y_control:
+        Control-only panel (N x T) used for placebo CV.
+    n_treated_units:
+        Number of placebo treated units to sample each trial.
+    treated_periods:
+        Number of placebo treated (post) periods (assumed final columns).
+    unit_grid:
+        Grid of candidate values for lambda_unit (unit-distance decay).
+    time_grid:
+        Grid of candidate values for lambda_time (time-distance decay).
+    nn_grid:
+        Grid of candidate values for lambda_nn (nuclear-norm penalty).
+    lambdas_init:
+        Optional initial values (lambda_unit, lambda_time, lambda_nn). If None, initializes
+        each lambda to the mean of its corresponding grid.
+    max_iter:
+        Maximum number of coordinate-descent iterations.
+    n_trials:
+        Number of placebo trials per grid point in each coordinate update.
+    n_jobs:
+        joblib parallelism. -1 uses all available cores.
+    prefer:
+        joblib backend preference. Use 'threads' by default for solver stability.
+    random_seed:
+        Seed for generating trial seeds (deterministic CV).
+    solver, verbose:
+        Passed through to TROP_TWFE_average.
+
+    Returns
+    -------
+    Tuple[float, float, float]
+        (lambda_unit, lambda_time, lambda_nn) at the converged fixed point.
+
+    Raises
+    ------
+    RuntimeError
+        If the procedure does not converge to a fixed point within `max_iter`.
+    """
+
+    Y = np.asarray(Y_control, dtype=float)
+    _validate_panel(Y, treated_periods, n_treated_units)
+
+    unit_grid_list = _as_list(unit_grid)
+    time_grid_list = _as_list(time_grid)
+    nn_grid_list = _as_list(nn_grid)
+
+    if lambdas_init is None:
+        lambda_unit = float(np.mean(unit_grid_list))
+        lambda_time = float(np.mean(time_grid_list))
+        lambda_nn = float(np.mean(nn_grid_list))
+    else:
+        lambda_unit, lambda_time, lambda_nn = map(float, lambdas_init)
+
+    for _ in range(max_iter):
+        old = (lambda_unit, lambda_time, lambda_nn)
+
+        lambda_unit = TROP_cv_single(
+            Y, n_treated_units, treated_periods,
+            fixed_lambdas=(lambda_time, lambda_nn),
+            lambda_grid=unit_grid_list,
+            lambda_cv="unit",
+            n_trials=n_trials, n_jobs=n_jobs, prefer=prefer,
+            random_seed=random_seed, solver=solver, verbose=verbose
+        )
+
+        lambda_time = TROP_cv_single(
+            Y, n_treated_units, treated_periods,
+            fixed_lambdas=(lambda_unit, lambda_nn),
+            lambda_grid=time_grid_list,
+            lambda_cv="time",
+            n_trials=n_trials, n_jobs=n_jobs, prefer=prefer,
+            random_seed=random_seed, solver=solver, verbose=verbose
+        )
+
+        lambda_nn = TROP_cv_single(
+            Y, n_treated_units, treated_periods,
+            fixed_lambdas=(lambda_unit, lambda_time),
+            lambda_grid=nn_grid_list,
+            lambda_cv="nn",
+            n_trials=n_trials, n_jobs=n_jobs, prefer=prefer,
+            random_seed=random_seed, solver=solver, verbose=verbose
+        )
+
+        new = (lambda_unit, lambda_time, lambda_nn)
+        if new == old:
+            return new
+
+    raise RuntimeError("TROP_cv_cycle did not converge (no fixed point) within max_iter.")
+
+
+def TROP_cv_joint(
+    Y_control: ArrayLike,
+    n_treated_units: int,
+    treated_periods: int,
+    unit_grid: Sequence[float],
+    time_grid: Sequence[float],
+    nn_grid: Sequence[float],
+    *,
+    n_trials: int = 200,
+    n_jobs: int = -1,
+    prefer: str = "threads",
+    random_seed: int = 0,
+    solver: Optional[str] = None,
+    verbose: bool = False,
+) -> Tuple[float, float, float]:
+    """
+    Joint grid search over (lambda_unit, lambda_time, lambda_nn).
+
+    Parameters
+    ----------
+    Y_control:
+        Control-only panel (N x T) used for placebo CV.
+    n_treated_units:
+        Number of placebo treated units to sample each trial.
+    treated_periods:
+        Number of placebo treated (post) periods (assumed final columns).
+    unit_grid:
+        Grid of candidate values for lambda_unit (unit-distance decay).
+    time_grid:
+        Grid of candidate values for lambda_time (time-distance decay).
+    nn_grid:
+        Grid of candidate values for lambda_nn (nuclear-norm penalty).
+    n_trials:
+        Number of placebo trials per (lambda_unit, lambda_time, lambda_nn) triple.
+    n_jobs:
+        joblib parallelism. -1 uses all available cores.
+    prefer:
+        joblib backend preference. Use 'threads' by default for solver stability.
+    random_seed:
+        Seed for generating trial seeds (deterministic CV).
+    solver, verbose:
+        Passed through to TROP_TWFE_average.
+
+    Returns
+    -------
+    Tuple[float, float, float]
+        (lambda_unit, lambda_time, lambda_nn) triple that minimizes the RMSE of placebo ATEs.
+    """
+
+    Y = np.asarray(Y_control, dtype=float)
+    _validate_panel(Y, treated_periods, n_treated_units)
+
+    unit_grid_list = _as_list(unit_grid)
+    time_grid_list = _as_list(time_grid)
+    nn_grid_list = _as_list(nn_grid)
+
+    base_rng = np.random.default_rng(random_seed)
+    seeds = base_rng.integers(0, 2**32 - 1, size=n_trials, dtype=np.uint32)
+
+    best_params: Optional[Tuple[float, float, float]] = None
+    best_score: float = float("inf")
+
+    for lambda_unit in unit_grid_list:
+        for lambda_time in time_grid_list:
+            for lambda_nn in nn_grid_list:
+                ates = Parallel(n_jobs=n_jobs, prefer=prefer)(
+                    delayed(_simulate_ate)(
+                        int(seed),
+                        Y,
+                        n_treated_units,
+                        treated_periods,
+                        float(lambda_unit),
+                        float(lambda_time),
+                        float(lambda_nn),
+                        solver,
+                        verbose,
+                    )
+                    for seed in seeds
+                )
+
+                ates_arr = np.asarray(ates, dtype=float)
+                ates_arr = ates_arr[np.isfinite(ates_arr)]
+                if ates_arr.size == 0:
+                    continue  # skip invalid setting
+
+                score = float(np.sqrt(np.mean(ates_arr**2)))
+                if score < best_score:
+                    best_score = score
+                    best_params = (float(lambda_unit), float(lambda_time), float(lambda_nn))
+
+    if best_params is None:
+        raise RuntimeError("All parameter combinations failed during joint CV. Check solver/settings.")
+    return best_params

trop-0.1.2/src/trop/estimator.py

@@ -0,0 +1,157 @@
+from __future__ import annotations
+
+import math
+from typing import Iterable, Optional, Sequence, Union
+
+import numpy as np
+import cvxpy as cp
+
+
+ArrayLike = Union[np.ndarray, Sequence[Sequence[float]]]
+
+
+def TROP_TWFE_average(
+    Y: ArrayLike,
+    W: ArrayLike,
+    treated_units: Sequence[int],
+    lambda_unit: float,
+    lambda_time: float,
+    lambda_nn: float,
+    treated_periods: int = 10,
+    solver: Optional[str] = None,
+    verbose: bool = False,
+) -> float:
+    """
+    Triply Robust Panel (TROP) estimator in a TWFE framework with:
+      - distance-based unit weights (lambda_unit)
+      - distance-based time weights (lambda_time)
+      - optional low-rank regression adjustment (nuclear norm penalty, lambda_nn)
+
+    Parameters
+    ----------
+    Y:
+        Outcome matrix of shape (N, T).
+    W:
+        Treatment indicator matrix of shape (N, T). Typically binary {0,1}.
+        Convention in the provided codebase: treated units are treated in the final
+        `treated_periods` columns, but the function will run for any W.
+    treated_units:
+        Indices of treated units (row indices into Y/W).
+    lambda_unit:
+        Nonnegative tuning parameter controlling unit-weight decay.
+    lambda_time:
+        Nonnegative tuning parameter controlling time-weight decay.
+    lambda_nn:
+        Nuclear norm penalty weight for the low-rank component L.
+        Use np.inf to disable low-rank adjustment (i.e., no L term).
+    treated_periods:
+        Number of treated (post) periods at the end of the panel used to define:
+          - the time-distance center
+          - the pre-period mask for unit distance computation
+    solver:
+        Optional CVXPY solver name. If None, chooses a reasonable default:
+          - if lambda_nn is finite: uses SCS (supports nuclear norm / SDP forms)
+          - if lambda_nn is infinite: uses OSQP (fast for pure quadratic problems)
+    verbose:
+        Passed to CVXPY solve().
+
+    Returns
+    -------
+    float
+        Estimated average treatment effect tau.
+
+    """
+    Y = np.asarray(Y, dtype=float)
+    W = np.asarray(W, dtype=float)
+
+    if Y.ndim != 2 or W.ndim != 2:
+        raise ValueError(f"Y and W must be 2D arrays. Got Y.ndim={Y.ndim}, W.ndim={W.ndim}.")
+    if Y.shape != W.shape:
+        raise ValueError(f"Y and W must have the same shape. Got Y={Y.shape}, W={W.shape}.")
+
+    N, T = Y.shape
+
+    if not isinstance(treated_periods, int) or treated_periods <= 0:
+        raise ValueError("treated_periods must be a positive integer.")
+    if treated_periods >= T:
+        raise ValueError(f"treated_periods must be < T. Got treated_periods={treated_periods}, T={T}.")
+
+    treated_units_arr = np.asarray(treated_units, dtype=int)
+    if treated_units_arr.size == 0:
+        raise ValueError("treated_units must contain at least one unit index.")
+    if np.any(treated_units_arr < 0) or np.any(treated_units_arr >= N):
+        raise ValueError(f"treated_units contains out-of-range indices for N={N}: {treated_units_arr}")
+
+    if lambda_unit < 0 or lambda_time < 0:
+        raise ValueError("lambda_unit and lambda_time should be nonnegative.")
+
+    # ---------------------------------------------------------------------
+    # Distance-based time weights
+    # ---------------------------------------------------------------------
+    # Distance to the center of the treated block near the end of the panel.
+    # dist_time = abs(arange(T) - (T - treated_periods/2))
+    center = T - treated_periods / 2.0
+    dist_time = np.abs(np.arange(T, dtype=float) - center)
+
+    # ---------------------------------------------------------------------
+    # Distance-based unit weights
+    # ---------------------------------------------------------------------
+    average_treated = np.mean(Y[treated_units_arr, :], axis=0)
+
+    # Pre-period mask: 1 in pre, 0 in treated/post
+    mask = np.ones((N, T), dtype=float)
+    mask[:, -treated_periods:] = 0.0
+
+    # RMS distance to average treated trajectory over pre-periods
+    # dist_unit[i] = sqrt( sum_pre (avg_tr - Y_i)^2 / (#pre) )
+    A = np.sum(((average_treated - Y) ** 2) * mask, axis=1)
+    B = np.sum(mask, axis=1)
+
+    if np.any(B == 0):
+        raise ValueError(
+            "Pre-period mask has zero pre-periods for at least one unit."
+        )
+
+    dist_unit = np.sqrt(A / B)
+
+    # Convert distances to weights
+    delta_unit = np.exp(-lambda_unit * dist_unit)          # shape (N,)
+    delta_time = np.exp(-lambda_time * dist_time)          # shape (T,)
+    delta = np.outer(delta_unit, delta_time)               # shape (N, T)
+
+    # ---------------------------------------------------------------------
+    # CVXPY problem: weighted TWFE
+    # ---------------------------------------------------------------------
+    unit_effects = cp.Variable((1, N))
+    time_effects = cp.Variable((1, T))
+    mu = cp.Variable()     # intercept
+    tau = cp.Variable()    # treatment effect
+
+    # Broadcast TWFE components to N x T
+    unit_factor = cp.kron(np.ones((T, 1)), unit_effects).T
+    time_factor = cp.kron(np.ones((N, 1)), time_effects)
+
+    is_low_rank = not math.isinf(float(lambda_nn))
+
+    if is_low_rank:
+        L = cp.Variable((N, T))
+        residual = Y - mu - unit_factor - time_factor - L - W * tau
+        loss = cp.sum_squares(cp.multiply(residual, delta)) + float(lambda_nn) * cp.norm(L, "nuc")
+        default_solver = "SCS"  # robust choice for nuclear norm problems
+    else:
+        residual = Y - mu - unit_factor - time_factor - W * tau
+        loss = cp.sum_squares(cp.multiply(residual, delta))
+        default_solver = "OSQP"  # fast for pure quadratic objective
+
+    prob = cp.Problem(cp.Minimize(loss))
+
+    chosen_solver = solver or default_solver
+    prob.solve(solver=chosen_solver, verbose=verbose)
+
+    if tau.value is None or not np.isfinite(tau.value):
+        raise RuntimeError(
+            "Optimization did not return a valid tau. "
+            f"Solver={chosen_solver}, status={prob.status}."
+        )
+
+    return float(tau.value)

trop-0.1.2/src/trop.egg-info/PKG-INFO

@@ -0,0 +1,53 @@
+Metadata-Version: 2.4
+Name: trop
+Version: 0.1.2
+Summary: Triply Robust Panel (TROP) estimator: weighted TWFE with optional low-rank adjustment.
+Author: Susan Athey, Guido Imbens, Zhaonan Qu, Davide Viviano
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/zhaonanq/TROP
+Project-URL: Repository, https://github.com/zhaonanq/TROP
+Project-URL: Issues, https://github.com/zhaonanq/TROP/issues
+Keywords: causal-inference,panel-data,factor-models,difference-in-differences,synthetic-control,synthetic-controls,trop,twfe
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.23
+Requires-Dist: cvxpy>=1.4
+Requires-Dist: joblib>=1.2
+Requires-Dist: osqp>=0.6.5
+Requires-Dist: scs>=3.2.4
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: ruff>=0.5.0; extra == "dev"
+Requires-Dist: black>=24.0.0; extra == "dev"
+Dynamic: license-file
+
+# TROP: Triply Robust Panel Estimator
+
+`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.
+
+
+Reference:
+
+> Susan Athey, Guido Imbens, Zhaonan Qu, Davide Viviano (2025).  
+> *Triply Robust Panel Estimators*.  
+> arXiv:2508.21536.
+
+---
+
+## Installation
+
+```
+pip install trop
+```

trop-0.1.2/src/trop.egg-info/SOURCES.txt

@@ -0,0 +1,12 @@
+LICENSE
+README.md
+pyproject.toml
+src/trop/__init__.py
+src/trop/cv.py
+src/trop/estimator.py
+src/trop.egg-info/PKG-INFO
+src/trop.egg-info/SOURCES.txt
+src/trop.egg-info/dependency_links.txt
+src/trop.egg-info/requires.txt
+src/trop.egg-info/top_level.txt
+tests/test_smoke_trop.py

trop-0.1.2/src/trop.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

trop-0.1.2/src/trop.egg-info/requires.txt

@@ -0,0 +1,10 @@
+numpy>=1.23
+cvxpy>=1.4
+joblib>=1.2
+osqp>=0.6.5
+scs>=3.2.4
+
+[dev]
+pytest>=7.0
+ruff>=0.5.0
+black>=24.0.0

trop-0.1.2/src/trop.egg-info/top_level.txt

@@ -0,0 +1 @@
+trop

trop-0.1.2/tests/test_smoke_trop.py

@@ -0,0 +1,46 @@
+import numpy as np
+import pytest
+
+from trop import TROP_TWFE_average
+
+
+def test_trop_twfe_average_smoke():
+    # Small synthetic panel
+    N, T = 8, 20
+    treated_periods = 5
+    treated_units = [0, 1]
+
+    rng = np.random.default_rng(0)
+    Y = rng.normal(size=(N, T))
+
+    # Treatment: treated_units get treated in last treated_periods
+    W = np.zeros((N, T))
+    W[treated_units, -treated_periods:] = 1.0
+
+    tau = TROP_TWFE_average(
+        Y=Y,
+        W=W,
+        treated_units=treated_units,
+        lambda_unit=0.1,
+        lambda_time=0.1,
+        lambda_nn=np.inf,   # simplest path: no nuclear-norm component
+        treated_periods=treated_periods,
+    )
+
+    assert np.isfinite(tau), "tau should be a finite scalar"
+    assert isinstance(tau, float)
+
+
+def test_invalid_shapes_raises():
+    Y = np.zeros((5, 10))
+    W = np.zeros((5, 9))  # mismatched
+    with pytest.raises(ValueError):
+        TROP_TWFE_average(
+            Y=Y,
+            W=W,
+            treated_units=[0],
+            lambda_unit=0.1,
+            lambda_time=0.1,
+            lambda_nn=np.inf,
+            treated_periods=2,
+        )

trop-0.1.0/PKG-INFO

@@ -1,24 +0,0 @@
-Metadata-Version: 2.4
-Name: trop
-Version: 0.1.0
-Summary: Package that implements the Triply Robust Panel Estimator (TROP)
-License: MIT
-Author: meganndare, zhaonanq
-Requires-Python: >=3.9
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Classifier: Programming Language :: Python :: 3.13
-Classifier: Programming Language :: Python :: 3.14
-Requires-Dist: cvxpy (==1.4.1)
-Requires-Dist: joblib (==1.3.2)
-Requires-Dist: numpy (==1.26.2)
-Requires-Dist: pandas (==2.1.3)
-Description-Content-Type: text/markdown
-
-# Triply Robust Panel Estimators (TROP)
-
-This package will soon contain the replication files and implementation of the TROP estimator.

trop-0.1.0/README.md

@@ -1,3 +0,0 @@
-# Triply Robust Panel Estimators (TROP)
-
-This package will soon contain the replication files and implementation of the TROP estimator.

trop-0.1.0/pyproject.toml

@@ -1,21 +0,0 @@
-[project]
-name = "trop"
-version = "0.1.0"
-description = "Package that implements the Triply Robust Panel Estimator (TROP)"
-authors = [
-    {name = "meganndare, zhaonanq"}
-]
-license = {text = "MIT"}
-readme = "README.md"
-requires-python = ">=3.9"
-dependencies = [
-    "numpy (==1.26.2)",
-    "pandas (==2.1.3)",
-    "cvxpy (==1.4.1)",
-    "joblib (==1.3.2)"
-]
-
-
-[build-system]
-requires = ["poetry-core>=2.0.0,<3.0.0"]
-build-backend = "poetry.core.masonry.api"

trop-0.1.0/trop/placeholder.py

@@ -1 +0,0 @@
-# placeholder to test package publishing