Zalmoxis.mixing

def parse_layer_components(config_value: str) -> LayerMixture:
    """Parse a layer EOS config string into a LayerMixture.

    Formats
    -------
    Single material (backward compat)::

        "PALEOS:iron"  ->  LayerMixture(["PALEOS:iron"], [1.0])

    Multi-material with mass fractions::

        "PALEOS:MgSiO3:0.85+PALEOS:H2O:0.15"
        ->  LayerMixture(["PALEOS:MgSiO3", "PALEOS:H2O"], [0.85, 0.15])

    Parameters
    ----------
    config_value : str
        Layer EOS config string from TOML.

    Returns
    -------
    LayerMixture

    Raises
    ------
    ValueError
        If no components found or fractions are invalid.
    """
    if not config_value or not config_value.strip():
        raise ValueError('Empty EOS config string.')

    segments = config_value.split('+')
    components = []
    fractions = []

    for seg in segments:
        eos_name, frac = _parse_single_component(seg)
        if not eos_name:
            raise ValueError(f'Empty EOS name in component: {seg!r}')
        components.append(eos_name)
        fractions.append(frac)

    if not components:
        raise ValueError(f'No components parsed from: {config_value!r}')

    # Normalize fractions if needed
    total = sum(fractions)
    if len(components) > 1 and abs(total - 1.0) > 1e-6:
        logger.warning(
            f'Mass fractions sum to {total:.6f}, normalizing to 1.0. '
            f'Components: {components}, fractions: {fractions}'
        )
        fractions = [f / total for f in fractions]
    elif len(components) == 1:
        fractions = [1.0]

    if any(f < 0 for f in fractions):
        raise ValueError(f'Negative mass fraction in {config_value!r}: {fractions}')

    return LayerMixture(components=components, fractions=fractions)

def parse_all_layer_mixtures(
    layer_eos_config: dict[str, str],
) -> dict[str, LayerMixture]:
    """Parse all layers in layer_eos_config into LayerMixture objects.

    Parameters
    ----------
    layer_eos_config : dict
        Per-layer EOS config, e.g.
        ``{"core": "PALEOS:iron", "mantle": "PALEOS:MgSiO3:0.85+PALEOS:H2O:0.15"}``.

    Returns
    -------
    dict
        Per-layer LayerMixture objects, keyed by layer name.
    """
    return {
        layer: parse_layer_components(eos_str)
        for layer, eos_str in layer_eos_config.items()
        if eos_str
    }

def calculate_mixed_density(
    pressure,
    temperature,
    mixture,
    material_dictionaries,
    solidus_func,
    liquidus_func,
    interpolation_functions,
    mushy_zone_factors=None,
    condensed_rho_min=CONDENSED_RHO_MIN_DEFAULT,
    condensed_rho_scale=CONDENSED_RHO_SCALE_DEFAULT,
    binodal_T_scale=BINODAL_T_SCALE_DEFAULT,
):
    """Compute volume-additive mixed density for a layer mixture.

    For a single component, delegates directly to ``calculate_density()``.
    For multiple components, evaluates each independently at (P, T) and
    returns a suppressed harmonic mean where each component is weighted
    by the product of:

    1. A density-based sigmoid (``_condensed_weight``): suppresses
       gas-like components by density.
    2. A binodal sigmoid (``_binodal_factor``): suppresses H2 when it
       is thermodynamically immiscible with its partner (below the
       binodal temperature). Only applies to H2 components in mixtures
       with silicate or water.

    Notes
    -----
    The suppressed harmonic mean does not conserve mass: when a component
    is suppressed (sigma near 0), its mass is excluded from the structural
    calculation. This is physically equivalent to treating vapor-phase
    volatiles as having negligible structural contribution. The
    approximation is valid when suppressed components are at low density
    (vapor/gas) and do not contribute meaningfully to hydrostatic support.

    Parameters
    ----------
    pressure : float
        Pressure in Pa.
    temperature : float
        Temperature in K.
    mixture : LayerMixture
        Layer mixture specification.
    material_dictionaries : dict
        EOS registry.
    solidus_func : callable or None
        Solidus melting curve function.
    liquidus_func : callable or None
        Liquidus melting curve function.
    interpolation_functions : dict
        Interpolation cache.
    mushy_zone_factors : dict or float or None
        Per-EOS mushy zone factors. Dict keyed by EOS name, a single
        float (applied to all), or None (default 1.0 for all).
    condensed_rho_min : float
        Sigmoid center for phase-aware suppression (kg/m^3).
    condensed_rho_scale : float
        Sigmoid width for phase-aware suppression (kg/m^3).
    binodal_T_scale : float
        Binodal sigmoid width in K. Controls the smooth transition at
        the miscibility boundary. Default 50 K.

    Returns
    -------
    float or None
        Mixed density in kg/m^3, or None if all components are gas-like.
    """
    from .eos_functions import calculate_density

    # Fast path: single component (no suppression)
    if mixture.is_single():
        mzf = _get_mushy_zone_factor(mixture.components[0], mushy_zone_factors)
        return calculate_density(
            pressure,
            material_dictionaries,
            mixture.components[0],
            temperature,
            solidus_func,
            liquidus_func,
            interpolation_functions,
            mzf,
        )

    # Multi-component suppressed harmonic mean
    w_eff_sum = 0.0
    inv_rho_sum = 0.0
    for eos_name, w_i in zip(mixture.components, mixture.fractions):
        if w_i <= 0:
            continue
        mzf = _get_mushy_zone_factor(eos_name, mushy_zone_factors)
        rho_i = calculate_density(
            pressure,
            material_dictionaries,
            eos_name,
            temperature,
            solidus_func,
            liquidus_func,
            interpolation_functions,
            mzf,
        )
        if rho_i is None or not np.isfinite(rho_i) or rho_i <= 0:
            # Abort the entire mixture rather than skipping this component.
            # A None from calculate_density indicates an EOS table coverage
            # gap, not a vapor-phase state (vapor has low but finite density).
            # Treating it as suppressed (continue) would silently hide table
            # errors. The Picard loop falls back to old_density for this shell.
            return None
        # Per-component sigmoid center and width: each volatile has its own
        # critical density. Fall back to the global config value if not listed.
        rho_min_i = _COMPONENT_RHO_MIN.get(eos_name, condensed_rho_min)
        rho_scale_i = _COMPONENT_RHO_SCALE.get(eos_name, condensed_rho_scale)
        sigma_i = _condensed_weight(rho_i, rho_min_i, rho_scale_i)
        sigma_i *= _binodal_factor(
            eos_name, w_i, mixture, pressure, temperature, binodal_T_scale
        )
        w_eff = w_i * sigma_i
        if w_eff <= 0:
            continue
        w_eff_sum += w_eff
        inv_rho_sum += w_eff / rho_i

    if w_eff_sum <= 0 or inv_rho_sum <= 0:
        return None
    return w_eff_sum / inv_rho_sum

def calculate_mixed_density_batch(
    pressures,
    temperatures,
    mixture,
    material_dictionaries,
    solidus_func,
    liquidus_func,
    interpolation_functions,
    mushy_zone_factors=None,
    condensed_rho_min=CONDENSED_RHO_MIN_DEFAULT,
    condensed_rho_scale=CONDENSED_RHO_SCALE_DEFAULT,
    binodal_T_scale=BINODAL_T_SCALE_DEFAULT,
):
    """Vectorized mixed density for a batch of (P, T) points in one layer.

    Parameters
    ----------
    pressures : numpy.ndarray
        1D array of pressures in Pa.
    temperatures : numpy.ndarray
        1D array of temperatures in K.
    mixture : LayerMixture
        Layer mixture specification (same for all points in the batch).
    material_dictionaries : dict
        EOS registry.
    solidus_func : callable or None
        Solidus melting curve function.
    liquidus_func : callable or None
        Liquidus melting curve function.
    interpolation_functions : dict
        Interpolation cache.
    mushy_zone_factors : dict or float or None
        Per-EOS mushy zone factors.
    condensed_rho_min : float
        Global sigmoid center fallback (kg/m^3).
    condensed_rho_scale : float
        Global sigmoid width fallback (kg/m^3).
    binodal_T_scale : float
        Binodal sigmoid width in K.

    Returns
    -------
    numpy.ndarray
        1D array of densities in kg/m^3. NaN where lookup fails or all
        components are suppressed.
    """
    from .eos_functions import calculate_density_batch

    n = len(pressures)

    # Fast path: single component (no suppression)
    if mixture.is_single():
        mzf = _get_mushy_zone_factor(mixture.components[0], mushy_zone_factors)
        return calculate_density_batch(
            pressures,
            temperatures,
            material_dictionaries,
            mixture.components[0],
            solidus_func,
            liquidus_func,
            interpolation_functions,
            mzf,
        )

    # Multi-component: evaluate each component, then suppressed harmonic mean
    w_eff_sum = np.zeros(n)
    inv_rho_sum = np.zeros(n)
    any_invalid = np.zeros(n, dtype=bool)

    for eos_name, w_i in zip(mixture.components, mixture.fractions):
        if w_i <= 0:
            continue
        mzf = _get_mushy_zone_factor(eos_name, mushy_zone_factors)
        rho_i = calculate_density_batch(
            pressures,
            temperatures,
            material_dictionaries,
            eos_name,
            solidus_func,
            liquidus_func,
            interpolation_functions,
            mzf,
        )
        # Mark shells where any component has invalid density
        bad = ~np.isfinite(rho_i) | (rho_i <= 0)
        any_invalid |= bad

        # Per-component sigmoid
        rho_min_i = _COMPONENT_RHO_MIN.get(eos_name, condensed_rho_min)
        rho_scale_i = _COMPONENT_RHO_SCALE.get(eos_name, condensed_rho_scale)
        sigma_i = _condensed_weight_batch(np.where(bad, 1.0, rho_i), rho_min_i, rho_scale_i)

        # Binodal suppression (scalar per shell, only for H2 components)
        if eos_name in _H2_EOS_NAMES:
            for j in range(n):
                if not bad[j]:
                    sigma_i[j] *= _binodal_factor(
                        eos_name,
                        w_i,
                        mixture,
                        pressures[j],
                        temperatures[j],
                        binodal_T_scale,
                    )

        w_eff = w_i * sigma_i
        w_eff = np.where(bad, 0.0, w_eff)

        ok = w_eff > 0
        w_eff_sum += np.where(ok, w_eff, 0.0)
        inv_rho_sum += np.where(ok, w_eff / np.where(ok, rho_i, 1.0), 0.0)

    result = np.where(
        (w_eff_sum > 0) & (inv_rho_sum > 0) & ~any_invalid,
        w_eff_sum / np.where(inv_rho_sum > 0, inv_rho_sum, 1.0),
        np.nan,
    )
    return result

def get_mixed_nabla_ad(
    pressure,
    temperature,
    mixture,
    material_dictionaries,
    interpolation_functions,
    solidus_func=None,
    liquidus_func=None,
    mushy_zone_factors=None,
    condensed_rho_min=CONDENSED_RHO_MIN_DEFAULT,
    condensed_rho_scale=CONDENSED_RHO_SCALE_DEFAULT,
    binodal_T_scale=BINODAL_T_SCALE_DEFAULT,
    precomputed_densities=None,
):
    """Compute mass-fraction-weighted nabla_ad for a mixture.

    For a single component, returns that component's nabla_ad directly.
    For multiple components, computes a weighted average where each
    component's contribution is scaled by the same sigmoid suppression
    used for density. This prevents vapor-phase components from
    contributing their (typically large) nabla_ad to the mixture.

    Components that do not support nabla_ad (e.g., Seager2007) are skipped
    and their fraction is redistributed among T-dependent components.

    Parameters
    ----------
    pressure : float
        Pressure in Pa.
    temperature : float
        Temperature in K.
    mixture : LayerMixture
        Layer mixture specification.
    material_dictionaries : dict
        EOS registry.
    interpolation_functions : dict
        Interpolation cache.
    solidus_func : callable or None
        Solidus function (for PALEOS-2phase).
    liquidus_func : callable or None
        Liquidus function (for PALEOS-2phase).
    mushy_zone_factors : dict or float or None
        Per-EOS mushy zone factors. Dict keyed by EOS name, a single
        float (applied to all), or None (default 1.0 for all).
    condensed_rho_min : float
        Sigmoid center for phase-aware suppression (kg/m^3).
    condensed_rho_scale : float
        Sigmoid width for phase-aware suppression (kg/m^3).
    binodal_T_scale : float
        Binodal sigmoid width in K for H2 miscibility suppression.
        Default 50 K.
    precomputed_densities : dict or None
        Optional pre-computed per-component densities, keyed by EOS name.
        When provided, skips the ``calculate_density()`` call for the
        sigmoid weight computation, avoiding redundant EOS table lookups
        when density was already computed in ``calculate_mixed_density()``.

    Returns
    -------
    float or None
        Dimensionless adiabatic gradient, or None if no component provides it.

    Notes
    -----
    Suppressing a vapor component's nabla_ad means the temperature profile
    ignores the volatile. This is consistent with the density suppression
    (both treat the vapor as structurally absent) but differs from reality
    where vapor affects heat transport. This is a known limitation.
    """
    from .eos_functions import calculate_density

    if mixture.is_single():
        return _nabla_ad_for_component(
            mixture.components[0],
            pressure,
            temperature,
            material_dictionaries,
            interpolation_functions,
            solidus_func,
            liquidus_func,
        )

    # Multi-component weighted average with condensed-phase suppression
    weighted_sum = 0.0
    weight_total = 0.0

    for eos_name, w_i in zip(mixture.components, mixture.fractions):
        if w_i <= 0:
            continue
        # Use pre-computed density if available, otherwise compute it
        if precomputed_densities is not None and eos_name in precomputed_densities:
            rho_i = precomputed_densities[eos_name]
        else:
            mzf = _get_mushy_zone_factor(eos_name, mushy_zone_factors)
            rho_i = calculate_density(
                pressure,
                material_dictionaries,
                eos_name,
                temperature,
                solidus_func,
                liquidus_func,
                interpolation_functions,
                mzf,
            )
        if rho_i is not None and np.isfinite(rho_i) and rho_i > 0:
            rho_min_i = _COMPONENT_RHO_MIN.get(eos_name, condensed_rho_min)
            rho_scale_i = _COMPONENT_RHO_SCALE.get(eos_name, condensed_rho_scale)
            sigma_i = _condensed_weight(rho_i, rho_min_i, rho_scale_i)
            sigma_i *= _binodal_factor(
                eos_name, w_i, mixture, pressure, temperature, binodal_T_scale
            )
        else:
            sigma_i = 0.0
        w_eff = w_i * sigma_i
        if w_eff <= 0:
            continue
        nabla = _nabla_ad_for_component(
            eos_name,
            pressure,
            temperature,
            material_dictionaries,
            interpolation_functions,
            solidus_func,
            liquidus_func,
        )
        if nabla is not None and np.isfinite(nabla) and nabla >= 0:
            weighted_sum += w_eff * nabla
            weight_total += w_eff

    if weight_total <= 0:
        return None

    # Renormalize by actual weight of contributing components
    return weighted_sum / weight_total

def any_component_is_tdep(
    layer_mixtures: dict[str, LayerMixture],
) -> bool:
    """Check if any component in any layer uses a T-dependent EOS.

    Parameters
    ----------
    layer_mixtures : dict
        Per-layer LayerMixture objects.

    Returns
    -------
    bool
    """
    return any(mix.has_tdep() for mix in layer_mixtures.values())