FlowMatching Estimator

Flow-matching posterior estimation with a truncated-prior proposal.

Overview

FlowMatching learns a conditional flow q_c(u | x) and a marginal flow q_m(u) over the standard-normal latent space u of the prior. Each flow is a velocity field v(w, t, s), an MLP trained by flow matching: it regresses the straight-line bridge from N(0, I) at t = 0 to the data at t = 1. Draws come from integrating the ODE forward with Euler steps. Densities come from the continuous change of variables, with the ODE run backward.

Key features:

• Exact per-round whitening. Each flow works in the frame w of a ZCA whitener that is refit to the round's training data, so the velocity field always sees a unit-scale target, however far the data has contracted. The latent-frame NLL includes the whitener's log-determinant, so rounds with different whiteners are still compared correctly.
• Moving averages. The evaluated, published and saved networks are exponential moving averages of the trained ones.
• Truncated-prior proposal. Proposals are prior samples truncated to a region read off the conditional flow. A hysteretic two-region ladder moves that region once per round (below).
• No importance weights for the posterior. The training data is the truncated prior, so the conditional flow itself estimates the posterior. Posterior samples are conditional-flow draws cut at its x_sigma contour.

Note

FlowMatching requires a TransformedPrior such as Product; it works in the "standard_normal" mode.

The proposal: a region ladder

A region is a level set {ln q_c(u) > thr} of a conditional flow at the observation, intersected with its parent region. A region is frozen when it is minted. It keeps copies of the flows it was read off, the embedded observation and its threshold, so it keeps meaning the same thing as the networks change.

Two nested regions are live:

• the inner region I must contain the x_sigma mass of the current conditional flow;
• the outer region O is what proposals sample, so the training data reaches beyond I.

Wider regions wait on a stack. After every round, the train actor measures m_out, the current conditional flow's mass outside I, and moves the ladder:

Condition	Action
m_out > leak(x_sigma)	EXPAND: I ← O, O ← pop the stack (or the prior)
otherwise	mint a candidate at leak(x_sigma + delta_sigma) inside I
candidate mass ≤ vratio × mass of I (with v_min_ess)	CONTRACT: push O, O ← I, I ← candidate
otherwise	hold

leak(x) is the two-sided normal tail mass: x = 3 means 2.7e-3 of the mass may lie outside I. A region is sampled by importance sampling. The sampler draws from the region's marginal flow, keeps the draws inside the region, weights them by prior over marginal flow, and resamples without replacement, so no simulation is spent twice. The pool of weighted draws is extended until its effective sample size is ess_factor times the draws taken from it, so proposals stay faithful samples of the truncated prior even where the marginal flow has a hole.

The ladder needs the node's conditions at the observation. The train actor gets them at launch; evidence derived from observed nodes is simulated once from the observation. Without them, e.g. in amortized runs, the proposal stays the prior.

The ladder logs [mint] and [ladder] lines to graph/<node>/train/output.log, and ladder:* metrics. The sample actor logs proposal:acceptance, proposal:eps (the ESS per accepted draw) and proposal:pool_ess.

Reference settings

The defaults are the settings of the reference run that produced the O1b region of the LDC MBHB study (t7b_s4k_b8k_max64k_reject_essgate: 9 parameters, 84 rounds, 3.3e5 simulations on one A100). Its buffer policy maps onto falcon's buffer as follows:

Reference (O1b)	falcon
drop samples outside the new outer region, oldest first, down to a floor of 8000	discard_samples: true (default) with buffer.min_samples: 8000
stop generating while 64000 samples are inside the region	buffer.max_samples: 64000, buffer.simulate_when_full: false
20% validation, handled like the training set	buffer.validation_fraction: 0.2
4000 new simulations per round, drawn from the new outer region	wait_for_sims: 4000 (see Training Loop)

Keeping the buffer inside the outer region matters: the whitener can only zoom in as far as the training data has contracted.

Configuration

estimator:
  _target_: falcon.estimators.FlowMatching
  max_epochs: 200
  val_every_epochs: 3        # validation solves the density ODE
  patience_epochs: 12
  lr: 3.0e-4
  embedding:
    _target_: model.E
    _input_: [x]
  hidden: 256
  layers: 4

See examples/01_minimal/config_flow_matching.yml for a complete example.

Configuration Reference

Networks and training

Parameter	Type	Default	Description
hidden	int	512	Velocity-field MLP width
layers	int	6	Velocity-field MLP hidden layers
time_dim	int	32	Number of Fourier time features (even)
layernorm	bool	false	LayerNorm after every hidden layer
lr	float	3e-4	Learning rate; reset at the start of every round
betas	tuple	(0.9, 0.9)	AdamW betas
grad_clip	float	1.0	Gradient-norm clip, per flow (0 = off)
ema_decay	float	0.995	Decay of the moving averages, per step
embedding_epochs	int	10	Train the embedding only in the first N epochs of each round, then keep it fixed (null = throughout)
time_late_k, time_late_mix	float	8.0, 0.5	Training times: a mix share drawn from k t^(k-1), which favours the sharp end of the path
sample_steps	int	128	Euler steps per draw
density_steps	int	32	ODE steps per density
div_probes	int	4	Hutchinson probes for proposal and posterior densities (0 = exact trace; validation is always exact)
eval_chunk	int	16384	Rows per forward pass when sampling or evaluating densities
wait_for_sims	int	0	After a round, wait until this many samples were simulated from the proposal it published before the next round
val_batch_size	int	16384	Batch size of validation and of the discard sweep. The density ODE is bound by its number of small steps, not by the rows, so one large batch costs about as much as one of 128
embed_chunk	int	512	Rows per forward pass of the embedding outside training steps

Region ladder

Parameter	Type	Default	Description
x_sigma	float	3.0	Mass the inner region must hold, in sigmas; also the posterior truncation
delta_sigma	float	1.0	Dead band: new inner regions are minted at x_sigma + delta_sigma
min_keep_frac	float	0.01	At least this share of a region's own marginal-flow draws fall inside it
vratio	float	0.8	Contract only if the candidate has at most this share of the inner region's prior mass
v_min_ess	float	1000	Effective sample size that mass ratio must carry
v_max_draws	int	1048576	Maximum marginal-flow draws spent on it
chain_depth	int	1	Ancestors tested for region membership (-1 = all)
n_region	int	65536	Draws per pass when minting or sampling a region
n_mout	int	65536	Draws for m_out
max_sample_passes	int	64	Maximum passes per proposal request
ess_factor	float	4.0	Extend the pool of region draws until its ESS is this many times the draws taken from it
readout_draws	int	65536	Draws that set the posterior truncation

The defaults are the O1b settings (above), sized for an A100; on a CPU, reduce hidden, layers and the draw counts.

The training loop parameters (max_epochs, patience_epochs, val_every_epochs, max_rounds, patience_rounds, prior_rounds, batch_size, lr_decay_factor, lr_patience_epochs, cache_on_device, max_cache_samples) are as in Flow; see Training Loop. discard_samples defaults to true here: accepted rounds discard samples outside the outer region.

Training details

• Network groups. The first group is the conditional flow together with the embedding, judged on nll. The second is the marginal flow, judged on nll_aux. Both are latent-frame NLLs with the exact trace. Early stopping watches loss, the conditional NLL without the rare rows whose ODE ran away. The training loss is the flow-matching loss, so the train_loss and val_loss in the log are on different scales; val_fm is the validation flow-matching loss.
• Per-sample results. The latent parameters of a sample never change, and its summary only changes while the embedding trains. Both are kept on the device under the sample's buffer id, so after embedding_epochs a training step no longer runs the embedding or touches the data arrays, and validation runs the embedding only for samples it has not seen since the embedding last moved.
• Whitener frames. Every round refits each flow's whitener to the round's training data before the first epoch. A flow that was not promoted in the previous round keeps the frame of the best network it resumes from. Once the proposal region has settled, rounds can alternate between a rejected round in a refit frame and an accepted round in the held one; the O1b reference run did this from round ~50 on.
• Discretisation. Densities and draws come from fixed-step Euler integration. When the conditional flow must contract its target strongly, the density at the default steps is biased high by a fraction of a nat. The whitener's zoom, as the proposal region contracts, keeps that contraction moderate.

Class Reference

FlowMatching

FlowMatching(*, max_epochs=200, lr=0.0003, embedding=None, device=None, batch_size=128, val_batch_size=16384, embed_chunk=512, wait_for_sims=0, patience_epochs=12, val_every_epochs=3, max_rounds=None, patience_rounds=10, prior_rounds=0, cache_on_device=False, max_cache_samples=0, discard_samples=True, hidden=512, layers=6, time_dim=32, layernorm=False, betas=(0.9, 0.9), grad_clip=1.0, lr_decay_factor=1.0, lr_patience_epochs=8, ema_decay=0.995, embedding_epochs=10, time_late_k=8.0, time_late_mix=0.5, sample_steps=128, density_steps=32, div_probes=4, eval_chunk=16384, x_sigma=3.0, delta_sigma=1.0, min_keep_frac=0.01, vratio=0.8, v_min_ess=1000, v_max_draws=1048576, chain_depth=1, n_region=65536, n_mout=65536, max_sample_passes=64, ess_factor=4.0, readout_draws=65536)

Bases: TorchModel

Flow-matching posterior estimation with a truncated-prior proposal.

Works in the standard-normal latent space of a TransformedPrior (e.g. Product). Training runs in rounds (see RoundTrainer); every round refits the whitener of each flow to the round's training data, except for a flow that was not promoted in the previous round, which keeps the frame it was trained in.

Network groups: the conditional flow with the embedding (primary, judged on the validation NLL nll) and the marginal flow (nll_aux). Both NLLs are latent-frame densities, comparable across whitener refits. Early stopping watches loss, the conditional NLL without the rare rows whose ODE ran away. The training loss is the flow-matching loss. The evaluated, published and saved networks are exponential moving averages of the trained ones, and every round resumes training from them.

The proposal is the prior truncated to the outer region of the region ladder (see falcon.estimators.region_ladder), which moves once per round in the train actor. It needs the node's conditions at the observation; without them (e.g. amortized runs) it stays the prior.

The defaults are the settings of the reference run that produced the O1b region of the LDC MBHB study (t7b_s4k_b8k_max64k_reject_essgate). Its buffer policy corresponds to discard_samples: true with buffer.min_samples as the floor, buffer.max_samples as the cap, buffer.simulate_when_full: false and buffer.validation_fraction: 0.2; its 4000 simulations per round, drawn from the new outer region, correspond to wait_for_sims: 4000.

Parameters:

Name	Type	Description	Default
max_epochs	int	Maximum epochs per round.	200
lr	float	Learning rate; reset at the start of every round.	0.0003
embedding		Embedding config dict (with _target_ etc.).	None
device	Optional[str]	Device string (e.g. "cuda:0"); auto-detected if None.	None
batch_size	int	Mini-batch size.	128
val_batch_size	Optional[int]	Batch size of validation and of the discard sweep (None = batch_size). The density ODE costs about the same for 128 rows as for thousands, so large batches make it cheap.	16384
embed_chunk	int	Rows per forward pass of the embedding outside training steps.	512
wait_for_sims	int	After each round, wait until this many samples have been simulated since the round published its proposal, so the next round trains on data from the new proposal region (0 = start the next round at once).	0
patience_epochs	int	End the round once the best validation loss is this many epochs old (checked at validations).	12
val_every_epochs	int	Validate after every N-th epoch (and at the last epoch); validation solves the density ODE, so it is not cheap.	3
max_rounds	Optional[int]	Maximum number of rounds (None = unlimited).	None
patience_rounds	int	Stop training after this many rejected rounds in a row.	10
prior_rounds	int	Rounds that simulate from the prior before switching to the learned proposal.	0
cache_on_device	bool	Cache training data on the estimator device.	False
max_cache_samples	int	Cap on cached training samples (0 = all).	0
discard_samples	bool	After accepted rounds, discard samples outside the outer region (oldest first, never below buffer.min_samples).	True
hidden	int	Width of the velocity-field MLP.	512
layers	int	Hidden layers of the velocity-field MLP.	6
time_dim	int	Number of Fourier time features (even).	32
layernorm	bool	LayerNorm after every hidden layer.	False
betas	tuple	AdamW beta coefficients.	(0.9, 0.9)
grad_clip	float	Gradient-norm clip, separately for the conditional flow with the embedding and for the marginal flow (0 = off).	1.0
lr_decay_factor	float	LR decay factor for the plateau scheduler (1.0 = no decay).	1.0
lr_patience_epochs	int	Epochs without validation improvement before LR decay.	8
ema_decay	float	Decay of the exponential moving averages, per step.	0.995
embedding_epochs	Optional[int]	Train the embedding only in the first N epochs of each round, then keep it fixed (None = train it throughout).	10
time_late_k	float	Exponent of the late-time tilt of the training times.	8.0
time_late_mix	float	Fraction of training times drawn from the tilt.	0.5
sample_steps	int	Euler steps when drawing from a flow.	128
density_steps	int	ODE steps when evaluating a density.	32
div_probes	int	Hutchinson probes for the divergence in proposal and posterior densities (0 = exact trace); validation always uses the exact trace.	4
eval_chunk	int	Rows per forward pass when sampling or evaluating densities.	16384
x_sigma	float	The inner region must contain this many sigmas of the conditional flow's mass (as two-sided normal tail mass: 3 means 2.7e-3 may lie outside); also the posterior truncation.	3.0
delta_sigma	float	Half-width of the ladder's dead band; new inner regions are minted at x_sigma + delta_sigma.	1.0
min_keep_frac	float	At least this fraction of a region's own marginal-flow draws must fall inside it.	0.01
vratio	float	Contract only if the candidate has at most this fraction of the inner region's prior mass.	0.8
v_min_ess	float	Effective sample size the mass ratio must carry to contract.	1000
v_max_draws	int	Maximum marginal-flow draws spent on that ratio.	1048576
chain_depth	int	Ancestors tested when checking membership in a region (-1 = all).	1
n_region	int	Draws per pass when minting or sampling a region.	65536
n_mout	int	Draws for the mass outside the inner region.	65536
max_sample_passes	int	Maximum passes when filling a proposal request.	64
ess_factor	float	Extend the pool of weighted region draws until its effective sample size is this many times the draws taken from it.	4.0
readout_draws	int	Draws used to set the posterior truncation.	65536

Source code in src/falcon/estimators/flow_matching.py

def __init__(
    self,
    *,
    # Most commonly changed
    max_epochs: int = 200,
    lr: float = 3e-4,
    embedding=None,
    device: Optional[str] = None,
    # Training loop
    batch_size: int = 128,
    val_batch_size: Optional[int] = 16384,
    embed_chunk: int = 512,
    wait_for_sims: int = 0,
    patience_epochs: int = 12,
    val_every_epochs: int = 3,
    max_rounds: Optional[int] = None,
    patience_rounds: int = 10,
    prior_rounds: int = 0,
    cache_on_device: bool = False,
    max_cache_samples: int = 0,
    discard_samples: bool = True,
    # Network
    hidden: int = 512,
    layers: int = 6,
    time_dim: int = 32,
    layernorm: bool = False,
    # Optimizer
    betas: tuple = (0.9, 0.9),
    grad_clip: float = 1.0,
    lr_decay_factor: float = 1.0,
    lr_patience_epochs: int = 8,
    ema_decay: float = 0.995,
    embedding_epochs: Optional[int] = 10,
    # Flow matching
    time_late_k: float = 8.0,
    time_late_mix: float = 0.5,
    sample_steps: int = 128,
    density_steps: int = 32,
    div_probes: int = 4,
    eval_chunk: int = 16384,
    # Region ladder
    x_sigma: float = 3.0,
    delta_sigma: float = 1.0,
    min_keep_frac: float = 0.01,
    vratio: float = 0.8,
    v_min_ess: float = 1000,
    v_max_draws: int = 1048576,
    chain_depth: int = 1,
    n_region: int = 65536,
    n_mout: int = 65536,
    max_sample_passes: int = 64,
    ess_factor: float = 4.0,
    readout_draws: int = 65536,
):
    if embedding_epochs is not None and embedding_epochs < 0:
        raise ValueError(f"embedding_epochs must be None or >= 0, got {embedding_epochs}")
    self.max_epochs = max_epochs
    self.lr = lr
    self.embedding = embedding
    self.device = device
    self.batch_size = batch_size
    self.val_batch_size = val_batch_size
    self.embed_chunk = embed_chunk
    self.wait_for_sims = wait_for_sims
    self.patience_epochs = patience_epochs
    self.val_every_epochs = val_every_epochs
    self.max_rounds = max_rounds
    self.patience_rounds = patience_rounds
    self.prior_rounds = prior_rounds
    self.cache_on_device = cache_on_device
    self.max_cache_samples = max_cache_samples
    self.discard_samples = discard_samples
    self.hidden = hidden
    self.layers = layers
    self.time_dim = time_dim
    self.layernorm = layernorm
    self.betas = betas
    self.grad_clip = grad_clip
    self.lr_decay_factor = lr_decay_factor
    self.lr_patience_epochs = lr_patience_epochs
    self.ema_decay = ema_decay
    self.embedding_epochs = embedding_epochs
    self.time_late_k = time_late_k
    self.time_late_mix = time_late_mix
    self.sample_steps = sample_steps
    self.density_steps = density_steps
    self.div_probes = div_probes
    self.eval_chunk = eval_chunk
    self.x_sigma = x_sigma
    self.delta_sigma = delta_sigma
    self.min_keep_frac = min_keep_frac
    self.vratio = vratio
    self.v_min_ess = v_min_ess
    self.v_max_draws = v_max_draws
    self.chain_depth = chain_depth
    self.n_region = n_region
    self.n_mout = n_mout
    self.max_sample_passes = max_sample_passes
    self.ess_factor = ess_factor
    self.readout_draws = readout_draws

build

build(init_tree)

Source code in src/falcon/estimators/flow_matching.py

def build(self, init_tree) -> None:
    debug("Initializing networks...")
    conditions = {k: self._to_tensor(v, self.device) for k, v in init_tree["conditions"].items()}
    embedding = instantiate_embedding(self.embedding).to(self.device)
    embedding.eval()
    with torch.no_grad():
        self._cond_dim = embedding(conditions).shape[1]
    self._frozen_params = {name for name, p in embedding.named_parameters() if not p.requires_grad}

    self._embedding = embedding
    self._flow_c = self._new_flow()
    self._flow_m = self._new_flow()
    for module in (self._embedding, self._flow_c, self._flow_m):
        module.requires_grad_(False)
    self._set_groups({
        "conditional": NetworkGroup(
            {"flow": self._flow_c, "embedding": self._embedding}, "nll",
            embedding=self._embedding,
        ),
        "marginal": NetworkGroup({"flow": self._flow_m}, "nll_aux"),
    })
    debug(f"FlowMatching networks built: param_dim={self.param_dim}, cond_dim={self._cond_dim}")

train_step

train_step(batch)

Source code in src/falcon/estimators/flow_matching.py

def train_step(self, batch) -> Dict[str, float]:
    self._ensure_training()
    nets = self._train_nets
    u = self._latent(batch)
    w_c = self._flow_c.whitener.whiten(u).float()
    w_m = self._flow_m.whitener.whiten(u).float()

    joint = self._embedding_trains()
    if joint:
        nets["embedding"].train(True)
        s = nets["embedding"](self._conditions(batch)).float()
    else:
        with torch.no_grad():
            s = self._fixed_summary(batch)

    loss_c = fm_loss_late(nets["velocity_c"], w_c, s, self.time_late_k, self.time_late_mix)
    loss_m = fm_loss_late(nets["velocity_m"], w_m, torch.zeros_like(s), self.time_late_k, self.time_late_mix)
    self._optimizer.zero_grad(set_to_none=True)
    (loss_c + loss_m).backward()
    if self.grad_clip:
        torch.nn.utils.clip_grad_norm_(self._params_c, self.grad_clip)
        torch.nn.utils.clip_grad_norm_(self._params_m, self.grad_clip)
    self._optimizer.step()

    ema_update(self._flow_c.velocity, nets["velocity_c"], self.ema_decay)
    ema_update(self._flow_m.velocity, nets["velocity_m"], self.ema_decay)
    if joint:
        ema_update(self._embedding, nets["embedding"], self.ema_decay)
        self._summary_stale = True
    self._steps += 1
    return {"loss": loss_c.item(), "loss_aux": loss_m.item()}

val_step

val_step(batch)

Source code in src/falcon/estimators/flow_matching.py

def val_step(self, batch) -> Dict[str, float]:
    u = self._latent(batch)
    s = self._fixed_summary(batch)
    zeros = torch.zeros_like(s)
    kw = self._density_kw(exact=True)
    log_prob_c = self._flow_c.log_prob(u, s, **kw)
    log_prob_m = self._flow_m.log_prob(u, zeros, **kw)
    return {
        "loss": _robust_nll(log_prob_c),
        "nll": float(-log_prob_c.mean()),
        "nll_aux": float(-log_prob_m.mean()),
        "fm": val_fm_loss(self._flow_c.velocity, self._flow_c.whitener.whiten(u).float(), s),
        "fm_aux": val_fm_loss(self._flow_m.velocity, self._flow_m.whitener.whiten(u).float(), zeros),
    }

on_train_start

on_train_start(batches)

Source code in src/falcon/estimators/flow_matching.py

def on_train_start(self, batches) -> None:
    self._ensure_training()
    u = torch.cat([self._latent(batch).cpu() for batch in batches])
    refit = []
    for name, flow in (("conditional", self._flow_c), ("marginal", self._flow_m)):
        if not self._hold_gauge.get(name, False):
            flow.whitener.refit(u)
            refit.append(name)
    held = [name for name in ("conditional", "marginal") if name not in refit]
    info(
        f"Whitener refit on {len(u)} samples"
        + (f" (frame held: {', '.join(held)})" if held else "")
        + f" | logdet conditional={float(self._flow_c.whitener.logdet):.3f}"
        f" marginal={float(self._flow_m.whitener.logdet):.3f}"
    )
    log({
        "whitener:logdet_conditional": float(self._flow_c.whitener.logdet),
        "whitener:logdet_marginal": float(self._flow_m.whitener.logdet),
    })
    self._steps = 0
    self._steps_per_epoch = len(u) // self.batch_size

on_round_end

on_round_end(promoted)

Source code in src/falcon/estimators/flow_matching.py

def on_round_end(self, promoted: Dict[str, bool]) -> bool:
    # A flow that was not promoted resumes from the best network next
    # round, in the frame that network was trained in
    self._hold_gauge = {name: not p for name, p in promoted.items()}
    if self._observation is None:
        if not self._warned_no_observation:
            warning(
                f"The conditions of '{self.theta_key}' are not fixed by the observations, "
                "so the proposal stays the prior"
            )
            self._warned_no_observation = True
        return False
    conditions = {k: self._to_tensor(v, self.device) for k, v in self._observation.items()}
    with torch.no_grad():
        s_obs = self._summary("conditional", conditions).float()[:1]
    current = FlowPair(self._flow_c, self._flow_m, s_obs, self._sample_kw(), self._density_kw())
    trees = {
        "c": module_state(self._flow_c),
        "m": module_state(self._flow_m),
        "s_obs": s_obs.cpu().numpy(),
    }
    action = self._ladder.step(current, trees)
    return action in ("EXPAND", "CONTRACT")

RegionLadder

RegionLadder(config, load_pair, param_dim, device)

Inner region I, outer region O and a stack of wider regions.

Regions are records (plain values and numpy trees) keyed by id; the prior is the permanent region "prior", which is never stored.

Parameters:

Name	Type	Description	Default
config	LadderConfig	Ladder parameters.	required
load_pair	Callable[[Dict[str, Any]], FlowPair]	Builds the FlowPair of a region record.	required
param_dim	int	Dimension of the latent space.	required
device		Device of the flows.	required

volume

volume(rid)

Prior mass of a region, including the cut by its parent.

cut

cut(rid, u, depth=None)

Membership of latent points in a region, including its parents up to depth.

A region is the intersection of its own level set with its parent, so a point can only be in a child if it is in the parent; depth (the chain_depth by default, -1 for all) limits how many ancestors are tested. Ancestors are evaluated only on the survivors.

mass_outside

mass_outside(current, rid)

The current conditional flow's mass outside a region.

The draws come from the mixture Q = (q_c + q_m) / 2 and are weighted by q_c / Q (at most 2): the marginal half populates the outside densely, so a tail mass of 1e-6 is an average of many small terms instead of a count of rare ones.

mint

mint(current, target_leak, parent)

Threshold and prior mass of a candidate region at target_leak inside parent.

The threshold is the target_leak quantile of ln q_c over draws from the mixture Q = (q_c + q_m) / 2, weighted by q_c / Q. It is then lowered, if necessary, so that at least min_keep_frac of the marginal flow's draws survive the cut, which widens the region (the safe direction) and keeps sampling it affordable.

The prior mass V of the level set and the ratio R of the candidate's mass inside the parent to the parent's mass are estimated on marginal-flow draws weighted by prior / q_m; more draws are added until R carries v_min_ess or v_max_draws is reached.

Returns:

Type	Description
Optional[Dict[str, float]]	The candidate's statistics, or None if no draw had a finite density.

step

step(current, trees)

Move the ladder once, with the current best flows; returns the action.

Parameters:

Name	Type	Description	Default
current	FlowPair	The current best flows at the observation.	required
trees	Dict[str, Any]	What a region minted now stores to rebuild current (c, m and s_obs).	required

Returns:

Type	Description
str	"EXPAND", "CONTRACT" or "hold (...)".

sample

sample(n)

n latent draws from the prior truncated to O, with their log density.

Draws are taken without replacement from a pool of weighted draws of O, which lives until O changes. Before taking them, the pool is extended until its effective sample size is ess_factor times everything taken from it so far, this request included (at most max_sample_passes passes; a pool that misses the target stops enforcing it).

pool_ess

pool_ess()

Effective sample size of every draw the current proposal pool has held.

export

export(full)

The ladder as a state tree: everything, or only what sampling needs.

load

load(tree)

Install a tree from export().