API Reference

Detailed technical documentation for the SPLIT pipeline’s modules. Only names with docstrings are listed; private helpers are intentionally omitted.

SPLIT Core (split.split)

class split.split.SPLIT(emri_config_path: str, sample_config_path: str, out_dir: str, custom_injection_func: Callable | None = None, custom_analysis_func: Callable | None = None)

Bases: object

Segmented Posteriors for Long-Inspiral Templates (SPLIT).

An orchestrator class designed to manage loosely-coherent Bayesian parameter estimation for Extreme-Mass-Ratio Inspirals (EMRIs) across multi-GPU hardware. It seamlessly handles the loading of user configurations, the generation and slicing of injected waveform data (including custom/non-Kerr waveforms), prior boundary construction, and the execution of the Eryn-based sampler.

Some choices, such as the Ensemble Sampler moves are currently baked in.

build_priors() → None

Evaluates the physical trajectory and dynamically constructs the Eryn prior bounds.

This method executes the true EMRI trajectory to find the exact injection values at the start of every block. It then constructs Eryn ProbDistContainer objects for the ‘static’ and ‘evolving’ parameter branches based strictly on the parameters requested in the user’s sample configuration JSON, ensuring periodic variables (like phases) are properly registered.

generate_injection_data() → None

Generates the target waveform, scales it to the desired SNR, and processes it for the loosely-coherent likelihood.

This method handles: 1. Injecting either standard Kerr or custom modified waveforms. 2. Windowing the full dataset to compute an accurate 1-year Network SNR. 3. Scaling the luminosity distance to exactly match the requested SNR. 4. Slicing the raw data into independent time blocks. 5. Applying localized Tukey windows to each block to prevent spectral leakage. 6. Computing the frequency-domain FFTs and masked Noise PSDs for the workers.

run_sampler() → None

Configures and executes the Eryn MCMC sampler.

This method acts as the engine of the SPLIT pipeline. It performs: 1. File management, including resuming safely from old HDF5 backends. 2. Dynamic index mapping for fixed vs. actively sampled parameters. 3. Initializing and jittering the start_state array across walkers and temperatures. 4. Setting up custom MCMC moves (Blocked Gibbs and Stretch). 5. Booting the spawned multi-GPU Python multiprocessing.Pool for likelihood evaluations. 6. Running the MCMC loop, checking Gelman-Rubin convergence, and saving diagnostic plots.

split.split.init_worker(d_fft_cpu: ndarray, PSD_cpu: ndarray, freq_mask_cpu: ndarray, Tobs_block_padded: float, dt: float, Nblocks: int, slice_length: int, use_response: bool, wave_gen_func: Any, tdi_kwargs_base: Dict | None = None, index_lambda: int | None = None, index_beta: int | None = None, t_buffer: float | None = None) → None

Initializes the completely isolated environment for each multiprocessing worker. Each worker gets its own GPU context, its own copy of the static data on that GPU, and its own set of ResponseWrapper instances built natively on that GPU.

split.split.window_gen_block_worker(*pars: Any, resp_instance: Any, T: float, dt: float, slice_length: int, alpha: float = 0.05) → Any

Generates padded waveform, slices valid middle block, applies Tukey window.

Moves (split.moves)

class split.moves.BlockedGibbsGaussianMove(*args: Any, **kwargs: Any)

Bases: MHMove

Probabilistic Blocked Gibbs sampler with fixed Gaussian proposals. On each step either the static (hyper) branch or one randomly chosen evolving leaf is updated with a multivariate Gaussian kernel. The branch to update is selected with probability prob_hyper for the static branch and 1 - prob_hyper for the evolving branch. Pre-supplied covariance matrices are scaled by the parallel-tempering inverse temperature to preserve detailed balance across temperature levels.

get_proposal(branches_coords, random, branches_inds=None, **kwargs)

Generate the proposed state.

Parameters:

• branches_coords (dict) – Keys are branch names. Values are arrays of shape (ntemps, nwalkers, nleaves_max, ndim).

• random – Random number generator instance.

class split.moves.BlockedGibbsStretchMove(*args: Any, **kwargs: Any)

Bases: RedBlueMove

Probabilistic Blocked Gibbs sampler using affine-invariant stretch proposals. Adapts the Goodman–Weare stretch move (RedBlue ensemble) to the two-branch SPLIT structure. On each step either the static branch or one randomly chosen evolving leaf is updated via a stretch proposal along the line connecting an active and a complementary walker. The correct log-factor (ndim - 1) * log(z) is accumulated for detailed balance.

get_proposal(s_all, c_all, random, **kwargs)

Generate the proposal.

Parameters:

• s_all (dict) – Keys are branch_names and values are coordinates for which a proposal is to be generated.

• c_all (dict) – Keys are branch_names and values are lists. These lists contain all the complement array values.

Notes

self.xp comes from the parent RedBlueMove class.

class split.moves.SequentialAdaptiveBlockedGibbsGaussianMove(*args: Any, **kwargs: Any)

Bases: RedBlueMove

Deterministic sequential Blocked Gibbs sampler with an adaptive empirical-covariance kernel. Combines sequential block scheduling (via SharedState) with an online covariance estimate built from the complementary walker ensemble. A burn-in phase uses a heavily compressed scale factor for broad initial exploration, then switches to the theoretically optimal Gelman–Roberts–Gilks scaling (2.38^2 / d) once burn_in_steps iterations have elapsed. A diagonal regularizer reg prevents covariance collapse during early burn-in when walkers are tightly clustered.

class split.moves.SequentialBlockedGibbsGaussianMove(*args: Any, **kwargs: Any)

Bases: MHMove

Deterministic sequential Blocked Gibbs sampler with fixed Gaussian proposals. Cycles through all evolving blocks and the static branch, updating exactly one block per call. A SharedState object keeps the schedule synchronized when this move is mixed with other sequential moves (e.g. SequentialBlockedGibbsStretchMove). The fixed covariance matrices are pre-scaled by the inverse temperature at each step.

class split.moves.SequentialBlockedGibbsStretchMove(*args: Any, **kwargs: Any)

Bases: RedBlueMove

Deterministic sequential Blocked Gibbs sampler using affine-invariant stretch proposals. Cycles through all evolving blocks and the static branch. Both halves of the RedBlue ensemble update the same block within each full iteration (controlled via floor-division on SharedState.step). Compatible with SequentialBlockedGibbsGaussianMove when both share the same SharedState instance.

Priors (split.priors)

class split.priors.MarkovStudenttPrior(prior_ev: Any, prior_st: Any, dt_block: float, nu: float, sigma_dict: Dict[str, float], samp_config: Dict, emri_config: Dict, all_param_names: List[str], true_evolving_dict: Dict[str, ndarray], traj_instance: Any, custom_arg_names: List[str], analysis_add_args_dict: Dict[str, float])

Bases: object

Custom Eryn Prior evaluating a heavy-tailed Student-t transition penalty for loosely-coherent EMRI blocks against theoretical vacuum-GR trajectories.

This class bypasses the assumption that loosely-coherent blocks are fully independent. It applies a base prior (uniform) to all blocks, and then evaluates a Markovian transition penalty: it forward-evolves the parameters from block (i-1) using pure GR, and applies a Student-t penalty based on how far the proposed parameters in block (i) deviate from that prediction.

logpdf(coords: Dict[str, ndarray], inds: Dict[str, ndarray], supps: Any | None = None, branch_supps: Any | None = None) → ndarray

Evaluates the joint log-prior probability for the entire MCMC ensemble.

This method is called natively by Eryn when “all_models_together” is specified. It first evaluates the independent base bounds. If a walker is within bounds, it computes the autoregressive transition penalty block-by-block.

Parameters:

• coords (dict) – Dictionary containing the proposed coordinates for all walkers. - coords[“evolving”] shape: (ntemps, nwalkers, Nblocks, ndim_ev) - coords[“static”] shape: (ntemps, nwalkers, 1, ndim_st)

• inds (dict) – Dictionary of boolean arrays indicating which leaves/blocks are active. - inds[“evolving”] shape: (ntemps, nwalkers, Nblocks)

• supps (object, optional) – Supplemental overall information passed by Eryn. (Unused here)

• branch_supps (dict, optional) – Supplemental branch-specific information passed by Eryn. (Unused here)

Returns:

total_logP + penalty – A 2D array of shape (ntemps, nwalkers) containing the final computed log-prior probability (base bounds + trajectory penalty) for each walker. Walkers proposing unphysical jumps or values outside the base bounds will return -inf.

Return type:

np.ndarray

Diagnostics (split.diagnostics)

split.diagnostics.check_convergence(sampler, Nblocks, min_autocorr_iters=50, autocorr_threshold=50, gelmanrubin_threshold=1.05)

Evaluates sampler convergence via Autocorrelation (tau), and Gelman-Rubin (R-hat). Logs all sampler statistics, including PT swaps.

Returns:

True if all convergence criteria are met, False otherwise.

Return type:

bool

split.diagnostics.update_diagnostic_plots(sampler, diagnostics_dir, Nblocks, static_in_names, ev_in_names, val_samp_st, val_samp_ev, true_pars_all, traj_config, min_autocorr_iters=10, autocorr_threshold=50, discard_frac=0.5, max_plot=10000, corner_kwargs=None)

Extract multi-branch chains, plot 1D walks, static posteriors, and t=0 projections.

traj_config (dict): Contains all trajectory mapping parameters:

dt, slice_length, idx_st_in, idx_ev_in, idx_st_fix, idx_ev_fix, val_st_fix, val_ev_fix, kerr_traj_instance, traj_indices, total_pars_len.

Utilities (split.utils)

split.utils.compute_rhat(x)

Custom EMRIs (split.customEMRIs)

split.customEMRIs is a namespace package; its submodules are documented individually below.

split.customEMRIs.AccEccEqPn5AAK

class split.customEMRIs.AccEccEqPn5AAK.AccEccEqPN5Traj(*args: Any, **kwargs: Any)

Bases: PN5

PN5 EMRI trajectory with a power-law modification to the p and e fluxes induced by planetary migration effect in EccEq EMRIs due to an accretion disk. Model from Duque et al. (2024) https://arxiv.org/pdf/2411.03436.

add_fixed_parameters(m1: float, m2: float, a: float, additional_args=None)

Additional trajectory parameters are initialized here. For the accretion model, valid list of

modify_rhs(ydot: ndarray, y: ndarray, **kwargs) → ndarray

This function allows the user to modify the right-hand side of the ODE after any required Jacobian transforms have been applied.

By default, this function returns the input right-hand side unchanged.

class split.customEMRIs.AccEccEqPn5AAK.AccEccEqPn5AAKWaveform(*args: Any, **kwargs: Any)

Bases: AAKWaveformBase

Waveform generation class for AAK with AccEccEqPN5Traj trajectory as defined above.

This class generates waveforms based on the Augmented Analytic Kludge given in the EMRI Kludge Suite. However, here the trajectory is vastly improved by employing the 5PN fluxes for generic Kerr orbits from Fujita & Shibata 2020.

The 5PN trajectory produces orbital and phase trajectories. The trajectory is calculated until the orbit reaches within 0.2 of the separatrix, determined from arXiv:1912.07609. The fundamental frequencies along the trajectory at each point are then calculated from the orbital parameters and the spin value given by (Schmidt 2002).

These frequencies along the trajectory are then used to map to the frequency basis of the Analytic Kludge. This mapping takes the form of time evolving large mass and spin parameters, as well as the use of phases and frequencies in \((alpha, \Phi, \gamma)\):

\[\Phi = \Phi_\phi,\]

\[\gamma = \Phi_\phi + \Phi_\Theta,\]

\[alpha = \Phi_\phi + \Phi_\Theta + \Phi_r.\]

The frequencies in that basis are found by taking the time derivatives of each equation above.

This class has GPU capabilities and works from the sparse trajectory methodoligy with cubic spine interpolation of the smoothly varying waveform quantities. This waveform does not have the freedom in terms of user-chosen quantitites that few.waveform.base.SphericalHarmonicWaveformBase contains. This is mainly due to the specific waveform constructions particular to the AAK/AK.

Please note: the 5PN trajectory and AAK waveform take the parameter \(Y\equiv\cos{\iota}=L/\sqrt{L^2 + Q}\) rather than \(x_I\) as is accepted for relativistic waveforms and in the generic waveform interface discussed above. The generic waveform interface directly converts \(x_I\) to \(Y\).

Parameters:

• inspiral_kwargs – Optional kwargs to pass to the inspiral generator. Important Note: These kwargs are passed online, not during instantiation like other kwargs here. Default is {}. This is stored as an attribute.

• sum_kwargs – Optional kwargs to pass to the sum module during instantiation. Default is {}.

split.customEMRIs.FastKerrEccentricEquatorialAccretionFlux

class split.customEMRIs.FastKerrEccentricEquatorialAccretionFlux.FastKerrEccentricEquatorialAccretionFlux(*args: Any, **kwargs: Any)

Bases: SphericalHarmonicWaveformBase, KerrEccentricEquatorial

Prebuilt model for fast Kerr eccentric equatorial flux-based waveforms.

This model combines the most efficient modules to produce the fastest accurate EMRI waveforms. It leverages GPU hardware for maximal acceleration, but is also available on for CPUs.

The trajectory module used here is few.trajectory.inspiral for a flux-based, sparse trajectory. This returns approximately 100 points.

The amplitudes are then determined with few.amplitude.ampinterp2d.AmpInterp2D along these sparse trajectories. This gives complex amplitudes for all modes in this model at each point in the trajectory. These are then filtered with few.utils.modeselector.ModeSelector.

The modes that make it through the filter are then summed by few.summation.interpolatedmodesum.InterpolatedModeSum.

See few.waveform.base.SphericalHarmonicWaveformBase for information on inputs. See examples as well.

Parameters:

• inspiral_kwargs – Optional kwargs to pass to the inspiral generator. Important Note: These kwargs are passed online, not during instantiation like other kwargs here. Default is {}.

• amplitude_kwargs – Optional kwargs to pass to the amplitude generator during instantiation. Default is {}.

• sum_kwargs – Optional kwargs to pass to the sum module during instantiation. Default is {}.

• Ylm_kwargs – Optional kwargs to pass to the Ylm generator during instantiation. Default is {}.

• *args – args for waveform model.

• **kwargs – kwargs for waveform model.

class split.customEMRIs.FastKerrEccentricEquatorialAccretionFlux.KerrEccEqAccFlux(*args: Any, **kwargs: Any)

Bases: KerrEccEqFlux

KerrEccEq EMRI trajectory with a power-law modification to the p and e fluxes induced by planetary migration effect in EccEq EMRIs due to an accretion disk. Model from Duque et al. (2024) https://arxiv.org/pdf/2411.03436.

add_fixed_parameters(m1: float, m2: float, a: float, additional_args=None)

Additional trajectory parameters are initialized here. For the accretion model, valid list of

modify_rhs(ydot: ndarray, y: ndarray, **kwargs) → ndarray

This function allows the user to modify the right-hand side of the ODE after any required Jacobian transforms have been applied.

By default, this function returns the input right-hand side unchanged.