vskywatcher

vskywatcher (Virtual Skywatcher) computes weighted line-of-sight integrals over 3D cosmological fields. It is built upon JAX python library, and provides support for automatic differentiation.

Main features

• Simple Python interface, C++/CUDA under the hood via Kokkos
• Integrates a 3D field along a large number of view directions in parallel
• Runs on CPU (OpenMP) and GPU (CUDA) without requiring low-level code
• Fully differentiable via JAX — use jax.grad directly on the integration
• Designed for simulation-based inference workflows

What it computes

For each line of sight defined by a direction \hat{n}, vskywatcher evaluates:

I(\hat{n}) = \int_{\chi_{\min}}^{\chi_{\max}} \delta\bigl(\mathbf{r}_{\rm obs} + \chi\,\hat{n}\bigr)\, W(\chi)\, d\chi

where \delta is a 3D scalar field sampled on a regular cubic grid, \chi is the co-moving distance along the ray, and W(\chi) is a user-defined weighting function (e.g. a photometric redshift distribution, a lensing kernel, or a selection function).

Why Gaussian mixtures?

An arbitrary weighting function requires numerical quadrature, which is slow and introduces discretization errors. vskywatcher instead requires W to be expressed as a linear combination of Gaussians:

W(\chi) = \sum_{g=1}^{G} A_g \exp\!\left(-\frac{(\chi - \mu_g)^2}{2\sigma_g^2}\right)

Within each grid cell, the trilinearly-interpolated field is a degree-3 polynomial in \chi. The integral of a polynomial against a Gaussian is known analytically, so the full integral is computed exactly and efficiently, cell by cell. Any smooth window function can be approximated to arbitrary precision by a Gaussian mixture.

It is up to the user to build such a mixture for their desired weighting function. They provide the mixture as a JAX array of shape 3*G, where each Gaussian is described, in order, by its amplitude, mean, and standard deviation. Another option is to provide an N*3*G table, where each of the N lines of sight has its own Gaussian mixture.

Under the hood, if you use vskywatcher with a GPU as computing device, you should be aware that:

• Gaussians can be specific to one line of sight or shared by all. In low-level code, this does not make any difference: Gaussians are always handled as line-of-sight-specific.
• Gaussians are stored in GPU shared memory to be reused by all threads assigned to a given line of sight. Above a GPU-architecture-dependent number of Gaussians, this does not just slow computation down — the scratch memory reserved per team may become too small for the actual mixture size, which can lead to incorrect results or a crash. As a rule of thumb, staying under 1000 Gaussians in double precision or 2000 in single precision is safe across all supported GPU architectures.

Installation

vskywatcher is split into two packages: the pure Python frontend and a compiled backend that matches your hardware.

CPU (OpenMP):

pip install vskywatcher[openmp]

GPU (CUDA 12):

pip install vskywatcher[cuda12]

JAX is required and will be installed automatically as a dependency.

Supported GPU architectures

The CUDA package is compiled for the following NVIDIA architectures:

Architecture	GPU examples	Compute capability
Volta	V100	sm_70
Ampere	A100	sm_80
Ada Lovelace	L40S, RTX 4090	sm_89
Hopper	H100, H200	sm_90
Blackwell	B100	sm_100
Blackwell	B200, GB200	sm_120

If you need support for a different architecture, please let us know.

Quick start

The main function one may want to use in vskywatcher is integrate_field_with_weight, taking as input:

• The field (Nx*Ny*Nz) containing all density contrasts on the grid vertices (or any other field you want to integrate upon)
• The line of sight directions as an array of 3D vectors (N*3)
• The Gaussian mixtures for each direction (N*3*G)
• The resolution, i.e. the edge length of a grid cell in Mpc/h (float/double)
• The observer position in the grid (array size 3)
• The position of the origin of the grid (array size 3)
• The list of the lower bound of integration for each line of sight
• The list of the upper bound of integration for each line of sight

It is advised to use the helper function truncate_to_domain_with_weight that performs a number of corrections and conversions to bring your inputs to the expected format:

• Normalizing directions
• Converting a 3*G Gaussian mixture into the N*3*G required input if you want to use the same weighting function for all lines of sight.
• Defaulting the observer to (0,0,0) and, if unspecified, computing the origin so that the observer sits at the center of the domain.
• Truncating lower and upper bounds to the domain boundaries.

Note that the observer can be outside the domain.

import jax.numpy as jnp
from vskywatcher import integrate_field_with_weight
from vskywatcher.truncate_to_domain import truncate_to_domain_with_weight

# 3D field on a cubic grid (e.g. dark matter density)
field = jnp.ones((256, 256, 256), dtype=jnp.float32)

# Line-of-sight directions (N x 3)
directions = jnp.array([
    [1.0, 0.0, 0.0],
    [0.0, 1.0, 0.0],
    [1.0, 1.0, 0.0],
])

# Weighting function: G Gaussians per direction, described by (amplitude, mean, sigma)
# Shape: (3, G) — amplitudes, means, sigmas stacked along axis 0
n_gaussians = 5
weighting_function = jnp.array([
    [1.0] * n_gaussians,    # amplitudes
    [0.5] * n_gaussians,    # means (co-moving distance in Mpc/h)
    [0.1] * n_gaussians,    # standard deviations (Mpc/h)
])

# Clip the field and Gaussians to the overlapping domain
resolution = 1.0 / 256  # co-moving size of one cell edge in Mpc/h
field, weighting_function, directions, resolution, observer, origin, starts, ends = \
    truncate_to_domain_with_weight(field, weighting_function, directions, resolution)

# Integrate — returns one value per direction
result = integrate_field_with_weight(
    field, weighting_function, directions, resolution, observer, origin, starts, ends
)
print(result)  # shape (N,)

Automatic differentiation

integrate_field_with_weight supports jax.grad out of the box:

import jax

def loss(field):
    return jnp.sum(integrate_field_with_weight(
        field, weighting_function, directions, resolution, observer, origin, starts, ends
    ))

grad = jax.grad(loss)(field)  # shape identical to field

Documentation

Full API documentation is available in docs/. To build it locally:

pip install sphinx
sphinx-build -b html docs/source docs/build/html