zephyrus.fractionation

def solve_closure(phi, X, m, T, g0, b, return_diag=False):
    """Active-set solution of the N-species fractionation closure.

    Parameters
    ----------
    phi : float
        Total mass flux at the base [g cm^-2 s^-1], non-negative.
    X : array
        Mole fractions, non-negative, summing to 1.
    m : array
        Particle masses [g].
    T : float
        Temperature [K].
    g0 : float
        Gravity at the base [cm s^-2].
    b : array
        Symmetric matrix of binary diffusion parameters [cm^-1 s^-1],
        diagonal ignored (conventionally np.inf).
    return_diag : bool
        When True, also return the multiplier ``1/Hbar`` and the active set.

    Returns
    -------
    Phi : array
        Number fluxes [cm^-2 s^-1], guaranteed non-negative and satisfying
        ``sum m_i Phi_i = phi``. For ``phi = 0`` the fluxes are zero, the
        active set empty, and ``1/Hbar`` the continuous limit
        ``min_j m_j g0 / kT``.
    """
    flux, inv_h_bar, escaping, _failed = _solve(phi, X, m, T, g0, b)
    if return_diag:
        return flux, inv_h_bar, escaping
    return flux

def solve_fixed_active(phi, X, m, T, g0, b, active):
    """Solve the linear closure on a fixed active set.

    Unknowns: ``w_j`` for j in ``active`` and the inverse ``1/Hbar`` of the
    shared density scale height. Returns ``(w_full, inv_h_bar)`` with
    ``w = 0`` for retained species. The single-active case is solved
    analytically, the general case with two-sided diagonal equilibration,
    which controls the roughly 25 decades the matrix entries can span.
    """
    kT = kb_cgs * T
    act = sorted(active)
    ret = [k for k in range(len(X)) if k not in active]
    na = len(act)
    if na == 1:
        j = act[0]
        w = np.zeros(len(X))
        w[j] = phi / (m[j] * X[j])
        inv_h_bar = m[j] * g0 / kT + w[j] * sum(X[k] / b[j, k] for k in ret)
        return w, inv_h_bar
    mat = np.zeros((na + 1, na + 1))
    rhs = np.zeros(na + 1)
    for row, j in enumerate(act):
        diag = 0.0
        for col, i in enumerate(act):
            if i == j:
                continue
            mat[row, col] += X[i] / b[i, j]
            diag += X[i] / b[i, j]
        for k in ret:
            diag += X[k] / b[j, k]
        mat[row, row] -= diag
        mat[row, na] = 1.0  # the +1/Hbar column
        rhs[row] = m[j] * g0 / kT
    for col, j in enumerate(act):
        mat[na, col] = m[j] * X[j]
    rhs[na] = phi
    # Two-sided equilibration: D_r M D_c y = D_r rhs, solution = D_c y.
    dr = 1.0 / np.max(np.abs(mat), axis=1)
    ms = mat * dr[:, None]
    dc = 1.0 / np.max(np.abs(ms), axis=0)
    ms = ms * dc[None, :]
    sol = dc * np.linalg.solve(ms, rhs * dr)
    w = np.zeros(len(X))
    for col, j in enumerate(act):
        w[j] = sol[col]
    return w, sol[na]

def first_threshold(X, m, T, g0, b):
    """Mass flux at which the first heavy species entrains, in g cm^-2 s^-1.

    Below this flux only the lightest species escapes; the value is the
    smallest crossover over the heavier species of the multi-background
    drag balance.
    """
    kT = kb_cgs * T
    light = int(np.argmin(m))
    best = np.inf
    for k in range(len(X)):
        if k == light:
            continue
        denom = X[light] / b[light, k] + sum(
            X[kk] / b[light, kk] for kk in range(len(X)) if kk != light
        )
        w_star = (m[k] - m[light]) * g0 / kT / denom
        best = min(best, m[light] * X[light] * w_star)
    return best

def closure_per_species(
    mdot: float, element_fractions: dict, T_wind: float, M_p: float, r_base: float
) -> tuple[dict, dict, dict]:
    """Per-element escape rates [kg/s] from the closure at the wind base.

    SI boundary around the cgs solver: the bulk rate ``mdot`` [kg/s]
    converts to the base mass flux, the closure partitions it over the
    atomized composition with the binary-diffusion library coefficients,
    and the number fluxes convert back to per-element mass rates that sum
    to the bulk rate. Rock-forming species (Na, Mg, Si, Fe) raise the
    ``rock_former_bij`` flag: their coefficients sit in the widest
    provenance class, and Na and Mg ionize where rock vapor exists while
    these are neutral-gas coefficients. Elements with no kinetic diameter
    (Al, P, Cl, K, Ca, Ti) take the coefficients of their nearest-mass
    covered substitute and raise the ``substituted_bij`` flag; their own
    masses still enter the closure.

    Returns ``(per_element, diagnostics, flags)`` with the diagnostics
    carrying the active set, the retained species, the multiplier, the
    relative mass-conservation residual, and the per-pair coefficient
    provenance.
    """
    flags: dict = {}
    species = sorted(element_fractions, key=lambda el: ELEMENT_AMU[el])
    X = np.array([element_fractions[el] for el in species], float)
    X = X / X.sum()
    m_g = masses_g(species)
    rows = build_rows(species)
    b = bmatrix(species, T_wind, rows=rows)
    prov = {f'{r.i}-{r.j}': (r.provenance, r.uncertainty) for r in rows}
    if any(el in ROCK_FORMERS for el in species):
        flags['rock_former_bij'] = [el for el in species if el in ROCK_FORMERS]
    covered = substitutable()
    substituted = [el for el in species if el not in covered and len(species) > 1]
    if substituted:
        flags['substituted_bij'] = substituted

    g0_cgs = (G * M_p / r_base**2) * 1e2  # m/s^2 -> cm/s^2
    phi_cgs = (mdot / (4.0 * math.pi * r_base**2)) * 0.1  # kg/m^2/s -> g/cm^2/s

    flux, inv_h_bar, active, failed = _solve(phi_cgs, X, m_g, T_wind, g0_cgs, b)
    if failed:
        flags['closure_trace_resolve_failed'] = True
    area_cm2 = 4.0 * math.pi * (r_base * 1e2) ** 2
    per_element = {
        el: float(flux[k]) * area_cm2 * float(m_g[k]) * 1e-3 for k, el in enumerate(species)
    }

    total = sum(per_element.values())
    conservation = abs(total - mdot) / mdot if mdot > 0 else 0.0
    diag = dict(
        active_set=sorted(species[k] for k in active),
        retained=[el for el in species if el not in {species[k] for k in active}],
        inv_H_bar_cgs=float(inv_h_bar),
        mass_conservation_rel=float(conservation),
        b_provenance=prov,
    )
    return per_element, diag, flags

def unfractionated_split(
    mdot: float, reservoirs: dict | None, element_fractions: dict
) -> tuple[dict, dict]:
    """Split a bulk rate over elements without fractionation.

    The protocol of the energy-limited path: reservoir mass fractions when
    reservoir masses are supplied; otherwise the mass fractions of the
    atomized wind-base composition, with the ``split_from_base_composition``
    flag recording the substitution. A reservoir that has run dry carries no
    proportions to split by, so it falls back to the same composition, which
    is the state the end of an evolutionary track reaches. Negative reservoir
    masses are malformed input. The split conserves the bulk rate exactly.
    """
    flags: dict = {}
    if reservoirs and any(m < 0.0 for m in reservoirs.values()):
        raise ValueError('reservoir masses must be non-negative')
    if reservoirs and sum(reservoirs.values()) > 0.0:
        tot = sum(reservoirs.values())
        fracs = {el: mass / tot for el, mass in reservoirs.items()}
    else:
        flags['split_from_base_composition'] = True
        mass = {el: x * ELEMENT_AMU[el] for el, x in element_fractions.items()}
        tot = sum(mass.values())
        fracs = {el: mm / tot for el, mm in mass.items()}
    return {el: mdot * f for el, f in fracs.items()}, flags