tramdag.plots

tramdag.plots

Figures of a TRAM-DAG: the DAG, the marginals, the training curve, a VC's beta.

matplotlib is an optional dependency: pip install "tramdag[plots]". It is imported on the first call, so importing tramdag never needs it.

plot_dag

plot_dag(
    spec_or_flow,
    *,
    labels=True,
    legend=True,
    node_kind=True,
    modifiers="edge",
    style=None,
    ax=None,
    path=None,
    title=None
)

Draw the labelled DAG of a spec (or of a fitted flow).

Layers run left to right, a node one layer past its parents. Continuous nodes are ellipses, ordinal nodes rounded boxes with their level count. Each edge is drawn by the term that owns it: LS thin gray, CS thick blue, a complex intercept (CI) dashed orange, a VC treatment edge red; a multi-parent CS/CI is labelled joint. A VC modifier is a dotted red arrow onto the treatment edge it modifies. A term without an entry in EDGE_STYLE draws dotted gray. EDGE_LABEL maps the to_matrix tag VCm to the label VC modifier.

The layout units are NODE_H for the node height and ROW_DY for the distance between rows. A new figure measures INCH_PER_UNIT inches per unit, and the fonts fit it. An edge that skips a layer bends out by BULGE per skipped layer. Given axes shrink the fonts with the space they have per layout unit. A modifier points at the midpoint of its treatment edge, a position of its own.

Parameters:

Name	Type	Description	Default
spec_or_flow	dict[str, NodeSpec] | CausalFlowDAG	The DAG to draw. A flow draws its spec.	required
labels	bool	Write the term name on each edge, by default True.	True
legend	bool	Add a legend of the terms used, by default True.	True
node_kind	bool	Write the node kind under each name, by default True.	True
modifiers	str	"edge" (default) points a VC modifier at the midpoint of its treatment edge; "node" draws it as an edge into the outcome node.	'edge'
style	str | dict | None	"light" (default) or "dark", or a dict over the light preset with any of the keys of tramdag.plots.STYLES["light"]: text, muted, node_edge, node_face (per kind), edge (per term), label_box (None for no box) and font_scale.	None
ax	Axes | None	Draw into this axes; by default a new figure sized to the layout.	None
path	str | Path | None	Save the figure here (150 dpi) after drawing, by default None (not saved).	None
title	str | None	Axes title, by default none.	None

Returns:

Type	Description
Axes	The axes drawn into.

Raises:

Type	Description
ValueError	If the spec has no node.

Source code in src/tramdag/plots.py

def plot_dag(
    spec_or_flow,
    *,
    labels: bool = True,
    legend: bool = True,
    node_kind: bool = True,
    modifiers: str = "edge",
    style=None,
    ax=None,
    path=None,
    title=None,
):
    """Draw the labelled DAG of a spec (or of a fitted flow).

    Layers run left to right, a node one layer past its parents. Continuous
    nodes are ellipses, ordinal nodes rounded boxes with their level count.
    Each edge is drawn by the term that owns it: ``LS`` thin gray, ``CS`` thick
    blue, a complex intercept (``CI``) dashed orange, a ``VC`` treatment edge
    red; a multi-parent CS/CI is labelled ``joint``. A VC modifier is a dotted
    red arrow onto the treatment edge it modifies. A term without an entry in
    ``EDGE_STYLE`` draws dotted gray. ``EDGE_LABEL`` maps the ``to_matrix`` tag
    ``VCm`` to the label ``VC modifier``.

    The layout units are ``NODE_H`` for the node height and ``ROW_DY`` for the
    distance between rows. A new figure measures ``INCH_PER_UNIT`` inches per
    unit, and the fonts fit it. An edge that skips a layer bends out by
    ``BULGE`` per skipped layer. Given axes shrink the fonts with the space
    they have per layout unit. A modifier points at the midpoint of its
    treatment edge, a position of its own.

    Parameters
    ----------
    spec_or_flow : dict[str, NodeSpec] | CausalFlowDAG
        The DAG to draw. A flow draws its ``spec``.
    labels : bool, optional
        Write the term name on each edge, by default True.
    legend : bool, optional
        Add a legend of the terms used, by default True.
    node_kind : bool, optional
        Write the node kind under each name, by default True.
    modifiers : str, optional
        ``"edge"`` (default) points a VC modifier at the midpoint of its
        treatment edge; ``"node"`` draws it as an edge into the outcome node.
    style : str | dict | None, optional
        ``"light"`` (default) or ``"dark"``, or a dict over the light preset
        with any of the keys of ``tramdag.plots.STYLES["light"]``: ``text``,
        ``muted``, ``node_edge``, ``node_face`` (per kind), ``edge`` (per
        term), ``label_box`` (``None`` for no box) and ``font_scale``.
    ax : matplotlib.axes.Axes | None, optional
        Draw into this axes; by default a new figure sized to the layout.
    path : str | Path | None, optional
        Save the figure here (150 dpi) after drawing, by default None (not
        saved).
    title : str | None, optional
        Axes title, by default none.

    Returns
    -------
    matplotlib.axes.Axes
        The axes drawn into.

    Raises
    ------
    ValueError
        If the spec has no node.
    """
    plt = _plt()
    spec = getattr(spec_or_flow, "spec", spec_or_flow)
    if not spec:
        raise ValueError("plot_dag needs a spec with at least one node")
    st = _style(style)
    pos, layer_dx = _layout(spec)
    edges, mods = _dag_edges(spec, modifiers)
    bulges = _bulges(spec, pos, layer_dx, edges)
    pos_all = pos | {
        (mod, t_edge): tuple(_arc(pos, t_edge, bulges[t_edge], [0.5])[0])
        for mod, t_edge in mods
    }
    mod_bulges = {
        (mod, (mod, t_edge)): _clear(spec, pos_all, (mod, (mod, t_edge)), 0.0)
        for mod, t_edge in mods
    }
    x_lo, x_hi, y_lo, y_hi = _limits(spec, pos_all, bulges | mod_bulges)
    created = ax is None
    if created:
        _, ax = plt.subplots(
            figsize=(
                INCH_PER_UNIT * (x_hi - x_lo),
                INCH_PER_UNIT * (y_hi - y_lo) + 0.5,
            )
        )
    else:
        box = ax.get_position()
        w, h = ax.figure.get_size_inches()
        per_unit = min(box.width * w / (x_hi - x_lo), box.height * h / (y_hi - y_lo))
        st = st | dict(font_scale=st["font_scale"] * per_unit / INCH_PER_UNIT)
    patches = {
        name: _draw_node(ax, name, spec[name], xy, st, node_kind)
        for name, xy in pos.items()
    }
    for edge, bulge in bulges.items():
        _draw_edge(ax, patches, spec, pos, edge, labels, bulge, st)
    for edge, bulge in mod_bulges.items():
        _draw_modifier(ax, patches, pos_all, edge, bulge, st)
    if legend:
        names = {e for _, _, e, _ in bulges} | ({"VC modifier"} if mods else set())
        _legend(ax, names, st)
    ax.set_xlim(x_lo, x_hi)
    ax.set_ylim(y_lo, y_hi)
    ax.set_aspect("equal")
    ax.set_axis_off()
    if title:
        ax.set_title(title, color=st["text"])
    _finish(ax.figure, path, created)
    return ax

plot_marginals

plot_marginals(
    flow,
    df,
    *,
    ncols=3,
    seed=None,
    colors=("C0", "C1"),
    legend="axes",
    ax=None,
    path=None,
    title="observed vs sampled marginals"
)

Observed vs sampled marginal of every node, one panel each.

Ordinal nodes compare level proportions side by side; continuous nodes a density histogram of the data with the flow's sample as a step outline. The sample has as many rows as df.

Parameters:

Name	Type	Description	Default
flow	CausalFlowDAG	The fitted flow.	required
df	DataFrame	The data to compare against (the validation split, typically).	required
ncols	int	Panels per row of a new figure, by default 3.	3
seed	int | None	Seed of the flow's sample, by default None.	None
colors	tuple[str, str]	Colours of the data and of the flow's sample, by default ("C0", "C1").	('C0', 'C1')
legend	str | None	"axes" (default) puts a legend in every panel, "figure" one legend for the figure, None none.	'axes'
ax	sequence of matplotlib.axes.Axes | None	Draw into these panels, one per node in the flow's order; by default a new figure.	None
path	str | Path | None	Save the figure here (150 dpi) after drawing, by default None (not saved).	None
title	str | None	Title of a new figure, by default "observed vs sampled marginals"; None for none. A figure passed through ax keeps its own title.	'observed vs sampled marginals'

Returns:

Type	Description
numpy.ndarray of matplotlib.axes.Axes	The panels, in the flow's node order (unused panels of a new figure are switched off).

Source code in src/tramdag/plots.py

def plot_marginals(
    flow,
    df: pd.DataFrame,
    *,
    ncols: int = 3,
    seed=None,
    colors=("C0", "C1"),
    legend: str | None = "axes",
    ax=None,
    path=None,
    title="observed vs sampled marginals",
):
    """Observed vs sampled marginal of every node, one panel each.

    Ordinal nodes compare level proportions side by side; continuous nodes a
    density histogram of the data with the flow's sample as a step outline.
    The sample has as many rows as ``df``.

    Parameters
    ----------
    flow : CausalFlowDAG
        The fitted flow.
    df : pd.DataFrame
        The data to compare against (the validation split, typically).
    ncols : int, optional
        Panels per row of a new figure, by default 3.
    seed : int | None, optional
        Seed of the flow's sample, by default None.
    colors : tuple[str, str], optional
        Colours of the data and of the flow's sample, by default
        ``("C0", "C1")``.
    legend : str | None, optional
        ``"axes"`` (default) puts a legend in every panel, ``"figure"`` one
        legend for the figure, ``None`` none.
    ax : sequence of matplotlib.axes.Axes | None, optional
        Draw into these panels, one per node in the flow's order; by default
        a new figure.
    path : str | Path | None, optional
        Save the figure here (150 dpi) after drawing, by default None (not
        saved).
    title : str | None, optional
        Title of a new figure, by default ``"observed vs sampled
        marginals"``; ``None`` for none. A figure passed through ``ax`` keeps
        its own title.

    Returns
    -------
    numpy.ndarray of matplotlib.axes.Axes
        The panels, in the flow's node order (unused panels of a new figure
        are switched off).
    """
    plt = _plt()
    sample = flow.sample(len(df), seed=seed)
    created = ax is None
    if created:
        nrows = -(-len(flow.order) // ncols)
        fig, axes = plt.subplots(
            nrows, ncols, figsize=(3.6 * ncols, 2.8 * nrows), squeeze=False
        )
        for unused in axes.flat[len(flow.order) :]:
            unused.set_axis_off()
    else:
        axes = np.asarray(ax, dtype=object)
        fig = axes.flat[0].figure
    for panel, name in zip(axes.flat, flow.order, strict=False):
        _marginal_panel(panel, name, flow.spec[name], df, sample, colors)
        panel.set_title(name)
        if legend == "axes":
            panel.legend(fontsize=8, frameon=False)
    if legend == "figure":
        fig.legend(
            *axes.flat[0].get_legend_handles_labels(),
            loc="upper center",
            bbox_to_anchor=(0.5, 0.0),
            ncol=2,
            frameon=False,
        )
    if title and created:
        fig.suptitle(title)
    _finish(fig, path, created)
    return axes

plot_training

plot_training(
    flow, *, stops=False, ax=None, path=None, title=None
)

Draw the summed train and validation NLL per epoch.

flow.history accumulates across fit calls, so the curves cover every epoch the flow has trained, not only the last call. Each validation entry is drawn at the epoch it was measured in, which is what keeps the two curves aligned when one fit validated and another did not: an unvalidated fit in between leaves a gap instead of shifting the curve back to epoch 1.

The y axis zooms past the initial drop: its top keeps each curve after 10 % of its epochs, so a validation curve that rises later stays in view. The limits use finite values only, so a diverged epoch does not blow them up. The legend sits below the axes, off the curves, and the title stays left of the stop names. The stop marks come after the zoom, so they hang from the visible top.

Parameters:

Name	Type	Description	Default
flow	CausalFlowDAG	The fitted flow; flow.history is read.	required
stops	bool	Mark each node's stop epoch, the length of its history["train"], with a dashed line, and its history["best_epoch"], where a fit recorded one, with a triangle at the top. By default no marks. The triangles come from the last fit that ran EarlyStopping; a fit without it removes history["best_epoch"].	False
ax	Axes | None	Draw into this axes; by default a new figure.	None
path	str | Path | None	Save the figure here (150 dpi) after drawing, by default None (not saved).	None
title	str | None	Axes title, by default "training".	None

Returns:

Type	Description
Axes	The axes drawn into.

Raises:

Type	Description
TypeError	If stops is not a bool.
ValueError	If flow.history is empty.

Source code in src/tramdag/plots.py

def plot_training(flow, *, stops=False, ax=None, path=None, title=None):
    """Draw the summed train and validation NLL per epoch.

    ``flow.history`` accumulates across ``fit`` calls, so the curves cover
    every epoch the flow has trained, not only the last call. Each validation
    entry is drawn at the epoch it was measured in, which is what keeps the two
    curves aligned when one ``fit`` validated and another did not: an
    unvalidated fit in between leaves a gap instead of shifting the curve back
    to epoch 1.

    The y axis zooms past the initial drop: its top keeps each curve after 10 %
    of its epochs, so a validation curve that rises later stays in view. The
    limits use finite values only, so a diverged epoch does not blow them up.
    The legend sits below the axes, off the curves, and the title stays left of
    the stop names. The stop marks come after the zoom, so they hang
    from the visible top.

    Parameters
    ----------
    flow : CausalFlowDAG
        The fitted flow; ``flow.history`` is read.
    stops : bool, optional
        Mark each node's stop epoch, the length of its ``history["train"]``,
        with a dashed line, and its ``history["best_epoch"]``, where a fit
        recorded one, with a triangle at the top. By default no marks. The
        triangles come from the last fit that ran ``EarlyStopping``; a fit
        without it removes ``history["best_epoch"]``.
    ax : matplotlib.axes.Axes | None, optional
        Draw into this axes; by default a new figure.
    path : str | Path | None, optional
        Save the figure here (150 dpi) after drawing, by default None (not
        saved).
    title : str | None, optional
        Axes title, by default ``"training"``.

    Returns
    -------
    matplotlib.axes.Axes
        The axes drawn into.

    Raises
    ------
    TypeError
        If ``stops`` is not a bool.
    ValueError
        If ``flow.history`` is empty.
    """
    if not isinstance(stops, bool):
        raise TypeError(f"stops must be a bool, got {stops!r}")
    plt = _plt()
    hist = flow.history
    if not hist.get("train"):
        raise ValueError("plot_training needs a fitted flow; its history is empty")
    train = np.array([sum(d.values()) for d in hist["train"]])
    curves = {"train": (np.arange(1, len(train) + 1), train)}
    if hist.get("val"):
        val = np.array([sum(d.values()) for d in hist["val"]])
        curves["val"] = (np.asarray(hist["val_epoch"], dtype=float), val)
    created = ax is None
    if created:
        _, ax = plt.subplots(figsize=(7.5, 3.6))
    for label, (x, curve) in curves.items():
        ax.plot(x, curve, label=f"{label} NLL (total)")
    lo = min(c[np.isfinite(c)].min() for _, c in curves.values())
    hi = max(
        c[len(c) // 10 :][np.isfinite(c[len(c) // 10 :])].max()
        for _, c in curves.values()
    )
    if hi > lo:
        ax.set_ylim(lo - 0.05 * (hi - lo), hi)
    if stops:
        _mark_stops(ax, flow)
    ax.set_xlabel("epoch")
    ax.set_ylabel("NLL")
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.18), ncol=2, frameon=False)
    ax.set_title(title or "training", loc="left")
    _finish(ax.figure, path, created)
    return ax

plot_varying_coef

plot_varying_coef(
    flow, df, node, *, by, t=None, ax=None, path=None
)

Draw a VC term's effect \(\beta(x)\) along one modifier.

The grid spans by over its range in df, or its levels for an ordinal by; every other node column is held at an observed median in df, so an ordinal value stays a level. The values come from varying_coef, on the node's latent scale. A dotted line marks \(\beta = 0\).

Parameters:

Name	Type	Description	Default
flow	CausalFlowDAG	The fitted flow.	required
df	DataFrame	Rows that give the range of by and the medians of the other modifiers.	required
node	str	The node that carries the VC term.	required
by	str	The modifier along which to draw.	required
t	str | None	The treatment of the VC term; optional when the node has one.	None
ax	Axes | None	Draw into this axes; by default a new figure.	None
path	str | Path | None	Save the figure here (150 dpi) after drawing, by default None (not saved).	None

Returns:

Type	Description
Axes	The axes drawn into.

Raises:

Type	Description
ValueError	If by is not a modifier of the node's VC term.

Source code in src/tramdag/plots.py

def plot_varying_coef(
    flow, df: pd.DataFrame, node: str, *, by: str, t=None, ax=None, path=None
):
    r"""Draw a VC term's effect $\beta(x)$ along one modifier.

    The grid spans ``by`` over its range in ``df``, or its levels for an
    ordinal ``by``; every other node column is held at an observed median
    in ``df``, so an ordinal value stays a level. The values come from
    [`varying_coef`][tramdag.flow.CausalFlowDAG.varying_coef], on the node's
    latent scale. A dotted line marks $\beta = 0$.

    Parameters
    ----------
    flow : CausalFlowDAG
        The fitted flow.
    df : pd.DataFrame
        Rows that give the range of ``by`` and the medians of the other
        modifiers.
    node : str
        The node that carries the VC term.
    by : str
        The modifier along which to draw.
    t : str | None, optional
        The treatment of the VC term; optional when the node has one.
    ax : matplotlib.axes.Axes | None, optional
        Draw into this axes; by default a new figure.
    path : str | Path | None, optional
        Save the figure here (150 dpi) after drawing, by default None (not
        saved).

    Returns
    -------
    matplotlib.axes.Axes
        The axes drawn into.

    Raises
    ------
    ValueError
        If ``by`` is not a modifier of the node's VC term.
    """
    plt = _plt()
    _, t_name, mods = flow._vc_term(node, t)
    if by not in mods:
        raise ValueError(
            f"{by!r} is not a modifier of the VC term on {t_name!r} of node "
            f"{node!r} (has {list(mods)})"
        )
    if flow.nodes[by].kind == "ordinal":
        grid = np.unique(df[by])
    else:
        grid = np.linspace(df[by].min(), df[by].max(), 200)
    nodes = [c for c in flow.order if c in df.columns]
    medians = {c: df[c].quantile(0.5, interpolation="nearest") for c in nodes}
    rows = pd.DataFrame({c: np.full(len(grid), m) for c, m in medians.items()})
    rows[by] = grid
    beta = flow.varying_coef(rows, node, t=t)
    created = ax is None
    if created:
        _, ax = plt.subplots(figsize=(5, 3.4))
    ax.plot(grid, beta)
    ax.axhline(0.0, ls=":", lw=1, color="gray")
    ax.set_xlabel(by)
    ax.set_ylabel(r"$\beta$" + f"({by})")
    _finish(ax.figure, path, created)
    return ax