tramdag.flow

tramdag.flow

CausalFlowDAG: a single triangular normalizing flow on a user-defined DAG.

The flow maps iid standard-logistic latents U to the observed variables X in topological order. sample, abduct, pmf and density answer the observational, interventional (do=) and counterfactual queries; the read-outs that only read fitted weights live in readouts.py.

CausalFlowDAG

Bases: FitMixin, ReadoutsMixin, Module

A causal normalizing flow defined by a DAG specification.

Parameters:

Name	Type	Description	Default
spec	dict[str, NodeSpec]	The DAG specification, {name: ContinuousNode | OrdinalNode}.	required
device	str	Torch device, by default "cpu".	'cpu'
seed	int | None	If given, seeds the weight initialization. Each node draws its weights from a stream keyed on this seed and the node's name, so a node starts from the same weights in every spec that contains it. Construction does not change torch's global RNG state. Without a seed, the weights come from torch's global RNG. fit, fit_jointly, sample and abduct take their own seed=.	None
init	str	Weight initialization of every linear layer (the LS weights and the CI/CS/VC networks): "torch" (default, nn.Linear's Kaiming-uniform), "glorot" (glorot-uniform weights, zero biases) or "normal" (N(0, 0.05^2) weights and biases). A VC head's output layer stays zero either way. Stored in the checkpoint.	'torch'
zero_shift_networks	bool	Zero the output layer of every CS network after init, so each shift starts at 0, by default False. It draws no random numbers, so the other weights stay those of the seed. A VC head starts at zero anyway. Not stored: a checkpoint holds the weights.	False

Attributes:

Name	Type	Description
nodes	ModuleDict	One Node per variable, in topological order.
order	list[str]	The node names in topological order.
history	dict	The nodes' fit histories, see history.
meta	dict	The provenance that load fills.

Source code in src/tramdag/flow.py

class CausalFlowDAG(FitMixin, ReadoutsMixin, nn.Module):
    """A causal normalizing flow defined by a DAG specification.

    Parameters
    ----------
    spec : dict[str, NodeSpec]
        The DAG specification, ``{name: ContinuousNode | OrdinalNode}``.
    device : str, optional
        Torch device, by default ``"cpu"``.
    seed : int | None, optional
        If given, seeds the weight initialization. Each node draws its
        weights from a stream keyed on this seed and the node's name, so a
        node starts from the same weights in every spec that contains it.
        Construction does not change torch's global RNG state. Without a
        seed, the weights come from torch's global RNG. ``fit``,
        ``fit_jointly``, ``sample`` and ``abduct`` take their own ``seed=``.
    init : str, optional
        Weight initialization of every linear layer (the LS weights and the
        CI/CS/VC networks): ``"torch"`` (default, ``nn.Linear``'s
        Kaiming-uniform), ``"glorot"`` (glorot-uniform weights, zero biases)
        or ``"normal"`` (N(0, 0.05^2) weights and biases). A VC head's
        output layer stays zero either way. Stored in the checkpoint.
    zero_shift_networks : bool, optional
        Zero the output layer of every ``CS`` network after ``init``, so
        each shift starts at 0, by default ``False``. It draws no random
        numbers, so the other weights stay those of the seed. A ``VC`` head
        starts at zero anyway. Not stored: a checkpoint holds the weights.

    Attributes
    ----------
    nodes : nn.ModuleDict
        One [`Node`][tramdag.nodes.Node] per variable, in topological order.
    order : list[str]
        The node names in topological order.
    history : dict
        The nodes' fit histories, see
        [`history`][tramdag.flow.CausalFlowDAG.history].
    meta : dict
        The provenance that [`load`][tramdag.flow.CausalFlowDAG.load] fills.
    """

    def __init__(
        self,
        spec: dict[str, NodeSpec],
        device: str = "cpu",
        seed: int | None = None,
        init: str = "torch",
        zero_shift_networks: bool = False,
    ):
        super().__init__()
        if init not in ("torch", "glorot", "normal"):
            raise ValueError(
                f"init must be 'torch', 'glorot' or 'normal', got {init!r}"
            )
        self.spec = spec
        self.order = validate_and_sort(spec)
        self.init = init
        self.nodes = nn.ModuleDict(
            {name: self._build_node(name, seed) for name in self.order}
        )
        if zero_shift_networks:
            for m in self.modules():
                if isinstance(m, ComplexShiftModule):
                    nn.init.zeros_(m.net[-1].weight)
        self.device = torch.device(device)
        self.meta: dict = {}
        self.to(self.device)

    def outcome(self) -> str:
        """Give the one node that no other node reads.

        Returns
        -------
        str
            The name of the only node without children.

        Raises
        ------
        ValueError
            If the DAG has more than one such node; the message names them.
        """
        parents = {p for node in self.spec.values() for p in node_parents(node)}
        sinks = sorted(set(self.order) - parents)
        if len(sinks) != 1:
            raise ValueError(f"outcome needs exactly one sink, the DAG has {sinks!r}")
        return sinks[0]

    def _schema(self, name: str) -> dict[str, str | int]:
        """Give a node's parent schema, ``{parent: "continuous" | n_levels}``."""
        return {p: schema_entry(self.spec[p]) for p in node_parents(self.spec[name])}

    def _build_node(self, name: str, seed: int | None) -> Node:
        """Build and initialize one node from ``(seed, name)``.

        Without a seed the node draws from torch's global stream.
        """
        with seeded(seed, name):
            node = Node(name, self.spec[name], self._schema(name))
            if self.init != "torch":
                for m in node.modules():
                    if isinstance(m, nn.Linear):
                        _init_linear(m, self.init)
                for m in node.modules():
                    if isinstance(m, ShiftModule):
                        m.post_init()
        return node

    @property
    def _dtype(self) -> torch.dtype:
        """Current model dtype: float32, unless the caller converts the flow.

        Every tensor built from a frame takes this dtype, so the read-outs work
        in either dtype without carrying a dtype argument.
        """
        return next(self.parameters()).dtype

    def _tensorize(
        self,
        df: pd.DataFrame,
        cols: list[str] | tuple[str, ...] | None = None,
        *,
        levels: bool = True,
    ) -> dict[str, Tensor]:
        """DataFrame columns -> one ``(n,)`` tensor each, in the model dtype.

        ``cols=None`` takes every node, in topological order. ``levels=False``
        skips the ordinal level check, for a frame of *latents* whose columns
        carry node names but real values.

        Raises
        ------
        KeyError
            If the frame lacks one of the columns, by name. A spec/data
            mismatch would otherwise surface deep inside a tensor op.
        ValueError
            If an ordinal column is not a level index of its node. The
            one-hot encoding would otherwise truncate 1.5 to level 1 in
            silence, or fail inside ``one_hot`` without naming the node.
        """
        cols = self.order if cols is None else cols
        ordinal = {n: nd.levels for n, nd in self.nodes.items() if nd.kind == "ordinal"}
        return tensorize(df, cols, ordinal if levels else {}, self._dtype, self.device)

    def _to_frame(self, values: dict[str, Tensor], index=None) -> pd.DataFrame:
        """Tensors -> DataFrame; an ordinal column goes back as a level index."""
        out = {}
        for k, v in values.items():
            arr = v.cpu().numpy()
            out[k] = arr.astype(np.int64) if self.nodes[k].kind == "ordinal" else arr
        return pd.DataFrame(out, index=index)

    def _generator(self, seed: int | None) -> torch.Generator | None:
        """Give a seeded generator on this flow's device, or None for unseeded."""
        if seed is None:
            return None
        return torch.Generator(device=self.device).manual_seed(seed)

    def _node(self, name: str) -> Node:
        """Look a node up by name, with the same error everywhere."""
        if name not in self.nodes:
            raise KeyError(f"unknown node {name!r}")
        return self.nodes[name]

    def _check_do(self, do: dict[str, float]) -> None:
        """Refuse an unknown ``do`` node or an ordinal value off the levels."""
        for name, value in do.items():
            if self._node(name).kind == "ordinal":
                check_level_values(name, [value], self.spec[name].levels)

    def _features(self, values: dict[str, Tensor]) -> dict[str, Tensor]:
        return {
            name: self.nodes[name].encode(v)
            for name, v in values.items()
            if name in self.spec
        }

    def _theta_shift(
        self, nd: Node, feats: dict[str, Tensor], values: dict[str, Tensor], n: int
    ) -> tuple[Tensor, Tensor]:
        """Evaluate one node's transform parameters and shift.

        The side columns join here, so every query that reaches a node's
        parameters gets them. A query that composed the call itself would drop
        a centered term's propensity column by omission, which is the failure
        ``VaryingCoefficientModule`` has to catch at run time.

        ``feats`` is passed in rather than derived, because a caller looping
        over the nodes encodes the parents once for all of them.
        """
        return nd.theta_shift(feats | self._side_feats(nd, values, n), n)

    def _side_feats(
        self, nd: Node, values: dict[str, Tensor], n: int
    ) -> dict[str, Tensor]:
        """Give the node's side columns: frozen from ``values``, else live.

        Training frames carry them as ordinary columns; queries recompute
        them from the fitted flow (``ShiftModule.live_side``).
        """
        out = {}
        for m in nd.shifts.values():
            cols = m.side_columns()
            if all(c in values for c in cols):
                out.update({c: values[c] for c in cols})
            else:
                out.update(m.live_side(self, values, n))
        return out

    def _query_side_columns(self, nd: Node) -> list[str]:
        """List the extra columns a query needs to recompute live side inputs.

        These are the columns beyond ``nd.parents``: the parents of the
        treatment nodes (which cannot be centered themselves, so one level
        is all there is).
        """
        cols = [p for m in nd.shifts.values() for p in m.extra_columns(self)]
        return [c for c in dict.fromkeys(cols) if c not in nd.parents]

    @torch.no_grad()
    def _propensity(self, nd: Node, values: dict[str, Tensor], n: int) -> Tensor:
        r"""Give $P(\text{node} = 1 \mid \mathrm{pa})$ for a binary ordinal treatment.

        $P(x \le 0) = \sigma(\vartheta_0 - s)$, so the answer is
        $\sigma(s - \vartheta_0)$. No side columns: chained centering is refused
        by the spec, so a treatment node never carries a centered term itself.
        """
        theta, shift = nd.theta_shift(nd.features(values), n)
        return torch.sigmoid(shift - theta[:, 0])

    def _conditional(
        self, df: pd.DataFrame, node: str, do: dict[str, float] | None
    ) -> tuple[Node, Tensor, Tensor, int]:
        """Evaluate one node's conditional at the rows of ``df``.

        ``do`` overrides columns before the parents are read, which is what
        makes [`pmf`][] and [`density`][] interventional. Gives the node,
        its transform parameters, its shift and the row count.
        """
        nd = self._node(node)
        do = do or {}
        self._check_do(do)
        columns = list(nd.parents) + self._query_side_columns(nd)
        for name in do:
            if name not in columns:
                raise ValueError(
                    f"do key {name!r} is not a parent of node {node!r}; "
                    f"its parents are {columns!r}"
                )
        df = df.assign(**do)
        n = len(df)
        values = self._tensorize(df, columns)
        theta, shift = self._theta_shift(nd, nd.features(values), values, n)
        return nd, theta, shift, n

    def _calibrate(
        self, train_df: pd.DataFrame, *, marginal_init: bool = False
    ) -> CausalFlowDAG:
        """Calibrate every node on the training rows ([`Node.calibrate`][]).

        The columns are checked for all nodes first. When a node refuses the
        frame, the nodes this call calibrated lose their flag again, so a
        corrected frame calibrates the whole flow anew.
        """
        check_columns(train_df, self.order)
        fresh = [nd for nd in self.nodes.values() if not bool(nd.calibrated)]
        try:
            for nd in fresh:
                nd.calibrate(train_df, marginal_init=marginal_init)
        except ValueError:
            for nd in fresh:
                nd.calibrated.fill_(False)
            raise
        return self

    def _node_log_prob(
        self,
        values: dict[str, Tensor],
        nodes: list[str] | None = None,
    ) -> dict[str, Tensor]:
        """Compute the per-node log-likelihood contributions.

        Parameters
        ----------
        values : dict[str, Tensor]
            Raw node values, keyed by node name, each shape ``(n,)``.
        nodes : list[str] | None, optional
            Restrict the computation to these nodes. A subset is exact
            because the per-node losses
            are independent. ``None`` (default) computes every node.

        Returns
        -------
        dict[str, Tensor]
            One log-likelihood tensor per node, each shape ``(n,)``.
        """
        feats = self._features(values)
        n = next(iter(values.values())).shape[0]
        out = {}
        for name in self.order if nodes is None else nodes:
            nd = self.nodes[name]
            theta, shift = self._theta_shift(nd, feats, values, n)
            out[name] = nd.log_prob(theta, shift, values[name])
        return out

    @torch.no_grad()
    def log_prob(self, df: pd.DataFrame, *, nodes: list[str] | None = None) -> Tensor:
        """Compute the joint log-likelihood per row.

        Parameters
        ----------
        df : pd.DataFrame
            Observations, one column per node.
        nodes : list[str] | None, optional
            Sum only these nodes' contributions. A subset is exact, because
            the per-node losses are independent: ``nodes=["Y"]`` is the
            conditional log-likelihood of ``Y`` given its parents, per row,
            in log space (safer than the log of [`pmf`][], which
            underflows in the tail). ``None`` (default) is the joint.

        Returns
        -------
        Tensor
            ``log p(x)`` per row, shape ``(n,)``.

        Raises
        ------
        ValueError
            If ``nodes`` is empty, or an ordinal value is not a level index.
        KeyError
            If ``nodes`` names an unknown node, or ``df`` lacks a column.
        """
        if nodes is not None and not nodes:
            raise ValueError("nodes=[] sums nothing; omit it for the joint")
        for name in nodes or ():
            self._node(name)
        per_node = self._node_log_prob(self._tensorize(df), nodes)
        return torch.stack(list(per_node.values()), dim=0).sum(dim=0)

    @torch.no_grad()
    def nll(self, df: pd.DataFrame) -> dict[str, float]:
        """Compute the mean negative log-likelihood per node (a diagnostic).

        Parameters
        ----------
        df : pd.DataFrame
            Observations, one column per node.

        Returns
        -------
        dict[str, float]
            The mean NLL, keyed by node name.
        """
        per_node = self._node_log_prob(self._tensorize(df))
        return {k: float(-v.mean()) for k, v in per_node.items()}

    @torch.no_grad()
    def sample(
        self,
        n: int | None = None,
        *,
        do: dict[str, float] | None = None,
        u: pd.DataFrame | None = None,
        seed: int | None = None,
    ) -> pd.DataFrame:
        """Sample from the flow, with optional interventions.

        Parameters
        ----------
        n : int | None, optional
            Number of samples. Ignored if ``u`` is given.
        do : dict[str, float] | None, optional
            Interventions, as ``{node: value}``. An intervened node is
            clamped and its parent dependence removed (graph mutilation).
        u : pd.DataFrame | None, optional
            Latent variables, as returned by [`abduct`][]. If given,
            they are pushed through the flow. With ``do``, this gives a
            counterfactual.
        seed : int | None, optional
            If given, seeds the latent draw. Ignored if ``u`` is given.

        Returns
        -------
        pd.DataFrame
            The samples, one column per node.

        Raises
        ------
        ValueError
            If both ``n`` and ``u`` are omitted, or a ``do`` value of an
            ordinal node is not a level index.

        Notes
        -----
        Under ``do``, a centered VC re-derives ``t_do - e_hat(x)``, so the
        features of a node are never cached. With ``u`` given, the result takes
        the index of ``u``, so a counterfactual lines up with its factual row.
        """
        do = do or {}
        self._check_do(do)
        if u is not None:
            n = len(u)
            u_vals = self._tensorize(u, levels=False)
        elif n is not None:
            gen = self._generator(seed)
            u_vals = {
                name: StandardLogistic.sample((n,), device=self.device, generator=gen)
                for name in self.order
            }
        else:
            raise ValueError("sample() needs n or u")

        values: dict[str, Tensor] = {}
        for name in self.order:
            if name in do:
                values[name] = torch.full(
                    (n,), float(do[name]), dtype=self._dtype, device=self.device
                )
                continue
            nd = self.nodes[name]
            feats = nd.features(values)
            theta, shift = self._theta_shift(nd, feats, values, n)
            values[name] = nd.sample(theta, shift, u_vals[name])
        return self._to_frame(values, index=None if u is None else u.index)

    @torch.no_grad()
    def abduct(self, df: pd.DataFrame, *, seed: int | None = None) -> pd.DataFrame:
        r"""Recover the latent variables ``u`` from observations (the abduction step).

        A continuous node inverts exactly: $u = h(x) + s$. This is
        $\operatorname{logit} P(X \le x \mid \text{parents})$, an
        interventional quantity at the given parent values. For an
        ordinal node the latent is only interval-identified, so it is
        sampled from the standard logistic truncated to the observed
        level's interval.

        Parameters
        ----------
        df : pd.DataFrame
            Observations, one column per node.
        seed : int | None, optional
            If given, seeds the truncated draw for the ordinal nodes.

        Returns
        -------
        pd.DataFrame
            The latents, one column per node, aligned with the rows of
            ``df``.
        """
        gen = self._generator(seed)
        values = self._tensorize(df)
        feats = self._features(values)
        n = len(df)
        u = {}
        for name in self.order:
            nd = self.nodes[name]
            theta, shift = self._theta_shift(nd, feats, values, n)
            u[name] = nd.abduct(theta, shift, values[name], generator=gen)
        return pd.DataFrame({k: v.cpu().numpy() for k, v in u.items()}, index=df.index)

    @torch.no_grad()
    def pmf(
        self, df: pd.DataFrame, node: str, *, do: dict[str, float] | None = None
    ) -> np.ndarray:
        """Give the class probabilities of an ordinal node in closed form.

        Parameters
        ----------
        df : pd.DataFrame
            Rows that supply the parent values of the node.
        node : str
            Name of the ordinal node.
        do : dict[str, float] | None, optional
            Parent values ``{parent: value}``, set in ``df`` before the node
            reads its parents. A key that is not a parent raises.

        Returns
        -------
        np.ndarray
            The class probabilities, shape ``(n, levels)``.

        Raises
        ------
        ValueError
            If ``node`` is continuous, or a ``do`` key is a node that is not
            a parent.
        KeyError
            If ``node`` or a ``do`` key is not a node.
        """
        if self._node(node).kind != "ordinal":
            raise ValueError(
                f"pmf() requires an ordinal node, {node!r} is continuous; use density()"
            )
        _, theta, shift, _ = self._conditional(df, node, do)
        return ordinal_pmf(theta, shift).cpu().numpy()

    @torch.no_grad()
    def density(
        self,
        df: pd.DataFrame,
        node: str,
        grid,
        *,
        do: dict[str, float] | None = None,
    ) -> np.ndarray:
        r"""Give the analytic conditional density of a continuous node on a grid.

        The continuous counterpart of [`pmf`][]: for every row of ``df``
        the density $p(\text{node} = g \mid \mathrm{pa})$ at each grid value $g$, in
        closed form from the transform, with no sampling. The evaluation takes one
        (row, grid value) pair per entry: the rows repeat and the grid tiles.

        Parameters
        ----------
        df : pd.DataFrame
            Rows that supply the parent values of the node.
        node : str
            Name of the continuous node.
        grid : array-like
            Values of ``node`` at which to evaluate the density, shape ``(m,)``.
        do : dict[str, float] | None, optional
            Parent values ``{parent: value}``, set in ``df`` before the node
            reads its parents. A key that is not a parent raises.

        Returns
        -------
        np.ndarray
            The densities, shape ``(n, m)``.

        Raises
        ------
        ValueError
            If ``node`` is ordinal (use ``pmf``), or a ``do`` key is a node
            that is not a parent.
        KeyError
            If ``node`` or a ``do`` key is not a node.
        """
        if self._node(node).kind != "continuous":
            raise ValueError(
                f"density() requires a continuous node, {node!r} is ordinal; use pmf()"
            )
        nd, theta, shift, n = self._conditional(df, node, do)
        y = torch.tensor(np.array(grid), dtype=self._dtype, device=self.device)
        m = y.numel()
        u0, ladj = nd.ut.forward(theta.repeat_interleave(m, 0), y.repeat(n))
        log_p = StandardLogistic.log_prob(u0 + shift.repeat_interleave(m)) + ladj
        return log_p.exp().view(n, m).cpu().numpy()

    @torch.no_grad()
    def scores(self, df: pd.DataFrame, node: str) -> pd.DataFrame:
        r"""Give the scores $\psi_i = \partial \ell_i / \partial \beta$.

        The method form of [`node_scores`][tramdag.scores.node_scores], which
        documents the arguments and the column naming.
        """
        return scores.node_scores(self, df, node)

    @torch.no_grad()
    def effect_modifier_scan(
        self,
        df: pd.DataFrame,
        node: str,
        *,
        t: str,
        candidates: list[str] | None = None,
        column: str | None = None,
    ) -> pd.DataFrame:
        """Rank candidate effect modifiers with a fluctuation scan.

        The method form of
        [`effect_modifier_scan`][tramdag.scores.effect_modifier_scan], which
        documents the method, the arguments and the result columns.
        """
        return scores.effect_modifier_scan(
            self, df, node, t=t, candidates=candidates, column=column
        )

    def save(self, path: str | Path) -> None:
        """Write the model, its history and its provenance to a checkpoint.

        The file holds the spec and the weights, the nodes' training
        ``history``, and a ``meta`` block with the tramdag version, the save
        time and the device.

        Parameters
        ----------
        path : str | Path
            Target file. Parent directories are created when missing.
        """
        write_checkpoint(
            path,
            {
                "spec": spec_to_dict(self.spec),
                "init": self.init,
                "state_dict": self.state_dict(),
                "history": {n: nd.history for n, nd in self.nodes.items()},
            },
            self.device,
        )

    @classmethod
    def load(cls, path: str | Path, device: str = "cpu") -> CausalFlowDAG:
        """Restore a model from a checkpoint.

        The nodes' ``history`` and ``flow.meta`` are refilled.

        Parameters
        ----------
        path : str | Path
            Checkpoint file written by [`save`][].
        device : str, optional
            Torch device to load onto, by default ``"cpu"``.

        Returns
        -------
        CausalFlowDAG
            The restored model, in eval mode.
        """
        ckpt = torch.load(path, map_location=device, weights_only=False)
        flow = cls(
            spec_from_dict(ckpt["spec"]),
            device=device,
            init=ckpt["init"],
        )
        load_weights(flow, ckpt["state_dict"])
        for name, history in ckpt["history"].items():
            flow.nodes[name].history = history
        flow.meta = ckpt["meta"]
        flow.eval()
        return flow

history property

history

The nodes' histories, one {node: value} dict per epoch per key.

Nodes stop at different epochs. A node that stopped earlier repeats its last entry, so a summed curve stays defined. A node without entries for a key shows NaN there. "val_epoch" is the longest node's, which is exact while every fit validates. Each node's best epoch is in node.history["best_epoch"].

outcome

outcome()

Give the one node that no other node reads.

Returns:

Type	Description
str	The name of the only node without children.

Raises:

Type	Description
ValueError	If the DAG has more than one such node; the message names them.

Source code in src/tramdag/flow.py

def outcome(self) -> str:
    """Give the one node that no other node reads.

    Returns
    -------
    str
        The name of the only node without children.

    Raises
    ------
    ValueError
        If the DAG has more than one such node; the message names them.
    """
    parents = {p for node in self.spec.values() for p in node_parents(node)}
    sinks = sorted(set(self.order) - parents)
    if len(sinks) != 1:
        raise ValueError(f"outcome needs exactly one sink, the DAG has {sinks!r}")
    return sinks[0]

log_prob

log_prob(df, *, nodes=None)

Compute the joint log-likelihood per row.

Parameters:

Name	Type	Description	Default
df	DataFrame	Observations, one column per node.	required
nodes	list[str] | None	Sum only these nodes' contributions. A subset is exact, because the per-node losses are independent: nodes=["Y"] is the conditional log-likelihood of Y given its parents, per row, in log space (safer than the log of pmf, which underflows in the tail). None (default) is the joint.	None

Returns:

Type	Description
Tensor	log p(x) per row, shape (n,).

Raises:

Type	Description
ValueError	If nodes is empty, or an ordinal value is not a level index.
KeyError	If nodes names an unknown node, or df lacks a column.

Source code in src/tramdag/flow.py

@torch.no_grad()
def log_prob(self, df: pd.DataFrame, *, nodes: list[str] | None = None) -> Tensor:
    """Compute the joint log-likelihood per row.

    Parameters
    ----------
    df : pd.DataFrame
        Observations, one column per node.
    nodes : list[str] | None, optional
        Sum only these nodes' contributions. A subset is exact, because
        the per-node losses are independent: ``nodes=["Y"]`` is the
        conditional log-likelihood of ``Y`` given its parents, per row,
        in log space (safer than the log of [`pmf`][], which
        underflows in the tail). ``None`` (default) is the joint.

    Returns
    -------
    Tensor
        ``log p(x)`` per row, shape ``(n,)``.

    Raises
    ------
    ValueError
        If ``nodes`` is empty, or an ordinal value is not a level index.
    KeyError
        If ``nodes`` names an unknown node, or ``df`` lacks a column.
    """
    if nodes is not None and not nodes:
        raise ValueError("nodes=[] sums nothing; omit it for the joint")
    for name in nodes or ():
        self._node(name)
    per_node = self._node_log_prob(self._tensorize(df), nodes)
    return torch.stack(list(per_node.values()), dim=0).sum(dim=0)

nll

nll(df)

Compute the mean negative log-likelihood per node (a diagnostic).

Parameters:

Name	Type	Description	Default
df	DataFrame	Observations, one column per node.	required

Returns:

Type	Description
dict[str, float]	The mean NLL, keyed by node name.

Source code in src/tramdag/flow.py

@torch.no_grad()
def nll(self, df: pd.DataFrame) -> dict[str, float]:
    """Compute the mean negative log-likelihood per node (a diagnostic).

    Parameters
    ----------
    df : pd.DataFrame
        Observations, one column per node.

    Returns
    -------
    dict[str, float]
        The mean NLL, keyed by node name.
    """
    per_node = self._node_log_prob(self._tensorize(df))
    return {k: float(-v.mean()) for k, v in per_node.items()}

sample

sample(n=None, *, do=None, u=None, seed=None)

Sample from the flow, with optional interventions.

Parameters:

Name	Type	Description	Default
n	int | None	Number of samples. Ignored if u is given.	None
do	dict[str, float] | None	Interventions, as {node: value}. An intervened node is clamped and its parent dependence removed (graph mutilation).	None
u	DataFrame | None	Latent variables, as returned by abduct. If given, they are pushed through the flow. With do, this gives a counterfactual.	None
seed	int | None	If given, seeds the latent draw. Ignored if u is given.	None

Returns:

Type	Description
DataFrame	The samples, one column per node.

Raises:

Type	Description
ValueError	If both n and u are omitted, or a do value of an ordinal node is not a level index.

Notes

Under do, a centered VC re-derives t_do - e_hat(x), so the features of a node are never cached. With u given, the result takes the index of u, so a counterfactual lines up with its factual row.

Source code in src/tramdag/flow.py

@torch.no_grad()
def sample(
    self,
    n: int | None = None,
    *,
    do: dict[str, float] | None = None,
    u: pd.DataFrame | None = None,
    seed: int | None = None,
) -> pd.DataFrame:
    """Sample from the flow, with optional interventions.

    Parameters
    ----------
    n : int | None, optional
        Number of samples. Ignored if ``u`` is given.
    do : dict[str, float] | None, optional
        Interventions, as ``{node: value}``. An intervened node is
        clamped and its parent dependence removed (graph mutilation).
    u : pd.DataFrame | None, optional
        Latent variables, as returned by [`abduct`][]. If given,
        they are pushed through the flow. With ``do``, this gives a
        counterfactual.
    seed : int | None, optional
        If given, seeds the latent draw. Ignored if ``u`` is given.

    Returns
    -------
    pd.DataFrame
        The samples, one column per node.

    Raises
    ------
    ValueError
        If both ``n`` and ``u`` are omitted, or a ``do`` value of an
        ordinal node is not a level index.

    Notes
    -----
    Under ``do``, a centered VC re-derives ``t_do - e_hat(x)``, so the
    features of a node are never cached. With ``u`` given, the result takes
    the index of ``u``, so a counterfactual lines up with its factual row.
    """
    do = do or {}
    self._check_do(do)
    if u is not None:
        n = len(u)
        u_vals = self._tensorize(u, levels=False)
    elif n is not None:
        gen = self._generator(seed)
        u_vals = {
            name: StandardLogistic.sample((n,), device=self.device, generator=gen)
            for name in self.order
        }
    else:
        raise ValueError("sample() needs n or u")

    values: dict[str, Tensor] = {}
    for name in self.order:
        if name in do:
            values[name] = torch.full(
                (n,), float(do[name]), dtype=self._dtype, device=self.device
            )
            continue
        nd = self.nodes[name]
        feats = nd.features(values)
        theta, shift = self._theta_shift(nd, feats, values, n)
        values[name] = nd.sample(theta, shift, u_vals[name])
    return self._to_frame(values, index=None if u is None else u.index)

abduct

abduct(df, *, seed=None)

Recover the latent variables u from observations (the abduction step).

A continuous node inverts exactly: \(u = h(x) + s\). This is \(\operatorname{logit} P(X \le x \mid \text{parents})\), an interventional quantity at the given parent values. For an ordinal node the latent is only interval-identified, so it is sampled from the standard logistic truncated to the observed level's interval.

Parameters:

Name	Type	Description	Default
df	DataFrame	Observations, one column per node.	required
seed	int | None	If given, seeds the truncated draw for the ordinal nodes.	None

Returns:

Type	Description
DataFrame	The latents, one column per node, aligned with the rows of df.

Source code in src/tramdag/flow.py

@torch.no_grad()
def abduct(self, df: pd.DataFrame, *, seed: int | None = None) -> pd.DataFrame:
    r"""Recover the latent variables ``u`` from observations (the abduction step).

    A continuous node inverts exactly: $u = h(x) + s$. This is
    $\operatorname{logit} P(X \le x \mid \text{parents})$, an
    interventional quantity at the given parent values. For an
    ordinal node the latent is only interval-identified, so it is
    sampled from the standard logistic truncated to the observed
    level's interval.

    Parameters
    ----------
    df : pd.DataFrame
        Observations, one column per node.
    seed : int | None, optional
        If given, seeds the truncated draw for the ordinal nodes.

    Returns
    -------
    pd.DataFrame
        The latents, one column per node, aligned with the rows of
        ``df``.
    """
    gen = self._generator(seed)
    values = self._tensorize(df)
    feats = self._features(values)
    n = len(df)
    u = {}
    for name in self.order:
        nd = self.nodes[name]
        theta, shift = self._theta_shift(nd, feats, values, n)
        u[name] = nd.abduct(theta, shift, values[name], generator=gen)
    return pd.DataFrame({k: v.cpu().numpy() for k, v in u.items()}, index=df.index)

pmf

pmf(df, node, *, do=None)

Give the class probabilities of an ordinal node in closed form.

Parameters:

Name	Type	Description	Default
df	DataFrame	Rows that supply the parent values of the node.	required
node	str	Name of the ordinal node.	required
do	dict[str, float] | None	Parent values {parent: value}, set in df before the node reads its parents. A key that is not a parent raises.	None

Returns:

Type	Description
ndarray	The class probabilities, shape (n, levels).

Raises:

Type	Description
ValueError	If node is continuous, or a do key is a node that is not a parent.
KeyError	If node or a do key is not a node.

Source code in src/tramdag/flow.py

@torch.no_grad()
def pmf(
    self, df: pd.DataFrame, node: str, *, do: dict[str, float] | None = None
) -> np.ndarray:
    """Give the class probabilities of an ordinal node in closed form.

    Parameters
    ----------
    df : pd.DataFrame
        Rows that supply the parent values of the node.
    node : str
        Name of the ordinal node.
    do : dict[str, float] | None, optional
        Parent values ``{parent: value}``, set in ``df`` before the node
        reads its parents. A key that is not a parent raises.

    Returns
    -------
    np.ndarray
        The class probabilities, shape ``(n, levels)``.

    Raises
    ------
    ValueError
        If ``node`` is continuous, or a ``do`` key is a node that is not
        a parent.
    KeyError
        If ``node`` or a ``do`` key is not a node.
    """
    if self._node(node).kind != "ordinal":
        raise ValueError(
            f"pmf() requires an ordinal node, {node!r} is continuous; use density()"
        )
    _, theta, shift, _ = self._conditional(df, node, do)
    return ordinal_pmf(theta, shift).cpu().numpy()

density

density(df, node, grid, *, do=None)

Give the analytic conditional density of a continuous node on a grid.

The continuous counterpart of pmf: for every row of df the density \(p(\text{node} = g \mid \mathrm{pa})\) at each grid value \(g\), in closed form from the transform, with no sampling. The evaluation takes one (row, grid value) pair per entry: the rows repeat and the grid tiles.

Parameters:

Name	Type	Description	Default
df	DataFrame	Rows that supply the parent values of the node.	required
node	str	Name of the continuous node.	required
grid	array - like	Values of node at which to evaluate the density, shape (m,).	required
do	dict[str, float] | None	Parent values {parent: value}, set in df before the node reads its parents. A key that is not a parent raises.	None

Returns:

Type	Description
ndarray	The densities, shape (n, m).

Raises:

Type	Description
ValueError	If node is ordinal (use pmf), or a do key is a node that is not a parent.
KeyError	If node or a do key is not a node.

Source code in src/tramdag/flow.py

@torch.no_grad()
def density(
    self,
    df: pd.DataFrame,
    node: str,
    grid,
    *,
    do: dict[str, float] | None = None,
) -> np.ndarray:
    r"""Give the analytic conditional density of a continuous node on a grid.

    The continuous counterpart of [`pmf`][]: for every row of ``df``
    the density $p(\text{node} = g \mid \mathrm{pa})$ at each grid value $g$, in
    closed form from the transform, with no sampling. The evaluation takes one
    (row, grid value) pair per entry: the rows repeat and the grid tiles.

    Parameters
    ----------
    df : pd.DataFrame
        Rows that supply the parent values of the node.
    node : str
        Name of the continuous node.
    grid : array-like
        Values of ``node`` at which to evaluate the density, shape ``(m,)``.
    do : dict[str, float] | None, optional
        Parent values ``{parent: value}``, set in ``df`` before the node
        reads its parents. A key that is not a parent raises.

    Returns
    -------
    np.ndarray
        The densities, shape ``(n, m)``.

    Raises
    ------
    ValueError
        If ``node`` is ordinal (use ``pmf``), or a ``do`` key is a node
        that is not a parent.
    KeyError
        If ``node`` or a ``do`` key is not a node.
    """
    if self._node(node).kind != "continuous":
        raise ValueError(
            f"density() requires a continuous node, {node!r} is ordinal; use pmf()"
        )
    nd, theta, shift, n = self._conditional(df, node, do)
    y = torch.tensor(np.array(grid), dtype=self._dtype, device=self.device)
    m = y.numel()
    u0, ladj = nd.ut.forward(theta.repeat_interleave(m, 0), y.repeat(n))
    log_p = StandardLogistic.log_prob(u0 + shift.repeat_interleave(m)) + ladj
    return log_p.exp().view(n, m).cpu().numpy()

scores

scores(df, node)

Give the scores \(\psi_i = \partial \ell_i / \partial \beta\).

The method form of node_scores, which documents the arguments and the column naming.

Source code in src/tramdag/flow.py

@torch.no_grad()
def scores(self, df: pd.DataFrame, node: str) -> pd.DataFrame:
    r"""Give the scores $\psi_i = \partial \ell_i / \partial \beta$.

    The method form of [`node_scores`][tramdag.scores.node_scores], which
    documents the arguments and the column naming.
    """
    return scores.node_scores(self, df, node)

effect_modifier_scan

effect_modifier_scan(
    df, node, *, t, candidates=None, column=None
)

Rank candidate effect modifiers with a fluctuation scan.

The method form of effect_modifier_scan, which documents the method, the arguments and the result columns.

Source code in src/tramdag/flow.py

@torch.no_grad()
def effect_modifier_scan(
    self,
    df: pd.DataFrame,
    node: str,
    *,
    t: str,
    candidates: list[str] | None = None,
    column: str | None = None,
) -> pd.DataFrame:
    """Rank candidate effect modifiers with a fluctuation scan.

    The method form of
    [`effect_modifier_scan`][tramdag.scores.effect_modifier_scan], which
    documents the method, the arguments and the result columns.
    """
    return scores.effect_modifier_scan(
        self, df, node, t=t, candidates=candidates, column=column
    )

save

save(path)

Write the model, its history and its provenance to a checkpoint.

The file holds the spec and the weights, the nodes' training history, and a meta block with the tramdag version, the save time and the device.

Parameters:

Name	Type	Description	Default
path	str | Path	Target file. Parent directories are created when missing.	required

Source code in src/tramdag/flow.py

def save(self, path: str | Path) -> None:
    """Write the model, its history and its provenance to a checkpoint.

    The file holds the spec and the weights, the nodes' training
    ``history``, and a ``meta`` block with the tramdag version, the save
    time and the device.

    Parameters
    ----------
    path : str | Path
        Target file. Parent directories are created when missing.
    """
    write_checkpoint(
        path,
        {
            "spec": spec_to_dict(self.spec),
            "init": self.init,
            "state_dict": self.state_dict(),
            "history": {n: nd.history for n, nd in self.nodes.items()},
        },
        self.device,
    )

load classmethod

load(path, device='cpu')

Restore a model from a checkpoint.

The nodes' history and flow.meta are refilled.

Parameters:

Name	Type	Description	Default
path	str | Path	Checkpoint file written by save.	required
device	str	Torch device to load onto, by default "cpu".	'cpu'

Returns:

Type	Description
CausalFlowDAG	The restored model, in eval mode.

Source code in src/tramdag/flow.py

@classmethod
def load(cls, path: str | Path, device: str = "cpu") -> CausalFlowDAG:
    """Restore a model from a checkpoint.

    The nodes' ``history`` and ``flow.meta`` are refilled.

    Parameters
    ----------
    path : str | Path
        Checkpoint file written by [`save`][].
    device : str, optional
        Torch device to load onto, by default ``"cpu"``.

    Returns
    -------
    CausalFlowDAG
        The restored model, in eval mode.
    """
    ckpt = torch.load(path, map_location=device, weights_only=False)
    flow = cls(
        spec_from_dict(ckpt["spec"]),
        device=device,
        init=ckpt["init"],
    )
    load_weights(flow, ckpt["state_dict"])
    for name, history in ckpt["history"].items():
        flow.nodes[name].history = history
    flow.meta = ckpt["meta"]
    flow.eval()
    return flow

shift_curve

shift_curve(node, parent, grid)

Evaluate one fitted shift term on a 1-D grid of parent values.

The parent runs over grid with the term's other inputs absent, so the term must read exactly this one parent (an LS or one-parent CS). The evaluation runs through the term's own shift_value, so custom terms work too. The grid is copied, as in _tensorize: torch warns on a non-writable numpy array, which a broadcast or read-only grid is.

Parameters:

Name	Type	Description	Default
node	str	Name of the node that carries the term.	required
parent	str	The term's key in the node's shifts: the parent's name.	required
grid	array - like	Parent values at which to evaluate the term, shape (m,).	required

Returns:

Type	Description
ndarray	The shift values, shape (m,): the curve a replication plots against the data-generating truth.

Raises:

Type	Description
KeyError	If node is unknown, or the node has no shift term keyed parent.
ValueError	If parent is an ordinal node.

Source code in src/tramdag/readouts.py

@torch.no_grad()
def shift_curve(self, node: str, parent: str, grid) -> np.ndarray:
    """Evaluate one fitted shift term on a 1-D grid of parent values.

    The parent runs over ``grid`` with the term's other inputs absent, so the
    term must read exactly this one parent (an ``LS`` or one-parent ``CS``).
    The evaluation runs through the term's own ``shift_value``, so custom
    terms work too. The grid is copied, as in ``_tensorize``: torch warns on
    a non-writable numpy array, which a broadcast or read-only grid is.

    Parameters
    ----------
    node : str
        Name of the node that carries the term.
    parent : str
        The term's key in the node's shifts: the parent's name.
    grid : array-like
        Parent values at which to evaluate the term, shape ``(m,)``.

    Returns
    -------
    np.ndarray
        The shift values, shape ``(m,)``: the curve a replication plots
        against the data-generating truth.

    Raises
    ------
    KeyError
        If ``node`` is unknown, or the node has no shift term keyed
        ``parent``.
    ValueError
        If ``parent`` is an ordinal node.
    """
    nd = self._node(node)
    if parent not in nd.shifts:
        raise KeyError(
            f"node {node!r} has no shift term keyed {parent!r}; "
            f"available: {sorted(nd.shifts)}"
        )
    if self.spec[parent].kind == "ordinal":
        raise ValueError(
            f"{parent!r} is an ordinal parent, so it enters the term "
            "one-hot over all its levels and a 1-D grid of level indices "
            "is not its input. Read its effect with ls_coefficients() "
            "instead, whose weights are the level contrasts."
        )
    x = torch.as_tensor(np.array(grid), dtype=self._dtype, device=self.device)
    x = x.view(-1, 1)
    curve = nd.shifts[parent].shift_value(nd, {parent: x})
    return curve.cpu().numpy().ravel()

varying_coef

varying_coef(df, node, *, t=None)

Evaluate the fitted effect function \(\beta(x)\) of a VC term.

The value is \(\beta_0 + b_\Theta(\text{modifiers})\), in closed form from the fitted term; only the modifier columns of df are read. It lives on the node's latent scale.

Parameters:

Name	Type	Description	Default
df	DataFrame	Rows at which to evaluate \(\beta\). Must contain every modifier column of the term.	required
node	str	Name of the node that carries the VC term.	required
t	str | None	Treatment name of the VC term. Optional when the node has exactly one VC term.	None

Returns:

Type	Description
ndarray	The \(\beta\) values, shape (n,). Constant when the term has no modifiers.

Raises:

Type	Description
KeyError	If node is unknown, if the node has no VC term on t, or if a modifier column is missing from df.
ValueError	If the node has no VC term, or if t is omitted while the node has several VC terms.

Source code in src/tramdag/readouts.py

@torch.no_grad()
def varying_coef(
    self, df: pd.DataFrame, node: str, *, t: str | None = None
) -> np.ndarray:
    r"""Evaluate the fitted effect function $\beta(x)$ of a ``VC`` term.

    The value is $\beta_0 + b_\Theta(\text{modifiers})$, in closed form from the
    fitted term; only the modifier columns of ``df`` are read. It lives on
    the node's latent scale.

    Parameters
    ----------
    df : pd.DataFrame
        Rows at which to evaluate $\beta$. Must contain every modifier
        column of the term.
    node : str
        Name of the node that carries the VC term.
    t : str | None, optional
        Treatment name of the VC term. Optional when the node has
        exactly one VC term.

    Returns
    -------
    np.ndarray
        The $\beta$ values, shape ``(n,)``. Constant when the term has
        no modifiers.

    Raises
    ------
    KeyError
        If ``node`` is unknown, if the node has no VC term on ``t``, or
        if a modifier column is missing from ``df``.
    ValueError
        If the node has no VC term, or if ``t`` is omitted while the
        node has several VC terms.
    """
    nd, t, mods = self._vc_term(node, t)
    mod_feat = None
    if mods:
        feats = self._features(self._tensorize(df, mods))
        mod_feat = nd.net_input(feats, mods, t)
    return nd.shifts[t].beta(mod_feat, len(df)).cpu().numpy()

ls_coefficients

ls_coefficients()

Give the per-node linear-shift weights.

A continuous parent has one weight. An ordinal parent has one weight per level (its one-hot encoding), identified only up to a common constant, so read them as differences w[k] - w[0]. Only LS terms have a weight; CS and VC shifts are skipped.

Returns:

Type	Description
dict[str, dict[str, ndarray]]	The weights, as {node: {parent: array}}. A node without linear-shift terms is absent.

Source code in src/tramdag/readouts.py

@torch.no_grad()
def ls_coefficients(self) -> dict[str, dict[str, np.ndarray]]:
    """Give the per-node linear-shift weights.

    A continuous parent has one weight. An ordinal parent has one weight
    per level (its one-hot encoding), identified only up to a common
    constant, so read them as differences ``w[k] - w[0]``. Only ``LS``
    terms have a weight; ``CS`` and ``VC`` shifts are skipped.

    Returns
    -------
    dict[str, dict[str, np.ndarray]]
        The weights, as ``{node: {parent: array}}``. A node without
        linear-shift terms is absent.
    """
    weights = {name: self.nodes[name].ls_coefficients() for name in self.order}
    return {name: w for name, w in weights.items() if w}

to_matrix

to_matrix()

Give the labeled adjacency matrix of term effects.

Returns:

Type	Description
DataFrame	Rows are parents and columns are children. A cell holds the term tag: "LS", "CS", "CI", "VC" for a VC treatment, or "VCm" for a VC modifier. An empty cell means there is no edge. A multi-parent term carries its parent group as a suffix. When several terms share a cell, their tags join with "+".

Source code in src/tramdag/readouts.py

def to_matrix(self) -> pd.DataFrame:
    """Give the labeled adjacency matrix of term effects.

    Returns
    -------
    pd.DataFrame
        Rows are parents and columns are children. A cell holds the
        term tag: ``"LS"``, ``"CS"``, ``"CI"``, ``"VC"`` for a VC
        treatment, or ``"VCm"`` for a VC modifier. An empty cell means
        there is no edge. A multi-parent term carries its parent group
        as a suffix. When several terms share a cell, their tags join
        with ``"+"``.
    """
    m = pd.DataFrame("", index=list(self.order), columns=list(self.order))
    for child in self.order:
        for term in self.spec[child].terms:
            for p, tag, joint in term.cells():
                if joint:
                    tag = f"{tag}{list(term.parents)}"
                cur = m.loc[p, child]
                m.loc[p, child] = f"{cur}+{tag}" if cur else tag
    return m

intercept_contributions

intercept_contributions(df, node)

Decompose a complex intercept into mean-centered per-parent parts.

Each network's output over the rows of df is centered to mean zero per parameter; the removed means collect into baseline, so \(\vartheta(\mathrm{pa}) = \text{baseline} + \sum \text{contributions}\) exactly. Post-hoc only: it reads the fitted weights and changes nothing.

Parameters:

Name	Type	Description	Default
df	DataFrame	Rows over which to center and at which to evaluate the contributions. Must contain every intercept-parent column.	required
node	str	Name of a node whose intercept term has parents.	required

Returns:

Type

Description

dict

Three keys. "baseline" is the absorbed constant, a (P,) array: the sum of the per-term means. P is the node's transform-parameter count: ut.n_params for a continuous node, levels - 1 cutpoint parameters for an ordinal node. "contributions" is {term_label: (n, P) array}: each term's mean-centered contribution at each row, columns summing to about zero over the rows. term_label is the term's parents joined by "+". "parents" is {term_label: tuple(parent_names)}.

Raises:

Type	Description
KeyError	If node is unknown, or if an intercept-parent column is missing from df.
ValueError	If the node's intercept term has no parents.

Notes

An additive intercept holds one network per parent group in its nets; a single intercept term, joint or not, is itself the one network. The contributions live in the transform's unconstrained parameter space, before the monotonicity constraint.

Source code in src/tramdag/readouts.py

@torch.no_grad()
def intercept_contributions(self, df: pd.DataFrame, node: str) -> dict:
    r"""Decompose a complex intercept into mean-centered per-parent parts.

    Each network's output over the rows of ``df`` is centered to mean zero
    per parameter; the removed means collect into ``baseline``, so
    $\vartheta(\mathrm{pa}) = \text{baseline} + \sum \text{contributions}$ exactly.
    Post-hoc only: it reads the fitted weights and changes nothing.

    Parameters
    ----------
    df : pd.DataFrame
        Rows over which to center and at which to evaluate the
        contributions. Must contain every intercept-parent column.
    node : str
        Name of a node whose intercept term has parents.

    Returns
    -------
    dict
        Three keys. ``"baseline"`` is the absorbed constant, a ``(P,)``
        array: the sum of the per-term means. ``P`` is the node's
        transform-parameter count: ``ut.n_params`` for a continuous
        node, ``levels - 1`` cutpoint parameters for an ordinal node.
        ``"contributions"`` is ``{term_label: (n, P) array}``: each
        term's mean-centered contribution at each row, columns summing
        to about zero over the rows. ``term_label`` is the term's
        parents joined by ``"+"``. ``"parents"`` is
        ``{term_label: tuple(parent_names)}``.

    Raises
    ------
    KeyError
        If ``node`` is unknown, or if an intercept-parent column is
        missing from ``df``.
    ValueError
        If the node's intercept term has no parents.

    Notes
    -----
    An additive intercept holds one network per parent group in its
    ``nets``; a single intercept term, joint or not, is itself the one
    network. The contributions live in the transform's unconstrained
    parameter space, before the monotonicity constraint.
    """
    nd = self._node(node)
    groups = nd.intercept.groups
    if not groups:
        raise ValueError(
            f"node {node!r} has no complex-intercept (I) terms with parents. "
            "Its intercept is unconditional, so there is nothing to decompose."
        )
    feats = self._features(self._tensorize(df, nd.intercept.ci_parents))
    nets = getattr(nd.intercept, "nets", [nd.intercept])

    contributions: dict[str, np.ndarray] = {}
    parents: dict[str, tuple] = {}
    baseline = None
    for net, grp in zip(nets, groups, strict=True):
        raw = net(nd.net_input(feats, grp, "@I"))
        mean = raw.mean(dim=0, keepdim=True)
        label = "+".join(grp)
        contributions[label] = (raw - mean).cpu().numpy()
        parents[label] = grp
        baseline = mean if baseline is None else baseline + mean
    return {
        "baseline": baseline.cpu().numpy().ravel(),
        "contributions": contributions,
        "parents": parents,
    }

design_matrix

design_matrix(df, node, *, drop_first=False)

Encode a node's parents the way the flow feeds them to its shifts.

A continuous parent stays raw in one column named after it. An ordinal parent becomes one column per level, named "{parent}[{k}]", the same one-hot the flow builds internally.

Use drop_first=True to get the design a classical reference expects (statsmodels OrderedModel, R polr). It drops the level-0 column of each ordinal parent, so the remaining coefficients match w[k] - w[0] on the flow side.

Parameters:

Name	Type	Description	Default
df	DataFrame	Rows to encode. Must contain every parent column of node.	required
node	str	Name of the node whose parents are encoded.	required
drop_first	bool	Drop each ordinal parent's level-0 column, by default False.	False

Returns:

Type	Description
DataFrame	One column per encoded feature, indexed like df.

Raises:

Type	Description
KeyError	If node is unknown, or df lacks a parent column.
ValueError	If an ordinal parent value is not a level index.

Source code in src/tramdag/readouts.py

@torch.no_grad()
def design_matrix(
    self, df: pd.DataFrame, node: str, *, drop_first: bool = False
) -> pd.DataFrame:
    """Encode a node's parents the way the flow feeds them to its shifts.

    A continuous parent stays raw in one column named after it. An
    ordinal parent becomes one column per level, named
    ``"{parent}[{k}]"``, the same one-hot the flow builds internally.

    Use ``drop_first=True`` to get the design a classical reference
    expects (``statsmodels`` ``OrderedModel``, R ``polr``). It drops the
    level-0 column of each ordinal parent, so the remaining coefficients
    match ``w[k] - w[0]`` on the flow side.

    Parameters
    ----------
    df : pd.DataFrame
        Rows to encode. Must contain every parent column of ``node``.
    node : str
        Name of the node whose parents are encoded.
    drop_first : bool, optional
        Drop each ordinal parent's level-0 column, by default ``False``.

    Returns
    -------
    pd.DataFrame
        One column per encoded feature, indexed like ``df``.

    Raises
    ------
    KeyError
        If ``node`` is unknown, or ``df`` lacks a parent column.
    ValueError
        If an ordinal parent value is not a level index.
    """
    nd = self._node(node)
    feats = nd.features(self._tensorize(df, nd.parents))
    cols: dict[str, np.ndarray] = {}
    for p in nd.parents:
        arr = feats[p].cpu().numpy()
        if self.spec[p].kind == "continuous":
            cols[p] = arr[:, 0]
        else:
            for k in range(1 if drop_first else 0, arr.shape[1]):
                cols[f"{p}[{k}]"] = arr[:, k]
    return pd.DataFrame(cols, index=df.index)

warm_start

warm_start(train_df)

Warm-start every node that has a VC term.

Each such node runs Node.warm_start. Call it before fit.

Parameters:

Name	Type	Description	Default
train_df	DataFrame	Training rows with every node's column.	required

Returns:

Type	Description
dict[str, dict]	The reports of the nodes with a VC term, keyed by node name. The dict is empty when no node has one.

Source code in src/tramdag/fitting.py

def warm_start(self, train_df: pd.DataFrame) -> dict[str, dict]:
    """Warm-start every node that has a ``VC`` term.

    Each such node runs [`Node.warm_start`][tramdag.nodes.Node.warm_start].
    Call it before ``fit``.

    Parameters
    ----------
    train_df : pd.DataFrame
        Training rows with every node's column.

    Returns
    -------
    dict[str, dict]
        The reports of the nodes with a ``VC`` term, keyed by node name.
        The dict is empty when no node has one.
    """
    return {
        name: self.nodes[name].warm_start(train_df)
        for name, node_spec in self.spec.items()
        if any(isinstance(t, VaryingCoefficient) for t in node_spec.terms)
    }

fit

fit(
    train_df,
    *,
    epochs,
    learning_rate=None,
    batch_size=512,
    validation_data=None,
    validation_split=None,
    verbose=0,
    seed=None,
    marginal_init=False,
    optimizer=None,
    callbacks=None,
    n_jobs=1
)

Fit every node by maximum likelihood, node by node.

The joint likelihood is a sum of per-node terms with disjoint parameters, so fitting the nodes one by one is exact. Each node runs Node.fit, whose docstring describes the loop, the history and the callback contract. The validation split happens once here, so all nodes see the same rows. The nodes fit in topological order; with validation data, the nodes with a centered VC fit last, so their validation NLL uses the fitted treatment node.

epochs, learning_rate, optimizer and callbacks each take one value for every node or a dict keyed by node name. The node gets value[name]. A missing key gives the argument's default (None), so a collections.defaultdict gives its own fallback. epochs has no default, so a missing key raises. A key that names no node raises too.

Parameters:

Name	Type	Description	Default
train_df	DataFrame	Training data, one column per node, plus the side columns.	required
epochs	int | dict	Passes over the data, one value or one per node. Required; there is no default.	required
learning_rate	float | dict	Adam step size of the default optimizer, one value or one per node. The default (None) is 1e-2. A node cannot get both a learning_rate and an optimizer.	None
batch_size	int	Rows per gradient step, by default 512.	512
validation_data	DataFrame | None	Validation rows, with the same columns as train_df, by default None.	None
validation_split	float | None	Keras' rule: the LAST fraction of train_df becomes the validation set, without shuffling, and only the remaining rows train and calibrate. Shuffle the frame first if its row order means anything. Mutually exclusive with validation_data.	None
verbose	int	0 (default) is silent; N >= 1 prints every N epochs.	0
seed	int | None	Seeds the minibatch shuffling, by default None. Every node fits with this seed, so a node shuffles as Node.fit with the same seed. The result does not depend on the other nodes or on n_jobs. Unseeded with n_jobs > 1, one seed is drawn from torch's global RNG; see the Notes.	None
marginal_init	bool	Start every simple intercept at its column's marginal, by default False. Applied only by the fit that calibrates a node. See Node.fit.	False
optimizer	callable | dict | None	A factory f(node) -> Optimizer, called once per node fit, for example lambda node: torch.optim.AdamW(node.parameters()), or a dict of factories or optimizer instances by node name. The default (None) is Adam(lr=1e-2). One optimizer instance cannot serve every node, so the dict holds the instances.	None
callbacks	Callback | list | callable | dict | None	The callbacks of every node, or a dict of them by node name, by default None. A bare callable is an on_epoch_end hook, as in Node.fit. A shared shipped callback resets at every node's fit begin, so its attributes describe the last node only; a dict keeps one instance per node.	None
n_jobs	int	Number of processes, by default 1 (serial, topological order). n_jobs > 1 forks n_jobs workers, one task per node (Linux only). A worker runs torch on one thread. With the same seed, its result equals a serial fit on one thread; with more threads a network can differ in the last bits. Callback state stays in the workers; read node.history instead.	1

Returns:

Type	Description
CausalFlowDAG	self, fitted, in eval mode.

Notes

Without a seed, n_jobs > 1 draws one from torch's RNG, because forked workers would otherwise replay one RNG state. Every worker gets this one seed. A centered VC validates against the live propensity of its fitted treatment node, as every query does. Those nodes therefore fit second. The callbacks normalize once here, so an iterator of callbacks survives a single pass.

Raises:

Type	Description
TypeError	If optimizer is one optimizer instance instead of a factory or a dict.
KeyError	If a frame lacks a node column, a dict epochs lacks a node, or a dict key names no node.
ValueError	See Node.fit. Every frame and every node's callbacks are checked before the first node fits. Also if a node gets both a learning_rate and an optimizer.

Source code in src/tramdag/fitting.py

def fit(
    self,
    train_df: pd.DataFrame,
    *,
    epochs: int | dict[str, int],
    learning_rate: float | dict[str, float] | None = None,
    batch_size: int = 512,
    validation_data: pd.DataFrame | None = None,
    validation_split: float | None = None,
    verbose: int = 0,
    seed: int | None = None,
    marginal_init: bool = False,
    optimizer: Callable[[Node], torch.optim.Optimizer] | dict | None = None,
    callbacks=None,
    n_jobs: int = 1,
) -> CausalFlowDAG:
    """Fit every node by maximum likelihood, node by node.

    The joint likelihood is a sum of per-node terms with disjoint
    parameters, so fitting the nodes one by one is exact. Each node runs
    [`Node.fit`][tramdag.nodes.Node.fit], whose docstring describes the
    loop, the history and the callback contract. The validation split
    happens once here, so all nodes see the same rows. The nodes fit in
    topological order; with validation data, the nodes with a centered
    ``VC`` fit last, so their validation NLL uses the fitted treatment
    node.

    ``epochs``, ``learning_rate``, ``optimizer`` and ``callbacks`` each
    take one value for every node or a dict keyed by node name. The node
    gets ``value[name]``. A missing key gives the argument's default
    (None), so a ``collections.defaultdict`` gives its own fallback.
    ``epochs`` has no default, so a missing key raises. A key that names
    no node raises too.

    Parameters
    ----------
    train_df : pd.DataFrame
        Training data, one column per node, plus the side columns.
    epochs : int | dict
        Passes over the data, one value or one per node. Required; there
        is no default.
    learning_rate : float | dict, optional
        Adam step size of the default optimizer, one value or one per
        node. The default (None) is 1e-2. A node cannot get both a
        ``learning_rate`` and an ``optimizer``.
    batch_size : int, optional
        Rows per gradient step, by default 512.
    validation_data : pd.DataFrame | None, optional
        Validation rows, with the same columns as ``train_df``, by
        default None.
    validation_split : float | None, optional
        Keras' rule: the LAST fraction of ``train_df`` becomes the
        validation set, without shuffling, and only the remaining rows
        train and calibrate. Shuffle the frame first if its row order
        means anything. Mutually exclusive with ``validation_data``.
    verbose : int, optional
        0 (default) is silent; ``N >= 1`` prints every ``N`` epochs.
    seed : int | None, optional
        Seeds the minibatch shuffling, by default None. Every node fits
        with this ``seed``, so a node shuffles as
        [`Node.fit`][tramdag.nodes.Node.fit] with the same ``seed``. The
        result does not depend on the other nodes or on ``n_jobs``.
        Unseeded with ``n_jobs > 1``, one seed is drawn from torch's
        global RNG; see the Notes.
    marginal_init : bool, optional
        Start every simple intercept at its column's marginal, by default
        ``False``. Applied only by the fit that calibrates a node. See
        [`Node.fit`][tramdag.nodes.Node.fit].
    optimizer : callable | dict | None, optional
        A factory ``f(node) -> Optimizer``, called once per node fit, for
        example ``lambda node: torch.optim.AdamW(node.parameters())``, or
        a dict of factories or optimizer instances by node name. The
        default (None) is ``Adam(lr=1e-2)``. One optimizer instance
        cannot serve every node, so the dict holds the instances.
    callbacks : Callback | list | callable | dict | None, optional
        The callbacks of every node, or a dict of them by node name, by
        default None. A bare callable is an ``on_epoch_end`` hook, as in
        [`Node.fit`][tramdag.nodes.Node.fit]. A shared shipped callback
        resets at every node's fit begin, so its attributes describe the
        last node only; a dict keeps one instance per node.
    n_jobs : int, optional
        Number of processes, by default 1 (serial, topological order).
        ``n_jobs > 1`` forks ``n_jobs`` workers, one task per node (Linux
        only). A worker runs torch on one thread. With the same ``seed``,
        its result equals a serial fit on one thread; with more threads a
        network can differ in the last bits. Callback state stays in the workers;
        read ``node.history`` instead.

    Returns
    -------
    CausalFlowDAG
        ``self``, fitted, in eval mode.

    Notes
    -----
    Without a ``seed``, ``n_jobs > 1`` draws one from torch's RNG, because
    forked workers would otherwise replay one RNG state. Every worker
    gets this one seed. A centered VC
    validates against the live propensity of its fitted treatment node, as
    every query does. Those nodes therefore fit second. The callbacks
    normalize once here, so an iterator of callbacks survives a single pass.

    Raises
    ------
    TypeError
        If ``optimizer`` is one optimizer instance instead of a factory
        or a dict.
    KeyError
        If a frame lacks a node column, a dict ``epochs`` lacks a node, or
        a dict key names no node.
    ValueError
        See [`Node.fit`][tramdag.nodes.Node.fit]. Every frame and every
        node's callbacks are checked before the first node fits. Also if
        a node gets both a ``learning_rate`` and an ``optimizer``.
    """
    if isinstance(optimizer, torch.optim.Optimizer):
        raise TypeError(
            "optimizer= of CausalFlowDAG.fit is a factory f(node) -> "
            "Optimizer or a dict of optimizers by node name; got one "
            "optimizer instance for every node"
        )
    train_df, validation_data = _split_validation(
        train_df, validation_data, validation_split
    )
    self._check_frames(train_df, validation_data)
    if seed is None and n_jobs > 1:
        seed = int(torch.randint(2**62, ()))
    if not isinstance(callbacks, dict):
        callbacks = _normalize_callbacks(callbacks)
    jobs = _node_jobs(
        self.order,
        dict(
            epochs=epochs,
            learning_rate=learning_rate,
            optimizer=optimizer,
            callbacks=callbacks,
        ),
        dict(
            batch_size=batch_size,
            validation_data=validation_data,
            verbose=verbose,
            seed=seed,
            marginal_init=marginal_init,
        ),
    )
    self._calibrate(train_df, marginal_init=marginal_init)
    later = [n for n in self.order if self.nodes[n].side_columns()]
    if validation_data is None:
        later = []
    self._fit_nodes(train_df, {n: jobs[n] for n in jobs if n not in later}, n_jobs)
    for name in later:
        jobs[name]["validation_data"] = self._with_live_side(name, validation_data)
    self._fit_nodes(train_df, {n: jobs[n] for n in later}, n_jobs)
    return self.eval()

fit_jointly

fit_jointly(
    train_df,
    *,
    epochs,
    learning_rate=None,
    batch_size=512,
    validation_data=None,
    validation_split=None,
    verbose=0,
    seed=None,
    marginal_init=False,
    optimizer=None,
    callbacks=None
)

Fit all nodes together: one optimizer, one minibatch loop.

Each step draws one batch of rows and minimizes the sum of the nodes' mean NLLs. The optimum is that of fit; the training dynamics differ. One learning-rate schedule and one stop rule act on the summed NLL, which is the protocol of the reference implementation. Each node's history gets its own train and validation NLL per epoch, so flow.history, history_frame and plot_training read a joint fit as they read a per-node one.

Parameters:

Name	Type	Description	Default
train_df	DataFrame	Training data, one column per node, plus the side columns.	required
epochs	int	Passes over the data. Required; there is no default.	required
learning_rate	float	Adam step size of the default optimizer. The default (None) is 1e-2. Exclusive with optimizer.	None
batch_size	int	Rows per gradient step, by default 512.	512
validation_data	DataFrame | None	Validation rows, one column per node, by default None. A centered VC validates against the live propensity of its treatment node.	None
validation_split	float | None	Keras' rule, as in fit. Mutually exclusive with validation_data.	None
verbose	int	0 (default) is silent; N >= 1 prints the summed NLLs every N epochs and on the final epoch.	0
seed	int | None	Seeds the minibatch shuffling, by default None. The row order of every epoch is that of fit with the same seed.	None
marginal_init	bool	Start every simple intercept at its column's marginal, by default False. Applied only by the fit that calibrates a node. See Node.fit.	False
optimizer	Optimizer | callable | None	A torch optimizer over flow.parameters(), or a factory f(flow) -> Optimizer. The default (None) is Adam(lr=1e-2).	None
callbacks	Callback | callable | list | None	As in Node.fit, but every hook gets the flow instead of a node. EarlyStopping then tracks the summed validation NLL and restores the whole flow.	None

Returns:

Type	Description
CausalFlowDAG	self, fitted, in eval mode.

Raises:

Type	Description
ValueError	See Node.fit.
KeyError	If a frame lacks a node column.

Source code in src/tramdag/fitting.py

def fit_jointly(
    self,
    train_df: pd.DataFrame,
    *,
    epochs: int,
    learning_rate: float | None = None,
    batch_size: int = 512,
    validation_data: pd.DataFrame | None = None,
    validation_split: float | None = None,
    verbose: int = 0,
    seed: int | None = None,
    marginal_init: bool = False,
    optimizer: torch.optim.Optimizer
    | Callable[[CausalFlowDAG], torch.optim.Optimizer]
    | None = None,
    callbacks=None,
) -> CausalFlowDAG:
    """Fit all nodes together: one optimizer, one minibatch loop.

    Each step draws one batch of rows and minimizes the sum of the nodes'
    mean NLLs. The optimum is that of [`fit`][tramdag.flow.CausalFlowDAG.fit];
    the training dynamics differ. One learning-rate schedule and one stop
    rule act on the summed NLL, which is the protocol of the reference
    implementation. Each node's history gets its own train and validation
    NLL per epoch, so ``flow.history``, ``history_frame`` and
    ``plot_training`` read a joint fit as they read a per-node one.

    Parameters
    ----------
    train_df : pd.DataFrame
        Training data, one column per node, plus the side columns.
    epochs : int
        Passes over the data. Required; there is no default.
    learning_rate : float, optional
        Adam step size of the default optimizer. The default (None) is
        1e-2. Exclusive with ``optimizer``.
    batch_size : int, optional
        Rows per gradient step, by default 512.
    validation_data : pd.DataFrame | None, optional
        Validation rows, one column per node, by default None. A centered
        ``VC`` validates against the live propensity of its treatment
        node.
    validation_split : float | None, optional
        Keras' rule, as in ``fit``. Mutually exclusive with
        ``validation_data``.
    verbose : int, optional
        0 (default) is silent; ``N >= 1`` prints the summed NLLs every
        ``N`` epochs and on the final epoch.
    seed : int | None, optional
        Seeds the minibatch shuffling, by default None. The row order of
        every epoch is that of ``fit`` with the same ``seed``.
    marginal_init : bool, optional
        Start every simple intercept at its column's marginal, by default
        ``False``. Applied only by the fit that calibrates a node. See
        [`Node.fit`][tramdag.nodes.Node.fit].
    optimizer : torch.optim.Optimizer | callable | None, optional
        A torch optimizer over ``flow.parameters()``, or a factory
        ``f(flow) -> Optimizer``. The default (None) is
        ``Adam(lr=1e-2)``.
    callbacks : Callback | callable | list | None, optional
        As in [`Node.fit`][tramdag.nodes.Node.fit], but every hook gets
        the flow instead of a node. ``EarlyStopping`` then tracks the
        summed validation NLL and restores the whole flow.

    Returns
    -------
    CausalFlowDAG
        ``self``, fitted, in eval mode.

    Raises
    ------
    ValueError
        See [`Node.fit`][tramdag.nodes.Node.fit].
    KeyError
        If a frame lacks a node column.
    """
    _check_learning_rate(learning_rate, optimizer)
    cbs = _normalize_callbacks(callbacks)
    train_df, validation_data = _split_validation(
        train_df, validation_data, validation_split
    )
    self._check_frames(train_df, validation_data)
    self._calibrate(train_df, marginal_init=marginal_init)
    nodes = [self.nodes[name] for name in self.order]
    side = [c for nd in nodes for c in nd.side_columns()]
    values = self._tensorize(train_df, [*self.order, *dict.fromkeys(side)])
    generator = None if seed is None else torch.Generator().manual_seed(seed)
    opt = _make_optimizer(self, optimizer, learning_rate)
    _train_loop(
        self,
        "flow",
        nodes,
        values,
        lambda: None if validation_data is None else self.nll(validation_data),
        opt,
        cbs,
        epochs=epochs,
        batch_size=batch_size,
        verbose=verbose,
        generator=generator,
    )
    return self

fit_classical

fit_classical(train_df, *, max_iter=5000, history_size=50)

Fit an all-ls model the classical way, node by node.

Each node runs Node.fit_classical: float64, full batch, L-BFGS and Newton. The joint objective is a sum of per-node terms with disjoint parameters, so the per-node optima are the joint one. Any term other than a simple intercept or LS raises before a node is fitted.

Parameters:

Name	Type	Description	Default
train_df	DataFrame	Training data, one column per node.	required
max_iter	int	Upper limit on L-BFGS iterations and Newton steps per node, by default 5000.	5000
history_size	int	L-BFGS memory, by default 50.	50

Returns:

Type	Description
dict	converged (every node converged), final_nll (the sum of the per-node mean NLLs), seconds, the fitted coefficients as {node: {parent: array}} (a node without LS terms is absent), and nodes, the per-node reports.

Notes

The report holds the float64 values of the fit. The flow keeps float32 parameters, so ls_coefficients, pmf and log_prob use these values rounded to float32, about 1e-7 relative. Compare the report with other software.

Warns:

Type	Description
UserWarning	For each node that did not converge.

Raises:

Type	Description
ValueError	If a term is not classical, an ordinal value is not a level index, or a node's calibration refuses the frame.
KeyError	If the frame lacks a node column.

Source code in src/tramdag/fitting.py

def fit_classical(
    self,
    train_df: pd.DataFrame,
    *,
    max_iter: int = 5000,
    history_size: int = 50,
) -> dict:
    """Fit an all-``ls`` model the classical way, node by node.

    Each node runs [`Node.fit_classical`][tramdag.nodes.Node.fit_classical]:
    float64, full batch, L-BFGS and Newton. The joint objective is a sum of per-node
    terms with disjoint parameters, so the per-node optima are the joint
    one. Any term other than a simple intercept or ``LS`` raises before a
    node is fitted.

    Parameters
    ----------
    train_df : pd.DataFrame
        Training data, one column per node.
    max_iter : int, optional
        Upper limit on L-BFGS iterations and Newton steps per node, by
        default 5000.
    history_size : int, optional
        L-BFGS memory, by default 50.

    Returns
    -------
    dict
        ``converged`` (every node converged), ``final_nll`` (the sum of
        the per-node mean NLLs), ``seconds``, the fitted ``coefficients``
        as ``{node: {parent: array}}`` (a node without ``LS`` terms is
        absent), and ``nodes``, the per-node reports.

    Notes
    -----
    The report holds the float64 values of the fit. The flow keeps float32
    parameters, so ``ls_coefficients``, ``pmf`` and ``log_prob`` use these
    values rounded to float32, about 1e-7 relative. Compare the report with
    other software.

    Warns
    -----
    UserWarning
        For each node that did not converge.

    Raises
    ------
    ValueError
        If a term is not classical, an ordinal value is not a level index,
        or a node's calibration refuses the frame.
    KeyError
        If the frame lacks a node column.
    """
    for nd in self.nodes.values():
        nd._check_classical()
    self._tensorize(train_df)
    self._calibrate(train_df)
    reports = {
        name: self.nodes[name].fit_classical(
            train_df, max_iter=max_iter, history_size=history_size
        )
        for name in self.order
    }
    return {
        "converged": all(r["converged"] for r in reports.values()),
        "final_nll": sum(r["final_nll"] for r in reports.values()),
        "seconds": sum(r["seconds"] for r in reports.values()),
        "coefficients": {
            n: r["coefficients"] for n, r in reports.items() if r["coefficients"]
        },
        "nodes": reports,
    }