antecedent

240class AcceptedGraph:
241    """Versioned accepted CPDAG/PAG/DAG/temporal completion for estimate-only clicks.
242
243    Estimate / prepare / refresh never call discovery. Only :meth:`rediscover`
244    or constructing a new handle may replace structure and bump :attr:`version`.
245    """
246
247    __slots__ = ("_graph", "_version", "_algorithm_id")
248
249    def __init__(
250        self,
251        graph: _GraphTypes,
252        *,
253        version: int = 1,
254        algorithm_id: str | None = None,
255    ) -> None:
256        if version < 1:
257            raise ValueError("version must be >= 1")
258        self._graph = graph
259        self._version = int(version)
260        self._algorithm_id = algorithm_id
261
262    @property
263    def graph(self) -> _GraphTypes:
264        return self._graph
265
266    @property
267    def version(self) -> int:
268        return self._version
269
270    @property
271    def algorithm_id(self) -> str | None:
272        return self._algorithm_id
273
274    @classmethod
275    def from_graph(
276        cls,
277        graph: _GraphTypes,
278        *,
279        algorithm_id: str | None = None,
280        version: int = 1,
281    ) -> AcceptedGraph:
282        """Hold a reviewed or hand-authored graph artifact."""
283        return cls(graph, version=version, algorithm_id=algorithm_id)
284
285    @classmethod
286    def from_discovery(
287        cls,
288        result: DiscoveryResult,
289        *,
290        algorithm_id: str,
291        version: int = 1,
292    ) -> AcceptedGraph:
293        """Accept a standalone ``discover_*`` result into a session artifact."""
294        if not algorithm_id:
295            raise ValueError("algorithm_id is required for discovery provenance")
296        algo = algorithm_id.lower().replace("pcmci_plus", "pcmci+")
297        if algo in ("pcmci", "pcmci+", "lpcmci"):
298            graph: _GraphTypes = _result_to_temporal_graph(result, algo)
299        else:
300            graph = _result_to_graph(result, algo)
301        return cls(graph, version=version, algorithm_id=algo)
302
303    def replace(
304        self,
305        graph: _GraphTypes,
306        *,
307        algorithm_id: str | None = None,
308    ) -> AcceptedGraph:
309        """Explicit structure replace (bumps version). Returns a new handle."""
310        return AcceptedGraph(
311            graph,
312            version=self._version + 1,
313            algorithm_id=algorithm_id if algorithm_id is not None else self._algorithm_id,
314        )
315
316    def rediscover(
317        self,
318        data: Any,
319        discovery: _AnyDiscovery,
320        *,
321        seed: int = 1,
322        threads: int = 1,
323    ) -> AcceptedGraph:
324        """User-triggered rediscovery; never called by estimate / prepare."""
325        if isinstance(discovery, (PCMCI, PCMCIPlus, LPCMCI)):
326            result, algo = _run_temporal_discovery(
327                data, discovery, seed=seed, threads=threads
328            )
329            graph: _GraphTypes = _result_to_temporal_graph(result, algo)
330        else:
331            result, algo = _run_static_discovery(
332                data, discovery, seed=seed, threads=threads
333            )
334            graph = _result_to_graph(result, algo)
335        return AcceptedGraph(graph, version=self._version + 1, algorithm_id=algo)
336
337    def analyze(self, data: Any, query: Any, **kwargs: Any) -> Any:
338        """Estimate on the held graph (default ``latency="interactive"``).
339
340        Rejects caller ``discovery=``. Does not bump :attr:`version`.
341        """
342        from .estimation import analyze
343
344        if "discovery" in kwargs and kwargs["discovery"] is not None:
345            raise CausalUnsupportedError(
346                "AcceptedGraph.analyze rejects discovery=; structure is already accepted "
347                "(call rediscover() for an explicit structure refresh)"
348            )
349        kwargs.pop("discovery", None)
350        kwargs.setdefault("latency", "interactive")
351        kwargs["graph"] = self._graph
352        return analyze(data, query=query, **kwargs)
353
354    def prepare(self, data: Any, *, query: Any, **kwargs: Any) -> Any:
355        """Compile-once prepared handle on the held static DAG/edges."""
356        from .estimation import PreparedAnalysis
357
358        if isinstance(self._graph, (Cpdag, Pag, Admg, TemporalCpdag, TemporalPag)):
359            raise CausalUnsupportedError(
360                "PreparedAnalysis requires a fully oriented Dag/TemporalDag (or edge list); "
361                "complete CPDAG/PAG review first, then AcceptedGraph.from_graph(...)"
362            )
363        kwargs.setdefault("latency", "interactive")
364        return PreparedAnalysis.prepare(data, query=query, graph=self._graph, **kwargs)  # type: ignore[arg-type]
365
366    def to_json(self) -> str:
367        """Serialize for durable hold (JSON interchange, not CBOR wire)."""
368        kind, payload = _encode_graph(self._graph)
369        return json.dumps(
370            {
371                "format": "causal.AcceptedGraph/v1",
372                "version": self._version,
373                "algorithm_id": self._algorithm_id,
374                "kind": kind,
375                "payload": payload,
376            },
377            separators=(",", ":"),
378        )
379
380    @classmethod
381    def from_json(cls, s: str) -> AcceptedGraph:
382        """Restore from :meth:`to_json`."""
383        obj = json.loads(s)
384        if obj.get("format") != "causal.AcceptedGraph/v1":
385            raise ValueError(f"unsupported AcceptedGraph format: {obj.get('format')!r}")
386        graph = _decode_graph(obj["kind"], obj["payload"])
387        return cls(
388            graph,
389            version=int(obj["version"]),
390            algorithm_id=obj.get("algorithm_id"),
391        )

219@dataclass(frozen=True)
220class AnalysisResult:
221    """Nested analysis result matching the Rust facade sections."""
222
223    identification: IdentificationView
224    estimate: EstimateView
225    posterior: PosteriorView | None
226    validation: ValidationView
227    performance: PerformanceView
228    diagnostics: list[str]
229    provenance: dict[str, Any]
230    mediation: MediationView | None = None
231    plan: PlanView | None = None
232    _raw: Any = None
233    _prepared: Any = None
234
235    @property
236    def effect(self) -> float:
237        """Primary scalar effect (mediation total when present, else estimate ATE/mean)."""
238        if self.mediation is not None and self.mediation.total is not None:
239            return float(self.mediation.total)
240        if self.estimate.mediation is not None and self.estimate.mediation.total is not None:
241            return float(self.estimate.mediation.total)
242        return self.estimate.ate
243
244    @property
245    def ate(self) -> float:
246        """Alias for :attr:`effect` (prefer ``effect`` for non-ATE queries)."""
247        return self.effect
248
249    def refresh(
250        self,
251        data: Mapping[str, Any] | Any,
252        *,
253        seed: int = 1,
254        threads: int = 1,
255    ) -> AnalysisResult:
256        """Re-estimate on new data via the retained prepared handle.
257
258        Only results from :meth:`PreparedAnalysis.estimate` / ``refresh`` support
259        this. One-shot :func:`analyze` results raise ``TypeError``.
260        """
261        if self._prepared is None:
262            raise TypeError(
263                "AnalysisResult.refresh requires a result from PreparedAnalysis; "
264                "use PreparedAnalysis.prepare(...) then estimate/refresh"
265            )
266        return self._prepared.estimate(data, seed=seed, threads=threads)
267
268    def refute(
269        self,
270        data: Mapping[str, Any] | Any,
271        suite: Refute | Literal["placebo", "full", "cheap"] | bool | str = "placebo",
272        *,
273        seed: int = 1,
274        threads: int = 1,
275        cancel: Any | None = None,
276    ) -> AnalysisResult:
277        """Second-click refute via the retained prepared handle."""
278        if self._prepared is None:
279            raise TypeError(
280                "AnalysisResult.refute requires a result from PreparedAnalysis; "
281                "use PreparedAnalysis.prepare(...) then estimate"
282            )
283        if isinstance(suite, Refute):
284            suite = str(suite)
285        return self._prepared.refute(
286            data, suite, seed=seed, threads=threads, cancel=cancel
287        )

10@dataclass(frozen=True)
11class AverageEffect:
12    """Average treatment effect (static tabular)."""
13
14    treatment: str
15    outcome: str
16    control_level: float = 0.0
17    active_level: float = 1.0
18    target_population: object | None = None
19    kind: Literal["average"] = "average"

20@dataclass(frozen=True)
21class Bayesian:
22    """Bayesian g-computation (Laplace / conjugate / HMC backends).
23
24    Parameters
25    ----------
26    n_draws:
27        Posterior draw count.
28    prior_scale:
29        Isotropic Gaussian coefficient prior scale when ``prior_from`` is unset.
30        Ignored when ``prior_from`` is provided.
31    prior_from:
32        Posterior artifact bytes from a previous ``result.posterior.artifact``,
33        or a ``ComposedPrior`` from ``compose_external_priors``.
34        Artifact hydrate is deferred until the target design is prepared.
35    mapping:
36        How to map an artifact into the target prior. ``None`` auto-selects:
37        identical coefficient subspace when designs match (sequential Bayes),
38        or ``PriorMapping.effect_functional(...)`` when designs differ and the
39        artifact has an effect quantity. Never silent ``coef_i → coef_i`` across
40        heterogeneous designs. Ignored when ``prior_from`` is a ``ComposedPrior``.
41    backend:
42        Inference backend: ``laplace`` (default), ``conjugate``, or ``hmc``.
43    """
44
45    n_draws: int = 1000
46    prior_scale: float = 10.0
47    prior_from: bytes | ComposedPrior | None = None
48    mapping: PriorMapping | None = None
49    backend: Literal["laplace", "conjugate", "hmc"] = "laplace"
50    kind: Literal["bayesian"] = "bayesian"

70@dataclass(frozen=True)
71class ConditionalEffect:
72    """Conditional / context average effect with a single effect modifier."""
73
74    treatment: str
75    outcome: str
76    modifier: str
77    control_level: float = 0.0
78    active_level: float = 1.0
79    kind: Literal["conditional"] = "conditional"

 95@dataclass(frozen=True)
 96class Counterfactual:
 97    """Unit-level ITE via GCM abduction–action–prediction."""
 98
 99    treatment: str
100    outcome: str
101    control_level: float = 0.0
102    active_level: float = 1.0
103    kind: Literal["counterfactual"] = "counterfactual"

24class Estimator(StrEnum):
25    """Estimation strategies (wire ids match Rust ``EstimatorId``)."""
26
27    LINEAR_ADJUSTMENT_ATE = "linear.adjustment.ate"
28    PROPENSITY_WEIGHTING = "propensity.weighting"
29    PROPENSITY_MATCHING = "propensity.matching"
30    PROPENSITY_STRATIFICATION = "propensity.stratification"
31    DISTANCE_MATCHING = "distance.matching"
32    AIPW = "aipw"
33    GLM_ADJUSTMENT = "glm.adjustment"
34    FRONTDOOR_TWO_STAGE = "frontdoor.two_stage"
35    IV_WALD = "iv.wald"
36    IV_2SLS = "iv.2sls"
37    RD_SHARP = "rd.sharp"
38    BAYESIAN_GCOMP = "bayesian.gcomp"
39    TEMPORAL_LINEAR_ADJUSTMENT = "temporal.linear.adjustment"
40    FUNCTIONAL_DISTRIBUTION = "functional.distribution"
41    FUNCTIONAL_EFFECT = "functional.effect"
42    CONDITIONAL_LINEAR_ADJUSTMENT = "conditional.linear.adjustment"
43    TEMPORAL_MEDIATION = "temporal.mediation"