tud_lbm.pipeline.setup

JAX-friendly SimulationSetup and build_setup factory.

SimulationSetup is an immutable typing.NamedTuple (i.e. a valid JAX pytree) that holds everything needed by the jitted step function. It is built from a validated SimulationConfig via build_setup().

Design rationale

  • Closed over, not ``static_argnums``: SimulationSetup is captured by functools.partial (or a closure) when constructing the scan body. JAX treats it as a compile-time constant and caches the compilation. If the setup changes, a new compilation is triggered — the correct behaviour.

  • No mutable class instances: All fields are Python scalars, tuples, or jax.Array values. No operator objects live here.

Usage:

from config.simulation_config import SimulationConfig
from setup.simulation_setup import build_setup

cfg = SimulationConfig(grid_shape=(64, 64), tau=0.8)
setup = build_setup(cfg)

Classes

SimulationSetup

Immutable operator container — only what the jitted step needs.

Functions

build_setup(→ SimulationSetup)

Construct a JAX-friendly SimulationSetup from a config.

Module Contents

class tud_lbm.pipeline.setup.SimulationSetup[source]

Bases: NamedTuple

Immutable operator container — only what the jitted step needs.

SimulationSetup stores built artifacts (operators, masks, closures) and physics scalars that the step function reads at JIT time. IO and initialisation metadata live on the original SimulationConfig, accessible via the config reference.

config[source]

The original SimulationConfig — for IO, init, and any metadata that does not enter the JIT boundary.

lattice[source]

The Lattice pytree.

grid_shape[source]

Spatial dimensions, e.g. (64, 64).

tau[source]

Relaxation time (> 0.5).

collision_scheme[source]

Name of the collision model ("bgk" / "mrt").

k_diag[source]

MRT relaxation rates (None for BGK).

bc_masks[source]

Pre-computed boundary-condition masks (BCMasks).

forces[source]

Pre-built force setup (ForceSetup) containing specs and source-term callable, or None if no forces are active.

multiphase_params[source]

None for single-phase runs.

gradient_standard[source]

Standard gradient ∇μ (chemical potential). Always used for chemical-potential gradient. Never wetting-corrected.

gradient_density[source]

Density gradient ∇ρ used in source term. Wetting-corrected when applicable.

laplacian_density[source]

Laplacian of density ∇²ρ in chemical-potential computation. Wetting-corrected when applicable.

gradient_density_wetting[source]

Parametric density gradient for hysteresis optimisation. Populated for hysteresis runs. Signature: (grid, phi_l, phi_r, d_rho_l, d_rho_r) -> result. None for non-hysteresis cases.

laplacian_density_wetting[source]

Parametric Laplacian of density for hysteresis optimisation. Populated for hysteresis runs. Signature: (grid, phi_l, phi_r, d_rho_l, d_rho_r) -> result. None for non-hysteresis cases.

step_fn[source]

The unbound step operator resolved from the registry, implementing StepOperator. Signature: (setup, state) → state_next.

wetting_fn[source]

The hysteresis operator for updating wetting state, implementing HysteresisOperator. Built when hysteresis_config is present; None otherwise.

extra_state_plugins[source]

Active plugin tuple used to initialise and update operation-specific extra state (e.g. electric potential, wetting state).

collision_fn[source]

Pre-built collision operator, resolved at setup time.

equilibrium_fn[source]

Pre-built equilibrium operator, resolved at setup time.

macroscopic_fn[source]

Pre-built macroscopic operator, resolved at setup time.

streaming_fn[source]

Pre-built streaming operator, resolved at setup time.

bc_fn[source]

Pre-built boundary-condition operator, resolved at setup time.

initial_f_fn[source]

Pre-built initial population function, resolved at setup time.

config: tud_lbm.config.simulation_config.SimulationConfig[source]
lattice: tud_lbm.lattice.lattice.Lattice[source]
grid_shape: tuple[int, Ellipsis][source]
tau: float[source]
collision_scheme: str[source]
k_diag: tuple[float, Ellipsis] | None = None[source]
bc_masks: tud_lbm.operators.boundary.BCMasks | None = None[source]
forces: tud_lbm.operators.force.ForceSetup | None = None[source]
multiphase_params: tud_lbm.operators.macroscopic.MultiphaseParams | None = None[source]
gradient_standard: tud_lbm.operators.protocols.DifferentialOperator | None = None[source]
gradient_density: tud_lbm.operators.protocols.DifferentialOperator | None = None[source]
laplacian_density: tud_lbm.operators.protocols.DifferentialOperator | None = None[source]
gradient_density_wetting: tud_lbm.operators.protocols.DifferentialOperator | None = None[source]
laplacian_density_wetting: tud_lbm.operators.protocols.DifferentialOperator | None = None[source]
step_fn: tud_lbm.operators.protocols.StepOperator | None = None[source]
wetting_fn: tud_lbm.operators.protocols.HysteresisOperator | None = None[source]
extra_state_plugins: tuple[tud_lbm.operators.protocols.ExtraStatePlugin, Ellipsis] = ()[source]
collision_fn: tud_lbm.operators.protocols.CollisionOperator | None = None[source]
equilibrium_fn: tud_lbm.operators.protocols.EquilibriumOperator | None = None[source]
macroscopic_fn: tud_lbm.operators.protocols.MacroscopicOperator[Ellipsis, tuple[jax.numpy.ndarray, Ellipsis]] | None = None[source]
streaming_fn: tud_lbm.operators.protocols.StreamingOperator | None = None[source]
bc_fn: tud_lbm.operators.protocols.BoundaryOperator | None = None[source]
initial_f_fn: tud_lbm.operators.protocols.InitialPopulationOperator[Ellipsis, jax.numpy.ndarray] | None = None[source]
tud_lbm.pipeline.setup.build_setup(config: tud_lbm.config.simulation_config.SimulationConfig) → SimulationSetup[source]

Construct a JAX-friendly SimulationSetup from a config.

Parameters:

config – A validated SimulationConfig.

Returns:

An immutable SimulationSetup NamedTuple ready for the jitted step function.

Raises:

ValueError – If wetting configuration is present but sim_type is not “multiphase”.