zephyrus.nozzle

def curvature_a(q: float) -> float:
    """Dimensionless L1 curvature A(q) of Jackson et al. (2017) Eq. (10).

    ``q`` is the donor-over-accretor mass ratio. The fit is symmetric under
    ``q`` to ``1/q``, equals 8 at equal masses, and tends to 4 at extreme
    ratios; it reproduces the numerical root of their Eq. (7) to 0.3% for
    all ``q`` (their Figure 2).
    """
    if q <= 0.0 or not math.isfinite(q):
        raise ValueError(f'q must be a positive finite mass ratio, got {q!r}')
    return 4.0 + _B1 / (_B2 + q ** (1.0 / 3.0) + q ** (-1.0 / 3.0))

def eggleton_lobe_radius(q: float, separation: float) -> float:
    """Volume-equivalent Roche lobe radius of the donor, in m.

    The Eggleton (1983, ApJ 268, 368) fit as printed in Jackson et al.
    (2017) Section 2.1, accurate to 1% for all mass ratios ``q`` (donor
    over accretor); ``separation`` is the orbital separation [m].
    """
    if q <= 0.0 or not math.isfinite(q):
        raise ValueError(f'q must be a positive finite mass ratio, got {q!r}')
    _positive('separation', separation)
    q13 = q ** (1.0 / 3.0)
    return separation * 0.49 * q13**2 / (0.6 * q13**2 + math.log(1.0 + q13))

def volume_averaged_potential(r_v: float, M_d: float, M_a: float, separation: float) -> float:
    """Volume-averaged Roche potential at volume-equivalent radius ``r_v``.

    Jackson et al. (2017) Eq. (14): the potential of the equipotential
    surface enclosing the same volume as a sphere of radius ``r_v`` around
    the donor ``M_d``, with accretor ``M_a`` at ``separation`` [m], in
    J/kg. The expansion converges inside the donor's Roche lobe and not
    outside it, and disagrees with a direct numerical evaluation by a few
    percent as ``r_v`` approaches the lobe radius, which the primary
    quantifies as about a factor of two in the final rate.
    """
    _positive('r_v', r_v)
    _positive('M_d', M_d)
    _positive('M_a', M_a)
    _positive('separation', separation)
    m_t = M_d + M_a
    x = r_v / separation
    bracket = (
        1.0
        + (m_t / M_d) * x**3 / 3.0
        + (4.0 / 45.0) * ((m_t**2 + 9.0 * M_a**2 + 3.0 * M_a * m_t) / M_d**2) * x**6
    )
    return (
        -(G * M_a / separation + G * M_a**2 / (2.0 * separation * m_t))
        - (G * M_d / r_v) * bracket
    )

def nozzle_candidate(
    M_p: float,
    M_star: float,
    a: float,
    e: float,
    rho_ph: float,
    r_ph: float,
    T: float,
    mu_kg: float,
    n_phase: int = 64,
    sigma_m2: float | None = None,
    kn_crit: float = 1.0,
) -> tuple[float, dict]:
    """Orbit-averaged Roche-lobe overflow rate through the L1 nozzle.

    The rate is Jackson et al. (2017) Eq. (3) at each orbital separation,
    averaged in time over the orbit and duty-cycled over the arc where the
    overflow description applies.

    Parameters
    ----------
    M_p, M_star : float
        Donor (planet) and accretor (star) masses [kg].
    a, e : float
        Semi-major axis [m] and eccentricity.
    rho_ph, r_ph : float
        Density [kg m^-3] and radius [m] of the launch level. The profile
        radius stands in for the volume-equivalent photospheric radius
        without the primary's Appendix distortion conversion. Along an
        isothermal column rho_ph exp(Phi_ph / v_th^2) is level-invariant,
        so the level choice cancels there and not on a non-isothermal one.
    T, mu_kg : float
        Temperature [K] and mean particle mass [kg] evaluating the
        isothermal sound speed and the exponential barrier.
    n_phase : int
        Midpoint nodes in eccentric anomaly for the orbit average; the
        quadrature is converged well below the default. At ``e = 0`` every
        node holds the same value and the average is the instantaneous rate.
    sigma_m2 : float or None
        Collision cross section [m^2] of the gas at L1. When given, a phase
        applies only where the L1 Knudsen number is at most ``kn_crit``;
        ``None`` skips that test, which leaves the closed form comparable
        with the primary's published rates.
    kn_crit : float
        Knudsen threshold at L1, used only with ``sigma_m2``.

    Returns
    -------
    (rate, detail)
        The orbit-averaged Eq. (3) rate [kg/s], unguarded so that it stays
        comparable with the primary's published rates, and a detail dict.
        The rate a caller should compete is
        ``detail['rate_applicable_kg_s']``, the same average duty-cycled
        over the arc where the overflow description applies. The detail
        dict holds the phase-independent sound speed, sonic radius, mass
        ratio, and curvature; the geometry (lobe radius, both potentials,
        the applied exponent, the nozzle area, ``saturated``) at periapsis;
        the applicable and saturated orbit fractions; the periapsis and
        apoapsis rates; the L1 density, Gaussian width, and Knudsen number at
        periapsis (``n_L1_m3``, ``H_L1_m``, ``kn_L1``); and the lift power,
        duty-cycled
        (``power_lift_W``) and over the full orbit
        (``power_lift_full_orbit_W``).
    """
    if n_phase < 1:
        raise ValueError(f'n_phase must be at least 1, got {n_phase!r}')
    for name, value in (
        ('M_p', M_p),
        ('M_star', M_star),
        ('a', a),
        ('r_ph', r_ph),
        ('T', T),
        ('mu_kg', mu_kg),
    ):
        _positive(name, value)
    if not math.isfinite(rho_ph) or rho_ph < 0.0:
        raise ValueError(f'rho_ph must be non-negative and finite, got {rho_ph!r}')
    if not math.isfinite(e) or not 0.0 <= e < 1.0:
        raise ValueError(f'e must satisfy 0 <= e < 1, got {e!r}')
    if sigma_m2 is not None:
        _positive('sigma_m2', sigma_m2)
    _positive('kn_crit', kn_crit)
    v_th = math.sqrt(kb * T / mu_kg)
    q = M_p / M_star
    a_curv = curvature_a(q)
    # The Figure 9 crossover quantity: the overflow description applies
    # where this sonic radius reaches the L1 distance (see module notes).
    r_sonic = G * M_p / (2.0 * v_th**2)

    # A secular caller needs mass per unit time, so each phase takes the
    # circular formula at its separation a (1 - e cos E), time-weighted by
    # Kepler's dt proportional to (1 - e cos E) dE.
    w_sum = rate_sum = power_sum = duty_rate_sum = duty_power_sum = 0.0
    applicable_sum = saturated_sum = 0.0
    for i in range(n_phase):
        ecc_anomaly = (i + 0.5) * 2.0 * math.pi / n_phase
        w = 1.0 - e * math.cos(ecc_anomaly)
        st = _phase_state(M_p, M_star, a * w, rho_ph, r_ph, v_th, q, a_curv)
        w_sum += w
        rate_sum += w * st['rate']
        power_sum += w * st['power']
        # Off the applicable arc, which surrounds periapsis because the L1
        # distance grows with separation, the gas chokes at its own sonic
        # surface first and the nozzle carries nothing.
        if r_sonic >= st['R_L1'] and _collisional(st, v_th, mu_kg, sigma_m2, kn_crit):
            duty_rate_sum += w * st['rate']
            duty_power_sum += w * st['power']
            applicable_sum += w
        if st['saturated']:
            saturated_sum += w

    peri = _phase_state(M_p, M_star, a * (1.0 - e), rho_ph, r_ph, v_th, q, a_curv)
    apo = _phase_state(M_p, M_star, a * (1.0 + e), rho_ph, r_ph, v_th, q, a_curv)
    n_l1, h_l1 = _l1_density_and_width(peri, v_th, mu_kg)
    return rate_sum / w_sum, dict(
        T_K=T,
        mu_kg=mu_kg,
        v_th=v_th,
        R_sonic=r_sonic,
        q=q,
        A=a_curv,
        n_phase=n_phase,
        a_periapsis=a * (1.0 - e),
        r_lobe=peri['r_lobe'],
        phi_L1=peri['phi_L1'],
        phi_ph=peri['phi_ph'],
        delta_phi=peri['delta_phi'],
        exponent_applied=peri['exponent_applied'],
        area_m2=peri['area_m2'],
        saturated=peri['saturated'],
        R_L1=peri['R_L1'],
        R_sonic_over_R_L1=r_sonic / peri['R_L1'],
        R_sonic_over_R_L1_apoapsis=r_sonic / apo['R_L1'],
        rate_periapsis_kg_s=peri['rate'],
        rate_apoapsis_kg_s=apo['rate'],
        rate_applicable_kg_s=duty_rate_sum / w_sum,
        n_L1_m3=n_l1,
        H_L1_m=h_l1,
        kn_L1=_kn_l1(n_l1, h_l1, sigma_m2),
        applicable=applicable_sum > 0.0,
        applicable_orbit_fraction=applicable_sum / w_sum,
        saturated_orbit_fraction=saturated_sum / w_sum,
        power_lift_W=duty_power_sum / w_sum,
        power_lift_full_orbit_W=power_sum / w_sum,
    )