tud_lbm.pipeline.state
======================

.. py:module:: tud_lbm.pipeline.state

.. autoapi-nested-parse::

   Dynamic simulation state for TUD-LBM.

   Provides :class:`State` and :class:`WettingState` — the pytree-compatible
   carry objects used inside ``jax.lax.scan``.

   Also provides state builder functions :func:`build_optional_fields` and
   :func:`build_extra_state` that orchestrate state composition.

   Public API::

       from state import State, WettingState, build_optional_fields, build_extra_state



Submodules
----------

.. toctree::
   :maxdepth: 1

   /autoapi/tud_lbm/pipeline/state/extra_state/index
   /autoapi/tud_lbm/pipeline/state/state/index


Functions
---------

.. autoapisummary::

   tud_lbm.pipeline.state.build_optional_fields
   tud_lbm.pipeline.state.build_extra_state
   tud_lbm.pipeline.state.update_extra_state


Package Contents
----------------

.. py:function:: build_optional_fields(setup: tud_lbm.pipeline.setup.SimulationSetup, nx: int, ny: int, nz: int, d: int) -> tuple[jax.numpy.ndarray | None, jax.numpy.ndarray | None]

   Build the optional ``force`` and ``force_ext`` fields.

   Returns zero-filled arrays for fields that will be written by the
   step function, and ``None`` for fields that remain unused.
   JAX's ``lax.scan`` requires the pytree structure to be constant
   across iterations, so any field that transitions from ``None`` →
   array must start as zeros.

   :param setup: :class:`~setup.simulation_setup.SimulationSetup`.
   :param nx: Grid size in x.
   :param ny: Grid size in y.
   :param nz: Grid size in z.
   :param d: Lattice dimension (e.g. 2 for D2Q9).

   :returns:

             * ``force``: Zeros ``(nx, ny, nz, 1, d)`` for multiphase runs,
               ``None`` otherwise.
             * ``force_ext``: Zeros ``(nx, ny, nz, 1, d)`` when forces are active,
               ``None`` otherwise.
   :rtype: A tuple ``(force, force_ext)``

   .. seealso:: :func:`state._optional_fields._build_optional_fields`


.. py:function:: build_extra_state(setup: tud_lbm.pipeline.setup.SimulationSetup) -> dict[str, Any]

   Collect extra State fields initialised by active extra-state plugins.

   :param setup: :class:`~setup.simulation_setup.SimulationSetup`.

   :returns: A dictionary mapping field names to initialised arrays.
             Empty when no plugins are active.

   .. seealso:: :func:`state._extra_state._build_extra_state`


.. py:function:: update_extra_state(setup: tud_lbm.pipeline.setup.SimulationSetup, prev_state: state.State, new_state: state.State, **context: dict[str, Any]) -> state.State

   Apply plugin-driven extra-state updates after one step.