trop 0.1.1__py3-none-any.whl → 0.1.2__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,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.1.dist-info → trop-0.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trop
-Version: 0.1.1
+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
@@ -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.2.dist-info/RECORD

@@ -0,0 +1,8 @@
+trop/__init__.py,sha256=B94vrDZevg2l6ijN4lut7wo0MYTEtBTTZhqmMQtq7Qg,205
+trop/cv.py,sha256=7tumpAaAiiK_F8nUUAcAwnecNAfc2XghftgJsNWVsAQ,14017
+trop/estimator.py,sha256=FWMO39GbL6k3Vz5g1V7SpR6t5wP3N81V5gGSFIe65Xw,6001
+trop-0.1.2.dist-info/licenses/LICENSE,sha256=VqjvjioQz04uLYBj4ye0x-_Ss77-WTIuEWWCW_awEz8,1065
+trop-0.1.2.dist-info/METADATA,sha256=YSyeONhxn4JO_NZBKT1ubVa5J96BhlWGtmRi-GJ0rLk,1997
+trop-0.1.2.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
+trop-0.1.2.dist-info/top_level.txt,sha256=jaqQZFm3D5B4vPBAKZtXfEAYnpl9FKsNHqlM49kcwTI,5
+trop-0.1.2.dist-info/RECORD,,

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

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.1)
 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,,