zephyrus.collision

def mass_loss(
    v_c: float,
    M_i: float,
    M_t: float,
    rho_i: float,
    rho_t: float,
    R_i: float,
    R_t: float,
    b: float,
) -> float:
    r"""Fractional atmospheric mass loss of the target in a giant impact.

    Implements the scaling law of Kegerreis et al. (2020), their Eqn. 1:

    $X \approx 0.64 \left[ \left(\frac{v_c}{v_{esc}}\right)^2
    \left(\frac{M_i}{M_{tot}}\right)^{1/2}
    \left(\frac{\rho_i}{\rho_t}\right)^{1/2} f_M(b) \right]^{0.65}$

    capped at 1 for total erosion, where subscript (i) is the impactor,
    (t) the target, and $M_{tot} = M_i + M_t$. The mutual escape speed is
    $v_{esc} = \sqrt{2 G (M_t + M_i) / (R_t + R_i)}$, and $f_M(b)$ is the
    fractional interacting mass of their Eqn. B1, built from the
    density-weighted spherical caps of common height
    $d = (R_t + R_i)(1 - b)$. The common-height caps are a linearised
    bookkeeping, so outside the fitted geometry (a much denser, much
    smaller impactor near head-on) the raw $f_M$ can leave $[0, 1]$, and
    is clamped to it here, and can vary non-monotonically with $b$; at
    equal bulk densities $f_M$ reduces exactly to the interacting volume
    $f_V$ of their Eqn. B2.

    Conventions the caller must honour (Kegerreis et al. 2020, Sect. 2):
    $v_c$ is the speed at first contact, not at infinity; the masses and
    radii exclude any atmosphere, with the radii taken at its base; and
    $b \equiv \sin\beta$ for impact angle $\beta$ (0 head-on, 1 grazing).

    The fit is constrained for target masses of roughly 0.3 to 3 Earth
    masses, impactors down to about 0.05 Earth masses, bulk densities of
    about half to double Earth's, speeds of 1 to 3 $v_{esc}$, any angle,
    and thin atmospheres of order 1 percent of the planet mass. The
    median deviation of the simulations from the law is 9 percent,
    rising to about 20 percent for slow, head-on impacts.

    Parameters
    ----------
    v_c : float
        Collision speed at first contact between impactor and target [m/s].
    M_i : float
        Mass of the impactor, excluding any atmosphere [kg].
    M_t : float
        Mass of the target, excluding any atmosphere [kg].
    rho_i : float
        Bulk density of the impactor, excluding any atmosphere [kg/m^3].
    rho_t : float
        Bulk density of the target, excluding any atmosphere [kg/m^3].
    R_i : float
        Radius of the impactor, at the base of any atmosphere [m].
    R_t : float
        Radius of the target, at the base of any atmosphere [m].
    b : float
        Dimensionless impact parameter, the sine of the impact angle,
        in [0, 1]: 0 is head-on, 1 is fully grazing.

    Returns
    -------
    float
        Fractional mass loss of the target body's atmosphere, in [0, 1].

    Raises
    ------
    ValueError
        If ``b`` lies outside [0, 1], if any mass, radius, or density is
        not strictly positive and finite, or if ``v_c`` is negative,
        not sub-luminal (>= c), or not finite. Inputs are scalar;
        arrays are not supported.

    References
    ----------
    1. Kegerreis J.A., Eke V.R., Catling D.C., Massey R.J., Teodoro
       L.F.A., Zahnle K.J. (2020). Atmospheric Erosion by Giant Impacts
       onto Terrestrial Planets: A Scaling Law for any Speed, Angle,
       Mass, and Density. ApJL 901, L31. doi:10.3847/2041-8213/abb5fb
    """
    if not 0.0 <= b <= 1.0:
        raise ValueError(f'Impact parameter b must be in [0, 1], got {b!r}')
    for name, value in (
        ('M_i', M_i),
        ('M_t', M_t),
        ('rho_i', rho_i),
        ('rho_t', rho_t),
        ('R_i', R_i),
        ('R_t', R_t),
    ):
        if not (value > 0.0 and np.isfinite(value)):
            raise ValueError(f'{name} must be strictly positive and finite, got {value!r}')
    if not (0.0 <= v_c < c and np.isfinite(v_c)):
        raise ValueError(
            f'Collision speed v_c must be non-negative, sub-luminal (< c), and finite, got {v_c!r}'
        )

    # Mutual escape speed of the pair at contact
    v_esc = np.sqrt((2.0 * G * (M_t + M_i)) / (R_t + R_i))

    # Fractional interacting mass f_M (Kegerreis et al. 2020, Eqn. B1).
    f_m = _interacting_mass_fraction_kegerreis(R_t, R_i, rho_t, rho_i, b)

    m_tot = M_i + M_t
    bracket = (v_c / v_esc) ** 2 * (M_i / m_tot) ** 0.5 * (rho_i / rho_t) ** 0.5 * f_m

    # Fractional atmospheric mass loss of the target, capped at 1 for
    # total erosion (Kegerreis et al. 2020, Eqn. 1)
    x_loss = 0.64 * bracket**0.65
    return min(max(x_loss, 0.0), 1.0)

def mass_loss_roche2026(
    v_c: float,
    M_i: float,
    M_t: float,
    R_i: float,
    R_t: float,
    b: float,
    f_atm: float,
) -> float:
    r"""Fractional atmospheric mass loss of the target in a giant impact.

    Implements the scaling law of Roche et al. (2026).

    Parameters
    ----------
    v_c : float
        Collision speed at first contact [m/s].
    M_i : float
        Impactor refractory mass, excluding atmosphere [kg].
    M_t : float
        Target refractory mass, excluding atmosphere [kg].
    R_i : float
        Impactor refractory radius, at base of atmosphere [m].
    R_t : float
        Target refractory radius, at base of atmosphere [m].
    b : float
        Dimensionless impact parameter $\sin\beta$ in [0, 1].
    f_atm : float
        Target atmospheric mass fraction $M_{\text{atm}} / (M_t + M_{\text{atm}})$ in [0, 1).

    Returns
    -------
    float
        Fractional atmospheric mass loss of the target in [0, 1].

    Raises
    ------
    TypeError
        If any input is a string, bytes, or non-scalar sequence/array.
    ValueError
        If inputs violate physical domain constraints.

    Notes
    -----
    Target mass M_t is the refractory core-plus-mantle mass; mutual
    escape speed v_esc and specific impact energy Q'_R evaluate with total
    target mass M_t / (1 - f_atm). Stability clamps apply to inputs outside
    the numerical stability bounds. This function returns only the loss
    fraction; use ``impact_loss`` if validity flags or diagnostics are needed.
    Inputs are scalar; arrays are not supported.

    References
    ----------
    1. Roche M.J., Lock S.J., Carter P.J., Leinhardt Z.M. (2026).
       Giant impacts preferentially remove low-mass atmospheres:
       a generalised scaling law for impact-driven atmospheric loss.
       arXiv:2610.06077 (accepted, ApJL); Zenodo doi:10.5281/zenodo.23192423.
    """
    frac, _, _ = _eval_roche2026(v_c, M_i, M_t, R_i, R_t, b, f_atm)
    return frac

def impact_loss(
    law: str,
    *,
    v_c: float,
    M_i: float,
    M_t: float,
    R_i: float,
    R_t: float,
    b: float,
    rho_i: float | None = None,
    rho_t: float | None = None,
    f_atm: float | None = None,
) -> ImpactLossResult:
    r"""Evaluate atmospheric erosion scaling law for a giant impact.

    Parameters
    ----------
    law : str
        Scaling law name ('kegerreis2020' or 'roche2026').
    v_c : float
        Collision speed at first contact [m/s].
    M_i : float
        Impactor mass, excluding atmosphere [kg].
    M_t : float
        Target mass, excluding atmosphere [kg].
    R_i : float
        Impactor radius, at base of atmosphere [m].
    R_t : float
        Target radius, at base of atmosphere [m].
    b : float
        Dimensionless impact parameter $\sin\beta$ in [0, 1].
    rho_i : float, optional
        Impactor bulk density [kg/m^3] (required for 'kegerreis2020').
    rho_t : float, optional
        Target bulk density [kg/m^3] (required for 'kegerreis2020').
    f_atm : float, optional
        Target atmospheric mass fraction in [0, 1) (required for 'roche2026').

    Returns
    -------
    ImpactLossResult
        Container holding law name, loss fraction, flags, and diagnostics.

    Raises
    ------
    TypeError
        If any input is a string, bytes, or non-scalar sequence/array.
    ValueError
        If law is unsupported or required law-specific inputs are missing.

    Notes
    -----
    For 'roche2026', target mass M_t is the refractory mass, while mutual escape
    speed v_esc and specific impact energy Q'_R evaluate with total target mass
    M_t / (1 - f_atm). Non-zero far-field loss at zero impact energy is recorded in
    diagnostics['X_FF_zero_energy'] and flagged as 'X_FF_zero_energy'. Arguments
    not used by the selected scaling law are ignored (e.g. bulk densities for
    'roche2026', or f_atm for 'kegerreis2020'). An airless target (f_atm = 0)
    returns no diagnostic flags. Inputs are scalar; arrays are not supported.
    """
    if law == 'kegerreis2020':
        if rho_i is None or rho_t is None:
            missing = [k for k, v in (('rho_i', rho_i), ('rho_t', rho_t)) if v is None]
            raise ValueError(f'kegerreis2020 requires {", ".join(missing)}')
        vals = _as_floats(
            v_c=v_c, M_i=M_i, M_t=M_t, rho_i=rho_i, rho_t=rho_t, R_i=R_i, R_t=R_t, b=b
        )
        frac = mass_loss(**vals)
        v_esc = mutual_escape_speed(vals['M_t'], vals['M_i'], vals['R_t'], vals['R_i'])
        diag = {
            'v_esc': v_esc,
            'v_ratio': vals['v_c'] / v_esc,
            'gamma': vals['M_i'] / (vals['M_t'] + vals['M_i']),
            'mass_ratio': vals['M_i'] / vals['M_t'],
            'f_M': _interacting_mass_fraction_kegerreis(
                vals['R_t'], vals['R_i'], vals['rho_t'], vals['rho_i'], vals['b']
            ),
        }
        return ImpactLossResult(law=law, fraction=frac, flags=(), diagnostics=diag)

    if law == 'roche2026':
        if f_atm is None:
            raise ValueError('roche2026 requires f_atm')
        frac, flags, diag = _eval_roche2026(v_c, M_i, M_t, R_i, R_t, b, f_atm)
        return ImpactLossResult(law=law, fraction=frac, flags=flags, diagnostics=diag)

    raise ValueError(f"Unknown scaling law {law!r}, expected 'kegerreis2020' or 'roche2026'")

def mutual_escape_speed(
    M_1: float,
    M_2: float,
    R_1: float,
    R_2: float,
) -> float:
    r"""Mutual escape speed of two spherical bodies at contact [m/s].

    Parameters
    ----------
    M_1 : float
        Mass of the first body [kg].
    M_2 : float
        Mass of the second body [kg].
    R_1 : float
        Radius of the first body [m].
    R_2 : float
        Radius of the second body [m].

    Returns
    -------
    float
        Mutual escape speed $\sqrt{2 G (M_1 + M_2) / (R_1 + R_2)}$ [m/s].

    Raises
    ------
    TypeError
        If any input is a string, bytes, or non-scalar sequence/array.
    ValueError
        If any mass or radius is not strictly positive and finite, or if
        floating-point conversion overflows.
    """
    vals = _as_floats(M_1=M_1, M_2=M_2, R_1=R_1, R_2=R_2)
    _check_strictly_positive(**vals)
    return float(np.sqrt(2.0 * G * (vals['M_1'] + vals['M_2']) / (vals['R_1'] + vals['R_2'])))

def specific_impact_energy(
    v_c: float,
    M_i: float,
    M_t_tot: float,
    R_i: float,
    R_t: float,
    b: float,
) -> float:
    r"""Modified specific impact energy Q'_R in MJ/kg.

    Parameters
    ----------
    v_c : float
        Impact velocity at surface contact [m/s].
    M_i : float
        Impactor mass [kg].
    M_t_tot : float
        Target mass including its atmosphere ($M_\mathrm{t}^\mathrm{tot}$ in
        Roche et al. (2026), Eq. 1 and $Q'_\mathrm{R}$) [kg].
    R_i : float
        Impactor refractory radius [m].
    R_t : float
        Target refractory radius [m].
    b : float
        Impact parameter $\sin\beta$ in [0, 1].

    Returns
    -------
    float
        Modified specific impact energy $Q'_R$ in MJ/kg.

    Raises
    ------
    TypeError
        If any input is a string, bytes, or non-scalar sequence/array.
    ValueError
        If b is outside [0, 1], if masses or radii are not strictly positive
        and finite, or if v_c is negative, non-finite, or >= c.

    References
    ----------
    1. Roche, M. J., Lock, S. J., Dou, J., Carter, P. J., Kegerreis, J. A.,
       & Leinhardt, Z. M. (2025), "Atmospheric Loss during Giant Impacts:
       Mechanisms and Scaling of Near- and Far-field Loss", The Planetary
       Science Journal, 6, 149, doi:10.3847/PSJ/add929, Eqns. 13 to 18.
    2. Leinhardt, Z. M., & Stewart, S. T. (2012), "Collisions between
       gravity-dominated bodies. I. Outcome regimes and scaling laws",
       ApJ 745, 79, doi:10.1088/0004-637X/745/1/79.
    """
    v_c, M_i, M_t_tot, R_i, R_t, b = _as_floats(
        v_c=v_c, M_i=M_i, M_t_tot=M_t_tot, R_i=R_i, R_t=R_t, b=b
    ).values()
    _check_strictly_positive(M_i=M_i, M_t_tot=M_t_tot, R_i=R_i, R_t=R_t)
    _check_b_and_speed(b, v_c)

    impact_param = (R_t + R_i) * b
    if impact_param + R_i <= R_t:
        alpha = 1.0
    else:
        interact_len = R_t + R_i - impact_param
        alpha = (3.0 * R_i * interact_len**2 - interact_len**3) / (4.0 * R_i**3)

    m_tot = M_t_tot + M_i
    mu_alpha = (alpha * M_i) / (1.0 + (alpha * M_i) / M_t_tot)
    return float(mu_alpha * v_c**2 / (2.0 * m_tot) / 1.0e6)