tud_lbm.io.analysis.stability
=============================

.. py:module:: tud_lbm.io.analysis.stability

.. autoapi-nested-parse::

   Stability diagnostics for the ``lax.scan`` time loop.

   Enabled via the ``--debug-stability`` CLI flag, which sets
   ``DEBUG_FLAG_STABILITY`` in :mod:`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 :class:`StabilityAbortError` when any metric is NaN.



Attributes
----------

.. autoapisummary::

   tud_lbm.io.analysis.stability.logger


Exceptions
----------

.. autoapisummary::

   tud_lbm.io.analysis.stability.StabilityAbortError


Functions
---------

.. autoapisummary::

   tud_lbm.io.analysis.stability.checkerboard_amplitude
   tud_lbm.io.analysis.stability.wake_mask
   tud_lbm.io.analysis.stability.compute_stability_metrics
   tud_lbm.io.analysis.stability.make_stability_callback


Module Contents
---------------

.. py:data:: logger

.. py:exception:: StabilityAbortError

   Bases: :py:obj:`FloatingPointError`


   Raised by the stability NaN guard; aborts the scan.


.. py:function:: checkerboard_amplitude(rho: jax.numpy.ndarray, mask: jax.numpy.ndarray) -> jax.numpy.ndarray

   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.

   :param rho: Density field, shape ``(nx, ny, nz, 1, 1)``.
   :param mask: Boolean mask, broadcastable to ``rho``.

   :returns: Scalar checkerboard amplitude.


.. py:function:: 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

   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).

   :param rho: Density field, shape ``(nx, ny, nz, 1, 1)``.
   :param grad_rho: Density gradient, shape ``(nx, ny, nz, 1, d)``.
   :param rho_l: Liquid coexistence density.
   :param rho_v: Vapor coexistence density.
   :param vapor_frac: Vapor threshold — ``rho < rho_v + vapor_frac * (rho_l - rho_v)``.
   :param grad_frac: Interface exclusion — ``|grad rho| < grad_frac * (rho_l - rho_v)``.

   :returns: Boolean mask, shape ``(nx, ny, nz, 1, 1)``.


.. py:function:: 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

   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]``.


.. py:function:: 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

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

   Mirrors :func:`~tud_lbm.io.callbacks.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.

   :param out_dir: Directory for ``stability_log.csv``.
   :param gradient_density: ``setup.gradient_density`` (``None`` for single-phase).
   :param mp: ``setup.multiphase_params`` (``None`` for single-phase).
   :param log_interval: Steps between samples (the run's ``save_interval``).
   :param vapor_frac: See :func:`wake_mask`.
   :param grad_frac: See :func:`wake_mask`.

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