trop 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

trop/cv.py

@@ -13,25 +13,7 @@ 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.
-    """
+    """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
@@ -42,24 +24,7 @@ def _validate_panel(Y: np.ndarray, treated_periods: int, n_treated_units: int) -
 
 
 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.
-    """
+    """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.")
@@ -118,41 +83,48 @@ def TROP_cv_single(
     verbose: bool = False,
 ) -> float:
     """
-    Cross-validate one lambda parameter while keeping the other two fixed.
+    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:
-        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.
+    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
-        Lambda value that minimizes RMSE of placebo ATEs.
+        Selected lambda value minimizing the RMSE of placebo estimates.
     """
     Y = np.asarray(Y_control, dtype=float)
     _validate_panel(Y, treated_periods, n_treated_units)
@@ -234,53 +206,54 @@ def TROP_cv_cycle(
     verbose: bool = False,
 ) -> Tuple[float, float, float]:
     """
-    Coordinate-descent style cross-validation for (lambda_unit, lambda_time, lambda_nn).
+    Tune (lambda_unit, lambda_time, lambda_nn) by coordinate-descent placebo cross-validation.
 
-    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.
+    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:
-        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:
+    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:
+    n_trials : int, default=200
         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.
+    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[float, float, float]
-        (lambda_unit, lambda_time, lambda_nn) at the converged fixed point.
+    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 to a fixed point within `max_iter`.
+        If the procedure does not converge within `max_iter`.
     """
-
     Y = np.asarray(Y_control, dtype=float)
     _validate_panel(Y, treated_periods, n_treated_units)
 
@@ -348,39 +321,50 @@ def TROP_cv_joint(
     verbose: bool = False,
 ) -> Tuple[float, float, float]:
     """
-    Joint grid search over (lambda_unit, lambda_time, lambda_nn).
+    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:
-        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.
+    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[float, float, float]
-        (lambda_unit, lambda_time, lambda_nn) triple that minimizes the RMSE of placebo ATEs.
-    """
+    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)
 

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.2.dist-info → trop-0.1.4.dist-info}/METADATA

@@ -1,12 +1,13 @@
 Metadata-Version: 2.4
 Name: trop
-Version: 0.1.2
+Version: 0.1.4
 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
+Project-URL: Homepage, https://github.com/ostasovskyi/TROP-Estimator
+Project-URL: Repository, https://github.com/ostasovskyi/TROP-Estimator
+Project-URL: Issues, https://github.com/ostasovskyi/TROP-Estimator/issues
+Project-URL: Documentation, https://ostasovskyi.github.io/TROP-Estimator/
 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
@@ -46,6 +47,13 @@ Reference:
 
 ---
 
+## Links
+
+- **Documentation**: https://ostasovskyi.github.io/TROP-Estimator/
+- **Source code**: https://github.com/ostasovskyi/TROP-Estimator
+
+---
+
 ## Installation
 
 ```

trop-0.1.4.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.4.dist-info/licenses/LICENSE,sha256=VqjvjioQz04uLYBj4ye0x-_Ss77-WTIuEWWCW_awEz8,1065
+trop-0.1.4.dist-info/METADATA,sha256=LD3UoSdXXTq_gx_jhVy55OoKEcOvi5Ws7f1J_IQxMmw,2258
+trop-0.1.4.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+trop-0.1.4.dist-info/top_level.txt,sha256=jaqQZFm3D5B4vPBAKZtXfEAYnpl9FKsNHqlM49kcwTI,5
+trop-0.1.4.dist-info/RECORD,,

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

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

trop-0.1.2.dist-info/RECORD

@@ -1,8 +0,0 @@
-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,,