"""NPE pipeline for discrete normalizing flows (parallel track).
The flow IS the density model: no ``ConditionalWrapper``, no ``GenerativeMethod``.
Trains ``q(obs | cond)`` by max-likelihood. NPE convention: ``obs = theta``,
``cond = x`` (mirrors ``ConditionalPipeline`` so the diagnostics run unchanged).
"""
import warnings
import jax
import jax.numpy as jnp
from tqdm.auto import tqdm
from gensbi.recipes.pipeline import AbstractPipeline, _chunked_draw, _validate_chunk_size
from gensbi.recipes.utils import _require_channel, _single_obs
[docs]
def _warn_unused_kwargs(kwargs):
"""Warn that solver-style kwargs are ignored by the (solver-free) flow.
The flow pipeline mirrors the flow-matching surface (which accepts
``**sampler_kwargs``), but a normalizing flow has no ODE/SDE solver, so
arguments like ``step_size``/``nsteps``/``solver`` do not apply and are
silently ignored apart from this warning.
"""
if kwargs:
keys = ", ".join(sorted(kwargs))
warnings.warn(
f"flow pipeline ignores unsupported keyword argument(s): {keys}. "
"A normalizing flow has no solver, so these have no effect.",
UserWarning, stacklevel=3,
)
[docs]
class ConditionalFlowPipeline(AbstractPipeline):
"""Max-likelihood NPE pipeline wrapping an ``MAFlow``.
Parameters
----------
model : MAFlow
A pre-built flow (e.g. ``MAFlow(MAFlowParams(rngs=rngs, dim=dim_obs, cond_dim=dim_cond))``).
train_dataset, val_dataset : iterable
Yield ``(obs, cond)`` batches. Shape is ``(B, dim, C)`` for each
variable (``C = 1`` for the tabular path; see ``ch_obs``/``ch_cond``).
dim_obs, dim_cond : int
ch_obs, ch_cond : int, optional
Channel count for the obs and cond variables. Default 1 (tabular SBI).
Values > 1 enable the ``(B, dim, C)`` channel-passthrough path: the
channel axis is preserved and forwarded to the flow unchanged (the flow
must be built with matching ``channels``/``cond_channels`` in
:class:`~gensbi.models.MAFlowParams`).
structured_obs, structured_cond : bool, optional
If ``True``, the modeled variable / condition keeps its native
structured shape (the model owns it) instead of the tabular
``(B, dim, 1)`` layout. Default ``False``.
Notes
-----
Every single-observation method (:meth:`sample`, :meth:`log_prob`,
:meth:`get_sampler`, :meth:`get_log_prob_fn`) expects ``x_o`` to carry a
**leading batch axis** (size 1 for one observation) **and a channel axis**:
shape ``(1, dim_cond, C)`` for tabular, or ``(1,) + per_obs_shape`` for
structured. A bare ``(B, dim)`` tensor is rejected — add ``[..., None]``
for ``C = 1``. A batch axis > 1 raises ``ValueError`` — pass a batch to
:meth:`sample_batched` instead.
"""
def __init__(self, model, train_dataset, val_dataset, dim_obs, dim_cond,
ch_obs=1, ch_cond=1, params=None, training_config=None,
structured_obs=False, structured_cond=False):
super().__init__(
model, train_dataset, val_dataset, dim_obs, dim_cond,
ch_obs=ch_obs, ch_cond=ch_cond, params=params,
training_config=training_config)
[docs]
self._standardized = False
[docs]
self.structured_obs = structured_obs
[docs]
self.structured_cond = structured_cond
[docs]
def _prep_obs(self, x):
x = jnp.asarray(x)
return x if self.structured_obs else _require_channel(x, "obs")
[docs]
def _prep_cond(self, x):
x = jnp.asarray(x)
return x if self.structured_cond else _require_channel(x, "cond")
# --- abstract methods the flow pipeline does not use (mirror ConditionalPipeline) ---
@classmethod
[docs]
def init_pipeline_from_config(cls, *args, **kwargs):
"""Not implemented: the flow pipeline requires a pre-built model.
Raises
------
NotImplementedError
Always. Construct an ``MAFlow`` and pass it as ``model=`` to the
pipeline constructor instead.
"""
raise NotImplementedError(
"ConditionalFlowPipeline takes a pre-built flow; build a `MAFlow` "
"and pass it as model=.")
[docs]
def _make_model(self, params):
raise NotImplementedError(
"Pass a pre-built MAFlow as model=; the flow pipeline does not build "
"models from params.")
@classmethod
[docs]
def get_default_params(cls, *args, **kwargs):
"""Not implemented: the flow pipeline takes a pre-built ``MAFlow``.
Raises
------
NotImplementedError
Always. There are no default model params to return; construct an
``MAFlow`` directly and pass it as ``model=``.
"""
raise NotImplementedError(
"ConditionalFlowPipeline takes a pre-built MAFlow; there are no model "
"params to default.")
# --- the flow IS the model: no wrapper ---
[docs]
def _wrap_model(self):
self.model_wrapped = self.model
self.ema_model_wrapped = self.ema_model
# --- methods implemented in later tasks (Tasks 3, 5, 6) ---
[docs]
def get_loss_fn(self):
"""Return the max-likelihood loss function for training.
Returns a closure ``loss_fn(model, batch, key) -> Array`` that
computes the mean negative log-likelihood
``-mean(log q(obs | cond))``. ``batch = (obs, cond)`` with each
element of shape ``(B, dim, 1)``. NPE convention: ``obs = theta``,
``cond = x``. The ``key`` argument is accepted for interface
compatibility but is unused.
Returns
-------
loss_fn : Callable
A function ``(model, batch, key) -> Array`` returning the scalar
mean negative log-likelihood.
"""
def loss_fn(model, batch, key):
obs, cond = batch
obs = self._prep_obs(obs)
cond = self._prep_cond(cond)
# Loss is always computed in fp32 regardless of the model's
# compute dtype (defense-in-depth on top of the models-emit-fp32
# contract).
return jnp.asarray(-jnp.mean(model.log_prob(obs, cond)), jnp.float32)
return loss_fn
[docs]
def fit_standardization(self, obs_data, axis=0):
"""Fit the Standardize bijection buffers from training observations.
Computes per-dimension mean and standard deviation of ``obs_data``
and stores them as buffers on both the live model and the EMA model.
EMA only averages ``Param`` variables, so the non-Param buffers must
be set explicitly here. Must be called before :meth:`train` when
input standardization is desired.
Parameters
----------
obs_data : Array
Training observations of shape ``(N, dim_obs)`` or
``(N, dim_obs, 1)`` (the autoregressive target; e.g. theta for
NPE). For multichannel flows (``ch_obs > 1``) the shape is
``(N, dim_obs, C)`` and ``axis=(0, 1)`` yields per-channel stats.
axis : int or tuple of int, optional
Reduction axis or axes for the mean/std computation. Default
is ``0`` (per-dimension stats over the batch), which is the
correct choice for the tabular (``C == 1``) path. Pass
``axis=(0, 1)`` for per-channel standardization when ``C > 1``.
"""
obs = jnp.asarray(obs_data)
mean = jnp.mean(obs, axis=axis)
std = jnp.std(obs, axis=axis)
std = jnp.where(std < 1e-6, 1.0, std) # guard zero-variance dims
self.model.set_standardization(mean, std)
self.ema_model.set_standardization(mean, std)
self._standardized = True
[docs]
def train(self, rngs, nsteps=None, save_model=True):
"""Train the flow model, warning if standardization was skipped.
Delegates to :meth:`AbstractPipeline.train` after checking that
:meth:`fit_standardization` was called.
Parameters
----------
rngs : nnx.Rngs
Random number generators for training and validation steps.
nsteps : int or None, optional
Number of training steps. If ``None``, taken from
``training_config["nsteps"]``. Default is ``None``.
save_model : bool, optional
If ``True`` (default), serialise the model to disk after
training.
Returns
-------
loss_array : list
Per-step training losses.
val_loss_array : list
Validation losses recorded at each validation checkpoint.
"""
if not self._standardized:
warnings.warn(
"fit_standardization() was not called before train(); the "
"Standardize bijection stays at identity. Call "
"pipeline.fit_standardization(theta_train) first if you want "
"input standardization.",
UserWarning, stacklevel=2)
return super().train(rngs, nsteps=nsteps, save_model=save_model)
[docs]
def get_sampler(self, x_o, use_ema=True, **kwargs):
"""Return a sampler closure for a single conditioning observation.
Parameters
----------
x_o : Array
Single conditioning observation. Must carry a leading batch axis
and a channel axis for tabular cond: shape ``(1, dim_cond, C)``.
For structured cond: ``(1,) + per_observation_shape``.
A leading batch axis > 1 raises ``ValueError`` (use
:meth:`sample_batched` for many conditions).
use_ema : bool, optional
If ``True`` (default), use the EMA model; otherwise use the
live model.
Returns
-------
sampler : Callable
A function ``(key, nsamples) -> Array`` returning the model's
native output shape ``(nsamples, dim_obs, C)`` (channel always
carried).
"""
_warn_unused_kwargs(kwargs)
flow = self.ema_model if use_ema else self.model
mode = "none" if self.structured_cond else "require"
cond = _single_obs(x_o, channel=mode)[0] # (cond_dim, C_cond) or structured per-obs shape
def sampler(key, nsamples):
cond_b = jnp.broadcast_to(cond, (nsamples,) + cond.shape)
return flow.sample(key, cond=cond_b) # model owns (nsamples, dim, C)
return sampler
[docs]
def sample(self, key, x_o, nsamples=10_000, use_ema=True,
chunk_size=None, show_progress_bars=True, **kwargs):
"""Draw posterior samples for a single conditioning observation.
Parameters
----------
key : jax.random.PRNGKey
Random key.
x_o : Array
Single conditioning observation carrying a leading batch axis of
size 1 (see :meth:`get_sampler` for the shape convention). A leading
batch axis > 1 raises ``ValueError``.
nsamples : int, optional
Number of posterior samples to draw. Default is 10 000.
use_ema : bool, optional
If ``True`` (default), use the EMA model.
chunk_size : int, optional
Maximum number of samples drawn per device call. ``None``
(default) draws everything in one call — identical to the
historical behavior. Set it to bound memory when drawing many
samples from a deep flow.
show_progress_bars : bool, optional
Show a progress bar over chunks (only when chunking is
active). Default is True.
Returns
-------
samples : Array
Posterior samples of shape ``(nsamples, dim_obs, 1)`` for the
tabular default (``C = 1``), or ``(nsamples, dim_obs, C)`` for
``ch_obs = C`` — the channel axis is always carried for a
vector-modeled variable regardless of ``structured_cond`` (a
structured condition changes only ``x_o``'s expected shape, not
the modeled variable's). When ``structured_obs=True``, samples
instead have shape ``(nsamples,) + per_obs_shape``, the model's
native structured output.
"""
sampler = self.get_sampler(x_o, use_ema=use_ema, **kwargs)
return _chunked_draw(
sampler, key, nsamples, chunk_size,
show_progress_bars=show_progress_bars,
)
[docs]
def sample_batched(self, key, x_o, nsamples=10_000, *, use_ema=True,
chunk_size=None, show_progress_bars=True, **kwargs):
"""Draw posterior samples for a batch of conditioning observations.
Each condition is repeated ``nsamples`` times and concatenated
into a single flattened ``(B * nsamples, ...)`` batch. Without
``chunk_size`` the whole batch runs in **one** autoregressive
pass (memory scales with ``B * nsamples``); with ``chunk_size``
the flattened batch is sliced into pieces of at most
``chunk_size`` rows per ``flow.sample`` call — chunk boundaries
may fall inside a condition, which is fine because every row is
independent.
Parameters
----------
key : jax.random.PRNGKey
Random key for the batched sampling pass.
x_o : Array
Batch of observations. For tabular cond: shape
``(B, dim_cond, C)`` (a bare ``(B, dim_cond)`` raises
``ValueError`` — add a trailing channel axis). For structured
cond: ``(B,) + per_obs_shape``.
nsamples : int, optional
Number of posterior samples per observation. Default is
10 000.
use_ema : bool, optional
If ``True`` (default), use the EMA model.
chunk_size : int, optional
Maximum number of rows of the flattened ``B * nsamples``
batch per device call. ``None`` (default) keeps the
historical single-pass behavior.
show_progress_bars : bool, optional
Show a progress bar over chunks (only when chunking is
active). Default is True.
**kwargs : dict, optional
Extra keyword arguments accepted for interface compatibility
and ignored with a warning (e.g. solver arguments from
:class:`~gensbi.recipes.pipeline.AbstractPipeline`).
Returns
-------
samples : Array
Posterior samples of shape ``(nsamples, B, dim_obs, 1)`` for the
tabular default (``C = 1``), or ``(nsamples, B, dim_obs, C)`` for
``ch_obs = C``. When ``structured_obs=True``, samples instead
have shape ``(nsamples, B) + per_obs_shape``. In both cases
``out[:, i]`` is the samples for condition ``i``.
"""
_warn_unused_kwargs(kwargs)
_validate_chunk_size(chunk_size)
flow = self.ema_model if use_ema else self.model
x_o = jnp.asarray(x_o)
if not self.structured_cond:
x_o = _require_channel(x_o, "x_o")
B = x_o.shape[0]
cond = jnp.repeat(x_o, nsamples, axis=0) # (B*nsamples, ...): c0 x nsamples, c1 x nsamples, ...
total = B * nsamples
if chunk_size is None or chunk_size >= total:
samples = flow.sample(key, cond=cond) # ONE batched AR pass
else:
n_chunks = (total + chunk_size - 1) // chunk_size
keys = jax.random.split(key, n_chunks)
loop = range(n_chunks)
if show_progress_bars:
loop = tqdm(loop, desc="Sampling")
chunks = []
for i in loop:
sl = slice(i * chunk_size, min((i + 1) * chunk_size, total))
chunk = flow.sample(keys[i], cond=cond[sl])
chunks.append(jax.block_until_ready(chunk))
samples = jnp.concatenate(chunks, axis=0)
samples = samples.reshape((B, nsamples) + samples.shape[1:])
return jnp.moveaxis(samples, 0, 1) # (nsamples, B, dim_obs, C)
[docs]
def get_log_prob_fn(self, x_o, use_ema=True, **kwargs):
"""Return a log-probability closure for a single conditioning observation.
Parameters
----------
x_o : Array
Single conditioning observation carrying a leading batch axis of
size 1 (see :meth:`get_sampler` for the shape convention). A leading
batch axis > 1 raises ``ValueError``.
use_ema : bool, optional
If ``True`` (default), use the EMA model.
Returns
-------
log_prob_fn : Callable
A function ``(x_1) -> Array`` of shape ``(B,)`` evaluating
the conditional log-probability ``log q(x_1 | x_o)`` for a
batch of ``B`` parameter vectors. ``x_1`` has shape
``(B, dim_obs)`` or ``(B, dim_obs, 1)`` on the tabular path,
or ``(B, dim_obs, C)`` when ``ch_obs > 1``
(channel-passthrough).
"""
_warn_unused_kwargs(kwargs)
flow = self.ema_model if use_ema else self.model
mode = "none" if self.structured_cond else "require"
cond = _single_obs(x_o, channel=mode)[0] # (cond_dim, C_cond) or structured per-obs shape
def log_prob_fn(x_1):
obs = self._prep_obs(x_1)
cond_b = jnp.broadcast_to(cond, (obs.shape[0],) + cond.shape)
return flow.log_prob(obs, cond_b) # (B,)
return log_prob_fn
[docs]
def log_prob(self, x_1, x_o, use_ema=True, **kwargs):
"""Evaluate the conditional log-probability for a batch of samples.
Parameters
----------
x_1 : Array
Batch of parameter vectors of shape ``(B, dim_obs)`` or
``(B, dim_obs, 1)``.
x_o : Array
Single conditioning observation carrying a leading batch axis of
size 1 (see :meth:`get_sampler` for the shape convention). A leading
batch axis > 1 raises ``ValueError``.
use_ema : bool, optional
If ``True`` (default), use the EMA model.
Returns
-------
log_prob : Array
Log-probabilities of shape ``(B,)``.
"""
return self.get_log_prob_fn(x_o, use_ema=use_ema, **kwargs)(x_1)
