zephyrus.composition

def parse_formula(name: str) -> dict[str, int]:
    """Element counts of a simple molecular formula string.

    Parses formulas of the kind atmospheric chemistry codes emit, for
    example ``'CO2' -> {'C': 1, 'O': 2}``. Trailing annotations after an
    underscore are stripped, because species names such as ``'O_1'``,
    ``'N_2D'``, or ``'S8_l_s'`` denote excited states or condensates of the
    same stoichiometry.

    Parameters
    ----------
    name : str
        Molecular formula, e.g., ``'H2O'``, ``'CO2'``, ``'He'``.

    Returns
    -------
    dict
        Element symbol to integer count.

    Raises
    ------
    ValueError
        If the string cannot be parsed as a formula or contains an element
        outside ``ELEMENT_AMU``.
    """
    base = name.split('_')[0]
    out: dict[str, int] = {}
    pos = 0
    for m in _FORMULA_TOKEN.finditer(base):
        if m.start() != pos:
            raise ValueError(f'cannot parse formula {name!r}')
        pos = m.end()
        el, cnt = m.group(1), int(m.group(2) or 1)
        if el not in ELEMENT_AMU:
            raise ValueError(f'unknown element {el!r} in formula {name!r}')
        out[el] = out.get(el, 0) + cnt
    if pos != len(base) or not out:
        raise ValueError(f'cannot parse formula {name!r}')
    return out

def species_mass_amu(name: str) -> float:
    """Molecular mass of a formula string, in atomic mass units."""
    return sum(ELEMENT_AMU[el] * n for el, n in parse_formula(name).items())

def atomize(vmr: dict[str, float]) -> dict[str, float]:
    """Element mole fractions of an atomized composition.

    Reduces molecular volume mixing ratios to the mole fractions of their
    constituent atoms, which is the composition a fully dissociated gas
    would have. Used where escape physics operates on atoms: at the XUV
    wind base, where molecules are photodissociated well below the level
    where the wind is launched (Murray-Clay et al. 2009, ApJ 693, 23).

    Parameters
    ----------
    vmr : dict
        Species name to mole fraction. Need not sum to 1; the result is
        renormalized.

    Returns
    -------
    dict
        Element symbol to mole fraction, summing to 1.

    Raises
    ------
    ValueError
        If no species carries a positive mole fraction.
    """
    counts: dict[str, float] = {}
    for sp, x in vmr.items():
        if x <= 0.0:
            continue
        for el, n in parse_formula(sp).items():
            counts[el] = counts.get(el, 0.0) + x * n
    tot = sum(counts.values())
    if tot <= 0.0:
        raise ValueError('empty composition: no species has a positive mole fraction')
    return {el: c / tot for el, c in counts.items()}

def mean_particle_mass(element_fractions: dict[str, float]) -> float:
    """Mean particle mass of an atomized composition, in kg.

    ``element_fractions`` maps element symbols to mole fractions summing
    to 1, as returned by :func:`atomize`.
    """
    return sum(x * ELEMENT_AMU[el] for el, x in element_fractions.items()) * amu