"""

Continuous Difference-in-Differences estimator.

Implements Callaway, Goodman-Bacon & Sant'Anna (2024),

"Difference-in-Differences with a Continuous Treatment" (NBER WP 32117).

Estimates dose-response curves ATT(d) and ACRT(d), as well as summary

parameters ATT^{glob} and ACRT^{glob}, with optional multiplier bootstrap

inference.

"""

import dataclasses

import warnings

from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple

import numpy as np

import pandas as pd

from diff_diff._base import BaseEstimator

from diff_diff._deprecation import NOT_SUPPLIED, warn_deprecated_kwarg

from diff_diff.aggregation import AggregationKit

from diff_diff.bootstrap_utils import (

compute_effect_bootstrap_stats,

generate_bootstrap_weights_batch,

)

from diff_diff.continuous_did_aggregation import _ContinuousDiDAggregationMixin

from diff_diff.continuous_did_bspline import (

SATURATED_TOL,

bspline_derivative_design_matrix,

bspline_design_matrix,

build_bspline_basis,

default_dose_grid,

saturated_derivative_design_matrix,

saturated_design_matrix,

saturated_dose_levels,

)

from diff_diff.continuous_did_results import (

ContinuousDiDResults,

DoseResponseCurve,

)

from diff_diff.linalg import _rank_guarded_inv, solve_logit, solve_ols

from diff_diff.survey import (

_resolve_survey_for_fit,

_validate_unit_constant_survey,

build_unit_first_row_index,

compute_survey_vcov,

)

from diff_diff.utils import safe_inference

if TYPE_CHECKING:

from diff_diff.survey import ResolvedSurveyDesign, SurveyDesign

__all__ = ["ContinuousDiD", "ContinuousDiDResults", "DoseResponseCurve"]

#: Pruned per-cell payload keys the post-fit event-study recompute reads

#: (row M-025). Exactly the ``_bootstrap_info`` subset

#: ``_compute_event_study_inference`` consumes - the K-dimensional spline

#: machinery (bread, ee_treated, Psi_eval, dPsi_*, beta_pred) is

#: deliberately NOT retained and dies with fit(). ``w_treated``/

#: ``w_control``/``w_treated_arr`` are copied only when present because

#: the consumer's survey-mass branch keys on ``"w_treated" in b_info``.

_ES_PAYLOAD_KEYS = (

"treated_indices",

"control_indices",

"n_treated",

"n_control",

"att_glob",

"mu_0",

"delta_y_treated",

"ee_control",

"w_treated",

"w_control",

"w_treated_arr",

)

def _build_continuous_aggregation_kit(

estimator: "ContinuousDiD",

gt_results: Dict[Tuple, Dict],

gt_bootstrap_info: Dict[Tuple, Dict],

precomp: Dict[str, Any],

resolved_survey: Optional["ResolvedSurveyDesign"],

has_post_cells: bool,

survey_df: Optional[int],

survey_metadata: Optional[Any],

) -> AggregationKit:

"""Build the post-fit aggregation kit for ContinuousDiD (row M-025).

``influence`` is empty BY DESIGN: the event-study recompute reads the

pruned per-cell payload + unit-level arrays in ``bookkeeping``, not a

per-unit EIF dict on the kit's influence contract; the ``simple`` /

``dose`` levels are pure views over stored public results fields and

never read the kit at all.

On bootstrap fits (``n_bootstrap > 0``) the kit is SCALARS-ONLY:

``aggregate('event_study')`` fails closed before reading any payload

and the views never read the kit, so a populated payload there would

be pure dead retention. The scalars still distinguish the

bootstrap-NotImplementedError gate from the legacy-pickle no-kit

ValueError.

"""

is_bootstrap = estimator.n_bootstrap > 0

gt_summary: Dict[Tuple, Dict[str, Any]] = {}

gt_es_payload: Dict[Tuple, Dict[str, Any]] = {}

if not is_bootstrap:

for gt, r in gt_results.items():

gt_summary[gt] = {

"att_glob": float(r["att_glob"]),

"n_treated": int(r["n_treated"]),

}

b_info = gt_bootstrap_info.get(gt, {})

if not b_info:

gt_es_payload[gt] = {}

continue

pruned = {k: b_info[k] for k in _ES_PAYLOAD_KEYS if k in b_info}

cov_if = b_info.get("cov_if")

pruned["cov_if"] = (

{

"cell_indices": cov_if["cell_indices"],

"if_att_glob": cov_if["if_att_glob"],

}

if cov_if is not None

else None

)

gt_es_payload[gt] = pruned

bookkeeping: Dict[str, Any] = {

"gt_summary": gt_summary,

"gt_es_payload": gt_es_payload,

"n_units": None if is_bootstrap else precomp["n_units"],

"unit_cohorts": None if is_bootstrap else precomp["unit_cohorts"],

"unit_survey_weights": (None if is_bootstrap else precomp.get("unit_survey_weights")),

"unit_first_panel_row": (None if is_bootstrap else precomp["unit_first_panel_row"]),

# PANEL-LEVEL design ref (locked decision: the unit-level collapse

# stays inside the verbatim recompute body; on replicate designs

# this carries the (n_obs x R) replicate matrix - documented in

# the REGISTRY memory contract).

"resolved_survey": None if is_bootstrap else resolved_survey,

"has_post_cells": has_post_cells,

"survey_df": survey_df,

"n_bootstrap": estimator.n_bootstrap,

"base_period": estimator.base_period,

# Fit-final COPY - the ES carrier metadata source; never the

# mutable public field (post-fit mutation of replicate_method /

# df_survey must not reach post-fit provenance).

"survey_metadata": (

dataclasses.replace(survey_metadata) if survey_metadata is not None else None

),

}

return AggregationKit(

bookkeeping=bookkeeping,

influence={},

alpha=estimator.alpha,

anticipation=estimator.anticipation,

cband=False,

bootstrap=None,

)

class ContinuousDiD(_ContinuousDiDAggregationMixin, BaseEstimator):

"""

Continuous Difference-in-Differences estimator.

Implements the methodology from Callaway, Goodman-Bacon & Sant'Anna (2024)

for estimating dose-response curves when treatment has a continuous intensity.

Parameters

----------

degree : int, default=3

B-spline degree (3 = cubic).

num_knots : int, default=0

Number of interior knots for the B-spline basis.

dvals : array-like, optional

Custom dose evaluation grid. If None, uses quantile-based default.

control_group : str, default="never_treated"

``"never_treated"``, ``"not_yet_treated"``, or ``"lowest_dose"``.

``"lowest_dose"`` implements Remark 3.1 (CGBS 2024) for settings with no

never-treated / zero-dose units (``P(D=0) = 0``): the lowest-dose group

``d_L`` becomes the comparison and the estimand is ``ATT(d) − ATT(d_L)``.

Requires a genuine lowest-dose group (``>= 2`` units at ``d_L``, i.e.

``P(D=d_L) > 0``) and no never-treated units present. Single-cohort only

(multi-cohort and ``covariates=`` raise ``NotImplementedError``).

anticipation : int, default=0

Number of periods of treatment anticipation.

base_period : str, default="varying"

``"varying"`` or ``"universal"``.

alpha : float, default=0.05

Significance level for confidence intervals.

n_bootstrap : int, default=0

Number of multiplier bootstrap iterations. 0 for analytical SEs only.

bootstrap_weights : str, default="rademacher"

Bootstrap weight type: ``"rademacher"``, ``"mammen"``, or ``"webb"``.

seed : int, optional

Random seed for reproducibility.

rank_deficient_action : str, default="warn"

Action for rank-deficient B-spline OLS: ``"warn"``, ``"error"``, or ``"silent"``.

covariates : list of str, optional

DEPRECATED constructor home (row M-084; warns with

``FutureWarning``, removed in 4.0) - pass ``covariates=`` to

``fit()`` instead (the sklearn hyperparameter/data split).

Column names of covariates for **conditional** parallel trends

(``E[ΔY(0) | D=d, X] = E[ΔY(0) | D=0, X]``). When ``None`` (default) the

estimator uses unconditional parallel trends. Covariates enter through a

covariate-adjusted per-cell control counterfactual (see ``estimation_method``).

Covariates are read from the base period of each ``(g, t)`` comparison.

Not currently composable with ``survey_design=`` (raises ``NotImplementedError``).

estimation_method : str, default="dr"

Covariate-adjustment method (only used when ``covariates`` is set):

``"reg"`` (outcome regression) or ``"dr"`` (doubly-robust, the default).

``"ipw"`` is **not supported on the dose / event-study aggregation** — pure

IPW's covariate adjustment is a single scalar level shift, so it cannot

adjust the dose-response *shape* (ACRT(d) would be identical to the

unconditional fit); it raises ``NotImplementedError``. ``reg`` and ``dr``

share the dose-response shape and ACRT(d); ``dr`` differs only in the

``overall_att`` / ATT(d) level and in its doubly-robust standard errors.

pscore_trim : float, default=0.01

Propensity-score trimming bound for the ``dr`` path (scores clipped to

``[pscore_trim, 1 - pscore_trim]``).

epv_threshold : float, default=10.0

Events-per-variable threshold for the ``dr`` propensity logit diagnostics.

pscore_fallback : str, default="error"

Action when ``dr`` propensity estimation raises (the logit IRLS fails

with a ``LinAlgError`` / ``ValueError``, e.g. perfect separation or rank

deficiency): ``"error"`` (re-raise — the default, fail-closed so a `dr`

fit never silently degrades to a non-DR estimate) or ``"unconditional"``

(fall back to an unconditional propensity with a warning; the affected

cells are then reg-like — use only when you knowingly accept that). Note:

low events-per-variable emits a diagnostic warning but does not itself

trigger the fallback.

treatment_type : str, default="continuous"

Dose-response model: ``"continuous"`` (B-spline sieve, the default) or

``"discrete"`` (saturated per-dose-level regression, CGBS 2024 Eq. 4.1).

On the discrete path each distinct dose level gets its own effect

coefficient — ``ATT(d_j) = mean_{D=d_j}(ΔY) − control`` (a per-level 2×2

DiD) — and ``ACRT(d_j)`` is the paper's backward finite difference on the

grid ``{0, d_1, ..., d_J}`` (``ACRT(d_1) = ATT(d_1)/d_1``, so a binary

dose ``D in {0, 1}`` gives ``ACRT = ATT``). It composes with

``covariates`` and ``survey_design`` and reduces to the per-level 2×2 DiD

standard error.

Multi-cohort fits must share the same dose support across cohorts (else

``NotImplementedError``); an off-support ``dvals`` value raises

``ValueError``.

Examples

--------

>>> from diff_diff import ContinuousDiD, generate_continuous_did_data

>>> data = generate_continuous_did_data(n_units=200, seed=42)

>>> est = ContinuousDiD(n_bootstrap=199, seed=42)

>>> results = est.fit(data, outcome="outcome", unit="unit",

... time="period", first_treat="first_treat",

... dose="dose")

>>> results.overall_att # doctest: +SKIP

>>> results.aggregate("dose") # doctest: +SKIP

"""

_VALID_CONTROL_GROUPS = {"never_treated", "not_yet_treated", "lowest_dose"}

_VALID_BASE_PERIODS = {"varying", "universal"}

_VALID_ESTIMATION_METHODS = {"reg", "dr", "ipw"}

_VALID_TREATMENT_TYPES = {"continuous", "discrete"}

def __init__(

self,

degree: int = 3,

num_knots: int = 0,

dvals: Optional[np.ndarray] = None,

control_group: str = "never_treated",

anticipation: int = 0,

base_period: str = "varying",

alpha: float = 0.05,

n_bootstrap: int = 0,

bootstrap_weights: str = "rademacher",

seed: Optional[int] = None,

rank_deficient_action: str = "warn",

covariates: Optional[List[str]] = None,

estimation_method: str = "dr",

pscore_trim: float = 0.01,

epv_threshold: float = 10.0,

pscore_fallback: str = "error",

treatment_type: str = "continuous",

):

self.degree = degree

self.num_knots = num_knots

self.dvals = np.asarray(dvals, dtype=float) if dvals is not None else None

self.control_group = control_group

self.anticipation = anticipation

self.base_period = base_period

self.alpha = alpha

self.n_bootstrap = n_bootstrap

self.bootstrap_weights = bootstrap_weights

self.seed = seed

self.rank_deficient_action = rank_deficient_action

# M-084: constructor covariates= is deprecated (removed in 4.0);

# the design-matrix column spec moves to fit() per the sklearn

# hyperparameter/data split. Raw-keep storage: the value still

# routes exactly as before, and get_params round-trips it.

if covariates is not None:

warn_deprecated_kwarg(

type(self).__name__,

"covariates",

"pass covariates to fit() instead",

)

self.covariates = covariates

self.estimation_method = estimation_method

self.pscore_trim = pscore_trim

self.epv_threshold = epv_threshold

self.pscore_fallback = pscore_fallback

self.treatment_type = treatment_type

self._validate_constrained_params()

def _validate_constrained_params(self) -> None:

"""Validate control_group, base_period, and estimation_method values."""

if self.control_group not in self._VALID_CONTROL_GROUPS:

raise ValueError(

f"Invalid control_group: '{self.control_group}'. "

f"Must be one of {self._VALID_CONTROL_GROUPS}."

)

if self.base_period not in self._VALID_BASE_PERIODS:

raise ValueError(

f"Invalid base_period: '{self.base_period}'. "

f"Must be one of {self._VALID_BASE_PERIODS}."

)

if self.estimation_method not in self._VALID_ESTIMATION_METHODS:

raise ValueError(

f"Invalid estimation_method: '{self.estimation_method}'. "

f"Must be one of {self._VALID_ESTIMATION_METHODS}."

)

if self.pscore_fallback not in {"unconditional", "error"}:

raise ValueError(

f"Invalid pscore_fallback: '{self.pscore_fallback}'. "

"Must be 'unconditional' or 'error'."

)

if not (np.isfinite(self.pscore_trim) and 0.0 <= self.pscore_trim < 0.5):

raise ValueError(

f"Invalid pscore_trim: {self.pscore_trim}. " "Must be finite and in [0, 0.5)."

)

if not (np.isfinite(self.epv_threshold) and self.epv_threshold > 0):

raise ValueError(

f"Invalid epv_threshold: {self.epv_threshold}. Must be finite and > 0."

)

if self.treatment_type not in self._VALID_TREATMENT_TYPES:

raise ValueError(

f"Invalid treatment_type: '{self.treatment_type}'. "

f"Must be one of {self._VALID_TREATMENT_TYPES}."

)

if self.control_group == "lowest_dose" and self.covariates is not None:

# The covariate estimand under lowest-dose-as-control shifts to

# conditional PT *relative to d_L* (E[ΔY(0)|D=d,X] = E[ΔY(0)|d_L,X]).

# Deferred (see the TODO.md ContinuousDiD row) rather than silently estimated with the wrong

# identifying assumption.

raise NotImplementedError(

"control_group='lowest_dose' does not yet compose with covariates= "

"(the conditional-parallel-trends estimand relative to the lowest "

"dose d_L is deferred). Use covariates=None for the unconditional "

"lowest-dose fit."

)

# get_params/set_params come from BaseEstimator.

# ------------------------------------------------------------------

# Main fit

# ------------------------------------------------------------------

def fit(

self,

data: pd.DataFrame,

outcome: str,

unit: str,

time: str,

first_treat: str,

dose: str,

aggregate: Any = NOT_SUPPLIED,

survey_design: Optional["SurveyDesign"] = None,

covariates: Optional[List[str]] = None,

) -> ContinuousDiDResults:

"""

Fit the continuous DiD estimator.

Parameters

----------

data : pd.DataFrame

Panel data.

outcome : str

Outcome column name.

unit : str

Unit identifier column.

time : str

Time period column.

first_treat : str

First treatment period column (0 or inf for never-treated).

dose : str

Continuous dose column.

aggregate : str, optional

DEPRECATED (row M-025, removed in 4.0) - aggregate as a

post-fit step instead: ``results.aggregate('event_study')``

for the binarized event study (underscored - the

``"eventstudy"`` spelling dies with this parameter), or

``results.aggregate('dose')`` / ``results.aggregate('simple')``

views. The dose-response curves and overall ATT/ACRT are

always computed by ``fit()``, so ``aggregate="dose"`` was

already a no-op. Supplying ANY value (including ``None``)

warns ``FutureWarning``; supplied values still run the legacy

routing unchanged until 4.0.

survey_design : SurveyDesign, optional

Survey design specification for design-based inference.

Supports weighted estimation and Taylor series linearization

variance with strata, PSU, and FPC.

covariates : list of str, optional

Covariate column names for the conditional-parallel-trends

estimand (the canonical fit-level home - row M-084). The

deprecated constructor ``covariates=`` still routes and warns;

supplying both raises ``ValueError``.

Returns

-------

ContinuousDiDResults

"""

# M-025 deprecation shim: a plain fit() never warns; supplying

# aggregate= with ANY value (None included) warns once, then the

# legacy routing below runs unchanged - "eventstudy" still

# computes the fit-time surface and invalid strings still reach

# the pre-existing ValueError. Only the SENTINEL normalizes to

# None (it would otherwise fail the _VALID_AGGREGATES check on

# every plain fit). The post-fit successor validates its own

# (unified) vocabulary.

if aggregate is not NOT_SUPPLIED:

warnings.warn(

"ContinuousDiD.fit(aggregate=) is deprecated and will be "

"removed in 4.0. Fit once, then aggregate as a post-fit "

"step: results = ContinuousDiD().fit(...); "

"results.aggregate('event_study') (note the underscore - "

"the 'eventstudy' spelling dies with this parameter) / "

".aggregate('dose') / .aggregate('simple'). The "

"dose-response curves and overall ATT/ACRT are always "

"computed by fit(), so aggregate='dose' was already "

"redundant.",

FutureWarning,

stacklevel=2,

)

else:

aggregate = None

# 1. Validate & prepare

_VALID_AGGREGATES = (None, "dose", "eventstudy")

if aggregate not in _VALID_AGGREGATES:

raise ValueError(

f"Invalid aggregate: '{aggregate}'. " f"Must be one of {_VALID_AGGREGATES}."

)

# Resolve survey design if provided

resolved_survey, survey_weights, survey_weight_type, survey_metadata = (

_resolve_survey_for_fit(survey_design, data, "analytical")

)

# Validate within-unit constancy for panel survey designs

if resolved_survey is not None:

_validate_unit_constant_survey(data, unit, survey_design)

# Bootstrap + survey supported via PSU-level multiplier bootstrap.

# M-084: fit-time covariates= is the canonical home; the deprecated

# constructor spec still routes (raw-keep). Supplying both is

# ambiguous and fails loudly.

if covariates is not None and self.covariates is not None:

raise ValueError(

"covariates= was supplied both to the constructor "

"(deprecated, row M-084) and to fit(); pass it to fit() only."

)

effective_covariates = covariates if covariates is not None else self.covariates

if self.control_group == "lowest_dose" and effective_covariates is not None:

# Mirror of the constructor-time guard for the fit-level spec.

raise NotImplementedError(

"control_group='lowest_dose' does not yet compose with covariates= "

"(the conditional-parallel-trends estimand relative to the lowest "

"dose d_L is deferred). Use covariates=None for the unconditional "

"lowest-dose fit."

)

df = data.copy()

cov_cols = list(effective_covariates) if effective_covariates else []

for col in [outcome, unit, time, first_treat, dose, *cov_cols]:

if col not in df.columns:

raise ValueError(f"Column '{col}' not found in data.")

# Covariate-path guards (conditional parallel trends).

if cov_cols:

if survey_design is not None:

raise NotImplementedError(

"ContinuousDiD does not yet support covariates= together with "

"survey_design= (weighted covariate outcome-regression / "

"propensity influence functions are a follow-up). Use one or "

"the other for now."

)

if self.estimation_method == "ipw":

raise NotImplementedError(

"estimation_method='ipw' is not supported with covariates on the "

"dose-response / event-study aggregation. Pure IPW's covariate "

"adjustment is a single scalar (a propensity-reweighted control "

"mean), which shifts only the ATT(d) level and leaves ACRT(d) "

"identical to the unconditional fit — it cannot adjust the "

"dose-response shape. Use estimation_method='reg' or 'dr'."

)

# Fail closed on missing/non-finite covariates: a per-cell fallback to

# unconditional estimation would silently mix conditional-PT and

# unconditional-PT cells in the aggregate (no-silent-failures).

cov_nonfinite = ~np.isfinite(df[cov_cols].to_numpy(dtype=float))

if cov_nonfinite.any():

n_bad = int(cov_nonfinite.any(axis=1).sum())

raise ValueError(

f"{n_bad} row(s) have missing/non-finite covariate values. "

"ContinuousDiD requires complete covariates (a per-cell fallback "

"would mix conditional and unconditional estimands). Drop or "

"impute the affected rows, or fit without covariates."

)

# Verify dose is time-invariant

dose_nunique = df.groupby(unit)[dose].nunique()

if dose_nunique.max() > 1:

bad_units = dose_nunique[dose_nunique > 1].index.tolist()

raise ValueError(

f"Dose must be time-invariant. Units with varying dose: {bad_units[:5]}"

)

# Normalize first_treat: +inf → 0 (R-style never-treated encoding).

# Count rows recategorized so users can see how many units just

# crossed from "treated at some point" to "never treated" — silent

# recategorization here would shift the control composition (axis-E

# silent coercion). Only positive infinity is recoded (to match the

# existing `.replace([np.inf, float("inf")], 0)` semantics on the

# next line).

first_treat_vals = df[first_treat].values

# Reject NaN first_treat explicitly. NaN survives preprocessing but

# satisfies neither the treated (g > 0) nor never-treated (g == 0)

# mask, so affected units would be silently excluded from the

# estimator (same silent-failure shape as `first_treat < 0`).

nan_mask = pd.isna(df[first_treat])

n_nan_first_treat = int(nan_mask.sum())

if n_nan_first_treat > 0:

raise ValueError(

f"{n_nan_first_treat} row(s) have NaN '{first_treat}' "

f"values. Valid values are 0 (never-treated) or a positive "

f"treatment period; such units would otherwise be silently "

f"excluded from both treated and control pools."

)

inf_mask = np.isposinf(first_treat_vals)

n_inf_first_treat = int(inf_mask.sum())

if n_inf_first_treat > 0:

warnings.warn(

f"{n_inf_first_treat} row(s) have inf in '{first_treat}'; "

f"treating the corresponding units as never-treated. Pass an "

f"explicit never-treated marker (0) if this is not intended.",

UserWarning,

stacklevel=2,

)

# Reject negative first_treat values (including -inf) explicitly.

# Without this guard they would survive preprocessing but fall out of

# both the treated (g > 0) and never-treated (g == 0) masks, silently

# excluding the affected units.

negative_mask = first_treat_vals < 0

n_negative_first_treat = int(negative_mask.sum())

if n_negative_first_treat > 0:

raise ValueError(

f"{n_negative_first_treat} row(s) have negative '{first_treat}' "

f"values (including -inf). Valid values are 0 (never-treated) "

f"or a positive treatment period; such units would otherwise "

f"be silently excluded from both treated and control pools."

)

df[first_treat] = df[first_treat].replace([np.inf, float("inf")], 0)

# Drop units with positive first_treat but zero dose (R convention)

unit_info = df.groupby(unit).first()[[first_treat, dose]]

drop_units = unit_info[(unit_info[first_treat] > 0) & (unit_info[dose] == 0)].index

if len(drop_units) > 0:

warnings.warn(

f"Dropping {len(drop_units)} units with positive first_treat but zero dose.",

UserWarning,

stacklevel=2,

)

df = df[~df[unit].isin(drop_units)]

# Validate no negative doses among treated units

treated_doses = df.loc[df[first_treat] > 0, dose]

if (treated_doses < 0).any():

n_neg = int((treated_doses < 0).sum())

raise ValueError(

f"Found {n_neg} treated unit(s) with negative dose. "

f"Dose must be strictly positive for treated units (D > 0)."

)

# Discrete-dose handling / detection.

unit_doses = df.loc[df[first_treat] > 0].groupby(unit)[dose].first()

treated_unit_doses = unit_doses[unit_doses > 0]

unique_pos_doses = treated_unit_doses.unique()

if self.treatment_type == "discrete":

# Saturated regression: warn if the fit is over-parameterized

# (near-continuous / degenerate per-level SE) so the user can see

# that a saturated basis is a poor fit for near-continuous dose.

n_levels = len(unique_pos_doses)

n_treated_total = int(len(treated_unit_doses))

min_per_level = int(treated_unit_doses.value_counts().min()) if n_levels else 0

if n_levels and (min_per_level < 2 or n_levels > n_treated_total / 2):

warnings.warn(

f"treatment_type='discrete' with {n_levels} dose level(s) over "

f"{n_treated_total} treated unit(s) (min {min_per_level} unit(s) per "

"level). The saturated regression is over-parameterized / "

"near-continuous; per-level standard errors are degenerate when a "

"level has fewer than 2 units. Consider treatment_type='continuous' "

"(B-spline) if the dose is effectively continuous.",

UserWarning,

stacklevel=2,

)

else:

# Continuous B-spline path: flag an integer-valued dose so the user

# knows the saturated regression is available.

is_integer = len(unique_pos_doses) > 0 and np.allclose(

unique_pos_doses, np.round(unique_pos_doses)

)

if is_integer:

warnings.warn(

f"Dose appears discrete ({len(unique_pos_doses)} unique integer "

"values). B-spline smoothing may be inappropriate for discrete "

"treatments; pass treatment_type='discrete' for a saturated "

"(per-dose-level) regression.",

UserWarning,

stacklevel=2,

)

# Force dose=0 for never-treated units with nonzero dose. Report the

# affected row count via UserWarning so users can see whether their

# never-treated rows had unintended nonzero doses — silent zeroing

# here would quietly shift part of the control trajectory (axis-E

# silent coercion, paired with the `first_treat=inf -> 0` fix above).

never_treated_mask = df[first_treat] == 0

nonzero_dose_rows = never_treated_mask & (df[dose] != 0)

n_nonzero_dose_never_treated = int(nonzero_dose_rows.sum())

if n_nonzero_dose_never_treated > 0:

warnings.warn(

f"{n_nonzero_dose_never_treated} row(s) have '{first_treat}'=0 "

f"(never-treated) but nonzero '{dose}'; zeroing the dose. Pass "

f"dose=0 for never-treated rows to avoid this coercion.",

UserWarning,

stacklevel=2,

)

df.loc[never_treated_mask, dose] = 0.0

# Verify balanced panel

all_periods = set(df[time].unique())

unit_periods = df.groupby(unit)[time].apply(set)

is_unbalanced = unit_periods.apply(lambda s: s != all_periods)

if is_unbalanced.any():

n_bad = int(is_unbalanced.sum())

raise ValueError(

"Unbalanced panel detected. ContinuousDiD requires a balanced panel. "

f"{n_bad} unit(s) have missing periods."

)

# Identify groups and time periods

unit_cohort = df.groupby(unit)[first_treat].first()

treatment_groups = sorted([g for g in unit_cohort.unique() if g > 0])

time_periods = sorted(df[time].unique())

if len(treatment_groups) == 0:

raise ValueError("No treated units found (all first_treat == 0).")

n_control = int((unit_cohort == 0).sum())

if self.control_group == "never_treated" and n_control == 0:

raise ValueError(

"No never-treated units found. Use control_group='not_yet_treated' "

"or add never-treated units."

)

if self.control_group == "not_yet_treated" and n_control == 0:

raise ValueError(

"No never-treated (D=0) units found. With control_group='not_yet_treated', "

"dose-response curve identification requires P(D=0) > 0. For settings "

"with no untreated group, use control_group='lowest_dose' (Remark 3.1: "

"the lowest-dose group becomes the comparison, estimand ATT(d)-ATT(d_L)). "

"Otherwise add never-treated units or use a dataset with D=0 observations."

)

# Remark 3.1 (control_group="lowest_dose"): the lowest-dose group d_L is

# the comparison. Compute d_L ONCE here (from the treated unit doses,

# before precompute) and thread it via precomp -> every d_L-referencing

# consumer runs after this, and fit() stays config-idempotent (no fitted

# self attr). The d_L cluster (|dose - d_L| <= SATURATED_TOL) is the

# single source of truth for both the mask and the modelled dose set.

lowest_dose: Optional[float] = None

if self.control_group == "lowest_dose":

if n_control > 0:

raise ValueError(

"control_group='lowest_dose' is for settings with no never-treated "

f"units (Remark 3.1, P(D=0)=0), but {n_control} never-treated unit(s) "

"were found; they would be silently dropped. Use "

"control_group='never_treated' or 'not_yet_treated', or remove the "

"never-treated units."

)

if len(treatment_groups) > 1:

# NOTE (deferred multi-cohort follow-up): a future multi-cohort

# lowest_dose must use a WITHIN-cohort d_L reference and a

# support-aware cross-cohort aggregation, and must exclude the d_L

# controls from the survey group/bin mass sums (which key off

# unit_cohorts==g and would otherwise double-count them). Harmless

# today because this path is fenced off here.

raise NotImplementedError(

"control_group='lowest_dose' with multiple treatment cohorts is not "

f"yet implemented ({len(treatment_groups)} cohorts found). Remark 3.1 "

"is defined for a single treatment date; use a single-cohort panel "

"(multi-period single-cohort is supported)."

)

dose_arr = treated_unit_doses.to_numpy(dtype=float)

d_L = float(np.min(dose_arr))

n_dL = int(np.sum(np.abs(dose_arr - d_L) <= SATURATED_TOL))

if n_dL < 2:

msg = (

f"control_group='lowest_dose' requires a lowest-dose *group* — a mass "

f"point at the minimum dose d_L={d_L:g} with >= 2 units (P(D=d_L) > 0), "

f"but only {n_dL} unit is at d_L. The reference group must have enough "

"units to form its own control variance; a singleton minimum is not a "

"lowest-dose group."

)

if self.treatment_type != "discrete":

msg += (

" On a truly continuous dose without a mass point at the minimum, "

"Remark 3.1 does not apply."

)

raise ValueError(msg)

above = dose_arr[dose_arr - d_L > SATURATED_TOL]

if len(saturated_dose_levels(above)) < 1:

raise ValueError(

f"control_group='lowest_dose': no treated dose above the lowest dose "

f"d_L={d_L:g}. The estimand ATT(d)-ATT(d_L) needs at least one dose "

"level above d_L, but all treated units share the same dose."

)

# A lowest modelled dose d_1 very close to d_L makes the boundary

# ACRT(d_1)=ATT(d_1)/(d_1-d_L) and its SE explode; warn (not an error).

d_1 = float(np.min(above))

dose_span = float(np.max(dose_arr) - d_L)

if dose_span > 0 and (d_1 - d_L) < 0.01 * dose_span:

warnings.warn(

f"control_group='lowest_dose': the lowest modelled dose d_1={d_1:g} is "

f"very close to the reference d_L={d_L:g} (gap {d_1 - d_L:g}, "

f"{100 * (d_1 - d_L) / dose_span:.2g}% of the dose range); the boundary "

"ACRT(d_1)=ATT(d_1)/(d_1-d_L) and its standard error may be very large.",

UserWarning,

stacklevel=2,

)

lowest_dose = d_L

# Re-resolve survey design on filtered df if rows were dropped

# (survey arrays must align with df, not the original data)

if resolved_survey is not None and len(df) < len(data):

resolved_survey, survey_weights, survey_weight_type, survey_metadata = (

_resolve_survey_for_fit(survey_design, df, "analytical")

)

# 2. Precompute structures

precomp = self._precompute_structures(

df,

outcome,

unit,

time,

first_treat,

dose,

time_periods,

survey_weights=survey_weights,

covariates=effective_covariates,

)

# Thread the lowest-dose reference d_L (Remark 3.1) to the per-cell

# dose-response so it swaps the control group and shifts the discrete

# ACRT reference. None on the never/not-yet-treated paths.

precomp["lowest_dose"] = lowest_dose

# Compute dvals (evaluation grid); for discrete treatment, also the

# saturated dose levels (the global basis support). Under lowest_dose the

# lowest-dose group d_L is the reference (not modelled): the basis / grid

# / levels span only the *modelled* doses strictly above d_L.

all_treated_doses = precomp["dose_vector"][precomp["dose_vector"] > 0]

if lowest_dose is not None:

modelled_doses = all_treated_doses[all_treated_doses - lowest_dose > SATURATED_TOL]

if self.dvals is not None:

bad = self.dvals[self.dvals <= lowest_dose + SATURATED_TOL]

if bad.size:

raise ValueError(

f"control_group='lowest_dose': dvals contain {int(bad.size)} "

f"value(s) <= the reference dose d_L={lowest_dose:g}. d_L is the "

"omitted reference (ATT(d_L)=0 by construction); evaluate the "

"dose-response only at doses strictly above d_L."

)

# Survey subpopulation weights could reduce the d_L control group to

# zero or a single positive-weight unit, leaving no identified

# reference to difference against: with one positive-weight d_L unit,

# mu_0 equals that unit's dY so its ee_control = w*(dY - mu_0) = 0 and

# the reference contributes zero variance (understated SE). The raw

# >= 2 guard runs before survey weights, so enforce the same effective

# >= 2 positive-weight requirement here. Fail closed.

usw0 = precomp.get("unit_survey_weights")

if usw0 is not None:

dv0 = precomp["dose_vector"]

dL_w = usw0[np.abs(dv0 - lowest_dose) <= SATURATED_TOL]

n_pos_dL = int(np.count_nonzero(dL_w > 0))

if n_pos_dL < 2:

raise ValueError(

f"control_group='lowest_dose': the lowest-dose group d_L="

f"{lowest_dose:g} has {n_pos_dL} positive-weight unit(s) after "

"survey/subpopulation weighting (< 2 needed for an identified "

"reference variance; a single positive-weight reference unit "

"contributes zero control-side variance). Widen the subpopulation "

"or use a different dose grid."

)

else:

modelled_doses = all_treated_doses

levels: Optional[np.ndarray] = None

if self.treatment_type == "discrete":

levels = saturated_dose_levels(modelled_doses)

if self.dvals is not None:

# A saturated model can only be evaluated at observed dose

# levels; reject an off-support request (no silent snapping).

off_support = np.array(

[not np.any(np.abs(levels - d) <= SATURATED_TOL) for d in self.dvals]

)

if off_support.any():

raise ValueError(

f"treatment_type='discrete': requested dvals contain "

f"{int(off_support.sum())} value(s) that are not observed dose "

f"levels {levels.tolist()}. The saturated basis can only be "

"evaluated at observed dose levels."

)

dvals = self.dvals

else:

dvals = levels

# Multi-cohort heterogeneous dose support would produce a silent-zero

# aggregation bias: a cohort missing a global level yields a dropped

# zero column -> that cell's att_d[level] = 0 -> the plain-sum dose

# aggregation biases that dose toward zero. Fence it off; support-aware

# aggregation (average each dose only over the cohorts that observe it)

# is a deferred follow-up. Single-cohort (incl. multi-period),

# 2-period, and shared-support multi-cohort are all allowed.

if len(treatment_groups) > 1:

for g in treatment_groups:

g_doses = precomp["dose_vector"][

(precomp["unit_cohorts"] == g) & (precomp["dose_vector"] > 0)

]

g_levels = saturated_dose_levels(g_doses)

if g_levels.shape != levels.shape or not np.allclose(

g_levels, levels, atol=SATURATED_TOL

):

raise NotImplementedError(

"treatment_type='discrete' with multiple treatment cohorts "

"requires every cohort to share the same dose support. "

f"Cohort {g} covers {g_levels.tolist()} but the global dose "

f"levels are {levels.tolist()}. Support-aware aggregation "

"(averaging each dose only over the cohorts that observe it) "

"is not yet implemented; use a single cohort or ensure a "

"shared dose support."

)

# Survey subpopulation weights could zero out every treated unit at a

# dose level, leaving that level in the (unweighted) basis support but

# with zero effective mass -> a dropped column -> a silent-zero ATT(d).

# Fail closed if any level has no positive treated weight anywhere.

usw = precomp.get("unit_survey_weights")

if usw is not None:

dv = precomp["dose_vector"]

empty = [

float(d) for d in levels if not np.any(usw[np.abs(dv - d) <= SATURATED_TOL] > 0)

]

if empty:

raise ValueError(

"treatment_type='discrete': dose level(s) "

f"{empty} have zero positive survey weight among treated units "

"(e.g. removed by a subpopulation filter). The saturated model "

"cannot estimate an unweighted level; drop the level from the "

"dose grid or widen the subpopulation."

)

elif self.dvals is not None:

dvals = self.dvals

else:

dvals = default_dose_grid(modelled_doses)

# Build B-spline knots from the modelled treated doses (excludes the d_L

# reference group under lowest_dose; unused on the discrete branch, but

# harmless to construct).

knots, degree = build_bspline_basis(

modelled_doses, degree=self.degree, num_knots=self.num_knots

)

# 3. Iterate over (g,t) cells

gt_results = {}

gt_bootstrap_info = {}

for g in treatment_groups:

for t in time_periods:

result = self._compute_dose_response_gt(

precomp,

g,

t,

knots,

degree,

dvals,

survey_weights=precomp.get("unit_survey_weights"),

resolved_survey=resolved_survey,

levels=levels,

)

if result is not None:

gt_results[(g, t)] = result

gt_bootstrap_info[(g, t)] = result.get("_bootstrap_info", {})

# Filter out NaN cells (e.g., from zero effective survey mass)

gt_results = {

gt: r for gt, r in gt_results.items() if np.isfinite(r.get("att_glob", np.nan))

}

if len(gt_results) == 0:

raise ValueError("No valid (g,t) cells computed.")

# 4. Aggregate

post_gt = {(g, t): r for (g, t), r in gt_results.items() if t >= g - self.anticipation}

# Dose-response aggregation

n_grid = len(dvals)

# NaN-initialized SE/CI fields (used when post_gt is empty or as defaults)

att_d_se = np.full(n_grid, np.nan)

att_d_ci_lower = np.full(n_grid, np.nan)

att_d_ci_upper = np.full(n_grid, np.nan)

acrt_d_se = np.full(n_grid, np.nan)

acrt_d_ci_lower = np.full(n_grid, np.nan)

acrt_d_ci_upper = np.full(n_grid, np.nan)

overall_att_se = np.nan

overall_att_t = np.nan

overall_att_p = np.nan

overall_att_ci = (np.nan, np.nan)

overall_acrt_se = np.nan

overall_acrt_t = np.nan

overall_acrt_p = np.nan

overall_acrt_ci = (np.nan, np.nan)

att_d_p = None

acrt_d_p = None

# Event study aggregation (binarized) — runs on ALL (g,t) cells

event_study_effects = None

if aggregate == "eventstudy":

event_study_effects = self._aggregate_event_study(

gt_results,

gt_bootstrap_info=gt_bootstrap_info,

unit_survey_weights=precomp.get("unit_survey_weights"),

unit_cohorts=precomp["unit_cohorts"],

anticipation=self.anticipation,

)

_survey_df = None # Set by analytical branch when survey is active

if len(post_gt) == 0:

warnings.warn(

"No post-treatment (g,t) cells available for aggregation. "

"This can occur when all treatments start after the last observed "

"period or all cells were skipped due to insufficient data.",

UserWarning,

stacklevel=2,

)

overall_att = np.nan

overall_acrt = np.nan

agg_att_d = np.full(n_grid, np.nan)

agg_acrt_d = np.full(n_grid, np.nan)

else:

# Compute cell weights: group-proportional (matching R's contdid convention).

# Each group g gets weight proportional to its number of treated units.

# When survey weights present, use sum(w_g) / sum(w) instead of n_g / N.

# Within each group, weight is divided equally among post-treatment cells.

group_n_treated = {}

group_n_post_cells = {}

unit_sw = precomp.get("unit_survey_weights")

for (g, t), r in post_gt.items():

if g not in group_n_treated:

if unit_sw is not None:

# Survey-weighted group size: sum of weights for treated units in g

g_mask = precomp["unit_cohorts"] == g

group_n_treated[g] = float(np.sum(unit_sw[g_mask]))

else:

group_n_treated[g] = float(r["n_treated"])

group_n_post_cells[g] = 0

group_n_post_cells[g] += 1

total_treated = sum(group_n_treated.values())

cell_weights = {}

if total_treated > 0:

for (g, t), r in post_gt.items():

pg = group_n_treated[g] / total_treated

cell_weights[(g, t)] = pg / group_n_post_cells[g]

agg_att_d = np.zeros(n_grid)

agg_acrt_d = np.zeros(n_grid)

overall_att = 0.0

overall_acrt = 0.0

for gt, w in cell_weights.items():

r = post_gt[gt]

agg_att_d += w * r["att_d"]

agg_acrt_d += w * r["acrt_d"]

overall_att += w * r["att_glob"]