trop 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

trop/__init__.py

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

trop/cv.py

@@ -0,0 +1,411 @@
+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 placebo CV inputs."""
+    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 an iterable of grid values to a non-empty list of floats."""
+    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:
+    """
+    Tune one TROP tuning parameter via placebo cross-validation on a control-only panel.
+
+    For each candidate value in `lambda_grid`, this routine repeatedly assigns a placebo
+    treatment to random units in the last `treated_periods` columns, computes the
+    corresponding TROP estimate, and selects the lambda that minimizes the RMSE of
+    placebo effects.
+
+    Parameters
+    ----------
+    Y_control : array_like of shape (N, T)
+        Control-only outcome panel used for placebo cross-validation.
+    n_treated_units : int
+        Number of placebo treated units sampled (without replacement) per trial.
+    treated_periods : int
+        Number of placebo treated (post) periods, taken as the final columns.
+    fixed_lambdas : tuple of float, default=(0.0, 0.0)
+        Values held fixed for the two lambdas not being tuned. Interpretation depends on
+        `lambda_cv`:
+        - 'unit': (lambda_time, lambda_nn)
+        - 'time': (lambda_unit, lambda_nn)
+        - 'nn'  : (lambda_unit, lambda_time)
+    lambda_grid : iterable of float or None, default=None
+        Candidate values for the lambda being tuned. If None, uses ``np.arange(0, 2, 0.2)``.
+    lambda_cv : {'unit', 'time', 'nn'}, default='unit'
+        Which lambda to tune.
+    n_trials : int, default=200
+        Number of placebo trials per candidate lambda.
+    n_jobs : int, default=-1
+        Number of parallel jobs for placebo trials. ``-1`` uses all available cores.
+    prefer : {'threads', 'processes'}, default='threads'
+        joblib backend preference.
+    random_seed : int, default=0
+        Seed for generating trial seeds (deterministic tuning).
+    solver : str or None, default=None
+        CVXPY solver passed to ``TROP_TWFE_average``.
+    verbose : bool, default=False
+        Verbosity flag passed to ``TROP_TWFE_average``.
+
+    Returns
+    -------
+    float
+        Selected lambda value minimizing the RMSE of placebo estimates.
+    """
+    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]:
+    """
+    Tune (lambda_unit, lambda_time, lambda_nn) by coordinate-descent placebo cross-validation.
+
+    Iteratively updates one tuning parameter at a time using `TROP_cv_single` (holding the
+    other two fixed) until the selected triplet stops changing or `max_iter` is reached.
+    Each update minimizes the RMSE of placebo effects on a control-only panel.
+
+    Parameters
+    ----------
+    Y_control : array_like of shape (N, T)
+        Control-only outcome panel used for placebo cross-validation.
+    n_treated_units : int
+        Number of placebo treated units sampled (without replacement) per trial.
+    treated_periods : int
+        Number of placebo treated (post) periods, taken as the final columns.
+    unit_grid : sequence of float
+        Candidate values for `lambda_unit` (unit-distance decay).
+    time_grid : sequence of float
+        Candidate values for `lambda_time` (time-distance decay).
+    nn_grid : sequence of float
+        Candidate values for `lambda_nn` (nuclear-norm penalty).
+    lambdas_init : tuple of float or None, default=None
+        Initial values (lambda_unit, lambda_time, lambda_nn). If None, initializes each
+        parameter to the mean of its grid.
+    max_iter : int, default=50
+        Maximum number of coordinate-descent iterations.
+    n_trials : int, default=200
+        Number of placebo trials per grid point in each coordinate update.
+    n_jobs : int, default=-1
+        Number of parallel jobs for placebo trials. ``-1`` uses all available cores.
+    prefer : {'threads', 'processes'}, default='threads'
+        joblib backend preference.
+    random_seed : int, default=0
+        Seed for generating trial seeds (deterministic tuning).
+    solver : str or None, default=None
+        CVXPY solver passed to ``TROP_TWFE_average``.
+    verbose : bool, default=False
+        Verbosity flag passed to ``TROP_TWFE_average``.
+
+    Returns
+    -------
+    tuple of float
+        (lambda_unit, lambda_time, lambda_nn) at the fixed point of the coordinate updates.
+
+    Raises
+    ------
+    RuntimeError
+        If the procedure does not converge 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]:
+    """
+    Select (lambda_unit, lambda_time, lambda_nn) by joint placebo cross-validation.
+
+    Performs a full grid search over `unit_grid` × `time_grid` × `nn_grid`. For each
+    candidate triple, repeatedly assigns a placebo treatment to random units in the
+    last `treated_periods` columns and selects the triple that minimizes the RMSE of
+    placebo effects on the control-only panel.
+
+    Parameters
+    ----------
+    Y_control : array_like of shape (N, T)
+        Control-only outcome panel used for placebo cross-validation.
+    n_treated_units : int
+        Number of placebo treated units sampled (without replacement) per trial.
+    treated_periods : int
+        Number of placebo treated (post) periods, taken as the final columns.
+    unit_grid : sequence of float
+        Candidate values for `lambda_unit` (unit-distance decay).
+    time_grid : sequence of float
+        Candidate values for `lambda_time` (time-distance decay).
+    nn_grid : sequence of float
+        Candidate values for `lambda_nn` (nuclear-norm penalty).
+    n_trials : int, default=200
+        Number of placebo trials per candidate triple.
+    n_jobs : int, default=-1
+        Number of parallel jobs for placebo trials. ``-1`` uses all available cores.
+    prefer : {'threads', 'processes'}, default='threads'
+        joblib backend preference.
+    random_seed : int, default=0
+        Seed for generating trial seeds (deterministic tuning).
+    solver : str or None, default=None
+        CVXPY solver passed to ``TROP_TWFE_average``.
+    verbose : bool, default=False
+        Verbosity flag passed to ``TROP_TWFE_average``.
+
+    Returns
+    -------
+    tuple of float
+        (lambda_unit, lambda_time, lambda_nn) minimizing the RMSE of placebo estimates.
+
+    Raises
+    ------
+    RuntimeError
+        If all parameter combinations fail (e.g., solver failures for every triple).
+    """
+    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/estimator.py

@@ -22,44 +22,48 @@ def TROP_TWFE_average(
     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)
+    Compute the TROP treatment effect with unit/time weighting and optional low-rank
+    outcome model.
 
     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().
+    Y : array_like of shape (N, T)
+        Outcome matrix.
+    W : array_like of shape (N, T)
+        Treatment indicator (often binary). The estimator uses ``W`` as provided;
+        ``treated_periods`` is used only to construct weights/masks, not to infer
+        treatment timing.
+    treated_units : sequence of int
+        Row indices of treated units used to form the reference (average) treated
+        trajectory for unit-distance weighting.
+    lambda_unit : float
+        Nonnegative decay parameter for unit weights: ``exp(-lambda_unit * dist_unit)``.
+    lambda_time : float
+        Nonnegative decay parameter for time weights: ``exp(-lambda_time * dist_time)``.
+    lambda_nn : float
+        Nuclear-norm penalty weight for the low-rank component ``L``. Use
+        ``np.inf`` to disable the low-rank adjustment (i.e., omit ``L``).
+    treated_periods : int, default=10
+        Number of final columns treated as the "post/tail block" for constructing
+        (a) the pre-period mask (all but last ``treated_periods`` columns) used in
+        unit distances, and (b) the time-distance center.
+    solver : str or None, default=None
+        CVXPY solver name. If None, uses "SCS" when ``lambda_nn`` is finite and
+        "OSQP" when ``lambda_nn`` is infinite.
+    verbose : bool, default=False
+        Passed to ``cvxpy.Problem.solve``.
 
     Returns
     -------
     float
-        Estimated average treatment effect tau.
-
+        Estimated treatment-effect parameter ``tau`` from the weighted TWFE objective.
+
+    Raises
+    ------
+    ValueError
+        If input shapes are inconsistent or tuning parameters are invalid.
+    RuntimeError
+        If the optimization fails to produce a finite ``tau``.
     """
     Y = np.asarray(Y, dtype=float)
     W = np.asarray(W, dtype=float)

{trop-0.1.1.dist-info → trop-0.1.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trop
-Version: 0.1.1
+Version: 0.1.3
 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
@@ -24,6 +24,7 @@ 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
@@ -34,18 +35,15 @@ Dynamic: license-file
 
 # TROP: Triply Robust Panel Estimator
 
-This package provides a Python implementation of the **Triply Robust Panel (TROP)** estimator introduced in:
+`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.
 
-The initial release (v0.1.0) exposes the function:
-
-- `TROP_TWFE_average(Y, W, treated_units, lambda_unit, lambda_time, lambda_nn, treated_periods=..., solver=...)`
-
-which estimates an average treatment effect `tau` in panel settings using a weighted TWFE objective with optional low-rank adjustment.
-
 ---
 
 ## Installation

trop-0.1.3.dist-info/RECORD

@@ -0,0 +1,8 @@
+trop/__init__.py,sha256=B94vrDZevg2l6ijN4lut7wo0MYTEtBTTZhqmMQtq7Qg,205
+trop/cv.py,sha256=tbT2mvPO_gQ28GT5CIPZR1NqL11J3PjxsTHYxNWmVRk,14908
+trop/estimator.py,sha256=iZFH9D664OsEZascxAEdSqOBPj26DofXHk6jWgz_42Q,6257
+trop-0.1.3.dist-info/licenses/LICENSE,sha256=VqjvjioQz04uLYBj4ye0x-_Ss77-WTIuEWWCW_awEz8,1065
+trop-0.1.3.dist-info/METADATA,sha256=paPD2_iwXAHPc-y6Dlu3CmpgrnGRHZplRqgjsV-t2y8,1997
+trop-0.1.3.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+trop-0.1.3.dist-info/top_level.txt,sha256=jaqQZFm3D5B4vPBAKZtXfEAYnpl9FKsNHqlM49kcwTI,5
+trop-0.1.3.dist-info/RECORD,,

{trop-0.1.1.dist-info → trop-0.1.3.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

trop-0.1.1.dist-info/RECORD

@@ -1,7 +0,0 @@
-trop/__init__.py,sha256=DW6eDmMyaY1tQ6wb-EP48fTNtkeUOuvqE5l88d8SnrA,73
-trop/estimator.py,sha256=FWMO39GbL6k3Vz5g1V7SpR6t5wP3N81V5gGSFIe65Xw,6001
-trop-0.1.1.dist-info/licenses/LICENSE,sha256=VqjvjioQz04uLYBj4ye0x-_Ss77-WTIuEWWCW_awEz8,1065
-trop-0.1.1.dist-info/METADATA,sha256=nM0XBF9nad4XGbpHiBmtLnPXs-mKxXM3sDxvt_ALl6Y,2069
-trop-0.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-trop-0.1.1.dist-info/top_level.txt,sha256=jaqQZFm3D5B4vPBAKZtXfEAYnpl9FKsNHqlM49kcwTI,5
-trop-0.1.1.dist-info/RECORD,,