tramdag.fitting

tramdag.fitting

The fitting paths: NodeFitMixin for one node, FitMixin for the flow.

The joint likelihood is a sum of per-node terms with disjoint parameters, so every node can fit on its own. NodeFitMixin.fit is one minibatch Adam loop over one node (validation, verbose printing and the callback hooks included); NodeFitMixin.fit_classical is the float64 full-batch L-BFGS exact-MLE route for a simple intercept with LS terms. FitMixin runs them for every node of the flow, in topological order or in processes; FitMixin.fit_jointly trains all nodes in one loop with one optimizer.

NodeFitMixin

The two fitting paths of one node, mixed into Node.

Source code in src/tramdag/fitting.py

class NodeFitMixin:
    """The two fitting paths of one node, mixed into [`Node`][tramdag.nodes.Node]."""

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

    def _check_side_columns(self, train_df: pd.DataFrame) -> list[str]:
        r"""Check the terms' side columns in the frame; give their names.

        A centered ``VC`` needs its propensity column
        $P(t = 1 \mid \mathrm{pa}_t)$ per row, merged into the frame as an
        ordinary column. ``docs/varying-coefficients.md`` says how to compute
        the column out of fold.
        """
        for m in self.shifts.values():
            for col in m.side_columns():
                if col not in train_df.columns:
                    raise ValueError(
                        f"the centered VC on node {self.name!r} needs its "
                        f"propensity column {col!r} in the training "
                        "frame; compute P(t=1|pa_t) out of fold and "
                        "merge it as a column"
                    )
                m.check_column(self.name, col, train_df[col].to_numpy())
        return self.side_columns()

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

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

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

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

    def fit(
        self,
        train_df: pd.DataFrame,
        *,
        epochs: int,
        learning_rate: float | None = None,
        batch_size: int = 512,
        validation_data: pd.DataFrame | None = None,
        validation_split: float | None = None,
        verbose: int = 0,
        seed: int | None = None,
        marginal_init: bool = False,
        optimizer: torch.optim.Optimizer
        | Callable[[Node], torch.optim.Optimizer]
        | None = None,
        callbacks=None,
    ) -> Node:
        """Fit the node by maximum likelihood with one minibatch loop, Adam by default.

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

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

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

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

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

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

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

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

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

        _train_loop(
            self,
            self.name,
            [self],
            values,
            validate,
            opt,
            cbs,
            epochs=epochs,
            batch_size=batch_size,
            verbose=verbose,
            generator=generator,
        )
        return self

    def _recenter(self, values: dict[str, Tensor]) -> None:
        r"""Run every shift term's post-fit ``finalize`` (the VC re-centering).

        A VC term re-splits $\beta_0$ and $b_\Theta$ so the head sums to zero
        over the train rows; the modelled function does not change.
        """
        feats = self.features(values)
        for m in self.shifts.values():
            m.finalize(self, feats)

    def fit_classical(
        self,
        train_df: pd.DataFrame,
        *,
        max_iter: int = 5000,
        history_size: int = 50,
    ) -> dict:
        """Fit a simple intercept with ``LS`` terms the classical way.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    def _check_classical(self) -> None:
        """Refuse a term the classical fit cannot take."""
        other = sorted({t.name for t in self.node_spec.terms if not t.classical})
        if other:
            raise ValueError(
                "fit_classical requires an all-`ls` spec, that is a simple "
                f"intercept and LS terms only; node {self.name!r} has {other} "
                "terms. Use fit() for flexible models."
            )

side_columns

side_columns()

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

Source code in src/tramdag/fitting.py

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

row_log_prob

row_log_prob(values)

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

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

Source code in src/tramdag/fitting.py

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

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

nll

nll(df)

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

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

Source code in src/tramdag/fitting.py

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

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

fit

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

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

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

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

Parameters:

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

Returns:

Type	Description
Node	self, fitted, in eval mode.

Notes

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

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

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

Raises:

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

Source code in src/tramdag/fitting.py

def fit(
    self,
    train_df: pd.DataFrame,
    *,
    epochs: int,
    learning_rate: float | None = None,
    batch_size: int = 512,
    validation_data: pd.DataFrame | None = None,
    validation_split: float | None = None,
    verbose: int = 0,
    seed: int | None = None,
    marginal_init: bool = False,
    optimizer: torch.optim.Optimizer
    | Callable[[Node], torch.optim.Optimizer]
    | None = None,
    callbacks=None,
) -> Node:
    """Fit the node by maximum likelihood with one minibatch loop, Adam by default.

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

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

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

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

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

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

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

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

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

    _train_loop(
        self,
        self.name,
        [self],
        values,
        validate,
        opt,
        cbs,
        epochs=epochs,
        batch_size=batch_size,
        verbose=verbose,
        generator=generator,
    )
    return self

fit_classical

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

Fit a simple intercept with LS terms the classical way.

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

Parameters:

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

Returns:

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

Warns:

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

Raises:

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

Notes

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

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

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

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

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

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

Source code in src/tramdag/fitting.py

def fit_classical(
    self,
    train_df: pd.DataFrame,
    *,
    max_iter: int = 5000,
    history_size: int = 50,
) -> dict:
    """Fit a simple intercept with ``LS`` terms the classical way.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

FitMixin

The flow's fitting paths, mixed into CausalFlowDAG.

fit and fit_classical loop over the nodes; fit_jointly trains them in one loop. Each node fits itself (Node.fit).

Source code in src/tramdag/fitting.py

class FitMixin:
    """The flow's fitting paths, mixed into [`CausalFlowDAG`][tramdag.CausalFlowDAG].

    `fit` and `fit_classical` loop over the nodes; `fit_jointly` trains them
    in one loop. Each node fits itself ([`Node.fit`][tramdag.nodes.Node.fit]).
    """

    @property
    def history(self) -> dict[str, list]:
        """The nodes' histories, one ``{node: value}`` dict per epoch per key.

        Nodes stop at different epochs. A node that stopped earlier repeats
        its last entry, so a summed curve stays defined. A node without
        entries for a key shows NaN there. ``"val_epoch"`` is the longest
        node's, which is exact while every fit validates. Each node's best
        epoch is in ``node.history["best_epoch"]``.
        """
        out = {}
        for key in ("train", "val", "lr"):
            rows = {name: nd.history.get(key, []) for name, nd in self.nodes.items()}
            if key == "train" or any(rows.values()):
                out[key] = _pad_last(rows)
        if "val" in out:
            out["val_epoch"] = max(
                (nd.history.get("val_epoch", []) for nd in self.nodes.values()),
                key=len,
            )
        return out

    def history_frame(self) -> pd.DataFrame:
        """Give the nodes' histories as one row per node and train epoch.

        Returns
        -------
        pd.DataFrame
            Columns ``node``, ``epoch`` (from 1), ``train``, ``val`` (NaN in an
            epoch without validation) and one ``lr[k]`` per optimizer group
            ``k`` (NaN in a node without that group), and the boolean
            ``best``, true in the epoch that ``EarlyStopping`` restored
            (``node.history["best_epoch"]``) and false everywhere for a node
            without one. The minimum of ``val`` is not the restored epoch
            when ``min_delta`` is above 0. A node that no ``fit`` trained has
            no rows, and a node that stopped early does not repeat its last
            entry.
        """
        frames = []
        for name, nd in self.nodes.items():
            h = nd.history
            epochs = range(1, len(h["train"]) + 1)
            val = dict(zip(h.get("val_epoch", []), h.get("val", []), strict=True))
            frame = pd.DataFrame(
                {
                    "node": name,
                    "epoch": epochs,
                    "train": h["train"],
                    "val": [val.get(e, math.nan) for e in epochs],
                    "best": [e == h.get("best_epoch") for e in epochs],
                }
            )
            lr = pd.DataFrame(h.get("lr", [])).rename(columns="lr[{}]".format)
            frames.append(pd.concat([frame, lr], axis=1))
        return pd.concat(frames, ignore_index=True)

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

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

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

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

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

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

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

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

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

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

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

    def _check_frames(
        self, train_df: pd.DataFrame, validation_data: pd.DataFrame | None
    ) -> None:
        """Check every frame for every node, so a bad one fails before any fit."""
        self._tensorize(train_df)
        if validation_data is not None:
            self._tensorize(validation_data)
        for nd in self.nodes.values():
            nd._check_side_columns(train_df)

    def _fit_nodes(self, train_df: pd.DataFrame, jobs: dict, n_jobs: int) -> None:
        """Run ``Node.fit`` for each job, serially or in forked workers.

        The parent calibrates every node before the fork, so the shapes of the
        returned states match.
        """
        global _FORK_JOB
        if n_jobs == 1 or len(jobs) < 2:
            for name, kwargs in jobs.items():
                self.nodes[name].fit(train_df, **kwargs)
            return
        _FORK_JOB = (self, train_df, jobs)
        try:
            with multiprocessing.get_context("fork").Pool(n_jobs) as pool:
                results = pool.map(_fit_in_child, list(jobs))
        finally:
            _FORK_JOB = None
        for name, (state, history) in zip(jobs, results, strict=True):
            self.nodes[name].load_state_dict(state)
            self.nodes[name].history = history
            self.nodes[name].eval()

    def _with_live_side(self, name: str, df: pd.DataFrame) -> pd.DataFrame:
        """Set the node's side columns in ``df`` to their live values."""
        nd = self.nodes[name]
        cols = dict.fromkeys([*nd.parents, *self._query_side_columns(nd)])
        live = self._side_feats(nd, self._tensorize(df, list(cols)), len(df))
        return df.assign(**{c: live[c].cpu().numpy() for c in nd.side_columns()})

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

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

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

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

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

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

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

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

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

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

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

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

history property

history

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

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

history_frame

history_frame()

Give the nodes' histories as one row per node and train epoch.

Returns:

Type

Description

DataFrame

Columns node, epoch (from 1), train, val (NaN in an epoch without validation) and one lr[k] per optimizer group k (NaN in a node without that group), and the boolean best, true in the epoch that EarlyStopping restored (node.history["best_epoch"]) and false everywhere for a node without one. The minimum of val is not the restored epoch when min_delta is above 0. A node that no fit trained has no rows, and a node that stopped early does not repeat its last entry.

Source code in src/tramdag/fitting.py

def history_frame(self) -> pd.DataFrame:
    """Give the nodes' histories as one row per node and train epoch.

    Returns
    -------
    pd.DataFrame
        Columns ``node``, ``epoch`` (from 1), ``train``, ``val`` (NaN in an
        epoch without validation) and one ``lr[k]`` per optimizer group
        ``k`` (NaN in a node without that group), and the boolean
        ``best``, true in the epoch that ``EarlyStopping`` restored
        (``node.history["best_epoch"]``) and false everywhere for a node
        without one. The minimum of ``val`` is not the restored epoch
        when ``min_delta`` is above 0. A node that no ``fit`` trained has
        no rows, and a node that stopped early does not repeat its last
        entry.
    """
    frames = []
    for name, nd in self.nodes.items():
        h = nd.history
        epochs = range(1, len(h["train"]) + 1)
        val = dict(zip(h.get("val_epoch", []), h.get("val", []), strict=True))
        frame = pd.DataFrame(
            {
                "node": name,
                "epoch": epochs,
                "train": h["train"],
                "val": [val.get(e, math.nan) for e in epochs],
                "best": [e == h.get("best_epoch") for e in epochs],
            }
        )
        lr = pd.DataFrame(h.get("lr", [])).rename(columns="lr[{}]".format)
        frames.append(pd.concat([frame, lr], axis=1))
    return pd.concat(frames, ignore_index=True)

warm_start

warm_start(train_df)

Warm-start every node that has a VC term.

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

Parameters:

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

Returns:

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

Source code in src/tramdag/fitting.py

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

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

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

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

fit

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

Fit every node by maximum likelihood, node by node.

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

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

Parameters:

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

Returns:

Type	Description
CausalFlowDAG	self, fitted, in eval mode.

Notes

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

Raises:

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

Source code in src/tramdag/fitting.py

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

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

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

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

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

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

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

fit_jointly

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

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

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

Parameters:

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

Returns:

Type	Description
CausalFlowDAG	self, fitted, in eval mode.

Raises:

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

Source code in src/tramdag/fitting.py

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

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

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

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

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

fit_classical

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

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

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

Parameters:

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

Returns:

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

Notes

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

Warns:

Type	Description
UserWarning	For each node that did not converge.

Raises:

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

Source code in src/tramdag/fitting.py

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

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

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

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

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

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

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