Interpolators (Surrogates)

Gaussian process and surrogate modeling for efficient parameter inference.

Overview

These modules provide surrogate models and interpolation methods that accelerate likelihood evaluation. Instead of computing expensive waveform models for each point, the integrators can query these surrogate models for fast approximations.

See also

• Integrators (MC Sampling) - MC integration modules that use these surrogates

• Likelihood and Priors - Full likelihood evaluation (for comparison)

Bayesian Least Squares

RIFT.interpolators.BayesianLeastSquares.fit_quadratic(x, y, x0=None, variable_symmetry_list=None, gamma_x=None, prior_x_gamma=None, prior_quadratic_gamma=None, verbose=False, n_digits=None, hard_regularize_negative=False, hard_regularize_scale=1)

Simple least squares to a quadratic. Written out in long form so I can impose priors as needed on (a) the fit coefficients [=regularization] and (b) the ‘x’ coordinates. INPUT:

x = array so x[0] , x[1], x[2] are points. y = array of y values

OUTPUT

peak_val_est, best_val_est, my_fisher_est, linear_term_est, fit_here : python function providing best fit. Usually what you want, but cannot be pickled/saved to text

OPTIONAL

variable_symmetry_list = list of length x, indicating symmetry under ONE discrete symmetry (so far)

RIFT.interpolators.BayesianLeastSquares.fit_quadratic_and_resample(x, y, npts, rho_fac=1, x0=None, gamma_x=None, prior_x_gamma=None, prior_quadratic_gamma=None, verbose=False, n_digits=None, hard_regularize_negative=False, hard_regularize_scale=1, force_include_mean=False)

Simple least squares to a quadratic, and resamples from the quadratic derived from the fit. Critical for iterative evaluation of

• Fisher matrix

• lnLmarg (ILE)

TO DO:

• implement non-stochastic placement as option (e.g., like effectiveFisher.py)

Constrained Quadratic Likelihood

RIFT.interpolators.ConstrainedQuadraticLikelihood.fit_grid(sample, values, extra_guess_mu=None, extra_guess_cov=None, **kwargs)

Fit multivariate normal object to samples with values. Note the fit assumes the data is PERFECTLY quadratic,

and ENFORCES a nonnegative-definite parameterization, so make sure you curate the data to make that possible

Parameters

sample: array like, shape = (npts, ndim)

Input sample locations

values: array like, shape = (npts,)

Input sample values

limits: array like, shape = (ndim, 2)

Input bounded interval limits

seed: int, optional

Input seed for random numbers

nwalk: int, optional

Input number of random walkers

nstep: int, optional

Input number of walker steps

carryover: float, optional

Input fraction of best guesses carried over through genetic algorithm

sig_factor: float, optional

Input controls jump sizes

RIFT.interpolators.ConstrainedQuadraticLikelihood.fit_quadratic_alt(x, y, y_err=None, x0=None, symmetry_list=None, verbose=False, hard_regularize_negative=True)

RIFT.interpolators.ConstrainedQuadraticLikelihood.quad_residuals(x, yvals, lnL_offset, mu, icov)

RIFT.interpolators.ConstrainedQuadraticLikelihood.quad_residuals_vector(x, yvals, lnL_offset, mu, icov)

Efficient SKLearn GP Save

RIFT.interpolators.efficient_save_sklearn_gp.export_gp_compact(fname_base, gp)

RIFT.interpolators.efficient_save_sklearn_gp.gp2json(gp)

Converts gp object to (a) json object and (b) data for X_train_ and y_train

RIFT.interpolators.efficient_save_sklearn_gp.load_gp(fname_base)

Basic GP

class RIFT.interpolators.gp.GaussianProcess1d(dat1d, sigma0, sigmab=1e-05, h=None)

Bases: object

Very simple gaussian process model. Does not take into account errors that vary with position.

predict(x_new)

RIFT.interpolators.gp.squared_exponential_kernel(dx, h, sigma0, sigmab)

GPyTorch Wrapper

Internal GP

GPR_Fit.py python 2.7.13 Vera Delfavero Date updated: 02/02/2019

The function of this script is to provide the function, “fit_gp” to the PE code (utils_ConstructIntrinsicPosterior_GenericCoordinates.py)

That function will fit training data, using a sparse gaussian proces, and return a function which takes only test data as an input, and returns only the estimated mean. (Because that’s all the PE code uses) This saves a lot of time, because calculating the covariance of the test data is expensive. This function is also a sparse gaussian proces, which should make it at least an order of magnitude faster than the sklearn gaussian process code.

The basic algorithm for this can be found on page 88 of R&W’s book on Gaussian Processes and Machine Learning.

RIFT.interpolators.internal_GP.GP_fit_function(x_train, y_train, thetas, white_noise, kernel_function=<function compact_kernel>, **kwargs)

Objective: return a function which can sample the distribution

with only the test inputs as an input. This function only needs to return a mean.

RIFT.interpolators.internal_GP.compact_kernel(Xp, Xq, thetas, white_noise)

compact_kernel: Return a sparse representation of the

compact kernel described on page 88 of R&W. Most of the functions which generate the sparse kernel involve non-sparse matricies.

Inputs:

Xp: Dense array containing input vector Xq: Dense array containing additional input vector

Note: if Xq is given, the returned kernel is the

transpose of the true kernel

thetas: Hyperparameters

thetas[0]: Scale length used to construct sparse matrix.

Outputs:

K: Sparse kernel representation (csc matrix)

RIFT.interpolators.internal_GP.fit_gp(x_train, y_train, y_errors=None, sparse_coeff=3.0, white_noise=0.01)

This function can directly replace fit_gp in utils_ConstructIntrinsicPosterior_GenericCoordinates.py

It is an implimentation of my sparse fitting code, however most of this function was copied from the original.

Inputs:

x_train (numpy.ndarray):

Training data for model shape: (m, n) where m is the number of points,

and n is the dimensionality

y_train (numpy.ndarray):

Training data for model shape: (m, ) where m is the number of points

sparse_coeff (float):

This hyperparameter determines the sparseness of the matrix. A sparse_coeff of 1. is not sparse. A sparse_coeff of inf is diagonal. What value you can use for this is largely dependent on how

much training data you have, as well as the dimensionality of your data

About each dimension in parameter space, the scale length is

calculated for that dimension as h/hs, where h is the range of the training data in that dimension.

white_noise (float):

The white noise hyperparemeter

Outputs:

my_fit (function(x_sample)):

A function which will sample our model on x_sample,

for an x_sample like x_train

GPU Interpolation

SENNI (Surrogate Enabled Nested Nested Importance Sampling)

Differentiable (JAX) surrogates: jax_gp

RIFT.interpolators.jax_gp is a pure-JAX framework for fitting a continuous, jax.grad-able surrogate of the marginalized log-likelihood lnL(theta) over the ILE intrinsic grid, and exporting it as a small, portable artifact that reloads with no RIFT/lalsimutils dependency. It is the differentiable counterpart to the GP/RF fits the legacy CIP uses, intended for AD use cases (gradient-based sampling, population inference) and for shipping a self-contained likelihood surrogate.

Interpolators

Each interpolator fits in a per-dimension whitened space and exposes a pure-JAX lnL_physical(theta) (differentiable in the fit coordinates).

Artifact export / import

Self-contained, differentiable export for jax_gp interpolators.

This replaces “dump the lnL grid and hope”: the artifact written here is a small, portable bundle that reconstructs a pure-JAX, differentiable lnL(theta) – jax.grad / jax.value_and_grad work out of the box, which is what downstream users need.

Layout (two files sharing a base path):

<base>.npz        whitening vectors + method parameters (NumPy arrays)
<base>.meta.json  schema/method/dimension/target-scaling + coordinate names

coord_names records which fit coordinates the axes of theta are (e.g. ["mc", "eta", "chi_eff", ...]). The exported lnL is differentiable in those fit coordinates – the same space the GP was trained on. (Pushing the derivative all the way back to raw physical parameters would require a JAX reimplementation of CIP’s coordinate transforms; that is deliberately out of scope here and noted as future work.)

RIFT.interpolators.jax_gp.export.exists(path)

RIFT.interpolators.jax_gp.export.load(path)

Reconstruct a (predict + differentiate) interpolator from an exported bundle.

Returns an interpolator instance exposing the usual contract: predict(X), predict_callable(), lnL_physical(theta) (pure JAX), lnL_and_grad(theta), grad_fn().

RIFT.interpolators.jax_gp.export.save(model, path, coord_names=None, extra_meta=None)

Write model to <path>.npz + <path>.meta.json; return the base path.

Applications

export_artifact packages a single ILE .net into a differentiable artifact; jax_cip is a standalone pure-JAX intrinsic-posterior path; export_at_scale points at a real run directory, exports the artifact, and validates it against the run’s own CIP posterior (apples-to-apples, with RIFT’s priors), locally or over HTCondor. See applications/EXPORT_AT_SCALE.md and applications/ARTIFACT.md.

export_at_scale – point at a real RIFT run directory, ship a differentiable lnL artifact for its all.net, and validate it against the run’s own posterior.

This is the “do it for real, do it for many” wrapper around the jax_gp primitives. export_artifact/jax_cip already turn a single all.net into a pure-JAX lnL(theta) surrogate; this tool adds the three things a production sweep needs:

1. Discovery. Read a run directory the way RIFT left it – detect the all.net column layout (tides / no-tides / precessing all differ), parse the active args_cip_list.txt for the fit coordinates and the prior box (--mc-range / --eta-range / --chi-max), and locate the run’s own final CIP posterior (posterior_samples-<N>.dat) to validate against.

2. Export + validate the *deliverable*. Fit the surrogate in dimension-agnostic physical fit coordinates – [mc, delta_mc] plus whichever spin/tidal columns actually vary, so it covers aligned, BNS and precessing runs without needing a hand-written coordinate transform – export it, reload the saved artifact, draw a posterior from the reloaded bytes (importance sampling over the run’s prior box), and report the Jensen-Shannon divergence of the mc/q/ chi_eff marginals against the run’s CIP posterior.

3. Scale. batch discovers many run directories and either runs them locally or emits an HTCondor DAG – one node per run, the submit file templated from the run’s own CIP.sub (accounting group, singularity image, requirements) so the jobs land in the same place the run itself ran.

Nothing is written back into the run directory: every artifact + report goes under a separate --workroot (default ./export_at_scale_out), one subdirectory per run.

At this stage we interpolate only ``all.net`` (the existing intrinsic ILE deliverable). The distance-grid export is a separate track.

CLI:

# one run, immediately
python -m RIFT.interpolators.jax_gp.applications.export_at_scale one \
    --run-dir /path/to/rundir --workroot /path/to/out

# inspect what discovery found, without doing any work
python -m RIFT.interpolators.jax_gp.applications.export_at_scale discover \
    --run-dir /path/to/rundir

# many runs -> a condor DAG (submit it with condor_submit_dag)
python -m RIFT.interpolators.jax_gp.applications.export_at_scale batch \
    --runs '/data/*/S*/rift*/' --workroot /path/to/out --condor

RIFT.interpolators.jax_gp.applications.export_at_scale.ALL_COMPARE_PARAMS = ('mc', 'eta', 'q', 'chi_eff', 'chiMinus', 's1z', 'chi1_perp', 'phi1', 's2z', 'chi2_perp', 'phi2')

every intrinsic parameter the validation reports a JS on. Masses, the aligned-spin combinations (chi_eff, chiMinus), and the cylindrical-polar spin of each body (aligned component s{b}z, in-plane magnitude chi{b}_perp, azimuth phi{b}).

RIFT.interpolators.jax_gp.applications.export_at_scale.build_sampling(spec, fit_names, spin_modes=None)

Return the RIFT-measure sampling spec for the artifact’s fit_names.

spin_modes (per body {"1": m, "2": m} with m in {"sph", "cart_z", None}) fixes how each body’s spin is sampled, derived from the raw physics (which spin columns vary), NOT from fit_names – because the fit may use chi_eff/chiMinus instead of s1z/s2z. When omitted, it is inferred from fit_names (back-compat: Cartesian fit coords).

Returns a dict with names/lo/hi (prior box), ln_prior(theta), to_fit(theta) (sampling coords -> the artifact’s fit_names vector), raw_to_sample (ILE data -> sampling coords) and to_compare (samples -> every physical comparison parameter).

RIFT.interpolators.jax_gp.applications.export_at_scale.detect_net_layout(path, max_probe=200)

Infer the column map of a RIFT all.net / .composite file.

Returns (cols, intrinsic_names, has_neff) where cols is the name->index dict load_ile_net() consumes and intrinsic_names are the physical columns of the intrinsic block (e.g. ['m1','m2','s1x',...]).

The detection is tail-anchored (... lnL sigma_lnL ntot [neff]) and sanity-checked on the candidate sigma_lnL column (a small positive MC error), so it is robust to the aligned/precessing/tidal width differences without a hard-coded table.

RIFT.interpolators.jax_gp.applications.export_at_scale.discover_run(run_dir)

Inspect a RIFT run directory and return a dict describing how to export + validate it. Raises if there is no usable all.net.

RIFT.interpolators.jax_gp.applications.export_at_scale.expand_runs(patterns)

Expand globs/paths to run directories that actually contain a usable all.net.

RIFT.interpolators.jax_gp.applications.export_at_scale.fit_and_export(spec, out_base, method='svgp', sigma_cut=0.6, lnL_offset=40.0, cap_points=8000, n_features=256, n_opt_steps=300, seed=0, quadgp_residual='svgp', keep_curv_frac=0.01, ls_lo_frac=0.2, ls_hi_frac=1.0, mass_coord='eta', spin_coord='auto')

Build, persist and cold-reload-verify a differentiable lnL artifact for the run’s all.net. Returns a metadata dict (also the fit-coord arrays needed by validation, under private keys).

spin_coord="auto" (default) fits BOTH aligned_eff (chi_eff/chiMinus principal axes – fixes the sharp low-mass aligned spin) and cartesian (s1z/s2z) and keeps whichever has the lower peak-region holdout RMSE. This is the “never regress” guard: aligned_eff is a large net win but worse on a minority of events, and holdout RMSE reliably picks the better of the two (it selects cartesian exactly where aligned_eff would regress). Pass an explicit aligned_eff / cartesian to skip the selection (half the fit cost).

RIFT.interpolators.jax_gp.applications.export_at_scale.fit_prior_box(fit_names, spec)

Per-coordinate uniform-prior box [lo, hi] in fit coords, from the run’s CIP ranges (mc, eta->delta_mc, chi-max per spin component).

RIFT.interpolators.jax_gp.applications.export_at_scale.fit_to_physical(Xfit, fit_names)

Inverse of the mass part + spin passthrough: fit coords -> physical dict with arrays m1, m2, q, chi_eff and any spin components present.

RIFT.interpolators.jax_gp.applications.export_at_scale.main(argv=None)

RIFT.interpolators.jax_gp.applications.export_at_scale.raw_to_fit(X_raw, raw_names, mass_coord='eta', spin_coord='aligned_eff')

(m1,m2,spin...) -> (mc, <mass2>, <aligned spin>, <in-plane spin>) columns.

The surrogate must be fit in the coordinates the lnL Fisher is quadratic in and axis-aligned with (the quadratic core whitens per-dimension and the GP residual uses axis-aligned ARD lengthscales – neither can represent a sharp ridge along a diagonal direction).

mass_coord: "eta" (default) or "delta_mc". eta is the Fisher-quadratic variable; delta_mc hides that curvature near equal mass (it is quadratic in delta_mc), so eta fixes the mass-ratio (q) marginal.

spin_coord: "aligned_eff" (default) replaces the aligned components (s1z, s2z) with (chi_eff, chiMinus) – the principal axes of the aligned-spin Fisher. In (s1z, s2z) the well-measured chi_eff is a diagonal ridge an axis-aligned ARD GP over-smooths (the low-mass aligned-spin failure mode); rotating to (chi_eff, chiMinus) makes it axis-aligned and resolvable. The in-plane components (s1x, s1y, s2x, s2y) are kept as-is (weakly constrained). "cartesian" keeps the raw components.

RIFT.interpolators.jax_gp.applications.export_at_scale.run_one(run_dir, workroot, method='quadgp', n_samples=40000, seed=0, cap_points=8000, n_features=256, n_opt_steps=300, lnL_offset=40.0, sigma_cut=0.6, sampler='auto', keep_curv_frac=0.01, ls_lo_frac=0.2, ls_hi_frac=1.0, mass_coord='eta', spin_coord='auto', write_plot=True)

Discover -> fit+export -> validate one run; write all artifacts under workroot/<label>/ and return the full report.

RIFT.interpolators.jax_gp.applications.export_at_scale.validate_artifact(spec, fit_meta, out_dir, n_samples=40000, inflate=1.2, seed=0, sampler='auto', compare_params=('mc', 'eta', 'q', 'chi_eff', 'chiMinus', 's1z', 'chi1_perp', 'phi1', 's2z', 'chi2_perp', 'phi2'))

Draw a posterior from the reloaded artifact and JS-compare its marginals to the run’s CIP posterior. Writes posterior_interp.dat and returns the report.

The target is the run’s actual posterior lnL(theta) + ln prior(theta) – NOT a flat-prior caricature – so the comparison is apples-to-apples with CIP. Sampling is done in RIFT’s own coordinates (build_sampling()): spins in (chi, cos_theta, phi), where RIFT’s isotropic prior is flat, plus the non-uniform mass-prior shape (mc_prior ~ mc, eta_prior ~ eta^-6/5). This both matches CIP’s measure and avoids the Cartesian-spin 1/chi^2 singularity.

sampler: "gaussian" – fast mu-matched importance sampling (great in low dimension); "nuts" – gradient-based NUTS preconditioned with the data covariance, which exploits the artifact’s AD gradients and explores the curved, high-dimensional precessing posterior far better than a single Gaussian proposal; "auto" (default) picks nuts when there are >3 sampling dimensions, else gaussian.

RIFT.interpolators.jax_gp.applications.export_at_scale.write_condor_batch(run_dirs, workroot, method='quadgp', n_samples=40000, sampler='auto', python=None, pythonpath=None, accounting_group=None)

Emit a one-node-per-run HTCondor DAG under workroot/condor/.

The submit file is templated from the first run’s own CIP.sub (accounting group / singularity image / requirements) so the validation jobs land in the same pool the run itself used. Returns the DAG path.

Package a RIFT ILE .net file into a self-contained, differentiable lnL artifact.

This is the “ship it” end of the jax_gp pipeline: it takes raw ILE output (the per-point Monte-Carlo lnL evaluations CIP normally consumes) and produces a small, portable bundle (<base>.npz + <base>.meta.json) that reconstructs a pure-JAX, jax.grad-able lnL(theta) – with no RIFT/lalsimutils dependency at load time.

The surrogate is fit in fit coordinates (the decorrelated space CIP itself fits in), so the exported lnL is differentiable in those coordinates – recorded as coord_names in the meta – not in the raw physical parameters. See ARTIFACT.md.

Example

python export_artifact.py --net /home/oshaughn/all.net \
    --out /tmp/gw170817_rff --coords bns

RIFT.interpolators.jax_gp.applications.export_artifact.RAW_COORDS = ('m1', 'm2', 's1z', 's2z', 'lambda1', 'lambda2')

raw 6-parameter intrinsic columns as loaded from an ILE .net file

RIFT.interpolators.jax_gp.applications.export_artifact.build_artifact(net_path, out_base, coords='bns', method='rff', sigma_cut=0.6, lnL_offset=40.0, cap_points=8000, n_features=512, n_opt_steps=300, seed=0, quadgp_residual='svgp')

Build and persist a differentiable lnL artifact from an ILE .net file.

The pipeline mirrors CIP’s fit preparation: load + de-dupe + sigma_cut the ILE points, transform to fit coordinates, keep only the high-lnL region (lnL > max - lnL_offset), optionally subsample to cap_points, then fit the chosen interpolator using the per-point Monte-Carlo errors. The model is exported via RIFT.interpolators.jax_gp.export.save(), reloaded, and checked for predict agreement and finite jax.grad.

Parameters

net_pathstr

Path to the RIFT ILE .net file.

out_basestr

Output base path; writes <out_base>.npz and <out_base>.meta.json.

coords{“bns”, “raw”}, optional

Fit-coordinate system (see _build_fit_coordinates()).

methodstr, optional

jax_gp interpolator name ("rff", "exact", "svgp", "quadgp"). "quadgp" is the PE-grade quadratic-core + GP-residual surrogate; its residual backend is selected by quadgp_residual.

quadgp_residualstr, optional

Residual GP for method="quadgp" ("svgp", "exact", "rff"); ignored otherwise.

sigma_cutfloat, optional

Drop ILE points whose reported sigma_lnL exceeds this (CIP default 0.6).

lnL_offsetfloat, optional

Keep only points with lnL > max(lnL) - lnL_offset.

cap_pointsint or None, optional

If set and there are more surviving points, random-subsample down to this many (like CIP --cap-points).

n_features, n_opt_steps, seedoptional

Passed to the interpolator constructor (where applicable).

Returns

dict

Metadata: n_train, coord_names, lnL_max, holdout_rmse, plus the resolved build settings.

RIFT.interpolators.jax_gp.applications.export_artifact.main(argv=None)

argparse CLI: build an artifact and print its metadata as JSON.

Compare intrinsic-posterior 1D marginals via Jensen-Shannon divergence.

The success metric for jax_cip is NOT sampler ESS – it is how close the recovered posterior is to the production CIP+RF benchmark. We quantify that with the JS divergence of each 1D marginal (start with mc, the best-measured parameter), in bits (log base 2, so JS in [0, 1]).

Usage:

python -m RIFT.interpolators.jax_gp.applications.compare –a jax_out.xml.gz –b cip_rf_out.xml.gz –param mc

Both inputs are RIFT ChooseWaveformParams XML (what CIP and jax_cip both write).

Caveat (per design notes): a reliable JS needs enough independent samples in BOTH posteriors. jax_cip’s per-run effective sample count is modest, so accumulate over many seeds/instances (Condor) before trusting a small JS. js_divergence_1d here also returns a bootstrap stderr so you can see when you are statistics-limited.

RIFT.interpolators.jax_gp.applications.compare.js_divergence_1d(a, b, bins=80, value_range=None)

Jensen-Shannon divergence (bits) between two 1D sample sets via histograms.

RIFT.interpolators.jax_gp.applications.compare.js_with_stderr(a, b, bins=80, n_boot=200, seed=0)

JS plus a bootstrap stderr (so you can tell when you’re statistics-limited).

RIFT.interpolators.jax_gp.applications.compare.load_param_from_xml(fname, param='mc')

Load a 1D parameter array (default mc, in Msun) from a RIFT samples XML.

RIFT.interpolators.jax_gp.applications.compare.load_param_pooled(fnames, param='mc')

Concatenate a parameter across several XMLs (pool e.g. 10 benchmark runs).

RIFT.interpolators.jax_gp.applications.compare.main(argv=None)