sequince.smc ============ .. py:module:: sequince.smc Attributes ---------- .. autoapisummary:: sequince.smc.logger sequince.smc.IterationCallable sequince.smc.StepCallable Classes ------- .. autoapisummary:: sequince.smc.SMCConfig sequince.smc.SMCRuntime sequince.smc.SequentialMonteCarlo Module Contents --------------- .. py:data:: logger .. py:type:: IterationCallable :canonical: Callable[[ParticleState, StepContext, Any], tuple[ParticleState, StepContext, Any, tuple[Any, ...]]] .. py:type:: StepCallable :canonical: Callable[[ParticleState, StepContext, 'RandomGenerator'], tuple[ParticleState, StepContext, Any]] .. py:class:: SMCConfig Read-only snapshot of runner settings, separate from runtime state. :ivar steps: Ordered pipeline. Step objects retain their identity and mutable state. :vartype steps: tuple of callables :ivar callbacks: Host callbacks invoked after each iteration. :vartype callbacks: tuple of callables :ivar checkpoint_writers: Writers invoked at checkpoint boundaries. :vartype checkpoint_writers: tuple of CheckpointWriter :ivar auto_normalize: Whether the runner normalizes weights after each iteration. :vartype auto_normalize: bool .. rubric:: Notes Containers and fields are immutable; referenced steps, callbacks, and writers are not copied or frozen. This is a host-side configuration view, not a checkpoint payload or compiled carry. .. py:attribute:: steps :type: tuple[StepCallable, ...] .. py:attribute:: callbacks :type: tuple[Callable[..., sequince.steps.base.StepContext | None], ...] .. py:attribute:: checkpoint_writers :type: tuple[sequince.checkpoint.CheckpointWriter, ...] .. py:attribute:: auto_normalize :type: bool .. py:class:: SMCRuntime Prepared SMC runtime state for fresh or resumed execution. ``rng`` contains either an ``orng.RandomGenerator`` for eager execution or an explicit backend random state for functional execution. .. py:attribute:: state :type: sequince.state.ParticleState .. py:attribute:: produced_states :type: list[sequince.state.ParticleState] .. py:attribute:: initial_state :type: sequince.state.ParticleState | None .. py:attribute:: t :type: int .. py:attribute:: next_observation_index :type: int .. py:attribute:: shared_context :type: collections.abc.Mapping[str, Any] .. py:attribute:: metadata :type: collections.abc.Mapping[str, Any] .. py:attribute:: evaluation_counts :type: sequince.target.EvaluationCounts .. py:attribute:: outputs :type: list[tuple[Any, ...]] .. py:attribute:: rng :type: Any .. py:class:: SequentialMonteCarlo(steps, *, callbacks = None, checkpoint_writers = None, auto_normalize = True) Composable Sequential Monte Carlo runner. Users supply an ordered collection of SMC steps that will be called in sequence at each iteration. Steps return the updated particle state, context, and an explicit output describing the completed work. :param steps: A sequence of callables implementing SMC steps. Each callable should accept the current ``ParticleState``, current ``StepContext``, and an ``orng.RandomGenerator``, and return ``(state, context, output)``. :param callbacks: Optional callbacks invoked after each completed iteration with the current state, context, and partial ``SMCResult``. A callback may return an updated context, or ``None`` when observational. :param checkpoint_writers: Optional persistence hooks invoked with a committed checkpoint after callbacks have completed. :param auto_normalize: If ``True``, normalize particle weights after each iteration. .. py:property:: config :type: SMCConfig Return a read-only snapshot of the current runner settings. .. py:method:: from_config(config) :classmethod: Construct a runner from a configuration snapshot. :param config: Runner settings, such as those returned by ``runner.config``. :type config: SMCConfig :returns: A new runner constructed through ``cls`` with the supplied settings. :rtype: SequentialMonteCarlo .. rubric:: Notes Steps, callbacks, and checkpoint writers are shared by reference. This does not clone their mutable state or restore a checkpoint. Subclasses must accept the same constructor arguments. .. py:method:: state_dict() Return mutable state for the configured step pipeline. .. py:method:: load_state_dict(state) Restore mutable state for the configured step pipeline. .. py:method:: iteration(state, context, rng) Apply one eager SMC iteration. The returned RNG is the same mutable generator supplied by the caller. Returning it gives eager and functional iterations a common internal calling convention. .. py:method:: functional_iteration(state, context, random_state) Apply one pure, explicitly RNG-threaded SMC iteration. .. rubric:: Notes After each step, ``context.stop_immediately`` guards subsequent steps, preserving terminal particles and RNG state. Guarded JAX steps require initialized carries with matching structures, shapes, and dtypes. Ordinary stop requests are handled after the complete iteration. Skipped JAX outputs contain zero placeholders with their usual structure; eager backend skips return None as the step output. .. py:method:: run(initial_state = None, observations = None, num_steps = None, rng = None, shared_context = None, metadata = None, outputs = None, runtime = None) Run the SMC algorithm using a mutable random generator. .. py:method:: initialize_functional(initial_state, context) Return shape-complete state and context for functional execution. .. py:method:: run_functional(initial_state = None, observations = None, num_steps = None, random_state = None, shared_context = None, metadata = None, outputs = None, runtime = None, iteration = None) Run the SMC algorithm with explicit random-state threading. ``iteration`` may be an externally transformed version of :meth:`functional_iteration`, for example ``jax.jit(smc.functional_iteration)``. :returns: * **result** (*SMCResult*) -- The completed SMC result. * **random_state** (*object*) -- The final backend-native random state. .. py:method:: resume_from(checkpoint, *, rng = None, shared_context = None) Reconstruct eager runtime state from a checkpoint. .. py:method:: resume_functional_from(checkpoint, *, random_state = None, shared_context = None) Reconstruct functional runtime state from a checkpoint.