tud_lbm.pipeline.setup ====================== .. py:module:: tud_lbm.pipeline.setup .. autoapi-nested-parse:: JAX-friendly SimulationSetup and ``build_setup`` factory. :class:`SimulationSetup` is an immutable :class:`typing.NamedTuple` (i.e. a valid JAX pytree) that holds everything needed by the jitted step function. It is built from a validated :class:`~config.simulation_config.SimulationConfig` via :func:`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 ------- .. autoapisummary:: tud_lbm.pipeline.setup.SimulationSetup Functions --------- .. autoapisummary:: tud_lbm.pipeline.setup.build_setup Module Contents --------------- .. py:class:: SimulationSetup Bases: :py:obj:`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 :class:`~config.simulation_config.SimulationConfig`, accessible via the :attr:`config` reference. .. attribute:: config The original :class:`SimulationConfig` — for IO, init, and any metadata that does not enter the JIT boundary. .. attribute:: lattice The :class:`~setup.lattice.Lattice` pytree. .. attribute:: grid_shape Spatial dimensions, e.g. ``(64, 64)``. .. attribute:: tau Relaxation time (> 0.5). .. attribute:: collision_scheme Name of the collision model (``"bgk"`` / ``"mrt"``). .. attribute:: k_diag MRT relaxation rates (``None`` for BGK). .. attribute:: bc_masks Pre-computed boundary-condition masks (:class:`BCMasks`). .. attribute:: forces Pre-built force setup (:class:`~operators.force.ForceSetup`) containing specs and source-term callable, or ``None`` if no forces are active. .. attribute:: multiphase_params ``None`` for single-phase runs. .. attribute:: gradient_standard Standard gradient ``∇μ`` (chemical potential). Always used for chemical-potential gradient. Never wetting-corrected. .. attribute:: gradient_density Density gradient ``∇ρ`` used in source term. Wetting-corrected when applicable. .. attribute:: laplacian_density Laplacian of density ``∇²ρ`` in chemical-potential computation. Wetting-corrected when applicable. .. attribute:: gradient_density_wetting 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. .. attribute:: laplacian_density_wetting 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. .. attribute:: step_fn The unbound step operator resolved from the registry, implementing :class:`~operators.protocols.StepOperator`. Signature: ``(setup, state) → state_next``. .. attribute:: wetting_fn The hysteresis operator for updating wetting state, implementing :class:`~operators.protocols.HysteresisOperator`. Built when ``hysteresis_config`` is present; ``None`` otherwise. .. attribute:: extra_state_plugins Active plugin tuple used to initialise and update operation-specific extra state (e.g. electric potential, wetting state). .. attribute:: collision_fn Pre-built collision operator, resolved at setup time. .. attribute:: equilibrium_fn Pre-built equilibrium operator, resolved at setup time. .. attribute:: macroscopic_fn Pre-built macroscopic operator, resolved at setup time. .. attribute:: streaming_fn Pre-built streaming operator, resolved at setup time. .. attribute:: bc_fn Pre-built boundary-condition operator, resolved at setup time. .. attribute:: initial_f_fn Pre-built initial population function, resolved at setup time. .. py:attribute:: config :type: tud_lbm.config.simulation_config.SimulationConfig .. py:attribute:: lattice :type: tud_lbm.lattice.lattice.Lattice .. py:attribute:: grid_shape :type: tuple[int, Ellipsis] .. py:attribute:: tau :type: float .. py:attribute:: collision_scheme :type: str .. py:attribute:: k_diag :type: tuple[float, Ellipsis] | None :value: None .. py:attribute:: bc_masks :type: tud_lbm.operators.boundary.BCMasks | None :value: None .. py:attribute:: forces :type: tud_lbm.operators.force.ForceSetup | None :value: None .. py:attribute:: multiphase_params :type: tud_lbm.operators.macroscopic.MultiphaseParams | None :value: None .. py:attribute:: gradient_standard :type: tud_lbm.operators.protocols.DifferentialOperator | None :value: None .. py:attribute:: gradient_density :type: tud_lbm.operators.protocols.DifferentialOperator | None :value: None .. py:attribute:: laplacian_density :type: tud_lbm.operators.protocols.DifferentialOperator | None :value: None .. py:attribute:: gradient_density_wetting :type: tud_lbm.operators.protocols.DifferentialOperator | None :value: None .. py:attribute:: laplacian_density_wetting :type: tud_lbm.operators.protocols.DifferentialOperator | None :value: None .. py:attribute:: step_fn :type: tud_lbm.operators.protocols.StepOperator | None :value: None .. py:attribute:: wetting_fn :type: tud_lbm.operators.protocols.HysteresisOperator | None :value: None .. py:attribute:: extra_state_plugins :type: tuple[tud_lbm.operators.protocols.ExtraStatePlugin, Ellipsis] :value: () .. py:attribute:: collision_fn :type: tud_lbm.operators.protocols.CollisionOperator | None :value: None .. py:attribute:: equilibrium_fn :type: tud_lbm.operators.protocols.EquilibriumOperator | None :value: None .. py:attribute:: macroscopic_fn :type: tud_lbm.operators.protocols.MacroscopicOperator[Ellipsis, tuple[jax.numpy.ndarray, Ellipsis]] | None :value: None .. py:attribute:: streaming_fn :type: tud_lbm.operators.protocols.StreamingOperator | None :value: None .. py:attribute:: bc_fn :type: tud_lbm.operators.protocols.BoundaryOperator | None :value: None .. py:attribute:: initial_f_fn :type: tud_lbm.operators.protocols.InitialPopulationOperator[Ellipsis, jax.numpy.ndarray] | None :value: None .. py:function:: build_setup(config: tud_lbm.config.simulation_config.SimulationConfig) -> SimulationSetup Construct a JAX-friendly :class:`SimulationSetup` from a config. :param config: A validated :class:`SimulationConfig`. :returns: An immutable :class:`SimulationSetup` NamedTuple ready for the jitted step function. :raises ValueError: If wetting configuration is present but sim_type is not "multiphase".