tud_lbm.operators.wetting.hysteresis.hysteresis

Wetting hysteresis optimisation — pure functions.

Ported from update_timestep.UpdateMultiphaseHysteresis.

The legacy class stores mutable wetting parameters on self and uses @partial(jit, static_argnums=(0,)) which causes JIT cache bloat. This module replaces it with pure functions that operate on the WettingState NamedTuple carried through jax.lax.scan.

All inner optimisation loops use optax + jax.lax.while_loop with early convergence exit and are fully jittable.

Design

update_wetting_state is the top-level entry point. It:

  1. Measures contact angles and contact-line locations from rho_t_plus1.

  2. Checks whether each side is inside the hysteresis window.

  3. Builds per-side objectives (CLL-pin or CA-target, selected via jnp.where).

  4. Runs two sequential jax.lax.while_loop optimisations — one for each side — masking the other side’s parameters.

  5. Returns an updated WettingState — no mutation.

The inner _evaluate_with_params closure performs a single LBM step with trial wetting parameters so that jax.value_and_grad can differentiate through it.

Functions

update_wetting_state(...)

Pure JAX update of wetting / hysteresis parameters.

update_wetting_state_chemical_step(...)

Hysteresis update where CA targets are determined per-side by chemical step position.

Module Contents

tud_lbm.operators.wetting.hysteresis.hysteresis.update_wetting_state(wetting: tud_lbm.pipeline.state.state.WettingState, rho_t_plus1: jax.numpy.ndarray, setup: tud_lbm.pipeline.setup.SimulationSetup, *, trial_step_fn: collections.abc.Callable[[tud_lbm.operators.wetting._params.WettingParams], tuple[jax.numpy.ndarray, jax.numpy.ndarray]] | None = None) → tud_lbm.pipeline.state.state.WettingState[source]

Pure JAX update of wetting / hysteresis parameters.

This replaces the mutable __call__ method. It operates entirely on the WettingState NamedTuple and returns a new instance — no side-effects.

Each side (left and right) is optimised independently in its own jax.lax.while_loop, masking out the other side’s parameters. This gives each side clean gradients and clean Adam state at the cost of two trial-step evaluations per outer iteration instead of one.

Hysteresis regime is selected strictly from the current contact angle: pinned when ca_receding <= ca <= ca_advancing; otherwise CA-targeting toward the exceeded bound.

Parameters:

  • wetting – Current WettingState.

  • rho_t_plus1 – Density field, shape (nx, ny, nz, 1, 1).

  • setup – SimulationSetup (closed-over, not traced).

  • f_t – Pre-step populations, shape (nx, ny, nz, q, 1).

  • force_tot – Total force from the current step (optional).

  • trial_step_fn – Callable (WettingParams) → (f_out, rho_out) that evaluates a single multiphase physics pass with trial wetting parameters. Required; no default is provided.

Returns:

Updated WettingState.

tud_lbm.operators.wetting.hysteresis.hysteresis.update_wetting_state_chemical_step(wetting: tud_lbm.pipeline.state.state.WettingState, rho_t_plus1: jax.numpy.ndarray, setup: tud_lbm.pipeline.setup.SimulationSetup, *, trial_step_fn: collections.abc.Callable[[tud_lbm.operators.wetting._params.WettingParams], tuple[jax.numpy.ndarray, jax.numpy.ndarray]]) → tud_lbm.pipeline.state.state.WettingState[source]

Hysteresis update where CA targets are determined per-side by chemical step position.

Identical to update_wetting_state except (ca_advancing, ca_receding) for each side is selected by comparing that side’s contact-line location against the chemical step position, rather than using a single global hysteresis window.

Parameters:

  • wetting – Current WettingState. wetting.cll_left and wetting.cll_right are used to determine which region each contact line occupies.

  • rho_t_plus1 – Post-step density field.

  • setup – SimulationSetup. Must carry chemical_step_config with keys: chemical_step_location, ca_advancing_pre_step, ca_receding_pre_step, ca_advancing_post_step, ca_receding_post_step.

  • trial_step_fn – Callable (WettingParams) -> (f_out, rho_out). Provided by the step function via partial application.

Returns:

Updated WettingState with optimised wetting parameters and measured CA/CLL.