fix(lwdid): native time-scale pre-period identification + review polish

Author: gorgeousfish

diff_diff/lwdid_trend_diagnostics.py

@@ -27,7 +27,7 @@
 
 import warnings
 from dataclasses import dataclass
-from typing import List, Optional
+from typing import Any, List, Optional
 
 import numpy as np
 import pandas as pd
@@ -55,8 +55,9 @@ class PreTrendEstimate:
 
     Attributes
     ----------
-    period : int
-        Calendar period (pseudo-post) used for this estimate.
+    period : scalar
+        Calendar period (pseudo-post) used for this estimate, on the
+        native scale of the ``time`` column (integer, datetime64, ...).
     att : float
         Estimated average treatment effect on the treated.
     se : float
@@ -67,7 +68,7 @@ class PreTrendEstimate:
         Two-sided p-value for testing H0: ATT = 0.
     """
 
-    period: int
+    period: Any
     att: float
     se: float
     t_stat: float
@@ -216,14 +217,16 @@ def _identify_pre_periods(data: pd.DataFrame, time: str, treatment: str, unit: s
 
     Returns
     -------
-    tuple of (list, int)
-        (pre_periods sorted, first_treat_time)
+    tuple of (list, scalar)
+        (pre_periods sorted, first_treat_time). The first-treatment time
+        is kept on the native scale of the ``time`` column (integer,
+        datetime64, ...), matching the LWDiD estimator.
     """
     treated_times = data.loc[data[treatment] == 1, time].unique()
     if len(treated_times) == 0:
         raise ValueError("No treated observations found in the data.")
 
-    first_treat = int(min(treated_times))
+    first_treat = min(treated_times)
     all_times = sorted(data[time].unique())
     pre_periods = [t for t in all_times if t < first_treat]
 
@@ -372,7 +375,7 @@ def _placebo_pre_trends(
     pre_periods, first_treat = _identify_pre_periods(data, time, treatment, unit)
 
     if len(pre_periods) < 2:
-        raise ValueError(
+        raise InsufficientPrePeriodsError(
             f"Need at least 2 pre-treatment periods for parallel trends test, "
             f"got {len(pre_periods)}."
         )
@@ -416,7 +419,7 @@ def _placebo_pre_trends(
             pval = 2 * (1 - stats.norm.cdf(abs(t_stat)))
             pre_effects.append(
                 PreTrendEstimate(
-                    period=int(pseudo_post_start),
+                    period=pseudo_post_start,
                     att=float(result.att),
                     se=float(result.se),
                     t_stat=float(t_stat),

docs/api/lwdid.rst

@@ -269,6 +269,25 @@ Results container returned by :meth:`~diff_diff.LWDiD.fit`.
       ~LWDiDResults.to_dataframe
       ~LWDiDResults.to_dict
 
+Input Contract
+--------------
+
+:meth:`~diff_diff.LWDiD.fit` validates the treatment design before any
+transformation is applied. Three requirements are enforced:
+
+- **Absorbing treatment** — within each unit the ``treatment`` indicator
+  must be non-decreasing over time: once a unit switches from 0 to 1 it
+  must remain treated. Units that revert to 0 raise ``ValueError``.
+- **Common timing** — when ``first_treat`` is not supplied, all treated
+  units must first switch on in the same period. Heterogeneous onsets
+  are rejected with a ``ValueError`` pointing to the staggered interface
+  (pass ``first_treat``).
+- **Staggered consistency** — when ``first_treat`` is supplied, the
+  ``treatment`` indicator must satisfy :math:`D_{it} = 1[t \ge g_i]`,
+  where :math:`g_i` is the unit's first-treatment period. Units that are
+  never treated (``first_treat`` coded NaN or 0) must have no treated
+  rows.
+
 Example Usage
 -------------
 

tests/test_lwdid_trend_diagnostics.py

@@ -169,6 +169,22 @@ def test_recommendation_summary(self, panel_data):
         assert isinstance(s, str)
         assert "RECOMMENDATION" in s
 
+    def test_datetime_time_scale(self):
+        """A datetime64 time column is handled on its native scale."""
+        rng = np.random.default_rng(42)
+        years = pd.date_range("2000-01-01", periods=8, freq="YS")
+        records = []
+        for i in range(80):
+            d = int(i < 25)
+            for k, t in enumerate(years, start=1):
+                y = 1.0 + 0.1 * k + rng.normal(0, 0.3)
+                if d and k > 4:
+                    y += 2.0
+                records.append({"unit": i, "time": t, "y": y, "treat": d * int(k > 4)})
+        df = pd.DataFrame(records)
+        rec = recommend_transformation(df, outcome="y", unit="unit", time="time", treatment="treat")
+        assert isinstance(rec, TransformationRecommendation)
+
 
 # ---------------------------------------------------------------------------
 # Edge cases