202class CausalFlowDAG(nn.Module):
 203    """A causal normalizing flow defined by ``spec = {name: NodeSpec}``."""
 204
 205    def __init__(
 206        self, spec: dict[str, NodeSpec], device: str = "cpu", seed: int | None = None
 207    ):
 208        """Build the flow from ``spec``.
 209
 210        Args:
 211            seed: if given, seeds weight initialisation deterministically
 212                (``torch.manual_seed`` is called before the nodes are
 213                constructed). Because init happens here, this is the single
 214                obvious knob for a reproducible model — ``fit(seed=...)`` only
 215                controls minibatch shuffling.
 216        """
 217        super().__init__()
 218        if seed is not None:
 219            torch.manual_seed(seed)
 220        self.spec = spec
 221        self.order = validate_and_sort(spec)
 222        self.nodes = nn.ModuleDict(
 223            {name: _Node(name, spec[name], spec) for name in self.order}
 224        )
 225        self.device = torch.device(device)
 226        self.history: dict = {"train": [], "val": [], "lr": [], "time": []}
 227        self.meta: dict = {}  # provenance attached at save() (machine, versions)
 228        self.vc_center_info: dict = {}  # OOF bookkeeping of centered VC terms (fit)
 229        self.to(self.device)
 230
 231    # ------------------------------------------------------------------ data
 232    def _encode_parent(self, name: str, values: Tensor) -> Tensor:
 233        """Encode the values of a node for use as a parent feature.
 234
 235        This follows the original TRAM-DAG convention. A continuous parent stays
 236        raw, shape ``(n, 1)``. An ordinal parent is one-hot encoded, shape
 237        ``(n, levels)``.
 238        """
 239        node = self.spec[name]
 240        if isinstance(node, OrdinalNode):
 241            return torch.nn.functional.one_hot(
 242                values.long(), num_classes=node.levels
 243            ).to(values.dtype)
 244        return values.view(-1, 1)
 245
 246    @property
 247    def _dtype(self) -> torch.dtype:
 248        """Current model dtype (float32 normally; float64 inside fit_classical)."""
 249        return next(self.parameters()).dtype
 250
 251    def _tensorize(self, df: pd.DataFrame) -> dict[str, Tensor]:
 252        np_dtype = np.float64 if self._dtype == torch.float64 else np.float32
 253        out = {}
 254        for name in self.order:
 255            vals = torch.as_tensor(
 256                df[name].to_numpy(dtype=np_dtype), device=self.device
 257            )
 258            out[name] = vals
 259        return out
 260
 261    def _features(self, values: dict[str, Tensor]) -> dict[str, Tensor]:
 262        return {name: self._encode_parent(name, vals) for name, vals in values.items()}
 263
 264    # ------------------------------------------- centered-VC propensity (e_hat)
 265    def _vc_ehat_live(
 266        self, nd: _Node, values: dict[str, Tensor], n: int
 267    ) -> dict[str, Tensor] | None:
 268        """Recompute ``e_hat(pa_on) = P(on = 1 | pa_on)`` for the centered VC terms.
 269
 270        The value comes from this flow's own fitted ``on`` node, as a full-data
 271        propensity fit. That is the DML prediction convention. Training uses
 272        frozen out-of-fold values instead, see :meth:`fit`.
 273
 274        The result is detached, so no gradient reaches the ``on`` node from the
 275        loss of this node.
 276
 277        The function derives the value from the current parent values, so
 278        ``do``-mutilated sampling uses ``t - e_hat(x)`` with the intervened ``t``
 279        and the observed ``x``. It never reads a cached value.
 280        """
 281        out = {}
 282        for g in nd._vc_groups:
 283            if not g.center:
 284                continue
 285            on_nd = self.nodes[g.on]
 286            feats = self._features({p: values[p] for p in on_nd.parents})
 287            theta, shift = on_nd.theta_shift(
 288                feats, n, vc_ehat=self._vc_ehat_live(on_nd, values, n)
 289            )
 290            # binary ordinal on: P(on <= 0) = sigmoid(theta_0 - s),
 291            # so e = sigmoid(s - theta_0)
 292            out[g.on] = torch.sigmoid(shift - theta[:, 0]).detach()
 293        return out or None
 294
 295    def _vc_ehat_columns(self, nd: _Node) -> list[str]:
 296        """List the extra columns needed for the centered VC terms of ``nd``.
 297
 298        These are the columns beyond ``nd.parents``, namely the parents of the
 299        treatment nodes, found recursively.
 300        """
 301        cols: list[str] = []
 302        for g in nd._vc_groups:
 303            if not g.center:
 304                continue
 305            on_nd = self.nodes[g.on]
 306            cols += [p for p in on_nd.parents] + self._vc_ehat_columns(on_nd)
 307        return [c for c in dict.fromkeys(cols) if c not in nd.parents]
 308
 309    # ------------------------------------------------------------- likelihood
 310    def node_log_prob(
 311        self,
 312        values: dict[str, Tensor],
 313        nodes: list[str] | None = None,
 314        vc_ehat: dict[str, dict[str, Tensor]] | None = None,
 315    ) -> dict[str, Tensor]:
 316        """Per-node log-likelihood contributions, each (n,).
 317
 318        ``nodes`` restricts computation to a subset (used to skip frozen nodes
 319        during training — valid because the per-node losses are independent).
 320        ``vc_ehat`` ({node: {on: e_hat}}) overrides the propensity used by
 321        centered VC terms — ``fit`` passes the frozen **out-of-fold** values for
 322        the training rows; when omitted, the live full-fit propensity is
 323        recomputed from the flow's own treatment node.
 324        """
 325        feats = self._features(values)
 326        n = next(iter(values.values())).shape[0]
 327        out = {}
 328        for name in self.order if nodes is None else nodes:
 329            node = self.nodes[name]
 330            ehat = (
 331                vc_ehat.get(name)
 332                if vc_ehat is not None
 333                else self._vc_ehat_live(node, values, n)
 334            )
 335            theta, shift = node.theta_shift(feats, n, vc_ehat=ehat)
 336            x = values[name]
 337            if node.kind == "continuous":
 338                z0, ladj = node.ut.forward(theta, x)
 339                z = z0 + shift
 340                out[name] = StandardLogistic.log_prob(z) + ladj
 341            else:
 342                out[name] = ordinal_log_prob(theta, shift, x)
 343        return out
 344
 345    def log_prob(self, df: pd.DataFrame) -> Tensor:
 346        """Joint log-likelihood log p(x) per row, shape (n,)."""
 347        per_node = self.node_log_prob(self._tensorize(df))
 348        return torch.stack(list(per_node.values()), dim=0).sum(dim=0)
 349
 350    def nll(self, df: pd.DataFrame) -> dict[str, float]:
 351        """Mean negative log-likelihood per node (diagnostic)."""
 352        with torch.no_grad():
 353            per_node = self.node_log_prob(self._tensorize(df))
 354        return {k: float(-v.mean()) for k, v in per_node.items()}
 355
 356    # ------------------------------------------------------------------- fit
 357    def _set_ranges(self, train_df: pd.DataFrame, marginal_init: bool = False) -> None:
 358        """Map the train 5%/95% quantiles onto the transform domain.
 359
 360        This is the min-max scaling of the original implementation.
 361
 362        ``marginal_init``: opt-in calibrated Bernstein init (see ``fit``). Applied only
 363        on the first fit (the same ``not ut._fitted`` guard as range-setting), so a
 364        multi-phase fit does not reset a partially-trained intercept.
 365        """
 366        from .transforms import BernsteinUT, ordinal_marginal_init_theta
 367
 368        for name in self.order:
 369            node = self.nodes[name]
 370            if node.kind == "continuous" and not node.ut._fitted:
 371                q = train_df[name].quantile([0.05, 0.95])
 372                node.ut.set_range(q.iloc[0], q.iloc[1])
 373                if (
 374                    marginal_init
 375                    and isinstance(node.ut, BernsteinUT)
 376                    and isinstance(node.intercept, SimpleIntercept)
 377                ):
 378                    with torch.no_grad():
 379                        node.intercept.theta.copy_(node.ut.marginal_init_theta())
 380            elif (
 381                node.kind == "ordinal"
 382                and marginal_init
 383                and isinstance(node.intercept, SimpleIntercept)
 384                and not getattr(node.intercept, "_marginal_inited", False)
 385            ):
 386                # calibrate unconditional cutpoints to the marginal class log-odds
 387                counts = np.bincount(
 388                    train_df[name].to_numpy().astype(np.int64),
 389                    minlength=self.spec[name].levels,
 390                )
 391                with torch.no_grad():
 392                    node.intercept.theta.copy_(ordinal_marginal_init_theta(counts))
 393                node.intercept._marginal_inited = True
 394
 395    def fit(
 396        self,
 397        train_df: pd.DataFrame,
 398        val_df: pd.DataFrame | None = None,
 399        epochs: int = 500,
 400        learning_rate: float = 1e-2,
 401        batch_size: int = 512,
 402        verbose: int = 50,
 403        seed: int | None = None,
 404        restore_best: bool = False,
 405        schedule: str | None = None,
 406        plateau_patience: int = 15,
 407        freeze_patience: int | None = None,
 408        min_delta: float = 1e-4,
 409        marginal_init: bool = False,
 410        vc_warm_start: bool = True,
 411    ) -> CausalFlowDAG:
 412        """Jointly fit all nodes by maximum likelihood.
 413
 414        By default training keeps the **final** (converged) weights, so an
 415        all-``ls`` model trained to convergence reproduces the classical maximum
 416        likelihood estimate exactly (e.g. matches ``statsmodels``/``polr``).
 417
 418        The optimizer holds one parameter group per node. Because the joint NLL
 419        decomposes per node with independent gradients, per-node learning rates
 420        and freezing are exactly equivalent to independent per-node training.
 421
 422        Args:
 423            val_df: optional held-out set, used only for monitoring (and for
 424                ``restore_best``, ``schedule="plateau"`` and ``freeze_patience``).
 425                If omitted, the training set is used for the validation metric.
 426            restore_best: if True, snapshot each node's best-validation weights
 427                during training and restore them at the end. This is a mild
 428                early-stopping regularization and the convention of the original
 429                implementation. The fit is then *not* the training-data MLE, so
 430                leave it False for an exact classical comparison. Default False.
 431            schedule: learning-rate schedule. ``None`` = constant (the classic
 432                behavior); ``"onecycle"`` = ``OneCycleLR`` (warmup to
 433                ``learning_rate``, then anneal; stepped per batch);
 434                ``"cosine"`` = ``CosineAnnealingLR`` over ``epochs``;
 435                ``"plateau"`` = **per-node** decay: a node's lr is multiplied by
 436                0.3 whenever its own validation NLL hasn't improved by
 437                ``min_delta`` for ``plateau_patience`` epochs (floor 1e-3 ×
 438                ``learning_rate``).
 439            freeze_patience: if set, a node whose validation NLL hasn't improved
 440                by ``min_delta`` for this many epochs is **frozen** — excluded
 441                from the loss and backward pass (a real compute saving, since
 442                per-node losses are independent). When every node is frozen the
 443                fit returns early. Freeze epochs are recorded in
 444                ``history["frozen"]``.
 445            marginal_init: if True, calibrate each *unconditional* node's intercept
 446                to its marginal at init, instead of zuko's default zero init.
 447                Bernstein continuous nodes -> the linear map of the pre-scaled
 448                domain onto the standard-logistic 5%/95% quantiles (default is
 449                ~2.5x too steep); ordinal nodes -> cutpoints set to the empirical
 450                class log-odds (default zeros = near-uniform). Pure init — the
 451                converged MLE is unchanged — applied once (first fit only).
 452                Opt-in; default off. Affects only ``SimpleIntercept`` nodes
 453                (conditional ci intercepts are left untouched).
 454            vc_warm_start: if True (default), each ``VC`` term's ``beta0`` is
 455                initialised from the classical all-``ls`` solution of its node's
 456                conditional (deterministic L-BFGS on a throwaway proxy) before
 457                training, so the penalized head starts at the classical answer
 458                and only learns deviations. Applied once per term (a buffer that
 459                survives ``save``/``load`` guards re-runs). No-op without VC terms.
 460
 461        For ``VC`` terms the objective is the **penalized** NLL on the
 462        total-likelihood scale — each term adds ``penalty * ||b_theta weights||^2``
 463        to the summed NLL, i.e. ``penalty * ||w||^2 / n_train`` to the mean loss
 464        (a fixed Gaussian prior: the shrinkage vanishes as n grows, the classical
 465        penalized-likelihood convention; ``beta0`` unpenalized). The recorded
 466        ``history`` NLLs stay pure likelihoods. After training, each ``b_theta``
 467        is re-centered to mean zero over the training data (function-preserving;
 468        the constant moves into ``beta0``). ``VC(center=...)`` terms run a
 469        stage-1 out-of-fold propensity computation before the loop
 470        (:meth:`_vc_oof_stage`); the training loss uses those frozen OOF values,
 471        while the epoch-level validation monitor (and every post-fit query)
 472        uses the live full-fit treatment node.
 473
 474        Calling ``fit`` again continues training (e.g. a second phase with a
 475        lower learning rate); freezing state does not carry across calls.
 476        """
 477        if schedule not in (None, "onecycle", "cosine", "plateau"):
 478            raise ValueError(f"unknown schedule {schedule!r}")
 479        if seed is not None:
 480            torch.manual_seed(seed)
 481        self._set_ranges(train_df, marginal_init=marginal_init)
 482        if vc_warm_start:
 483            self._vc_warm_start(train_df)
 484        # VC effect heads whose L2 penalty joins the loss, per owning node
 485        vc_penalized = {
 486            name: [
 487                self.nodes[name].shifts[g.on]
 488                for g in self.nodes[name]._vc_groups
 489                if g.mods and self.nodes[name].shifts[g.on].penalty > 0
 490            ]
 491            for name in self.order
 492        }
 493        # stage 1 for centered VC terms (issue #30): frozen OUT-OF-FOLD e_hat
 494        # for the training rows — a plain tensor, so the Y-node loss has no
 495        # gradient path into the treatment node (per-node factorization intact).
 496        vc_ehat_train = self._vc_oof_stage(train_df)
 497
 498        train_vals = self._tensorize(train_df)
 499        val_vals = self._tensorize(val_df) if val_df is not None else train_vals
 500        n = len(train_df)
 501        steps_per_epoch = (n + batch_size - 1) // batch_size
 502
 503        opt = torch.optim.Adam(
 504            [
 505                {
 506                    "params": list(self.nodes[name].parameters()),
 507                    "lr": learning_rate,
 508                    "node": name,
 509                }
 510                for name in self.order
 511            ]
 512        )
 513        sched = None
 514        if schedule == "onecycle":
 515            sched = torch.optim.lr_scheduler.OneCycleLR(
 516                opt, max_lr=learning_rate, total_steps=epochs * steps_per_epoch
 517            )
 518        elif schedule == "cosine":
 519            sched = torch.optim.lr_scheduler.CosineAnnealingLR(
 520                opt, T_max=epochs, eta_min=learning_rate * 1e-3
 521            )
 522
 523        if restore_best and not hasattr(self, "_best"):
 524            self._best = {name: (float("inf"), None) for name in self.order}
 525        best = self._best if restore_best else None
 526        # per-node plateau/freeze bookkeeping (local to this fit call)
 527        node_best = {name: float("inf") for name in self.order}
 528        node_bad = {name: 0 for name in self.order}
 529        frozen: set[str] = set()
 530        t0 = time.perf_counter()
 531        t_offset = self.history["time"][-1] if self.history.get("time") else 0.0
 532        prev_train: dict[str, float] = {}
 533
 534        for epoch in range(epochs):
 535            self.train()
 536            active = [name for name in self.order if name not in frozen]
 537            perm = torch.randperm(n, device=self.device)
 538            train_acc = {name: prev_train.get(name, float("nan")) for name in frozen}
 539            train_acc.update({name: 0.0 for name in active})
 540            for start in range(0, n, batch_size):
 541                idx = perm[start : start + batch_size]
 542                batch = {k: v[idx] for k, v in train_vals.items()}
 543                ehat_batch = (
 544                    None
 545                    if vc_ehat_train is None
 546                    else {
 547                        nm: {on: e[idx] for on, e in d.items()}
 548                        for nm, d in vc_ehat_train.items()
 549                    }
 550                )
 551                per_node = self.node_log_prob(batch, nodes=active, vc_ehat=ehat_batch)
 552                node_nlls = {k: -v.mean() for k, v in per_node.items()}
 553                loss = torch.stack(list(node_nlls.values())).sum()
 554                for name in active:  # VC penalty (excluded from history)
 555                    for m in vc_penalized[name]:
 556                        loss = loss + m.penalty * m.l2() / n
 557                opt.zero_grad()
 558                loss.backward()
 559                opt.step()
 560                if schedule == "onecycle":
 561                    sched.step()
 562                w = len(idx) / n
 563                for k, v in node_nlls.items():
 564                    train_acc[k] += float(v.detach()) * w
 565            if schedule == "cosine":
 566                sched.step()
 567            prev_train = train_acc
 568
 569            self.eval()
 570            with torch.no_grad():
 571                val_per_node = {
 572                    k: float(-v.mean()) for k, v in self.node_log_prob(val_vals).items()
 573                }
 574            self.history["train"].append(train_acc)
 575            self.history["val"].append(val_per_node)
 576            self.history.setdefault("lr", []).append(
 577                max(g["lr"] for g in opt.param_groups)
 578            )
 579            self.history.setdefault("time", []).append(
 580                t_offset + time.perf_counter() - t0
 581            )
 582
 583            # per-node improvement tracking (plateau decay + freezing)
 584            for g in opt.param_groups:
 585                name = g["node"]
 586                if name in frozen:
 587                    continue
 588                if val_per_node[name] < node_best[name] - min_delta:
 589                    node_best[name] = val_per_node[name]
 590                    node_bad[name] = 0
 591                else:
 592                    node_bad[name] += 1
 593                if (
 594                    schedule == "plateau"
 595                    and node_bad[name] > 0
 596                    and node_bad[name] % plateau_patience == 0
 597                ):
 598                    g["lr"] = max(g["lr"] * 0.3, learning_rate * 1e-3)
 599                # under "plateau", only freeze nodes whose lr has already been
 600                # decayed substantially — otherwise a node can freeze while a
 601                # smaller lr would still make progress toward the optimum
 602                lr_decayed = schedule != "plateau" or g[
 603                    "lr"
 604                ] <= learning_rate * 1e-2 * (1 + 1e-9)
 605                if (
 606                    freeze_patience is not None
 607                    and lr_decayed
 608                    and node_bad[name] >= freeze_patience
 609                ):
 610                    frozen.add(name)
 611                    self.history.setdefault("frozen", {}).setdefault(
 612                        name, len(self.history["val"])
 613                    )  # 1-based global epoch
 614
 615            if restore_best:
 616                for name in self.order:
 617                    if val_per_node[name] < best[name][0]:
 618                        best[name] = (
 619                            val_per_node[name],
 620                            copy.deepcopy(self.nodes[name].state_dict()),
 621                        )
 622
 623            if verbose and (epoch % verbose == 0 or epoch == epochs - 1):
 624                tot_t = sum(train_acc.values())
 625                tot_v = sum(val_per_node.values())
 626                print(
 627                    f"[epoch {epoch + 1:5d}/{epochs}] train NLL {tot_t:.4f}  "
 628                    f"val NLL {tot_v:.4f}"
 629                    + (f"  frozen {sorted(frozen)}" if frozen else "")
 630                )
 631
 632            if len(frozen) == len(self.order):  # everything converged
 633                if verbose:
 634                    print(f"[epoch {epoch + 1:5d}] all nodes frozen — stopping.")
 635                break
 636
 637        if restore_best:  # restore per-node best-validation weights
 638            for name, (_, state) in best.items():
 639                if state is not None:
 640                    self.nodes[name].load_state_dict(state)
 641        self._recenter_vc(train_vals)
 642        self.eval()
 643        return self
 644
 645    # ------------------------------------------------- varying-coefficient (VC)
 646    def _vc_warm_start(self, train_df: pd.DataFrame) -> None:
 647        """Initialise every VC term's ``beta0`` from the classical solution.
 648
 649        The value comes from the all-``ls`` solution of the node's conditional.
 650        Issue #28 recommends this warm start.
 651
 652        A throwaway proxy of the node (same kind/transform, every parent an LS
 653        term, parent marginals irrelevant to the conditional because the joint
 654        NLL decomposes per node) is fitted with the deterministic
 655        :meth:`fit_classical`, and the ``on`` coefficient copied into ``beta0``
 656        (for a binary ordinal treatment, the identified one-hot difference
 657        ``w[1] - w[0]``). ``b_theta`` already starts at the zero function
 658        (zero-initialised output layer). Runs once per term — the
 659        ``warm_started`` buffer survives ``save``/``load``.
 660        """
 661        for name in self.order:
 662            nd = self.nodes[name]
 663            todo = [g for g in nd._vc_groups if not bool(nd.shifts[g[0]].warm_started)]
 664            if not todo:
 665                continue
 666            node_spec = self.spec[name]
 667            proxy_spec: dict[str, NodeSpec] = {}
 668            for p in nd.parents:
 669                pn = self.spec[p]
 670                proxy_spec[p] = (
 671                    OrdinalNode(levels=pn.levels)
 672                    if isinstance(pn, OrdinalNode)
 673                    else ContinuousNode(transform="affine")
 674                )
 675            ls_terms = [LS(p) for p in nd.parents]
 676            if isinstance(node_spec, OrdinalNode):
 677                proxy_spec[name] = OrdinalNode(levels=node_spec.levels, terms=ls_terms)
 678            else:
 679                proxy_spec[name] = ContinuousNode(
 680                    transform=node_spec.transform,
 681                    transform_kwargs=dict(node_spec.transform_kwargs),
 682                    terms=ls_terms,
 683                )
 684            proxy = CausalFlowDAG(proxy_spec, device=str(self.device))
 685            proxy.fit_classical(train_df[list(nd.parents) + [name]], verbose=False)
 686            for g in todo:
 687                w = proxy.nodes[name].shifts[g.on].weight.detach()
 688                b0 = float(w[-1] - w[0]) if g.on_is_ord else float(w[0])
 689                m = nd.shifts[g.on]
 690                with torch.no_grad():
 691                    m.beta0.fill_(b0)
 692                m.warm_started.fill_(True)
 693
 694    @torch.no_grad()
 695    def _predict_p1(self, on: str, df: pd.DataFrame) -> np.ndarray:
 696        """Give ``P(on = 1 | pa_on)`` from this flow's ``on`` node.
 697
 698        The treatment is binary ordinal, so the value is
 699        ``sigmoid(shift - theta_0)``.
 700        """
 701        nd = self.nodes[on]
 702        np_dtype = np.float64 if self._dtype == torch.float64 else np.float32
 703        values = {
 704            p: torch.as_tensor(df[p].to_numpy(dtype=np_dtype), device=self.device)
 705            for p in nd.parents
 706        }
 707        feats = self._features(values)
 708        theta, shift = nd.theta_shift(
 709            feats, len(df), vc_ehat=self._vc_ehat_live(nd, values, len(df))
 710        )
 711        return torch.sigmoid(shift - theta[:, 0]).cpu().numpy()
 712
 713    def _vc_oof_stage(
 714        self, train_df: pd.DataFrame
 715    ) -> dict[str, dict[str, Tensor]] | None:
 716        """Compute stage 1 of the two-stage centered-VC design, issue #30.
 717
 718        The result holds the frozen training-time propensities, as
 719        ``{node: {on: (n,) tensor}}``.
 720
 721        For ``center=True`` the values are **out-of-fold** — K refits of the
 722        treatment node only, each predicting its held-out fold (the DML
 723        cross-fitting requirement; in-sample e_hat reintroduces the
 724        own-observation bias and can be worse than no centering). For
 725        ``center="col"`` the user-supplied cross-fitted column is taken as-is.
 726        Bookkeeping lands in ``self.vc_center_info[(node, on)]`` (``e_oof``,
 727        ``fold_id``, ``folds``, ``source``) so tests can assert the fold
 728        structure — a later "simplification" to in-sample e_hat fails CI.
 729        """
 730        jobs = [
 731            (name, g)
 732            for name in self.order
 733            for g in self.nodes[name]._vc_groups
 734            if g.center
 735        ]
 736        if not jobs:
 737            return None
 738        self.vc_center_info = {}
 739        np_dtype = np.float64 if self._dtype == torch.float64 else np.float32
 740        out: dict[str, dict[str, Tensor]] = {}
 741        rng_state = torch.get_rng_state()  # proxies reseed; keep fit reproducible
 742        try:
 743            for name, g in jobs:
 744                if isinstance(g.center, str):  # user-supplied cross-fitted col
 745                    if g.center not in train_df.columns:
 746                        raise KeyError(f"center column {g.center!r} not in train_df.")
 747                    e = train_df[g.center].to_numpy(dtype=np.float64)
 748                    if not ((e > 0.0) & (e < 1.0)).all():
 749                        raise ValueError(
 750                            f"center column {g.center!r} must hold propensities "
 751                            "strictly inside (0, 1)."
 752                        )
 753                    fold_id = None
 754                else:
 755                    e, fold_id = self._vc_oof_propensity(g.on, train_df, g.folds)
 756                out.setdefault(name, {})[g.on] = torch.as_tensor(
 757                    e.astype(np_dtype), device=self.device
 758                )
 759                self.vc_center_info[(name, g.on)] = {
 760                    "source": g.center if isinstance(g.center, str) else "oof-refit",
 761                    "folds": None if fold_id is None else int(g.folds),
 762                    "fold_id": fold_id,
 763                    "e_oof": e.copy(),
 764                    "n": len(train_df),
 765                }
 766        finally:
 767            torch.set_rng_state(rng_state)
 768        return out
 769
 770    def _vc_oof_propensity(
 771        self, on: str, train_df: pd.DataFrame, k: int
 772    ) -> tuple[np.ndarray, np.ndarray]:
 773        """Compute the out-of-fold ``P(on=1|pa_on)``.
 774
 775        The function refits the ``on`` node K times, and only that node. Each
 776        refit uses a single-node proxy whose parents are sources, because their
 777        marginals cannot influence the conditional. Each refit then predicts the
 778        fold it never saw.
 779
 780        A treatment with all-``ls`` terms uses :meth:`fit_classical`, which is
 781        deterministic and takes seconds. Any other treatment uses a
 782        fixed-budget Adam fit.
 783        """
 784        on_nd = self.nodes[on]
 785        if any(g.center for g in on_nd._vc_groups):
 786            raise NotImplementedError(
 787                f"treatment node {on!r} itself has a centered VC term. "
 788                "Chained centering is not supported."
 789            )
 790        node_spec = self.spec[on]
 791        proxy_spec: dict[str, NodeSpec] = {}
 792        for p in on_nd.parents:
 793            pn = self.spec[p]
 794            proxy_spec[p] = (
 795                OrdinalNode(levels=pn.levels)
 796                if isinstance(pn, OrdinalNode)
 797                else ContinuousNode(transform="affine")
 798            )
 799        terms = list(node_spec.terms) if node_spec.terms else None
 800        proxy_spec[on] = OrdinalNode(levels=2, terms=terms)
 801        all_ls = all(t.effect == "LS" for t in (terms or []))
 802        cols = list(on_nd.parents) + [on]
 803
 804        n = len(train_df)
 805        fold_id = np.random.default_rng(0).permutation(n) % k
 806        e = np.empty(n, dtype=np.float64)
 807        for j in range(k):
 808            proxy = CausalFlowDAG(proxy_spec, device=str(self.device), seed=0)
 809            held_in = train_df.iloc[fold_id != j][cols]
 810            if all_ls:
 811                proxy.fit_classical(held_in, verbose=False)
 812            else:
 813                proxy.fit(
 814                    held_in,
 815                    epochs=300,
 816                    learning_rate=1e-2,
 817                    verbose=0,
 818                    seed=0,
 819                    restore_best=False,
 820                )
 821            e[fold_id == j] = proxy._predict_p1(on, train_df.iloc[fold_id == j])
 822        return e, fold_id
 823
 824    @torch.no_grad()
 825    def _recenter_vc(self, values: dict[str, Tensor]) -> None:
 826        """Re-split every VC term so ``b_theta`` sums to zero over the train rows.
 827
 828        The removed constant moves into ``beta0``, so the modelled function does
 829        not change.
 830        """
 831        feats: dict[str, Tensor] | None = None
 832        for name in self.order:
 833            nd = self.nodes[name]
 834            for g in nd._vc_groups:
 835                if not g.mods:
 836                    continue
 837                if feats is None:
 838                    feats = self._features(values)
 839                nd.shifts[g.on].recenter(torch.cat([feats[p] for p in g.mods], dim=1))
 840
 841    @torch.no_grad()
 842    def varying_coef(
 843        self, node: str, data: pd.DataFrame, on: str | None = None
 844    ) -> np.ndarray:
 845        """Evaluate the fitted effect function ``beta(x)`` of a ``VC`` term.
 846
 847        The function reads the rows of ``data``. It is the first-class read-out
 848        of issue #28.
 849
 850        The value comes in closed form from the fitted term, as
 851        ``beta0 + b_theta(modifiers)``. It is deterministic, it is free of ``y``
 852        because only the modifier columns of ``data`` are read, and it needs no
 853        abduction. For a binary treatment it is identical to the abduction
 854        difference ``u(x, t=1, y) - u(x, t=0, y)``.
 855
 856        The value lives on the latent, log-odds scale of the node. A continuous
 857        node adds it. An ordinal node subtracts it from the cutpoints.
 858
 859        For a centered term, that is ``center=...``, the form of the returned
 860        ``beta`` does not change, but ``beta0`` then reads as the effect at the
 861        treatment margin, which is the observed propensities.
 862
 863        Args:
 864            node: name of the node carrying the VC term.
 865            on: the VC term's treatment name; optional when the node has exactly
 866                one VC term.
 867
 868        Returns an ``(n,)`` array of ``beta`` values (constant when the term has
 869        no modifiers).
 870        """
 871        if node not in self.nodes:
 872            raise KeyError(f"unknown node {node!r}")
 873        nd = self.nodes[node]
 874        vcs = {g.on: g.mods for g in nd._vc_groups}
 875        if not vcs:
 876            raise ValueError(f"node {node!r} has no VC term.")
 877        if on is None:
 878            if len(vcs) > 1:
 879                raise ValueError(
 880                    f"node {node!r} has several VC terms ({sorted(vcs)}). "
 881                    "Pass on=<treatment name>."
 882                )
 883            on = next(iter(vcs))
 884        if on not in vcs:
 885            raise KeyError(
 886                f"node {node!r} has no VC term on {on!r} (has {sorted(vcs)})."
 887            )
 888        mods = vcs[on]
 889        missing = [p for p in mods if p not in data.columns]
 890        if missing:
 891            raise KeyError(f"data is missing modifier column(s): {missing}")
 892        mod_feat = None
 893        if mods:
 894            np_dtype = np.float64 if self._dtype == torch.float64 else np.float32
 895            vals = {
 896                p: torch.as_tensor(data[p].to_numpy(dtype=np_dtype), device=self.device)
 897                for p in mods
 898            }
 899            feats = self._features(vals)
 900            mod_feat = torch.cat([feats[p] for p in mods], dim=1)
 901        return nd.shifts[on].beta(mod_feat, len(data)).cpu().numpy()
 902
 903    # --------------------------------------------------------- classical fit
 904    def _is_all_ls(self) -> bool:
 905        return all(
 906            term.effect == "LS"
 907            for node in self.spec.values()
 908            for term in node_terms(node)
 909        )
 910
 911    def ls_coefficients(self) -> dict[str, dict[str, np.ndarray]]:
 912        """Give the per-node linear-shift weights, as ``{node: {parent: array}}``.
 913
 914        For an all-``ls`` model these are the interpretable log-odds-ratio
 915        coefficients.
 916        """
 917        out: dict[str, dict[str, np.ndarray]] = {}
 918        for name in self.order:
 919            shifts = self.nodes[name].shifts
 920            if shifts:
 921                out[name] = {
 922                    p: m.weight.detach().cpu().numpy().ravel().copy()
 923                    for p, m in shifts.items()
 924                }
 925        return out
 926
 927    def to_matrix(self) -> pd.DataFrame:
 928        """Give the labelled adjacency matrix of term effects.
 929
 930        Rows are parents and columns are children. This is the meta-adjacency
 931        view of the paper. A cell holds ``"LS"``, ``"CS"`` or ``"CI"``, and an
 932        empty cell means there is no edge. A multi-parent term carries its parent
 933        group as a suffix.
 934        """
 935        labels = {"I": "CI", "LS": "LS", "CS": "CS"}
 936        m = pd.DataFrame("", index=list(self.order), columns=list(self.order))
 937        for child in self.order:
 938            for term in node_terms(self.spec[child]):
 939                if term.effect == "VC":  # treatment cell "VC", modifiers "VCm"
 940                    cells = [(term.parents[0], "VC")] + [
 941                        (p, "VCm") for p in term.parents[1:]
 942                    ]
 943                else:
 944                    tag = labels[term.effect]
 945                    if len(term.parents) > 1:
 946                        tag = f"{tag}{list(term.parents)}"
 947                    cells = [(p, tag) for p in term.parents]
 948                for p, tag in cells:  # a VC modifier may share its cell with
 949                    cur = m.loc[p, child]  # a prognostic term -> join with "+"
 950                    m.loc[p, child] = f"{cur}+{tag}" if cur else tag
 951        return m
 952
 953    @torch.no_grad()
 954    def intercept_contributions(self, node: str, data: pd.DataFrame) -> dict:
 955        """Decompose a complex intercept into mean-centered per-term parts.
 956
 957        The parts are contributions to the transform parameters of the node. Use
 958        them to plot additive partial effects.
 959
 960        An additive complex intercept ``terms=[I("x1"), I("x2")]`` builds one
 961        network per ``I``-term and **sums their outputs in unconstrained
 962        parameter space**: ``theta(pa) = net_1(x1) + net_2(x2)``. The sum is
 963        identified (so every L1/L2/L3 query is correct), but each term's output is
 964        identified only up to a constant — a constant moves freely between the
 965        nets. This makes the *raw* per-term outputs not directly comparable.
 966
 967        Following the usual additive-model / GAM convention, this resolves the
 968        ambiguity by a **sum-to-zero (mean-centering) constraint applied over the
 969        rows of** ``data``: each term's contribution is centered to mean zero
 970        (per parameter), and the removed constants are collected into a single
 971        ``baseline``. The decomposition is exact —
 972
 973            ``theta(pa) = baseline + sum_terms contribution_term(pa)``
 974
 975        — so ``baseline`` plus the (uncentered) row sum of the contributions
 976        reproduces the model's transform parameters. This is **post-hoc only**:
 977        it reads the fitted weights and changes nothing about the model or any
 978        frozen number (issue #20, Option A). Shift terms (``LS``/``CS``) are a
 979        separate, already-interpretable slot — see :meth:`ls_coefficients`.
 980
 981        Args:
 982            node: name of a node with at least one complex-intercept (``I``) term
 983                that has parents.
 984            data: rows over which to center (and at which to evaluate the
 985                contributions); must contain every intercept-parent column.
 986
 987        Returns a dict with:
 988            ``"baseline"``: ``(P,)`` array — the absorbed constant (sum of the
 989                per-term means), where ``P`` is the node's transform-parameter
 990                count: ``ut.n_params`` for a continuous node (e.g. Bernstein
 991                coefficients, the [widths|heights|derivatives] block of an RQ
 992                spline, or the 2 affine parameters), ``levels - 1`` cutpoint
 993                parameters for an ordinal node. The contributions live in the
 994                transform's **unconstrained** parameter space (where the model
 995                sums the additive terms, before the monotonicity constraint), so
 996                they are exact partial effects on those parameters but not, in
 997                general, an additive shift of the curve itself.
 998            ``"contributions"``: ``{term_label: (n, P) array}`` — each term's
 999                mean-centered contribution at each row (columns sum to ~0 over
1000                rows). ``term_label`` is the term's parents joined by ``"+"``.
1001            ``"parents"``: ``{term_label: tuple(parent_names)}``.
1002        """
1003        if node not in self.nodes:
1004            raise KeyError(f"unknown node {node!r}")
1005        nd = self.nodes[node]
1006        groups = nd._intercept_groups
1007        if not groups:
1008            raise ValueError(
1009                f"node {node!r} has no complex-intercept (I) terms with parents. "
1010                "Its intercept is unconditional, so there is nothing to decompose."
1011            )
1012        missing = [p for p in nd.ci_parents if p not in data.columns]
1013        if missing:
1014            raise KeyError(f"data is missing intercept-parent column(s): {missing}")
1015
1016        np_dtype = np.float64 if self._dtype == torch.float64 else np.float32
1017        vals = {
1018            p: torch.as_tensor(data[p].to_numpy(dtype=np_dtype), device=self.device)
1019            for p in nd.ci_parents
1020        }
1021        feats = self._features(vals)
1022        # one net per group: the additive case stores them in intercept_nets;
1023        # a single (possibly joint) I-term is the lone `intercept` network.
1024        nets = (
1025            list(nd.intercept_nets) if nd.intercept_nets is not None else [nd.intercept]
1026        )
1027
1028        contributions: dict[str, np.ndarray] = {}
1029        parents: dict[str, tuple] = {}
1030        baseline = None
1031        for net, grp in zip(nets, groups):
1032            raw = net(torch.cat([feats[p] for p in grp], dim=1))  # (n, P)
1033            mean = raw.mean(dim=0, keepdim=True)  # (1, P)
1034            label = "+".join(grp)
1035            contributions[label] = (raw - mean).cpu().numpy()
1036            parents[label] = grp
1037            baseline = mean if baseline is None else baseline + mean
1038        return {
1039            "baseline": baseline.cpu().numpy().ravel(),
1040            "contributions": contributions,
1041            "parents": parents,
1042        }
1043
1044    def fit_classical(
1045        self,
1046        train_df: pd.DataFrame,
1047        *,
1048        max_iter: int = 400,
1049        tol: float = 1e-6,
1050        verbose: bool = True,
1051    ) -> dict:
1052        """Fit an all-``ls`` model the classical way.
1053
1054        The fit uses full batches, float64, and L-BFGS with a strong-Wolfe line
1055        search. There are no minibatches, no schedule and no early stopping, so
1056        the fit is deterministic and bit-reproducible. It lands on the exact
1057        maximum-likelihood estimate and matches classical software, that is
1058        ``statsmodels`` ``OrderedModel`` and R ``polr`` or ``Colr``. It is much
1059        faster than minibatch Adam.
1060
1061        This method is valid only when every edge is ``ls``, because each
1062        node-conditional is then a classical transformation model. Any other
1063        model raises. For a ``cs`` or ``ci`` model use :meth:`fit`, where the
1064        minibatch noise also regularizes the MLPs.
1065
1066        float64 is a transient compute mode. The model is upcast for the fit,
1067        and ``self.double()`` converts the parameters and the range buffers of
1068        the transforms in one call. Afterwards the model returns to float32, so
1069        the stored model and ``save``/``load`` stay float32. Double precision is
1070        what lets the line search resolve the optimum cleanly.
1071
1072        Convergence is judged by **NLL flatness** (relative change < ``tol``
1073        between L-BFGS rounds). Note that ``|grad|`` and individual coefficients
1074        do *not* settle to machine precision: a continuous node's Bernstein
1075        intercept, and weakly-identified directions like rare one-hot levels or a
1076        flat treatment-effect ridge, keep drifting along near-zero-curvature
1077        valleys long after the likelihood (and the well-identified coefficients)
1078        have reached the MLE. Correctness is therefore verified by comparison to
1079        classical software (see ``experiments/validate_ls.py``), not by this flag.
1080
1081        Returns a convergence report (iterations, final NLL, gradient norm,
1082        max coefficient change at the last round, wall-time, and the fitted
1083        :meth:`ls_coefficients`).
1084        """
1085        if not self._is_all_ls():
1086            raise ValueError(
1087                "fit_classical requires an all-`ls` spec, that is every edge "
1088                "term 'ls'. This spec has cs, ci or vc terms. Use fit() for "
1089                "flexible models."
1090            )
1091        self._set_ranges(train_df)
1092
1093        self.double()  # parameters + buffers (xmin/xmax) -> float64, one call
1094        assert next(self.parameters()).dtype == torch.float64
1095        t0 = time.perf_counter()
1096        chunk = 25  # inner L-BFGS iterations per round; we stop on NLL change
1097        try:
1098            vals = self._tensorize(train_df)
1099            self.train()
1100            opt = torch.optim.LBFGS(
1101                self.parameters(),
1102                lr=1.0,
1103                max_iter=chunk,
1104                history_size=50,
1105                tolerance_grad=0.0,
1106                tolerance_change=0.0,
1107                line_search_fn="strong_wolfe",
1108            )
1109
1110            def closure():
1111                opt.zero_grad()
1112                nll = torch.stack(
1113                    [-lp.mean() for lp in self.node_log_prob(vals).values()]
1114                ).sum()
1115                nll.backward()
1116                return nll
1117
1118            def flat_coefs() -> np.ndarray:
1119                cs = self.ls_coefficients()
1120                return (
1121                    np.concatenate([w for node in cs.values() for w in node.values()])
1122                    if cs
1123                    else np.zeros(1)
1124                )
1125
1126            prev_nll, prev_c, final_nll, n_iter, converged, coef_delta = (
1127                float("inf"),
1128                flat_coefs(),
1129                float("nan"),
1130                0,
1131                False,
1132                float("inf"),
1133            )
1134            for _ in range(max(1, max_iter // chunk)):
1135                final_nll = float(opt.step(closure))
1136                n_iter += chunk
1137                cur_c = flat_coefs()
1138                coef_delta = float(np.abs(cur_c - prev_c).max())
1139                prev_c = cur_c
1140                if abs(prev_nll - final_nll) < tol * (1.0 + abs(final_nll)):
1141                    converged = True
1142                    break
1143                prev_nll = final_nll
1144            grad_norm = float(
1145                torch.cat(
1146                    [
1147                        p.grad.reshape(-1)
1148                        for p in self.parameters()
1149                        if p.grad is not None
1150                    ]
1151                ).norm()
1152            )
1153            coefs = self.ls_coefficients()  # read while still float64
1154        finally:
1155            self.float()  # restore canonical float32 (lossy ~1e-7, harmless)
1156        self.eval()
1157
1158        report = {
1159            "converged": converged,
1160            "n_iter": n_iter,
1161            "final_nll": final_nll,
1162            "grad_norm": grad_norm,
1163            "coef_delta": coef_delta,
1164            "seconds": time.perf_counter() - t0,
1165            "coefficients": coefs,
1166        }
1167        if verbose:
1168            print(
1169                f"fit_classical: {n_iter} L-BFGS iters, NLL {final_nll:.6f}, "
1170                f"{report['seconds']:.2f}s"
1171                + ("" if converged else f"  (NLL still moving at {max_iter} iters)")
1172            )
1173        return report
1174
1175    # ------------------------------------------------------- causal queries
1176    @torch.no_grad()
1177    def sample(
1178        self,
1179        n: int | None = None,
1180        *,
1181        do: dict[str, float] | None = None,
1182        u: pd.DataFrame | None = None,
1183        seed: int | None = None,
1184    ) -> pd.DataFrame:
1185        """Sample from the (optionally mutilated) flow.
1186
1187        Args:
1188            n: number of samples (ignored if ``u`` is given).
1189            do: interventions {node: value}; intervened nodes are clamped and
1190                their parent dependence removed (graph mutilation).
1191            u: latent variables (as returned by :meth:`abduct`). If given, they
1192                are pushed through the flow — together with ``do`` this yields
1193                counterfactuals (Pearl's abduction -> action -> prediction).
1194        """
1195        do = do or {}
1196        gen = None
1197        if seed is not None:
1198            gen = torch.Generator(device=self.device).manual_seed(seed)
1199
1200        np_dtype = np.float64 if self._dtype == torch.float64 else np.float32
1201        if u is not None:
1202            n = len(u)
1203            u_vals = {
1204                name: torch.as_tensor(
1205                    u[name].to_numpy(dtype=np_dtype, copy=True), device=self.device
1206                )
1207                for name in self.order
1208            }
1209        elif n is not None:
1210            u_vals = {
1211                name: StandardLogistic.sample((n,), device=self.device)
1212                if gen is None
1213                else StandardLogistic.icdf(
1214                    torch.rand((n,), device=self.device, generator=gen)
1215                )
1216                for name in self.order
1217            }
1218        else:
1219            raise ValueError("Provide either n or u.")
1220
1221        values: dict[str, Tensor] = {}
1222        for name in self.order:
1223            if name in do:
1224                values[name] = torch.full(
1225                    (n,), float(do[name]), dtype=self._dtype, device=self.device
1226                )
1227                continue
1228            node = self.nodes[name]
1229            feats = self._features({p: values[p] for p in node.parents})
1230            # centered VC: e_hat(pa_on) is re-derived from the already-sampled
1231            # ancestor values — under do the regressor is t_do - e_hat(x), never
1232            # a cached training value
1233            theta, shift = node.theta_shift(
1234                feats, n, vc_ehat=self._vc_ehat_live(node, values, n)
1235            )
1236            z = u_vals[name]
1237            if node.kind == "continuous":
1238                values[name] = node.ut.inverse(theta, z - shift)
1239            else:
1240                values[name] = ordinal_sample(theta, shift, z)
1241        return pd.DataFrame({k: v.cpu().numpy() for k, v in values.items()})
1242
1243    @torch.no_grad()
1244    def abduct(self, df: pd.DataFrame, seed: int | None = None) -> pd.DataFrame:
1245        """Pearl abduction: recover the latent variables ``u`` from observations.
1246
1247        Continuous nodes are inverted exactly (``u = h(x) + shift``); for ordinal
1248        nodes the latent is only interval-identified, so it is sampled from the
1249        standard logistic truncated to the observed level's interval.
1250        """
1251        gen = None
1252        if seed is not None:
1253            gen = torch.Generator(device=self.device).manual_seed(seed)
1254        values = self._tensorize(df)
1255        feats = self._features(values)
1256        n = len(df)
1257        u = {}
1258        for name in self.order:
1259            node = self.nodes[name]
1260            theta, shift = node.theta_shift(
1261                feats, n, vc_ehat=self._vc_ehat_live(node, values, n)
1262            )
1263            x = values[name]
1264            if node.kind == "continuous":
1265                z0, _ = node.ut.forward(theta, x)
1266                u[name] = z0 + shift
1267            else:
1268                u[name] = ordinal_abduct(theta, shift, x, generator=gen)
1269        return pd.DataFrame({k: v.cpu().numpy() for k, v in u.items()})
1270
1271    @torch.no_grad()
1272    def pmf(
1273        self, df: pd.DataFrame, node: str, do: dict[str, float] | None = None
1274    ) -> np.ndarray:
1275        """Give the analytic class probabilities of an ordinal node.
1276
1277        The result has shape ``(n, levels)``. The parents of the node come from
1278        ``df``, after the ``do`` overrides are applied.
1279        """
1280        if not isinstance(self.spec[node], OrdinalNode):
1281            raise ValueError(f"pmf() requires an ordinal node, '{node}' is continuous.")
1282        df_local = df.copy()
1283        for col, val in (do or {}).items():
1284            df_local[col] = val
1285        nd = self.nodes[node]
1286        np_dtype = np.float64 if self._dtype == torch.float64 else np.float32
1287        cols = list(nd.parents) + self._vc_ehat_columns(nd)  # + e_hat inputs
1288        values = {
1289            p: torch.as_tensor(df_local[p].to_numpy(dtype=np_dtype), device=self.device)
1290            for p in cols
1291        }
1292        feats = self._features({p: values[p] for p in nd.parents})
1293        theta, shift = nd.theta_shift(
1294            feats, len(df_local), vc_ehat=self._vc_ehat_live(nd, values, len(df_local))
1295        )
1296        return ordinal_pmf(theta, shift).cpu().numpy()
1297
1298    # ------------------------------------------------------------------ scores
1299    @torch.no_grad()
1300    def scores(
1301        self, df: pd.DataFrame, node: str, params: str = "shift"
1302    ) -> pd.DataFrame:
1303        """Give the per-observation scores ``psi_i = d l_i / d theta``, issue #29.
1304
1305        The scores belong to the interpretable shift coefficients of a node and
1306        are analytic and exact, see ``tramdag.scores``. ``params="shift"`` is the
1307        only option and covers every ``LS`` weight and the ``beta0`` of every
1308        ``VC`` term.
1309
1310        At a fitted MLE each column sums to about zero. Order the rows by a
1311        covariate that truly modifies the treatment effect and the cumulative sum
1312        of the treatment column drifts. :meth:`effect_modifier_scan` measures
1313        that drift.
1314
1315        This is a pure read-out. It touches no fitting or sampling code path.
1316        """
1317        if params != "shift":
1318            raise ValueError(f"params='shift' is the only option, got {params!r}.")
1319        from .scores import node_scores
1320
1321        return node_scores(self, df, node)
1322
1323    @torch.no_grad()
1324    def effect_modifier_scan(
1325        self, df: pd.DataFrame, node: str, on: str, candidates: list[str] | None = None
1326    ) -> pd.DataFrame:
1327        """Rank candidate effect modifiers with a Zeileis-Hornik fluctuation scan.
1328
1329        Issue #29 describes the method. Each candidate covariate is ranked by how
1330        strongly the scores of the ``on`` coefficient drift when the rows are
1331        ordered by it. A cheap all-``ls`` fit is enough, so this gives a measured
1332        shortlist for ``VC`` modifiers.
1333
1334        Returns
1335        -------
1336        pd.DataFrame
1337            One row per candidate, with ``stat``, ``p_value``, ``crit_5pct`` and
1338            ``flag``. See ``tramdag.scores.effect_modifier_scan``.
1339        """
1340        from .scores import effect_modifier_scan
1341
1342        return effect_modifier_scan(self, df, node, on, candidates=candidates)
1343
1344    # ------------------------------------------------------------------- io
1345    def save(self, path: str | Path) -> None:
1346        """Write the model, its history and its provenance to a checkpoint.
1347
1348        The file holds the spec and the weights, the training ``history``, and a
1349        ``meta`` block with the tramdag version, the save time, the device, and
1350        the machine that trained the model. A cached run therefore stays
1351        self-describing: the file alone is enough to rebuild a training-curve
1352        plot or to compare timings.
1353        """
1354        from datetime import datetime, timezone
1355
1356        from . import __version__
1357        from .env import machine_info
1358
1359        path = Path(path)
1360        path.parent.mkdir(parents=True, exist_ok=True)
1361        meta = {
1362            "tramdag_version": __version__,
1363            "saved_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
1364            "device": str(self.device),
1365            "machine": machine_info(),
1366        }
1367        torch.save(
1368            {
1369                "spec": spec_to_dict(self.spec),
1370                "state_dict": self.state_dict(),
1371                "history": self.history,
1372                "meta": meta,
1373            },
1374            path,
1375        )
1376
1377    @classmethod
1378    def load(cls, path: str | Path, device: str = "cpu") -> CausalFlowDAG:
1379        """Restore a model from a checkpoint.
1380
1381        ``flow.history`` and ``flow.meta`` are refilled, so a cached model can
1382        still produce training and diagnostic plots, and can report the machine
1383        that trained it.
1384        """
1385        ckpt = torch.load(path, map_location=device, weights_only=False)
1386        flow = cls(spec_from_dict(ckpt["spec"]), device=device)
1387        for name in flow.order:  # mark transforms as fitted before loading buffers
1388            node = flow.nodes[name]
1389            if node.kind == "continuous":
1390                node.ut._fitted = True
1391        flow.load_state_dict(ckpt["state_dict"])
1392        flow.history = ckpt.get(
1393            "history", {"train": [], "val": [], "lr": [], "time": []}
1394        )
1395        flow.meta = ckpt.get("meta", {})
1396        flow.eval()
1397        return flow