BaseEstimator

The model contract that every estimator implements, independent of the network framework.

Overview

An estimator node runs the same model class in two Ray actors:

• the train actor (<node>/train) drives it through falcon.core.round_trainer.RoundTrainer, which owns the round loop, the round counters, the best state and the checkpoint;
• the sample actor (<node>) drives it through falcon.core.model_sampler.ModelSampler, which installs the best state and falls back to the prior until one exists.

The model keeps its networks as mutable state on self and knows nothing about rounds, publishing or Ray. The two sides exchange the best state as a state tree: a nested dict of numpy arrays and plain values (falcon.core.state_io). The same tree is written to graph/<node>/best_state.npz.

Methods, grouped by who calls them:

Called by	Methods
both	setup, groups, built, build, import_state
train actor	init_from_batch, train_step, evaluate, discard_mask, on_round_start, on_train_start, on_validation_end, on_round_end, set_observations, snapshot, restore, export_state, export_proposal
sample actor	sample_prior, sample
both, optional	import_proposal

Rules that let any framework implement the contract:

• snapshot() / restore() are fast in-memory copies with opaque handles;
• export_state() returns copies as numpy; nothing outside the model keeps references to the model's own arrays;
• sampling draws its randomness from the rng argument.

A model whose proposal is more than its networks keeps a proposal state. FlowMatching, for example, keeps its region ladder there. The optional hooks for it are all no-ops by default:

• on_train_start(batches) receives the round's training set after the best network's baseline validation, before the first epoch;
• on_round_end(promoted) runs after the acceptance test, with the best state installed, and returns True when the proposal state changed;
• set_observations(conditions) gives the train actor the node's conditions at the observation;
• export_proposal(full) / import_proposal(tree) move the proposal state as a state tree. The trainer publishes the sampling subset (full=False) whenever it changes and saves the full state in best_state.npz.

Torch estimators derive from falcon.estimators.TorchModel, which implements the state handling for nn.Module network groups; Flow, FlowMatching and GaussianFullCov are examples.

Writing a JAX estimator

JAX functions are pure, so a JAX model keeps its latest state on self and replaces it after each compiled step. With buffer donation the step reuses the old buffers, so memory does not double; old arrays then become invalid, and snapshot() must copy:

import jax, jax.numpy as jnp, numpy as np, optax
from falcon.core.base_estimator import BaseEstimator

class JaxGaussian(BaseEstimator):
    def build(self, init_tree):
        self.params = init_params(init_tree)              # pytree on the device
        self.opt = optax.adamw(self.lr)
        self.opt_state = self.opt.init(self.params)
        self._step = jax.jit(make_step(self.opt), donate_argnums=(0, 1))

    def groups(self):
        return {"model": "loss"}

    @property
    def built(self):
        return hasattr(self, "params")

    def train_step(self, batch):
        batch = {k: jnp.asarray(v.cpu().numpy()) for k, v in batch.items()}
        self.params, self.opt_state, loss = self._step(self.params, self.opt_state, batch)
        return {"loss": float(loss)}

    def snapshot(self, group):
        return jax.tree.map(jnp.copy, self.params)       # donation invalidates old arrays

    def restore(self, group, handle):
        self.params = jax.tree.map(jnp.copy, handle)

    def export_state(self, group):
        return {"params": flatten_to_numpy(self.params)}   # nested dict of np.ndarray

    def import_state(self, group, tree):
        self.params = unflatten_from_numpy(tree["params"])

    def sample(self, rng, num_samples, conditions, mode):
        key = jax.random.key(int(rng.integers(2**63 - 1)))
        ...

Set XLA_PYTHON_CLIENT_PREALLOCATE=false through the node's ray.runtime_env.env_vars when a JAX actor shares its GPU with other actors; JAX otherwise reserves most of the GPU memory at start-up.

Class Reference

BaseEstimator

Bases: ABC

Model of one estimator family, used unchanged by both actors of a node.

The train actor (<node>/train) drives it through RoundTrainer, the sample actor (<node>) through ModelSampler. The model holds its networks as mutable state on self and knows nothing about rounds, publishing or Ray; each actor calls only the methods it needs.

Lifecycle:

1. __init__: pure config storage (flat YAML kwargs). The round-loop parameters are read from attributes named like the RoundConfig fields.
2. setup(): runtime wiring inside the actor (simulator, keys, device).
3. build(init_tree): creates the networks. The trainer gets the tree from init_from_batch() on its first batch; the sampler gets it from the published state.

State rules, so that any framework can implement this:

• snapshot() / restore() keep and restore fast in-memory copies of a network group; the handles are opaque.
• export_state() returns copies as numpy trees; import_state() installs such a tree. Nothing outside the model keeps references to the model's own arrays.
• Randomness in sampling comes from the rng argument.

Example::

graph.add_node("z", estimator=Flow(max_epochs=300, net_type="nsf"))

cache_device property

cache_device

Device for the training data cache (None = CPU).

history property

history

Model-specific training history, saved next to the loop's history.

built abstractmethod property

built

Whether build() has been called.

setup

setup(simulator_instance, theta_key, condition_keys)

Wire up runtime components inside the actor.

Parameters:

Name	Type	Description	Default
simulator_instance		Live prior/simulator of the node.	required
theta_key	str	Name of the node being estimated.	required
condition_keys	List[str]	Evidence and scaffold node names.	required

Source code in src/falcon/core/base_estimator.py

def setup(self, simulator_instance, theta_key: str, condition_keys: List[str]) -> None:
    """Wire up runtime components inside the actor.

    Args:
        simulator_instance: Live prior/simulator of the node.
        theta_key: Name of the node being estimated.
        condition_keys: Evidence and scaffold node names.
    """
    self.simulator_instance = simulator_instance
    self.theta_key = theta_key
    self.condition_keys = list(condition_keys or [])

set_num_threads

set_num_threads(num_threads)

Limit the CPU threads of the network framework in this process.

Called by the actor, whose CPU budget is shared with the node's other actor. Frameworks that honour OMP_NUM_THREADS at import need nothing here.

Source code in src/falcon/core/base_estimator.py

def set_num_threads(self, num_threads: int) -> None:
    """Limit the CPU threads of the network framework in this process.

    Called by the actor, whose CPU budget is shared with the node's other
    actor. Frameworks that honour ``OMP_NUM_THREADS`` at import need
    nothing here.
    """

batch_keys

batch_keys()

Buffer keys the training batches must contain.

Source code in src/falcon/core/base_estimator.py

def batch_keys(self) -> List[str]:
    """Buffer keys the training batches must contain."""
    return [f"{self.theta_key}.value", f"{self.theta_key}.log_prob",
            *[f"{k}.value" for k in self.condition_keys]]

groups abstractmethod

groups()

Network groups mapped to their validation metric; the first is primary.

Groups are compared and promoted independently; a round is accepted when the primary group improves.

Source code in src/falcon/core/base_estimator.py

@abstractmethod
def groups(self) -> Dict[str, str]:
    """Network groups mapped to their validation metric; the first is primary.

    Groups are compared and promoted independently; a round is accepted
    when the primary group improves.
    """

init_from_batch abstractmethod

init_from_batch(batch)

What build() needs, taken from the first training batch.

Source code in src/falcon/core/base_estimator.py

@abstractmethod
def init_from_batch(self, batch) -> StateTree:
    """What ``build()`` needs, taken from the first training batch."""

build abstractmethod

build(init_tree)

Create the networks.

Source code in src/falcon/core/base_estimator.py

@abstractmethod
def build(self, init_tree: StateTree) -> None:
    """Create the networks."""

snapshot abstractmethod

snapshot(group)

Opaque in-memory copy of a group's current state.

Source code in src/falcon/core/base_estimator.py

@abstractmethod
def snapshot(self, group: str) -> Any:
    """Opaque in-memory copy of a group's current state."""

restore abstractmethod

restore(group, handle)

Restore a group from snapshot().

Source code in src/falcon/core/base_estimator.py

@abstractmethod
def restore(self, group: str, handle: Any) -> None:
    """Restore a group from ``snapshot()``."""

export_state abstractmethod

export_state(group)

Copy of a group's state as a numpy tree.

Source code in src/falcon/core/base_estimator.py

@abstractmethod
def export_state(self, group: str) -> StateTree:
    """Copy of a group's state as a numpy tree."""

import_state abstractmethod

import_state(group, tree)

Install a group's state from a numpy tree.

Source code in src/falcon/core/base_estimator.py

@abstractmethod
def import_state(self, group: str, tree: StateTree) -> None:
    """Install a group's state from a numpy tree."""

load_legacy

load_legacy(node_dir)

Read a checkpoint written before best_state.npz existed.

Returns a tree in the published format, or None.

Source code in src/falcon/core/base_estimator.py

def load_legacy(self, node_dir: Path) -> Optional[StateTree]:
    """Read a checkpoint written before ``best_state.npz`` existed.

    Returns a tree in the published format, or None.
    """
    return None

train_step abstractmethod

train_step(batch)

One optimizer update on one training batch; returns metrics incl. "loss".

Source code in src/falcon/core/base_estimator.py

@abstractmethod
def train_step(self, batch) -> Dict[str, float]:
    """One optimizer update on one training batch; returns metrics incl. "loss"."""

evaluate abstractmethod

evaluate(batch)

Validation metrics for one batch, without updating anything.

Must include "loss" (early stopping) and the metric of every group (acceptance test).

Source code in src/falcon/core/base_estimator.py

@abstractmethod
def evaluate(self, batch) -> Dict[str, float]:
    """Validation metrics for one batch, without updating anything.

    Must include "loss" (early stopping) and the metric of every group
    (acceptance test).
    """

discard_mask abstractmethod

discard_mask(batch)

Discard test for one batch with the primary group; True = discard.

Called after an accepted round, when the primary group equals the best network.

Source code in src/falcon/core/base_estimator.py

@abstractmethod
def discard_mask(self, batch) -> np.ndarray:
    """Discard test for one batch with the primary group; True = discard.

    Called after an accepted round, when the primary group equals the
    best network.
    """

on_round_start

on_round_start()

Called at the start of every round after the first, once the best state has been restored (e.g. reset the learning rate).

Source code in src/falcon/core/base_estimator.py

def on_round_start(self) -> None:
    """Called at the start of every round after the first, once the best
    state has been restored (e.g. reset the learning rate)."""

on_validation_end

on_validation_end(epoch, val_metrics)

Called after each validation within a round (e.g. LR scheduler step).

Returns:

Type	Description
Optional[Dict[str, float]]	Optional dict of extra metrics for the epoch summary line.

Source code in src/falcon/core/base_estimator.py

def on_validation_end(self, epoch: int, val_metrics: Dict[str, float]) -> Optional[Dict[str, float]]:
    """Called after each validation within a round (e.g. LR scheduler step).

    Returns:
        Optional dict of extra metrics for the epoch summary line.
    """
    return None

on_train_start

on_train_start(batches)

Called once per round before the first epoch, with the round's training batches.

Runs after the best network was validated on the round's data, so data-dependent set-up done here (e.g. refitting a normalisation to the whole training set) does not affect that baseline. The batches are not shuffled.

Source code in src/falcon/core/base_estimator.py

def on_train_start(self, batches: Iterable) -> None:
    """Called once per round before the first epoch, with the round's training batches.

    Runs after the best network was validated on the round's data, so
    data-dependent set-up done here (e.g. refitting a normalisation to the
    whole training set) does not affect that baseline. The batches are
    not shuffled.
    """

on_round_end

on_round_end(promoted)

Called after the acceptance test, with the best state of every group installed.

Parameters:

Name	Type	Description	Default
promoted	Dict[str, bool]	Whether each group was promoted in this round.	required

Returns:

Type	Description
bool	True if the proposal state changed and must be published.

Source code in src/falcon/core/base_estimator.py

def on_round_end(self, promoted: Dict[str, bool]) -> bool:
    """Called after the acceptance test, with the best state of every group installed.

    Args:
        promoted: Whether each group was promoted in this round.

    Returns:
        True if the proposal state changed and must be published.
    """
    return False

set_observations

set_observations(conditions)

Condition values at the observation, given to the train actor before training.

Only called when every condition of the node is fixed by the observations; evidence derived from them is already simulated.

Parameters:

Name	Type	Description	Default
conditions	Conditions	Condition arrays (numpy) with a leading dimension of 1.	required

Source code in src/falcon/core/base_estimator.py

def set_observations(self, conditions: Conditions) -> None:
    """Condition values at the observation, given to the train actor before training.

    Only called when every condition of the node is fixed by the
    observations; evidence derived from them is already simulated.

    Args:
        conditions: Condition arrays (numpy) with a leading dimension of 1.
    """

export_proposal

export_proposal(full)

Proposal state that is not part of any network group; None if there is none.

Parameters:

Name	Type	Description	Default
full	bool	Everything needed to resume training (for the checkpoint) instead of only what sampling needs (for the sample actors).	required

Source code in src/falcon/core/base_estimator.py

def export_proposal(self, full: bool) -> Optional[StateTree]:
    """Proposal state that is not part of any network group; None if there is none.

    Args:
        full: Everything needed to resume training (for the checkpoint)
            instead of only what sampling needs (for the sample actors).
    """
    return None

import_proposal

import_proposal(tree)

Install a tree from export_proposal().

Source code in src/falcon/core/base_estimator.py

def import_proposal(self, tree: StateTree) -> None:
    """Install a tree from ``export_proposal()``."""

sample_prior abstractmethod

sample_prior(rng, num_samples)

Prior samples: {"value": ndarray, "log_prob": ndarray}.

Source code in src/falcon/core/base_estimator.py

@abstractmethod
def sample_prior(self, rng: np.random.Generator, num_samples: int) -> dict:
    """Prior samples: ``{"value": ndarray, "log_prob": ndarray}``."""

sample abstractmethod

sample(rng, num_samples, conditions, mode)

Samples from the networks: {"value": ndarray, "log_prob": ndarray}.

Parameters:

Name	Type	Description	Default
rng	Generator	Source of randomness for this call.	required
num_samples	int	Number of samples.	required
conditions	Conditions	Condition arrays (numpy); a leading dimension of 1 is broadcast over all samples.	required
mode	str	"proposal" (tempered, used for simulations) or "posterior".	required

Source code in src/falcon/core/base_estimator.py

@abstractmethod
def sample(self, rng: np.random.Generator, num_samples: int,
           conditions: Conditions, mode: str) -> dict:
    """Samples from the networks: ``{"value": ndarray, "log_prob": ndarray}``.

    Args:
        rng: Source of randomness for this call.
        num_samples: Number of samples.
        conditions: Condition arrays (numpy); a leading dimension of 1 is
            broadcast over all samples.
        mode: ``"proposal"`` (tempered, used for simulations) or ``"posterior"``.
    """

RoundConfig dataclass

RoundConfig(max_epochs=100, patience_epochs=16, val_every_epochs=1, max_rounds=None, patience_rounds=10, prior_rounds=0, batch_size=128, val_batch_size=None, wait_for_sims=0, max_cache_samples=0, cache_on_device=False, discard_samples=False)

Parameters of the round-based training loop (see docs/training.md).

TorchModel

Bases: BaseEstimator

Implements the state and sampling plumbing of BaseEstimator for torch.

Subclasses register their networks with _set_groups() in build() and implement:

• init_from_batch() / build()
• train_step() / val_step() / discard_test()
• _sample_prior() / _sample()

val_step, discard_test and _sample run without gradients, and sampling runs with torch's random state seeded from the given rng.

discard_test

discard_test(batch)

Boolean mask (tensor or array) of length len(batch); True = discard.