tud_lbm.pipeline.state

Dynamic simulation state for TUD-LBM.

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

Also provides state builder functions build_optional_fields() and build_extra_state() that orchestrate state composition.

Public API:

from state import State, WettingState, build_optional_fields, build_extra_state

Submodules

Functions

`build_optional_fields`(...)	Build the optional `force` and `force_ext` fields.
`build_extra_state`(→ dict[str, Any])	Collect extra State fields initialised by active extra-state plugins.
`update_extra_state`(→ state.State)	Apply plugin-driven extra-state updates after one step.

Package Contents

tud_lbm.pipeline.state.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][source]

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.

Parameters:

setup – SimulationSetup.
nx – Grid size in x.
ny – Grid size in y.
nz – Grid size in z.
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.

Return type:

A tuple (force, force_ext)