Calliope.solve

def atmosphere_mass(pin, ddict):
    """Atmospheric mass of volatiles and totals for H, C, and N.

    CALLIOPE stores pressures in bar throughout; the only conversion
    to SI Pa happens here (factor 1e5) when computing column mass
    `kg = p_Pa * 4 pi R^2 / g`.
    """

    p_d = get_partial_pressures(pin, ddict)
    mu_atm = atmosphere_mean_molar_mass(p_d)

    mass_atm_d = {}
    for key, value in p_d.items():
        mass_atm_d[key] = value * 1.0e5 / ddict['gravity']
        mass_atm_d[key] *= 4.0 * np.pi * ddict['radius'] ** 2.0
        mass_atm_d[key] *= molar_mass[key] / mu_atm

    # Per-element kg accumulators are built in mol-of-element first
    # (atom counts: H2O=2H, CH4=4H, H2S=2H, NH3=3H+1N, SO2=2O+1S, ...)
    # then multiplied by the element molar mass at the end.
    mass_atm_d['H'] = 2 * mass_atm_d['H2O'] / molar_mass['H2O']
    if is_included('H2', ddict):
        mass_atm_d['H'] += 2 * mass_atm_d['H2'] / molar_mass['H2']
    if is_included('CH4', ddict):
        mass_atm_d['H'] += 4 * mass_atm_d['CH4'] / molar_mass['CH4']
    if is_included('H2S', ddict):
        mass_atm_d['H'] += 2 * mass_atm_d['H2S'] / molar_mass['H2S']
    if is_included('NH3', ddict):
        mass_atm_d['H'] += 3 * mass_atm_d['NH3'] / molar_mass['NH3']
    mass_atm_d['H'] *= molar_mass['H']

    mass_atm_d['C'] = mass_atm_d['CO2'] / molar_mass['CO2']
    if is_included('CO', ddict):
        mass_atm_d['C'] += mass_atm_d['CO'] / molar_mass['CO']
    if is_included('CH4', ddict):
        mass_atm_d['C'] += mass_atm_d['CH4'] / molar_mass['CH4']
    mass_atm_d['C'] *= molar_mass['C']

    mass_atm_d['N'] = 2 * mass_atm_d['N2'] / molar_mass['N2']
    if is_included('NH3', ddict):
        mass_atm_d['N'] += mass_atm_d['NH3'] / molar_mass['NH3']
    mass_atm_d['N'] *= molar_mass['N']

    mass_atm_d['O'] = mass_atm_d['H2O'] / molar_mass['H2O']
    mass_atm_d['O'] += 2 * mass_atm_d['O2'] / molar_mass['O2']
    if is_included('CO', ddict):
        mass_atm_d['O'] += mass_atm_d['CO'] / molar_mass['CO']
    if is_included('CO2', ddict):
        mass_atm_d['O'] += mass_atm_d['CO2'] / molar_mass['CO2'] * 2.0
    if is_included('SO2', ddict):
        mass_atm_d['O'] += mass_atm_d['SO2'] / molar_mass['SO2'] * 2.0
    mass_atm_d['O'] *= molar_mass['O']

    mass_atm_d['S'] = 2 * mass_atm_d['S2'] / molar_mass['S2']
    if is_included('SO2', ddict):
        mass_atm_d['S'] += mass_atm_d['SO2'] / molar_mass['SO2']
    if is_included('H2S', ddict):
        mass_atm_d['S'] += mass_atm_d['H2S'] / molar_mass['H2S']
    mass_atm_d['S'] *= molar_mass['S']

    for e in element_list:
        mass_atm_d[e] = max(0.0, mass_atm_d[e])

    return mass_atm_d

def atmosphere_mean_molar_mass(p_d):
    """Mean molar mass of the atmosphere"""

    ptot = get_total_pressure(p_d)

    mu_atm = 0
    for key, value in p_d.items():
        mu_atm += molar_mass[key] * value
    mu_atm /= ptot

    return mu_atm

def dissolved_mass(pin, ddict):
    """Volatile masses in the (molten) mantle"""

    mass_int_d = {}

    p_d = get_partial_pressures(pin, ddict)
    ptot = get_total_pressure(p_d)

    # Henry's-law / power-law solubility laws return ppmw of the
    # species in the melt; prefactor converts ppmw -> kg dissolved.
    prefactor = 1e-6 * ddict['M_mantle'] * ddict['Phi_global']

    sol_H2O = SolubilityH2O()
    ppmw_H2O = sol_H2O(p_d['H2O'])
    mass_int_d['H2O'] = prefactor * ppmw_H2O

    sol_CO2 = SolubilityCO2()
    ppmw_CO2 = sol_CO2(p_d['CO2'], ddict['T_magma'])
    mass_int_d['CO2'] = prefactor * ppmw_CO2

    if is_included('CO', ddict):
        sol_CO = SolubilityCO()
        ppmw_CO = sol_CO(p_d['CO'], ptot)
        mass_int_d['CO'] = prefactor * ppmw_CO
    else:
        mass_int_d['CO'] = 0.0

    if is_included('CH4', ddict):
        sol_CH4 = SolubilityCH4()
        ppmw_CH4 = sol_CH4(p_d['CH4'], ptot)
        mass_int_d['CH4'] = prefactor * ppmw_CH4
    else:
        mass_int_d['CH4'] = 0.0

    # Override class default 'libourel'; dasgupta carries fO2 + p_total
    # dependence which the linear Libourel law does not.
    sol_N2 = SolubilityN2('dasgupta')
    ppmw_N2 = sol_N2(p_d['N2'], ptot, ddict['T_magma'], ddict['fO2_shift_IW'])
    mass_int_d['N2'] = prefactor * ppmw_N2

    sol_S2 = SolubilityS2()
    ppmw_S2 = sol_S2(p_d['S2'], ddict['T_magma'], ddict['fO2_shift_IW'])
    mass_int_d['S2'] = prefactor * ppmw_S2

    # No SolubilityH2S / NH3 / SO2 / O2 / H2 in CALLIOPE — these
    # species do not partition into the melt phase and contribute
    # zero to the dissolved-element tallies below.
    mass_int_d['H'] = mass_int_d['H2O'] * 2 / molar_mass['H2O']
    if is_included('CH4', ddict):
        mass_int_d['H'] += mass_int_d['CH4'] * 4 / molar_mass['CH4']
    mass_int_d['H'] *= molar_mass['H']

    mass_int_d['C'] = mass_int_d['CO2'] / molar_mass['CO2']
    if is_included('CO', ddict):
        mass_int_d['C'] += mass_int_d['CO'] / molar_mass['CO']
    if is_included('CH4', ddict):
        mass_int_d['C'] += mass_int_d['CH4'] / molar_mass['CH4']
    mass_int_d['C'] *= molar_mass['C']

    mass_int_d['N'] = mass_int_d['N2']

    mass_int_d['O'] = mass_int_d['H2O'] / molar_mass['H2O']
    mass_int_d['O'] += mass_int_d['CO2'] / molar_mass['CO2'] * 2.0
    if is_included('CO', ddict):
        mass_int_d['O'] += mass_int_d['CO'] / molar_mass['CO']
    mass_int_d['O'] *= molar_mass['O']

    mass_int_d['S'] = mass_int_d['S2']

    for e in element_list:
        mass_int_d[e] = max(0.0, mass_int_d[e])

    return mass_int_d

def equilibrium_atmosphere(
    target_d,
    ddict,
    hide_warnings=True,
    rtol=1e-5,
    atol=1e10,
    xtol=1e-8,
    p_guess=None,
    nsolve=1500,
    nguess=7500,
    print_result=True,
    opt_solver=True,
):
    """Solve for surface partial pressures assuming melt-vapour equilibrium.

    Parameters
    ----------
    target_d : dict
        Target elemental mass inventories [kg], with keys 'H', 'C', 'N', 'S'.
    ddict : dict
        Dictionary of coupler options variables (planet, magma state, inclusion flags).
    hide_warnings : bool, default True
        Hide floating point runtime warnings raised by `scipy` for poor guesses.
    rtol : float, default 1e-5
        Relative tolerance for mass conservation.
    atol : float, default 1e10
        Absolute tolerance for mass conservation [kg].
    xtol : float, default 1e-8
        Relative tolerance for fsolve.
    p_guess : dict or None, default None
        Dictionary of initial guess for primary-species partial pressures [bar].
        Must contain the keys 'H2O', 'CO2', 'N2', 'S2', each mapping to a
        finite real number. If None, an internal Monte-Carlo log-uniform
        draw is used. A non-dict value raises TypeError; missing keys or
        non-finite values raise ValueError.
    nsolve : int, default 1500
        Maximum number of inner-solver iterations per attempt.
    nguess : int, default 7500
        Maximum number of Monte-Carlo restarts before giving up.
    print_result : bool, default True
        If True, log final outgassed partial pressures at INFO level.
    opt_solver : bool, default True
        If True, alternate between fsolve and trust-constr on each restart.

    Returns
    -------
    partial_pressures : dict
        Volatile partial pressures [bar] keyed `<species>_bar`, plus per-species
        reservoir masses [kg], elemental totals, residuals, and atmospheric
        diagnostics (P_surf, M_atm, atm_kg_per_mol, ratios).
    """

    if print_result:
        log.info('Solving for equilibrium partial pressures at surface')
    log.debug('    target masses: %s' % str(target_d))

    # Hard ub = 1e7 bar is well above any realistic magma-ocean
    # surface pressure; it prevents trust-constr from exploring
    # non-physical regions during a poor cold start.
    lb = [0.0] * 4
    ub = [1e7] * 4

    if p_guess is None:
        x0 = get_initial_pressures(target_d)
    else:
        # Validate up front so a missing key surfaces as a clear ValueError
        # rather than a bare KeyError from the tuple construction below.
        # The isinstance check handles cases where a caller passes a list,
        # a pandas Series, or accidentally a non-dict object — without it,
        # the membership test below would raise an opaque TypeError.
        if not isinstance(p_guess, dict):
            raise TypeError(f'p_guess must be a dict or None, got {type(p_guess).__name__}.')
        required = ('H2O', 'CO2', 'N2', 'S2')
        missing = [k for k in required if k not in p_guess]
        if missing:
            raise ValueError(
                f'p_guess is missing required keys: {missing}. '
                f'Expected all of {list(required)}.'
            )
        x0 = (p_guess['H2O'], p_guess['CO2'], p_guess['N2'], p_guess['S2'])

        # Reject non-finite values up front. NaN > 1e-10 evaluates False,
        # which would silently collapse ub to 1.0 and propagate NaN through
        # the solver to produce garbage output with no error signal.
        for k, v in zip(required, x0):
            if not np.isfinite(v):
                raise ValueError(f'p_guess[{k!r}] must be a finite real number, got {v!r}.')

        # Zero or near-zero guess collapses ub from 1e7 to 1.0 to keep
        # trust-constr from wandering inside a degenerate slot.
        for i in range(4):
            ub[i] = ub[i] if (x0[i] > 1e-10) else 1.0

    bounds = opt.Bounds(lb=lb, ub=ub)

    tolerance = np.amax(list(target_d.values())) * rtol + atol + TRUNC_MASS
    log.debug('Required tolerance: %g' % tolerance)

    with warnings.catch_warnings():
        # Solver Monte-Carlo restarts produce poor guesses that trip
        # RuntimeWarning / UserWarning; the bad attempts are discarded
        # so the warnings should not reach the caller.
        if hide_warnings:
            warnings.filterwarnings('ignore', category=RuntimeWarning)
            warnings.filterwarnings('ignore', category=UserWarning)

        solver: int = 0
        for count in range(nguess):
            if solver == 0:
                sol, _, ier, _ = opt.fsolve(
                    func, x0, args=(ddict, target_d), maxfev=nsolve, xtol=xtol, full_output=True
                )
                success = bool(ier == 1)
            else:
                result = opt.minimize(
                    obj,
                    x0,
                    args=(ddict, target_d),
                    method='trust-constr',
                    bounds=bounds,
                    options={'maxiter': nsolve, 'xtol': xtol},
                )
                success = result.success
                sol = result.x

            # Reject solver-claimed successes that miss the residual
            # tolerance: fsolve and trust-constr converge by their own
            # criteria, which are not the kg-mass-balance criterion.
            this_resid = func(sol, ddict, target_d)
            loss = np.amax(np.abs(this_resid))
            if loss > tolerance:
                if success:
                    log.debug('Solution rejected by residual')
                    log.debug('    d(i=%d) = %.2e kg' % (np.argmax(this_resid), loss))
                success = False

            if success:
                break

            x0 = get_initial_pressures(target_d)

            # Alternate fsolve <-> trust-constr on each restart so a
            # basin one solver cannot escape gets a chance from the
            # other. Disable via opt_solver=False to pin to fsolve.
            if opt_solver:
                solver = 1 - solver

    if not success:
        raise RuntimeError(
            'Could not find solution for volatile abundances (max attempts, %d)' % nguess
        )

    log.debug('    Initial guess attempt number = %d' % count)

    res_l = func(sol, ddict, target_d)
    log.debug('    Residuals: %s' % res_l)

    sol_dict = {'H2O': sol[0], 'CO2': sol[1], 'N2': sol[2], 'S2': sol[3]}
    p_d = get_partial_pressures(sol_dict, ddict)

    # Pass the 4-key primary dict (sol_dict), not the expanded p_d:
    # atmosphere_mass and dissolved_mass each call get_partial_pressures
    # internally, so feeding them p_d would re-run gas-phase chemistry
    # on a dict that already contains the secondaries (read but ignored).
    mass_atm_d = atmosphere_mass(sol_dict, ddict)
    mass_int_d = dissolved_mass(sol_dict, ddict)

    # CALLIOPE has no solid-mantle reservoir; every `_kg_solid` field
    # is 0.0 and exists only for schema parity with the downstream
    # PROTEUS hf_row which carries solid/melt/atmosphere splits.
    outdict = {'M_atm': 0.0, 'P_surf': 0.0}
    for s in volatile_species:
        outdict[s + '_bar'] = 0.0
        outdict[s + '_kg_atm'] = 0.0
        outdict[s + '_kg_liquid'] = 0.0
        outdict[s + '_kg_solid'] = 0.0
        outdict[s + '_kg_total'] = 0.0

        if s in p_d.keys():
            outdict[s + '_bar'] = p_d[s]
            outdict['P_surf'] += outdict[s + '_bar']

    for s in volatile_species:
        outdict[s + '_vmr'] = outdict[s + '_bar'] / outdict['P_surf']

        if print_result:
            log.info(
                '    %-6s : %-8.2f bar (%.2e VMR)'
                % (s, outdict[s + '_bar'], outdict[s + '_vmr'])
            )

    all = [s for s in volatile_species]
    all.extend(['H', 'C', 'N', 'S', 'O'])
    for s in all:
        tot_kg = 0.0

        if s in mass_atm_d.keys():
            outdict[s + '_kg_atm'] = mass_atm_d[s]
            tot_kg += mass_atm_d[s]

        if s in mass_int_d.keys():
            outdict[s + '_kg_liquid'] = mass_int_d[s]
            outdict[s + '_kg_solid'] = 0.0
            tot_kg += mass_int_d[s]

        outdict[s + '_kg_total'] = tot_kg

    for s in volatile_species:
        outdict['M_atm'] += outdict[s + '_kg_atm']

    outdict['atm_kg_per_mol'] = 0.0
    for s in volatile_species:
        outdict[s + '_mol_atm'] = outdict[s + '_kg_atm'] / molar_mass[s]
        outdict[s + '_mol_solid'] = outdict[s + '_kg_solid'] / molar_mass[s]
        outdict[s + '_mol_liquid'] = outdict[s + '_kg_liquid'] / molar_mass[s]
        outdict[s + '_mol_total'] = (
            outdict[s + '_mol_atm'] + outdict[s + '_mol_solid'] + outdict[s + '_mol_liquid']
        )

        outdict['atm_kg_per_mol'] += outdict[s + '_vmr'] * molar_mass[s]

    for e1 in element_list:
        for e2 in element_list:
            if e1 == e2:
                continue
            em1 = outdict[e1 + '_kg_atm']
            em2 = outdict[e2 + '_kg_atm']
            if em2 == 0:
                continue
            outdict['%s/%s_atm' % (e1, e2)] = em1 / em2

    outdict['H_res'] = res_l[0]
    outdict['C_res'] = res_l[1]
    outdict['N_res'] = res_l[2]
    outdict['S_res'] = res_l[3]

    return outdict

def func(pin_arr, ddict, mass_target_d):
    """Mass-balance residual [kg per element] for the four primary partial pressures [bar]."""

    pin_dict = {'H2O': pin_arr[0], 'CO2': pin_arr[1], 'N2': pin_arr[2], 'S2': pin_arr[3]}

    mass_atm_d = atmosphere_mass(pin_dict, ddict)
    mass_int_d = dissolved_mass(pin_dict, ddict)

    res_l = [0.0] * 4
    for i, vol in enumerate(['H', 'C', 'N', 'S']):
        res_l[i] = mass_atm_d[vol] + mass_int_d[vol] - mass_target_d[vol]

    return res_l

def get_initial_pressures(target_d):
    """Cold-start guesses for the four primary partial pressures [bar].

    Log-uniform draw over [1e-12, 1e5] bar, covering ~17 orders of
    magnitude from trace-volatile undersaturation up to ~100 kbar
    (the upper end of magma-ocean surface-pressure regimes).
    `target_d` is accepted for API stability but not consulted.
    """
    pH2O = 10 ** np.random.uniform(low=-12, high=5)
    pCO2 = 10 ** np.random.uniform(low=-12, high=5)
    pN2 = 10 ** np.random.uniform(low=-12, high=5)
    pS2 = 10 ** np.random.uniform(low=-12, high=5)

    return pH2O, pCO2, pN2, pS2

def get_partial_pressures(pin, ddict):
    """Partial pressures [bar] of all 11 species from the 4 primaries.

    `pin` provides H2O, CO2, N2, S2; the other 7 species are derived
    via equilibrium constants and the fO2 buffer.
    """

    fO2_shift = ddict['fO2_shift_IW']

    p_d = {s: 0.0 for s in volatile_species}

    p_d['H2O'] = pin['H2O']

    if is_included('H2', ddict):
        gamma = ModifiedKeq('janaf_H2')
        gamma = gamma(ddict['T_magma'], fO2_shift)
        p_d['H2'] = gamma * pin['H2O']

    p_d['CO2'] = pin['CO2']

    if is_included('CO', ddict):
        gamma = ModifiedKeq('janaf_CO')
        gamma = gamma(ddict['T_magma'], fO2_shift)
        p_d['CO'] = gamma * pin['CO2']

    if is_included('H2', ddict) and is_included('CH4', ddict):
        gamma = ModifiedKeq('schaefer_CH4')
        gamma = gamma(ddict['T_magma'], fO2_shift)
        p_d['CH4'] = gamma * pin['CO2'] * p_d['H2'] ** 2.0

    p_d['N2'] = pin['N2']

    if is_included('NH3', ddict):
        gamma = ModifiedKeq('janaf_NH3')
        gamma = gamma(ddict['T_magma'], fO2_shift)
        p_d['NH3'] = (gamma * pin['N2'] * p_d['H2'] ** 3) ** 0.5

    # O2 is set unconditionally from the fO2 buffer regardless of the
    # `O2_included` flag, because every sulfur and carbon equilibrium
    # below reads p_O2 (SO2 reaction, CO2-CO couple, CH4 couple).
    fO2_model = OxygenFugacity()
    p_d['O2'] = 10.0 ** fO2_model(ddict['T_magma'], fO2_shift)

    p_d['S2'] = pin['S2']

    if is_included('SO2', ddict):
        gamma = ModifiedKeq('janaf_SO2')
        gamma = gamma(ddict['T_magma'], fO2_shift)
        p_d['SO2'] = (gamma * pin['S2'] * p_d['O2'] ** 2) ** 0.5

    if is_included('H2S', ddict) and is_included('H2', ddict):
        gamma = ModifiedKeq('janaf_H2S')
        gamma = gamma(ddict['T_magma'], fO2_shift)
        p_d['H2S'] = (gamma * pin['S2'] * p_d['H2'] ** 2) ** 0.5

    # Silent clip: solver Monte-Carlo restarts can produce negative
    # `pin`, which then propagates through the sqrt expressions; the
    # downstream mass tallies require non-negative pressures.
    for k in p_d.keys():
        p_d[k] = max(0.0, p_d[k])

    return p_d

def get_total_pressure(p_d):
    """Sum partial pressures to get total pressure"""
    return sum(p_d.values())

def obj(pin_arr, ddict, mass_target_d):
    """Function to compute the residual of the mass balance given the partial pressures [bar]"""

    res_l = func(pin_arr, ddict, mass_target_d)
    return np.dot(res_l, res_l) ** 0.5