proteus.outgas.trapping¶

def critical_melt_fraction(config: Config) -> float:
    """Disaggregation melt fraction ``phi_c`` [1], the solver's rheological transition.

    The single definition of ``phi_c`` for the trapping step, so the top of the
    drainage front and the no-drainage bound cannot drift apart from each other
    or from the interior rheology.
    """
    return float(config.interior_energetics.rfront_loc)

def crystallised_mass_from_phi(m_mantle: float, phi_prev: float, phi_now: float) -> float:
    """Mass crystallised over one step [kg], from the melt fraction.

    ``dM_RM = M_mantle * [Phi(t-1) - Phi(t)]``. Phi_global is a mass fraction
    and therefore intensive, so a Zalmoxis re-solve that changes ``M_mantle``
    does not inject a spurious increment the way differencing
    ``M_mantle_solid`` would. It also matches the melt reservoir CALLIOPE
    dissolves into, which is sized from Phi_global rather than from
    ``M_mantle_liquid``; the prevent-warming clamp moves Phi_global without
    resynchronising those masses, so the two disagree on clamped steps.

    A negative increment means the mantle remelted that much mass, which
    releases trapped volatiles back to the melt.

    Parameters
    ----------
    m_mantle : float
        Mantle mass [kg]. A non-positive or non-finite value gives 0.
    phi_prev, phi_now : float
        Global melt fraction at the previous and the current step [1]. A
        non-finite value gives 0.
    """
    if not np.isfinite(m_mantle) or m_mantle <= 0.0:
        return 0.0
    if not (np.isfinite(phi_prev) and np.isfinite(phi_now)):
        return 0.0
    return float(m_mantle) * (float(phi_prev) - float(phi_now))

def derived_total_elements(config: Config) -> tuple[str, ...]:
    """Elements whose whole-planet total the chemistry recomputes each step.

    Oxygen under ``planet.fO2_source = 'user_constant'`` is the only case.
    Holding the oxygen fugacity at a fixed buffer offset requires the chemistry
    to move oxygen, so ``O_kg_total`` is an output of the solve rather than a
    conserved budget the escape chain debits; ``outgas/common.py`` says as much
    where it lists the ``_kg_total`` slot as escape-owned for every element
    except O. Under ``'from_O_budget'`` the wrapper restores the authoritative
    budget after the copy, so oxygen is conserved state there and is absent
    from this tuple.

    Trapping still treats oxygen like every other element around a solve: the
    trapped oxygen is taken out before it and added back to the total the
    chemistry writes, so ``O_kg_total`` is the whole planet's oxygen under
    either source. The per-element closure is not asserted for it, because the
    chemistry total closes only to the solver's own convention.
    """
    if getattr(getattr(config, 'planet', None), 'fO2_source', '') == 'user_constant':
        return ('O',)
    return ()

def effective_partition(f_tl: float | np.ndarray, d_z: float) -> float | np.ndarray:
    """Effective crystal/melt partition coefficient including trapped melt.

    Reduces to ``f_tl`` for a perfectly incompatible species (``d_z = 0``) and
    to 1 for a species that does not fractionate (``d_z = 1``). ``f_tl`` may be
    a scalar or a per-step array; the algebra is the same either way.

    Parameters
    ----------
    f_tl : float or ndarray
        Trapped melt fraction, a mass fraction in [0, 1].
    d_z : float
        Crystal/melt partition coefficient [1], non-negative.

    Returns
    -------
    float or ndarray
        ``(1 - f_tl) * d_z + f_tl`` [1].

    Raises
    ------
    ValueError
        If ``f_tl`` is not finite or lies outside [0, 1], or ``d_z`` is negative.
    """
    fraction = np.asarray(f_tl, dtype=float)
    if not np.all(np.isfinite(fraction)):
        raise ValueError('F_tl must be finite at every step')
    if np.any((fraction < 0.0) | (fraction > 1.0)):
        outside = fraction[(fraction < 0.0) | (fraction > 1.0)]
        raise ValueError(f'F_tl must lie in [0, 1], got {outside.ravel()[0]!r}')
    if d_z < 0.0:
        raise ValueError(f'D_Z must be non-negative, got {d_z!r}')
    return (1.0 - f_tl) * d_z + f_tl

def escapable_inventory(hf_row: dict, element: str, config: Config | None = None) -> float:
    """Whole-planet inventory of one element that escape can still reach [kg]."""
    total = float(hf_row.get(f'{element}_kg_total', 0.0))
    if not np.isfinite(total):
        return total
    return max(0.0, total - unreachable_mass(hf_row, element, config))

def locked_condensed_mass(hf_row: dict, element: str, config: Config | None = None) -> float:
    """Mass of one element held in condensates that no longer re-vapourise [kg].

    The ``{element}_kg_condensed`` reservoir with ``outgas.prevent_desaturation``
    on, and 0 otherwise (also without a config), since then each solve may
    re-vapourise the condensate.
    """
    if config is None or getattr(config.outgas, 'prevent_desaturation', False) is not True:
        return 0.0
    mass = float(hf_row.get(f'{element}_kg_condensed', 0.0))
    return max(0.0, mass) if np.isfinite(mass) else 0.0

def locked_solid_mass(hf_row: dict, element: str) -> float:
    """Whole-planet mass of one element locked in the solid mantle [kg].

    The single definition of what trapping has put beyond reach, read from
    ``{element}_kg_trapped``. Outgassing condensates such as atmodeller's graphite
    are held in ``{element}_kg_condensed``, not in ``_kg_solid``; see
    :func:`locked_condensed_mass` for when they are locked too.
    """
    return _trapped(hf_row, element)

def melt_concentration(
    kg_liquid: float | np.ndarray, melt_mass: float | np.ndarray
) -> float | np.ndarray:
    """Species mass fraction in the melt [1]; zero wherever no melt remains.

    A fully crystallised mantle carries no melt to draw from, so the quotient is
    replaced by zero instead of dividing by zero.

    Parameters
    ----------
    kg_liquid : float or ndarray
        Mass of the species dissolved in the melt [kg]; negative values count as 0.
    melt_mass : float or ndarray
        Mass of the melt [kg].
    """
    kg = np.maximum(np.asarray(kg_liquid, dtype=float), 0.0)
    melt = np.asarray(melt_mass, dtype=float)
    c_z = np.zeros(np.shape(kg), dtype=float)
    np.divide(kg, melt, out=c_z, where=melt > 0.0)
    if np.ndim(kg_liquid) == 0:
        return float(c_z)
    return c_z

def partition_coefficients(config: Config) -> dict[str, float]:
    """Crystal/melt partition coefficient per species, from config.

    Water is the only species with appreciable lattice incorporation; the rest
    default to zero and are buried through the interstitial-melt term alone.

    Returns
    -------
    dict
        ``outgas.D_const_<species>`` [1] per trapped species; 0 for a species,
        such as a noble gas, that has no such field.
    """
    return {s: float(getattr(config.outgas, f'D_const_{s}', 0.0)) for s in TRAPPED_SPECIES}

def released_mass(hf_row: dict, hf_all, phi_prev: float, phi_now: float) -> dict[str, float]:
    """Mass each trapped species and element returns to the melt on a remelt [kg].

    The mass the burial ledger holds below ``phi_now``, rebuilt from the
    completed rows in ``hf_all``. Trapped mass the ledger does not hold returns
    in proportion to the solid that remelted; a complete remelt returns all of
    it. Each release is capped at what is trapped.
    """
    current = np.array([_trapped(hf_row, name) for name in TRAPPED_COLUMNS])
    share = remelted_fraction(phi_prev, phi_now)
    if share >= 1.0:
        return dict(zip(TRAPPED_COLUMNS, current))
    history = np.column_stack(
        [
            hf_all[f'{name}_kg_trapped'].to_numpy(dtype=float)
            if f'{name}_kg_trapped' in hf_all
            else np.zeros(len(hf_all))
            for name in TRAPPED_COLUMNS
        ]
    )
    layers = build_ledger(hf_all['Phi_global'].to_numpy(dtype=float), history)
    held = ledger_total(layers, len(TRAPPED_COLUMNS))
    from_ledger = remelt_to(layers, phi_now) if layers else np.zeros(len(TRAPPED_COLUMNS))
    unledgered = np.maximum(current - held, 0.0)
    release = np.minimum(from_ledger + share * unledgered, current)
    return dict(zip(TRAPPED_COLUMNS, release))

def remelt_returned_mass(hf_row: dict) -> bool:
    """Whether this step's trapping released buried mass back to the melt."""
    branch = float(hf_row.get('trap_branch', 0.0))
    step = float(hf_row.get('trap_kg_step', 0.0))
    return bool(np.isfinite(branch) and round(branch) == BRANCH_REMELT and step < 0.0)

def remelted_fraction(phi_prev: float, phi_now: float) -> float:
    """Share of the solid mantle that remelted over one step [1].

    ``(Phi(t) - Phi(t-1)) / (1 - Phi(t-1))``, clipped to ``[0, 1]``: 1 when the
    mantle is fully molten again, 0 when it froze further or did not change.

    Parameters
    ----------
    phi_prev, phi_now : float
        Global melt fraction at the previous and the current step [1]. A
        non-finite value gives 0.
    """
    if not (np.isfinite(phi_prev) and np.isfinite(phi_now)):
        return 0.0
    solid_prev = 1.0 - float(phi_prev)
    if solid_prev <= 0.0:
        return 1.0 if phi_now > phi_prev else 0.0
    return float(np.clip((float(phi_now) - float(phi_prev)) / solid_prev, 0.0, 1.0))

def restore_trapped_mass(hf_row: dict) -> None:
    """Put the trapped mass back after a chemistry solve.

    Adds it to whatever the backend wrote into ``_kg_solid``, and restores each
    element's whole-planet ``_kg_total``. A species total is an output of the
    solve, the backend's atmosphere plus melt, so for every trapped species it is
    rebuilt as ``atm + liquid + solid``, and the solid and total moles follow the
    masses. Putting back the species total saved before the solve would instead
    overwrite the partition the solve has just produced.
    """
    for species in vol_list:
        mass = _trapped(hf_row, species)
        if mass <= 0.0:
            continue
        _shift(hf_row, f'{species}_kg_solid', mass)
        hf_row[f'{species}_kg_total'] = _reservoir_sum(hf_row, species, 'kg')
        mmw = eval_gas_mmw(species)
        hf_row[f'{species}_mol_solid'] = float(hf_row[f'{species}_kg_solid']) / mmw
        hf_row[f'{species}_mol_total'] = _reservoir_sum(hf_row, species, 'mol')
    for element in (*vol_element_list, *noble_gases):
        mass = _trapped(hf_row, element)
        if mass <= 0.0:
            continue
        _shift(hf_row, f'{element}_kg_solid', mass)
        _shift(hf_row, f'{element}_kg_total', mass)
        if element in noble_gases:
            # A noble gas is its own element; its moles follow the masses.
            mmw = eval_gas_mmw(element)
            hf_row[f'{element}_mol_solid'] = float(hf_row[f'{element}_kg_solid']) / mmw
            hf_row[f'{element}_mol_total'] = float(hf_row[f'{element}_kg_total']) / mmw

def run_trapping(
    config: Config,
    hf_row: dict,
    hf_all,
    interior_o=None,
    *,
    init_stage: bool = False,
) -> TrappingStep | None:
    """Bury volatiles into the solid mantle over this crystallisation step.

    Called once per iteration, after the interior has advanced and any structure
    re-solve has finished, and before escape and outgassing read the
    inventories. Returns ``None`` when trapping is inactive: during the
    initialisation stage, whose interior steps do not advance the mantle,
    when ``hf_all`` holds no earlier row or the time is not above 0, where
    there is no previous state to difference, and whenever the mode is
    ``none``. The first step after the initialisation stage differences
    against the last row of that stage.

    Parameters
    ----------
    config : Config
        Model configuration.
    hf_row : dict
        Current helpfile row, modified in place.
    hf_all : pandas.DataFrame or None
        Completed rows. The last one supplies the previous melt fraction.
    interior_o : Interior_t or None
        Interior state. The front scheme reads its melt-fraction, density and
        pressure profiles to locate the freezing front.
    init_stage : bool
        True during the initialisation stage, which traps nothing.
    """
    mode = getattr(config.outgas, 'trap_mode', 'none')
    # A config that never declared a mode traps nothing: the schema restricts the
    # field at load, so a non-string means a config built in code without trapping,
    # whose other attributes nobody chose as trapping parameters.
    if not isinstance(mode, str):
        return None
    if mode not in TRAPPING_MODES:
        raise ValueError(f'Unknown trapping mode {mode!r}; expected one of {TRAPPING_MODES}')
    if mode == 'none':
        return None
    if init_stage:
        # The initialisation iterations advance the clock but not the mantle,
        # and the time is reset to zero after each, so there is nothing to bury.
        return None
    if hf_all is None or len(hf_all) < 1:
        return None
    if float(hf_row.get('Time', 0.0)) <= 0.0:
        return None

    prev = hf_all.iloc[-1].to_dict()
    m_mantle = float(hf_row.get('M_mantle', 0.0))
    phi_now = float(hf_row.get('Phi_global', float('nan')))
    phi_prev = float(prev.get('Phi_global', float('nan')))
    dm_rm = crystallised_mass_from_phi(m_mantle, phi_prev, phi_now)
    # The dissolved masses in hf_row are still those of the previous chemistry
    # solve, so the melt they were dissolved into is the previous one. Dividing
    # by it recovers the concentration the solver computed.
    m_mantle_prev = float(prev.get('M_mantle', m_mantle))
    melt_mass = m_mantle_prev * max(0.0, min(1.0, phi_prev)) if np.isfinite(phi_prev) else 0.0
    if dm_rm < 0.0:
        released = released_mass(hf_row, hf_all, phi_prev, phi_now)
        return _release(hf_row, mode, dm_rm, melt_mass, released)
    if dm_rm == 0.0 or melt_mass <= 0.0:
        step = TrappingStep(mode=mode, dm_rm=dm_rm, f_tl=0.0, melt_mass=melt_mass)
        _record(hf_row, step)
        return step

    step = _drainage_fraction(config, hf_row, prev, interior_o, dm_rm)
    if step is None:
        # Without the interior profiles there is no front; the crystals still
        # take up D_Z, so the step buries at that alone, on a branch of its own.
        step = TrappingStep(mode=mode, dm_rm=0.0, f_tl=0.0, melt_mass=0.0)
        step.branch = BRANCH_FALLBACK
    step.mode = mode
    step.dm_rm = dm_rm
    step.melt_mass = melt_mass

    d_z = partition_coefficients(config)
    for species in TRAPPED_SPECIES:
        kg_liquid = float(hf_row.get(f'{species}_kg_liquid', 0.0))
        if kg_liquid <= 0.0 or not _carried(hf_row, species):
            continue
        c_z = melt_concentration(kg_liquid, melt_mass)
        mass, capped = trapped_mass(d_z[species], step.f_tl, c_z, dm_rm, kg_liquid)
        if mass <= 0.0:
            continue
        if capped:
            step.supply_capped.append(species)
        step.trapped_kg[species] = mass

    _apply_to_reservoirs(hf_row, step.trapped_kg)
    _record(hf_row, step)
    return step

def trapped_mass(
    d_z: float,
    f_tl: float,
    c_z: float,
    dm_rm: float,
    available_kg: float,
) -> tuple[float, bool]:
    """Mass of one species buried over a crystallisation increment [kg].

    ``[(1 - F_tl) D_Z + F_tl] C_Z dM_RM``, Sim et al. (2024) Eq. 6, capped at
    ``available_kg``, the mass the melt holds. A non-positive increment buries
    nothing. Returns the mass and whether the cap bound.

    Parameters
    ----------
    d_z : float
        Crystal/melt partition coefficient [1], non-negative.
    f_tl : float
        Trapped melt fraction, a mass fraction in [0, 1].
    c_z : float
        Mass fraction of the species in the melt [1].
    dm_rm : float
        Mantle mass crystallised over the increment [kg].
    available_kg : float
        Mass of the species the melt holds [kg].

    Returns
    -------
    tuple
        Buried mass [kg], at least 0, and whether the cap at ``available_kg`` bound.
    """
    raw = float(effective_partition(f_tl, d_z)) * float(c_z) * max(float(dm_rm), 0.0)
    available = max(float(available_kg), 0.0)
    return max(min(raw, available), 0.0), raw > available

@contextmanager
def trapped_mass_withheld(hf_row: dict):
    """Hide the trapped mass from the chemistry for the duration of a solve.

    Withholds on entry and restores on exit, also when the solve raises, so the
    row never carries the reachable inventory as the whole planet's.
    """
    withhold_trapped_mass(hf_row)
    try:
        yield
    finally:
        restore_trapped_mass(hf_row)

def trapping_active(config: Config) -> bool:
    """Whether the run traps volatiles; with ``trap_mode = 'none'`` nothing
    this module owns may change what the run does."""
    mode = getattr(config.outgas, 'trap_mode', 'none')
    return isinstance(mode, str) and mode != 'none'

def unreachable_mass(hf_row: dict, element: str, config: Config | None = None) -> float:
    """Whole-planet mass of one element beyond the reach of escape and outgassing [kg].

    The mass trapped in the solid mantle, plus the condensed reservoir when
    condensates are kept from re-vapourising. The escape step subtracts this mass
    from the mass available to escape, and the desiccation gate subtracts it from
    the inventory it tests, so a planet whose atmosphere and melt have both emptied
    still registers as desiccated.
    """
    return locked_solid_mass(hf_row, element) + locked_condensed_mass(hf_row, element, config)

def withhold_trapped_mass(hf_row: dict) -> None:
    """Take the trapped mass out of the reservoirs a chemistry solve owns.

    Every outgassing backend partitions a whole-planet inventory between melt and
    atmosphere, and rewrites the ``_kg_solid`` columns, with zero for every backend;
    atmodeller writes its condensates to ``_kg_condensed`` instead. Trapped mass
    belongs to neither. Left in an element total, it would be re-dissolved by the
    solve; left in a solid column, it would be overwritten by the backend. It is
    therefore subtracted from each element's
    ``_kg_total`` and from each species' and element's ``_kg_solid``, which leaves
    the backend exactly the share it owns.

    Pair with :func:`restore_trapped_mass`, or use :func:`trapped_mass_withheld`,
    which pairs the two. An element whose total the chemistry recomputes, such
    as oxygen under a buffered fO2, is treated the same way: the solve writes its
    own total over the withheld one and the restore adds the trapped share to
    it, and a solve that writes nothing leaves the total where it started.
    """
    for species in vol_list:
        mass = _trapped(hf_row, species)
        if mass > 0.0:
            _shift(hf_row, f'{species}_kg_solid', -mass)
    for element in (*vol_element_list, *noble_gases):
        mass = _trapped(hf_row, element)
        if mass <= 0.0:
            continue
        _shift(hf_row, f'{element}_kg_solid', -mass)
        _shift(hf_row, f'{element}_kg_total', -mass)