Invariants

def check_accepted(result, config_params, rtol, atol, density_at=None, ice=False):
    """Raise StructureInvariantError when an accepted result breaks an invariant.

    See ``deviations`` for the parameters; the mass limit is ``result['accept_mass_limit']``,
    which the accept path writes. A result with ``converged`` False is not checked.

    Raises
    ------
    StructureInvariantError
        Names the first broken invariant in the order of ``deviations``, its deviation and
        tolerance, the node for the profile invariants, the values that enter it and the
        ``accepted_on`` label; or the profile arrays or layer masses that the result lacks; or
        a malformed result (arrays that are not numeric, of more than one shape or dimension,
        fewer than 2 nodes, a layer mass that is not one real number); or an
        ``accept_mass_limit``, ``rtol`` or ``atol`` that is not a finite positive number.
    """
    if not result.get('converged'):
        return
    accepted_on = result.get('accepted_on')
    missing = [k for k in _ARRAYS + _MASSES if result.get(k) is None]
    if missing:
        raise StructureInvariantError(
            f'Accepted structure ({accepted_on=}) has no {", ".join(missing)}.'
        )
    try:
        shapes = [np.asarray(result[k], float).shape for k in _ARRAYS]
    except (TypeError, ValueError) as err:
        raise StructureInvariantError(
            f'Accepted structure ({accepted_on=}) is malformed: {err}'
        ) from err
    scalars = [_real_scalar(result[k]) for k in _MASSES]
    if len(set(shapes)) != 1 or len(shapes[0]) != 1 or shapes[0][0] < 2 or not all(scalars):
        raise StructureInvariantError(
            f'Accepted structure ({accepted_on=}) is malformed: the profile arrays need one '
            f'1-D shape of at least 2 nodes and cmb_mass, core_mantle_mass one real number '
            f'each; got shapes {shapes}, cmb_mass {result["cmb_mass"]!r}, core_mantle_mass '
            f'{result["core_mantle_mass"]!r}.'
        )
    limit = result.get('accept_mass_limit')
    if not _finite_positive(limit):
        raise StructureInvariantError(
            f'Accepted structure ({accepted_on=}) has no valid accept_mass_limit: {limit!r}.'
        )
    if not (_finite_positive(rtol) and _finite_positive(atol)):
        raise StructureInvariantError(
            f'Accepted structure ({accepted_on=}) has invalid integration tolerances: '
            f'rtol={rtol!r}, atol={atol!r}.'
        )
    details = {}
    for name, (dev, i, tol) in deviations(
        result, config_params, limit, rtol, atol, density_at, ice, details
    ).items():
        if not dev <= tol:
            node = '' if name in _SCALAR or i < 0 else f' at node {i}'
            err = StructureInvariantError(
                f'Accepted structure ({accepted_on=}) breaks {name}: deviation {dev:.3e} > '
                f'tolerance {tol:.3e}{node}; {details[name]}.'
            )
            cause = details.get(f'{name} cause')
            if cause is not None:
                raise err from cause
            raise err

def deviations(
    result, config_params, mass_limit, rtol, atol, density_at=None, ice=False, details=None
):
    """Return the deviation of each invariant of an accepted result.

    Parameters
    ----------
    result : dict
        Result of ``main`` with ``converged`` True.
    config_params : dict
        Configuration; reads ``planet_mass``, ``core_mass_fraction`` and
        ``mantle_mass_fraction``, and ``target_surface_pressure`` for the density check of
        padded nodes.
    mass_limit : float
        Largest relative mass error that the accept path of ``result`` allows.
    rtol, atol : float
        ``relative_tolerance`` and ``absolute_tolerance`` of the structure integration of
        ``result``. The gravity tolerance at node i is ``TOLERANCES['gravity']`` times
        ``rtol + atol (1 / M_i + 1 / g_i)``, with ``g_i = G M_i / r_i^2``. The RK45 dense
        output keeps ``g r^2 - G M`` per step to about 19.9 rtol + 13.6 atol (1/M_i + 1/g_i)
        relative to ``G M``, an estimate from the Butcher tableau. ``g r^2 - G M`` is constant
        under the exact equations for any density, so an error left at a step end stays, and
        relative to ``G M`` it falls outward as ``M_i / M``. The largest node deviation is the
        dense-output error inside a long step: on the PALEOS-2phase adiabat below it is
        ``+6.4e-6 G M`` at r/R 0.22 in the core, 20.8 units, and ``1.4e-7 G M`` at the last node
        with P > 0.
        The factor 100 is therefore an empirical bound: accepted real solves through ``main``
        of 26 setups, 0.1 to 10 M_E, reach 2.83 to 20.78 times ``rtol + atol (1 / M_i +
        1 / g_i)``, a margin of at least 4.8. They cover the analytic, Wolf & Bower (isothermal
        3000 and 4000 K, linear), PALEOS, PALEOS-2phase (linear, adiabatic), PALEOS with 3 %
        Chabrier H and 3-layer PALEOS ice EOS, 150 and 2000 nodes, rtol 1e-5 and 1e-9, Picard,
        Newton and the JAX kernel. The largest is ``input/default.toml`` at 1 M_E with the
        layers of the PALEOS smoke tests (``Seager2007:iron``, ``PALEOS-2phase:MgSiO3``), the
        adiabatic profile and the default tolerances; the smallest is a 3 % Chabrier H mantle
        at rtol 1e-9. At rtol 1e-5 and atol 1e-6 the tolerance at the surface is about 1e-3
        for any node count, so there the check sees gravity faults of about 0.1 % or more;
        at the first nodes of a fine mesh it is several percent (2.1e-2 at node 1 of 2000).
    density_at : callable or None
        ``(P, T, layer) -> rho`` for the arrays of the nodes of one layer, by the EOS path
        that the solver used at the returned nodes; called once per layer that has nodes.
        A reply that is not a finite number breaks the check. The density check is skipped
        when ``density_at`` is None or when ``result['density_from_eos']`` is False.
    ice : bool
        True when the planet has an ice layer.
    details : dict or None
        If given, filled with {invariant: text} naming the values that enter the deviation and
        their expected values, for the error message, and with ``'density_eos cause'``: the
        exception of a density lookup that raised.

    Returns
    -------
    dict
        {invariant: (deviation, node index, tolerance)}. ``pressure_monotone`` uses the nodes
        before the first node with P <= 0, and ``padding`` the nodes from there on, the surface
        past a stop of the integration, which must have P = 0, the mass of that first node
        and, since no mass lies above the stop, the gravity of that mass at their own radius
        (Gauss), not the gravity at the stop; ``gravity`` checks every node but the centre.
        ``positive`` and ``density_eos`` read every node, ``density_eos`` with the lookup of a
        padded node at ``target_surface_pressure``, the surface pressure that the solve targets;
        ``mass_monotone``, the layer masses, the node checks and ``mass_closure`` read all nodes
        or the last one, and ``centre`` reads node 0.
        ``cmb_node`` and ``core_mantle_node`` mirror the PROTEUS reading of the profile: the
        layer boundary is the first node with m >= cmb_mass (core_mantle_mass with an ice
        layer), compared exactly, and a result where no node reaches it would give the caller
        R_core = 0; its node index is then -1. So a standalone core_mass_fraction 1 run passes
        only with a mass surplus, and with an ice layer so does a run with core_mass_fraction
        plus mantle_mass_fraction 1: a mass deficit is a fault even where no caller reads
        that node. A density lookup that raises breaks ``density_eos`` with an infinite
        deviation and the error in its text.
    """
    P, rho, m, r, g, T = (
        np.asarray(result[k], float)
        for k in ('pressure', 'density', 'mass_enclosed', 'radii', 'gravity', 'temperature')
    )
    ok = np.isfinite(P) & (P > 0)
    k = len(P) if ok.all() else int(np.argmin(ok))
    out = {}

    def put(name, dev, i, tol=None, text=''):
        out[name] = (float(dev), int(i), TOLERANCES[name] if tol is None else tol)
        if details is not None:
            details[name] = text

    put(
        'centre',
        np.count_nonzero([r[0], m[0], g[0]]),
        0,
        text=f'radius, mass, gravity {r[0]:.6e}, {m[0]:.6e}, {g[0]:.6e}, expected 0',
    )
    bad = ~(np.isfinite(rho) & (rho > 0) & np.isfinite(T) & (T > 0))
    i = int(np.argmax(bad)) if bad.any() else k
    put(
        'positive',
        np.count_nonzero(bad) + (k < 2),
        i,
        text=f'density {rho[i]:.6e}, temperature {T[i]:.6e}'
        if bad.any()
        else f'{k} nodes with P > 0',
    )
    # the solver pads a surface stop with zero pressure, the stop mass and its gravity at r
    with np.errstate(divide='ignore', invalid='ignore'):
        gauss = np.abs(g[k:] * r[k:] ** 2 / (G * m[k : k + 1]) - 1) <= _RTOL_GAUSS
    off = (P[k:] != 0) | (m[k:] != m[k : k + 1]) | ~gauss
    i = k + int(np.argmax(off)) if off.any() else k
    put(
        'padding',
        np.count_nonzero(P[k:] != 0)
        + np.count_nonzero(m[k:] != m[k : k + 1])
        + np.count_nonzero(~gauss),
        i,
        text=f'pressure, mass, gravity {P[i]:.6e}, {m[i]:.6e}, {g[i]:.6e}; expected 0, '
        f'the mass of node {k} and G m / r^2 {G * m[k] / r[i] ** 2:.6e}'
        if off.any()
        else '',
    )
    for name, x, sign in (('pressure_monotone', P[:k], 1), ('mass_monotone', m, -1)):
        rise = sign * np.diff(x) / np.abs(x).max(initial=1.0)
        i = int(np.argmax(rise)) + 1 if len(rise) else 0
        put(
            name,
            rise.max(initial=0.0),
            i,
            text=f'{x[i - 1]:.6e} at node {i - 1}, {x[i]:.6e} here' if i else '',
        )
    M, cmf = config_params['planet_mass'], config_params['core_mass_fraction']
    mmf = config_params.get('mantle_mass_fraction', 0.0)
    for name, f in (('cmb_mass', cmf), ('core_mantle_mass', cmf + mmf)):
        put(
            name,
            abs(result[name] / (f * M) - 1) if f > 0 else float(result[name] != 0),
            len(m) - 1,
            text=f'{name} {result[name]:.13e} kg, expected {f * M:.13e} kg',
        )
    # PROTEUS takes the CMB (and with an ice layer the mantle top) at argmax(m >= mass)
    for name, f, key in (
        ('cmb_node', cmf, 'cmb_mass'),
        ('core_mantle_node', cmf + mmf, 'core_mantle_mass'),
    ):
        if f > 0 and (name == 'cmb_node' or ice):
            reach = m >= result[key]
            i = int(np.argmax(reach)) if reach.any() else -1
            put(
                name,
                not (reach[i] and r[i] > 0),
                i,
                text=f'{key} {result[key]:.6e} kg, largest mass_enclosed {m.max():.6e} kg, '
                f'radius of the first node at or above it {r[i] if reach[i] else "none"}',
            )
    with np.errstate(divide='ignore', invalid='ignore'):
        # the bound takes G M / r^2, not the gravity under test, which could widen it
        g_ref = G * m[1:] / r[1:] ** 2
        dg = np.abs(g[1:] / g_ref - 1)
        tol_g = TOLERANCES['gravity'] * (rtol + atol * (1 / m[1:] + 1 / g_ref))
        # a node without mass breaks the check
        ok_g = (m[1:] > 0) & np.isfinite(dg)
        ratio = np.where(ok_g, dg / tol_g, np.inf)
        i = int(np.argmax(ratio)) if len(ratio) else -1
        dev, tol = (dg[i], tol_g[i]) if i >= 0 and ok_g[i] else (np.inf, 0.0)
        j = i + 1
        text = (
            f'gravity {g[j]:.6e} m/s^2, expected {G * m[j] / r[j] ** 2:.6e} (G m / r^2)'
            if i >= 0
            else ''
        )
    put('gravity', dev, j, tol, text=text)
    # the form of the accept test
    put(
        'mass_closure',
        abs(m[-1] - M) / M,
        len(m) - 1,
        mass_limit,
        text=f'mass_enclosed[-1] {m[-1]:.6e} kg, target {M:.6e} kg',
    )
    if density_at is not None and result.get('density_from_eos'):
        n = len(P)
        lookup = P.copy()
        if k < n:  # a padded node is the surface, at the target surface pressure
            lookup[k:] = config_params['target_surface_pressure']
        layer = _layer(m, result['cmb_mass'], result['core_mantle_mass'], ice)
        ref = np.full(n, np.nan)
        for name in np.unique(layer):
            sel = layer == name
            try:
                ref[sel] = np.asarray(density_at(lookup[sel], T[sel], str(name)), dtype=float)
            except Exception as err:  # reported as the density_eos deviation
                failed = (
                    f'density lookup of the {name} layer failed: {type(err).__name__}: {err}'
                )
                put('density_eos', np.inf, int(np.argmax(sel)), text=failed)
                if details is not None:
                    details['density_eos cause'] = err
                return out
        with np.errstate(divide='ignore', invalid='ignore'):
            d = np.abs(rho / ref - 1)
        d[~np.isfinite(d)] = np.inf
        i = int(np.argmax(d))
        put(
            'density_eos',
            d.max(initial=0.0),
            i,
            text=f'density {rho[i]:.6e}, expected {ref[i]:.6e} ({layer[i]} EOS)',
        )
    return out