ps_eor: Power-spectrum analysis and signal separation for 21-cm cosmology

ps_eor provides data containers, power-spectrum estimators, simulations, and signal-separation tools for low-frequency interferometric 21-cm experiments. It supports the path from gridded visibility data to spatial, cylindrical, and spherical power spectra, including Gaussian-process foreground and systematic separation.

📖 Documentation: https://ps-eor.readthedocs.io/en/latest/ — getting started, topic-based user guide, pstool command reference, and full API reference.

Upgrading from ps_eor 0.x to 1.0? Read the ML-GPR migration guide.

What the package does

Data handling — build analysis-ready visibility cubes from calibrated images:

• read frequency-ordered FITS images (Jy/PSF), using matching PSF files for an accurate Kelvin conversion, or inferring the PSF area from image metadata;
• select the field of view and baseline range, Fourier-transform to the uv plane, and store visibilities, weights, frequencies, and metadata as HDF5.

Power spectra — measure 21-cm power spectra:

• spatial, cylindrical, and spherical spectra, plus cross-spectra between independent observations, with the instrumental and cosmological normalizations applied.

Signal separation — suppress foregrounds and systematics, preserve the 21-cm signal:

• polynomial and PCA foreground models;
• ML-GPR, a Gaussian-process framework that infers posterior foreground, systematic, noise, and 21-cm components and propagates them to power spectra — for one observation or jointly across many, separating repeatable sky from observation-dependent excess variance.

Simulation and sensitivity — generate test data and forecast performance:

• draw Gaussian-process component cubes (foreground, 21-cm, systematic; single- or multi-epoch) from covariance kernels or an input power spectrum;
• generate thermal-noise data cubes from measured weights or a simulated telescope layout and observing strategy;
• forecast UV coverage and power-spectrum sensitivity, with built-in models for LOFAR, NenuFAR, MWA, SKA-Low, HERA, AARTFAAC, and OVRO-LWA.

The GPR approach follows Mertens et al. (2018); the multi-observation extension relates to Munshi et al. (2025).

Installation

The base package contains the data and power-spectrum functionality:

pip install ps-eor

Install the Gaussian-process backend and its standard samplers, including NUTS through Pyro, with:

pip install "ps-eor[ml-gpr]"

UltraNest is available separately:

pip install "ps-eor[ml-gpr,ultranest]"

For development:

git clone https://gitlab.com/flomertens/ps_eor.git
cd ps_eor
pip install -e ".[ml-gpr,dev]"

Python 3.10 or newer is required.

Creating a visibility cube from FITS images

Given one calibrated image and one matching PSF image per frequency, ps_eor can perform the Jy/PSF-to-Kelvin conversion and construct a weighted CartDataCube. The files are sorted using their FITS frequency metadata:

from pathlib import Path

import numpy as np

from ps_eor import datacube, psutil

image_files = [str(path) for path in Path("fits/images").glob("*.fits")]
psf_files = [str(path) for path in Path("fits/psfs").glob("*.fits")]

image_files = psutil.sort_by_fits_key(image_files, "CRVAL3")
psf_files = psutil.sort_by_fits_key(psf_files, "CRVAL3")

cube = datacube.CartDataCube.load_from_fits_image_and_psf(
    image_files,
    psf_files,
    umin=50,
    umax=250,
    theta_fov=np.deg2rad(4),
    int_time=10,
    total_time=10 * 3600,
)
cube.save("visibilities.h5")

Here, umin and umax are in wavelengths, theta_fov is in radians, and the integration and total observing times are in seconds. The same operation is available as pstool gen_vis_cube.

Basic power-spectrum workflow

import numpy as np

from ps_eor import datacube, pspec

cube = datacube.CartDataCube.load("visibilities.h5")

config = pspec.PowerSpectraConfig(
    el=2 * np.pi * np.arange(cube.ru.min(), cube.ru.max(), 10),
    window_fct="hann",
)
ps_gen = pspec.PowerSpectraBuilder(config).get(cube)

spatial = ps_gen.get_ps(cube)
cylindrical = ps_gen.get_ps2d(cube)

kbins = np.logspace(np.log10(ps_gen.kmin), np.log10(0.5), 10)
spherical = ps_gen.get_ps3d(kbins, cube)

ML-GPR workflow

The high-level fitter is configured from TOML:

from ps_eor.ml_gpr import MLGPRConfigFile, MLGPRForegroundFitter

config = MLGPRConfigFile.get_defaults()
fitter = MLGPRForegroundFitter(config)
noise_for_fit = fitter.process_noise_cube(noise_cube)
result = fitter.run(data_cube, noise_for_fit)

Advanced users can work directly with the lower-level building blocks:

from ps_eor.ml_gpr.kernels import UVScaledKernel
from ps_eor.ml_gpr.multidata import MultiData
from ps_eor.ml_gpr.regressor import MultiGPRegressor

Command-line interface

The package installs pstool, which can run the main analysis and sensitivity workflows without writing Python code:

Each command documents its inputs and options:

pstool --help
pstool COMMAND --help

ML-GPR tutorials

The numbered notebooks form a focused ML-GPR tutorial series:

1. Quick start
2. Simulating data
3. The covariance model
4. Running a fit
5. Training the VAE
6. ML-GPR with the VAE kernel
7. Config-driven runs
8. Multi-epoch GPR
9. Data in-painting
10. Validation and failure modes

Specialist examples, including Cross-GPR and sampler-prior comparisons, live in the advanced notebooks.

Tests

pytest

Frozen numerical fixtures used to compare the GPyTorch implementation with the former GPy backend are stored under tests/data/gpr_golden/.

License

ps_eor is distributed under the GNU General Public License v3 or later.