sequince.steps.control
======================

.. py:module:: sequince.steps.control


Attributes
----------

.. autoapisummary::

   sequince.steps.control.logger


Classes
-------

.. autoapisummary::

   sequince.steps.control.ConditionalStepOutput
   sequince.steps.control.ContinueWhile
   sequince.steps.control.StopIf
   sequince.steps.control.PeriodicStep
   sequince.steps.control.RepeatStep
   sequince.steps.control.SequentialStep


Functions
---------

.. autoapisummary::

   sequince.steps.control.apply_functional_if


Module Contents
---------------

.. py:data:: logger

.. py:function:: apply_functional_if(predicate, apply, state, context, random_state)

   Conditionally execute a functional step through backend control flow.

   :param predicate: Execute the step when true; otherwise preserve the entire input carry.
   :type predicate: bool or scalar array
   :param apply: Pure, traceable ``apply(state, context, random_state)`` operation.
   :type apply: callable
   :param state: Initialized particle state, including any evaluation cache.
   :type state: ParticleState
   :param context: Initialized context with all metadata required by the operation.
   :type context: StepContext
   :param random_state: Explicit backend random state.
   :type random_state: object

   :returns: **state, context, random_state, output** -- Executed step results, or unchanged inputs. Skipped output is None on
             eager backends. On JAX it contains zero array placeholders in the step's
             output structure, with static fields unchanged.
   :rtype: tuple

   .. rubric:: Notes

   Uses acai.cond: eager backends execute only the selected branch. On JAX,
   abstract evaluation traces ``apply`` even when skipped, but does not execute
   transitions or consume random draws. Both branches must preserve the carry's
   pytree structure, shapes, and dtypes. Initialize missing leaves with the
   functional initialization hooks before calling this helper. Output leaves
   must support JAX zero construction; placeholders are not step diagnostics.
   PyTorch execution is eager; structured torch.compile is not guaranteed.


.. py:class:: ConditionalStepOutput

   Output from a step that may skip its child.


   .. py:attribute:: executed
      :type:  bool


   .. py:attribute:: child_output
      :type:  Any
      :value: None



.. py:class:: ContinueWhile(predicate)

   Bases: :py:obj:`sequince.steps.base.SMCStep`


   Ensure the SMC loop continues while a predicate holds.

   This step clears any stop requests raised by earlier steps when the supplied
   predicate evaluates to ``True`` for the current ``ParticleState`` and
   ``StepContext``.  When the predicate returns ``False`` the stop flags remain
   untouched, allowing the run to terminate.

   :param predicate: A callable receiving the current ``ParticleState`` and ``StepContext``
                     and returning a boolean indicating whether to continue the SMC loop.


   .. py:method:: apply(state, context, rng)


.. py:class:: StopIf(predicate, *, immediate = False)

   Bases: :py:obj:`sequince.steps.base.SMCStep`


   Request the SMC loop to stop when a predicate holds.

   This step sets a stop request flag in the context when the supplied predicate
   evaluates to ``True`` for the current ``ParticleState`` and ``StepContext``.
   When the predicate returns ``False`` the stop flag remains untouched, allowing
   the run to continue.

   :param predicate: A callable receiving the current ``ParticleState`` and ``StepContext``
                     and returning a boolean indicating whether to request stopping the SMC
                     loop.
   :param immediate: When ``True``, halt the current iteration immediately. When ``False``,
                     stop after the current iteration finishes.


   .. py:method:: apply_functional(state, context, random_state)

      Evaluate a traceable predicate without consuming random state.



   .. py:method:: apply(state, context, rng)


.. py:class:: PeriodicStep(step, n, offset = 0)

   Bases: :py:obj:`sequince.steps.base.SMCStep`


   Apply a step every N iterations.

   This step applies the supplied ``step`` only on iterations where the current
   iteration number is a multiple of ``n``.  On other iterations, it simply
   returns the input state unchanged.

   :param step: The step to apply every N iterations.
   :param n: The number of iterations between applications of the step.
   :param offset: An optional offset to apply to the iteration number before checking for
                  divisibility.  This allows you to, for example, apply a step every 5
                  iterations starting from iteration 2 by setting ``n=5`` and ``offset=2``.


   .. py:attribute:: step


   .. py:attribute:: n


   .. py:attribute:: offset
      :value: 0



   .. py:method:: apply(state, context, rng)


   .. py:method:: state_dict()

      Return mutable state required by subsequent SMC iterations.



   .. py:method:: load_state_dict(state)

      Restore state produced by :meth:`state_dict`.



.. py:class:: RepeatStep(step, n)

   Bases: :py:obj:`sequince.steps.base.SMCStep`


   Repeat a step or sequence of steps a fixed number of times.


   .. py:attribute:: step


   .. py:attribute:: n


   .. py:method:: apply(state, context, rng)


   .. py:method:: state_dict()

      Return mutable state required by subsequent SMC iterations.



   .. py:method:: load_state_dict(state)

      Restore state produced by :meth:`state_dict`.



.. py:class:: SequentialStep(steps)

   Bases: :py:obj:`sequince.steps.base.SMCStep`


   Apply several SMC steps in order.

   This is useful for grouping multiple steps together into a single step,
   which can then be used in other contexts (e.g., as part of a
   `RepeatStep` or `PeriodicStep`).


   .. py:attribute:: steps


   .. py:method:: apply(state, context, rng)


   .. py:method:: state_dict()

      Return mutable state required by subsequent SMC iterations.



   .. py:method:: load_state_dict(state)

      Restore state produced by :meth:`state_dict`.