tud_lbm.io.analysis.stability

Stability diagnostics for the lax.scan time loop.

Enabled via the --debug-stability CLI flag, which sets DEBUG_FLAG_STABILITY in tud_lbm.config.config_overview. Every save_interval steps the runner samples a small vector of on-device metrics and ships it to a host callback:

  • max|u| — maximum velocity magnitude

  • max|grad mu| — maximum chemical-potential gradient magnitude, derived from the stored interaction force (grad mu = -F_int / rho)

  • min rho / max rho — density range

  • checkerboard amplitude — L2 norm of rho - 3x3-smoothed rho restricted to the droplet wake (vapor phase, interface excluded)

The host callback appends one row per sample to stability_log.csv (flushed per write, so the curve survives a crash), prints a status line, and raises StabilityAbortError when any metric is NaN.

Attributes

logger

Exceptions

StabilityAbortError

Raised by the stability NaN guard; aborts the scan.

Functions

checkerboard_amplitude(→ jax.numpy.ndarray)

L2 norm of rho minus its 3x3 xy-uniform-smoothed version, on mask.

wake_mask(→ jax.numpy.ndarray)

Boolean mask for the droplet wake: vapor phase, interface excluded.

compute_stability_metrics(→ jax.numpy.ndarray)

Compute the stability-metric vector for one state.

make_stability_callback(→ collections.abc.Callable)

Build a stability-diagnostics callback for a lax.scan body.

Module Contents

tud_lbm.io.analysis.stability.logger[source]
exception tud_lbm.io.analysis.stability.StabilityAbortError[source]

Bases: FloatingPointError

Raised by the stability NaN guard; aborts the scan.

tud_lbm.io.analysis.stability.checkerboard_amplitude(rho: jax.numpy.ndarray, mask: jax.numpy.ndarray) → jax.numpy.ndarray[source]

L2 norm of rho minus its 3x3 xy-uniform-smoothed version, on mask.

The smoothing is a separable 3-point mean along the x and y axes (z untouched — for 2D simulations nz=1). Rolls wrap at domain edges; the wake mask excludes wall/interface layers where that matters.

Parameters:

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

  • mask – Boolean mask, broadcastable to rho.

Returns:

Scalar checkerboard amplitude.

tud_lbm.io.analysis.stability.wake_mask(rho: jax.numpy.ndarray, grad_rho: jax.numpy.ndarray, rho_l: float, rho_v: float, vapor_frac: float, grad_frac: float) → jax.numpy.ndarray[source]

Boolean mask for the droplet wake: vapor phase, interface excluded.

Recomputed at every sample, so the region follows the droplet. A cell is in the wake when its density is close to the vapor density AND the local density gradient is flat (excluding the diffuse interface, whose steep but smooth variation would dominate the checkerboard residual).

Parameters:

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

  • grad_rho – Density gradient, shape (nx, ny, nz, 1, d).

  • rho_l – Liquid coexistence density.

  • rho_v – Vapor coexistence density.

  • vapor_frac – Vapor threshold — rho < rho_v + vapor_frac * (rho_l - rho_v).

  • grad_frac – Interface exclusion — |grad rho| < grad_frac * (rho_l - rho_v).

Returns:

Boolean mask, shape (nx, ny, nz, 1, 1).

tud_lbm.io.analysis.stability.compute_stability_metrics(state: tud_lbm.pipeline.state.state.State, *, gradient_density: tud_lbm.operators.protocols.DifferentialOperator | None = None, mp: tud_lbm.operators.macroscopic.MultiphaseParams | None = None, vapor_frac: float = 0.2, grad_frac: float = 0.05) → jax.numpy.ndarray[source]

Compute the stability-metric vector for one state.

max|grad mu| is derived from the stored force rather than recomputed: the multiphase step stores force = -rho*grad(mu) + force_ext with force_ext kept separately, so grad mu = -(force - force_ext) / rho. Single-phase states (force is None) report 0.

Without multiphase parameters the checkerboard mask falls back to the whole domain.

Returns:

Flat vector [max_u, max_grad_mu, rho_min, rho_max, checkerboard_amp, n_wake_cells].

tud_lbm.io.analysis.stability.make_stability_callback(out_dir: str | pathlib.Path, *, gradient_density: tud_lbm.operators.protocols.DifferentialOperator | None, mp: tud_lbm.operators.macroscopic.MultiphaseParams | None, log_interval: int, vapor_frac: float, grad_frac: float) → collections.abc.Callable[source]

Build a stability-diagnostics callback for a lax.scan body.

Mirrors make_save_callback(): interval checking runs on-device via jax.lax.cond and the metric reductions live inside the true branch, so off-interval steps cost one integer compare.

Parameters:

  • out_dir – Directory for stability_log.csv.

  • gradient_density – setup.gradient_density (None for single-phase).

  • mp – setup.multiphase_params (None for single-phase).

  • log_interval – Steps between samples (the run’s save_interval).

  • vapor_frac – See wake_mask().

  • grad_frac – See wake_mask().

Returns:

do_check(state, t) — no return value; logs as a side effect and aborts the scan on NaN.