Source code for setu.io

"""Serialization API for saving and loading trained estimators."""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING, TypeVar

import equinox as eqx

if TYPE_CHECKING:
    from setu.discrete_nle import DiscreteNLE
    from setu.mixed_nle import MixedNLE
    from setu.nle import NLE

T = TypeVar("T")




def save_estimator(estimator: NLE | MixedNLE | DiscreteNLE, path: str | Path) -> None:
    """Save a trained NLE, MixedNLE, or DiscreteNLE to disk.

    Uses equinox serialization to save all array leaves of the estimator pytree.

    Args:
        estimator: Trained NLE or MixedNLE instance.
        path: File path to save to (e.g., "model.eqx").

    Example:
        >>> nle = create_nle(x_dim=3, theta_dim=3, key=key)
        >>> result = fit_nle(nle, dataset, key=key)
        >>> save_estimator(result.nle, "trained_nle.eqx")
    """
    eqx.tree_serialise_leaves(str(path), estimator)





def load_estimator(path: str | Path, like: T) -> T:
    """Load a trained NLE, MixedNLE, or DiscreteNLE from disk.

    Uses equinox deserialization. The `like` parameter provides the pytree
    structure for deserialization -- it must be an estimator with the **exact
    same structure** as the saved model. This means:

    - Same architecture (x_dim, theta_dim, flow_type, hidden_dims, n_layers)
    - Same fitted state: if the saved model was trained (has transforms fitted),
      `like` must also have fitted transforms. The simplest way to ensure this
      is to train a model with the same architecture on dummy data, or to keep
      a copy of the untrained-but-fitted template.

    For convenience, use the pattern: save the trained estimator, and when
    loading, create a new estimator with the same architecture and train it
    briefly (even 1 epoch) to match the pytree structure.

    Args:
        path: File path to load from.
        like: An estimator with matching pytree structure. Must have the same
            hyperparameters and fitted state as the saved model.

    Returns:
        Estimator with weights loaded from disk.

    Example:
        >>> # Save after training
        >>> save_estimator(result.nle, "trained_nle.eqx")
        >>>
        >>> # Load: create matching template (train briefly on dummy data)
        >>> template_nle = create_nle(x_dim=3, theta_dim=3, key=key)
        >>> template_result = fit_nle(
        ...     template_nle, dummy_dataset, max_epochs=1, key=key
        ... )
        >>> loaded = load_estimator("trained_nle.eqx", like=template_result.nle)
    """
    return eqx.tree_deserialise_leaves(str(path), like)