tud_lbm.operators.force

Force operators — composite builder and ForceParams.

Public API: build_forces(), ForceParams, ForceSetup

Implementation modules are internal; use the factory to access.

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

ForceParams	One pre-built force contribution.
ForceSetup	Container for force definitions and the resolved source-term callable.

Functions

build_forces(→ ForceSetup)	Discover *_force fields on config, build ForceSetup with specs and source term.
compute_total_force_ext(...)	Compute the summed external force contribution.

Package Contents

class tud_lbm.operators.force.ForceParams

Bases: NamedTuple

One pre-built force contribution.

name

Registry key, e.g. "gravity_force".

compute_fn

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

precomputed

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

name: str

compute_fn: Any

precomputed: Any | None = None

class tud_lbm.operators.force.ForceSetup

Bases: 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.

specs

Tuple of ForceParams for each active force.

source_term

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

specs: tuple[ForceParams, Ellipsis]

source_term: collections.abc.Callable[[Any, Any, Any, Any], Any]

tud_lbm.operators.force.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

Parameters:

• config – A validated configuration object with *_force fields.

• grid_shape – Spatial dimensions, e.g. (64, 64, 1) or (64, 64, 32) for (nx, ny, nz).

• lattice – The simulation lattice.

Returns:

A ForceSetup containing force specs and the source-term callable.

tud_lbm.operators.force.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.

Parameters:

• setup – The SimulationSetup.

• state – Current State.

• force_setup – The 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).

Return type:

Tuple of (total_force, updated_state) where