tramdag.nodes¶

tramdag.nodes ¶

The node model: one TRAM regression per variable, alone or in a flow.

Node bundles a variable's intercept (transform parameters), monotone transform and shift modules, and fits itself; CausalFlowDAG holds one per node and the DAG lives in which parents each node reads. scores.py, the read-outs and the tests read kind, parents, shifts, intercept, ut, net_input and theta_shift by name.

Node ¶

Bases: NodeFitMixin, Module

One TRAM regression of a variable on its parents, alone or in a flow.

It is an intercept plus additive shift terms. The intercept produces the transform parameters \(\vartheta\). The shift terms add up on the latent scale. The node-kind maths lives in the five methods log_prob, sample, abduct, marginal_theta and encode, in the module function encode and in the score derivative of scores. The rest of the package reads kind or the schema entry to check, convert or display values.

The node builds one module per shift term in formula order, because the draws of one seeded stream are pinned to it. Each shift names its own key: the parent, "a+b" for a joint CS, the treatment for a VC. The intercept_module of the first term picks the intercept slot: the free theta_0, one joint net, or one net per parent summed in coefficient space.

Parameters:

Name	Type	Description	Default
`name`	`str`	The node's column name.	required
`node_spec`	`NodeSpec`	Specification of the node.	required
`schema`	`dict[str, str \| int]`	The parent schema, `{parent: "continuous" \| n_levels}`. It sizes the parent features: a continuous parent enters raw, an ordinal one one-hot over its levels.	required
`seed`	`int \| None`	If given, seeds the weight initialization from `(seed, name)`. The weights are then those of this node in `CausalFlowDAG(spec, seed=seed)` with the default `init`. Construction does not change torch's global RNG state. Without a seed, the weights come from torch's global RNG. A lone node always uses torch's default init; `init=` exists on `CausalFlowDAG` only.	`None`

Attributes:

Name	Type	Description
`history`	`dict`	The fit history, see `fit`. `"train"` holds the mean train NLL per epoch.
`calibrated`	`Tensor`	Boolean buffer, set by the first fit's calibration and kept in the checkpoint. It is a buffer, not a Python bool, so the flag rides in the state dict and a loaded node does not recalibrate on its next fit.
`meta`	`dict`	The provenance that `load` fills.

Source code in src/tramdag/nodes.py

class Node(NodeFitMixin, nn.Module):
    r"""One TRAM regression of a variable on its parents, alone or in a flow.

    It is an intercept plus additive shift terms. The intercept produces the
    transform parameters $\vartheta$. The shift terms add up on the latent
    scale. The node-kind maths lives in the five methods ``log_prob``,
    ``sample``, ``abduct``, ``marginal_theta`` and ``encode``, in the module
    function ``encode`` and in the score derivative of ``scores``. The
    rest of the package reads ``kind`` or the schema entry to check, convert
    or display values.

    The node builds one module per shift term in formula order, because the
    draws of one seeded stream are pinned to it. Each shift names its own key: the
    parent, ``"a+b"`` for a joint CS, the treatment for a VC. The
    ``intercept_module`` of the first term picks the intercept slot: the free
    theta_0, one joint net, or one net per parent summed in coefficient space.

    Parameters
    ----------
    name : str
        The node's column name.
    node_spec : NodeSpec
        Specification of the node.
    schema : dict[str, str | int]
        The parent schema, ``{parent: "continuous" | n_levels}``. It sizes
        the parent features: a continuous parent enters raw, an ordinal one
        one-hot over its levels.
    seed : int | None, optional
        If given, seeds the weight initialization from ``(seed, name)``. The
        weights are then those of this node in ``CausalFlowDAG(spec,
        seed=seed)`` with the default ``init``. Construction does not change
        torch's global RNG state. Without a seed, the weights come from
        torch's global RNG. A lone node always uses torch's default init;
        ``init=`` exists on [`CausalFlowDAG`][tramdag.CausalFlowDAG] only.

    Attributes
    ----------
    history : dict
        The fit history, see [`fit`][tramdag.nodes.Node.fit]. ``"train"`` holds
        the mean train NLL per epoch.
    calibrated : Tensor
        Boolean buffer, set by the first fit's calibration and kept in the
        checkpoint. It is a buffer, not a Python bool, so the flag rides in the
        state dict and a loaded node does not recalibrate on its next fit.
    meta : dict
        The provenance that [`load`][tramdag.nodes.Node.load] fills.
    """

    def __init__(
        self,
        name: str,
        node_spec: NodeSpec,
        schema: dict[str, str | int],
        seed: int | None = None,
    ):
        super().__init__()
        with seeded(seed, name):
            self.name = name
            self.node_spec = node_spec
            self.kind = node_spec.kind
            terms = node_spec.terms
            self.parents = tuple(node_parents(node_spec))
            self.schema = {p: schema[p] for p in self.parents}
            if node_spec.kind == "continuous":
                self.ut = make_univariate_transform(
                    node_spec.transform, **node_spec.transform_kwargs
                )
                n_params = self.ut.n_params
            else:
                self.ut = None
                self.levels = node_spec.levels
                n_params = node_spec.levels - 1
            self.encoding = schema_entry(node_spec)
            self.intercept = terms[0].module(terms[0], self.schema, n_params)
            self.shifts = nn.ModuleDict()
            for term in terms[1:]:
                m = term.module(term, self.schema)
                self.shifts[m.key] = m
        self.register_buffer("calibrated", torch.tensor(False))
        self.history: dict = {"train": []}
        self.meta: dict = {}

    def encode(self, values: Tensor) -> Tensor:
        """Encode this node's values for use as a parent feature.

        A continuous parent stays raw, shape ``(n, 1)``; an ordinal parent is
        one-hot encoded, shape ``(n, levels)``.
        """
        return encode(values, self.encoding)

    def features(self, values: dict[str, Tensor]) -> dict[str, Tensor]:
        """Encode this node's parents out of a raw tensor dict."""
        return {p: encode(values[p], self.schema[p]) for p in self.parents}

    def tensorize(self, df: pd.DataFrame, cols=None) -> dict[str, Tensor]:
        """DataFrame columns -> one ``(n,)`` tensor each, in the node's dtype.

        ``cols=None`` takes the node's own column and its parents. Ordinal
        columns, own or parent, are checked to hold level indices.

        Raises
        ------
        KeyError
            If the frame lacks one of the columns.
        ValueError
            If an ordinal column is not a level index.
        """
        cols = (self.name, *self.parents) if cols is None else cols
        entries = self.schema | {self.name: self.encoding}
        levels = {c: e for c, e in entries.items() if e != "continuous"}
        p = next(self.parameters())
        return tensorize(df, cols, levels, p.dtype, p.device)

    def calibrate(self, train_df: pd.DataFrame, *, marginal_init: bool = False) -> Node:
        r"""Take the data-dependent state from the training rows, once.

        Every term calibrates itself: the intercept term maps the node's
        train ``range_q``/``1 - range_q`` quantiles (a transform option, set in
        ``transform_kwargs``, default 5%/95%; ``0.0`` is the min/max) onto the
        transform's pre-scaled domain, and every term with an ``input_transform=``
        freezes its statistics (minmax lo/hi, standardize mean/std, a
        callable's frozen train columns).

        The first ``fit`` or ``fit_classical`` calls this when it has not run
        yet; a loaded node is already calibrated, and later fits on other rows
        reuse this state. Data on a new scale needs a new node.

        ``marginal_init`` additionally starts a simple intercept at the
        column's marginal: a Bernstein intercept at the Bernstein
        approximation of $\operatorname{logit} \hat F(y)$, an ordinal one at
        the marginal class log-odds; spline/affine intercepts and intercepts
        with parents are untouched. The start rides on this method's guard, so
        a second ``fit`` (the next phase of a schedule) continues training
        instead of discarding the intercept it just trained.

        Parameters
        ----------
        train_df : pd.DataFrame
            Training rows with the node's column and its parents' columns.
        marginal_init : bool, optional
            Also set the calibrated start, by default ``False``: an
            uninitialized intercept starts at zuko's zero instead.

        Returns
        -------
        Node
            ``self``.
        """
        if bool(self.calibrated):
            return self
        self.tensorize(train_df)
        self.intercept.calibrate_intercept(train_df, train_df[self.name], self.ut)
        for m in self.shifts.values():
            m.calibrate(train_df)
        if marginal_init:
            theta = self.marginal_theta(train_df[self.name].to_numpy())
            if theta is not None:
                self.intercept.marginal_start(theta)
        self.calibrated.fill_(True)
        return self

    def log_prob(self, theta: Tensor, shift: Tensor, x: Tensor) -> Tensor:
        """``log p(x | pa)`` from the transform parameters and the shift."""
        if self.kind == "continuous":
            u0, ladj = self.ut.forward(theta, x)
            return StandardLogistic.log_prob(u0 + shift) + ladj
        return ordinal_log_prob(theta, shift, x)

    def sample(self, theta: Tensor, shift: Tensor, u: Tensor) -> Tensor:
        """Push the latent ``u`` forward to an observed value."""
        if self.kind == "continuous":
            return self.ut.inverse(theta, u - shift)
        return ordinal_sample(theta, shift, u)

    def abduct(self, theta: Tensor, shift: Tensor, x: Tensor, generator=None) -> Tensor:
        """Recover the latent: exact (continuous) or truncated-sampled (ordinal)."""
        if self.kind == "continuous":
            u0, _ = self.ut.forward(theta, x)
            return u0 + shift
        return ordinal_abduct(theta, shift, x, generator=generator)

    def marginal_theta(self, column: np.ndarray):
        r"""Give the node's marginal-start $\vartheta$, or ``None`` when there is none.

        Ordinal: the empirical class log-odds. Continuous: the transform's own
        marginal start over the same column (``None`` for spline/affine,
        nothing to set). Only a free intercept applies it.
        """
        if self.kind == "ordinal":
            counts = np.bincount(column.astype(np.int64), minlength=self.levels)
            return ordinal_marginal_init_theta(counts)
        return self.ut.marginal_init_theta(column)

    def net_input(self, feats: dict[str, Tensor], parents, key: str) -> Tensor:
        """Concatenate parent features for one term's network.

        ``key`` names the term ("@I" for the intercept, the shift key
        otherwise); a term with an ``input_transform`` gets its continuous
        parent columns transformed with the statistics frozen at
        calibration. Every network input goes through here, training and
        the read-outs (``varying_coef``, ``intercept_contributions``) alike,
        so the model seen at inference is the model that was fitted. Linear
        shifts and the VC treatment column are not network inputs and never
        pass through.
        """
        term = self.intercept if key == "@I" else self.shifts[key]
        tr = term.input_transform
        cols = []
        for p in parents:
            x = feats[p]
            if tr is not None and p in tr.cols:
                x = tr(x, tr.cols.index(p))
            cols.append(x)
        return torch.cat(cols, dim=1)

    def theta_shift(self, feats: dict[str, Tensor], n: int) -> tuple[Tensor, Tensor]:
        """Compute the transform parameters and the total shift of the node.

        The shifts add in ``order``, plain shifts first and VC last. The stable
        sort keeps the pinned order.

        Parameters
        ----------
        feats : dict[str, Tensor]
            Encoded parent features keyed by parent name, plus the terms'
            side columns (a centered VC's propensities, frozen from the
            training frame, injected live by the flow at query time).
        n : int
            Batch size.

        Returns
        -------
        tuple[Tensor, Tensor]
            The transform parameters, shape ``(n, P)``, and the total
            shift, shape ``(n,)``.

        Raises
        ------
        RuntimeError
            If a centered VC term is evaluated without its propensity
            column in ``feats``.
        """
        theta = self.intercept.theta_value(self, feats, n)
        shift = torch.zeros(n, dtype=theta.dtype, device=theta.device)
        for m in sorted(self.shifts.values(), key=lambda m: m.order):
            shift = shift + m.shift_value(self, feats)
        return theta, shift

    def ls_coefficients(self) -> dict[str, np.ndarray]:
        """Give the ``LS`` weights as ``{parent: array}``; empty without ``LS``."""
        return {
            parent: m.weight.detach().cpu().numpy().ravel().copy()
            for parent, m in self.shifts.items()
            if isinstance(m, LinearShiftModule)
        }

    def parameter_groups(self) -> dict[str, list[nn.Parameter]]:
        r"""Give every parameter of the node once, by its kind.

        An optimizer factory builds its parameter groups from this, for
        example weight decay on ``"networks"`` only.

        Returns
        -------
        dict[str, list[nn.Parameter]]
            ``"intercept"``: the free $\vartheta$ of a simple intercept.
            ``"linear"``: the ``LS`` weights and each ``VC``'s $\beta_0$.
            ``"networks"``: every network, the ``CI``, ``CS`` and $b_\Theta$
            of a ``VC``. Each term module gives its own through
            ``parameter_groups``.
        """
        groups = {"intercept": [], "linear": [], "networks": []}
        for m in (self.intercept, *self.shifts.values()):
            for role, params in m.parameter_groups().items():
                groups[role] += params
        return groups

    def warm_start(self, train_df: pd.DataFrame) -> dict:
        r"""Start each ``VC`` $\beta_0$ at the classical linear treatment effect.

        An all-``LS`` twin of the node, a simple intercept plus one ``LS``
        term per parent, fits classically on ``train_df``. Each ``VC``
        $\beta_0$ takes the twin's coefficient of its treatment. A binary
        ordinal treatment takes the level-1-against-0 contrast ``w[1] - w[0]``,
        a continuous one its weight ``w[0]``. Call it before ``fit``; the twin
        ignores a centering column. The twin starts from seed 0, so the result
        does not depend on torch's global RNG.

        Parameters
        ----------
        train_df : pd.DataFrame
            Training rows: the node's column and its parents' columns.

        Returns
        -------
        dict
            The twin's [`fit_classical`][tramdag.nodes.Node.fit_classical]
            report.

        Raises
        ------
        ValueError
            If the node has no ``VC`` term.
        """
        vcs = [
            m for m in self.shifts.values() if isinstance(m, VaryingCoefficientModule)
        ]
        if not vcs:
            raise ValueError(f"node {self.name!r} has no VC term to warm-start")
        spec = copy.copy(self.node_spec)
        spec.terms = [Intercept(), *(LinearShift(p) for p in self.parents)]
        twin = Node(self.name, spec, self.schema, seed=0).to(
            next(self.parameters()).device
        )
        report = twin.fit_classical(train_df)
        with torch.no_grad():
            for m in vcs:
                w = report["coefficients"][m.key]
                m.beta0.fill_(float(w[1] - w[0] if m.t_is_ord else w[0]))
        return report

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

        The file holds the node's name, spec, parent schema and weights, its
        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,
            {
                "name": self.name,
                "node_spec": spec_to_dict({self.name: self.node_spec})[self.name],
                "schema": self.schema,
                "state_dict": self.state_dict(),
                "history": self.history,
            },
            next(self.parameters()).device,
        )

    @classmethod
    def load(cls, path: str | Path, device: str = "cpu") -> Node:
        """Restore a node from a checkpoint written by [`save`][].

        The node's ``history`` and ``meta`` are refilled.

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

        Returns
        -------
        Node
            The restored node, in eval mode.
        """
        ckpt = torch.load(path, map_location=device, weights_only=False)
        name = ckpt["name"]
        spec = spec_from_dict({name: ckpt["node_spec"]})[name]
        node = cls(name, spec, ckpt["schema"]).to(device)
        load_weights(node, ckpt["state_dict"])
        node.history = ckpt["history"]
        node.meta = ckpt["meta"]
        return node.eval()

encode ¶

encode(values)

Encode this node's values for use as a parent feature.

A continuous parent stays raw, shape (n, 1); an ordinal parent is one-hot encoded, shape (n, levels).

Source code in src/tramdag/nodes.py

def encode(self, values: Tensor) -> Tensor:
    """Encode this node's values for use as a parent feature.

    A continuous parent stays raw, shape ``(n, 1)``; an ordinal parent is
    one-hot encoded, shape ``(n, levels)``.
    """
    return encode(values, self.encoding)

features ¶

features(values)

Encode this node's parents out of a raw tensor dict.

Source code in src/tramdag/nodes.py

def features(self, values: dict[str, Tensor]) -> dict[str, Tensor]:
    """Encode this node's parents out of a raw tensor dict."""
    return {p: encode(values[p], self.schema[p]) for p in self.parents}

tensorize ¶

tensorize(df, cols=None)

DataFrame columns -> one (n,) tensor each, in the node's dtype.

cols=None takes the node's own column and its parents. Ordinal columns, own or parent, are checked to hold level indices.

Raises:

Type	Description
`KeyError`	If the frame lacks one of the columns.
`ValueError`	If an ordinal column is not a level index.

Source code in src/tramdag/nodes.py

def tensorize(self, df: pd.DataFrame, cols=None) -> dict[str, Tensor]:
    """DataFrame columns -> one ``(n,)`` tensor each, in the node's dtype.

    ``cols=None`` takes the node's own column and its parents. Ordinal
    columns, own or parent, are checked to hold level indices.

    Raises
    ------
    KeyError
        If the frame lacks one of the columns.
    ValueError
        If an ordinal column is not a level index.
    """
    cols = (self.name, *self.parents) if cols is None else cols
    entries = self.schema | {self.name: self.encoding}
    levels = {c: e for c, e in entries.items() if e != "continuous"}
    p = next(self.parameters())
    return tensorize(df, cols, levels, p.dtype, p.device)

calibrate ¶

calibrate(train_df, *, marginal_init=False)

Take the data-dependent state from the training rows, once.

Every term calibrates itself: the intercept term maps the node's train range_q/1 - range_q quantiles (a transform option, set in transform_kwargs, default 5%/95%; 0.0 is the min/max) onto the transform's pre-scaled domain, and every term with an input_transform= freezes its statistics (minmax lo/hi, standardize mean/std, a callable's frozen train columns).

The first fit or fit_classical calls this when it has not run yet; a loaded node is already calibrated, and later fits on other rows reuse this state. Data on a new scale needs a new node.

marginal_init additionally starts a simple intercept at the column's marginal: a Bernstein intercept at the Bernstein approximation of \(\operatorname{logit} \hat F(y)\), an ordinal one at the marginal class log-odds; spline/affine intercepts and intercepts with parents are untouched. The start rides on this method's guard, so a second fit (the next phase of a schedule) continues training instead of discarding the intercept it just trained.

Parameters:

Name	Type	Description	Default
`train_df`	`DataFrame`	Training rows with the node's column and its parents' columns.	required
`marginal_init`	`bool`	Also set the calibrated start, by default `False`: an uninitialized intercept starts at zuko's zero instead.	`False`

Returns:

Type	Description
`Node`	`self`.

Source code in src/tramdag/nodes.py

def calibrate(self, train_df: pd.DataFrame, *, marginal_init: bool = False) -> Node:
    r"""Take the data-dependent state from the training rows, once.

    Every term calibrates itself: the intercept term maps the node's
    train ``range_q``/``1 - range_q`` quantiles (a transform option, set in
    ``transform_kwargs``, default 5%/95%; ``0.0`` is the min/max) onto the
    transform's pre-scaled domain, and every term with an ``input_transform=``
    freezes its statistics (minmax lo/hi, standardize mean/std, a
    callable's frozen train columns).

    The first ``fit`` or ``fit_classical`` calls this when it has not run
    yet; a loaded node is already calibrated, and later fits on other rows
    reuse this state. Data on a new scale needs a new node.

    ``marginal_init`` additionally starts a simple intercept at the
    column's marginal: a Bernstein intercept at the Bernstein
    approximation of $\operatorname{logit} \hat F(y)$, an ordinal one at
    the marginal class log-odds; spline/affine intercepts and intercepts
    with parents are untouched. The start rides on this method's guard, so
    a second ``fit`` (the next phase of a schedule) continues training
    instead of discarding the intercept it just trained.

    Parameters
    ----------
    train_df : pd.DataFrame
        Training rows with the node's column and its parents' columns.
    marginal_init : bool, optional
        Also set the calibrated start, by default ``False``: an
        uninitialized intercept starts at zuko's zero instead.

    Returns
    -------
    Node
        ``self``.
    """
    if bool(self.calibrated):
        return self
    self.tensorize(train_df)
    self.intercept.calibrate_intercept(train_df, train_df[self.name], self.ut)
    for m in self.shifts.values():
        m.calibrate(train_df)
    if marginal_init:
        theta = self.marginal_theta(train_df[self.name].to_numpy())
        if theta is not None:
            self.intercept.marginal_start(theta)
    self.calibrated.fill_(True)
    return self

log_prob ¶

log_prob(theta, shift, x)

log p(x | pa) from the transform parameters and the shift.

Source code in src/tramdag/nodes.py

def log_prob(self, theta: Tensor, shift: Tensor, x: Tensor) -> Tensor:
    """``log p(x | pa)`` from the transform parameters and the shift."""
    if self.kind == "continuous":
        u0, ladj = self.ut.forward(theta, x)
        return StandardLogistic.log_prob(u0 + shift) + ladj
    return ordinal_log_prob(theta, shift, x)

sample ¶

sample(theta, shift, u)

Push the latent u forward to an observed value.

Source code in src/tramdag/nodes.py

def sample(self, theta: Tensor, shift: Tensor, u: Tensor) -> Tensor:
    """Push the latent ``u`` forward to an observed value."""
    if self.kind == "continuous":
        return self.ut.inverse(theta, u - shift)
    return ordinal_sample(theta, shift, u)

abduct ¶

abduct(theta, shift, x, generator=None)

Recover the latent: exact (continuous) or truncated-sampled (ordinal).

Source code in src/tramdag/nodes.py

def abduct(self, theta: Tensor, shift: Tensor, x: Tensor, generator=None) -> Tensor:
    """Recover the latent: exact (continuous) or truncated-sampled (ordinal)."""
    if self.kind == "continuous":
        u0, _ = self.ut.forward(theta, x)
        return u0 + shift
    return ordinal_abduct(theta, shift, x, generator=generator)

marginal_theta ¶

marginal_theta(column)

Give the node's marginal-start \(\vartheta\), or None when there is none.

Ordinal: the empirical class log-odds. Continuous: the transform's own marginal start over the same column (None for spline/affine, nothing to set). Only a free intercept applies it.

Source code in src/tramdag/nodes.py

def marginal_theta(self, column: np.ndarray):
    r"""Give the node's marginal-start $\vartheta$, or ``None`` when there is none.

    Ordinal: the empirical class log-odds. Continuous: the transform's own
    marginal start over the same column (``None`` for spline/affine,
    nothing to set). Only a free intercept applies it.
    """
    if self.kind == "ordinal":
        counts = np.bincount(column.astype(np.int64), minlength=self.levels)
        return ordinal_marginal_init_theta(counts)
    return self.ut.marginal_init_theta(column)

net_input ¶

net_input(feats, parents, key)

Concatenate parent features for one term's network.

key names the term ("@I" for the intercept, the shift key otherwise); a term with an input_transform gets its continuous parent columns transformed with the statistics frozen at calibration. Every network input goes through here, training and the read-outs (varying_coef, intercept_contributions) alike, so the model seen at inference is the model that was fitted. Linear shifts and the VC treatment column are not network inputs and never pass through.

Source code in src/tramdag/nodes.py

def net_input(self, feats: dict[str, Tensor], parents, key: str) -> Tensor:
    """Concatenate parent features for one term's network.

    ``key`` names the term ("@I" for the intercept, the shift key
    otherwise); a term with an ``input_transform`` gets its continuous
    parent columns transformed with the statistics frozen at
    calibration. Every network input goes through here, training and
    the read-outs (``varying_coef``, ``intercept_contributions``) alike,
    so the model seen at inference is the model that was fitted. Linear
    shifts and the VC treatment column are not network inputs and never
    pass through.
    """
    term = self.intercept if key == "@I" else self.shifts[key]
    tr = term.input_transform
    cols = []
    for p in parents:
        x = feats[p]
        if tr is not None and p in tr.cols:
            x = tr(x, tr.cols.index(p))
        cols.append(x)
    return torch.cat(cols, dim=1)

theta_shift ¶

theta_shift(feats, n)

Compute the transform parameters and the total shift of the node.

The shifts add in order, plain shifts first and VC last. The stable sort keeps the pinned order.

Parameters:

Name	Type	Description	Default
`feats`	`dict[str, Tensor]`	Encoded parent features keyed by parent name, plus the terms' side columns (a centered VC's propensities, frozen from the training frame, injected live by the flow at query time).	required
`n`	`int`	Batch size.	required

Returns:

Type	Description
`tuple[Tensor, Tensor]`	The transform parameters, shape `(n, P)`, and the total shift, shape `(n,)`.

Raises:

Type	Description
`RuntimeError`	If a centered VC term is evaluated without its propensity column in `feats`.

Source code in src/tramdag/nodes.py

def theta_shift(self, feats: dict[str, Tensor], n: int) -> tuple[Tensor, Tensor]:
    """Compute the transform parameters and the total shift of the node.

    The shifts add in ``order``, plain shifts first and VC last. The stable
    sort keeps the pinned order.

    Parameters
    ----------
    feats : dict[str, Tensor]
        Encoded parent features keyed by parent name, plus the terms'
        side columns (a centered VC's propensities, frozen from the
        training frame, injected live by the flow at query time).
    n : int
        Batch size.

    Returns
    -------
    tuple[Tensor, Tensor]
        The transform parameters, shape ``(n, P)``, and the total
        shift, shape ``(n,)``.

    Raises
    ------
    RuntimeError
        If a centered VC term is evaluated without its propensity
        column in ``feats``.
    """
    theta = self.intercept.theta_value(self, feats, n)
    shift = torch.zeros(n, dtype=theta.dtype, device=theta.device)
    for m in sorted(self.shifts.values(), key=lambda m: m.order):
        shift = shift + m.shift_value(self, feats)
    return theta, shift

ls_coefficients ¶

ls_coefficients()

Give the LS weights as {parent: array}; empty without LS.

Source code in src/tramdag/nodes.py

def ls_coefficients(self) -> dict[str, np.ndarray]:
    """Give the ``LS`` weights as ``{parent: array}``; empty without ``LS``."""
    return {
        parent: m.weight.detach().cpu().numpy().ravel().copy()
        for parent, m in self.shifts.items()
        if isinstance(m, LinearShiftModule)
    }

parameter_groups ¶

parameter_groups()

Give every parameter of the node once, by its kind.

An optimizer factory builds its parameter groups from this, for example weight decay on "networks" only.

Returns:

Type	Description
`dict[str, list[Parameter]]`	`"intercept"`: the free \(\vartheta\) of a simple intercept. `"linear"`: the `LS` weights and each `VC`'s \(\beta_0\). `"networks"`: every network, the `CI`, `CS` and \(b_\Theta\) of a `VC`. Each term module gives its own through `parameter_groups`.

Source code in src/tramdag/nodes.py

def parameter_groups(self) -> dict[str, list[nn.Parameter]]:
    r"""Give every parameter of the node once, by its kind.

    An optimizer factory builds its parameter groups from this, for
    example weight decay on ``"networks"`` only.

    Returns
    -------
    dict[str, list[nn.Parameter]]
        ``"intercept"``: the free $\vartheta$ of a simple intercept.
        ``"linear"``: the ``LS`` weights and each ``VC``'s $\beta_0$.
        ``"networks"``: every network, the ``CI``, ``CS`` and $b_\Theta$
        of a ``VC``. Each term module gives its own through
        ``parameter_groups``.
    """
    groups = {"intercept": [], "linear": [], "networks": []}
    for m in (self.intercept, *self.shifts.values()):
        for role, params in m.parameter_groups().items():
            groups[role] += params
    return groups

warm_start ¶

warm_start(train_df)

Start each VC \(\beta_0\) at the classical linear treatment effect.

An all-LS twin of the node, a simple intercept plus one LS term per parent, fits classically on train_df. Each VC \(\beta_0\) takes the twin's coefficient of its treatment. A binary ordinal treatment takes the level-1-against-0 contrast w[1] - w[0], a continuous one its weight w[0]. Call it before fit; the twin ignores a centering column. The twin starts from seed 0, so the result does not depend on torch's global RNG.

Parameters:

Name	Type	Description	Default
`train_df`	`DataFrame`	Training rows: the node's column and its parents' columns.	required

Returns:

Type	Description
`dict`	The twin's `fit_classical` report.

Raises:

Type	Description
`ValueError`	If the node has no `VC` term.

Source code in src/tramdag/nodes.py

def warm_start(self, train_df: pd.DataFrame) -> dict:
    r"""Start each ``VC`` $\beta_0$ at the classical linear treatment effect.

    An all-``LS`` twin of the node, a simple intercept plus one ``LS``
    term per parent, fits classically on ``train_df``. Each ``VC``
    $\beta_0$ takes the twin's coefficient of its treatment. A binary
    ordinal treatment takes the level-1-against-0 contrast ``w[1] - w[0]``,
    a continuous one its weight ``w[0]``. Call it before ``fit``; the twin
    ignores a centering column. The twin starts from seed 0, so the result
    does not depend on torch's global RNG.

    Parameters
    ----------
    train_df : pd.DataFrame
        Training rows: the node's column and its parents' columns.

    Returns
    -------
    dict
        The twin's [`fit_classical`][tramdag.nodes.Node.fit_classical]
        report.

    Raises
    ------
    ValueError
        If the node has no ``VC`` term.
    """
    vcs = [
        m for m in self.shifts.values() if isinstance(m, VaryingCoefficientModule)
    ]
    if not vcs:
        raise ValueError(f"node {self.name!r} has no VC term to warm-start")
    spec = copy.copy(self.node_spec)
    spec.terms = [Intercept(), *(LinearShift(p) for p in self.parents)]
    twin = Node(self.name, spec, self.schema, seed=0).to(
        next(self.parameters()).device
    )
    report = twin.fit_classical(train_df)
    with torch.no_grad():
        for m in vcs:
            w = report["coefficients"][m.key]
            m.beta0.fill_(float(w[1] - w[0] if m.t_is_ord else w[0]))
    return report

save ¶

save(path)

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

The file holds the node's name, spec, parent schema and weights, its 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/nodes.py

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

    The file holds the node's name, spec, parent schema and weights, its
    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,
        {
            "name": self.name,
            "node_spec": spec_to_dict({self.name: self.node_spec})[self.name],
            "schema": self.schema,
            "state_dict": self.state_dict(),
            "history": self.history,
        },
        next(self.parameters()).device,
    )

load `classmethod` ¶

load(path, device='cpu')

Restore a node from a checkpoint written by save.

The node's history and meta are refilled.

Parameters:

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

Returns:

Type	Description
`Node`	The restored node, in eval mode.

Source code in src/tramdag/nodes.py

@classmethod
def load(cls, path: str | Path, device: str = "cpu") -> Node:
    """Restore a node from a checkpoint written by [`save`][].

    The node's ``history`` and ``meta`` are refilled.

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

    Returns
    -------
    Node
        The restored node, in eval mode.
    """
    ckpt = torch.load(path, map_location=device, weights_only=False)
    name = ckpt["name"]
    spec = spec_from_dict({name: ckpt["node_spec"]})[name]
    node = cls(name, spec, ckpt["schema"]).to(device)
    load_weights(node, ckpt["state_dict"])
    node.history = ckpt["history"]
    node.meta = ckpt["meta"]
    return node.eval()

side_columns ¶

side_columns()

Name the frame columns the node's terms need beyond the parents.

Source code in src/tramdag/fitting.py

def side_columns(self) -> list[str]:
    """Name the frame columns the node's terms need beyond the parents."""
    cols = [c for m in self.shifts.values() for c in m.side_columns()]
    return list(dict.fromkeys(cols))

row_log_prob ¶

row_log_prob(values)

Give log p(x | pa) per row of already tensorized columns.

values holds the node's column, its parents' columns and its side columns, each shape (n,).

Source code in src/tramdag/fitting.py

def row_log_prob(self, values: dict[str, Tensor]) -> Tensor:
    """Give ``log p(x | pa)`` per row of already tensorized columns.

    ``values`` holds the node's column, its parents' columns and its side
    columns, each shape ``(n,)``.
    """
    n = len(values[self.name])
    feats = self.features(values) | {c: values[c] for c in self.side_columns()}
    theta, shift = self.theta_shift(feats, n)
    return self.log_prob(theta, shift, values[self.name])

nll ¶

nll(df)

Give the node's mean negative log-likelihood over the rows of df.

df holds the node's column, its parents' columns and its side columns (a centered VC's propensity).

Source code in src/tramdag/fitting.py

@torch.no_grad()
def nll(self, df: pd.DataFrame) -> float:
    """Give the node's mean negative log-likelihood over the rows of ``df``.

    ``df`` holds the node's column, its parents' columns and its side
    columns (a centered VC's propensity).
    """
    values = self.tensorize(df, (self.name, *self.parents, *self.side_columns()))
    return float(-self.row_log_prob(values).mean())

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
)

Fit the node by maximum likelihood with one minibatch loop, Adam by default.

The loss is the node's mean NLL plus its own VC penalties. The loop keeps the final weights, and a second fit call continues the training. The loop computes the validation NLL and prints the progress; learning-rate schedules, early stopping and best-weight restoration are the caller's, through optimizer and callbacks; callbacks ships EarlyStopping. The VC penalty never enters history["train"], and a VC term is re-centered after the loop.

node.history accumulates across fits: "train" and "val" hold one mean NLL per epoch, "val_epoch" the train epoch each validation entry belongs to, "lr" the optimizer's rates per epoch as {group index: lr}. "best_epoch" is the best epoch that an EarlyStopping of the last fit found, counted like "val_epoch"; a fit without one removes it.

Parameters:

Name	Type	Description	Default
`train_df`	`DataFrame`	Training rows: the node's column, its parents' columns and its side columns. Other columns are ignored.	required
`epochs`	`int`	Number of 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. `len(train_df)` is one full-batch step per epoch.	`512`
`validation_data`	`DataFrame \| None`	Validation rows, with the same columns as `train_df`, by default None. When given (or split off), `fit` appends the validation NLL to `node.history["val"]` after every epoch; `EarlyStopping` reads it there. A centered `VC` reads its propensity column from this frame; the flow's `fit` puts the live propensity of the fitted treatment node there.	`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. Mutually exclusive with `validation_data`.	`None`
`verbose`	`int`	0 (default) is silent. `N >= 1` prints one line every `N` epochs and on the final epoch: node, epoch counter, train NLL, validation NLL when validation is configured.	`0`
`seed`	`int \| None`	Seeds the minibatch shuffling with a private generator, by default None (torch's global RNG). `CausalFlowDAG.fit` and `fit_jointly` shuffle from the same generator, so with the same `seed` all three see the same row order in every epoch. `Node(seed=)` and `CausalFlowDAG(seed=)` seed the weight initialization.	`None`
`marginal_init`	`bool`	Start a simple intercept at its column's marginal, by default `False`. Applied only by the fit that calibrates the node. A spline or affine transform has no start and skips it.	`False`
`optimizer`	`Optimizer \| callable \| None`	A torch optimizer over `node.parameters()`, or a factory `f(node) -> Optimizer` (for example to split the parameters into groups). The default (None) is `Adam(lr=1e-2)`.	`None`
`callbacks`	`Callback \| callable \| list \| None`	One entry or a list, by default None. A `Callback` hooks all three points of the fit; its docstring is the contract. A bare callable is an `on_epoch_end` hook, `cb(node, epoch, optimizer)`; a `True` return stops the fit. Two callbacks that restore weights are refused.	`None`

Returns:

Type	Description
`Node`	`self`, fitted, in eval mode.

Notes

fit validates the frames before it calibrates, so a malformed frame cannot half-mutate the node. A batch of one row is skipped: batch norm needs two rows and one row gives no gradient.

history["val_epoch"] holds the train epoch of each validation entry. history accumulates across fit calls, so an unvalidated fit would otherwise shift every later validation entry away from its own epoch.

Every callback runs at each epoch end, so a stop does not skip a monitoring callback. The learning rate is read after the callbacks, so a scheduler's decision for this epoch shows. on_fit_end runs before the VC re-centering, so weights that a callback restores (such as EarlyStopping) still take part in it.

Raises:

Type	Description
`ValueError`	If `learning_rate` and `optimizer` are both given, the validation arguments clash, a centered VC term's propensity column is missing or out of [0, 1], two callbacks restore weights, an ordinal value is not a level index, the node's calibration refuses the frame, or `marginal_init` meets a `range_q=0` Bernstein transform.
`KeyError`	If a frame lacks a column the node reads.

Source code in src/tramdag/fitting.py

def fit(
    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[[Node], torch.optim.Optimizer]
    | None = None,
    callbacks=None,
) -> Node:
    """Fit the node by maximum likelihood with one minibatch loop, Adam by default.

    The loss is the node's mean NLL plus its own ``VC`` penalties. The
    loop keeps the **final** weights, and a second ``fit`` call continues
    the training. The loop computes the validation NLL and prints the
    progress; learning-rate schedules, early stopping and best-weight
    restoration are the caller's, through ``optimizer`` and
    ``callbacks``; [`callbacks`][tramdag.callbacks] ships ``EarlyStopping``.
    The ``VC`` penalty never enters ``history["train"]``, and a ``VC``
    term is re-centered after the loop.

    ``node.history`` accumulates across fits: ``"train"`` and ``"val"``
    hold one mean NLL per epoch, ``"val_epoch"`` the train epoch each
    validation entry belongs to, ``"lr"`` the optimizer's rates per epoch
    as ``{group index: lr}``. ``"best_epoch"`` is the best epoch that an
    ``EarlyStopping`` of the last fit found, counted like ``"val_epoch"``;
    a fit without one removes it.

    Parameters
    ----------
    train_df : pd.DataFrame
        Training rows: the node's column, its parents' columns and its
        side columns. Other columns are ignored.
    epochs : int
        Number of 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. ``len(train_df)`` is one
        full-batch step per epoch.
    validation_data : pd.DataFrame | None, optional
        Validation rows, with the same columns as ``train_df``, by
        default None. When given (or split off), ``fit`` appends the
        validation NLL to ``node.history["val"]`` after every epoch;
        ``EarlyStopping`` reads it there. A centered ``VC`` reads its
        propensity column from this frame; the flow's ``fit`` puts the
        live propensity of the fitted treatment node there.
    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. Mutually exclusive with ``validation_data``.
    verbose : int, optional
        0 (default) is silent. ``N >= 1`` prints one line every ``N``
        epochs and on the final epoch: node, epoch counter, train NLL,
        validation NLL when validation is configured.
    seed : int | None, optional
        Seeds the minibatch shuffling with a private generator, by
        default None (torch's global RNG). ``CausalFlowDAG.fit`` and
        ``fit_jointly`` shuffle from the same generator, so with the same
        ``seed`` all three see the same row order in every epoch.
        ``Node(seed=)`` and ``CausalFlowDAG(seed=)`` seed the weight
        initialization.
    marginal_init : bool, optional
        Start a simple intercept at its column's marginal, by default
        ``False``. Applied only by the fit that calibrates the node. A
        spline or affine transform has no start and skips it.
    optimizer : torch.optim.Optimizer | callable | None, optional
        A torch optimizer over ``node.parameters()``, or a factory
        ``f(node) -> Optimizer`` (for example to split the parameters into
        groups). The default (None) is ``Adam(lr=1e-2)``.
    callbacks : Callback | callable | list | None, optional
        One entry or a list, by default None. A
        [`Callback`][tramdag.callbacks.Callback]
        hooks all three points of the fit; its docstring is the contract.
        A bare callable is an ``on_epoch_end`` hook, ``cb(node, epoch,
        optimizer)``; a ``True`` return stops the fit. Two callbacks that
        restore weights are refused.

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

    Notes
    -----
    ``fit`` validates the frames before it calibrates, so a malformed frame
    cannot half-mutate the node. A batch of one row is skipped: batch norm
    needs two rows and one row gives no gradient.

    ``history["val_epoch"]`` holds the train epoch of each validation
    entry. ``history`` accumulates across ``fit`` calls, so an unvalidated
    fit would otherwise shift every later validation entry away from its
    own epoch.

    Every callback runs at each epoch end, so a stop does not skip a
    monitoring callback. The learning rate is read after the callbacks, so
    a scheduler's decision for this epoch shows. ``on_fit_end`` runs before
    the VC re-centering, so weights that a callback restores (such as
    ``EarlyStopping``) still take part in it.

    Raises
    ------
    ValueError
        If ``learning_rate`` and ``optimizer`` are both given, the
        validation arguments clash, a centered VC term's
        propensity column is missing or out of [0, 1], two callbacks
        restore weights, an ordinal value is not a level index, the
        node's calibration refuses the frame, or ``marginal_init`` meets
        a ``range_q=0`` Bernstein transform.
    KeyError
        If a frame lacks a column the node reads.
    """
    _check_learning_rate(learning_rate, optimizer)
    cbs = _normalize_callbacks(callbacks)
    train_df, validation_data = _split_validation(
        train_df, validation_data, validation_split
    )
    cols = (self.name, *self.parents, *self._check_side_columns(train_df))
    values = self.tensorize(train_df, cols)
    val_values = (
        self.tensorize(validation_data, cols)
        if validation_data is not None
        else None
    )
    self.calibrate(train_df, marginal_init=marginal_init)
    generator = None if seed is None else torch.Generator().manual_seed(seed)
    opt = _make_optimizer(self, optimizer, learning_rate)

    def validate():
        if val_values is None:
            return None
        with torch.no_grad():
            return {self.name: float(-self.row_log_prob(val_values).mean())}

    _train_loop(
        self,
        self.name,
        [self],
        values,
        validate,
        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 a simple intercept with LS terms the classical way.

The fit uses full batches, float64, and L-BFGS with a strong-Wolfe line search, then Newton steps on the exact Hessian. There are no minibatches, no schedule and no early stopping, so the fit is deterministic; the report says whether it reached the maximum-likelihood estimate. Any other term raises. docs/fitting.md says why.

Parameters:

Name	Type	Description	Default
`train_df`	`DataFrame`	Training rows: the node's column and its parents' columns.	required
`max_iter`	`int`	Upper limit on L-BFGS iterations and Newton steps together, by default 5000.	`5000`
`history_size`	`int`	L-BFGS memory, by default 50.	`50`

Returns:

Type	Description
`dict`	A convergence report: `converged`, `stop_reason`, `n_iter`, `final_nll`, `grad_norm`, `decrement`, `seconds`, and the fitted `LS` weights as `coefficients`, `{parent: array}`.

Warns:

Type	Description
`UserWarning`	If the report is not `converged`, with its `stop_reason` and `decrement`.

Raises:

Type	Description
`ValueError`	If a term is not classical: anything but the simple intercept and `LS`, if an ordinal value is not a level index, or if the node's calibration refuses the frame.
`KeyError`	If the frame lacks a column the node reads.

Notes

L-BFGS stops when one iteration lowers the NLL by less than 1e-9. On an ill-conditioned NLL, for example a parent on a large raw scale next to a rare level, this happens short of the optimum. Newton steps on the exact Hessian then finish the fit. The coefficients are read while the node is float64. The cast back to float32 loses about 1e-7 relative.

float64 is a transient compute mode. The node is upcast for the fit, and self.double() converts the parameters and the range buffers of the transform in one call. Afterwards the node returns to float32, so the stored model and save/load stay float32. Double precision is what lets the line search resolve the optimum cleanly.

stop_reason is what ended the run: "tolerance" when L-BFGS and the Newton steps stopped on their own, "max_iter" when the iteration budget ran out, and "max_eval" when torch's evaluation budget (max_iter * 5 // 4 closure calls, line searches included) ran out first. n_iter counts the L-BFGS iterations and the Newton steps. tolerance_grad is 0, so the gradient never ends the run. Only a "tolerance" stop of L-BFGS gets the Newton steps.

The Newton steps stop when the decrement is at most DECREMENT_TOL (1e-9), or when no step lowers the NLL. The decrement is the decrease of the mean NLL that a Newton step still expects. It does not depend on the scale of the parents, which the gradient norm does. A Bernstein NLL converges only linearly in its weakly curved tail coefficients, so this can take many steps. A spline NLL is only piecewise smooth: at a kink its Hessian is indefinite, no step lowers the NLL, and the decrement is above the bound.

converged is stop_reason == "tolerance": a stall short of the optimum is not converged, because the Newton steps continue from it until they stop on their own or spend the budget.

The Newton step skips the eigenvalues of the Hessian at float64 noise level. So it does not move along a direction with zero curvature, such as the one that the one-hot of an ordinal parent shares with the cutpoints. Compare the contrasts w[k] - w[0] with other software.

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 a simple intercept with ``LS`` terms the classical way.

    The fit uses full batches, float64, and L-BFGS with a strong-Wolfe
    line search, then Newton steps on the exact Hessian. There are no
    minibatches, no schedule and no early stopping, so the fit is
    deterministic; the report says whether it reached the
    maximum-likelihood estimate. Any other term raises.
    ``docs/fitting.md`` says why.

    Parameters
    ----------
    train_df : pd.DataFrame
        Training rows: the node's column and its parents' columns.
    max_iter : int, optional
        Upper limit on L-BFGS iterations and Newton steps together, by
        default 5000.
    history_size : int, optional
        L-BFGS memory, by default 50.

    Returns
    -------
    dict
        A convergence report: ``converged``, ``stop_reason``, ``n_iter``,
        ``final_nll``, ``grad_norm``, ``decrement``, ``seconds``, and the
        fitted ``LS`` weights as ``coefficients``, ``{parent: array}``.

    Warns
    -----
    UserWarning
        If the report is not ``converged``, with its ``stop_reason`` and
        ``decrement``.

    Raises
    ------
    ValueError
        If a term is not classical: anything but the simple intercept and
        ``LS``, if an ordinal value is not a level index, or if the node's
        calibration refuses the frame.
    KeyError
        If the frame lacks a column the node reads.

    Notes
    -----
    L-BFGS stops when one iteration lowers the NLL by less than 1e-9. On an
    ill-conditioned NLL, for example a parent on a large raw scale next to a
    rare level, this happens short of the optimum. Newton steps on the exact
    Hessian then finish the fit. The coefficients are read while the node
    is float64. The cast back to float32 loses about 1e-7 relative.

    float64 is a transient compute mode. The node is upcast for the fit,
    and ``self.double()`` converts the parameters and the range buffers of
    the transform in one call. Afterwards the node returns to float32, so
    the stored model and ``save``/``load`` stay float32. Double precision
    is what lets the line search resolve the optimum cleanly.

    ``stop_reason`` is what ended the run: ``"tolerance"`` when L-BFGS and
    the Newton steps stopped on their own, ``"max_iter"`` when the
    iteration budget ran out, and ``"max_eval"`` when torch's evaluation
    budget (``max_iter * 5 // 4`` closure calls, line searches included)
    ran out first. ``n_iter`` counts the L-BFGS iterations and the Newton
    steps. ``tolerance_grad`` is 0, so the gradient never ends the run.
    Only a ``"tolerance"`` stop of L-BFGS gets the Newton steps.

    The Newton steps stop when the ``decrement`` is at most
    ``DECREMENT_TOL`` (1e-9), or when no step lowers the NLL. The decrement
    is the decrease of the mean NLL that a Newton step still expects. It
    does not depend on the scale of the parents, which the gradient norm
    does. A Bernstein NLL converges only linearly in its weakly curved
    tail coefficients, so this can take many steps. A spline NLL
    is only piecewise smooth: at a kink its Hessian is indefinite, no step
    lowers the NLL, and the decrement is above the bound.

    ``converged`` is ``stop_reason == "tolerance"``: a stall short of the
    optimum is not converged, because the Newton steps continue from it
    until they stop on their own or spend the budget.

    The Newton step skips the eigenvalues of the Hessian at float64 noise
    level. So it does not move along a direction with zero curvature,
    such as the one that the one-hot of an ordinal parent shares with the
    cutpoints. Compare the contrasts ``w[k] - w[0]`` with other software.
    """
    self._check_classical()
    self.calibrate(train_df)
    self.double()
    t0 = time.perf_counter()
    try:
        values = self.tensorize(train_df)
        self.train()
        opt = torch.optim.LBFGS(
            self.parameters(),
            lr=1.0,
            max_iter=max_iter,
            history_size=history_size,
            tolerance_grad=0.0,
            tolerance_change=1e-9,
            line_search_fn="strong_wolfe",
        )

        def closure():
            opt.zero_grad()
            loss = -self.row_log_prob(values).mean()
            loss.backward()
            return loss

        opt.step(closure)
        state = next(iter(opt.state.values()))
        n_iter = state["n_iter"]
        if n_iter >= max_iter:
            stop_reason = "max_iter"
        elif state["func_evals"] >= opt.param_groups[0]["max_eval"]:
            stop_reason = "max_eval"
        else:
            stop_reason = "tolerance"
        if stop_reason == "tolerance":
            decrement, steps, settled = _newton_polish(
                self, values, max_iter - n_iter
            )
            n_iter += steps
            if not settled:
                stop_reason = "max_iter"
        else:
            decrement = _newton(self, values)[1]
        final_nll = float(closure().detach())
        grad_norm = float(
            torch.nn.utils.get_total_norm(
                [p.grad for p in self.parameters() if p.grad is not None]
            )
        )
        converged = stop_reason == "tolerance"
        coefs = self.ls_coefficients()
    finally:
        self.float()
    self.eval()

    if not converged:
        warnings.warn(
            f"fit_classical did not converge on node {self.name!r}: "
            f"stop_reason={stop_reason!r} after {n_iter} iterations, "
            f"decrement={decrement:.2e} (bound {DECREMENT_TOL}). Raise max_iter "
            "or read the report.",
            UserWarning,
            stacklevel=2,
        )
    return {
        "converged": converged,
        "stop_reason": stop_reason,
        "n_iter": n_iter,
        "final_nll": final_nll,
        "grad_norm": grad_norm,
        "decrement": decrement,
        "seconds": time.perf_counter() - t0,
        "coefficients": coefs,
    }

encode ¶

encode(values, entry)

Encode one column as a parent feature.

entry is a schema entry: "continuous" stays raw, shape (n, 1); a level count one-hot encodes, shape (n, levels).

Source code in src/tramdag/nodes.py

def encode(values: Tensor, entry: str | int) -> Tensor:
    """Encode one column as a parent feature.

    ``entry`` is a schema entry: ``"continuous"`` stays raw, shape ``(n, 1)``;
    a level count one-hot encodes, shape ``(n, levels)``.
    """
    if entry == "continuous":
        return values.view(-1, 1)
    one_hot = nn.functional.one_hot(values.long(), num_classes=entry)
    return one_hot.to(values.dtype)

schema_entry ¶

schema_entry(node_spec)

Give a node's schema entry: "continuous" or its level count.

Source code in src/tramdag/nodes.py

def schema_entry(node_spec: NodeSpec) -> str | int:
    """Give a node's schema entry: ``"continuous"`` or its level count."""
    return "continuous" if node_spec.kind == "continuous" else node_spec.levels

check_columns ¶

check_columns(df, cols)

Name the columns df lacks, before any tensor op would.

Source code in src/tramdag/nodes.py

def check_columns(df: pd.DataFrame, cols) -> None:
    """Name the columns ``df`` lacks, before any tensor op would."""
    missing = [c for c in cols if c not in df.columns]
    if missing:
        raise KeyError(
            f"the data frame lacks the column(s) {missing}; this needs "
            f"{list(cols)}, the frame has {list(df.columns)}"
        )

check_level_values ¶

check_level_values(name, values, levels)

Reject ordinal values that are not level indices of their node.

bincount, the cutpoint likelihood and the one-hot parent encoding all take the values as 0..levels-1; a 1-based or non-integer value would silently be truncated instead of failing.

Source code in src/tramdag/nodes.py

def check_level_values(name: str, values, levels: int) -> None:
    """Reject ordinal values that are not level indices of their node.

    ``bincount``, the cutpoint likelihood and the one-hot parent
    encoding all take the values as ``0..levels-1``; a 1-based or
    non-integer value would silently be truncated instead of failing.
    """
    v = np.asarray(values, dtype=np.float64)
    if v.size == 0:
        return
    fractional = bool((v != np.round(v)).any())
    if fractional or v.min() < 0 or v.max() >= levels:
        raise ValueError(
            f"node {name!r}: an ordinal column holds the level indices "
            f"0..{levels - 1}, got values in [{v.min()}, {v.max()}]"
            f"{' (non-integer)' if fractional else ''}"
        )

tensorize ¶

tensorize(df, cols, levels, dtype, device)

DataFrame columns -> one (n,) tensor each, in dtype on device.

A column named in levels is checked to hold level indices first, any other column to hold finite values.

Source code in src/tramdag/nodes.py

def tensorize(
    df: pd.DataFrame, cols, levels: dict[str, int], dtype, device
) -> dict[str, Tensor]:
    """DataFrame columns -> one ``(n,)`` tensor each, in ``dtype`` on ``device``.

    A column named in ``levels`` is checked to hold level indices first,
    any other column to hold finite values.
    """
    check_columns(df, cols)
    out = {}
    for c in cols:
        values = np.ascontiguousarray(df[c].to_numpy(dtype=float))
        if c in levels:
            check_level_values(c, values, levels[c])
        elif not np.isfinite(values).all():
            raise ValueError(f"column {c!r} holds nan or inf values")
        out[c] = torch.tensor(values, dtype=dtype, device=device)
    return out

write_checkpoint ¶

write_checkpoint(path, payload, device)

Save payload plus a meta block with torch.save.

meta holds the tramdag version, the save time and the device.

Raises:

Type	Description
`ValueError`	If the spec does not pickle (a callable `input_transform` that is not a module-level function).

Source code in src/tramdag/nodes.py

def write_checkpoint(path: str | Path, payload: dict, device) -> None:
    """Save ``payload`` plus a ``meta`` block with ``torch.save``.

    ``meta`` holds the tramdag version, the save time and the device.

    Raises
    ------
    ValueError
        If the spec does not pickle (a callable ``input_transform`` that is not
        a module-level function).
    """
    from . import __version__

    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    meta = {
        "tramdag_version": __version__,
        "saved_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "device": str(device),
    }
    try:
        torch.save(payload | {"meta": meta}, path)
    except (pickle.PicklingError, AttributeError) as err:
        raise ValueError(
            "the spec does not serialize: a callable input_transform "
            "must be a picklable module-level function; use "
            "'minmax'/'standardize', or def the function at module level"
        ) from err

load_weights ¶

load_weights(module, state_dict)

Load a checkpoint's state dict into a freshly built module.

A callable input transform's frozen train columns take the checkpoint's shape first; their buffer is empty until calibration.

Source code in src/tramdag/nodes.py

def load_weights(module: nn.Module, state_dict: dict) -> None:
    """Load a checkpoint's state dict into a freshly built module.

    A callable input transform's frozen train columns take the checkpoint's
    shape first; their buffer is empty until calibration.
    """
    for name, t in state_dict.items():
        if not name.endswith(".train_cols"):
            continue
        buf = module.get_buffer(name)
        if buf.shape != t.shape:
            mod_path, _, buf_name = name.rpartition(".")
            module.get_submodule(mod_path).register_buffer(
                buf_name, torch.empty_like(t)
            )
    module.load_state_dict(state_dict)

seeded ¶

seeded(seed, name)

Draw the block's random numbers from a stream keyed on (seed, name).

The block runs in torch.random.fork_rng, so torch's global RNG state is the same after the block as before it. Without a seed the block draws from torch's global stream.

Source code in src/tramdag/nodes.py

@contextlib.contextmanager
def seeded(seed: int | None, name: str):
    """Draw the block's random numbers from a stream keyed on ``(seed, name)``.

    The block runs in ``torch.random.fork_rng``, so torch's global RNG state
    is the same after the block as before it. Without a seed the block draws
    from torch's global stream.
    """
    if seed is None:
        yield
        return
    with torch.random.fork_rng(devices=[]):
        key = [seed, zlib.crc32(name.encode())]
        state = np.random.SeedSequence(key).generate_state(1)[0]
        torch.default_generator.manual_seed(int(state))
        yield

tramdag.nodes¶

tramdag.nodes ¶

Node ¶

encode ¶

features ¶

tensorize ¶

calibrate ¶

log_prob ¶

sample ¶

abduct ¶

marginal_theta ¶

net_input ¶

theta_shift ¶

ls_coefficients ¶

parameter_groups ¶

warm_start ¶

save ¶

load classmethod ¶

side_columns ¶

row_log_prob ¶

nll ¶

fit ¶

fit_classical ¶

encode ¶

schema_entry ¶

check_columns ¶

check_level_values ¶

tensorize ¶

write_checkpoint ¶

load_weights ¶

seeded ¶

load `classmethod` ¶