tud_lbm.operators.force
=======================

.. py:module:: tud_lbm.operators.force

.. autoapi-nested-parse::

   Force operators — composite builder and ForceParams.

   Public API: build_forces(), ForceParams, ForceSetup

   Implementation modules are internal; use the factory to access.

   .. rubric:: Example

   from operators.force import build_forces

   force_setup = build_forces(config, (64, 64))
   specs = force_setup.specs
   source_fn = force_setup.source_term



Classes
-------

.. autoapisummary::

   tud_lbm.operators.force.ForceParams
   tud_lbm.operators.force.ForceSetup


Functions
---------

.. autoapisummary::

   tud_lbm.operators.force.build_forces
   tud_lbm.operators.force.compute_total_force_ext


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

.. py:class:: ForceParams

   Bases: :py:obj:`NamedTuple`


   One pre-built force contribution.

   .. attribute:: name

      Registry key, e.g. ``"gravity_force"``.

   .. attribute:: compute_fn

      Pure function ``(state, precomputed, lattice) → jnp.ndarray``
      of shape ``(nx, ny, nz, 1, d)``.
      Returns the force contribution for this physics.

   .. attribute:: precomputed

      Optional pre-computed data (e.g. gravity template array).


   .. py:attribute:: name
      :type:  str


   .. py:attribute:: compute_fn
      :type:  Any


   .. py:attribute:: precomputed
      :type:  Any | None
      :value: None



.. py:class:: ForceSetup

   Bases: :py:obj:`NamedTuple`


   Container for force definitions and the resolved source-term callable.

   Bundles force specifications with the source-term function to provide
   a unified interface for the rest of the codebase.

   .. attribute:: specs

      Tuple of :class:`ForceParams` for each active force.

   .. attribute:: source_term

      Callable that computes the well-balanced forcing source term,
      signature ``(rho, u, force, lattice, *, gradient) → jnp.ndarray``.


   .. py:attribute:: specs
      :type:  tuple[ForceParams, Ellipsis]


   .. py:attribute:: source_term
      :type:  collections.abc.Callable[[Any, Any, Any, Any], Any]


.. py:function:: build_forces(config: tud_lbm.config.SimulationConfig, grid_shape: tuple[int, Ellipsis], lattice: Any) -> ForceSetup

   Discover ``*_force`` fields on config, build ForceSetup with specs and source term.

   Each force operator in the registry must expose:

   - ``build(params, grid_shape, config, lattice)`` → precomputed data (or None)
   - ``compute(state, precomputed, **kwargs)`` → force array

   :param config: A validated configuration object with ``*_force`` fields.
   :param grid_shape: Spatial dimensions, e.g. ``(64, 64, 1)`` or ``(64, 64, 32)`` for (nx, ny, nz).
   :param lattice: The simulation lattice.

   :returns: A :class:`ForceSetup` containing force specs and the source-term callable.


.. py:function:: compute_total_force_ext(setup: tud_lbm.setup.SimulationSetup, state: tud_lbm.pipeline.state.State, force_setup: ForceSetup | None) -> tuple[jax.numpy.ndarray | None, Any]

   Compute the summed external force contribution.

   Iterates over all force specs in the setup, calls each force's `compute_fn`,
   and accumulates contributions.

   :param setup: The :class:`~setup.simulation_setup.SimulationSetup`.
   :param state: Current :class:`~state.state.State`.
   :param force_setup: The :class:`ForceSetup` containing force specs, or None if no forces.

   :returns:

             - *total_force* is the summed force array, or None if no forces are active.
             - *updated_state* is the unchanged state (extra-state plugins handle updates).
   :rtype: Tuple of ``(total_force, updated_state)`` where