Source code for hwoutils.conversions

"""Unit conversion functions using centralized constants.

Pure JAX implementations — no astropy dependency. Functions are intentionally
NOT JIT-compiled so JAX can fuse them into larger computation graphs.
"""

import jax.numpy as jnp

from hwoutils import constants as const

# ---------------------------------------------------------------------------
# Flux conversions
# ---------------------------------------------------------------------------




def jy_to_photons_per_nm_per_m2(flux_jy, wavelength_nm):
    """Convert flux density from Janskys to photons/s/nm/m².

    Args:
        flux_jy: Flux density in Janskys.
        wavelength_nm: Wavelength in nanometers.

    Returns:
        Flux density in photons/s/nm/m².
    """
    return flux_jy * const.Jy / (wavelength_nm * const.h)





def photons_per_nm_per_m2_to_jy(flux_phot, wavelength_nm):
    """Convert flux density from photons/s/nm/m² to Janskys.

    Args:
        flux_phot: Flux density in photons/s/nm/m².
        wavelength_nm: Wavelength in nanometers.

    Returns:
        Flux density in Janskys.
    """
    return flux_phot * (wavelength_nm * const.h) / const.Jy





def mag_per_arcsec2_to_jy_per_arcsec2(mag_per_arcsec2):
    """Convert surface brightness from mag/arcsec² to Jy/arcsec² (AB).

    Args:
        mag_per_arcsec2: Surface brightness in magnitudes per arcsec².

    Returns:
        Surface brightness in Jy/arcsec².
    """
    # Astropy calculates `to(u.Jy)` as `10**( -0.4 * (mag + 48.600000) )`
    # and `1 Jy = 1e-23 erg/s/cm^2/Hz`
    # which comes out to 3630.7805477010028 Jy, not the nominal 3631.0 Jy.
    f0_jy = 3630.7805477010028
    return f0_jy * 10 ** (-0.4 * mag_per_arcsec2)



# ---------------------------------------------------------------------------
# Length conversions
# ---------------------------------------------------------------------------




def nm_to_um(length_nm):
    """Convert nanometers to micrometers."""
    return length_nm * const.nm2um





def um_to_nm(length_um):
    """Convert micrometers to nanometers."""
    return length_um * const.um2nm





def nm_to_wavenumber_cm(wavelength_nm):
    """Convert wavelength [nm] to wavenumber [cm^-1] (nu = 1 / lambda).

    Args:
        wavelength_nm: Wavelength in nanometers.

    Returns:
        Wavenumber in inverse centimeters.
    """
    return const.cm2nm / wavelength_nm





def wavenumber_cm_to_nm(wavenumber_cm):
    """Convert wavenumber [cm^-1] to wavelength [nm].

    Args:
        wavenumber_cm: Wavenumber in inverse centimeters.

    Returns:
        Wavelength in nanometers.
    """
    return const.cm2nm / wavenumber_cm





def au_to_m(length_au):
    """Convert AU to meters."""
    return length_au * const.AU2m





def m_to_au(length_m):
    """Convert meters to AU."""
    return length_m * const.m2AU





def Rearth_to_m(length_Rearth):
    """Convert Earth radii to meters."""
    return length_Rearth * const.Rearth2m



# ---------------------------------------------------------------------------
# Velocity conversions
# ---------------------------------------------------------------------------




def au_per_yr_to_m_per_s(velocity_au_per_yr):
    """Convert AU/yr to m/s."""
    return velocity_au_per_yr * const.AU2m / const.yr2s



# ---------------------------------------------------------------------------
# Angular conversions
# ---------------------------------------------------------------------------




def arcsec_to_rad(angle_arcsec):
    """Convert arcseconds to radians."""
    return angle_arcsec * const.arcsec2rad





def rad_to_arcsec(angle_rad):
    """Convert radians to arcseconds."""
    return angle_rad * const.rad2arcsec





def mas_to_arcsec(angle_mas):
    """Convert milliarcseconds to arcseconds."""
    return angle_mas * const.mas2arcsec





def arcsec_to_mas(angle_arcsec):
    """Convert arcseconds to milliarcseconds."""
    return angle_arcsec * const.arcsec2mas





def arcsec_to_lambda_d(angle_arcsec, wavelength_nm, diameter_m):
    """Convert angular separation to lambda/D units.

    Args:
        angle_arcsec: Angular separation in arcseconds.
        wavelength_nm: Wavelength in nanometers.
        diameter_m: Telescope diameter in meters.

    Returns:
        Angular separation in lambda/D.
    """
    angle_rad = angle_arcsec * const.arcsec2rad
    wavelength_m = wavelength_nm * const.nm2m
    lambda_d_rad = wavelength_m / diameter_m
    return angle_rad / lambda_d_rad





def lambda_d_to_arcsec(angle_lambda_d, wavelength_nm, diameter_m):
    """Convert lambda/D units to angular separation in arcseconds.

    Args:
        angle_lambda_d: Angular separation in lambda/D.
        wavelength_nm: Wavelength in nanometers.
        diameter_m: Telescope diameter in meters.

    Returns:
        Angular separation in arcseconds.
    """
    wavelength_m = wavelength_nm * const.nm2m
    lambda_d_rad = wavelength_m / diameter_m
    angle_rad = angle_lambda_d * lambda_d_rad
    return angle_rad * const.rad2arcsec



# ---------------------------------------------------------------------------
# Mass conversions
# ---------------------------------------------------------------------------




def Msun_to_kg(mass_solar):
    """Convert solar masses to kilograms."""
    return mass_solar * const.Msun2kg





def Mearth_to_kg(mass_earth):
    """Convert Earth masses to kilograms."""
    return mass_earth * const.Mearth2kg



# ---------------------------------------------------------------------------
# Distance conversions
# ---------------------------------------------------------------------------




def au_to_arcsec(distance_au, distance_pc):
    """Convert physical distance in AU to angular separation.

    Args:
        distance_au: Physical distance in AU.
        distance_pc: Distance to system in parsecs.

    Returns:
        Angular separation in arcseconds.
    """
    return distance_au / distance_pc





def arcsec_to_au(angle_arcsec, distance_pc):
    """Convert angular separation to physical distance.

    Args:
        angle_arcsec: Angular separation in arcseconds.
        distance_pc: Distance to system in parsecs.

    Returns:
        Physical distance in AU.
    """
    return angle_arcsec * distance_pc



# ---------------------------------------------------------------------------
# Time conversions
# ---------------------------------------------------------------------------




def years_to_days(time_years):
    """Convert years to days."""
    return time_years * 365.25





def days_to_years(time_days):
    """Convert days to years."""
    return time_days / 365.25





def is_leap_year(year):
    """Determine if a year is a leap year.

    Args:
        year: The year to check.

    Returns:
        True if the year is a leap year.
    """
    return (year % 4 == 0) & ((year % 100 != 0) | (year % 400 == 0))





def days_in_year(year):
    """Return the number of days in a year (365 or 366).

    Args:
        year: The year to check.

    Returns:
        Number of days.
    """
    return 365 + is_leap_year(year)





def gregorian_to_jd(year, month, day):
    """Convert a Gregorian date to a Julian day.

    Args:
        year: The year.
        month: The month.
        day: The day.

    Returns:
        The Julian day.
    """
    a = jnp.floor((14 - month) / 12)
    y = year + 4800 - a
    m = month + 12 * a - 3
    jdn = (
        day
        + jnp.floor((153 * m + 2) / 5)
        + 365 * y
        + jnp.floor(y / 4)
        - jnp.floor(y / 100)
        + jnp.floor(y / 400)
        - 32045
    )
    return jdn - 0.5





def jd_to_decimal_year(jd):
    """Convert a Julian day to a decimal year.

    Args:
        jd: The Julian day.

    Returns:
        The decimal year.
    """
    year_approx = 1970.0 + (jd - 2440587.5) / 365.2425
    year = jnp.floor(year_approx)

    jd_start = gregorian_to_jd(year, 1, 1)
    jd_end = gregorian_to_jd(year + 1, 1, 1)

    year = jnp.where(jd < jd_start, year - 1, year)
    jd_start = gregorian_to_jd(year, 1, 1)
    jd_end = gregorian_to_jd(year + 1, 1, 1)

    return year + (jd - jd_start) / (jd_end - jd_start)





def decimal_year_to_jd(decimal_year):
    """Convert a decimal year to a Julian day.

    Args:
        decimal_year: The decimal year.

    Returns:
        The Julian day.
    """
    year = jnp.floor(decimal_year)
    year_fraction = decimal_year - year

    jd_start = gregorian_to_jd(year, 1, 1)
    jd_end = gregorian_to_jd(year + 1, 1, 1)

    return jd_start + year_fraction * (jd_end - jd_start)



# ---------------------------------------------------------------------------
# Albedo conventions
# ---------------------------------------------------------------------------

# Lambertian-sphere ratio of geometric to spherical (Bond) albedo.
# Seager 2010, eq 3.36: a uniformly Lambertian sphere has A_g = (2/3) A_S.
SPHERICAL_TO_GEOMETRIC_ALBEDO_LAMBERTIAN: float = 2.0 / 3.0




def spherical_to_geometric_albedo(A_S):
    """Convert spherical albedo to geometric albedo (Lambertian sphere).

    Args:
        A_S: Spherical (Bond) albedo, dimensionless.

    Returns:
        Geometric albedo A_g, dimensionless.

    Notes:
        Exact for a uniformly Lambertian-reflecting sphere
        (Seager 2010, eq 3.36): A_g = (2/3) * A_S. For atmospheres
        with anisotropic scattering (Mie clouds, etc.) the actual
        A_g / A_S ratio depends on the scattering phase function;
        this is the leading-order approximation that the ExoJaxPhysicalModel
        uses to convert its 2-stream plane-parallel reflectivity to a
        disk-integrated geometric albedo.
    """
    return SPHERICAL_TO_GEOMETRIC_ALBEDO_LAMBERTIAN * A_S





def geometric_to_spherical_albedo(A_g):
    """Convert geometric albedo to spherical albedo (Lambertian sphere).

    Inverse of :func:`spherical_to_geometric_albedo`:
    A_S = (3/2) * A_g.
    """
    return A_g / SPHERICAL_TO_GEOMETRIC_ALBEDO_LAMBERTIAN