chore(linting): add Black + Ruff PEP 8 tooling and apply project-wide formatting

# -*- coding: utf-8 -*-

"""

Convenience package initializer.

The actual library code resides in the inner ``PyAETNA`` directory generated by

in test environments where only the repository root is on ``PYTHONPATH``, we

expose the inner package's contents at the top level.

"""

import importlib.util, os, sys

import importlib.util

import os

import sys

# Resolve the path to the sibling directory containing the real implementation.

_current_dir = os.path.abspath(os.path.dirname(__file__))

_inner_path = os.path.join(_current_dir, "PyAETNA")

# After adjusting ``sys.path`` we import everything from the inner package so

# that attributes are available directly under the top‑level ``PyAETNA`` name.

from .version import __version__

_spec = importlib.util.find_spec("PyAETNA")

if _spec is not None and _spec.origin != __file__:

    _module = importlib.import_module("PyAETNA")

still detected and preferred if available. The --xla_force_host_platform_device_count

flag does not prevent JAX from using GPUs when present on the system.

"""

import os

# Configure host CPU device count for JAX before importing it.

# via env var PYAETNA_XLA_FORCE_HOST_PLATFORM_DEVICE_COUNT.

# NOTE: This setting does NOT affect GPU detection or usage.

_cpu_count = os.cpu_count() or 1

_override = os.environ.get('PYAETNA_XLA_FORCE_HOST_PLATFORM_DEVICE_COUNT')

_override = os.environ.get("PYAETNA_XLA_FORCE_HOST_PLATFORM_DEVICE_COUNT")

try:

    _n_devices = max(1, int(_override)) if _override is not None else _cpu_count

except ValueError:

    _n_devices = _cpu_count

_flag = f'--xla_force_host_platform_device_count={_n_devices}'

_existing = os.environ.get('XLA_FLAGS', '')

if '--xla_force_host_platform_device_count=' not in _existing:

    os.environ['XLA_FLAGS'] = (_existing + ' ' + _flag).strip()

_flag = f"--xla_force_host_platform_device_count={_n_devices}"

_existing = os.environ.get("XLA_FLAGS", "")

if "--xla_force_host_platform_device_count=" not in _existing:

    os.environ["XLA_FLAGS"] = (_existing + " " + _flag).strip()

from jaxoplanet.orbits.keplerian import System, Central, Body

from jaxoplanet.orbits.ttv import TTVOrbit, compute_expected_transit_times

import jax

import jax.numpy as jnp

import matplotlib.pyplot as plt

import numpy as np

from astropy.constants import G, M_sun, R_sun, c

from jax import jit

from jaxoplanet.light_curves import limb_dark_light_curve

from jaxoplanet.light_curves.emission import light_curve as emission_lc

from jaxoplanet.starry.orbit import SurfaceSystem

from jaxoplanet.starry.surface import Surface

from jaxoplanet.orbits.keplerian import Body, Central, System

from jaxoplanet.orbits.ttv import TTVOrbit

from jaxoplanet.starry.core.basis import (

    A1 as _jxp_A1,

)

from jaxoplanet.starry.core.basis import (

    A2_inv as _jxp_A2_inv,

)

# Internal jaxoplanet functions used for the dense RM reimplementation below.

# We avoid surface_radial_velocity directly because its internal sparse matrix

# operations (BCOO) fail in the backward pass when used inside jax.vmap

# (see compute_rossiter_mclaughlin docstring for details).

from jaxoplanet.starry.core.basis import U as _jxp_U, A1 as _jxp_A1, A2_inv as _jxp_A2_inv

from jaxoplanet.starry.core.solution import rT as _jxp_rT, solution_vector as _jxp_solution_vector

from jaxoplanet.starry.core.rotation import dot_rz as _jxp_dot_rz

from jaxoplanet.starry.core.basis import (

    U as _jxp_U,

)

from jaxoplanet.starry.core.polynomials import Pijk as _jxp_Pijk

from jaxoplanet.starry.core.rotation import dot_rz as _jxp_dot_rz

from jaxoplanet.starry.core.solution import (

    rT as _jxp_rT,

)

from jaxoplanet.starry.core.solution import (

    solution_vector as _jxp_solution_vector,

)

from jaxoplanet.starry.doppler import rv_map_expansion as _jxp_rv_map_expansion

import jax

import jax.numpy as jnp

import matplotlib.pyplot as plt

from jax import jit

from astropy.constants import M_sun, R_sun, G, c

import numpy as np

from jaxoplanet.starry.orbit import SurfaceSystem

from jaxoplanet.starry.surface import Surface

## Use 64-bit precision for accurate orbital mechanics

# jax.config.update("jax_enable_x64", True) #accuracy is already managed at higher level

# Physical constants from astropy.constants (module-level for use in JIT functions)

RHO_SUN_CGS = float(3.0 * M_sun.cgs.value / (4.0 * np.pi * R_sun.cgs.value**3))  # Solar mean density [g/cm³]

RHO_SUN_CGS = float(

    3.0 * M_sun.cgs.value / (4.0 * np.pi * R_sun.cgs.value**3)

)  # Solar mean density [g/cm³]

G_CGS = float(G.cgs.value)  # Gravitational constant [cm³/(g·s²)]

SECONDS_PER_DAY = 86400.0  # Seconds per day (exact by definition)

_RM_rT_UDEG = np.array(_jxp_rT(_RM_UDEG), dtype=np.float32)  # (9,)

_RM_rT_TOTAL = np.array(_jxp_rT(_RM_TOTAL_DEG), dtype=np.float32)  # (16,)

@jit

def reparametrize_LD(q1, q2):

    """Convert (q1, q2) to (u1, u2) limb darkening coefficients. JIT-compiled."""

    u2 = jnp.sqrt(q1) * (1 - 2 * q2)

    return (u1, u2)

@jit

def reparametrize_LD_inverse(u1, u2):

    """Convert (u1, u2) to (q1, q2) limb darkening parameters. JIT-compiled."""

    q2 = 0.5 * u1 / (u1 + u2)

    return (q1, q2)

@jit

def reparametrize_ecc_omega(ec, es):

    """

    omega = jnp.arctan2(es, ec)

    return (ecc, omega)

@jit

def compute_semimajor_axis_scaled(rho_star, period):

    """

    return a_over_rstar

@jit

def compute_inclination(b_param, a_over_rstar, ecc=0.0, omega=0.0):

    """

    rho_c = (1.0 - ecc**2) / (1.0 + ecc * jnp.sin(omega))

    return jnp.arccos(jnp.clip(b_param / (a_over_rstar * rho_c), 0.0, 1.0 - 1e-10))

@jit

def compute_transit_duration(period, rho_star, b_param, rp_by_rs, ecc=0.0, omega=0.0, T23=False):

def compute_transit_duration(

    period, rho_star, b_param, rp_by_rs, ecc=0.0, omega=0.0, T23=False

):

    """

    Compute transit duration (T14 or T23) from orbital parameters and stellar density.

@jit

def compute_eclipse_duration(period, rho_star, b_param, rp_by_rs, ecc=0.0, omega=0.0, T23=False):

def compute_eclipse_duration(

    period, rho_star, b_param, rp_by_rs, ecc=0.0, omega=0.0, T23=False

):

    """

    Compute secondary eclipse duration (total T14 or full T23).

    return -jnp.arctan2(x, -z)

@jit

def getLambertPhaseCurve(alpha):

    """

    """

    return (jnp.sin(alpha) + (jnp.pi - alpha) * jnp.cos(alpha)) / jnp.pi

@jit

def getIsotropicPhaseCurve(alpha):

    """

    # Per-planet corrected positions  (loop unrolled at JIT trace time)

    corrected = [

        body.relative_position(times + ltt[i])

        for i, body in enumerate(system.bodies)

        body.relative_position(times + ltt[i]) for i, body in enumerate(system.bodies)

    ]

    return (

        jnp.stack([pos[0] for pos in corrected]),

    _f = lambda x: jnp.asarray(x, dtype=jnp.float32)

    return _compute_rossiter_mclaughlin_jit(

        jnp.asarray(times, dtype=jnp.float32),

        _f(rp_by_rs), _f(b_param), _f(ecosomega), _f(esinomega),

        _f(period), _f(rho_star), _f(v_sin_i), _f(lambda_spin),

        _f(inc_star), _f(q1), _f(q2),

        _f(rp_by_rs),

        _f(b_param),

        _f(ecosomega),

        _f(esinomega),

        _f(period),

        _f(rho_star),

        _f(v_sin_i),

        _f(lambda_spin),

        _f(inc_star),

        _f(q1),

        _f(q2),

    )

# Transit Timing Variation (TTV) lightcurve model

# ──────────────────────────────────────────────────────────────────────────────

def compute_transit_ttv(times, transit_times, duration, impact_param, radius_ratio,

                        u_limb_dark):

def compute_transit_ttv(

    times, transit_times, duration, impact_param, radius_ratio, u_limb_dark

):

    """Compute the transit light curve for a single planet using TTV parameterization.

    Unlike :func:`compute_lightcurve`, which is parameterized by a global

@jit

def compute_lightcurve(times, q_limb_dark, rho_star,

                    t0s, periods, b_params, rp_by_rs, ecosomega, esinomega,

                    amp_night, amp_day_minus_amp_night, amp_beam, amp_ellip, phase_offsets, tidal_lag, K_amp):

def compute_lightcurve(

    times,

    q_limb_dark,

    rho_star,

    t0s,

    periods,

    b_params,

    rp_by_rs,

    ecosomega,

    esinomega,

    amp_night,

    amp_day_minus_amp_night,

    amp_beam,

    amp_ellip,

    phase_offsets,

    tidal_lag,

    K_amp,

):

    """

    Compute multi-planet lightcurve with phase curves, beaming, and ellipsoidal variations.

            impact_param=b_safe[i],

            eccentricity=eccentricities[i],

            omega_peri=omegas[i],

            radial_velocity_semiamplitude=K_amp[i]

            radial_velocity_semiamplitude=K_amp[i],

        )

    # ==========================================

    # --- Component 1: Primary Transits ---

    u_limb_dark = reparametrize_LD(q_limb_dark[0], q_limb_dark[1])

    # Raw transit computed with b_safe — always finite

    lc_transit_raw = limb_dark_light_curve(system, u_limb_dark)(times)  # (n_times, n_planets)

    lc_transit_raw = limb_dark_light_curve(system, u_limb_dark)(

        times

    )  # (n_times, n_planets)

    # Zero out planets that have no transit (b >= 1 + k)

    lc_transit_raw = lc_transit_raw * transit_occurs[None, :]  # broadcast mask

    lc_transit = jnp.sum(lc_transit_raw, axis=-1)

            radius=rp_by_rs[i],

            impact_param=b_safe[i],

            eccentricity=eccentricities[i],

            omega_peri=omegas[i]

            omega_peri=omegas[i],

        )

        # Define planet surface with thermal emission amplitude

        # normalize=False is CRITICAL to preserve amplitude scaling

        planet_surface = Surface(

            amplitude=amp_day[i] * 1e-6,  # Convert ppm to dimensionless flux

            u=(),  # No limb darkening for planet emission

            normalize=False  # Preserve amplitude (default True would rescale)

            normalize=False,  # Preserve amplitude (default True would rescale)

        )

        planet_bodies_with_surfaces.append((planet_body, planet_surface))

    # Build emission system for smooth eclipse calculation

    emission_system = SurfaceSystem(

        central=star,

        central_surface=star_surface,

        bodies=planet_bodies_with_surfaces

        central=star, central_surface=star_surface, bodies=planet_bodies_with_surfaces

    )

    # Compute emission light curves with smooth secondary eclipse profiles

    # where φ = SIGNED orbital phase, δ_lag = tidal lag angle

    # We SUBTRACT the lag because the bulge lags BEHIND the planet's position

    # Factor of 2: two maxima per orbit (both quadratures)

    lc_ellip = -amp_ellip[:, None] * jnp.cos(2 * (orbital_phase_signed - tidal_lag[:, None]))

    lc_ellip = -amp_ellip[:, None] * jnp.cos(

        2 * (orbital_phase_signed - tidal_lag[:, None])

    )

    # ==========================================

    # 4. Combine and Return (Vectorized Summation)

@jit

def compute_lightcurve_supersampled(

    times, exptime_seconds,

    q_limb_dark, rho_star,

    t0s, periods, b_params, rp_by_rs, ecosomega, esinomega,

    amp_night, amp_day_minus_amp_night, amp_beam, amp_ellip,

    phase_offsets, tidal_lag, K_amp,

    times,

    exptime_seconds,

    q_limb_dark,

    rho_star,

    t0s,

    periods,

    b_params,

    rp_by_rs,

    ecosomega,

    esinomega,

    amp_night,

    amp_day_minus_amp_night,

    amp_beam,

    amp_ellip,

    phase_offsets,

    tidal_lag,

    K_amp,

):

    """

    Compute multi-planet lightcurve with **per-point** exposure-time integration.

    t_flat = t_sub.reshape(-1)  # (N * ns,)

    # Evaluate model at all sub-sample times in a single vectorised call

    (planet_offset_flat, lc_phase_flat, lc_beam_flat,

     lc_ellip_flat, lc_transit_flat, rvs_flat) = compute_lightcurve(

        t_flat, q_limb_dark, rho_star,

        t0s, periods, b_params, rp_by_rs, ecosomega, esinomega,

        amp_night, amp_day_minus_amp_night, amp_beam, amp_ellip,

        phase_offsets, tidal_lag, K_amp,

    (

        planet_offset_flat,

        lc_phase_flat,

        lc_beam_flat,

        lc_ellip_flat,

        lc_transit_flat,

        rvs_flat,

    ) = compute_lightcurve(

        t_flat,

        q_limb_dark,

        rho_star,

        t0s,

        periods,

        b_params,

        rp_by_rs,

        ecosomega,

        esinomega,

        amp_night,

        amp_day_minus_amp_night,

        amp_beam,

        amp_ellip,

        phase_offsets,

        tidal_lag,

        K_amp,

    )

    # Average over sub-samples — 1-D arrays (n_times * ns,) → (n_times,)

    t0s = jnp.array([0.0, 0.5])

    periods = jnp.array([1.5, 4.0])  # Orbital periods [days]

    b_params = jnp.array([0.1, 0.0])

    rp_by_rs = jnp.array([0.01, .1])

    rp_by_rs = jnp.array([0.01, 0.1])

    ecosomega = jnp.array([0.0, 0.0])  # between -1 and 1

    esinomega = jnp.array([0.0, 0.0])  # between -1 and 1

    amp_night = jnp.array([0, 0])  # in ppm

    # Method 1: Use stellar density (can be sampled in MCMC)

    rho_star = 1.4  # g/cm³ (solar-like)

    planet_offset1, phases1, beams1, ellips1, transits1, rvs1 = compute_lightcurve(

        t_grid, q_limb_dark, rho_star,

        t0s, periods, b_params, rp_by_rs, ecosomega, esinomega,

        amp_night, amp_day_minus_amp_night, amp_beam, amp_ellip, phase_offsets, tidal_lag, K_amp

        t_grid,

        q_limb_dark,

        rho_star,

        t0s,

        periods,

        b_params,

        rp_by_rs,

        ecosomega,

        esinomega,

        amp_night,

        amp_day_minus_amp_night,

        amp_beam,

        amp_ellip,

        phase_offsets,

        tidal_lag,

        K_amp,

    )

    flux1 = 1.0 * 0 + planet_offset1

    fig, axes = plt.subplots(6, 1, figsize=(10, 12), sharex=True)

    # 1. Total Lightcurve (with flexible stellar density)

    axes[0].plot(t_grid, flux1, 'k', lw=1.5)

    axes[0].set_title(f"Total Lightcurve (ρ_star = {rho_star:.2f} g/cm³, Dimensionless & Flexible)", fontsize=12, weight='bold')

    axes[0].plot(t_grid, flux1, "k", lw=1.5)

    axes[0].set_title(

        f"Total Lightcurve (ρ_star = {rho_star:.2f} g/cm³, Dimensionless & Flexible)",

        fontsize=12,

        weight="bold",

    )

    axes[0].set_ylabel("Relative Flux")

    # 2. Phase Curves (Thermal/Reflected) - use first model's components

    # Sum over planets to show total signal for that component

    axes[1].plot(t_grid, jnp.sum(phases1, axis=0), 'C1') 

    axes[1].plot(t_grid, jnp.sum(phases1, axis=0), "C1")

    axes[1].set_ylabel("Phase Curve\n(ppm)")

    axes[1].text(0.02, 0.85, "Note: Cuts to 0 during Secondary Eclipse", transform=axes[1].transAxes, fontsize=9)

    axes[1].text(

        0.02,

        0.85,

        "Note: Cuts to 0 during Secondary Eclipse",

        transform=axes[1].transAxes,

        fontsize=9,

    )

    # 3. Doppler Beaming

    axes[2].plot(t_grid, jnp.sum(beams1, axis=0), 'C2')

    axes[2].plot(t_grid, jnp.sum(beams1, axis=0), "C2")

    axes[2].set_ylabel("Doppler Beaming\n(ppm)")

    # 4. Ellipsoidal Variations

    axes[3].plot(t_grid, jnp.sum(ellips1, axis=0), 'C3')

    axes[3].plot(t_grid, jnp.sum(ellips1, axis=0), "C3")

    axes[3].set_ylabel("Ellipsoidal\n(ppm)")

    # 5. Transits

    axes[4].plot(t_grid, transits1 * 1e6, 'C0')

    axes[4].plot(t_grid, transits1 * 1e6, "C0")

    axes[4].set_ylabel("Transits\n(ppm)")

    axes[4].set_xlabel("Time [days]")

    # 6. RVs

    axes[5].plot(t_grid, jnp.sum(rvs1, axis=0), 'C4')

    axes[5].plot(t_grid, jnp.sum(rvs1, axis=0), "C4")

    axes[5].set_ylabel("Radial Velocity\n(m/s)")

    axes[5].set_xlabel("Time [days]")

    # Test 1: Circular orbit, different phase offsets

    offsets_to_test = [0.0, jnp.pi / 6, jnp.pi / 4, jnp.pi / 3]

    colors = ['C0', 'C1', 'C2', 'C3']

    colors = ["C0", "C1", "C2", "C3"]

    for offset, color in zip(offsets_to_test, colors):

        _, phases_test, _, _, _, _ = compute_lightcurve(

            t_test, q_limb_dark, rho_star,

            jnp.array([0.0]), jnp.array([4.0]), jnp.array([0.1]), jnp.array([0.1]),

            jnp.array([0.0]), jnp.array([0.0]),  # Circular

            jnp.array([100.0]), jnp.array([800.0]), jnp.array([0.0]), jnp.array([0.0]),

            jnp.array([offset]), jnp.array([0.0]), jnp.array([100.0])

            t_test,

            q_limb_dark,

            rho_star,

            jnp.array([0.0]),

            jnp.array([4.0]),

            jnp.array([0.1]),

            jnp.array([0.1]),

            jnp.array([0.0]),

            jnp.array([0.0]),  # Circular

            jnp.array([100.0]),

            jnp.array([800.0]),

            jnp.array([0.0]),

            jnp.array([0.0]),

            jnp.array([offset]),

            jnp.array([0.0]),

            jnp.array([100.0]),

        )

        label = f'offset = {offset*180/jnp.pi:.0f}°'

        label = f"offset = {offset*180/jnp.pi:.0f}°"

        axes2[0].plot(t_test, phases_test[0], color=color, label=label, lw=1.5)

    axes2[0].axvline(2.0, color='k', ls='--', alpha=0.5, label='Secondary Eclipse')

    axes2[0].set_ylabel('Phase Curve (ppm)')

    axes2[0].set_title('Circular Orbit: Phase Offset Comparison', fontsize=12, weight='bold')

    axes2[0].legend(loc='upper right')

    axes2[0].axvline(2.0, color="k", ls="--", alpha=0.5, label="Secondary Eclipse")

    axes2[0].set_ylabel("Phase Curve (ppm)")

    axes2[0].set_title(

        "Circular Orbit: Phase Offset Comparison", fontsize=12, weight="bold"

    )

    axes2[0].legend(loc="upper right")

    axes2[0].grid(True, alpha=0.3)