flatmaps — Flat-sky Fourier utilities

flatmaps is the mathematical foundation of the analysis pipeline. Every function that works in harmonic space — power spectra, Gaussian realisations, filtering, polarisation rotation — lives here.

All functions operate on flat-sky maps: 2D NumPy arrays parameterised by:

flatskymapparams = [nx, ny, dx, dy]   # dx, dy in arcminutes/pixel

For the 6° × 6° training patches used in this project the standard value is [256, 256, 1.40625, 1.40625]. Pass this list to every function that accepts a flatskymapparams argument.

Common patterns

Compute a mean power spectrum over N maps:

from foregrounds_diffusion.flatmaps import map2cl
from foregrounds_diffusion.moments import mean_cls

mapparams = [256, 256, 1.40625, 1.40625]
el, mean_cl, std_cl = mean_cls(maps_nhw, mapparams, lmin=300, lmax=4000, binsize=60)

Apply a low-pass filter at ℓ = 7000:

from foregrounds_diffusion.flatmaps import get_lpf_hpf, bandpass_filter

lpf = get_lpf_hpf(mapparams, 7000, filter_type=0)
filtered = bandpass_filter(raw_map, lpf)

Draw a correlated Gaussian realisation:

from foregrounds_diffusion.flatmaps import make_gaussian_realisation

sim = make_gaussian_realisation(mapparams, el, cl_cib, cl2=cl_tsz, cl12=cl_cross)
# sim[0] = CIB-like field, sim[1] = correlated tSZ-like field

GPU acceleration

map2cl_torch() computes all N power spectra in a single batched torch.fft.fft2 call, giving a substantial speedup over calling map2cl() in a Python loop. See 13 — Profiling and Benchmarking Baseline for timing comparisons.

API

Flat-sky Fourier utilities for CMB map analysis.

This module provides the core mathematical primitives used throughout the pipeline. All functions operate on flat-sky maps: 2D NumPy arrays where each pixel represents a fixed angular area on the sky, parameterised by:

flatskymapparams = [nx, ny, dx, dy]   # dx, dy in arcminutes/pixel

For the training patches used in this project: [256, 256, 1.40625, 1.40625] (6° × 6° at 1.40625 arcmin/pixel).

Functions

Power spectrum estimation:

map2cl() — auto- or cross-power spectrum of one or two maps mean_cls() — mean spectrum over a stack (see moments) cl2map() — draw a single Gaussian realisation from a 1D C_ℓ make_gaussian_realisation() — correlated two-field realisation

GPU-accelerated:

build_lbin_idx_fft2() — precompute ℓ-bin indices for PyTorch map2cl_torch() — batched power spectrum on GPU

Filtering:

get_lpf_hpf() — low-pass, high-pass, or band-pass filter mask bandpass_filter() — apply a 2D filter to a map

Polarisation:

convert_eb_qu() — rotate between E/B and Q/U representations

Utilities:

get_lxly() — 2D Fourier wavenumber grids radial_profile() — azimuthally averaged radial profile cl_to_cl2d() — interpolate 1D C_ℓ onto the 2D Fourier grid

foregrounds_diffusion.flatmaps.get_lxly(flatskymapparams)

Return 2D Fourier wavenumber grids lx and ly.

Parameters:

flatskymapparams (list) – [nx, ny, dx, dy] where ny, nx = flatskymap.shape and dx, dy are the pixel resolution in arcminutes. Example: [100, 100, 0.5, 0.5] gives a 50’ x 50’ map at 0.5’ resolution.

Returns:

lx, ly (ndarray) – 2D arrays of Fourier wavenumbers.

foregrounds_diffusion.flatmaps.get_lxly_az_angle(lx, ly)

Return the azimuthal angle in Fourier space.

Parameters:

lx, ly (ndarray) – 2D Fourier wavenumber arrays from get_lxly().

Returns:

ndarray – Azimuthal angle array.

foregrounds_diffusion.flatmaps.get_lpf_hpf(flatskymapparams, lmin_lmax, filter_type=0)

Build a 2D Fourier-space filter on the flat-sky grid.

Parameters:

• flatskymapparams (list) – [nx, ny, dx, dy] — map dimensions and pixel size in arcmin.

• lmin_lmax (float or tuple of float) – For filter_type 0 or 1: a single ℓ cutoff. For filter_type 2: a (lmin, lmax) pair defining the passband.

• filter_type (int, optional) – 0 — low-pass (default): zero all ℓ > lmin_lmax. 1 — high-pass: zero all ℓ < lmin_lmax. 2 — band-pass: pass only lmin ≤ ℓ ≤ lmax.

Returns:

fft_filter (ndarray, shape (ny, nx)) – Binary 2D mask (0 or 1) in Fourier space. Multiply by np.fft.fft2(map) then np.fft.ifft2 to apply.

foregrounds_diffusion.flatmaps.bandpass_filter(fmap, bp)

Apply a 2D bandpass filter to a real-valued flat-sky map.

Parameters:

• fmap (ndarray, shape (ny, nx)) – Input real-space map.

• bp (ndarray, shape (ny, nx)) – 2D filter in Fourier space (e.g. from get_lpf_hpf() with filter_type=2).

Returns:

ndarray – Filtered real-space map.

foregrounds_diffusion.flatmaps.cl_to_cl2d(el, cl, flatskymapparams)

Interpolate a 1D power spectrum onto a 2D Fourier grid.

Parameters:

• el (array_like) – Multipole values at which cl is defined.

• cl (array_like) – 1D power spectrum C_ℓ.

• flatskymapparams (list) – [nx, ny, dx, dy] — see get_lxly().

Returns:

ndarray – 2D power spectrum on the Fourier grid.

foregrounds_diffusion.flatmaps.map2cl(flatskymapparams, flatskymap1, flatskymap2=None, binsize=None, minbin=100, maxbin=10000, _ell_bin_cache=None)

Compute auto- or cross-power spectrum of flat-sky map(s).

Parameters:

• flatskymapparams (list) – [nx, ny, dx, dy] — see get_lxly().

• flatskymap1 (ndarray, shape (ny, nx)) – First map.

• flatskymap2 (ndarray, shape (ny, nx), optional) – Second map for cross-spectrum. Auto-spectrum computed when None.

• binsize (float, optional) – Bin width in ℓ. Computed automatically when None.

• minbin, maxbin (float) – Minimum and maximum ℓ bins.

• _ell_bin_cache (tuple, optional) – Pre-computed (binarr, bin_idx, valid_ell, binsize) from _build_ell_bin_cache(). When supplied, skips get_lxly, sqrt, and np.digitize — pass through mean_cls / mean_cross_cls to amortise across N maps.

Returns:

el, cl (ndarray) – Binned multipoles and power spectrum.

foregrounds_diffusion.flatmaps.cl2map(flatskymapparams, cl, el=None)

Generate a Gaussian realisation of a flat-sky map from a 1D C_ℓ.

For correlated multi-field realisations see make_gaussian_realisation().

Parameters:

• flatskymapparams (list) – [nx, ny, dx, dy] — see get_lxly().

• cl (array_like) – 1D power spectrum.

• el (array_like, optional) – Multipoles. Defaults to np.arange(len(cl)).

Returns:

ndarray – Simulated flat-sky map.

foregrounds_diffusion.flatmaps.make_gaussian_realisation(mapparams, el, cl, cl2=None, cl12=None, bl=None, qu_or_eb='qu')

Generate a (possibly correlated two-field) Gaussian flat-sky realisation.

Parameters:

• mapparams (list) – [nx, ny, dx, dy] — see get_lxly().

• el (array_like) – Multipoles.

• cl (array_like) – Auto-spectrum of field 1 (or the only field when cl2 is None).

• cl2 (array_like, optional) – Auto-spectrum of field 2. Required together with cl12.

• cl12 (array_like, optional) – Cross-spectrum between field 1 and field 2.

• bl (array_like, optional) – Beam transfer function (1D or 2D). Applied to the output if given.

• qu_or_eb ({‘qu’, ‘eb’}) – Whether polarisation output should be in Q/U or E/B convention.

Returns:

ndarray – Simulated map (1D or 3-component array for polarisation).

foregrounds_diffusion.flatmaps.convert_eb_qu(map1, map2, flatskymapparams, eb_to_qu=1)

Convert between E/B and Q/U polarisation representations.

Parameters:

• map1, map2 (ndarray) – Input polarisation maps.

• flatskymapparams (list) – [nx, ny, dx, dy] — see get_lxly().

• eb_to_qu (int) – If 1 convert E/B → Q/U; if 0 convert Q/U → E/B.

Returns:

map1_mod, map2_mod (ndarray) – Rotated polarisation maps.

foregrounds_diffusion.flatmaps.radial_profile(z, xy=None, bin_size=1.0, minbin=0.0, maxbin=10.0, to_arcmins=1)

Compute the radial profile of a real- or Fourier-space image.

Parameters:

• z (ndarray) – 2D image.

• xy (tuple of ndarray, optional) – Pre-computed (x, y) coordinate arrays. Computed from z when None.

• bin_size (float) – Radial bin width.

• minbin, maxbin (float) – Radial range.

• to_arcmins (int) – If 1, multiply radius by 60 (convert degrees to arcminutes).

Returns:

ndarray, shape (nbins, 3) – Columns: bin centre, mean value, error on mean.

foregrounds_diffusion.flatmaps.build_lbin_idx_fft2(mapparams, binsize=None, minbin=100, maxbin=10000)

Pre-compute ℓ-bin index tensor for map2cl_torch().

Uses the same full fft2 frequency grid as _build_ell_bin_cache() so that map2cl_torch() produces per-bin means that are numerically identical (within float32 tolerance) to the CPU map2cl().

Parameters:

• mapparams (list) – [nx, ny, dx, dy] — see get_lxly().

• binsize (float, optional) – Bin width in ℓ. Defaults to the smallest ℓ spacing.

• minbin, maxbin (float) – Multipole range.

Returns:

• lbin_idx (torch.Tensor, shape (ny * nx,), dtype long) – 0-indexed bin assignment per fft2 pixel. Out-of-range pixels receive the sentinel value n_bins.

• bin_counts (torch.Tensor, shape (n_bins,), dtype float32) – Number of non-zero fft2 pixels per bin (for normalisation).

• n_bins (int) – Number of ℓ bins.

foregrounds_diffusion.flatmaps.map2cl_torch(maps_nhw, lbin_idx, bin_counts, n_bins, dx_arcmin)

Batched auto-power spectrum on GPU via torch.fft.fft2.

Computes all N power spectra in a single batched FFT call, replacing the Python loop over N maps in map2cl(). Uses the full fft2 grid (same as the CPU implementation) so per-bin means match to within float32 accumulation differences (rtol ≈ 1e-3).

Parameters:

• maps_nhw (torch.Tensor, shape (N, H, W)) – Batch of flat-sky maps on any device (CPU or CUDA).

• lbin_idx (torch.Tensor, shape (H * W,), dtype long) – Bin assignments from build_lbin_idx_fft2(), on the same device.

• bin_counts (torch.Tensor, shape (n_bins,), dtype float32) – Pixel counts per bin from build_lbin_idx_fft2(), on the same device.

• n_bins (int) – Number of ℓ bins.

• dx_arcmin (float) – Pixel resolution in arcminutes (same for x and y).

Returns:

torch.Tensor, shape (N, n_bins), dtype float32 – Mean power per ℓ bin for each of the N input maps.

Examples

>>> lbin_idx, bin_counts, n_bins = build_lbin_idx_fft2(mapparams)
>>> maps_t = torch.from_numpy(maps_np)          # (N, H, W)
>>> cl_batch = map2cl_torch(maps_t, lbin_idx, bin_counts, n_bins, 1.40625)