sigway.ms_solver

Solve a model of inflation to get the primordial spectrum \(\mathcal{P}_\zeta(k)\), rather than writing it down by hand.

The Mukhanov–Sasaki equation

During inflation the comoving curvature perturbation \(\zeta\) is most cleanly evolved through the gauge-invariant Mukhanov–Sasaki (MS) variable \(v_k = z\,\zeta_k\), where \(z = a\,\dot\phi/H\). Each Fourier mode obeys a forced-oscillator equation,

\[ \frac{\mathrm{d}^2 v_k}{\mathrm{d}\eta^2} + \left(k^2 - \frac{1}{z}\frac{\mathrm{d}^2 z}{\mathrm{d}\eta^2}\right) v_k = 0 , \]

with \(\eta\) the conformal time. Deep inside the horizon (\(k \gg aH\)) the \(k^2\) term dominates and the mode oscillates like a flat-space (Bunch–Davies) vacuum; as the mode crosses the horizon the \(z''/z\) term takes over and \(\zeta_k\) freezes to a constant. The curvature power spectrum is read off from the frozen amplitude,

\[ \mathcal{P}_\zeta(k) = \frac{k^3}{2\pi^2}\,\bigl|\zeta_k\bigr|^2 = \frac{k^3}{2\pi^2}\,\left|\frac{v_k}{z}\right|^2 . \]

How SingleFieldSolver works

For a given potential \(V(\phi)\) it proceeds in three steps (the code integrates in e-folds \(N=\ln a\) rather than conformal time):

1. Background — integrate the inflaton equation of motion together with the Friedmann equation to obtain \(\phi(N)\) and \(H(N)\), and hence \(z(N)\).
2. Modes — for each comoving wavenumber \(k\), evolve the MS equation from a sub-horizon Bunch–Davies initial condition (a few e-folds before horizon crossing) through to well after horizon crossing, where \(\zeta_k\) has frozen.
3. Spectrum — evaluate \(\mathcal{P}_\zeta(k)\) from the frozen modes.

This turns a model of inflation — e.g. an ultra-slow-roll quasi-inflection-point potential — directly into the small-scale \(\mathcal{P}_\zeta\) that sources scalar-induced GWs. Wrap the solver in SingleFieldPerturbations to feed it to an OmegaGW model.

Solver

sigway.ms_solver.SingleFieldSolver

SingleFieldSolver(V, phi0=0.0, pi0=0.0, N_CMB_to_end=65.0, max_efolds=1000.0, cmb_bounds=CMB_BOUNDS, check_consistency=False, N_subhorizon=3.0, N_suphorizon=7.0, k=None, upsample=False, background_solver_opts={}, perturbation_solver_opts={}, error_on_fail=False)

Compute the primordial scalar power spectrum \(\mathcal{P}_\zeta(k)\) for a single inflaton field.

Given a potential \(V(\phi, *\mathrm{params})\), this class:

1. Integrates the background Klein-Gordon and Friedmann equations in e-fold time \(N\) from the initial field value \(\phi_0\) (on the slow-roll attractor) until the end of inflation (\(\epsilon_H = 1\)).
2. For each comoving wavenumber \(k\) in the requested grid, integrates the Mukhanov-Sasaki equation from Bunch-Davies initial conditions (set \(N_{\rm sub}\) e-folds before horizon crossing) to well after horizon crossing (\(N_{\rm sup}\) e-folds after), then reads off the frozen amplitude to obtain \(\mathcal{P}_\zeta(k)\).

The potential is normalised internally as \(U(\phi) = V(\phi)/V(\phi_0)\), so the overall energy scale \(V_0 \equiv V(\phi_0)\) is restored at the end. The rescaled background variables are \(\pi \equiv \mathrm{d}\phi/\mathrm{d}N\) and \(h = H / \sqrt{V_0/3}\).

The primary entry points are:

• solver.run(k, *params) — runs the full calculation and returns a callable interpolant \(\mathcal{P}_\zeta(k_{\rm new})\).
• solver(k, *params) — same calculation, returns the raw \(\mathcal{P}_\zeta\) array instead.

Parameters:

Name	Type	Description	Default
V	callable	Inflaton potential \(V(\phi, *\mathrm{params})\). Must be written in JAX (using jax.numpy) so that automatic differentiation can compute \(V'\) and \(V''\) internally.	required
phi0	float	Initial field value \(\phi_0\). Also the normalisation point: \(U(\phi) \equiv V(\phi)/V(\phi_0)\). Default 0.0.	0.0
pi0	float	Initial field velocity \(\pi_0 = \mathrm{d}\phi/\mathrm{d}N\big|_0\). Currently stored for reference; the background integration always starts from the slow-roll attractor value derived from \(V'(\phi_0)\). Default 0.0.	0.0
N_CMB_to_end	float	Number of e-folds from the CMB pivot scale \(k_* \approx 0.05\,\mathrm{Mpc}^{-1}\) to the end of inflation, assuming instantaneous reheating. This sets the absolute \(k\)-to-\(N\) mapping. Default 65.0.	65.0
max_efolds	float	Hard upper limit on the number of e-folds integrated for the background. Increase this for models with very long inflationary phases. Default 1000.0.	1000.0
cmb_bounds	dict	Gaussian CMB observational priors on \([\ln(10^{10}\mathcal{P}_\zeta),\, n_s,\, r]\), supplied as a dict with keys "means" (shape (3,)) and "cov" (shape (3, 3)). Used only when check_consistency=True. Defaults to Planck 2018 best-fit values.	CMB_BOUNDS
check_consistency	bool	If True, compute the Gaussian log-likelihood of the slow-roll CMB observables against cmb_bounds after the background run. Informational only; does not affect the perturbation calculation. Default False.	False
N_subhorizon	float	Number of e-folds before horizon crossing (\(k = aH\)) at which to impose Bunch-Davies vacuum initial conditions on each mode. Default 3.0.	3.0
N_suphorizon	float	Number of e-folds after horizon crossing at which to read off the frozen super-horizon mode amplitude. Default 7.0.	7.0
k	array - like or callable or None	Wavenumber grid (in \(\mathrm{s}^{-1}\)) onto which the spectrum is interpolated when upsample=True. Ignored otherwise. Default None.	None
upsample	bool	If True, evaluate the returned power spectrum on the dense k grid rather than on the wavenumbers passed to run. Requires k to be set. Default False.	False
background_solver_opts	dict	Override any field of the default SolverOptions for the background integrator. Defaults: rtol=1e-8, atol=1e-8, max_steps=100000, dt0=1e-3, saveat=SaveAt(steps=True). Tighten tolerances if the background trajectory looks noisy for potentials with sharp features.	{}
perturbation_solver_opts	dict	Override any field of the default SolverOptions for the perturbation integrator. Defaults: rtol=1e-6, atol=1e-6, max_steps=1000000, dt0=1e-3, saveat=SaveAt(t1=True) (only the final mode amplitude is stored, saving memory). Setting saveat=SaveAt(steps=True) retains the full mode history but is considerably slower.	{}
error_on_fail	bool	If True, raise an error when the perturbation integrator fails for a given mode. Default False.	False

Examples:

Ultra-slow-roll model with a near-inflection-point potential:

>>> import jax.numpy as jnp
>>> from sigway.ms_solver import SingleFieldSolver
>>> def V(phi, a, lam, v, nfac):
...     b = (1+nfac)*(1 - a**2/3 + a**2/3*(9/(2*a**2)-1)**(2/3))
...     x = phi/v
...     return lam*v**4/12 * x**2*(6-4*a*x+3*x**2)/(1+b*x**2)**2
>>> solver = SingleFieldSolver(V, phi0=3.0, pi0=0.0, N_CMB_to_end=58.0,
...                            k=jnp.geomspace(1e-5, 10.0, 200))

After construction, call solver.run(k, *params) to obtain a callable interpolant of \(\mathcal{P}_\zeta(k)\), or solver(k, *params) for the raw spectrum array.

__call__

__call__(k, *params)

Compute and return the raw \(\mathcal{P}_\zeta\) array at k.

Identical to run but returns the power-spectrum values directly instead of a callable interpolant. Convenient for quickly plotting or fitting the spectrum at a fixed parameter set.

Parameters:

Name	Type	Description	Default
k	array - like	Comoving wavenumber grid in \(\mathrm{s}^{-1}\).	required
*params	float	Scalar potential parameters forwarded to \(V(\phi, *\mathrm{params})\).	()

Returns:

Type	Description
ndarray	Dimensionful scalar power spectrum \(\mathcal{P}_\zeta(k)\) evaluated at each wavenumber in k.

epsilon_h

epsilon_h(y)

First Hubble slow-roll parameter \(\epsilon_H = \pi^2/2\).

\(\epsilon_H\) measures how quickly the Hubble rate is changing: \(\epsilon_H = -\dot{H}/H^2\). Inflation ends when \(\epsilon_H = 1\).

Parameters:

Name	Type	Description	Default
y	float or array - like	Field velocity \(\pi = \mathrm{d}\phi/\mathrm{d}N\).	required

Returns:

Type	Description
float or ndarray	\(\epsilon_H = \pi^2/2\).

eta_h

eta_h(phi, y, h, params)

Second Hubble slow-roll parameter \(\eta_H\).

\(\eta_H\) characterises the curvature of the inflationary trajectory. In slow roll \(|\eta_H| \ll 1\); large \(|\eta_H|\) signals departures from slow roll (e.g. an inflection point).

Parameters:

Name	Type	Description	Default
phi	float or array - like	Inflaton field value \(\phi\).	required
y	float or array - like	Field velocity \(\pi = \mathrm{d}\phi/\mathrm{d}N\), same shape as phi.	required
h	float or array - like	Rescaled Hubble parameter \(h\), same shape as phi.	required
params	tuple	Extra parameters forwarded to \(V(\phi, *\mathrm{params})\).	required

Returns:

Type	Description
float or ndarray	\(\eta_H = -6 + 2\epsilon_H - U'(\phi)\,\pi / (\epsilon_H\,h^2)\).

n_s

n_s(epsilon_h, eta_h)

Scalar spectral index \(n_s\) in the slow-roll approximation.

A scale-invariant spectrum has \(n_s = 1\); Planck measures \(n_s \approx 0.965\).

Parameters:

Name	Type	Description	Default
epsilon_h	float or array - like	First Hubble slow-roll parameter \(\epsilon_H\).	required
eta_h	float or array - like	Second Hubble slow-roll parameter \(\eta_H\).	required

Returns:

Type	Description
float or ndarray	\(n_s = 1 - 2\epsilon_H - \eta_H\).

p_at_cmb

p_at_cmb(N, phi, y, h, params)

Evaluate the CMB log-likelihood at the pivot scale using slow-roll observables.

Interpolates the background to \(N_{\rm CMB} = N_{\rm end} - N_{\rm CMB\_to\_end}\), computes \([\ln(10^{10}\mathcal{P}_\zeta),\, n_s,\, r]\) in the slow-roll approximation, and returns their multivariate-Gaussian log-probability against the stored CMB priors.

Parameters:

Name	Type	Description	Default
N	array - like	E-fold array from run_background.	required
phi	array - like	Inflaton field values \(\phi(N)\).	required
y	array - like	Field velocity \(\pi(N) = \mathrm{d}\phi/\mathrm{d}N\).	required
h	array - like	Rescaled Hubble parameter \(h(N)\).	required
params	tuple	Extra parameters forwarded to the potential \(V(\phi, *\mathrm{params})\).	required

Returns:

Name	Type	Description
p	float	Multivariate-Gaussian log-likelihood \(\ln\mathcal{L}[\ln(10^{10}\mathcal{P}_\zeta),\, n_s,\, r]\).
params_at_cmb	ndarray	Array \([\ln(10^{10}\mathcal{P}_\zeta),\, n_s,\, r]\) evaluated at the CMB pivot scale.

plot_evolution

plot_evolution(k, params)

Plot the background evolution and scalar power spectrum.

Produces a five-panel figure (e-folds \(N - N_{\rm end}\) on the shared x-axis) showing, from top to bottom:

1. \(\mathcal{P}_\zeta(k)\): slow-roll approximation vs. full Mukhanov-Sasaki result.
2. Slow-roll parameters \(\epsilon_H\) and \(|\eta_H|\).
3. Inflaton field \(\phi\).
4. Field velocity \(-\pi = -\mathrm{d}\phi/\mathrm{d}N\).
5. Rescaled Hubble parameter \(h\).

A vertical dashed line marks the CMB pivot scale; the top panel has a secondary x-axis labelled in \(k\,[\mathrm{s}^{-1}]\).

Parameters:

Name	Type	Description	Default
k	array - like	Comoving wavenumber grid in \(\mathrm{s}^{-1}\) for the full perturbation run.	required
params	tuple	Extra parameters forwarded to the potential \(V(\phi, *\mathrm{params})\).	required

Returns:

Type	Description
Figure	Five-panel figure of the background and perturbation evolution.

plot_potential

plot_potential(params, phi_range=None, n_points=1000, relative=False)

Plot the inflaton potential \(V(\phi)\).

Parameters:

Name	Type	Description	Default
params	tuple	Extra parameters forwarded to \(V(\phi, *\mathrm{params})\).	required
phi_range	tuple of float	(phi_min, phi_max) field range for the plot. Defaults to \((0.1\,\phi_0,\; 2\,\phi_0)\).	None
n_points	int	Number of field-value samples at which to evaluate \(V\). Default 1000.	1000
relative	bool	If True, plot \(V/V(\phi_0)\) vs. \(\phi/\phi_0\) (dimensionless axes). Default False.	False

Returns:

Type	Description
Figure	Figure object containing the potential plot.

pzeta_sr

pzeta_sr(y, h, params)

Slow-roll estimate of the scalar power spectrum \(\mathcal{P}_\zeta\).

Uses the standard slow-roll formula \(\mathcal{P}_\zeta \approx V(\phi_0)\,h^2 / (8\pi^2\,\epsilon_H)\) where \(\epsilon_H = \pi^2/2\). This is fast to evaluate along the whole background trajectory and useful for comparison with the full Mukhanov-Sasaki result.

Parameters:

Name	Type	Description	Default
y	float or array - like	Field velocity \(\pi = \mathrm{d}\phi/\mathrm{d}N\).	required
h	float or array - like	Rescaled Hubble parameter \(h\), same shape as y.	required
params	tuple	Extra parameters forwarded to \(V(\phi, *\mathrm{params})\).	required

Returns:

Type	Description
float or ndarray	Slow-roll approximation to \(\mathcal{P}_\zeta\).

r

r(epsilon_h)

Tensor-to-scalar ratio \(r\) in the slow-roll approximation.

Current CMB upper limits place \(r \lesssim 0.036\) (95% CL).

Parameters:

Name	Type	Description	Default
epsilon_h	float or array - like	First Hubble slow-roll parameter \(\epsilon_H\).	required

Returns:

Type	Description
float or ndarray	\(r = 16\,\epsilon_H\).

run

run(k, *params)

Compute \(\mathcal{P}_\zeta(k)\) and return a callable interpolant.

Runs the full background + perturbation calculation over the wavenumber grid k and fits a cubic-Hermite spline to the result. The returned callable can then be evaluated at any \(k\) inside the original grid. This is the entry point used by SingleFieldPerturbations.

Parameters:

Name	Type	Description	Default
k	array - like	Comoving wavenumber grid in \(\mathrm{s}^{-1}\) at which to solve the Mukhanov-Sasaki equation.	required
*params	float	Scalar potential parameters forwarded to \(V(\phi, *\mathrm{params})\).	()

Returns:

Type	Description
callable	A function P_zeta_interpolation(k_new) that evaluates the cubic-Hermite interpolant of \(\mathcal{P}_\zeta\) at wavenumbers k_new within the original k range.

run_background

run_background(params)

Integrate the inflationary background equations.

Evolves the inflaton \(\phi(N)\) and Hubble parameter \(H(N)\) from the slow-roll attractor at \(\phi_0\) until \(\epsilon_H \geq 1\) (end of inflation) or max_efolds is reached. The returned arrays use the internally rescaled variables: \(h = H/\sqrt{V(\phi_0)/3}\).

Parameters:

Name	Type	Description	Default
params	tuple	Extra parameters forwarded to the potential \(V(\phi, *\mathrm{params})\).	required

Returns:

Name	Type	Description
N	ndarray	E-fold number at each saved integration step.
phi	ndarray	Inflaton field \(\phi(N)\) at each saved step.
y	ndarray	Field velocity \(\pi(N) = \mathrm{d}\phi/\mathrm{d}N\) at each saved step.
h	ndarray	Rescaled Hubble parameter \(h(N) = H(N)/\sqrt{V(\phi_0)/3}\) at each saved step.

run_perturbations

run_perturbations(k, N, phi, y, h, params)

Solve the Mukhanov-Sasaki equation for all modes and return \(\mathcal{P}_\zeta(k)\).

For each wavenumber \(k\) the integration window is \([N_k - N_{\rm sub},\, N_k + N_{\rm sup}]\) where \(N_k\) is the horizon crossing e-fold (\(k = a H\)). All modes are integrated simultaneously using jax.vmap.

Parameters:

Name	Type	Description	Default
k	array - like	Comoving wavenumber grid in \(\mathrm{s}^{-1}\).	required
N	array - like	E-fold array from run_background.	required
phi	array - like	Inflaton field values \(\phi(N)\), same shape as N.	required
y	array - like	Field velocity \(\pi(N) = \mathrm{d}\phi/\mathrm{d}N\), same shape as N.	required
h	array - like	Rescaled Hubble parameter \(h(N)\), same shape as N.	required
params	tuple	Extra parameters forwarded to the potential \(V(\phi, *\mathrm{params})\).	required

Returns:

Type	Description
ndarray	Dimensionful scalar power spectrum \(\mathcal{P}_\zeta(k)\), shape (len(k),).

Raises:

Type	Description
ConsistencyError	If the background trajectory is too short to reach the CMB pivot scale, or if the horizon-crossing e-fold \(N_k\) leaves insufficient room to set Bunch-Davies initial conditions.

run_single_k

run_single_k(k, N, phi, y, h, params)

Solve the Mukhanov-Sasaki equation for a single mode, retaining the full history.

Identical to run_perturbations for one wavenumber, but the complete mode-function trajectory \((\phi, \pi, h, \mathrm{Re}\,\Delta\phi, \mathrm{Re}\,\Delta\phi', \mathrm{Im}\,\Delta\phi, \mathrm{Im}\,\Delta\phi')\) is stored at every adaptive step. Useful for diagnosing oscillations, horizon crossing, or super-horizon freeze-out behaviour.

Parameters:

Name	Type	Description	Default
k	float	Single comoving wavenumber in \(\mathrm{s}^{-1}\).	required
N	array - like	E-fold array from run_background.	required
phi	array - like	Inflaton field values \(\phi(N)\).	required
y	array - like	Field velocity \(\pi(N) = \mathrm{d}\phi/\mathrm{d}N\).	required
h	array - like	Rescaled Hubble parameter \(h(N)\).	required
params	tuple	Extra parameters forwarded to the potential \(V(\phi, *\mathrm{params})\).	required

Returns:

Name	Type	Description
sol	Solution	Full ODE solution; sol.ys has shape (steps, 7) with columns \((\phi, \pi, h, \mathrm{Re}\,\Delta\phi, \mathrm{Re}\,\Delta\phi', \mathrm{Im}\,\Delta\phi, \mathrm{Im}\,\Delta\phi')\) at each saved step.
lograt	float	\(\ln(k/aH)\) at the initial e-fold \(N_{\rm in}\).

Raises:

Type	Description
ConsistencyError	If the background trajectory is too short to reach the CMB pivot scale, or if \(N_k\) leaves insufficient room for Bunch-Davies initial conditions.

Options & errors

sigway.ms_solver.SolverOptions

Bases: namedtuple('SolverOptions', ['rtol', 'atol', 'max_steps', 'dt0', 'saveat'])

Numerical settings for the adaptive ODE integrator (diffrax Tsit5).

One instance is used for the background solver and a separate one for the perturbation solver; both are configured via the background_solver_opts and perturbation_solver_opts arguments of SingleFieldSolver.

Attributes:

Name	Type	Description
rtol	float	Relative error tolerance for the adaptive (PID) step-size controller. Tighter values increase accuracy at the cost of more ODE steps.
atol	float	Absolute error tolerance for the adaptive step-size controller.
max_steps	int	Maximum number of ODE steps allowed before the integrator gives up. Increase this for potentials with sharp features or very long inflation.
dt0	float	Initial step size in e-folds \(\Delta N\).
saveat	SaveAt	Specification of which solution points to store. The default SaveAt(steps=True) retains every adaptive step (needed for the background); SaveAt(t1=True) stores only the final value (used for the perturbation solver to save memory).

sigway.ms_solver.ConsistencyError

Bases: Exception

Raised when the model fails a physical self-consistency check.

Typical causes: the background trajectory produces too few e-folds to reach the CMB pivot scale, or the horizon-crossing e-fold \(N_k\) is so early that there is insufficient room to initialise the Bunch-Davies vacuum \(N_{\mathrm{sub}}\) e-folds before crossing.

Helper functions

sigway.ms_solver.interpolation_inner

interpolation_inner(knew, k, coeff)