Code Overview#

spotgp models stellar photometric variability caused by rotating starspots as a Gaussian Process. The GP covariance kernel is derived analytically from a physical spot evolution model, enabling fast, gradient-based inference of stellar rotation periods, differential rotation, spot lifetimes, and inclination.

Module architecture#

Module descriptions#

Foundation#

Module	Key exports	Role
`params.py`	`resolve_hparam`, `KERNEL_HPARAM_KEYS`	Validates and normalizes raw hyperparameter dicts. Single source of truth for parameter names, envelope detection, and amplitude modes.
`banded_cholesky.py`	`banded_cholesky_compact`, `banded_solve_compact`	O(n·b) memory Cholesky factorization and solve for banded symmetric positive-definite matrices. Used internally by `GPSolver`.
`psd.py`	`compute_psd`	Lomb–Scargle PSD of unevenly sampled time series via `astropy`.

Data layer#

Module	Key exports	Role
`observations.py`	`TimeSeriesData`	Container for observed time series (x, y, yerr). Handles NaN masking, normalization, sigma clipping. Provides `compute_psd()`, `compute_acf()`, `from_lightkurve()`, and plotting methods.
`plotting.py`	`crb_corner_plot`	Corner plots for Cramér–Rao bound and posterior visualization.

Model layer#

The model layer defines the physics of spot evolution and stellar geometry. It is composed of three independent components — defined in separate modules — that are combined into a single SpotEvolutionModel.

`envelope.py` — Spot size evolution#

EnvelopeFunction is an abstract base class. Subclass it and implement tau_spot (property) and Gamma(t) to define a new spot shape. Everything else has working defaults:

Method	Default	Purpose
`Gamma(t)`	required	Normalized spot-size envelope, peak = 1
`tau_spot`	required	Characteristic timescale [days]
`R_Gamma(lag)`	FFT interpolation	Autocorrelation ∫ Γ(t) Γ(t+lag) dt
`Gamma_hat(omega)`	FFT interpolation	Fourier transform magnitude \|FT[Γ]\|(ω)
`lspot`	`0.0`	Plateau duration [days]
`param_dict`	`{}`	Free parameters exposed to `GPSolver`
`kernel_support()`	`lspot + 6·tau_spot`	Lag beyond which R_Γ ≈ 0

Use check_functions() to compare an analytic override of R_Gamma or Gamma_hat against the FFT baseline.

Built-in envelopes:

Class	Parameters	Notes
`TrapezoidSymmetricEnvelope`	`lspot`, `tau_spot`	Analytic `Gamma_hat` and `R_Gamma`
`TrapezoidAsymmetricEnvelope`	`lspot`, `tau_em`, `tau_dec`	Analytic `R_Gamma`; rise ≠ decay
`SkewedGaussianEnvelope`	`sigma_sn`, `n_sn`	Skew-normal shape
`ExponentialEnvelope`	`tau_spot`	Analytic `Gamma_hat`, `R_Gamma`

`visibility.py` — Stellar visibility function#

VisibilityFunction encodes how much flux a spot at latitude φ contributes as the star rotates under differential rotation. The contribution is expanded as a Fourier series in rotation harmonics with coefficients cₙ(inc, φ).

Parameters: peq (equatorial period), kappa (differential rotation shear), inc (stellar inclination).

Subclass	Description
`EdgeOnVisibilityFunction`	Closed-form coefficients for edge-on (I = π/2), solid-body rotation
`FullGeometryVisibilityFunction`	Exact piecewise projected area (Eq. 5) without the small-spot approximation; coefficients computed numerically via DFT

`latitude.py` — Latitude distribution#

LatitudeDistributionFunction defines the probability density p(φ) over stellar latitude. The default is uniform over [−π/2, π/2]. Subclass and override __call__ and/or lat_range to define a custom distribution.

Attribute / method	Purpose
`lat_range`	`(min, max)` latitude bounds used for spot placement and kernel integration
`__call__(phi)`	Unnormalized PDF evaluated at φ; normalization is handled internally

The PDF weights the latitude integral inside AnalyticKernel, and lat_range sets the uniform sampling bounds for spot placement in LightcurveModel.

`spot_model.py` — Model assembly#

SpotEvolutionModel assembles the three components:

model = SpotEvolutionModel(
    envelope=TrapezoidSymmetricEnvelope(lspot=15.0, tau_spot=5.0),
    visibility=VisibilityFunction(peq=10.0, kappa=0.3, inc=np.pi / 3),
    sigma_k=0.01,
    latitude_distribution=GaussianLatitude(sigma_deg=20.0),  # optional
)

param_keys exposes the full ordered parameter vector (peq, kappa, inc, <envelope params>, sigma_k) used by GPSolver.

Kernel layer#

`analytic_kernel.py` — `AnalyticKernel`#

Computes the GP covariance kernel by integrating the per-latitude kernel contributions over the latitude distribution:

\[K(\tau) = \sigma_k^2 \, R_\Gamma(\tau) \int p(\phi)\, \left[ c_0^2(\phi) + 2\sum_{n=1}^{N} c_n^2(\phi) \cos(n\,\omega_0(\phi)\,\tau) \right] d\phi\]

where Rᵧ(τ) is the autocorrelation of the spot envelope and cₙ are the visibility Fourier coefficients.

The latitude integral uses jax.lax.scan for O(M) memory regardless of the number of latitude points. Call build_jax() once after construction to pre-compile the XLA kernels.

Key parameters: n_harmonics (default 3), n_lat (default 64), quadrature ("trapezoid" or "gauss-legendre").

`numerical_kernel.py` — `NumericalKernel`#

Estimates the kernel empirically by simulating many lightcurves with LightcurveModel and averaging their autocovariance. Used for benchmarking and validating the analytic kernel against Monte Carlo simulations.

Simulation#

`lightcurve.py` — `LightcurveModel`#

Simulates a stellar lightcurve by summing the flux deficit of nspot independently evolving spots. Each spot is placed at a random longitude, latitude (drawn uniformly within lat_range), and emergence time; its angular size follows the envelope Gamma(t).

Spot positions rotate at the latitude-dependent rate ω₀(φ) = 2π(1 − κ sin²φ)/Peq.

lc = LightcurveModel.from_spot_model(model, nspot=30, tsim=150, tsamp=0.5)

Includes plot_lightcurve() and animate_lightcurve() for visualization.

Inference layer#

`gp_solver.py` — `GPSolver`#

Builds the GP covariance matrix from AnalyticKernel, factorises it via Cholesky (full or banded), and evaluates the marginal log-posterior. Two functions are JIT-compiled: log_posterior (the log-density handed to samplers) and value_and_grad_log_posterior (used by fit_map and the gradient accessors); neg_log_posterior, grad_log_posterior, and grad_neg_log_posterior are free Python wrappers around them.

Call build_jax() once before fitting to pre-compile both. The banded Cholesky solver (default) uses the kernel support to determine bandwidth and achieves O(n·b) memory and O(n·b²) time. On uniformly sampled data the covariance is Toeplitz, and the kernel is automatically evaluated only once per distinct lag (b+1 values banded, N values full) instead of once per matrix entry. Uniform-cadence data with gaps (the common Kepler/TESS case) gets the same treatment via integer cadence offsets: distinct lags are evaluated once and gathered through a precomputed index table. Genuinely irregular sampling falls back to the general per-entry evaluation.

The bandwidth is derived from the prior upper bounds of the envelope parameters divided by the cadence, so wide priors on lspot/tau_spot at fine cadence can push b toward N — a warning is emitted when b ≥ N/2, since the banded solver then has no advantage over dense Cholesky; tighten the bounds, downsample, or pass matrix_solver="cholesky_full".

Multi-start fits (fit_map(nopt=N), fit_map_parallel, fit_acf_parallel) with method="L-BFGS-B" and batch=True run all restarts as a single vmapped jaxopt.LBFGSB program instead of one scipy optimizer per thread (requires the optional jaxopt dependency; falls back to the thread pool otherwise). The compiled program is cached on the solver: the first call pays a large one-off XLA compilation, so this pays off when the same solver configuration is fit repeatedly in a session or on GPU — for a single fit the default thread pool is usually faster end-to-end.

data = TimeSeriesData(t, flux, flux_err)
gp = GPSolver(data, model, bounds=bounds).build_jax()
theta_map, result = gp.fit_map(nopt=5)

Key methods:

Method	Purpose
`fit_map(nopt=N)`	L-BFGS-B MAP optimization, N random restarts
`predict()`	GP posterior mean and variance at new times
`plot_prediction()`	Posterior mean ± σ bands over the data
`plot_acf()`	Empirical ACF vs analytic kernel
`plot_psd()`	Lomb–Scargle PSD vs analytic PSD
`plot_covariance_matrix()`	Banded covariance matrix with sparsity annotation
`mass_matrix_hessian_map()`	Hessian-based posterior covariance at MAP

`mcmc.py` — `MCMCSampler` / `BlackJAXSampler`#

Wraps GPSolver with MCMC sampling. BlackJAXSampler uses the BlackJAX NUTS sampler with gradient information from grad_log_posterior. Provides posterior summaries, corner plots, convergence diagnostics, and checkpointing.

Data flow: fitting a lightcurve#

Observed flux (t, y, yerr)
         │
         ▼
   TimeSeriesData               ← normalize, sigma_clip
   (observations.py)            ← compute_psd, compute_acf
         │
         ▼
   SpotEvolutionModel           ← EnvelopeFunction  (envelope.py)
   (spot_model.py)              ← VisibilityFunction (visibility.py)
                                ← LatitudeDistributionFunction (latitude.py)
         │
         ▼
   AnalyticKernel.kernel(lag)
         │
         ▼
   GPSolver(data, model)
         │
         ├─ fit_map()           → theta_MAP (point estimate)
         │
         └─ BlackJAXSampler     → posterior samples

Extending the library#

Four extension points allow custom physics without modifying core code:

Extension point	Base class	Minimum required
Custom spot shape	`EnvelopeFunction`	`tau_spot`, `Gamma(t)`
Custom visibility geometry	`VisibilityFunction`	`cn_squared(phi, n_harmonics)`
Custom latitude distribution	`LatitudeDistributionFunction`	`__call__(phi)`
Custom amplitude parameterization	pass `sigma_k` directly	—

See the tutorials for worked examples:

Code Overview

Contents

Code Overview#

Module architecture#

Module descriptions#

Foundation#

Data layer#

Model layer#

envelope.py — Spot size evolution#

visibility.py — Stellar visibility function#

latitude.py — Latitude distribution#

spot_model.py — Model assembly#