tramdag.modules

tramdag.modules

The term modules: one nn.Module per term, built from the term's spec class.

A spec term (Term subclass: LS, CS, VC, I) is plain data and carries the spec-level rules, and its module attribute is the class here that trains it (LinearShift.module is LinearShiftModule). A module is constructed from its term and the node's parent schema, module(term, schema) (the intercept slot adds n_params), holds the term's network and owns the runtime behaviour: shift_value/theta_value, post_init, regularizer, finalize, score_columns and the side-input contract. This module imports nothing from spec: it reads the schema, {parent: "continuous" | n_levels}, and a term's parents and options, so spec can import it.

A custom term is two classes: a ShiftModule subclass with __init__(term, schema) and shift_value, and a Term subclass for the options and checks whose module is that class.

The widths and the activation of every network are options of the term classes, written once in the signatures in spec; the modules here take what they are given.

Module	Network	Term
LinearShiftModule	Linear(n, 1, bias=False)	LS
ComplexShiftModule	hidden stack to 1 output, no bias	CS
ComplexInterceptModule	hidden stack to n_params outputs, no bias	I(...)
AdditiveInterceptModule	one stack per parent, outputs summed	additive I
SimpleInterceptModule	free parameter vector, no parent	I()
VaryingCoefficientModule	beta0 + penalized hidden stack to 1	VC

Parent features: a continuous parent enters raw, in one column; an ordinal parent one-hot, in levels columns. ACTIVATIONS maps the activation names a term may name to their torch.nn classes.

TermModule

Bases: Module

What every term module shares: the input transform and its calibration.

Source code in src/tramdag/modules.py

class TermModule(nn.Module):
    """What every term module shares: the input transform and its calibration."""

    @property
    def input_transform(self):
        """The term's frozen network-input transform, or ``None``.

        A module's ``__init__`` registers one (``_attach_input_transform``) when
        the term declares ``input_transform=`` over continuous parents; a plain
        class attribute would shadow the registered submodule.
        """
        return self._modules.get("_input_transform")

    def parameter_groups(self) -> dict[str, list[nn.Parameter]]:
        """Give the term's parameters by role; all of them are ``"networks"`` here.

        A term with a free intercept or an interpretable coefficient overrides
        this and gives those as ``"intercept"`` or ``"linear"``.
        """
        return {"networks": list(self.parameters())}

    def calibrate(self, train_df: pd.DataFrame) -> None:
        """Freeze this term's data-dependent state: the input-transform stats.

        ``Node.calibrate`` calls this once per term; a term without
        an ``input_transform`` has nothing to freeze. The intercept slot has a
        second step on top of this one, see
        [`calibrate_intercept`][tramdag.modules.InterceptModule.calibrate_intercept].
        """
        tr = self.input_transform
        if tr is None:
            return
        cols = torch.stack(
            [
                torch.as_tensor(train_df[p].to_numpy(dtype=np.float32).copy())
                for p in tr.cols
            ],
            dim=1,
        )
        tr.set_stats(cols)

input_transform property

input_transform

The term's frozen network-input transform, or None.

A module's __init__ registers one (_attach_input_transform) when the term declares input_transform= over continuous parents; a plain class attribute would shadow the registered submodule.

parameter_groups

parameter_groups()

Give the term's parameters by role; all of them are "networks" here.

A term with a free intercept or an interpretable coefficient overrides this and gives those as "intercept" or "linear".

Source code in src/tramdag/modules.py

def parameter_groups(self) -> dict[str, list[nn.Parameter]]:
    """Give the term's parameters by role; all of them are ``"networks"`` here.

    A term with a free intercept or an interpretable coefficient overrides
    this and gives those as ``"intercept"`` or ``"linear"``.
    """
    return {"networks": list(self.parameters())}

calibrate

calibrate(train_df)

Freeze this term's data-dependent state: the input-transform stats.

Node.calibrate calls this once per term; a term without an input_transform has nothing to freeze. The intercept slot has a second step on top of this one, see calibrate_intercept.

Source code in src/tramdag/modules.py

def calibrate(self, train_df: pd.DataFrame) -> None:
    """Freeze this term's data-dependent state: the input-transform stats.

    ``Node.calibrate`` calls this once per term; a term without
    an ``input_transform`` has nothing to freeze. The intercept slot has a
    second step on top of this one, see
    [`calibrate_intercept`][tramdag.modules.InterceptModule.calibrate_intercept].
    """
    tr = self.input_transform
    if tr is None:
        return
    cols = torch.stack(
        [
            torch.as_tensor(train_df[p].to_numpy(dtype=np.float32).copy())
            for p in tr.cols
        ],
        dim=1,
    )
    tr.set_stats(cols)

ShiftModule

Bases: TermModule, ABC

A shift term's behavior hooks, on top of its network.

__init__(term, schema) builds the network from the term's options and the parents' widths in the schema, and sets key (the node's ModuleDict key) and parents (the term's written parents); a subclass may keep more (VaryingCoefficientModule keeps mods/t_is_ord/propensity_col). Layers are built in a fixed order under fixed attribute names, so state-dict paths and the seeded RNG stream stay bit-stable.

scored is true when score_columns gives coefficients. order sets the summing order of the shifts; VC comes last, which is the pinned order.

Source code in src/tramdag/modules.py

class ShiftModule(TermModule, ABC):
    """A shift term's behavior hooks, on top of its network.

    ``__init__(term, schema)`` builds the network from the term's options and the
    parents' widths in the schema, and sets ``key`` (the node's ModuleDict key)
    and ``parents`` (the term's written parents); a subclass may keep more
    (``VaryingCoefficientModule`` keeps ``mods``/``t_is_ord``/``propensity_col``).
    Layers are built in a fixed order under fixed attribute names, so
    state-dict paths and the seeded RNG stream stay bit-stable.

    ``scored`` is true when ``score_columns`` gives coefficients. ``order``
    sets the summing order of the shifts; VC comes last, which is the pinned
    order.
    """

    scored = False
    order = 0

    key: str
    parents: tuple[str, ...]

    @abstractmethod
    def shift_value(self, node: Node, feats: dict) -> Tensor:
        """Give this term's contribution to the node's shift, shape ``(n,)``.

        ``feats`` holds the node's encoded parents plus this term's side
        columns (frozen from the data frame during training, injected live
        by the flow at query time).
        """

    def post_init(self) -> None:
        """Re-apply construction-time invariants after a global weight init."""

    def regularizer(self) -> Tensor | None:
        """Give the term's penalty on the total-likelihood scale, or ``None``."""
        return None

    def finalize(self, node: Node, feats: dict) -> None:
        """Run the term's post-fit step (after the after-fit callbacks)."""

    def score_columns(self, node: Node, feats: dict, dlds) -> dict:
        """Give the per-observation score columns of this term's coefficients.

        Empty for a term with no interpretable coefficient (``CS``).
        """
        return {}

    def side_columns(self) -> tuple[str, ...]:
        """Name the data-frame columns this term needs beyond the parents.

        Training reads them from ``train_df`` like any other column (so the
        validation split and the minibatch slicing come for free).
        """
        return ()

    def check_column(self, node_name: str, col: str, values) -> None:
        """Validate one side column's values at fit time."""

    def live_side(self, flow, values: dict, n: int) -> dict:
        """Recompute the side columns from the fitted flow, at query time."""
        return {}

    def extra_columns(self, flow) -> list[str]:
        """List the extra columns queries must tensorize for ``live_side``."""
        return []

shift_value abstractmethod

shift_value(node, feats)

Give this term's contribution to the node's shift, shape (n,).

feats holds the node's encoded parents plus this term's side columns (frozen from the data frame during training, injected live by the flow at query time).

Source code in src/tramdag/modules.py

@abstractmethod
def shift_value(self, node: Node, feats: dict) -> Tensor:
    """Give this term's contribution to the node's shift, shape ``(n,)``.

    ``feats`` holds the node's encoded parents plus this term's side
    columns (frozen from the data frame during training, injected live
    by the flow at query time).
    """

post_init

post_init()

Re-apply construction-time invariants after a global weight init.

Source code in src/tramdag/modules.py

def post_init(self) -> None:
    """Re-apply construction-time invariants after a global weight init."""

regularizer

regularizer()

Give the term's penalty on the total-likelihood scale, or None.

Source code in src/tramdag/modules.py

def regularizer(self) -> Tensor | None:
    """Give the term's penalty on the total-likelihood scale, or ``None``."""
    return None

finalize

finalize(node, feats)

Run the term's post-fit step (after the after-fit callbacks).

Source code in src/tramdag/modules.py

def finalize(self, node: Node, feats: dict) -> None:
    """Run the term's post-fit step (after the after-fit callbacks)."""

score_columns

score_columns(node, feats, dlds)

Give the per-observation score columns of this term's coefficients.

Empty for a term with no interpretable coefficient (CS).

Source code in src/tramdag/modules.py

def score_columns(self, node: Node, feats: dict, dlds) -> dict:
    """Give the per-observation score columns of this term's coefficients.

    Empty for a term with no interpretable coefficient (``CS``).
    """
    return {}

side_columns

side_columns()

Name the data-frame columns this term needs beyond the parents.

Training reads them from train_df like any other column (so the validation split and the minibatch slicing come for free).

Source code in src/tramdag/modules.py

def side_columns(self) -> tuple[str, ...]:
    """Name the data-frame columns this term needs beyond the parents.

    Training reads them from ``train_df`` like any other column (so the
    validation split and the minibatch slicing come for free).
    """
    return ()

check_column

check_column(node_name, col, values)

Validate one side column's values at fit time.

Source code in src/tramdag/modules.py

def check_column(self, node_name: str, col: str, values) -> None:
    """Validate one side column's values at fit time."""

live_side

live_side(flow, values, n)

Recompute the side columns from the fitted flow, at query time.

Source code in src/tramdag/modules.py

def live_side(self, flow, values: dict, n: int) -> dict:
    """Recompute the side columns from the fitted flow, at query time."""
    return {}

extra_columns

extra_columns(flow)

List the extra columns queries must tensorize for live_side.

Source code in src/tramdag/modules.py

def extra_columns(self, flow) -> list[str]:
    """List the extra columns queries must tensorize for ``live_side``."""
    return []

InterceptModule

Bases: TermModule, ABC

The intercept slot's behavior hooks, on top of its network.

A node has exactly one intercept term (normalization guarantees node.terms[0]); it produces the transform parameters theta. groups carries the parent groups (empty for a simple intercept, one tuple for a joint net, one per parent for an additive one) and ci_parents their flat order. intercept_module picks the class for an I term; each class constructs from (term, schema, n_params).

Source code in src/tramdag/modules.py

class InterceptModule(TermModule, ABC):
    """The intercept slot's behavior hooks, on top of its network.

    A node has exactly one intercept term (normalization guarantees
    ``node.terms[0]``); it produces the transform parameters ``theta``.
    ``groups`` carries the parent groups (empty for a simple intercept,
    one tuple for a joint net, one per parent for an additive one) and
    ``ci_parents`` their flat order. [`intercept_module`][] picks the class
    for an ``I`` term; each class constructs from ``(term, schema, n_params)``.
    """

    groups: list[tuple[str, ...]]
    ci_parents: list[str]

    def calibrate_intercept(self, train_df: pd.DataFrame, own, ut) -> None:
        """Freeze the input stats and set the transform's domain, once.

        The intercept slot calibrates one thing more than every other term, so
        it says so in its own name rather than widening ``calibrate``: ``own``
        is the node's training column and ``ut`` its monotone transform
        (``None`` for an ordinal node: cutpoints have no domain), whose
        ``range_q``/``1 - range_q`` quantiles map onto the pre-scaled domain.
        """
        self.calibrate(train_df)
        if ut is None:
            return
        q = own.quantile([ut.range_q, 1.0 - ut.range_q])
        if q.iloc[1] <= q.iloc[0]:
            raise ValueError(
                f"node {own.name!r}: the {ut.range_q:.0%} and "
                f"{1 - ut.range_q:.0%} quantiles coincide at {q.iloc[0]}. A "
                "continuous node needs a spread of values; a level index "
                "needs OrdinalNode()."
            )
        ut.set_range(q.iloc[0], q.iloc[1])

    @abstractmethod
    def theta_value(self, node: Node, feats: dict, n: int) -> Tensor:
        """Give the transform parameters, shape ``(n, P)``."""

    def marginal_start(self, theta: Tensor) -> None:
        """Set the calibrated marginal start; only a free intercept has one."""

calibrate_intercept

calibrate_intercept(train_df, own, ut)

Freeze the input stats and set the transform's domain, once.

The intercept slot calibrates one thing more than every other term, so it says so in its own name rather than widening calibrate: own is the node's training column and ut its monotone transform (None for an ordinal node: cutpoints have no domain), whose range_q/1 - range_q quantiles map onto the pre-scaled domain.

Source code in src/tramdag/modules.py

def calibrate_intercept(self, train_df: pd.DataFrame, own, ut) -> None:
    """Freeze the input stats and set the transform's domain, once.

    The intercept slot calibrates one thing more than every other term, so
    it says so in its own name rather than widening ``calibrate``: ``own``
    is the node's training column and ``ut`` its monotone transform
    (``None`` for an ordinal node: cutpoints have no domain), whose
    ``range_q``/``1 - range_q`` quantiles map onto the pre-scaled domain.
    """
    self.calibrate(train_df)
    if ut is None:
        return
    q = own.quantile([ut.range_q, 1.0 - ut.range_q])
    if q.iloc[1] <= q.iloc[0]:
        raise ValueError(
            f"node {own.name!r}: the {ut.range_q:.0%} and "
            f"{1 - ut.range_q:.0%} quantiles coincide at {q.iloc[0]}. A "
            "continuous node needs a spread of values; a level index "
            "needs OrdinalNode()."
        )
    ut.set_range(q.iloc[0], q.iloc[1])

theta_value abstractmethod

theta_value(node, feats, n)

Give the transform parameters, shape (n, P).

Source code in src/tramdag/modules.py

@abstractmethod
def theta_value(self, node: Node, feats: dict, n: int) -> Tensor:
    """Give the transform parameters, shape ``(n, P)``."""

marginal_start

marginal_start(theta)

Set the calibrated marginal start; only a free intercept has one.

Source code in src/tramdag/modules.py

def marginal_start(self, theta: Tensor) -> None:
    """Set the calibrated marginal start; only a free intercept has one."""

SimpleInterceptModule

Bases: InterceptModule

The free simple intercept: one theta vector, no parents.

Parameters:

Name	Type	Description	Default
term	Intercept	The parentless intercept term.	required
schema	dict[str, str | int]	The parent schema (unused: no parents to size).	required
n_params	int	Number of transform parameters.	required

Source code in src/tramdag/modules.py

class SimpleInterceptModule(InterceptModule):
    """The free simple intercept: one theta vector, no parents.

    Parameters
    ----------
    term : Intercept
        The parentless intercept term.
    schema : dict[str, str | int]
        The parent schema (unused: no parents to size).
    n_params : int
        Number of transform parameters.
    """

    def __init__(self, term: Term, schema: dict[str, str | int], n_params: int):
        super().__init__()
        self.theta = nn.Parameter(torch.zeros(n_params))
        self.groups, self.ci_parents = [], []

    def forward(self, n: int) -> Tensor:
        """Broadcast the parameters over a batch of ``n`` rows, shape ``(n, P)``."""
        return self.theta.unsqueeze(0).expand(n, -1)

    def theta_value(self, node: Node, feats: dict, n: int) -> Tensor:
        """Broadcast the free theta over the batch."""
        return self(n)

    def parameter_groups(self) -> dict[str, list[nn.Parameter]]:
        """Give the free theta as ``"intercept"``."""
        return {"intercept": [self.theta]}

    def marginal_start(self, theta: Tensor) -> None:
        """Start at the node's data marginal."""
        with torch.no_grad():
            self.theta.copy_(theta)

forward

forward(n)

Broadcast the parameters over a batch of n rows, shape (n, P).

Source code in src/tramdag/modules.py

def forward(self, n: int) -> Tensor:
    """Broadcast the parameters over a batch of ``n`` rows, shape ``(n, P)``."""
    return self.theta.unsqueeze(0).expand(n, -1)

theta_value

theta_value(node, feats, n)

Broadcast the free theta over the batch.

Source code in src/tramdag/modules.py

def theta_value(self, node: Node, feats: dict, n: int) -> Tensor:
    """Broadcast the free theta over the batch."""
    return self(n)

parameter_groups

parameter_groups()

Give the free theta as "intercept".

Source code in src/tramdag/modules.py

def parameter_groups(self) -> dict[str, list[nn.Parameter]]:
    """Give the free theta as ``"intercept"``."""
    return {"intercept": [self.theta]}

marginal_start

marginal_start(theta)

Start at the node's data marginal.

Source code in src/tramdag/modules.py

def marginal_start(self, theta: Tensor) -> None:
    """Start at the node's data marginal."""
    with torch.no_grad():
        self.theta.copy_(theta)

ComplexInterceptModule

Bases: InterceptModule

A single (possibly joint multi-parent) complex intercept net.

Several parents given to one term feed a single network, so they interact. The term's units, activation and batch_norm size the network.

Parameters:

Name	Type	Description	Default
term	Intercept	The intercept term, with at least one parent.	required
schema	dict[str, str | int]	The parent schema, {parent: "continuous" | n_levels}, for the parents' feature widths.	required
n_params	int	Number of transform parameters to produce.	required

Source code in src/tramdag/modules.py

class ComplexInterceptModule(InterceptModule):
    """A single (possibly joint multi-parent) complex intercept net.

    Several parents given to one term feed a single network, so they interact.
    The term's ``units``, ``activation`` and ``batch_norm`` size the network.

    Parameters
    ----------
    term : Intercept
        The intercept term, with at least one parent.
    schema : dict[str, str | int]
        The parent schema, ``{parent: "continuous" | n_levels}``, for the
        parents' feature widths.
    n_params : int
        Number of transform parameters to produce.
    """

    def __init__(self, term: Term, schema: dict[str, str | int], n_params: int):
        super().__init__()
        self.net = _nn(
            feat_width(schema, term.parents),
            term.units,
            n_params,
            activation=term.activation,
            batch_norm=term.batch_norm,
        )
        self.groups = [tuple(term.parents)]
        self.ci_parents = list(term.parents)
        _attach_input_transform(self, term, schema, tuple(term.parents))

    def forward(self, x: Tensor) -> Tensor:
        """Map parent features ``(n, n_features)`` to parameters ``(n, n_params)``."""
        return self.net(x)

    def theta_value(self, node: Node, feats: dict, n: int) -> Tensor:
        """Run the one net over the joint parent features."""
        return self(node.net_input(feats, self.ci_parents, "@I"))

forward

forward(x)

Map parent features (n, n_features) to parameters (n, n_params).

Source code in src/tramdag/modules.py

def forward(self, x: Tensor) -> Tensor:
    """Map parent features ``(n, n_features)`` to parameters ``(n, n_params)``."""
    return self.net(x)

theta_value

theta_value(node, feats, n)

Run the one net over the joint parent features.

Source code in src/tramdag/modules.py

def theta_value(self, node: Node, feats: dict, n: int) -> Tensor:
    """Run the one net over the joint parent features."""
    return self(node.net_input(feats, self.ci_parents, "@I"))

AdditiveInterceptModule

Bases: InterceptModule

allow_interaction=False: one net per parent, outputs summed.

Each parent reshapes the transform independently, in unconstrained coefficient space. The submodule stays named nets, because tests/test_statedict_stability.py pins the state-dict path.

Parameters:

Name	Type	Description	Default
term	Intercept	The intercept term with allow_interaction=False and two or more parents.	required
schema	dict[str, str | int]	The parent schema, {parent: "continuous" | n_levels}, for the parents' feature widths.	required
n_params	int	Number of transform parameters to produce.	required

Source code in src/tramdag/modules.py

class AdditiveInterceptModule(InterceptModule):
    """``allow_interaction=False``: one net per parent, outputs summed.

    Each parent reshapes the transform independently, in unconstrained
    coefficient space. The submodule stays named ``nets``, because
    ``tests/test_statedict_stability.py`` pins the state-dict path.

    Parameters
    ----------
    term : Intercept
        The intercept term with ``allow_interaction=False`` and two or more
        parents.
    schema : dict[str, str | int]
        The parent schema, ``{parent: "continuous" | n_levels}``, for the
        parents' feature widths.
    n_params : int
        Number of transform parameters to produce.
    """

    def __init__(self, term: Term, schema: dict[str, str | int], n_params: int):
        super().__init__()
        self.groups = [(p,) for p in term.parents]
        self.ci_parents = list(term.parents)
        self.nets = nn.ModuleList(
            _nn(
                feat_width(schema, grp),
                term.units,
                n_params,
                activation=term.activation,
                batch_norm=term.batch_norm,
            )
            for grp in self.groups
        )
        _attach_input_transform(self, term, schema, tuple(term.parents))

    def theta_value(self, node: Node, feats: dict, n: int) -> Tensor:
        """Sum the per-parent nets in coefficient space."""
        return sum(
            net(node.net_input(feats, grp, "@I"))
            for net, grp in zip(self.nets, self.groups, strict=True)
        )

theta_value

theta_value(node, feats, n)

Sum the per-parent nets in coefficient space.

Source code in src/tramdag/modules.py

def theta_value(self, node: Node, feats: dict, n: int) -> Tensor:
    """Sum the per-parent nets in coefficient space."""
    return sum(
        net(node.net_input(feats, grp, "@I"))
        for net, grp in zip(self.nets, self.groups, strict=True)
    )

LinearShiftModule

Bases: ShiftModule

LS: one raw-unit coefficient per feature of the single parent, no bias.

Keyed by the parent's name.

Parameters:

Name	Type	Description	Default
term	LinearShift	The term, with its one parent.	required
schema	dict[str, str | int]	The parent schema, for the parent's feature width.	required

Source code in src/tramdag/modules.py

class LinearShiftModule(ShiftModule):
    r"""``LS``: one raw-unit coefficient per feature of the single parent, no bias.

    Keyed by the parent's name.

    Parameters
    ----------
    term : LinearShift
        The term, with its one parent.
    schema : dict[str, str | int]
        The parent schema, for the parent's feature width.
    """

    scored = True

    def __init__(self, term: Term, schema: dict[str, str | int]):
        super().__init__()
        self.fc = nn.Linear(feat_width(schema, term.parents), 1, bias=False)
        self.key = term.parents[0]
        self.parents = tuple(term.parents)

    @property
    def weight(self) -> Tensor:
        """Tensor: the shift coefficients, shape ``(n_features,)``."""
        return self.fc.weight.squeeze(0)

    def forward(self, x: Tensor) -> Tensor:
        """Give the shift ``(n,)`` from the encoded features ``(n, n_features)``."""
        return self.fc(x).squeeze(-1)

    def parameter_groups(self) -> dict[str, list[nn.Parameter]]:
        """Give the weight as ``"linear"``."""
        return {"linear": [self.fc.weight]}

    def shift_value(self, node: Node, feats: dict) -> Tensor:
        """Give the raw parent column times the weight, with no input transform."""
        return self(torch.cat([feats[p] for p in self.parents], dim=1))

    def score_columns(self, node: Node, feats: dict, dlds) -> dict:
        r"""One column per weight: the parent (continuous) or its one-hot levels.

        $\partial \ell_i / \partial \beta = (\partial \ell_i / \partial s_i)\, x_i$,
        analytic and exact.
        """
        (parent,) = self.parents
        psi = (dlds.unsqueeze(1) * feats[parent]).cpu().numpy()
        if node.schema[parent] != "continuous":
            return {f"{parent}[{k}]": psi[:, k] for k in range(psi.shape[1])}
        return {self.key: psi[:, 0]}

weight property

weight

Tensor: the shift coefficients, shape (n_features,).

forward

forward(x)

Give the shift (n,) from the encoded features (n, n_features).

Source code in src/tramdag/modules.py

def forward(self, x: Tensor) -> Tensor:
    """Give the shift ``(n,)`` from the encoded features ``(n, n_features)``."""
    return self.fc(x).squeeze(-1)

parameter_groups

parameter_groups()

Give the weight as "linear".

Source code in src/tramdag/modules.py

def parameter_groups(self) -> dict[str, list[nn.Parameter]]:
    """Give the weight as ``"linear"``."""
    return {"linear": [self.fc.weight]}

shift_value

shift_value(node, feats)

Give the raw parent column times the weight, with no input transform.

Source code in src/tramdag/modules.py

def shift_value(self, node: Node, feats: dict) -> Tensor:
    """Give the raw parent column times the weight, with no input transform."""
    return self(torch.cat([feats[p] for p in self.parents], dim=1))

score_columns

score_columns(node, feats, dlds)

One column per weight: the parent (continuous) or its one-hot levels.

\(\partial \ell_i / \partial \beta = (\partial \ell_i / \partial s_i)\, x_i\), analytic and exact.

Source code in src/tramdag/modules.py

def score_columns(self, node: Node, feats: dict, dlds) -> dict:
    r"""One column per weight: the parent (continuous) or its one-hot levels.

    $\partial \ell_i / \partial \beta = (\partial \ell_i / \partial s_i)\, x_i$,
    analytic and exact.
    """
    (parent,) = self.parents
    psi = (dlds.unsqueeze(1) * feats[parent]).cpu().numpy()
    if node.schema[parent] != "continuous":
        return {f"{parent}[{k}]": psi[:, k] for k in range(psi.shape[1])}
    return {self.key: psi[:, 0]}

ComplexShiftModule

Bases: ShiftModule

CS: an additive network shift g(x) over its parents.

One net over the concatenated parents, keyed 'a' or 'a+b'. The term's units, activation and batch_norm size the network.

Parameters:

Name	Type	Description	Default
term	ComplexShift	The term, with its parents.	required
schema	dict[str, str | int]	The parent schema, {parent: "continuous" | n_levels}, for the parents' feature widths.	required

Source code in src/tramdag/modules.py

class ComplexShiftModule(ShiftModule):
    """``CS``: an additive network shift ``g(x)`` over its parents.

    One net over the concatenated parents, keyed ``'a'`` or ``'a+b'``. The
    term's ``units``, ``activation`` and ``batch_norm`` size the network.

    Parameters
    ----------
    term : ComplexShift
        The term, with its parents.
    schema : dict[str, str | int]
        The parent schema, ``{parent: "continuous" | n_levels}``, for the
        parents' feature widths.
    """

    def __init__(self, term: Term, schema: dict[str, str | int]):
        super().__init__()
        self.parents = tuple(term.parents)
        self.net = _nn(
            feat_width(schema, self.parents),
            term.units,
            1,
            activation=term.activation,
            batch_norm=term.batch_norm,
        )
        self.key = "+".join(self.parents)
        _attach_input_transform(self, term, schema, self.parents)

    def forward(self, x: Tensor) -> Tensor:
        """Give the shift ``(n,)`` from the encoded features ``(n, n_features)``."""
        return self.net(x).squeeze(-1)

    def shift_value(self, node: Node, feats: dict) -> Tensor:
        """Give the net over this term's (possibly input-transformed) features."""
        return self(node.net_input(feats, self.parents, self.key))

forward

forward(x)

Give the shift (n,) from the encoded features (n, n_features).

Source code in src/tramdag/modules.py

def forward(self, x: Tensor) -> Tensor:
    """Give the shift ``(n,)`` from the encoded features ``(n, n_features)``."""
    return self.net(x).squeeze(-1)

shift_value

shift_value(node, feats)

Give the net over this term's (possibly input-transformed) features.

Source code in src/tramdag/modules.py

def shift_value(self, node: Node, feats: dict) -> Tensor:
    """Give the net over this term's (possibly input-transformed) features."""
    return self(node.net_input(feats, self.parents, self.key))

VaryingCoefficientModule

Bases: ShiftModule

VC: \(\beta(\text{mod})\, x_t\) with \(\beta(x) = \beta_0 + b_\Theta(x)\).

The weights of \(b_\Theta\) carry the L2 penalty (l2; fit adds penalty * l2() / n to the node's mean NLL, which equals adding penalty * l2() to the summed NLL); \(\beta_0\) is not penalized. The output layer starts at zero, so \(\beta(x) = \beta_0\) at construction. With n_features == 0 there is no network and the term is LS(t). After the fit, recenter moves the training-mean of b_theta into beta0 through the center buffer, a reparameterization that leaves the modelled function unchanged. Keyed by the treatment's name.

Parameters:

Name	Type	Description	Default
term	VaryingCoefficient	The term: treatment first in parents, then the modifiers.	required
schema	dict[str, str | int]	The parent schema, for the modifiers' feature widths and the treatment's kind.	required

Source code in src/tramdag/modules.py

class VaryingCoefficientModule(ShiftModule):
    r"""``VC``: $\beta(\text{mod})\, x_t$ with $\beta(x) = \beta_0 + b_\Theta(x)$.

    The weights of $b_\Theta$ carry the L2 ``penalty`` (``l2``; ``fit`` adds
    ``penalty * l2() / n`` to the node's mean NLL, which equals adding
    ``penalty * l2()`` to the summed NLL); $\beta_0$ is not penalized. The
    output layer starts at zero, so $\beta(x) = \beta_0$ at construction. With
    ``n_features == 0`` there is no network and the term is ``LS(t)``. After
    the fit, ``recenter`` moves the training-mean of ``b_theta`` into
    ``beta0`` through the ``center`` buffer, a reparameterization that leaves
    the modelled function unchanged. Keyed by the treatment's name.

    Parameters
    ----------
    term : VaryingCoefficient
        The term: treatment first in ``parents``, then the modifiers.
    schema : dict[str, str | int]
        The parent schema, for the modifiers' feature widths and the
        treatment's kind.
    """

    scored = True
    order = 1

    def __init__(self, term: Term, schema: dict[str, str | int]):
        super().__init__()
        t, mods = term.parents[0], tuple(term.parents[1:])
        self.penalty = float(term.penalty)
        self.beta0 = nn.Parameter(torch.zeros(()))
        self.register_buffer("center", torch.zeros(()))
        n_features = feat_width(schema, mods)
        if n_features > 0:
            self.net = _nn(
                n_features,
                term.units,
                1,
                activation=term.activation,
                batch_norm=term.batch_norm,
                zero_init_last=True,
            )
        else:
            self.net = None
        self.key = t
        self.parents = tuple(term.parents)
        self.mods = mods
        self.t_is_ord = schema[t] != "continuous"
        self.propensity_col = term.propensity
        _attach_input_transform(self, term, schema, mods)

    def beta(self, mod_feats: Tensor | None, n: int) -> Tensor:
        r"""Give the effect values $\beta(x)$, shape ``(n,)``.

        ``mod_feats`` is ``None`` if, and only if, the term has no modifiers;
        ``n`` is the batch size, used in that case.
        """
        if self.net is None:
            return (self.beta0 - self.center).expand(n)
        return self.beta0 + self.net(mod_feats).squeeze(-1) - self.center

    def forward(self, t: Tensor, mod_feats: Tensor | None) -> Tensor:
        r"""Give the shift $\beta(\text{modifiers})\, t$, shape ``(n,)``.

        ``t`` is the raw treatment column ``(n, 1)``, ``mod_feats`` the encoded
        modifier features or ``None`` without modifiers.
        """
        return self.beta(mod_feats, t.shape[0]) * t.squeeze(-1)

    def l2(self) -> Tensor:
        """Sum the squared ``b_theta`` weights, the penalized quantity.

        ``beta0`` is never included; 0 without modifiers.
        """
        if self.net is None:
            return torch.zeros((), device=self.beta0.device, dtype=self.beta0.dtype)
        return sum(p.pow(2).sum() for p in self.net.parameters())

    @torch.no_grad()
    def recenter(self, mod_feats: Tensor | None) -> None:
        """Re-split ``beta0`` and ``b_theta`` so ``b_theta`` has mean zero.

        The mean is taken over ``mod_feats``. The removed constant moves into
        ``beta0``, so the modelled function does not change. Without
        modifiers this does nothing.
        """
        if self.net is None:
            return
        delta = (self.net(mod_feats).squeeze(-1) - self.center).mean()
        self.center += delta
        self.beta0 += delta

    def regressor(self, feats: dict) -> Tensor:
        r"""Give the ``(n, 1)`` column $\beta$ multiplies: the treatment, raw.

        The one-hot level-1 indicator for a binary ordinal treatment, the
        value itself for a continuous one; a centered term subtracts its
        propensity column (the regressor ``t - e_hat(x)`` of Robinson, 1988). It is
        also the score of ``beta0``, so ``score_columns`` reads it here.
        """
        if self.propensity_col and self.propensity_col not in feats:
            raise RuntimeError(
                f"centered VC term on {self.key!r} needs its propensity "
                f"column {self.propensity_col!r}. Internal callers inject it; "
                "never evaluate a centered term without its propensity."
            )
        t = feats[self.key][:, -1:] if self.t_is_ord else feats[self.key]
        if self.propensity_col:
            t = t - feats[self.propensity_col].view(-1, 1)
        return t

    def shift_value(self, node: Node, feats: dict) -> Tensor:
        r"""$\beta(\text{modifiers})$ times the regressor, guarding a centered term."""
        t = self.regressor(feats)
        mod_feat = node.net_input(feats, self.mods, self.key) if self.mods else None
        return self(t, mod_feat)

    def post_init(self) -> None:
        r"""Re-zero the head's output layer: $\beta(x) = \beta_0$ at start."""
        if self.net is not None:
            nn.init.zeros_(self.net[-1].weight)

    def regularizer(self) -> Tensor | None:
        r"""Give the penalty $\lambda \lVert b_\Theta \rVert^2$ on the NLL scale.

        ``None`` without a head to shrink (no modifiers, or penalty 0).
        """
        if self.net is None or self.penalty == 0:
            return None
        return self.penalty * self.l2()

    def parameter_groups(self) -> dict[str, list[nn.Parameter]]:
        r"""Give $\beta_0$ as ``"linear"`` and $b_\Theta$ as ``"networks"``."""
        nets = [] if self.net is None else list(self.net.parameters())
        return {"linear": [self.beta0], "networks": nets}

    def finalize(self, node: Node, feats: dict) -> None:
        """Re-split ``beta0``/``b_theta``: the head sums to zero over train."""
        if self.mods:
            self.recenter(node.net_input(feats, self.mods, self.key))

    def score_columns(self, node: Node, feats: dict, dlds) -> dict:
        r"""One column, keyed by the treatment: the ``beta0`` score.

        $\partial s / \partial \beta_0$ is the term's own ``regressor``, so forward
        and score share one definition by construction.
        """
        t = self.regressor(feats)
        return {self.key: (dlds * t.squeeze(-1)).cpu().numpy()}

    def side_columns(self) -> tuple[str, ...]:
        """Name the propensity column a centered term reads from the frame."""
        return (self.propensity_col,) if self.propensity_col else ()

    def check_column(self, node_name: str, col: str, values) -> None:
        """Propensities are probabilities."""
        if not ((values >= 0) & (values <= 1)).all():
            raise ValueError(
                f"column {col!r} (centered VC on node {node_name!r}) must "
                "hold probabilities in [0, 1]"
            )

    def live_side(self, flow, values: dict, n: int) -> dict:
        """Give the full-data propensity from the flow's own treatment node.

        Detached (no gradient reaches the treatment node from this node's
        loss) and derived from the current parent values, so
        ``do``-mutilated sampling centers with the intervened ``t``. Training
        uses the frozen column.
        """
        if not self.propensity_col:
            return {}
        p1 = flow._propensity(flow.nodes[self.key], values, n).detach()
        return {self.propensity_col: p1}

    def extra_columns(self, flow) -> list[str]:
        """Give the treatment's parents (a treatment cannot be centered itself)."""
        return list(flow.nodes[self.key].parents) if self.propensity_col else []

beta

beta(mod_feats, n)

Give the effect values \(\beta(x)\), shape (n,).

mod_feats is None if, and only if, the term has no modifiers; n is the batch size, used in that case.

Source code in src/tramdag/modules.py

def beta(self, mod_feats: Tensor | None, n: int) -> Tensor:
    r"""Give the effect values $\beta(x)$, shape ``(n,)``.

    ``mod_feats`` is ``None`` if, and only if, the term has no modifiers;
    ``n`` is the batch size, used in that case.
    """
    if self.net is None:
        return (self.beta0 - self.center).expand(n)
    return self.beta0 + self.net(mod_feats).squeeze(-1) - self.center

forward

forward(t, mod_feats)

Give the shift \(\beta(\text{modifiers})\, t\), shape (n,).

t is the raw treatment column (n, 1), mod_feats the encoded modifier features or None without modifiers.

Source code in src/tramdag/modules.py

def forward(self, t: Tensor, mod_feats: Tensor | None) -> Tensor:
    r"""Give the shift $\beta(\text{modifiers})\, t$, shape ``(n,)``.

    ``t`` is the raw treatment column ``(n, 1)``, ``mod_feats`` the encoded
    modifier features or ``None`` without modifiers.
    """
    return self.beta(mod_feats, t.shape[0]) * t.squeeze(-1)

l2

l2()

Sum the squared b_theta weights, the penalized quantity.

beta0 is never included; 0 without modifiers.

Source code in src/tramdag/modules.py

def l2(self) -> Tensor:
    """Sum the squared ``b_theta`` weights, the penalized quantity.

    ``beta0`` is never included; 0 without modifiers.
    """
    if self.net is None:
        return torch.zeros((), device=self.beta0.device, dtype=self.beta0.dtype)
    return sum(p.pow(2).sum() for p in self.net.parameters())

recenter

recenter(mod_feats)

Re-split beta0 and b_theta so b_theta has mean zero.

The mean is taken over mod_feats. The removed constant moves into beta0, so the modelled function does not change. Without modifiers this does nothing.

Source code in src/tramdag/modules.py

@torch.no_grad()
def recenter(self, mod_feats: Tensor | None) -> None:
    """Re-split ``beta0`` and ``b_theta`` so ``b_theta`` has mean zero.

    The mean is taken over ``mod_feats``. The removed constant moves into
    ``beta0``, so the modelled function does not change. Without
    modifiers this does nothing.
    """
    if self.net is None:
        return
    delta = (self.net(mod_feats).squeeze(-1) - self.center).mean()
    self.center += delta
    self.beta0 += delta

regressor

regressor(feats)

Give the (n, 1) column \(\beta\) multiplies: the treatment, raw.

The one-hot level-1 indicator for a binary ordinal treatment, the value itself for a continuous one; a centered term subtracts its propensity column (the regressor t - e_hat(x) of Robinson, 1988). It is also the score of beta0, so score_columns reads it here.

Source code in src/tramdag/modules.py

def regressor(self, feats: dict) -> Tensor:
    r"""Give the ``(n, 1)`` column $\beta$ multiplies: the treatment, raw.

    The one-hot level-1 indicator for a binary ordinal treatment, the
    value itself for a continuous one; a centered term subtracts its
    propensity column (the regressor ``t - e_hat(x)`` of Robinson, 1988). It is
    also the score of ``beta0``, so ``score_columns`` reads it here.
    """
    if self.propensity_col and self.propensity_col not in feats:
        raise RuntimeError(
            f"centered VC term on {self.key!r} needs its propensity "
            f"column {self.propensity_col!r}. Internal callers inject it; "
            "never evaluate a centered term without its propensity."
        )
    t = feats[self.key][:, -1:] if self.t_is_ord else feats[self.key]
    if self.propensity_col:
        t = t - feats[self.propensity_col].view(-1, 1)
    return t

shift_value

shift_value(node, feats)

\(\beta(\text{modifiers})\) times the regressor, guarding a centered term.

Source code in src/tramdag/modules.py

def shift_value(self, node: Node, feats: dict) -> Tensor:
    r"""$\beta(\text{modifiers})$ times the regressor, guarding a centered term."""
    t = self.regressor(feats)
    mod_feat = node.net_input(feats, self.mods, self.key) if self.mods else None
    return self(t, mod_feat)

post_init

post_init()

Re-zero the head's output layer: \(\beta(x) = \beta_0\) at start.

Source code in src/tramdag/modules.py

def post_init(self) -> None:
    r"""Re-zero the head's output layer: $\beta(x) = \beta_0$ at start."""
    if self.net is not None:
        nn.init.zeros_(self.net[-1].weight)

regularizer

regularizer()

Give the penalty \(\lambda \lVert b_\Theta \rVert^2\) on the NLL scale.

None without a head to shrink (no modifiers, or penalty 0).

Source code in src/tramdag/modules.py

def regularizer(self) -> Tensor | None:
    r"""Give the penalty $\lambda \lVert b_\Theta \rVert^2$ on the NLL scale.

    ``None`` without a head to shrink (no modifiers, or penalty 0).
    """
    if self.net is None or self.penalty == 0:
        return None
    return self.penalty * self.l2()

parameter_groups

parameter_groups()

Give \(\beta_0\) as "linear" and \(b_\Theta\) as "networks".

Source code in src/tramdag/modules.py

def parameter_groups(self) -> dict[str, list[nn.Parameter]]:
    r"""Give $\beta_0$ as ``"linear"`` and $b_\Theta$ as ``"networks"``."""
    nets = [] if self.net is None else list(self.net.parameters())
    return {"linear": [self.beta0], "networks": nets}

finalize

finalize(node, feats)

Re-split beta0/b_theta: the head sums to zero over train.

Source code in src/tramdag/modules.py

def finalize(self, node: Node, feats: dict) -> None:
    """Re-split ``beta0``/``b_theta``: the head sums to zero over train."""
    if self.mods:
        self.recenter(node.net_input(feats, self.mods, self.key))

score_columns

score_columns(node, feats, dlds)

One column, keyed by the treatment: the beta0 score.

\(\partial s / \partial \beta_0\) is the term's own regressor, so forward and score share one definition by construction.

Source code in src/tramdag/modules.py

def score_columns(self, node: Node, feats: dict, dlds) -> dict:
    r"""One column, keyed by the treatment: the ``beta0`` score.

    $\partial s / \partial \beta_0$ is the term's own ``regressor``, so forward
    and score share one definition by construction.
    """
    t = self.regressor(feats)
    return {self.key: (dlds * t.squeeze(-1)).cpu().numpy()}

side_columns

side_columns()

Name the propensity column a centered term reads from the frame.

Source code in src/tramdag/modules.py

def side_columns(self) -> tuple[str, ...]:
    """Name the propensity column a centered term reads from the frame."""
    return (self.propensity_col,) if self.propensity_col else ()

check_column

check_column(node_name, col, values)

Propensities are probabilities.

Source code in src/tramdag/modules.py

def check_column(self, node_name: str, col: str, values) -> None:
    """Propensities are probabilities."""
    if not ((values >= 0) & (values <= 1)).all():
        raise ValueError(
            f"column {col!r} (centered VC on node {node_name!r}) must "
            "hold probabilities in [0, 1]"
        )

live_side

live_side(flow, values, n)

Give the full-data propensity from the flow's own treatment node.

Detached (no gradient reaches the treatment node from this node's loss) and derived from the current parent values, so do-mutilated sampling centers with the intervened t. Training uses the frozen column.

Source code in src/tramdag/modules.py

def live_side(self, flow, values: dict, n: int) -> dict:
    """Give the full-data propensity from the flow's own treatment node.

    Detached (no gradient reaches the treatment node from this node's
    loss) and derived from the current parent values, so
    ``do``-mutilated sampling centers with the intervened ``t``. Training
    uses the frozen column.
    """
    if not self.propensity_col:
        return {}
    p1 = flow._propensity(flow.nodes[self.key], values, n).detach()
    return {self.propensity_col: p1}

extra_columns

extra_columns(flow)

Give the treatment's parents (a treatment cannot be centered itself).

Source code in src/tramdag/modules.py

def extra_columns(self, flow) -> list[str]:
    """Give the treatment's parents (a treatment cannot be centered itself)."""
    return list(flow.nodes[self.key].parents) if self.propensity_col else []

feat_width

feat_width(schema, parents)

Total feature width of the parents (ordinal one-hot, continuous raw).

Source code in src/tramdag/modules.py

def feat_width(schema: dict[str, str | int], parents: tuple[str, ...]) -> int:
    """Total feature width of the parents (ordinal one-hot, continuous raw)."""
    return sum(1 if schema[p] == "continuous" else schema[p] for p in parents)

intercept_module

intercept_module(term, schema, n_params)

Construct the node's intercept module from its I term.

One I term becomes one of three modules: the free theta without parents, one joint net over all parents, or one net per parent with the coefficient vectors summed (allow_interaction=False).

Source code in src/tramdag/modules.py

def intercept_module(
    term: Term, schema: dict[str, str | int], n_params: int
) -> InterceptModule:
    """Construct the node's intercept module from its ``I`` term.

    One ``I`` term becomes one of three modules: the free theta without
    parents, one joint net over all parents, or one net per parent with the
    coefficient vectors summed (``allow_interaction=False``).
    """
    if not term.parents:
        return SimpleInterceptModule(term, schema, n_params)
    if term.allow_interaction:
        return ComplexInterceptModule(term, schema, n_params)
    return AdditiveInterceptModule(term, schema, n_params)