API Reference (spotgp)

Lightcurve Model

class spotgp.lightcurve.LightcurveModel(peq=4.0, kappa=0.0, inc=1.5707963267948966, nspot=None, tau_spot=None, tem=2, tdec=2, alpha_max=0.1, fspot=0, lspot=5, long=[0, 6.283185307179586], lat=[-1.5707963267948966, 1.5707963267948966], tsim=28, tsamp=0.02, limb_darkening=False, tmax=None, rotate=True, grow=True, nspot_rate=None)

Bases: AnimationMixin

JAX-accelerated star with spots and its lightcurve.

Same interface as the numpy version but uses JAX for vectorized computation across all spots simultaneously.

Parameters:

• peq (float) – Equatorial period of the star.

• kappa (float) – Differential rotation shear.

• inc (float) – Inclination of the star.

• nspot (int) – Number of spots.

• tau_spot (float, optional) – Timescale for both emergence and decay of the spots. Defaults to None.

• tem (float, optional) – Emergence timescale of the spots. Defaults to 2.

• tdec (float, optional) – Decay timescale of the spots. Defaults to 2.

• alpha_max (float, optional) – Maximum angular area of the spots. Defaults to 0.1.

• fspot (float, optional) – Spot contrast fraction. Defaults to 0.

• lspot (float, optional) – Spot lifetime. Defaults to 5.

• long (list, optional) – Range of spot longitudes. Defaults to [0, 2*pi].

• lat (list, optional) – Range of spot latitudes. Defaults to [0, pi].

• tsim (float, optional) – End simulation time. Defaults to 28.

• tsamp (float, optional) – Sampling cadence. Defaults to 0.02.

• limb_darkening (bool, optional) – Flag to enable limb darkening. Defaults to False.

classmethod from_spot_model(spot_model: SpotEvolutionModel, nspot: int = None, *, nspot_rate: float = None, **kwargs)

Construct a LightcurveModel from a SpotEvolutionModel.

Parameters:

• spot_model (SpotEvolutionModel) – Fully configured spot evolution model.

• nspot (int, optional) – Total number of spots to simulate.

• nspot_rate (float, optional) – Spot emergence rate [spots/day]. The actual number of spots is max(1, int(nspot_rate * tsim)). Exactly one of nspot or nspot_rate must be provided.

• **kwargs – Forwarded to LightcurveModel.__init__ (e.g. tsim, tsamp, lat, long).

Return type:

LightcurveModel

classmethod from_hparam(hparam: dict, nspot: int = None, *, nspot_rate: float = None, **kwargs)

Construct a LightcurveModel from a GPSolver-compatible hparam dict.

Accepts the same raw hparam dict that GPSolver/AnalyticKernel take, including all amplitude modes (sigma_k, nspot_rate, or nspot), and both symmetric (tau) and asymmetric (tau_em + tau_dec) envelopes. This removes the need to manually decompose the dict in scripts.

Parameters:

• hparam (dict) – Raw hyperparameter dict. Must contain peq, kappa, inc, lspot, tau_spot (or tau_em/tau_dec), and an amplitude specification.

• nspot (int, optional) – Total number of spots to simulate.

• nspot_rate (float, optional) – Spot emergence rate [spots/day]. Exactly one of nspot or nspot_rate must be provided.

• **kwargs – Forwarded to LightcurveModel.__init__ (e.g. tsim, tsamp, lat, long).

Return type:

LightcurveModel

Flux(teval)

Compute the full lightcurve using JAX vmap over all spots.

Instead of a Python loop over nspot, all spots are computed in parallel via JAX’s vmap.

plot_lightcurve(show_spots=True, show_title=True)

Plot the lightcurve.

Analytic Kernel

class spotgp.analytic_kernel.AnalyticKernel(model_or_hparam, n_harmonics=3, n_lat=64, lat_range=None, quadrature='trapezoid')

Bases: object

JAX-accelerated analytic GP kernel for stellar rotation variability.

Parameters:

• model_or_hparam (SpotEvolutionModel or dict) – Either a SpotEvolutionModel instance (new API) or a raw hparam dict (backward-compatible old API).

• n_harmonics (int) – Number of Fourier harmonics for the visibility function (default 3).

• n_lat (int) – Number of latitude quadrature points (default 64).

• lat_range (tuple) – (min, max) latitude in radians (default (-pi/2, pi/2)).

• quadrature (str) – Latitude integration method: “trapezoid” or “gauss-legendre”.

omega0(phi)

Latitude-dependent rotation angular frequency [rad/day].

R_Gamma(lag)

Autocorrelation of the squared envelope (delegates to envelope).

cn_squared(phi)

Squared Fourier visibility coefficients at latitude phi.

kernel_single_latitude(lag, phi)

Single-spot kernel at a fixed latitude.

kernel(lag, lat_dist=None)

Full GP kernel averaged over latitude.

Uses jax.lax.scan for memory-efficient accumulation: only one lag-sized buffer is live at a time — O(M) instead of O(n_lat·M).

When the visibility function is an EdgeOnVisibilityFunction, the latitude-averaged |c_n|^2 are known constants and the latitude loop is bypassed entirely.

Parameters:

• lag (array_like) – Time lags [days]. Can be 1D or 2D.

• lat_dist (callable or None) – Latitude probability density. If None, uniform.

Returns:

K

Return type:

ndarray, same shape as lag input.

kernel_solid_body(lag, lat_dist=None)

Kernel for solid-body rotation (kappa=0).

compute_psd(omega, lat_dist=None)

Analytic power spectral density.

Parameters:

• omega (array_like) – Angular frequencies [rad/day].

• lat_dist (callable or None) – Latitude probability density.

Returns:

• freq (ndarray [cycles/day])

• power (ndarray)

build_jax(n_lag=256)

Pre-compile and warm up JAX JIT computation for this kernel.

jax.lax.scan (used inside kernel()) triggers XLA compilation on its first call for a given array shape. That compilation can take several seconds and is easy to mistake for slow runtime. Call build_jax() once after constructing the kernel to pay that cost upfront — subsequent calls to kernel() and compute_psd() with the same shape will be fast.

Parameters:

n_lag (int) – Length of the dummy lag array used to drive compilation (default 256). The actual value does not matter as long as it is representative of the sizes you will use at runtime.

Returns:

self – Returns self so the call can be chained: ak = AnalyticKernel(model).build_jax().

Return type:

AnalyticKernel

Numerical Kernel

class spotgp.numerical_kernel.NumericalKernel(model_or_hparam, tsim=20, tsamp=0.05, nsim=1000.0, verbose=True)

Bases: object

Gaussian Process for Stellar Rotation

Parameters:

• hparam (dict) –

  Dict of hyperparameters. Required keys: peq, kappa, inc, lspot, tau_spot, alpha_max. For the kernel amplitude, provide EITHER:

  • sigma_k : overall amplitude prefactor, OR

  • nspot + fspot : number of spots and spot contrast, from which sigma_k is computed as sqrt(N_spot) * (1 - f_spot) / pi.

  Note: nspot is always required for the numerical simulations.

• tsim (float) – Simulation time (default: 20).

• tsamp (float) – Time sampling (default: 0.05).

• nsim (int) – Number of simulations (default: 1e3).

• verbose (bool) – Whether to print verbose output (default: True).

get_acf()

compute_psd(tarr=None, freq_min=None, freq_max=None, normalization='psd', nsims=100)

plot_autocorrelation()

Plot the autocorrelation function.

Returns: - fig: matplotlib Figure object

spotgp.numerical_kernel.generate_sims(theta, nsim=1000.0, **kwargs)

Generate synthetic lightcurves for a given set of parameters.

Parameters:

• theta (tuple) – Tuple containing peq, kappa, inc, and nspot parameters.

• nsim (int, optional) – Number of simulations. Defaults to 1000.

• **kwargs – Additional arguments for the StarSpot class.

Returns:

Array of synthetic lightcurves.

Return type:

numpy.ndarray

spotgp.numerical_kernel.avg_covariance_tlag(K)

GP Solver

class spotgp.gp_solver.GPSolver(data_or_x, y=None, yerr=None, model_or_hparam=None, kernel_type='analytic', mean=None, fit_sigma_n=False, bounds=None, log_prior=None, matrix_solver='cholesky_banded', bandwidth=None, save_dir=None, **kernel_kwargs)

Bases: FittingMixin, GPPlotsMixin, MassMatrixMixin

JAX-accelerated Gaussian Process solver for stellar lightcurves.

Handles covariance matrix construction, Cholesky factorization, log-likelihood evaluation, prediction, MAP estimation, and mass matrix computation.

Parameters:

• data_or_x (TimeSeriesData or array_like, shape (N,)) – Either a TimeSeriesData object, or observation times [days]. When a TimeSeriesData is passed, y and yerr must be None (they are read from the data object).

• y (array_like, shape (N,), optional) – Observed flux values. Required when data_or_x is an array.

• yerr (array_like, shape (N,) or float, optional) – Measurement uncertainties (1-sigma). Required when data_or_x is an array.

• model_or_hparam (SpotEvolutionModel or dict) – Either a SpotEvolutionModel (new API) or a raw hparam dict (backward-compatible old API).

• kernel_type ({"analytic"}) – Which kernel to use (default: “analytic”).

• mean (float or callable or None) – Mean function.

• fit_sigma_n (bool) – If True, include white noise amplitude sigma_n as a free parameter for optimization/sampling (default False).

• bounds (dict or None) – Parameter bounds for optimization. If None, uses defaults.

• log_prior (callable or None) – Custom log-prior function f(theta_arr) -> scalar. If None, uses soft uniform within bounds.

• kernel_kwargs (dict) – Extra kwargs forwarded to the kernel constructor.

DEFAULT_BOUNDS = {'inc': (0.01, 3.1315926535897933), 'kappa': (0.001, 0.999), 'lat_max': (0.0, 1.5707963267948966), 'lat_min': (0.0, 1.5707963267948966), 'lspot': (0.1, 20.0), 'n_sn': (-10.0, 10.0), 'peq': (0.5, 50.0), 'sigma_k': (1e-06, 1.0), 'sigma_n': (1e-06, 0.1), 'sigma_sn': (0.05, 10.0), 'tau_dec': (0.05, 10.0), 'tau_em': (0.05, 10.0), 'tau_spot': (0.05, 10.0)}

build_jax(recompute=True)

Pre-compile and warm up the solver’s JAX JIT functions.

Exactly two XLA compilations are triggered: the value-only log_posterior (the log-density handed to samplers) and value_and_grad_log_posterior (used by fit_map and the gradient accessors). The remaining public functions (neg_log_posterior, grad_log_posterior, grad_neg_log_posterior) are Python wrappers around these two and need no compilation of their own.

Call build_jax() once after constructing the solver to pay the compilation cost upfront instead of inside the first fit or MCMC step.

Returns:

self – Returns self so the call can be chained: gp = GPSolver(...).build_jax().

Return type:

GPSolver

log_likelihood()

Marginal log-likelihood of the data under the GP.

Returns:

logL – log p(y | X, theta)

Return type:

float

predict(xpred, return_cov=False)

Predictive distribution at new input locations.

Parameters:

• xpred (array_like, shape (M,)) – Prediction times.

• return_cov (bool) – If True, return full predictive covariance.

Returns:

• mu_pred (ndarray, shape (M,))

• var_pred (ndarray, shape (M,) or (M, M))

sample_prior(xpred, n_samples=1, rng=None)

Draw samples from the GP prior.

sample_posterior(xpred, n_samples=1, rng=None)

Draw samples from the GP posterior.

sample_lightcurves(theta=None, xpred=None, n_samples=5, n_points=2000, source='prior', rng=None)

Sample lightcurves from the GP prior or posterior.

Parameters:

• theta (dict or array_like, optional) – Kernel parameters. Accepts a physical dict with keys from param_keys, a sampling-space dict with log_-prefixed keys, or a flat array matching param_keys. If None, uses the current internal hyperparameters.

• xpred (array_like, optional) – Times at which to evaluate the samples. If None, uses n_points evenly spaced times spanning the data baseline.

• n_samples (int) – Number of lightcurve samples to draw (default 5).

• n_points (int) – Number of prediction points when xpred is None (default 2000).

• source ({'prior', 'posterior'}) – Whether to sample from the GP prior or posterior (default ‘prior’).

• rng (numpy.random.Generator, optional) – Random number generator for reproducibility.

Returns:

• xpred (ndarray, shape (M,)) – Prediction times.

• samples (ndarray, shape (n_samples, M)) – Sampled lightcurves.

compute_acf(tlags=None, n_bins=50, normalize=True)

Compute the empirical autocorrelation function of the data.

Delegates to self.data.compute_acf().

Parameters:

• tlags (array_like, optional) – Bin edges for time lags [days]. If provided, n_bins is inferred as len(tlags) - 1. If None, n_bins linearly spaced bins from 0 to half the baseline are used.

• n_bins (int) – Number of lag bins (used when tlags is None, default 50).

• normalize (bool) – If True (default), normalize so ACF(0) ~ 1.

Returns:

• lag_centers (ndarray, shape (n_bins,)) – Bin centers.

• acf (ndarray, shape (n_bins,)) – Empirical ACF at each bin center.

compute_kernel(tlags)

Evaluate the analytic kernel at the given time lags.

Parameters:

tlags (array_like, shape (M,)) – Time lags [days].

Returns:

K – Kernel values at each lag.

Return type:

ndarray, shape (M,)

get_theta()

Return the current kernel hyperparameters as a dictionary.

Returns:

theta – Keys and values for all kernel (and optionally noise) hyperparameters, e.g. {“peq”: 5.0, “kappa”: 0.2, …}.

Return type:

dict

update_hparam(hparam)

Update hyperparameters and rebuild kernel and covariance.

Accepts a SpotEvolutionModel, an hparam dict (legacy keys like lspot), or a theta-style dict whose keys match self.spot_model.param_keys (e.g. tau_em, lat_min).

save(path)

Save solver state to an HDF5 file.

Writes data, model configuration, bounds, and any completed fit results (MAP, ACF, mass matrix). Can be called repeatedly — each call overwrites only the groups whose data changed.

Parameters:

path (str) – File path (should end in .h5 or .hdf5).

classmethod load(path)

Reconstruct a GPSolver from an HDF5 file.

The returned solver has fit results restored but is NOT JIT-compiled. Call .build_jax() before running fits or sampling.

Parameters:

path (str) – Path to an HDF5 file previously written by save().

Return type:

GPSolver

Power Spectral Density

spotgp.psd.compute_psd(y, t=None, dt=None, normalization='psd', freq_min=None, freq_max=None, n_freq=None, samples_per_peak=5)

Compute the Power Spectral Density of a time series using astropy.timeseries.LombScargle.

Works for both evenly and unevenly sampled data.

Parameters:

• y (array-like, shape (N,)) – Time series values.

• t (array-like, shape (N,), optional) – Sample times. If None, integer indices scaled by dt are used.

• dt (float, optional) – Sampling interval. Used only when t is None (default: 1).

• normalization ({"psd", "standard", "model", "log"}) – Passed directly to LombScargle.autopower / power.

• freq_min (float, optional) – Minimum frequency to evaluate.

• freq_max (float, optional) – Maximum frequency to evaluate.

• n_freq (int, optional) – Number of frequency grid points.

• samples_per_peak (float, optional) – Controls the frequency grid density (default 5).

Returns:

• freq (ndarray) – Frequencies in cycles per unit time.

• power (ndarray) – PSD evaluated at each frequency.

MCMC Sampler

class spotgp.mcmc.MCMCSampler(gp)

Bases: object

Base MCMC sampler for GP hyperparameters.

Wraps a GPSolver object and provides shared storage, diagnostics, summary statistics, corner plots, and dict conversion. Subclasses implement specific sampling algorithms (e.g. NUTS).

Parameters:

gp (GPSolver) – A configured GPSolver instance.

property param_keys

property n_params

summary()

Print summary statistics of the posterior samples.

Returns:

stats – Parameter names mapped to (mean, std, 16%, 50%, 84%).

Return type:

dict

plot_covariance(method='fisher', theta_map=None, n_sigma=2, n_grid=200, samples=None, figsize=None, color='C0', alpha=0.3, true_params=None, savefig=None, **corner_kwargs)

Corner plot of 2D covariance ellipses from the Hessian or Fisher matrix, with 1D marginal Gaussians on the diagonal.

Uses corner.corner to lay out the figure when MCMC samples are provided, and overlays the Laplace/Fisher Gaussian approximation (ellipses + 1D marginals).

Parameters:

• method ({"fisher", "hessian_map", "laplace"}) – Which matrix to use for the Gaussian approximation.

• theta_map (array_like, optional) – Center of the ellipses. If None, uses MAP estimate.

• n_sigma (float) – Number of sigma for the ellipse contours (default 2).

• n_grid (int) – Grid resolution for the ellipse curves (default 200).

• samples (array_like, optional) – If provided, plotted as the corner histogram/contours. If None, the figure is created with empty axes and only the Gaussian approximation is drawn.

• figsize (tuple, optional) – Figure size.

• color (str) – Color for Gaussian ellipses and marginals (default “C0”).

• alpha (float) – Fill alpha for the ellipse interiors (default 0.3).

• true_params (dict or array_like, optional) – True parameter values to mark with crosshairs.

• savefig (str, optional) – If provided, save figure to this path.

• **corner_kwargs – Extra keyword arguments forwarded to corner.corner (e.g. quantiles, show_titles, hist_kwargs).

Returns:

fig, axes

Return type:

matplotlib Figure and 2D array of Axes.

plot_corner_map(samples=None, checkpoint_path=None, cmap='viridis', marker_size=40, savefig=None, true_params=None, **corner_kwargs)

Corner plot of MCMC samples with MAP solutions overlaid as scatter points colored by their log-likelihood.

Parameters:

• samples (array_like, optional) – Shape (n_samples, n_params). If None, loads from the checkpoint file.

• checkpoint_path (str, optional) – Path to checkpoint .npz file containing MAP solutions. If None, uses the default checkpoint file.

• cmap (str) – Colormap for the MAP scatter points (default “viridis”).

• marker_size (float) – Marker size for scatter points (default 40).

• savefig (str, optional) – If provided, save figure to this path.

• true_params (dict or array_like, optional) – True parameter values to mark with crosshairs.

• **corner_kwargs – Extra keyword arguments forwarded to corner.corner.

Returns:

fig, axes

Return type:

matplotlib Figure and 2D array of Axes.

to_dict(samples=None)

Convert samples array to a dict keyed by parameter name.

Parameters:

samples (jnp.ndarray, optional) – Shape (n_samples, n_params). If None, uses self.samples.

Returns:

d – {param_name: array of shape (n_samples,)}

Return type:

dict

class spotgp.mcmc.BlackJAXSampler(gp, save_dir='results', checkpoint_file='mcmc_checkpoint.npz')

Bases: MCMCSampler

NUTS sampler using the BlackJAX library.

Inherits diagnostics, summary, plotting, and dict conversion from MCMCSampler. Adds run_map, run_warmup, and run_sampling for gradient-based No-U-Turn sampling with dual-averaging step-size adaptation.

When multiple chains are requested, sampling is parallelized across available devices via jax.pmap. Chains are distributed evenly across devices (n_chains must be divisible by jax.device_count()). On a single GPU this behaves identically to the previous jax.vmap implementation.

Parameters:

• gp (GPSolver) – A configured GPSolver instance.

• save_dir (str, optional) – Directory for all outputs produced by this sampler (corner plots, covariance plots, etc.). Created automatically if it does not exist. When set, save_checkpoint will default to saving the checkpoint inside this directory.

• checkpoint_file (str, optional) – Path to the checkpoint file. When provided, overrides the default save_dir/mcmc_checkpoint.npz. If neither checkpoint_file nor save_dir is given, no checkpoint file is set until one is passed to a later method.

run_map(nopt=10, keys=None, checkpoint_file=None, theta0=None, **kwargs)

Find MAP solutions via parallel multi-start optimization.

Runs GPSolver.fit_map_parallel and stores the results. If the checkpoint file already contains MAP data, loads from it instead of re-running the optimization.

Parameters:

• nopt (int) – Number of independent optimization restarts (default 10).

• keys (list of str, optional) – Parameter names to optimize. If None, uses all bounded parameters from GPSolver.

• theta0 (dict, optional) – Initial parameter guess to include as one of the optimization starting points. Replaces one random start so the total number of restarts stays nopt.

• checkpoint_file (str, optional) – Path to save/load MAP solutions. If provided, also updates the sampler’s default checkpoint path. Defaults to self._checkpoint_file.

• **kwargs – Additional keyword arguments passed to GPSolver.fit_map_parallel (e.g. method, maxiter).

Returns:

all_theta_maps – All MAP solutions sorted by objective (best first).

Return type:

list of dict

run_warmup(n_warmup=500, theta_init=None, mass_matrix_method='hessian_map', step_size=None, rng_key=None, target_accept=0.8, progress_bar=False, n_chains=1, checkpoint_file=None, warmup_method='window_adaptation', pathfinder_maxiter=100, pathfinder_maxcor=10, pathfinder_num_elbo=200)

Run warmup phase: adapt step size and mass matrix.

Supports three warmup strategies:

• "window_adaptation" (default): BlackJAX’s standard dual-averaging window adaptation of both step size and mass matrix.

• "pathfinder": multi-path Pathfinder via L-BFGS.

• "dual_averaging": fixes the mass matrix (from Hessian at MAP) and only adapts the step size.

After warmup, adapted parameters are stored on the sampler and a checkpoint is saved (if checkpoint_file is set).

Parameters:

• n_warmup (int) – Number of warmup steps (default 500).

• theta_init (dict or array_like, optional) – Initial position. If None, uses GPSolver’s MAP estimate. Can also be a list of dicts or 2-D array for per-chain starting points.

• mass_matrix_method ({"hessian_map", "fisher", "laplace", "diagonal", None}) – Method to estimate the mass matrix.

• step_size (float, optional) – Initial NUTS step size. If None, a heuristic is used.

• rng_key (jax.random.PRNGKey, optional) – Random key. Default: PRNGKey(0).

• target_accept (float) – Target acceptance rate (default 0.8).

• progress_bar (bool) – If True, show progress during window adaptation.

• n_chains (int) – Number of chains (used to validate device count and store per-chain init positions).

• checkpoint_file (str, optional) – Override the default checkpoint file path. When set, updates self._checkpoint_file for all subsequent save/load operations. Defaults to save_dir/mcmc_checkpoint.npz when save_dir is set.

• warmup_method ({"window_adaptation", "pathfinder", "dual_averaging"}) – Warmup strategy.

• pathfinder_maxiter (int) – Max L-BFGS iterations for Pathfinder (default 100).

• pathfinder_maxcor (int) – L-BFGS history size for Pathfinder (default 10).

• pathfinder_num_elbo (int) – Number of ELBO samples for Pathfinder (default 200).

run_sampling(n_samples=1000)

Run NUTS sampling using adapted parameters from run_warmup.

Must be called after run_warmup (or will use parameters restored from a checkpoint).

Parameters:

n_samples (int) – Number of post-warmup samples per chain (default 1000).

Returns:

• samples (jnp.ndarray) – Shape (n_samples, n_params) when n_chains=1, or (n_chains, n_samples, n_params) when n_chains > 1.

• info (dict) – Sampling diagnostics (arrays have a leading chain dimension when n_chains > 1).

save_checkpoint(path=None, append_samples=True, plot_corner=False)

Save sampler state to disk for later resumption.

When append_samples=True (the default), new samples are appended to any existing samples already stored in path, and self.samples is cleared from memory. This enables a sample-checkpoint-clear loop that keeps memory usage constant.

Parameters:

• path (str, optional) – File path (saved as .npz). If None, uses the checkpoint_file set in run_warmup, or save_dir/checkpoint.npz if save_dir was set.

• append_samples (bool) – If True, append current self.samples to any samples already on disk, then clear self.samples from memory. If False, overwrite with only the current in-memory samples.

• plot_corner (bool) – If True, load all samples currently on disk after saving and write a corner plot to save_dir/corner_plot.png (or alongside the checkpoint file if save_dir is not set).

load_checkpoint(checkpoint_file=None)

Restore sampler state from a checkpoint file.

Loads only the NUTS state and adapted kernel parameters needed to resume sampling. Samples stored in the file are not loaded into memory — use load_samples to read them later.

Parameters:

checkpoint_file (str, optional) – Path to a .npz checkpoint file. If provided, also updates the sampler’s default checkpoint path. If None, uses the default save_dir/mcmc_checkpoint.npz.

run_smc(n_particles=500, n_mcmc_steps=10, n_adapt_steps=25, target_ess=0.5, target_accept=0.6, rng_key=None, step_size=None, mass_matrix_method='hessian_map', theta_init=None, max_tempering_steps=200, checkpoint_every=10, checkpoint_file=None, particle_batch_size=None, max_num_doublings=10)

Run adaptive tempered Sequential Monte Carlo.

Starts from the prior and anneals toward the full posterior using an adaptive temperature schedule. At each tempering step, particles are resampled and rejuvenated with NUTS moves. The NUTS step size is re-adapted via dual averaging at each tempering stage using a representative particle.

Parameters:

• n_particles (int) – Number of SMC particles (default 500).

• n_mcmc_steps (int) – NUTS rejuvenation steps per tempering stage (default 10).

• n_adapt_steps (int) – Dual-averaging warmup steps to adapt the NUTS step size at each tempering stage (default 25).

• target_ess (float) – Target effective sample size as a fraction of n_particles (default 0.5).

• target_accept (float) – Target NUTS acceptance rate for dual averaging (default 0.6).

• rng_key (jax.random.PRNGKey, optional) – Random key. Default: PRNGKey(42).

• step_size (float, optional) – Initial NUTS step size. If None, a heuristic from the mass matrix is used.

• mass_matrix_method (str, optional) – Method to estimate the inverse mass matrix (default "hessian_map"). Set to None to use an identity matrix.

• theta_init (dict or array_like, optional) – Reference point for mass matrix estimation. If None, the MAP estimate is used.

• max_tempering_steps (int) – Safety limit on the number of tempering stages (default 200).

• checkpoint_every (int) – Save a checkpoint every this many tempering steps (default 10). Set to 0 to disable periodic checkpointing.

• checkpoint_file (str, optional) – Override the default checkpoint file path.

• particle_batch_size (int, optional) – Process particles in batches of this size to limit GPU memory usage. When multiple GPUs are visible the batches are distributed across devices via jax.pmap. n_particles must be divisible by this value (and by batch_size * n_devices for multi-GPU). If None, all particles are evaluated at once (original blackjax behavior).

• max_num_doublings (int, optional) – Maximum NUTS tree depth (default 10). Lower values (e.g. 5-6) reduce peak GPU memory per particle at the cost of shorter trajectories.

Returns:

• samples (np.ndarray, shape (n_particles, n_params)) – Weighted posterior particles at the final temperature.

• info (dict) – Diagnostics including tempering schedule and log evidence estimate.

static load_samples(path, flatten_chains=True)

Read all samples from a checkpoint file without loading the sampler state.

Parameters:

• path (str) – Path to a .npz checkpoint file.

• flatten_chains (bool) – If True (default), collapse the chain dimension so the returned array is always (n_total, n_params). Set to False to get the raw (n_chains, n_samples, n_params) array for per-chain diagnostics (e.g. R-hat).

Returns:

samples – Shape (n_total, n_params) when flatten_chains=True, or (n_chains, n_samples, n_params) otherwise.

Return type:

np.ndarray