Skip to content

tramdag.readouts

tramdag.readouts

Read-outs of a fitted flow: ReadoutsMixin, composed into CausalFlowDAG.

Coefficients (ls_coefficients), effect curves (shift_curve, varying_coef), the intercept decomposition (intercept_contributions), the meta-adjacency view (to_matrix) and the classical design matrix (design_matrix).

ReadoutsMixin

The stateless read-outs, mixed into CausalFlowDAG.

Source code in src/tramdag/readouts.py
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
class ReadoutsMixin:
    """The stateless read-outs, mixed into [`CausalFlowDAG`][tramdag.CausalFlowDAG]."""

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

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

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

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

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

    def _vc_term(self, node: str, t: str | None) -> tuple:
        """Find the node, treatment name and modifiers of one VC term."""
        nd = self._node(node)
        vcs = {
            m.key: m.mods
            for m in nd.shifts.values()
            if isinstance(m, VaryingCoefficientModule)
        }
        if not vcs:
            raise ValueError(f"node {node!r} has no VC term")
        if t is None:
            if len(vcs) > 1:
                raise ValueError(
                    f"node {node!r} has several VC terms ({sorted(vcs)}); "
                    "pass t=<treatment name>"
                )
            t = next(iter(vcs))
        if t not in vcs:
            raise KeyError(f"node {node!r} has no VC term on {t!r} (has {sorted(vcs)})")
        return nd, t, vcs[t]

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

shift_curve

shift_curve(node, parent, grid)

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

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

Parameters:

Name Type Description Default
node str

Name of the node that carries the term.

required
parent str

The term's key in the node's shifts: the parent's name.

required
grid array - like

Parent values at which to evaluate the term, shape (m,).

required

Returns:

Type Description
ndarray

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

Raises:

Type Description
KeyError

If node is unknown, or the node has no shift term keyed parent.

ValueError

If parent is an ordinal node.

Source code in src/tramdag/readouts.py
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
@torch.no_grad()
def shift_curve(self, node: str, parent: str, grid) -> np.ndarray:
    """Evaluate one fitted shift term on a 1-D grid of parent values.

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

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

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

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

varying_coef

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

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

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

Parameters:

Name Type Description Default
df DataFrame

Rows at which to evaluate \(\beta\). Must contain every modifier column of the term.

required
node str

Name of the node that carries the VC term.

required
t str | None

Treatment name of the VC term. Optional when the node has exactly one VC term.

None

Returns:

Type Description
ndarray

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

Raises:

Type Description
KeyError

If node is unknown, if the node has no VC term on t, or if a modifier column is missing from df.

ValueError

If the node has no VC term, or if t is omitted while the node has several VC terms.

Source code in src/tramdag/readouts.py
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
@torch.no_grad()
def varying_coef(
    self, df: pd.DataFrame, node: str, *, t: str | None = None
) -> np.ndarray:
    r"""Evaluate the fitted effect function $\beta(x)$ of a ``VC`` term.

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

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

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

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

ls_coefficients

ls_coefficients()

Give the per-node linear-shift weights.

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

Returns:

Type Description
dict[str, dict[str, ndarray]]

The weights, as {node: {parent: array}}. A node without linear-shift terms is absent.

Source code in src/tramdag/readouts.py
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
@torch.no_grad()
def ls_coefficients(self) -> dict[str, dict[str, np.ndarray]]:
    """Give the per-node linear-shift weights.

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

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

to_matrix

to_matrix()

Give the labeled adjacency matrix of term effects.

Returns:

Type Description
DataFrame

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

Source code in src/tramdag/readouts.py
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
def to_matrix(self) -> pd.DataFrame:
    """Give the labeled adjacency matrix of term effects.

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

intercept_contributions

intercept_contributions(df, node)

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

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

Parameters:

Name Type Description Default
df DataFrame

Rows over which to center and at which to evaluate the contributions. Must contain every intercept-parent column.

required
node str

Name of a node whose intercept term has parents.

required

Returns:

Type Description
dict

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

Raises:

Type Description
KeyError

If node is unknown, or if an intercept-parent column is missing from df.

ValueError

If the node's intercept term has no parents.

Notes

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

Source code in src/tramdag/readouts.py
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
@torch.no_grad()
def intercept_contributions(self, df: pd.DataFrame, node: str) -> dict:
    r"""Decompose a complex intercept into mean-centered per-parent parts.

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

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

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

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

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

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

design_matrix

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

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

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

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

Parameters:

Name Type Description Default
df DataFrame

Rows to encode. Must contain every parent column of node.

required
node str

Name of the node whose parents are encoded.

required
drop_first bool

Drop each ordinal parent's level-0 column, by default False.

False

Returns:

Type Description
DataFrame

One column per encoded feature, indexed like df.

Raises:

Type Description
KeyError

If node is unknown, or df lacks a parent column.

ValueError

If an ordinal parent value is not a level index.

Source code in src/tramdag/readouts.py
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
@torch.no_grad()
def design_matrix(
    self, df: pd.DataFrame, node: str, *, drop_first: bool = False
) -> pd.DataFrame:
    """Encode a node's parents the way the flow feeds them to its shifts.

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

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

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

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

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