API reference

API reference#

DISCO-EB and the interface the CMB code (DISCO-CMB) may use.

DISCO-EB solves the background, the thermal history (RECFAST with reionization) and the linear perturbation hierarchy of each wave mode on the GPU. discoeb.component holds the component model, discoeb.eb_model the species and their equations, discoeb.recombination the thermal history and discoeb.perturbations the layouts, batching and modax launches. The CMB code lives in the separate package discocmb (DISCO-CMB) and turns this output into angular power spectra: line-of-sight sources, grids, transfer functions, projection and lensing. It imports only from this package itself and from the generic numerics of discoeb.spline_interpolation. DISCO-EB does not depend on it.

The contract. A caller asks for named quantities at chosen conformal times of a batch of cosmologies and wave modes. The kernel launch evaluates them from the solver’s dense output and returns those values, never the state history:

  • solve_eb() (cosmologies, k, tau, keep=names, extra=components) returns an EBSolution whose values are (n_cosmologies, len(keep), n_k, n_tau), with each cosmology’s conformal age tau0. Past its own tau0, a cosmology’s trajectories hold their state today.

  • A name is a state element (state_names(), including the thermal block’s lna, thermo_0..2 and the optical depth kappa) or any quantity the model provides (adotoa, opacity, sigma, dgrho, clxg, polter, R …). extra extends the vocabulary: stateless Component subclasses whose providers read the model’s quantities by name and add their own, compiled into the same kernel. The CMB’s line-of-sight sources work this way, so both parts share one component model and nothing of the CMB enters the hierarchy’s right-hand side.

  • Error weights go by state name ({"theta_2": f(k)}), so the caller needs neither the state’s index order nor the launch order.

  • Differentiation: pass DIFFERENTIABLE_PARAMETERS as JAX values, or the nine DIRECTED_NAMES inputs (directed_inputs()) as inputs, and values becomes differentiable (jax.jvp, jax.grad) through forward sensitivities the kernel computes. The derivatives hold tau0 and the reionization redshift fixed.

background_history(), density_coefficients() and conformal_ages() give the background for building grids. A cosmology is a parameter dictionary or an object as_cosmology() returns.

discoeb.solve_eb(cosmologies, k_values, tau, *, keep=None, extra=(), inputs=None, rtol: float = 0.0001, atol: float = 0.0001, first_step: float = 0.01, max_steps: int = 20000, error_weights=None, thermo_weight: float = 1.0) → EBSolution[source]#

Solve the Einstein-Boltzmann system of cosmologies and keep keep at tau.

Like solve_perturbation_history_batch(), this runs one modax Rodas5P launch over every (wave mode, cosmology) pair. Unlike it, the launch records only the quantities named in keep at the requested times, through modax’s save hook. The kernel evaluates them from the dense-output state at each save and stores len(keep) numbers, so the device holds n_k * n_cosmologies * len(tau) * len(keep) values and never the history. Keeping only the comoving curvature perturbation at 1000 times for 128 modes and 128 cosmologies stores 130 MB instead of the 7 GB of the 55-variable history.

Parameters:

  • cosmologies – One cosmology (a parameter dictionary or a cosmology object) or a sequence of them, all with one perturbation layout. The values of DIFFERENTIABLE_PARAMETERS may be JAX arrays, traced or not. The others must be plain numbers.

  • k_values – Wave numbers in 1/Mpc, in any order.

  • tau – Strictly increasing conformal times in Mpc at which to record, starting at or after TAU_START. Every launch starts at TAU_START, so a first time equal to it records the initial state. A cosmology’s trajectory ends at its conformal age tau0, and later times hold its state today. A last time past every tau0 therefore returns the present-day values of the whole batch.

  • keep – Names of the quantities to record: state elements (layout_state_names(), or the recombination block’s lna, thermo_0, thermo_1, thermo_2, kappa) or quantities the joint model provides, such as the comoving curvature perturbation R, the shear sigma, the expansion rate adotoa or the opacity. Default: every element of the perturbation state, which amounts to the hierarchy’s full history.

  • extra – Components to append to the model for the hook: stateless Component subclasses whose providers read the model’s quantities and add their own (the line-of-sight sources of discocmb.cmb). The hook cache tells them apart by class.

  • inputs – The (n_cosmologies, 9) directed inputs of DIRECTED_NAMES (the five background densities and the recombination block’s YHe, Tcmb, H0, Omegab), replacing those the cosmologies imply (directed_inputs()). They are differentiable in their own right: the solution becomes a function of them, at fixed background endpoints of the cosmologies.

  • rtol – Solver settings, as in solve_perturbation_history_batch(). error_weights is either {state name: weight} or modax’s array (n_k * n_cosmologies, nvar) in launch order (wave modes ascending, cosmologies within). In the dictionary, each weight is a number or a function of the trajectories’ wave numbers that returns one per trajectory, and unnamed states keep weight one.

  • atol – Solver settings, as in solve_perturbation_history_batch(). error_weights is either {state name: weight} or modax’s array (n_k * n_cosmologies, nvar) in launch order (wave modes ascending, cosmologies within). In the dictionary, each weight is a number or a function of the trajectories’ wave numbers that returns one per trajectory, and unnamed states keep weight one.

  • first_step – Solver settings, as in solve_perturbation_history_batch(). error_weights is either {state name: weight} or modax’s array (n_k * n_cosmologies, nvar) in launch order (wave modes ascending, cosmologies within). In the dictionary, each weight is a number or a function of the trajectories’ wave numbers that returns one per trajectory, and unnamed states keep weight one.

  • max_steps – Solver settings, as in solve_perturbation_history_batch(). error_weights is either {state name: weight} or modax’s array (n_k * n_cosmologies, nvar) in launch order (wave modes ascending, cosmologies within). In the dictionary, each weight is a number or a function of the trajectories’ wave numbers that returns one per trajectory, and unnamed states keep weight one.

  • error_weights – Solver settings, as in solve_perturbation_history_batch(). error_weights is either {state name: weight} or modax’s array (n_k * n_cosmologies, nvar) in launch order (wave modes ascending, cosmologies within). In the dictionary, each weight is a number or a function of the trajectories’ wave numbers that returns one per trajectory, and unnamed states keep weight one.

  • thermo_weight – Solver settings, as in solve_perturbation_history_batch(). error_weights is either {state name: weight} or modax’s array (n_k * n_cosmologies, nvar) in launch order (wave modes ascending, cosmologies within). In the dictionary, each weight is a number or a function of the trajectories’ wave numbers that returns one per trajectory, and unnamed states keep weight one.

Notes

The hook compiles once per batch shape and set of names, like the right-hand side. The returned values are differentiable in the cosmologies’ Omegam, Omegab, H0, Tcmb and Neff under jax.jvp, jax.jacfwd and jax.grad. They are not under jit, since the host finds the background endpoints from the concrete values. The derivative runs through the five background densities and the recombination block’s Tcmb, H0 and Omegab. The forward-sensitivity launch of _kept_values() carries it, with a hook that stores each kept quantity’s derivative beside it, so it also costs one joint launch and no history. It is the derivative at fixed tau0 and reionization redshift.

class discoeb.EBSolution(names: tuple[str, ...], k: ndarray, tau: ndarray, tau0: ndarray, values: Array, layout: PerturbationLayout)[source]#

Bases: object

What solve_eb() kept: the selected quantities on the requested grid.

values is a device array (n_cosmologies, len(names), n_k, n_tau) with wave modes and times in the caller’s order. solution[name] is that quantity’s (n_cosmologies, n_k, n_tau) block. tau0 is each cosmology’s conformal age, past which its column holds the state today. values is differentiable in the parameters of DIFFERENTIABLE_PARAMETERS given as JAX values.

discoeb.state_names(cosmology) → tuple[str, ...][source]#

Return the perturbation state names of a cosmology’s layout, for keep and error_weights.

discoeb.conformal_ages(cosmologies) → ndarray[source]#

Return each cosmology’s conformal age tau0 (Mpc), where its trajectories stop.

Reads the same cached background endpoints as solve_eb(), so a solve of the same cosmologies right after reuses them.

discoeb.background_history(cosmology)[source]#

Return (a, tau) of one cosmology’s background quadrature, host arrays ending at tau0.

discoeb.density_coefficients(cosmology) → tuple[float, float, float, float, float][source]#

Return (grhog, grhornomass, grhoc, grhob, grhov) of discoeb.eb_model.density_parameters().

discoeb.directed_inputs(cosmologies) → ndarray[source]#

Return the (n_cosmologies, 9) directed inputs of DIRECTED_NAMES the cosmologies imply.

These are the five background densities and the recombination block’s YHe, Tcmb, H0 and Omegab: the variables solve_eb() takes as inputs and differentiates with respect to.

discoeb.as_cosmology(cosmology)#

Normalise a cosmology given as a parameter dictionary (Omegam, Omegab, H0, Tcmb, YHe, Neff, A_s, n_s, … ) to the solver’s cosmology object. Returns one already normalised as it is.

discoeb.DIFFERENTIABLE_PARAMETERS#

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable’s items.

If the argument is a tuple, the return value is the same object.

discoeb.DIRECTED_NAMES#

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable’s items.

If the argument is a tuple, the return value is the same object.