refactor: phase 1 - centralized physical constants in utils/constants.py

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 astropy.constants import 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.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,

)

from jaxoplanet.starry.core.basis import A1 as _jxp_A1

from jaxoplanet.starry.core.basis import A2_inv as _jxp_A2_inv

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.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

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³]

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

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

# Import physical constants from centralized location

from .utils.constants import G_CGS, RHO_SUN_CGS, SECONDS_PER_DAY

# Speed of light in R_sun per day.

# Use astropy constants for the speed of light and solar radius so that

    # BCOO internally, but ONLY for BCOO.todense() (scatter-to-dense), which

    # DOES support vmap (it's a scatter with static indices and dynamic values).

    def _rm_per_step(xi, yi, zi):

    def _rm_per_step(xi: jax.Array, yi: jax.Array, zi: jax.Array) -> jax.Array:

        """Compute RM RV [R_sun/day] for a single time step.  All ops are dense."""

        # Angular position of the planet on the stellar disk (azimuth in sky plane).

        # This determines which part of the rotating stellar surface is occulted.

) -> jax.Array:

    """Public wrapper: converts scalar inputs to JAX float32 arrays, then calls the JIT core."""

    def _f(x):

    def _f(x: Any) -> jax.Array:

        return jnp.asarray(x, dtype=jnp.float32)

    return _compute_rossiter_mclaughlin_jit(

import matplotlib.pyplot as plt

import numpy as np

from astropy.constants import G, M_earth, M_jup, M_sun, R_earth, R_jup, R_sun

from matplotlib.backends.backend_pdf import PdfPages

from matplotlib.patches import Patch, Rectangle

    reparametrize_ecc_omega,

    reparametrize_LD,

)

# Constants for stellar physics (from astropy)

M_SUN_CGS = float(M_sun.cgs.value)

R_SUN_CGS = float(R_sun.cgs.value)

G_CGS = float(G.cgs.value)

RHO_SUN_CGS = M_SUN_CGS / (4.0 / 3.0 * np.pi * R_SUN_CGS**3)  # [g/cm^3]

# Planetary radius constants (CGS)

R_JUP_CGS = float(R_jup.cgs.value)

R_EARTH_CGS = float(R_earth.cgs.value)

# Planetary mass constants (CGS)

M_JUP_CGS = float(M_jup.cgs.value)

M_EARTH_CGS = float(M_earth.cgs.value)

# Conversion factors

RSUN_TO_RJUP = R_SUN_CGS / R_JUP_CGS

RSUN_TO_REARTH = R_SUN_CGS / R_EARTH_CGS

MSUN_TO_MJUP = M_SUN_CGS / M_JUP_CGS

MSUN_TO_MEARTH = M_SUN_CGS / M_EARTH_CGS

MJUP_TO_MEARTH = M_JUP_CGS / M_EARTH_CGS

from .utils.constants import (

    G_CGS,

    M_EARTH_CGS,

    M_JUP_CGS,

    M_SUN_CGS,

    R_SUN_CGS,

    RHO_SUN_CGS,

    RSUN_TO_REARTH,

    RSUN_TO_RJUP,

)

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

# Passband / blackbody / stellar-spectrum helpers  (phase-curve physics)

    derived_order = []

    derived_circular_note_lines: list[str] = []

    def _get_val(name):

    def _get_val(name: str) -> tuple[Any, Any, bool]:

        if name in free_params:

            return samples_kept[name], map_params[name], True

        val = fixed_params.get(name, 0.0)

        if "radius_star" in derived_samples:

            rstar_vals = derived_samples["radius_star"]

            rstar_map = derived_map["radius_star"]

            rstar_is_free = (

            rstar_is_free_derived: bool = bool(

                (np.std(rstar_vals) > 1e-12) if np.size(rstar_vals) > 1 else False

            )

            rstar_is_free = rstar_is_free_derived

        if rstar_is_free or rp_free:

            rp_jup_vals = rstar_vals * rp_vals * RSUN_TO_RJUP

                derived_samples["mass_star"],

                derived_map["mass_star"],

            )

            mstar_is_free = (

            mstar_is_free_derived: bool = bool(

                (np.std(mstar_vals) > 1e-12) if np.size(mstar_vals) > 1 else False

            )

            mstar_is_free = mstar_is_free_derived

        has_mstar = mstar_vals is not None and mstar_map > 0

    )

    data_types = param_config.get("data_types", {})

    def _get_val(name):

    def _get_val(name: str) -> tuple[Any, Any, bool]:

        if name in free_params:

            return samples_kept[name], map_params[name], True

        val = fixed_params.get(name, 0.0)

    if not derived_samples and not derived_fixed and not physics_priors:

        return

    def _fmt(val, width=12, sig=6):

    def _fmt(val: Any, width: int = 12, sig: int = 6) -> str:

        try:

            if val is None or (isinstance(val, float) and np.isnan(val)):

                return f"{'-':>{width}}"

            logg_nominal = None

            logg_source = None

            def _scalar_from_spec(spec):

            def _scalar_from_spec(spec: Any) -> Any:

                """Return a single representative scalar from a prior spec."""

                if spec is None:

                    return None

# Utilities package for PyAETNA

# This package contains repository-wide utilities and constants

from .constants import *  # noqa: F403

"""Centralized physical constants and conversion factors for PyAETNA."""

import numpy as np

from astropy.constants import G, M_earth, M_jup, M_sun, R_earth, R_jup, R_sun

# Physical constants (CGS)

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

M_SUN_CGS = float(M_sun.cgs.value)  # Solar mass [g]

R_SUN_CGS = float(R_sun.cgs.value)  # Solar radius [cm]

M_JUP_CGS = float(M_jup.cgs.value)  # Jupiter mass [g]

R_JUP_CGS = float(R_jup.cgs.value)  # Jupiter radius [cm]

M_EARTH_CGS = float(M_earth.cgs.value)  # Earth mass [g]

R_EARTH_CGS = float(R_earth.cgs.value)  # Earth radius [cm]

# Derived constants

RHO_SUN_CGS = M_SUN_CGS / (

    4.0 / 3.0 * np.pi * R_SUN_CGS**3

)  # Solar mean density [g/cm³]

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

# Conversion factors

RSUN_TO_RJUP = R_SUN_CGS / R_JUP_CGS

RSUN_TO_REARTH = R_SUN_CGS / R_EARTH_CGS

MSUN_TO_MJUP = M_SUN_CGS / M_JUP_CGS

MSUN_TO_MEARTH = M_SUN_CGS / M_EARTH_CGS

MJUP_TO_MEARTH = M_JUP_CGS / M_EARTH_CGS