From b2623ed79b79bbabc1d041a4d3388ac4cc690fd7 Mon Sep 17 00:00:00 2001 From: jgslunde Date: Wed, 20 May 2026 23:57:53 +0200 Subject: [PATCH 01/23] Changed the DetectorTOD dataclass to support compressed flags entry. --- src/commander4/data_models/detector_TOD.py | 54 +++++++++++++++------- 1 file changed, 38 insertions(+), 16 deletions(-) diff --git a/src/commander4/data_models/detector_TOD.py b/src/commander4/data_models/detector_TOD.py index 2558710..337e94a 100644 --- a/src/commander4/data_models/detector_TOD.py +++ b/src/commander4/data_models/detector_TOD.py @@ -47,6 +47,7 @@ def __init__( tod_is_compressed: bool = True, pix_is_compressed: bool = True, psi_is_compressed: bool = True, + flag_is_compressed: bool = True, ): """Construct a DetectorTOD. @@ -94,7 +95,13 @@ def __init__( self._tod_is_compressed = tod_is_compressed self._pix_is_compressed = pix_is_compressed self._psi_is_compressed = psi_is_compressed - self._processing_mask_TOD = np.packbits(processing_mask_map[self.pix]) + self._flag_is_compressed = flag_is_compressed + processing_mask = processing_mask_map[self.pix] + self._processing_mask_TOD = np.packbits(processing_mask) + if flag_encoded is not None and flag_bitmask is not None: + bad_data_mask = ~(self.flag & flag_bitmask) + self._bad_data_mask = np.packbits(bad_data_mask) + self._full_mask = np.packbits(bad_data_mask & processing_mask) if orb_dir_vec is not None: log.logassert_np(orb_dir_vec.size == 3, "orb_dir_vec must be a vector of size 3.", logger) self._orb_dir_vec = orb_dir_vec.astype(np.float32, copy=False) @@ -163,6 +170,18 @@ def psi(self) -> NDArray[np.floating]: psi = self._psi_encoded return psi[:self.ntod] # Crop to actual size (might be cut to fast FFT length) + @property + def flag(self) -> NDArray[np.integer]: + if self._flag_is_compressed: + flag = np.zeros(self.ntod_original, dtype=np.int64) + flag = cpp_utils.huffman_decode(np.frombuffer(self._flag_encoded, dtype=np.uint8), + self._huffman_tree, self._huffman_symbols, flag) + # TODO: Uncomment when SO flags are fixed. + # flag = np.cumsum(flag) + else: + flag = self._flag_encoded + return flag[:self.ntod] + @property def processing_mask_TOD(self) -> NDArray[np.bool_]: """Boolean mask selecting valid (unmasked) TOD samples. @@ -176,24 +195,27 @@ def processing_mask_TOD(self) -> NDArray[np.bool_]: raise ValueError(f"Mask size {mask.size} doesn't match TOD size {self.tod.size}.") return mask[:self.tod.size] + @property - def flags(self) -> NDArray[np.floating]: - """ - Returns the uncompressed flag array. - """ - flags = np.zeros(self.ntod_original, dtype=np.int64) - flags = cpp_utils.huffman_decode(np.frombuffer(self._flag_encoded, dtype=np.uint8), - self._huffman_tree, self._huffman_symbols, flags) - flags = np.cumsum(flags) - flags = flags[:self.ntod] - return flags + def full_mask(self) -> NDArray[np.bool_]: + mask = np.unpackbits(self._full_mask).view(bool) + if mask.size > self.tod.size + 7 or mask.size < self.tod.size: + raise ValueError(f"Mask size {mask.size} doesn't match TOD size {self.tod.size}.") + return mask[:self.tod.size] @property - def excluded_tod_mask(self) -> NDArray[np.bool_]: - """ - Returns a mask given by the intersection between the flag array and the flag bitmask. - """ - return (self.flags & self._flag_bitmask).astype(np.bool_) + def bad_data_mask(self) -> NDArray[np.bool_]: + mask = np.unpackbits(self._bad_data_mask).view(bool) + if mask.size > self.tod.size + 7 or mask.size < self.tod.size: + raise ValueError(f"Mask size {mask.size} doesn't match TOD size {self.tod.size}.") + return mask[:self.tod.size] + + # @property + # def excluded_tod_mask(self) -> NDArray[np.bool_]: + # """ + # Returns a mask given by the intersection between the flag array and the flag bitmask. + # """ + # return (self.flags & self._flag_bitmask).astype(np.bool_) @property From 70cdd88da86f81c5f5f480be55839af11477cde0 Mon Sep 17 00:00:00 2001 From: jgslunde Date: Mon, 25 May 2026 11:05:23 +0200 Subject: [PATCH 02/23] Implemented live pointing construction, and abstracted away pointing models into classes. Also added an SO-SAT reader --- src/commander4/data_models/detector_TOD.py | 157 ++++++----- src/commander4/data_models/pointing.py | 249 ++++++++++++++++++ .../experiments/SO/tod_reader_SO_SAT.py | 191 ++++++++++++++ src/commander4/tod_processing.py | 159 +++++------ 4 files changed, 599 insertions(+), 157 deletions(-) create mode 100644 src/commander4/data_models/pointing.py create mode 100644 src/commander4/experiments/SO/tod_reader_SO_SAT.py diff --git a/src/commander4/data_models/detector_TOD.py b/src/commander4/data_models/detector_TOD.py index 337e94a..fa82f26 100644 --- a/src/commander4/data_models/detector_TOD.py +++ b/src/commander4/data_models/detector_TOD.py @@ -1,20 +1,23 @@ import numpy as np import logging from numpy.typing import NDArray -import ducc0 -import os + from commander4.cmdr4_support import utils as cpp_utils +from commander4.data_models.pointing import PixelPointing, DetectorBoresightPointing import commander4.output.log as log +from commander4.logging.performance_logger import benchmark, bench_summary, start_bench,\ + stop_bench, log_memory, increment_count, bench_reset logger = logging.getLogger(__name__) + class DetectorTOD: """Holds time-ordered data (TOD) for a single detector within a scan. - The raw pixel and polarization angle arrays are stored in Huffman-compressed form - and decompressed on demand via the ``pix`` and ``psi`` properties. The TOD array, - evaluation nside, data nside, and sampling frequency are stored as plain public - attributes. + The pointing information is stored in a ``PixelPointing`` or + ``DetectorBoresightPointing`` object and exposed via ``get_pix()``, + ``get_psi()``, and ``get_pix_psi()``. The TOD array, evaluation nside, + data nside, and sampling frequency are stored as plain public attributes. Attributes: tod (NDArray[np.floating]): 1-D array of calibrated time-ordered samples. @@ -27,15 +30,11 @@ class DetectorTOD: def __init__( self, tod: NDArray[np.floating], - pix_encoded: NDArray[np.integer] | bytes, - psi_encoded: NDArray[np.integer] | NDArray[np.floating] | bytes, - nside: int, - data_nside: int, + pointing: PixelPointing | DetectorBoresightPointing, fsamp: float, orb_dir_vec: NDArray[np.floating] | None, - huffman_tree: NDArray, - huffman_symbols: NDArray, - npsi: int, + huffman_tree: NDArray | None, + huffman_symbols: NDArray | None, processing_mask_map: NDArray[np.bool_], ntod_original: int, ntod_optimal: int, @@ -45,30 +44,24 @@ def __init__( flag_bitmask: int | None = None, init_scalars: NDArray | None = None, tod_is_compressed: bool = True, - pix_is_compressed: bool = True, - psi_is_compressed: bool = True, flag_is_compressed: bool = True, + det_response: NDArray | None = None, ): """Construct a DetectorTOD. Args: tod: 1-D floating-point array of calibrated time samples. - pix_encoded: Huffman-encoded (or raw) pixel index array. - psi_encoded: Huffman-encoded (or raw) polarization angle array. - nside: HEALPix nside for map evaluation. - data_nside: HEALPix nside the pixel indices are stored at on disk. + pointing: Pointing representation for this detector. Must be a + ``PixelPointing`` or ``DetectorBoresightPointing`` instance. fsamp: Sampling frequency in Hz. orb_dir_vec: Unit vector of the spacecraft orbital velocity (size 3), or None if orbital dipole is not used. - huffman_tree: Huffman decoding tree (passed to C++ decoder). - huffman_symbols: Huffman symbol table (passed to C++ decoder). - npsi: Number of discretised polarization angle bins. + huffman_tree: Huffman decoding tree for the flag stream, or None. + huffman_symbols: Huffman symbol table for the flag stream, or None. processing_mask_map: Boolean HEALPix map selecting valid pixels. ntod_original: Original TOD length before Fourier-length cropping. flag_encoded: Huffman-encoded flag array, or None. flag_bitmask: Bitmask applied to flags to identify excluded samples. - pix_is_compressed: Whether ``pix_encoded`` is Huffman-compressed. - psi_is_compressed: Whether ``psi_encoded`` is Huffman-compressed. """ if not tod_is_compressed: log.logassert_np(tod.ndim==1, "'value' must be a 1D array", logger) @@ -76,28 +69,46 @@ def __init__( f"type, is {tod.dtype}", logger) log.logassert_np(processing_mask_map.dtype == bool, "Processing mask is not boolean type", logger) + log.logassert_np( + isinstance(pointing, (PixelPointing, DetectorBoresightPointing)), + "pointing must be a PixelPointing or DetectorBoresightPointing instance.", + logger, + ) + log.logassert_np( + pointing.ntod_original == ntod_original, + "Pointing ntod_original does not match DetectorTOD ntod_original.", + logger, + ) + log.logassert_np( + pointing.ntod == ntod_optimal, + "Pointing ntod does not match DetectorTOD ntod_optimal.", + logger, + ) + if flag_encoded is not None and flag_is_compressed: + log.logassert_np( + huffman_tree is not None and huffman_symbols is not None, + "Compressed flags require Huffman metadata.", + logger, + ) self._tod = tod self.ntod_original = ntod_original self.ntod = ntod_optimal - self.nside = nside - self.data_nside = data_nside + self.nside = pointing.nside + self.data_nside = pointing.data_nside self.fsamp = fsamp self.init_scalars = init_scalars - self._pix_encoded = pix_encoded - self._psi_encoded = psi_encoded self._flag_encoded = flag_encoded self._flag_bitmask = flag_bitmask self._huffman_tree = huffman_tree self._huffman_symbols = huffman_symbols self._huffman_tree2 = huffman_tree2 self._huffman_symbols2 = huffman_symbols2 - self._npsi = npsi self._tod_is_compressed = tod_is_compressed - self._pix_is_compressed = pix_is_compressed - self._psi_is_compressed = psi_is_compressed self._flag_is_compressed = flag_is_compressed - processing_mask = processing_mask_map[self.pix] + self.pointing = pointing + processing_mask = processing_mask_map[self.get_pix()] self._processing_mask_TOD = np.packbits(processing_mask) + self.det_response = det_response if flag_encoded is not None and flag_bitmask is not None: bad_data_mask = ~(self.flag & flag_bitmask) self._bad_data_mask = np.packbits(bad_data_mask) @@ -112,7 +123,7 @@ def __init__( @property def tod(self) -> NDArray[np.floating]: if self._tod_is_compressed: - tod = np.zeros(self.ntod_original, dtype=np.int32) + tod = np.zeros(self.ntod_original, dtype=self._huffman_symbols2.dtype) tod[:] = cpp_utils.huffman_decode(np.frombuffer(self._tod, dtype=np.uint8), self._huffman_tree2, self._huffman_symbols2, tod)[:self.ntod] tod[:] = np.cumsum(tod) @@ -122,62 +133,34 @@ def tod(self) -> NDArray[np.floating]: return tod - @property - def pix(self) -> NDArray[np.integer]: - """Decompressed HEALPix pixel indices at the evaluation nside. - - If the stored pixel array is Huffman-compressed, it is decoded and - cumulative-summed on each access. When ``data_nside != nside`` the - indices are re-projected to the evaluation resolution. - """ - if self._pix_is_compressed: - pix = np.zeros(self.ntod_original, dtype=np.int64) - pix = cpp_utils.huffman_decode(np.frombuffer(self._pix_encoded, dtype=np.uint8), - self._huffman_tree, self._huffman_symbols, pix) - #TODO: Include cumsum in the C++ decode, so it can't be forgotten? - pix = np.cumsum(pix) - else: - pix = self._pix_encoded - # The TOD was cropped to an ideal Fourier length, but because the pix entry is compressed, - # we need to unpack the entire original array, and then crop it to the correct length. - pix = pix[:self.ntod] - - if self.nside != self.data_nside: - # If the data nside does not match the specified evaluation nside, we convert to it. - # pix = hp.ang2pix(self.nside, *hp.pix2ang(self.data_nside, pix)) - nthreads = int(os.environ["OMP_NUM_THREADS"]) - geom_from = ducc0.healpix.Healpix_Base(self.data_nside, "RING") - geom_to = ducc0.healpix.Healpix_Base(self.nside, "RING") - ang = geom_from.pix2ang(pix, nthreads=nthreads) - pix = geom_to.ang2pix(ang, nthreads=nthreads) + def get_pix(self, nside: int | None = None) -> NDArray[np.integer]: + start_bench("pointing") + pix = self.pointing.get_pix(nside) + stop_bench("pointing") return pix - @property - def psi(self) -> NDArray[np.floating]: - """Decompressed polarization angles in radians. + def get_psi(self, nside: int | None = None) -> NDArray[np.integer] | NDArray[np.floating]: + start_bench("pointing") + psi = self.pointing.get_psi(nside) + stop_bench("pointing") + return psi - If the stored array is Huffman-compressed, it is decoded, cumulative-summed, - and converted from integer bins to radians on each access. - """ - if self._psi_is_compressed: - psi = np.zeros(self.ntod_original, dtype=np.int64) - psi = cpp_utils.huffman_decode(np.frombuffer(self._psi_encoded, dtype=np.uint8), - self._huffman_tree, self._huffman_symbols, psi) - psi = np.cumsum(psi) - psi = psi[:self.ntod] - psi = 2*np.pi * psi.astype(np.float32, copy=False)/self._npsi - else: - psi = self._psi_encoded - return psi[:self.ntod] # Crop to actual size (might be cut to fast FFT length) + def get_pix_psi( + self, + nside: int | None = None, + ) -> tuple[NDArray[np.integer], NDArray[np.integer] | NDArray[np.floating]]: + start_bench("pointing") + pix_psi = self.pointing.get_pix_psi(nside) + stop_bench("pointing") + return pix_psi @property def flag(self) -> NDArray[np.integer]: if self._flag_is_compressed: - flag = np.zeros(self.ntod_original, dtype=np.int64) + flag = np.zeros(self.ntod_original, dtype=self._huffman_symbols.dtype) flag = cpp_utils.huffman_decode(np.frombuffer(self._flag_encoded, dtype=np.uint8), self._huffman_tree, self._huffman_symbols, flag) - # TODO: Uncomment when SO flags are fixed. - # flag = np.cumsum(flag) + flag = np.cumsum(flag) else: flag = self._flag_encoded return flag[:self.ntod] @@ -228,4 +211,16 @@ def orb_dir_vec(self) -> NDArray[np.floating]: if self._orb_dir_vec is not None: return self._orb_dir_vec else: - raise ValueError("Attempted to access self.orb_dir_vec, which is not set.") \ No newline at end of file + raise ValueError("Attempted to access self.orb_dir_vec, which is not set.") + + def IQU_response(self, psi: NDArray | None = None): + # psi can be passed as an argument to avoid re-calculating it if already available. + if psi is None: + psi = self.get_psi() + response = np.zeros((3, psi.shape[-1])) + if self.det_response[0] != 0: + response[0,:] = self.det_response[0] + if self.det_response[1] != 0: + response[1,:] = np.cos(2.0*psi)*self.det_response[1] + response[2,:] = np.sin(2.0*psi)*self.det_response[1] + return response \ No newline at end of file diff --git a/src/commander4/data_models/pointing.py b/src/commander4/data_models/pointing.py new file mode 100644 index 0000000..842824f --- /dev/null +++ b/src/commander4/data_models/pointing.py @@ -0,0 +1,249 @@ +import numpy as np +import logging +from numpy.typing import NDArray +import ducc0 +import os +import healpy as hp +from pixell.bunch import Bunch +from pixell import coordsys +from commander4.cmdr4_support import utils as cpp_utils +import commander4.output.log as log + +logger = logging.getLogger(__name__) + + +class ScanBoresightPointing: + """Scan-level boresight pointing evaluated once for all detectors.""" + + def __init__(self, + time_start_mjd: float, + time_end_mjd: float, + ntod_original: int, + site: NDArray, + bore: NDArray, + detoffs: NDArray, + polangs: NDArray | float, + nside: int, + ntod: int | None = None): + self.site = Bunch( + lat = site[0], + lon = site[1], + alt = site[2], + weather = "toco") + self.detoffs = np.asarray(detoffs) + self.polangs = np.asarray(polangs) + self.nside = nside + self.data_nside = nside + self.ntod_original = ntod_original + self.ntod = ntod_original if ntod is None else ntod + self.ndet = self.detoffs.shape[0] + log.logassert_np(self.ntod <= self.ntod_original, "ntod cannot exceed ntod_original.", logger) + log.logassert_np(self.detoffs.ndim == 2, "detoffs must be a 2D array.", logger) + log.logassert_np(self.detoffs.shape[1] == 2, "detoffs must have shape (ndet, 2).", logger) + log.logassert_np( + self.polangs.size == self.ndet, + "polangs must contain one polarization angle per detector.", + logger, + ) + time_start_unix = (time_start_mjd - 40587.0) * 86400.0 + time_end_unix = (time_end_mjd - 40587.0) * 86400.0 + time_unix = np.linspace(time_start_unix, time_end_unix, ntod_original) + + self.bore_point = self.initialize_boresight(time_unix, bore, site=self.site) + + + def initialize_boresight( + self, + ctime: NDArray[np.floating], + bore: NDArray, + sys: str = "cel", + site=None, + weather: str = "typical", + ): + icoord = coordsys.Coords(az=bore[0], el=bore[1], roll=bore[2]) + return coordsys.transform("hor", sys, icoord, ctime=ctime, site=site, weather=weather) + + + def get_det_point( + self, + idet: int, + ) -> tuple[NDArray[np.floating], NDArray[np.floating], NDArray[np.floating]]: + log.logassert_np(0 <= idet < self.ndet, f"Detector index {idet} out of range.", logger) + # By slicing instead of indexing we keep the 1-sized detector dimension. + detoff = self.detoffs[idet:idet+1] + polang = self.polangs[idet:idet+1] + qdet = coordsys.rotation_xieta(detoff[:, 0], detoff[:, 1], polang) + ocoord = self.bore_point * qdet[:, None] + # TODO: The lines below are absurdly slow, taking 95% of the runtime of this function, + # being almost as time-consuming as a full hp.ang2pix call. I tried replacing the call with + # a Numba function, but couldn't achieve a speedup. Should be looked into. + dec = np.asarray(ocoord.dec[0, :self.ntod]) + ra = np.asarray(ocoord.ra[0, :self.ntod]) + psi = np.asarray(ocoord.psi[0, :self.ntod]) + return dec, ra, psi + + + def get_pix_psi( + self, + idet: int, + nside: int | None = None, + ) -> tuple[NDArray[np.integer], NDArray[np.floating]]: + target_nside = self.nside if nside is None else nside + dec, ra, psi = self.get_det_point(idet) + theta = np.pi/2.0 - dec + pix = hp.ang2pix(target_nside, theta, ra) + psi = psi.astype(np.float32, copy=False)[:self.ntod] + return pix, psi + + + def get_pix(self, idet: int, nside: int | None = None) -> NDArray[np.integer]: + return self.get_pix_psi(idet, nside)[0] + + + def get_psi(self, idet: int, nside: int | None = None) -> NDArray[np.floating]: + return self.get_pix_psi(idet, nside)[1] + + + +class DetectorBoresightPointing: + def __init__(self, scan_pointing: ScanBoresightPointing, idet: int): + self.scan_pointing = scan_pointing + self.idet = int(idet) + log.logassert_np( + 0 <= self.idet < self.scan_pointing.ndet, + f"Detector index {self.idet} out of range.", + logger, + ) + self.nside = scan_pointing.nside + self.data_nside = scan_pointing.data_nside + self.ntod_original = scan_pointing.ntod_original + self.ntod = scan_pointing.ntod + + def get_pix(self, nside: int | None = None) -> NDArray[np.integer]: + return self.scan_pointing.get_pix(self.idet, nside) + + def get_psi(self, nside: int | None = None) -> NDArray[np.floating]: + return self.scan_pointing.get_psi(self.idet, nside) + + def get_pix_psi(self, nside: int | None = None) -> tuple[NDArray[np.integer], NDArray[np.floating]]: + return self.scan_pointing.get_pix_psi(self.idet, nside) + + + +class PixelPointing: + """Joint pixel and polarization-angle representation for one detector TOD.""" + + def __init__(self, + pix: bytes | NDArray[np.integer], + psi: bytes | NDArray[np.integer] | NDArray[np.floating], + huffman_tree: NDArray | None, + huffman_symbols: NDArray | None, + npsi: int | None, + nside: int, + data_nside: int, + ntod_original: int, + ntod: int, + ): + self.nside = nside + self.data_nside = data_nside + self.ntod_original = ntod_original + self.ntod = ntod + self.pix_encoded = pix + self.psi_encoded = psi + self.huffman_tree = huffman_tree + self.huffman_symbols = huffman_symbols + self.npsi = npsi + self.pix_is_compressed = isinstance(pix, bytes) + self.psi_is_compressed = isinstance(psi, bytes) + self._test_input() + + + def _test_input(self): + log.logassert_np(self.ntod <= self.ntod_original, "ntod cannot exceed ntod_original.", logger) + log.logassert_np( + self.pix_is_compressed or isinstance(self.pix_encoded, np.ndarray), + "'pix' must be provided as bytes or a numpy array.", + logger, + ) + log.logassert_np( + self.psi_is_compressed or isinstance(self.psi_encoded, np.ndarray), + "'psi' must be provided as bytes or a numpy array.", + logger, + ) + if self.pix_is_compressed: + log.logassert_np( + self.huffman_tree is not None and self.huffman_symbols is not None, + "Compressed pix requires Huffman metadata.", + logger, + ) + if self.psi_is_compressed: + log.logassert_np( + self.huffman_tree is not None and self.huffman_symbols is not None, + "Compressed psi requires Huffman metadata.", + logger, + ) + log.logassert_np(self.npsi is not None, "Compressed psi requires npsi.", logger) + if not self.pix_is_compressed: + pix_array = np.asarray(self.pix_encoded) + log.logassert_np(pix_array.ndim == 1, "'pix' must be a 1D array", logger) + log.logassert_np( + np.issubdtype(pix_array.dtype, np.integer), + "'pix' array must have integer dtype.", + logger, + ) + log.logassert_np( + pix_array.size >= self.ntod, + f"'pix' length {pix_array.size} is shorter than ntod {self.ntod}.", + logger, + ) + if not self.psi_is_compressed: + psi_array = np.asarray(self.psi_encoded) + log.logassert_np(psi_array.ndim == 1, "'psi' must be a 1D array", logger) + log.logassert_np( + np.issubdtype(psi_array.dtype, np.integer) + or np.issubdtype(psi_array.dtype, np.floating), + "'psi' array must have numeric dtype.", + logger, + ) + log.logassert_np( + psi_array.size >= self.ntod, + f"'psi' length {psi_array.size} is shorter than ntod {self.ntod}.", + logger, + ) + + def get_pix(self, nside: int | None = None) -> NDArray[np.integer]: + """Return HEALPix pixel indices at the requested output nside.""" + target_nside = self.nside if nside is None else nside + if self.pix_is_compressed: + pix = np.zeros(self.ntod_original, dtype=np.int64) + pix = cpp_utils.huffman_decode(np.frombuffer(self.pix_encoded, dtype=np.uint8), + self.huffman_tree, self.huffman_symbols, pix) + pix = np.cumsum(pix) + else: + pix = self.pix_encoded + + pix = pix[:self.ntod] + if target_nside != self.data_nside: + nthreads = int(os.environ["OMP_NUM_THREADS"]) + geom_from = ducc0.healpix.Healpix_Base(self.data_nside, "RING") + geom_to = ducc0.healpix.Healpix_Base(target_nside, "RING") + ang = geom_from.pix2ang(pix, nthreads=nthreads) + pix = geom_to.ang2pix(ang, nthreads=nthreads) + return pix + + def get_psi(self, nside: int | None = None) -> NDArray[np.floating]: + """Return polarization angles, cropped to the active TOD length.""" + if self.psi_is_compressed: + psi = np.zeros(self.ntod_original, dtype=np.int64) + psi = cpp_utils.huffman_decode(np.frombuffer(self.psi_encoded, dtype=np.uint8), + self.huffman_tree, self.huffman_symbols, psi) + psi = np.cumsum(psi) + psi = psi[:self.ntod] + psi = 2 * np.pi * psi.astype(np.float32, copy=False) / self.npsi + else: + psi = self.psi_encoded + return psi[:self.ntod] + + def get_pix_psi(self, nside: int | None = None) -> tuple[NDArray, NDArray]: + return self.get_pix(nside), self.get_psi(nside) + diff --git a/src/commander4/experiments/SO/tod_reader_SO_SAT.py b/src/commander4/experiments/SO/tod_reader_SO_SAT.py new file mode 100644 index 0000000..48f9b9b --- /dev/null +++ b/src/commander4/experiments/SO/tod_reader_SO_SAT.py @@ -0,0 +1,191 @@ +import logging +import numpy as np +import healpy as hp +from astropy.io import fits +import h5py +import gc +from numpy.typing import NDArray +from pixell.bunch import Bunch +from mpi4py import MPI + +from commander4.cmdr4_support import utils as cpp_utils +from commander4.data_models.detector_TOD import DetectorTOD +from commander4.data_models.detector_group_TOD import DetGroupTOD +from commander4.data_models.scan_TOD import ScanTOD +from commander4.simulations.inplace_litebird_sim import replace_tod_with_sim +from commander4.output.log import logassert +from commander4.noise_sampling.noise_psd import NoisePSD, NoisePSDOof +from commander4.logging.performance_logger import benchmark, bench_summary, start_bench,\ + stop_bench, log_memory, increment_count, bench_reset +from commander4.data_models.pointing import DetectorBoresightPointing, ScanBoresightPointing +logger = logging.getLogger(__name__) + +def get_processing_mask(my_band: Bunch) -> DetectorTOD: + """ Finds and returns the processing mask for the relevant band. + """ + hdul = fits.open(my_band.processing_mask) + mask = hdul[1].data["TEMPERATURE"].flatten().astype(bool) + nside = np.sqrt(mask.size//12) + if nside != my_band.eval_nside: + mask = hp.ud_grade(mask.astype(np.float64), my_band.eval_nside) == 1 + return mask + +def find_good_Fourier_time(Fourier_times:NDArray, ntod:int) -> int: + if ntod <= 10_000 or ntod >= 400_000: + return ntod + search_start = int(0.99*ntod) # Consider sizes up to 1% smaller than ntod. + best_ntod = np.argmin(Fourier_times[search_start:ntod+1]) + best_ntod += search_start + assert(best_ntod <= ntod) + return best_ntod + + +def tod_reader(band_comm: MPI.Comm, my_experiment: str, my_band: Bunch, det_names: list[str], + params: Bunch, scan_idx_start: int, + scan_idx_stop: int) -> DetGroupTOD: + start_bench("reader-startup") + oids = [] + pids = [] + filepaths = [] + bandname = my_band._name + expname = my_experiment._name + + with open(my_band.filelist) as infile: + infile.readline() + for line in infile: + pid, filepath, _, _, _ = line.split() + pids.append(f"{int(pid):06d}") + filepaths.append(filepath[1:-1]) + oids.append(filepath.split(".")[0].split("_")[-1]) + if "processing_mask" in my_band: + processing_mask_map = np.ones(12*my_band.eval_nside**2, dtype=bool) + if band_comm.Get_rank() == 0: + processing_mask_map[:] = get_processing_mask(my_band) + band_comm.Bcast(processing_mask_map, root=0) + else: + processing_mask_map = np.ones(12*my_band.eval_nside**2, dtype=bool) + + if "bad_PIDs_path" in my_experiment: + bad_PIDs = np.load(my_experiment.bad_PIDs_path) + else: + bad_PIDs = np.array([]) + + Fourier_times = np.load(my_experiment.Fourier_times_path) + + # Attempting to reduce fragmentation by allocating buffers. + ntod_upper_bound = int(100*100*3600) # 10 hour scan. + + ntod_sum_original = 0 + ntod_sum_final = 0 + scan_list = [] + num_included = 0 + stop_bench("reader-startup") + for i_pid in range(scan_idx_start, scan_idx_stop+1): + pid = pids[i_pid] + filepath = filepaths[i_pid] + start_bench("fileread") + if pid in bad_PIDs: + continue + good_scan = True + with h5py.File(filepath, "r") as f: + ntod = int(f[f"/{pid}/common/ntod"][()].item()) + ntod_optimal = find_good_Fourier_time(Fourier_times, ntod) + huffman_tree = f[f"/{pid}/common/hufftree"][()] + huffman_symbols = f[f"/{pid}/common/huffsymb"][()] + # Second Huffman set might not exist. + if f"/{pid}/common/hufftree2" in f and f"/{pid}/common/huffsymb2" in f: + huffman_tree2 = f[f"/{pid}/common/hufftree2"][()] + huffman_symbols2 = f[f"/{pid}/common/huffsymb2"][()] + else: + huffman_tree2 = None + huffman_symbols2 = None + fsamp = float(f["/common/fsamp/"][()].item()) + det_responses = f["/common/resp/"][()] + + # The detector names are stored as a single "Bytes-like" string, formatted like a + # Python list. We extract the string from the Bytes, and then re-create the list with .split(","). + det_names_file = f["/common/det"].asstr()[()].split(",") + det_names_file = [det.strip() for det in det_names_file] + + processing_mask_nside = hp.npix2nside(processing_mask_map.size) + logassert(my_band.eval_nside == processing_mask_nside, + f"Processing mask (band {bandname}) " + f"has nside {processing_mask_nside} while eval_nside = {my_band.eval_nside} " + "(NB: eval_nside can be set different from native data nside)", logger) + + if ntod > ntod_upper_bound: + raise ValueError(f"{ntod_upper_bound} {ntod}") + + all_detector_offsets = f["/common/detoff/"][()] + all_polarization_angles = f["/common/polang/"][()] + site_location = f["/common/site/"][()] + boresight = f[f"/{pid}/common/bore/"][()] + time_start_mjd = f[f"/{pid}/common/time/"][0] + time_end_mjd = f[f"/{pid}/common/time_end/"][0] + + scan_pointing = ScanBoresightPointing(time_start_mjd, time_end_mjd, ntod, site_location, + boresight, all_detector_offsets, all_polarization_angles, + my_band.eval_nside, ntod_optimal) + + detector_list = [] + for idet, det_name in enumerate(det_names): + # Find the index of the current detector in the file order of detectors. + det_file_idx = det_names_file.index(det_name) + + if my_experiment.tod_is_compressed: + tod = f[f"/{pid}/{det_name}/ztod/"][()] + else: + tod = f[f"/{pid}/{det_name}/tod/"][:ntod_optimal].astype(np.float32) + + pointing = DetectorBoresightPointing(scan_pointing, det_file_idx) + det_response = det_responses[det_file_idx] + + flag_encoded = f[f"/{pid}/{det_name}/flag/"][()] + gain_init, sigma0_init, fknee_init, alpha_init = f[f"/{pid}/{det_name}/scalars"][()] + + detector = DetectorTOD(tod, pointing, fsamp, np.zeros(3), huffman_tree, huffman_symbols, + processing_mask_map, ntod, ntod_optimal, + huffman_tree2=huffman_tree2, + huffman_symbols2=huffman_symbols2, + flag_encoded=flag_encoded, + flag_bitmask=my_experiment.flag_bitmask, + tod_is_compressed=my_experiment.tod_is_compressed, + det_response=det_response) + if np.sum(detector.full_mask) == 0 or (detector.tod == 0).all(): + continue + detector_list.append(detector) + ntod_sum_original += ntod + ntod_sum_final += ntod_optimal + + if good_scan: + scanID = int(pid) + scan = ScanTOD(detector_list, 0., scanID) + scan_list.append(scan) + num_included += 1 + if i_pid % 10 == 0: + gc.collect() + + ndet = len(det_names) + noise_model = NoisePSDOof() + band_tod = DetGroupTOD(scan_list, expname, bandname, my_band.eval_nside, my_band.freq, + my_band.fwhm, fsamp, ndet, my_band.polarization, noise_model) + + ### Collect some info on master rank of each detector and print it ### + local_tot_scans = scan_idx_stop - scan_idx_start + local_stats = np.array([num_included, local_tot_scans, ntod_sum_final, ntod_sum_original]) + global_stats = np.zeros_like(local_stats) + band_comm.Reduce(local_stats, global_stats, op=MPI.SUM, root=0) + if band_comm.Get_rank() == 0: + total_included, total_scans, total_ntod_final, total_ntod_original = global_stats + frac_included = 0.0 + if total_scans > 0: + frac_included = total_included / total_scans * 100.0 + avg_scan_remaining = 0.0 + if total_ntod_original > 0: + avg_scan_remaining = total_ntod_final / total_ntod_original * 100.0 + logger.info(f"Band {bandname} finished reading TODs from file.") + logger.info(f"Fraction of scans included for {bandname}: {frac_included:.1f} %") + logger.info(f"Fraction of TODs left after Fourier cut for {bandname}: "\ + f"{avg_scan_remaining:.1f} %") + + return band_tod \ No newline at end of file diff --git a/src/commander4/tod_processing.py b/src/commander4/tod_processing.py index aa61792..96c6372 100644 --- a/src/commander4/tod_processing.py +++ b/src/commander4/tod_processing.py @@ -8,11 +8,11 @@ from numpy.typing import NDArray import healpy as hp -import matplotlib.pyplot as plt from pixell.bunch import Bunch from astropy.io import fits import pysm3.units as pysm3_u +from commander4.output.log import logassert from commander4.data_models.detector_map import DetectorMap from commander4.data_models.detector_group_TOD import DetGroupTOD from commander4.data_models.TOD_samples import TODSamples @@ -86,12 +86,14 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output mapmaker_invvar = WeightsMapmakerIQU(band_comm, experiment_data.nside) for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): - pix = det.pix - psi = det.psi + mask_bad = det.bad_data_mask + pix, psi = det.get_pix_psi() + pix = pix[mask_bad] + psi = psi[mask_bad] sigma0, fknee, alpha = tod_samples.noise_params[iscan, idet, :] gain = tod_samples.gain(iscan, idet) inv_var = (gain/sigma0)**2 - mapmaker_invvar.accumulate_to_map(inv_var, pix, psi) + mapmaker_invvar.accumulate_to_map(inv_var, pix, psi, response=det.det_response) mapmaker_invvar.gather_map() mapmaker_invvar.normalize_map() if ismaster: @@ -105,7 +107,8 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output mapmaker_invvar = WeightsMapmaker(band_comm, experiment_data.nside) for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): - pix = det.pix + mask_bad = det.bad_data_mask + pix = det.get_pix()[mask_bad] sigma0, fknee, alpha = tod_samples.noise_params[iscan, idet, :] gain = tod_samples.gain(iscan, idet) inv_var = (gain/sigma0)**2 @@ -136,16 +139,20 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): d_sky = det.tod.copy() - pix = det.pix - psi = None if pols=="I" else det.psi + pix, psi = det.get_pix_psi() sigma0, fknee, alpha = tod_samples.noise_params[iscan, idet, :] gain = tod_samples.gain(iscan, idet) inv_var = (gain/sigma0)**2 + response = det.det_response if pols == "IQU" else None ### ORBITAL DIPOLE ### sky_orb_dipole = get_s_orb_TOD(det, experiment_data, pix) d_sky -= gain*sky_orb_dipole - mapmaker_orbdipole.accumulate_to_map(sky_orb_dipole, inv_var, pix, psi) + if pols == "IQU": + mapmaker_orbdipole.accumulate_to_map(sky_orb_dipole, inv_var, pix, psi, + response=response) + else: + mapmaker_orbdipole.accumulate_to_map(sky_orb_dipole, inv_var, pix, psi) ### CORRELATED NOISE SAMPLING ### if do_ncorr_sampling: @@ -156,17 +163,21 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output Ntod = sky_subtracted_TOD.shape[0] Nfft = Ntod + 1 # mirrored FFT: nfft=2*Ntod, n=nfft/2+1=Ntod+1 freq = rfftfreq(2 * Ntod, d = 1/det.fsamp) - mask = det.processing_mask_TOD - sigma0_ncorr = calc_sigma0_robust(sky_subtracted_TOD, mask) + mask_full = det.full_mask + sigma0_ncorr = calc_sigma0_robust(sky_subtracted_TOD, mask_full) C_1f_inv = np.zeros(Nfft) C_1f_inv[1:] = 1.0 / (sigma0_ncorr**2*(freq[1:]/fknee)**alpha) # fill_all_masked(sky_subtracted_TOD, mask, sigma0_ncorr) err_tol = 1e-6 n_corr_est, residual = corr_noise_realization_with_gaps(sky_subtracted_TOD, - mask, sigma0_ncorr, C_1f_inv, - err_tol=err_tol) - mapmaker_ncorr.accumulate_to_map((n_corr_est/gain).astype(np.float32, copy=False), - inv_var, pix, psi) + mask_full, sigma0_ncorr, + C_1f_inv, err_tol=err_tol) + if pols == "IQU": + mapmaker_ncorr.accumulate_to_map((n_corr_est/gain).astype(np.float32, copy=False), + inv_var, pix, psi, response=response) + else: + mapmaker_ncorr.accumulate_to_map((n_corr_est/gain).astype(np.float32, copy=False), + inv_var, pix, psi) if residual > err_tol: num_failed_convergences_ncorr += 1 worst_residual_ncorr = max(worst_residual_ncorr, residual) @@ -334,13 +345,15 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu mapmaker_invvar = WeightsMapmakerIQU(band_comm, experiment_data.nside) for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): - pix = det.pix - psi = det.psi + mask_bad = det.bad_data_mask + pix, psi = det.get_pix_psi() + pix = pix[mask_bad] + psi = psi[mask_bad] sigma0, fknee, alpha = tod_samples.noise_params[iscan,idet,:] gain = tod_samples.gain(iscan, idet) # sigma0 is in detector-units, transform into uK_RJ by dividing it by the gain. inv_var = (gain/sigma0)**2 - mapmaker_invvar.accumulate_to_map(inv_var, pix, psi) + mapmaker_invvar.accumulate_to_map(inv_var, pix, psi, response=det.det_response) mapmaker_invvar.gather_map() mapmaker_invvar.normalize_map() @@ -362,23 +375,26 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): start_bench("binned-mapmaker") - d_sky = det.tod.copy() - pix = det.pix - psi = det.psi + mask_bad = det.bad_data_mask + d_sky = det.tod.copy()[mask_bad] + pix, psi = det.get_pix_psi() + pix = pix[mask_bad] + psi = psi[mask_bad] + response = det.det_response gain = tod_samples.gain(iscan, idet) sigma0, fknee, alpha = tod_samples.noise_params[iscan,idet,:] # sigma0 is in detector-units, transform into uK_RJ by dividing it by the gain. inv_var = (gain/sigma0)**2 ### ORBITAL DIPOLE ### - sky_orb_dipole = get_s_orb_TOD(det, experiment_data, pix) + sky_orb_dipole = get_s_orb_TOD(det, experiment_data, pix)[mask_bad] d_sky -= gain*sky_orb_dipole stop_bench("binned-mapmaker", increment_count=False) ### CORRELATED NOISE SAMPLING ### if do_ncorr_sampling: start_bench("ncorr-sampling") - s_tot = get_static_sky_TOD(compsep_output, pix, psi) + s_tot = get_static_sky_TOD(compsep_output, pix, psi)[mask_bad] s_tot += sky_orb_dipole sky_subtracted_TOD = det.tod.copy() sky_subtracted_TOD -= gain*s_tot @@ -386,39 +402,32 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu Nfft = Ntod + 1 # mirrored FFT: nfft=2*Ntod, n=nfft/2+1=Ntod+1 freq = rfftfreq(2 * Ntod, d = 1/det.fsamp) - mask = det.processing_mask_TOD - sigma0_ncorr = calc_sigma0_robust(sky_subtracted_TOD, mask) + mask_full = det.full_mask + sigma0_ncorr = calc_sigma0_robust(sky_subtracted_TOD, mask_full) C_1f_inv = np.zeros(Nfft) C_1f_inv[1:] = 1.0 / (sigma0_ncorr**2*(freq[1:]/fknee)**alpha) err_tol = 1e-6 # Inpaint masked regions with linear slope + white noise. # In the CG solver this is only used to define the starting guess, # but if the CG fails it is also used to generate the fallback solution. - fill_all_masked(sky_subtracted_TOD, mask, sigma0_ncorr) + fill_all_masked(sky_subtracted_TOD, mask_full, sigma0_ncorr) n_corr_est, residual, niter, did_conv = corr_noise_realization_with_gaps(sky_subtracted_TOD, - mask, sigma0_ncorr, C_1f_inv, + mask_full, sigma0_ncorr, C_1f_inv, err_tol=err_tol, max_iter=200) - - # if band_comm.Get_rank() == 0 and idet == 0 and chain == 1: - # if iscan == 300 or iscan == 600 or iscan == 900: - # np.save(f"corrdata/mirrorfft_ncorr_signal_{experiment_data.band_name}_{iscan}_{iter}.npy", sky_subtracted_TOD) - # np.save(f"corrdata/mirrorfft_ncorr_ncorr_{experiment_data.band_name}_{iscan}_{iter}.npy", n_corr_est) - # np.save(f"corrdata/mirrorfft_ncorr_mask_{experiment_data.band_name}_{iscan}_{iter}.npy", mask) - # np.save(f"corrdata/mirrorfft_ncorr_C_1f_inv_{experiment_data.band_name}_{iscan}_{iter}.npy", C_1f_inv) - resid = (sky_subtracted_TOD - n_corr_est) * mask + resid = (sky_subtracted_TOD - n_corr_est) * mask_full var_resid = np.dot(resid, resid) - var_data = np.dot(sky_subtracted_TOD * mask, sky_subtracted_TOD * mask) + var_data = np.dot(sky_subtracted_TOD * mask_full, sky_subtracted_TOD * mask_full) # If either of the two tests failed, use fallback for n_corr. if var_resid > var_data or not did_conv: # Direcly solve constrained realization system without a mask. n_corr_est, _, _, _ = corr_noise_realization_with_gaps(sky_subtracted_TOD, - np.ones_like(mask, dtype=bool), sigma0_ncorr, C_1f_inv) + np.ones_like(mask_full, dtype=bool), sigma0_ncorr, C_1f_inv) # if band_comm.Get_rank() == 0 and idet == 0 and chain == 1: # if iscan == 300 or iscan == 600 or iscan == 900: # np.save(f"corrdata/mirrorfft_corrected_ncorr_signal_{experiment_data.band_name}_{iscan}_{iter}.npy", sky_subtracted_TOD) # np.save(f"corrdata/mirrorfft_corrected_ncorr_ncorr_{experiment_data.band_name}_{iscan}_{iter}.npy", n_corr_est) mapmaker_ncorr.accumulate_to_map((n_corr_est/gain).astype(np.float32, copy=False), - inv_var, pix, psi) + inv_var, pix, psi, response=response) if not did_conv: num_failed_convergences_ncorr += 1 if var_resid > var_data: @@ -441,8 +450,9 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu log_memory("ncorr-sampling") start_bench("binned-mapmaker") - mapmaker.accumulate_to_map(d_sky/gain, inv_var, pix, psi) - mapmaker_orbdipole.accumulate_to_map(sky_orb_dipole, inv_var, pix, psi) + mapmaker.accumulate_to_map(d_sky/gain, inv_var, pix, psi, response=response) + mapmaker_orbdipole.accumulate_to_map(sky_orb_dipole, inv_var, pix, psi, + response=response) stop_bench("binned-mapmaker", increment_count=False) ### PRINT NOISE SAMPLING STATS ### @@ -473,18 +483,18 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu residuals = residuals[residuals != 0] residuals = np.array([0]) if len(residuals) == 0 else residuals niters = np.concatenate(niters) - logger.info(f"{experiment_data.nu}GHz: fknees {np.min(fknees):.4f} "\ - f"{np.percentile(fknees, 1):.4f} {np.mean(fknees):.4f} {np.percentile(fknees, 99):.4f}"\ - f" {np.max(fknees):.4f}") - logger.info(f"{experiment_data.nu}GHz: alphas {np.min(alphas):.4f} "\ - f"{np.percentile(alphas, 1):.4f} {np.mean(alphas):.4f} {np.percentile(alphas, 99):.4f}"\ - f" {np.max(alphas):.4f}") - logger.info(f"{experiment_data.nu}GHz: residuals {np.min(residuals):.2e} "\ - f"{np.percentile(residuals, 1):.2e} {np.mean(residuals):.2e} {np.percentile(residuals, 99):.2e}"\ - f" {np.max(residuals):.2e}") - logger.info(f"{experiment_data.nu}GHz: iterations {np.min(niters):.4f} "\ - f"{np.percentile(niters, 1):.4f} {np.mean(niters):.4f} {np.percentile(niters, 99):.4f}"\ - f" {np.max(niters):.4f}") + logger.info(f"{experiment_data.nu}GHz: fknees {np.nanmin(fknees):.4f} "\ + f"{np.nanpercentile(fknees, 1):.4f} {np.nanmean(fknees):.4f} {np.nanpercentile(fknees, 99):.4f}"\ + f" {np.nanmax(fknees):.4f}") + logger.info(f"{experiment_data.nu}GHz: alphas {np.nanmin(alphas):.4f} "\ + f"{np.nanpercentile(alphas, 1):.4f} {np.nanmean(alphas):.4f} {np.nanpercentile(alphas, 99):.4f}"\ + f" {np.nanmax(alphas):.4f}") + logger.info(f"{experiment_data.nu}GHz: residuals {np.nanmin(residuals):.2e} "\ + f"{np.nanpercentile(residuals, 1):.2e} {np.nanmean(residuals):.2e} {np.nanpercentile(residuals, 99):.2e}"\ + f" {np.nanmax(residuals):.2e}") + logger.info(f"{experiment_data.nu}GHz: iterations {np.nanmin(niters):.4f} "\ + f"{np.nanpercentile(niters, 1):.4f} {np.nanmean(niters):.4f} {np.nanpercentile(niters, 99):.4f}"\ + f" {np.nanmax(niters):.4f}") start_bench("binned-mapmaker") @@ -664,17 +674,17 @@ def estimate_white_noise(experiment_data: DetGroupTOD, tod_samples: TODSamples, tod_samples (TODSamples): Updated TOD samples with sigma0 estimates. """ for iscan, scan in enumerate(experiment_data.scans): - for idetector, detector in enumerate(scan.detectors): - pix = detector.pix - psi = detector.psi if "QU" in experiment_data.pols else None + for idet, det in enumerate(scan.detectors): + pix, psi = det.get_pix_psi() # FIXME: Should maybe n_corr be subtracted here as well? - gain = tod_samples.gain(iscan, idetector) - sky_subtracted_tod = detector.tod.copy() + gain = tod_samples.gain(iscan, idet) + sky_subtracted_tod = det.tod.copy() sky_subtracted_tod -= gain*get_static_sky_TOD(det_compsep_map, pix, psi=psi) - sky_subtracted_tod -= gain*get_s_orb_TOD(detector, experiment_data, pix) - mask = detector.processing_mask_TOD + sky_subtracted_tod -= gain*get_s_orb_TOD(det, experiment_data, pix) + mask = det.full_mask sigma0 = calc_sigma0_robust(sky_subtracted_tod, mask) - tod_samples.noise_params[iscan,idetector,0] = sigma0 + logassert(sigma0 != 0, "sigma0 is 0, which should never happen.", logger) + tod_samples.noise_params[iscan,idet,0] = sigma0 if iscan == len(experiment_data.scans) - 1: log_memory("sigma0-est") return tod_samples @@ -712,8 +722,7 @@ def sample_absolute_gain(band_comm: MPI.Comm, experiment_data: DetGroupTOD, tod_ assert((ntod_down+1)*down_factor >= det.tod.shape[0]) - pix = det.pix # Only decompressing pix once for efficiency. - psi = det.psi + pix, psi = det.get_pix_psi() pix = pix[indices_centers] psi = psi[indices_centers] @@ -738,7 +747,7 @@ def sample_absolute_gain(band_comm: MPI.Comm, experiment_data: DetGroupTOD, tod_ residual_tod += tod_samples.abs_gain*s_orb # Now we can add back in the orbital dipole. ### Solving Equation 16 from BP7 ### - mask = det.processing_mask_TOD[indices_centers] + mask = det.full_mask[indices_centers] # White noise level, adjusted for downsampling. sigma0_effective = tod_samples.noise_params[iscan,idet,0] * np.sqrt(1.0/f_samp) # Inpaint masked regions with the calib signal times the absolute gain (+ white noise). @@ -815,8 +824,7 @@ def sample_relative_gain(band_comm: MPI.Comm, experiment_data: DetGroupTOD, ntod_down = indices_centers.size # Define the residual for this sampling step, as per Eq. (17) - pix = det.pix - psi = det.psi + pix, psi = det.get_pix_psi() pix = pix[indices_centers] psi = psi[indices_centers] @@ -828,7 +836,7 @@ def sample_relative_gain(band_comm: MPI.Comm, experiment_data: DetGroupTOD, residual_tod = det.tod[:ntod_down*down_factor].reshape((ntod_down, down_factor)) residual_tod = np.mean(residual_tod, axis=-1) residual_tod -= gain*s_cal - mask = det.processing_mask_TOD[indices_centers] + mask = det.full_mask[indices_centers] # sigma0 = calc_sigma0_robust(residual_tod, mask) # Setup FFT-based calculation for N^-1 operations @@ -942,8 +950,7 @@ def sample_temporal_gain_variations(band_comm: MPI.Comm, experiment_data: DetGro ntod_down = indices_centers.size # Per Eq. (26), the residual is d - (g0 + Delta_g)*s - pix = det.pix - psi = det.psi + pix, psi = det.get_pix_psi() pix = pix[indices_centers] psi = psi[indices_centers] @@ -955,7 +962,7 @@ def sample_temporal_gain_variations(band_comm: MPI.Comm, experiment_data: DetGro residual_tod = np.mean(residual_tod, axis=-1) residual_tod -= gain*s_cal - mask = det.processing_mask_TOD[indices_centers] + mask = det.full_mask[indices_centers] # FFT-based N^-1 operation setup # Ntod = residual_tod.shape[0] @@ -1176,14 +1183,14 @@ def process_tod(mpi_info: Bunch, experiment_data: DetGroupTOD, logger.info(f"Chain {chain} iter{iter} {experiment_data.nu}GHz: Finished temporal "\ f"gain estimation in {timing_dict['temp-gain']:.1f}s.") - ### WHITE NOISE ESTIMATION ### - t0 = time.time() - with benchmark("sigma0-est"): - tod_samples = estimate_white_noise(experiment_data, tod_samples, compsep_output, params) - timing_dict["wn-est-2"] = time.time() - t0 - if band_comm.Get_rank() == 0: - logger.info(f"Chain {chain} iter{iter} {experiment_data.nu}GHz: Finished white noise "\ - f"estimation in {timing_dict['wn-est-2']:.1f}s.") + # ### WHITE NOISE ESTIMATION ### + # t0 = time.time() + # with benchmark("sigma0-est"): + # tod_samples = estimate_white_noise(experiment_data, tod_samples, compsep_output, params) + # timing_dict["wn-est-2"] = time.time() - t0 + # if band_comm.Get_rank() == 0: + # logger.info(f"Chain {chain} iter{iter} {experiment_data.nu}GHz: Finished white noise "\ + # f"estimation in {timing_dict['wn-est-2']:.1f}s.") ### MAPMAKING ### do_ncorr_sampling = params.general.sample_corr_noise and iter >=\ From 6d640e023272332be1ab504a7042349e2ff366de Mon Sep 17 00:00:00 2001 From: jgslunde Date: Mon, 25 May 2026 11:09:33 +0200 Subject: [PATCH 03/23] Added dependencies --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 680dea1..ece3313 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,6 +29,8 @@ dependencies = [ "healpy", "tqdm", "camb", + "numpy-quaternion", + "qpoint", ] # Add new Python runtime dependencies here. [project.optional-dependencies] From e8ade143d55977877420254a03795df59ec8b0c9 Mon Sep 17 00:00:00 2001 From: jgslunde Date: Mon, 25 May 2026 11:11:04 +0200 Subject: [PATCH 04/23] Added support for a 'response' in the mapmaker, such that the I- and QU-sensitivity can be adjusted per-detector. --- src/commander4/utils/mapmaker.py | 80 +++++++++++++++++++-------- src/lib_cpp/ctypes/mapmaker.cpp | 95 ++++++++++++++++++++++++++++++++ 2 files changed, 152 insertions(+), 23 deletions(-) diff --git a/src/commander4/utils/mapmaker.py b/src/commander4/utils/mapmaker.py index 0d9a9d8..1f48976 100644 --- a/src/commander4/utils/mapmaker.py +++ b/src/commander4/utils/mapmaker.py @@ -147,6 +147,11 @@ def __init__(self, map_comm:MPI.Comm, nside:int, dtype=np.float32): self.maplib.map_accumulator_IQU_f64.argtypes = [ct_f64_dim2, ct_f64_dim1, ct.c_double, ct_i64_dim1, ct_f64_dim1, ct.c_int64, ct.c_int64] + self.maplib.map_accumulator_IQU_response_f64.argtypes = [ct_f64_dim2, ct_f64_dim1, + ct.c_double, ct_i64_dim1, + ct_f64_dim1, ct.c_double, + ct.c_double, ct.c_int64, + ct.c_int64] self.maplib.map_solve_IQU_f64.argtypes = [ct_f64_dim2, ct_f64_dim2, ct_f64_dim2, ct.c_int64] @@ -158,7 +163,8 @@ def final_map(self): return self._finalized_map - def accumulate_to_map(self, tod:NDArray, weights:NDArray, pix:NDArray, psi:NDArray): + def accumulate_to_map(self, tod:NDArray, weights:NDArray, pix:NDArray, psi:NDArray, + response: NDArray | None = None): """Accumulate I,Q,U signal into the local map buffer.""" # Check that we are still in business, and haven't already called "gather_map". logassert(self._map_signal is not None, "Tried accumulating to finalized map", logger) @@ -166,11 +172,19 @@ def accumulate_to_map(self, tod:NDArray, weights:NDArray, pix:NDArray, psi:NDArr tod_f64 = np.ascontiguousarray(tod, dtype=np.float64) weight_f64 = float(weights) psi_f64 = np.ascontiguousarray(psi, dtype=np.float64) - self.maplib.map_accumulator_IQU_f64(self._map_signal, tod_f64, weight_f64, - pix.astype(np.int64, copy=False), psi_f64, ntod, - self.npix) - - def accumulate_to_map_Python(self, tod:NDArray, weights:NDArray, pix:NDArray, psi:NDArray): + pix_i64 = pix.astype(np.int64, copy=False) + if response is None: + self.maplib.map_accumulator_IQU_f64(self._map_signal, tod_f64, weight_f64, + pix_i64, psi_f64, ntod, self.npix) + else: + response_I = float(response[0]) + response_QU = float(response[1]) + self.maplib.map_accumulator_IQU_response_f64(self._map_signal, tod_f64, weight_f64, + pix_i64, psi_f64, response_I, + response_QU, ntod, self.npix) + + def accumulate_to_map_Python(self, tod:NDArray, weights:NDArray, pix:NDArray, psi:NDArray, + response: NDArray | None = None): """Reference accumulator matching the ctypes IQU implementation.""" # Reference implementation matching the ctypes IQU accumulator. logassert(self._map_signal is not None, "Tried accumulating to finalized map", logger) @@ -179,9 +193,14 @@ def accumulate_to_map_Python(self, tod:NDArray, weights:NDArray, pix:NDArray, ps ang = 2.0 * np.ascontiguousarray(psi, dtype=np.float64) c2 = np.cos(ang) s2 = np.sin(ang) - np.add.at(self._map_signal[0], pix_idx, w_tod) - np.add.at(self._map_signal[1], pix_idx, w_tod * c2) - np.add.at(self._map_signal[2], pix_idx, w_tod * s2) + if response is None: + response_I, response_QU = 1.0, 1.0 + else: + response_I = float(response[0]) + response_QU = float(response[1]) + np.add.at(self._map_signal[0], pix_idx, w_tod * response_I) + np.add.at(self._map_signal[1], pix_idx, w_tod * response_QU * c2) + np.add.at(self._map_signal[2], pix_idx, w_tod * response_QU * s2) def gather_map(self): """Reduce the local IQU buffers across MPI ranks into the root map.""" @@ -270,8 +289,10 @@ def __init__(self, map_comm:MPI.Comm, nside:int, dtype=np.float32): ct_f64_dim1 = np.ctypeslib.ndpointer(dtype=ct.c_double, ndim=1, flags="contiguous") ct_f64_dim2 = np.ctypeslib.ndpointer(dtype=ct.c_double, ndim=2, flags="contiguous") self.maplib.map_weight_accumulator_IQU_f64.argtypes = [ct_f64_dim2, ct.c_double, - ct_i64_dim1, ct_f64_dim1, ct.c_int64, - ct.c_int64] + ct_i64_dim1, ct_f64_dim1, ct.c_int64, ct.c_int64] + self.maplib.map_weight_accumulator_IQU_response_f64.argtypes = [ct_f64_dim2, + ct.c_double, ct_i64_dim1, ct_f64_dim1, ct.c_double, + ct.c_double, ct.c_int64, ct.c_int64] self.maplib.map_invdiag_IQU_f64.argtypes = [ct_f64_dim2, ct_f64_dim2, ct.c_int64] @property @@ -288,18 +309,26 @@ def final_cov_map(self): logger) return self._gathered_map - def accumulate_to_map(self, weight:float, pix:NDArray, psi:NDArray): + def accumulate_to_map(self, weight:float, pix:NDArray, psi:NDArray, response:NDArray | None = None): """Accumulate IQU weight/covariance elements into the local buffer.""" # Check that we are still in business, and haven't already called "gather_map". logassert(self._map_signal is not None, "Tried accumulating to finalized map", logger) ntod = pix.shape[0] weight_f64 = float(weight) psi_f64 = np.ascontiguousarray(psi, dtype=np.float64) - self.maplib.map_weight_accumulator_IQU_f64(self._map_signal, weight_f64, - pix.astype(np.int64, copy=False), psi_f64, ntod, - self.npix) - - def accumulate_to_map_Python(self, weight:float, pix:NDArray, psi:NDArray): + pix_i64 = pix.astype(np.int64, copy=False) + if response is None: + self.maplib.map_weight_accumulator_IQU_f64(self._map_signal, weight_f64, + pix_i64, psi_f64, ntod, self.npix) + else: + response_I = float(response[0]) + response_QU = float(response[1]) + self.maplib.map_weight_accumulator_IQU_response_f64(self._map_signal, weight_f64, + pix_i64, psi_f64, response_I, + response_QU, ntod, self.npix) + + def accumulate_to_map_Python(self, weight:float, pix:NDArray, psi:NDArray, + response: NDArray | None = None): """Reference accumulator matching the ctypes IQU weights implementation.""" # Reference implementation matching the ctypes IQU weight accumulator. logassert(self._map_signal is not None, "Tried accumulating to finalized map", logger) @@ -308,12 +337,17 @@ def accumulate_to_map_Python(self, weight:float, pix:NDArray, psi:NDArray): c2 = np.cos(ang) s2 = np.sin(ang) weight_f64 = float(weight) - np.add.at(self._map_signal[0], pix_idx, weight_f64) - np.add.at(self._map_signal[1], pix_idx, weight_f64 * c2) - np.add.at(self._map_signal[2], pix_idx, weight_f64 * s2) - np.add.at(self._map_signal[3], pix_idx, weight_f64 * c2 * c2) - np.add.at(self._map_signal[4], pix_idx, weight_f64 * s2 * c2) - np.add.at(self._map_signal[5], pix_idx, weight_f64 * s2 * s2) + if response is None: + response_I, response_QU = 1.0, 1.0 + else: + response_I = float(response[0]) + response_QU = float(response[1]) + np.add.at(self._map_signal[0], pix_idx, weight_f64 * response_I * response_I) + np.add.at(self._map_signal[1], pix_idx, weight_f64 * response_I * response_QU * c2) + np.add.at(self._map_signal[2], pix_idx, weight_f64 * response_I * response_QU * s2) + np.add.at(self._map_signal[3], pix_idx, weight_f64 * response_QU * response_QU * c2 * c2) + np.add.at(self._map_signal[4], pix_idx, weight_f64 * response_QU * response_QU * s2 * c2) + np.add.at(self._map_signal[5], pix_idx, weight_f64 * response_QU * response_QU * s2 * s2) @property def inv_N_diag(self): diff --git a/src/lib_cpp/ctypes/mapmaker.cpp b/src/lib_cpp/ctypes/mapmaker.cpp index a40fc2a..afd6eea 100644 --- a/src/lib_cpp/ctypes/mapmaker.cpp +++ b/src/lib_cpp/ctypes/mapmaker.cpp @@ -67,6 +67,28 @@ void _map_accumulator_IQU_T(T *map, const T *tod, const T weight, int64_t *pix, } } +/** Response-scaled version of _map_accumulator_IQU_T. + * + * Uses the effective pointing row [response_I, response_QU*cos(2 psi), + * response_QU*sin(2 psi)] for each sample. + */ +template +void _map_accumulator_IQU_response_T(T *map, const T *tod, const T weight, + int64_t *pix, const double *psi, + const double response_I, + const double response_QU, + int64_t scan_len, int64_t num_pix){ + const T response_I_T = static_cast(response_I); + const T response_QU_T = static_cast(response_QU); + for(int64_t itod=0; itod(std::cos(2.0 * psi[itod])); + const T sin2psi = static_cast(std::sin(2.0 * psi[itod])); + map[pix[itod]] += tod[itod] * weight * response_I_T; // I + map[pix[itod] + num_pix] += tod[itod] * weight * response_QU_T * cos2psi; // Q + map[pix[itod] + 2*num_pix] += tod[itod] * weight * response_QU_T * sin2psi; // U + } +} + /** Simple serial transpose of the mapmaker operator for only I, which accumulates the values on the TOD, given a certain pointing on a map. * * Notes: @@ -137,6 +159,35 @@ void _map_weight_accumulator_IQU_T(T *map, const T weight, int64_t *pix, const d } } +/** Response-scaled version of _map_weight_accumulator_IQU_T. + * + * Uses the effective pointing row [response_I, response_QU*cos(2 psi), + * response_QU*sin(2 psi)] and accumulates its weighted outer product. + */ +template +void _map_weight_accumulator_IQU_response_T(T *map, const T weight, + int64_t *pix, const double *psi, + const double response_I, + const double response_QU, + int64_t scan_len, + int64_t num_pix){ + const T response_I_T = static_cast(response_I); + const T response_QU_T = static_cast(response_QU); + const T response_I_sq = response_I_T * response_I_T; + const T response_I_QU = response_I_T * response_QU_T; + const T response_QU_sq = response_QU_T * response_QU_T; + for(int64_t itod=0; itod(std::cos(2.0 * psi[itod])); + const T sin2psi = static_cast(std::sin(2.0 * psi[itod])); + map[pix[itod]] += weight * response_I_sq; // II + map[pix[itod] + num_pix] += weight * response_I_QU * cos2psi; // IQ + map[pix[itod] + 2*num_pix] += weight * response_I_QU * sin2psi; // IU + map[pix[itod] + 3*num_pix] += weight * response_QU_sq * cos2psi*cos2psi; // QQ + map[pix[itod] + 4*num_pix] += weight * response_QU_sq * sin2psi*cos2psi; // QU + map[pix[itod] + 5*num_pix] += weight * response_QU_sq * sin2psi*sin2psi; // UU + } +} + /** Invert a symmetric positive definite (SPD) 3x3 matrix. * Accepts the 6 (out of 9) unique elements of such a matrix as arguments, @@ -331,6 +382,26 @@ void map_accumulator_IQU_f64(double *map, double *tod, double weight, int64_t *p _map_accumulator_IQU_T(map, tod, weight, pix, psi, scan_len, num_pix); } +extern "C" +void map_accumulator_IQU_response_f32(float *map, float *tod, float weight, + int64_t *pix, double *psi, + double response_I, double response_QU, + int64_t scan_len, int64_t num_pix){ + _map_accumulator_IQU_response_T(map, tod, weight, pix, psi, + response_I, response_QU, + scan_len, num_pix); +} + +extern "C" +void map_accumulator_IQU_response_f64(double *map, double *tod, double weight, + int64_t *pix, double *psi, + double response_I, double response_QU, + int64_t scan_len, int64_t num_pix){ + _map_accumulator_IQU_response_T(map, tod, weight, pix, psi, + response_I, response_QU, + scan_len, num_pix); +} + extern "C" void map2tod_f64(double *map, double *tod, int64_t *pix, int64_t scan_len){ _map2tod_T(map, tod, pix, scan_len); @@ -361,6 +432,30 @@ void map_weight_accumulator_IQU_f64(double *map, double weight, int64_t *pix, do _map_weight_accumulator_IQU_T(map, weight, pix, psi, scan_len, num_pix); } +extern "C" +void map_weight_accumulator_IQU_response_f32(float *map, float weight, + int64_t *pix, double *psi, + double response_I, + double response_QU, + int64_t scan_len, + int64_t num_pix){ + _map_weight_accumulator_IQU_response_T(map, weight, pix, psi, + response_I, response_QU, + scan_len, num_pix); +} + +extern "C" +void map_weight_accumulator_IQU_response_f64(double *map, double weight, + int64_t *pix, double *psi, + double response_I, + double response_QU, + int64_t scan_len, + int64_t num_pix){ + _map_weight_accumulator_IQU_response_T(map, weight, pix, psi, + response_I, response_QU, + scan_len, num_pix); +} + extern "C" void map_solve_IQU_f32(float *map_out, float *map_rhs, float *norm_map, int64_t num_pix){ _map_solve_IQU_T(map_out, map_rhs, norm_map, num_pix); From 4ce7265d886c92c589f8b8b826dd55b0f6567dfd Mon Sep 17 00:00:00 2001 From: jgslunde Date: Mon, 25 May 2026 11:11:44 +0200 Subject: [PATCH 05/23] Added Huffman decompression support across a wider range of types. --- src/lib_cpp/utils_pymod.cc | 34 ++++++++++++++++++++++++++++------ 1 file changed, 28 insertions(+), 6 deletions(-) diff --git a/src/lib_cpp/utils_pymod.cc b/src/lib_cpp/utils_pymod.cc index 7889c95..5370b2f 100644 --- a/src/lib_cpp/utils_pymod.cc +++ b/src/lib_cpp/utils_pymod.cc @@ -298,14 +298,35 @@ template static NpArr Py2_huffman_decode(const CNpArr &bytes_, return out_; } +template static NpArr Py3_huffman_decode(const CNpArr &bytes, + const CNpArr &tree, const CNpArr &symb, const NpArr &out, + const char *dtype_descr) + { + MR_assert(isPyarr(out), "type mismatch: 'out' must have the same dtype as 'symb' (", + dtype_descr, ")"); + return Py2_huffman_decode(bytes, tree, symb, out); + } + static NpArr Py_huffman_decode(const CNpArr &bytes, const CNpArr &tree, const CNpArr &symb, const NpArr &out) { - if (isPyarr(symb)) - return Py2_huffman_decode (bytes, tree, symb, out); + if (isPyarr(symb)) + return Py3_huffman_decode(bytes, tree, symb, out, "i1"); + if (isPyarr(symb)) + return Py3_huffman_decode(bytes, tree, symb, out, "u1"); + if (isPyarr(symb)) + return Py3_huffman_decode(bytes, tree, symb, out, "i2"); + if (isPyarr(symb)) + return Py3_huffman_decode(bytes, tree, symb, out, "u2"); if (isPyarr(symb)) - return Py2_huffman_decode (bytes, tree, symb, out); - MR_fail("type matching failed: 'symb' has neither type 'i8' nor 'i4'"); + return Py3_huffman_decode(bytes, tree, symb, out, "i4"); + if (isPyarr(symb)) + return Py3_huffman_decode(bytes, tree, symb, out, "u4"); + if (isPyarr(symb)) + return Py3_huffman_decode(bytes, tree, symb, out, "i8"); + if (isPyarr(symb)) + return Py3_huffman_decode(bytes, tree, symb, out, "u8"); + MR_fail("type mismatch: 'symb' must have an integer dtype among 'i1', 'u1', 'i2', 'u2', 'i4', 'u4', 'i8', or 'u8'"); } constexpr const char *Py_huffman_decode_DS = R"""( @@ -317,11 +338,12 @@ bytes: numpy.ndarray((nbytes,), dtype=np.uint8) the bit stream tree: numpy.ndarray((ntree,), dtype=np.int64) the tree array -symb: numpy.ndarray((nsymb,), dtype=np.int64 or np.int32) +symb: numpy.ndarray((nsymb,), dtype any signed or unsigned 8/16/32/64-bit integer type) the array of possible symbols in the stream out: numpy.ndarray((ndata,), dtype identical to that of symb) the array into which the uncopressed data is written - The size of this array *must* match the number of decoded symbols! + The size of this array *must* match the number of decoded symbols! + The dtype of this array *must* be identical to that of symb. Returns ------- From ef778a22fc128b257754519945541ebbde094603 Mon Sep 17 00:00:00 2001 From: jgslunde Date: Mon, 25 May 2026 13:53:19 +0200 Subject: [PATCH 06/23] Added SO-LAT to list of TOD readers --- src/commander4/tod_reader.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/commander4/tod_reader.py b/src/commander4/tod_reader.py index 3368993..04ae15f 100644 --- a/src/commander4/tod_reader.py +++ b/src/commander4/tod_reader.py @@ -11,6 +11,7 @@ from commander4.experiments.planck.tod_reader_planck_sim import tod_reader as tod_reader_planck_sim from commander4.experiments.akari.tod_reader_akari import tod_reader as tod_reader_akari from commander4.experiments.SO.tod_reader_SO_LAT import tod_reader as tod_reader_SO_LAT +from commander4.experiments.SO.tod_reader_SO_SAT import tod_reader as tod_reader_SO_SAT # Dictionary containing known experiments and the location of their TOD reading scripts. # The `experiment_id`` field in the parameter file decides what TOD reader is used in this dict. @@ -22,6 +23,7 @@ "litebird_sim_spawndetectors" : tod_reader_litebird_sim_spawndetectors, "akari" : tod_reader_akari, "SO_LAT" : tod_reader_SO_LAT, + "SO_SAT" : tod_reader_SO_SAT, } def read_tods_from_file(band_comm: MPI.Comm, my_experiment: Bunch, my_band: Bunch, my_det: Bunch, From e99ac6826df8fb44558e1fc08282a7f14fe262c4 Mon Sep 17 00:00:00 2001 From: jgslunde Date: Mon, 25 May 2026 13:54:13 +0200 Subject: [PATCH 07/23] Removed depricated test and added a new one. --- tests/test_huffman_decode.py | 40 ++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) create mode 100644 tests/test_huffman_decode.py diff --git a/tests/test_huffman_decode.py b/tests/test_huffman_decode.py new file mode 100644 index 0000000..b7b8626 --- /dev/null +++ b/tests/test_huffman_decode.py @@ -0,0 +1,40 @@ +import numpy as np +import pytest +from numpy.testing import assert_array_equal + +from commander4.cmdr4_support import utils as cpp_utils +from commander4.compression import huffman + + +def test_huffman_decode_roundtrip_int8() -> None: + values = np.array([-8, -3, 7, -3, 2, 7, -8, 2], dtype=np.int8) + tree, symb, sym_codes, sym_lengths = huffman.build_huffman_tree([values]) + encoded = np.frombuffer(huffman.huffman_compress_array(values, sym_codes, sym_lengths), dtype=np.uint8) + out = np.empty(values.size, dtype=values.dtype) + + decoded = cpp_utils.huffman_decode(encoded, tree, symb, out) + + assert decoded.dtype == values.dtype + assert_array_equal(decoded, values) + + +def test_huffman_decode_roundtrip_uint8() -> None: + values = np.array([0, 255, 3, 255, 1, 0, 7, 1], dtype=np.uint8) + tree, symb, sym_codes, sym_lengths = huffman.build_huffman_tree([values]) + encoded = np.frombuffer(huffman.huffman_compress_array(values, sym_codes, sym_lengths), dtype=np.uint8) + out = np.empty(values.size, dtype=values.dtype) + + decoded = cpp_utils.huffman_decode(encoded, tree, symb, out) + + assert decoded.dtype == values.dtype + assert_array_equal(decoded, values) + + +def test_huffman_decode_requires_out_dtype_to_match_symb() -> None: + values = np.array([0, 255, 3, 255, 1, 0, 7, 1], dtype=np.uint8) + tree, symb, sym_codes, sym_lengths = huffman.build_huffman_tree([values]) + encoded = np.frombuffer(huffman.huffman_compress_array(values, sym_codes, sym_lengths), dtype=np.uint8) + out = np.empty(values.size, dtype=np.int16) + + with pytest.raises(RuntimeError, match="'out' must have the same dtype as 'symb'"): + cpp_utils.huffman_decode(encoded, tree, symb, out) \ No newline at end of file From 525ced54c65aa4f453a0985ee8f0b03e02371c23 Mon Sep 17 00:00:00 2001 From: jgslunde Date: Mon, 25 May 2026 16:12:41 +0200 Subject: [PATCH 08/23] Renamed some stuff --- src/commander4/data_models/detector_TOD.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/commander4/data_models/detector_TOD.py b/src/commander4/data_models/detector_TOD.py index fa82f26..2fe7017 100644 --- a/src/commander4/data_models/detector_TOD.py +++ b/src/commander4/data_models/detector_TOD.py @@ -41,7 +41,7 @@ def __init__( huffman_tree2: NDArray | None = None, huffman_symbols2: NDArray | None = None, flag_encoded: NDArray[np.integer] | bytes | None = None, - flag_bitmask: int | None = None, + bad_data_bitmask: int | None = None, init_scalars: NDArray | None = None, tod_is_compressed: bool = True, flag_is_compressed: bool = True, @@ -98,19 +98,20 @@ def __init__( self.fsamp = fsamp self.init_scalars = init_scalars self._flag_encoded = flag_encoded - self._flag_bitmask = flag_bitmask + self._bad_data_bitmask = bad_data_bitmask self._huffman_tree = huffman_tree self._huffman_symbols = huffman_symbols self._huffman_tree2 = huffman_tree2 self._huffman_symbols2 = huffman_symbols2 self._tod_is_compressed = tod_is_compressed self._flag_is_compressed = flag_is_compressed + self.processing_mask_map = processing_mask_map self.pointing = pointing processing_mask = processing_mask_map[self.get_pix()] - self._processing_mask_TOD = np.packbits(processing_mask) + self._processing_mask = np.packbits(processing_mask) self.det_response = det_response - if flag_encoded is not None and flag_bitmask is not None: - bad_data_mask = ~(self.flag & flag_bitmask) + if flag_encoded is not None and bad_data_bitmask is not None: + bad_data_mask = ~(self.flag & bad_data_bitmask) self._bad_data_mask = np.packbits(bad_data_mask) self._full_mask = np.packbits(bad_data_mask & processing_mask) if orb_dir_vec is not None: @@ -166,12 +167,12 @@ def flag(self) -> NDArray[np.integer]: return flag[:self.ntod] @property - def processing_mask_TOD(self) -> NDArray[np.bool_]: + def processing_mask(self) -> NDArray[np.bool_]: """Boolean mask selecting valid (unmasked) TOD samples. Stored internally as a packed bit array and unpacked on each access. """ - mask = np.unpackbits(self._processing_mask_TOD).view(bool) + mask = np.unpackbits(self._processing_mask).view(bool) if mask.size > self.tod.size + 7 or mask.size < self.tod.size: # The bytearray is stored in multiples of 8, so it can be up to 7 elements # longer than the TOD. If it's even longer or shorter, something is wrong. From 424bbea676ed892f469ff4c379d90401c4dd5519 Mon Sep 17 00:00:00 2001 From: jgslunde Date: Wed, 27 May 2026 19:16:21 +0200 Subject: [PATCH 09/23] Fixed a bug where number of threads was incorrectly set in FFT --- src/commander4/utils/math_operations.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/commander4/utils/math_operations.py b/src/commander4/utils/math_operations.py index f089cc0..ea20319 100644 --- a/src/commander4/utils/math_operations.py +++ b/src/commander4/utils/math_operations.py @@ -248,7 +248,7 @@ def complist_norm(comp_list:list[Component]) -> float: ###### GENERAL MATH STUFF ###### -def forward_rfft(data:NDArray[np.floating], nthreads:int = 1): +def forward_rfft(data:NDArray[np.floating], nthreads:int = None) -> NDArray: """ Forward real Fourier transform, equivalent to scipy.fft.rfft. Args: data (np.array): Real-valued data array to be Fourier transformed. @@ -260,7 +260,7 @@ def forward_rfft(data:NDArray[np.floating], nthreads:int = 1): nthreads = int(os.environ.get("OMP_NUM_THREADS", 1)) if nthreads is None else nthreads return ducc0.fft.r2c(data, nthreads=nthreads) -def backward_rfft(data_f:NDArray, ntod:int, nthreads:int = None) -> NDArray[np.floating]: +def backward_rfft(data_f:NDArray, ntod:int, nthreads:int = None) -> NDArray: """ Backward real Fourier transform, equivalent to scipy.fft.irfft. Args: data_f (np.array): Complex Fourier coefficients to be converted back to real data. @@ -278,7 +278,7 @@ def backward_rfft(data_f:NDArray, ntod:int, nthreads:int = None) -> NDArray[np.f return ducc0.fft.c2r(data_f, lastsize=ntod, forward=False, nthreads=nthreads, inorm=2) -def forward_rfft_mirrored(data: NDArray, nthreads: int = 1): +def forward_rfft_mirrored(data: NDArray, nthreads: int = None) -> NDArray: """Forward real FFT on a mirrored (reflected) copy of the input. The input is mirrored so that ``dt[0:ntod] = data``, From db948de93f27a240fe84b38b40a1b0f13b6c8463 Mon Sep 17 00:00:00 2001 From: jgslunde Date: Wed, 27 May 2026 22:17:00 +0200 Subject: [PATCH 10/23] Few different changes to support variable amount of detectors (some detectors are discarded on read-in). --- src/commander4/data_models/TOD_samples.py | 75 ++++++++++--------- src/commander4/data_models/detector_TOD.py | 86 +++++++++++++++++----- src/commander4/tod_processing.py | 4 +- 3 files changed, 107 insertions(+), 58 deletions(-) diff --git a/src/commander4/data_models/TOD_samples.py b/src/commander4/data_models/TOD_samples.py index 2cb76ad..cccf0fe 100644 --- a/src/commander4/data_models/TOD_samples.py +++ b/src/commander4/data_models/TOD_samples.py @@ -68,6 +68,7 @@ class TODSamples: def __init__(self, experiment_data: DetGroupTOD, params: Bunch, + my_band: Bunch, band_comm: MPI.Comm, chain: int, noise_psd_class: str = "oof", @@ -92,45 +93,43 @@ def __init__(self, if self.band_comm.Get_rank() == 0: logger.info(f"Band {self.band_name} initializing TOD samples from default values.") - all_det_gains = [] - myband_noise_params = None - - for exp_name in params.experiments: - experiment = params.experiments[exp_name] - for iband, band_name in enumerate(experiment.bands): - band = experiment.bands[band_name] - # Fixed typo here: self.band.name -> self.band_name - if self.experiment_name == exp_name and self.band_name == band_name: - # Decide how to set initial noise parmams. - if "initial_noise_params" in band: - # Option 1: They are specified in the parameter file. - myband_noise_params = np.array(band.initial_noise_params) - elif experiment_data.scans[0].detectors[0].init_scalars is not None: - # Option 2: There were entries in the read-in files. - myband_noise_params = np.array([[det.init_scalars[1:4] for det in scan.detectors] for scan in experiment_data.scans]) - else: - # Option 3: Fallback to sensible defaults. - myband_noise_params = np.array([1e-3, 0.1, -1.0]) - # Loop over all detectors to record their initial gain values. - for idet, det_name in enumerate(band.detectors): - detector = band.detectors[det_name] - # Decide how to set initial gain values. - if "gain_est" in detector: - all_det_gains.append(detector.gain_est) - # TODO: Make this work - # elif experiment_data.scans[0].detectors[0].init_scalars is not None: - # # FIXME: This entry seemed to be off by 1e-6 compared to my - # # estimates, but is this always true? Needs to be checked! - # all_det_gains.append(1e-6*experiment_data.scans[0].detectors[].init_scalars[0]) - else: - all_det_gains.append(1.0) - - all_det_gains = np.array(all_det_gains) - abs_gain = float(np.mean(all_det_gains)) - self.abs_gain = abs_gain - self.rel_gain = all_det_gains - abs_gain + self.noise_params = np.zeros((self.nscans, self.ndet, 3)) + np.nan + self.abs_gain = 0.0 + self.rel_gain = np.zeros((self.ndet)) self.temporal_gain = np.zeros((self.nscans, self.ndet)) - self.noise_params = np.full((self.nscans, self.ndet, 3), myband_noise_params) + + all_det_gains = np.zeros((self.nscans, self.ndet)) + # all_det_gains = [] + # myband_noise_params = None + + if "initial_noise_params" in my_band: + # Option 1: They are specified in the parameter file. + self.noise_params[:] = np.array(my_band.initial_noise_params) + elif experiment_data.scans[0].detectors[0].init_scalars is not None: + # Option 2: There were entries in the read-in files. + for iscan, scan in enumerate(experiment_data.scans): + for idet, det in enumerate(scan.detectors): + self.noise_params[iscan,idet] = det.init_scalars[1:] + else: + # Option 3: Fallback to sensible defaults. + logger.warning("Did not find initial noise parameters, falling back to sensible defaults.") + self.noise_params[:] = np.array([1e-3, 0.1, -1.0]) + + if "gain" in my_band.detectors[experiment_data.scans[0].detectors[0].name]: + for iscan, scan in enumerate(experiment_data.scans): + for idet, det in enumerate(scan.detectors): + all_det_gains[iscan,idet] = my_band[det.name].gain + elif experiment_data.scans[0].detectors[0].init_scalars is not None: + for iscan, scan in enumerate(experiment_data.scans): + for idet, det in enumerate(scan.detectors): + all_det_gains[iscan,idet] = det.init_scalars[0] + else: + raise ValueError("Did not find initial gain value in input files.") + + all_det_gains = np.array(all_det_gains) + self.abs_gain = float(np.nanmean(all_det_gains)) + self.rel_gain = np.nanmean(all_det_gains, axis=0) - self.abs_gain + self.temporal_gain = all_det_gains - self.rel_gain - self.abs_gain else: # --------------------------------------------------------- diff --git a/src/commander4/data_models/detector_TOD.py b/src/commander4/data_models/detector_TOD.py index 2fe7017..7fcf189 100644 --- a/src/commander4/data_models/detector_TOD.py +++ b/src/commander4/data_models/detector_TOD.py @@ -20,16 +20,19 @@ class DetectorTOD: data nside, and sampling frequency are stored as plain public attributes. Attributes: + name (str): Unique name of this detector. tod (NDArray[np.floating]): 1-D array of calibrated time-ordered samples. ntod (int): Number of time samples (after any Fourier-length cropping). nside (int): HEALPix nside at which this detector should be evaluated. data_nside (int): HEALPix nside at which the pixel indices are stored on disk. fsamp (float): Sampling frequency in Hz. - init_scalars (array): 4-element array of initial guesses for gain, sigma0, fknee, and alpha. """ def __init__( self, - tod: NDArray[np.floating], + name: str, + det_idx_fullband: int, + det_idx_local: int, + tod: NDArray[np.floating] | bytes | np.void, pointing: PixelPointing | DetectorBoresightPointing, fsamp: float, orb_dir_vec: NDArray[np.floating] | None, @@ -40,7 +43,7 @@ def __init__( ntod_optimal: int, huffman_tree2: NDArray | None = None, huffman_symbols2: NDArray | None = None, - flag_encoded: NDArray[np.integer] | bytes | None = None, + flag_encoded: NDArray[np.integer] | bytes | np.void | None = None, bad_data_bitmask: int | None = None, init_scalars: NDArray | None = None, tod_is_compressed: bool = True, @@ -50,7 +53,12 @@ def __init__( """Construct a DetectorTOD. Args: - tod: 1-D floating-point array of calibrated time samples. + name: Unique name of the current detector. + det_idx_fullband: Unique detector-index among all the detectors on the relevant band. + det_idx_local: Unique detector-index among all detectors in the current scan, where + some detectors might be missing from the full set of detectors on the band. + tod: Calibrated time samples, either as a decoded 1-D floating-point + array or as a compressed binary payload. pointing: Pointing representation for this detector. Must be a ``PixelPointing`` or ``DetectorBoresightPointing`` instance. fsamp: Sampling frequency in Hz. @@ -60,10 +68,22 @@ def __init__( huffman_symbols: Huffman symbol table for the flag stream, or None. processing_mask_map: Boolean HEALPix map selecting valid pixels. ntod_original: Original TOD length before Fourier-length cropping. - flag_encoded: Huffman-encoded flag array, or None. + flag_encoded: Flag samples, either decoded or Huffman-encoded, or None. flag_bitmask: Bitmask applied to flags to identify excluded samples. """ - if not tod_is_compressed: + if tod_is_compressed: + log.logassert_np( + isinstance(tod, (bytes, np.void)), + "Compressed TOD must be provided as bytes or numpy.void.", + logger, + ) + log.logassert_np( + huffman_tree2 is not None and huffman_symbols2 is not None, + "Compressed TOD requires Huffman metadata.", + logger, + ) + else: + log.logassert_np(isinstance(tod, np.ndarray), "'tod' must be a numpy array.", logger) log.logassert_np(tod.ndim==1, "'value' must be a 1D array", logger) log.logassert_np(tod.dtype in [np.float64,np.float32], "TOD dtype must be floating "\ f"type, is {tod.dtype}", logger) @@ -84,20 +104,50 @@ def __init__( "Pointing ntod does not match DetectorTOD ntod_optimal.", logger, ) - if flag_encoded is not None and flag_is_compressed: - log.logassert_np( - huffman_tree is not None and huffman_symbols is not None, - "Compressed flags require Huffman metadata.", - logger, - ) - self._tod = tod + if flag_encoded is not None: + if flag_is_compressed: + log.logassert_np( + isinstance(flag_encoded, (bytes, np.void)), + "Compressed flags must be provided as bytes or numpy.void.", + logger, + ) + log.logassert_np( + huffman_tree is not None and huffman_symbols is not None, + "Compressed flags require Huffman metadata.", + logger, + ) + else: + log.logassert_np( + isinstance(flag_encoded, np.ndarray), + "Decoded flags must be provided as a numpy array.", + logger, + ) + log.logassert_np(flag_encoded.ndim == 1, "'flag' must be a 1D array.", logger) + log.logassert_np( + np.issubdtype(flag_encoded.dtype, np.integer), + "Decoded flags must have integer dtype.", + logger, + ) + log.logassert_np( + flag_encoded.size >= ntod_optimal, + f"'flag' length {flag_encoded.size} is shorter than ntod {ntod_optimal}.", + logger, + ) + self.name = name + self.det_idx_fullband = det_idx_fullband + self.det_idx_local = det_idx_local + self._tod = np.frombuffer(tod, dtype=np.uint8) if tod_is_compressed else tod self.ntod_original = ntod_original self.ntod = ntod_optimal self.nside = pointing.nside self.data_nside = pointing.data_nside self.fsamp = fsamp self.init_scalars = init_scalars - self._flag_encoded = flag_encoded + self._flag_encoded = ( + np.frombuffer(flag_encoded, dtype=np.uint8) + if flag_encoded is not None and flag_is_compressed + else flag_encoded + ) self._bad_data_bitmask = bad_data_bitmask self._huffman_tree = huffman_tree self._huffman_symbols = huffman_symbols @@ -125,13 +175,13 @@ def __init__( def tod(self) -> NDArray[np.floating]: if self._tod_is_compressed: tod = np.zeros(self.ntod_original, dtype=self._huffman_symbols2.dtype) - tod[:] = cpp_utils.huffman_decode(np.frombuffer(self._tod, dtype=np.uint8), - self._huffman_tree2, self._huffman_symbols2, tod)[:self.ntod] + tod[:] = cpp_utils.huffman_decode(self._tod, + self._huffman_tree2, self._huffman_symbols2, tod) tod[:] = np.cumsum(tod) tod = tod.astype(np.float32) else: tod = self._tod - return tod + return tod[:self.ntod] def get_pix(self, nside: int | None = None) -> NDArray[np.integer]: @@ -159,7 +209,7 @@ def get_pix_psi( def flag(self) -> NDArray[np.integer]: if self._flag_is_compressed: flag = np.zeros(self.ntod_original, dtype=self._huffman_symbols.dtype) - flag = cpp_utils.huffman_decode(np.frombuffer(self._flag_encoded, dtype=np.uint8), + flag = cpp_utils.huffman_decode(self._flag_encoded, self._huffman_tree, self._huffman_symbols, flag) flag = np.cumsum(flag) else: diff --git a/src/commander4/tod_processing.py b/src/commander4/tod_processing.py index 96c6372..a366a90 100644 --- a/src/commander4/tod_processing.py +++ b/src/commander4/tod_processing.py @@ -628,8 +628,8 @@ def init_tod_processing(mpi_info: Bunch, params: Bunch) -> tuple[Bunch, str, Det if mpi_info.tod.is_master: logger.info(f"TOD: Finished reading all files in {time.time()-t0:.1f}s.") - tod_samples_chain1 = TODSamples(experiment_data, params, band_comm, 1) - tod_samples_chain2 = TODSamples(experiment_data, params, band_comm, 2) + tod_samples_chain1 = TODSamples(experiment_data, params, my_band, band_comm, 1) + tod_samples_chain2 = TODSamples(experiment_data, params, my_band, band_comm, 2) # Creating "tod_band_masters", an array which maps the band index to the rank of the master # of that band. From 9bb0be6f3a2ea6859a721ca6067a34cbbb712783 Mon Sep 17 00:00:00 2001 From: jgslunde Date: Wed, 27 May 2026 22:22:50 +0200 Subject: [PATCH 11/23] Some tweaks to how compressed entries are stored, and some changes to masking calculations --- src/commander4/data_models/detector_TOD.py | 15 ++++-- src/commander4/data_models/pointing.py | 60 +++++++++++++++++----- src/commander4/data_models/scan_TOD.py | 3 -- src/commander4/tod_processing.py | 60 ++++++++++++---------- 4 files changed, 91 insertions(+), 47 deletions(-) diff --git a/src/commander4/data_models/detector_TOD.py b/src/commander4/data_models/detector_TOD.py index 7fcf189..37423a2 100644 --- a/src/commander4/data_models/detector_TOD.py +++ b/src/commander4/data_models/detector_TOD.py @@ -155,15 +155,18 @@ def __init__( self._huffman_symbols2 = huffman_symbols2 self._tod_is_compressed = tod_is_compressed self._flag_is_compressed = flag_is_compressed + # The Huffman decoder expects uint8 arrays; for bytes and HDF5-backed + # numpy.void payloads the internal storage is rewritten as a zero-copy + # uint8 view once at construction. self.processing_mask_map = processing_mask_map self.pointing = pointing processing_mask = processing_mask_map[self.get_pix()] self._processing_mask = np.packbits(processing_mask) self.det_response = det_response if flag_encoded is not None and bad_data_bitmask is not None: - bad_data_mask = ~(self.flag & bad_data_bitmask) - self._bad_data_mask = np.packbits(bad_data_mask) - self._full_mask = np.packbits(bad_data_mask & processing_mask) + good_data_mask = (self.flag & bad_data_bitmask) == 0 + self._good_data_mask = np.packbits(good_data_mask) + self._full_mask = np.packbits(good_data_mask & processing_mask) if orb_dir_vec is not None: log.logassert_np(orb_dir_vec.size == 3, "orb_dir_vec must be a vector of size 3.", logger) self._orb_dir_vec = orb_dir_vec.astype(np.float32, copy=False) @@ -232,14 +235,16 @@ def processing_mask(self) -> NDArray[np.bool_]: @property def full_mask(self) -> NDArray[np.bool_]: + """Boolean mask keeping samples that pass both flag and processing cuts.""" mask = np.unpackbits(self._full_mask).view(bool) if mask.size > self.tod.size + 7 or mask.size < self.tod.size: raise ValueError(f"Mask size {mask.size} doesn't match TOD size {self.tod.size}.") return mask[:self.tod.size] @property - def bad_data_mask(self) -> NDArray[np.bool_]: - mask = np.unpackbits(self._bad_data_mask).view(bool) + def good_data_mask(self) -> NDArray[np.bool_]: + """Boolean mask keeping samples that pass the bad-data flag cut.""" + mask = np.unpackbits(self._good_data_mask).view(bool) if mask.size > self.tod.size + 7 or mask.size < self.tod.size: raise ValueError(f"Mask size {mask.size} doesn't match TOD size {self.tod.size}.") return mask[:self.tod.size] diff --git a/src/commander4/data_models/pointing.py b/src/commander4/data_models/pointing.py index 842824f..2b44b54 100644 --- a/src/commander4/data_models/pointing.py +++ b/src/commander4/data_models/pointing.py @@ -13,7 +13,14 @@ class ScanBoresightPointing: - """Scan-level boresight pointing evaluated once for all detectors.""" + """Evaluate one scan's boresight once and reuse it for all detectors. + + The scan boresight is propagated for the full original TOD length in sky + coordinates and kept as a shared object. Individual detector pointings are + then obtained by rotating that common boresight with per-detector xi/eta + offsets and polarization angles, which avoids recomputing the expensive + time-dependent coordinate transform for every detector. + """ def __init__(self, time_start_mjd: float, @@ -45,10 +52,13 @@ def __init__(self, "polangs must contain one polarization angle per detector.", logger, ) + # pixell's time-dependent coordinate transforms use Unix seconds. time_start_unix = (time_start_mjd - 40587.0) * 86400.0 time_end_unix = (time_end_mjd - 40587.0) * 86400.0 time_unix = np.linspace(time_start_unix, time_end_unix, ntod_original) + # Build the boresight for the full native scan once; shorter requests + # are handled later by slicing to self.ntod. self.bore_point = self.initialize_boresight(time_unix, bore, site=self.site) @@ -60,6 +70,7 @@ def initialize_boresight( site=None, weather: str = "typical", ): + """Transform boresight az/el/roll samples into the requested sky frame.""" icoord = coordsys.Coords(az=bore[0], el=bore[1], roll=bore[2]) return coordsys.transform("hor", sys, icoord, ctime=ctime, site=site, weather=weather) @@ -72,6 +83,8 @@ def get_det_point( # By slicing instead of indexing we keep the 1-sized detector dimension. detoff = self.detoffs[idet:idet+1] polang = self.polangs[idet:idet+1] + # Apply the detector's focal-plane offset and polarization rotation on + # top of the shared boresight quaternion. qdet = coordsys.rotation_xieta(detoff[:, 0], detoff[:, 1], polang) ocoord = self.bore_point * qdet[:, None] # TODO: The lines below are absurdly slow, taking 95% of the runtime of this function, @@ -90,6 +103,7 @@ def get_pix_psi( ) -> tuple[NDArray[np.integer], NDArray[np.floating]]: target_nside = self.nside if nside is None else nside dec, ra, psi = self.get_det_point(idet) + # healpy expects co-latitude theta rather than declination. theta = np.pi/2.0 - dec pix = hp.ang2pix(target_nside, theta, ra) psi = psi.astype(np.float32, copy=False)[:self.ntod] @@ -106,6 +120,13 @@ def get_psi(self, idet: int, nside: int | None = None) -> NDArray[np.floating]: class DetectorBoresightPointing: + """Detector-specific view onto a shared ScanBoresightPointing. + + This wrapper stores only the detector index and forwards all queries to the + shared scan-level object. That keeps the per-detector interface simple while + avoiding duplication of boresight and site state. + """ + def __init__(self, scan_pointing: ScanBoresightPointing, idet: int): self.scan_pointing = scan_pointing self.idet = int(idet) @@ -131,11 +152,18 @@ def get_pix_psi(self, nside: int | None = None) -> tuple[NDArray[np.integer], ND class PixelPointing: - """Joint pixel and polarization-angle representation for one detector TOD.""" + """Store pixel and polarization-angle pointing for one detector TOD. + + The pointing can be supplied either as decoded 1D arrays or as Huffman- + compressed binary payloads. Compressed payloads are kept compact in memory + and decoded only on demand in `get_pix()` and `get_psi()`. Pixel samples are + stored at `data_nside` and optionally remapped to another output `nside` + after decompression. + """ def __init__(self, - pix: bytes | NDArray[np.integer], - psi: bytes | NDArray[np.integer] | NDArray[np.floating], + pix: bytes | np.void | NDArray[np.integer], + psi: bytes | np.void | NDArray[np.integer] | NDArray[np.floating], huffman_tree: NDArray | None, huffman_symbols: NDArray | None, npsi: int | None, @@ -153,8 +181,12 @@ def __init__(self, self.huffman_tree = huffman_tree self.huffman_symbols = huffman_symbols self.npsi = npsi - self.pix_is_compressed = isinstance(pix, bytes) - self.psi_is_compressed = isinstance(psi, bytes) + self.pix_is_compressed = isinstance(pix, (bytes, np.void)) + self.psi_is_compressed = isinstance(psi, (bytes, np.void)) + # The Huffman decoder consumes uint8 arrays; for HDF5-backed np.void + # inputs, frombuffer gives a zero-copy view over the stored payload. + self.pix_compressed_u8 = np.frombuffer(pix, dtype=np.uint8) if self.pix_is_compressed else None + self.psi_compressed_u8 = np.frombuffer(psi, dtype=np.uint8) if self.psi_is_compressed else None self._test_input() @@ -162,12 +194,12 @@ def _test_input(self): log.logassert_np(self.ntod <= self.ntod_original, "ntod cannot exceed ntod_original.", logger) log.logassert_np( self.pix_is_compressed or isinstance(self.pix_encoded, np.ndarray), - "'pix' must be provided as bytes or a numpy array.", + "'pix' must be provided as bytes, numpy.void, or a numpy array.", logger, ) log.logassert_np( self.psi_is_compressed or isinstance(self.psi_encoded, np.ndarray), - "'psi' must be provided as bytes or a numpy array.", + "'psi' must be provided as bytes, numpy.void, or a numpy array.", logger, ) if self.pix_is_compressed: @@ -216,8 +248,10 @@ def get_pix(self, nside: int | None = None) -> NDArray[np.integer]: target_nside = self.nside if nside is None else nside if self.pix_is_compressed: pix = np.zeros(self.ntod_original, dtype=np.int64) - pix = cpp_utils.huffman_decode(np.frombuffer(self.pix_encoded, dtype=np.uint8), - self.huffman_tree, self.huffman_symbols, pix) + pix = cpp_utils.huffman_decode(self.pix_compressed_u8, self.huffman_tree, + self.huffman_symbols, pix) + # The compressed stream stores first differences, so reconstruct the + # absolute pixel indices with a cumulative sum. pix = np.cumsum(pix) else: pix = self.pix_encoded @@ -235,8 +269,10 @@ def get_psi(self, nside: int | None = None) -> NDArray[np.floating]: """Return polarization angles, cropped to the active TOD length.""" if self.psi_is_compressed: psi = np.zeros(self.ntod_original, dtype=np.int64) - psi = cpp_utils.huffman_decode(np.frombuffer(self.psi_encoded, dtype=np.uint8), - self.huffman_tree, self.huffman_symbols, psi) + psi = cpp_utils.huffman_decode(self.psi_compressed_u8, self.huffman_tree, + self.huffman_symbols, psi) + # psi is compressed as differences of digitized angle bins; first + # recover the bin index stream, then convert bins back to radians. psi = np.cumsum(psi) psi = psi[:self.ntod] psi = 2 * np.pi * psi.astype(np.float32, copy=False) / self.npsi diff --git a/src/commander4/data_models/scan_TOD.py b/src/commander4/data_models/scan_TOD.py index 12a66a8..cf781e5 100644 --- a/src/commander4/data_models/scan_TOD.py +++ b/src/commander4/data_models/scan_TOD.py @@ -7,9 +7,6 @@ class ScanTOD: detectors (list[DetectorTOD]): Per-detector TOD data for this scan. start_time (float): Start time of the scan (mission elapsed time). scanID (int): Unique integer identifier for this scan. - scan_num (int): Compact - scan_idx_start (int): Global start index of this scan within the full TOD. - scan_idx_stop (int): Global stop index (exclusive) of this scan. """ def __init__ (self, detlist: list[DetectorTOD], start_time: float, scan_id: int): self.detectors = detlist diff --git a/src/commander4/tod_processing.py b/src/commander4/tod_processing.py index a366a90..6ff9258 100644 --- a/src/commander4/tod_processing.py +++ b/src/commander4/tod_processing.py @@ -86,10 +86,10 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output mapmaker_invvar = WeightsMapmakerIQU(band_comm, experiment_data.nside) for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): - mask_bad = det.bad_data_mask + good_data_mask = det.good_data_mask pix, psi = det.get_pix_psi() - pix = pix[mask_bad] - psi = psi[mask_bad] + pix = pix[good_data_mask] + psi = psi[good_data_mask] sigma0, fknee, alpha = tod_samples.noise_params[iscan, idet, :] gain = tod_samples.gain(iscan, idet) inv_var = (gain/sigma0)**2 @@ -107,8 +107,8 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output mapmaker_invvar = WeightsMapmaker(band_comm, experiment_data.nside) for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): - mask_bad = det.bad_data_mask - pix = det.get_pix()[mask_bad] + good_data_mask = det.good_data_mask + pix = det.get_pix()[good_data_mask] sigma0, fknee, alpha = tod_samples.noise_params[iscan, idet, :] gain = tod_samples.gain(iscan, idet) inv_var = (gain/sigma0)**2 @@ -140,6 +140,9 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output for idet, det in enumerate(scan.detectors): d_sky = det.tod.copy() pix, psi = det.get_pix_psi() + good_data_mask = det.good_data_mask + pix_masked = pix[good_data_mask] + psi_masked = psi[good_data_mask] sigma0, fknee, alpha = tod_samples.noise_params[iscan, idet, :] gain = tod_samples.gain(iscan, idet) inv_var = (gain/sigma0)**2 @@ -173,11 +176,13 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output mask_full, sigma0_ncorr, C_1f_inv, err_tol=err_tol) if pols == "IQU": - mapmaker_ncorr.accumulate_to_map((n_corr_est/gain).astype(np.float32, copy=False), - inv_var, pix, psi, response=response) + mapmaker_ncorr.accumulate_to_map( + (n_corr_est/gain).astype(np.float32, copy=False), + inv_var, pix, psi, response=response) else: - mapmaker_ncorr.accumulate_to_map((n_corr_est/gain).astype(np.float32, copy=False), - inv_var, pix, psi) + mapmaker_ncorr.accumulate_to_map( + (n_corr_est/gain).astype(np.float32, copy=False), + inv_var, pix, psi) if residual > err_tol: num_failed_convergences_ncorr += 1 worst_residual_ncorr = max(worst_residual_ncorr, residual) @@ -191,12 +196,14 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output d_sky -= n_corr_est + d_sky_masked = d_sky[good_data_mask] + cg_mapmaker.accum_to_RHS( scan_tod=det, sigma0=sigma0, - pix=pix, - psi=psi, - scan_tod_arr=d_sky/gain + pix=pix_masked, + psi=psi_masked, + scan_tod_arr=d_sky_masked/gain ) ### PRINT NOISE SAMPLING STATS ### @@ -345,10 +352,10 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu mapmaker_invvar = WeightsMapmakerIQU(band_comm, experiment_data.nside) for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): - mask_bad = det.bad_data_mask + good_data_mask = det.good_data_mask pix, psi = det.get_pix_psi() - pix = pix[mask_bad] - psi = psi[mask_bad] + pix = pix[good_data_mask] + psi = psi[good_data_mask] sigma0, fknee, alpha = tod_samples.noise_params[iscan,idet,:] gain = tod_samples.gain(iscan, idet) # sigma0 is in detector-units, transform into uK_RJ by dividing it by the gain. @@ -375,11 +382,11 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): start_bench("binned-mapmaker") - mask_bad = det.bad_data_mask - d_sky = det.tod.copy()[mask_bad] + good_data_mask = det.good_data_mask + d_sky = det.tod.copy() pix, psi = det.get_pix_psi() - pix = pix[mask_bad] - psi = psi[mask_bad] + pix_masked = pix[good_data_mask] + psi_masked = psi[good_data_mask] response = det.det_response gain = tod_samples.gain(iscan, idet) sigma0, fknee, alpha = tod_samples.noise_params[iscan,idet,:] @@ -387,14 +394,14 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu inv_var = (gain/sigma0)**2 ### ORBITAL DIPOLE ### - sky_orb_dipole = get_s_orb_TOD(det, experiment_data, pix)[mask_bad] + sky_orb_dipole = get_s_orb_TOD(det, experiment_data, pix) d_sky -= gain*sky_orb_dipole stop_bench("binned-mapmaker", increment_count=False) ### CORRELATED NOISE SAMPLING ### if do_ncorr_sampling: start_bench("ncorr-sampling") - s_tot = get_static_sky_TOD(compsep_output, pix, psi)[mask_bad] + s_tot = get_static_sky_TOD(compsep_output, pix, psi) s_tot += sky_orb_dipole sky_subtracted_TOD = det.tod.copy() sky_subtracted_TOD -= gain*s_tot @@ -406,14 +413,13 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu sigma0_ncorr = calc_sigma0_robust(sky_subtracted_TOD, mask_full) C_1f_inv = np.zeros(Nfft) C_1f_inv[1:] = 1.0 / (sigma0_ncorr**2*(freq[1:]/fknee)**alpha) - err_tol = 1e-6 # Inpaint masked regions with linear slope + white noise. # In the CG solver this is only used to define the starting guess, # but if the CG fails it is also used to generate the fallback solution. fill_all_masked(sky_subtracted_TOD, mask_full, sigma0_ncorr) n_corr_est, residual, niter, did_conv = corr_noise_realization_with_gaps(sky_subtracted_TOD, mask_full, sigma0_ncorr, C_1f_inv, - err_tol=err_tol, max_iter=200) + err_tol=1e-4, max_iter=20) resid = (sky_subtracted_TOD - n_corr_est) * mask_full var_resid = np.dot(resid, resid) var_data = np.dot(sky_subtracted_TOD * mask_full, sky_subtracted_TOD * mask_full) @@ -445,14 +451,14 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu niters.append(niter) d_sky -= n_corr_est - stop_bench("ncorr-sampling", increment_count=False) + stop_bench("ncorr-sampling") if iscan == len(experiment_data.scans) - 1: log_memory("ncorr-sampling") + d_sky_masked = d_sky[good_data_mask] start_bench("binned-mapmaker") - mapmaker.accumulate_to_map(d_sky/gain, inv_var, pix, psi, response=response) - mapmaker_orbdipole.accumulate_to_map(sky_orb_dipole, inv_var, pix, psi, - response=response) + mapmaker.accumulate_to_map(d_sky_masked/gain, inv_var, pix_masked, psi_masked, response=response) + mapmaker_orbdipole.accumulate_to_map(sky_orb_dipole, inv_var, pix, psi, response=response) stop_bench("binned-mapmaker", increment_count=False) ### PRINT NOISE SAMPLING STATS ### From 98e7cb48747e83a469fd0e3749df38879e8214e0 Mon Sep 17 00:00:00 2001 From: jgslunde Date: Fri, 29 May 2026 15:32:34 +0200 Subject: [PATCH 12/23] Added a TODView class to serve as a frontend for retrieving TOD-related data, combining information from both TODSamples and Detector_TOD --- src/commander4/data_models/TOD_samples.py | 46 ++- src/commander4/data_models/tod_view.py | 458 ++++++++++++++++++++++ src/commander4/tod_processing.py | 261 ++++-------- 3 files changed, 584 insertions(+), 181 deletions(-) create mode 100644 src/commander4/data_models/tod_view.py diff --git a/src/commander4/data_models/TOD_samples.py b/src/commander4/data_models/TOD_samples.py index cccf0fe..cbe06cb 100644 --- a/src/commander4/data_models/TOD_samples.py +++ b/src/commander4/data_models/TOD_samples.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import os import numpy as np import h5py @@ -7,8 +9,11 @@ from mpi4py import MPI from pixell.bunch import Bunch import logging +import typing -from commander4.data_models.detector_group_TOD import DetGroupTOD +from commander4.data_models.jump_corrections import JumpCatalog +if typing.TYPE_CHECKING: + from commander4.data_models.detector_group_TOD import DetGroupTOD logger = logging.getLogger(__name__) @@ -57,6 +62,29 @@ def _gather_scan_distributed_array(band_comm: MPI.Comm, local_array: NDArray, return global_array +def _gather_variable_length_1d_array(band_comm: MPI.Comm, local_array: NDArray) -> NDArray | None: + """Gather a 1-D array with rank-dependent length onto the root rank.""" + local_array = np.ascontiguousarray(local_array) + local_count = np.array([local_array.size], dtype=np.int64) + + if band_comm.Get_rank() == 0: + counts = np.zeros(band_comm.Get_size(), dtype=np.int64) + else: + counts = None + band_comm.Gather(local_count, counts, root=0) + + recvbuf = None + global_array = None + if band_comm.Get_rank() == 0: + displacements = np.cumsum(counts) - counts + global_array = np.empty(np.sum(counts), dtype=local_array.dtype) + mpi_type = MPI._typedict[local_array.dtype.char] + recvbuf = (global_array, counts, displacements, mpi_type) + + band_comm.Gatherv(local_array, recvbuf=recvbuf, root=0) + return global_array + + class TODSamples: """ A class for holding all sampled TOD-quantities, such as gains and correlated noise parameters, for one MPI rank. Quantities that vary with detectors and/or scans are stored as @@ -65,6 +93,7 @@ class TODSamples: Quantities that are the same across a band (like absolute gain) have identical copies for all ranks on the same band. """ + def __init__(self, experiment_data: DetGroupTOD, params: Bunch, @@ -84,6 +113,7 @@ def __init__(self, self.scan_idx_start = experiment_data.scan_idx_start self.scan_idx_stop = experiment_data.scan_idx_stop self.scan_ids = np.array([scan.scan_id for scan in experiment_data.scans]) + self.jumps = JumpCatalog.empty(self.nscans, self.ndet) # Gibbs-sampled quantities if not params.general.init_from_chain: @@ -173,6 +203,7 @@ def __init__(self, # 4. Load and slice Per-Scan arrays (Distributed across ranks) self.temporal_gain = f["temporal_gain"][local_indices, :] if "temporal_gain" in f else None self.noise_params = f["noise_params"][local_indices, ...] if "noise_params" in f else None + self.jumps = JumpCatalog.from_hdf5(f, local_indices, self.ndet) if self.band_comm.Get_rank() == 0: logger.debug(f"Initial absolute gain estimate for {self.band_name}: {self.abs_gain:.3e}.") @@ -240,6 +271,13 @@ def write_chain_to_file(self, itr: int): if self.noise_params is not None: noise_params_global = _gather_scan_distributed_array(band_comm, self.noise_params, scans_per_rank) + + # 5. Jump corrections (per-scan per-detector ragged quantity) + jump_counts_local, jump_locations_local, jump_offsets_local = self.jumps.pack() + jump_counts_global = _gather_scan_distributed_array(band_comm, jump_counts_local, + scans_per_rank) + jump_locations_global = _gather_variable_length_1d_array(band_comm, jump_locations_local) + jump_offsets_global = _gather_variable_length_1d_array(band_comm, jump_offsets_local) #################################################################### # Write results to file. #################################################################### @@ -261,4 +299,8 @@ def write_chain_to_file(self, itr: int): if temporal_gain_global is not None: file["temporal_gain"] = temporal_gain_global if noise_params_global is not None: - file["noise_params"] = noise_params_global \ No newline at end of file + file["noise_params"] = noise_params_global + if jump_counts_global is not None: + file["jump_counts"] = jump_counts_global + file["jump_locations"] = jump_locations_global + file["jump_offsets"] = jump_offsets_global \ No newline at end of file diff --git a/src/commander4/data_models/tod_view.py b/src/commander4/data_models/tod_view.py new file mode 100644 index 0000000..45d7a62 --- /dev/null +++ b/src/commander4/data_models/tod_view.py @@ -0,0 +1,458 @@ +import numpy as np +import logging +from numpy.typing import NDArray +from typing import Literal + +from pixell.bunch import Bunch + +from commander4.data_models.detector_group_TOD import DetGroupTOD +from commander4.data_models.TOD_samples import TODSamples +from commander4.utils.map_utils import get_static_sky_TOD, get_s_orb_TOD + +logger = logging.getLogger(__name__) + + +class TODView: + """Materialize one detector at a time and build derived TOD views on demand. + + The view keeps only the currently focused detector decoded in memory. Calling + ``focus()`` always discards the previous detector's cached arrays, which matches the + one-detector-at-a-time TOD-processing architecture. + """ + + _ALL_GAIN_TERMS = ("abs", "rel", "temp") + + def __init__( + self, + experiment_data: DetGroupTOD, + tod_samples: TODSamples, + compsep_output: NDArray | None = None, + ): + """Initialize a detector-local view over one band's TOD data. + + Args: + experiment_data: Static TOD container for the current band. + tod_samples: Sampled gain and noise parameters for the current chain state. + compsep_output: Optional default sky model used by sky-subtraction helpers. + """ + self.experiment_data = experiment_data + self.tod_samples = tod_samples + self.compsep_output = compsep_output + self._iscan: int | None = None + self._idet: int | None = None + self._det = None + self._clear_cache() + + def _clear_cache(self): + """Drop all arrays materialized for the current detector.""" + self._tod = None + self._corrected_tod = None + self._pix = None + self._psi = None + self._flag = None + self._processing_mask = None + self._good_data_mask = None + self._full_mask = None + self._static_sky = None + self._orbital_dipole = None + self._downsampled: dict[int, Bunch] = {} + + def focus(self, iscan: int, idet: int) -> "TODView": + """Focus the view on one detector and discard any previous materialization.""" + self._det = self.experiment_data.scans[iscan].detectors[idet] + self._iscan = iscan + self._idet = idet + self._clear_cache() + return self + + def _require_focus(self): + """Return the current detector or raise if the view was not focused yet.""" + if self._iscan is None or self._idet is None or self._det is None: + raise ValueError("Attempted to use TODView before calling TODView.focus().") + return self._det + + @property + def iscan(self) -> int: + self._require_focus() + return self._iscan + + @property + def idet(self) -> int: + self._require_focus() + return self._idet + + @property + def detector(self): + return self._require_focus() + + @property + def fsamp(self) -> float: + return self.detector.fsamp + + @property + def det_response(self) -> NDArray | None: + return self.detector.det_response + + @property + def noise_params(self) -> NDArray: + self._require_focus() + return self.tod_samples.noise_params[self._iscan, self._idet] + + @property + def sigma0(self) -> float: + return float(self.noise_params[0]) + + def get_gain(self, gain_terms: tuple[str, ...] | None = _ALL_GAIN_TERMS) -> float: + """Return the selected subset of the current detector gain model.""" + if gain_terms is None: + gain_terms = () + elif "all" in gain_terms: + gain_terms = self._ALL_GAIN_TERMS + + gain = 0.0 + for term in gain_terms: + if term == "abs": + gain += self.tod_samples.abs_gain + elif term == "rel": + gain += self.tod_samples.rel_gain[self.idet] + elif term == "temp": + gain += self.tod_samples.temporal_gain[self.iscan, self.idet] + return float(gain) + + @property + def tod(self) -> NDArray[np.floating]: + if self._tod is None: + self._tod = self.detector.tod + return self._tod + + @property + def corrected_tod(self) -> NDArray[np.floating]: + """Return the raw TOD with the stored jump offsets applied in detector units.""" + if self._corrected_tod is None: + jump = self.tod_samples.jumps.get(self.iscan, self.idet) + self._corrected_tod = self.tod if jump.is_empty() else jump.apply(self.tod) + return self._corrected_tod + + def _get_fullres_pix_psi(self) -> tuple[NDArray[np.integer], NDArray[np.floating] | NDArray[np.integer]]: + """Decode full-resolution pointing once and reuse it for the current detector.""" + if self._pix is None or self._psi is None: + self._pix, self._psi = self.detector.get_pix_psi() + return self._pix, self._psi + + @property + def pix(self) -> NDArray[np.integer]: + return self._get_fullres_pix_psi()[0] + + @property + def psi(self) -> NDArray[np.floating] | NDArray[np.integer]: + return self._get_fullres_pix_psi()[1] + + @property + def flag(self) -> NDArray[np.integer]: + if self._flag is None: + self._flag = self.detector.flag + return self._flag + + def _unpack_mask(self, attr_name: str) -> NDArray[np.bool_]: + """Unpack one of DetectorTOD's packed bit masks for the focused detector.""" + det = self.detector + packed = getattr(det, attr_name, None) + if packed is None: + # Some datasets may not define every cut explicitly; fall back to permissive masks. + if attr_name == "_processing_mask": + return np.ones(det.ntod, dtype=bool) + if attr_name == "_good_data_mask": + return np.ones(det.ntod, dtype=bool) + if attr_name == "_full_mask": + return self.processing_mask.copy() + raise ValueError(f"Detector mask '{attr_name}' is unavailable.") + mask = np.unpackbits(packed).view(bool) + if mask.size > det.ntod + 7 or mask.size < det.ntod: + raise ValueError(f"Mask size {mask.size} doesn't match ntod {det.ntod}.") + return mask[:det.ntod] + + @property + def processing_mask(self) -> NDArray[np.bool_]: + if self._processing_mask is None: + self._processing_mask = self._unpack_mask("_processing_mask") + return self._processing_mask + + @property + def good_data_mask(self) -> NDArray[np.bool_]: + if self._good_data_mask is None: + self._good_data_mask = self._unpack_mask("_good_data_mask") + return self._good_data_mask + + @property + def full_mask(self) -> NDArray[np.bool_]: + if self._full_mask is None: + self._full_mask = self._unpack_mask("_full_mask") + return self._full_mask + + def _downsample_factor_or_default(self, downsample_factor: int | None) -> int: + """Validate a downsampling factor and replace ``None`` with unity.""" + if downsample_factor is None: + downsample_factor = 1 + if downsample_factor < 1: + raise ValueError("downsample_factor must be >= 1.") + return int(downsample_factor) + + def _materialize_downsampled(self, downsample_factor: int) -> Bunch: + """Cache a downsampled detector view for the requested averaging factor.""" + factor = self._downsample_factor_or_default(downsample_factor) + cached = self._downsampled.get(factor) + if cached is not None: + return cached + + if factor == 1: + data = Bunch( + tod=self.corrected_tod, + pix=self.pix, + psi=self.psi, + processing_mask=self.processing_mask, + good_data_mask=self.good_data_mask, + full_mask=self.full_mask, + indices=np.arange(self.detector.ntod, dtype=np.int64), + ) + else: + # Average the jump-corrected TOD blocks, while keeping pointing and masks at the block + # centers to match the existing gain-calibration logic. + indices_edges = np.arange(0, self.detector.ntod, factor) + indices = (indices_edges[1:] + indices_edges[:-1]) // 2 + ntod_down = indices.size + tod = self.corrected_tod[:ntod_down*factor].reshape((ntod_down, factor)) + data = Bunch( + tod=np.mean(tod, axis=-1), + pix=self.pix[indices], + psi=self.psi[indices], + processing_mask=self.processing_mask[indices], + good_data_mask=self.good_data_mask[indices], + full_mask=self.full_mask[indices], + indices=indices, + ) + + self._downsampled[factor] = data + return data + + def get_mask( + self, + mask: Literal["none", "processing", "good", "full"] = "none", + downsample_factor: int | None = None, + ) -> NDArray[np.bool_] | None: + """Return one of the cached detector masks, optionally at downsampled resolution.""" + if mask == "none": + return None + data = self._materialize_downsampled(self._downsample_factor_or_default(downsample_factor)) + if mask == "processing": + return data.processing_mask + if mask == "good": + return data.good_data_mask + if mask == "full": + return data.full_mask + raise ValueError(f"Unknown mask mode '{mask}'.") + + def _require_compsep_output(self, compsep_output: NDArray | None) -> NDArray: + """Resolve the sky model to use for static-sky subtraction.""" + sky_model = self.compsep_output if compsep_output is None else compsep_output + if sky_model is None: + raise ValueError("A component-separation sky map must be provided for sky subtraction.") + return sky_model + + def get_static_sky_tod( + self, + compsep_output: NDArray | None = None, + downsample_factor: int | None = None, + ) -> NDArray[np.floating]: + """Evaluate the static sky model along the focused detector pointing.""" + factor = self._downsample_factor_or_default(downsample_factor) + sky_model = self._require_compsep_output(compsep_output) + if factor == 1 and compsep_output is None: + if self._static_sky is None: + # Reuse the full-resolution sky TOD when both pointing and sky model match. + self._static_sky = get_static_sky_TOD(sky_model, self.pix, psi=self.psi) + return self._static_sky + data = self._materialize_downsampled(factor) + return get_static_sky_TOD(sky_model, data.pix, psi=data.psi) + + def get_orbital_dipole_tod(self, downsample_factor: int | None = None) -> NDArray[np.floating]: + """Evaluate the orbital dipole for the focused detector.""" + factor = self._downsample_factor_or_default(downsample_factor) + if factor == 1: + if self._orbital_dipole is None: + self._orbital_dipole = get_s_orb_TOD(self.detector, self.experiment_data, self.pix) + return self._orbital_dipole + data = self._materialize_downsampled(factor) + return get_s_orb_TOD(self.detector, self.experiment_data, data.pix) + + def _normalize_signal_name(self, signal_name: str) -> str: + """Map user-facing TOD component names onto internal canonical names.""" + normalized = signal_name.lower() + aliases = { + "sky": "static_sky", + "static_sky": "static_sky", + "orb": "orbital_dipole", + "orbital_dipole": "orbital_dipole", + } + if normalized not in aliases: + raise ValueError(f"Unknown TOD signal '{signal_name}'.") + return aliases[normalized] + + def _get_signal_tod( + self, + signal_name: str, + compsep_output: NDArray | None = None, + downsample_factor: int | None = None, + ) -> NDArray[np.floating]: + """Return one named model TOD evaluated for the focused detector.""" + normalized = self._normalize_signal_name(signal_name) + if normalized == "static_sky": + return self.get_static_sky_tod(compsep_output=compsep_output, + downsample_factor=downsample_factor) + if normalized == "orbital_dipole": + return self.get_orbital_dipole_tod(downsample_factor=downsample_factor) + raise ValueError(f"Unhandled TOD signal '{signal_name}'.") + + def get_tod( + self, + *, + subtract: tuple[tuple[str, tuple[str, ...]], ...] | None = None, + divide_by_gain: tuple[str, ...] | None = None, + downsample_factor: int = 1, + mask: Literal["none", "processing", "good", "full"] = "none", + compsep_output: NDArray | None = None, + ) -> NDArray[np.floating]: + """Return a jump-corrected detector-local TOD after subtracting selected model terms. + + Args: + subtract: Sequence of ``(signal_name, gain_terms)`` pairs. Each signal is evaluated + and subtracted after multiplying it by the selected gain subset. + divide_by_gain: Gain terms to divide the final TOD by, or ``None``. + downsample_factor: Average the TOD in contiguous chunks before subtraction. + mask: Optional mask to apply to the output TOD. + compsep_output: Optional sky model override for static-sky subtraction. + """ + factor = self._downsample_factor_or_default(downsample_factor) + tod = np.array(self._materialize_downsampled(factor).tod, copy=True) + + if subtract is not None: + for signal_name, gain_terms in subtract: + # All supported residuals in this class are linear combinations of named model TODs. + signal = self._get_signal_tod(signal_name, compsep_output=compsep_output, + downsample_factor=factor) + tod -= self.get_gain(gain_terms) * signal + + if divide_by_gain is not None: + gain = self.get_gain(divide_by_gain) + if gain == 0: + raise ValueError("Cannot divide TOD by a zero gain.") + tod /= gain + + # Apply masking last so callers can request either the full residual or the cut samples. + mask_arr = self.get_mask(mask, downsample_factor=factor) + return tod if mask_arr is None else tod[mask_arr] + + def _fill_masked_calibration_samples( + self, + tod: NDArray[np.floating], + mask: NDArray[np.bool_], + signal: NDArray[np.floating], + gain_terms: tuple[str, ...], + downsample_factor: int, + rng: np.random.Generator | None, + ) -> NDArray[np.floating]: + """Fill masked calibration samples with signal plus white noise in detector units.""" + sigma0_effective = self.sigma0 * np.sqrt(1.0 / downsample_factor) + filled = np.array(tod, copy=True) + if (~mask).any(): + normal = np.random.normal if rng is None else rng.normal + noise = normal(0.0, sigma0_effective, signal[~mask].shape) + # The masked regions retain only the target gain term times the calibrator signal. + filled[~mask] = self.get_gain(gain_terms) * signal[~mask] + noise + return filled + + def get_abs_calib_tod( + self, + *, + compsep_output: NDArray | None = None, + downsample_factor: int | None = None, + calibrate_on_full_sky: bool | None = None, + preserve_target_gain: bool = True, + fill_masked: bool = True, + rng: np.random.Generator | None = None, + ) -> Bunch: + """Return the residual, calibrator signal, and mask used for absolute-gain sampling.""" + factor = int(self.fsamp) if downsample_factor is None else int(downsample_factor) + data = self._materialize_downsampled(factor) + mask = self.get_mask("full", downsample_factor=factor) + s_sky = self.get_static_sky_tod(compsep_output=compsep_output, downsample_factor=factor) + s_orb = self.get_orbital_dipole_tod(downsample_factor=factor) + + if calibrate_on_full_sky is None: + calibrate_on_full_sky = self.experiment_data.nu > 380.0 + + if calibrate_on_full_sky: + s_cal = s_sky + s_orb + # For the full-sky branch, optionally preserve the absolute term in the residual so + # callers can choose between the current implementation and the algebraically cleaner + # target-gain-preserving form. + abs_subtract = ("rel", "temp") if preserve_target_gain else self._ALL_GAIN_TERMS + subtract = (("sky", abs_subtract), ("orbital_dipole", abs_subtract)) + else: + s_cal = s_orb + # Low-frequency absolute calibration keeps only the absolute orbital-dipole term. + subtract = (("sky", self._ALL_GAIN_TERMS), ("orbital_dipole", ("rel", "temp"))) + + tod = self.get_tod(subtract=subtract, downsample_factor=factor, compsep_output=compsep_output) + if fill_masked: + tod = self._fill_masked_calibration_samples(tod, mask, s_cal, ("abs",), factor, rng) + return Bunch(tod=tod, s_cal=s_cal, pix=data.pix, psi=data.psi, mask=mask) + + def get_rel_calib_tod( + self, + *, + compsep_output: NDArray | None = None, + downsample_factor: int | None = None, + fill_masked: bool = True, + rng: np.random.Generator | None = None, + ) -> Bunch: + """Return the residual, calibrator signal, and mask used for relative-gain sampling.""" + factor = int(self.fsamp) if downsample_factor is None else int(downsample_factor) + data = self._materialize_downsampled(factor) + mask = self.get_mask("full", downsample_factor=factor) + s_sky = self.get_static_sky_tod(compsep_output=compsep_output, downsample_factor=factor) + s_orb = self.get_orbital_dipole_tod(downsample_factor=factor) + s_cal = s_sky + s_orb + # Relative-gain sampling removes the absolute and temporal gain terms, leaving only the + # detector-dependent residual gain multiplying the calibrator signal. + subtract = (("sky", ("abs", "temp")), ("orbital_dipole", ("abs", "temp"))) + + tod = self.get_tod(subtract=subtract, downsample_factor=factor, compsep_output=compsep_output) + if fill_masked: + tod = self._fill_masked_calibration_samples(tod, mask, s_cal, ("rel",), factor, rng) + return Bunch(tod=tod, s_cal=s_cal, pix=data.pix, psi=data.psi, mask=mask) + + def get_temp_calib_tod( + self, + *, + compsep_output: NDArray | None = None, + downsample_factor: int | None = None, + fill_masked: bool = True, + rng: np.random.Generator | None = None, + ) -> Bunch: + """Return the residual, calibrator signal, and mask used for temporal-gain sampling.""" + factor = int(self.fsamp) if downsample_factor is None else int(downsample_factor) + data = self._materialize_downsampled(factor) + mask = self.get_mask("full", downsample_factor=factor) + s_sky = self.get_static_sky_tod(compsep_output=compsep_output, downsample_factor=factor) + s_orb = self.get_orbital_dipole_tod(downsample_factor=factor) + s_cal = s_sky + s_orb + # Temporal-gain sampling removes the band-mean and detector-mean gains so the residual + # only carries the scan-dependent gain fluctuation. + subtract = (("sky", ("abs", "rel")), ("orbital_dipole", ("abs", "rel"))) + + tod = self.get_tod(subtract=subtract, downsample_factor=factor, compsep_output=compsep_output) + if fill_masked: + tod = self._fill_masked_calibration_samples(tod, mask, s_cal, ("temp",), factor, rng) + return Bunch(tod=tod, s_cal=s_cal, pix=data.pix, psi=data.psi, mask=mask) + + diff --git a/src/commander4/tod_processing.py b/src/commander4/tod_processing.py index 6ff9258..509dfd1 100644 --- a/src/commander4/tod_processing.py +++ b/src/commander4/tod_processing.py @@ -16,12 +16,13 @@ from commander4.data_models.detector_map import DetectorMap from commander4.data_models.detector_group_TOD import DetGroupTOD from commander4.data_models.TOD_samples import TODSamples +from commander4.data_models.jump_corrections import JumpCorrection +from commander4.data_models.tod_view import TODView from commander4.utils.mapmaker import MapmakerIQU, WeightsMapmakerIQU, WeightsMapmaker, Mapmaker from commander4.utils.CG_mapmaker import CGMapmakerI, CGMapmakerIQU from commander4.solvers.preconditioners import InvNPreconditionerI, InvNPreconditionerIQU from commander4.noise_sampling.noise_sampling import sample_noise_PS_params, fill_all_masked from commander4.noise_sampling.sample_ncorr import corr_noise_realization_with_gaps -from commander4.utils.map_utils import get_static_sky_TOD, get_s_orb_TOD from commander4.utils.math_operations import forward_rfft, backward_rfft from commander4.noise_sampling.sigma0 import calc_sigma0_robust from commander4.tod_reader import read_tods_from_file @@ -82,18 +83,19 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output # This is purely to reduce the maximum concurrent memory requirement, and is slightly slower # as we have to de-compress pix and psi twice. pols = experiment_data.pols + scan_view = TODView(experiment_data, tod_samples, compsep_output=compsep_output) if pols == "IQU": mapmaker_invvar = WeightsMapmakerIQU(band_comm, experiment_data.nside) for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): - good_data_mask = det.good_data_mask - pix, psi = det.get_pix_psi() - pix = pix[good_data_mask] - psi = psi[good_data_mask] - sigma0, fknee, alpha = tod_samples.noise_params[iscan, idet, :] - gain = tod_samples.gain(iscan, idet) + view = scan_view.focus(iscan, idet) + good_data_mask = view.good_data_mask + pix = view.pix[good_data_mask] + psi = view.psi[good_data_mask] + sigma0, fknee, alpha = view.noise_params + gain = view.get_gain() inv_var = (gain/sigma0)**2 - mapmaker_invvar.accumulate_to_map(inv_var, pix, psi, response=det.det_response) + mapmaker_invvar.accumulate_to_map(inv_var, pix, psi, response=view.det_response) mapmaker_invvar.gather_map() mapmaker_invvar.normalize_map() if ismaster: @@ -107,10 +109,11 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output mapmaker_invvar = WeightsMapmaker(band_comm, experiment_data.nside) for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): - good_data_mask = det.good_data_mask - pix = det.get_pix()[good_data_mask] - sigma0, fknee, alpha = tod_samples.noise_params[iscan, idet, :] - gain = tod_samples.gain(iscan, idet) + view = scan_view.focus(iscan, idet) + good_data_mask = view.good_data_mask + pix = view.pix[good_data_mask] + sigma0, fknee, alpha = view.noise_params + gain = view.get_gain() inv_var = (gain/sigma0)**2 mapmaker_invvar.accumulate_to_map(inv_var, pix) mapmaker_invvar.gather_map() @@ -138,19 +141,19 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output ### MAIN SCAN LOOP ### for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): - d_sky = det.tod.copy() - pix, psi = det.get_pix_psi() - good_data_mask = det.good_data_mask + view = scan_view.focus(iscan, idet) + pix, psi = view.pix, view.psi + good_data_mask = view.good_data_mask pix_masked = pix[good_data_mask] psi_masked = psi[good_data_mask] - sigma0, fknee, alpha = tod_samples.noise_params[iscan, idet, :] - gain = tod_samples.gain(iscan, idet) + sigma0, fknee, alpha = view.noise_params + gain = view.get_gain() inv_var = (gain/sigma0)**2 - response = det.det_response if pols == "IQU" else None + response = view.det_response if pols == "IQU" else None ### ORBITAL DIPOLE ### - sky_orb_dipole = get_s_orb_TOD(det, experiment_data, pix) - d_sky -= gain*sky_orb_dipole + sky_orb_dipole = view.get_orbital_dipole_tod() + d_sky = view.get_tod(subtract=(("orbital_dipole", TODView._ALL_GAIN_TERMS),)) if pols == "IQU": mapmaker_orbdipole.accumulate_to_map(sky_orb_dipole, inv_var, pix, psi, response=response) @@ -159,14 +162,14 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output ### CORRELATED NOISE SAMPLING ### if do_ncorr_sampling: - s_tot = get_static_sky_TOD(compsep_output, pix, psi=psi) - s_tot += sky_orb_dipole - sky_subtracted_TOD = det.tod.copy() - sky_subtracted_TOD -= gain*s_tot + sky_subtracted_TOD = view.get_tod( + subtract=(("sky", TODView._ALL_GAIN_TERMS), + ("orbital_dipole", TODView._ALL_GAIN_TERMS)), + ) Ntod = sky_subtracted_TOD.shape[0] Nfft = Ntod + 1 # mirrored FFT: nfft=2*Ntod, n=nfft/2+1=Ntod+1 - freq = rfftfreq(2 * Ntod, d = 1/det.fsamp) - mask_full = det.full_mask + freq = rfftfreq(2 * Ntod, d=1/view.fsamp) + mask_full = view.full_mask sigma0_ncorr = calc_sigma0_robust(sky_subtracted_TOD, mask_full) C_1f_inv = np.zeros(Nfft) C_1f_inv[1:] = 1.0 / (sigma0_ncorr**2*(freq[1:]/fknee)**alpha) @@ -188,7 +191,7 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output worst_residual_ncorr = max(worst_residual_ncorr, residual) ### CORRELATED NOISE POWER SPECTRUM PARAMETERS SAMPLING ### - fknee, alpha = sample_noise_PS_params(n_corr_est, sigma0_ncorr, det.fsamp, alpha, + fknee, alpha = sample_noise_PS_params(n_corr_est, sigma0_ncorr, view.fsamp, alpha, freq_max=2.0, n_grid=150, n_burnin=4) tod_samples.noise_params[iscan, idet, :] = sigma0, fknee, alpha alphas.append(alpha) @@ -349,18 +352,19 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu nscans = len(experiment_data.scans) ndet = experiment_data.ndet pols = experiment_data.pols + scan_view = TODView(experiment_data, tod_samples, compsep_output=compsep_output) mapmaker_invvar = WeightsMapmakerIQU(band_comm, experiment_data.nside) for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): - good_data_mask = det.good_data_mask - pix, psi = det.get_pix_psi() - pix = pix[good_data_mask] - psi = psi[good_data_mask] - sigma0, fknee, alpha = tod_samples.noise_params[iscan,idet,:] - gain = tod_samples.gain(iscan, idet) + view = scan_view.focus(iscan, idet) + good_data_mask = view.good_data_mask + pix = view.pix[good_data_mask] + psi = view.psi[good_data_mask] + sigma0, fknee, alpha = view.noise_params + gain = view.get_gain() # sigma0 is in detector-units, transform into uK_RJ by dividing it by the gain. inv_var = (gain/sigma0)**2 - mapmaker_invvar.accumulate_to_map(inv_var, pix, psi, response=det.det_response) + mapmaker_invvar.accumulate_to_map(inv_var, pix, psi, response=view.det_response) mapmaker_invvar.gather_map() mapmaker_invvar.normalize_map() @@ -382,34 +386,34 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): start_bench("binned-mapmaker") - good_data_mask = det.good_data_mask - d_sky = det.tod.copy() - pix, psi = det.get_pix_psi() + view = scan_view.focus(iscan, idet) + good_data_mask = view.good_data_mask + pix, psi = view.pix, view.psi pix_masked = pix[good_data_mask] psi_masked = psi[good_data_mask] - response = det.det_response - gain = tod_samples.gain(iscan, idet) - sigma0, fknee, alpha = tod_samples.noise_params[iscan,idet,:] + response = view.det_response + gain = view.get_gain() + sigma0, fknee, alpha = view.noise_params # sigma0 is in detector-units, transform into uK_RJ by dividing it by the gain. inv_var = (gain/sigma0)**2 ### ORBITAL DIPOLE ### - sky_orb_dipole = get_s_orb_TOD(det, experiment_data, pix) - d_sky -= gain*sky_orb_dipole + sky_orb_dipole = view.get_orbital_dipole_tod() + d_sky = view.get_tod(subtract=(("orbital_dipole", TODView._ALL_GAIN_TERMS),)) stop_bench("binned-mapmaker", increment_count=False) ### CORRELATED NOISE SAMPLING ### if do_ncorr_sampling: start_bench("ncorr-sampling") - s_tot = get_static_sky_TOD(compsep_output, pix, psi) - s_tot += sky_orb_dipole - sky_subtracted_TOD = det.tod.copy() - sky_subtracted_TOD -= gain*s_tot + sky_subtracted_TOD = view.get_tod( + subtract=(("sky", TODView._ALL_GAIN_TERMS), + ("orbital_dipole", TODView._ALL_GAIN_TERMS)), + ) Ntod = sky_subtracted_TOD.shape[0] Nfft = Ntod + 1 # mirrored FFT: nfft=2*Ntod, n=nfft/2+1=Ntod+1 - freq = rfftfreq(2 * Ntod, d = 1/det.fsamp) + freq = rfftfreq(2 * Ntod, d=1/view.fsamp) - mask_full = det.full_mask + mask_full = view.full_mask sigma0_ncorr = calc_sigma0_robust(sky_subtracted_TOD, mask_full) C_1f_inv = np.zeros(Nfft) C_1f_inv[1:] = 1.0 / (sigma0_ncorr**2*(freq[1:]/fknee)**alpha) @@ -441,7 +445,7 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu worst_residual_ncorr = max(worst_residual_ncorr, residual) ### CORRELATED NOISE POWER SPECTRUM PARAMETERS SAMPLING ### - fknee, alpha = sample_noise_PS_params(n_corr_est, sigma0_ncorr, det.fsamp, alpha, + fknee, alpha = sample_noise_PS_params(n_corr_est, sigma0_ncorr, view.fsamp, alpha, freq_max=2.0, n_grid=150, n_burnin=4) tod_samples.noise_params[iscan,idet,:] = sigma0, fknee, alpha @@ -717,53 +721,21 @@ def sample_absolute_gain(band_comm: MPI.Comm, experiment_data: DetGroupTOD, tod_ # Calibrate on the full sky at high frequencies, as the orbital dipole is too faint. calibrate_on_full_sky = experiment_data.nu > 380.0 + scan_view = TODView(experiment_data, tod_samples, compsep_output=det_compsep_map) for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): - f_samp = det.fsamp - down_factor = int(f_samp) - indices_edges = np.arange(0, det.ntod, down_factor) - indices_centers = (indices_edges[1:] + indices_edges[:-1])//2 - ntod_down = indices_centers.size - - assert((ntod_down+1)*down_factor >= det.tod.shape[0]) - - pix, psi = det.get_pix_psi() - pix = pix[indices_centers] - psi = psi[indices_centers] - - s_orb = get_s_orb_TOD(det, experiment_data, pix) - sky_model_TOD = get_static_sky_TOD(det_compsep_map, pix, psi=psi) - gain = tod_samples.gain(iscan,idet) - - if calibrate_on_full_sky: - # Calibrate on the full sky model (static sky + orbital dipole), - # analogous to sample_relative_gain / sample_temporal_gain_variations. - s_cal = sky_model_TOD + s_orb - residual_tod = det.tod[:ntod_down*down_factor].reshape((ntod_down, down_factor)) - residual_tod = np.mean(residual_tod, axis=-1) - residual_tod -= gain*s_cal - else: - # Default: calibrate on the orbital dipole only. - s_cal = s_orb - residual_tod = det.tod[:ntod_down*down_factor].reshape((ntod_down, down_factor)) - residual_tod = np.mean(residual_tod, axis=-1) - residual_tod -= gain*sky_model_TOD # Subtracting sky signals. - residual_tod -= gain*s_orb - residual_tod += tod_samples.abs_gain*s_orb # Now we can add back in the orbital dipole. - - ### Solving Equation 16 from BP7 ### - mask = det.full_mask[indices_centers] - # White noise level, adjusted for downsampling. - sigma0_effective = tod_samples.noise_params[iscan,idet,0] * np.sqrt(1.0/f_samp) - # Inpaint masked regions with the calib signal times the absolute gain (+ white noise). - # TODO: This should ideally be some constrained realization with correlated noise, - # but I'm not 100% sure how to do so safely without potentially biasing the signal. - residual_tod[~mask] = tod_samples.abs_gain*s_cal[~mask]\ - + np.random.normal(0, sigma0_effective, s_cal[~mask].shape) - - N_inv_s = experiment_data.apply_N_inv(s_cal, tod_samples.noise_params[iscan,idet], samprate=1.0) - N_inv_d = experiment_data.apply_N_inv(residual_tod, tod_samples.noise_params[iscan,idet], samprate=1.0) + view = scan_view.focus(iscan, idet) + calib = view.get_abs_calib_tod( + downsample_factor=int(view.fsamp), + calibrate_on_full_sky=calibrate_on_full_sky, + preserve_target_gain=False, + ) + s_cal = calib.s_cal + residual_tod = calib.tod + + N_inv_s = experiment_data.apply_N_inv(s_cal, view.noise_params, samprate=1.0) + N_inv_d = experiment_data.apply_N_inv(residual_tod, view.noise_params, samprate=1.0) # Add to the numerator and denominator. # sum_s_T_N_inv_d += np.dot(s_cal[mask], N_inv_d[mask]) @@ -819,50 +791,16 @@ def sample_relative_gain(band_comm: MPI.Comm, experiment_data: DetGroupTOD, # local_r_T_N_inv_s = 0.0 local_r_T_N_inv_s = np.zeros(ndet, dtype=np.float32) + scan_view = TODView(experiment_data, tod_samples, compsep_output=det_compsep_map) for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): - f_samp = det.fsamp - down_factor = int(f_samp) - indices_edges = np.arange(0, det.ntod, down_factor) - indices_centers = (indices_edges[1:] + indices_edges[:-1])//2 - ntod_down = indices_centers.size - - # Define the residual for this sampling step, as per Eq. (17) - pix, psi = det.get_pix_psi() - pix = pix[indices_centers] - psi = psi[indices_centers] - - s_cal = get_static_sky_TOD(det_compsep_map, pix, psi) - - s_cal += get_s_orb_TOD(det, experiment_data, pix) - - gain = tod_samples.abs_gain + tod_samples.temporal_gain[iscan,idet] - residual_tod = det.tod[:ntod_down*down_factor].reshape((ntod_down, down_factor)) - residual_tod = np.mean(residual_tod, axis=-1) - residual_tod -= gain*s_cal - mask = det.full_mask[indices_centers] - # sigma0 = calc_sigma0_robust(residual_tod, mask) - - # Setup FFT-based calculation for N^-1 operations - # Ntod = residual_tod.shape[0] - # Nrfft = Ntod // 2 + 1 - # freqs = rfftfreq(Ntod, 1.0) - # inv_power_spectrum = np.zeros(Nrfft) - # sigma0, fknee, alpha = tod_samples.noise_params[iscan,idet,:] - # inv_power_spectrum[1:] = 1.0 / (sigma0**2 * (1 + (freqs[1:] / fknee))**alpha) - - # s_fft = forward_rfft(s_tot) - # N_inv_s_fft = s_fft * inv_power_spectrum - # N_inv_s = backward_rfft(N_inv_s_fft, Ntod) - - sigma0_effective = tod_samples.noise_params[iscan,idet,0] * np.sqrt(1.0/f_samp) - N_inv_s = experiment_data.apply_N_inv(s_cal, tod_samples.noise_params[iscan,idet], samprate=1.0) - - # Inpaint on the masked regions the sky signal times only the detector-residual gain. - residual_tod[~mask] = tod_samples.rel_gain[idet]*s_cal[~mask]\ - + np.random.normal(0, sigma0_effective, s_cal[~mask].shape) + view = scan_view.focus(iscan, idet) + calib = view.get_rel_calib_tod(downsample_factor=int(view.fsamp)) + s_cal = calib.s_cal + residual_tod = calib.tod + N_inv_s = experiment_data.apply_N_inv(s_cal, view.noise_params, samprate=1.0) s_T_N_inv_s_scan = np.dot(s_cal, N_inv_s) r_T_N_inv_s_scan = np.dot(residual_tod, N_inv_s) @@ -939,6 +877,7 @@ def sample_temporal_gain_variations(band_comm: MPI.Comm, experiment_data: DetGro # Local calculations on each rank A_qq_local = np.zeros((ndet, nscans_local), dtype=np.float64) b_q_local = np.zeros((ndet, nscans_local), dtype=np.float64) + scan_view = TODView(experiment_data, tod_samples, compsep_output=det_compsep_map) for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): @@ -949,47 +888,13 @@ def sample_temporal_gain_variations(band_comm: MPI.Comm, experiment_data: DetGro # (simply passing the full data through the FFTs seems like a bad idea because of # ringing from the large residual in the galactic plane). - f_samp = det.fsamp - down_factor = int(f_samp) - indices_edges = np.arange(0, det.ntod, down_factor) - indices_centers = (indices_edges[1:] + indices_edges[:-1])//2 - ntod_down = indices_centers.size - - # Per Eq. (26), the residual is d - (g0 + Delta_g)*s - pix, psi = det.get_pix_psi() - pix = pix[indices_centers] - psi = psi[indices_centers] - - s_cal = get_static_sky_TOD(det_compsep_map, pix, psi) - s_cal += get_s_orb_TOD(det, experiment_data, pix) - - gain = tod_samples.abs_gain + tod_samples.rel_gain[idet] - residual_tod = det.tod[:ntod_down*down_factor].reshape((ntod_down, down_factor)) - residual_tod = np.mean(residual_tod, axis=-1) - residual_tod -= gain*s_cal - - mask = det.full_mask[indices_centers] - - # FFT-based N^-1 operation setup - # Ntod = residual_tod.shape[0] - # Nrfft = Ntod // 2 + 1 - # freqs = rfftfreq(Ntod, 1.0) - # inv_power_spectrum = np.zeros(Nrfft) - # sigma0, fknee, alpha = tod_samples.noise_params[iscan,idet,:] - # inv_power_spectrum[1:] = 1.0 / (sigma0**2 * (1 + (freqs[1:] / fknee))**alpha) - - sigma0_effective = tod_samples.noise_params[iscan,idet,0] * np.sqrt(1.0/f_samp) - - # In the masked regions, inpaint the total sky model times only the temporal gain estimate. - residual_tod[~mask] = tod_samples.temporal_gain[iscan,idet]*s_cal[~mask]\ - + np.random.normal(0, sigma0_effective, s_cal[~mask].shape) - - # Calculate N^-1 * s_tot and N^-1 * residual_tod - # N_inv_s = backward_rfft(forward_rfft(s_tot) * inv_power_spectrum, Ntod) - # N_inv_r = backward_rfft(forward_rfft(residual_tod) * inv_power_spectrum, Ntod) + view = scan_view.focus(iscan, idet) + calib = view.get_temp_calib_tod(downsample_factor=int(view.fsamp)) + s_cal = calib.s_cal + residual_tod = calib.tod - N_inv_s = experiment_data.apply_N_inv(s_cal, tod_samples.noise_params[iscan,idet], samprate=1.0) - N_inv_r = experiment_data.apply_N_inv(residual_tod, tod_samples.noise_params[iscan,idet], samprate=1.0) + N_inv_s = experiment_data.apply_N_inv(s_cal, view.noise_params, samprate=1.0) + N_inv_r = experiment_data.apply_N_inv(residual_tod, view.noise_params, samprate=1.0) # Calculate elements for the linear system A_qq = np.dot(s_cal, N_inv_s) @@ -1129,13 +1034,11 @@ def process_tod(mpi_info: Bunch, experiment_data: DetGroupTOD, tod_samples (TODSamples): Updated sampled TOD parameters. """ # Steps: - # 1. Initialize n_corr_est and alpha/fknee values. - # 2. Sample the gain from the sky-subtracted TOD (Skipped on iter==1 because we don't have - # a reliable sky-subtracted TOD). - # 3. Estimate White noise from the sky-subtracted TOD. - # 4. Sample correlated noise and PS parameters (skipped on iter==1). + # 1. Detect and store jump corrections from the flag stream. + # 2. Estimate white noise from the jump-corrected, sky-subtracted TOD. + # 3. Sample the gain from the jump-corrected, sky-subtracted TOD. + # 4. Sample correlated noise and PS parameters. # 5. Mapmaking on TOD - corr_noise_TOD - orb_dipole_TOD. - # (In other words, on iteration 1 we do just do White noise estimation -> Mapmaking.) timing_dict = {} waittime_dict = {} From 42158ffc64a2b8a44193c29a75df19aa2cabb074 Mon Sep 17 00:00:00 2001 From: jgslunde Date: Fri, 29 May 2026 15:43:46 +0200 Subject: [PATCH 13/23] Added jump-corrections as a sampling step. The corrections are applied to the TOD upon reads. --- .../data_models/jump_corrections.py | 205 ++++++++++++++++++ src/commander4/tod_processing.py | 100 ++++++++- 2 files changed, 299 insertions(+), 6 deletions(-) create mode 100644 src/commander4/data_models/jump_corrections.py diff --git a/src/commander4/data_models/jump_corrections.py b/src/commander4/data_models/jump_corrections.py new file mode 100644 index 0000000..a243445 --- /dev/null +++ b/src/commander4/data_models/jump_corrections.py @@ -0,0 +1,205 @@ +from dataclasses import dataclass + +import h5py +import numpy as np +from numpy.typing import NDArray + + +@dataclass(slots=True) +class JumpCorrection: + """Additive offsets that are applied after one or more detected jump locations.""" + + locations: NDArray[np.int64] + offsets: NDArray[np.float32] + + def __post_init__(self): + """Normalize storage and validate that the jump metadata is well-formed.""" + self.locations = np.asarray(self.locations, dtype=np.int64) + self.offsets = np.asarray(self.offsets, dtype=np.float32) + if self.locations.ndim != 1 or self.offsets.ndim != 1: + raise ValueError("JumpCorrection expects 1-D locations and offsets arrays.") + if self.locations.size != self.offsets.size: + raise ValueError("jump locations and offsets must have the same length.") + + @classmethod + def empty(cls) -> "JumpCorrection": + """Return an empty correction object for detectors without any detected jumps.""" + return cls(np.empty(0, dtype=np.int64), np.empty(0, dtype=np.float32)) + + @property + def size(self) -> int: + """Number of jump discontinuities represented by this correction.""" + return int(self.locations.size) + + def is_empty(self) -> bool: + """Return whether this correction contains any offsets.""" + return self.size == 0 + + def apply( + self, + tod: NDArray[np.floating], + *, + inplace: bool = False, + ) -> NDArray[np.floating]: + """Apply the stored offsets to a TOD array in detector units.""" + corrected_tod = tod if inplace else np.array(tod, copy=True) + for jump_location, jump_offset in zip(self.locations, self.offsets): + corrected_tod[jump_location:] += jump_offset + return corrected_tod + + @classmethod + def detect( + cls, + tod: NDArray[np.floating], + flag: NDArray[np.integer], + valid_mask: NDArray[np.bool_], + n_window: int, + *, + jump_bitmask: int, + ) -> tuple["JumpCorrection", int]: + """Estimate jump offsets from flagged regions and neighboring valid samples. + + Args: + tod: Raw detector TOD in detector units. + flag: Per-sample flag stream. Contiguous regions with a non-zero + ``flag & jump_bitmask`` mark jumps. + valid_mask: Boolean mask defining which samples are allowed in the pre/post windows. + n_window: Number of valid samples to average on each side of a jump. + jump_bitmask: Integer bitmask used to tag jumps in the flag stream. + + Returns: + A ``JumpCorrection`` object plus the number of flagged jump regions that were skipped + because either side lacked enough valid samples. + """ + if n_window < 1: + raise ValueError("n_window must be >= 1.") + if jump_bitmask < 1: + raise ValueError("jump_bitmask must be >= 1.") + + jump_indices = np.flatnonzero((flag & jump_bitmask) != 0) + if jump_indices.size == 0: + return cls.empty(), 0 + + breaks = np.flatnonzero(np.diff(jump_indices) > 1) + jump_starts = np.concatenate(([jump_indices[0]], jump_indices[breaks + 1])) + jump_stops = np.concatenate((jump_indices[breaks] + 1, [jump_indices[-1] + 1])) + valid_indices = np.flatnonzero(valid_mask) + corrected_tod = np.array(tod, copy=True) + jump_locations = [] + jump_offsets = [] + num_skipped = 0 + + for jump_start, jump_stop in zip(jump_starts, jump_stops): + before_stop = np.searchsorted(valid_indices, jump_start, side="left") + after_start = np.searchsorted(valid_indices, jump_stop, side="left") + before_indices = valid_indices[max(0, before_stop - n_window):before_stop] + after_indices = valid_indices[after_start:after_start + n_window] + if before_indices.size < n_window or after_indices.size < n_window: + num_skipped += 1 + continue + + mean_before = np.mean(corrected_tod[before_indices], dtype=np.float64) + mean_after = np.mean(corrected_tod[after_indices], dtype=np.float64) + jump_offset = float(mean_before - mean_after) + + # Later jumps should be estimated relative to the already corrected baseline. + corrected_tod[jump_stop:] += jump_offset + jump_locations.append(int(jump_stop)) + jump_offsets.append(jump_offset) + + return cls(jump_locations, jump_offsets), num_skipped + + +class JumpCatalog: + """Per-scan and per-detector container for jump corrections.""" + + def __init__(self, entries: NDArray): + """Wrap a 2-D object array of ``JumpCorrection`` instances.""" + if entries.ndim != 2: + raise ValueError("JumpCatalog expects a 2-D object array.") + self._entries = entries + self.nscans, self.ndet = entries.shape + + @classmethod + def empty(cls, nscans: int, ndet: int) -> "JumpCatalog": + """Allocate an empty jump catalog for one MPI rank's local scans.""" + entries = np.empty((nscans, ndet), dtype=object) + for iscan in range(nscans): + for idet in range(ndet): + entries[iscan, idet] = JumpCorrection.empty() + return cls(entries) + + @classmethod + def from_hdf5( + cls, + file: h5py.File, + local_indices: list[int], + ndet: int, + ) -> "JumpCatalog": + """Reconstruct a local jump catalog from packed HDF5 datasets when present.""" + catalog = cls.empty(len(local_indices), ndet) + dataset_names = {"jump_counts", "jump_locations", "jump_offsets"} + if not dataset_names.issubset(file.keys()): + return catalog + + counts_global = file["jump_counts"][:] + locations_global = file["jump_locations"][:] + offsets_global = file["jump_offsets"][:] + counts_flat = counts_global.reshape(-1) + starts_flat = np.cumsum(counts_flat, dtype=np.int64) - counts_flat + + for iscan_local, iscan_global in enumerate(local_indices): + row_start = iscan_global * ndet + for idet in range(ndet): + flat_index = row_start + idet + count = int(counts_flat[flat_index]) + start = int(starts_flat[flat_index]) + stop = start + count + catalog.set( + iscan_local, + idet, + JumpCorrection(locations_global[start:stop], offsets_global[start:stop]), + ) + return catalog + + def get(self, iscan: int, idet: int) -> JumpCorrection: + """Return the correction object for one detector in one scan.""" + jump = self._entries[iscan, idet] + return JumpCorrection.empty() if jump is None else jump + + def set(self, iscan: int, idet: int, jump: JumpCorrection | None): + """Store one detector-local correction, replacing ``None`` by an empty correction.""" + self._entries[iscan, idet] = JumpCorrection.empty() if jump is None else jump + + def apply( + self, + tod: NDArray[np.floating], + iscan: int, + idet: int, + *, + inplace: bool = False, + ) -> NDArray[np.floating]: + """Apply the stored correction for one detector and scan to a TOD array.""" + return self.get(iscan, idet).apply(tod, inplace=inplace) + + def pack(self) -> tuple[NDArray[np.int64], NDArray[np.int64], NDArray[np.float32]]: + """Pack ragged corrections into counts plus flat arrays for MPI/HDF5 output.""" + counts = np.zeros((self.nscans, self.ndet), dtype=np.int64) + flat_locations = [] + flat_offsets = [] + + for iscan in range(self.nscans): + for idet in range(self.ndet): + jump = self.get(iscan, idet) + counts[iscan, idet] = jump.size + if not jump.is_empty(): + flat_locations.append(jump.locations) + flat_offsets.append(jump.offsets) + + packed_locations = ( + np.concatenate(flat_locations) if flat_locations else np.empty(0, dtype=np.int64) + ) + packed_offsets = ( + np.concatenate(flat_offsets) if flat_offsets else np.empty(0, dtype=np.float32) + ) + return counts, packed_locations, packed_offsets \ No newline at end of file diff --git a/src/commander4/tod_processing.py b/src/commander4/tod_processing.py index 509dfd1..b7684b7 100644 --- a/src/commander4/tod_processing.py +++ b/src/commander4/tod_processing.py @@ -683,23 +683,99 @@ def estimate_white_noise(experiment_data: DetGroupTOD, tod_samples: TODSamples, Output: tod_samples (TODSamples): Updated TOD samples with sigma0 estimates. """ + scan_view = TODView(experiment_data, tod_samples, compsep_output=det_compsep_map) for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): - pix, psi = det.get_pix_psi() + view = scan_view.focus(iscan, idet) # FIXME: Should maybe n_corr be subtracted here as well? - gain = tod_samples.gain(iscan, idet) - sky_subtracted_tod = det.tod.copy() - sky_subtracted_tod -= gain*get_static_sky_TOD(det_compsep_map, pix, psi=psi) - sky_subtracted_tod -= gain*get_s_orb_TOD(det, experiment_data, pix) - mask = det.full_mask + sky_subtracted_tod = view.get_tod( + subtract=(("sky", TODView._ALL_GAIN_TERMS), + ("orbital_dipole", TODView._ALL_GAIN_TERMS)), + ) + mask = view.full_mask sigma0 = calc_sigma0_robust(sky_subtracted_tod, mask) logassert(sigma0 != 0, "sigma0 is 0, which should never happen.", logger) + logassert(sigma0 != np.inf, "sigma0 is inf, which should never happen.", logger) tod_samples.noise_params[iscan,idet,0] = sigma0 if iscan == len(experiment_data.scans) - 1: log_memory("sigma0-est") return tod_samples +def sample_jump_detection(band_comm: MPI.Comm, experiment_data: DetGroupTOD, + tod_samples: TODSamples, params: Bunch) -> TODSamples: + """Detect jump discontinuities from the flag stream and store additive post-jump offsets. + + A jump is identified by a contiguous region with a non-zero + ``flag & experiments.[experiment_name].jump_bitmask``. For each region, the offset is + estimated from the last ``N`` valid samples before the jump and the first ``N`` valid samples + after it, where validity is defined by ``full_mask``. The correction is then applied to all + later samples when a TOD is requested through ``TODView.get_tod()``. + """ + n_window = int(getattr(params.general, "jump_detection_window", 10)) + if n_window < 1: + raise ValueError("jump_detection_window must be >= 1.") + experiment_params = params.experiments[experiment_data.experiment_name] + if "jump_bitmask" not in experiment_params or experiment_params.jump_bitmask is None: + raise ValueError( + "Jump sampling is enabled, but " + f"experiments.{experiment_data.experiment_name}.jump_bitmask is not specified." + ) + jump_bitmask = int(experiment_params.jump_bitmask) + + scan_view = TODView(experiment_data, tod_samples) + num_applied_local = 0 + num_skipped_local = 0 + offsets_local = [] + jump_counts_local = [] + + for iscan, scan in enumerate(experiment_data.scans): + for idet, det in enumerate(scan.detectors): + view = scan_view.focus(iscan, idet) + if getattr(view.detector, "_flag_encoded", None) is None or not hasattr(view.detector, "_full_mask"): + tod_samples.jumps.set(iscan, idet, None) + jump_counts_local.append(0) + continue + jump, num_skipped = JumpCorrection.detect( + view.tod, + view.flag, + view.full_mask, + n_window, + jump_bitmask=jump_bitmask, + ) + tod_samples.jumps.set(iscan, idet, jump) + jump_counts_local.append(jump.size) + num_skipped_local += num_skipped + if not jump.is_empty(): + offsets_local.extend(jump.offsets.astype(np.float64, copy=False)) + num_applied_local += jump.size + + num_applied = band_comm.reduce(num_applied_local, op=MPI.SUM, root=0) + num_skipped = band_comm.reduce(num_skipped_local, op=MPI.SUM, root=0) + gathered_offsets = band_comm.gather(np.asarray(offsets_local, dtype=np.float64), root=0) + gathered_jump_counts = band_comm.gather(np.asarray(jump_counts_local, dtype=np.int32), root=0) + + if band_comm.Get_rank() == 0: + all_jump_counts = np.concatenate(gathered_jump_counts) if gathered_jump_counts else np.empty(0) + if all_jump_counts.size > 0: + logger.debug( + f"Band {experiment_data.band_name} jump counts per detector-scan: " + f"min={np.min(all_jump_counts)}, avg={np.mean(all_jump_counts):.2f}, " + f"max={np.max(all_jump_counts)} over {all_jump_counts.size} samples." + ) + if num_applied > 0: + all_offsets = np.concatenate([arr for arr in gathered_offsets if arr.size > 0]) + logger.info(f"Band {experiment_data.band_name} jump detection: applied {num_applied} " + f"offsets, skipped {num_skipped}, median |offset| = " + f"{np.median(np.abs(all_offsets)):.3e}.") + elif num_skipped > 0: + logger.info(f"Band {experiment_data.band_name} jump detection skipped {num_skipped} " + f"flagged regions because there were not enough valid samples around them.") + + log_memory("jump-detect") + return tod_samples + + def sample_absolute_gain(band_comm: MPI.Comm, experiment_data: DetGroupTOD, tod_samples: TODSamples, det_compsep_map: NDArray): @@ -1046,6 +1122,18 @@ def process_tod(mpi_info: Bunch, experiment_data: DetGroupTOD, det_comm = mpi_info.det.comm band_comm = mpi_info.band.comm TOD_comm = mpi_info.tod.comm + ### JUMP DETECTION ### + if getattr(params.general, "sample_jump_detection", True) and iter >= int( + getattr(params.general, "sample_jump_detection_from_iter_num", 1) + ): + t0 = time.time() + with benchmark("jump-detect"): + tod_samples = sample_jump_detection(band_comm, experiment_data, tod_samples, params) + timing_dict["jump-detect"] = time.time() - t0 + if mpi_info.band.is_master: + logger.info(f"Chain {chain} iter{iter} {experiment_data.nu}GHz: Finished jump " + f"detection in {timing_dict['jump-detect']:.1f}s.") + ### WHITE NOISE ESTIMATION ### t0 = time.time() with benchmark("sigma0-est"): From e5472966f76e96b984a237fbe864dedf44d753b1 Mon Sep 17 00:00:00 2001 From: jgslunde Date: Tue, 2 Jun 2026 11:26:49 +0200 Subject: [PATCH 14/23] Refactored many things. --- params/Planck_TODs/param_PlanckTODs_CG.yml | 7 +- .../Planck_TODs/param_PlanckTODs_perpix.yml | 2 +- src/commander4/communication.py | 119 ++--- src/commander4/compsep_processing.py | 133 ++--- src/commander4/data_models/detector_TOD.py | 8 +- .../experiments/SO/tod_reader_SO_SAT.py | 2 +- .../experiments/planck/tod_reader_planck.py | 35 +- src/commander4/mpi_management.py | 21 +- src/commander4/output/write_chains_files.py | 20 +- src/commander4/sky_models/component.py | 466 ++++++++++++------ src/commander4/sky_models/sky_model.py | 36 +- .../solvers/perpix_compsep_solver.py | 2 +- src/commander4/tod_processing.py | 48 +- src/commander4/utils/CG_mapmaker.py | 4 +- src/commander4/utils/math_operations.py | 3 +- 15 files changed, 539 insertions(+), 367 deletions(-) diff --git a/params/Planck_TODs/param_PlanckTODs_CG.yml b/params/Planck_TODs/param_PlanckTODs_CG.yml index 06f4ed5..48010c4 100644 --- a/params/Planck_TODs/param_PlanckTODs_CG.yml +++ b/params/Planck_TODs/param_PlanckTODs_CG.yml @@ -331,11 +331,11 @@ experiments: tod_files_prefix: "LFI_030_" num_MPI_tasks: 64 #24 eval_nside: 512 - data_nside: 512 fwhm: 32.4 # [arcmin] freq: 28.4 # [GHz] fsamp: 32.5079365079365 # [Hz] polarization: "IQU" + noise_psd: "NoisePSDOof" processing_mask: "/mn/stornext/d5/data/duncanwa/WMAP/data/mask_proc_030_res_v5.fits" filelist: "/mn/stornext/d16/cmbco/bp/mathew/test/filelist_30.txt" detectors: @@ -360,10 +360,8 @@ experiments: tod_files_prefix: "LFI_044_" num_MPI_tasks: 40 eval_nside: 512 - data_nside: 512 fwhm: 27.1 # [arcmin] freq: 44.1 # [GHz] - fsamp: 46.5454545454545 polarization: "IQU" processing_mask: "/mn/stornext/d5/data/duncanwa/WMAP/data/mask_proc_044_res_v5.fits" filelist: "/mn/stornext/d16/cmbco/bp/mathew/test/filelist_44.txt" @@ -397,14 +395,11 @@ experiments: tod_files_prefix: "LFI_070_" num_MPI_tasks: 64 eval_nside: 1024 - data_nside: 1024 fwhm: 13.3 # [arcmin] freq: 70.1 # [GHz] - fsamp: 78.7692307692308 polarization: "IQU" processing_mask: "/mn/stornext/d5/data/duncanwa/WMAP/data/mask_proc_070_res_v5.fits" filelist: "/mn/stornext/d16/cmbco/bp/mathew/test/filelist_70.txt" - detectors: detectors: 18M: name: "18M" diff --git a/params/Planck_TODs/param_PlanckTODs_perpix.yml b/params/Planck_TODs/param_PlanckTODs_perpix.yml index 0e82688..3db2c3e 100644 --- a/params/Planck_TODs/param_PlanckTODs_perpix.yml +++ b/params/Planck_TODs/param_PlanckTODs_perpix.yml @@ -21,7 +21,7 @@ general: sample_abs_gain_from_iter_num: 2 sample_rel_gain: True sample_rel_gain_from_iter_num: 2 - sample_temporal_gain: False + sample_temporal_gain: True sample_temporal_gain_from_iter_num: 2 perform_compsep: true diff --git a/src/commander4/communication.py b/src/commander4/communication.py index 7dae1a0..bace09c 100644 --- a/src/commander4/communication.py +++ b/src/commander4/communication.py @@ -1,15 +1,47 @@ import numpy as np import logging +from collections.abc import Mapping from numpy.typing import NDArray from pixell.bunch import Bunch from commander4.data_models.detector_map import DetectorMap from commander4.data_models.detector_group_TOD import DetGroupTOD from commander4.maps_from_file import read_data_map_from_file -from commander4.output import log +from commander4.utils.execution_ids import get_execution_band_id logger = logging.getLogger(__name__) + +def _get_compsep_sender_id_for_tod_band(todproc_my_band_id: str, + senders: Mapping[str, int]) -> str: + """Return the CompSep execution-view ID that sends the sky model back to a TOD band. + + For bands that have both intensity and QU execution views, the intensity rank sends the + already reassembled `SkyModel`. QU-only bands fall back to their `_QU` identifier. + """ + intensity_key = get_execution_band_id(todproc_my_band_id, "I") + if intensity_key in senders: + return intensity_key + pol_key = get_execution_band_id(todproc_my_band_id, "QU") + if pol_key in senders: + return pol_key + raise KeyError(f"No CompSep sender found for TOD band '{todproc_my_band_id}'.") + + +def _should_send_compsep_result(compsep_my_band_id: str, + destinations: Mapping[str, int] | None) -> bool: + """Return whether this CompSep execution view should send the realized sky model to TOD.""" + if destinations is None or compsep_my_band_id not in destinations: + return False + if compsep_my_band_id.endswith("_QU"): + paired_intensity_id = get_execution_band_id( + compsep_my_band_id.removesuffix("_QU"), + "I", + ) + if paired_intensity_id in destinations: + return False + return True + ########################################################### # ON TOD SIDE ########################################################### @@ -18,27 +50,27 @@ # (uppercase) whereever possible (e.g. where arrays are communicated). def receive_compsep(mpi_info: Bunch, experiment_data: DetGroupTOD, todproc_my_band_id: str, senders: dict[str, int]) -> NDArray[np.floating]: - """ MPI-receive the results from compsep (used in conjunction with send_compsep). + """Receive the CompSep sky model, realize the local band map, and broadcast it within TOD. - Input: + Args: mpi_info (Bunch): The data structure containing all MPI relevant data. experiment_data (DetGroupTOD): The experiment TOD data, used to determine band frequency, resolution, and polarization for evaluating the sky model. - todproc_my_band_id (str): The string uniquely indentifying the experiment+band of this rank - (example: 'PlanckLFI$$$30GHz'). + todproc_my_band_id (str): The string uniquely indentifying the band of this rank + (example: '30GHz'). senders (dict[str, int]): A dictionary mapping the string in 'todproc_my_band_id' to the - world rank of the sender task (on the CompSep side). + world rank of the sender task (on the CompSep side), keyed by execution-view band ID. Returns: - detector_map_arr (NDArray): The sky model evaluated at the band frequency and resolution, - broadcast to all processes on the band communicator. + NDArray: The sky model realized at the local band frequency and resolution, broadcast to + all processes on the band communicator. """ world_comm = mpi_info.world.comm band_comm = mpi_info.band.comm is_band_master = mpi_info.band.is_master if is_band_master: - # As only the compsep rank holding I will send everything at once. - sky_model = world_comm.recv(source=senders[todproc_my_band_id+'_I']) + source_band_id = _get_compsep_sender_id_for_tod_band(todproc_my_band_id, senders) + sky_model = world_comm.recv(source=senders[source_band_id]) else: sky_model = None # Currently all TOD MPI ranks need a copy of the relevant detector map, @@ -58,44 +90,21 @@ def send_tod(mpi_info: Bunch, tod_map_dict: dict[DetectorMap], todproc_my_band_i mpi_info (Bunch): The data structure containing all MPI relevant data. tod_map_dict (dict[str, DetectorMap]): The output maps keyed by polarization component (e.g. 'I', 'QU') from process_tod for the band belonging to this process. - todproc_my_band_id (str): The string uniquely indentifying the experiment+band of this rank, - regardless of polarization (example: 'PlanckLFI$$$30GHz'). + todproc_my_band_id (str): The string uniquely indentifying the band of this rank, + regardless of polarization (example: '30GHz'). receivers (Bunch): Maps a band identifier to the band master on the compsep side. """ if mpi_info.tod.is_master: logger.info(f"Compsep band masters: {mpi_info.world.compsep_band_masters}") if mpi_info.band.is_master: - # my_pol = mpi_info.world.tod_band_pols[todproc_my_band_id] - #I - my_todproc_pols = tod_map_dict.keys() - for pol in my_todproc_pols: - target_band = todproc_my_band_id+'_'+pol - if target_band in receivers.keys(): #Compsep wants an I pol for todproc_my_band_id - mpi_info.world.comm.send(tod_map_dict[pol], dest=receivers[target_band]) + for pol, detector_map in tod_map_dict.items(): + target_band = get_execution_band_id(todproc_my_band_id, pol) + if target_band in receivers.keys(): + mpi_info.world.comm.send(detector_map, dest=receivers[target_band]) else: logger.info(f"Pol-{pol} TOD-processing result discarded, "\ f"as band {todproc_my_band_id} does not require it on compsep side.") - # target_band = todproc_my_band_id+'_I' - # if target_band in receivers.keys(): #Compsep wants an I pol for todproc_my_band_id - # log.logassert("I" in my_todproc_pols, - # "Polarization 'I' requested from compsep but not found in TOD processing", - # logger) - # mpi_info.world.comm.send(tod_map_dict["I"], dest=receivers[target_band]) - # else: - # logger.info(f"Intensity TOD-processing result discarded, as band {todproc_my_band_id} "\ - # "is only defined for QU.") - # #QU - # target_band = todproc_my_band_id+'_QU' - # if target_band in receivers.keys(): #Compsep wants a QU pol for todproc_my_band_id - # log.logassert("QU" in my_pol, - # "Polarization 'QU' requested from compsep but not found in TOD processing", - # logger) - # mpi_info.world.comm.send(tod_map_dict["QU"], dest=receivers[target_band]) - # else: - # logger.info(f"QU TOD-processing result discarded, as band {todproc_my_band_id} is "\ - # "only defined for Intensity.") - ########################################################### # ON COMPSEP SIDE @@ -108,11 +117,11 @@ def receive_tod(mpi_info: Bunch, senders: dict[str,int], my_band: Bunch, compsep Input: mpi_info (Bunch): The data structure containing all MPI relevant data. senders: (dict[str->int]): A dictionary mapping a string uniquely identifying each - experiment+band to the world rank of the sender task (on the CompSep side). + band to the world rank of the sender task (on the CompSep side). my_band (Bunch): The section of the parameter file corresponding to this CompSep band, as a "Bunch" type, it also has an 'identifier' field. - compsep_band_id (str): The string uniquely indentifying the experiment+band+pol of this - rank (example: 'PlanckLFI$$$30GHz_I'). + compsep_band_id (str): The string uniquely indentifying the band+pol of this + rank (example: '30GHz_I'). curr_tod_output (DetectorMap): The current map output from the TOD process. Should be None unless map is read from file already in a previous iteration. @@ -131,27 +140,25 @@ def receive_tod(mpi_info: Bunch, senders: dict[str,int], my_band: Bunch, compsep f" from TOD process with global rank {senders[compsep_band_id]}") curr_tod_output = mpi_info.world.comm.recv(source=senders[compsep_band_id]) + return curr_tod_output -def send_compsep(mpi_info: Bunch, compsep_my_band_id: str, band_sky_map, +def send_compsep(mpi_info: Bunch, compsep_my_band_id: str, sky_model, destinations: dict[str, int]|None) -> None: - """ MPI-send the results from compsep to a destinations on the TOD processing side - (used in conjunction with receive_compsep). Assumes the COMM_WORLD communicator. + """Send the CompSep sky model back to TOD when this execution view owns the return path. + + For split IQU bands, the QU master first transfers its alms to the I master inside the CompSep + communicator. Only the matching `_I` execution view then sends the fully reassembled sky model + back to TOD. Pure QU bands send directly from their `_QU` execution view. Args: mpi_info (Bunch): The data structure containing all MPI relevant data. - compsep_my_band_id (str): The string uniquely indentifying the experiment+band+pol of this - rank (example: 'PlanckLFI$$$30GHz_I'). - band_sky_map (NDArray): A sky realization at a given band frequency. + compsep_my_band_id (str): The string uniquely indentifying the band+pol of this + rank (example: '30GHz_I'). + sky_model: The realized sky model object for this Gibbs sample. destinations (dict[str->int]): A dictionary mapping the string in 'compsep_my_band_id' to the world rank of the destination task (on the TOD side) (This is the same as is found in mpi_info) """ - if destinations is not None: - # If the band our rank is holding is not in "destinations", it means it did not come from - # TOD-processing, and should not be sent back either. - if compsep_my_band_id in destinations: - # Currently in the end of compsep_processing the polarization ranks send over to the - # Intensity ones which then in turn broadcast and evaluate the sky. - if compsep_my_band_id.endswith('_I'): - mpi_info.world.comm.send(band_sky_map, dest=destinations[compsep_my_band_id]) + if _should_send_compsep_result(compsep_my_band_id, destinations): + mpi_info.world.comm.send(sky_model, dest=destinations[compsep_my_band_id]) diff --git a/src/commander4/compsep_processing.py b/src/commander4/compsep_processing.py index 9900f23..5282099 100644 --- a/src/commander4/compsep_processing.py +++ b/src/commander4/compsep_processing.py @@ -1,5 +1,6 @@ import numpy as np import logging +from copy import deepcopy from pixell.bunch import Bunch from commander4.output.log import logassert @@ -9,14 +10,17 @@ from commander4.solvers.CG_compsep_solver import CompSepSolver from commander4.solvers.perpix_compsep_solver import solve_compsep_perpix from commander4.output.write_chains_files import write_compsep_chain_to_file +from commander4.utils.execution_ids import get_execution_band_id logger = logging.getLogger(__name__) - def init_compsep_processing(mpi_info: Bunch, params: Bunch)\ -> tuple[CompList, str, dict[str, int], Bunch]: - """ To be run once before starting component separation processing. - Determines whether the process is compsep master, and the number of bands. + """Set up the rank-local execution view for component separation. + + Each CompSep rank owns exactly one execution view of one band. The global CompSep rank space + is split into a contiguous intensity block followed by a contiguous QU block, and we match the + current rank against those two logical streams here. Args: mpi_info (Bunch): The data structure containing all MPI relevant data. @@ -25,48 +29,12 @@ def init_compsep_processing(mpi_info: Bunch, params: Bunch)\ Returns: mpi_info (Bunch): The data structure containing all MPI relevant data, modified to contain also the band masters dictionary. - band_identifier (str): Unique string for the experiment+band this rank is working on. + band_identifier (str): Unique string for the band execution view this rank is working on. my_band (Bunch): A subset of the full parameter file for the band this rank is working on. """ logger.info(f"CompSep: Hello from CompSep-rank {mpi_info.compsep.rank} (on machine "\ f"{mpi_info.processor_name}), dedicated to band {mpi_info.compsep.rank}.") - ### Creating list of all components ### - # comp_list = [] - # for component_str in params.components: - # component = params.components[component_str] - # if component.enabled: - # if component.params.lmax == "full": - # component.params.lmax = (params.general.nside*5)//2 - # if component.params.polarization == "I": #I-only - # # 'getattr' loads the class specified by "component_class" from model.component. - # # This class is then instantiated with the "params" specified, and appended to - # # the components list. - # comp_list.append(getattr(component_lib, component.component_class)(component.params, - # params.general, allocate_empty_alms=True)) - # elif component.params.polarization == "QU": #QU-only - # comp_list.append(getattr(component_lib, component.component_class)(component.params, - # params.general, allocate_empty_alms=True)) - # elif component.params.polarization == "IQU": - # #I - # comp_list.append(getattr(component_lib, component.component_class)( - # component.params, - # params.general, - # allocate_empty_alms=True, - # longname = component.params.longname+"_Instensity", - # shortname = component.params.longname+"_I", - # eval_pol="I")) - # #QU - # comp_list.append(getattr(component_lib, component.component_class)( - # component.params, - # params.general, - # allocate_empty_alms=True, - # longname = component.params.longname+"_Polarization", - # shortname = component.params.longname+"_QU", - # eval_pol="QU")) - # else: - # raise ValueError(f"Unrecognized polarization in parameter file for component {component_str}") - comp_list = CompList.init_from_params(params.components, params) ### Setting up info for each band, including where to get the data from ### @@ -75,51 +43,40 @@ def init_compsep_processing(mpi_info: Bunch, params: Bunch)\ current_band_idx_I = 0 current_band_idx_QU = mpi_info.compsep.QU_master band_identifier = None + my_band = None for band_str in params.CompSep_bands: # Intensity band = params.CompSep_bands[band_str] if band.enabled: if band.polarization == "I": if current_band_idx_I == mpi_info.compsep.rank: - my_band = band - if my_band.get_from != "file": - band_identifier = f"{my_band.get_from}$$${band_str}" - else: - band_identifier = band_str + my_band = deepcopy(band) + band_identifier = get_execution_band_id(band_str, "I") logger.info(f"Rank {mpi_info.compsep.rank} just matched band {band_identifier}") my_band.identifier = band_identifier + my_band.polarization = "I" current_band_idx_I += 1 elif band.polarization == "QU": if current_band_idx_QU == mpi_info.compsep.rank: - my_band = band - if my_band.get_from != "file": - band_identifier = f"{my_band.get_from}$$${band_str}" - else: - band_identifier = band_str + my_band = deepcopy(band) + band_identifier = get_execution_band_id(band_str, "QU") logger.info(f"Rank {mpi_info.compsep.rank} matched band {band_identifier}") my_band.identifier = band_identifier + my_band.polarization = "QU" current_band_idx_QU += 1 elif band.polarization == "IQU": - #if the band is defined as IQU we split it in two. - #I + # Each IQU band occupies one rank in the intensity block and one in the QU block. if current_band_idx_I == mpi_info.compsep.rank: - my_band = band - if my_band.get_from != "file": - band_identifier = f"{my_band.get_from}$$${band_str}_I" - else: - band_identifier = band_str+"_I" + my_band = deepcopy(band) + band_identifier = get_execution_band_id(band_str, "I") logger.info(f"Rank {mpi_info.compsep.rank} just matched band {band_identifier}") my_band.identifier = band_identifier my_band.polarization = "I" current_band_idx_I += 1 - #QU if current_band_idx_QU == mpi_info.compsep.rank: - my_band = band - if my_band.get_from != "file": - band_identifier = f"{my_band.get_from}$$${band_str}_QU" - else: - band_identifier = band_str+"_QU" + my_band = deepcopy(band) + band_identifier = get_execution_band_id(band_str, "QU") logger.info(f"Rank {mpi_info.compsep.rank} matched band {band_identifier}") my_band.identifier = band_identifier my_band.polarization = "QU" @@ -134,6 +91,11 @@ def init_compsep_processing(mpi_info: Bunch, params: Bunch)\ logassert(current_band_idx_QU == mpi_info.compsep.size, "Number of acquired QU bands "\ f"({current_band_idx_QU}) do not match number of MPI tasks assigned to QU "\ f"({mpi_info.compsep.QU_master})", logger) + if my_band is None or band_identifier is None: + logassert(False, + f"CompSep rank {mpi_info.compsep.rank} was not assigned to any enabled band. " + "Check that CompSep_bands matches the configured I/QU rank counts.", + logger) data_world = (band_identifier, mpi_info.world.rank) data_compsep = (band_identifier, mpi_info.compsep.rank) @@ -169,8 +131,8 @@ def process_compsep(mpi_info: Bunch, detector_data: DetectorMap, iter: int, chai compsep_rank = mpi_info.compsep.rank subcolor = mpi_info.compsep.subcolor #Subcolor splits the compsep ranks into: Pol -> 1, Int -> 0 compsep_subcomm = mpi_info.compsep.subcomm - # comp_sublist = split_complist(comp_list, subcolor) - comp_sublist = comp_list.split(subcolor) + target_pol = "I" if subcolor == 0 else "QU" + comp_sublist = comp_list.split_for_eval_pol(target_pol) ### 2. SOLVE COMPSEP: band maps -> component alms (either by per-pixel or CG solver) ### if params.general.pixel_compsep_sampling: @@ -181,34 +143,15 @@ def process_compsep(mpi_info: Bunch, detector_data: DetectorMap, iter: int, chai comp_sublist = compsep_solver.solve(comp_sublist) ### 3. CLEANUP: Gather I+QU alm solutions and make plots. ### - # Pol master sends the portion of list to the Intensity master rank, - # and then it will broadcast through the compsep_comm - if mpi_info.compsep.is_QU_master: - t=0 - for comp in comp_sublist: - logger.debug(f"[MPI Comm] Sending {comp.shortname} from QU {comp._data.shape} "\ - f"{comp._data.dtype} {t} to {mpi_info.compsep.I_master}") - compsep_comm.Send(comp._data, dest=mpi_info.compsep.I_master, tag=t) - t+=1 - - if mpi_info.compsep.is_I_master: - t_pol=0 - t_int=0 - for comp in comp_list: - if comp.is_pol: #if it is a pol component receive it from the QU_master - logger.debug(f"[MPI Comm] Receiving {comp.shortname} from QU {comp._data.shape} "\ - f"{comp._data.dtype} {t_pol} from {mpi_info.compsep.QU_master}") - compsep_comm.Recv(comp._data, source=mpi_info.compsep.QU_master, tag=t_pol) - t_pol+=1 - else: # Otherwise it copy it over from the local intensity sublist held on I_master - logger.debug(f"[MPI Comm] Copying {comp.shortname} from I {comp._data.shape} "\ - f"{comp._data.dtype} {t_int} from local I") - comp._data = comp_sublist[t_int]._data - t_int+=1 - - # In any case the component, received from QU or computed and accumulated on I, is bcasted - for comp in comp_list: - comp.bcast_data_blocking(compsep_comm, root=mpi_info.compsep.master) + comp_list.reassemble_from_split_solution( + comp_sublist, + compsep_comm, + is_I_master=mpi_info.compsep.is_I_master, + is_QU_master=mpi_info.compsep.is_QU_master, + I_master=mpi_info.compsep.I_master, + QU_master=mpi_info.compsep.QU_master, + root=mpi_info.compsep.master, + ) sky_model = SkyModel(comp_list) sky_model_at_band = sky_model.get_sky_at_nu(detector_data.nu, detector_data.nside, "IQU", @@ -224,8 +167,8 @@ def process_compsep(mpi_info: Bunch, detector_data: DetectorMap, iter: int, chai compsep_comm.Barrier() - if compsep_comm.Get_rank() == 0: - write_compsep_chain_to_file(comp_list, params, chain, iter) + if compsep_rank == mpi_info.compsep.master: + write_compsep_chain_to_file(comp_list.joined(), params, chain, iter) return sky_model # Return the full sky realization for my band. diff --git a/src/commander4/data_models/detector_TOD.py b/src/commander4/data_models/detector_TOD.py index 37423a2..2c168e3 100644 --- a/src/commander4/data_models/detector_TOD.py +++ b/src/commander4/data_models/detector_TOD.py @@ -46,7 +46,7 @@ def __init__( flag_encoded: NDArray[np.integer] | bytes | np.void | None = None, bad_data_bitmask: int | None = None, init_scalars: NDArray | None = None, - tod_is_compressed: bool = True, + tod_is_compressed: bool = False, flag_is_compressed: bool = True, det_response: NDArray | None = None, ): @@ -225,7 +225,9 @@ def processing_mask(self) -> NDArray[np.bool_]: Stored internally as a packed bit array and unpacked on each access. """ + start_bench("numpy-unpack") mask = np.unpackbits(self._processing_mask).view(bool) + stop_bench("numpy-unpack") if mask.size > self.tod.size + 7 or mask.size < self.tod.size: # The bytearray is stored in multiples of 8, so it can be up to 7 elements # longer than the TOD. If it's even longer or shorter, something is wrong. @@ -236,7 +238,9 @@ def processing_mask(self) -> NDArray[np.bool_]: @property def full_mask(self) -> NDArray[np.bool_]: """Boolean mask keeping samples that pass both flag and processing cuts.""" + start_bench("numpy-unpack") mask = np.unpackbits(self._full_mask).view(bool) + stop_bench("numpy-unpack") if mask.size > self.tod.size + 7 or mask.size < self.tod.size: raise ValueError(f"Mask size {mask.size} doesn't match TOD size {self.tod.size}.") return mask[:self.tod.size] @@ -244,7 +248,9 @@ def full_mask(self) -> NDArray[np.bool_]: @property def good_data_mask(self) -> NDArray[np.bool_]: """Boolean mask keeping samples that pass the bad-data flag cut.""" + start_bench("numpy-unpack") mask = np.unpackbits(self._good_data_mask).view(bool) + stop_bench("numpy-unpack") if mask.size > self.tod.size + 7 or mask.size < self.tod.size: raise ValueError(f"Mask size {mask.size} doesn't match TOD size {self.tod.size}.") return mask[:self.tod.size] diff --git a/src/commander4/experiments/SO/tod_reader_SO_SAT.py b/src/commander4/experiments/SO/tod_reader_SO_SAT.py index 48f9b9b..9216e3a 100644 --- a/src/commander4/experiments/SO/tod_reader_SO_SAT.py +++ b/src/commander4/experiments/SO/tod_reader_SO_SAT.py @@ -148,7 +148,7 @@ def tod_reader(band_comm: MPI.Comm, my_experiment: str, my_band: Bunch, det_name huffman_tree2=huffman_tree2, huffman_symbols2=huffman_symbols2, flag_encoded=flag_encoded, - flag_bitmask=my_experiment.flag_bitmask, + bad_data_bitmask=my_experiment.bad_data_bitmask, tod_is_compressed=my_experiment.tod_is_compressed, det_response=det_response) if np.sum(detector.full_mask) == 0 or (detector.tod == 0).all(): diff --git a/src/commander4/experiments/planck/tod_reader_planck.py b/src/commander4/experiments/planck/tod_reader_planck.py index a6ad2f0..789d74e 100644 --- a/src/commander4/experiments/planck/tod_reader_planck.py +++ b/src/commander4/experiments/planck/tod_reader_planck.py @@ -13,6 +13,9 @@ from commander4.data_models.scan_TOD import ScanTOD from commander4.data_models.detector_group_TOD import DetGroupTOD from commander4.noise_sampling.noise_psd import NoisePSD, NoisePSDOof +from commander4.data_models.pointing import PixelPointing +from commander4.logging.performance_logger import benchmark, bench_summary, start_bench,\ + stop_bench, log_memory, increment_count, bench_reset logger = logging.getLogger(__name__) @@ -37,8 +40,9 @@ def find_good_Fourier_time(Fourier_times:NDArray, ntod:int) -> int: return best_ntod -def tod_reader(band_comm: MPI.Comm, my_experiment: Bunch, my_band: Bunch, det_names: list[str], - params: Bunch, scan_idx_start: int, scan_idx_stop: int) -> DetectorTOD: +def tod_reader(band_comm: MPI.Comm, my_experiment: str, my_band: Bunch, all_det_names: list[str], + params: Bunch, scan_idx_start: int, + scan_idx_stop: int) -> DetGroupTOD: oids = [] pids = [] filepaths = [] @@ -79,6 +83,8 @@ def tod_reader(band_comm: MPI.Comm, my_experiment: Bunch, my_band: Bunch, det_na num_included = 0 ntod_sum_original = 0 ntod_sum_final = 0 + ndet = len(all_det_names) + det_init_scalars = np.zeros((ndet, 4)) + np.nan # Small de-sycnronization sleep. # time.sleep(5.0 * (band_comm.Get_rank() / band_comm.Get_size())) for i_pid in range(scan_idx_start, scan_idx_stop): @@ -100,12 +106,15 @@ def tod_reader(band_comm: MPI.Comm, my_experiment: Bunch, my_band: Bunch, det_na fsamp = float(f["/common/fsamp/"][()].item()) npsi = int(f["/common/npsi/"][()].item()) detector_list = [] - for det_name in det_names: + idet_accepted = 0 + for idet, det_name in enumerate(all_det_names): tod = f[f"/{pid}/{det_name}/tod/"][:ntod_optimal].astype(np.float32, copy=False) pix_encoded = f[f"/{pid}/{det_name}/pix/"][()] psi_encoded = f[f"/{pid}/{det_name}/psi/"][()] if "QU" in my_band.polarization else [] flag_encoded = f[f"/{pid}/{det_name}/flag/"][()] init_scalars = f[f"/{pid}/{det_name}/scalars/"][()] + # Data format has this weird thing were gain seems to be in "micro-gain"... + init_scalars[0] *= 1e-6 flag_buffer[:ntod] = cpp_utils.huffman_decode(np.frombuffer(flag_encoded, dtype=np.uint8), huffman_tree, huffman_symbols, flag_buffer[:ntod]) @@ -117,16 +126,21 @@ def tod_reader(band_comm: MPI.Comm, my_experiment: Bunch, my_band: Bunch, det_na # Check for crazy data. if np.mean(np.abs(tod)) > 0.001 or np.std(tod) > 0.001: good_scan = False - detector = DetectorTOD(tod, pix_encoded, psi_encoded, my_band.eval_nside, - data_nside, fsamp, vsun, huffman_tree, huffman_symbols, - npsi, processing_mask_map, ntod, - init_scalars = init_scalars, - pix_is_compressed=my_experiment.pix_is_compressed, - psi_is_compressed=my_experiment.psi_is_compressed \ - if "QU" in my_band.polarization else False) + + det_init_scalars[idet] = init_scalars + det_pointing = PixelPointing(pix_encoded, psi_encoded, huffman_tree, + huffman_symbols, npsi, my_band.eval_nside, data_nside, + ntod, ntod_optimal) + + detector = DetectorTOD(det_name, idet, idet_accepted, tod, det_pointing, fsamp, + vsun, huffman_tree, huffman_symbols, processing_mask_map, + ntod, ntod_optimal, + bad_data_bitmask = 6111232, + init_scalars = init_scalars) detector_list.append(detector) ntod_sum_original += ntod ntod_sum_final += ntod_optimal + idet_accepted += 1 if good_scan: scan = ScanTOD(detector_list, 0., scanID) scan_list.append(scan) @@ -136,7 +150,6 @@ def tod_reader(band_comm: MPI.Comm, my_experiment: Bunch, my_band: Bunch, det_na f"{i_pid-scan_idx_start}/{nscans}") if i_pid % 10 == 0: gc.collect() - ndet = len(det_names) # Number of detectors *should* be the same for all scans. noise_model = NoisePSDOof() diff --git a/src/commander4/mpi_management.py b/src/commander4/mpi_management.py index cd9a077..693fb6b 100644 --- a/src/commander4/mpi_management.py +++ b/src/commander4/mpi_management.py @@ -90,7 +90,6 @@ def init_mpi(params): # Testing revealed 24 to be a good number (regardless of nside), but I tested this on the # new 384-core nodes, the optimal number is probably slightly lower on the older owls. my_num_threads_numba = min(24,my_num_threads) - my_num_threads_numba = min(24,my_num_threads) else: raise ValueError("My rank ({worldrank}) exceeds the combined number of allocated tasks to" f"both TOD ({global_params.MPI_config.ntask_tod}) and compsep " \ @@ -108,7 +107,6 @@ def init_mpi(params): # resulted in some weirdeties, like many duplicate open file handles even when x=1. if False: # This code should enter production, but threadpoolctl is not yet a dependency. - import numpy import numba from threadpoolctl import threadpool_info @@ -208,6 +206,11 @@ def init_mpi_tod(mpi_info, params): # handling all communication with CompSep. TOD_rank = 0 current_detector_id = 0 # A unique number identifying every detector of every band. + my_band_name = None + my_band_id = None + my_experiment_name = None + my_det_id = None + my_detector_name = None for exp_name in params.experiments: experiment = params.experiments[exp_name] if not experiment.enabled: @@ -224,14 +227,10 @@ def init_mpi_tod(mpi_info, params): my_band_id = iband my_experiment_name = exp_name for idet, det_name in enumerate(band.detectors): - num_ranks_this_detector = len(TOD_ranks_per_detector[idet]) - detector = band.detectors[det_name] # Check if our rank belongs to this detector if MPIrank_tod in TOD_ranks_per_detector[idet]: # What is my rank number among the ranks processing this detector? my_det_id = idet - # Setting our unique detector id. Note that this is a global, not per band. - my_detector_id = current_detector_id my_detector_name = det_name current_detector_id += 1 # Update detector counter. else: @@ -245,6 +244,14 @@ def init_mpi_tod(mpi_info, params): if is_tod_master: logger.info(f"TOD: {MPIsize_tod} tasks successfully allocated to TOD proc.") + if my_band_id is None or my_det_id is None or my_experiment_name is None: + log.lograise( + RuntimeError, + f"TOD rank {MPIrank_tod} was not assigned to an enabled band/detector. " + "Check experiment enable flags and the TOD rank allocation in the parameter file.", + logger, + ) + # Create communicators for each different band. band_comm = mpi_info.tod.comm.Split(my_band_id, key=MPIrank_tod) # Get my local rank, and the total size of, the band-communicator IvsQU'm on. @@ -301,7 +308,7 @@ def init_mpi_compsep(mpi_info, params): for the 'compsep' context. """ - MPIsize_compsep, MPIrank_compsep = mpi_info.compsep.size, mpi_info.compsep.rank + MPIsize_compsep = mpi_info.compsep.size ### Setting up info for each band, including where to get the data from ### ###(map from file, or receive from TOD processing) ### diff --git a/src/commander4/output/write_chains_files.py b/src/commander4/output/write_chains_files.py index 64776d8..3274e55 100644 --- a/src/commander4/output/write_chains_files.py +++ b/src/commander4/output/write_chains_files.py @@ -3,11 +3,9 @@ import numpy as np import healpy as hp import datetime -from mpi4py import MPI from pixell.bunch import Bunch -from commander4.data_models.TOD_samples import TODSamples -from commander4.sky_models.component import Component +from commander4.sky_models.component import Component, CompList def write_map_chain_to_file(params: Bunch, chain: int, iter: int, exp_name:str, @@ -32,13 +30,23 @@ def write_map_chain_to_file(params: Bunch, chain: int, iter: int, exp_name:str, file[key] = value -def write_compsep_chain_to_file(comp_list: list[Component], params: Bunch, chain: int, iter: int): +def write_compsep_chain_to_file(comp_list: list[Component] | CompList, params: Bunch, + chain: int, iter: int): chain_dir = os.path.join(params.general.output_paths.chains, "compsep") chain_file = os.path.join(chain_dir, f"chain{chain:02d}_iter{iter:04d}.h5") + components = comp_list.components if isinstance(comp_list, CompList) else comp_list with h5py.File(chain_file, "w") as file: file["metadata/datetime"] = datetime.datetime.now().isoformat() file["metadata/parameter_file_as_string"] = params.parameter_file_as_string - for comp in comp_list: + seen_shortnames = set() + for comp in components: + if comp.shortname in seen_shortnames: + raise ValueError(f"Duplicate component shortname '{comp.shortname}' in compsep chain.") + seen_shortnames.add(comp.shortname) file[f"comps/{comp.shortname}/alms"] = comp.alms file[f"comps/{comp.shortname}/longname"] = comp.longname - file[f"comps/{comp.shortname}/shortname"] = comp.shortname \ No newline at end of file + file[f"comps/{comp.shortname}/shortname"] = comp.shortname + if comp.defined_pol is not None: + file[f"comps/{comp.shortname}/defined_pol"] = comp.defined_pol + if comp.eval_pol is not None: + file[f"comps/{comp.shortname}/eval_pol"] = comp.eval_pol \ No newline at end of file diff --git a/src/commander4/sky_models/component.py b/src/commander4/sky_models/component.py index f814512..202acd6 100644 --- a/src/commander4/sky_models/component.py +++ b/src/commander4/sky_models/component.py @@ -3,6 +3,7 @@ import numpy as np import pysm3.units as pysm3u import healpy as hp +import inspect import logging from copy import deepcopy from scipy.interpolate import interp1d @@ -42,100 +43,133 @@ def __init__(self, comp_params: Bunch, global_params: Bunch): self.global_params = global_params self.longname = comp_params.longname if "longname" in comp_params else "Unknown Component" self.shortname = comp_params.shortname if "shortname" in comp_params else "comp" + self.defined_pol = comp_params.polarization if "polarization" in comp_params else None + self.eval_pol = self.defined_pol self.double_prec = False if global_params.CG_float_precision == "single" else True self._data = None - def __add__(self, other): + @property + def logical_id(self) -> str: + return self.longname + + @property + def logical_key(self) -> tuple[type["Component"], str]: + return (type(self), self.logical_id) + + @property + def execution_key(self) -> tuple[type["Component"], str, str | None]: + return (type(self), self.logical_id, self.eval_pol) + + @property + def is_split_view(self) -> bool: + return self.defined_pol == "IQU" and self.eval_pol in ("I", "QU") + + @property + def execution_label(self) -> str: + if self.eval_pol is None or not self.is_split_view: + return self.shortname + return f"{self.shortname}[{self.eval_pol}]" + + def _assert_consistent_comp(self, other: "Component") -> None: + if not isinstance(other, Component): + raise TypeError("Both operands must be Component objects.") if type(self) is not type(other): raise TypeError("Both operands must be of the same Component type.") + mismatched = [ + attr for attr in ( + "longname", + "shortname", + "defined_pol", + "eval_pol", + ) + if getattr(self, attr) != getattr(other, attr) + ] + if mismatched: + raise ValueError( + "Components must represent the same execution view. " + f"Mismatched fields: {', '.join(mismatched)}" + ) if self._data is None or other._data is None: - raise ValueError("Cannot add Components with no data.") + raise ValueError("Cannot operate on Components with no data.") if self._data.shape != other._data.shape: raise ValueError("Data arrays of the two Components must match in size.") + + def join_split_views(self, other: "Component") -> "Component": + if not isinstance(other, Component): + raise TypeError("Can only join Component objects.") + if type(self) is not type(other): + raise TypeError("Split views must be of the same Component type.") + if self.defined_pol != "IQU" or other.defined_pol != "IQU": + raise ValueError("Only IQU-defined components can be joined.") + if not self.is_split_view or not other.is_split_view: + raise ValueError("Only split component views can be joined.") + if {self.eval_pol, other.eval_pol} != {"I", "QU"}: + raise ValueError("Joining requires one intensity view and one QU view.") + mismatched = [ + attr for attr in ("longname", "shortname", "defined_pol") + if getattr(self, attr) != getattr(other, attr) + ] + if mismatched: + raise ValueError( + "Split views must refer to the same logical component. " + f"Mismatched fields: {', '.join(mismatched)}" + ) + if self._data is None or other._data is None: + raise ValueError("Cannot join split views with no data.") + intensity_comp, pol_comp = (self, other) if self.eval_pol == "I" else (other, self) + if intensity_comp._data.shape[1:] != pol_comp._data.shape[1:]: + raise ValueError("Split views must have compatible alm dimensions.") + joined = deepcopy(intensity_comp) + joined.eval_pol = joined.defined_pol + joined._data = np.concatenate((intensity_comp._data, pol_comp._data), axis=0) + return joined + + def __add__(self, other): + self._assert_consistent_comp(other) out = deepcopy(self) inplace_arr_add(out._data, other._data) return out def __iadd__(self, other): - if type(self) is not type(other): - raise TypeError("Both operands must be of the same Component type.") - if self._data is None or other._data is None: - raise ValueError("Cannot add Components with no data.") - if self._data.shape != other._data.shape: - raise ValueError("Data arrays of the two Components must match in size.") + self._assert_consistent_comp(other) inplace_arr_add(self._data, other._data) return self def __sub__(self, other): - if type(self) is not type(other): - raise TypeError("Both operands must be of the same Component type.") - if self._data is None or other._data is None: - raise ValueError("Cannot add Components with no data.") - if self._data.shape != other._data.shape: - raise ValueError("Data arrays of the two Components must match in size.") + self._assert_consistent_comp(other) out = deepcopy(self) inplace_arr_sub(out._data, other._data) return out def __isub__(self, other): - if type(self) is not type(other): - raise TypeError("Both operands must be of the same Component type.") - if self._data is None or other._data is None: - raise ValueError("Cannot add Components with no data.") - if self._data.shape != other._data.shape: - raise ValueError("Data arrays of the two Components must match in size.") + self._assert_consistent_comp(other) inplace_arr_sub(self._data, other._data) return self def __mul__(self, other): - if type(self) is not type(other): - raise TypeError("Both operands must be of the same Component type.") - if self._data is None or other._data is None: - raise ValueError("Cannot add Components with no data.") - if self._data.shape != other._data.shape: - raise ValueError("Data arrays of the two Components must match in size.") + self._assert_consistent_comp(other) out = deepcopy(self) inplace_arr_prod(out._data, other._data) return out def __imul__(self, other): - if type(self) is not type(other): - raise TypeError("Both operands must be of the same Component type.") - if self._data is None or other._data is None: - raise ValueError("Cannot add Components with no data.") - if self._data.shape != other._data.shape: - raise ValueError("Data arrays of the two Components must match in size.") + self._assert_consistent_comp(other) inplace_arr_prod(self._data, other._data) return self def __truediv__(self, other): - if type(self) is not type(other): - raise TypeError("Both operands must be of the same Component type.") - if self._data is None or other._data is None: - raise ValueError("Cannot add Components with no data.") - if self._data.shape != other._data.shape: - raise ValueError("Data arrays of the two Components must match in size.") + self._assert_consistent_comp(other) out = deepcopy(self) inplace_arr_truediv(out._data, other._data) return out def __itruediv__(self, other): - if type(self) is not type(other): - raise TypeError("Both operands must be of the same Component type.") - if self._data is None or other._data is None: - raise ValueError("Cannot add Components with no data.") - if self._data.shape != other._data.shape: - raise ValueError("Data arrays of the two Components must match in size.") - inplace_arr_add(self._data, other._data) + self._assert_consistent_comp(other) + inplace_arr_truediv(self._data, other._data) return self def __matmul__(self, other): - if type(self) is not type(other): - raise TypeError("Both operands must be of the same Component type.") - if self._data is None or other._data is None: - raise ValueError("Cannot add Components with no data.") - if self._data.shape != other._data.shape: - raise ValueError("Data arrays of the two Components must match in size.") + self._assert_consistent_comp(other) return dot(self._data, other._data) def bcast_data_blocking(self, comm:MPI.Comm, root=0): @@ -223,7 +257,11 @@ def npol(self): @property def spin(self): - return 2 if self.is_pol else 0 + if self.eval_pol == "I": + return 0 + if self.eval_pol in ("QU", "IQU"): + return 2 + raise ValueError(f"Unsupported polarization '{self.eval_pol}'.") @property def is_pol(self): @@ -298,15 +336,24 @@ def apply_smoothing_prior_sqrt(self): almxfl(self._data[ipol], smooth_p_sqrt, inplace=True) return self._data + def _realize_alms_as_map(self, component_alms, nside: int, fwhm: float = 0): + """Realize component alms as a map. + + Joined `IQU` components still need separate intensity and spin-2 synthesis calls, since + DUCC does not accept a 3-row alm block in one call. + """ + component_alms = hp.smoothalm(component_alms, fwhm, inplace=False) + if self.eval_pol != "IQU": + return alm_to_map(component_alms, nside, self.lmax, spin=self.spin) + intensity_map = alm_to_map(component_alms[:1], nside, self.lmax, spin=0) + pol_map = alm_to_map(component_alms[1:], nside, self.lmax, spin=2) + return np.concatenate((intensity_map, pol_map), axis=0) + def get_component_map(self, nside:int, fwhm:int=0): component_alms = self.alms if component_alms is None: raise ValueError("component_alms property not set.") - if fwhm == 0: - return alm_to_map(component_alms, nside, self.lmax, spin = self.spin) - else: - return alm_to_map(hp.smoothalm(component_alms, fwhm, inplace=False), nside, self.lmax, - spin = self.spin) + return self._realize_alms_as_map(component_alms, nside, fwhm) def get_sky(self, nu, nside, fwhm=0): return self.get_component_map(nside, fwhm)*self.get_sed(nu) @@ -316,12 +363,7 @@ def get_sed(self, nu): #overwrite of the dot product as the diffuse component will have alm _data with complex encoding def __matmul__(self, other): - if type(self) is not type(other): - raise TypeError("Both operands must be of the same Component type.") - if self._data is None or other._data is None: - raise ValueError("Cannot add Components with no data.") - if self._data.shape != other._data.shape: - raise ValueError("Data arrays of the two Components must match in size.") + self._assert_consistent_comp(other) res = 0.0 for ipol in range(self.npol): res += _dot_complex_alm_1D_arrays(self._data[ipol], other._data[ipol], self.lmax) @@ -431,11 +473,7 @@ def get_sky_anisotropies(self, nu, nside, fwhm=0): # Zero out the quadrupole (l=2) for m in range(3): # m = 0, 1, 2 component_alms[:,hp.Alm.getidx(self.lmax, 2, m)] = 0.0 + 0.0j - if fwhm == 0: - return alm_to_map(component_alms, nside, self.lmax, spin = self.spin)*self.get_sed(nu) - else: - return alm_to_map(hp.smoothalm(component_alms, fwhm, inplace=False), nside, self.lmax, - spin = self.spin)*self.get_sed(nu) + return self._realize_alms_as_map(component_alms, nside, fwhm) * self.get_sed(nu) class CMBRelQuad(TemplateComponent): pass @@ -607,6 +645,8 @@ def __init__(self, comp_params: Bunch, global_params: Bunch): self.longname = comp_params.longname if "longname" in comp_params\ else "Unknown PointSourceComp" self.shortname = comp_params.shortname if "shortname" in comp_params else "pscomp" + self.defined_pol = "I" + self.eval_pol = "I" @property def is_pol(self) -> bool: @@ -757,8 +797,6 @@ def _eval_from_band_map(self, map, nu): """ mJysr_to_uKRJ = (pysm3u.mJy / pysm3u.steradian).to(pysm3u.uK_RJ, equivalencies=pysm3u.cmb_equivalencies(nu*pysm3u.GHz)) - uKRJ_to_mJysr = (pysm3u.uK_RJ).to(pysm3u.mJy / pysm3u.steradian, - equivalencies=pysm3u.cmb_equivalencies(nu*pysm3u.GHz)) sed_s = self.get_sed(nu) _numba_eval_from_map(map[0,:], self.pix_disc_idx_list, self.beam_disc_val_list, self._data[0,:], sed_s = sed_s) @@ -836,9 +874,9 @@ def split_complist(comp_list: list[Component], color:int, if color not in IvsQU_colors: logging.warning(f"Color {color} not in colors assigned to I or QU ({IvsQU_colors})!") else: + target_pol = "I" if color == IvsQU_colors[0] else "QU" for comp in comp_list: - if comp.is_pol == (color == IvsQU_colors[1]): - # print("Comp", comp.shortname, comp.pol) + if comp.eval_pol == target_pol: out_comp_list.append(comp) return out_comp_list @@ -846,8 +884,82 @@ def split_complist(comp_list: list[Component], color:int, class CompList: def __init__(self, comp_list:list[Component]): + self._validate_comp_list(comp_list) self.comp_list = comp_list + @staticmethod + def _validate_comp_list(comp_list: list[Component]) -> None: + """Check that a component list has a coherent logical and execution-view layout.""" + if not isinstance(comp_list, list): + raise TypeError("comp_list must be a list of Component objects.") + + grouped_by_longname = {} + shortname_to_longname = {} + for idx, comp in enumerate(comp_list): + if not isinstance(comp, Component): + raise TypeError(f"comp_list[{idx}] must be a Component.") + if comp.defined_pol is not None: + assert_pol_supported(comp.defined_pol) + if comp.eval_pol is not None: + assert_pol_supported(comp.eval_pol) + + prev_longname = shortname_to_longname.get(comp.shortname) + if prev_longname is not None and prev_longname != comp.longname: + raise ValueError( + f"Shortname {comp.shortname!r} is used for both {prev_longname!r} and " + f"{comp.longname!r}." + ) + shortname_to_longname[comp.shortname] = comp.longname + grouped_by_longname.setdefault(comp.longname, []).append(comp) + + for longname, group in grouped_by_longname.items(): + component_types = {type(comp) for comp in group} + if len(component_types) > 1: + raise ValueError(f"Longname {longname!r} is shared across multiple component classes.") + + shortnames = {comp.shortname for comp in group} + if len(shortnames) > 1: + raise ValueError( + f"Longname {longname!r} is associated with multiple shortnames: " + f"{sorted(shortnames)!r}." + ) + + split_views = [comp for comp in group if comp.is_split_view] + unsplit_views = [comp for comp in group if not comp.is_split_view] + if split_views and unsplit_views: + raise ValueError( + f"Longname {longname!r} mixes split and unsplit execution views." + ) + if len(unsplit_views) > 1: + raise ValueError(f"Duplicate logical component {longname!r}.") + if len(split_views) > 2: + raise ValueError(f"Component {longname!r} has too many split execution views.") + split_pols = [comp.eval_pol for comp in split_views] + if len(set(split_pols)) != len(split_pols): + raise ValueError(f"Component {longname!r} repeats a split execution view.") + + @staticmethod + def _instantiate_component(component: Bunch, global_params: Bunch, + eval_pol: str | None = None) -> Component: + """Instantiate one execution-view component from a parameter-file component entry. + + Diffuse IQU components opt into splitting via the `eval_pol` constructor argument, while + other component classes can ignore that detail entirely by omitting the parameter. + """ + component_cls = getattr(component_lib, component.component_class) + init_params = inspect.signature(component_cls.__init__).parameters + kwargs = {} + if "allocate_empty_alms" in init_params: + kwargs["allocate_empty_alms"] = True + if eval_pol is not None: + log.logassert( + "eval_pol" in init_params, + f"Component class '{component.component_class}' does not support polarization splitting.", + logger, + ) + kwargs["eval_pol"] = eval_pol + return component_cls(component.params, global_params, **kwargs) + @classmethod def init_from_params(cls, components:Bunch, params:Bunch): # Determine whether any of the bands actually have polarization by checking whether any @@ -857,67 +969,152 @@ def init_from_params(cls, components:Bunch, params:Bunch): for component_str in components: component = components[component_str] if component.enabled: - if component.params.lmax == "full": + if "lmax" in component.params and component.params.lmax == "full": component.params.lmax = (params.general.nside*5)//2 - if component.params.polarization == "I": #I-only - # 'getattr' loads the class specified by "component_class" from model.component. - # This class is then instantiated with the "params" specified, and appended to - # the components list. - comp_list.append(getattr(component_lib, component.component_class)(component.params, - params.general, allocate_empty_alms=True)) - elif component.params.polarization == "QU": #QU-only + component_pol = component.params.polarization if "polarization" in component.params else "I" + if component_pol == "I": + comp_list.append(cls._instantiate_component(component, params.general)) + elif component_pol == "QU": if not pol_bands_exist: logging.warning(f"Component '{component_str}' is specified as QU-only but " f"ntask_compsep_QU=0 (no polarized bands). Skipping.") continue - comp_list.append(getattr(component_lib, component.component_class)(component.params, - params.general, allocate_empty_alms=True)) - elif component.params.polarization == "IQU": - # I - comp_list.append(getattr(component_lib, component.component_class)( - component.params, - params.general, - allocate_empty_alms=True, - longname = component.params.longname+"_Instensity", - shortname = component.params.longname+"_I", - eval_pol="I")) - # QU + comp_list.append(cls._instantiate_component(component, params.general)) + elif component_pol == "IQU": + intensity_comp = cls._instantiate_component(component, params.general, eval_pol="I") + comp_list.append(intensity_comp) if not pol_bands_exist: logging.warning(f"Component '{component_str}' is specified as IQU but " f"ntask_compsep_QU=0 (no polarized bands). " f"Only the intensity (I) part will be used.") else: - comp_list.append(getattr(component_lib, component.component_class)( - component.params, - params.general, - allocate_empty_alms=True, - longname = component.params.longname+"_Polarization", - shortname = component.params.longname+"_QU", - eval_pol="QU")) + pol_comp = cls._instantiate_component(component, params.general, + eval_pol="QU") + comp_list.append(pol_comp) else: - raise ValueError(f"Unrecognized polarization in parameter file for component {component_str}") + raise ValueError( + f"Unrecognized polarization in parameter file for component {component_str}" + ) return cls(comp_list) + def _assert_consistent_comps(self, other: "CompList") -> None: + if not isinstance(other, CompList): + raise TypeError("Both operands must be CompList objects.") + if len(self.comp_list) != len(other.comp_list): + raise ValueError("Component lists must match in length.") + self_keys = [comp.execution_key for comp in self.comp_list] + other_keys = [comp.execution_key for comp in other.comp_list] + if self_keys != other_keys: + raise ValueError("Component lists must contain the same execution views in the same order.") + + def components_for_eval_pol(self, target_pol: str) -> list[Component]: + assert_pol_supported(target_pol) + return [comp for comp in self.comp_list if comp.eval_pol == target_pol] + + def split_for_eval_pol(self, target_pol: str) -> "CompList": + """Return the execution-view subset evaluated for one polarization stream.""" + return CompList(self.components_for_eval_pol(target_pol)) + + def copy_matching_data_from(self, other: "CompList") -> None: + if not isinstance(other, CompList): + raise TypeError("Input must be a CompList.") + other_by_key = {} + for comp in other.comp_list: + if comp.execution_key in other_by_key: + raise ValueError(f"Duplicate component execution key {comp.execution_key!r}.") + other_by_key[comp.execution_key] = comp + self_keys = {comp.execution_key for comp in self.comp_list} + extra_keys = [key for key in other_by_key if key not in self_keys] + if extra_keys: + raise ValueError(f"Found unknown components in source CompList: {extra_keys!r}") + for comp in self.comp_list: + other_comp = other_by_key.get(comp.execution_key) + if other_comp is None: + continue + comp._assert_consistent_comp(other_comp) + np.copyto(comp._data, other_comp._data) + + def reassemble_from_split_solution(self, local_solution: "CompList", comm: MPI.Comm, + *, is_I_master: bool, is_QU_master: bool, + I_master: int, QU_master: int, root: int = 0) -> None: + """Collect split I/QU solver results back into the full execution list on the CompSep root. + + The intensity master already owns the local I solution and receives the QU views from the + QU master. Once the full execution list is assembled on the root, every rank gets a copy + through a blocking broadcast so `SkyModel` realization can happen locally. + """ + if is_I_master: + self.copy_matching_data_from(local_solution) + + pol_components = self.components_for_eval_pol("QU") + if pol_components: + if is_QU_master: + local_pol_by_key = {comp.execution_key: comp for comp in local_solution.comp_list} + for tag, comp in enumerate(pol_components): + local_comp = local_pol_by_key.get(comp.execution_key) + if local_comp is None: + raise ValueError( + f"Missing QU component {comp.execution_label} on QU master." + ) + comm.Send(local_comp._data, dest=I_master, tag=tag) + if is_I_master: + for tag, comp in enumerate(pol_components): + comm.Recv(comp._data, source=QU_master, tag=tag) + + for comp in self.comp_list: + comp.bcast_data_blocking(comm, root=root) + + def joined(self) -> "CompList": + """Collapse split execution views back to one logical component per `longname`.""" + grouped_components = {} + logical_order = [] + for comp in self.comp_list: + if comp.logical_key not in grouped_components: + grouped_components[comp.logical_key] = [] + logical_order.append(comp.logical_key) + grouped_components[comp.logical_key].append(comp) + + joined_components = [] + for logical_key in logical_order: + group = grouped_components[logical_key] + split_views = [comp for comp in group if comp.is_split_view] + unsplit_views = [comp for comp in group if not comp.is_split_view] + if unsplit_views and split_views: + raise ValueError( + f"Logical component {logical_key[1]!r} mixes split and unsplit execution views." + ) + if len(unsplit_views) > 1: + raise ValueError(f"Duplicate unsplit component {logical_key[1]!r}.") + if unsplit_views: + joined_components.append(deepcopy(unsplit_views[0])) + continue + if len(split_views) == 1: + joined_components.append(deepcopy(split_views[0])) + continue + if len(split_views) != 2: + raise ValueError( + f"Expected one or two execution views for {logical_key[1]!r}, got {len(group)}." + ) + joined_components.append(split_views[0].join_split_views(split_views[1])) + + return CompList(joined_components) + def split(self, color:int, IvsQU_colors:tuple = (0,1)): """ Extracts from `comp_list` only the components containing the correct Stokes parameter based on the passed `color` of the local MPI rank. By default, color=0 will treat Intensity and color=1 QU. A list with the relevant components is returned. """ - out_comp_list = [] IvsQU_colors = IvsQU_colors[:2] #cut off eventual elements in excess if color not in IvsQU_colors: logging.warning(f"Color {color} not in colors assigned to I or QU ({IvsQU_colors})!") + return CompList([]) elif color == IvsQU_colors[0]: target_pol = "I" elif color == IvsQU_colors[1]: target_pol = "QU" - for comp in self.comp_list: - if comp.eval_pol == target_pol: - out_comp_list.append(comp) - - return CompList(out_comp_list) + return self.split_for_eval_pol(target_pol) @property def components(self): @@ -932,69 +1129,60 @@ def __matmul__(self, other) -> float: components with alms. It will automatically handle the correct dot product definition for each type of Component. """ - if len(self.comp_list) != len(other): - raise ValueError("Component lists must match in length.") + self._assert_consistent_comps(other) res = 0.0 for c1, c2 in zip(self.components, other.components): res += float(c1 @ c2) return res def __add__(self, other): - if len(self.comp_list) != len(other): - raise ValueError("Component lists must match in length.") + self._assert_consistent_comps(other) out = deepcopy(self) - for o, c in zip(out.components, other.components): - o += c - return o + for out_comp, other_comp in zip(out.components, other.components): + out_comp += other_comp + return out def __iadd__(self, other): - if len(self.comp_list) != len(other): - raise ValueError("Component lists must match in length.") + self._assert_consistent_comps(other) for c1, c2 in zip(self.components, other.components): c1 += c2 return self def __sub__(self, other): - if len(self.comp_list) != len(other): - raise ValueError("Component lists must match in length.") + self._assert_consistent_comps(other) out = deepcopy(self) - for o, c in zip(out.components, other.components): - o -= c - return o + for out_comp, other_comp in zip(out.components, other.components): + out_comp -= other_comp + return out def __isub__(self, other): - if len(self.comp_list) != len(other): - raise ValueError("Component lists must match in length.") + self._assert_consistent_comps(other) for c1, c2 in zip(self.components, other.components): c1 -= c2 return self def __mul__(self, other): - if len(self.comp_list) != len(other): - raise ValueError("Component lists must match in length.") + self._assert_consistent_comps(other) out = deepcopy(self) - for o, c in zip(out.components, other.components): - o *= c - return o + for out_comp, other_comp in zip(out.components, other.components): + out_comp *= other_comp + return out def __imul__(self, other): - if len(self.comp_list) != len(other): - raise ValueError("Component lists must match in length.") + self._assert_consistent_comps(other) for c1, c2 in zip(self.components, other.components): c1 *= c2 return self def __truediv__(self, other): - if len(self.comp_list) != len(other): - raise ValueError("Component lists must match in length.") + self._assert_consistent_comps(other) out = deepcopy(self) - for o, c in zip(out.components, other.components): - o /= c - return o + for out_comp, other_comp in zip(out.components, other.components): + out_comp /= other_comp + return out def __itruediv__(self, other): - if len(self.comp_list) != len(other): - raise ValueError("Component lists must match in length.") + self._assert_consistent_comps(other) for c1, c2 in zip(self.components, other.components): c1 /= c2 return self @@ -1069,8 +1257,7 @@ def apply_smoothing_prior_sqrt(self): def inplace_add_scaled(self, list_other, scalar): """ `list_inplace += scalar*list_other` """ - if len(self) != len(list_other): - raise ValueError("Component lists must match in length.") + self._assert_consistent_comps(list_other) for ci, co in zip(self.comp_list, list_other.comp_list): inplace_add_scaled_vec(ci._data, co._data, scalar) @@ -1078,8 +1265,7 @@ def inplace_add_scaled(self, list_other, scalar): def inplace_scale_and_add(self, list_other, scalar): """ `list_inplace = scalar*list_inplace + list_other` """ - if len(self) != len(list_other): - raise ValueError("Component lists must match in length.") + self._assert_consistent_comps(list_other) for ci, co in zip(self.comp_list, list_other.comp_list): inplace_scale_add(ci._data, co._data, scalar) \ No newline at end of file diff --git a/src/commander4/sky_models/sky_model.py b/src/commander4/sky_models/sky_model.py index 23cd15a..cffc32a 100644 --- a/src/commander4/sky_models/sky_model.py +++ b/src/commander4/sky_models/sky_model.py @@ -13,29 +13,41 @@ def get_sky(self, band): raise NotImplementedError def get_sky_at_nu(self, nu, nside, pols_required, fwhm=None): - """ Get sky at specific frequency. + """Get the realized sky at one frequency. + + The component list may be either the split execution list used during CompSep (`I` and + `QU` views) or a joined logical list containing `IQU` components. """ npix = 12*nside**2 if pols_required == "I": skymap = np.zeros((1, npix), dtype=np.float32) - for component in self._components: - if not component.is_pol: - skymap[0] += component.get_sky(nu, nside, fwhm)[0] elif pols_required == "QU": skymap = np.zeros((2, npix), dtype=np.float32) - for component in self._components: - if component.is_pol: - skymap[0:] += component.get_sky(nu, nside, fwhm) elif pols_required == "IQU": skymap = np.zeros((3, npix), dtype=np.float32) - for component in self._components: - if component.is_pol: - skymap[1:] += component.get_sky(nu, nside, fwhm) - else: - skymap[0] += component.get_sky(nu, nside, fwhm)[0] else: raise ValueError("Unrecognized polarization string") + + for component in self._components: + if component.eval_pol == "I": + if pols_required in ("I", "IQU"): + skymap[0] += component.get_sky(nu, nside, fwhm)[0] + elif component.eval_pol == "QU": + if pols_required == "QU": + skymap += component.get_sky(nu, nside, fwhm) + elif pols_required == "IQU": + skymap[1:] += component.get_sky(nu, nside, fwhm) + elif component.eval_pol == "IQU": + component_sky = component.get_sky(nu, nside, fwhm) + if pols_required == "I": + skymap[0] += component_sky[0] + elif pols_required == "QU": + skymap += component_sky[1:] + else: + skymap += component_sky + else: + raise ValueError(f"Unsupported component polarization '{component.eval_pol}'.") return skymap # class SkyModel: diff --git a/src/commander4/solvers/perpix_compsep_solver.py b/src/commander4/solvers/perpix_compsep_solver.py index 9c5f149..67e72b1 100644 --- a/src/commander4/solvers/perpix_compsep_solver.py +++ b/src/commander4/solvers/perpix_compsep_solver.py @@ -148,7 +148,7 @@ def solve_compsep_perpix(proc_comm: MPI.Comm, detector_data: DetectorMap, ctypes_lib.solve_compsep(npix, nband, ncomp, maps_sky.astype(np.float64, copy=False), maps_rms.astype(np.float64, copy=False), M, rand, comp_maps[ipol]) logger.info(f"Finished pixel-by-pixel component separation in {time.time()-t0:.2f}s "\ - f"for polarization {ipol+1} of 3.") + f"for polarization {ipol+1} of {npol}.") comp_maps = proc_comm.bcast(comp_maps, root=0) for icomp in range(ncomp): diff --git a/src/commander4/tod_processing.py b/src/commander4/tod_processing.py index b7684b7..ad3744e 100644 --- a/src/commander4/tod_processing.py +++ b/src/commander4/tod_processing.py @@ -24,6 +24,7 @@ from commander4.noise_sampling.noise_sampling import sample_noise_PS_params, fill_all_masked from commander4.noise_sampling.sample_ncorr import corr_noise_realization_with_gaps from commander4.utils.math_operations import forward_rfft, backward_rfft +from commander4.utils.execution_ids import get_execution_band_ids from commander4.noise_sampling.sigma0 import calc_sigma0_robust from commander4.tod_reader import read_tods_from_file from commander4.output.write_chains_files import write_map_chain_to_file @@ -421,10 +422,14 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu # In the CG solver this is only used to define the starting guess, # but if the CG fails it is also used to generate the fallback solution. fill_all_masked(sky_subtracted_TOD, mask_full, sigma0_ncorr) - n_corr_est, residual, niter, did_conv = corr_noise_realization_with_gaps(sky_subtracted_TOD, - mask_full, sigma0_ncorr, C_1f_inv, - err_tol=1e-4, max_iter=20) - resid = (sky_subtracted_TOD - n_corr_est) * mask_full + # n_corr_est, residual, niter, did_conv = corr_noise_realization_with_gaps(sky_subtracted_TOD, + # mask_full, sigma0_ncorr, C_1f_inv, + # err_tol=1e-4, max_iter=1) + residual = np.inf + resid = np.inf + niter = 0 + did_conv = False + # resid = (sky_subtracted_TOD - n_corr_est) * mask_full var_resid = np.dot(resid, resid) var_data = np.dot(sky_subtracted_TOD * mask_full, sky_subtracted_TOD * mask_full) # If either of the two tests failed, use fallback for n_corr. @@ -584,7 +589,6 @@ def init_tod_processing(mpi_info: Bunch, params: Bunch) -> tuple[Bunch, str, Det # handling all communication with CompSep. # All the non-master ranks will have None values, and receive info from master further down. det_names = [] - my_experiment_name = None my_band_name = None my_experiment = None my_band = None @@ -607,7 +611,6 @@ def init_tod_processing(mpi_info: Bunch, params: Bunch) -> tuple[Bunch, str, Det my_band_pol = band.polarization my_band_id = iband # What is my rank number among the ranks processing this detector? - my_experiment_name = exp_name my_experiment = experiment # Setting our unique detector id. Note that this is a global, not per band. tot_num_scans = experiment.num_scans @@ -643,31 +646,24 @@ def init_tod_processing(mpi_info: Bunch, params: Bunch) -> tuple[Bunch, str, Det # Creating "tod_band_masters", an array which maps the band index to the rank of the master # of that band. - todproc_my_band_id = f"{my_experiment_name}$$${my_band_name}" - data_world = (todproc_my_band_id, mpi_info.world.rank) if mpi_info.band.is_master else None + todproc_my_band_id = my_band_name + data_world = ( + todproc_my_band_id, + mpi_info.world.rank, + my_band_pol, + ) if mpi_info.band.is_master else None data_tod = (todproc_my_band_id, mpi_info.tod.rank) if mpi_info.band.is_master else None - # pols_tod_bands = (todproc_my_band_id, my_band_pol) if mpi_info.band.is_master else None all_data_world = mpi_info.tod.comm.allgather(data_world) all_data_tod = mpi_info.tod.comm.allgather(data_tod) - # all_pol_data = mpi_info.tod.comm.allgather(pols_tod_bands) - - world_band_masters_dict = {} - if "I" in my_band_pol: - world_band_masters_dict.update({item[0]+'_I': # First I: - item[1] for item in all_data_world if item is not None}) - if "QU" in my_band_pol: - world_band_masters_dict.update({item[0]+'_QU': # Then QU: - item[1] for item in all_data_world if item is not None}) + + world_band_masters_dict = { + execution_band_id: item[1] + for item in all_data_world if item is not None + for execution_band_id in get_execution_band_ids(item[0], item[2]) + } tod_band_masters_dict = {item[0]: item[1] for item in all_data_tod if item is not None} - # tod_band_pol_dict = {item[0]: item[1] for item in all_pol_data if item is not None} - # logger.info(f"world_band_masters_dict: {world_band_masters_dict}") - # logger.info(f"tod_band_masters_dict: {tod_band_masters_dict}") - # logger.info(f"tod_band_pol_dict: {tod_band_pol_dict}") - # logger.info(f"TOD: Rank {mpi_info.tod.rank:4} assigned scans {my_scans_start:6} - "\ - # f"{my_scans_stop:6} on band {my_band_id:4}.") mpi_info['world']['tod_band_masters'] = world_band_masters_dict mpi_info['tod']['tod_band_masters'] = tod_band_masters_dict - # mpi_info['world']['tod_band_pols'] = tod_band_pol_dict return mpi_info, todproc_my_band_id, experiment_data, tod_samples_chain1, tod_samples_chain2 @@ -1123,7 +1119,7 @@ def process_tod(mpi_info: Bunch, experiment_data: DetGroupTOD, band_comm = mpi_info.band.comm TOD_comm = mpi_info.tod.comm ### JUMP DETECTION ### - if getattr(params.general, "sample_jump_detection", True) and iter >= int( + if getattr(params.general, "sample_jump_detection", False) and iter >= int( getattr(params.general, "sample_jump_detection_from_iter_num", 1) ): t0 = time.time() diff --git a/src/commander4/utils/CG_mapmaker.py b/src/commander4/utils/CG_mapmaker.py index 31faa74..96f1c76 100644 --- a/src/commander4/utils/CG_mapmaker.py +++ b/src/commander4/utils/CG_mapmaker.py @@ -438,7 +438,7 @@ def apply_P_adjoint(self, in_scan: ScanTOD, out_map:NDArray, pix=None, psi=None, assert npix_out == hp.nside2npix(in_scan.nside), "out_map size must match scan's eval nside." assert pix.shape == scan_tod_arr.shape, "pix shape must match scan_tod_arr." pix = in_scan.pix if pix is None else pix - ntod = in_scan.tod.shape[-1] + ntod = scan_tod_arr.shape[-1] self.map_accumulator(out_map, scan_tod_arr, 1, pix.astype(np.int64, copy=False), ntod) return out_map @@ -555,7 +555,7 @@ def apply_P_adjoint(self, in_scan: ScanTOD, out_map:NDArray, pix=None, psi=None, # assert psi.shape == scan_tod_arr.shape, "psi shape must match scan_tod_arr." pix = in_scan.pix if pix is None else pix psi = in_scan.psi if psi is None else psi - ntod = in_scan.tod.shape[-1] + ntod = scan_tod_arr.shape[-1] self.map_accumulator_IQU(out_map, scan_tod_arr, 1, pix.astype(np.int64, copy=False), psi.astype(np.float64, copy=False), ntod, npix_out) return out_map diff --git a/src/commander4/utils/math_operations.py b/src/commander4/utils/math_operations.py index ea20319..d0f28d7 100644 --- a/src/commander4/utils/math_operations.py +++ b/src/commander4/utils/math_operations.py @@ -230,8 +230,7 @@ def complist_dot(comp_list1:CompList, comp_list2:CompList) -> float: components with alms. It will automatically handle the correct dot product definition for each type of Component. """ - if len(comp_list1) != len(comp_list2): - raise ValueError("Component lists must match in length.") + comp_list1._assert_consistent_comps(comp_list2) if len(comp_list1) == 0: print("WARNING dot prod between empty comp list") res = 0.0 From 28b3dfc901238bd6e2c3a6ecc4f300619119be7b Mon Sep 17 00:00:00 2001 From: jgslunde Date: Wed, 3 Jun 2026 17:05:49 +0200 Subject: [PATCH 15/23] Parameter file dumped to chain is no longer re-ordered to be alphabetical. --- src/commander4/parse_params.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/commander4/parse_params.py b/src/commander4/parse_params.py index f775693..c1b3ea2 100644 --- a/src/commander4/parse_params.py +++ b/src/commander4/parse_params.py @@ -134,7 +134,7 @@ def as_bunch_recursive(dict_of_dicts, name=None): # For reproducability, create custom entries in the parameter object which holds the entire # parameter file as a fully resolved YAML string (with any !inc directives expanded). -params.parameter_file_as_string = yaml.dump(params_dict) +params.parameter_file_as_string = yaml.dump(params_dict, sort_keys=False) # Storing Commander4 version number or git commit. # params.metadata.version_number = # print(get_version_info("commander4", __file__)) \ No newline at end of file From be705db4efcb296382d0095cb879934964dea1cb Mon Sep 17 00:00:00 2001 From: jgslunde Date: Wed, 3 Jun 2026 23:55:26 +0200 Subject: [PATCH 16/23] Implemented compsep sampling groups, and refactored a bunch of other stuff --- src/commander4/compsep_processing.py | 285 ++++++---- src/commander4/output/write_chains_files.py | 2 +- src/commander4/sky_models/component.py | 515 +++++++++--------- src/commander4/sky_models/sky_model.py | 52 +- src/commander4/solvers/CG_compsep_solver.py | 35 +- .../solvers/perpix_compsep_solver.py | 9 +- 6 files changed, 441 insertions(+), 457 deletions(-) diff --git a/src/commander4/compsep_processing.py b/src/commander4/compsep_processing.py index 5282099..76333f3 100644 --- a/src/commander4/compsep_processing.py +++ b/src/commander4/compsep_processing.py @@ -1,6 +1,7 @@ import numpy as np import logging from copy import deepcopy +from mpi4py import MPI from pixell.bunch import Bunch from commander4.output.log import logassert @@ -10,93 +11,139 @@ from commander4.solvers.CG_compsep_solver import CompSepSolver from commander4.solvers.perpix_compsep_solver import solve_compsep_perpix from commander4.output.write_chains_files import write_compsep_chain_to_file -from commander4.utils.execution_ids import get_execution_band_id +from commander4.utils.execution_ids import get_execution_band_id, EXECUTION_POLS logger = logging.getLogger(__name__) + +def _sampling_group_selects_band(selected_bands: list[str] | None, band_name: str, + band_identifier: str) -> bool: + """Whether a sampling group acts on a band, matched by base name or execution-view identifier. + + `selected_bands` of None means "all bands". + """ + if selected_bands is None: + return True + return band_name in selected_bands or band_identifier in selected_bands + + +def _filter_sampling_group_components(comp_list: CompList, + selected_components: list[str] | None) -> CompList: + """Subset of `comp_list` whose component names are selected by the sampling group. + + `selected_components` of None means "all components". The returned `CompList` shares the + underlying `Component` objects with `comp_list` (it is a view, not a copy). + """ + if selected_components is None: + return CompList(list(comp_list)) + selected_names = set(selected_components) + return CompList([comp for comp in comp_list if comp.comp_name in selected_names]) + + +def _validate_sampling_groups(sampling_groups: Bunch, comp_list: CompList, params: Bunch) -> None: + """Fail fast if any enabled sampling group references a non-existent component or band. + + `comps` and `bands` are expected to be lists of strings naming existing components and bands + (bands may be given either as a base name or as an execution-view identifier). + """ + known_comp_names = {comp.comp_name for comp in comp_list.joined()} + known_band_names = set() + for band_str in params.CompSep_bands: + band = params.CompSep_bands[band_str] + if not band.enabled: + continue + known_band_names.add(band_str) + for eval_pol in EXECUTION_POLS[band.polarization]: + known_band_names.add(get_execution_band_id(band_str, eval_pol)) + + for group_name in sampling_groups: + group = sampling_groups[group_name] + if "enabled" in group and not group.enabled: + continue + if "comps" in group: + unknown = sorted(set(group.comps) - known_comp_names) + logassert(not unknown, + f"Sampling group {group_name!r} references unknown component(s) {unknown}. " + f"Known components: {sorted(known_comp_names)}.", logger) + if "bands" in group: + unknown = sorted(set(group.bands) - known_band_names) + logassert(not unknown, + f"Sampling group {group_name!r} references unknown band(s) {unknown}. " + f"Known bands: {sorted(known_band_names)}.", logger) + + def init_compsep_processing(mpi_info: Bunch, params: Bunch)\ - -> tuple[CompList, str, dict[str, int], Bunch]: + -> tuple[CompList, Bunch, str, Bunch]: """Set up the rank-local execution view for component separation. - Each CompSep rank owns exactly one execution view of one band. The global CompSep rank space - is split into a contiguous intensity block followed by a contiguous QU block, and we match the - current rank against those two logical streams here. + Each CompSep rank owns exactly one execution view of one band. The global CompSep rank space is + split into a contiguous intensity block (ranks ``[0, QU_master)``) followed by a contiguous QU + block (ranks ``[QU_master, size)``), and we match the current rank against those two streams. Args: mpi_info (Bunch): The data structure containing all MPI relevant data. params (Bunch): The parameters from the input parameter file. Returns: - mpi_info (Bunch): The data structure containing all MPI relevant data, - modified to contain also the band masters dictionary. + comp_list (CompList): The full execution-view component list (identical on all CompSep ranks). + mpi_info (Bunch): `mpi_info`, extended with this rank's band name/identifier and the + band-master dictionaries. band_identifier (str): Unique string for the band execution view this rank is working on. - my_band (Bunch): A subset of the full parameter file for the band this rank is working on. + my_band (Bunch): The parameter-file subset for the band this rank is working on. """ logger.info(f"CompSep: Hello from CompSep-rank {mpi_info.compsep.rank} (on machine "\ f"{mpi_info.processor_name}), dedicated to band {mpi_info.compsep.rank}.") comp_list = CompList.init_from_params(params.components, params) + comp_names = [comp.comp_name for comp in comp_list.joined()] + logassert(len(comp_names) == len(set(comp_names)), + f"Duplicate component names found in CompSep setup: {comp_names}", logger) - ### Setting up info for each band, including where to get the data from ### - ### (map from file, or receive from TOD processing) ### - - current_band_idx_I = 0 - current_band_idx_QU = mpi_info.compsep.QU_master + ### Match this rank to its band execution view. Intensity views fill the I-rank block in band + ### order, QU views fill the QU-rank block; the two cursors track those contiguous layouts. ### + band_cursor = {"I": 0, "QU": mpi_info.compsep.QU_master} band_identifier = None + band_name = None my_band = None - for band_str in params.CompSep_bands: # Intensity + for band_str in params.CompSep_bands: band = params.CompSep_bands[band_str] - if band.enabled: - if band.polarization == "I": - if current_band_idx_I == mpi_info.compsep.rank: - my_band = deepcopy(band) - band_identifier = get_execution_band_id(band_str, "I") - logger.info(f"Rank {mpi_info.compsep.rank} just matched band {band_identifier}") - my_band.identifier = band_identifier - my_band.polarization = "I" - current_band_idx_I += 1 - - elif band.polarization == "QU": - if current_band_idx_QU == mpi_info.compsep.rank: - my_band = deepcopy(band) - band_identifier = get_execution_band_id(band_str, "QU") - logger.info(f"Rank {mpi_info.compsep.rank} matched band {band_identifier}") - my_band.identifier = band_identifier - my_band.polarization = "QU" - current_band_idx_QU += 1 - - elif band.polarization == "IQU": - # Each IQU band occupies one rank in the intensity block and one in the QU block. - if current_band_idx_I == mpi_info.compsep.rank: - my_band = deepcopy(band) - band_identifier = get_execution_band_id(band_str, "I") - logger.info(f"Rank {mpi_info.compsep.rank} just matched band {band_identifier}") - my_band.identifier = band_identifier - my_band.polarization = "I" - current_band_idx_I += 1 - if current_band_idx_QU == mpi_info.compsep.rank: - my_band = deepcopy(band) - band_identifier = get_execution_band_id(band_str, "QU") - logger.info(f"Rank {mpi_info.compsep.rank} matched band {band_identifier}") - my_band.identifier = band_identifier - my_band.polarization = "QU" - current_band_idx_QU += 1 - else: - raise ValueError(f"Unrecognized polarization in parameter file for band {band_str}") - - #sanity check: - logassert(current_band_idx_I == mpi_info.compsep.QU_master, "Number of acquired Intensity "\ - f"bands ({current_band_idx_I}) do not match number of MPI tasks assigned to "\ - f"Intensity ({mpi_info.compsep.QU_master})", logger) - logassert(current_band_idx_QU == mpi_info.compsep.size, "Number of acquired QU bands "\ - f"({current_band_idx_QU}) do not match number of MPI tasks assigned to QU "\ - f"({mpi_info.compsep.QU_master})", logger) + if not band.enabled: + continue + if band.polarization not in EXECUTION_POLS: + raise ValueError(f"Unrecognized polarization in parameter file for band {band_str}") + for eval_pol in EXECUTION_POLS[band.polarization]: + if band_cursor[eval_pol] == mpi_info.compsep.rank: + my_band = deepcopy(band) + band_name = band_str + band_identifier = get_execution_band_id(band_str, eval_pol) + my_band.identifier = band_identifier + my_band.polarization = eval_pol + logger.info(f"Rank {mpi_info.compsep.rank} matched band {band_identifier}") + band_cursor[eval_pol] += 1 + + # Sanity checks: the I cursor must have consumed exactly the I-rank block [0, QU_master), and + # the QU cursor exactly the QU-rank block [QU_master, size). + n_I_ranks = mpi_info.compsep.QU_master + n_QU_ranks = mpi_info.compsep.size - mpi_info.compsep.QU_master + logassert(band_cursor["I"] == mpi_info.compsep.QU_master, + f"Number of enabled Intensity band views ({band_cursor['I']}) does not match the " + f"number of CompSep ranks assigned to Intensity ({n_I_ranks}).", logger) + logassert(band_cursor["QU"] == mpi_info.compsep.size, + f"Number of enabled QU band views ({band_cursor['QU'] - mpi_info.compsep.QU_master}) " + f"does not match the number of CompSep ranks assigned to QU ({n_QU_ranks}).", logger) if my_band is None or band_identifier is None: logassert(False, f"CompSep rank {mpi_info.compsep.rank} was not assigned to any enabled band. " "Check that CompSep_bands matches the configured I/QU rank counts.", logger) - + + sampling_groups = params.sampling_groups_compsep if "sampling_groups_compsep" in params \ + else Bunch() + _validate_sampling_groups(sampling_groups, comp_list, params) + + mpi_info.compsep.band_name = band_name + mpi_info.compsep.band_identifier = band_identifier + data_world = (band_identifier, mpi_info.world.rank) data_compsep = (band_identifier, mpi_info.compsep.rank) all_data_world = mpi_info.compsep.comm.allgather(data_world) @@ -111,64 +158,100 @@ def init_compsep_processing(mpi_info: Bunch, params: Bunch)\ def process_compsep(mpi_info: Bunch, detector_data: DetectorMap, iter: int, chain: int, params: Bunch, comp_list: CompList) -> SkyModel: - """ Performs a single component separation iteration. - Called by each compsep process, which are each responsible for a single band. - + """Perform a single component-separation iteration. + + Called by every CompSep rank, each of which owns one band execution view. Loops over the + configured sampling groups; for each group the participating ranks (those whose band and at + least one of whose components are selected) form a solver sub-communicator, run the requested + sampler, and then broadcast the updated components so that `comp_list` is identical on every + CompSep rank again before the next group. + Args: mpi_info (Bunch): The data structure containing all MPI relevant data. - detector_data (DetectorMap): The detector map for this MPI rank's band, cleaned of all - "TOD" components (correlated noise and orbital dipole). + detector_data (DetectorMap): The detector map for this rank's band, cleaned of all "TOD" + components (correlated noise and orbital dipole). iter (int): The current Gibbs iteration (used only for printing and seeding). chain (int): The current chain (used only for printing and seeding). params (Bunch): The parameters from the input parameter file. + comp_list (CompList): The full execution-view component list, updated in place. Returns: - detector_maps (np.array): The band-integrated total sky. + sky_model (SkyModel): The full sky realization, wrapping the updated `comp_list`. """ - - ### 1. MPI SETUP: Split into I and QU ### compsep_comm = mpi_info.compsep.comm compsep_rank = mpi_info.compsep.rank - subcolor = mpi_info.compsep.subcolor #Subcolor splits the compsep ranks into: Pol -> 1, Int -> 0 + # subcomm splits the CompSep ranks by polarization (subcolor 0 -> I, 1 -> QU). compsep_subcomm = mpi_info.compsep.subcomm - target_pol = "I" if subcolor == 0 else "QU" - comp_sublist = comp_list.split_for_eval_pol(target_pol) - - ### 2. SOLVE COMPSEP: band maps -> component alms (either by per-pixel or CG solver) ### - if params.general.pixel_compsep_sampling: - comp_sublist = solve_compsep_perpix(compsep_subcomm, detector_data, comp_sublist, params) - else: - compsep_solver = CompSepSolver(detector_data, params, compsep_subcomm) - - comp_sublist = compsep_solver.solve(comp_sublist) - - ### 3. CLEANUP: Gather I+QU alm solutions and make plots. ### - comp_list.reassemble_from_split_solution( - comp_sublist, - compsep_comm, - is_I_master=mpi_info.compsep.is_I_master, - is_QU_master=mpi_info.compsep.is_QU_master, - I_master=mpi_info.compsep.I_master, - QU_master=mpi_info.compsep.QU_master, - root=mpi_info.compsep.master, - ) + target_pol = "I" if mpi_info.compsep.subcolor == 0 else "QU" + sampling_groups = params.sampling_groups_compsep if "sampling_groups_compsep" in params \ + else Bunch() + # SkyModel wraps `comp_list` by reference, so a single instance reflects all in-place updates + # made below, and is what we ultimately return. sky_model = SkyModel(comp_list) - sky_model_at_band = sky_model.get_sky_at_nu(detector_data.nu, detector_data.nside, "IQU", - fwhm=np.deg2rad(detector_data.fwhm/60.0)) - pol_names = ["Q", "U"] if detector_data.pol else ["I"] - pol_offset = 1 if detector_data.pol else 0 - for ipol in range(detector_data.npol): - chi2 = np.mean(np.abs(detector_data.map_sky[ipol] - - sky_model_at_band[ipol + pol_offset])/detector_data.map_rms[ipol]) - logger.info(f"Reduced chi2 on rank {compsep_rank} for pol={pol_names[ipol]} "\ - f"({detector_data.nu}GHz): {chi2:.3f}") + for sampling_group_name in sampling_groups: + sampling_group = sampling_groups[sampling_group_name] + if "enabled" in sampling_group and not sampling_group.enabled: + continue + + sampled_components = sampling_group.comps if "comps" in sampling_group else None + sampled_bands = sampling_group.bands if "bands" in sampling_group else None + band_is_active = _sampling_group_selects_band(sampled_bands, mpi_info.compsep.band_name, + mpi_info.compsep.band_identifier) + # This rank's components (for its own polarization stream) that take part in this group. + # `active_sublist` shares Component objects with `comp_list`, so copying the solver result + # into it updates `comp_list` on this rank; the broadcast below propagates it to all ranks. + active_sublist = _filter_sampling_group_components( + comp_list.split_for_eval_pol(target_pol), sampled_components) + should_solve = band_is_active and len(active_sublist) > 0 + + # The solving ranks of each polarization form their own solver communicator. + solver_comm = compsep_subcomm.Split(0 if should_solve else MPI.UNDEFINED, key=compsep_rank) + if should_solve: + sample_class = sampling_group.sample_class + if sample_class == "amplitude_sampler_perpix": + solved_sublist = solve_compsep_perpix(solver_comm, detector_data, active_sublist, + params) + elif sample_class == "amplitude_sampler_CG": + solved_sublist = CompSepSolver(detector_data, params, solver_comm).solve( + active_sublist) + else: + raise ValueError( + f"Unknown compsep sampling class {sample_class!r} for sampling group " + f"{sampling_group_name!r}.") + active_sublist.copy_matching_data_from(solved_sublist) + solver_comm.Free() - compsep_comm.Barrier() + any_active = compsep_comm.allreduce(1 if should_solve else 0, op=MPI.SUM) + if not any_active: + if compsep_rank == mpi_info.compsep.master: + logger.info( + f"Sampling group {sampling_group_name!r} had no active band/component overlap.") + continue + + # Restore global consistency: for each polarization that was solved this group, the + # lowest-ranked solver (which holds the authoritative result) broadcasts its component views + # to all ranks. A polarization that no rank solved is already identical everywhere. + for eval_pol in ("I", "QU"): + solved_here = should_solve and target_pol == eval_pol + source = compsep_comm.allreduce(compsep_rank if solved_here else compsep_comm.size, + op=MPI.MIN) + if source < compsep_comm.size: + comp_list.broadcast_pol_views(compsep_comm, eval_pol=eval_pol, source=source) + + # Print new per-band chi2s against the updated sky model. + sky_model_at_band = sky_model.get_sky_at_nu(detector_data.nu, detector_data.nside, "IQU", + fwhm=np.deg2rad(detector_data.fwhm/60.0)) + pol_names = ["Q", "U"] if detector_data.pol else ["I"] + pol_offset = 1 if detector_data.pol else 0 + for ipol in range(detector_data.npol): + chi2 = np.mean(np.abs(detector_data.map_sky[ipol] - + sky_model_at_band[ipol + pol_offset])/detector_data.map_rms[ipol]) + logger.info(f"Reduced chi2 on rank {compsep_rank} for pol={pol_names[ipol]} "\ + f"({detector_data.nu}GHz): {chi2:.3f}") if compsep_rank == mpi_info.compsep.master: write_compsep_chain_to_file(comp_list.joined(), params, chain, iter) return sky_model # Return the full sky realization for my band. - diff --git a/src/commander4/output/write_chains_files.py b/src/commander4/output/write_chains_files.py index 3274e55..632928f 100644 --- a/src/commander4/output/write_chains_files.py +++ b/src/commander4/output/write_chains_files.py @@ -44,7 +44,7 @@ def write_compsep_chain_to_file(comp_list: list[Component] | CompList, params: B raise ValueError(f"Duplicate component shortname '{comp.shortname}' in compsep chain.") seen_shortnames.add(comp.shortname) file[f"comps/{comp.shortname}/alms"] = comp.alms - file[f"comps/{comp.shortname}/longname"] = comp.longname + file[f"comps/{comp.shortname}/comp_name"] = comp.comp_name file[f"comps/{comp.shortname}/shortname"] = comp.shortname if comp.defined_pol is not None: file[f"comps/{comp.shortname}/defined_pol"] = comp.defined_pol diff --git a/src/commander4/sky_models/component.py b/src/commander4/sky_models/component.py index 202acd6..66b4e39 100644 --- a/src/commander4/sky_models/component.py +++ b/src/commander4/sky_models/component.py @@ -3,7 +3,6 @@ import numpy as np import pysm3.units as pysm3u import healpy as hp -import inspect import logging from copy import deepcopy from scipy.interpolate import interp1d @@ -19,6 +18,7 @@ _dot_complex_alm_1D_arrays, _numba_proj2map, _numba_eval_from_map, inplace_scale_add from commander4.utils.map_utils import gauss_beam, get_gauss_beam_radius, get_npol, assert_pol_supported from commander4.data_models.band import Band +from commander4.utils.execution_ids import EXECUTION_POLS logger = logging.getLogger(__name__) @@ -38,19 +38,53 @@ def g(nu): # First tier component classes class Component: - def __init__(self, comp_params: Bunch, global_params: Bunch): + default_shortname = "comp" + legal_pols: tuple[str, ...] = ("I", "QU", "IQU") + requires_defined_pol = False + + @classmethod + def _assert_legal_pol(cls, pol: str | None, *, role: str, required: bool = False) -> None: + if pol is None: + log.logassert( + not required, + f"{cls.__name__} requires a defined polarization mode.", + logger, + ) + return + assert_pol_supported(pol) + log.logassert( + pol in cls.legal_pols, + f"{cls.__name__} does not support {role} polarization {pol!r}. " + f"Allowed polarizations: {cls.legal_pols!r}.", + logger, + ) + + def __init__(self, comp_params: Bunch, global_params: Bunch, *, + shortname: str | None = None, comp_name: str | None = None, + eval_pol: str | None = None, allocate_empty_alms: bool = False): self.comp_params = comp_params self.global_params = global_params - self.longname = comp_params.longname if "longname" in comp_params else "Unknown Component" - self.shortname = comp_params.shortname if "shortname" in comp_params else "comp" + self.shortname = ( + shortname + if shortname is not None + else comp_params.shortname if "shortname" in comp_params + else self.default_shortname + ) + self.comp_name = comp_params._name if comp_name is None else comp_name self.defined_pol = comp_params.polarization if "polarization" in comp_params else None - self.eval_pol = self.defined_pol + type(self)._assert_legal_pol( + self.defined_pol, + role="defined", + required=type(self).requires_defined_pol, + ) + self.eval_pol = self.defined_pol if eval_pol is None else eval_pol + type(self)._assert_legal_pol(self.eval_pol, role="evaluation") self.double_prec = False if global_params.CG_float_precision == "single" else True self._data = None @property def logical_id(self) -> str: - return self.longname + return self.comp_name @property def logical_key(self) -> tuple[type["Component"], str]: @@ -77,7 +111,7 @@ def _assert_consistent_comp(self, other: "Component") -> None: raise TypeError("Both operands must be of the same Component type.") mismatched = [ attr for attr in ( - "longname", + "comp_name", "shortname", "defined_pol", "eval_pol", @@ -106,7 +140,7 @@ def join_split_views(self, other: "Component") -> "Component": if {self.eval_pol, other.eval_pol} != {"I", "QU"}: raise ValueError("Joining requires one intensity view and one QU view.") mismatched = [ - attr for attr in ("longname", "shortname", "defined_pol") + attr for attr in ("comp_name", "shortname", "defined_pol") if getattr(self, attr) != getattr(other, attr) ] if mismatched: @@ -124,49 +158,35 @@ def join_split_views(self, other: "Component") -> "Component": joined._data = np.concatenate((intensity_comp._data, pol_comp._data), axis=0) return joined - def __add__(self, other): + def _apply_array_op(self, other: "Component", arr_op, *, inplace: bool) -> "Component": self._assert_consistent_comp(other) - out = deepcopy(self) - inplace_arr_add(out._data, other._data) - return out + target = self if inplace else deepcopy(self) + arr_op(target._data, other._data) + return target + + def __add__(self, other): + return self._apply_array_op(other, inplace_arr_add, inplace=False) def __iadd__(self, other): - self._assert_consistent_comp(other) - inplace_arr_add(self._data, other._data) - return self + return self._apply_array_op(other, inplace_arr_add, inplace=True) def __sub__(self, other): - self._assert_consistent_comp(other) - out = deepcopy(self) - inplace_arr_sub(out._data, other._data) - return out + return self._apply_array_op(other, inplace_arr_sub, inplace=False) def __isub__(self, other): - self._assert_consistent_comp(other) - inplace_arr_sub(self._data, other._data) - return self + return self._apply_array_op(other, inplace_arr_sub, inplace=True) def __mul__(self, other): - self._assert_consistent_comp(other) - out = deepcopy(self) - inplace_arr_prod(out._data, other._data) - return out + return self._apply_array_op(other, inplace_arr_prod, inplace=False) def __imul__(self, other): - self._assert_consistent_comp(other) - inplace_arr_prod(self._data, other._data) - return self + return self._apply_array_op(other, inplace_arr_prod, inplace=True) def __truediv__(self, other): - self._assert_consistent_comp(other) - out = deepcopy(self) - inplace_arr_truediv(out._data, other._data) - return out + return self._apply_array_op(other, inplace_arr_truediv, inplace=False) def __itruediv__(self, other): - self._assert_consistent_comp(other) - inplace_arr_truediv(self._data, other._data) - return self + return self._apply_array_op(other, inplace_arr_truediv, inplace=True) def __matmul__(self, other): self._assert_consistent_comp(other) @@ -233,21 +253,24 @@ def _zeros_like(self, other, dtype=None, order='K', subok=True, shape=None): # Second tier component classes class DiffuseComponent(Component): + requires_defined_pol = True + def __init__(self, comp_params: Bunch, global_params: Bunch, - allocate_empty_alms=False, eval_pol:None|str=None): - super().__init__(comp_params, global_params) + allocate_empty_alms=False, eval_pol:None|str=None, + comp_name: str | None = None, shortname: str | None = None): + super().__init__( + comp_params, + global_params, + shortname=shortname, + comp_name=comp_name, + eval_pol=eval_pol, + allocate_empty_alms=allocate_empty_alms, + ) self.spatially_varying_MM = comp_params.spatially_varying_MM self.lmax = comp_params.lmax self.smoothing_prior_FWHM = comp_params.smoothing_prior_FWHM self.smoothing_prior_amplitude = comp_params.smoothing_prior_amplitude self._data = None # Alm data is not allocated by default. - assert_pol_supported(comp_params.polarization) - self.defined_pol = comp_params.polarization #polarization as defined on the parameter file - if eval_pol is not None: - assert_pol_supported(eval_pol) - self.eval_pol = eval_pol - else: - self.eval_pol = self.defined_pol #if not passed, eval_pol defaults to defined_pol if allocate_empty_alms: self.allocate_empty_alms() @@ -441,14 +464,18 @@ class TemplateComponent(Component): # Third tier component classes class CMB(DiffuseComponent): + default_shortname = "cmb" + def __init__(self, comp_params: Bunch, global_params: Bunch, allocate_empty_alms=False, - shortname = None, longname = None, eval_pol = None): - super().__init__(comp_params, global_params, allocate_empty_alms, eval_pol) - #this gives priority: 1) arg, 2) param and 3) default - self.longname = longname if longname is not None else \ - comp_params.longname if "longname" in comp_params else "CMB" - self.shortname = shortname if shortname is not None else \ - comp_params.shortname if "shortname" in comp_params else "cmb" + shortname = None, eval_pol = None, comp_name: str | None = None): + super().__init__( + comp_params, + global_params, + allocate_empty_alms=allocate_empty_alms, + eval_pol=eval_pol, + comp_name=comp_name, + shortname=shortname, + ) def get_sed(self, nu): """Calculates the spectral energy distribution (SED) for CMB emission. @@ -479,17 +506,22 @@ class CMBRelQuad(TemplateComponent): pass class ThermalDust(DiffuseComponent): + default_shortname = "term-dust" + def __init__(self, comp_params: Bunch, global_params: Bunch, allocate_empty_alms=False, - shortname = None, longname = None, eval_pol = None): - super().__init__(comp_params, global_params, allocate_empty_alms, eval_pol) + shortname = None, eval_pol = None, comp_name: str | None = None): + super().__init__( + comp_params, + global_params, + allocate_empty_alms=allocate_empty_alms, + eval_pol=eval_pol, + comp_name=comp_name, + shortname=shortname, + ) self.beta = comp_params.beta self.T = comp_params.T self.nu0 = comp_params.nu0 self.prior_l_power_law = 2.5 - self.longname = longname if longname is not None else \ - comp_params.longname if "longname" in comp_params else "Thermal Dust" - self.shortname = shortname if shortname is not None else \ - comp_params.shortname if "shortname" in comp_params else "term-dust" def get_sed(self, nu): """Calculates the spectral energy distribution (SED) for Thermal Dust emission. @@ -506,17 +538,22 @@ def get_sed(self, nu): class Synchrotron(DiffuseComponent): + default_shortname = "sync" + def __init__(self, comp_params: Bunch, global_params: Bunch, allocate_empty_alms=False, - shortname = None, longname = None, eval_pol = None): - super().__init__(comp_params, global_params, allocate_empty_alms, eval_pol) + shortname = None, eval_pol = None, comp_name: str | None = None): + super().__init__( + comp_params, + global_params, + allocate_empty_alms=allocate_empty_alms, + eval_pol=eval_pol, + comp_name=comp_name, + shortname=shortname, + ) self.beta = comp_params.beta self.nu0 = comp_params.nu0 self.nside_comp_map = 512 self.prior_l_power_law = -3 - self.longname = longname if longname is not None else \ - comp_params.longname if "longname" in comp_params else "Synchrotron" - self.shortname = shortname if shortname is not None else \ - comp_params.shortname if "shortname" in comp_params else "sync" def get_sed(self, nu): """Calculates the spectral energy distribution (SED) for Synchrotron emission. @@ -530,15 +567,20 @@ def get_sed(self, nu): class FreeFree(DiffuseComponent): + default_shortname = "ff" + def __init__(self, comp_params: Bunch, global_params: Bunch, allocate_empty_alms=False, - shortname = None, longname = None, eval_pol = None): - super().__init__(comp_params, global_params, allocate_empty_alms, eval_pol) + shortname = None, eval_pol = None, comp_name: str | None = None): + super().__init__( + comp_params, + global_params, + allocate_empty_alms=allocate_empty_alms, + eval_pol=eval_pol, + comp_name=comp_name, + shortname=shortname, + ) self.T = comp_params.T # Electron temperature in K self.nu0 = comp_params.nu0 # Reference frequency in GHz - self.longname = longname if longname is not None else \ - comp_params.longname if "longname" in comp_params else "Free-Free" - self.shortname = shortname if shortname is not None else \ - comp_params.shortname if "shortname" in comp_params else "ff" def _gaunt_factor(self, nu, T): """Calculates the Gaunt factor for free-free emission, as per Eq. 18 in BP1. @@ -569,6 +611,8 @@ def get_sed(self, nu): return sed class SpinningDust(DiffuseComponent): + default_shortname = "spin-dust" + """ Spinning Dust component spectral model, based on spinning dust. The SED is derived from the SpDust2 code template for the Cold Neutral Medium. @@ -578,14 +622,21 @@ class SpinningDust(DiffuseComponent): # Columns: Frequency (GHz), Emissivity (proportional to Intensity) def __init__(self, comp_params: Bunch, global_params: Bunch, allocate_empty_alms=False, - shortname = None, longname = None, eval_pol = None): + shortname = None, eval_pol = None, comp_name: str | None = None): """ Args: nu_peak (float): The peak frequency of the spinning dust component in GHz. nu_0 (float): The reference frequency of the spinning dust template in GHz. This will not impact the shape of the SED, just the absolute scaling. """ - super().__init__(comp_params, global_params, allocate_empty_alms, eval_pol) + super().__init__( + comp_params, + global_params, + allocate_empty_alms=allocate_empty_alms, + eval_pol=eval_pol, + comp_name=comp_name, + shortname=shortname, + ) # Read SpDust2 template data. This is a simulation of what the spectral shape of # spinning dust emission should look like if it happens to peak at 30 GHz. @@ -593,10 +644,6 @@ def __init__(self, comp_params: Bunch, global_params: Bunch, allocate_empty_alms self.nu_peak_ref = 30.0 # The reference peak frequency of 30 GHz. self.nu_peak_eval = comp_params.nu_peak self.nu_0 = comp_params.nu_0 # Reference frequency for the amplitude map in GHz - self.longname = longname if longname is not None else \ - comp_params.longname if "longname" in comp_params else "Spinning Dust" - self.shortname = shortname if shortname is not None else \ - comp_params.shortname if "shortname" in comp_params else "spin-dust" # Create an logarithmic interpolation function from the SpDust2 template log_nu = np.log(freqs) @@ -640,11 +687,20 @@ def get_sed(self, nu: float|NDArray[np.floating]): # NON DIFFUSE COMPONENTS class PointSourcesComponent(Component): - def __init__(self, comp_params: Bunch, global_params: Bunch): - super().__init__(comp_params, global_params) - self.longname = comp_params.longname if "longname" in comp_params\ - else "Unknown PointSourceComp" - self.shortname = comp_params.shortname if "shortname" in comp_params else "pscomp" + default_shortname = "pscomp" + legal_pols: tuple[str, ...] = ("I",) + + def __init__(self, comp_params: Bunch, global_params: Bunch, *, + shortname: str | None = None, comp_name: str | None = None, + eval_pol: str | None = None, allocate_empty_alms: bool = False): + super().__init__( + comp_params, + global_params, + shortname=shortname, + comp_name=comp_name, + eval_pol="I" if eval_pol is None else eval_pol, + allocate_empty_alms=allocate_empty_alms, + ) self.defined_pol = "I" self.eval_pol = "I" @@ -657,10 +713,19 @@ def npol(self) -> int: return 1 class RadioSources(PointSourcesComponent): - def __init__(self, comp_params: Bunch, global_params: Bunch): - super().__init__(comp_params, global_params) - self.longname = comp_params.longname if "longname" in comp_params else "RadioPointSources" - self.shortname = comp_params.shortname if "shortname" in comp_params else "radsources" + default_shortname = "radsources" + + def __init__(self, comp_params: Bunch, global_params: Bunch, *, + shortname: str | None = None, comp_name: str | None = None, + eval_pol: str | None = None, allocate_empty_alms: bool = False): + super().__init__( + comp_params, + global_params, + shortname=shortname, + comp_name=comp_name, + eval_pol=eval_pol, + allocate_empty_alms=allocate_empty_alms, + ) #reference frequency self.nu0 = comp_params.nu_0 #tabulated data @@ -860,41 +925,39 @@ def apply_smoothing_prior_sqrt(self): def __repr__(self): return f"Radio Source \n amps: {self._data}" - -#FIXME: this will go within ComponentList object when implemented -def split_complist(comp_list: list[Component], color:int, - IvsQU_colors:tuple = (0,1)) -> list[Component]: - """ - Extracts from `comp_list` only the components containing the correct Stokes parameter based - on the passed `color` of the local MPI rank. By default, color=0 will treat Intensity and - color=1 polarization. A list with the relevant components is returned. - """ - out_comp_list = [] - IvsQU_colors = IvsQU_colors[:2] #cut off eventual elements in excess - if color not in IvsQU_colors: - logging.warning(f"Color {color} not in colors assigned to I or QU ({IvsQU_colors})!") - else: - target_pol = "I" if color == IvsQU_colors[0] else "QU" - for comp in comp_list: - if comp.eval_pol == target_pol: - out_comp_list.append(comp) - - return out_comp_list - - class CompList: def __init__(self, comp_list:list[Component]): self._validate_comp_list(comp_list) self.comp_list = comp_list + @staticmethod + def _group_by_logical_key( + comp_list: list[Component], + ) -> list[tuple[tuple[type["Component"], str], list[Component]]]: + grouped_components = {} + logical_order = [] + for comp in comp_list: + if comp.logical_key not in grouped_components: + grouped_components[comp.logical_key] = [] + logical_order.append(comp.logical_key) + grouped_components[comp.logical_key].append(comp) + return [(logical_key, grouped_components[logical_key]) for logical_key in logical_order] + + @staticmethod + def _partition_execution_views( + group: list[Component], + ) -> tuple[list[Component], list[Component]]: + split_views = [comp for comp in group if comp.is_split_view] + unsplit_views = [comp for comp in group if not comp.is_split_view] + return split_views, unsplit_views + @staticmethod def _validate_comp_list(comp_list: list[Component]) -> None: """Check that a component list has a coherent logical and execution-view layout.""" if not isinstance(comp_list, list): raise TypeError("comp_list must be a list of Component objects.") - grouped_by_longname = {} - shortname_to_longname = {} + shortname_to_comp_name = {} for idx, comp in enumerate(comp_list): if not isinstance(comp, Component): raise TypeError(f"comp_list[{idx}] must be a Component.") @@ -903,98 +966,80 @@ def _validate_comp_list(comp_list: list[Component]) -> None: if comp.eval_pol is not None: assert_pol_supported(comp.eval_pol) - prev_longname = shortname_to_longname.get(comp.shortname) - if prev_longname is not None and prev_longname != comp.longname: + prev_comp_name = shortname_to_comp_name.get(comp.shortname) + if prev_comp_name is not None and prev_comp_name != comp.comp_name: raise ValueError( - f"Shortname {comp.shortname!r} is used for both {prev_longname!r} and " - f"{comp.longname!r}." + f"Shortname {comp.shortname!r} is used for both {prev_comp_name!r} and " + f"{comp.comp_name!r}." ) - shortname_to_longname[comp.shortname] = comp.longname - grouped_by_longname.setdefault(comp.longname, []).append(comp) + shortname_to_comp_name[comp.shortname] = comp.comp_name - for longname, group in grouped_by_longname.items(): + for logical_key, group in CompList._group_by_logical_key(comp_list): + comp_name = logical_key[1] component_types = {type(comp) for comp in group} if len(component_types) > 1: - raise ValueError(f"Longname {longname!r} is shared across multiple component classes.") + raise ValueError( + f"Component name {comp_name!r} is shared across multiple component classes." + ) shortnames = {comp.shortname for comp in group} if len(shortnames) > 1: raise ValueError( - f"Longname {longname!r} is associated with multiple shortnames: " + f"Component name {comp_name!r} is associated with multiple shortnames: " f"{sorted(shortnames)!r}." ) - split_views = [comp for comp in group if comp.is_split_view] - unsplit_views = [comp for comp in group if not comp.is_split_view] + split_views, unsplit_views = CompList._partition_execution_views(group) if split_views and unsplit_views: raise ValueError( - f"Longname {longname!r} mixes split and unsplit execution views." + f"Component name {comp_name!r} mixes split and unsplit execution views." ) if len(unsplit_views) > 1: - raise ValueError(f"Duplicate logical component {longname!r}.") + raise ValueError(f"Duplicate logical component {comp_name!r}.") if len(split_views) > 2: - raise ValueError(f"Component {longname!r} has too many split execution views.") + raise ValueError(f"Component {comp_name!r} has too many split execution views.") split_pols = [comp.eval_pol for comp in split_views] if len(set(split_pols)) != len(split_pols): - raise ValueError(f"Component {longname!r} repeats a split execution view.") - - @staticmethod - def _instantiate_component(component: Bunch, global_params: Bunch, - eval_pol: str | None = None) -> Component: - """Instantiate one execution-view component from a parameter-file component entry. - - Diffuse IQU components opt into splitting via the `eval_pol` constructor argument, while - other component classes can ignore that detail entirely by omitting the parameter. - """ - component_cls = getattr(component_lib, component.component_class) - init_params = inspect.signature(component_cls.__init__).parameters - kwargs = {} - if "allocate_empty_alms" in init_params: - kwargs["allocate_empty_alms"] = True - if eval_pol is not None: - log.logassert( - "eval_pol" in init_params, - f"Component class '{component.component_class}' does not support polarization splitting.", - logger, - ) - kwargs["eval_pol"] = eval_pol - return component_cls(component.params, global_params, **kwargs) + raise ValueError(f"Component {comp_name!r} repeats a split execution view.") @classmethod def init_from_params(cls, components:Bunch, params:Bunch): - # Determine whether any of the bands actually have polarization by checking whether any - # MPI ranks are dedicated to QU-processing (maybe slightly hacky but works fine). - pol_bands_exist = params.general.MPI_config.ntask_compsep_QU > 0 + # An execution view for a given polarization is only created if there are CompSep MPI ranks + # assigned to process it. This lets the same component configuration run in intensity-only, + # QU-only, or joint IQU setups; any requested polarization with no ranks is dropped (with a + # warning), and a component left with no views at all is skipped entirely. + pol_has_ranks = { + "I": params.general.MPI_config.ntask_compsep_I > 0, + "QU": params.general.MPI_config.ntask_compsep_QU > 0, + } comp_list = [] for component_str in components: component = components[component_str] - if component.enabled: - if "lmax" in component.params and component.params.lmax == "full": - component.params.lmax = (params.general.nside*5)//2 - component_pol = component.params.polarization if "polarization" in component.params else "I" - if component_pol == "I": - comp_list.append(cls._instantiate_component(component, params.general)) - elif component_pol == "QU": - if not pol_bands_exist: - logging.warning(f"Component '{component_str}' is specified as QU-only but " - f"ntask_compsep_QU=0 (no polarized bands). Skipping.") - continue - comp_list.append(cls._instantiate_component(component, params.general)) - elif component_pol == "IQU": - intensity_comp = cls._instantiate_component(component, params.general, eval_pol="I") - comp_list.append(intensity_comp) - if not pol_bands_exist: - logging.warning(f"Component '{component_str}' is specified as IQU but " - f"ntask_compsep_QU=0 (no polarized bands). " - f"Only the intensity (I) part will be used.") - else: - pol_comp = cls._instantiate_component(component, params.general, - eval_pol="QU") - comp_list.append(pol_comp) - else: - raise ValueError( - f"Unrecognized polarization in parameter file for component {component_str}" - ) + if not component.enabled: + continue + component_cls = getattr(component_lib, component.component_class) + if "lmax" in component.params and component.params.lmax == "full": + component.params.lmax = (params.general.nside*5)//2 + component_pol = component.params.polarization if "polarization" in component.params \ + else "I" + if component_pol not in EXECUTION_POLS: + raise ValueError( + f"Unrecognized polarization in parameter file for component {component_str}") + requested_pols = EXECUTION_POLS[component_pol] + active_pols = [eval_pol for eval_pol in requested_pols if pol_has_ranks[eval_pol]] + if not active_pols: + logging.warning(f"Component '{component_str}' is specified as {component_pol} but " + f"no CompSep ranks are assigned to its polarization(s) " + f"({'/'.join(requested_pols)}). Skipping.") + continue + if len(active_pols) < len(requested_pols): + skipped = [pol for pol in requested_pols if pol not in active_pols] + logging.warning(f"Component '{component_str}' is specified as {component_pol} but " + f"no CompSep ranks are assigned to {'/'.join(skipped)}; only the " + f"{'/'.join(active_pols)} part will be used.") + for eval_pol in active_pols: + comp_list.append(component_cls(component.params, params.general, eval_pol=eval_pol, + comp_name=component._name, allocate_empty_alms=True)) return cls(comp_list) def _assert_consistent_comps(self, other: "CompList") -> None: @@ -1034,51 +1079,21 @@ def copy_matching_data_from(self, other: "CompList") -> None: comp._assert_consistent_comp(other_comp) np.copyto(comp._data, other_comp._data) - def reassemble_from_split_solution(self, local_solution: "CompList", comm: MPI.Comm, - *, is_I_master: bool, is_QU_master: bool, - I_master: int, QU_master: int, root: int = 0) -> None: - """Collect split I/QU solver results back into the full execution list on the CompSep root. + def broadcast_pol_views(self, comm: MPI.Comm, *, eval_pol: str, source: int) -> None: + """Broadcast all execution views of `eval_pol` from `source` to every rank in `comm`. - The intensity master already owns the local I solution and receives the QU views from the - QU master. Once the full execution list is assembled on the root, every rank gets a copy - through a blocking broadcast so `SkyModel` realization can happen locally. + Used after a sampling step: only the ranks that actually solved a given polarization hold + the updated component data, so broadcasting that polarization's views from one authoritative + `source` rank restores a globally consistent component list. """ - if is_I_master: - self.copy_matching_data_from(local_solution) - - pol_components = self.components_for_eval_pol("QU") - if pol_components: - if is_QU_master: - local_pol_by_key = {comp.execution_key: comp for comp in local_solution.comp_list} - for tag, comp in enumerate(pol_components): - local_comp = local_pol_by_key.get(comp.execution_key) - if local_comp is None: - raise ValueError( - f"Missing QU component {comp.execution_label} on QU master." - ) - comm.Send(local_comp._data, dest=I_master, tag=tag) - if is_I_master: - for tag, comp in enumerate(pol_components): - comm.Recv(comp._data, source=QU_master, tag=tag) - - for comp in self.comp_list: - comp.bcast_data_blocking(comm, root=root) + for comp in self.components_for_eval_pol(eval_pol): + comp.bcast_data_blocking(comm, root=source) def joined(self) -> "CompList": - """Collapse split execution views back to one logical component per `longname`.""" - grouped_components = {} - logical_order = [] - for comp in self.comp_list: - if comp.logical_key not in grouped_components: - grouped_components[comp.logical_key] = [] - logical_order.append(comp.logical_key) - grouped_components[comp.logical_key].append(comp) - + """Collapse split execution views back to one logical component per `comp_name`.""" joined_components = [] - for logical_key in logical_order: - group = grouped_components[logical_key] - split_views = [comp for comp in group if comp.is_split_view] - unsplit_views = [comp for comp in group if not comp.is_split_view] + for logical_key, group in self._group_by_logical_key(self.comp_list): + split_views, unsplit_views = self._partition_execution_views(group) if unsplit_views and split_views: raise ValueError( f"Logical component {logical_key[1]!r} mixes split and unsplit execution views." @@ -1098,23 +1113,6 @@ def joined(self) -> "CompList": joined_components.append(split_views[0].join_split_views(split_views[1])) return CompList(joined_components) - - def split(self, color:int, IvsQU_colors:tuple = (0,1)): - """ - Extracts from `comp_list` only the components containing the correct Stokes parameter based - on the passed `color` of the local MPI rank. By default, color=0 will treat Intensity and - color=1 QU. A list with the relevant components is returned. - """ - IvsQU_colors = IvsQU_colors[:2] #cut off eventual elements in excess - if color not in IvsQU_colors: - logging.warning(f"Color {color} not in colors assigned to I or QU ({IvsQU_colors})!") - return CompList([]) - elif color == IvsQU_colors[0]: - target_pol = "I" - elif color == IvsQU_colors[1]: - target_pol = "QU" - - return self.split_for_eval_pol(target_pol) @property def components(self): @@ -1134,58 +1132,37 @@ def __matmul__(self, other) -> float: for c1, c2 in zip(self.components, other.components): res += float(c1 @ c2) return res + + def _apply_componentwise_op(self, other: "CompList", component_op, *, inplace: bool) -> "CompList": + self._assert_consistent_comps(other) + target = self if inplace else deepcopy(self) + for target_comp, other_comp in zip(target.components, other.components): + component_op(target_comp, other_comp) + return target def __add__(self, other): - self._assert_consistent_comps(other) - out = deepcopy(self) - for out_comp, other_comp in zip(out.components, other.components): - out_comp += other_comp - return out + return self._apply_componentwise_op(other, Component.__iadd__, inplace=False) def __iadd__(self, other): - self._assert_consistent_comps(other) - for c1, c2 in zip(self.components, other.components): - c1 += c2 - return self + return self._apply_componentwise_op(other, Component.__iadd__, inplace=True) def __sub__(self, other): - self._assert_consistent_comps(other) - out = deepcopy(self) - for out_comp, other_comp in zip(out.components, other.components): - out_comp -= other_comp - return out + return self._apply_componentwise_op(other, Component.__isub__, inplace=False) def __isub__(self, other): - self._assert_consistent_comps(other) - for c1, c2 in zip(self.components, other.components): - c1 -= c2 - return self + return self._apply_componentwise_op(other, Component.__isub__, inplace=True) def __mul__(self, other): - self._assert_consistent_comps(other) - out = deepcopy(self) - for out_comp, other_comp in zip(out.components, other.components): - out_comp *= other_comp - return out + return self._apply_componentwise_op(other, Component.__imul__, inplace=False) def __imul__(self, other): - self._assert_consistent_comps(other) - for c1, c2 in zip(self.components, other.components): - c1 *= c2 - return self + return self._apply_componentwise_op(other, Component.__imul__, inplace=True) def __truediv__(self, other): - self._assert_consistent_comps(other) - out = deepcopy(self) - for out_comp, other_comp in zip(out.components, other.components): - out_comp /= other_comp - return out + return self._apply_componentwise_op(other, Component.__itruediv__, inplace=False) def __itruediv__(self, other): - self._assert_consistent_comps(other) - for c1, c2 in zip(self.components, other.components): - c1 /= c2 - return self + return self._apply_componentwise_op(other, Component.__itruediv__, inplace=True) def __getitem__(self, index): return self.comp_list[index] diff --git a/src/commander4/sky_models/sky_model.py b/src/commander4/sky_models/sky_model.py index cffc32a..46148b5 100644 --- a/src/commander4/sky_models/sky_model.py +++ b/src/commander4/sky_models/sky_model.py @@ -19,7 +19,7 @@ def get_sky_at_nu(self, nu, nside, pols_required, fwhm=None): `QU` views) or a joined logical list containing `IQU` components. """ npix = 12*nside**2 - + if pols_required == "I": skymap = np.zeros((1, npix), dtype=np.float32) elif pols_required == "QU": @@ -49,53 +49,3 @@ def get_sky_at_nu(self, nu, nside, pols_required, fwhm=None): else: raise ValueError(f"Unsupported component polarization '{component.eval_pol}'.") return skymap - -# class SkyModel: -# def __init__(self, components_I, components_Q=None, components_U=None): -# # components = list of Component objects -# self.components_I = components_I -# self.components_Q = components_Q -# self.components_U = components_U -# self.all_components = [components_I, components_Q, components_U] -# self.comp_names = ["I", "Q", "U"] - -# def get_sky(self, band): -# """ Get sky from a bandpass. -# """ -# raise NotImplementedError - -# def get_sky_at_nu(self, nu, nside, fwhm=None, pol=(True,True,True)): -# """ Get sky at specific frequency. -# """ -# npix = 12*nside**2 -# skymap = np.zeros((np.sum(pol), npix)) -# idx = 0 -# for ipol in range(3): -# if pol[ipol]: -# if self.all_components[ipol] is None: -# raise ValueError("Attempted to create sky model containing " -# f"{self.comp_names[ipol]} but this component is not set.") -# for component in self.all_components[ipol]: -# skymap[idx] += component.get_sky(nu, nside, ipol, fwhm) -# idx += 1 -# return skymap - - # def get_foreground_sky_at_nu(self, nu, nside, fwhm=None): - # """ Get sky, excluding the cmb, at specific frequency. - # """ - # npix = 12*nside**2 - # skymap = np.zeros((npix)) - # for component in self.components: - # if not isinstance(component, CMB): - # skymap += component.get_sky(nu, nside, fwhm) - # return skymap - - # def get_cmb_sky_at_nu(self, nu, nside, fwhm=None): - # """ Get sky, excluding the cmb, at specific frequency. - # """ - # npix = 12*nside**2 - # skymap = np.zeros((npix)) - # for component in self.components: - # if isinstance(component, CMB): - # skymap += component.get_sky(nu, nside, fwhm) - # return skymap \ No newline at end of file diff --git a/src/commander4/solvers/CG_compsep_solver.py b/src/commander4/solvers/CG_compsep_solver.py index b670fe7..8d2d277 100644 --- a/src/commander4/solvers/CG_compsep_solver.py +++ b/src/commander4/solvers/CG_compsep_solver.py @@ -8,7 +8,7 @@ from pixell.bunch import Bunch from commander4.data_models.detector_map import DetectorMap -from commander4.sky_models.component import Component, CompList +from commander4.sky_models.component import CompList from commander4.utils.math_operations import alm_to_map_adjoint, gaussian_random_alm, almxfl,\ complist_dot, complist_norm from commander4.solvers.dense_matrix_math import DenseMatrix @@ -79,9 +79,6 @@ def project_all_comps_to_band(self, comp_list_in: CompList, In Commander4 notation, applies A matrix, from comp list to band alms. """ - # band_out.alms = np.zeros_like(band_out.alms) - # for comp in comp_list_in: - # comp.project_comp_to_band(band_out, nthreads=self.nthreads) comp_list_in.project_comp_to_band(band_out, nthreads=self.nthreads) # B Y^-1 M Y a alm_out = self.det_map.apply_B(band_out.alms) @@ -102,8 +99,6 @@ def eval_all_comps_from_band(self, band_in:Band, self.det_map.apply_B(band_in.alms) # Y^T M^T Y^-1^T B^T a - # for comp in comp_list_out: - # comp.eval_comp_from_band(band_in, nthreads=self.nthreads) comp_list_out.eval_comp_from_band(band_in, nthreads=self.nthreads) return comp_list_out @@ -125,10 +120,6 @@ def apply_LHS_matrix(self, comp_list_in: CompList) -> CompList: comp_list = deepcopy(comp_list_in) if myrank == 0: # this task actually holds a component - # for comp in comp_list: - # - # comp.apply_smoothing_prior_sqrt() - # S^{1/2} a comp_list.apply_smoothing_prior_sqrt() @@ -137,8 +128,6 @@ def apply_LHS_matrix(self, comp_list_in: CompList) -> CompList: # and are therefore limited to <2GB arrays... We have to fallback to blocking communication # for >2GB arrays. In the future we should probably implement chunking instead. - # for comp in comp_list: - # comp.bcast_data_blocking(self.CompSep_comm) comp_list.bcast_data_blocking(self.CompSep_comm) # B Y^-1 M Y S^{1/2} a @@ -155,15 +144,9 @@ def apply_LHS_matrix(self, comp_list_in: CompList) -> CompList: use_blocking = biggest_size_bytes > MPI_LIMIT_32BIT if use_blocking: print(f"Fallback to blocking comm (array size = {biggest_size_bytes:.2e}B)") - # for comp in comp_list: - # comp.accum_data_blocking(self.CompSep_comm) comp_list.accum_data_blocking(self.CompSep_comm) else: - # requests = [] - # for comp in comp_list: - # req = comp.accum_data_non_blocking(self.CompSep_comm) - # requests.append(req) requests = comp_list.accum_data_non_blocking(self.CompSep_comm) if myrank == 0: @@ -217,9 +200,6 @@ def calc_RHS_mean(self, comp_list: CompList) -> CompList: if myrank == 0: for comp in comp_list: logger.info(f"RHS1 comp-{comp.shortname}: {np.mean(np.abs(comp._data)):.2e}") - # b = comp_list - # else: - # b = [] return comp_list @@ -257,11 +237,9 @@ def calc_RHS_fluct(self, comp_list: CompList) -> CompList: if myrank == 0: for comp in comp_list: - logger.info(f"RHS1 comp-{comp.shortname}: {np.mean(np.abs(comp._data)):.2e}") - else: - b = [] + logger.info(f"RHS2 comp-{comp.shortname}: {np.mean(np.abs(comp._data)):.2e}") - return b + return comp_list def calc_RHS_prior_mean(self, comp_list: CompList) -> CompList: @@ -278,7 +256,7 @@ def calc_RHS_prior_mean(self, comp_list: CompList) -> CompList: for ipol in range(self.npol): almxfl(mu[ipol], comp.P_smoothing_prior.astype(self.float_dtype, copy=False), inplace=True) - logger.info(f"RHS3 comp-{comp.longname}: {np.mean(np.abs(mu)):.2e}") + logger.info(f"RHS3 comp-{comp.comp_name}: {np.mean(np.abs(mu)):.2e}") mu_s.append(mu) else: mu_s = [] @@ -296,7 +274,7 @@ def calc_RHS_prior_fluct(self, comp_list: CompList) -> CompList: eta2 = np.zeros((self.npol, comp.alm_len_complex), dtype=self.complex_dtype) for ipol in range(self.npol): eta2[ipol] = gaussian_random_alm(comp.lmax, comp.lmax, self.spin, 1) - logger.info(f"RHS4 comp-{comp.longname}: {np.mean(np.abs(eta2)):.2e}") + logger.info(f"RHS4 comp-{comp.comp_name}: {np.mean(np.abs(eta2)):.2e}") eta2_s.append(eta2) else: eta2_s = [] @@ -363,7 +341,7 @@ def solve_CG(self, LHS: Callable, RHS: CompList, x0 = None, M = None, # A-norm error is only defined for the full vector. logger.info(f"CG iter {iter:3d} - True A-norm error: {CG_Anorm_error:.3e}") # We can print the individual component L2 errors. - logger.info(f"CG iter {iter:3d} - {self.comp_list[mycomp].longname} - "\ + logger.info(f"CG iter {iter:3d} - {self.comp_list[mycomp].comp_name} - "\ f"True L2 error: {CG_errors_true:.3e}") else: if x_true is not None: @@ -402,7 +380,6 @@ def solve(self, comp_list:CompList, seed=None) -> CompList: # RHS4 = self.calc_RHS_prior_fluct() # RHS = [_R1 + _R2 + _R3 + _R4 for _R1, _R2, _R3, _R4 in zip(RHS1, RHS2, RHS3, RHS4)] RHS = RHS1 - #del(self.map_sky) # Initialize the precondidioner class, which is in the module "solvers.preconditioners", # and has a name specified by self.params.compsep.preconditioner. diff --git a/src/commander4/solvers/perpix_compsep_solver.py b/src/commander4/solvers/perpix_compsep_solver.py index 67e72b1..ca25f03 100644 --- a/src/commander4/solvers/perpix_compsep_solver.py +++ b/src/commander4/solvers/perpix_compsep_solver.py @@ -136,12 +136,9 @@ def solve_compsep_perpix(proc_comm: MPI.Comm, detector_data: DetectorMap, nband = len(freqs) comp_maps[ipol] = np.zeros((ncomp, npix)) M = np.empty((nband, ncomp)) - idx = 0 - for i in range(ncomp): - # if ipol == 0 or comp_list[i].polarized: - M[:,idx] = comp_list[i].get_sed(freqs) - idx += 1 - rand = np.random.randn(npix,nband) + for icomp in range(ncomp): + M[:, icomp] = comp_list[icomp].get_sed(freqs) + rand = np.random.randn(npix, nband) # TODO: Write unit tests that confirm Python and C gives same answers. # TODO: Should scale M to make solution more well-conditioned, and then adjust # solution with the scaling factor used. From 0b79b58f771961a83eae9d7f0e42044be09da9f9 Mon Sep 17 00:00:00 2001 From: jgslunde Date: Wed, 3 Jun 2026 23:55:58 +0200 Subject: [PATCH 17/23] Updated parameter files with sampling groups support --- params/Planck_TODs/param_PlanckTODs_CG.yml | 16 +++++++++++++--- params/Planck_TODs/param_PlanckTODs_perpix.yml | 13 ++++++++++--- 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/params/Planck_TODs/param_PlanckTODs_CG.yml b/params/Planck_TODs/param_PlanckTODs_CG.yml index 48010c4..fc79224 100644 --- a/params/Planck_TODs/param_PlanckTODs_CG.yml +++ b/params/Planck_TODs/param_PlanckTODs_CG.yml @@ -54,15 +54,25 @@ general: chain_maps_interval: 5 output_paths: - stats: "../cmdr4_stats/PlanckTODs_CG/" - chains: "../cmdr4_chains/PlanckTODs_CG/" + stats: "../cmdr4_stats/PlanckTODs_CG_samplegroups/" + chains: "../cmdr4_chains/PlanckTODs_CG_samplegroups/" logging: console: level: 'info' file: level: 'info' - filename: '../PlanckTODs_CG.log' + filename: '../PlanckTODs_CG_samplegroups.log' + + + +sampling_groups_compsep: + sample_amps_perpix: + enabled: true + sample_class: "amplitude_sampler_CG" + comps: [CMB, Synchrotron, ThermalDust, FreeFree] + bands: [Planck30GHz, Haslam408MHz, WMAPKa, WMAPQ1, WMAPQ2, WMAPV1, WMAPV2, Planck857GHz, Planck353GHz] + components: CMB: diff --git a/params/Planck_TODs/param_PlanckTODs_perpix.yml b/params/Planck_TODs/param_PlanckTODs_perpix.yml index 3db2c3e..e9032d1 100644 --- a/params/Planck_TODs/param_PlanckTODs_perpix.yml +++ b/params/Planck_TODs/param_PlanckTODs_perpix.yml @@ -44,15 +44,22 @@ general: init_from_chain: false output_paths: - chains: "../cmdr4_chains/PlanckTODs_perpix_newclass_v5_abs+rel/" - stats: "../cmdr4_stats/PlanckTODs_perpix_newclass_v5_abs+rel/" + chains: "../cmdr4_chains/PlanckTODs_perpix_sampgroups/" + stats: "../cmdr4_stats/PlanckTODs_perpix_sampgroups/" logging: console: level: 'debug' file: level: 'info' - filename: '../cmdr4_logs/output_PlanckTODs_perpix_newclass_v1.log' + filename: '../cmdr4_logs/output_PlanckTODs_perpix_sampgroups' + +sampling_groups_compsep: + sample_amps_perpix: + enabled: true + sample_class: "amplitude_sampler_perpix" + comps: [CMB, Synchrotron, ThermalDust, FreeFree] + bands: [Planck30GHz, Haslam408MHz, WMAPKa, WMAPQ1, WMAPQ2, WMAPV1, WMAPV2, Planck857GHz, Planck353GHz] components: From 1913663ebd951ea3f9496a36eace1d7d2a6bfc09 Mon Sep 17 00:00:00 2001 From: jgslunde Date: Wed, 3 Jun 2026 23:56:12 +0200 Subject: [PATCH 18/23] Added colors to log output --- src/commander4/output/log.py | 46 ++++++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/src/commander4/output/log.py b/src/commander4/output/log.py index 495e986..2d2f354 100644 --- a/src/commander4/output/log.py +++ b/src/commander4/output/log.py @@ -1,5 +1,6 @@ import logging import logging.config +import sys # --- Debug levels and their numeric values, including the Python defaults (commented out) --- # DEBUG (10) # >100 per iter, intermediate processing steps and per-rank or per-detector details. @@ -14,6 +15,38 @@ logging.VERBOSE = VERBOSE +class ColorFormatter(logging.Formatter): + """ + Logging formatter that prepends ANSI color codes to console output based on + log level. + + Color scheme: + DEBUG - dim/faint (subtle, lower visual weight) + VERBOSE - dark gray + INFO - default terminal color + QUIET - default terminal color + WARNING - yellow + ERROR - red + CRITICAL - bold red + """ + + _RESET = '\033[0m' + _COLORS = { + logging.DEBUG: '\033[2m', # Dim / faint + VERBOSE: '\033[90m', # Dark gray (bright-black) + logging.INFO: '', # Default – no color code + QUIET: '', # Default – no color code + logging.WARNING: '\033[33m', # Yellow + logging.ERROR: '\033[31m', # Red + logging.CRITICAL: '\033[1;31m', # Bold red + } + + def format(self, record: logging.LogRecord) -> str: + message = super().format(record) + color = self._COLORS.get(record.levelno, '') + return f'{color}{message}{self._RESET}' if color else message + + def init_loggers(logger_params): """ Intended usage: This function is called once at the very beginning of the @@ -109,6 +142,19 @@ def level_val(name): return getattr(logging, name, 20) logging.config.dictConfig(config_dict) + # Apply ColorFormatter to the console handler manually. dictConfig's '()' + # factory passes keyword arguments directly to the constructor, but + # logging.Formatter.__init__ uses 'fmt' not 'format', so the factory path + # silently falls back to an unconfigured formatter. Setting it here avoids + # that problem entirely. + if 'console' in logger_params: + _fmt = '{asctime} - {name} - {levelname} - {message}' + _color_fmt = ColorFormatter(fmt=_fmt, style='{') + for _handler in logging.root.handlers: + if isinstance(_handler, logging.StreamHandler) and not isinstance(_handler, logging.FileHandler): + _handler.setFormatter(_color_fmt) + break + # Configure logging to redirect warnings from py.warning. Note that this will *prevent* these # from being sent to sys.stderr, to avoid duplication. logging.captureWarnings(True) From 2e31ad197b4f3b42ac8c78c4b32089ece73ec58f Mon Sep 17 00:00:00 2001 From: jgslunde Date: Thu, 4 Jun 2026 00:03:03 +0200 Subject: [PATCH 19/23] tiny fix --- src/commander4/mpi_management.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/commander4/mpi_management.py b/src/commander4/mpi_management.py index 693fb6b..858d421 100644 --- a/src/commander4/mpi_management.py +++ b/src/commander4/mpi_management.py @@ -133,7 +133,7 @@ def init_mpi(params): compsep_master = global_params.MPI_config.ntask_tod world_comm.barrier() - time.sleep(worldrank*1e-2) # Small sleep to get prints in nice order. + time.sleep(worldrank*1e-5) # Small sleep to get prints in nice order. mpi_info['world'] = Bunch() mpi_info['world']['comm'] = world_comm From 4b40dd4da434c71b19f313ba564e47ccd8d6b50b Mon Sep 17 00:00:00 2001 From: jgslunde Date: Sun, 7 Jun 2026 12:33:37 +0200 Subject: [PATCH 20/23] Added support for initializing from a sky model (also without spawning any compsep-ranks), and refactored some stuff --- src/commander4/cli.py | 39 ++-- src/commander4/communication.py | 18 ++ src/commander4/compsep_processing.py | 53 ++++-- src/commander4/data_models/TOD_samples.py | 28 +-- src/commander4/mpi_management.py | 28 ++- src/commander4/sky_models/component.py | 219 ++++++++++++++++++---- src/commander4/sky_models/sky_model.py | 14 ++ 7 files changed, 317 insertions(+), 82 deletions(-) diff --git a/src/commander4/cli.py b/src/commander4/cli.py index 21a0cd5..ece8aff 100755 --- a/src/commander4/cli.py +++ b/src/commander4/cli.py @@ -64,9 +64,11 @@ def run_commander4(params: Bunch, params_dict: dict): # as Numpy will not respect a change in thread count after it has been loaded. import numpy as np import commander4.output.log as log - from commander4.tod_processing import process_tod, init_tod_processing, get_initial_sky - from commander4.compsep_processing import process_compsep, init_compsep_processing - from commander4.communication import receive_tod, send_tod, receive_compsep, send_compsep + from commander4.tod_processing import process_tod, init_tod_processing + from commander4.compsep_processing import process_compsep, init_compsep_processing,\ + get_initial_sky_model + from commander4.communication import receive_tod, send_tod, receive_compsep, send_compsep,\ + get_local_initial_sky # Unique seed per worldrank. Hash used for slightly improved entropy. Modulus because seed needs # to be 32 bit. Numpy recommends instead carrying around an instance of 'np.random.default_rng', @@ -103,21 +105,32 @@ def run_commander4(params: Bunch, params_dict: dict): mpi_info.world.compsep_band_masters, root=mpi_info.world.compsep_master) mpi_info['world']['compsep_band_masters'] = world_compsep_band_masters_dict - ###### Sending empty data back and forth ###### + ###### Exchanging the initial sky model ###### + # Component separation is active iff CompSep ranks were allocated (compsep_master is then a + # valid world rank). This single flag replaces the old `perform_compsep` parameter. + compsep_active = mpi_info.world.compsep_master is not None curr_tod_output = None if mpi_info.world.color == 0: - # Chain #1 do TOD processing, resulting in maps_chain1 (we start with a fake output of - # component separation, containing a completely empty sky). - compsep_output_black = get_initial_sky(experiment_data) + # The initial sky model is built from each component's init_from / init_chain_path (else + # zeros). If CompSep ranks exist they build and send it (as for every later iteration); + # otherwise we build it locally so a sensible fixed sky is available with no CompSep ranks. + if compsep_active: + curr_compsep_output = receive_compsep(mpi_info, experiment_data, my_band_tod_id, + mpi_info.world.compsep_band_masters) + else: + curr_compsep_output = get_local_initial_sky(mpi_info, experiment_data, params) curr_tod_output, tod_samples = process_tod(mpi_info, experiment_data, tod_samples_chain1, - compsep_output_black, params, 1, 1) - if params.general.perform_compsep: + curr_compsep_output, params, 1, 1) + if compsep_active: send_tod(mpi_info, curr_tod_output, my_band_tod_id, mpi_info.world.compsep_band_masters) - curr_compsep_output = compsep_output_black elif mpi_info.world.color == 1: + # Send the initial sky model to TOD before receiving the first TOD output, mirroring the + # process_compsep -> send_compsep -> receive_tod order used inside the main loop. + send_compsep(mpi_info, my_band_compsep_id, get_initial_sky_model(components), + mpi_info.world.tod_band_masters) curr_tod_output = receive_tod(mpi_info, mpi_info.world.tod_band_masters, my_band, my_band_compsep_id, curr_tod_output) @@ -150,7 +163,7 @@ def run_commander4(params: Bunch, params_dict: dict): logger.info(f"TOD: Rank {mpi_info.tod.rank} finished chain {chain_num}, iter "\ f"{iter_num} in {time.time()-t0:.2f}s. Receiving compsep results.") t0 = time.time() - if params.general.perform_compsep: + if compsep_active: curr_compsep_output = receive_compsep(mpi_info, experiment_data, my_band_tod_id, mpi_info.world.compsep_band_masters) @@ -159,7 +172,7 @@ def run_commander4(params: Bunch, params_dict: dict): f"results for chain {chain_num}, iter {iter_num} "\ f"(time spent waiting+receiving = "\ f"{time.time()-t0:.1f}s).") - if params.general.perform_compsep: + if compsep_active: send_tod(mpi_info, curr_tod_output, my_band_tod_id, mpi_info.world.compsep_band_masters) if mpi_info.tod.is_master: @@ -191,7 +204,7 @@ def run_commander4(params: Bunch, params_dict: dict): f" chain {chain_num}, iter {iter_num} (time spent waiting+receiving = "\ f"{time.time()-t0:.1f}s).") # stop compsep machinery - if mpi_info.world.is_master: + if mpi_info.world.is_master and compsep_active: logger.info("TOD: sending STOP signal to compsep") mpi_info.world.comm.send(True, dest=mpi_info.world.compsep_master) diff --git a/src/commander4/communication.py b/src/commander4/communication.py index bace09c..0433ab4 100644 --- a/src/commander4/communication.py +++ b/src/commander4/communication.py @@ -8,6 +8,7 @@ from commander4.data_models.detector_group_TOD import DetGroupTOD from commander4.maps_from_file import read_data_map_from_file from commander4.utils.execution_ids import get_execution_band_id +from commander4.sky_models.sky_model import build_initial_sky_model logger = logging.getLogger(__name__) @@ -81,6 +82,23 @@ def receive_compsep(mpi_info: Bunch, experiment_data: DetGroupTOD, todproc_my_ba return detector_map_arr +def get_local_initial_sky(mpi_info: Bunch, experiment_data: DetGroupTOD, + params: Bunch) -> NDArray[np.floating]: + """Build the initial sky model locally and realize it at this TOD band. + + Used when there are no CompSep ranks: the band master builds the SkyModel from the component + parameters and init files, broadcasts it within the band communicator, and every rank realizes + it at the band frequency/resolution. Mirrors `receive_compsep`, minus the cross-world receive. + """ + if mpi_info.band.is_master: + sky_model = build_initial_sky_model(params) + else: + sky_model = None + sky_model = mpi_info.band.comm.bcast(sky_model, root=0) + return sky_model.get_sky_at_nu(experiment_data.nu, experiment_data.nside, experiment_data.pols, + fwhm=np.deg2rad(experiment_data.fwhm/60.0)) + + def send_tod(mpi_info: Bunch, tod_map_dict: dict[DetectorMap], todproc_my_band_id: str, receivers: Bunch) -> None: """ MPI-send the results from a single band TOD processing to a task on the CompSep side diff --git a/src/commander4/compsep_processing.py b/src/commander4/compsep_processing.py index 76333f3..c8d9485 100644 --- a/src/commander4/compsep_processing.py +++ b/src/commander4/compsep_processing.py @@ -16,6 +16,20 @@ logger = logging.getLogger(__name__) +def _sampling_group_selection(sampling_group: Bunch, key: str) -> list[str] | None: + """Return the names a sampling group selects for `key` ('comps' or 'bands'), or None for "all". + + Both a missing entry and the literal string ``"all"`` select everything (returned as None); + otherwise the entry is expected to be a list of names. + """ + if key not in sampling_group: + return None + value = sampling_group[key] + if isinstance(value, str) and value == "all": + return None + return value + + def _sampling_group_selects_band(selected_bands: list[str] | None, band_name: str, band_identifier: str) -> bool: """Whether a sampling group acts on a band, matched by base name or execution-view identifier. @@ -44,7 +58,8 @@ def _validate_sampling_groups(sampling_groups: Bunch, comp_list: CompList, param """Fail fast if any enabled sampling group references a non-existent component or band. `comps` and `bands` are expected to be lists of strings naming existing components and bands - (bands may be given either as a base name or as an execution-view identifier). + (bands may be given either as a base name or as an execution-view identifier), the string + "all", or omitted. The latter two select everything and are not checked against names. """ known_comp_names = {comp.comp_name for comp in comp_list.joined()} known_band_names = set() @@ -60,13 +75,15 @@ def _validate_sampling_groups(sampling_groups: Bunch, comp_list: CompList, param group = sampling_groups[group_name] if "enabled" in group and not group.enabled: continue - if "comps" in group: - unknown = sorted(set(group.comps) - known_comp_names) + selected_comps = _sampling_group_selection(group, "comps") + if selected_comps is not None: + unknown = sorted(set(selected_comps) - known_comp_names) logassert(not unknown, f"Sampling group {group_name!r} references unknown component(s) {unknown}. " f"Known components: {sorted(known_comp_names)}.", logger) - if "bands" in group: - unknown = sorted(set(group.bands) - known_band_names) + selected_bands = _sampling_group_selection(group, "bands") + if selected_bands is not None: + unknown = sorted(set(selected_bands) - known_band_names) logassert(not unknown, f"Sampling group {group_name!r} references unknown band(s) {unknown}. " f"Known bands: {sorted(known_band_names)}.", logger) @@ -141,6 +158,10 @@ def init_compsep_processing(mpi_info: Bunch, params: Bunch)\ else Bunch() _validate_sampling_groups(sampling_groups, comp_list, params) + # Load the initial component alms (from each component's init_from / init_chain_path, else + # zeros). Done identically on every CompSep rank so comp_list starts globally consistent. + comp_list.load_initial_alms(params) + mpi_info.compsep.band_name = band_name mpi_info.compsep.band_identifier = band_identifier @@ -156,6 +177,15 @@ def init_compsep_processing(mpi_info: Bunch, params: Bunch)\ return comp_list, mpi_info, band_identifier, my_band +def get_initial_sky_model(comp_list: CompList) -> SkyModel: + """Wrap the freshly-initialized `comp_list` as a SkyModel for the pre-loop initial send to TOD. + + `comp_list` already holds its initial alms (set in `init_compsep_processing`), so this is just + the same `SkyModel(comp_list)` that `process_compsep` produces in later iterations. + """ + return SkyModel(comp_list) + + def process_compsep(mpi_info: Bunch, detector_data: DetectorMap, iter: int, chain: int, params: Bunch, comp_list: CompList) -> SkyModel: """Perform a single component-separation iteration. @@ -195,8 +225,8 @@ def process_compsep(mpi_info: Bunch, detector_data: DetectorMap, iter: int, chai if "enabled" in sampling_group and not sampling_group.enabled: continue - sampled_components = sampling_group.comps if "comps" in sampling_group else None - sampled_bands = sampling_group.bands if "bands" in sampling_group else None + sampled_components = _sampling_group_selection(sampling_group, "comps") + sampled_bands = _sampling_group_selection(sampling_group, "bands") band_is_active = _sampling_group_selects_band(sampled_bands, mpi_info.compsep.band_name, mpi_info.compsep.band_identifier) # This rank's components (for its own polarization stream) that take part in this group. @@ -240,14 +270,15 @@ def process_compsep(mpi_info: Bunch, detector_data: DetectorMap, iter: int, chai if source < compsep_comm.size: comp_list.broadcast_pol_views(compsep_comm, eval_pol=eval_pol, source=source) - # Print new per-band chi2s against the updated sky model. - sky_model_at_band = sky_model.get_sky_at_nu(detector_data.nu, detector_data.nside, "IQU", + # Print new per-band chi2s against the updated sky model. Realize only this band's own + # polarization, so inert (unsolved) component views are not synthesized needlessly. + band_pol = "QU" if detector_data.pol else "I" + sky_model_at_band = sky_model.get_sky_at_nu(detector_data.nu, detector_data.nside, band_pol, fwhm=np.deg2rad(detector_data.fwhm/60.0)) pol_names = ["Q", "U"] if detector_data.pol else ["I"] - pol_offset = 1 if detector_data.pol else 0 for ipol in range(detector_data.npol): chi2 = np.mean(np.abs(detector_data.map_sky[ipol] - - sky_model_at_band[ipol + pol_offset])/detector_data.map_rms[ipol]) + sky_model_at_band[ipol])/detector_data.map_rms[ipol]) logger.info(f"Reduced chi2 on rank {compsep_rank} for pol={pol_names[ipol]} "\ f"({detector_data.nu}GHz): {chi2:.3f}") diff --git a/src/commander4/data_models/TOD_samples.py b/src/commander4/data_models/TOD_samples.py index cbe06cb..4c1b085 100644 --- a/src/commander4/data_models/TOD_samples.py +++ b/src/commander4/data_models/TOD_samples.py @@ -115,13 +115,15 @@ def __init__(self, self.scan_ids = np.array([scan.scan_id for scan in experiment_data.scans]) self.jumps = JumpCatalog.empty(self.nscans, self.ndet) + init_chain_path = getattr(params.general, "init_chain_path", False) + init_from_chain = bool(init_chain_path) # Gibbs-sampled quantities - if not params.general.init_from_chain: + if not init_from_chain: # --------------------------------------------------------- # Standard Initialization (No file provided) # --------------------------------------------------------- if self.band_comm.Get_rank() == 0: - logger.info(f"Band {self.band_name} initializing TOD samples from default values.") + logger.info("No previous chain provided. Starting fresh Gibbs chain.") self.noise_params = np.zeros((self.nscans, self.ndet, 3)) + np.nan self.abs_gain = 0.0 @@ -166,24 +168,24 @@ def __init__(self, # Disk Initialization (Read from previous chain) # --------------------------------------------------------- # 1. Find the latest iteration for chain 01 - init_dir = params.general.init_chain_dir - pattern = f"tod/{self.experiment_name}_{self.band_name}_chain{self.chain:02d}_iter*.h5" - search_path = os.path.join(init_dir, pattern) - files = glob.glob(search_path) + # init_dir = params.general.init_chain_dir + # pattern = f"tod/{self.experiment_name}_{self.band_name}_chain{self.chain:02d}_iter*.h5" + # search_path = os.path.join(init_dir, pattern) + # files = glob.glob(search_path) - if not files: - raise FileNotFoundError(f"No chain files found matching: {search_path}") + # if not files: + # raise FileNotFoundError(f"No chain files found matching: {search_path}") - # Sorting alphabetically naturally sorts by the zero-padded iteration number - files.sort() - latest_file = files[-1] + # # Sorting alphabetically naturally sorts by the zero-padded iteration number + # files.sort() + # latest_file = files[-1] if self.band_comm.Get_rank() == 0: logger.info(f"Band {self.band_name} initializing TOD samples from existing chain: "\ - f"{latest_file}.") + f"{init_chain_path}.") # 2. Extract data mapping - with h5py.File(latest_file, "r") as f: + with h5py.File(init_chain_path, "r") as f: # Read the global scan_ids saved by the Gatherv operation global_scan_ids = f["scan_ids"][:] diff --git a/src/commander4/mpi_management.py b/src/commander4/mpi_management.py index 858d421..d0d8a2c 100644 --- a/src/commander4/mpi_management.py +++ b/src/commander4/mpi_management.py @@ -58,7 +58,8 @@ def init_mpi(params): if is_world_master: # Every rank doesn't need to throw an error. tot_num_Compsep_bands_I = 0 tot_num_Compsep_bands_QU = 0 - for band in params.CompSep_bands: + bands = getattr(params, "CompSep_bands", []) # Fallback to empty list if doesn't exist. + for band in bands: if params.CompSep_bands[band].enabled and "I" in params.CompSep_bands[band].polarization: tot_num_Compsep_bands_I += 1 if params.CompSep_bands[band].enabled and "QU" in params.CompSep_bands[band].polarization: @@ -68,10 +69,27 @@ def init_mpi(params): log.lograise(RuntimeError, f"Total number of MPI tasks ({worldsize}) must equal the sum " f"of tasks for TOD ({global_params.MPI_config.ntask_tod}) + CompSep I + QU" f"({global_params.MPI_config.ntask_compsep_I} + {global_params.MPI_config.ntask_compsep_QU}).", logger) - if tot_num_CompSep_ranks != tot_num_Compsep_bands: + # With CompSep ranks allocated, there must be exactly one per enabled band view. With none + # allocated, CompSep is simply off (TOD-only) and any enabled CompSep_bands are ignored. + if tot_num_CompSep_ranks > 0 and tot_num_CompSep_ranks != tot_num_Compsep_bands: log.lograise(RuntimeError, f"CompSep needs exactly as many MPI tasks " f"({tot_num_CompSep_ranks}) as there are bands " f"({tot_num_Compsep_bands}).", logger) + if tot_num_CompSep_ranks == 0 and tot_num_Compsep_bands > 0: + logger.info(f"No CompSep MPI ranks allocated; running TOD-only. The " + f"{tot_num_Compsep_bands} enabled CompSep band view(s) are ignored, and TOD " + f"ranks use the initial sky model built from the components.") + # Component separation is "on" iff CompSep ranks are allocated; enabled sampling groups + # therefore require at least one CompSep rank to run on. + sampling_groups = params.sampling_groups_compsep if "sampling_groups_compsep" in params \ + else {} + has_enabled_sampling_group = any( + "enabled" not in sampling_groups[group] or sampling_groups[group].enabled + for group in sampling_groups) + if has_enabled_sampling_group and tot_num_CompSep_ranks == 0: + log.lograise(RuntimeError, "Enabled compsep sampling groups are configured, but no " + "CompSep MPI ranks are allocated (ntask_compsep_I = ntask_compsep_QU = 0).", + logger) # Split the world communicator into a communicator for compsep and one for TOD (with "color" # being the keyword for the split). @@ -122,7 +140,7 @@ def init_mpi(params): if color == MPI.UNDEFINED: return -1 world_comm.barrier() - time.sleep(worldrank*1e-3) # Small sleep to get prints in nice order. + time.sleep(worldrank*1e-5) # Small sleep to get prints in nice order. logger.debug(f"MPI split performed, hi from worldrank {worldrank} (on machine "\ f"{MPI.Get_processor_name()}) subcomrank {proc_comm.Get_rank()} from color "\ f"{color} of size {proc_comm.Get_size()}.") @@ -130,7 +148,7 @@ def init_mpi(params): # Determine the world ranks of the respective master tasks for compsep and TOD # We ensured that this works by the "key=worldrank" in the split command. tod_master = 0 if global_params.MPI_config.ntask_tod > 0 else None - compsep_master = global_params.MPI_config.ntask_tod + compsep_master = global_params.MPI_config.ntask_tod if tot_num_CompSep_ranks > 0 else None world_comm.barrier() time.sleep(worldrank*1e-5) # Small sleep to get prints in nice order. @@ -263,7 +281,7 @@ def init_mpi_tod(mpi_info, params): is_det_master = MPIrank_det == 0 tod_comm.Barrier() - time.sleep(MPIrank_tod*1e-3) # Small sleep to get prints in nice order. + time.sleep(MPIrank_tod*1e-5) # Small sleep to get prints in nice order. logger.debug(f"TOD: Hello from TOD-rank {MPIrank_tod} (on machine {MPI.Get_processor_name()}) "\ f"dedicated to band {my_band_id}, with local rank {MPIrank_band} (local "\ diff --git a/src/commander4/sky_models/component.py b/src/commander4/sky_models/component.py index 66b4e39..ba1d874 100644 --- a/src/commander4/sky_models/component.py +++ b/src/commander4/sky_models/component.py @@ -3,6 +3,7 @@ import numpy as np import pysm3.units as pysm3u import healpy as hp +import h5py import logging from copy import deepcopy from scipy.interpolate import interp1d @@ -254,8 +255,12 @@ def _zeros_like(self, other, dtype=None, order='K', subok=True, shape=None): # Second tier component classes class DiffuseComponent(Component): requires_defined_pol = True + # The unit in which this component's amplitude (alms) is internally represented -- always uK_RJ + # for diffuse components, including the CMB. Init sky maps are converted to it from their own + # ``units`` (at the component's reference frequency); chain alms are already stored in it. + amplitude_unit = "uK_RJ" - def __init__(self, comp_params: Bunch, global_params: Bunch, + def __init__(self, comp_params: Bunch, global_params: Bunch, allocate_empty_alms=False, eval_pol:None|str=None, comp_name: str | None = None, shortname: str | None = None): super().__init__( @@ -270,10 +275,42 @@ def __init__(self, comp_params: Bunch, global_params: Bunch, self.lmax = comp_params.lmax self.smoothing_prior_FWHM = comp_params.smoothing_prior_FWHM self.smoothing_prior_amplitude = comp_params.smoothing_prior_amplitude + # Unit of an init_from sky map for this component (None -> assume it is already in + # `amplitude_unit`). Only used when reading FITS init maps, not compsep chains. + self.units = comp_params.units if "units" in comp_params else None self._data = None # Alm data is not allocated by default. if allocate_empty_alms: self.allocate_empty_alms() + def _reference_frequency(self, comp_params: Bunch) -> float: + """Reference frequency (GHz) for this view's polarization. + + ``nu_ref`` is either a scalar (shared by I and QU) or a 2-element list ``[nu_I, nu_QU]``. + """ + nu_ref = comp_params.nu_ref + if isinstance(nu_ref, (list, tuple)): + return nu_ref[0] if self.eval_pol == "I" else nu_ref[1] + return nu_ref + + def init_map_to_amplitude(self, sky_map: NDArray) -> NDArray: + """Convert an init sky map (in ``self.units``) to this component's amplitude unit. + + The conversion is done at the component's reference frequency (``self.nu_ref``) using pysm3's + CMB equivalencies. It is a no-op when the units are unspecified or already equal to the + amplitude unit. + """ + if self.units is None or self.units == self.amplitude_unit: + return sky_map + ref_freq = getattr(self, "nu_ref", None) + log.logassert(ref_freq is not None, + f"Component {self.comp_name!r}: converting an init map from {self.units!r} to " + f"{self.amplitude_unit!r} requires a reference frequency, but none is defined.", + logger) + factor = (1*pysm3u.Unit(self.units)).to( + pysm3u.Unit(self.amplitude_unit), + equivalencies=pysm3u.cmb_equivalencies(ref_freq*pysm3u.GHz)).value + return sky_map * factor + @property def npol(self): return get_npol(self.eval_pol) @@ -465,6 +502,9 @@ class TemplateComponent(Component): # Third tier component classes class CMB(DiffuseComponent): default_shortname = "cmb" + # Like all diffuse components, the CMB amplitude is stored internally in uK_RJ, referenced to + # `nu_ref` (default 1 GHz, where uK_RJ ~= uK_CMB). `get_sed` is therefore the *ratio* of the + # thermodynamic-to-RJ conversion at `nu` relative to `nu_ref` (it inherits amplitude_unit=uK_RJ). def __init__(self, comp_params: Bunch, global_params: Bunch, allocate_empty_alms=False, shortname = None, eval_pol = None, comp_name: str | None = None): @@ -476,17 +516,25 @@ def __init__(self, comp_params: Bunch, global_params: Bunch, allocate_empty_alms comp_name=comp_name, shortname=shortname, ) + # The CMB blackbody is polarization-independent, so a scalar reference suffices. The choice + # is arbitrary (the sky is invariant to it); 1 GHz keeps stored amplitudes ~= uK_CMB. + self.nu_ref = self._reference_frequency(comp_params) if "nu_ref" in comp_params else 1.0 def get_sed(self, nu): - """Calculates the spectral energy distribution (SED) for CMB emission. - The result is unitless, but meant to be multiplied by a RJ brightness temperature. + """SED for CMB emission: the thermodynamic-to-RJ conversion at `nu` relative to `nu_ref`. + + The CMB amplitude is stored in uK_RJ referenced to `nu_ref`, so multiplying by this ratio + yields the uK_RJ brightness at `nu`. The result is dimensionless. + Args: - nu (float or np.ndarray): Frequency in GHz at which to evaluate the SED. + nu (float or np.ndarray): Frequency in GHz at which to evaluate the SED. Returns: The SED scaling factor (float or np.ndarray). """ - return (np.ones_like(nu)*pysm3u.uK_CMB).to(pysm3u.uK_RJ,equivalencies= - pysm3u.cmb_equivalencies(nu*u.GHz)).value + def cmb_to_rj(f): + return (np.ones_like(f)*pysm3u.uK_CMB).to( + pysm3u.uK_RJ, equivalencies=pysm3u.cmb_equivalencies(f*u.GHz)).value + return cmb_to_rj(nu) / cmb_to_rj(self.nu_ref) def get_sky_anisotropies(self, nu, nside, fwhm=0): if self.alms is None: @@ -520,7 +568,7 @@ def __init__(self, comp_params: Bunch, global_params: Bunch, allocate_empty_alms ) self.beta = comp_params.beta self.T = comp_params.T - self.nu0 = comp_params.nu0 + self.nu_ref = self._reference_frequency(comp_params) self.prior_l_power_law = 2.5 def get_sed(self, nu): @@ -533,8 +581,8 @@ def get_sed(self, nu): """ # Modified blackbody, in uK_CMB x = (h_over_k*nu)/(self.T) - x0 = (h_over_k*self.nu0)/(self.T) - return (nu / self.nu0)**(self.beta + 1.0) * np.expm1(x0) / np.expm1(x) + x0 = (h_over_k*self.nu_ref)/(self.T) + return (nu / self.nu_ref)**(self.beta + 1.0) * np.expm1(x0) / np.expm1(x) class Synchrotron(DiffuseComponent): @@ -551,7 +599,7 @@ def __init__(self, comp_params: Bunch, global_params: Bunch, allocate_empty_alms shortname=shortname, ) self.beta = comp_params.beta - self.nu0 = comp_params.nu0 + self.nu_ref = self._reference_frequency(comp_params) self.nside_comp_map = 512 self.prior_l_power_law = -3 @@ -563,7 +611,7 @@ def get_sed(self, nu): Returns: The SED scaling factor (float or np.ndarray). """ - return (nu/self.nu0)**self.beta + return (nu/self.nu_ref)**self.beta class FreeFree(DiffuseComponent): @@ -580,7 +628,7 @@ def __init__(self, comp_params: Bunch, global_params: Bunch, allocate_empty_alms shortname=shortname, ) self.T = comp_params.T # Electron temperature in K - self.nu0 = comp_params.nu0 # Reference frequency in GHz + self.nu_ref = self._reference_frequency(comp_params) # Reference frequency in GHz def _gaunt_factor(self, nu, T): """Calculates the Gaunt factor for free-free emission, as per Eq. 18 in BP1. @@ -604,10 +652,10 @@ def get_sed(self, nu): The SED scaling factor (float or np.ndarray). """ gaunt_nu = self._gaunt_factor(nu, self.T) - gaunt_nu0 = self._gaunt_factor(self.nu0, self.T) - - # The scaling is proportional to nu^-2 * g_ff(nu), normalized to 1 at nu0. - sed = (self.nu0 / nu)**2 * (gaunt_nu / gaunt_nu0) + gaunt_nu_ref = self._gaunt_factor(self.nu_ref, self.T) + + # The scaling is proportional to nu^-2 * g_ff(nu), normalized to 1 at nu_ref. + sed = (self.nu_ref / nu)**2 * (gaunt_nu / gaunt_nu_ref) return sed class SpinningDust(DiffuseComponent): @@ -727,7 +775,7 @@ def __init__(self, comp_params: Bunch, global_params: Bunch, *, allocate_empty_alms=allocate_empty_alms, ) #reference frequency - self.nu0 = comp_params.nu_0 + self.nu_ref = comp_params.nu_0 #tabulated data ps_bunch = self.read_dat_to_bunch(comp_params.template_path) #per-source amplitudes @@ -795,10 +843,10 @@ def compute_pix_beams(self, band_fwhm_r, band_nside, recompute=False): def get_sed(self, nu:float): """ - Returns a list of sed's, one per `alpha_list`, evaluated at `nu`, with ref frequency `nu0`. + Returns a list of sed's, one per `alpha_list`, evaluated at `nu`, with ref frequency `nu_ref`. Freq. are in GHz. """ - return (nu/self.nu0)**(self.alpha_arr - 2) + return (nu/self.nu_ref)**(self.alpha_arr - 2) def get_sky(self, nu:float, nside:int, fwhm:float=0.0): """ @@ -925,6 +973,90 @@ def apply_smoothing_prior_sqrt(self): def __repr__(self): return f"Radio Source \n amps: {self._data}" +# Stokes channels stored, in order, for each polarization mode. Used to map the rows of a stored +# (npol, ...) array (whose layout follows its polarization mode) onto the rows an execution view +# needs. Applies equally to chain alm arrays and FITS maps, since both are laid out by polarization. +_POL_CHANNELS = {"I": ("I",), "QU": ("Q", "U"), "IQU": ("I", "Q", "U")} + + +def _pol_row_indices(data: NDArray, eval_pol: str, shortname: str, source_path: str): + """Row indices in a stored (npol, ...) array for `eval_pol`'s Stokes channels. + + The stored polarization mode is inferred from the number of rows (1=I, 2=QU, 3=IQU). Returns + None if the stored data does not contain all channels `eval_pol` needs, so the caller can leave + those alms at zero. Raises only if the row count is not a recognized polarization mode. + """ + nrows = data.shape[0] + stored_pol = {1: "I", 2: "QU", 3: "IQU"}.get(nrows) + log.logassert(stored_pol is not None, + f"Initial data for component {shortname!r} in {source_path!r} has an unexpected " + f"first dimension ({nrows}); expected 1 (I), 2 (QU) or 3 (IQU).", logger) + row_of = {channel: row for row, channel in enumerate(_POL_CHANNELS[stored_pol])} + if any(channel not in row_of for channel in _POL_CHANNELS[eval_pol]): + return None + return [row_of[channel] for channel in _POL_CHANNELS[eval_pol]] + + +def _read_view_alms_from_chain(comp: "DiffuseComponent", chain_path: str) -> NDArray | None: + """This view's alms from a compsep chain (``comps//alms``), or None if not present. + + A missing component is logged as an error (but not fatal); a component present without this + view's polarization is a benign partial initialization and only debug-logged. + """ + with h5py.File(chain_path, "r") as f: + group_path = f"comps/{comp.shortname}" + if group_path not in f or "alms" not in f[group_path]: + logger.error(f"Component {comp.comp_name!r} (shortname {comp.shortname!r}) not found in " + f"init chain {chain_path!r}; leaving its alms at zero.") + return None + stored_alms = f[f"{group_path}/alms"][()] + rows = _pol_row_indices(stored_alms, comp.eval_pol, comp.shortname, chain_path) + if rows is None: + logger.debug(f"Init chain {chain_path!r} has no {comp.eval_pol!r} data for component " + f"{comp.comp_name!r}; leaving those alms at zero.") + return None + return project_alms(np.ascontiguousarray(stored_alms[rows]), comp.lmax) + + +def _read_view_alms_from_fits(comp: "DiffuseComponent", fits_path: str) -> NDArray | None: + """This view's alms from a FITS sky map (transformed), or None if its polarization isn't present. + + The map's polarization content is inferred purely from its shape (npol, npix), so the column + names do not matter. The map is converted from its ``units`` to the component's amplitude unit + (at the component's reference frequency) before being transformed to alms. + """ + sky_map = np.atleast_2d(hp.read_map(fits_path, field=None)) + rows = _pol_row_indices(sky_map, comp.eval_pol, comp.shortname, fits_path) + if rows is None: + logger.debug(f"Init map {fits_path!r} has no {comp.eval_pol!r} data for component " + f"{comp.comp_name!r}; leaving those alms at zero.") + return None + view_map = np.ascontiguousarray(sky_map[rows], dtype=np.float64) + view_map = comp.init_map_to_amplitude(view_map) + nside = hp.npix2nside(view_map.shape[-1]) + return map_to_alm(view_map, nside, comp.lmax, spin=comp.spin) + + +def _load_component_alms(comp: "DiffuseComponent", source_path: str) -> None: + """Set `comp`'s initial alms from `source_path`, dispatching on its file type. + + ``.h5``/``.hd5`` files are read as compsep chains (alms taken directly); ``.fits`` files are + read as sky maps and transformed to alms. If the source does not contain this component or its + polarization, the alms are left at their initial value (zeros). + """ + lower_path = str(source_path).lower() + if lower_path.endswith((".h5", ".hd5")): + view_alms = _read_view_alms_from_chain(comp, source_path) + elif lower_path.endswith(".fits"): + view_alms = _read_view_alms_from_fits(comp, source_path) + else: + log.logassert(False, + f"Unsupported init file {source_path!r} for component {comp.comp_name!r}: " + f"expected a .h5/.hd5 chain or a .fits map.", logger) + if view_alms is not None: + comp.alms = view_alms.astype(comp.dtype, copy=False) + + class CompList: def __init__(self, comp_list:list[Component]): self._validate_comp_list(comp_list) @@ -1004,14 +1136,10 @@ def _validate_comp_list(comp_list: list[Component]) -> None: @classmethod def init_from_params(cls, components:Bunch, params:Bunch): - # An execution view for a given polarization is only created if there are CompSep MPI ranks - # assigned to process it. This lets the same component configuration run in intensity-only, - # QU-only, or joint IQU setups; any requested polarization with no ranks is dropped (with a - # warning), and a component left with no views at all is skipped entirely. - pol_has_ranks = { - "I": params.general.MPI_config.ntask_compsep_I > 0, - "QU": params.general.MPI_config.ntask_compsep_QU > 0, - } + # Build the full logical component list: every enabled component contributes one execution + # view per polarization it defines (I, QU, or both for an IQU component). Construction is + # deliberately independent of the MPI/compsep layout -- a view whose polarization is not + # actually solved or used in a given run simply stays inert at its initial value. comp_list = [] for component_str in components: component = components[component_str] @@ -1025,23 +1153,34 @@ def init_from_params(cls, components:Bunch, params:Bunch): if component_pol not in EXECUTION_POLS: raise ValueError( f"Unrecognized polarization in parameter file for component {component_str}") - requested_pols = EXECUTION_POLS[component_pol] - active_pols = [eval_pol for eval_pol in requested_pols if pol_has_ranks[eval_pol]] - if not active_pols: - logging.warning(f"Component '{component_str}' is specified as {component_pol} but " - f"no CompSep ranks are assigned to its polarization(s) " - f"({'/'.join(requested_pols)}). Skipping.") - continue - if len(active_pols) < len(requested_pols): - skipped = [pol for pol in requested_pols if pol not in active_pols] - logging.warning(f"Component '{component_str}' is specified as {component_pol} but " - f"no CompSep ranks are assigned to {'/'.join(skipped)}; only the " - f"{'/'.join(active_pols)} part will be used.") - for eval_pol in active_pols: + for eval_pol in EXECUTION_POLS[component_pol]: comp_list.append(component_cls(component.params, params.general, eval_pol=eval_pol, comp_name=component._name, allocate_empty_alms=True)) return cls(comp_list) + def load_initial_alms(self, params: Bunch) -> None: + """Populate each component's alms with an initial guess read from a file. + + For every component the source is its own ``init_from`` parameter (inside the component's + ``params`` block) if present, otherwise the global ``params.general.init_chain_path``. The + source may be a compsep chain (``.h5``/``.hd5``, alms read directly) or a FITS sky map + (``.fits``, transformed to alms); the type is decided by the file extension. If neither path + is set the alms are left at their allocated value (zeros). Only diffuse (alm-based) + components are supported for now. + """ + global_path = params.general.init_chain_path if "init_chain_path" in params.general else None + for comp in self.comp_list: + has_explicit_path = "init_from" in comp.comp_params + source_path = comp.comp_params.init_from if has_explicit_path else global_path + if not source_path: + continue # No initial guess requested; leave the allocated zeros. + if not isinstance(comp, DiffuseComponent): + log.logassert(not has_explicit_path, + f"Component {comp.comp_name!r}: 'init_from' is currently only " + f"supported for diffuse (alm-based) components.", logger) + continue + _load_component_alms(comp, source_path) + def _assert_consistent_comps(self, other: "CompList") -> None: if not isinstance(other, CompList): raise TypeError("Both operands must be CompList objects.") diff --git a/src/commander4/sky_models/sky_model.py b/src/commander4/sky_models/sky_model.py index 46148b5..f1833a6 100644 --- a/src/commander4/sky_models/sky_model.py +++ b/src/commander4/sky_models/sky_model.py @@ -1,7 +1,21 @@ import numpy as np +from pixell.bunch import Bunch from commander4.sky_models.component import CompList +def build_initial_sky_model(params: Bunch) -> "SkyModel": + """Construct the initial sky model directly from the parameter file. + + Builds the full component list, loads each component's initial alms (from its ``init_from`` or + the global ``init_chain_path``, else zeros), and wraps it in a SkyModel. This is rank-agnostic + and performs no MPI, so it is used both by CompSep ranks and -- when no CompSep ranks exist -- + by the TOD band masters to construct the initial sky locally. + """ + comp_list = CompList.init_from_params(params.components, params) + comp_list.load_initial_alms(params) + return SkyModel(comp_list) + + class SkyModel: def __init__(self, components:CompList): # components = list of Component objects From 366ea5e8daf1330f6c947f1b48e3947817b07e11 Mon Sep 17 00:00:00 2001 From: jgslunde Date: Sun, 14 Jun 2026 10:33:52 +0200 Subject: [PATCH 21/23] Refactored a bunch of stuff. --- src/commander4/data_models/TOD_samples.py | 101 ++- src/commander4/data_models/detector_TOD.py | 9 +- .../data_models/detector_group_TOD.py | 29 +- src/commander4/data_models/pointing.py | 7 +- src/commander4/data_models/tod_view.py | 163 ++--- .../experiments/planck/tod_reader_planck.py | 6 +- src/commander4/noise_sampling/noise_psd.py | 132 +++- .../noise_sampling/noise_sampling.py | 83 --- src/commander4/noise_sampling/sample_ncorr.py | 186 +++++- src/commander4/output/log.py | 40 +- .../simulations/inplace_litebird_sim.py | 26 +- src/commander4/tod_processing.py | 606 +++++++++++------- tests/test_components.py | 458 +++++++++++++ tests/test_compsep_framework.py | 232 +++++++ tests/test_gain_calibration.py | 283 ++++++++ tests/test_math_utils.py | 80 --- tests/test_noise_sampling.py | 239 +++++++ tests/test_tod_power_spectra.py | 150 +++++ 18 files changed, 2286 insertions(+), 544 deletions(-) create mode 100644 tests/test_components.py create mode 100644 tests/test_compsep_framework.py create mode 100644 tests/test_gain_calibration.py delete mode 100644 tests/test_math_utils.py create mode 100644 tests/test_tod_power_spectra.py diff --git a/src/commander4/data_models/TOD_samples.py b/src/commander4/data_models/TOD_samples.py index 4c1b085..be5da73 100644 --- a/src/commander4/data_models/TOD_samples.py +++ b/src/commander4/data_models/TOD_samples.py @@ -94,13 +94,14 @@ class TODSamples: all ranks on the same band. """ + TOD_PS_NBIN = 100 # Fixed bin count for the optional low-resolution TOD power spectra. + def __init__(self, experiment_data: DetGroupTOD, params: Bunch, my_band: Bunch, band_comm: MPI.Comm, chain: int, - noise_psd_class: str = "oof", ): # Meta-information self.params = params @@ -110,10 +111,38 @@ def __init__(self, self.band_name = experiment_data.band_name self.ndet = experiment_data.ndet self.nscans = experiment_data.nscans + # The noise model defines how many parameters per detector-scan (first entry is sigma0). + self.noise_model = experiment_data.noise_model + self.npar = self.noise_model.npar self.scan_idx_start = experiment_data.scan_idx_start self.scan_idx_stop = experiment_data.scan_idx_stop self.scan_ids = np.array([scan.scan_id for scan in experiment_data.scans]) self.jumps = JumpCatalog.empty(self.nscans, self.ndet) + # Per detector-scan acceptance flag. Currently always True (no scans rejected); kept as a + # chain-tracked quantity so scan rejection can become a sampled step in the future. + self.accept = np.ones((self.nscans, self.ndet), dtype=bool) + + # Low-resolution (log-binned) TOD power spectra, written to the chain by default: a shared + # binned frequency axis plus the binned periodograms of several per-detector-scan TOD views + # (all in detector units): the raw TOD, the correlated-noise realization, the TOD with only + # the correlated noise removed (sky + white noise retained), and the residual (sky model, + # orbital dipole, and correlated noise all subtracted). Filled during TOD processing. The + # binned frequency edges differ per scan (scans have different lengths), so freqs are stored. + ps_shape = (self.nscans, self.ndet, self.TOD_PS_NBIN) + self.tod_ps_freqs = np.full(ps_shape, np.nan, dtype=np.float32) + self.tod_ps_ncorr = np.full(ps_shape, np.nan, dtype=np.float32) + self.tod_ps_raw = np.full(ps_shape, np.nan, dtype=np.float32) + self.tod_ps_ncorrsub = np.full(ps_shape, np.nan, dtype=np.float32) + self.tod_ps_residual = np.full(ps_shape, np.nan, dtype=np.float32) + + # Optional DEBUG: the entire per-sample correlated-noise (n_corr) TODs, written to the chain + # only when explicitly requested (the data is very large). Collected ragged as one float32 + # array per detector-scan; ``None`` disables collection. + if bool(getattr(params.general, "write_ncorr_tods_to_chain", False)): + self.ncorr_tods: list[list[NDArray | None]] | None = \ + [[None] * self.ndet for _ in range(self.nscans)] + else: + self.ncorr_tods = None init_chain_path = getattr(params.general, "init_chain_path", False) init_from_chain = bool(init_chain_path) @@ -125,7 +154,7 @@ def __init__(self, if self.band_comm.Get_rank() == 0: logger.info("No previous chain provided. Starting fresh Gibbs chain.") - self.noise_params = np.zeros((self.nscans, self.ndet, 3)) + np.nan + self.noise_params = np.zeros((self.nscans, self.ndet, self.npar)) + np.nan self.abs_gain = 0.0 self.rel_gain = np.zeros((self.ndet)) self.temporal_gain = np.zeros((self.nscans, self.ndet)) @@ -143,9 +172,14 @@ def __init__(self, for idet, det in enumerate(scan.detectors): self.noise_params[iscan,idet] = det.init_scalars[1:] else: - # Option 3: Fallback to sensible defaults. - logger.warning("Did not find initial noise parameters, falling back to sensible defaults.") - self.noise_params[:] = np.array([1e-3, 0.1, -1.0]) + # Option 3: Fall back to the noise model's default parameters (ensuring a finite + # sigma0, which the model leaves as NaN to be estimated from the data). + logger.warning("Did not find initial noise parameters, falling back to the noise " + "model's default parameters.") + default_params = np.array(self.noise_model.params, dtype=np.float64) + if not np.isfinite(default_params[0]): + default_params[0] = 1e-3 + self.noise_params[:] = default_params if "gain" in my_band.detectors[experiment_data.scans[0].detectors[0].name]: for iscan, scan in enumerate(experiment_data.scans): @@ -196,7 +230,8 @@ def __init__(self, try: local_indices = [global_id_to_index[sid] for sid in self.scan_ids] except KeyError as e: - raise ValueError(f"Local scan ID {e} not found in the global chain file {latest_file}.") + raise ValueError(f"Local scan ID {e} not found in the global chain file " + f"{init_chain_path}.") from e # 3. Load Per-Band and Per-Detector arrays (Identical across ranks) self.abs_gain = float(f["abs_gain"][...]) if "abs_gain" in f else None @@ -205,6 +240,7 @@ def __init__(self, # 4. Load and slice Per-Scan arrays (Distributed across ranks) self.temporal_gain = f["temporal_gain"][local_indices, :] if "temporal_gain" in f else None self.noise_params = f["noise_params"][local_indices, ...] if "noise_params" in f else None + self.accept = f["accept"][local_indices, ...].astype(bool) self.jumps = JumpCatalog.from_hdf5(f, local_indices, self.ndet) if self.band_comm.Get_rank() == 0: @@ -235,6 +271,26 @@ def gain_all(self) -> NDArray[np.floating]: return gain + def _pack_ncorr_tods(self) -> tuple[NDArray, NDArray]: + """Pack the optional per-(scan, det) correlated-noise TODs for ragged chain storage. + + Returns a ``(nscans, ndet)`` int64 array of per-detector-scan lengths and a 1-D float32 + concatenation of all segments in scan-major, detector-minor order (matching how the + gather routines concatenate). The reader reconstructs each TOD by walking the lengths. + """ + lengths = np.zeros((self.nscans, self.ndet), dtype=np.int64) + segments = [] + for iscan in range(self.nscans): + for idet in range(self.ndet): + seg = self.ncorr_tods[iscan][idet] + if seg is not None: + seg = np.asarray(seg, dtype=np.float32).ravel() + lengths[iscan, idet] = seg.size + segments.append(seg) + flat = np.concatenate(segments) if segments else np.zeros(0, dtype=np.float32) + return lengths, flat + + def write_chain_to_file(self, itr: int): band_comm = self.band_comm params = self.params @@ -274,6 +330,30 @@ def write_chain_to_file(self, itr: int): noise_params_global = _gather_scan_distributed_array(band_comm, self.noise_params, scans_per_rank) + # 4b. Acceptance flags (per-scan per-detector; stored as int8 for MPI/HDF compatibility) + accept_global = _gather_scan_distributed_array(band_comm, self.accept.astype(np.int8), + scans_per_rank) + + # 4c. Low-resolution TOD power spectra (per-scan per-detector per-bin). + tod_ps_freqs_global = _gather_scan_distributed_array(band_comm, self.tod_ps_freqs, + scans_per_rank) + tod_ps_ncorr_global = _gather_scan_distributed_array(band_comm, self.tod_ps_ncorr, + scans_per_rank) + tod_ps_raw_global = _gather_scan_distributed_array(band_comm, self.tod_ps_raw, + scans_per_rank) + tod_ps_ncorrsub_global = _gather_scan_distributed_array(band_comm, self.tod_ps_ncorrsub, + scans_per_rank) + tod_ps_residual_global = _gather_scan_distributed_array(band_comm, self.tod_ps_residual, + scans_per_rank) + + # 4d. Optional DEBUG: full per-sample correlated-noise TODs (ragged per-scan per-detector). + ncorr_lengths_global = ncorr_flat_global = None + if self.ncorr_tods is not None: + ncorr_lengths_local, ncorr_flat_local = self._pack_ncorr_tods() + ncorr_lengths_global = _gather_scan_distributed_array(band_comm, ncorr_lengths_local, + scans_per_rank) + ncorr_flat_global = _gather_variable_length_1d_array(band_comm, ncorr_flat_local) + # 5. Jump corrections (per-scan per-detector ragged quantity) jump_counts_local, jump_locations_local, jump_offsets_local = self.jumps.pack() jump_counts_global = _gather_scan_distributed_array(band_comm, jump_counts_local, @@ -302,6 +382,15 @@ def write_chain_to_file(self, itr: int): file["temporal_gain"] = temporal_gain_global if noise_params_global is not None: file["noise_params"] = noise_params_global + file["accept"] = accept_global + file["tod_ps_freqs"] = tod_ps_freqs_global + file["tod_ps_ncorr"] = tod_ps_ncorr_global + file["tod_ps_raw"] = tod_ps_raw_global + file["tod_ps_ncorrsub"] = tod_ps_ncorrsub_global + file["tod_ps_residual"] = tod_ps_residual_global + if ncorr_lengths_global is not None: + file["ncorr_tod_lengths"] = ncorr_lengths_global + file["ncorr_tod_flat"] = ncorr_flat_global if jump_counts_global is not None: file["jump_counts"] = jump_counts_global file["jump_locations"] = jump_locations_global diff --git a/src/commander4/data_models/detector_TOD.py b/src/commander4/data_models/detector_TOD.py index 2c168e3..5f9f921 100644 --- a/src/commander4/data_models/detector_TOD.py +++ b/src/commander4/data_models/detector_TOD.py @@ -149,10 +149,15 @@ def __init__( else flag_encoded ) self._bad_data_bitmask = bad_data_bitmask - self._huffman_tree = huffman_tree self._huffman_symbols = huffman_symbols - self._huffman_tree2 = huffman_tree2 + self._huffman_tree = huffman_tree + # C++ decoder accepts only int64 for the tree. + if self._huffman_tree is not None: + self._huffman_tree = self._huffman_tree.astype(np.int64, copy=False) self._huffman_symbols2 = huffman_symbols2 + self._huffman_tree2 = huffman_tree2 + if self._huffman_tree2 is not None: + self._huffman_tree2 = self._huffman_tree2.astype(np.int64, copy=False) self._tod_is_compressed = tod_is_compressed self._flag_is_compressed = flag_is_compressed # The Huffman decoder expects uint8 arrays; for bytes and HDF5-backed diff --git a/src/commander4/data_models/detector_group_TOD.py b/src/commander4/data_models/detector_group_TOD.py index be7fe7e..e699c83 100644 --- a/src/commander4/data_models/detector_group_TOD.py +++ b/src/commander4/data_models/detector_group_TOD.py @@ -3,7 +3,7 @@ from commander4.data_models.scan_TOD import ScanTOD from commander4.noise_sampling.noise_psd import NoisePSD -from commander4.utils.math_operations import forward_rfft, backward_rfft +from commander4.utils.math_operations import forward_rfft_mirrored, backward_rfft_mirrored class DetGroupTOD: """Container for all scan TODs belonging to one detector group (experiment + band). @@ -43,20 +43,31 @@ def __init__(self, scans: list[ScanTOD], experiment_name: str, band_name: str, n def apply_N_inv(self, tod: NDArray, noise_params: NDArray, samprate: float|None = None, inplace=False) -> NDArray: - """ Modulates the input TOD with the noise model of this Det-Group, using the specified - noise parameters. If a sample rate is specified, the TOD is assumed to have been - downsampled, and the white noise level will be adjusted accordingly. + """ Applies the inverse noise covariance N^-1 of this Det-Group to the input TOD, using the + specified noise parameters. If a sample rate is specified, the TOD is assumed to have + been downsampled, and the noise level is scaled accordingly. The DC (mean) mode is + projected out, matching the Commander3 ``multiply_inv_N`` convention. """ - # TODO: Some noise-PS types might require different handling, such as flat PS. actual_samprate = samprate if samprate is not None else self.fsamp tod_out = tod if inplace else np.zeros_like(tod) + + # White-noise fast path: P(f) = sigma0^2 (flat), so N^-1 is a scalar and the FFT is skipped. + if self.noise_model.is_white: + scale = float(noise_params[0])**2 + if samprate is not None and samprate != self.fsamp: + scale *= samprate/self.fsamp + tod_out[:] = tod/scale + tod_out -= np.mean(tod_out) # Project out the DC mode (mean). + return tod_out + + # Mirrored FFT (length 2*ntod) reduces boundary/periodicity ringing. ntod = tod.shape[0] - freqs = np.fft.rfftfreq(ntod, d=1.0/actual_samprate) + freqs = np.fft.rfftfreq(2*ntod, d=1.0/actual_samprate) noise_PS = self.noise_model.eval_full(freqs, noise_params) - - tod_f = forward_rfft(tod) if samprate is not None and samprate != self.fsamp: noise_PS *= samprate/self.fsamp + tod_f = forward_rfft_mirrored(tod) tod_f /= noise_PS - tod_out[:] = backward_rfft(tod_f, ntod) + tod_f[0] = 0.0 # Project out the DC mode (mean), matching Commander multiply_inv_N. + tod_out[:] = backward_rfft_mirrored(tod_f, ntod) return tod_out \ No newline at end of file diff --git a/src/commander4/data_models/pointing.py b/src/commander4/data_models/pointing.py index 2b44b54..7ed1e91 100644 --- a/src/commander4/data_models/pointing.py +++ b/src/commander4/data_models/pointing.py @@ -178,7 +178,8 @@ def __init__(self, self.ntod = ntod self.pix_encoded = pix self.psi_encoded = psi - self.huffman_tree = huffman_tree + # C++ decoder accepts only int64 for the tree. + self.huffman_tree = huffman_tree.astype(np.int64, copy=False) self.huffman_symbols = huffman_symbols self.npsi = npsi self.pix_is_compressed = isinstance(pix, (bytes, np.void)) @@ -247,7 +248,7 @@ def get_pix(self, nside: int | None = None) -> NDArray[np.integer]: """Return HEALPix pixel indices at the requested output nside.""" target_nside = self.nside if nside is None else nside if self.pix_is_compressed: - pix = np.zeros(self.ntod_original, dtype=np.int64) + pix = np.zeros(self.ntod_original, dtype=self.huffman_symbols.dtype) pix = cpp_utils.huffman_decode(self.pix_compressed_u8, self.huffman_tree, self.huffman_symbols, pix) # The compressed stream stores first differences, so reconstruct the @@ -268,7 +269,7 @@ def get_pix(self, nside: int | None = None) -> NDArray[np.integer]: def get_psi(self, nside: int | None = None) -> NDArray[np.floating]: """Return polarization angles, cropped to the active TOD length.""" if self.psi_is_compressed: - psi = np.zeros(self.ntod_original, dtype=np.int64) + psi = np.zeros(self.ntod_original, dtype=self.huffman_symbols.dtype) psi = cpp_utils.huffman_decode(self.psi_compressed_u8, self.huffman_tree, self.huffman_symbols, psi) # psi is compressed as differences of digitized angle bins; first diff --git a/src/commander4/data_models/tod_view.py b/src/commander4/data_models/tod_view.py index 45d7a62..6f09c2f 100644 --- a/src/commander4/data_models/tod_view.py +++ b/src/commander4/data_models/tod_view.py @@ -21,6 +21,13 @@ class TODView: """ _ALL_GAIN_TERMS = ("abs", "rel", "temp") + # Calibration targets, mapped onto the model signals they span. Sampling a gain term against + # one of these reduces the calibration residual to (target gain term) * s_cal + noise. + _CALIB_TARGET_SIGNALS = { + "orbital_dipole": ("orbital_dipole",), + "sky": ("sky",), + "full_sky": ("sky", "orbital_dipole"), + } def __init__( self, @@ -102,6 +109,12 @@ def noise_params(self) -> NDArray: def sigma0(self) -> float: return float(self.noise_params[0]) + @property + def accept(self) -> bool: + """Whether the focused detector-scan is accepted (currently always True).""" + self._require_focus() + return bool(self.tod_samples.accept[self._iscan, self._idet]) + def get_gain(self, gain_terms: tuple[str, ...] | None = _ALL_GAIN_TERMS) -> float: """Return the selected subset of the current detector gain model.""" if gain_terms is None: @@ -215,8 +228,9 @@ def _materialize_downsampled(self, downsample_factor: int) -> Bunch: indices=np.arange(self.detector.ntod, dtype=np.int64), ) else: - # Average the jump-corrected TOD blocks, while keeping pointing and masks at the block - # centers to match the existing gain-calibration logic. + # Average the jump-corrected TOD over contiguous blocks. Pointing and masks are kept at + # the block centers; model TODs are not evaluated at this pointing but block-averaged at + # full rate (see get_static_sky_tod), so they share the data's downsampling transfer. indices_edges = np.arange(0, self.detector.ntod, factor) indices = (indices_edges[1:] + indices_edges[:-1]) // 2 ntod_down = indices.size @@ -258,31 +272,45 @@ def _require_compsep_output(self, compsep_output: NDArray | None) -> NDArray: raise ValueError("A component-separation sky map must be provided for sky subtraction.") return sky_model + def _block_average(self, tod: NDArray[np.floating], factor: int) -> NDArray[np.floating]: + """Average a full-rate array over the same contiguous blocks as the downsampled TOD.""" + ntod_down = self._materialize_downsampled(factor).tod.shape[0] + return tod[:ntod_down*factor].reshape((ntod_down, factor)).mean(axis=-1) + def get_static_sky_tod( self, compsep_output: NDArray | None = None, downsample_factor: int | None = None, ) -> NDArray[np.floating]: - """Evaluate the static sky model along the focused detector pointing.""" + """Evaluate the static sky model along the focused detector pointing. + + For ``downsample_factor > 1`` the model is evaluated at the full sampling rate and averaged + over the same sample blocks as the data, integrating the model over the scan path within + each block rather than sampling it at the block-center pixel. Model and data thereby see + the same downsampling transfer function, which keeps e.g. gain estimates unbiased. + """ factor = self._downsample_factor_or_default(downsample_factor) sky_model = self._require_compsep_output(compsep_output) - if factor == 1 and compsep_output is None: + if compsep_output is None: if self._static_sky is None: # Reuse the full-resolution sky TOD when both pointing and sky model match. self._static_sky = get_static_sky_TOD(sky_model, self.pix, psi=self.psi) - return self._static_sky - data = self._materialize_downsampled(factor) - return get_static_sky_TOD(sky_model, data.pix, psi=data.psi) + sky_tod = self._static_sky + else: + sky_tod = get_static_sky_TOD(sky_model, self.pix, psi=self.psi) + return sky_tod if factor == 1 else self._block_average(sky_tod, factor) def get_orbital_dipole_tod(self, downsample_factor: int | None = None) -> NDArray[np.floating]: - """Evaluate the orbital dipole for the focused detector.""" + """Evaluate the orbital dipole for the focused detector. + + Downsampling block-averages the full-rate dipole TOD, mirroring ``get_static_sky_tod``. + """ factor = self._downsample_factor_or_default(downsample_factor) + if self._orbital_dipole is None: + self._orbital_dipole = get_s_orb_TOD(self.detector, self.experiment_data, self.pix) if factor == 1: - if self._orbital_dipole is None: - self._orbital_dipole = get_s_orb_TOD(self.detector, self.experiment_data, self.pix) return self._orbital_dipole - data = self._materialize_downsampled(factor) - return get_s_orb_TOD(self.detector, self.experiment_data, data.pix) + return self._block_average(self._orbital_dipole, factor) def _normalize_signal_name(self, signal_name: str) -> str: """Map user-facing TOD component names onto internal canonical names.""" @@ -370,89 +398,68 @@ def _fill_masked_calibration_samples( filled[~mask] = self.get_gain(gain_terms) * signal[~mask] + noise return filled - def get_abs_calib_tod( + def get_calib_tod( self, + target_term: str, + calibrate_against: str, *, compsep_output: NDArray | None = None, downsample_factor: int | None = None, - calibrate_on_full_sky: bool | None = None, - preserve_target_gain: bool = True, fill_masked: bool = True, rng: np.random.Generator | None = None, ) -> Bunch: - """Return the residual, calibrator signal, and mask used for absolute-gain sampling.""" - factor = int(self.fsamp) if downsample_factor is None else int(downsample_factor) - data = self._materialize_downsampled(factor) - mask = self.get_mask("full", downsample_factor=factor) - s_sky = self.get_static_sky_tod(compsep_output=compsep_output, downsample_factor=factor) - s_orb = self.get_orbital_dipole_tod(downsample_factor=factor) - - if calibrate_on_full_sky is None: - calibrate_on_full_sky = self.experiment_data.nu > 380.0 + """Return the residual, calibrator signal, and mask used to sample one gain term. - if calibrate_on_full_sky: - s_cal = s_sky + s_orb - # For the full-sky branch, optionally preserve the absolute term in the residual so - # callers can choose between the current implementation and the algebraically cleaner - # target-gain-preserving form. - abs_subtract = ("rel", "temp") if preserve_target_gain else self._ALL_GAIN_TERMS - subtract = (("sky", abs_subtract), ("orbital_dipole", abs_subtract)) - else: - s_cal = s_orb - # Low-frequency absolute calibration keeps only the absolute orbital-dipole term. - subtract = (("sky", self._ALL_GAIN_TERMS), ("orbital_dipole", ("rel", "temp"))) + The detector model is ``d = (g_abs + g_rel + g_temp) * (s_sky + s_orb) + n``. To sample + ``target_term`` against a calibrator signal ``s_cal`` (the subset of {static sky, orbital + dipole} selected by ``calibrate_against``), each model signal is subtracted with the + appropriate gain terms so the residual reduces to ``g_target * s_cal + n``: + - signals making up the calibrator keep the target term (only the *other* terms are + subtracted), contributing ``g_target * s`` to the residual; + - signals outside the calibrator are subtracted in full and thus removed. - tod = self.get_tod(subtract=subtract, downsample_factor=factor, compsep_output=compsep_output) - if fill_masked: - tod = self._fill_masked_calibration_samples(tod, mask, s_cal, ("abs",), factor, rng) - return Bunch(tod=tod, s_cal=s_cal, pix=data.pix, psi=data.psi, mask=mask) - - def get_rel_calib_tod( - self, - *, - compsep_output: NDArray | None = None, - downsample_factor: int | None = None, - fill_masked: bool = True, - rng: np.random.Generator | None = None, - ) -> Bunch: - """Return the residual, calibrator signal, and mask used for relative-gain sampling.""" - factor = int(self.fsamp) if downsample_factor is None else int(downsample_factor) - data = self._materialize_downsampled(factor) - mask = self.get_mask("full", downsample_factor=factor) - s_sky = self.get_static_sky_tod(compsep_output=compsep_output, downsample_factor=factor) - s_orb = self.get_orbital_dipole_tod(downsample_factor=factor) - s_cal = s_sky + s_orb - # Relative-gain sampling removes the absolute and temporal gain terms, leaving only the - # detector-dependent residual gain multiplying the calibrator signal. - subtract = (("sky", ("abs", "temp")), ("orbital_dipole", ("abs", "temp"))) - - tod = self.get_tod(subtract=subtract, downsample_factor=factor, compsep_output=compsep_output) - if fill_masked: - tod = self._fill_masked_calibration_samples(tod, mask, s_cal, ("rel",), factor, rng) - return Bunch(tod=tod, s_cal=s_cal, pix=data.pix, psi=data.psi, mask=mask) + Args: + target_term: Gain term being sampled, one of ``_ALL_GAIN_TERMS`` ("abs", "rel", "temp"). + calibrate_against: Calibrator, one of "orbital_dipole", "full_sky", or "sky". + compsep_output: Optional sky-model override for the static-sky term. + downsample_factor: Block-averaging factor applied to both the data and the model + TODs; defaults to one second (``int(fsamp)``). + fill_masked: If True, fill masked samples with ``g_target * s_cal`` plus white noise. + rng: Optional NumPy generator for the masked-sample noise. + + Returns: + Bunch with ``tod`` (residual), ``s_cal``, ``pix``, ``psi``, and ``mask``. + """ + if target_term not in self._ALL_GAIN_TERMS: + raise ValueError(f"Unknown gain term '{target_term}'; expected one of " + f"{self._ALL_GAIN_TERMS}.") + if calibrate_against not in self._CALIB_TARGET_SIGNALS: + raise ValueError(f"Unknown calibrate_against '{calibrate_against}'; expected one of " + f"{tuple(self._CALIB_TARGET_SIGNALS)}.") - def get_temp_calib_tod( - self, - *, - compsep_output: NDArray | None = None, - downsample_factor: int | None = None, - fill_masked: bool = True, - rng: np.random.Generator | None = None, - ) -> Bunch: - """Return the residual, calibrator signal, and mask used for temporal-gain sampling.""" factor = int(self.fsamp) if downsample_factor is None else int(downsample_factor) data = self._materialize_downsampled(factor) mask = self.get_mask("full", downsample_factor=factor) s_sky = self.get_static_sky_tod(compsep_output=compsep_output, downsample_factor=factor) s_orb = self.get_orbital_dipole_tod(downsample_factor=factor) - s_cal = s_sky + s_orb - # Temporal-gain sampling removes the band-mean and detector-mean gains so the residual - # only carries the scan-dependent gain fluctuation. - subtract = (("sky", ("abs", "rel")), ("orbital_dipole", ("abs", "rel"))) - tod = self.get_tod(subtract=subtract, downsample_factor=factor, compsep_output=compsep_output) + calib_signals = self._CALIB_TARGET_SIGNALS[calibrate_against] + other_terms = tuple(t for t in self._ALL_GAIN_TERMS if t != target_term) + # Calibrator signals keep the target gain term (subtract only the others); non-calibrator + # signals are subtracted in full so they drop out of the residual entirely. + subtract = tuple((name, other_terms if name in calib_signals else self._ALL_GAIN_TERMS) + for name in ("sky", "orbital_dipole")) + s_cal = np.zeros_like(s_sky) + if "sky" in calib_signals: + s_cal = s_cal + s_sky + if "orbital_dipole" in calib_signals: + s_cal = s_cal + s_orb + + tod = self.get_tod(subtract=subtract, downsample_factor=factor, + compsep_output=compsep_output) if fill_masked: - tod = self._fill_masked_calibration_samples(tod, mask, s_cal, ("temp",), factor, rng) + tod = self._fill_masked_calibration_samples(tod, mask, s_cal, (target_term,), factor, + rng) return Bunch(tod=tod, s_cal=s_cal, pix=data.pix, psi=data.psi, mask=mask) diff --git a/src/commander4/experiments/planck/tod_reader_planck.py b/src/commander4/experiments/planck/tod_reader_planck.py index 789d74e..f3d4a93 100644 --- a/src/commander4/experiments/planck/tod_reader_planck.py +++ b/src/commander4/experiments/planck/tod_reader_planck.py @@ -151,7 +151,11 @@ def tod_reader(band_comm: MPI.Comm, my_experiment: str, my_band: Bunch, all_det_ if i_pid % 10 == 0: gc.collect() - noise_model = NoisePSDOof() + # Initialize noise model with defaults and uniform priors suited for LFI. + noise_model = NoisePSDOof(P_active_mean = [np.nan, 0.1, -1.0], + P_active_rms = [np.nan, np.inf, np.inf], + P_uni = [[np.nan, np.nan], [0.01, 0.5], [-2.5, -0.25]], + nu_fit = [[np.nan, np.nan], [0, 3.0], [0, 3.0]]) band_tod = DetGroupTOD(scan_list, expname, bandname, my_band.eval_nside, my_band.freq, my_band.fwhm, fsamp, ndet, my_band.polarization, noise_model) diff --git a/src/commander4/noise_sampling/noise_psd.py b/src/commander4/noise_sampling/noise_psd.py index 90c5609..42852f3 100644 --- a/src/commander4/noise_sampling/noise_psd.py +++ b/src/commander4/noise_sampling/noise_psd.py @@ -1,10 +1,32 @@ import logging from typing import Optional import numpy as np +import pixell +from scipy.fft import rfftfreq from numpy.typing import NDArray +from commander4.utils.math_operations import forward_rfft logger = logging.getLogger(__name__) + +def _inversion_sampler_1d(lnL: NDArray, grid_points: NDArray) -> float: + """ Performs 1D inversion sampling on a grid. This involves calculating the cumulative + log-likelihood, normalizing this to be contained in [0,1], drawing a random number in [0,1], + and matching that to a (interpolated) grid point. + Args: + lnL (np.ndarray): Array of log-likelihood values at each grid point. + grid_points (np.ndarray): The corresponding parameter values for each grid point. + Returns: + sample (float): A single random sample drawn from the provided distribution. + """ + lnL -= np.max(lnL) + L = np.exp(lnL) # Calculate the linear likelihood. + cdf = np.cumsum(L) # Cumulative likelihood. + cdf /= cdf[-1] # Constrain it to [0,1]. + u = np.random.uniform(0, 1) + sample = np.interp(u, cdf, grid_points) # Find the x-value that matches the y-value we drew. + return sample + # =================================================================== # Base class # =================================================================== @@ -22,10 +44,11 @@ class NoisePSD: P_active: Informative prior ``[mean, rms]`` per parameter, shape (npar, 2). P_lognorm: If True the informative prior is log-normal, else Gaussian, shape (npar,). nu_fit: Frequency range ``[f_min, f_max]`` (Hz) for fitting, shape (npar, 2). - apply_filter: Whether to multiply output by a modulation filter spline. + is_white: True for purely white models (enables an FFT-free fast path in ``apply_N_inv``). """ param_names: tuple[str, ...] = () + is_white: bool = False # Override to True in purely white-noise subclasses. _jit_model_id: int = -1 # Numba model ID; -1 = Python fallback def __init__(self, @@ -53,26 +76,52 @@ def npar(self) -> int: return len(self.params) # ---- interface (override in subclasses) -------------------------------- - def eval_full(self, freqs: NDArray) -> NDArray: + # All models receive the noise parameters explicitly (rather than reading ``self.params``) + # because the live per-scan/per-detector values are stored in ``TODSamples.noise_params``, + # while a single shared model instance is attached to each ``DetGroupTOD``. + def eval_full(self, freqs: NDArray, noise_params: NDArray) -> NDArray: """Evaluate the *full* PSD (white + correlated) at each frequency in *freqs* (Hz).""" raise NotImplementedError - def eval_corr(self, freqs: NDArray) -> NDArray: + def eval_corr(self, freqs: NDArray, noise_params: NDArray) -> NDArray: """Evaluate the *correlated-only* PSD at each frequency in *freqs* (Hz).""" raise NotImplementedError - def compute_inv_corr_spectrum(self, freqs: NDArray) -> NDArray: - """Return sigma0^2 / P_corr(f) for each frequency. + def compute_inv_corr_spectrum(self, freqs: NDArray, noise_params: NDArray) -> NDArray: + """Return the inverse correlated-noise power spectrum 1 / P_corr(f) for each frequency. - Used to build ``C_corr_inv`` in the noise-sampling CG. - Frequencies where ``P_corr <= 0`` get ``inv = 0``. + This is the quantity added to the inverse white-noise level in the correlated-noise CG + (see ``corr_noise_realization_with_gaps``), so it carries units of inverse variance. + Frequencies where ``P_corr <= 0`` (e.g. the zero-frequency mode) get ``inv = 0``. """ - P_corr = self.eval_corr(freqs) + P_corr = self.eval_corr(freqs, noise_params) out = np.zeros_like(P_corr) good = P_corr > 0 - out[good] = self.params[0] ** 2 / P_corr[good] + out[good] = 1.0 / P_corr[good] return out + def sample_params(self, residual: NDArray, noise_params: NDArray, fsamp: float, *, + nu_min: float = 0.0, nu_max: float = np.inf, bin_psd: bool = False) -> NDArray: + """Draw a new sample of the model parameters (except ``sigma0`` at index 0). + + The fit is performed against the periodogram of the sky-subtracted *residual* TOD (with + masked gaps inpainted by the caller using the correlated-noise realization plus white + noise), fitting the *full* PSD model. This mirrors the Commander3 ``sample_noise_psd`` + routine and keeps the Markov-chain correlation length short relative to fitting the + drawn ``n_corr`` realization directly. + + Args: + residual: Sky-subtracted residual TOD (white + correlated noise), gaps inpainted. + noise_params: Current parameter values; ``noise_params[0]`` (sigma0) is held fixed. + fsamp: Sampling rate of *residual* (Hz). + nu_min, nu_max: Frequency range (Hz) over which to fit the PSD parameters. + bin_psd: If True, fit a log-binned periodogram (each bin Whittle-weighted by its mode + count); if False, fit every Fourier mode in range (default). + Returns: + A new ``noise_params`` array with the updated parameter values. + """ + raise NotImplementedError + # ---- prior evaluation ------------------------------------------------- def log_prior(self, param_idx: int, x: float) -> float: """Evaluate the log-prior for parameter *param_idx* at value *x*. @@ -111,10 +160,10 @@ class NoisePSDOof(NoisePSD): param_names = ('sigma0', 'fknee', 'alpha') def __init__(self, - P_active_mean = [np.nan, 0.1, -1.0], + P_active_mean = [np.nan, 10.0, -2.7], P_active_rms = [np.nan, np.inf, np.inf], - P_uni = [[np.nan, np.nan], [0.01, 0.5], [-2.5, -0.25]], - nu_fit = [[np.nan, np.nan], [0, 3.0], [0, 3.0]], + P_uni = [[np.nan, np.nan], [1.0, 100.0], [-4.5, -1.0]], + nu_fit = [[np.nan, np.nan], [0, 10.0], [0, 10.0]], **kw): P_lognorm = np.array([False, True, False]) # sigma0, fknee, alpha super().__init__(P_active_mean=P_active_mean[:3], P_active_rms=P_active_rms[:3], @@ -132,4 +181,61 @@ def eval_corr(self, freqs: NDArray, noise_params: NDArray) -> NDArray: vals = np.zeros(len(freqs), dtype=np.float64) pos = freqs > 0 vals[pos] = s0 ** 2 * (freqs[pos] / fk) ** a - return vals \ No newline at end of file + return vals + + def sample_params(self, residual: NDArray, noise_params: NDArray, fsamp: float, *, + nu_min: float = 0.0, nu_max: float = np.inf, bin_psd: bool = False, + n_grid: int = 150, n_burnin: int = 4) -> NDArray: + """ Draw a sample of (fknee, alpha) for the 1/f model by fitting the *full* PSD + P(f) = sigma0^2 (1 + (f/fknee)^alpha) to the periodogram of the sky-subtracted + *residual* TOD, with sigma0 (= noise_params[0]) held fixed. Uses a short Gibbs loop + that alternately inversion-samples fknee and alpha on grids spanning the model's + uniform priors (``P_uni``). See the base-class docstring for argument meaning. + Returns: + A new [sigma0, fknee, alpha] array with updated fknee and alpha. + """ + sigma0_sq = float(noise_params[0])**2 + alpha_current = float(noise_params[2]) + fknee_current = float(noise_params[1]) + Ntod = len(residual) + freqs = rfftfreq(Ntod, 1.0/fsamp) + power = (1.0 / Ntod) * np.abs(forward_rfft(residual))**2 + + # Restrict to the requested fit range, always excluding the zero-frequency (DC) mode. + in_fit = (freqs > 0) & (freqs >= nu_min) & (freqs <= nu_max) + f = freqs[in_fit] + p = power[in_fit] + if bin_psd: + # Log-bin the periodogram; each bin is Whittle-weighted by its number of modes. + bins = pixell.utils.expbin(f.size, nbin=100, nmin=1) + weight = (bins[:, 1] - bins[:, 0]).astype(np.float64) + f = pixell.utils.bin_data(bins, f) + p = pixell.utils.bin_data(bins, p) + else: + weight = np.ones(f.size, dtype=np.float64) + log_p = np.log(p) + w = weight[:, np.newaxis] + + # Grids span the uniform priors (single source of truth for the hard parameter bounds). + fk_lo, fk_hi = float(self.P_uni[1, 0]), float(self.P_uni[1, 1]) + al_lo, al_hi = float(self.P_uni[2, 0]), float(self.P_uni[2, 1]) + fknee_grid = np.logspace(np.log10(fk_lo), np.log10(fk_hi), n_grid) + alpha_grid = np.linspace(al_lo, al_hi, n_grid) + + for _ in range(n_burnin + 1): + # Whittle log-likelihood sum_l w_l (log p_l - log S_l - p_l/S_l), with the full model + # S = sigma0^2 (1 + (f/fknee)^alpha). 1. Sample fknee given the current alpha. + S = sigma0_sq * (1.0 + (f[:, np.newaxis] / fknee_grid) ** alpha_current) + resid = log_p[:, np.newaxis] - np.log(S) + log_L_fknee = np.sum(w * (resid - np.exp(resid)), axis=0) + fknee_current = float(_inversion_sampler_1d(log_L_fknee, fknee_grid)) + # 2. Sample alpha given the new fknee. + S = sigma0_sq * (1.0 + (f[:, np.newaxis] / fknee_current) ** alpha_grid) + resid = log_p[:, np.newaxis] - np.log(S) + log_L_alpha = np.sum(w * (resid - np.exp(resid)), axis=0) + alpha_current = float(_inversion_sampler_1d(log_L_alpha, alpha_grid)) + + out = np.array(noise_params, dtype=np.float64, copy=True) + out[1] = fknee_current + out[2] = alpha_current + return out \ No newline at end of file diff --git a/src/commander4/noise_sampling/noise_sampling.py b/src/commander4/noise_sampling/noise_sampling.py index cc1fd5a..a27c879 100644 --- a/src/commander4/noise_sampling/noise_sampling.py +++ b/src/commander4/noise_sampling/noise_sampling.py @@ -1,89 +1,6 @@ import numpy as np -from scipy.fft import rfftfreq -import pixell from numba import njit from numpy.typing import NDArray -from commander4.utils.math_operations import forward_rfft, backward_rfft, forward_rfft_mirrored,\ - backward_rfft_mirrored - -def _inversion_sampler_1d(lnL: NDArray, grid_points: NDArray) -> float: - """ Performs 1D inversion sampling on a grid. This involves calculating the cumulative - log-likelihood, normalizing this to be contained in [0,1], drawing a random number in [0,1], - and matching that to a (interpolated) grid point. - Args: - lnL (np.ndarray): Array of log-likelihood values at each grid point. - grid_points (np.ndarray): The corresponding parameter values for each grid point. - Returns: - sample (float): A single random sample drawn from the provided distribution. - """ - lnL -= np.max(lnL) - L = np.exp(lnL) # Calculate the linear likelihood. - cdf = np.cumsum(L) # Cumulative likelihood. - cdf /= cdf[-1] # Constrain it to [0,1]. - u = np.random.uniform(0, 1) - sample = np.interp(u, cdf, grid_points) # Find the x-value that matches the y-value we drew. - return sample - - -def sample_noise_PS_params(n_corr: NDArray, sigma0: float, f_samp: float, alpha_start=-1.0, - freq_max=3.0, n_grid=100, n_burnin=5) -> tuple[float, float]: - """ Function for drawing a sample of the fknee and alpha parameters for the correlated noise - under the power spectrum data model PS = sigma0*(f/fknee)**alpha, where sigma0 is known. - Note that this relates *only* to the correlated noise, without the "flat" white noise. - Args: - n_corr (np.array): 1D array of the correlated noise time series. - sigma0 (float): White noise level of full data. Since this data model does not have a - white noise floor, this essentially just scales the resulting fnee value. - f_samp (float): The sampling rate of the data (n_corr), in Hertz. - alpha_start (float): Starting guess for alpha in Gibbs sampler. - freq_max (float): Maximum frequency to consider for the PS (all 1/f information is - contained at low freqs). - n_grid (int): Number of grid points used for the inverse sampling. - n_burnin (int): Number of burn-in samples before drawing the "actual" - sample of fknee and alpha. - Returns: - fknee_sample (float): A single sample of fknee. - alpha_sample (float): A single sample of alpha. - """ - Ntod = len(n_corr) - freqs = rfftfreq(Ntod, 1.0/f_samp)[1:] # [1:] to Exclude freq=0 mode (same on line below). - n_corr_power = (1.0 / Ntod) * np.abs(forward_rfft(n_corr))[1:]**2 - Nrfft = freqs.size - bins = pixell.utils.expbin(Nrfft, nbin=100, nmin=1) - binned_freqs = pixell.utils.bin_data(bins, freqs) - binned_n_corr_power = pixell.utils.bin_data(bins, n_corr_power) - freq_mask = (binned_freqs <= freq_max) - log_freqs_masked = np.log(binned_freqs[freq_mask]) - n_corr_power_masked = binned_n_corr_power[freq_mask] - log_n_corr_power = np.log(n_corr_power_masked) - # Set up a grid of possible fknee values: from 2 times the min frequency to the max frequency. - fknee_grid = np.logspace(np.log10(freqs[0] * 2), np.log10(freq_max), n_grid) - log_fknee_grid = np.log(fknee_grid) - - alpha_grid = np.linspace(-2.5, -0.25, n_grid) - alpha_current = alpha_start - - log_sigma0_sq = np.log(sigma0**2) - # --- Main Gibbs Loop --- - for _ in range(n_burnin + 1): - # 1. Sample f_knee, given a fixed alpha - log_N_corr_ps = log_sigma0_sq + alpha_current *\ - (log_freqs_masked[:, np.newaxis] - log_fknee_grid) - residual = log_n_corr_power[:, np.newaxis] - log_N_corr_ps - log_L_fknee = np.sum(residual - np.exp(residual), axis=0) - # A faster but slightly less statistically robust way of calculating the likelihood - # (~30% speedup of code): - # log_L_fknee = -0.5 * np.sum((log_n_corr_power[:, np.newaxis] - log_N_corr_ps)**2, axis=0) - fknee_sample = float(_inversion_sampler_1d(log_L_fknee, fknee_grid)) - - log_fknee_sample = np.log(fknee_sample) - log_N_corr_ps = log_sigma0_sq + alpha_grid\ - * (log_freqs_masked[:, np.newaxis] - log_fknee_sample) - residual = log_n_corr_power[:, np.newaxis] - log_N_corr_ps - log_L_alpha = np.sum(residual - np.exp(residual), axis=0) - # log_L_alpha = -0.5 * np.sum((log_n_corr_power[:, np.newaxis] - log_N_corr_ps)**2, axis=0) - alpha_current = float(_inversion_sampler_1d(log_L_alpha, alpha_grid)) - return fknee_sample, alpha_current @njit(fastmath=True) diff --git a/src/commander4/noise_sampling/sample_ncorr.py b/src/commander4/noise_sampling/sample_ncorr.py index 0a1a8c0..00352d5 100644 --- a/src/commander4/noise_sampling/sample_ncorr.py +++ b/src/commander4/noise_sampling/sample_ncorr.py @@ -1,8 +1,20 @@ +import logging import numpy as np import pixell +from scipy.fft import rfftfreq +from mpi4py import MPI from numpy.typing import NDArray +from pixell.bunch import Bunch from commander4.utils.math_operations import forward_rfft, backward_rfft, forward_rfft_mirrored,\ backward_rfft_mirrored +from commander4.noise_sampling.noise_sampling import fill_all_masked +from commander4.noise_sampling.noise_psd import NoisePSD +from commander4.noise_sampling.sigma0 import calc_sigma0_robust + +from commander4.logging.performance_logger import benchmark, bench_summary, start_bench,\ + stop_bench, log_memory, increment_count, bench_reset + +logger = logging.getLogger(__name__) def corr_noise_realization_with_gaps(TOD: NDArray, mask: NDArray[np.bool_], sigma0: float, @@ -28,7 +40,10 @@ def corr_noise_realization_with_gaps(TOD: NDArray, mask: NDArray[np.bool_], sigm x_final (np.array): The TOD realization of the correlated noise. """ def apply_filter(vec, Fourier_filter): - return backward_rfft_mirrored(forward_rfft_mirrored(vec) * Fourier_filter, ntod=len(vec)) + start_bench("FFT") + res = backward_rfft_mirrored(forward_rfft_mirrored(vec) * Fourier_filter, ntod=len(vec)) + stop_bench("FFT") + return res def apply_LHS_scaled(x_small): u_x = np.zeros(Ntod, dtype=x_small.dtype) @@ -133,4 +148,171 @@ def apply_LHS(x_full): else: break x_full = CG_solver.x - return x_full \ No newline at end of file + return x_full + + +def _estimate_sigma0(residual: NDArray, n_corr: NDArray, mask: NDArray[np.bool_], + dec: int) -> float: + """Robust white-noise level from the fully-cleaned residual (sky- and n_corr-subtracted).""" + return float(calc_sigma0_robust(residual - n_corr, mask, down_factor=int(dec))) + + +def sample_correlated_noise(tod: NDArray, mask: NDArray[np.bool_], noise_params: NDArray, + noise_model: NoisePSD, fsamp: float, *, cg_err_tol: float, + cg_max_iter: int, sample_params: bool, sample_sigma0: bool = True, + nomono: bool = False, onlymono: bool = False, sigma0_dec: int = 1, + psd_fit_nu_min: float = 0.0, psd_fit_nu_max: float = np.inf, + psd_bin: bool = False) -> Bunch: + """ Draw a correlated-noise realization for one detector-scan and optionally resample sigma0 and + the noise-model parameters. The inverse correlated-noise spectrum is supplied by + *noise_model*, so any NoisePSD subclass (parameters of any length) can be plugged in. + + Masked gaps are inpainted in-place in *tod* (linear slope + white noise). The filled values + seed the CG warm-start and, when the masked CG is skipped or rejected, define the stationary + (full-mask) Wiener fallback solution. + + Args: + tod: Sky-subtracted TOD for one detector-scan. Modified in-place by gap inpainting. + mask: Boolean validity mask (True = valid sample). + noise_params: Current noise parameters; ``noise_params[0]`` is sigma0 (used for the CG). + noise_model: NoisePSD model providing ``compute_inv_corr_spectrum`` and ``sample_params``. + fsamp: Sampling rate of the TOD (Hz). + cg_err_tol: CG convergence tolerance (relative residual). + cg_max_iter: Maximum CG iterations. If 0, the masked-gap CG is skipped entirely and the + stationary (full-mask) Wiener solution is returned directly. + sample_params: Whether to resample the noise-model parameters (fknee, alpha, ...). + sample_sigma0: Whether to re-estimate sigma0 (noise_params[0]) from the n_corr-subtracted + residual after drawing n_corr (Commander3-aligned). + nomono: If True, project the per-scan monopole out of the residual and of n_corr (Fortran + ``nomono``); otherwise the DC is left in n_corr. + onlymono: If True, model the correlated noise as only the per-scan offset, skipping the CG + and parameter sampling (Fortran ``onlymono``). Takes precedence over ``nomono``. + sigma0_dec: Decimation (block-average) factor for the sigma0 estimator. + psd_fit_nu_min, psd_fit_nu_max: Frequency range (Hz) for PSD-parameter fitting. + psd_bin: Whether the PSD-parameter fit uses a (mode-count-weighted) binned periodogram. + Returns: + Bunch with fields ``n_corr`` (realization), ``noise_params`` (with updated sigma0 and/or + parameters), ``residual`` (CG residual; 0 when no masked CG ran), ``niter`` (CG iterations), + ``converged`` (bool), and ``high_var`` (variance sanity check failed). + """ + noise_params = np.array(noise_params, dtype=np.float64, copy=True) + sigma0 = float(noise_params[0]) + Ntod = tod.shape[0] + + # "Only monopole" mode: model the correlated noise as just the per-scan offset. + if onlymono: + mono = float(np.mean(tod[mask])) if mask.any() else 0.0 + n_corr = np.full(Ntod, mono, dtype=tod.dtype) + if sample_sigma0: + noise_params[0] = _estimate_sigma0(tod, n_corr, mask, sigma0_dec) + return Bunch(n_corr=n_corr, noise_params=noise_params, residual=0.0, niter=0, + converged=True, high_var=False) + + freq = rfftfreq(2 * Ntod, d=1.0/fsamp) # Mirrored-FFT grid: nfft=2*Ntod -> length Ntod+1. + C_corr_inv = noise_model.compute_inv_corr_spectrum(freq, noise_params) + # Inpaint masked regions: seeds the CG warm-start and feeds the stationary fallback solve. + fill_all_masked(tod, mask, sigma0) + if nomono and mask.any(): + tod = tod - np.mean(tod[mask]) # Solve for a mean-zero correlated noise component. + + high_var = False + if cg_max_iter == 0: + # User requested no CG steps: use the stationary (full-mask) Wiener solution directly. + n_corr, residual, niter, converged = corr_noise_realization_with_gaps( + tod, np.ones_like(mask), sigma0, C_corr_inv) + else: + n_corr, residual, niter, converged = corr_noise_realization_with_gaps( + tod, mask, sigma0, C_corr_inv, err_tol=cg_err_tol, max_iter=cg_max_iter) + # Sanity check: the residual (data minus n_corr) should not carry more power than the data. + resid = (tod - n_corr) * mask + high_var = bool(np.dot(resid, resid) > np.dot(tod*mask, tod*mask)) + if high_var or not converged: + # Fall back to the stationary solution that ignores the gaps. + n_corr, _, _, _ = corr_noise_realization_with_gaps( + tod, np.ones_like(mask), sigma0, C_corr_inv) + + if nomono and mask.any(): + n_corr = n_corr - np.mean(n_corr[mask]) + + # Re-estimate sigma0 from the fully-cleaned residual (after subtracting the n_corr realization). + if sample_sigma0: + noise_params[0] = _estimate_sigma0(tod, n_corr, mask, sigma0_dec) + + # Fit the PSD parameters to the residual periodogram (full model), inpainting masked samples + # with the n_corr realization plus white noise (Commander3 sample_noise_psd convention). + if sample_params: + residual_tod = tod.copy() + ngap = int(np.count_nonzero(~mask)) + if ngap > 0: + residual_tod[~mask] = n_corr[~mask] + float(noise_params[0]) * np.random.randn(ngap) + noise_params = noise_model.sample_params(residual_tod, noise_params, fsamp, + nu_min=psd_fit_nu_min, nu_max=psd_fit_nu_max, + bin_psd=psd_bin) + + return Bunch(n_corr=n_corr, noise_params=noise_params, residual=float(residual), + niter=int(niter), converged=bool(converged), high_var=high_var) + + +def _log_distribution(nu: float, label: str, values: NDArray, fmt: str = ".4f") -> None: + """Log the min / 1st-pct / mean / 99th-pct / max of *values* for one band.""" + values = np.asarray(values, dtype=np.float64) + logger.info(f"{nu}GHz: {label} {np.nanmin(values):{fmt}} {np.nanpercentile(values, 1):{fmt}} " + f"{np.nanmean(values):{fmt}} {np.nanpercentile(values, 99):{fmt}} " + f"{np.nanmax(values):{fmt}}") + + +def log_corr_noise_stats(band_comm: MPI.Comm, nu: float, noise_model: NoisePSD, + sampled_params: list[NDArray], residuals: list[float], + niters: list[int], n_failed_conv: int, n_high_var: int, + worst_residual: float, n_local_scans: int) -> None: + """ Reduce per-detector-scan correlated-noise diagnostics across the band communicator and log + a summary on the band master. Parameter distributions are reported per model parameter name + (skipping sigma0), so the summary adapts to any NoisePSD model. + + Args: + band_comm: Band-level MPI communicator. + nu: Band centre frequency (GHz), for labelling. + noise_model: The band's NoisePSD model (used for ``param_names``). + sampled_params: Locally sampled ``noise_params`` arrays (empty if params were not sampled). + residuals: Local CG residuals (0 entries are excluded from the summary). + niters: Local CG iteration counts. + n_failed_conv: Local count of non-converged masked-CG solves. + n_high_var: Local count of variance-sanity-check failures. + worst_residual: Local worst CG residual. + n_local_scans: Local number of detector-scans (for the "out of N" message). + """ + n_failed_conv = band_comm.reduce(n_failed_conv, op=MPI.SUM) + n_high_var = band_comm.reduce(n_high_var, op=MPI.SUM) + worst_residual = band_comm.reduce(worst_residual, op=MPI.MAX) + n_total = band_comm.reduce(n_local_scans, op=MPI.SUM) + residuals = band_comm.gather(residuals) + niters = band_comm.gather(niters) + sampled_params = band_comm.gather(sampled_params) + if band_comm.Get_rank() != 0: + return + + logger.debug(f"Worst corr-noise sampling residual (band {nu}GHz) = {worst_residual:.2e}.") + if n_failed_conv > 0: + logger.warning(f"Band {nu}GHz failed noise CG for {n_failed_conv} out of {n_total} scans. " + f"Worst residual = {worst_residual:.3e}.") + if n_high_var > 0: + logger.warning(f"Band {nu}GHz failed variance sanity check for {n_high_var} out of " + f"{n_total} scans.") + + residuals = np.concatenate([np.asarray(r, dtype=np.float64) for r in residuals]) + residuals = residuals[residuals != 0] + if residuals.size == 0: + residuals = np.array([0.0]) + niters = np.concatenate([np.asarray(n, dtype=np.float64) for n in niters]) + if niters.size == 0: + niters = np.array([0.0]) + _log_distribution(nu, "residuals", residuals, fmt=".2e") + _log_distribution(nu, "iterations", niters, fmt=".4f") + + # Per-parameter distributions (model-agnostic; sigma0 at index 0 is reported elsewhere). + flat = [p for sub in sampled_params for p in sub] + if flat: + arr = np.asarray(flat, dtype=np.float64) # shape (n_sampled_scans, npar) + for j in range(1, arr.shape[1]): + name = noise_model.param_names[j] if j < len(noise_model.param_names) else f"p{j}" + _log_distribution(nu, name, arr[:, j], fmt=".4f") \ No newline at end of file diff --git a/src/commander4/output/log.py b/src/commander4/output/log.py index 2d2f354..0af91fe 100644 --- a/src/commander4/output/log.py +++ b/src/commander4/output/log.py @@ -15,7 +15,16 @@ logging.VERBOSE = VERBOSE -class ColorFormatter(logging.Formatter): +class C4Formatter(logging.Formatter): + """Formatter that strips the leading 'commander4.' from logger names, so that + e.g. 'commander4.tod_processing' is displayed as just 'tod_processing'.""" + + def format(self, record: logging.LogRecord) -> str: + record.name = record.name.removeprefix('commander4.') + return super().format(record) + + +class ColorFormatter(C4Formatter): """ Logging formatter that prepends ANSI color codes to console output based on log level. @@ -100,7 +109,8 @@ def level_val(name): return getattr(logging, name, 20) 'formatters': { 'standard': { 'style': '{', - 'format': '{asctime} - {name} - {levelname} - {message}' + 'format': '{asctime} - {name} - {levelname} - {message}', + 'datefmt': '%H:%M:%S' # Time only; drops date and sub-second precision. }, }, 'handlers': {}, @@ -142,18 +152,20 @@ def level_val(name): return getattr(logging, name, 20) logging.config.dictConfig(config_dict) - # Apply ColorFormatter to the console handler manually. dictConfig's '()' - # factory passes keyword arguments directly to the constructor, but - # logging.Formatter.__init__ uses 'fmt' not 'format', so the factory path - # silently falls back to an unconfigured formatter. Setting it here avoids - # that problem entirely. - if 'console' in logger_params: - _fmt = '{asctime} - {name} - {levelname} - {message}' - _color_fmt = ColorFormatter(fmt=_fmt, style='{') - for _handler in logging.root.handlers: - if isinstance(_handler, logging.StreamHandler) and not isinstance(_handler, logging.FileHandler): - _handler.setFormatter(_color_fmt) - break + # Apply ColorFormatter to the console handler, and the (non-color) C4Formatter + # to the file handler, both manually. This is also where prefix-stripping is + # wired in for the file handler, since dictConfig's 'standard' formatter is a + # plain logging.Formatter. dictConfig's '()' factory passes keyword arguments + # directly to the constructor, but logging.Formatter.__init__ uses 'fmt' not + # 'format', so the factory path silently falls back to an unconfigured + # formatter; setting the formatters here avoids that problem entirely. + _fmt = '{asctime} - {name} - {levelname} - {message}' + _datefmt = '%H:%M:%S' # Time only; drops date and sub-second precision. + for _handler in logging.root.handlers: + if isinstance(_handler, logging.FileHandler): + _handler.setFormatter(C4Formatter(fmt=_fmt, datefmt=_datefmt, style='{')) + elif isinstance(_handler, logging.StreamHandler): + _handler.setFormatter(ColorFormatter(fmt=_fmt, datefmt=_datefmt, style='{')) # Configure logging to redirect warnings from py.warning. Note that this will *prevent* these # from being sent to sys.stderr, to avoid duplication. diff --git a/src/commander4/simulations/inplace_litebird_sim.py b/src/commander4/simulations/inplace_litebird_sim.py index 482eb82..2f75eba 100644 --- a/src/commander4/simulations/inplace_litebird_sim.py +++ b/src/commander4/simulations/inplace_litebird_sim.py @@ -13,11 +13,18 @@ from commander4.data_models.detector_TOD import DetectorTOD from commander4.data_models.detector_group_TOD import DetGroupTOD -from commander4.sky_models.component import ThermalDust, Synchrotron, FreeFree +from commander4.sky_models.component import ThermalDust, Synchrotron, FreeFree, SpinningDust from commander4.logging.performance_logger import benchmark, bench_summary, start_bench,\ stop_bench, log_memory, increment_count, bench_reset +def _scalar_nu_ref(comp_params): + """Resolve a component's `nu_ref` (scalar or [nu_I, nu_QU] list) to a single scalar (the I + value), since the in-place simulator uses one reference frequency per component.""" + nu_ref = comp_params.nu_ref + return nu_ref[0] if isinstance(nu_ref, (list, tuple)) else nu_ref + + def generate_cmb(freq, fwhm, units, nside, lmax, params): H0 = 67.5 ombh2 = 0.022 @@ -85,7 +92,7 @@ def generate_cmb(freq, fwhm, units, nside, lmax, params): def generate_thermal_dust(freq, fwhm, units, nside, params): - nu_dust = params.components.ThermalDust.params.nu0 + nu_dust = _scalar_nu_ref(params.components.ThermalDust.params) #d0 = constant beta 1.54 and T = 20 dust = pysm3.Sky(nside=min(1024,nside), preset_strings=["d0"], output_unit=units) @@ -93,9 +100,10 @@ def generate_thermal_dust(freq, fwhm, units, nside, params): dust_ref_smoothed = hp.smoothing(dust_ref, fwhm=fwhm)*dust_ref.unit dust_params = deepcopy(params.components.ThermalDust.params) + dust_params.nu_ref = nu_dust # single reference frequency for the in-place sim dust_params.polarized = True - dust = ThermalDust(dust_params, params) + dust = ThermalDust(dust_params, params, comp_name="ThermalDust") dust_s = dust_ref_smoothed*dust.get_sed(freq) dust_s = hp.ud_grade(dust_s.value, nside)*dust_s.unit @@ -103,7 +111,7 @@ def generate_thermal_dust(freq, fwhm, units, nside, params): def generate_sync(freq, fwhm, units, nside, params): - nu_sync = params.components.Synchrotron.params.nu0 + nu_sync = _scalar_nu_ref(params.components.Synchrotron.params) # s5 = const beta -3.1 sync = pysm3.Sky(nside=min(1024,nside), preset_strings=["s5"], output_unit=units) @@ -111,9 +119,10 @@ def generate_sync(freq, fwhm, units, nside, params): sync_ref_smoothed = hp.smoothing(sync_ref, fwhm=fwhm)*sync_ref.unit sync_params = deepcopy(params.components.Synchrotron.params) + sync_params.nu_ref = nu_sync # single reference frequency for the in-place sim sync_params.polarized = True - sync = Synchrotron(sync_params, params) + sync = Synchrotron(sync_params, params, comp_name="Synchrotron") sync_s = sync_ref_smoothed*sync.get_sed(freq) sync_s = hp.ud_grade(sync_s.value, nside)*sync_s.unit @@ -121,16 +130,17 @@ def generate_sync(freq, fwhm, units, nside, params): def generate_ff(freq, fwhm, units, nside, params): - nu_ff = params.components.FreeFree.params.nu0 + nu_ff = _scalar_nu_ref(params.components.FreeFree.params) ff = pysm3.Sky(nside=min(1024,nside), preset_strings=["f1"], output_unit=units) ff_ref = ff.get_emission(nu_ff*u.GHz) ff_ref_smoothed = hp.smoothing(ff_ref, fwhm=fwhm)*ff_ref.unit ff_params = deepcopy(params.components.FreeFree.params) + ff_params.nu_ref = nu_ff # single reference frequency for the in-place sim ff_params.polarized = False - ff = FreeFree(ff_params, params) + ff = FreeFree(ff_params, params, comp_name="FreeFree") ff_s = ff_ref_smoothed*ff.get_sed(freq) ff_s = hp.ud_grade(ff_s.value, nside)*ff_s.unit @@ -147,7 +157,7 @@ def generate_spdust(freq, fwhm, units, nside, params): spdust_params = deepcopy(params.components.SpinningDust.params) spdust_params.polarized = False - spdust = Synchrotron(spdust_params, params) + spdust = SpinningDust(spdust_params, params, comp_name="SpinningDust") spdust_s = spdust_ref_smoothed*spdust.get_sed(freq) spdust_s = hp.ud_grade(spdust_s.value, nside)*spdust_s.unit diff --git a/src/commander4/tod_processing.py b/src/commander4/tod_processing.py index ad3744e..d7226b5 100644 --- a/src/commander4/tod_processing.py +++ b/src/commander4/tod_processing.py @@ -7,10 +7,7 @@ import time from numpy.typing import NDArray -import healpy as hp from pixell.bunch import Bunch -from astropy.io import fits -import pysm3.units as pysm3_u from commander4.output.log import logassert from commander4.data_models.detector_map import DetectorMap @@ -21,8 +18,7 @@ from commander4.utils.mapmaker import MapmakerIQU, WeightsMapmakerIQU, WeightsMapmaker, Mapmaker from commander4.utils.CG_mapmaker import CGMapmakerI, CGMapmakerIQU from commander4.solvers.preconditioners import InvNPreconditionerI, InvNPreconditionerIQU -from commander4.noise_sampling.noise_sampling import sample_noise_PS_params, fill_all_masked -from commander4.noise_sampling.sample_ncorr import corr_noise_realization_with_gaps +from commander4.noise_sampling.sample_ncorr import sample_correlated_noise, log_corr_noise_stats from commander4.utils.math_operations import forward_rfft, backward_rfft from commander4.utils.execution_ids import get_execution_band_ids from commander4.noise_sampling.sigma0 import calc_sigma0_robust @@ -34,33 +30,81 @@ logger = logging.getLogger(__name__) -def get_initial_sky(experiment_data: DetGroupTOD) -> NDArray[np.float32]: - """ Returns a sky realization from a set of components. The set of components are listed in - the provided DetGroupTOD object, originally specified in the parameter file. - """ - initial_sky = np.zeros((3, 12*experiment_data.nside**2), dtype=np.float32) - for skyfile in experiment_data.sky_init_files: - with fits.open(skyfile) as hdul: - fields = ["TEMPERATURE", "Q_POLARISATION", "U_POLARISATION"] - for i, field in enumerate(fields): - data = hdul[1].data[field].flatten() - nside = hp.npix2nside(data.size) - if nside != experiment_data.nside: - data = hp.ud_grade(data, experiment_data.nside) - initial_sky[i] += data - - # Convert from uK_CMB to uK_RJ - initial_sky *= (1*pysm3_u.uK_CMB).to(pysm3_u.uK_RJ, - equivalencies=pysm3_u.cmb_equivalencies(experiment_data.nu*pysm3_u.GHz)).value - return initial_sky - def called_on_non_master(arr): logger.debug("Dummy precond has been called") return np.copy(arr) + +def _binned_tod_power_spectrum(tod: NDArray, fsamp: float, nbin: int) -> tuple[NDArray, NDArray]: + """ Log-binned periodogram of a TOD, for low-resolution diagnostics written to the chain. + + Computes the one-sided periodogram ``|rfft(tod)|^2 / Ntod`` on the natural frequency grid, + then averages it into exponentially spaced bins (``pixell.utils.expbin`` with ``nmin=1``). + expbin returns at most ``nbin`` bins (fewer for short TODs), so the binned frequencies and + power are returned padded to length ``nbin`` with NaN, giving a fixed width for the + per-detector-scan chain arrays. + Args: + tod (NDArray): Time-ordered data (any units; e.g. the raw TOD or an n_corr realization). + fsamp (float): Sampling rate (Hz). + nbin (int): Fixed output length (the maximum number of bins). + Returns: + (freqs, power): Each a length-``nbin`` array, NaN-padded beyond the actual bin count. + """ + ntod = len(tod) + freqs = rfftfreq(ntod, 1.0 / fsamp) + power = (1.0 / ntod) * np.abs(forward_rfft(tod)) ** 2 + bins = pixell.utils.expbin(freqs.size, nbin=nbin, nmin=1) + nb = bins.shape[0] + freqs_binned = np.full(nbin, np.nan, dtype=np.float64) + power_binned = np.full(nbin, np.nan, dtype=np.float64) + freqs_binned[:nb] = pixell.utils.bin_data(bins, freqs) + power_binned[:nb] = pixell.utils.bin_data(bins, power) + return freqs_binned, power_binned + + +def _record_tod_diagnostics(tod_samples: TODSamples, iscan: int, idet: int, view: TODView, + n_corr: NDArray | None) -> None: + """ Record per-detector-scan TOD diagnostics into the chain arrays. + + Stores the low-resolution log-binned power spectra (sharing one binned frequency axis) of + four detector-unit TOD views: + * ``raw``: the raw detector TOD. + * ``ncorrsub``: the TOD with only the correlated noise subtracted (sky signal, orbital + dipole, and white noise retained); equals ``raw`` when no n_corr drawn. + * ``residual``: the noise residual, with the sky model, orbital dipole, and correlated + noise all subtracted. + * ``ncorr``: the correlated-noise realization itself, stored only when one was drawn. + ``ncorrsub`` and ``residual`` use the jump-corrected stream (matching mapmaking and n_corr + sampling). When the off-by-default DEBUG full-``n_corr`` collection is enabled, also stores + the entire ``n_corr`` TOD for this detector-scan. + """ + nbin = tod_samples.TOD_PS_NBIN + freqs_binned, raw_binned = _binned_tod_power_spectrum(view.tod, view.fsamp, nbin) + tod_samples.tod_ps_freqs[iscan, idet] = freqs_binned + tod_samples.tod_ps_raw[iscan, idet] = raw_binned + + # Sky+orbital-dipole-subtracted residual, and the TOD with only the correlated noise removed. + # Both are fresh writable copies, so n_corr (when present) is subtracted in place from each. + residual_tod = view.get_tod(subtract=(("sky", TODView._ALL_GAIN_TERMS), + ("orbital_dipole", TODView._ALL_GAIN_TERMS))) + ncorrsub_tod = view.get_tod() + if n_corr is not None: + residual_tod -= n_corr + ncorrsub_tod -= n_corr + _, residual_binned = _binned_tod_power_spectrum(residual_tod, view.fsamp, nbin) + _, ncorrsub_binned = _binned_tod_power_spectrum(ncorrsub_tod, view.fsamp, nbin) + tod_samples.tod_ps_residual[iscan, idet] = residual_binned + tod_samples.tod_ps_ncorrsub[iscan, idet] = ncorrsub_binned + + if n_corr is not None: + _, ncorr_binned = _binned_tod_power_spectrum(n_corr, view.fsamp, nbin) + tod_samples.tod_ps_ncorr[iscan, idet] = ncorr_binned + if tod_samples.ncorr_tods is not None: + tod_samples.ncorr_tods[iscan][idet] = n_corr.astype(np.float32, copy=False) + def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output: NDArray, tod_samples: TODSamples, params: Bunch, chain: int, iter: int, - do_ncorr_sampling: bool) -> dict[str, DetectorMap]: + ncorr_cfg: Bunch) -> dict[str, DetectorMap]: """ Commander4 CG mapmaking. All ranks on the provided MPI communicator collaborates on creating the band maps (sky signal, inverse variance, possibly also aux maps like orbital dipole). Args: @@ -72,7 +116,8 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output params (Bunch): Parameter file as 'Param' object. chain (int): Current chain number. iter (int): Current Gibbs iteration. - do_ncorr_sampling (bool): Perform correlated noise sampling or not. + ncorr_cfg (Bunch): Correlated-noise sampling config (do_ncorr, do_param, cg_err_tol, + cg_max_iter). Output: dict[str, DetectorMap]: Dictionary containing the solved detector maps, keyed by polarization component ('I', 'QU'). @@ -90,10 +135,12 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): view = scan_view.focus(iscan, idet) + if not view.accept: + continue good_data_mask = view.good_data_mask pix = view.pix[good_data_mask] psi = view.psi[good_data_mask] - sigma0, fknee, alpha = view.noise_params + sigma0 = view.sigma0 gain = view.get_gain() inv_var = (gain/sigma0)**2 mapmaker_invvar.accumulate_to_map(inv_var, pix, psi, response=view.det_response) @@ -111,9 +158,11 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): view = scan_view.focus(iscan, idet) + if not view.accept: + continue good_data_mask = view.good_data_mask pix = view.pix[good_data_mask] - sigma0, fknee, alpha = view.noise_params + sigma0 = view.sigma0 gain = view.get_gain() inv_var = (gain/sigma0)**2 mapmaker_invvar.accumulate_to_map(inv_var, pix) @@ -132,22 +181,26 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output # mapmaker = BinMapmaker(band_comm, experiment_data.nside) mapmaker_orbdipole = BinMapmaker(band_comm, experiment_data.nside) - if do_ncorr_sampling: + if ncorr_cfg.do_ncorr: mapmaker_ncorr = BinMapmaker(band_comm, experiment_data.nside) - fknees = [] - alphas = [] + sampled_params = [] + residuals = [] + niters = [] num_failed_convergences_ncorr = 0 + num_too_high_var_ncorr = 0 worst_residual_ncorr = 0 - + ### MAIN SCAN LOOP ### for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): view = scan_view.focus(iscan, idet) + if not view.accept: + continue pix, psi = view.pix, view.psi good_data_mask = view.good_data_mask pix_masked = pix[good_data_mask] psi_masked = psi[good_data_mask] - sigma0, fknee, alpha = view.noise_params + sigma0 = view.sigma0 gain = view.get_gain() inv_var = (gain/sigma0)**2 response = view.det_response if pols == "IQU" else None @@ -162,23 +215,20 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output mapmaker_orbdipole.accumulate_to_map(sky_orb_dipole, inv_var, pix, psi) ### CORRELATED NOISE SAMPLING ### - if do_ncorr_sampling: + if ncorr_cfg.do_ncorr: sky_subtracted_TOD = view.get_tod( subtract=(("sky", TODView._ALL_GAIN_TERMS), ("orbital_dipole", TODView._ALL_GAIN_TERMS)), ) - Ntod = sky_subtracted_TOD.shape[0] - Nfft = Ntod + 1 # mirrored FFT: nfft=2*Ntod, n=nfft/2+1=Ntod+1 - freq = rfftfreq(2 * Ntod, d=1/view.fsamp) - mask_full = view.full_mask - sigma0_ncorr = calc_sigma0_robust(sky_subtracted_TOD, mask_full) - C_1f_inv = np.zeros(Nfft) - C_1f_inv[1:] = 1.0 / (sigma0_ncorr**2*(freq[1:]/fknee)**alpha) - # fill_all_masked(sky_subtracted_TOD, mask, sigma0_ncorr) - err_tol = 1e-6 - n_corr_est, residual = corr_noise_realization_with_gaps(sky_subtracted_TOD, - mask_full, sigma0_ncorr, - C_1f_inv, err_tol=err_tol) + res = sample_correlated_noise( + sky_subtracted_TOD, view.full_mask, np.array(view.noise_params, copy=True), + experiment_data.noise_model, view.fsamp, cg_err_tol=ncorr_cfg.cg_err_tol, + cg_max_iter=ncorr_cfg.cg_max_iter, sample_params=ncorr_cfg.do_param, + sample_sigma0=ncorr_cfg.sample_sigma0, nomono=ncorr_cfg.nomono, + onlymono=ncorr_cfg.onlymono, + sigma0_dec=ncorr_cfg.sigma0_dec, psd_fit_nu_min=ncorr_cfg.psd_fit_nu_min, + psd_fit_nu_max=ncorr_cfg.psd_fit_nu_max, psd_bin=ncorr_cfg.psd_bin) + n_corr_est = res.n_corr if pols == "IQU": mapmaker_ncorr.accumulate_to_map( (n_corr_est/gain).astype(np.float32, copy=False), @@ -187,19 +237,22 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output mapmaker_ncorr.accumulate_to_map( (n_corr_est/gain).astype(np.float32, copy=False), inv_var, pix, psi) - if residual > err_tol: + tod_samples.noise_params[iscan, idet, :] = res.noise_params + if ncorr_cfg.do_param: + sampled_params.append(np.array(res.noise_params, copy=True)) + if not res.converged: num_failed_convergences_ncorr += 1 - worst_residual_ncorr = max(worst_residual_ncorr, residual) - - ### CORRELATED NOISE POWER SPECTRUM PARAMETERS SAMPLING ### - fknee, alpha = sample_noise_PS_params(n_corr_est, sigma0_ncorr, view.fsamp, alpha, - freq_max=2.0, n_grid=150, n_burnin=4) - tod_samples.noise_params[iscan, idet, :] = sigma0, fknee, alpha - alphas.append(alpha) - fknees.append(fknee) + if res.high_var: + num_too_high_var_ncorr += 1 + worst_residual_ncorr = max(worst_residual_ncorr, res.residual) + residuals.append(res.residual) + niters.append(res.niter) d_sky -= n_corr_est + _record_tod_diagnostics(tod_samples, iscan, idet, view, + n_corr_est if ncorr_cfg.do_ncorr else None) + d_sky_masked = d_sky[good_data_mask] cg_mapmaker.accum_to_RHS( @@ -211,27 +264,12 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output ) ### PRINT NOISE SAMPLING STATS ### - if do_ncorr_sampling: - num_failed_convergences_ncorr = band_comm.reduce(num_failed_convergences_ncorr, op=MPI.SUM) - worst_residual_ncorr = band_comm.reduce(worst_residual_ncorr, op=MPI.MAX) - if band_comm.Get_rank() == 0: - if num_failed_convergences_ncorr > 0: - logger.info(f"Band {experiment_data.nu}GHz failed noise CG for "\ - f"{num_failed_convergences_ncorr} scans. "\ - f"Worst residual = {worst_residual_ncorr:.3e}.") - - alphas = band_comm.gather(alphas, root=0) - fknees = band_comm.gather(fknees, root=0) - if band_comm.Get_rank() == 0: - alphas = np.concatenate(alphas) - fknees = np.concatenate(fknees) - logger.info(f"{experiment_data.nu}GHz: fknees {np.min(fknees):.4f} "\ - f"{np.percentile(fknees, 1):.4f} {np.mean(fknees):.4f} {np.percentile(fknees, 99):.4f}"\ - f" {np.max(fknees):.4f}") - logger.info(f"{experiment_data.nu}GHz: alphas {np.min(alphas):.4f} "\ - f"{np.percentile(alphas, 1):.4f} {np.mean(alphas):.4f} {np.percentile(alphas, 99):.4f}"\ - f" {np.max(alphas):.4f}") - + if ncorr_cfg.do_ncorr: + log_corr_noise_stats(band_comm, experiment_data.nu, experiment_data.noise_model, + sampled_params, residuals, niters, num_failed_convergences_ncorr, + num_too_high_var_ncorr, worst_residual_ncorr, + len(experiment_data.scans)*experiment_data.ndet) + ### GATHER AND NORMALIZE MAPS ### @@ -289,7 +327,7 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output ##################### - if do_ncorr_sampling: + if ncorr_cfg.do_ncorr: mapmaker_ncorr.gather_map() mapmaker_ncorr.normalize_map(map_cov) map_corrnoise = mapmaker_ncorr.final_map @@ -314,7 +352,7 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output maps_to_file["map_rms"] = map_rms if params.general.write_orb_dipole_maps_to_chain: maps_to_file["map_orbdipole"] = map_orbdipole - if params.general.write_corr_noise_maps_to_chain and do_ncorr_sampling: + if params.general.write_corr_noise_maps_to_chain and ncorr_cfg.do_ncorr: maps_to_file["map_corrnoise"] = map_corrnoise if params.general.write_sky_model_maps_to_chain: maps_to_file["map_skymodel"] = compsep_output @@ -327,7 +365,7 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output: NDArray, tod_samples: TODSamples, params: Bunch, chain: int, iter: int, - do_ncorr_sampling: bool) -> dict[str, DetectorMap]: + ncorr_cfg: Bunch) -> dict[str, DetectorMap]: """ Commander4 bin mapmaking. All ranks on the provided MPI communicator collaborates on creating the band maps (sky signal, inverse variance, possibly also aux maps like orbital dipole). Args: @@ -339,7 +377,8 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu params (Bunch): Parameter file as 'Param' object. chain (int): Current chain number. iter (int): Current Gibbs iteration. - do_ncorr_sampling (bool): Perform correlated noise sampling or not. + ncorr_cfg (Bunch): Correlated-noise sampling config (do_ncorr, do_param, cg_err_tol, + cg_max_iter). Output: dict[str, DetectorMap]: Dictionary containing the solved detector maps, keyed by polarization component ('I', 'QU'). @@ -358,10 +397,12 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): view = scan_view.focus(iscan, idet) + if not view.accept: + continue good_data_mask = view.good_data_mask pix = view.pix[good_data_mask] psi = view.psi[good_data_mask] - sigma0, fknee, alpha = view.noise_params + sigma0 = view.sigma0 gain = view.get_gain() # sigma0 is in detector-units, transform into uK_RJ by dividing it by the gain. inv_var = (gain/sigma0)**2 @@ -372,10 +413,9 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu mapmaker = MapmakerIQU(band_comm, experiment_data.nside) mapmaker_orbdipole = MapmakerIQU(band_comm, experiment_data.nside) - if do_ncorr_sampling: + if ncorr_cfg.do_ncorr: mapmaker_ncorr = MapmakerIQU(band_comm, experiment_data.nside) - fknees = [] - alphas = [] + sampled_params = [] residuals = [] niters = [] num_failed_convergences_ncorr = 0 @@ -386,6 +426,8 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu ### MAIN SCAN LOOP ### for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): + if not tod_samples.accept[iscan, idet]: + continue start_bench("binned-mapmaker") view = scan_view.focus(iscan, idet) good_data_mask = view.good_data_mask @@ -394,7 +436,7 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu psi_masked = psi[good_data_mask] response = view.det_response gain = view.get_gain() - sigma0, fknee, alpha = view.noise_params + sigma0 = view.sigma0 # sigma0 is in detector-units, transform into uK_RJ by dividing it by the gain. inv_var = (gain/sigma0)**2 @@ -404,66 +446,42 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu stop_bench("binned-mapmaker", increment_count=False) ### CORRELATED NOISE SAMPLING ### - if do_ncorr_sampling: + if ncorr_cfg.do_ncorr: start_bench("ncorr-sampling") sky_subtracted_TOD = view.get_tod( subtract=(("sky", TODView._ALL_GAIN_TERMS), ("orbital_dipole", TODView._ALL_GAIN_TERMS)), ) - Ntod = sky_subtracted_TOD.shape[0] - Nfft = Ntod + 1 # mirrored FFT: nfft=2*Ntod, n=nfft/2+1=Ntod+1 - freq = rfftfreq(2 * Ntod, d=1/view.fsamp) - - mask_full = view.full_mask - sigma0_ncorr = calc_sigma0_robust(sky_subtracted_TOD, mask_full) - C_1f_inv = np.zeros(Nfft) - C_1f_inv[1:] = 1.0 / (sigma0_ncorr**2*(freq[1:]/fknee)**alpha) - # Inpaint masked regions with linear slope + white noise. - # In the CG solver this is only used to define the starting guess, - # but if the CG fails it is also used to generate the fallback solution. - fill_all_masked(sky_subtracted_TOD, mask_full, sigma0_ncorr) - # n_corr_est, residual, niter, did_conv = corr_noise_realization_with_gaps(sky_subtracted_TOD, - # mask_full, sigma0_ncorr, C_1f_inv, - # err_tol=1e-4, max_iter=1) - residual = np.inf - resid = np.inf - niter = 0 - did_conv = False - # resid = (sky_subtracted_TOD - n_corr_est) * mask_full - var_resid = np.dot(resid, resid) - var_data = np.dot(sky_subtracted_TOD * mask_full, sky_subtracted_TOD * mask_full) - # If either of the two tests failed, use fallback for n_corr. - if var_resid > var_data or not did_conv: - # Direcly solve constrained realization system without a mask. - n_corr_est, _, _, _ = corr_noise_realization_with_gaps(sky_subtracted_TOD, - np.ones_like(mask_full, dtype=bool), sigma0_ncorr, C_1f_inv) - # if band_comm.Get_rank() == 0 and idet == 0 and chain == 1: - # if iscan == 300 or iscan == 600 or iscan == 900: - # np.save(f"corrdata/mirrorfft_corrected_ncorr_signal_{experiment_data.band_name}_{iscan}_{iter}.npy", sky_subtracted_TOD) - # np.save(f"corrdata/mirrorfft_corrected_ncorr_ncorr_{experiment_data.band_name}_{iscan}_{iter}.npy", n_corr_est) + res = sample_correlated_noise( + sky_subtracted_TOD, view.full_mask, np.array(view.noise_params, copy=True), + experiment_data.noise_model, view.fsamp, cg_err_tol=ncorr_cfg.cg_err_tol, + cg_max_iter=ncorr_cfg.cg_max_iter, sample_params=ncorr_cfg.do_param, + sample_sigma0=ncorr_cfg.sample_sigma0, nomono=ncorr_cfg.nomono, + onlymono=ncorr_cfg.onlymono, + sigma0_dec=ncorr_cfg.sigma0_dec, psd_fit_nu_min=ncorr_cfg.psd_fit_nu_min, + psd_fit_nu_max=ncorr_cfg.psd_fit_nu_max, psd_bin=ncorr_cfg.psd_bin) + n_corr_est = res.n_corr mapmaker_ncorr.accumulate_to_map((n_corr_est/gain).astype(np.float32, copy=False), inv_var, pix, psi, response=response) - if not did_conv: + tod_samples.noise_params[iscan,idet,:] = res.noise_params + if ncorr_cfg.do_param: + sampled_params.append(np.array(res.noise_params, copy=True)) + if not res.converged: num_failed_convergences_ncorr += 1 - if var_resid > var_data: + if res.high_var: num_too_high_var_ncorr += 1 - worst_residual_ncorr = max(worst_residual_ncorr, residual) - - ### CORRELATED NOISE POWER SPECTRUM PARAMETERS SAMPLING ### - fknee, alpha = sample_noise_PS_params(n_corr_est, sigma0_ncorr, view.fsamp, alpha, - freq_max=2.0, n_grid=150, n_burnin=4) - - tod_samples.noise_params[iscan,idet,:] = sigma0, fknee, alpha - alphas.append(alpha) - fknees.append(fknee) - residuals.append(residual) - niters.append(niter) + worst_residual_ncorr = max(worst_residual_ncorr, res.residual) + residuals.append(res.residual) + niters.append(res.niter) d_sky -= n_corr_est stop_bench("ncorr-sampling") if iscan == len(experiment_data.scans) - 1: log_memory("ncorr-sampling") + _record_tod_diagnostics(tod_samples, iscan, idet, view, + n_corr_est if ncorr_cfg.do_ncorr else None) + d_sky_masked = d_sky[good_data_mask] start_bench("binned-mapmaker") mapmaker.accumulate_to_map(d_sky_masked/gain, inv_var, pix_masked, psi_masked, response=response) @@ -471,45 +489,10 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu stop_bench("binned-mapmaker", increment_count=False) ### PRINT NOISE SAMPLING STATS ### - if do_ncorr_sampling: - num_failed_convergences_ncorr = band_comm.reduce(num_failed_convergences_ncorr, op=MPI.SUM) - num_too_high_var_ncorr = band_comm.reduce(num_too_high_var_ncorr, op=MPI.SUM) - worst_residual_ncorr = band_comm.reduce(worst_residual_ncorr, op=MPI.MAX) - nscans_global = band_comm.reduce(nscans*ndet, op=MPI.SUM) - if band_comm.Get_rank() == 0: - logger.debug(f"Worst corr-noise sampling residual (band {experiment_data.nu}GHz) = "\ - f"{worst_residual_ncorr:.2e}.") - if num_failed_convergences_ncorr > 0: - logger.warning(f"Band {experiment_data.nu}GHz failed noise CG for "\ - f"{num_failed_convergences_ncorr} out of {nscans_global} scans. "\ - f"Worst residual = {worst_residual_ncorr:.3e}.") - if num_too_high_var_ncorr > 0: - logger.warning(f"Band {experiment_data.nu}GHz failed variance sanity check for "\ - f"{num_too_high_var_ncorr} out of {nscans_global} scans. ") - - alphas = band_comm.gather(alphas, root=0) - fknees = band_comm.gather(fknees, root=0) - residuals = band_comm.gather(residuals, root=0) - niters = band_comm.gather(niters, root=0) - if band_comm.Get_rank() == 0: - alphas = np.concatenate(alphas) - fknees = np.concatenate(fknees) - residuals = np.concatenate(residuals) - residuals = residuals[residuals != 0] - residuals = np.array([0]) if len(residuals) == 0 else residuals - niters = np.concatenate(niters) - logger.info(f"{experiment_data.nu}GHz: fknees {np.nanmin(fknees):.4f} "\ - f"{np.nanpercentile(fknees, 1):.4f} {np.nanmean(fknees):.4f} {np.nanpercentile(fknees, 99):.4f}"\ - f" {np.nanmax(fknees):.4f}") - logger.info(f"{experiment_data.nu}GHz: alphas {np.nanmin(alphas):.4f} "\ - f"{np.nanpercentile(alphas, 1):.4f} {np.nanmean(alphas):.4f} {np.nanpercentile(alphas, 99):.4f}"\ - f" {np.nanmax(alphas):.4f}") - logger.info(f"{experiment_data.nu}GHz: residuals {np.nanmin(residuals):.2e} "\ - f"{np.nanpercentile(residuals, 1):.2e} {np.nanmean(residuals):.2e} {np.nanpercentile(residuals, 99):.2e}"\ - f" {np.nanmax(residuals):.2e}") - logger.info(f"{experiment_data.nu}GHz: iterations {np.nanmin(niters):.4f} "\ - f"{np.nanpercentile(niters, 1):.4f} {np.nanmean(niters):.4f} {np.nanpercentile(niters, 99):.4f}"\ - f" {np.nanmax(niters):.4f}") + if ncorr_cfg.do_ncorr: + log_corr_noise_stats(band_comm, experiment_data.nu, experiment_data.noise_model, + sampled_params, residuals, niters, num_failed_convergences_ncorr, + num_too_high_var_ncorr, worst_residual_ncorr, nscans*ndet) start_bench("binned-mapmaker") @@ -522,7 +505,7 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu map_signal = mapmaker.final_map mapmaker_orbdipole.normalize_map(map_cov) map_orbdipole = mapmaker_orbdipole.final_map - if do_ncorr_sampling: + if ncorr_cfg.do_ncorr: mapmaker_ncorr.gather_map() mapmaker_ncorr.normalize_map(map_cov) map_corrnoise = mapmaker_ncorr.final_map @@ -550,7 +533,7 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu maps_to_file["map_rms"] = map_rms if params.general.write_orb_dipole_maps_to_chain: maps_to_file["map_orbdipole"] = map_orbdipole - if params.general.write_corr_noise_maps_to_chain and do_ncorr_sampling: + if params.general.write_corr_noise_maps_to_chain and ncorr_cfg.do_ncorr: maps_to_file["map_corrnoise"] = map_corrnoise if params.general.write_sky_model_maps_to_chain: maps_to_file["map_skymodel"] = compsep_output @@ -613,7 +596,9 @@ def init_tod_processing(mpi_info: Bunch, params: Bunch) -> tuple[Bunch, str, Det # What is my rank number among the ranks processing this detector? my_experiment = experiment # Setting our unique detector id. Note that this is a global, not per band. - tot_num_scans = experiment.num_scans + # A per-band ``num_scans`` (bands can hold different numbers of scans) takes + # precedence over the experiment-level value, which is the shared default. + tot_num_scans = band.num_scans if "num_scans" in band else experiment.num_scans scans = np.arange(tot_num_scans) my_scans = np.array_split(scans, mpi_info.band.size)[mpi_info.band.rank] my_scans_start = my_scans[0] @@ -632,11 +617,6 @@ def init_tod_processing(mpi_info: Bunch, params: Bunch) -> tuple[Bunch, str, Det with benchmark("fileread-tod"): experiment_data = read_tods_from_file(band_comm, my_experiment, my_band, det_names, params, my_scans_start, my_scans_stop) - #FIXME: Make this not a hacky fix. - if "init_sky_from" in my_band: - experiment_data.sky_init_files = my_band.init_sky_from - else: - experiment_data.sky_init_files = [] mpi_info.tod.comm.Barrier() if mpi_info.tod.is_master: logger.info(f"TOD: Finished reading all files in {time.time()-t0:.1f}s.") @@ -773,16 +753,95 @@ def sample_jump_detection(band_comm: MPI.Comm, experiment_data: DetGroupTOD, +# Valid calibration targets for gain sampling, and the per-term defaults used when a parameter file +# leaves `calibrate_against` unspecified (preserving the historical low-frequency behavior). +_VALID_CALIB_TARGETS = ("orbital_dipole", "full_sky", "sky") +_DEFAULT_CALIB_TARGETS = {"abs_gain": "orbital_dipole", + "rel_gain": "full_sky", + "temporal_gain": "full_sky"} + + +def _resolve_calib_target(params: Bunch, experiment_data: DetGroupTOD, gain_block: str) -> str: + """ Resolve which signal a gain term is calibrated against for the current band. + + A per-band ``calibrate_against`` (under ``experiments..bands..``) + overrides ``general..calibrate_against``, which in turn falls back to the + term's default in ``_DEFAULT_CALIB_TARGETS``. + """ + general_block = params.general[gain_block] + target = (general_block["calibrate_against"] if "calibrate_against" in general_block + else _DEFAULT_CALIB_TARGETS[gain_block]) + band = params.experiments[experiment_data.experiment_name].bands[experiment_data.band_name] + if gain_block in band and "calibrate_against" in band[gain_block]: + target = band[gain_block]["calibrate_against"] + if target not in _VALID_CALIB_TARGETS: + raise ValueError(f"{gain_block}.calibrate_against='{target}' is invalid; must be one of " + f"{_VALID_CALIB_TARGETS}.") + return target + + +def _resolve_gain_downsample_factor(params: Bunch, experiment_data: DetGroupTOD) -> int: + """ Downsampling factor (in samples) used when building the gain-calibration TODs. + + Derived from ``general.gain_calib_downsample_time`` (a duration in seconds, shared by all + gain terms) and the band's sampling rate. A duration of 0 disables downsampling. + """ + return max(1, int(round(params.general.gain_calib_downsample_time * experiment_data.fsamp))) + + +def _solve_relative_gain_system(s_weights: NDArray, r_weights: NDArray, prev_rel_gain: NDArray, + rng=None) -> NDArray: + """ Draw relative-gain deviations Delta g_i from the BP7 Sec. 3.4 constrained Gaussian. + + Solves the bordered linear system enforcing ``sum(Delta g_i) = 0`` over the *active* + detectors only -- those with nonzero calibration weight ``s_weights``. A detector rejected + on every scan (or with a vanishing calibrator) has ``s_weights == 0``; it would contribute + an all-zero row/column and, if two or more are present, make the matrix singular. Such + detectors are excluded from the solve (shrinking the system to the active set) and held at + their current relative gain. + + Args: + s_weights (NDArray): Per-detector ``sum_scans s^T N^-1 s`` (calibration weight), shape (ndet,). + r_weights (NDArray): Per-detector ``sum_scans r^T N^-1 s`` (residual projection), shape (ndet,). + prev_rel_gain (NDArray): Current relative gains, kept for excluded detectors, shape (ndet,). + rng: Optional NumPy random generator for the fluctuation term (defaults to ``np.random``). + + Returns: + NDArray: Full-length (ndet,) float32 relative-gain vector with the active entries resampled. + + Raises: + np.linalg.LinAlgError: If the reduced system is singular (left for the caller to handle). + """ + rng = np.random if rng is None else rng + out = np.array(prev_rel_gain, dtype=np.float32) + idx = np.flatnonzero(np.asarray(s_weights) > 0.0) + n = idx.size + if n == 0: + return out + d = np.asarray(s_weights)[idx].astype(np.float64) + r = np.asarray(r_weights)[idx].astype(np.float64) + A = np.zeros((n + 1, n + 1)) + A[:n, :n] = np.diag(d) + A[:n, n] = 0.5 # Lagrange-multiplier column enforcing the zero-sum constraint. + A[n, :n] = 1.0 # Constraint row: sum of active Delta g_i = 0. + b = np.zeros(n + 1) + b[:n] = r + np.sqrt(d) * rng.standard_normal(n) + solution = np.linalg.solve(A, b) # Raises LinAlgError if singular. + out[idx] = solution[:n].astype(np.float32) + return out + + def sample_absolute_gain(band_comm: MPI.Comm, experiment_data: DetGroupTOD, tod_samples: TODSamples, - det_compsep_map: NDArray): + det_compsep_map: NDArray, calibrate_against: str, downsample_factor: int): """ Draw a realization of the absolute gain term, g0, which is constant across all - detectors and all scans within a band. For frequencies < 380.0 GHz this is done using - only the orbital dipole, and above it uses the full sky. + detectors and all scans within a band, calibrated against ``calibrate_against``. Args: band_comm (MPI.Comm): The band-level MPI communicator. experiment_data (DetGroupTOD): The object holding all the scan data. tod_samples (TODSamples): Current sampled TOD parameters (updated in-place with g0). det_compsep_map (NDArray): The component-separation sky map for the detector. + calibrate_against (str): Calibrator signal, one of "orbital_dipole" | "full_sky" | "sky". + downsample_factor (int): Block-averaging factor for the calibration TODs. Returns: tod_samples (TODSamples): Updated TOD samples with the new g0 estimate. wait_time (float): Time spent waiting at the MPI barrier. @@ -791,18 +850,15 @@ def sample_absolute_gain(band_comm: MPI.Comm, experiment_data: DetGroupTOD, tod_ sum_s_T_N_inv_d = 0 # Accumulators for the numerator and denominator of eqn 16. sum_s_T_N_inv_s = 0 - # Calibrate on the full sky at high frequencies, as the orbital dipole is too faint. - calibrate_on_full_sky = experiment_data.nu > 380.0 scan_view = TODView(experiment_data, tod_samples, compsep_output=det_compsep_map) for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): view = scan_view.focus(iscan, idet) - calib = view.get_abs_calib_tod( - downsample_factor=int(view.fsamp), - calibrate_on_full_sky=calibrate_on_full_sky, - preserve_target_gain=False, - ) + if not view.accept: # Skip detector-scans flagged as bad; they carry no gain info. + continue + calib = view.get_calib_tod("abs", calibrate_against, + downsample_factor=downsample_factor) s_cal = calib.s_cal residual_tod = calib.tod @@ -818,16 +874,20 @@ def sample_absolute_gain(band_comm: MPI.Comm, experiment_data: DetGroupTOD, tod_ # The g0 term is fully global, so we reduce across both all scans and all bands: sum_s_T_N_inv_d = band_comm.reduce(sum_s_T_N_inv_d, op=MPI.SUM, root=0) sum_s_T_N_inv_s = band_comm.reduce(sum_s_T_N_inv_s, op=MPI.SUM, root=0) - g_sampled = 0.0 + # Default to the current value so a skipped or ill-posed solve leaves the gain unchanged. + g_sampled = tod_samples.abs_gain # Rank 0 draws a sample of g0 from eq (16) from BP6, and bcasts it to the other ranks. if band_comm.Get_rank() == 0: - eta = np.random.randn() - g_mean = sum_s_T_N_inv_d / sum_s_T_N_inv_s - g_std = 1.0 / np.sqrt(sum_s_T_N_inv_s) - - g_sampled = g_mean + eta * g_std - logger.info(f"Band {experiment_data.band_name} g0: {tod_samples.abs_gain:.4e} "\ - f"-> {g_sampled:.4e} (+/- {g_std:.4e})") + if not np.isfinite(sum_s_T_N_inv_s) or sum_s_T_N_inv_s <= 0.0: + logger.error(f"Band {experiment_data.band_name} absolute gain has no calibration " + f"weight (all detector-scans rejected or zero calibrator): not updating.") + else: + eta = np.random.randn() + g_mean = sum_s_T_N_inv_d / sum_s_T_N_inv_s + g_std = 1.0 / np.sqrt(sum_s_T_N_inv_s) + g_sampled = g_mean + eta * g_std + logger.info(f"Band {experiment_data.band_name} g0: {tod_samples.abs_gain:.4e} "\ + f"-> {g_sampled:.4e} (+/- {g_std:.4e})") t0 = time.time() band_comm.Barrier() @@ -843,7 +903,8 @@ def sample_absolute_gain(band_comm: MPI.Comm, experiment_data: DetGroupTOD, tod_ def sample_relative_gain(band_comm: MPI.Comm, experiment_data: DetGroupTOD, - tod_samples: TODSamples, det_compsep_map: NDArray): + tod_samples: TODSamples, det_compsep_map: NDArray, calibrate_against: str, + downsample_factor: int): """ Samples the detector-dependent relative gain (Delta g_i). This function implements the logic from Sec. 3.4 of BP7. Args: @@ -851,6 +912,8 @@ def sample_relative_gain(band_comm: MPI.Comm, experiment_data: DetGroupTOD, experiment_data (DetGroupTOD): The object holding scan data for the band. tod_samples (TODSamples): Current sampled TOD parameters. det_compsep_map (NDArray): The component-separation sky map for the detector. + calibrate_against (str): Calibrator signal, one of "orbital_dipole" | "full_sky" | "sky". + downsample_factor (int): Block-averaging factor for the calibration TODs. Returns: tod_samples (TODSamples): Updated TOD samples with relative gain estimates. """ @@ -869,7 +932,10 @@ def sample_relative_gain(band_comm: MPI.Comm, experiment_data: DetGroupTOD, for iscan, scan in enumerate(experiment_data.scans): for idet, det in enumerate(scan.detectors): view = scan_view.focus(iscan, idet) - calib = view.get_rel_calib_tod(downsample_factor=int(view.fsamp)) + if not view.accept: # Skip detector-scans flagged as bad; they carry no gain info. + continue + calib = view.get_calib_tod("rel", calibrate_against, + downsample_factor=downsample_factor) s_cal = calib.s_cal residual_tod = calib.tod N_inv_s = experiment_data.apply_N_inv(s_cal, view.noise_params, samprate=1.0) @@ -890,34 +956,37 @@ def sample_relative_gain(band_comm: MPI.Comm, experiment_data: DetGroupTOD, band_comm.Allreduce(MPI.IN_PLACE, local_r_T_N_inv_s, op=MPI.SUM) ### 3. Solve Global System ### - delta_g_samples = np.zeros(ndet, dtype=np.float32) + # Solve the constrained system (sum of Delta g_i = 0) over the active detectors only; detectors + # rejected on every scan or with a vanishing calibrator carry zero weight, are held at their + # current value, and are excluded so the bordered matrix stays non-singular. + delta_g_samples = np.array(tod_samples.rel_gain, dtype=np.float32) # default: leave unchanged if band_comm.Get_rank() == 0: - A = np.zeros((ndet + 1, ndet + 1)) - b = np.zeros(ndet + 1) - diagonal = np.array(local_s_T_N_inv_s) - A[:ndet, :ndet] = np.diag(diagonal) - A[:ndet, ndet] = 0.5 - A[ndet, :ndet] = 1.0 - eta = np.random.randn(ndet) - fluctuation_term = np.sqrt(diagonal) * eta - - b[:ndet] = np.array(local_r_T_N_inv_s) + fluctuation_term - - try: - solution = np.linalg.solve(A, b) - delta_g_samples[:] = solution[:ndet] - logger.info(f"Solved global relative gains for {ndet} detectors.") - except np.linalg.LinAlgError: - logger.error("Failed to solve global linear system for relative gain: Not updating") + n_active = int(np.count_nonzero(local_s_T_N_inv_s > 0.0)) + n_excluded = ndet - n_active + if n_active == 0: + logger.error(f"Band {experiment_data.band_name}: no detectors with calibration weight " + f"for relative gain; not updating.") + else: + try: + delta_g_samples = _solve_relative_gain_system(local_s_T_N_inv_s, + local_r_T_N_inv_s, tod_samples.rel_gain) + msg = f"Solved relative gains for {n_active} active detectors" + if n_excluded: + msg += f" ({n_excluded} excluded: rejected on all scans or zero calibrator)" + logger.info(msg + ".") + except np.linalg.LinAlgError: + logger.error("Failed to solve linear system for relative gain: Not updating.") + # Broadcast and apply on every rank, so all band ranks hold the identical relative-gain vector. + prev_rel_gain = np.array(tod_samples.rel_gain) band_comm.Bcast(delta_g_samples, root=0) + tod_samples.rel_gain[:] = delta_g_samples log_memory("rel-gain") - + if band_comm.Get_rank() == 0: logger.info(f"Rel gain for band {experiment_data.band_name}: min = "\ f"{np.min(delta_g_samples):.3e} max = {np.max(delta_g_samples):.3e}") logger.debug(f"Rel gains for band {experiment_data.band_name}: {delta_g_samples}\n"\ - f"Average change = {np.mean(np.abs(tod_samples.rel_gain - delta_g_samples))}") - tod_samples.rel_gain[:] = delta_g_samples + f"Average change = {np.mean(np.abs(prev_rel_gain - delta_g_samples))}") wait_time = 0 return tod_samples, wait_time @@ -926,7 +995,8 @@ def sample_relative_gain(band_comm: MPI.Comm, experiment_data: DetGroupTOD, def sample_temporal_gain_variations(band_comm: MPI.Comm, experiment_data: DetGroupTOD, tod_samples: TODSamples, det_compsep_map: NDArray, - chain: int, iter: int, params: Bunch): + chain: int, iter: int, params: Bunch, calibrate_against: str, + downsample_factor: int): """ Samples the time-dependent relative gain variations (delta g_qi). This function implements the logic from Sec. 3.5 of the BP7 paper, using a Wiener filter to smooth the gain solution over time (PIDs). It solves a global system for all scans of a given detector, which are @@ -940,6 +1010,8 @@ def sample_temporal_gain_variations(band_comm: MPI.Comm, experiment_data: DetGro chain (int): Current chain number. iter (int): Current Gibbs iteration. params (Bunch): Parameters from the parameter file. + calibrate_against (str): Calibrator signal, one of "orbital_dipole" | "full_sky" | "sky". + downsample_factor (int): Block-averaging factor for the calibration TODs. """ band_rank = band_comm.Get_rank() band_size = band_comm.Get_size() @@ -961,7 +1033,10 @@ def sample_temporal_gain_variations(band_comm: MPI.Comm, experiment_data: DetGro # ringing from the large residual in the galactic plane). view = scan_view.focus(iscan, idet) - calib = view.get_temp_calib_tod(downsample_factor=int(view.fsamp)) + if not view.accept: # Rejected detector-scans contribute zero weight (A_qq = b_q = 0); + continue # the Wiener prior then fills their temporal gain from neighbors. + calib = view.get_calib_tod("temp", calibrate_against, + downsample_factor=downsample_factor) s_cal = calib.s_cal residual_tod = calib.tod @@ -1130,33 +1205,73 @@ def process_tod(mpi_info: Bunch, experiment_data: DetGroupTOD, logger.info(f"Chain {chain} iter{iter} {experiment_data.nu}GHz: Finished jump " f"detection in {timing_dict['jump-detect']:.1f}s.") - ### WHITE NOISE ESTIMATION ### - t0 = time.time() - with benchmark("sigma0-est"): - tod_samples = estimate_white_noise(experiment_data, tod_samples, compsep_output, params) - timing_dict["wn-est-1"] = time.time() - t0 - if mpi_info.tod.is_master: - logger.info(f"Chain {chain} iter{iter} {experiment_data.nu}GHz: Finished white noise "\ - f"estimation in {timing_dict['wn-est-1']:.1f}s.") + ### CORRELATED-NOISE SAMPLING CONFIG ### + # All correlated-noise settings live in the nested ``general.corr_noise`` block. + cn = params.general.corr_noise + sample_corr_noise = cn.sample_corr_noise + from_iter = cn.sample_corr_noise_from_iter_num + sample_noise_params = cn.sample_noise_params + # Parameter sampling consumes a correlated-noise realization, so it cannot run without it. + if sample_noise_params and not sample_corr_noise: + raise ValueError("general.corr_noise.sample_noise_params requires sample_corr_noise=True " + "(parameter sampling needs a correlated-noise realization).") + do_ncorr = sample_corr_noise and iter >= from_iter + # Per-scan monopole handling mirrors the Fortran flags; both set at once is contradictory. + nomono = getattr(cn, "nomono", False) + onlymono = getattr(cn, "onlymono", False) + if nomono and onlymono and mpi_info.band.is_master: + logger.error("general.corr_noise.nomono and onlymono are both True, which is contradictory; " + "onlymono takes precedence.") + ncorr_cfg = Bunch( + do_ncorr=do_ncorr, + do_param=do_ncorr and sample_noise_params, + sample_sigma0=getattr(cn, "sample_sigma0", True), + cg_err_tol=cn.CG_err_tol, + cg_max_iter=cn.CG_max_iter, + nomono=nomono, + onlymono=onlymono, + sigma0_dec=getattr(cn, "sigma0_decimation", 1), + psd_fit_nu_min=getattr(cn, "psd_fit_nu_min", 0.0), + psd_fit_nu_max=getattr(cn, "psd_fit_nu_max", float("inf")), + psd_bin=getattr(cn, "psd_bin", False), + ) - ### ABSOLUTE GAIN CALIBRATION ### - if params.general.sample_abs_gain and iter >= params.general.sample_abs_gain_from_iter_num: + ### WHITE NOISE ESTIMATION ### + # When correlated-noise sampling runs it re-estimates sigma0 after subtracting n_corr + # (Commander3-aligned), so only estimate sigma0 standalone here when that will not happen. + if not (ncorr_cfg.do_ncorr and ncorr_cfg.sample_sigma0): + t0 = time.time() + with benchmark("sigma0-est"): + tod_samples = estimate_white_noise(experiment_data, tod_samples, compsep_output, params) + timing_dict["wn-est-1"] = time.time() - t0 + if mpi_info.tod.is_master: + logger.info(f"Chain {chain} iter{iter} {experiment_data.nu}GHz: Finished white noise "\ + f"estimation in {timing_dict['wn-est-1']:.1f}s.") + + ### ABSOLUTE GAIN CALIBRATION ### + if params.general.abs_gain.sample and iter >= params.general.abs_gain.sample_from_iter_num: + calib_target = _resolve_calib_target(params, experiment_data, "abs_gain") + downsample_factor = _resolve_gain_downsample_factor(params, experiment_data) t0 = time.time() with benchmark("abs-gain"): tod_samples, wait_time = sample_absolute_gain(band_comm, experiment_data, tod_samples, - compsep_output) + compsep_output, calib_target, + downsample_factor) timing_dict["abs-gain"] = time.time() - t0 waittime_dict["abs-gain"] = wait_time if mpi_info.band.is_master: logger.info(f"Chain {chain} iter{iter} {experiment_data.nu}GHz: Finished absolute "\ f"gain estimation in {timing_dict['abs-gain']:.1f}s.") - ### RELATIVE GAIN CALIBRATION ### - if params.general.sample_rel_gain and iter >= params.general.sample_rel_gain_from_iter_num: + ### RELATIVE GAIN CALIBRATION ### + if params.general.rel_gain.sample and iter >= params.general.rel_gain.sample_from_iter_num: + calib_target = _resolve_calib_target(params, experiment_data, "rel_gain") + downsample_factor = _resolve_gain_downsample_factor(params, experiment_data) t0 = time.time() with benchmark("rel-gain"): tod_samples, wait_time = sample_relative_gain(band_comm, experiment_data, tod_samples, - compsep_output) + compsep_output, calib_target, + downsample_factor) timing_dict["rel-gain"] = time.time() - t0 waittime_dict["rel-gain"] = wait_time if mpi_info.band.is_master: @@ -1164,13 +1279,16 @@ def process_tod(mpi_info: Bunch, experiment_data: DetGroupTOD, f"gain estimation in {timing_dict['rel-gain']:.1f}s.") - ### TEMPORAL GAIN CALIBRATION ### - if params.general.sample_temporal_gain\ - and iter >= params.general.sample_temporal_gain_from_iter_num: + ### TEMPORAL GAIN CALIBRATION ### + if params.general.temporal_gain.sample\ + and iter >= params.general.temporal_gain.sample_from_iter_num: + calib_target = _resolve_calib_target(params, experiment_data, "temporal_gain") + downsample_factor = _resolve_gain_downsample_factor(params, experiment_data) t0 = time.time() with benchmark("temporal-gain"): - tod_samples = sample_temporal_gain_variations(band_comm, experiment_data, - tod_samples, compsep_output, chain, iter, params) + tod_samples = sample_temporal_gain_variations(band_comm, experiment_data, tod_samples, + compsep_output, chain, iter, params, calib_target, + downsample_factor) timing_dict["temp-gain"] = time.time() - t0 if mpi_info.band.is_master: logger.info(f"Chain {chain} iter{iter} {experiment_data.nu}GHz: Finished temporal "\ @@ -1186,8 +1304,6 @@ def process_tod(mpi_info: Bunch, experiment_data: DetGroupTOD, # f"estimation in {timing_dict['wn-est-2']:.1f}s.") ### MAPMAKING ### - do_ncorr_sampling = params.general.sample_corr_noise and iter >=\ - params.general.sample_corr_noise_from_iter_num t0 = time.time() if "mapmaker" in params.experiments[experiment_data.experiment_name].bands[experiment_data.band_name]: @@ -1200,10 +1316,10 @@ def process_tod(mpi_info: Bunch, experiment_data: DetGroupTOD, if mapmaker_str == "CG": detmap_dict = tod2map_CG(band_comm, experiment_data, compsep_output, tod_samples, params, chain, - iter, do_ncorr_sampling) + iter, ncorr_cfg) elif mapmaker_str == "bin": - detmap_dict = tod2map_bin(band_comm, experiment_data, compsep_output, tod_samples, params, - chain, iter, do_ncorr_sampling) + detmap_dict = tod2map_bin(band_comm, experiment_data, compsep_output, tod_samples, params, + chain, iter, ncorr_cfg) else: raise ValueError(f'Mapmaker must be either "CG" or "bin", but {mapmaker_str} was given for'\ f' experiment {experiment_data.experiment_name}, band {experiment_data.band_name}') diff --git a/tests/test_components.py b/tests/test_components.py new file mode 100644 index 0000000..30eeff4 --- /dev/null +++ b/tests/test_components.py @@ -0,0 +1,458 @@ +from copy import deepcopy + +import h5py +import healpy as hp +import numpy as np +import pytest +from pixell.bunch import Bunch + +from commander4.sky_models.component import CMB, CompList, PointSourcesComponent, ThermalDust +from commander4.sky_models.sky_model import build_initial_sky_model +from commander4.utils.math_operations import complist_dot, map_to_alm + + +def _make_general(ntask_compsep_qu: int = 1, ntask_compsep_i: int = 1) -> Bunch: + return Bunch( + nside=2, + CG_float_precision="single", + MPI_config=Bunch(ntask_compsep_I=ntask_compsep_i, ntask_compsep_QU=ntask_compsep_qu), + ) + + +def _make_component_cfg(polarization: str = "IQU") -> Bunch: + return Bunch( + enabled=True, + component_class="CMB", + params=Bunch( + lmax=1, + polarization=polarization, + shortname="cmb", + spatially_varying_MM=False, + smoothing_prior_FWHM=0.0, + smoothing_prior_amplitude=1.0, + ), + ) + + +def _make_comp_list(polarization: str = "IQU", ntask_compsep_qu: int = 1) -> CompList: + params = Bunch(general=_make_general(ntask_compsep_qu)) + cmb = _make_component_cfg(polarization) + object.__setattr__(cmb, "_name", "cmb") + components = Bunch({"cmb": cmb}) + return CompList.init_from_params(components, params) + + +def _make_named_component_cfg(shortname: str, polarization: str = "IQU") -> Bunch: + cfg = _make_component_cfg(polarization) + cfg.params.shortname = shortname + return cfg + + +def _make_multi_comp_list() -> CompList: + params = Bunch(general=_make_general()) + cmb = _make_named_component_cfg("cmb", "IQU") + object.__setattr__(cmb, "_name", "cmb") + dust = _make_named_component_cfg("dust", "IQU") + object.__setattr__(dust, "_name", "dust") + components = Bunch( + { + "cmb": cmb, + "dust": dust, + } + ) + return CompList.init_from_params(components, params) + + +def test_init_from_params_requires_component_name() -> None: + params = Bunch(general=_make_general()) + components = Bunch({"cmb": _make_component_cfg("I")}) + + with pytest.raises(AttributeError, match="_name"): + CompList.init_from_params(components, params) + + +def test_init_from_params_does_not_mutate_component_params_name() -> None: + params = Bunch(general=_make_general()) + component_cfg = _make_component_cfg("I") + object.__setattr__(component_cfg, "_name", "cmb") + components = Bunch({"cmb": component_cfg}) + + CompList.init_from_params(components, params) + + assert "_name" not in component_cfg.params + + +def test_init_from_params_builds_all_defined_pol_views() -> None: + # Construction is independent of the MPI/compsep layout: an IQU component always yields both an + # I and a QU view, even when zero compsep ranks are configured. + params = Bunch(general=_make_general(ntask_compsep_qu=0, ntask_compsep_i=0)) + cmb = _make_named_component_cfg("cmb", "IQU") + object.__setattr__(cmb, "_name", "cmb") + ff = _make_named_component_cfg("ff", "I") + object.__setattr__(ff, "_name", "ff") + components = Bunch({"cmb": cmb, "ff": ff}) + + comp_list = CompList.init_from_params(components, params) + + assert [(comp.comp_name, comp.eval_pol) for comp in comp_list] == [ + ("cmb", "I"), ("cmb", "QU"), ("ff", "I")] + + +def test_complist_split_preserves_names_and_join_restores_logical_component() -> None: + comp_list = _make_comp_list("IQU") + + assert [comp.comp_name for comp in comp_list] == ["cmb", "cmb"] + assert [comp.shortname for comp in comp_list] == ["cmb", "cmb"] + assert [comp.eval_pol for comp in comp_list] == ["I", "QU"] + assert [comp.is_split_view for comp in comp_list] == [True, True] + + comp_list[0].alms[:] = 1.0 + 0.0j + comp_list[1].alms[:] = 2.0 + 0.0j + joined = comp_list.joined() + + assert len(joined) == 1 + assert joined[0].comp_name == "cmb" + assert joined[0].shortname == "cmb" + assert joined[0].eval_pol == "IQU" + assert not joined[0].is_split_view + assert np.all(joined[0].alms[0] == 1.0 + 0.0j) + assert np.all(joined[0].alms[1:] == 2.0 + 0.0j) + + +def test_component_itruediv_divides_data() -> None: + comp_list = _make_comp_list("I") + comp = comp_list[0] + other = deepcopy(comp) + comp.alms[:] = 6.0 + 0.0j + other.alms[:] = 3.0 + 0.0j + + comp /= other + + assert np.all(comp.alms == 2.0 + 0.0j) + + +def test_complist_add_returns_full_complist() -> None: + comp_list = _make_comp_list("IQU") + for idx, comp in enumerate(comp_list, start=1): + comp.alms[:] = idx + 0.0j + + summed = comp_list + comp_list + + assert isinstance(summed, CompList) + assert len(summed) == 2 + assert np.all(summed[0].alms == 2.0 + 0.0j) + assert np.all(summed[1].alms == 4.0 + 0.0j) + + +def test_complist_ops_require_matching_execution_views() -> None: + comp_list = _make_comp_list("IQU") + other = deepcopy(comp_list) + other.comp_list.reverse() + + with pytest.raises(ValueError, match="same execution views"): + _ = comp_list + other + with pytest.raises(ValueError, match="same execution views"): + _ = complist_dot(comp_list, other) + + +def test_complist_split_for_eval_pol_rejects_invalid_polarization(caplog) -> None: + comp_list = _make_comp_list("IQU") + + with pytest.raises(AssertionError): + comp_list.split_for_eval_pol("bad") + assert "Unsupported polarization string bad" in caplog.text + + +def test_point_sources_component_rejects_non_intensity_eval_pol(caplog) -> None: + params = Bunch(shortname="ps") + + with pytest.raises(AssertionError): + PointSourcesComponent(params, _make_general(), comp_name="ps", eval_pol="QU") + assert "PointSourcesComponent does not support evaluation polarization 'QU'" in caplog.text + + +def test_complist_split_for_eval_pol_returns_requested_execution_view() -> None: + comp_list = _make_comp_list("IQU") + + qu_only = comp_list.split_for_eval_pol("QU") + + assert len(qu_only) == 1 + assert qu_only[0].eval_pol == "QU" + + +def test_complist_constructor_rejects_duplicate_unsplit_comp_names() -> None: + comp = _make_comp_list("I")[0] + duplicate = deepcopy(comp) + + with pytest.raises(ValueError, match="Duplicate logical component"): + CompList([comp, duplicate]) + + +def test_complist_constructor_rejects_reused_shortname_for_distinct_comp_names() -> None: + comp = _make_comp_list("I")[0] + other = deepcopy(comp) + other.comp_name = "dust" + + with pytest.raises(ValueError, match="Shortname"): + CompList([comp, other]) + + +def test_copy_matching_data_from_leaves_omitted_components_unchanged() -> None: + comp_list = _make_multi_comp_list() + intensity = comp_list.split_for_eval_pol("I") + original_other = intensity[1].alms.copy() + updated_subset = CompList([deepcopy(intensity[0])]) + updated_subset[0].alms[:] = 7.0 + 0.0j + + intensity.copy_matching_data_from(updated_subset) + + assert np.all(intensity[0].alms == 7.0 + 0.0j) + assert np.array_equal(intensity[1].alms, original_other) + + +def _write_chain_alms(path, alms_by_shortname: dict) -> None: + with h5py.File(path, "w") as f: + for shortname, alms in alms_by_shortname.items(): + f[f"comps/{shortname}/alms"] = alms + + +def test_load_initial_alms_reads_and_splits_from_chain(tmp_path) -> None: + nalm = (1 + 1) * (1 + 2) // 2 # lmax == 1, matching the default component config. + cmb_alms = np.arange(3 * nalm, dtype=np.float64).reshape(3, nalm).astype(np.complex64) + ff_alms = (np.arange(nalm, dtype=np.float64) + 100).reshape(1, nalm).astype(np.complex64) + chain = tmp_path / "init_chain.h5" + _write_chain_alms(chain, {"cmb": cmb_alms, "ff": ff_alms}) + + general = _make_general() + general.init_chain_path = str(chain) + cmb = _make_named_component_cfg("cmb", "IQU") + object.__setattr__(cmb, "_name", "cmb") + ff = _make_named_component_cfg("ff", "I") + object.__setattr__(ff, "_name", "ff") + params = Bunch(general=general, components=Bunch({"cmb": cmb, "ff": ff})) + + comp_list = CompList.init_from_params(params.components, params) + comp_list.load_initial_alms(params) + + views = {(comp.comp_name, comp.eval_pol): comp for comp in comp_list} + # The joined IQU alms get split into the I row and the two QU rows. + assert np.array_equal(views[("cmb", "I")].alms, cmb_alms[0:1]) + assert np.array_equal(views[("cmb", "QU")].alms, cmb_alms[1:3]) + assert np.array_equal(views[("ff", "I")].alms, ff_alms[0:1]) + + +def test_load_initial_alms_prefers_per_component_init_from(tmp_path) -> None: + nalm = (1 + 1) * (1 + 2) // 2 + global_chain = tmp_path / "global.h5" + special_chain = tmp_path / "special.h5" + _write_chain_alms(global_chain, {"cmb": np.zeros((3, nalm), dtype=np.complex64)}) + _write_chain_alms(special_chain, {"cmb": np.full((3, nalm), 5.0, dtype=np.complex64)}) + + general = _make_general() + general.init_chain_path = str(global_chain) + cmb = _make_named_component_cfg("cmb", "IQU") + object.__setattr__(cmb, "_name", "cmb") + cmb.params.init_from = str(special_chain) # Per-component path takes precedence over the global one. + params = Bunch(general=general, components=Bunch({"cmb": cmb})) + + comp_list = CompList.init_from_params(params.components, params) + comp_list.load_initial_alms(params) + + assert all(np.all(comp.alms == 5.0) for comp in comp_list) + + +def test_load_initial_alms_leaves_zeros_without_a_source() -> None: + general = _make_general() # No init_chain_path, and no per-component init_from. + cmb = _make_named_component_cfg("cmb", "IQU") + object.__setattr__(cmb, "_name", "cmb") + params = Bunch(general=general, components=Bunch({"cmb": cmb})) + + comp_list = CompList.init_from_params(params.components, params) + comp_list.load_initial_alms(params) + + assert all(np.all(comp.alms == 0) for comp in comp_list) + + +def test_load_initial_alms_from_fits_map(tmp_path) -> None: + nside = 2 + lmax = 3 + npix = 12 * nside**2 + iqu_map = np.zeros((3, npix), dtype=np.float64) + iqu_map[0] = 1.0 + np.arange(npix) # Distinct I, Q, U so a wrong row would be detectable. + iqu_map[1] = 2.0 + iqu_map[2] = 3.0 + fits_path = tmp_path / "init_map.fits" + hp.write_map(str(fits_path), iqu_map, overwrite=True, dtype=np.float64) + + general = _make_general() + general.CG_float_precision = "double" # So component alms match map_to_alm output exactly. + cmb = _make_named_component_cfg("cmb", "IQU") + cmb.params.lmax = lmax + object.__setattr__(cmb, "_name", "cmb") + cmb.params.init_from = str(fits_path) + params = Bunch(general=general, components=Bunch({"cmb": cmb})) + + comp_list = CompList.init_from_params(params.components, params) + comp_list.load_initial_alms(params) + + views = {(comp.comp_name, comp.eval_pol): comp for comp in comp_list} + expected_I = map_to_alm(np.ascontiguousarray(iqu_map[0:1]), nside, lmax, spin=0) + expected_QU = map_to_alm(np.ascontiguousarray(iqu_map[1:3]), nside, lmax, spin=2) + assert np.allclose(views[("cmb", "I")].alms, expected_I) + assert np.allclose(views[("cmb", "QU")].alms, expected_QU) + + +def test_load_initial_alms_rejects_unknown_extension(tmp_path) -> None: + general = _make_general() + cmb = _make_named_component_cfg("cmb", "IQU") + object.__setattr__(cmb, "_name", "cmb") + cmb.params.init_from = str(tmp_path / "init_map.txt") + params = Bunch(general=general, components=Bunch({"cmb": cmb})) + + comp_list = CompList.init_from_params(params.components, params) + with pytest.raises(AssertionError): + comp_list.load_initial_alms(params) + + +def test_load_initial_alms_partial_source_leaves_missing_pol_zero(tmp_path) -> None: + # An intensity-only chain initializes the I view; the IQU component's QU view stays at zero + # rather than erroring (so e.g. I-from-chain + QU-from-zero works). + nalm = (1 + 1) * (1 + 2) // 2 + cmb_intensity_only = (np.arange(nalm, dtype=np.float64) + 1).reshape(1, nalm).astype(np.complex64) + chain = tmp_path / "intensity_only.h5" + _write_chain_alms(chain, {"cmb": cmb_intensity_only}) + + general = _make_general() + general.init_chain_path = str(chain) + cmb = _make_named_component_cfg("cmb", "IQU") + object.__setattr__(cmb, "_name", "cmb") + params = Bunch(general=general, components=Bunch({"cmb": cmb})) + + comp_list = CompList.init_from_params(params.components, params) + comp_list.load_initial_alms(params) + + views = {(comp.comp_name, comp.eval_pol): comp for comp in comp_list} + assert np.array_equal(views[("cmb", "I")].alms, cmb_intensity_only) + assert np.all(views[("cmb", "QU")].alms == 0) + + +def test_load_initial_alms_missing_component_logs_error_and_continues(tmp_path, caplog) -> None: + nalm = (1 + 1) * (1 + 2) // 2 + chain = tmp_path / "other_components.h5" + _write_chain_alms(chain, {"dust": np.ones((3, nalm), dtype=np.complex64)}) # no "cmb" entry + + general = _make_general() + general.init_chain_path = str(chain) + cmb = _make_named_component_cfg("cmb", "IQU") + object.__setattr__(cmb, "_name", "cmb") + params = Bunch(general=general, components=Bunch({"cmb": cmb})) + + comp_list = CompList.init_from_params(params.components, params) + with caplog.at_level("ERROR"): + comp_list.load_initial_alms(params) # must not raise + + assert all(np.all(comp.alms == 0) for comp in comp_list) + assert "not found" in caplog.text + + +def test_build_initial_sky_model_returns_realizable_model() -> None: + general = _make_general() # No init paths -> zero alms -> zero sky. + cmb = _make_named_component_cfg("cmb", "IQU") + cmb.params.lmax = 2 # Spin-2 (QU) synthesis requires lmax >= 2. + object.__setattr__(cmb, "_name", "cmb") + params = Bunch(general=general, components=Bunch({"cmb": cmb})) + + sky = build_initial_sky_model(params) + realized = sky.get_sky_at_nu(100.0, 2, "IQU", fwhm=0.0) + + assert realized.shape == (3, 12 * 2**2) + assert np.all(realized == 0) + + +def _dust_params(**overrides) -> Bunch: + params = Bunch(polarization="IQU", shortname="dust", spatially_varying_MM=False, + smoothing_prior_FWHM=0.0, smoothing_prior_amplitude=1.0, lmax=2, + beta=1.5, T=20.0, nu_ref=[857.0, 353.0], units="uK_RJ") + for key, value in overrides.items(): + params[key] = value + object.__setattr__(params, "_name", "dust") + return params + + +def test_diffuse_component_resolves_per_pol_reference_frequency() -> None: + general = _make_general() + dust_I = ThermalDust(_dust_params(), general, eval_pol="I", comp_name="dust") + dust_QU = ThermalDust(_dust_params(), general, eval_pol="QU", comp_name="dust") + + # nu_ref = [nu_I, nu_QU]: each view picks its own entry. + assert dust_I.nu_ref == 857.0 + assert dust_QU.nu_ref == 353.0 + # The SED is normalized to 1 at each view's own reference frequency. + assert np.isclose(dust_I.get_sed(857.0), 1.0) + assert np.isclose(dust_QU.get_sed(353.0), 1.0) + + +def test_scalar_reference_frequency_is_shared_by_both_polarizations() -> None: + general = _make_general() + dust_I = ThermalDust(_dust_params(nu_ref=545.0), general, eval_pol="I", comp_name="dust") + dust_QU = ThermalDust(_dust_params(nu_ref=545.0), general, eval_pol="QU", comp_name="dust") + + # A scalar nu_ref applies to both I and QU. + assert dust_I.nu_ref == 545.0 and dust_QU.nu_ref == 545.0 + + +def test_init_map_to_amplitude_is_noop_when_units_match() -> None: + general = _make_general() + dust = ThermalDust(_dust_params(units="uK_RJ"), general, eval_pol="I", comp_name="dust") + arr = np.ones((1, 12)) + + # uK_RJ already equals the dust amplitude unit, so the same array is returned untouched. + assert dust.init_map_to_amplitude(arr) is arr + + +def test_init_map_to_amplitude_converts_to_amplitude_unit() -> None: + import pysm3.units as pysm3u + + general = _make_general() + dust = ThermalDust(_dust_params(units="uK_CMB", nu_ref=100.0), general, + eval_pol="I", comp_name="dust") + expected = (1 * pysm3u.Unit("uK_CMB")).to( + pysm3u.uK_RJ, equivalencies=pysm3u.cmb_equivalencies(100.0 * pysm3u.GHz)).value + + out = dust.init_map_to_amplitude(np.ones((1, 12))) + + assert np.allclose(out, expected) + + +def test_cmb_init_map_to_amplitude_converts_uK_CMB_at_nu_ref() -> None: + import pysm3.units as pysm3u + + general = _make_general() + cmb_params = _make_component_cfg("IQU").params + cmb_params.units = "uK_CMB" + cmb = CMB(cmb_params, general, eval_pol="I", comp_name="cmb") + + # CMB amplitudes are stored in uK_RJ referenced to nu_ref (default 1 GHz), so a uK_CMB disk map + # is converted to uK_RJ at nu_ref. + assert cmb.nu_ref == 1.0 + expected = (1 * pysm3u.Unit("uK_CMB")).to( + pysm3u.uK_RJ, equivalencies=pysm3u.cmb_equivalencies(1.0 * pysm3u.GHz)).value + out = cmb.init_map_to_amplitude(np.ones((1, 12))) + assert np.allclose(out, expected) + + +def test_cmb_get_sed_is_unity_at_nu_ref_and_ratio_elsewhere() -> None: + import pysm3.units as pysm3u + + general = _make_general() + cmb_params = _make_component_cfg("IQU").params + cmb_params.nu_ref = 100.0 + cmb = CMB(cmb_params, general, eval_pol="I", comp_name="cmb") + + # The SED is normalized to 1 at the reference frequency and is the ratio of the + # thermodynamic-to-RJ conversion elsewhere. + assert np.isclose(cmb.get_sed(100.0), 1.0) + def g(f): + return (1 * pysm3u.uK_CMB).to( + pysm3u.uK_RJ, equivalencies=pysm3u.cmb_equivalencies(f * pysm3u.GHz)).value + assert np.isclose(cmb.get_sed(353.0), g(353.0) / g(100.0)) \ No newline at end of file diff --git a/tests/test_compsep_framework.py b/tests/test_compsep_framework.py new file mode 100644 index 0000000..fc0974a --- /dev/null +++ b/tests/test_compsep_framework.py @@ -0,0 +1,232 @@ +import numpy as np +import pytest +from pixell.bunch import Bunch + +from commander4.compsep_processing import ( + _filter_sampling_group_components, + _sampling_group_selection, + _sampling_group_selects_band, + _validate_sampling_groups, + init_compsep_processing, +) +from commander4.communication import _get_compsep_sender_id_for_tod_band, _should_send_compsep_result +from commander4.sky_models.component import CompList +from commander4.sky_models.sky_model import SkyModel +from commander4.utils.execution_ids import get_execution_band_id, get_execution_band_ids + + +def _make_general(ntask_compsep_qu: int = 1, ntask_compsep_i: int = 1) -> Bunch: + return Bunch( + nside=2, + CG_float_precision="single", + MPI_config=Bunch(ntask_compsep_I=ntask_compsep_i, ntask_compsep_QU=ntask_compsep_qu), + ) + + +def _make_component_cfg(polarization: str = "IQU") -> Bunch: + return Bunch( + enabled=True, + component_class="CMB", + params=Bunch( + lmax=2, + polarization=polarization, + shortname="cmb", + spatially_varying_MM=False, + smoothing_prior_FWHM=0.0, + smoothing_prior_amplitude=1.0, + ), + ) + + +def _make_comp_list(polarization: str = "IQU", ntask_compsep_qu: int = 1) -> CompList: + params = Bunch(general=_make_general(ntask_compsep_qu)) + cmb = _make_component_cfg(polarization) + object.__setattr__(cmb, "_name", "cmb") + components = Bunch({"cmb": cmb}) + return CompList.init_from_params(components, params) + + +def _make_multi_comp_list() -> CompList: + params = Bunch(general=_make_general()) + cmb = _make_component_cfg("IQU") + object.__setattr__(cmb, "_name", "CMB") + dust = Bunch( + enabled=True, + component_class="CMB", + params=Bunch( + lmax=2, + polarization="IQU", + shortname="dust", + spatially_varying_MM=False, + smoothing_prior_FWHM=0.0, + smoothing_prior_amplitude=1.0, + ), + ) + object.__setattr__(dust, "_name", "ThermalDust") + ff = Bunch( + enabled=True, + component_class="CMB", + params=Bunch( + lmax=2, + polarization="I", + shortname="ff", + spatially_varying_MM=False, + smoothing_prior_FWHM=0.0, + smoothing_prior_amplitude=1.0, + ), + ) + object.__setattr__(ff, "_name", "FreeFree") + components = Bunch( + { + "CMB": cmb, + "ThermalDust": dust, + "FreeFree": ff, + } + ) + return CompList.init_from_params(components, params) + + +def test_execution_band_id_helpers_use_plain_band_names() -> None: + assert get_execution_band_id("90GHz", "I") == "90GHz_I" + assert get_execution_band_id("353GHz", "QU") == "353GHz_QU" + assert get_execution_band_ids("90GHz", "IQU") == ("90GHz_I", "90GHz_QU") + + +def test_tod_receive_source_prefers_intensity_and_falls_back_to_qu() -> None: + senders = {"30GHz_I": 3, "30GHz_QU": 4} + assert _get_compsep_sender_id_for_tod_band("30GHz", senders) == "30GHz_I" + + senders = {"Planck353GHz_QU": 8} + assert _get_compsep_sender_id_for_tod_band("Planck353GHz", senders) == "Planck353GHz_QU" + + with pytest.raises(KeyError, match="No CompSep sender"): + _get_compsep_sender_id_for_tod_band("MissingBand", {}) + + +def test_should_send_compsep_result_skips_qu_when_i_sender_exists() -> None: + destinations = {"LFT_I": 1, "LFT_QU": 2} + + assert _should_send_compsep_result("LFT_I", destinations) + assert not _should_send_compsep_result("LFT_QU", destinations) + assert _should_send_compsep_result("Planck353GHz_QU", {"Planck353GHz_QU": 5}) + assert not _should_send_compsep_result("Unused_I", None) + + +def test_joined_skymodel_realizes_iqu_components() -> None: + comp_list = _make_comp_list("IQU") + joined = comp_list.joined() + + sky = SkyModel(joined).get_sky_at_nu(30.0, 2, "IQU", fwhm=0.0) + + assert sky.shape == (3, 48) + assert np.all(np.isfinite(sky)) + + +def test_sampling_group_component_filter_matches_comp_names_and_preserves_pol_split() -> None: + comp_list = _make_multi_comp_list() + + intensity = _filter_sampling_group_components( + comp_list.split_for_eval_pol("I"), + ["ThermalDust", "FreeFree"], + ) + polarization = _filter_sampling_group_components( + comp_list.split_for_eval_pol("QU"), + ["ThermalDust", "FreeFree"], + ) + + assert [comp.comp_name for comp in intensity] == ["ThermalDust", "FreeFree"] + assert [comp.eval_pol for comp in intensity] == ["I", "I"] + assert [comp.comp_name for comp in polarization] == ["ThermalDust"] + assert [comp.eval_pol for comp in polarization] == ["QU"] + + +def test_sampling_group_band_filter_accepts_base_and_execution_ids() -> None: + assert _sampling_group_selects_band(["Planck30GHz"], "Planck30GHz", "Planck30GHz_I") + assert _sampling_group_selects_band(["Planck30GHz_I"], "Planck30GHz", "Planck30GHz_I") + assert not _sampling_group_selects_band(["Planck44GHz"], "Planck30GHz", "Planck30GHz_I") + + +def test_validate_sampling_groups_rejects_unknown_names(caplog) -> None: + comp_list = _make_multi_comp_list() # CMB, ThermalDust, FreeFree + params = Bunch(CompSep_bands=Bunch({"Planck30GHz": Bunch(enabled=True, polarization="IQU")})) + + # Valid references (component name + execution-view band id) pass silently. + _validate_sampling_groups( + Bunch(g=Bunch(comps=["CMB"], bands=["Planck30GHz_QU"])), comp_list, params) + + with pytest.raises(AssertionError): + _validate_sampling_groups(Bunch(g=Bunch(comps=["DoesNotExist"])), comp_list, params) + assert "unknown component" in caplog.text + + with pytest.raises(AssertionError): + _validate_sampling_groups(Bunch(g=Bunch(bands=["NoSuchBand"])), comp_list, params) + assert "unknown band" in caplog.text + + +def test_sampling_group_selection_resolves_all_and_missing() -> None: + # Missing entry and the literal "all" both mean "everything" (None); a list is returned as-is. + assert _sampling_group_selection(Bunch(), "comps") is None + assert _sampling_group_selection(Bunch(comps="all"), "comps") is None + assert _sampling_group_selection(Bunch(bands="all"), "bands") is None + assert _sampling_group_selection(Bunch(comps=["CMB"]), "comps") == ["CMB"] + assert _sampling_group_selection(Bunch(bands=["Planck30GHz"]), "bands") == ["Planck30GHz"] + + +def test_validate_sampling_groups_accepts_all_and_missing() -> None: + comp_list = _make_multi_comp_list() + params = Bunch(CompSep_bands=Bunch({"Planck30GHz": Bunch(enabled=True, polarization="IQU")})) + + # "all" and omitted entries select everything and must not be checked against names. + _validate_sampling_groups(Bunch(g=Bunch(comps="all", bands="all")), comp_list, params) + _validate_sampling_groups(Bunch(g=Bunch()), comp_list, params) + + +def test_validate_sampling_groups_skips_disabled_group() -> None: + comp_list = _make_multi_comp_list() + params = Bunch(CompSep_bands=Bunch({"Planck30GHz": Bunch(enabled=True, polarization="IQU")})) + + # A disabled group is not validated, so its bogus references are tolerated. + _validate_sampling_groups( + Bunch(g=Bunch(enabled=False, comps=["DoesNotExist"], bands=["NoSuchBand"])), + comp_list, params) + + +def test_comp_name_comes_from_component_bunch_name() -> None: + params = Bunch(general=_make_general()) + component_cfg = _make_component_cfg("IQU") + object.__setattr__(component_cfg, "_name", "CMBFromName") + components = Bunch({"cmb": component_cfg}) + + comp_list = CompList.init_from_params(components, params) + + assert [comp.comp_name for comp in comp_list] == ["CMBFromName", "CMBFromName"] + + +def test_init_compsep_processing_rejects_duplicate_component_names(monkeypatch, caplog) -> None: + class _FakeCompList: + def joined(self): + return [Bunch(comp_name="dup"), Bunch(comp_name="dup")] + + class _FakeComm: + def allgather(self, data): + return [data] + + monkeypatch.setattr(CompList, "init_from_params", classmethod(lambda cls, *_: _FakeCompList())) + + mpi_info = Bunch( + processor_name="test-node", + world=Bunch(rank=0), + compsep=Bunch(rank=0, QU_master=1, size=1, comm=_FakeComm()), + ) + params = Bunch( + components=Bunch(), + CompSep_bands=Bunch( + { + "BandA": Bunch(enabled=True, polarization="I", get_from="file"), + } + ), + ) + + with pytest.raises(AssertionError): + init_compsep_processing(mpi_info, params) + assert "Duplicate component names found in CompSep setup" in caplog.text \ No newline at end of file diff --git a/tests/test_gain_calibration.py b/tests/test_gain_calibration.py new file mode 100644 index 0000000..3dbbcc0 --- /dev/null +++ b/tests/test_gain_calibration.py @@ -0,0 +1,283 @@ +"""Tests for the gain-calibration configuration and the unified calibrator builder. + +Covers two new pieces introduced when the three gain-sampling procedures became nested +parameter-file blocks with a per-term ``calibrate_against`` target: + +* ``_resolve_calib_target`` - resolves a gain term's calibrator, with a per-band override + taking precedence over the general-block value, which falls back to the term default. +* ``TODView.get_calib_tod`` - builds the calibration residual for one gain term against a + chosen calibrator signal, replacing the former per-term ``get_*_calib_tod`` methods. +""" + +from types import SimpleNamespace + +import numpy as np +import pytest +from pixell.bunch import Bunch + +from commander4.data_models.tod_view import TODView +from commander4.tod_processing import (_resolve_calib_target, _DEFAULT_CALIB_TARGETS, + _VALID_CALIB_TARGETS, _solve_relative_gain_system, + _resolve_gain_downsample_factor) + + +# -------------------------------------------------------------------------------------- +# _resolve_calib_target +# -------------------------------------------------------------------------------------- +def _make_params(general_blocks: dict, band_blocks: dict) -> Bunch: + """Build a params Bunch with the given general and per-band gain blocks.""" + return Bunch( + general=Bunch(**{name: Bunch(**vals) for name, vals in general_blocks.items()}), + experiments=Bunch(EXP=Bunch(bands=Bunch( + BAND=Bunch(**{name: Bunch(**vals) for name, vals in band_blocks.items()})))), + ) + + +def _exp_data(band="BAND"): + return SimpleNamespace(experiment_name="EXP", band_name=band) + + +def test_defaults_when_calibrate_against_absent(): + # No calibrate_against anywhere -> each term falls back to its documented default. + params = _make_params({"abs_gain": {}, "rel_gain": {}, "temporal_gain": {}}, {}) + for block, default in _DEFAULT_CALIB_TARGETS.items(): + assert _resolve_calib_target(params, _exp_data(), block) == default + assert _DEFAULT_CALIB_TARGETS["abs_gain"] == "orbital_dipole" + assert _DEFAULT_CALIB_TARGETS["rel_gain"] == "full_sky" + + +def test_general_block_value_used(): + params = _make_params({"abs_gain": {"calibrate_against": "full_sky"}}, {}) + assert _resolve_calib_target(params, _exp_data(), "abs_gain") == "full_sky" + + +def test_band_override_beats_general_and_default(): + # General says full_sky, band overrides to sky -> band wins. + params = _make_params({"abs_gain": {"calibrate_against": "full_sky"}}, + {"abs_gain": {"calibrate_against": "sky"}}) + assert _resolve_calib_target(params, _exp_data(), "abs_gain") == "sky" + # A band with no override block falls back to the general value. + params2 = _make_params({"abs_gain": {"calibrate_against": "full_sky"}}, {}) + assert _resolve_calib_target(params2, _exp_data(), "abs_gain") == "full_sky" + + +def test_invalid_target_raises(): + params = _make_params({"abs_gain": {"calibrate_against": "bogus"}}, {}) + with pytest.raises(ValueError): + _resolve_calib_target(params, _exp_data(), "abs_gain") + + +def test_valid_targets_contents(): + assert set(_VALID_CALIB_TARGETS) == {"orbital_dipole", "full_sky", "sky"} + + +# -------------------------------------------------------------------------------------- +# TODView.get_calib_tod +# -------------------------------------------------------------------------------------- +class _StubView(TODView): + """A TODView whose data accessors are stubbed so get_calib_tod can be tested in + isolation: it records the ``subtract`` spec passed to ``get_tod`` and supplies fixed + sky / orbital-dipole signals.""" + + def __init__(self, s_sky, s_orb): + super().__init__(None, None) + self._s_sky = s_sky + self._s_orb = s_orb + self.captured_subtract = None + + def _materialize_downsampled(self, factor): + n = self._s_sky.size + return Bunch(tod=np.zeros(n), pix=np.zeros(n, dtype=int), psi=np.zeros(n)) + + def get_mask(self, which, downsample_factor=1): + return np.ones(self._s_sky.size, dtype=bool) + + def get_static_sky_tod(self, compsep_output=None, downsample_factor=None): + return self._s_sky + + def get_orbital_dipole_tod(self, downsample_factor=None): + return self._s_orb + + def get_tod(self, *, subtract=None, downsample_factor=1, compsep_output=None, **kw): + self.captured_subtract = subtract + return np.zeros(self._s_sky.size) + + +def _make_stub(): + return _StubView(np.array([1.0, 2.0, 3.0, 4.0]), np.array([10.0, 20.0, 30.0, 40.0])) + + +ALL = ("abs", "rel", "temp") + + +@pytest.mark.parametrize("target,calib,expected_subtract,scal", [ + # Absolute gain on the orbital dipole: sky removed entirely, dipole keeps the abs term. + ("abs", "orbital_dipole", + (("sky", ALL), ("orbital_dipole", ("rel", "temp"))), "orb"), + # Absolute gain on the full sky (clean target-gain-preserving form): both signals keep abs. + ("abs", "full_sky", + (("sky", ("rel", "temp")), ("orbital_dipole", ("rel", "temp"))), "sky+orb"), + # Relative gain on the full sky: both signals keep the rel term. + ("rel", "full_sky", + (("sky", ("abs", "temp")), ("orbital_dipole", ("abs", "temp"))), "sky+orb"), + # Temporal gain on the full sky: both signals keep the temp term. + ("temp", "full_sky", + (("sky", ("abs", "rel")), ("orbital_dipole", ("abs", "rel"))), "sky+orb"), + # Absolute gain on the static sky only: dipole removed entirely, sky keeps abs. + ("abs", "sky", + (("sky", ("rel", "temp")), ("orbital_dipole", ALL)), "sky"), +]) +def test_get_calib_tod_builds_residual(target, calib, expected_subtract, scal): + view = _make_stub() + out = view.get_calib_tod(target, calib, downsample_factor=1, fill_masked=False) + assert view.captured_subtract == expected_subtract + expected_scal = {"orb": view._s_orb, "sky": view._s_sky, + "sky+orb": view._s_sky + view._s_orb}[scal] + np.testing.assert_allclose(out.s_cal, expected_scal) + + +def test_get_calib_tod_rejects_bad_arguments(): + view = _make_stub() + with pytest.raises(ValueError): + view.get_calib_tod("bogus", "full_sky", downsample_factor=1) + with pytest.raises(ValueError): + view.get_calib_tod("abs", "bogus", downsample_factor=1) + + +# -------------------------------------------------------------------------------------- +# _solve_relative_gain_system (reduced constrained solve + bad-detector exclusion) +# -------------------------------------------------------------------------------------- +_ZERO_RNG = SimpleNamespace(standard_normal=lambda n: np.zeros(n)) # disables the fluctuation term + + +def test_relgain_recovers_constrained_mean(): + # With the fluctuation term zeroed, the solve must reproduce the analytic constrained mean + # g_i = (r_i - 0.5*lambda)/d_i with lambda set so the active gains sum to zero. + s = np.array([2.0, 4.0, 1.0, 3.0]) + r = np.array([1.0, -2.0, 0.5, 0.7]) + out = _solve_relative_gain_system(s, r, np.zeros(4), rng=_ZERO_RNG) + lam = 2.0 * np.sum(r / s) / np.sum(1.0 / s) + expected = (r - 0.5 * lam) / s + np.testing.assert_allclose(out, expected, rtol=1e-5, atol=1e-6) + assert abs(out.sum()) < 1e-5 # zero-sum constraint over all (active) detectors + + +def test_relgain_excludes_zero_weight_detectors(): + # Two detectors have zero calibration weight -> the *full* bordered system is singular, but the + # helper solves the reduced active system and holds the excluded detectors at their prev value. + s = np.array([0.0, 2.0, 0.0, 3.0]) + r = np.array([5.0, 1.0, 9.0, -1.0]) + prev = np.array([0.11, 0.22, 0.33, 0.44], dtype=np.float32) + + # Sanity: the un-reduced 4-detector bordered system really is singular. + n = 4 + A_full = np.zeros((n + 1, n + 1)) + A_full[:n, :n] = np.diag(s); A_full[:n, n] = 0.5; A_full[n, :n] = 1.0 + with pytest.raises(np.linalg.LinAlgError): + np.linalg.solve(A_full, np.ones(n + 1)) + + out = _solve_relative_gain_system(s, r, prev, rng=_ZERO_RNG) + assert np.all(np.isfinite(out)) + assert out[0] == np.float32(0.11) and out[2] == np.float32(0.33) # excluded held at prev + assert abs(out[[1, 3]].sum()) < 1e-5 # active subset sums to zero + + +def test_relgain_single_zero_detector_is_held(): + s = np.array([0.0, 2.0, 5.0]) + r = np.array([3.0, 1.0, -1.0]) + prev = np.array([0.9, 0.0, 0.0], dtype=np.float32) + out = _solve_relative_gain_system(s, r, prev, rng=_ZERO_RNG) + assert out[0] == np.float32(0.9) + assert abs(out[[1, 2]].sum()) < 1e-5 + + +def test_relgain_no_active_detectors_returns_prev_unchanged(): + prev = np.array([1.0, 2.0, 3.0], dtype=np.float32) + out = _solve_relative_gain_system(np.zeros(3), np.zeros(3), prev) + np.testing.assert_array_equal(out, prev) + + +def test_relgain_deterministic_with_seeded_rng(): + s = np.array([2.0, 3.0, 4.0]); r = np.array([0.5, -0.5, 0.1]); prev = np.zeros(3) + a = _solve_relative_gain_system(s, r, prev, rng=np.random.default_rng(5)) + b = _solve_relative_gain_system(s, r, prev, rng=np.random.default_rng(5)) + np.testing.assert_array_equal(a, b) + + +# -------------------------------------------------------------------------------------- +# _resolve_gain_downsample_factor (general.gain_calib_downsample_time, in seconds) +# -------------------------------------------------------------------------------------- +def test_gain_downsample_factor_from_time(): + make = lambda t: Bunch(general=Bunch(gain_calib_downsample_time=t)) + exp = SimpleNamespace(fsamp=200.0) + assert _resolve_gain_downsample_factor(make(1.0), exp) == 200 + assert _resolve_gain_downsample_factor(make(0.25), exp) == 50 + assert _resolve_gain_downsample_factor(make(0.0), exp) == 1 # 0 disables downsampling. + assert _resolve_gain_downsample_factor(make(0.001), exp) == 1 # Clamped to at least 1. + assert _resolve_gain_downsample_factor(make(1.0), SimpleNamespace(fsamp=32.51)) == 33 + + +# -------------------------------------------------------------------------------------- +# Downsampling: model TODs are block-averaged like the data, not block-center sampled +# -------------------------------------------------------------------------------------- +NTOD, FACTOR = 12, 4 +# arange(0, ntod, factor) defines the block edges and the midpoint construction keeps the +# ntod//factor - 1 leading complete blocks (the trailing block is dropped). +NBLOCKS = NTOD // FACTOR - 1 + + +def _make_real_view(monkeypatch): + """A TODView over a minimal fake detector, exercising the real downsampling code paths.""" + monkeypatch.setenv("OMP_NUM_THREADS", "1") # Required by get_s_orb_TOD. + rng = np.random.default_rng(7) + pix = rng.integers(0, 12, size=NTOD) # Valid pixels for the nside=1 experiment below. + psi = rng.uniform(0.0, np.pi, size=NTOD) + det = SimpleNamespace(tod=rng.normal(size=NTOD), ntod=NTOD, fsamp=float(FACTOR), + get_pix_psi=lambda: (pix, psi), + orb_dir_vec=np.array([1.0, 0.0, 0.0], dtype=np.float32)) + experiment_data = SimpleNamespace(scans=[SimpleNamespace(detectors=[det])], nside=1, nu=30.0) + no_jump = SimpleNamespace(is_empty=lambda: True) + tod_samples = SimpleNamespace(jumps=SimpleNamespace(get=lambda iscan, idet: no_jump), + abs_gain=2.0, rel_gain=np.array([0.5]), + temporal_gain=np.array([[0.25]]), + accept=np.ones((1, 1), dtype=bool)) + skymap = rng.normal(size=(3, 12)) + view = TODView(experiment_data, tod_samples, compsep_output=skymap).focus(0, 0) + s_full = skymap[0, pix] + np.cos(2*psi)*skymap[1, pix] + np.sin(2*psi)*skymap[2, pix] + return view, det, s_full + + +def _block_mean(arr): + return arr[:NBLOCKS*FACTOR].reshape(NBLOCKS, FACTOR).mean(axis=-1) + + +def test_static_sky_downsampling_is_block_average(monkeypatch): + view, _, s_full = _make_real_view(monkeypatch) + out = view.get_static_sky_tod(downsample_factor=FACTOR) + np.testing.assert_allclose(out, _block_mean(s_full), rtol=2e-5, atol=1e-6) + # Regression guard: must NOT be the model sampled at the block-center pixels. + block_centers = np.array([2, 6]) + assert not np.allclose(out, s_full[block_centers]) + + +def test_data_and_model_share_block_definition(monkeypatch): + view, det, _ = _make_real_view(monkeypatch) + np.testing.assert_allclose(view.get_tod(downsample_factor=FACTOR), _block_mean(det.tod)) + + +def test_orbital_dipole_downsampling_is_block_average(monkeypatch): + view, _, _ = _make_real_view(monkeypatch) + orb_full = view.get_orbital_dipole_tod() + np.testing.assert_allclose(view.get_orbital_dipole_tod(downsample_factor=FACTOR), + _block_mean(orb_full), rtol=2e-5, atol=1e-9) + + +def test_get_calib_tod_downsampled_end_to_end(monkeypatch): + # Absolute gain against the static sky: residual = - (g_rel+g_temp) - g_all*, + # with every term block-averaged with the same kernel. + view, det, s_full = _make_real_view(monkeypatch) + orb_full = view.get_orbital_dipole_tod() + out = view.get_calib_tod("abs", "sky", downsample_factor=FACTOR, fill_masked=False) + np.testing.assert_allclose(out.s_cal, _block_mean(s_full), rtol=2e-5, atol=1e-6) + expected = _block_mean(det.tod) - 0.75*_block_mean(s_full) - 2.75*_block_mean(orb_full) + np.testing.assert_allclose(out.tod, expected, rtol=2e-5, atol=1e-6) diff --git a/tests/test_math_utils.py b/tests/test_math_utils.py deleted file mode 100644 index 59b3519..0000000 --- a/tests/test_math_utils.py +++ /dev/null @@ -1,80 +0,0 @@ -import numpy as np -import pytest -import sys -import os - -# Ugly trick to be able to import. -module_root_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.append(module_root_path) - -from src.python.utils.math_operations import calculate_sigma0, inplace_scale_add, inplace_add_scaled_vec,\ - inplace_arr_prod, inplace_scale, dot - - -### calculate_sigma0(tod, mask) ### -def test_sigma0(): - arr = np.random.normal(2.0, 3.0, 10) - mask = np.random.randint(0, 2, 10).astype(bool) - masked_arr = arr[mask] - numpy_sol = np.std(masked_arr[1:] - masked_arr[:-1])/np.sqrt(2) - assert calculate_sigma0(arr, mask) == pytest.approx(numpy_sol) - -def test_sigma0_masked(): - arr = np.random.normal(2.0, 3.0, 10) - mask = np.random.randint(0, 2, 10).astype(bool) - mask = np.zeros(10).astype(bool) - assert calculate_sigma0(arr, mask) == np.inf - mask[3] = True - assert calculate_sigma0(arr, mask) == np.inf - - -### dot(arr1, arr2) ### -def test_dot(): - arr1 = np.random.normal(2.0, 3.0, 100) - arr2 = np.random.normal(-3.5, 1.0, 100) - assert dot(arr1, arr2) == pytest.approx(np.dot(arr1, arr2)) - -def test_dot_f32_f64(): - arr1 = np.random.normal(2.0, 3.0, 100).astype(np.float32, copy=False) - arr2 = np.random.normal(-3.5, 1.0, 100) - assert dot(arr1, arr2) == pytest.approx(np.dot(arr1, arr2)) - -def test_dot_ndim(): - arr1 = np.random.normal(2.0, 3.0, (10,20)) - arr2 = np.random.normal(-3.5, 1.0, (10,20)) - assert dot(arr1, arr2) == pytest.approx(np.dot(arr1.flatten(), arr2.flatten())) - - -### inplace_scale_add(arr_main, arr_add, float_mult) ### -def test_inplace_add_scaled_vec(): - arr1 = np.random.normal(2.0, 3.0, 100) - arr2 = np.random.normal(-3.5, 1.0, 100) - value = -12.5 - answer = arr1 + value*arr2 - inplace_add_scaled_vec(arr1, arr2, value) - assert arr1 == pytest.approx(answer) - -def test_inplace_add_scaled_vec_integer(): - arr1 = np.random.normal(2.0, 3.0, 100) - arr2 = np.random.normal(-3.5, 1.0, 100) - value = 5 - answer = arr1 + value*arr2 - inplace_add_scaled_vec(arr1, arr2, value) - assert arr1 == pytest.approx(answer) - -def test_inplace_add_scaled_vec_f32_f64(): - arr1 = np.random.normal(2.0, 3.0, 100).astype(np.float32, copy=False) - arr2 = np.random.normal(-3.5, 1.0, 100) - value = 5 - answer = arr1 + value*arr2 - inplace_add_scaled_vec(arr1, arr2, value) - assert arr1 == pytest.approx(answer) - - -### inplace_arr_prod(arr_main, arr_prod) ### -def test_inplace_arr_prod(): - arr1 = np.random.normal(2.0, 3.0, 100) - arr2 = np.random.normal(-3.5, 1.0, 100) - answer = arr1*arr2 - inplace_arr_prod(arr1, arr2) - assert arr1 == pytest.approx(answer) diff --git a/tests/test_noise_sampling.py b/tests/test_noise_sampling.py index 9771004..212b5ef 100644 --- a/tests/test_noise_sampling.py +++ b/tests/test_noise_sampling.py @@ -1,7 +1,45 @@ import numpy as np import pytest +from numba import njit +from scipy.fft import rfftfreq from commander4.noise_sampling.sigma0 import calc_sigma0_simple, calc_sigma0_robust +from commander4.noise_sampling.noise_psd import NoisePSDOof +from commander4.noise_sampling.noise_sampling import fill_all_masked +from commander4.noise_sampling.sample_ncorr import (sample_correlated_noise, + corr_noise_realization_with_gaps) +from commander4.data_models.detector_group_TOD import DetGroupTOD + + +@njit +def _seed_numba_rng(seed: int) -> None: + """Seed Numba's internal RNG (independent of NumPy's global RNG, which is used by njit code).""" + np.random.seed(seed) + + +def _seed_all_rng(seed: int) -> None: + """Seed both the NumPy global RNG and Numba's internal RNG for fully reproducible sampling.""" + np.random.seed(seed) + _seed_numba_rng(seed) + + +def _synth_corr_noise(seed: int, ntod: int, fsamp: float, sigma0: float, fknee: float, + alpha: float) -> np.ndarray: + """Synthesize a correlated-noise stream whose periodogram follows sigma0^2 (f/fknee)^alpha. + + Uses the rfft/irfft convention (forward unnormalized, inverse divided by ntod) matching the + project's FFT helpers, so that (1/ntod)|rfft(n_corr)|^2 has expectation P_corr(f). + """ + rng = np.random.default_rng(seed) + freqs = np.fft.rfftfreq(ntod, d=1.0/fsamp) + P = np.zeros_like(freqs) + P[1:] = sigma0**2 * (freqs[1:]/fknee)**alpha + spec = np.sqrt(ntod*P/2.0) * (rng.standard_normal(freqs.size) + + 1j*rng.standard_normal(freqs.size)) + spec[0] = 0.0 # zero-mean + if ntod % 2 == 0: + spec[-1] = spec[-1].real # real Nyquist mode + return np.fft.irfft(spec, n=ntod) # =================================================================== # sigma0 estimation @@ -60,3 +98,204 @@ def test_few_samples(self): # With < 100 valid pairs, should return inf est = calc_sigma0_robust(tod, mask) assert est == np.inf + + +# =================================================================== +# NoisePSD model (1/f / "oof") +# =================================================================== + +class TestNoisePSDOof: + def test_eval_matches_formula(self): + m = NoisePSDOof() + s0, fk, a = 2.0, 0.5, -2.3 + freqs = np.array([0.0, 0.1, 1.0, 5.0]) + full = m.eval_full(freqs, np.array([s0, fk, a])) + corr = m.eval_corr(freqs, np.array([s0, fk, a])) + # At f = 0: full PSD is the white-noise floor, correlated PSD is zero. + assert full[0] == pytest.approx(s0**2) + assert corr[0] == 0.0 + pos = freqs > 0 + assert np.allclose(corr[pos], s0**2 * (freqs[pos]/fk)**a) + assert np.allclose(full[pos], s0**2 * (1.0 + (freqs[pos]/fk)**a)) + + def test_inv_corr_spectrum_is_one_over_Pcorr(self): + """compute_inv_corr_spectrum must equal 1/P_corr (the quantity the noise CG adds to + 1/sigma0^2), matching the previously hardcoded C_1f_inv = 1/(sigma0^2 (f/fknee)^alpha).""" + m = NoisePSDOof() + s0, fk, a = 1.3, 0.4, -1.7 + freqs = rfftfreq(2048, d=1.0/10.0) + inv = m.compute_inv_corr_spectrum(freqs, np.array([s0, fk, a])) + assert inv[0] == 0.0 # zero-frequency mode excluded + expected = np.zeros_like(freqs) + expected[1:] = 1.0 / (s0**2 * (freqs[1:]/fk)**a) + assert np.allclose(inv, expected) + + def test_sample_params_keeps_sigma0_and_uses_Puni_bounds(self): + m = NoisePSDOof() + sigma0 = 1.234 + residual = _synth_corr_noise(1, 2**14, 10.0, sigma0, 0.5, -2.0) \ + + np.random.default_rng(2).normal(0.0, sigma0, 2**14) + _seed_all_rng(0) + params = np.array([sigma0, 0.3, -1.5]) + out = m.sample_params(residual, params, 10.0, nu_min=0.0, nu_max=2.0) + assert out.shape == params.shape + assert out[0] == params[0] # sigma0 held fixed + assert out is not params # returns a fresh array + assert params[1] == 0.3 and params[2] == -1.5 # input untouched + # Grids span the model's uniform priors (single source of truth for the hard bounds). + assert m.P_uni[1, 0] <= out[1] <= m.P_uni[1, 1] # fknee in P_uni + assert m.P_uni[2, 0] <= out[2] <= m.P_uni[2, 1] # alpha in P_uni ([-4.5, -0.25]) + + def test_sample_params_recovers_input(self): + # The sampler builds its (fknee, alpha) grids from the model's uniform priors, so construct + # a model whose P_uni spans the injected values rather than relying on the defaults. + m = NoisePSDOof(P_uni=[[np.nan, np.nan], [0.01, 10.0], [-4.5, -0.25]]) + sigma0, fknee, alpha = 1.0, 0.5, -2.0 + # Fit is performed on the residual (white + correlated noise) with the full PSD model. + residual = _synth_corr_noise(7, 2**16, 10.0, sigma0, fknee, alpha) \ + + np.random.default_rng(8).normal(0.0, sigma0, 2**16) + for bin_psd in (False, True): + _seed_all_rng(123) + fks, als = [], [] + for _ in range(8): + out = m.sample_params(residual, np.array([sigma0, 0.2, -1.0]), 10.0, + nu_min=0.0, nu_max=2.0, bin_psd=bin_psd) + fks.append(out[1]) + als.append(out[2]) + assert np.mean(fks) == pytest.approx(fknee, rel=0.4), f"bin_psd={bin_psd}" + assert np.mean(als) == pytest.approx(alpha, abs=0.4), f"bin_psd={bin_psd}" + + +# =================================================================== +# Correlated-noise orchestrator +# =================================================================== + +class TestSampleCorrelatedNoise: + @staticmethod + def _setup(ntod=2**13, fsamp=10.0, sigma0=1.0, fknee=0.3, alpha=-1.8, seed=3): + rng = np.random.default_rng(seed) + tod = _synth_corr_noise(seed, ntod, fsamp, sigma0, fknee, alpha) \ + + rng.normal(0.0, sigma0, ntod) + mask = np.ones(ntod, dtype=bool) + mask[1000:1050] = False # contiguous gap + mask[5000] = False # single-sample gap + return tod, mask, np.array([sigma0, fknee, alpha]), fsamp + + def test_zero_cg_steps_equals_direct_fallback(self): + """cg_max_iter=0 must skip the masked CG and return the stationary full-mask solution.""" + m = NoisePSDOof() + tod, mask, params, fsamp = self._setup() + tod_a, tod_b = tod.copy(), tod.copy() + _seed_all_rng(42) + res = sample_correlated_noise(tod_a, mask, params.copy(), m, fsamp, cg_err_tol=1e-6, + cg_max_iter=0, sample_params=False) + # Reproduce the fallback manually: inpaint, then solve with an all-valid mask. + _seed_all_rng(42) + fill_all_masked(tod_b, mask, params[0]) + C = m.compute_inv_corr_spectrum(rfftfreq(2*tod_b.size, d=1.0/fsamp), params) + ref, _, _, _ = corr_noise_realization_with_gaps(tod_b, np.ones_like(mask), params[0], C) + assert np.allclose(res.n_corr, ref) + assert res.niter == 0 + assert res.converged + assert not res.high_var + + def test_masked_cg_runs_and_converges(self): + m = NoisePSDOof() + tod, mask, params, fsamp = self._setup() + _seed_all_rng(0) + res = sample_correlated_noise(tod.copy(), mask, params.copy(), m, fsamp, cg_err_tol=1e-6, + cg_max_iter=200, sample_params=False) + assert res.n_corr.shape == tod.shape + assert np.all(np.isfinite(res.n_corr)) + assert res.converged # an easy, well-conditioned system should converge + assert res.niter > 0 # the masked CG actually ran (unlike cg_max_iter=0) + + def test_param_sampling_toggle(self): + m = NoisePSDOof() + tod, mask, params, fsamp = self._setup() + # With both sampling switches off, noise_params come back untouched. + _seed_all_rng(0) + res_off = sample_correlated_noise(tod.copy(), mask, params.copy(), m, fsamp, + cg_err_tol=1e-6, cg_max_iter=0, sample_params=False, + sample_sigma0=False) + assert np.array_equal(res_off.noise_params, params) + # Turning on parameter sampling (only) updates fknee/alpha but keeps sigma0. + _seed_all_rng(0) + res_on = sample_correlated_noise(tod.copy(), mask, params.copy(), m, fsamp, + cg_err_tol=1e-6, cg_max_iter=0, sample_params=True, + sample_sigma0=False, psd_fit_nu_max=2.0) + assert res_on.noise_params[0] == params[0] # sigma0 fixed + assert not np.array_equal(res_on.noise_params[1:], params[1:]) # fknee/alpha updated + + def test_sigma0_reestimated_from_residual(self): + """sample_sigma0=True replaces noise_params[0] with a data-driven estimate.""" + m = NoisePSDOof() + tod, mask, params, fsamp = self._setup(sigma0=1.0) + params_in = params.copy() + params_in[0] = 5.0 # deliberately wrong sigma0 going in + _seed_all_rng(1) + res = sample_correlated_noise(tod.copy(), mask, params_in, m, fsamp, cg_err_tol=1e-6, + cg_max_iter=50, sample_params=False, sample_sigma0=True) + assert res.noise_params[0] == pytest.approx(1.0, rel=0.2) # recovered true sigma0 + + def test_monopole_modes(self): + m = NoisePSDOof() + ntod, fsamp, sigma0 = 2**13, 10.0, 1.0 + rng = np.random.default_rng(11) + base = _synth_corr_noise(11, ntod, fsamp, sigma0, 0.3, -1.8) + rng.normal(0, sigma0, ntod) + offset = 7.0 + mask = np.ones(ntod, dtype=bool) + mask[2000:2050] = False + params = np.array([sigma0, 0.3, -1.8]) + + _seed_all_rng(5) + keep = sample_correlated_noise((base + offset).copy(), mask, params.copy(), m, fsamp, + cg_err_tol=1e-6, cg_max_iter=50, sample_params=False) + _seed_all_rng(5) + remove = sample_correlated_noise((base + offset).copy(), mask, params.copy(), m, fsamp, + cg_err_tol=1e-6, cg_max_iter=50, sample_params=False, + nomono=True) + _seed_all_rng(5) + only = sample_correlated_noise((base + offset).copy(), mask, params.copy(), m, fsamp, + cg_err_tol=1e-6, cg_max_iter=50, sample_params=False, + onlymono=True) + # Default leaves the offset in n_corr; nomono projects it out; onlymono returns the offset. + assert np.mean(keep.n_corr[mask]) == pytest.approx(offset, abs=0.3) + assert np.mean(remove.n_corr[mask]) == pytest.approx(0.0, abs=1e-6) + assert np.allclose(only.n_corr, np.mean((base + offset)[mask])) + assert only.niter == 0 + # onlymono takes precedence when both are set (contradictory; the caller logs an error). + _seed_all_rng(5) + both = sample_correlated_noise((base + offset).copy(), mask, params.copy(), m, fsamp, + cg_err_tol=1e-6, cg_max_iter=50, sample_params=False, + nomono=True, onlymono=True) + assert np.allclose(both.n_corr, only.n_corr) and both.niter == 0 + + +# =================================================================== +# Inverse-noise application (apply_N_inv) +# =================================================================== + +class TestApplyNInv: + @staticmethod + def _detgroup(noise_model, fsamp=10.0): + return DetGroupTOD(scans=[], experiment_name="x", band_name="b", nside=64, nu=100.0, + fwhm=30.0, fsamp=fsamp, ndet=1, pols="IQU", noise_model=noise_model) + + def test_projects_out_dc(self): + """apply_N_inv must remove the DC (mean) mode, matching Commander multiply_inv_N.""" + dg = self._detgroup(NoisePSDOof()) + tod = np.random.default_rng(0).normal(0.0, 1.0, 4096) + 12.0 # large offset + out = dg.apply_N_inv(tod, np.array([1.0, 0.5, -2.0])) + assert abs(np.mean(out)) < 1e-6 * np.max(np.abs(out)) + + def test_white_fast_path(self): + class _White(NoisePSDOof): + is_white = True + dg = self._detgroup(_White()) + tod = np.random.default_rng(1).normal(0.0, 2.0, 4096) + 3.0 + sigma0 = 2.0 + out = dg.apply_N_inv(tod, np.array([sigma0, 0.5, -2.0])) + expected = tod / sigma0**2 + expected -= np.mean(expected) + assert np.allclose(out, expected) diff --git a/tests/test_tod_power_spectra.py b/tests/test_tod_power_spectra.py new file mode 100644 index 0000000..18abe8f --- /dev/null +++ b/tests/test_tod_power_spectra.py @@ -0,0 +1,150 @@ +"""Tests for the low-resolution (log-binned) TOD power spectra written to the chain.""" + +import numpy as np +import pytest + +from commander4.tod_processing import _binned_tod_power_spectrum, _record_tod_diagnostics +from commander4.data_models.TOD_samples import TODSamples + + +def test_shape_and_nan_padding(): + # expbin returns fewer than nbin bins for a long TOD, so both outputs are NaN-padded to nbin, + # and the finite entries form a contiguous leading block identical between freqs and power. + tod = np.random.default_rng(0).standard_normal(100_000) + nbin = 100 + f, p = _binned_tod_power_spectrum(tod, fsamp=10.0, nbin=nbin) + assert f.shape == (nbin,) and p.shape == (nbin,) + finite = np.isfinite(f) + assert 0 < finite.sum() < nbin + np.testing.assert_array_equal(finite, np.isfinite(p)) + nb = int(finite.sum()) + assert finite[:nb].all() and not finite[nb:].any() + + +def test_freqs_monotonic_and_in_range(): + fsamp = 8.0 + tod = np.random.default_rng(1).standard_normal(50_000) + f, _ = _binned_tod_power_spectrum(tod, fsamp, 100) + ff = f[np.isfinite(f)] + assert np.all(np.diff(ff) > 0) # strictly increasing bin centers + assert ff[0] >= 0.0 and ff[-1] <= fsamp / 2 + 1e-6 + + +def test_sinusoid_power_localized(): + fsamp, ntod, f0 = 100.0, 100_000, 5.0 + t = np.arange(ntod) / fsamp + f, p = _binned_tod_power_spectrum(np.sin(2 * np.pi * f0 * t), fsamp, 100) + peak_bin = int(np.nanargmax(p)) + assert abs(f[peak_bin] - f0) < 0.5 # peak power bin sits at the injected frequency + + +def test_white_noise_roughly_flat(): + # Periodogram |rfft|^2 / N of unit-variance white noise has expectation ~1 per mode; the + # bin-averaged spectrum (skipping the few noisy low-mode bins) should sit near 1. + tod = np.random.default_rng(3).standard_normal(200_000) + _, p = _binned_tod_power_spectrum(tod, fsamp=10.0, nbin=100) + mid = p[np.isfinite(p)][5:] + assert 0.5 < np.median(mid) < 2.0 + + +def test_short_tod_uses_fewer_bins_without_error(): + # A short TOD yields well under nbin bins; everything beyond stays NaN. + f, p = _binned_tod_power_spectrum(np.arange(64.0), fsamp=1.0, nbin=100) + assert np.isfinite(f).sum() < 40 + assert np.isfinite(p).sum() == np.isfinite(f).sum() + + +# -------------------------------------------------------------------------------------- +# Optional DEBUG: full n_corr TOD ragged packing +# -------------------------------------------------------------------------------------- +def _bare_tod_samples(nscans, ndet, ncorr_tods): + ts = TODSamples.__new__(TODSamples) # bypass __init__ (no MPI / data needed) + ts.nscans, ts.ndet, ts.ncorr_tods = nscans, ndet, ncorr_tods + return ts + + +def test_pack_ncorr_tods_ragged(): + a = np.arange(3, dtype=np.float32) + b = np.arange(5, dtype=np.float32) + 10.0 + # scan0: det0=a (len 3), det1 missing; scan1: det0 missing, det1=b (len 5). + ts = _bare_tod_samples(2, 2, [[a, None], [None, b]]) + lengths, flat = ts._pack_ncorr_tods() + assert lengths.tolist() == [[3, 0], [0, 5]] + # Flat concatenation is scan-major, detector-minor (a before b). + np.testing.assert_array_equal(flat, np.concatenate([a, b])) + assert flat.dtype == np.float32 + + +def test_pack_ncorr_tods_empty(): + ts = _bare_tod_samples(2, 2, [[None, None], [None, None]]) + lengths, flat = ts._pack_ncorr_tods() + assert lengths.sum() == 0 + assert flat.size == 0 and flat.dtype == np.float32 + + +# -------------------------------------------------------------------------------------- +# _record_tod_diagnostics: which TOD view feeds which recorded spectrum +# -------------------------------------------------------------------------------------- +class _DiagStubView: + """Stand-in for TODView exposing only what _record_tod_diagnostics reads. + + ``get_tod()`` (no subtract) returns the jump-corrected raw TOD; ``get_tod(subtract=...)`` + returns the sky+orbital-dipole-subtracted base. Both are fresh copies so the diagnostics + helper can subtract n_corr in place, mirroring the real TODView contract. + """ + def __init__(self, raw, sky_orb_subtracted, fsamp=10.0): + self._raw = raw + self._sky_orb_subtracted = sky_orb_subtracted + self.fsamp = fsamp + + @property + def tod(self): + return self._raw + + def get_tod(self, *, subtract=None, **kwargs): + base = self._raw if subtract is None else self._sky_orb_subtracted + return np.array(base, copy=True) + + +def _diag_tod_samples(): + ts = TODSamples.__new__(TODSamples) # bypass __init__ (no MPI / data needed) + shape = (1, 1, TODSamples.TOD_PS_NBIN) + for name in ("tod_ps_freqs", "tod_ps_raw", "tod_ps_ncorr", "tod_ps_ncorrsub", + "tod_ps_residual"): + setattr(ts, name, np.full(shape, np.nan, dtype=np.float32)) + ts.ncorr_tods = None + return ts + + +def _ps(tod, fsamp=10.0): + return _binned_tod_power_spectrum(tod, fsamp, TODSamples.TOD_PS_NBIN)[1] + + +def test_record_diagnostics_routes_each_tod_to_its_spectrum(): + rng = np.random.default_rng(7) + raw = rng.standard_normal(2000) + sky_orb_sub = rng.standard_normal(2000) # stand-in for raw - gain*(sky+orb) + n_corr = rng.standard_normal(2000) + ts = _diag_tod_samples() + _record_tod_diagnostics(ts, 0, 0, _DiagStubView(raw, sky_orb_sub), n_corr) + + # Recorded spectra are stored as float32, so compare at single-precision tolerance. + np.testing.assert_allclose(ts.tod_ps_raw[0, 0], _ps(raw), rtol=1e-5) + np.testing.assert_allclose(ts.tod_ps_ncorr[0, 0], _ps(n_corr), rtol=1e-5) + # n_corr removed from raw -> sky + white noise retained. + np.testing.assert_allclose(ts.tod_ps_ncorrsub[0, 0], _ps(raw - n_corr), rtol=1e-5) + # sky, orbital dipole, and n_corr all removed -> noise residual. + np.testing.assert_allclose(ts.tod_ps_residual[0, 0], _ps(sky_orb_sub - n_corr), rtol=1e-5) + + +def test_record_diagnostics_without_ncorr(): + rng = np.random.default_rng(8) + raw = rng.standard_normal(2000) + sky_orb_sub = rng.standard_normal(2000) + ts = _diag_tod_samples() + _record_tod_diagnostics(ts, 0, 0, _DiagStubView(raw, sky_orb_sub), None) + + # With no n_corr drawn: ncorrsub == raw, residual is just the sky-subtracted TOD, ncorr stays NaN. + np.testing.assert_allclose(ts.tod_ps_ncorrsub[0, 0], _ps(raw), rtol=1e-5) + np.testing.assert_allclose(ts.tod_ps_residual[0, 0], _ps(sky_orb_sub), rtol=1e-5) + assert np.isnan(ts.tod_ps_ncorr[0, 0]).all() From f3f922e1214de68b54f3b5446a58e1ce414ade4a Mon Sep 17 00:00:00 2001 From: jgslunde Date: Mon, 15 Jun 2026 12:15:49 +0200 Subject: [PATCH 22/23] Changed the way detector-scans are looped over and how TODView is used, to make wrong-detector-indexing bugs (due to missing detectors on different ranks) less likely. --- src/commander4/data_models/TOD_samples.py | 64 ++- .../data_models/detector_group_TOD.py | 22 + src/commander4/data_models/tod_view.py | 46 +- .../experiments/SO/tod_reader_SO_SAT.py | 14 +- .../simulations/inplace_litebird_sim.py | 5 +- src/commander4/tod_processing.py | 493 ++++++++---------- src/commander4/utils/CG_mapmaker.py | 47 +- tests/test_gain_calibration.py | 4 +- 8 files changed, 368 insertions(+), 327 deletions(-) diff --git a/src/commander4/data_models/TOD_samples.py b/src/commander4/data_models/TOD_samples.py index be5da73..a8c95cb 100644 --- a/src/commander4/data_models/TOD_samples.py +++ b/src/commander4/data_models/TOD_samples.py @@ -117,9 +117,22 @@ def __init__(self, self.scan_idx_start = experiment_data.scan_idx_start self.scan_idx_stop = experiment_data.scan_idx_stop self.scan_ids = np.array([scan.scan_id for scan in experiment_data.scans]) + # Ordered per-band detector names. Their position is the ``idet`` axis shared by every + # per-detector array (rel_gain, noise_params, temporal_gain, accept, tod_ps_*, jumps), so + # writing them to the chain lets a reader map each array column back to a physical detector. + # Identical across all ranks of a band (taken from the band's detector list). + self.det_names = list(my_band.detectors) self.jumps = JumpCatalog.empty(self.nscans, self.ndet) - # Per detector-scan acceptance flag. Currently always True (no scans rejected); kept as a - # chain-tracked quantity so scan rejection can become a sampled step in the future. + # Two distinct per-detector-scan boolean masks over the dense (nscans, ndet) grid: + # * present: whether this detector actually has data in this scan. Scans hold only the + # detectors present in them (DetGroupTOD/ScanTOD are sparse), so a detector missing from + # a scan leaves a `present=False` hole in the dense arrays. Derived from the data, not + # sampled, so it is rebuilt here on every construction (chain init included). + # * accept: data-quality flag for present data that is *not* flagged as bad. Defaults to + # all True; a chain-tracked quantity so bad-data rejection can become a sampled step. + self.present = np.zeros((self.nscans, self.ndet), dtype=bool) + for iscan, det in experiment_data.iter_detector_scans(): + self.present[iscan, det.det_idx_fullband] = True self.accept = np.ones((self.nscans, self.ndet), dtype=bool) # Low-resolution (log-binned) TOD power spectra, written to the chain by default: a shared @@ -159,18 +172,17 @@ def __init__(self, self.rel_gain = np.zeros((self.ndet)) self.temporal_gain = np.zeros((self.nscans, self.ndet)) - all_det_gains = np.zeros((self.nscans, self.ndet)) - # all_det_gains = [] - # myband_noise_params = None + # NaN so detector-scans with no data are excluded from the gain means below. + all_det_gains = np.full((self.nscans, self.ndet), np.nan) if "initial_noise_params" in my_band: # Option 1: They are specified in the parameter file. self.noise_params[:] = np.array(my_band.initial_noise_params) elif experiment_data.scans[0].detectors[0].init_scalars is not None: - # Option 2: There were entries in the read-in files. - for iscan, scan in enumerate(experiment_data.scans): - for idet, det in enumerate(scan.detectors): - self.noise_params[iscan,idet] = det.init_scalars[1:] + # Option 2: There were entries in the read-in files. Index by the detector's + # full-band column; absent detector-scans stay NaN. + for iscan, det in experiment_data.iter_detector_scans(): + self.noise_params[iscan, det.det_idx_fullband] = det.init_scalars[1:] else: # Option 3: Fall back to the noise model's default parameters (ensuring a finite # sigma0, which the model leaves as NaN to be estimated from the data). @@ -178,24 +190,28 @@ def __init__(self, "model's default parameters.") default_params = np.array(self.noise_model.params, dtype=np.float64) if not np.isfinite(default_params[0]): - default_params[0] = 1e-3 + default_params[0] = 1.0 self.noise_params[:] = default_params if "gain" in my_band.detectors[experiment_data.scans[0].detectors[0].name]: - for iscan, scan in enumerate(experiment_data.scans): - for idet, det in enumerate(scan.detectors): - all_det_gains[iscan,idet] = my_band[det.name].gain + for iscan, det in experiment_data.iter_detector_scans(): + all_det_gains[iscan, det.det_idx_fullband] = my_band[det.name].gain elif experiment_data.scans[0].detectors[0].init_scalars is not None: - for iscan, scan in enumerate(experiment_data.scans): - for idet, det in enumerate(scan.detectors): - all_det_gains[iscan,idet] = det.init_scalars[0] + for iscan, det in experiment_data.iter_detector_scans(): + all_det_gains[iscan, det.det_idx_fullband] = det.init_scalars[0] else: raise ValueError("Did not find initial gain value in input files.") - all_det_gains = np.array(all_det_gains) self.abs_gain = float(np.nanmean(all_det_gains)) - self.rel_gain = np.nanmean(all_det_gains, axis=0) - self.abs_gain - self.temporal_gain = all_det_gains - self.rel_gain - self.abs_gain + # Relative gain only for detectors with data in >=1 local scan; detectors absent from + # every local scan get 0 (never used downstream) and are kept out of the empty-slice + # mean. temporal_gain holes (absent detector-scans) collapse to 0 likewise. + present_any = np.isfinite(all_det_gains).any(axis=0) + self.rel_gain = np.zeros(self.ndet) + self.rel_gain[present_any] = (np.nanmean(all_det_gains[:, present_any], axis=0) + - self.abs_gain) + self.temporal_gain = np.nan_to_num(all_det_gains - self.rel_gain - self.abs_gain, + nan=0.0) else: # --------------------------------------------------------- @@ -330,7 +346,10 @@ def write_chain_to_file(self, itr: int): noise_params_global = _gather_scan_distributed_array(band_comm, self.noise_params, scans_per_rank) - # 4b. Acceptance flags (per-scan per-detector; stored as int8 for MPI/HDF compatibility) + # 4b. Presence and acceptance flags (per-scan per-detector; int8 for MPI/HDF compatibility). + # `present` marks real vs. absent (dummy) detector-scans; `accept` marks good vs. bad data. + present_global = _gather_scan_distributed_array(band_comm, self.present.astype(np.int8), + scans_per_rank) accept_global = _gather_scan_distributed_array(band_comm, self.accept.astype(np.int8), scans_per_rank) @@ -374,6 +393,10 @@ def write_chain_to_file(self, itr: int): file["metadata/datetime"] = datetime.datetime.now().isoformat() file["metadata/parameter_file_as_string"] = params.parameter_file_as_string file["scan_ids"] = scan_ids_global + # Detector names (per-band, identical across ranks): the `idet` axis of every + # per-detector array below. Variable-length UTF-8 for a clean string round-trip. + file.create_dataset("det_names", + data=np.array(self.det_names, dtype=h5py.string_dtype())) if abs_gain_global is not None: file["abs_gain"] = abs_gain_global if rel_gain_global is not None: @@ -382,6 +405,7 @@ def write_chain_to_file(self, itr: int): file["temporal_gain"] = temporal_gain_global if noise_params_global is not None: file["noise_params"] = noise_params_global + file["present"] = present_global file["accept"] = accept_global file["tod_ps_freqs"] = tod_ps_freqs_global file["tod_ps_ncorr"] = tod_ps_ncorr_global diff --git a/src/commander4/data_models/detector_group_TOD.py b/src/commander4/data_models/detector_group_TOD.py index e699c83..0d5930b 100644 --- a/src/commander4/data_models/detector_group_TOD.py +++ b/src/commander4/data_models/detector_group_TOD.py @@ -41,6 +41,28 @@ def __init__(self, scans: list[ScanTOD], experiment_name: str, band_name: str, n self.nscans_allranks: int = 0 # Total number of scans across all ranks (on this band). self.noise_model = noise_model + def iter_detector_scans(self, accept: NDArray | None = None): + """Iterate over present detector-scans, yielding ``(iscan, det)`` pairs. + + ``ScanTOD.detectors`` is sparse -- each scan lists only the detectors actually present in it + -- so this nested walk is the canonical way to traverse detector-scans. The detector's + full-band column ``det.det_idx_fullband`` is the index into the dense ``(nscans, ndet)`` + per-detector sample arrays (gain, noise params, accept, ...); the per-scan enumerate position + must never be used for that, and this iterator deliberately never exposes one. + + Args: + accept: Optional ``(nscans, ndet)`` boolean mask. When given, detector-scans whose entry + is False are skipped, so callers process only accepted (good-quality) data. + + Yields: + tuple[int, DetectorTOD]: the local scan index ``iscan`` and the present detector ``det``. + """ + for iscan, scan in enumerate(self.scans): + for det in scan.detectors: + if accept is not None and not accept[iscan, det.det_idx_fullband]: + continue + yield iscan, det + def apply_N_inv(self, tod: NDArray, noise_params: NDArray, samprate: float|None = None, inplace=False) -> NDArray: """ Applies the inverse noise covariance N^-1 of this Det-Group to the input TOD, using the diff --git a/src/commander4/data_models/tod_view.py b/src/commander4/data_models/tod_view.py index 6f09c2f..e6cd450 100644 --- a/src/commander4/data_models/tod_view.py +++ b/src/commander4/data_models/tod_view.py @@ -6,6 +6,7 @@ from pixell.bunch import Bunch from commander4.data_models.detector_group_TOD import DetGroupTOD +from commander4.data_models.detector_TOD import DetectorTOD from commander4.data_models.TOD_samples import TODSamples from commander4.utils.map_utils import get_static_sky_TOD, get_s_orb_TOD @@ -64,14 +65,47 @@ def _clear_cache(self): self._orbital_dipole = None self._downsampled: dict[int, Bunch] = {} - def focus(self, iscan: int, idet: int) -> "TODView": - """Focus the view on one detector and discard any previous materialization.""" - self._det = self.experiment_data.scans[iscan].detectors[idet] + def focus(self, iscan: int, det: DetectorTOD) -> "TODView": + """Focus the view on one present detector and discard any previous materialization. + + Args: + iscan: Scan index, local to this rank. + det: The detector to focus on -- an element of ``scans[iscan].detectors`` (which holds + only the detectors actually present in that scan). Its full-band index + ``det.det_idx_fullband`` is the column used to address every per-detector sample + array, so detectors absent from a scan are simply skipped rather than misaligning + the dense ``(nscans, ndet)`` arrays. + """ + self._det = det self._iscan = iscan - self._idet = idet + self._idet = det.det_idx_fullband # full-band column in the (nscans, ndet) sample arrays self._clear_cache() return self + def iter_focused(self, *, accepted_only: bool = False): + """Focus on each present detector-scan in turn, yielding this re-focused view. + + Canonical detector-scan loop for TOD processing. It folds away the per-detector boilerplate + (``focus`` + the ``accept`` check) and, importantly, never exposes a per-scan detector + position that could be mistaken for a dense-array column: address the ``(nscans, ndet)`` + sample arrays through ``view.idet`` (the full-band column) and ``view.iscan`` only. + + The same view instance is re-focused and yielded on every iteration (matching the + one-detector-at-a-time design), so callers must consume each view within the loop body and + not retain it across iterations. + + Args: + accepted_only: If True, skip detector-scans whose ``accept`` flag is False (bad data); + if False (default), every present detector-scan is yielded (e.g. white-noise/jump + passes). + + Yields: + TODView: this view, focused on the current present detector-scan. + """ + accept = self.tod_samples.accept if accepted_only else None + for iscan, det in self.experiment_data.iter_detector_scans(accept): + yield self.focus(iscan, det) + def _require_focus(self): """Return the current detector or raise if the view was not focused yet.""" if self._iscan is None or self._idet is None or self._det is None: @@ -85,6 +119,7 @@ def iscan(self) -> int: @property def idet(self) -> int: + """Full-band detector index (``det_idx_fullband``): the column in per-detector arrays.""" self._require_focus() return self._idet @@ -111,7 +146,8 @@ def sigma0(self) -> float: @property def accept(self) -> bool: - """Whether the focused detector-scan is accepted (currently always True).""" + """Whether the focused detector-scan is accepted, i.e. present *and* not flagged as bad + data. ``accept`` (data quality) is distinct from ``present`` (data exists at all).""" self._require_focus() return bool(self.tod_samples.accept[self._iscan, self._idet]) diff --git a/src/commander4/experiments/SO/tod_reader_SO_SAT.py b/src/commander4/experiments/SO/tod_reader_SO_SAT.py index 9216e3a..ff79cfc 100644 --- a/src/commander4/experiments/SO/tod_reader_SO_SAT.py +++ b/src/commander4/experiments/SO/tod_reader_SO_SAT.py @@ -128,6 +128,10 @@ def tod_reader(band_comm: MPI.Comm, my_experiment: str, my_band: Bunch, det_name my_band.eval_nside, ntod_optimal) detector_list = [] + # idet is the detector's full-band column (its position in ``det_names``); idet_accepted + # is its position among the detectors actually kept in this scan (det_idx_local), and is + # only advanced when a detector passes the cuts below. + idet_accepted = 0 for idet, det_name in enumerate(det_names): # Find the index of the current detector in the file order of detectors. det_file_idx = det_names_file.index(det_name) @@ -141,14 +145,17 @@ def tod_reader(band_comm: MPI.Comm, my_experiment: str, my_band: Bunch, det_name det_response = det_responses[det_file_idx] flag_encoded = f[f"/{pid}/{det_name}/flag/"][()] - gain_init, sigma0_init, fknee_init, alpha_init = f[f"/{pid}/{det_name}/scalars"][()] + # gain_init, sigma0_init, fknee_init, alpha_init: + init_scalars = f[f"/{pid}/{det_name}/scalars"][()] - detector = DetectorTOD(tod, pointing, fsamp, np.zeros(3), huffman_tree, huffman_symbols, + detector = DetectorTOD(det_name, idet, idet_accepted, tod, pointing, fsamp, + np.zeros(3), huffman_tree, huffman_symbols, processing_mask_map, ntod, ntod_optimal, huffman_tree2=huffman_tree2, huffman_symbols2=huffman_symbols2, flag_encoded=flag_encoded, bad_data_bitmask=my_experiment.bad_data_bitmask, + init_scalars=init_scalars, tod_is_compressed=my_experiment.tod_is_compressed, det_response=det_response) if np.sum(detector.full_mask) == 0 or (detector.tod == 0).all(): @@ -156,7 +163,10 @@ def tod_reader(band_comm: MPI.Comm, my_experiment: str, my_band: Bunch, det_name detector_list.append(detector) ntod_sum_original += ntod ntod_sum_final += ntod_optimal + idet_accepted += 1 + if len(detector_list) == 0: + good_scan = False if good_scan: scanID = int(pid) scan = ScanTOD(detector_list, 0., scanID) diff --git a/src/commander4/simulations/inplace_litebird_sim.py b/src/commander4/simulations/inplace_litebird_sim.py index 2f75eba..6af9e80 100644 --- a/src/commander4/simulations/inplace_litebird_sim.py +++ b/src/commander4/simulations/inplace_litebird_sim.py @@ -229,11 +229,12 @@ def replace_tod_with_sim(band_comm: MPI.Comm, detector_data: DetGroupTOD, band_p for scan in detector_data.scans: for det in scan.detectors: start_bench("orbdip") + pix, psi = det.get_pix_psi() ntod = det.tod.size det.tod[:] = np.zeros(ntod, dtype=np.float32) - det.tod[:] = I[det.pix] + Q[det.pix]*np.cos(2*det.psi) + U[det.pix]*np.sin(2*det.psi) + det.tod[:] = I[pix] + Q[pix]*np.cos(2*psi) + U[pix]*np.sin(2*psi) if sim_params.include_OrbitalDipole: - det.tod[:] += get_orbital_dipole(det, det.pix, freq, units) + det.tod[:] += get_orbital_dipole(det, pix, freq, units) stop_bench("orbdip") start_bench("noise") diff --git a/src/commander4/tod_processing.py b/src/commander4/tod_processing.py index d7226b5..9dde519 100644 --- a/src/commander4/tod_processing.py +++ b/src/commander4/tod_processing.py @@ -132,18 +132,14 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output scan_view = TODView(experiment_data, tod_samples, compsep_output=compsep_output) if pols == "IQU": mapmaker_invvar = WeightsMapmakerIQU(band_comm, experiment_data.nside) - for iscan, scan in enumerate(experiment_data.scans): - for idet, det in enumerate(scan.detectors): - view = scan_view.focus(iscan, idet) - if not view.accept: - continue - good_data_mask = view.good_data_mask - pix = view.pix[good_data_mask] - psi = view.psi[good_data_mask] - sigma0 = view.sigma0 - gain = view.get_gain() - inv_var = (gain/sigma0)**2 - mapmaker_invvar.accumulate_to_map(inv_var, pix, psi, response=view.det_response) + for view in scan_view.iter_focused(accepted_only=True): + good_data_mask = view.good_data_mask + pix = view.pix[good_data_mask] + psi = view.psi[good_data_mask] + sigma0 = view.sigma0 + gain = view.get_gain() + inv_var = (gain/sigma0)**2 + mapmaker_invvar.accumulate_to_map(inv_var, pix, psi, response=view.det_response) mapmaker_invvar.gather_map() mapmaker_invvar.normalize_map() if ismaster: @@ -155,17 +151,13 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output CG_maxiter=params.general.CG_mapmaker.maxiter) elif pols == "I": mapmaker_invvar = WeightsMapmaker(band_comm, experiment_data.nside) - for iscan, scan in enumerate(experiment_data.scans): - for idet, det in enumerate(scan.detectors): - view = scan_view.focus(iscan, idet) - if not view.accept: - continue - good_data_mask = view.good_data_mask - pix = view.pix[good_data_mask] - sigma0 = view.sigma0 - gain = view.get_gain() - inv_var = (gain/sigma0)**2 - mapmaker_invvar.accumulate_to_map(inv_var, pix) + for view in scan_view.iter_focused(accepted_only=True): + good_data_mask = view.good_data_mask + pix = view.pix[good_data_mask] + sigma0 = view.sigma0 + gain = view.get_gain() + inv_var = (gain/sigma0)**2 + mapmaker_invvar.accumulate_to_map(inv_var, pix) mapmaker_invvar.gather_map() if ismaster: precond = InvNPreconditionerI(utils.without_nan(1./mapmaker_invvar.final_map)) @@ -191,84 +183,80 @@ def tod2map_CG(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_output worst_residual_ncorr = 0 ### MAIN SCAN LOOP ### - for iscan, scan in enumerate(experiment_data.scans): - for idet, det in enumerate(scan.detectors): - view = scan_view.focus(iscan, idet) - if not view.accept: - continue - pix, psi = view.pix, view.psi - good_data_mask = view.good_data_mask - pix_masked = pix[good_data_mask] - psi_masked = psi[good_data_mask] - sigma0 = view.sigma0 - gain = view.get_gain() - inv_var = (gain/sigma0)**2 - response = view.det_response if pols == "IQU" else None + for view in scan_view.iter_focused(accepted_only=True): + pix, psi = view.pix, view.psi + good_data_mask = view.good_data_mask + pix_masked = pix[good_data_mask] + psi_masked = psi[good_data_mask] + sigma0 = view.sigma0 + gain = view.get_gain() + inv_var = (gain/sigma0)**2 + response = view.det_response if pols == "IQU" else None + + ### ORBITAL DIPOLE ### + sky_orb_dipole = view.get_orbital_dipole_tod() + d_sky = view.get_tod(subtract=(("orbital_dipole", TODView._ALL_GAIN_TERMS),)) + if pols == "IQU": + mapmaker_orbdipole.accumulate_to_map(sky_orb_dipole, inv_var, pix, psi, + response=response) + else: + mapmaker_orbdipole.accumulate_to_map(sky_orb_dipole, inv_var, pix, psi) - ### ORBITAL DIPOLE ### - sky_orb_dipole = view.get_orbital_dipole_tod() - d_sky = view.get_tod(subtract=(("orbital_dipole", TODView._ALL_GAIN_TERMS),)) + ### CORRELATED NOISE SAMPLING ### + if ncorr_cfg.do_ncorr: + sky_subtracted_TOD = view.get_tod( + subtract=(("sky", TODView._ALL_GAIN_TERMS), + ("orbital_dipole", TODView._ALL_GAIN_TERMS)), + ) + res = sample_correlated_noise( + sky_subtracted_TOD, view.full_mask, np.array(view.noise_params, copy=True), + experiment_data.noise_model, view.fsamp, cg_err_tol=ncorr_cfg.cg_err_tol, + cg_max_iter=ncorr_cfg.cg_max_iter, sample_params=ncorr_cfg.do_param, + sample_sigma0=ncorr_cfg.sample_sigma0, nomono=ncorr_cfg.nomono, + onlymono=ncorr_cfg.onlymono, + sigma0_dec=ncorr_cfg.sigma0_dec, psd_fit_nu_min=ncorr_cfg.psd_fit_nu_min, + psd_fit_nu_max=ncorr_cfg.psd_fit_nu_max, psd_bin=ncorr_cfg.psd_bin) + n_corr_est = res.n_corr if pols == "IQU": - mapmaker_orbdipole.accumulate_to_map(sky_orb_dipole, inv_var, pix, psi, - response=response) + mapmaker_ncorr.accumulate_to_map( + (n_corr_est/gain).astype(np.float32, copy=False), + inv_var, pix, psi, response=response) else: - mapmaker_orbdipole.accumulate_to_map(sky_orb_dipole, inv_var, pix, psi) - - ### CORRELATED NOISE SAMPLING ### - if ncorr_cfg.do_ncorr: - sky_subtracted_TOD = view.get_tod( - subtract=(("sky", TODView._ALL_GAIN_TERMS), - ("orbital_dipole", TODView._ALL_GAIN_TERMS)), - ) - res = sample_correlated_noise( - sky_subtracted_TOD, view.full_mask, np.array(view.noise_params, copy=True), - experiment_data.noise_model, view.fsamp, cg_err_tol=ncorr_cfg.cg_err_tol, - cg_max_iter=ncorr_cfg.cg_max_iter, sample_params=ncorr_cfg.do_param, - sample_sigma0=ncorr_cfg.sample_sigma0, nomono=ncorr_cfg.nomono, - onlymono=ncorr_cfg.onlymono, - sigma0_dec=ncorr_cfg.sigma0_dec, psd_fit_nu_min=ncorr_cfg.psd_fit_nu_min, - psd_fit_nu_max=ncorr_cfg.psd_fit_nu_max, psd_bin=ncorr_cfg.psd_bin) - n_corr_est = res.n_corr - if pols == "IQU": - mapmaker_ncorr.accumulate_to_map( - (n_corr_est/gain).astype(np.float32, copy=False), - inv_var, pix, psi, response=response) - else: - mapmaker_ncorr.accumulate_to_map( - (n_corr_est/gain).astype(np.float32, copy=False), - inv_var, pix, psi) - tod_samples.noise_params[iscan, idet, :] = res.noise_params - if ncorr_cfg.do_param: - sampled_params.append(np.array(res.noise_params, copy=True)) - if not res.converged: - num_failed_convergences_ncorr += 1 - if res.high_var: - num_too_high_var_ncorr += 1 - worst_residual_ncorr = max(worst_residual_ncorr, res.residual) - residuals.append(res.residual) - niters.append(res.niter) - - d_sky -= n_corr_est - - _record_tod_diagnostics(tod_samples, iscan, idet, view, - n_corr_est if ncorr_cfg.do_ncorr else None) - - d_sky_masked = d_sky[good_data_mask] - - cg_mapmaker.accum_to_RHS( - scan_tod=det, - sigma0=sigma0, - pix=pix_masked, - psi=psi_masked, - scan_tod_arr=d_sky_masked/gain - ) + mapmaker_ncorr.accumulate_to_map( + (n_corr_est/gain).astype(np.float32, copy=False), + inv_var, pix, psi) + tod_samples.noise_params[view.iscan, view.idet, :] = res.noise_params + if ncorr_cfg.do_param: + sampled_params.append(np.array(res.noise_params, copy=True)) + if not res.converged: + num_failed_convergences_ncorr += 1 + if res.high_var: + num_too_high_var_ncorr += 1 + worst_residual_ncorr = max(worst_residual_ncorr, res.residual) + residuals.append(res.residual) + niters.append(res.niter) + + d_sky -= n_corr_est + + _record_tod_diagnostics(tod_samples, view.iscan, view.idet, view, + n_corr_est if ncorr_cfg.do_ncorr else None) + + d_sky_masked = d_sky[good_data_mask] + + cg_mapmaker.accum_to_RHS( + scan_tod=view.detector, + sigma0=sigma0, + pix=pix_masked, + psi=psi_masked, + scan_tod_arr=d_sky_masked/gain + ) ### PRINT NOISE SAMPLING STATS ### if ncorr_cfg.do_ncorr: log_corr_noise_stats(band_comm, experiment_data.nu, experiment_data.noise_model, sampled_params, residuals, niters, num_failed_convergences_ncorr, num_too_high_var_ncorr, worst_residual_ncorr, - len(experiment_data.scans)*experiment_data.ndet) + sum(len(s.detectors) for s in experiment_data.scans)) ### GATHER AND NORMALIZE MAPS ### @@ -389,24 +377,18 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu # This is purely to reduce the maximum concurrent memory requirement, and is slightly slower # as we have to de-compress pix and psi twice. start_bench("binned-mapmaker") - nscans = len(experiment_data.scans) - ndet = experiment_data.ndet pols = experiment_data.pols scan_view = TODView(experiment_data, tod_samples, compsep_output=compsep_output) - mapmaker_invvar = WeightsMapmakerIQU(band_comm, experiment_data.nside) - for iscan, scan in enumerate(experiment_data.scans): - for idet, det in enumerate(scan.detectors): - view = scan_view.focus(iscan, idet) - if not view.accept: - continue - good_data_mask = view.good_data_mask - pix = view.pix[good_data_mask] - psi = view.psi[good_data_mask] - sigma0 = view.sigma0 - gain = view.get_gain() - # sigma0 is in detector-units, transform into uK_RJ by dividing it by the gain. - inv_var = (gain/sigma0)**2 - mapmaker_invvar.accumulate_to_map(inv_var, pix, psi, response=view.det_response) + mapmaker_invvar = WeightsMapmakerIQU(band_comm, experiment_data.nside) + for view in scan_view.iter_focused(accepted_only=True): + good_data_mask = view.good_data_mask + pix = view.pix[good_data_mask] + psi = view.psi[good_data_mask] + sigma0 = view.sigma0 + gain = view.get_gain() + # sigma0 is in detector-units, transform into uK_RJ by dividing it by the gain. + inv_var = (gain/sigma0)**2 + mapmaker_invvar.accumulate_to_map(inv_var, pix, psi, response=view.det_response) mapmaker_invvar.gather_map() mapmaker_invvar.normalize_map() @@ -424,75 +406,72 @@ def tod2map_bin(band_comm: MPI.Comm, experiment_data: DetGroupTOD, compsep_outpu stop_bench("binned-mapmaker") ### MAIN SCAN LOOP ### - for iscan, scan in enumerate(experiment_data.scans): - for idet, det in enumerate(scan.detectors): - if not tod_samples.accept[iscan, idet]: - continue - start_bench("binned-mapmaker") - view = scan_view.focus(iscan, idet) - good_data_mask = view.good_data_mask - pix, psi = view.pix, view.psi - pix_masked = pix[good_data_mask] - psi_masked = psi[good_data_mask] - response = view.det_response - gain = view.get_gain() - sigma0 = view.sigma0 - # sigma0 is in detector-units, transform into uK_RJ by dividing it by the gain. - inv_var = (gain/sigma0)**2 - - ### ORBITAL DIPOLE ### - sky_orb_dipole = view.get_orbital_dipole_tod() - d_sky = view.get_tod(subtract=(("orbital_dipole", TODView._ALL_GAIN_TERMS),)) - - stop_bench("binned-mapmaker", increment_count=False) - ### CORRELATED NOISE SAMPLING ### - if ncorr_cfg.do_ncorr: - start_bench("ncorr-sampling") - sky_subtracted_TOD = view.get_tod( - subtract=(("sky", TODView._ALL_GAIN_TERMS), - ("orbital_dipole", TODView._ALL_GAIN_TERMS)), - ) - res = sample_correlated_noise( - sky_subtracted_TOD, view.full_mask, np.array(view.noise_params, copy=True), - experiment_data.noise_model, view.fsamp, cg_err_tol=ncorr_cfg.cg_err_tol, - cg_max_iter=ncorr_cfg.cg_max_iter, sample_params=ncorr_cfg.do_param, - sample_sigma0=ncorr_cfg.sample_sigma0, nomono=ncorr_cfg.nomono, - onlymono=ncorr_cfg.onlymono, - sigma0_dec=ncorr_cfg.sigma0_dec, psd_fit_nu_min=ncorr_cfg.psd_fit_nu_min, - psd_fit_nu_max=ncorr_cfg.psd_fit_nu_max, psd_bin=ncorr_cfg.psd_bin) - n_corr_est = res.n_corr - mapmaker_ncorr.accumulate_to_map((n_corr_est/gain).astype(np.float32, copy=False), - inv_var, pix, psi, response=response) - tod_samples.noise_params[iscan,idet,:] = res.noise_params - if ncorr_cfg.do_param: - sampled_params.append(np.array(res.noise_params, copy=True)) - if not res.converged: - num_failed_convergences_ncorr += 1 - if res.high_var: - num_too_high_var_ncorr += 1 - worst_residual_ncorr = max(worst_residual_ncorr, res.residual) - residuals.append(res.residual) - niters.append(res.niter) - - d_sky -= n_corr_est - stop_bench("ncorr-sampling") - if iscan == len(experiment_data.scans) - 1: - log_memory("ncorr-sampling") - - _record_tod_diagnostics(tod_samples, iscan, idet, view, - n_corr_est if ncorr_cfg.do_ncorr else None) - - d_sky_masked = d_sky[good_data_mask] - start_bench("binned-mapmaker") - mapmaker.accumulate_to_map(d_sky_masked/gain, inv_var, pix_masked, psi_masked, response=response) - mapmaker_orbdipole.accumulate_to_map(sky_orb_dipole, inv_var, pix, psi, response=response) - stop_bench("binned-mapmaker", increment_count=False) + for view in scan_view.iter_focused(accepted_only=True): + start_bench("binned-mapmaker") + good_data_mask = view.good_data_mask + pix, psi = view.pix, view.psi + pix_masked = pix[good_data_mask] + psi_masked = psi[good_data_mask] + response = view.det_response + gain = view.get_gain() + sigma0 = view.sigma0 + # sigma0 is in detector-units, transform into uK_RJ by dividing it by the gain. + inv_var = (gain/sigma0)**2 + + ### ORBITAL DIPOLE ### + sky_orb_dipole = view.get_orbital_dipole_tod() + d_sky = view.get_tod(subtract=(("orbital_dipole", TODView._ALL_GAIN_TERMS),)) + + stop_bench("binned-mapmaker", increment_count=False) + ### CORRELATED NOISE SAMPLING ### + if ncorr_cfg.do_ncorr: + start_bench("ncorr-sampling") + sky_subtracted_TOD = view.get_tod( + subtract=(("sky", TODView._ALL_GAIN_TERMS), + ("orbital_dipole", TODView._ALL_GAIN_TERMS)), + ) + res = sample_correlated_noise( + sky_subtracted_TOD, view.full_mask, np.array(view.noise_params, copy=True), + experiment_data.noise_model, view.fsamp, cg_err_tol=ncorr_cfg.cg_err_tol, + cg_max_iter=ncorr_cfg.cg_max_iter, sample_params=ncorr_cfg.do_param, + sample_sigma0=ncorr_cfg.sample_sigma0, nomono=ncorr_cfg.nomono, + onlymono=ncorr_cfg.onlymono, + sigma0_dec=ncorr_cfg.sigma0_dec, psd_fit_nu_min=ncorr_cfg.psd_fit_nu_min, + psd_fit_nu_max=ncorr_cfg.psd_fit_nu_max, psd_bin=ncorr_cfg.psd_bin) + n_corr_est = res.n_corr + mapmaker_ncorr.accumulate_to_map((n_corr_est/gain).astype(np.float32, copy=False), + inv_var, pix, psi, response=response) + tod_samples.noise_params[view.iscan, view.idet, :] = res.noise_params + if ncorr_cfg.do_param: + sampled_params.append(np.array(res.noise_params, copy=True)) + if not res.converged: + num_failed_convergences_ncorr += 1 + if res.high_var: + num_too_high_var_ncorr += 1 + worst_residual_ncorr = max(worst_residual_ncorr, res.residual) + residuals.append(res.residual) + niters.append(res.niter) + + d_sky -= n_corr_est + stop_bench("ncorr-sampling") + + _record_tod_diagnostics(tod_samples, view.iscan, view.idet, view, + n_corr_est if ncorr_cfg.do_ncorr else None) + + d_sky_masked = d_sky[good_data_mask] + start_bench("binned-mapmaker") + mapmaker.accumulate_to_map(d_sky_masked/gain, inv_var, pix_masked, psi_masked, response=response) + mapmaker_orbdipole.accumulate_to_map(sky_orb_dipole, inv_var, pix, psi, response=response) + stop_bench("binned-mapmaker", increment_count=False) + if ncorr_cfg.do_ncorr: + log_memory("ncorr-sampling") ### PRINT NOISE SAMPLING STATS ### if ncorr_cfg.do_ncorr: log_corr_noise_stats(band_comm, experiment_data.nu, experiment_data.noise_model, sampled_params, residuals, niters, num_failed_convergences_ncorr, - num_too_high_var_ncorr, worst_residual_ncorr, nscans*ndet) + num_too_high_var_ncorr, worst_residual_ncorr, + sum(len(s.detectors) for s in experiment_data.scans)) start_bench("binned-mapmaker") @@ -660,21 +639,18 @@ def estimate_white_noise(experiment_data: DetGroupTOD, tod_samples: TODSamples, tod_samples (TODSamples): Updated TOD samples with sigma0 estimates. """ scan_view = TODView(experiment_data, tod_samples, compsep_output=det_compsep_map) - for iscan, scan in enumerate(experiment_data.scans): - for idet, det in enumerate(scan.detectors): - view = scan_view.focus(iscan, idet) - # FIXME: Should maybe n_corr be subtracted here as well? - sky_subtracted_tod = view.get_tod( - subtract=(("sky", TODView._ALL_GAIN_TERMS), - ("orbital_dipole", TODView._ALL_GAIN_TERMS)), - ) - mask = view.full_mask - sigma0 = calc_sigma0_robust(sky_subtracted_tod, mask) - logassert(sigma0 != 0, "sigma0 is 0, which should never happen.", logger) - logassert(sigma0 != np.inf, "sigma0 is inf, which should never happen.", logger) - tod_samples.noise_params[iscan,idet,0] = sigma0 - if iscan == len(experiment_data.scans) - 1: - log_memory("sigma0-est") + for view in scan_view.iter_focused(): + # FIXME: Should maybe n_corr be subtracted here as well? + sky_subtracted_tod = view.get_tod( + subtract=(("sky", TODView._ALL_GAIN_TERMS), + ("orbital_dipole", TODView._ALL_GAIN_TERMS)), + ) + mask = view.full_mask + sigma0 = calc_sigma0_robust(sky_subtracted_tod, mask) + logassert(sigma0 != 0, "sigma0 is 0, which should never happen.", logger) + logassert(sigma0 != np.inf, "sigma0 is inf, which should never happen.", logger) + tod_samples.noise_params[view.iscan, view.idet, 0] = sigma0 + log_memory("sigma0-est") return tod_samples @@ -705,26 +681,24 @@ def sample_jump_detection(band_comm: MPI.Comm, experiment_data: DetGroupTOD, offsets_local = [] jump_counts_local = [] - for iscan, scan in enumerate(experiment_data.scans): - for idet, det in enumerate(scan.detectors): - view = scan_view.focus(iscan, idet) - if getattr(view.detector, "_flag_encoded", None) is None or not hasattr(view.detector, "_full_mask"): - tod_samples.jumps.set(iscan, idet, None) - jump_counts_local.append(0) - continue - jump, num_skipped = JumpCorrection.detect( - view.tod, - view.flag, - view.full_mask, - n_window, - jump_bitmask=jump_bitmask, - ) - tod_samples.jumps.set(iscan, idet, jump) - jump_counts_local.append(jump.size) - num_skipped_local += num_skipped - if not jump.is_empty(): - offsets_local.extend(jump.offsets.astype(np.float64, copy=False)) - num_applied_local += jump.size + for view in scan_view.iter_focused(): + if getattr(view.detector, "_flag_encoded", None) is None or not hasattr(view.detector, "_full_mask"): + tod_samples.jumps.set(view.iscan, view.idet, None) + jump_counts_local.append(0) + continue + jump, num_skipped = JumpCorrection.detect( + view.tod, + view.flag, + view.full_mask, + n_window, + jump_bitmask=jump_bitmask, + ) + tod_samples.jumps.set(view.iscan, view.idet, jump) + jump_counts_local.append(jump.size) + num_skipped_local += num_skipped + if not jump.is_empty(): + offsets_local.extend(jump.offsets.astype(np.float64, copy=False)) + num_applied_local += jump.size num_applied = band_comm.reduce(num_applied_local, op=MPI.SUM, root=0) num_skipped = band_comm.reduce(num_skipped_local, op=MPI.SUM, root=0) @@ -852,24 +826,19 @@ def sample_absolute_gain(band_comm: MPI.Comm, experiment_data: DetGroupTOD, tod_ scan_view = TODView(experiment_data, tod_samples, compsep_output=det_compsep_map) - for iscan, scan in enumerate(experiment_data.scans): - for idet, det in enumerate(scan.detectors): - view = scan_view.focus(iscan, idet) - if not view.accept: # Skip detector-scans flagged as bad; they carry no gain info. - continue - calib = view.get_calib_tod("abs", calibrate_against, - downsample_factor=downsample_factor) - s_cal = calib.s_cal - residual_tod = calib.tod + # Skip detector-scans flagged as bad (accepted_only); they carry no gain info. + for view in scan_view.iter_focused(accepted_only=True): + calib = view.get_calib_tod("abs", calibrate_against, + downsample_factor=downsample_factor) + s_cal = calib.s_cal + residual_tod = calib.tod - N_inv_s = experiment_data.apply_N_inv(s_cal, view.noise_params, samprate=1.0) - N_inv_d = experiment_data.apply_N_inv(residual_tod, view.noise_params, samprate=1.0) + N_inv_s = experiment_data.apply_N_inv(s_cal, view.noise_params, samprate=1.0) + N_inv_d = experiment_data.apply_N_inv(residual_tod, view.noise_params, samprate=1.0) - # Add to the numerator and denominator. - # sum_s_T_N_inv_d += np.dot(s_cal[mask], N_inv_d[mask]) - # sum_s_T_N_inv_s += np.dot(s_cal[mask], N_inv_s[mask]) - sum_s_T_N_inv_d += np.dot(s_cal, N_inv_d) - sum_s_T_N_inv_s += np.dot(s_cal, N_inv_s) + # Add to the numerator and denominator. + sum_s_T_N_inv_d += np.dot(s_cal, N_inv_d) + sum_s_T_N_inv_s += np.dot(s_cal, N_inv_s) # The g0 term is fully global, so we reduce across both all scans and all bands: sum_s_T_N_inv_d = band_comm.reduce(sum_s_T_N_inv_d, op=MPI.SUM, root=0) @@ -929,25 +898,20 @@ def sample_relative_gain(band_comm: MPI.Comm, experiment_data: DetGroupTOD, scan_view = TODView(experiment_data, tod_samples, compsep_output=det_compsep_map) - for iscan, scan in enumerate(experiment_data.scans): - for idet, det in enumerate(scan.detectors): - view = scan_view.focus(iscan, idet) - if not view.accept: # Skip detector-scans flagged as bad; they carry no gain info. - continue - calib = view.get_calib_tod("rel", calibrate_against, - downsample_factor=downsample_factor) - s_cal = calib.s_cal - residual_tod = calib.tod - N_inv_s = experiment_data.apply_N_inv(s_cal, view.noise_params, samprate=1.0) + # Skip detector-scans flagged as bad (accepted_only); they carry no gain info. + for view in scan_view.iter_focused(accepted_only=True): + calib = view.get_calib_tod("rel", calibrate_against, + downsample_factor=downsample_factor) + s_cal = calib.s_cal + residual_tod = calib.tod + N_inv_s = experiment_data.apply_N_inv(s_cal, view.noise_params, samprate=1.0) - s_T_N_inv_s_scan = np.dot(s_cal, N_inv_s) - r_T_N_inv_s_scan = np.dot(residual_tod, N_inv_s) - # s_T_N_inv_s_scan = np.dot(s_tot[mask], N_inv_s[mask]) - # r_T_N_inv_s_scan = np.dot(residual_tod[mask], N_inv_s[mask]) + s_T_N_inv_s_scan = np.dot(s_cal, N_inv_s) + r_T_N_inv_s_scan = np.dot(residual_tod, N_inv_s) - # Add the contribution from this scan to the local sum - local_s_T_N_inv_s[idet] += s_T_N_inv_s_scan - local_r_T_N_inv_s[idet] += r_T_N_inv_s_scan + # Add the contribution from this scan to the local sum (full-band detector column). + local_s_T_N_inv_s[view.idet] += s_T_N_inv_s_scan + local_r_T_N_inv_s[view.idet] += r_T_N_inv_s_scan ### 2. Intra-Detector Reduction ### # Sum the local values across all ranks that share the same detector using det_comm. @@ -1023,32 +987,29 @@ def sample_temporal_gain_variations(band_comm: MPI.Comm, experiment_data: DetGro b_q_local = np.zeros((ndet, nscans_local), dtype=np.float64) scan_view = TODView(experiment_data, tod_samples, compsep_output=det_compsep_map) - for iscan, scan in enumerate(experiment_data.scans): - for idet, det in enumerate(scan.detectors): - # I'm still not sure what way of dealing with the masked samples are best: - # 1. Replace masked values with 0s before FFT. - # 2. Replace masked values with n_corr realizations before FFT. - # 3. Remove masked values by reducing TOD size before FFTs. - # (simply passing the full data through the FFTs seems like a bad idea because of - # ringing from the large residual in the galactic plane). - - view = scan_view.focus(iscan, idet) - if not view.accept: # Rejected detector-scans contribute zero weight (A_qq = b_q = 0); - continue # the Wiener prior then fills their temporal gain from neighbors. - calib = view.get_calib_tod("temp", calibrate_against, - downsample_factor=downsample_factor) - s_cal = calib.s_cal - residual_tod = calib.tod - - N_inv_s = experiment_data.apply_N_inv(s_cal, view.noise_params, samprate=1.0) - N_inv_r = experiment_data.apply_N_inv(residual_tod, view.noise_params, samprate=1.0) - - # Calculate elements for the linear system - A_qq = np.dot(s_cal, N_inv_s) - b_q = np.dot(s_cal, N_inv_r) - - A_qq_local[idet, iscan] = A_qq - b_q_local[idet, iscan] = b_q + # I'm still not sure what way of dealing with the masked samples are best: + # 1. Replace masked values with 0s before FFT. + # 2. Replace masked values with n_corr realizations before FFT. + # 3. Remove masked values by reducing TOD size before FFTs. + # (simply passing the full data through the FFTs seems like a bad idea because of + # ringing from the large residual in the galactic plane). + # Rejected detector-scans (accepted_only) contribute zero weight (A_qq = b_q = 0); the Wiener + # prior then fills their temporal gain from neighbors. + for view in scan_view.iter_focused(accepted_only=True): + calib = view.get_calib_tod("temp", calibrate_against, + downsample_factor=downsample_factor) + s_cal = calib.s_cal + residual_tod = calib.tod + + N_inv_s = experiment_data.apply_N_inv(s_cal, view.noise_params, samprate=1.0) + N_inv_r = experiment_data.apply_N_inv(residual_tod, view.noise_params, samprate=1.0) + + # Calculate elements for the linear system + A_qq = np.dot(s_cal, N_inv_s) + b_q = np.dot(s_cal, N_inv_r) + + A_qq_local[view.idet, view.iscan] = A_qq + b_q_local[view.idet, view.iscan] = b_q # Gather scan counts on all ranks (needed for gather/scatter with varying roots) scan_counts = np.array(band_comm.allgather(nscans_local), dtype=int) diff --git a/src/commander4/utils/CG_mapmaker.py b/src/commander4/utils/CG_mapmaker.py index 96f1c76..4e22dd4 100644 --- a/src/commander4/utils/CG_mapmaker.py +++ b/src/commander4/utils/CG_mapmaker.py @@ -237,36 +237,23 @@ def apply_LHS(self, in_map: NDArray): self.map_comm.Bcast(in_map, root=0) out_map = np.zeros_like(in_map) - pri = True - for iscan, scan in enumerate(self.detector_tod.scans): - for idet, det in enumerate(scan.detectors): - pix = det.pix - psi = det.psi - sigma0 = self.detector_samples.noise_params[iscan, idet, 0] - scan_tod_arr_aux = np.zeros_like(det.tod, dtype=self.f_dtype) #aux array to not modify scan.tod - # if self.map_comm.Get_rank() == 0 and pri: - # self.logger.info(f"##LHS 1 mean: {np.mean(in_map)}") - #P m - scan_tod_arr_aux = self.apply_P(in_map, det, pix=pix, psi=psi, scan_tod_arr=scan_tod_arr_aux) - # if self.map_comm.Get_rank() == 0 and pri: - # self.logger.info(f"##LHS 2 mean: {np.mean(scan_tod_arr_aux)}") - #T P m - scan_tod_arr_aux = self.apply_T(scan_tod_arr_aux) - # if self.map_comm.Get_rank() == 0 and pri: - # self.logger.info(f"##LHS 3 mean: {np.mean(scan_tod_arr_aux)}") - #N^-1 T P m - scan_tod_arr_aux = self.apply_inv_N(scan_tod_arr_aux, sigma0) - # if self.map_comm.Get_rank() == 0 and pri: - # self.logger.info(f"##LHS 4 mean: {np.mean(scan_tod_arr_aux)}") - #T^T N^-1 T P m - scan_tod_arr_aux = self.apply_T_adjoint(scan_tod_arr_aux) - # if self.map_comm.Get_rank() == 0 and pri: - # self.logger.info(f"##LHS 5 mean: {np.mean(scan_tod_arr_aux)}") - #P^T T^T N^-1 T P - out_map = self.apply_P_adjoint(det, out_map, pix=pix, psi=psi, scan_tod_arr=scan_tod_arr_aux) - # if self.map_comm.Get_rank() == 0 and pri: - # self.logger.info(f"##LHS 6 mean: {np.mean(out_map)}") - pri=False + # Iterate accepted detector-scans only, matching accum_to_RHS: the LHS operator + # P^T T^T N^-1 T P and the RHS P^T T^T N^-1 d must span the same set of detector-scans, or + # the CG would solve an inconsistent (A, b) pair. det.det_idx_fullband is the dense column. + for iscan, det in self.detector_tod.iter_detector_scans(self.detector_samples.accept): + pix, psi = det.get_pix_psi() + sigma0 = self.detector_samples.noise_params[iscan, det.det_idx_fullband, 0] + scan_tod_arr_aux = np.zeros_like(det.tod, dtype=self.f_dtype) #aux array to not modify scan.tod + #P m + scan_tod_arr_aux = self.apply_P(in_map, det, pix=pix, psi=psi, scan_tod_arr=scan_tod_arr_aux) + #T P m + scan_tod_arr_aux = self.apply_T(scan_tod_arr_aux) + #N^-1 T P m + scan_tod_arr_aux = self.apply_inv_N(scan_tod_arr_aux, sigma0) + #T^T N^-1 T P m + scan_tod_arr_aux = self.apply_T_adjoint(scan_tod_arr_aux) + #P^T T^T N^-1 T P + out_map = self.apply_P_adjoint(det, out_map, pix=pix, psi=psi, scan_tod_arr=scan_tod_arr_aux) send, recv = (MPI.IN_PLACE, out_map) if self.map_comm.Get_rank() == 0 else (out_map, None) self.map_comm.Reduce(send, recv, op=MPI.SUM, root=0) if not ismaster: diff --git a/tests/test_gain_calibration.py b/tests/test_gain_calibration.py index 3dbbcc0..5c2b875 100644 --- a/tests/test_gain_calibration.py +++ b/tests/test_gain_calibration.py @@ -233,7 +233,7 @@ def _make_real_view(monkeypatch): pix = rng.integers(0, 12, size=NTOD) # Valid pixels for the nside=1 experiment below. psi = rng.uniform(0.0, np.pi, size=NTOD) det = SimpleNamespace(tod=rng.normal(size=NTOD), ntod=NTOD, fsamp=float(FACTOR), - get_pix_psi=lambda: (pix, psi), + det_idx_fullband=0, get_pix_psi=lambda: (pix, psi), orb_dir_vec=np.array([1.0, 0.0, 0.0], dtype=np.float32)) experiment_data = SimpleNamespace(scans=[SimpleNamespace(detectors=[det])], nside=1, nu=30.0) no_jump = SimpleNamespace(is_empty=lambda: True) @@ -242,7 +242,7 @@ def _make_real_view(monkeypatch): temporal_gain=np.array([[0.25]]), accept=np.ones((1, 1), dtype=bool)) skymap = rng.normal(size=(3, 12)) - view = TODView(experiment_data, tod_samples, compsep_output=skymap).focus(0, 0) + view = TODView(experiment_data, tod_samples, compsep_output=skymap).focus(0, det) s_full = skymap[0, pix] + np.cos(2*psi)*skymap[1, pix] + np.sin(2*psi)*skymap[2, pix] return view, det, s_full From 1ef3596c07c4b957d54b26aecf15b39c095138c3 Mon Sep 17 00:00:00 2001 From: jgslunde Date: Mon, 15 Jun 2026 12:16:53 +0200 Subject: [PATCH 23/23] Updated various experiments to be more up-to-date with recent SO implementations. --- params/param_AkariTODs_CG.yml | 48 ++++++++++++++----- .../experiments/akari/tod_reader_akari.py | 32 +++++++++---- .../litebird/tod_reader_litebird_sim.py | 18 ++++--- .../tod_reader_litebird_sim_spawndetectors.py | 17 +++++-- .../experiments/planck/tod_reader_planck.py | 39 ++++++++------- 5 files changed, 104 insertions(+), 50 deletions(-) diff --git a/params/param_AkariTODs_CG.yml b/params/param_AkariTODs_CG.yml index 8c9426d..0368383 100644 --- a/params/param_AkariTODs_CG.yml +++ b/params/param_AkariTODs_CG.yml @@ -28,14 +28,32 @@ general: noise_fknee: 0.1 noise_alpha: -0.9 - sample_corr_noise: True - sample_corr_noise_from_iter_num: 8 - sample_abs_gain: True - sample_abs_gain_from_iter_num: 2 - sample_rel_gain: True - sample_rel_gain_from_iter_num: 4 - sample_temporal_gain: True - sample_temporal_gain_from_iter_num: 4 + corr_noise: + sample_corr_noise: True + sample_corr_noise_from_iter_num: 8 + sample_noise_params: True # Resample noise-model params (needs sample_corr_noise). + sample_sigma0: True # Re-estimate sigma0 from the n_corr-subtracted residual. + CG_err_tol: 1.0e-4 # Correlated-noise CG convergence tolerance. + CG_max_iter: 60 # Max CG iterations; 0 -> skip CG, use stationary fallback. + nomono: False # Project the per-scan monopole out of n_corr (Fortran "nomono"). + onlymono: False # Model the correlated noise as only the per-scan offset. + sigma0_decimation: 1 # Block-average factor for the sigma0 estimator. + psd_fit_nu_min: 0.0 # Min frequency (Hz) for noise-PSD parameter fitting. + psd_fit_nu_max: 2.0 # Max frequency (Hz) for noise-PSD parameter fitting. + psd_bin: False # Fit every Fourier mode (False) or mode-weighted binned PSD. + gain_calib_downsample_time: 1.0 # [s] Averaging chunk duration for gain-calibration TODs (0 = off). + abs_gain: # Absolute (band-global) gain g0. + sample: true + sample_from_iter_num: 2 + calibrate_against: orbital_dipole # "orbital_dipole" | "full_sky" | "sky"; per-band overridable. + rel_gain: # Per-detector relative gain. + sample: true + sample_from_iter_num: 4 + calibrate_against: full_sky # "orbital_dipole" | "full_sky" | "sky"; per-band overridable. + temporal_gain: # Per-scan temporal gain variations. + sample: true + sample_from_iter_num: 4 + calibrate_against: full_sky # "orbital_dipole" | "full_sky" | "sky"; per-band overridable. tod_processing_steps: Akari: @@ -69,6 +87,7 @@ general: write_orb_dipole_maps_to_chain: True write_corr_noise_maps_to_chain: True write_sky_model_maps_to_chain: True + write_ncorr_tods_to_chain: False # DEBUG ONLY: full per-sample n_corr TODs (very large data). output_paths: stats: "../cmdr4_stats/PlanckTODs_CG/" @@ -107,7 +126,8 @@ components: smoothing_prior_amplitude: 2.0e-3 # uK_RJ^2 beta: 1.56 T: 20 - nu0: 545 + nu_ref: [857, 353] # [I, QU] + units: "uK_RJ" lmax: 1024 smoothing_scale: 0.0 spatially_varying_MM: False @@ -122,7 +142,8 @@ components: smoothing_prior_FWHM: 60.0 # arcmin smoothing_prior_amplitude: 2.0e-2 # uK_RJ^2 beta: -3.1 - nu0: 30 + nu_ref: [0.408, 30] # [I, QU] + units: "uK_RJ" lmax: 512 smoothing_scale: 0.0 spatially_varying_MM: False @@ -137,7 +158,8 @@ components: smoothing_prior_FWHM: 30.0 # arcmin smoothing_prior_amplitude: 1.0e+9 # uK_RJ^2 T: 7000 - nu0: 0.408 #40 + nu_ref: 40 + units: "uK_RJ" lmax: 512 smoothing_scale: 0.0 spatially_varying_MM: False @@ -443,12 +465,14 @@ experiments: pix_is_compressed: true data_type: "TOD" mapmaker: "CG" - flag_bitmask: 0b1111111111111111111111111111101 + bad_data_bitmask: 0b1111111111111111111111111111101 #bad_PIDs_path: "/mn/stornext/d23/cmbco/jonas/c4_testing/Commander4/badPIDs.npy" Fourier_times_path: "/mn/stornext/d23/cmbco/jonas/c4_testing/analysis/FFT_times_0_399999.npy" num_scans: 11395 bands: AkariN60: + abs_gain: + calibrate_against: full_sky # Preserve full-sky abs-gain calibration (legacy nu>380 rule). enabled: true tod_files_prefix: "AKARI_N60_v1_" polarization: "I" diff --git a/src/commander4/experiments/akari/tod_reader_akari.py b/src/commander4/experiments/akari/tod_reader_akari.py index 9924c25..1bd8f11 100644 --- a/src/commander4/experiments/akari/tod_reader_akari.py +++ b/src/commander4/experiments/akari/tod_reader_akari.py @@ -12,6 +12,8 @@ from commander4.data_models.detector_TOD import DetectorTOD from commander4.data_models.scan_TOD import ScanTOD from commander4.data_models.detector_group_TOD import DetGroupTOD +from commander4.data_models.pointing import PixelPointing +from commander4.noise_sampling.noise_psd import NoisePSDOof def get_processing_mask(my_band: Bunch) -> DetectorTOD: @@ -85,31 +87,43 @@ def tod_reader(band_comm: MPI.Comm, my_experiment: str, my_band: Bunch, det_name if ntod > ntod_upper_bound: raise ValueError(f"{ntod_upper_bound} {ntod}") vsun = np.ones(3) # dummy, we don't have that in Akari. + # Akari is intensity-only: the files carry no psi, so we hand PixelPointing a zero psi of + # the right length (psi is unused by I-only mapmaking, but PixelPointing requires one). + psi_zeros = np.zeros(ntod_optimal, dtype=np.float32) detector_list = [] - for det_name in det_names: + # All detectors are kept, so the full-band column idet and the per-scan column + # idet_accepted advance together here. + idet_accepted = 0 + for idet, det_name in enumerate(det_names): tod = f[f"/{pid}/{det_name}/tod/"][:ntod_optimal].astype(np.float32) pix_encoded = f[f"/{pid}/{det_name}/pix/"][()] flag_encoded = f[f"/{pid}/{det_name}/flag/"][()] - detector = DetectorTOD(tod, pix_encoded, [], my_band.eval_nside, - my_band.data_nside, fsamp, vsun, huffman_tree, - huffman_symbols, npsi, processing_mask_map, ntod, + # gain_init, sigma0_init, fknee_init, alpha_init: + init_scalars = f[f"/{pid}/{det_name}/scalars"][()] + det_pointing = PixelPointing(pix_encoded, psi_zeros, huffman_tree, huffman_symbols, + npsi, my_band.eval_nside, my_band.data_nside, ntod, + ntod_optimal) + detector = DetectorTOD(det_name, idet, idet_accepted, tod, det_pointing, fsamp, + vsun, huffman_tree, huffman_symbols, processing_mask_map, + ntod, ntod_optimal, flag_encoded=flag_encoded, - flag_bitmask=my_experiment.flag_bitmaks, - pix_is_compressed=my_experiment.pix_is_compressed, - psi_is_compressed=False) + bad_data_bitmask=my_experiment.bad_data_bitmask, + init_scalars=init_scalars) detector_list.append(detector) ntod_sum_original += ntod ntod_sum_final += ntod_optimal + idet_accepted += 1 scanID = int(pid) - scan = ScanTOD(detector_list, 0., scanID, scan_idx_start, scan_idx_stop) + scan = ScanTOD(detector_list, 0., scanID) scan_list.append(scan) num_included += 1 if i_pid % 10 == 0: gc.collect() ndet = len(det_names) + noise_model = NoisePSDOof() band_tod = DetGroupTOD(scan_list, expname, bandname, my_band.eval_nside, my_band.freq, - my_band.fwhm, ndet, my_band.polarizations) + my_band.fwhm, fsamp, ndet, my_band.polarization, noise_model) ### Collect some info on master rank of each band and print it ### local_tot_scans = scan_idx_stop - scan_idx_start diff --git a/src/commander4/experiments/litebird/tod_reader_litebird_sim.py b/src/commander4/experiments/litebird/tod_reader_litebird_sim.py index 8c5cf5a..cca8085 100644 --- a/src/commander4/experiments/litebird/tod_reader_litebird_sim.py +++ b/src/commander4/experiments/litebird/tod_reader_litebird_sim.py @@ -12,6 +12,7 @@ from commander4.data_models.detector_TOD import DetectorTOD from commander4.data_models.scan_TOD import ScanTOD from commander4.data_models.detector_group_TOD import DetGroupTOD +from commander4.data_models.pointing import PixelPointing from commander4.noise_sampling.noise_psd import NoisePSD, NoisePSDOof from commander4.simulations.inplace_litebird_sim import replace_tod_with_sim from commander4.output.log import logassert @@ -101,7 +102,10 @@ def tod_reader(band_comm: MPI.Comm, my_experiment: str, my_band: Bunch, det_name raise ValueError(f"{ntod_upper_bound} {ntod}") detector_list = [] - for det_name in det_names: + # All detectors are kept (whole-scan rejection via the flag check below), so the + # full-band column idet and the per-scan column idet_accepted advance together here. + idet_accepted = 0 + for idet, det_name in enumerate(det_names): tod = f[f"/{pid}/{det_name}/tod/"][:ntod_optimal].astype(np.float32, copy=False) pix_encoded = f[f"/{pid}/{det_name}/pix/"][()] psi_encoded = f[f"/{pid}/{det_name}/psi/"][()] @@ -122,14 +126,16 @@ def tod_reader(band_comm: MPI.Comm, my_experiment: str, my_band: Bunch, det_name if np.sum(flag_buffer[:ntod_optimal]) != 0: good_scan = False - detector = DetectorTOD(tod, pix_encoded, psi_encoded, my_band.eval_nside, - data_nside, fsamp, vsun, huffman_tree, huffman_symbols, - npsi, processing_mask_map, ntod, - pix_is_compressed=my_experiment.pix_is_compressed, - psi_is_compressed=my_experiment.psi_is_compressed) + det_pointing = PixelPointing(pix_encoded, psi_encoded, huffman_tree, + huffman_symbols, npsi, my_band.eval_nside, data_nside, + ntod, ntod_optimal) + detector = DetectorTOD(det_name, idet, idet_accepted, tod, det_pointing, fsamp, + vsun, huffman_tree, huffman_symbols, processing_mask_map, + ntod, ntod_optimal) detector_list.append(detector) ntod_sum_original += ntod ntod_sum_final += ntod_optimal + idet_accepted += 1 if good_scan: scanID = int(pid) scan = ScanTOD(detector_list, 0., scanID) diff --git a/src/commander4/experiments/litebird/tod_reader_litebird_sim_spawndetectors.py b/src/commander4/experiments/litebird/tod_reader_litebird_sim_spawndetectors.py index 778b31d..9ca0f6f 100644 --- a/src/commander4/experiments/litebird/tod_reader_litebird_sim_spawndetectors.py +++ b/src/commander4/experiments/litebird/tod_reader_litebird_sim_spawndetectors.py @@ -12,6 +12,7 @@ from commander4.data_models.detector_TOD import DetectorTOD from commander4.data_models.scan_TOD import ScanTOD from commander4.data_models.detector_group_TOD import DetGroupTOD +from commander4.data_models.pointing import PixelPointing from commander4.noise_sampling.noise_psd import NoisePSD, NoisePSDOof from commander4.simulations.inplace_litebird_sim import replace_tod_with_sim from commander4.output.log import logassert @@ -127,6 +128,7 @@ def tod_reader(band_comm: MPI.Comm, my_experiment: str, my_band: Bunch, det_name detector_list = [] tod_alldet = np.zeros((ndet, ntod_optimal), dtype=np.float32) for idet in range(ndet): + det_name = det_names[idet] tod = tod_alldet[idet] # Ideally we would do the huffman compression outside the detector-loop, but that's @@ -139,11 +141,16 @@ def tod_reader(band_comm: MPI.Comm, my_experiment: str, my_band: Bunch, det_name psi_encoded = huffman.huffman_compress_array(psi, sym_codes, sym_lengths) pix_encoded = huffman.huffman_compress_array(pix, sym_codes, sym_lengths) - detector = DetectorTOD(tod, pix_encoded, psi_encoded, my_band.eval_nside, - data_nside, fsamp, vsun, huffman_tree, huffman_symbols, - npsi, processing_mask_map, ntod_optimal, - pix_is_compressed=True, # Hard-coded to true since we compress manually. - psi_is_compressed=True) + # Detectors are spawned from one shared pointing, so every detector is present in every + # scan: det_idx_fullband == det_idx_local == idet. pix/psi were just compressed to bytes + # (PixelPointing auto-detects the compression) at the optimal length, so for this reader + # ntod_original == ntod == ntod_optimal. + det_pointing = PixelPointing(pix_encoded, psi_encoded, huffman_tree, huffman_symbols, + npsi, my_band.eval_nside, data_nside, ntod_optimal, + ntod_optimal) + detector = DetectorTOD(det_name, idet, idet, tod, det_pointing, fsamp, vsun, + huffman_tree, huffman_symbols, processing_mask_map, + ntod_optimal, ntod_optimal) detector_list.append(detector) ntod_sum_original += ntod ntod_sum_final += ntod_optimal diff --git a/src/commander4/experiments/planck/tod_reader_planck.py b/src/commander4/experiments/planck/tod_reader_planck.py index f3d4a93..b9ada5d 100644 --- a/src/commander4/experiments/planck/tod_reader_planck.py +++ b/src/commander4/experiments/planck/tod_reader_planck.py @@ -8,7 +8,6 @@ from astropy.io import fits from mpi4py import MPI from pixell.bunch import Bunch -from commander4.cmdr4_support import utils as cpp_utils from commander4.data_models.detector_TOD import DetectorTOD from commander4.data_models.scan_TOD import ScanTOD from commander4.data_models.detector_group_TOD import DetGroupTOD @@ -73,11 +72,6 @@ def tod_reader(band_comm: MPI.Comm, my_experiment: str, my_band: Bunch, all_det_ Fourier_times = np.load(my_experiment.Fourier_times_path) - # Attempting to reduce fragmentation by allocating buffers. - ntod_upper_bound = int(my_band.fsamp*100*3600) # 10 hour scan. - flag_buffer = np.zeros(ntod_upper_bound, dtype=np.int64) - # tod_buffer = np.zeros(ntod_upper_bound, dtype=np.float32) - scan_list = [] nscans = scan_idx_stop - scan_idx_start num_included = 0 @@ -106,27 +100,23 @@ def tod_reader(band_comm: MPI.Comm, my_experiment: str, my_band: Bunch, all_det_ fsamp = float(f["/common/fsamp/"][()].item()) npsi = int(f["/common/npsi/"][()].item()) detector_list = [] + # idet is the detector's full-band column (its position in ``all_det_names``); + # idet_accepted (det_idx_local) advances only when a detector survives the cuts below. idet_accepted = 0 for idet, det_name in enumerate(all_det_names): tod = f[f"/{pid}/{det_name}/tod/"][:ntod_optimal].astype(np.float32, copy=False) pix_encoded = f[f"/{pid}/{det_name}/pix/"][()] - psi_encoded = f[f"/{pid}/{det_name}/psi/"][()] if "QU" in my_band.polarization else [] + # Intensity-only bands have no psi in the files; feed a zero psi (unused by I-only + # mapmaking, but PixelPointing requires a length-matched array). + if "QU" in my_band.polarization: + psi_encoded = f[f"/{pid}/{det_name}/psi/"][()] + else: + psi_encoded = np.zeros(ntod_optimal, dtype=np.float32) flag_encoded = f[f"/{pid}/{det_name}/flag/"][()] init_scalars = f[f"/{pid}/{det_name}/scalars/"][()] # Data format has this weird thing were gain seems to be in "micro-gain"... init_scalars[0] *= 1e-6 - flag_buffer[:ntod] = cpp_utils.huffman_decode(np.frombuffer(flag_encoded, dtype=np.uint8), - huffman_tree, huffman_symbols, flag_buffer[:ntod]) - flag_buffer[:ntod_optimal] = np.cumsum(flag_buffer[:ntod_optimal]) - flag_buffer[:ntod_optimal] &= 6111232 - if np.sum(flag_buffer[:ntod_optimal]) != 0: - good_scan = False - # tod_buffer[:ntod_optimal] = np.abs(tod) - # Check for crazy data. - if np.mean(np.abs(tod)) > 0.001 or np.std(tod) > 0.001: - good_scan = False - det_init_scalars[idet] = init_scalars det_pointing = PixelPointing(pix_encoded, psi_encoded, huffman_tree, huffman_symbols, npsi, my_band.eval_nside, data_nside, @@ -135,12 +125,25 @@ def tod_reader(band_comm: MPI.Comm, my_experiment: str, my_band: Bunch, all_det_ detector = DetectorTOD(det_name, idet, idet_accepted, tod, det_pointing, fsamp, vsun, huffman_tree, huffman_symbols, processing_mask_map, ntod, ntod_optimal, + flag_encoded=flag_encoded, bad_data_bitmask = 6111232, init_scalars = init_scalars) + # Bad samples are handled per-sample by the detector's full_mask (flag & + # bad_data_bitmask), rather than dropping the whole scan. Following the SO readers, + # only skip detectors with too little usable data (or pathological TOD magnitudes). + unmasked_fraction = np.sum(detector.full_mask)/detector.full_mask.size + if unmasked_fraction < 0.99: + continue + if (detector.tod == 0).all(): + continue + if np.mean(np.abs(tod)) > 0.001 or np.std(tod) > 0.001: + continue detector_list.append(detector) ntod_sum_original += ntod ntod_sum_final += ntod_optimal idet_accepted += 1 + if len(detector_list) == 0: + good_scan = False if good_scan: scan = ScanTOD(detector_list, 0., scanID) scan_list.append(scan)