Enhance type annotations and refactor code for improved clarity and maintainability (9fdb1198) · Commits · Gaetano Scandariato / PyAETNA

PyAETNA/body_models.py

+124 −80

Changes for PyAETNA/body_models.py: 124 added lines, 80 removed lines.
Original line number	Diff line number	Diff line
		from __future__ import annotations

		from typing import Any

		import jax
		import jax.numpy as jnp
		import matplotlib.pyplot as plt
		@@ -116,7 +120,9 @@ _RM_rT_TOTAL = np.array(_jxp_rT(_RM_TOTAL_DEG), dtype=np.float32) # (16,)


		@jit
		def reparametrize_LD(q1, q2):
		def reparametrize_LD(
		q1: float \| jax.Array, q2: float \| jax.Array
		) -> tuple[jax.Array, jax.Array]:
		"""Convert (q1, q2) to (u1, u2) limb darkening coefficients. JIT-compiled."""
		u1 = 2.0 * jnp.sqrt(q1) * q2
		u2 = jnp.sqrt(q1) * (1 - 2 * q2)
		@@ -124,7 +130,9 @@ def reparametrize_LD(q1, q2):


		@jit
		def reparametrize_LD_inverse(u1, u2):
		def reparametrize_LD_inverse(
		u1: float \| jax.Array, u2: float \| jax.Array
		) -> tuple[float \| jax.Array, float \| jax.Array]:
		"""Convert (u1, u2) to (q1, q2) limb darkening parameters. JIT-compiled."""
		q1 = jnp.square(u1 + u2)
		q2 = 0.5 * u1 / (u1 + u2)
		@@ -132,7 +140,9 @@ def reparametrize_LD_inverse(u1, u2):


		@jit
		def reparametrize_ecc_omega(ec, es):
		def reparametrize_ecc_omega(
		ec: float \| jax.Array, es: float \| jax.Array
		) -> tuple[float \| jax.Array, float \| jax.Array]:
		"""
		Convert Cartesian eccentricity to (e, ω).

		@@ -149,7 +159,9 @@ def reparametrize_ecc_omega(ec, es):


		@jit
		def compute_semimajor_axis_scaled(rho_star, period):
		def compute_semimajor_axis_scaled(
		rho_star: float \| jax.Array, period: float \| jax.Array
		) -> jax.Array:
		"""
		Compute scaled semi-major axis (a/R_star) from stellar density and orbital period.

		@@ -178,7 +190,12 @@ def compute_semimajor_axis_scaled(rho_star, period):


		@jit
		def compute_inclination(b_param, a_over_rstar, ecc=0.0, omega=0.0):
		def compute_inclination(
		b_param: float \| jax.Array,
		a_over_rstar: float \| jax.Array,
		ecc: float \| jax.Array = 0.0,
		omega: float \| jax.Array = 0.0,
		) -> jax.Array:
		"""
		Compute orbital inclination from transit impact parameter.

		@@ -215,8 +232,14 @@ def compute_inclination(b_param, a_over_rstar, ecc=0.0, omega=0.0):

		@jit
		def compute_transit_duration(
		period, rho_star, b_param, rp_by_rs, ecc=0.0, omega=0.0, T23=False
		):
		period: float \| jax.Array,
		rho_star: float \| jax.Array,
		b_param: float \| jax.Array,
		rp_by_rs: float \| jax.Array,
		ecc: float \| jax.Array = 0.0,
		omega: float \| jax.Array = 0.0,
		T23: bool = False,
		) -> jax.Array:
		"""
		Compute transit duration (T14 or T23) from orbital parameters and stellar density.

		@@ -276,8 +299,14 @@ def compute_transit_duration(

		@jit
		def compute_eclipse_duration(
		period, rho_star, b_param, rp_by_rs, ecc=0.0, omega=0.0, T23=False
		):
		period: float \| jax.Array,
		rho_star: float \| jax.Array,
		b_param: float \| jax.Array,
		rp_by_rs: float \| jax.Array,
		ecc: float \| jax.Array = 0.0,
		omega: float \| jax.Array = 0.0,
		T23: bool = False,
		) -> jax.Array:
		"""
		Compute secondary eclipse duration (total T14 or full T23).

		@@ -336,7 +365,9 @@ def compute_eclipse_duration(


		@jit
		def compute_eclipse_phase(ecc, omega):
		def compute_eclipse_phase(
		ecc: float \| jax.Array, omega: float \| jax.Array
		) -> jax.Array:
		"""
		Compute orbital phase of secondary eclipse for eccentric orbits.

		@@ -391,7 +422,7 @@ def compute_eclipse_phase(ecc, omega):


		@jit
		def compute_phase_angle_fast(x, y, z):
		def compute_phase_angle_fast(x: jax.Array, y: jax.Array, z: jax.Array) -> jax.Array:
		"""
		Calculate phase angle for jaxoplanet coordinate system.

		@@ -439,7 +470,7 @@ def compute_phase_angle_fast(x, y, z):


		@jit
		def compute_orbital_phase_signed(x, y, z):
		def compute_orbital_phase_signed(x: jax.Array, y: jax.Array, z: jax.Array) -> jax.Array:
		"""
		Calculate SIGNED orbital phase for phase curve modeling with hotspot offsets.

		@@ -495,7 +526,7 @@ def compute_orbital_phase_signed(x, y, z):


		@jit
		def getLambertPhaseCurve(alpha):
		def getLambertPhaseCurve(alpha: jax.Array) -> jax.Array:
		"""
		Exact phase function for a reflective Lambertian sphere.

		@@ -517,7 +548,7 @@ def getLambertPhaseCurve(alpha):


		@jit
		def getIsotropicPhaseCurve(alpha):
		def getIsotropicPhaseCurve(alpha: jax.Array) -> jax.Array:
		"""
		Phase function for isotropic scattering.

		@@ -538,7 +569,9 @@ def getIsotropicPhaseCurve(alpha):
		return 0.5 * (1 + jnp.cos(alpha))


		def _compute_ltt_positions(system, times, z_obs):
		def _compute_ltt_positions(
		system: Any, times: jax.Array, z_obs: jax.Array
		) -> tuple[jax.Array, jax.Array, jax.Array]:
		"""Re-evaluate planet positions at light-travel-time corrected epochs.

		Photons observed at t_obs were emitted at t_em = t_obs + z/c where
		@@ -578,7 +611,14 @@ def _compute_ltt_positions(system, times, z_obs):


		@jit
		def compute_radial_velocity(times, t0s, periods, ecosomega, esinomega, K_rv):
		def compute_radial_velocity(
		times: jax.Array,
		t0s: jax.Array,
		periods: jax.Array,
		ecosomega: jax.Array,
		esinomega: jax.Array,
		K_rv: jax.Array,
		) -> tuple[jax.Array, jax.Array]:
		"""
		Compute Keplerian radial velocity model for multiple planets.

		@@ -641,19 +681,19 @@ def compute_radial_velocity(times, t0s, periods, ecosomega, esinomega, K_rv):

		@jax.jit
		def _compute_rossiter_mclaughlin_jit(
		times,
		rp_by_rs,
		b_param,
		ecosomega,
		esinomega,
		period,
		rho_star,
		v_sin_i,
		lambda_spin,
		inc_star,
		q1,
		q2,
		):
		times: jax.Array,
		rp_by_rs: jax.Array,
		b_param: jax.Array,
		ecosomega: jax.Array,
		esinomega: jax.Array,
		period: jax.Array,
		rho_star: jax.Array,
		v_sin_i: jax.Array,
		lambda_spin: jax.Array,
		inc_star: jax.Array,
		q1: jax.Array,
		q2: jax.Array,
		) -> jax.Array:
		"""JIT-compiled core of compute_rossiter_mclaughlin. All args must be JAX arrays."""
		"""
		Compute the Rossiter-McLaughlin (RM) radial-velocity anomaly for a single planet.
		@@ -939,19 +979,19 @@ def _compute_rossiter_mclaughlin_jit(


		def compute_rossiter_mclaughlin(
		times,
		rp_by_rs,
		b_param,
		ecosomega,
		esinomega,
		period,
		rho_star,
		v_sin_i,
		lambda_spin,
		inc_star,
		q1=0.5,
		q2=0.1,
		):
		times: Any,
		rp_by_rs: Any,
		b_param: Any,
		ecosomega: Any,
		esinomega: Any,
		period: Any,
		rho_star: Any,
		v_sin_i: Any,
		lambda_spin: Any,
		inc_star: Any,
		q1: Any = 0.5,
		q2: Any = 0.1,
		) -> jax.Array:
		"""Public wrapper: converts scalar inputs to JAX float32 arrays, then calls the JIT core."""

		def _f(x):
		@@ -979,8 +1019,13 @@ def compute_rossiter_mclaughlin(


		def compute_transit_ttv(
		times, transit_times, duration, impact_param, radius_ratio, u_limb_dark
		):
		times: jax.Array,
		transit_times: jax.Array,
		duration: jax.Array,
		impact_param: jax.Array,
		radius_ratio: jax.Array,
		u_limb_dark: jax.Array,
		) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
		"""Compute the transit light curve for a single planet using TTV parameterization.

		Unlike :func:`compute_lightcurve`, which is parameterized by a global
		@@ -1124,23 +1169,23 @@ def compute_transit_ttv(

		@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,
		):
		times: jax.Array,
		q_limb_dark: jax.Array,
		rho_star: float \| jax.Array,
		t0s: jax.Array,
		periods: jax.Array,
		b_params: jax.Array,
		rp_by_rs: jax.Array,
		ecosomega: jax.Array,
		esinomega: jax.Array,
		amp_night: jax.Array,
		amp_day_minus_amp_night: jax.Array,
		amp_beam: jax.Array,
		amp_ellip: jax.Array,
		phase_offsets: jax.Array,
		tidal_lag: jax.Array,
		K_amp: jax.Array,
		) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]:
		"""
		Compute multi-planet lightcurve with phase curves, beaming, and ellipsoidal variations.

		@@ -1518,24 +1563,24 @@ SUPERSAMPLE_NUM = 7

		@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: jax.Array,
		exptime_seconds: jax.Array,
		q_limb_dark: jax.Array,
		rho_star: float \| jax.Array,
		t0s: jax.Array,
		periods: jax.Array,
		b_params: jax.Array,
		rp_by_rs: jax.Array,
		ecosomega: jax.Array,
		esinomega: jax.Array,
		amp_night: jax.Array,
		amp_day_minus_amp_night: jax.Array,
		amp_beam: jax.Array,
		amp_ellip: jax.Array,
		phase_offsets: jax.Array,
		tidal_lag: jax.Array,
		K_amp: jax.Array,
		) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]:
		"""
		Compute multi-planet lightcurve with per-point exposure-time integration.

		@@ -1645,7 +1690,6 @@ def compute_lightcurve_supersampled(


		if __name__ == "__main__":

		q_limb_dark = jnp.array([0.1, 0.3]) # Limb darkening
		t0s = jnp.array([0.0, 0.5])
		periods = jnp.array([1.5, 4.0]) # Orbital periods [days]

PyAETNA/gp.py

+38 −8

Changes for PyAETNA/gp.py: 38 added lines, 8 removed lines.
Original line number	Diff line number	Diff line
		@@ -48,6 +48,11 @@ with moderate Q) the benefits are:
		See ``docs/gp_per_visit_analysis/GP_PER_VISIT_ANALYSIS.md`` for benchmarks.
		"""

		from __future__ import annotations

		from collections.abc import Callable
		from typing import Any

		import jax
		import jax.numpy as jnp
		import numpy as np
		@@ -77,7 +82,15 @@ __all__ = [
		# =============================================================================


		def gp_predict(kernel, x, y, dy_sq, x_test, log_func=None, log_target=None):
		def gp_predict(
		kernel: Any,
		x: Any,
		y: Any,
		dy_sq: Any,
		x_test: Any,
		log_func: Callable[..., None] \| None = None,
		log_target: Any = None,
		) -> tuple[jax.Array, jax.Array]:
		"""
		GP prediction wrapper. Handles shape validation and logging before
		calling the JIT-compiled core.
		@@ -121,7 +134,9 @@ def gp_predict(kernel, x, y, dy_sq, x_test, log_func=None, log_target=None):


		@jax.jit
		def _gp_predict_jit(kernel, x, y, dy_sq, x_test):
		def _gp_predict_jit(
		kernel: Any, x: jax.Array, y: jax.Array, dy_sq: jax.Array, x_test: jax.Array
		) -> tuple[jax.Array, jax.Array]:
		"""Core JIT-compiled GP prediction logic."""
		gp = GaussianProcess(kernel, x, diag=dy_sq)
		gp_cond = gp.condition(y, x_test)
		@@ -133,7 +148,9 @@ def _gp_predict_jit(kernel, x, y, dy_sq, x_test):
		# =============================================================================


		def get_gp_x_variable(plData, gp_cfg, inst_mask, use_jax=True):
		def get_gp_x_variable(
		plData: Any, gp_cfg: dict[str, Any], inst_mask: Any, use_jax: bool = True
		) -> jax.Array \| np.ndarray:
		"""
		Extract GP input variable(s) for an instrument.

		@@ -184,7 +201,13 @@ def get_gp_x_variable(plData, gp_cfg, inst_mask, use_jax=True):
		# =============================================================================


		def gp_log_probability_per_visit(kernel, x_inst, res_inst, diag_inst, visit_indices):
		def gp_log_probability_per_visit(
		kernel: Any,
		x_inst: jax.Array,
		res_inst: jax.Array,
		diag_inst: jax.Array,
		visit_indices: list[Any],
		) -> jax.Array:
		"""
		Compute the GP marginal log-probability using a per-visit block-diagonal
		factorisation. Hyperparameters are shared across all visits.
		@@ -304,8 +327,13 @@ def gp_log_probability_per_visit(kernel, x_inst, res_inst, diag_inst, visit_indi


		def gp_predict_per_visit(
		kernel, x_inst, res_inst, diag_inst, visit_indices, x_pred_inst=None
		):
		kernel: Any,
		x_inst: Any,
		res_inst: Any,
		diag_inst: Any,
		visit_indices: list[Any],
		x_pred_inst: Any = None,
		) -> tuple[jax.Array, jax.Array]:
		"""
		Compute GP posterior predictions using per-visit block-diagonal
		conditioning, with shared kernel hyperparameters.
		@@ -362,7 +390,9 @@ def gp_predict_per_visit(
		# =============================================================================


		def compute_kernel_theoretical_psd(kernel, freqs, n_acf_lags=8192):
		def compute_kernel_theoretical_psd(
		kernel: Any, freqs: np.ndarray, n_acf_lags: int = 8192
		) -> np.ndarray:
		"""
		Compute the one-sided power spectral density (PSD) of a tinygp kernel
		at a set of requested cyclic frequencies via numerical FFT of the
		@@ -470,7 +500,7 @@ def compute_kernel_theoretical_psd(kernel, freqs, n_acf_lags=8192):
		# =============================================================================


		def compute_windowed_psd(kernel, freqs, t):
		def compute_windowed_psd(kernel: Any, freqs: Any, t: Any) -> jax.Array:
		"""
		Compute the discrete one-sided PSD approximation of a tinygp kernel,
		fully JAX-compatible and differentiable.

PyAETNA/kernels/factory.py

+39 −23

Changes for PyAETNA/kernels/factory.py: 39 added lines, 23 removed lines.
Original line number	Diff line number	Diff line
		@@ -7,8 +7,11 @@ This module contains the low-level (`build_single_kernel`) and high-level
		throughout PyAETNA's GP sampling and reporting code (see `PyAETNA/gp.py`).
		"""

		from __future__ import annotations

		import re
		from functools import reduce
		from typing import Any

		import jax.numpy as jnp
		import numpy as np
		@@ -36,7 +39,12 @@ SUPPORTED_KERNELS = (
		# =============================================================================


		def build_single_kernel(kernel_type, param_values, subspace_idx=None, use_jax=True):
		def build_single_kernel(
		kernel_type: str,
		param_values: list[Any],
		subspace_idx: int \| None = None,
		use_jax: bool = True,
		) -> Any:
		"""
		Build a single tinygp kernel term from parameter values.

		@@ -310,8 +318,12 @@ def build_single_kernel(kernel_type, param_values, subspace_idx=None, use_jax=Tr


		def build_gp_kernel_from_specs(
		params, kernel_specs, inst_suffix=None, use_jax=True, logger=None
		):
		params: dict[str, Any],
		kernel_specs: list[dict[str, Any]],
		inst_suffix: str \| None = None,
		use_jax: bool = True,
		logger: Any = None,
		) -> Any:
		"""
		Build a composite tinygp kernel from specifications and parameter dictionary.

		@@ -466,18 +478,18 @@ def build_gp_kernel_from_specs(


		def _build_spectral_mixture_params(
		params,
		spec,
		inst_suffix,
		section_name,
		per_param_scopes,
		default_scope,
		use_jax=True,
		logger=None,
		):
		params: dict[str, Any],
		spec: dict[str, Any],
		inst_suffix: str \| None,
		section_name: str,
		per_param_scopes: dict[str, str],
		default_scope: str,
		use_jax: bool = True,
		logger: Any = None,
		) -> list[Any]:
		"""Build parameter list for SpectralMixture kernel with proper handling of component structure."""
		xp = jnp if use_jax else np
		param_values = []
		param_values: list[Any] = []

		# Determine parameters and their suffixes
		param_names = spec.get("params", [])
		@@ -632,14 +644,14 @@ def _build_spectral_mixture_params(


		def _count_spectral_mixture_components(
		params,
		spec,
		inst_suffix,
		section_name,
		per_param_scopes,
		default_scope,
		logger=None,
		):
		params: dict[str, Any],
		spec: dict[str, Any],
		inst_suffix: str \| None,
		section_name: str,
		per_param_scopes: dict[str, str],
		default_scope: str,
		logger: Any = None,
		) -> int:
		"""
		Count the number of components for a SpectralMixture kernel based on parameter expansion.

		@@ -769,7 +781,9 @@ def _count_spectral_mixture_components(
		return component_count


		def build_gp_kernel_for_plotting(params, kernel_specs, inst_label):
		def build_gp_kernel_for_plotting(
		params: dict[str, Any], kernel_specs: list[dict[str, Any]], inst_label: Any
		) -> Any:
		"""
		Convenience wrapper for plotting: builds kernel using numpy (not JAX).

		@@ -786,7 +800,9 @@ def build_gp_kernel_for_plotting(params, kernel_specs, inst_label):
		)


		def build_gp_kernel_for_sampling(params, kernel_specs, inst_label):
		def build_gp_kernel_for_sampling(
		params: dict[str, Any], kernel_specs: list[dict[str, Any]], inst_label: Any
		) -> Any:
		"""
		Convenience wrapper for sampling: builds kernel using JAX.

PyAETNA/logger.py

+7 −4

Changes for PyAETNA/logger.py: 7 added lines, 4 removed lines.
Original line number	Diff line number	Diff line
		@@ -5,13 +5,16 @@ Provides a unified logger `_log` supporting both implicit module-level log file
		destinations (configured via `set_log_file`) and explicit model/file target parameters.
		"""

		from __future__ import annotations

		import time as ttt
		from typing import Any

		_LOG_FILE = None
		_LOG_FILE: str \| None = None
		_EIV_TRIGGERED = False


		def set_log_file(path):
		def set_log_file(path: str \| None) -> None:
		"""Set the module-wide default log file path.

		All subsequent calls to `_log` without an explicit target will write to `path`.
		@@ -21,12 +24,12 @@ def set_log_file(path):
		_LOG_FILE = path


		def get_log_file():
		def get_log_file() -> str \| None:
		"""Return the current default log file path."""
		return _LOG_FILE


		def _log(target_or_msg, msg=None, level=None):
		def _log(target_or_msg: Any, msg: str \| None = None, level: str \| None = None) -> None:
		"""Write a timestamped message to a log file or stdout.

		Signatures supported:

PyAETNA/model_utils.py

+32 −11

Changes for PyAETNA/model_utils.py: 32 added lines, 11 removed lines.
Original line number	Diff line number	Diff line
		@@ -18,6 +18,11 @@ Where:
		GP = Gaussian Process correlated-noise model
		"""

		from __future__ import annotations

		from collections.abc import Sequence
		from typing import Any

		import jax.numpy as jnp

		# ---------------------------------------------------------------------------
		@@ -26,14 +31,14 @@ import jax.numpy as jnp


		def build_mean_function(
		norm_model,
		planet_model,
		trend_model,
		normalization_mode,
		instrument,
		uniqueInstruments,
		data_types,
		):
		norm_model: Any,
		planet_model: Any,
		trend_model: Any,
		normalization_mode: str,
		instrument: Any,
		uniqueInstruments: Sequence[Any],
		data_types: dict[Any, str],
		) -> Any:
		"""
		Assemble the deterministic model (excluding GP) per data point.

		@@ -91,7 +96,14 @@ def build_mean_function(
		# ---------------------------------------------------------------------------


		def detrend_data(y, norm, trend, gp, normalization_mode, data_type):
		def detrend_data(
		y: Any,
		norm: Any,
		trend: Any,
		gp: Any,
		normalization_mode: str,
		data_type: str,
		) -> Any:
		"""
		Remove normalization, trend, and GP from observed data for plotting.

		@@ -130,7 +142,14 @@ def detrend_data(y, norm, trend, gp, normalization_mode, data_type):
		return (y - gp) / norm - trend


		def build_composite_for_plot(norm, planet, trend, gp, normalization_mode, data_type):
		def build_composite_for_plot(
		norm: Any,
		planet: Any,
		trend: Any,
		gp: Any,
		normalization_mode: str,
		data_type: str,
		) -> Any:
		"""
		Build the composite model line for plotting (full model = mean + GP).

		@@ -156,7 +175,9 @@ def build_composite_for_plot(norm, planet, trend, gp, normalization_mode, data_t
		return norm * (1.0 + planet + trend) + gp


		def build_partial_for_plot(norm, planet, trend, normalization_mode, data_type):
		def build_partial_for_plot(
		norm: Any, planet: Any, trend: Any, normalization_mode: str, data_type: str
		) -> Any:
		"""
		Build a partial model line (without GP) for plotting.