zephyrus.diffusion

def diameters() -> dict[str, float]:
    """Kinetic diameters in pm: printed values plus the van der Waals rule."""
    d = dict(D_ZK23)
    d.update(D_STANDARD_EXTRA)
    for s in _VDW_SCALED:
        d[s] = D_ZK23['O'] * BONDI_VDW_RADIUS_A[s] / BONDI_VDW_RADIUS_A['O']
    return d

def b_from_sn88(f_value: float, T_of_table: float) -> float:
    """Convert a Sasaki & Nakazawa (1988) Table 1 entry to b = n D in cm^-1 s^-1."""
    return f_value / (KB_SI * T_of_table) / 100.0

def sn88_fit(species: str, partner: str, t_lo: float = 1000, t_hi: float = 10000):
    """Fit b = A T^s to two Sasaki & Nakazawa temperature points.

    Defaults to the 1000 and 10000 K points, both above the 242 K validity
    floor their footnote attaches to the coldest column. Returns ``(A, s)``.
    """
    row = SN88_TABLE1[(species, partner)]
    b_lo = b_from_sn88(row[t_lo], t_lo)
    b_hi = b_from_sn88(row[t_hi], t_hi)
    s = math.log(b_hi / b_lo) / math.log(t_hi / t_lo)
    return b_lo / t_lo**s, s

def b_zk86(species: str, partner: str, T: float) -> float:
    """Zahnle & Kasting (1986) Table I fit b = A T^s, in cm^-1 s^-1."""
    a, s = ZK86_TABLE1[(species, partner)]
    return a * T**s

def eq10(b_anchor: float, target: tuple, anchor: tuple, mass: dict, diam: dict) -> float:
    """Zahnle & Kasting (2023) Eq. (10): scale b between pairs.

    ``b_ij = b_kl sqrt((1/m_i + 1/m_j) / (1/m_k + 1/m_l))
    ((d_k + d_l) / (d_i + d_j))^2``, the reduced-mass and hard-sphere
    diameter scaling their own estimated rows are built with. Masses in amu
    and diameters in pm (the units cancel).
    """
    i, j = target
    k, l = anchor  # noqa: E741
    rm_t = 1.0 / mass[i] + 1.0 / mass[j]
    rm_a = 1.0 / mass[k] + 1.0 / mass[l]
    return b_anchor * math.sqrt(rm_t / rm_a) * ((diam[k] + diam[l]) / (diam[i] + diam[j])) ** 2

def build_rows(species: list) -> list[Row]:
    """Assemble one Row per unordered pair of ``species``.

    Source order: the Zahnle & Kasting (2023) Table 2 row where one is
    printed (measured or estimated, at its printed value); the Sasaki &
    Nakazawa (1988) noble-gas rows; otherwise the Eq. (10) scaling from the
    measured Table 2 anchors, with the provenance class recording that the
    row is scaled and whether it rests on an estimated diameter. A species
    with no kinetic diameter (Al, P, Cl, K, Ca, Ti) takes the row of its
    nearest-mass covered substitute, in the ``substituted`` class.
    """
    mass, diam = ALL_MASS, diameters()
    rows = []
    for a in range(len(species)):
        for c in range(a + 1, len(species)):
            i, j = species[a], species[c]
            hit = _zk23_lookup((i, j))
            if hit is not None:
                b1000, cls, src = hit
                rows.append(
                    Row(i, j, b1000, ZK23_EXPONENT, f'ZK23 T2 [{cls}]', CLASS_OF_ZK23[cls], src)
                )
                continue
            sn = None
            for a_, b_ in ((i, j), (j, i)):
                if (a_, b_) in SN88_TABLE1 and b_ in ('H2', 'He'):
                    sn = (a_, b_)
            if sn is not None:
                a_fit, s_fit = sn88_fit(*sn)
                rows.append(
                    Row(
                        i,
                        j,
                        a_fit * 1000.0**s_fit,
                        ZK23_EXPONENT,
                        'SN88 T1 [M]',
                        'measured',
                        f'Sasaki & Nakazawa (1988) Table 1; fitted exponent '
                        f'{s_fit:.3f}, entered at 0.75',
                    )
                )
                continue
            if any(s not in mass or s not in diam for s in (i, j)):
                rows.append(_proxy_row(i, j))
                continue
            scaled, submode = _anchors_for((i, j), mass, diam)
            vals = np.array([v for _p, v in scaled])
            b1000 = float(np.exp(np.mean(np.log(vals))))
            est = tuple(s for s in (i, j) if s in _VDW_SCALED)
            rows.append(
                Row(
                    i,
                    j,
                    b1000,
                    ZK23_EXPONENT,
                    f'ZK23 Eq.(10) [{submode}]',
                    'scaled*' if est else 'scaled',
                    'anchors ' + ', '.join(f'{p[0]}-{p[1]}' for p, _ in scaled),
                )
            )
    return rows

def bmatrix(species: list, T: float, rows: list[Row] | None = None) -> np.ndarray:
    """Symmetric b matrix in cm^-1 s^-1, with np.inf on the diagonal.

    The infinite diagonal is the convention the fractionation closure
    expects (a species does not diffuse against itself).
    """
    if rows is None:
        rows = build_rows(species)
    idx = {s: k for k, s in enumerate(species)}
    n = len(species)
    b = np.full((n, n), np.inf)
    for r in rows:
        if r.i in idx and r.j in idx:
            b[idx[r.i], idx[r.j]] = b[idx[r.j], idx[r.i]] = r.b(T)
    off = ~np.eye(n, dtype=bool)
    if not np.all(np.isfinite(b[off])):
        missing = [
            (species[a], species[c])
            for a in range(n)
            for c in range(n)
            if a != c and not np.isfinite(b[a, c])
        ]
        raise ValueError(f'diffusion library incomplete for pairs {missing}')
    return b

def masses_g(species: list) -> np.ndarray:
    """Particle masses in g for the closure, from the shared mass table."""
    missing = [s for s in species if s not in ALL_MASS]
    if missing:
        raise KeyError(f'no mass tabulated for {missing}')
    return np.array([ALL_MASS[s] for s in species]) * AMU_G

def b_pair(sp_i: str, sp_j: str, T: float) -> tuple[float, str]:
    """Binary diffusion parameter b = n D for one pair, in SI [m^-1 s^-1].

    The fallback order, most trusted source first: a measured Zahnle &
    Kasting (2023) Table 2 row; the molecular-background table; the rest
    of the atomic library (estimated Table 2 rows, the Sasaki & Nakazawa
    rows, then the Eq. 10 scaling); as a last resort, the library value of
    the nearest-mass covered species, with the substitution recorded in
    the provenance string. Trailing state annotations in species names are
    stripped.

    Returns ``(b [m^-1 s^-1], provenance string)``.
    """
    key = tuple(sorted((sp_i.split('_')[0], sp_j.split('_')[0])))
    if key in _PAIR_CACHE:
        b1000, s, prov = _PAIR_CACHE[key]
        return b1000 * (T / 1000.0) ** s * 100.0, prov
    a, b = key
    if a == b:
        hit = _zk23_lookup(key)
        if hit is not None:
            b1000, cls, _src = hit
            prov = f'ZK23 T2 [{cls}]'
            _PAIR_CACHE[key] = (b1000, ZK23_EXPONENT, prov)
            return b1000 * (T / 1000.0) ** ZK23_EXPONENT * 100.0, prov
    # Order by provenance class: a measured Table 2 row, then the background
    # compilation (same measured primaries), then anything estimated, so the
    # Eq. (10) scaling never preempts a measured row.
    hit = _zk23_lookup(key)
    if hit is not None and hit[1] == 'M':
        b1000, cls, _src = hit
        prov = f'ZK23 T2 [{cls}]'
        _PAIR_CACHE[key] = (b1000, ZK23_EXPONENT, prov)
        return b1000 * (T / 1000.0) ** ZK23_EXPONENT * 100.0, prov
    fit = MOLECULAR_BACKGROUND.get(key) or MOLECULAR_BACKGROUND.get(key[::-1])
    if fit is not None:
        a_fit, s_fit = fit
        b1000 = a_fit * 1000.0**s_fit
        prov = 'molecular-background table'
        _PAIR_CACHE[key] = (b1000, s_fit, prov)
        return b1000 * (T / 1000.0) ** s_fit * 100.0, prov
    r = build_rows([a, b])[0] if a in ALL_MASS and b in ALL_MASS else _proxy_row(a, b)
    _PAIR_CACHE[key] = (r.b1000, r.exponent, r.provenance)
    return r.b(T) * 100.0, r.provenance

def b_mixture(sp: str, comp: dict[str, float], T: float) -> tuple[float, dict]:
    """Blanc's-law mixture diffusion parameter for one species, SI [m^-1 s^-1].

    ``b_mix = (1 - X_sp) / sum_j (X_j / b_sp,j)`` over the other species of
    the mixture ``comp`` (mole fractions). Returns ``(b_mix, provenance
    dict per pair)``; a species alone in its mixture has nothing to diffuse
    against and returns infinity.
    """
    x_sp = comp.get(sp, 0.0)
    s = 0.0
    prov: dict[str, str] = {}
    for oj, xj in comp.items():
        if oj == sp or xj <= 0.0:
            continue
        b, p = b_pair(sp, oj, T)
        s += xj / b
        prov[f'{sp}-{oj}'] = p
    if s == 0.0:
        return math.inf, prov
    return (1.0 - x_sp) / s, prov