diff --git a/CHANGELOG.md b/CHANGELOG.md index 94fabcb0..1f68317d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,5 @@ # HEAD +- Extract a shared driver for the systematic-effect modules (`for_each_observation`, `for_each_observation_with_pointings`, `normalize_observations`) and rename `pointings_in_obs.py` to `observation_utilities.py`, which now collects all the helpers that operate on a collection of `Observation` objects. Noise, gain drift, non-linearity, HWP differential emission and the CMB dipole now share one place for observation normalization, per-detector RNG resolution and TOD-component indirection instead of re-implementing it each. No change to public effect interfaces. - Stop supporting unused MPI processes and drop `MPI_COMM_GRID` class and its attributes. Now throwing an error when `comm_size != n_blocks_det * n_blocks_time`. It fixes [#364](https://github.com/litebird/litebird_sim/issues/364), while incorporating selective changes from [#372](https://github.com/litebird/litebird_sim/pull/372) and [#427](https://github.com/litebird/litebird_sim/pull/427). Reference PR: [#539](https://github.com/litebird/litebird_sim/pull/539). diff --git a/docs/source/h_maps.rst b/docs/source/h_maps.rst index c9c25ea6..8c30f9bf 100644 --- a/docs/source/h_maps.rst +++ b/docs/source/h_maps.rst @@ -38,7 +38,7 @@ These maps capture how the detector orientation is distributed across the sky during the scanning strategy. Setting :math:`m=0` is equivalent to the definition of h maps in McCallum et al (2021) without HWP modulation. -The maps can be generated by using the :func:`make_h_maps` function with a list of observations, or by using the observations of a simulation through the method :func:`.make_h_maps` of :class:`simulation` +The maps can be generated by using the :func:`~litebird_sim.mapmaking.h_maps.make_h_maps` function with a list of observations, or by using the observations of a simulation through the :meth:`.Simulation.make_h_maps` method. Example ----------- @@ -100,7 +100,7 @@ Note: for the couple (0,0) the resulting map is not the one expected from the de Output: ``HMapsResult`` ----------------------- -:func:`make_h_maps` returns a :class:`HMapsResult` object, which contains: +:func:`~litebird_sim.mapmaking.h_maps.make_h_maps` returns a :class:`~litebird_sim.mapmaking.h_maps.HMapsResult` object, which contains: - ``h_maps``: a dictionary indexed by detector name, then by ``(n, m)`` tuple, each entry being a :class:`.h_map_Re_and_Im` object. @@ -149,7 +149,7 @@ Maps are saved in **HDF5** format, one file per detector: │ └── Im └── ... -To reload maps from disk, use :func:`.load_h_maps_from_file`:: +To reload maps from disk, use :func:`~litebird_sim.mapmaking.h_maps.load_h_maps_from_file`:: from litebird_sim.mapmaking.h_maps import load_h_maps_from_file @@ -165,7 +165,7 @@ and all time samples. MPI support ----------- -The current implementation of :func:`make_h_maps` only allows to distribute observation by detector, i.e. each MPI process computes the h maps for a subset of detectors, but using all time samples. +The current implementation of :func:`~litebird_sim.mapmaking.h_maps.make_h_maps` only allows to distribute observation by detector, i.e. each MPI process computes the h maps for a subset of detectors, but using all time samples. Example: .. code-block:: python diff --git a/docs/source/index.rst b/docs/source/index.rst index 0e490b59..0939b31b 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -20,6 +20,21 @@ This is the User’s Manual of the LiteBIRD Simulation Framework. part6.rst appendix.rst +How to cite +=========== + +If you use the LiteBIRD Simulation Framework in your work, please cite +:cite:`2025:litebird:tomasi`: + + M. Tomasi *et al.* (LiteBIRD Collaboration), "A simulation framework for the + LiteBIRD instruments," *JCAP* **11** (2025) 040, + doi:`10.1088/1475-7516/2025/11/040 + `_, + `arXiv:2507.04918 `_. + +The corresponding BibTeX entry (key ``2025:litebird:tomasi``) is available in +:doc:`the bibliography `. + Indices and tables ================== diff --git a/docs/source/part1.rst b/docs/source/part1.rst index 2dee4393..3325b2d2 100644 --- a/docs/source/part1.rst +++ b/docs/source/part1.rst @@ -1,6 +1,22 @@ Introduction ============ +`LiteBIRD `_ is a JAXA-led satellite mission that +will map the polarization of the Cosmic Microwave Background over the full sky, +searching for the primordial *B*-mode signature of cosmic inflation. Meeting its +requirements relies on detailed end-to-end simulations of the instruments and of +the observation process. + +The **LiteBIRD Simulation Framework** (LBS) is the collaboration's toolkit for +those simulations: it turns *sky models* and *instrument descriptions* into +*time-ordered detector data* (TOD), injects instrumental noise and systematic +effects, and reduces the TOD back into sky maps. The rest of this manual +documents these steps part by part. + +If you are new to the framework, start with the installation instructions below +and then work through the :doc:`tutorial `. If you use LBS in your +work, please cite it (:cite:`2025:litebird:tomasi`). + .. toctree:: :maxdepth: 1 diff --git a/docs/source/part2.rst b/docs/source/part2.rst index 5b16f5c4..d16457a4 100644 --- a/docs/source/part2.rst +++ b/docs/source/part2.rst @@ -1,6 +1,11 @@ Structure of the framework ========================== +This part describes the core building blocks of the framework: the central +``Simulation`` class, the objects it ties together (observations, detectors, the +instrument model), and the cross-cutting facilities shared by every module — +units, reports, reproducible random numbers, and MPI parallelism. + .. toctree:: :maxdepth: 1 diff --git a/docs/source/part3.rst b/docs/source/part3.rst index a96bf860..189131ed 100644 --- a/docs/source/part3.rst +++ b/docs/source/part3.rst @@ -1,6 +1,10 @@ Basic simulation modules ======================== +This part covers the modules that synthesize time-ordered data: how the +spacecraft scans the sky (pointing and quaternions), how an input sky is sampled +into a TOD, beam convolution, the CMB dipole, and instrumental noise. + .. toctree:: :maxdepth: 1 diff --git a/docs/source/part4.rst b/docs/source/part4.rst index 58531bd4..55df3c34 100644 --- a/docs/source/part4.rst +++ b/docs/source/part4.rst @@ -1,6 +1,11 @@ Systematic effects ================== +This part documents the instrumental systematic effects that can be injected +into the TOD — gain drifts, detector non-linearity, half-wave-plate systematics, +and pointing disturbances — together with the spin-based :math:`h`-maps that +encode the orientation information of the scanning strategy. + .. toctree:: :maxdepth: 1 diff --git a/docs/source/part5.rst b/docs/source/part5.rst index d55fc66d..170578e9 100644 --- a/docs/source/part5.rst +++ b/docs/source/part5.rst @@ -1,6 +1,9 @@ Data-reduction modules ====================== +This part describes how time-ordered data is reduced back into sky maps, +covering the binned, destriped, and pair-differenced map-making algorithms. + .. toctree:: :maxdepth: 1 diff --git a/docs/source/part6.rst b/docs/source/part6.rst index f8526797..c03d680c 100644 --- a/docs/source/part6.rst +++ b/docs/source/part6.rst @@ -1,6 +1,9 @@ Tutorials ========= +This part collects worked, end-to-end examples that walk through a complete +simulation, from setting up the instrument to producing maps. + .. toctree:: :maxdepth: 1 diff --git a/docs/source/refs.bib b/docs/source/refs.bib index 86d55596..fa6c4765 100644 --- a/docs/source/refs.bib +++ b/docs/source/refs.bib @@ -159,4 +159,21 @@ @ARTICLE{2019:duivenvoorden:beamconv adsnote = {Provided by the SAO/NASA Astrophysics Data System} } +@ARTICLE{2025:litebird:tomasi, + author = {{Tomasi}, M. and {LiteBIRD Collaboration}}, + title = "{A simulation framework for the LiteBIRD instruments}", + journal = {JCAP}, + keywords = {Astrophysics - Instrumentation and Methods for Astrophysics}, + year = 2025, + month = nov, + volume = {2025}, + number = {11}, + eid = {040}, + pages = {040}, + doi = {10.1088/1475-7516/2025/11/040}, +archivePrefix = {arXiv}, + eprint = {2507.04918}, + primaryClass = {astro-ph.IM} +} + diff --git a/docs/source/scanning.rst b/docs/source/scanning.rst index 8166a468..17eb0d6d 100644 --- a/docs/source/scanning.rst +++ b/docs/source/scanning.rst @@ -848,7 +848,7 @@ API reference :undoc-members: :show-inheritance: -.. automodule:: litebird_sim.pointings_in_obs +.. automodule:: litebird_sim.observation_utilities :members: :undoc-members: :show-inheritance: diff --git a/litebird_sim/__init__.py b/litebird_sim/__init__.py index a426a2d5..15424e83 100644 --- a/litebird_sim/__init__.py +++ b/litebird_sim/__init__.py @@ -146,6 +146,14 @@ apply_quadratic_nonlin_to_observations, ) from .observations import Observation, TodDescription +from .observation_utilities import ( + normalize_observations, + for_each_observation, + for_each_observation_with_pointings, + prepare_pointings, + precompute_pointings, + apply_hwp_to_obs, +) from .pointing_sys import ( get_detector_orientation, FocalplaneCoord, @@ -155,11 +163,6 @@ from .pointings import ( PointingProvider, ) -from .pointings_in_obs import ( - prepare_pointings, - precompute_pointings, - apply_hwp_to_obs, -) from .profiler import TimeProfiler, profile_list_to_speedscope from .quaternions import ( quat_rotation_x, @@ -363,7 +366,10 @@ "RNGHierarchy", # pointings.py "PointingProvider", - # pointings_in_obs.py + # observation_utilities.py + "normalize_observations", + "for_each_observation", + "for_each_observation_with_pointings", "prepare_pointings", "precompute_pointings", "apply_hwp_to_obs", diff --git a/litebird_sim/beam_convolution.py b/litebird_sim/beam_convolution.py index a02a0f5c..18d26952 100644 --- a/litebird_sim/beam_convolution.py +++ b/litebird_sim/beam_convolution.py @@ -9,7 +9,7 @@ from .hwp import HWP from .mueller_convolver import MuellerConvolver from .observations import Observation -from .pointings_in_obs import ( +from .observation_utilities import ( _get_hwp_angle, _normalize_observations_and_pointings, _get_pointings_array, diff --git a/litebird_sim/dipole.py b/litebird_sim/dipole.py index a9b5ed5f..d6bad0c7 100644 --- a/litebird_sim/dipole.py +++ b/litebird_sim/dipole.py @@ -8,6 +8,7 @@ from .maps_and_harmonics import SphericalHarmonics from .observations import Observation from .spacecraft import SpacecraftPositionAndVelocity +from .observation_utilities import for_each_observation_with_pointings # Updated imports to match the new constants.py structure from .constants import C_LIGHT_KM_OVER_S, H_OVER_K_B, T_CMB_K @@ -882,43 +883,14 @@ def add_dipole_to_observations( # For convolved types we keep the full (θ, φ, ψ) columns; otherwise strip to (θ, φ). ptg_cols = slice(None) if apply_convolution else slice(0, 2) - if pointings is None: - if isinstance(observations, Observation): - obs_list = [observations] - if hasattr(observations, "pointing_matrix"): - ptg_list = [observations.pointing_matrix[:, :, ptg_cols]] - else: - ptg_list = [observations.get_pointings] - else: - obs_list = observations - ptg_list = [] - for ob in observations: - if hasattr(ob, "pointing_matrix"): - ptg_list.append(ob.pointing_matrix[:, :, ptg_cols]) - else: - ptg_list.append(ob.get_pointings) - else: - if isinstance(observations, Observation): - assert isinstance(pointings, np.ndarray), ( - "You must pass a list of observations *and* a list " - + "of pointing matrices to add_dipole_to_observations" - ) - obs_list = [observations] - ptg_list = [pointings[:, :, ptg_cols]] - else: - assert isinstance(pointings, list), ( - "When you pass a list of observations to add_dipole_to_observations" - + ", you must do the same for `pointings`" - ) - assert len(observations) == len(pointings), ( - f"The list of observations has {len(observations)} elements, but " - + f"the list of pointings has {len(pointings)} elements" - ) - obs_list = observations - ptg_list = [point[:, :, ptg_cols] for point in pointings] - - for cur_obs, cur_ptg in zip(obs_list, ptg_list): - tod = getattr(cur_obs, component) + for cur_obs, tod, cur_ptg in for_each_observation_with_pointings( + observations, pointings, component + ): + # Callables (lazy pointings) are forwarded untouched, matching the + # behaviour of the underlying normalizer. Convolved calculations need + # the psi column too, hence ptg_cols rather than a hardcoded 0:2. + if isinstance(cur_ptg, np.ndarray): + cur_ptg = cur_ptg[:, :, ptg_cols] # Resolve the beam S-parameters when convolution is requested. # 1. Use the explicitly provided beam_alms if given. diff --git a/litebird_sim/gaindrifts.py b/litebird_sim/gaindrifts.py index 85047b34..ec7c3885 100644 --- a/litebird_sim/gaindrifts.py +++ b/litebird_sim/gaindrifts.py @@ -5,7 +5,7 @@ import numpy as np from .observations import Observation -from .seeding import regenerate_or_check_detector_generators +from .observation_utilities import for_each_observation class GainDriftType(IntEnum): @@ -548,23 +548,13 @@ def apply_gaindrift_to_observations( if drift_params is None: drift_params = GainDriftParams() - if isinstance(observations, Observation): - obs_list = [observations] - elif isinstance(observations, list): - obs_list = observations - else: - raise TypeError( - "The parameter `observations` must be an `Observation` or a list of `Observation`." - ) - dets_random = regenerate_or_check_detector_generators( - observations=obs_list, - comm=None, + for cur_obs, tod, dets_random in for_each_observation( + observations, + component, user_seed=user_seed, dets_random=dets_random, - ) - - for cur_obs in obs_list: - tod = getattr(cur_obs, component) + requires_rng=True, + ): sampling_freq_hz = cur_obs.sampling_rate_hz focalplane_attr = getattr(cur_obs, drift_params.focalplane_group) diff --git a/litebird_sim/hwp_diff_emiss.py b/litebird_sim/hwp_diff_emiss.py index 860cbd8f..e0859650 100644 --- a/litebird_sim/hwp_diff_emiss.py +++ b/litebird_sim/hwp_diff_emiss.py @@ -5,7 +5,8 @@ from .hwp import HWP from .observations import Observation -from .pointings_in_obs import _get_hwp_angle +from .observation_utilities import _get_hwp_angle +from .observation_utilities import for_each_observation # We calculate the additive signal coming from hwp harmonics. @@ -83,13 +84,8 @@ def add_2f_to_observations( # in `observations.hwp_2f_tod` add_2f_to_observations(sim.observations, component="hwp_2f_tod") """ - if isinstance(observations, Observation): - obs_list = [observations] - else: - obs_list = observations - # iterate through each observation - for cur_obs in obs_list: + for cur_obs, tod, _ in for_each_observation(observations, component): if amplitude_2f_k is None: amplitude_2f_k = getattr(cur_obs, "amplitude_2f_k", None) assert amplitude_2f_k is not None, ( @@ -104,7 +100,7 @@ def add_2f_to_observations( hwp_angle = _get_hwp_angle(obs=cur_obs, hwp=hwp, pointing_dtype=pointings_dtype) add_2f( - tod=getattr(cur_obs, component), + tod=tod, hwp_angle=hwp_angle, pol_angle_rad=cur_obs.pol_angle_rad, amplitude_2f_k=amplitude_2f_k, diff --git a/litebird_sim/hwp_harmonics/hwp_harmonics.py b/litebird_sim/hwp_harmonics/hwp_harmonics.py index 51484e97..9bcc0f37 100644 --- a/litebird_sim/hwp_harmonics/hwp_harmonics.py +++ b/litebird_sim/hwp_harmonics/hwp_harmonics.py @@ -17,7 +17,7 @@ from ..input_sky import SkyInput from ..maps_and_harmonics import HealpixMap, SphericalHarmonics, interpolate_alm from ..observations import Observation -from ..pointings_in_obs import ( +from ..observation_utilities import ( _get_pointings_array, ) from .jones_methods import ( diff --git a/litebird_sim/madam.py b/litebird_sim/madam.py index bfdc4026..de4b3967 100644 --- a/litebird_sim/madam.py +++ b/litebird_sim/madam.py @@ -15,7 +15,7 @@ from .hwp import HWP from .mapmaking import ExternalDestriperParameters from .observations import Observation -from .pointings_in_obs import _get_pointings_and_pol_angles_det +from .observation_utilities import _get_pointings_and_pol_angles_det from .simulations import MpiDistributionDescr, Simulation diff --git a/litebird_sim/mapmaking/binner.py b/litebird_sim/mapmaking/binner.py index 3961c408..8e0dfa75 100644 --- a/litebird_sim/mapmaking/binner.py +++ b/litebird_sim/mapmaking/binner.py @@ -21,7 +21,7 @@ from litebird_sim.coordinates import CoordinateSystem from litebird_sim.hwp import HWP from litebird_sim.observations import Observation -from litebird_sim.pointings_in_obs import ( +from litebird_sim.observation_utilities import ( _get_hwp_angle, _normalize_observations_and_pointings, ) diff --git a/litebird_sim/mapmaking/common.py b/litebird_sim/mapmaking/common.py index 23d25bc4..3432df24 100644 --- a/litebird_sim/mapmaking/common.py +++ b/litebird_sim/mapmaking/common.py @@ -9,7 +9,7 @@ from litebird_sim.coordinates import CoordinateSystem from litebird_sim.observations import Observation -from litebird_sim.pointings_in_obs import _get_pointings_array, _get_pol_angle +from litebird_sim.observation_utilities import _get_pointings_array, _get_pol_angle # The threshold on the conditioning number used to determine if a pixel # was really “seen” or not diff --git a/litebird_sim/mapmaking/destriper.py b/litebird_sim/mapmaking/destriper.py index 0c8c328e..ffe9cda1 100644 --- a/litebird_sim/mapmaking/destriper.py +++ b/litebird_sim/mapmaking/destriper.py @@ -16,7 +16,7 @@ from litebird_sim.hwp import HWP from litebird_sim.mpi import MPI_ENABLED, MPI_COMM_WORLD from litebird_sim.observations import Observation -from litebird_sim.pointings_in_obs import ( +from litebird_sim.observation_utilities import ( _get_hwp_angle, _normalize_observations_and_pointings, ) diff --git a/litebird_sim/mapmaking/h_maps.py b/litebird_sim/mapmaking/h_maps.py index 9d079a30..69fe41e0 100644 --- a/litebird_sim/mapmaking/h_maps.py +++ b/litebird_sim/mapmaking/h_maps.py @@ -14,7 +14,7 @@ from litebird_sim.hwp import HWP from litebird_sim.maps_and_harmonics import HealpixMap from litebird_sim.observations import Observation -from litebird_sim.pointings_in_obs import ( +from litebird_sim.observation_utilities import ( _get_hwp_angle, _normalize_observations_and_pointings, ) diff --git a/litebird_sim/mapmaking/pair_differencing.py b/litebird_sim/mapmaking/pair_differencing.py index 7cfe5c58..d1247859 100644 --- a/litebird_sim/mapmaking/pair_differencing.py +++ b/litebird_sim/mapmaking/pair_differencing.py @@ -21,7 +21,7 @@ from litebird_sim.coordinates import CoordinateSystem from litebird_sim.hwp import HWP from litebird_sim.observations import Observation -from litebird_sim.pointings_in_obs import ( +from litebird_sim.observation_utilities import ( _get_hwp_angle, _normalize_observations_and_pointings, ) diff --git a/litebird_sim/noise.py b/litebird_sim/noise.py index 38971e14..ff002786 100644 --- a/litebird_sim/noise.py +++ b/litebird_sim/noise.py @@ -3,8 +3,7 @@ from numba import njit from ducc0.misc import OofaNoise -from .observations import Observation -from .seeding import regenerate_or_check_detector_generators +from .observation_utilities import for_each_observation # --- TRANSFER FUNCTIONS (MODELS) --- @@ -855,18 +854,6 @@ def add_noise_to_observations( if noise_type not in ["white", "one_over_f", "correlated"]: raise ValueError("Unknown noise type " + noise_type) - if isinstance(observations, Observation): - obs_list = [observations] - else: - obs_list = observations - - dets_random = regenerate_or_check_detector_generators( - observations=obs_list, - comm=None, - user_seed=user_seed, - dets_random=dets_random, - ) - if noise_type == "correlated": if correlation is None: raise ValueError( @@ -877,7 +864,13 @@ def add_noise_to_observations( common_mode_type = correlation.get("common_mode_type", "one_over_f") group_by = correlation.get("group_by", None) - for cur_obs in obs_list: + for cur_obs, tod, dets_random in for_each_observation( + observations, + component, + user_seed=user_seed, + dets_random=dets_random, + requires_rng=True, + ): if corr_matrix is not None: groups = None elif "groups" in correlation: @@ -885,7 +878,7 @@ def add_noise_to_observations( else: groups = _build_detector_groups(cur_obs, group_by) add_correlated_noise( - tod=getattr(cur_obs, component), + tod=tod, sampling_rate_hz=cur_obs.sampling_rate_hz, net_ukrts=cur_obs.net_ukrts, fknee_mhz=getattr(cur_obs, "fknee_mhz"), @@ -902,9 +895,15 @@ def add_noise_to_observations( ) return - for cur_obs in obs_list: + for cur_obs, tod, dets_random in for_each_observation( + observations, + component, + user_seed=user_seed, + dets_random=dets_random, + requires_rng=True, + ): add_noise( - tod=getattr(cur_obs, component), + tod=tod, noise_type=noise_type, sampling_rate_hz=cur_obs.sampling_rate_hz, net_ukrts=cur_obs.net_ukrts * scale, diff --git a/litebird_sim/non_linearity.py b/litebird_sim/non_linearity.py index 9164d31d..d3d003e3 100644 --- a/litebird_sim/non_linearity.py +++ b/litebird_sim/non_linearity.py @@ -3,7 +3,7 @@ import numpy as np from .observations import Observation -from .seeding import regenerate_or_check_detector_generators +from .observation_utilities import for_each_observation, normalize_observations from .units import Units, UnitUtils @@ -202,23 +202,16 @@ def apply_quadratic_nonlin_to_observations( if nl_params is None: nl_params = NonLinParams() - if isinstance(observations, Observation): - obs_list = [observations] - elif isinstance(observations, list): - obs_list = observations - else: - raise TypeError( - "The parameter `observations` must be an `Observation` or a list of `Observation`." - ) - dets_random = regenerate_or_check_detector_generators( - observations=obs_list, - comm=obs_list[0].comm_time_block, - user_seed=user_seed, - ) + obs_list = normalize_observations(observations) # iterate through each observation - for cur_obs in obs_list: - tod = getattr(cur_obs, component) + for cur_obs, tod, dets_random in for_each_observation( + obs_list, + component, + user_seed=user_seed, + requires_rng=True, + comm=obs_list[0].comm_time_block, + ): bandcenter_ghz = getattr(cur_obs, "bandcenter_ghz") bandwidth_ghz = getattr(cur_obs, "bandwidth_ghz") diff --git a/litebird_sim/pointings_in_obs.py b/litebird_sim/observation_utilities.py similarity index 70% rename from litebird_sim/pointings_in_obs.py rename to litebird_sim/observation_utilities.py index edef0bc9..6bf09e79 100644 --- a/litebird_sim/pointings_in_obs.py +++ b/litebird_sim/observation_utilities.py @@ -1,4 +1,4 @@ -from collections.abc import Callable +from collections.abc import Callable, Iterator import astropy.time import numpy as np @@ -11,6 +11,7 @@ from .hwp import HWP from .observations import Observation from .scanning import RotQuaternion +from .seeding import regenerate_or_check_detector_generators def prepare_pointings( @@ -419,3 +420,150 @@ def _get_pointings_and_pol_angles_det( ) return pointings_det, pol_angle + + +# --------------------------------------------------------------------------- +# Shared driver for the systematic-effect modules +# +# Every module that injects an effect into the time-ordered data (noise, gain +# drift, non-linearity, HWP differential emission, the CMB dipole, …) used to +# re-implement the same preamble before getting to its own physics: coerce the +# ``observations`` argument into a list, optionally build one RNG per detector +# for the current MPI rank, then loop pulling ``getattr(obs, component)`` to +# find the TOD array to modify. That preamble is identical across the effects, +# so it lives here once. An effect module now writes only the part that +# differs — which detector parameters it reads and which kernel it calls. +# --------------------------------------------------------------------------- + + +def normalize_observations( + observations: Observation | list[Observation], +) -> list[Observation]: + """Coerce *observations* into a list of :class:`.Observation`. + + A single :class:`.Observation` is wrapped in a one-element list; a list is + copied (so callers may mutate the result without touching their input). + Anything else raises :class:`TypeError`. + """ + if isinstance(observations, Observation): + return [observations] + if isinstance(observations, list): + return list(observations) + raise TypeError( + "The parameter `observations` must be an `Observation` or a list of " + "`Observation`." + ) + + +def _get_tod_component(obs: Observation, component: str) -> np.ndarray: + """Fetch the named TOD component from *obs*, failing fast on a bad name. + + Effect modules pass ``component`` as a plain string. A typo would otherwise + reach ``getattr`` and blow up one iteration into the loop, after RNG streams + are already built, with a bare ``AttributeError``. Check up front instead + and raise with the observation's registered TOD names (``obs.tod_list``) as + a hint. Any existing attribute is accepted, not only registered TODs, so + callers that stash a scratch component on the observation keep working. + """ + if not hasattr(obs, component): + registered = [td.name for td in obs.tod_list] + raise ValueError( + f"'{component}' is not an attribute of this observation; " + f"registered TOD components are {registered}." + ) + return getattr(obs, component) + + +def for_each_observation( + observations: Observation | list[Observation], + component: str = "tod", + *, + user_seed: int | None = None, + dets_random: list[np.random.Generator] | None = None, + requires_rng: bool = False, + comm=None, +) -> Iterator[tuple[Observation, np.ndarray, list[np.random.Generator] | None]]: + """Iterate over observations, yielding ``(obs, tod, dets_random)``. + + This is the common driver for the systematic-effect modules. It performs + the bookkeeping every effect shares: + + - normalizes *observations* into a list + (see :func:`normalize_observations`); + - if *requires_rng* is ``True``, resolves one RNG per detector for the + current MPI rank via + :func:`.regenerate_or_check_detector_generators` and yields the same list + on every iteration; + - yields the ``component`` TOD array (validated against + ``obs.tod_list``) so the caller does not repeat the attribute + indirection. + + Parameters + ---------- + observations : Observation or list of Observation + The observation(s) to iterate over. + component : str, optional + Name of the TOD attribute to fetch from each observation. Defaults to + ``"tod"``. + user_seed : int, optional + Master seed used to build the per-detector RNGs when *dets_random* is + not supplied. Only consulted when *requires_rng* is ``True``. + dets_random : list of numpy.random.Generator, optional + Pre-built per-detector RNGs. Only consulted when *requires_rng* is + ``True``. + requires_rng : bool, optional + Whether the effect is stochastic. When ``True`` exactly one of + *user_seed* or *dets_random* must be provided. When ``False`` (the + default) no RNGs are built and ``None`` is yielded in their place. + comm : optional + MPI communicator forwarded to + :func:`.regenerate_or_check_detector_generators`. Only consulted + when *requires_rng* is ``True``. + + Yields + ------ + tuple + ``(cur_obs, tod, dets_random)`` for each observation. ``dets_random`` + is the resolved list of generators when *requires_rng* is ``True``, + otherwise ``None``. The same ``dets_random`` object is yielded on every + iteration. + """ + obs_list = normalize_observations(observations) + + if requires_rng: + dets_random = regenerate_or_check_detector_generators( + observations=obs_list, + comm=comm, + user_seed=user_seed, + dets_random=dets_random, + ) + else: + dets_random = None + + for cur_obs in obs_list: + yield cur_obs, _get_tod_component(cur_obs, component), dets_random + + +def for_each_observation_with_pointings( + observations: Observation | list[Observation], + pointings: np.ndarray | list[np.ndarray] | None, + component: str = "tod", +) -> Iterator[tuple[Observation, np.ndarray, np.ndarray]]: + """Iterate over observations paired with their pointing matrices. + + The pointing-aware counterpart of :func:`for_each_observation`, for effects + that need a pointing matrix per observation (e.g. the CMB dipole). It uses + :func:`_normalize_observations_and_pointings`, so a single observation may + be paired with a single pointing array, or a list with a list. When + *pointings* is ``None`` the pointing matrix is taken from each observation. + + Yields + ------ + tuple + ``(cur_obs, tod, cur_ptg)`` for each observation, where ``cur_ptg`` is + either a pointing array or the observation's ``get_pointings`` callable + (matching the behaviour of the underlying normalizer). + """ + obs_list, ptg_list = _normalize_observations_and_pointings(observations, pointings) + for cur_obs, cur_ptg in zip(obs_list, ptg_list): + yield cur_obs, _get_tod_component(cur_obs, component), cur_ptg diff --git a/litebird_sim/observations.py b/litebird_sim/observations.py index 31f85c3a..e45b4668 100644 --- a/litebird_sim/observations.py +++ b/litebird_sim/observations.py @@ -158,6 +158,9 @@ class Observation: # Dynamic attributes set by beam synthesis name: list + channel: list + bandcenter_ghz: npt.NDArray + bandwidth_ghz: npt.NDArray fwhm_arcmin: npt.NDArray ellipticity: npt.NDArray psi_rad: npt.NDArray diff --git a/litebird_sim/scan_map.py b/litebird_sim/scan_map.py index a5379beb..ff9f315a 100644 --- a/litebird_sim/scan_map.py +++ b/litebird_sim/scan_map.py @@ -12,7 +12,7 @@ from .input_sky import SkyInput from .maps_and_harmonics import HealpixMap, SphericalHarmonics, interpolate_alm from .observations import Observation -from .pointings_in_obs import ( +from .observation_utilities import ( _get_hwp_angle, _get_pointings_array, _get_pol_angle, diff --git a/litebird_sim/simulations.py b/litebird_sim/simulations.py index 10ee60ff..f48fa854 100644 --- a/litebird_sim/simulations.py +++ b/litebird_sim/simulations.py @@ -65,7 +65,7 @@ from .noise import add_noise_to_observations from .non_linearity import NonLinParams, apply_quadratic_nonlin_to_observations from .observations import Observation, TodDescription -from .pointings_in_obs import ( +from .observation_utilities import ( precompute_pointings, prepare_pointings, ) diff --git a/test/test_destriper.py b/test/test_destriper.py index 78d41b71..953afdca 100644 --- a/test/test_destriper.py +++ b/test/test_destriper.py @@ -375,7 +375,7 @@ def setup_simulation( def test_map_maker_parts(): - from litebird_sim.pointings_in_obs import _normalize_observations_and_pointings + from litebird_sim.observation_utilities import _normalize_observations_and_pointings from litebird_sim.mapmaking.destriper import ( _store_pixel_idx_and_pol_angle_in_obs, _build_mask_detector_split, diff --git a/test/test_hwp_diff_emiss.py b/test/test_hwp_diff_emiss.py index 28f8ed3c..d3ab8870 100644 --- a/test/test_hwp_diff_emiss.py +++ b/test/test_hwp_diff_emiss.py @@ -2,7 +2,7 @@ from astropy.time import Time import litebird_sim as lbs -from litebird_sim.pointings_in_obs import _get_hwp_angle +from litebird_sim.observation_utilities import _get_hwp_angle def test_add_2f(): diff --git a/test/test_mpi.py b/test/test_mpi.py index 9d1c5b3c..3b4bfc5c 100644 --- a/test/test_mpi.py +++ b/test/test_mpi.py @@ -729,8 +729,8 @@ def test_non_linearity_seeding(): different MPI tasks share the same detector in different time samples """ - if lbs.MPI_COMM_WORLD.size < 2: - return + if lbs.MPI_COMM_WORLD.size != 2: + pytest.skip("This test requires exactly 2 MPI processes") rank = lbs.MPI_COMM_WORLD.rank diff --git a/test/test_observation_utilities.py b/test/test_observation_utilities.py new file mode 100644 index 00000000..08fe7256 --- /dev/null +++ b/test/test_observation_utilities.py @@ -0,0 +1,139 @@ +"""Tests for the shared systematic-effect driver in ``observation_utilities.py``. + +The driver concentrates the bookkeeping that every effect module used to +duplicate (observation normalization, per-detector RNG resolution, TOD +component indirection), so it can be exercised directly here without going +through any particular effect. +""" + +import numpy as np +import pytest +from astropy.time import Time + +import litebird_sim as lbs + + +def _make_two_observations(): + """Build a two-detector simulation with one local observation.""" + dets = [ + lbs.DetectorInfo(name="det_A", sampling_rate_hz=1.0), + lbs.DetectorInfo(name="det_B", sampling_rate_hz=1.0), + ] + sim = lbs.Simulation( + start_time=Time("2025-02-02T00:00:00"), + duration_s=10.0, + random_seed=12345, + ) + sim.create_observations(detectors=dets, split_list_over_processes=False) + return sim, dets + + +def test_normalize_observations_single_is_wrapped(): + sim, _ = _make_two_observations() + obs = sim.observations[0] + + result = lbs.normalize_observations(obs) + + assert result == [obs] + assert result[0] is obs + + +def test_normalize_observations_list_is_copied(): + sim, _ = _make_two_observations() + + result = lbs.normalize_observations(sim.observations) + + # Same elements, but a fresh list the caller may mutate safely. + assert result == sim.observations + assert result is not sim.observations + + +def test_normalize_observations_rejects_other_types(): + with pytest.raises(TypeError): + lbs.normalize_observations("not an observation") + + +def test_for_each_observation_without_rng_yields_tod_and_none(): + sim, _ = _make_two_observations() + + yielded = list(lbs.for_each_observation(sim.observations)) + + assert len(yielded) == len(sim.observations) + for (cur_obs, tod, dets_random), expected_obs in zip(yielded, sim.observations): + assert cur_obs is expected_obs + # The driver hands back the actual TOD array, not a copy. + assert tod is expected_obs.tod + assert dets_random is None + + +def test_for_each_observation_respects_component(): + sim, _ = _make_two_observations() + for obs in sim.observations: + obs.custom_tod = np.zeros_like(obs.tod) + + for cur_obs, tod, _ in lbs.for_each_observation( + sim.observations, component="custom_tod" + ): + assert tod is cur_obs.custom_tod + + +def test_for_each_observation_rejects_unknown_component(): + sim, _ = _make_two_observations() + with pytest.raises(ValueError, match="registered TOD components"): + list(lbs.for_each_observation(sim.observations, component="toed")) + + +def test_for_each_observation_resolves_dets_random_once(): + sim, dets = _make_two_observations() + hierarchy = lbs.RNGHierarchy(12345, num_detectors_per_rank=len(dets)) + dets_random = hierarchy.get_detector_level_generators_on_rank(0) + + seen = [ + rng + for _, _, rng in lbs.for_each_observation( + sim.observations, dets_random=dets_random, requires_rng=True + ) + ] + + # Same generator list is handed out on every iteration. + assert all(rng is dets_random for rng in seen) + assert len(dets_random) >= sim.observations[0].n_detectors + + +def test_for_each_observation_requires_seed_or_generators_when_stochastic(): + sim, _ = _make_two_observations() + + with pytest.raises(ValueError): + list(lbs.for_each_observation(sim.observations, requires_rng=True)) + + +def test_for_each_observation_with_pointings_pairs_arrays(): + sim, _ = _make_two_observations() + obs = sim.observations[0] + n_det, n_samp = obs.tod.shape + pointings = np.zeros((n_det, n_samp, 3)) + + yielded = list(lbs.for_each_observation_with_pointings(obs, pointings)) + + assert len(yielded) == 1 + cur_obs, tod, cur_ptg = yielded[0] + assert cur_obs is obs + assert tod is obs.tod + assert cur_ptg is pointings + + +def test_for_each_observation_with_pointings_pairs_lists(): + sim, _ = _make_two_observations() + obs_list = sim.observations + pointings_list = [ + np.zeros((obs.tod.shape[0], obs.tod.shape[1], 3)) for obs in obs_list + ] + + yielded = list(lbs.for_each_observation_with_pointings(obs_list, pointings_list)) + + assert len(yielded) == len(obs_list) + for (cur_obs, _, cur_ptg), exp_obs, exp_ptg in zip( + yielded, obs_list, pointings_list + ): + assert cur_obs is exp_obs + assert cur_ptg is exp_ptg diff --git a/test/test_scan_map.py b/test/test_scan_map.py index 9977b024..526e23f8 100644 --- a/test/test_scan_map.py +++ b/test/test_scan_map.py @@ -433,7 +433,7 @@ def test_scan_map_algebras(): pointings, _ = obs.get_pointings(idet) pixind = hpx.ang2pix(pointings[:, 0:2]) # Note: calling internal method _get_hwp_angle - hwp_angle = lbs.pointings_in_obs._get_hwp_angle(obs, hwp_new) + hwp_angle = lbs.observation_utilities._get_hwp_angle(obs, hwp_new) angle = 2 * pointings[:, 2] - 2 * obs.pol_angle_rad[idet] + 4 * hwp_angle tod[idet, :] = ( maps[0, pixind] @@ -462,7 +462,7 @@ def test_scan_map_algebras(): for idet in range(obs.n_detectors): pointings = obs.pointing_matrix[idet] pixind = hpx.ang2pix(pointings[:, 0:2]) - hwp_angle = lbs.pointings_in_obs._get_hwp_angle(obs, hwp_new) + hwp_angle = lbs.observation_utilities._get_hwp_angle(obs, hwp_new) angle = 2 * pointings[:, 2] - 2 * obs.pol_angle_rad[idet] + 4 * hwp_angle tod[idet, :] = ( maps[0, pixind]