Reference¶

cosmotile¶

Cosmotile.

class cosmotile.SliceInterpolator(*, coordinates, order, weights=None, origin=None)¶

A coeval box interpolated onto one fixed set of shell coordinates.

Built by make_lightcone_slice_interpolator(), and callable on any box of the right shape: the coordinates, order and sub-sample weights are fixed at construction, so tiling many fields onto the same shell costs one geometry setup, not one per field.

This was a functools.partial() before version 2.0, with the origin stapled on as an attribute. It is a class now so that the contract is inspectable and typed, but keywords and origin are kept so that code written against the old object keeps working.

Parameters:

coordinates (np.ndarray)
order (int)
weights (np.ndarray | None)
origin (np.ndarray | None)

property keywords: dict[str, Any]¶: The interpolation parameters, as functools.partial() used to expose them.

origin¶: The shell centre, kept because it is what turns pixel coordinates back into lines of sight – see make_lightcone_slice_vector_field().

cosmotile.apply_rsds(field, los_displacement, distance, n_subcells=4, *, outside='edge')¶

Apply redshift-space distortions to a field.

Notes

To ensure that we cover all the slices in the field after the velocities have been applied, the grid is padded on either end by the largest displacement that can reach it. What that padding holds is outside’s business – see below. Then, to ensure we don’t pick up cells with zero particles (after displacement), we interpolate the slices onto a finer regular grid (in comoving distance) and then displace the field on that grid.

The displaced fine grid is finally integrated over the radial extent of each output slice, rather than sampled at its centre. Every parcel therefore lands in exactly one output cell (in the proportions in which it straddles them), mass is conserved, and n_subcells is a real convergence parameter: raising it shrinks the cloud-in-cell kernel without anything falling between the output slices.

Convention. Unlike the tiling functions, this one works throughout in terms of radial cell averages: a slice is a cell running from edge to edge, and its value is the mean of the field over that cell. That is what makes the displacement mass-conserving and a zero displacement an exact round trip, and it is not optional – moving material around is only meaningful for a quantity with extent.

That is consistent with a lightcone built by point-sampling shells provided the slice spacing is comparable to the cell size of the coeval box, because the box’s own cell window then already supplies about one slice of radial smoothing (the cubic cell window is within 3% of isotropic even at Nyquist, so it acts as a radial top-hat of one cell). If your slices are much coarser than your cells, a point-sampled lightcone is not a radial cell average and this function will assume smoothing that is not there: build the lightcone with radial_width first, using residual_radial_width() to choose it.

Parameters:

field (ndarray) – The field to apply redshift-space distortions to, shape (nslices, ncoords). Taken as the mean over each radial slice, and returned the same way – see the convention note above.
los_displacement (ndarray) – The line-of-sight “apparent” displacement of the field, in pixel coordinates. Equal to v / H(z) / cell_size. Positive values are towards the observer, shape (nslices, ncoords). This is the same sign convention as the output of make_lightcone_slice_vector_field(), so the two can be chained directly. A parcel at comoving distance d with displacement u is observed at apparent distance d - u.
distance (ndarray) – The comoving distance to each slice in the field, in units of the cell size. shape (nslices,).
n_subcells (int) – The number of sub-cells per (smallest) output slice used for the displacement. Larger values resolve the displacement field more finely and give a more accurate answer, at proportionally greater cost.
outside (Literal['empty', 'edge']) –
What the field does beyond the range distance covers.

"edge" (the default) continues the field at its first and last slice values, moving with the boundary displacement. Material flows in as well as out, so on average nothing is lost. This is usually what you want: a set of slices is normally a window cut out of a larger field.

"empty" takes the field to be zero outside, so material displaced off either end is gone and the total can only fall. Choose it when your slices really are the whole field.

Returns:

The field in redshift space, same shape as field, again as radial cell averages.

Return type:

distorted

cosmotile.cell_window(shape, width=1.0, rfft=False)¶

Fourier response of a top-hat of width cells, on the mode grid of shape.

A simulation cell almost always holds the mean of the field over that cell, which is the same thing as a point sample of the field convolved with a top-hat of one cell on a side. In Fourier space that is

\[\tilde{T}(\mathbf{k}) = \prod_i \mathrm{sinc}(k_i w / 2),\]

which is what this returns, broadcast over the axes of shape rather than materialised. See deconvolve_cell_window() for why you might want it.

Parameters:

shape (Sequence[int]) – Shape of the grid whose Fourier modes the window is evaluated on.
width (float) – Width of the top-hat, in cells. One cell by default.
rfft (bool) – Whether the last axis is laid out for numpy.fft.rfftn() (half-spectrum) rather than numpy.fft.fftn().

Return type:

ndarray

cosmotile.deconvolve_cell_window(coeval, width=1.0)¶

Undo the cell-averaging built into a gridded simulation.

A coeval box almost always holds cell averages, g_n = (T * f)(x_n), where f is the underlying field and T a top-hat one cell wide. That single number is both “the mean of f over the cell” and “a point sample of the smoothed field s = T * f” – the same quantity under two names. Everything else in cosmotile reconstructs and evaluates whatever the box holds, so by default a lightcone carries T whether you want it or not.

This divides T out, turning cell averages into point samples of f. It is exact for a periodic field band-limited to the grid Nyquist, and the output window of an averaged lightcone is then only the window you asked for, rather than that window convolved with T.

You usually do not want this. The cell average is the field at the resolution the simulation actually has; deconvolving sharpens structure it never resolved and amplifies whatever aliased power and noise sit near Nyquist (a factor of 1 / prod_i sinc(k_i / 2), which reaches 3.9 at the Nyquist corner of a 3D grid, and about 1.6 in RMS on a white-noise box). Reach for it only when you need a lightcone whose window is exactly a quantity you have specified – for instance when comparing a line-of-sight power spectrum against P(k) times a known channel response.

It is a standalone function rather than a keyword on the interpolator because it need only be done once per box: applying it per shell would repeat it for every slice of a lightcone, in the same way the spline pre-filter used to be before prefilter_coeval(). If you use both, deconvolve first and pre-filter the result – the pre-filter is computed from whatever values it is handed.

Parameters:

coeval (ndarray) – The gridded field, of any dimensionality. Assumed periodic, as everywhere else.
width (float) – Width of the cell top-hat, in cells. One cell by default; must be less than two, since sinc(k w / 2) has its first zero at k = pi, w = 2 and a wider window is not invertible on the grid.

Returns:

Point samples of the underlying field, same shape and dtype-kind as the input.

Return type:

deconvolved

cosmotile.get_distance_to_shell_from_redshift(z, cell_size, cosmo=FlatLambdaCDM(name='Planck18', H0=<Quantity 67.66 km / (Mpc s)>, Om0=0.30966, Tcmb0=<Quantity 2.7255 K>, Neff=3.046, m_nu=<Quantity [0., 0., 0.06] eV>, Ob0=0.04897))¶

Get a distance to a shell, in units of cell size, from a given redshift.

Parameters:

z (float) – The redshift
cell_size (Annotated[Quantity, PhysicalType('length')]) – The resolution of the coeval simulation, in comoving units.
cosmo (FLRW) – The astropy cosmology.

Returns:

The distance, in units of pixels, to the shell.

Return type:

distance

cosmotile.healpix_subpixel_lonlat(nside, order='ring', subsample_level=0)¶

Angular coordinates of the HEALPix pixels, optionally sub-divided.

With subsample_level = 0 this returns the (npix,) pixel centres. With subsample_level = k it returns (4**k, npix) arrays holding the centres of the 4**k sub-pixels of side nside * 2**k that exactly tile each pixel. Since HEALPix sub-pixels are equal-area and tile their parent exactly, the unweighted mean over them is an unbiased estimate of the mean of the field over the pixel.

Parameters:

nside (int) – The Nside parameter of the healpix map.
order (Literal['ring', 'nested']) – The ordering of the pixels in the healpix map.
subsample_level (int) – How many times to halve the pixel side. Each level costs four times as many interpolations.

Returns:

In radians, of shape (npix,) or (4**k, npix).

Return type:

latitude, longitude

cosmotile.make_healpix_lightcone_slice(nside, order='ring', subsample_level=0, **kwargs)¶

Create a healpix lightcone slice in angular coordinates.

This is a simple wrapper around make_lightcone_slice() that sets up angular co-ordinates from a healpix grid.

Parameters:

nside (int) – The Nside parameter of the healpix map.
order (Literal['ring', 'nested']) – The ordering of the pixels in the healpix map.
subsample_level (int) –
If zero (the default), each pixel takes a single sample of the reconstructed field at its centre, which is the historical behaviour. Such a map carries the input box’s own cell window but not the HEALPix pixel window, so pixwin must not be divided out of its angular power spectrum.

If k > 0, each pixel is instead averaged over the 4**k sub-pixels of a map with nside * 2**k. The pixel window then does apply on top of the cell window, and power above l ~ 2 Nside is suppressed rather than aliased down. Cost grows as 4**k; k = 2 is usually enough.
kwargs (Any)

Return type:

Generator

:param All other parameters are passed through to make_lightcone_slice().:

cosmotile.make_lightcone_slice(*, coevals, **kwargs)¶

Create a lightcone slice in angular coordinates from two coeval simulations.

Interpolates the input coeval box to angular coordinates.

Whatever a cell of coevals holds is what comes out. A simulation cell almost always holds the mean of the field over that cell, so the output carries that cell window – plus the reconstruction kernel set by interpolation_order, plus any output window you request. See make_lightcone_slice_interpolator() for the full chain, and deconvolve_cell_window() if you need the cell window gone.

Parameters:

coevals (Sequence[ndarray] | ndarray | PrefilteredCoeval) – An iterable of rectangular coeval simulations to interpolate to the angular coordinates. Must have three dimensions (not necessarily the same size). Each box must have the same shape, and all are assumed to be at the same coordinates. Each coeval box can be a different simulated field. If you are tiling the same box onto many shells at interpolation_order >= 2, pass boxes pre-filtered by prefilter_coeval(), so that the spline pre-filter is computed once rather than once per shell.
kwargs (Any)

Return type:

Iterator[ndarray]

:param All other parameters are passed to make_lightcone_slice_interpolator().:

Yields:

field – Each interpolated field on the angular coordinates.

Parameters:

coevals (Sequence[ndarray] | ndarray | PrefilteredCoeval)
kwargs (Any)

Return type:

Iterator[ndarray]

cosmotile.make_lightcone_slice_interpolator(*, latitude, longitude, distance_to_shell, interpolation_order=1, origin=None, rotation=None, radial_width=1.0, coeval_cell_width=1.0, n_radial_samples=1)¶

Create a callable interpolator for a lightcone slice.

By default each output pixel is a single point sample of the reconstructed field, at the pixel centre and at exactly the shell radius. Both directions can instead be averaged over the extent the pixel really subtends – see latitude and radial_width below. Averaging costs one interpolation per sub-sample, so it is opt-in; nothing changes unless you ask for it.

What the output value is, precisely, is Q * L * T * f evaluated at the pixel:

f is the underlying continuous field;
T is the cell window the input box already carries, because a simulation cell holds the mean over that cell. It is always there. Remove it beforehand with deconvolve_cell_window() if – and only if – you need it gone;
L is the reconstruction kernel, set by interpolation_order;
Q is the output window: nothing by default, the sub-samples given in latitude and/or a radial top-hat of radial_width if you ask.

Because T never leaves, radial_width is the total window you want and only the remainder is applied on top of the box’s own. See “What a cell holds, and what comes out” in the accuracy documentation.

Parameters:

latitude (ndarray) –
An array of latitude coordinates onto which to tile the box. In radians from -pi/2 to pi/2.

May instead be two-dimensional, with shape (n_subsamples, npix). The box is then interpolated at every sub-sample and the results averaged, so that each output pixel is an average over the directions given rather than a sample at one of them. make_healpix_lightcone_slice() builds such an array from HEALPix sub-pixels.
longitude (ndarray) – An array, same shape as latitude, of longitude coordinates onto which to tile the box. In radians from 0 to 2pi.
distance_to_shell (float) – The distance to the spherical shell onto which to interpolate, in units of the cell-size of the coeval box(es) you wish to interpolate.
interpolation_order (int) – The order of interpolation. Must be in the range 0-5.
origin (ndarray | tuple[float, float, float] | None) – Define the location of the centre of the spherical shell, assuming that the (0,0,0) pixel of the coeval box is at (0,0,0) in cartesian coordinates.
rotation (Rotation | None) – The rotation by which to rotate the spherical coordinates before interpolation. This is done before shifting the origin, and is equivalent to rotating the coeval box beforing tiling it.
radial_width (float) –
The total radial top-hat you want the output to carry, in units of the cell size – in a lightcone, the slice spacing or the width of a frequency channel.

The box already supplies coeval_cell_width of radial smoothing, so only the remainder is applied here: the extra top-hat is sqrt(radial_width**2 - coeval_cell_width**2), since top-hat widths add in quadrature to leading order (residual_radial_width()). The default of 1.0 is therefore a no-op on a cell-averaged box – you already have a one-cell window. Values below coeval_cell_width are rejected: averaging cannot sharpen, and no lightcone can resolve better than the box it came from.

Only used when n_radial_samples > 1.
coeval_cell_width (float) – The radial top-hat the input box already carries, in units of the cell size. One by default, since a simulation cell holds the mean over that cell. Pass 0.0 if your box genuinely holds point samples, or if you have removed the cell window with deconvolve_cell_window().
n_radial_samples (int) –
The number of Gauss-Legendre nodes used to average over the extra width, weighted by the r^2 volume element – the volume element is always applied, since it is what makes the result the mean over the shell. The default of 1 samples the shell radius itself and does no averaging at all.

The nodes interpolate the coeval box at their own radii, so this genuinely averages the input cells that fall in the window rather than rescaling by a window function. Gauss-Legendre with n nodes integrates polynomials of degree 2n - 1 exactly, so a window spanning many cells of a field with power near Nyquist needs roughly n >= pi * width / 2 nodes; four is ample for a window of a cell or two.

Returns:

A callable that takes a 3D array of coeval values and returns a 2D array of interpolated values on a redshift slice.

Return type:

interpolator

cosmotile.make_lightcone_slice_vector_field(coeval_vector_fields, interpolator)¶

Interpolate a 3D vector field to a lightcone slice as a line-of-sight component.

This takes a sequence of 3D vector fields, eg. the velocity field, and interpolates each component to the lightcone slice. It then computes the line-of-sight component of each interpolated vector field, where positive values are oriented towards the observer.

Parameters:

coeval_vector_fields (Sequence[Sequence[ndarray]]) – An iterable of 3D vector fields to interpolate to the lightcone slice. Each vector field must be an iterable of 3 3D arrays, each of the same shape.
interpolator (SliceInterpolator) – A callable that takes a 3D array of coeval values and returns a 2D array of interpolated values on a redshift slice. This should be created by make_lightcone_slice_interpolator() using the properties of the coeval vector fields.

Yields:

los_component – The line-of-sight component of each interpolated vector field.

Return type:

Iterator[ndarray]

cosmotile.prefilter_coeval(coeval, order)¶

Apply the spline pre-filter to a coeval box once, for re-use across many shells.

Interpolating at order >= 2 requires the coeval box to be converted to B-spline coefficients first (see _interpolate_coeval()). Those coefficients depend only on the box and the order – not on the shell radius, rotation or origin – so for a lightcone of many shells the filter need only be computed once.

filtered = cosmotile.prefilter_coeval(coeval, order=3)
for radius in radii:
    (shell,) = cosmotile.make_lightcone_slice(
        coevals=filtered,
        latitude=lat,
        longitude=lon,
        distance_to_shell=radius,
        interpolation_order=3,
    )

Parameters:

coeval (ndarray) – The coeval box to pre-filter.
order (int) – The interpolation order the box is being prepared for. Must be in the range 0-5, and must match the interpolation_order it is later tiled with. Orders 0 and 1 use interpolating kernels and need no filter, so for them this only tags the box.

Returns:

The pre-filtered box, tagged with order. Pass it to make_lightcone_slice() (or make_lightcone_slice_interpolator()) in place of the raw box; no further flag is needed, since the tag is what tells the interpolator the filter has already been applied.

Return type:

prefiltered

Notes

The result of tiling a pre-filtered box is bit-identical to tiling the raw box at the same order; this only moves the work out of the per-shell loop.

Mutating coeval after calling this does not update the returned array for order > 1 (it is a fresh array); re-run this function if the box changes.

cosmotile.recommended_subsample_level(nside, distance_to_shell, cell_size=1.0, tolerance=0.01)¶

Choose subsample_level for a wanted accuracy on the pixel average.

Averaging over 4**k sub-pixels is a quadrature rule for the mean of the field over the pixel, so its error is set by how finely the sub-pixels sample the scale on which the field varies – the cell size. Empirically the fractional error on the angular power is about 0.15 (a / D)^2 for a sub-pixel arc a and cell size D, which inverts to the k returned here.

Note that this does not saturate once the sub-pixels reach the cell size. The reconstructed field is smooth, not structureless below a cell, so a finer rule keeps paying: measured errors run 6% at a sub-pixel arc of 0.65 cells, 1.3% at 0.33 and 0.3% at 0.16.

Parameters:

nside (int) – The Nside parameter of the healpix map.
distance_to_shell (float) – Shell radius, in cells – a HEALPix pixel subtends roughly 0.52 / nside radians, so its arc is about 0.52 r / nside cells.
cell_size (float) – Cell size of the coeval box, in cells. One by default.
tolerance (float) – Wanted fractional accuracy on the pixel average.

Returns:

The subsample_level to pass to make_healpix_lightcone_slice(). Zero when the pixel is already small enough that its centre is a good enough estimate of its mean.

Return type:

level

Notes

A large answer is telling you something: cost grows as 4**level, and needing more than two or three means the pixel is many cells across, i.e. nside is too coarse to resolve the box at this radius in the first place. Raising nside is then both cheaper and more useful than averaging a huge pixel very accurately.

cosmotile.transform_to_pixel_coords(*, comoving_radius, latitude, longitude, origin=None, rotation=None)¶

Transform input spherical coordinates to pixel coordinates wrt a coeval box.

Parameters:

comoving_radius (Annotated[Quantity, Unit("pix")]) – The radius of the spherical coordinates (in units of the cell size).
latitude (ndarray) – An array of latitude coordinates onto which to tile the box. In radians from -pi/2 to pi/2
longitude (ndarray) – An array, same size as latitude, of longitude coordinates onto which to tile the box. In radians from 0 to 2pi.
origin (Annotated[Quantity, Unit("pix")] | tuple[float, float, float] | None) – Define the location of the centre of the spherical shell, assuming that the (0,0,0) pixel of the coeval box is at (0,0,0) in cartesian coordinates. In units of the cell size.
rotation (Rotation | None) – The rotation by which to rotate the spherical coordinates before interpolation. This is done before shifting the origin, and is equivalent to rotating the coeval box beforing tiling it.

Return type:

ndarray

cosmotile.theory¶

Closed-form predictions for the statistics of a tiled shell.

cosmotile cuts a spherical shell out of a periodically-tiled coeval box. Given the three-dimensional power spectrum of the coeval box, the angular power spectrum of the shell is fixed, and this module evaluates it.

Three functions are provided.

continuum_angular_power(): What the shell would have if the box were infinite,

\[C_\ell = \frac{2}{\pi} \int \mathrm{d}k\, k^2 P(k) j_\ell^2(kr).\]
discrete_angular_power(): What a periodic box of finite size actually gives, since it contains only the discrete modes \(\mathbf{k} = 2\pi\mathbf{j}/L\),

\[C_\ell = \frac{4\pi}{V} \sum_{\mathbf{k}} P(k) j_\ell^2(kr).\]
interpolation_window(): The Fourier-space response \(W(\mathbf{k})\) of the spline interpolation that cosmotile uses to reconstruct the field between grid points, which multiplies the box modes before they are summed.

The difference between the first two is the power a finite box is missing; the third is the power the interpolation suppresses at the small-scale end. Together they bracket the window of validity documented in Accuracy and Limitations.

This is also how to check that cosmotile is working. Measure the angular power spectrum of a shell you have actually tiled – one realisation, or better an average over several – and compare it against discrete_angular_power() with interpolation_window() squared folded into the weights. That is the prediction the tiling is supposed to reproduce, and inside the window of validity it does so to a few percent. Comparing instead against continuum_angular_power() mixes in the finite box’s missing power, and comparing against P(l / r) / r**2 is the Limber approximation, which a geometrically thin shell does not have.

Notes

Everything here works in cell units: the coeval box is n cells on a side, the cell size is unity, so the box length is L = n, the volume is V = n**3, wavenumbers are k = 2 pi j / n and distances to shells are in cells.

The Fourier convention is

\[\delta(\mathbf{x}) = \sum_{\mathbf{k}} \delta_{\mathbf{k}} e^{i \mathbf{k}\cdot\mathbf{x}}, \qquad \langle |\delta_{\mathbf{k}}|^2 \rangle = P(k) / V,\]

which is the convention in which the angular power spectrum of a thin shell takes the familiar form above. It is also powerbox’s default (a = b = 1, vol_normalised_power=True) once boxlength is given in cells.

These are exact statements about a geometrically thin shell; none of them is the Limber approximation, which has no thin-shell limit because it needs a radial kernel of finite width to integrate over.

cosmotile.theory.continuum_angular_power(pk, radius, ells, kmax, nk=40000)¶

Evaluate the continuum integral C_l = (2/pi) int dk k^2 P(k) j_l^2(kr).

This is the infinite-box limit of discrete_angular_power(): what a tiled box would give if it contained every mode. The difference between the two is exactly the large-scale power a finite box is missing.

Parameters:

pk (Callable[[ndarray], ndarray]) – The three-dimensional power spectrum, a callable of the angular wavenumber k in inverse cell units.
radius (float) – Shell radius, in cells.
ells (ndarray) – Multipoles at which to evaluate.
kmax (float) – Upper limit of the integral, in inverse cell units. For a band-limited spectrum set this to the cut-off; otherwise take it well above max(ells) / radius.
nk (int) – Number of trapezoidal integration points. The integrand oscillates like j_ell^2, so this must comfortably resolve kmax * radius / pi periods.

Returns:

The angular power spectrum, one value per entry of ells.

Return type:

cosmotile.jax¶

The optional JAX backend. See Differentiable and GPU tiling for what it is for and how to use it.

A JAX backend for cosmotile.

Requires jax, an optional dependency: pip install cosmotile[jax]. The NumPy API in cosmotile never imports it.

The functions here take arrays and return arrays, so you can wrap them in jax.jit(), jax.vmap() or jax.grad(). They compute the same thing as the NumPy API, which returns lazy iterators and carries astropy units – neither of which survives being traced – so this is a smaller, flatter surface rather than a mirror of it.

prefilter_coeval(): Convert a box to B-spline coefficients. Needed once, before tiling at order 2 or above.
shell(): Interpolate a box onto one spherical shell.
shell_from_coordinates(): The same, when you already have the pixel coordinates.
lightcone_scan(): Build many shells one at a time, accumulating a result instead of keeping them all.
apply_rsds(): Apply redshift-space distortions.

See the backend guide for what it is for, what you can differentiate, and how to keep a large lightcone in memory.

class cosmotile.jax.PrefilteredCoeval(coefficients, order)¶

A coeval box converted to B-spline coefficients for one interpolation order.

Produced by prefilter_coeval() (or its JAX counterpart). The order is carried explicitly rather than hidden on the array, so the same object works for a NumPy array, a JAX array or anything else, and so a static type checker can see it.

The tag is what tells make_lightcone_slice() to skip the filter it would otherwise apply, and what lets it reject a box filtered for the wrong order.

Deliberately not an array: it supports neither arithmetic nor indexing, because the pre-filter of a slice is not the slice of the pre-filter and the pre-filter of twice a box is not twice the pre-filter. Both would be silently wrong, so both are a loud TypeError. Use numpy.asarray() to get the coefficients out, and re-filter the result if you derive a new box from them.

Parameters:

coefficients (Any)
order (int)

coefficients: Any¶: The B-spline coefficients. For orders 0 and 1 these are the box values unchanged.

property dtype: Any¶: Dtype of the underlying coefficients.

property ndim: int¶: Number of dimensions of the underlying coefficients.

order: int¶: The interpolation order the box was filtered for.

property shape: tuple[int, ...]¶: Shape of the underlying coefficients.

property spline_order: int¶: Deprecated alias for order.

Deprecated since version 2.0: Use order.

class cosmotile.jax.RsdPlan(fine_widths, refine_index, source_mask, interp_index, interp_weight, rebin_index, rebin_frac, out_widths, fine_cumulative, nslice, n_near, n_far, n_subcells, outside)¶

Everything about an RSD application that does not depend on the field values.

Built by make_rsd_plan(). Reusable across as many fields as share a radial grid, which for a lightcone is all of them.

Parameters:

fine_widths (Any)
refine_index (Any)
source_mask (Any)
interp_index (Any)
interp_weight (Any)
rebin_index (Any)
rebin_frac (Any)
out_widths (Any)
fine_cumulative (Any)
nslice (int)
n_near (int)
n_far (int)
n_subcells (int)
outside (str)

fine_cumulative: Any¶: (nfine,) cumulative width up to each sub-cell’s near edge.

fine_widths: Any¶: (nfine,) width of each sub-cell, in the units of distance.

interp_index: Any¶: (nfine, k) radial stencil for the displacement. k is 2 on a regular grid, and nslice on an irregular one, where the spline is global.

interp_weight: Any¶: (nfine, k) matching stencil weights.

n_near: int¶: Sub-cells of padding at the near and far ends.

n_subcells: int¶: Sub-cells per output slice.

nslice: int¶: Number of output slices.

out_widths: Any¶: (nslice,) width of each output slice.

outside: str¶

"empty" or "edge".

Type:: What the field does beyond the grid

rebin_frac: Any¶: (nslice + 1,) how far across that sub-cell it falls.

rebin_index: Any¶: (nslice + 1,) sub-cell each output edge falls in, for the re-binning.

refine_index: Any¶: (nfine,) index of the input slice each sub-cell takes its value from.

source_mask: Any¶

zero for outside="empty" and one for outside="edge". One inside the grid either way.

Type:: (nfine,) how much of the padding is real field

class cosmotile.jax.ShellSampling(directions, nodes, node_weights, half_width, npix, n_angular, order)¶

Everything about a shell’s geometry that does not depend on its radius.

Built by make_shell_sampling(). One of these is shared by every shell of a lightcone: the unit vectors are the expensive part and they are radius-independent, so a thousand-shell lightcone carries one of these plus a thousand scalars, rather than a thousand coordinate arrays.

The integer fields are kept separate from the array fields because cosmotile.jax registers this as a pytree with the integers as static metadata – order in particular cannot be a tracer, since the gather unrolls over it.

Parameters:

directions (Any)
nodes (Any)
node_weights (Any)
half_width (float)
npix (int)
n_angular (int)
order (int)

directions: Any¶: (3, n_angular * npix) unit vectors, angular-sub-sample-major.

half_width: float¶: Half the extra radial top-hat, in cells. Zero for point sampling.

n_angular: int¶: Number of angular sub-samples per output pixel.

node_weights: Any¶: The matching quadrature weights, before the r^2 volume element.

nodes: Any¶: Gauss-Legendre nodes on [-1, 1] across the shell’s radial extent.

npix: int¶: Number of output pixels.

order: int¶

the gather unrolls over it.

Type:: The interpolation order. Static

cosmotile.jax.apply_rsds(field, los_displacement, plan)¶

Apply redshift-space distortions to a field, on a fixed plan.

Computes the same thing as cosmotile.apply_rsds(), but with every shape fixed by plan rather than by the data, so it can be jitted, vmapped and differentiated. It is differentiable in both arguments: linearly in field, and piecewise-linearly in los_displacement through the cloud-in-cell kernel.

Parameters:

field (jax.Array) – (nslice, nangles) radial cell averages – see the convention note on cosmotile.apply_rsds(), which applies here unchanged.
los_displacement (jax.Array) – (nslice, nangles) apparent displacement in cells, positive towards the observer, matching cosmotile.make_lightcone_slice_vector_field().
plan (RsdPlan) – From cosmotile._rsd.make_rsd_plan(), built for this radial grid and a bound on the displacement.

Returns:

(nslice, nangles), again as radial cell averages. Mass that leaves the padded grid is gone, as it should be.

Return type:

distorted

cosmotile.jax.lightcone_scan(coeval, sampling, radii, body, init, *, remat=False, **kwargs)¶

Build the shells of a lightcone one at a time, accumulating a result.

A whole lightcone rarely fits in GPU memory – a thousand shells at nside=256 is 3.1 GB, and building all their coordinates at once needs a further 19 GB. This makes one shell at a time and hands it to body, which combines it with a running result and returns the updated one. Only a single shell is ever in memory.

Use it when what you want out is a summary of the lightcone rather than the lightcone itself – a likelihood, a power spectrum, a sum of squares, the maximum brightness – which is the usual case when you are differentiating through it. If you want the shells themselves and they fit, just call shell() in a loop.

Parameters:

coeval (Array | PrefilteredCoeval) – A PrefilteredCoeval for order >= 2, else a 3D array.
sampling (ShellSampling) – The shared, radius-independent geometry.
radii (Array) – (nshell,) shell radii in cells.
body (Callable[[Carry, Array], Carry]) – Called as body(result, shell) for each shell, and must return the updated result. Both must have the same shape and dtype every time.
init (Carry) – The starting value of the result, before any shell has been seen.
remat (bool) – Recompute each shell during the backward pass rather than storing it. Only worth it when body is expensive; the interpolation itself stores almost nothing.
**kwargs (Any) – Passed to shell() – rotation and origin.

Returns:

What body returned after the last shell.

Return type:

result

Examples

The mean squared brightness of every shell in a lightcone, differentiated with respect to the box it came from:

import jax, jax.numpy as jnp
from cosmotile import jax as cjax

radii = jnp.linspace(100.0, 400.0, 1000)

def total_power(box):
    coeff = cjax.prefilter_coeval(box, order=3)
    return cjax.lightcone_scan(
        coeff, sampling, radii, lambda acc, shell: acc + jnp.mean(shell**2), 0.0
    )

gradient = jax.grad(total_power)(box)

cosmotile.jax.make_rsd_plan(distance, *, n_subcells=4, max_displacement, outside='edge')¶

Precompute everything about an RSD application that the field values do not set.

Parameters:

distance (ndarray) – (nslice,) comoving distance to each slice, in cells. Plain numbers, not a Quantity: units do not survive a traced computation, so strip them before you get here.
n_subcells (int) – Sub-cells per output slice used for the displacement. A genuine convergence parameter: raising it shrinks the cloud-in-cell kernel.
max_displacement (float) –
The largest line-of-sight displacement the plan must accommodate, in cells, and the reason a JAX version is possible: it replaces the original’s measurement of the actual displacement, which made the padded shapes depend on the data.

Pass a physical bound – max |v| / H / cell_size – or, if you have the field already, float(np.abs(los_displacement).max()). Too large merely wastes memory.

With outside="empty" too small a bound does not corrupt the interior: material simply leaves the grid sooner than it would have, which is already what happens at the true edges. With outside="edge" it does matter, because the padding is also where material flows in from – make it comfortably larger than the displacement at the first and last slices.
outside (Literal['empty', 'edge']) – What the field does beyond the range distance covers: "edge" (the default) continues it at the first and last slice values, "empty" takes it to be zero. See cosmotile.apply_rsds(), whose keyword this mirrors.

Returns:

Pass to cosmotile.jax.apply_rsds().

Return type:

plan

cosmotile.jax.make_shell_sampling(*, latitude, longitude, order=1, radial_width=1.0, coeval_cell_width=1.0, n_radial_samples=1)¶

Build the radius-independent half of a shell’s geometry.

The arguments mirror make_lightcone_slice_interpolator(), minus everything that varies from shell to shell (radius, rotation, origin), so that one of these serves a whole lightcone.

Parameters:

latitude (ndarray) – Angular coordinates in radians. Either 1D, one per output pixel, or 2D with shape (n_angular, npix) to average over sub-samples of each pixel – see healpix_subpixel_lonlat().
longitude (ndarray) – Angular coordinates in radians. Either 1D, one per output pixel, or 2D with shape (n_angular, npix) to average over sub-samples of each pixel – see healpix_subpixel_lonlat().
order (int) – Interpolation order, in the range 0-5.
radial_width (float) – The total radial top-hat the output should carry, in cells. The box already supplies coeval_cell_width, so only the remainder is applied; see residual_radial_width().
coeval_cell_width (float) – The radial top-hat the box already carries, in cells.
n_radial_samples (int) – Number of Gauss-Legendre nodes across the residual width.

Returns:

Pass to cosmotile.jax.shell(), with a radius.

Return type:

sampling

cosmotile.jax.prefilter_coeval(coeval, order)¶

Convert a coeval box to B-spline coefficients, once, for re-use across shells.

Interpolating at order >= 2 requires the box to be converted to B-spline coefficients first, so that the reconstruction passes through the grid values rather than merely being attracted to them.

scipy does this with a recursive IIR filter, which is inherently sequential and has no JAX equivalent. Under the periodic boundary condition cosmotile uses, though, the filter is a circular deconvolution, so it is exactly a division by the Fourier response of the sampled kernel,

\[c = \mathcal{F}^{-1}\left[ \frac{\mathcal{F}[g]}{\prod_i b_p(k_i)} \right],\]

with \(b_p\) from bspline_dtft(). This is not an approximation to scipy’s filter: it agrees with it to roundoff (about 7e-15 relative at 384^3 in double precision), and being an FFT it is very much faster on a GPU – around fifteen to thirty times scipy’s filter on a CPU.

Unlike cosmotile.prefilter_coeval(), this is not optional: the JAX backend will not filter a box for you inside shell(), because that would redo the filter for every shell of a lightcone.

Parameters:

coeval (jax.typing.ArrayLike) – The box to filter. Any array JAX accepts.
order (int) – The interpolation order the box is being prepared for, in the range 0-5. Orders 0 and 1 use interpolating kernels and need no filter, so for them this only tags the box.

Returns:

The coefficients, tagged with order. Pass to shell().

Return type:

prefiltered

Notes

In single precision the deconvolution is a global operation, so its error is not the local error of the gather: expect a relative error around 3e-6 rather than the 1e-15 of double precision. See the precision discussion in the documentation.

cosmotile.jax.shell(coeval, sampling, radius, *, rotation=None, origin=None)¶

Interpolate a box onto one spherical shell.

Pure and differentiable in coeval, radius and origin. The coordinates are built inside the computation rather than passed in, which is what lets a whole lightcone be folded with lightcone_scan() while carrying only one scalar per shell.

Parameters:

coeval (Array | PrefilteredCoeval) – A PrefilteredCoeval for order >= 2, else a 3D array.
sampling (ShellSampling) – The radius-independent geometry, from make_shell_sampling().
radius (Array | float) – Shell radius in cells. May be a tracer.
rotation (Array | None) – A (3, 3) rotation matrix, or None. Call scipy.spatial.transform.Rotation.as_matrix() on the host to get one.
origin (Array | None) – A (3,) shell centre in cells, or None.

Returns:

(sampling.npix,) interpolated values.

Return type:

values

cosmotile.jax.shell_from_coordinates(coeval, coordinates, *, order, weights=None, npix=None)¶

Interpolate a box at arbitrary pixel coordinates.

The primitive the rest of the backend is built on. Pure, so it may be wrapped in jax.jit(), jax.vmap() or jax.grad() directly; it is deliberately not jit-ed itself, so that jit-ing the caller does not compile it twice.

Parameters:

coeval (Array | PrefilteredCoeval) – A PrefilteredCoeval for order >= 2, or a raw 3D array for orders 0 and 1.
coordinates (Array) – (3, nsample) coordinates in cells, sub-sample-major. Need not lie inside the box: the grid is periodic.
order (int) – Interpolation order, in the range 0-5. Static under jit.
weights (Array | None) – Sub-sample weights, or None to return one value per sample.
npix (int | None) – Number of output pixels. Required when weights is given.

Returns:

(npix,) if weights is given, else (nsample,).

Return type:

values