SMC steps

Common interface

Before functional initialization, the runner calls each step’s validate_functional_configuration(config, state) hook, passing its smc.config snapshot and the initial particle state. SMCConfig exposes steps, callbacks, checkpoint_writers, and auto_normalize as read-only fields with tuple containers. Referenced step and callback objects retain their identity and mutable state. The config is host-side configuration, separate from runtime metadata and checkpoints. Steps can reject incompatible pipeline settings here, before any target evaluation or cache preparation. The default implementation does nothing. Metadata and cache preparation remain in initialize_functional_metadata and initialize_functional_state_and_context. Algorithm-specific checks belong to the relevant steps, rather than the runner.

Functional conditional execution

sequince.steps.apply_functional_if(predicate, apply, state, context, key) is a reusable control-flow helper for functional steps, using acai.cond. It executes apply(state, context, key) only when the scalar predicate is true. Otherwise it returns the unchanged state, context, and RNG. On NumPy, CuPy, and PyTorch, only the selected branch executes and skipped output is None. PyTorch control flow is eager, without structured torch.compile guarantees. On JAX, skipped outputs contain zero array placeholders in the normal output structure, with static fields retained. Both JAX branches are traced; the operation must be pure and traceable even when skipped, and initialization must establish matching carry structures, shapes, and dtypes. Placeholder outputs do not describe executed transitions.

After each functional step, the runner guards subsequent steps using the runtime context.stop_immediately flag. StopIf(immediate=True), NS-SMC shell steps, and custom steps returning context.request_stop_immediately() all stop execution without an additional capability declaration. An ordinary stop request still permits the rest of the iteration to execute. Use initialize_functional before directly invoking a compiled iteration so all guarded steps have compatible carries; run_functional and compile_functional_iteration perform this initialization automatically. As in eager execution, runner-level automatic normalization still runs after the step sequence when enabled.

Weighting strategies

Nested-sampling SMC

Use make_nested_sampling_smc to construct the composable pipeline and wrap its result in NestedSamplingSMCResult for shell evidence and posterior samples. NestedSamplingLevel provides the per-shell diagnostics. The ordinary SMCResult.log_evidence() and SMCResult.posterior_samples() methods raise an error for NS-SMC results and direct callers to this wrapper. SMCResult.log_normalizing_constant() remains available for the constrained prior normalization.

Custom threshold schedules must implement schedule(log_likelihoods, state, context) -> float | scalar array. All three arguments are passed positionally. Raise StopIteration to integrate the remaining live tail and finish; other exceptions propagate to the caller. Fixed, weighted quantile, and worst-particle schedules implement this same signature.

During execution, context.metadata["nested_sampling"] holds the current threshold and bookkeeping as a plain dictionary. Steps replace this dictionary rather than mutate an earlier context. Shell expectations and evidence increments are stored in NestedSamplingLevel records. State metadata records snapshots after survivor weighting or terminal-shell integration, preserving result histories. Resume through the standard checkpoint API to restore both particles and context.

Functional and compiled execution

run_functional supports JAX arrays with unconditional resampling (always_resample=True) and an unchanged population size. BlackJAX and MiniPCN provide functional mutation kernels; host-only kernels such as emcee, Eryn, and the current Metropolis-Hastings implementations are rejected. Set logging_frequency=0 and use runner callbacks for host logging. Existing functional restrictions on transforms and mutation options still apply.

For example, using JAX-compatible density functions and particles:

import jax
from sequince import ParticleState, make_nested_sampling_smc
from sequince.jax import compile_functional_iteration

smc, target, metadata = make_nested_sampling_smc(
    log_prior, log_likelihood,
    threshold_schedule=[-10.0, -5.0, -1.0],
    evidence_tolerance=None,
    mutation="blackjax",
    mutation_kwargs={"algorithm": "random-walk",
                     "algorithm_kwargs": {"sigma": 0.1}},
    particle_transform=None,
    logging_frequency=0,
)
state = ParticleState(particles)
key = jax.random.key(42)
iteration = compile_functional_iteration(smc, state, key, metadata=metadata)
result, key = smc.run_functional(
    state, random_state=key, metadata=metadata, iteration=iteration,
)
nested = result.as_nested_sampling_result()
samples, weights = nested.posterior_samples()

Omit iteration for uncompiled functional execution, or supply jax.jit(smc.functional_iteration). The runner initializes numeric metadata and a complete target cache before the first iteration. Call initialize_functional yourself when invoking functional_iteration directly. Context metadata remains authoritative; state metadata contains independent historical snapshots. Stop reasons are numeric during functional execution; nested.stop_reason and nested.summary() expose readable strings. Shapes, dtypes, dictionary keys, and cache components must remain fixed.

Compiled schedules use the same three-argument signature. They must use JAX operations without host conversions, Python branching on traced values, mutable iterator state, or data-dependent exceptions. Return positive infinity to integrate the terminal tail. The built-in fixed schedule uses this sentinel on exhaustion while retaining StopIteration in eager execution. An explicit infinite threshold (including one returned by a custom schedule) is reported as no_survivors; exhausting a built-in fixed schedule is reported as schedule_exhausted. Ordinary user exceptions propagate during execution or tracing. Target functions and transformations must also be pure and traceable; log densities must be finite or negative infinity. Shape and cache ownership checks run during tracing; density values and threshold errors are checked at initialization and after each host-controlled iteration. Direct compiled calls must invoke the selection step’s validate_functional_state_and_context at the host boundary. Rejected internal mutation proposals are not inspected by these host checks.

After terminal-shell integration, later steps are conditionally skipped, including weighting, resampling, and mutation. Their carry (particles, caches, and RNG) is unchanged. Skipped outputs retain the usual structure with zero array placeholders; they do not represent executed transitions. Functional shell outputs use population-sized arrays and a particle mask. Optional prior-mass tail outputs also carry a valid flag, allowing two shells in one iteration. The result wrapper filters these flags and masks when reconstructing levels and posterior samples. Eager shell storage remains compact, and store_shell_particles=False avoids retaining shell particle arrays in both modes.

Evidence remains separate from constrained-prior normalizing constants. A num_steps limit suspends an incomplete run without integrating its live tail. Use resume_functional_from(checkpoint) and then run_functional(runtime=runtime, ...) to continue. With a newly constructed target, initialization refreshes incompatible cached evaluations before compilation/execution; RNG and historical shells are restored. Portable result serialization retains shell masks, diagnostics, and posterior interpretation.

Resampling

Mutation kernels

Proposals and dimensionality control

Control flow and utilities