EOS table consistency

def check(table, eps):
    """Test each cell against the Richardson bound of each identity.

    A derivative identity fails where |R_x(h)| > floor_x + FACTOR * est, with R_x(kh) =
    (4 R_k - R_2k) / 3, floor_x = (4 floor_1 + floor_2) / 3 and est = |R_x(2h) - R_x(h)| / 15,
    the measured h^4 term of R_x(h) at the cell. nabla_ad fails where |R| > floor. At 9 stored
    digits (eps 5e-9) the exact ideal gas and Murnaghan tables give no failing cell, and one
    planted value is found at the cells one and two steps from it.

    The floor uses eps_eff = max(eps, EPS_FLOOR): every table is checked at the 9-digit
    floor at best, so a single value off by about 1e-8 (relative) or less is out of reach on any
    table. Each identity reads only its own columns (``INPUTS``). A derivative identity is
    tested only where 4 cells on each side along its axis hold its differentiated column in one
    phase; NaN marks a missing value, and a stored 0 or inf is a value. Where all values an
    identity reads are present but its residual or bound is not finite, the cell fails; an
    identity is skipped only when no cell holds all values it reads, on the full stencil of a
    derivative identity (4 cells on each side in one phase). On a coarse grid a cell can lie
    outside the range where the Richardson estimate holds (R_x(2h) / R_x(h) far from 16), and
    an exact table can fail there: 1 cp cell of the Murnaghan table at 10 points per decade over
    P 1e5 to 10^11.5 Pa and T 300 to 7.5e3 K, and 29 cp_cv cells of the 60 x 60 Murnaghan table
    over P 1e5 to 1e11 Pa and T 500 to 5000 K, where cp and cv are positive.

    Parameters
    ----------
    table : dict
        As for ``residuals``, on a grid uniform in log10 P and log10 T: each axis within
        ``GRID_FIT`` of a linear fit, as 9-digit coordinates are.
    eps : float
        Relative precision of the stored columns (0.5 x 10^(1 - n) for n significant digits).

    Returns
    -------
    dict
        Identity name -> dict with ``rx`` (the Richardson residual, defined also at untested
        cells), ``bound``, ``present`` (all values the identity reads on its stencil),
        ``tested`` (both finite), ``blind`` (tested with a bound of 1 or more, where a missing
        derivative passes), ``fail`` (tested and above the bound, or present with a residual
        or bound that is not finite) arrays, and ``skipped``, True when no cell is present.

    Raises
    ------
    ValueError
        If eps is not a real number in [0, 1), an axis has fewer than 2 points or is not
        finite, increasing and uniform, a column is off the shape of the grid, or every
        identity is skipped.
    """
    eps = _eps(eps)
    r = {k: residuals(table, k) for k in (1, 2, 4)}
    const = {**table}
    for col, value in _PRESENT.items():
        const[col] = np.where(np.isnan(np.asarray(table[col], float)), np.nan, value)
    p = {k: residuals(const, k) for k in (1, 2, 4)}
    f1, f2 = floors(table, eps, 1), floors(table, eps, 2)
    with np.errstate(invalid='ignore'):
        return _nonempty(
            {
                n: _verdict(r[1][n], f1[n], np.isfinite(p[1][n]))
                if n == 'nabla_ad'
                else _richardson(
                    r[1][n], r[2][n], r[4][n], f1[n], f2[n], p[1][n], p[2][n], p[4][n]
                )
                for n in IDENTITIES
            }
        )

def check_pt(tables, eps):
    """``check`` for the Aragog P-T tables that eos_export writes (no u, cv or nabla_ad).

    Parameters
    ----------
    tables : dict
        ``PT_QUANTITIES`` name -> dict with ``P`` (Pa) and ``T`` (K) grid vectors and ``values``
        of shape (n_P, n_T); the fill value ``PT_FILL`` counts as a missing cell.
    eps : float
        Relative precision of the stored values; raised to ``EPS_FLOOR`` at least.

    Returns
    -------
    dict
        As for ``check``, for the identities cp = T (ds/dT)_P, (ds/dP)_T = -alpha / rho and
        alpha = -(dln rho/dT)_P, each with ``fill``: quantity -> number of ``PT_FILL`` cells.

    Raises
    ------
    ValueError
        If eps is not a real number in [0, 1), an axis is not 1-D, has fewer than 2 points, is
        not finite and increasing, or departs from a uniform or geometric grid by more than
        ``PT_GRID_FIT`` of its step plus its stored precision, T is not positive or P is
        negative, the four tables are not on one grid or off its shape, or every identity is
        skipped.

    Notes
    -----
    The floor includes the rounding of the stored P and T coordinates: a central difference
    over x[i-k], x[i+k] gains |D| (e[i+k] + e[i-k]) / (x[i+k] - x[i-k]), with e the departure
    of each coordinate from the uniform or geometric fit of its axis plus 4 ulp.
    """
    eps = _eps(eps)
    parts = {k: _pt_parts(tables, k, eps) for k in (1, 2, 4)}
    const = {}
    for q, t in tables.items():
        v = np.asarray(t['values'], float)
        const[q] = {**t, 'values': np.where(np.isnan(v) | (v == PT_FILL), PT_FILL, 1.0)}
    present = {k: _pt_parts(const, k, eps) for k in (1, 2, 4)}
    fill = {
        q: int((np.asarray(tables[q]['values'], float) == PT_FILL).sum()) for q in PT_QUANTITIES
    }
    out = {}
    with np.errstate(divide='ignore', invalid='ignore'):
        for n in PT_IDENTITIES:
            p = [present[k][n][0] / present[k][n][2] for k in (1, 2, 4)]
            r = {k: parts[k][n][0] / parts[k][n][2] - 1 for k in (1, 2, 4)}
            fl = {
                k: (parts[k][n][1] / np.abs(parts[k][n][2]) + np.abs(r[k] + 1) * parts[k][n][3])
                * eps
                + np.abs(r[k] + 1) * parts[k][n][4]
                for k in (1, 2)
            }
            out[n] = {**_richardson(r[1], r[2], r[4], fl[1], fl[2], *p), 'fill': dict(fill)}
    return _nonempty(out)

def floors(table, eps, k=1):
    """Round-off bound of each residual at step ``k`` for columns stored to relative ``eps``.

    For R + 1 = D / Q the bound is (E / |Q| + |D / Q| q) eps + |D / Q| c, with E the error of
    the central difference from its two stored ends, q the relative error of Q, and c the
    relative error from the stored P and T coordinates: in Q (e_P and e_T, ln 10 times the
    departure of log10 P and log10 T from their linear fit) and in the mean step of the
    derivative axis (2 eps / span, span = (n - 1) h in ln units). nabla_ad, with no
    derivative, gets |R + 1| (4 eps + e_P) from its four stored factors and P.
    """
    hx, hy, P, _, rho, _, _, cp, _, alpha, nab = _grid(table)
    sP, sT = 2 * eps / (hx * (rho.shape[0] - 1)), 2 * eps / (hy * (rho.shape[1] - 1))
    eP = np.log(10.0) * _departure(np.asarray(table['unique_log_p'], float))[:, None]
    eT = np.log(10.0) * _departure(np.asarray(table['unique_log_t'], float))[None, :]
    with np.errstate(divide='ignore', invalid='ignore'):
        a = P * alpha / rho
        c = {
            'cp': sT,
            'maxwell': eP + sP,
            'alpha': eT + sT,
            'dudt': eT + sT + eP * np.abs(a / (cp - a)),
            'cp_cv': eP + eT + sP,
        }
        out = {
            n: (E / np.abs(Q) + np.abs(D / Q) * q) * eps + np.abs(D / Q) * c[n]
            for n, (D, E, Q, q) in _parts(table, k).items()
        }
        out['nabla_ad'] = np.abs(nab * rho * cp / (P * alpha)) * (4 * eps + eP)
    return {n: out[n] for n in IDENTITIES}

def signs(table, eps):
    """Cells that break a sign condition beyond the stored precision ``eps``.

    Parameters
    ----------
    table : dict
        As for ``residuals``.
    eps : float
        Relative precision of the stored columns; raised to ``EPS_FLOOR`` at least.

    Returns
    -------
    dict
        Condition name -> dict with ``tested`` and ``fail`` boolean arrays. ``rho``, ``cp``
        and ``cv`` not positive; ``cp_ge_cv``: cv above cp by more than eps (|cp| + |cv|);
        ``kappa_T``: rho falls with P at fixed T by more than the rounding of its two
        neighbours (central difference over one cell on each side, one phase). A condition
        tests the cells where its values are not NaN; an infinite value breaks the rho, cp, cv
        and cp_ge_cv conditions, and kappa_T does not test it.

    Raises
    ------
    ValueError
        If eps is not a real number in [0, 1), or the axes or shapes fail as for ``check``.
    """
    eps = _eps(eps)
    _axes(table)
    rho, cp, cv = (np.asarray(table[c], float) for c in ('rho', 'cp', 'cv'))
    seg = segments(table, 0, ('rho',))
    rise = _stencil(rho, seg, 0, 1, 0.5, -1.0)  # rho[i+1] - rho[i-1]
    noise = _stencil(np.abs(rho), seg, 0, 1, 0.5, 1.0) * eps
    with np.errstate(invalid='ignore'):
        cond = {
            'rho': (rho, rho <= 0),
            'cp': (cp, cp <= 0),
            'cv': (cv, cv <= 0),
            'cp_ge_cv': (cp - cv, cv - cp > eps * (np.abs(cp) + np.abs(cv))),
            'kappa_T': (rise, rise < -noise),
        }
    return {
        n: {'tested': ~np.isnan(v), 'fail': ~np.isnan(v) & (f | np.isinf(v))}
        for n, (v, f) in cond.items()
    }

def residuals(table, k=1):
    """Relative residuals of the identities at stencil step ``k`` cells.

    Parameters
    ----------
    table : dict
        ``unique_log_p``, ``unique_log_t`` (uniform, log10 Pa and K) and the columns rho, u, s,
        cp, cv, alpha, nabla_ad (SI) and phase on the (P, T) grid, as from
        ``zalmoxis.eos_export.load_paleos_all_properties``.
    k : int
        Stencil step in grid cells.

    Returns
    -------
    dict
        Identity name -> residual array on the grid (NaN where it cannot be formed):
        ``cp``: (ds/dlnT)_P / cp - 1; ``maxwell``: (ds/dlnP)_T / (-P alpha / rho) - 1;
        ``alpha``: -(dln rho/dlnT)_P / (T alpha) - 1; ``dudt``: (du/dlnT)_P / T / (cp - P alpha
        / rho) - 1; ``nabla_ad``: nabla_ad rho cp / (P alpha) - 1 (no derivative);
        ``cp_cv``: (cp - cv) rho kappa_T / (T alpha^2) - 1 with kappa_T = (dln rho/dlnP)_T / P.
    """
    _, _, P, _, rho, _, _, cp, _, alpha, nab = _grid(table)
    parts = _parts(table, k)
    with np.errstate(divide='ignore', invalid='ignore'):
        out = {n: D / Q - 1 for n, (D, _, Q, _) in parts.items()}
        out['nabla_ad'] = nab * rho * cp / (P * alpha) - 1
    return {n: out[n] for n in IDENTITIES}

def derivative(f, seg, axis, k, h):
    """d f / d(axis) with step ``k h``, NaN where the stencil leaves one segment."""
    return _stencil(f, seg, axis, k, h, -1.0)

def segments(table, axis, cols=COLUMNS, valid=np.isfinite, code=None):
    """Label runs of cells of one phase along ``axis``; cells where ``valid`` (a function of a
    column) is false for one of ``cols`` get -1. ``code`` is ``_phase_code(table)`` if given."""
    ok = np.logical_and.reduce([valid(np.asarray(table[c], float)) for c in cols])
    code = np.where(ok, _phase_code(table) if code is None else code, -1)
    start = np.ones(code.shape, bool)
    inner = tuple(slice(1, None) if a == axis else slice(None) for a in range(code.ndim))
    start[inner] = np.diff(code, axis=axis) != 0
    return np.where(ok, np.cumsum(start, axis=axis), -1)