diff --git a/docs/source/h_maps.rst b/docs/source/h_maps.rst new file mode 100644 index 00000000..a6808ab0 --- /dev/null +++ b/docs/source/h_maps.rst @@ -0,0 +1,206 @@ +.. _h-maps: + +:math:`h`-maps +====================== + +.. contents:: Table of Contents + :depth: 2 + :local: + +Overview +-------- + +The ``h_maps`` module provides tools to generate **complex harmonic maps** +:math:`h_{n,m}` from LiteBIRD observations. These maps allow to perform a spin-based mapmaking, they encode the scanning +strategy information. The spin-based mapmaking can be performed with the `SMARTIES `_ package, it has an interface with LBS :math:`h` maps through the functions :func:`read_lbs_h_maps` and :func:`build_mask_from_lbs_h_maps`. + +This implementation follows the formalism introduced in McCallum et al (2021) +(`arXiv:2109.05038 `_), extended to the HWP case in Takase et al (2024) (`arXiv:2408.03040 `_). + +Mathematical formalism +---------------------- +For a scanning strategy with HWP modulation, the harmonic map of spin +order :math:`n,m` at pixel :math:`p` is defined as: + +.. math:: + :label: eq-h-maps + + h_{n,m}(p) = \frac{1}{N_p} \sum_{t \in p} e^{i n \psi_t + i m \chi_t} + +where: + +- :math:`N_p` is the number of observations (hits) in pixel :math:`p` +- :math:`\psi_t` is the parralactic angle of the detector at time sample :math:`t` +- :math:`\chi_t` is the HWP rotation angle at time :math:`t` +- The sum runs over all time samples :math:`t` that fall within pixel :math:`p` + +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` + +Example +----------- + +.. testcode:: + + import litebird_sim as lbs + import numpy as np + from astropy.time import Time + + #load an IMO + imo = lbs.IMO.from_file("/path/to/imo") + # Create a simple simulation + sim = lbs.Simulation( + base_path=:"./output_sim", + start_time=0, + duration_s=24 * 3600.0, + random_seed=12345, + imo=imo, + ) + + # ... (set up detectors, scanning strategy, etc.) ... + + #Create an observation + + sim.create_observations( + detectors=dets, + n_blocks_det=1, + n_blocks_time=1, + tod_dtype=precision, + split_list_over_processes=False, + allocate_tod=False, #We only need the pointing information to compute h maps, so we can save memory by not allocating the full TOD + ) + sim.prepare_pointings() + + result = lbs.make_h_maps( + observations=sim.observations, + nside=64, + n_m_couples=np.array([[0, 0], [2, 0], [4, 0]]), + output_coordinate_system=lbs.CoordinateSystem.Galactic, + save_to_file=True, + output_directory="./h_n_maps", + ) + +The ``n_m_couples`` parameter +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The spin orders to compute are specified as a 2D array of ``(n, m)`` pairs:: + + # Compute h_{0,0}, h_{2,0} and h_{4,0} + n_m_couples = np.array([ + [0, 0], + [2, 0], + [4, 0], + ]) + +Note: for the couple (0,0) the resulting map is not the one expected from the definition of eq :eq:`eq-h-maps`. This definition yield :math:`h_{0,0} = \frac{1}{N_p} \sum_{t \in p} e^{i 0 \psi_t + i 0 \chi_t} = 1` for all observed pixels, which is not very useful. Instead the :math:`h_{0,0}` is an hitmap, i.e. :math:`h_{0,0}(p) = N_p`. + +Output: ``HnMapResult`` +----------------------- + +:func:`make_h_maps` returns a :class:`HnMapResult` 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. +- ``coordinate_system``: the output coordinate system + (a :class:`.CoordinateSystem` object). +- ``detector_split``: the detector split used. +- ``time_split``: the time split used. +- ``duration_s``: the duration of the observation in seconds. +- ``sampling_rate_Hz``: the sampling rate of the observation in Hz. + +Accessing individual maps +~~~~~~~~~~~~~~~~~~~~~~~~~ + +Each entry in ``h_maps`` is a :class:`.h_map_Re_and_Im` object with: + +- ``real``: the real part of :math:`h_{n,m}` (NumPy array of shape ``(Npix,)``) +- ``imag``: the imaginary part (NumPy array of shape ``(Npix,)``) +- ``norm``: property returning :math:`\sqrt{\mathrm{Re}^2 + \mathrm{Im}^2}`, + with unobserved pixels set to :attr:`.UNSEEN_PIXEL_VALUE` +- ``n``, ``m``: the spin orders +- ``det_info``: the detector name + +Example:: + + h_2_0 = result.h_maps["detector_name"][2, 0] + print(h_2_0.real) # real part + print(h_2_0.imag) # imaginary part + print(h_2_0.norm) # amplitude + +Saving and loading maps +----------------------- + +Maps are saved in **HDF5** format, one file per detector: + +.. code-block:: text + + output_directory/ + └── h_maps_det_{detector_name}.h5 + ├── attrs: coordinate_system, det, detector_split, + │ time_split, duration, sampling_rate + ├── 0,0/ + │ ├── Re + │ └── Im + ├── 2,0/ + │ ├── Re + │ └── Im + └── ... + +To reload maps from disk, use :func:`.load_h_map_from_file`:: + + from litebird_sim.mapmaking.h_maps import load_h_map_from_file + + result = load_h_map_from_file("./h_n_maps/h_maps_det_mydetector.h5") + +Detector and time splits +------------------------ + +The ``detector_split`` and ``time_split`` parameters follow the same +conventions as the rest of the litebird_sim mapmaking framework (see +:ref:`mapmaking`). Use ``"full"`` (default) to include all detectors +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. +Example: +.. testcode:: + sim.create_observations( + detectors=dets, + n_blocks_det=2, # The detectors can be split over several MPI tasks + n_blocks_time=1, # But the time samples cannot be split, all are used for each detector + tod_dtype=precision, + split_list_over_processes=False, + allocate_tod=False, + ) + sim.prepare_pointings() + + result = lbs.make_h_maps( + observations=sim.observations, + nside=64, + n_m_couples=np.array([[0, 0], [2, 0], [4, 0]]), + output_coordinate_system=lbs.CoordinateSystem.Galactic, + save_to_file=True, + output_directory="./h_n_maps", + ) + + +API reference +------------- + +.. autoclass:: litebird_sim.mapmaking.h_maps.h_map_Re_and_Im + :members: norm + :undoc-members: + +.. autoclass:: litebird_sim.mapmaking.h_maps.HnMapResult + :undoc-members: + +.. autofunction:: litebird_sim.mapmaking.h_maps.make_h_maps + +.. autofunction:: litebird_sim.mapmaking.h_maps.save_hn_maps + +.. autofunction:: litebird_sim.mapmaking.h_maps.load_h_map_from_file + diff --git a/docs/source/part4.rst b/docs/source/part4.rst index 35f1baa0..58531bd4 100644 --- a/docs/source/part4.rst +++ b/docs/source/part4.rst @@ -8,3 +8,4 @@ Systematic effects non_linearity.rst hwp_sys.rst pointing_sys.rst + h_maps.rst \ No newline at end of file diff --git a/litebird_sim/__init__.py b/litebird_sim/__init__.py index 83a5fcce..2169dd11 100644 --- a/litebird_sim/__init__.py +++ b/litebird_sim/__init__.py @@ -9,6 +9,9 @@ DestriperParameters, DestriperResult, ExternalDestriperParameters, + load_h_map_from_file, + make_h_maps, + HnMapResult, make_pair_differenced_map, PairDifferencingResult, ) diff --git a/litebird_sim/mapmaking/__init__.py b/litebird_sim/mapmaking/__init__.py index d5097557..593dbda4 100644 --- a/litebird_sim/mapmaking/__init__.py +++ b/litebird_sim/mapmaking/__init__.py @@ -1,5 +1,6 @@ from .common import ExternalDestriperParameters from .binner import make_binned_map, check_valid_splits, BinnerResult +from .h_maps import HnMapResult, make_h_maps, load_h_map_from_file from .brahmap_gls import make_brahmap_gls_map from .destriper import ( make_destriped_map, @@ -23,6 +24,10 @@ "BinnerResult", "make_binned_map", "check_valid_splits", + # h_n.py + "HnMapResult", + "make_h_maps", + "load_h_map_from_files", # brahmap_gls "make_brahmap_gls_map", # destriper.py diff --git a/litebird_sim/mapmaking/binner.py b/litebird_sim/mapmaking/binner.py index 6b358648..837e7427 100644 --- a/litebird_sim/mapmaking/binner.py +++ b/litebird_sim/mapmaking/binner.py @@ -11,10 +11,10 @@ from dataclasses import dataclass from typing import Any -import healpy as hp import numpy as np import numpy.typing as npt from ducc0.healpix import Healpix_Base +from litebird_sim.healpix import UNSEEN_PIXEL_VALUE from numba import njit from litebird_sim import mpi @@ -86,8 +86,8 @@ def _solve_binning(nobs_matrix, atd): atd[ipix] = np.linalg.solve(nobs_matrix[ipix], atd[ipix]) nobs_matrix[ipix] = np.linalg.inv(nobs_matrix[ipix]) else: - nobs_matrix[ipix].fill(hp.UNSEEN) - atd[ipix].fill(hp.UNSEEN) + nobs_matrix[ipix].fill(UNSEEN_PIXEL_VALUE) + atd[ipix].fill(UNSEEN_PIXEL_VALUE) @njit diff --git a/litebird_sim/mapmaking/common.py b/litebird_sim/mapmaking/common.py index cc25e208..a0c35b78 100644 --- a/litebird_sim/mapmaking/common.py +++ b/litebird_sim/mapmaking/common.py @@ -131,6 +131,7 @@ def _compute_pixel_indices( hwp_angle: npt.NDArray | None, output_coordinate_system: CoordinateSystem, pointings_dtype=np.float64, + hmap_generation: bool = False, ) -> tuple[npt.NDArray, npt.NDArray]: """Compute the index of each pixel and its attack angle @@ -144,12 +145,15 @@ def _compute_pixel_indices( ``N_d`` the number of detectors and ``N_t`` the number of samples in the TOD, and the last rank represents the θ and φ angles (in radians) expressed in the Ecliptic reference frame. + + The option `hmap_generation` returns only the telescope orientation instead of + the polarization angle. """ assert len(pol_angle_detectors) == num_of_detectors - pixidx_all = np.empty((num_of_detectors, num_of_samples), dtype=int) - polang_all = np.empty((num_of_detectors, num_of_samples), dtype=np.float64) + pixidx_all = np.empty((num_of_detectors, num_of_samples), dtype=np.int32) + polang_all = np.empty((num_of_detectors, num_of_samples), dtype=pointings_dtype) for idet in range(num_of_detectors): curr_pointings_det, hwp_angle = _get_pointings_array( @@ -159,14 +163,60 @@ def _compute_pixel_indices( output_coordinate_system=output_coordinate_system, pointings_dtype=pointings_dtype, ) + if hmap_generation: + polang_all[idet] = curr_pointings_det[:, 2] + else: + polang_all[idet] = _get_pol_angle( + curr_pointings_det=curr_pointings_det, + hwp_angle=hwp_angle, + pol_angle_detectors=pol_angle_detectors[idet], + ) + + pixidx_all[idet] = hpx.ang2pix(curr_pointings_det[:, :2]) - polang_all[idet] = _get_pol_angle( + if output_coordinate_system == CoordinateSystem.Galactic: + # Free curr_pointings_det if the output map is already in Galactic coordinates + try: + del curr_pointings_det + except UnboundLocalError: + pass + + return pixidx_all, polang_all + + +def _compute_pixel_indices_single_detector( + hpx: Healpix_Base, + pointings: npt.NDArray | Callable, + pol_angle_detector: float, + num_of_samples: int, + detector_index: int, + hwp_angle: npt.NDArray | None, + output_coordinate_system: CoordinateSystem, + pointings_dtype=np.float64, + hmap_generation: bool = False, +) -> tuple[npt.NDArray, npt.NDArray]: + """ + Same as _compute_pixel_indices but for a single detector, thus returning only the pixel indices and polarization angles for that detector. + """ + pixidx = np.empty((num_of_samples), dtype=np.int32) + polang = np.empty((num_of_samples), dtype=pointings_dtype) + curr_pointings_det, hwp_angle = _get_pointings_array( + detector_idx=detector_index, + pointings=pointings, + hwp_angle=hwp_angle, + output_coordinate_system=output_coordinate_system, + pointings_dtype=pointings_dtype, + ) + + if hmap_generation: + polang = curr_pointings_det[:, 2] + else: + polang = _get_pol_angle( curr_pointings_det=curr_pointings_det, hwp_angle=hwp_angle, - pol_angle_detectors=pol_angle_detectors[idet], + pol_angle_detectors=pol_angle_detector, ) - - pixidx_all[idet] = hpx.ang2pix(curr_pointings_det[:, :2]) + pixidx = hpx.ang2pix(curr_pointings_det[:, :2]) if output_coordinate_system == CoordinateSystem.Galactic: # Free curr_pointings_det if the output map is already in Galactic coordinates @@ -175,7 +225,7 @@ def _compute_pixel_indices( except UnboundLocalError: pass - return pixidx_all, polang_all + return pixidx, polang def _cholesky_plain(A: npt.NDArray, dest_L: npt.NDArray) -> None: diff --git a/litebird_sim/mapmaking/destriper.py b/litebird_sim/mapmaking/destriper.py index 60587fec..a440d268 100644 --- a/litebird_sim/mapmaking/destriper.py +++ b/litebird_sim/mapmaking/destriper.py @@ -6,10 +6,11 @@ from typing import Any, cast from collections.abc import Callable -import healpy as hp import numpy as np import numpy.typing as npt from ducc0.healpix import Healpix_Base +from litebird_sim.healpix import UNSEEN_PIXEL_VALUE + from numba import njit, prange from litebird_sim.coordinates import CoordinateSystem, coord_sys_to_healpix_string @@ -20,6 +21,8 @@ _get_hwp_angle, _normalize_observations_and_pointings, ) + +from litebird_sim.maps_and_harmonics import HealpixMap from .common import ( _compute_pixel_indices, COND_THRESHOLD, @@ -2142,7 +2145,7 @@ def _load_rank0_destriper_results(file_path: Path) -> DestriperResult: valid_pixel=np.array(inpf["MASK"].data.field("VALID"), dtype=bool), is_cholesky=bool(inpf["NOBSMATR"].header["ISCHOL"]), ) - nside = hp.npix2nside(len(nobs_matrix.valid_pixel)) + nside = HealpixMap.npix_to_nside(len(nobs_matrix.valid_pixel)) if "GAL" in str(inpf["HITMAP"].header["COORDSYS"]).upper(): coord_sys = CoordinateSystem.Galactic diff --git a/litebird_sim/mapmaking/h_maps.py b/litebird_sim/mapmaking/h_maps.py new file mode 100644 index 00000000..3c742f45 --- /dev/null +++ b/litebird_sim/mapmaking/h_maps.py @@ -0,0 +1,408 @@ +# WIP +# ------- +# ------- + +from dataclasses import dataclass + +import numpy as np +import numpy.typing as npt +from numba import njit +import h5py +import os +from collections.abc import Callable +from litebird_sim.observations import Observation +from litebird_sim.coordinates import CoordinateSystem +from litebird_sim.pointings_in_obs import ( + _get_hwp_angle, + _normalize_observations_and_pointings, +) +from litebird_sim.hwp import HWP +from litebird_sim import mpi +from ducc0.healpix import Healpix_Base +from litebird_sim.healpix import UNSEEN_PIXEL_VALUE +from litebird_sim.maps_and_harmonics import HealpixMap +from litebird_sim.detectors import DetectorInfo +from .common import ( + _compute_pixel_indices_single_detector, + _build_mask_detector_split, + _build_mask_time_split, +) +import gc +import logging + +log = logging.getLogger(__name__) + + +@dataclass +class h_map_Re_and_Im: + """Single :math:`h_{n,m}` map component for one detector. + + :ivar real: Real part of :math:`h_{n,m}`. + :type real: npt.NDArray + :ivar imag: Imaginary part of :math:`h_{n,m}`. + :type imag: npt.NDArray + :ivar n: Spin index :math:`n`. + :type n: int + :ivar m: Spin index :math:`m`. + :type m: int + :ivar det_info: Detector name associated with this map. + :type det_info: str + """ + + real: npt.NDArray + imag: npt.NDArray + n: int + m: int + det_info: str + + @property + def norm(self): + """Return :math:`|h_{n,m}|` per pixel. + + :returns: Pixel-wise magnitude computed from ``real`` and ``imag``. + :rtype: npt.NDArray + """ + output = np.full(np.shape(self.real), UNSEEN_PIXEL_VALUE) + mask = self.real != UNSEEN_PIXEL_VALUE + output[mask] = np.sqrt(self.real[mask] ** 2 + self.imag[mask] ** 2) + return output + + +@dataclass +class HnMapResult: + """Result of :func:`make_h_maps`. + + :ivar h_maps: Mapping ``detector -> {(n, m): h_map_Re_and_Im}`` containing + all computed :math:`h_{n,m}` maps. + :type h_maps: dict[str, dict[tuple[int, int], h_map_Re_and_Im]] + :ivar duration_s: Duration of the simulated observation in seconds. + :type duration_s: float + :ivar sampling_rate_Hz: Sampling rate in hertz. + :type sampling_rate_Hz: float + :ivar coordinate_system: Coordinate system of the output maps. + :type coordinate_system: CoordinateSystem + :ivar detector_split: Detector split used to build the maps. + :type detector_split: str + :ivar time_split: Time split used to build the maps. + :type time_split: str + """ + + h_maps: dict[list[str], list[h_map_Re_and_Im]] + duration_s: float + sampling_rate_Hz: float + coordinate_system: CoordinateSystem = CoordinateSystem.Ecliptic + + detector_split: str = "full" + time_split: str = "full" + + +def load_h_map_from_file( + filepath: str, +) -> HnMapResult: + """Load :math:`h_{n,m}` maps from an HDF5 file. + + :param filepath: Path to an HDF5 file produced by :func:`save_hn_maps`. + :type filepath: str + :returns: Loaded maps and metadata. + :rtype: HnMapResult + """ + h_maps = {} + log.info(f"Loading h maps from file:{filepath}") + with h5py.File(filepath, "r") as f: + det_info = f.attrs["det"] + h_maps[det_info] = {} + for hn_map_key in f.keys(): + n, m = map(int, hn_map_key.split(",")) + log.debug(f"Loading h map n,m ={n},{m}") + group = f[hn_map_key] + h_maps[det_info][n, m] = h_map_Re_and_Im( + real=np.array(group["Re"]), + imag=np.array(group["Im"]), + det_info=det_info, + n=n, + m=m, + ) + + return HnMapResult( + h_maps=h_maps, + coordinate_system=CoordinateSystem[f.attrs["coordinate_system"]], + detector_split=f.attrs["detector_split"], + time_split=f.attrs["time_split"], + duration_s=f.attrs["duration_s"], + sampling_rate_Hz=f.attrs["sampling_rate_Hz"], + ) + + +@njit +def _solve_binning(nobs_matrix, atd): + # Solve the map-making equation + # + # This method alters the parameter `nobs_matrix`, so that after its completion + # each 3×3 matrix in nobs_matrix[idx, :, :] will be the *inverse*. + + # Expected shape: + # - `nobs_matrix`: array of shape (Ndet,N_p,3), where + # N_p is the number of pixels in the map and Ndet the number of detectors + # - `atd`: (Ndet,N_p, 2) + npix = atd.shape[0] + for ipix in range(npix): + if nobs_matrix[ipix, 0] != 0: + atd[ipix, 0] = atd[ipix, 0] / nobs_matrix[ipix, 0] + atd[ipix, 1] = atd[ipix, 1] / nobs_matrix[ipix, 0] + else: + atd[ipix] = 0 #If there is no hit we set the pixel to 0 + + + +@njit +def _accumulate_spin_terms_and_build_nobs_matrix( + n: int, + m: int, + pix: npt.ArrayLike, + psi: npt.ArrayLike, + hwp_angle: npt.ArrayLike | None, + t_mask: npt.ArrayLike, + nobs_matrix: npt.ArrayLike, +) -> None: + assert pix.shape == psi.shape + + for cur_pix_idx, cur_psi, cur_hwp_angle in zip(pix, psi, hwp_angle): + info_pix = nobs_matrix[cur_pix_idx] + + cos_n_m = np.cos(n * cur_psi + m * cur_hwp_angle) + sin_n_m = np.sin(n * cur_psi + m * cur_hwp_angle) + if (n, m) == (0, 0): + info_pix[0] = ( + 1 # if n=m=0 we compute the hit count, beceause it is useful to combine h maps later on + ) + else: + info_pix[0] += 1 + info_pix[1] += cos_n_m + info_pix[2] += sin_n_m + + +@njit +def _numba_extract_rhs(nobs_matrix: npt.ArrayLike, rhs: npt.ArrayLike) -> None: + # This is used internally by _extract_map_and_fill_info. + for idx in range(nobs_matrix.shape[0]): + # Extract the vector from the lower left triangle of the 3×3 matrix + # nobs_matrix[idx, :, :] + rhs[idx, 0] = nobs_matrix[idx, 1] + rhs[idx, 1] = nobs_matrix[idx, 2] + + +def _extract_rhs(n_obs: npt.ArrayLike) -> npt.ArrayLike: + # Extract the RHS of the mapmaking equation from the lower triangle of info + # The RHS has a shape (Ndet,Np,2) + rhs = np.empty((n_obs.shape[0], 2), dtype=n_obs.dtype) + + # The implementation in Numba of this code is ~5 times faster than the older + # implementation that used NumPy. + _numba_extract_rhs(n_obs, rhs) + + return rhs + + +def _build_nobs_matrix( + n: int, + m: int, + nside: int, + pixidx: npt.ArrayLike, + polang: npt.ArrayLike, + hwp_angle: npt.ArrayLike | None, + time_mask: npt.ArrayLike, +) -> npt.ArrayLike: + """Build the nobs matrix for all detectors and pixels, it has shape (Npix,3) and contains the accumulated spin terms and hit counts of each pixel for the considered detector.""" + n_pix = HealpixMap.nside_to_npix(nside) + + nobs_matrix = np.zeros((n_pix, 3)) + if hwp_angle is None: + hwp_angle = np.zeros_like(polang) + _accumulate_spin_terms_and_build_nobs_matrix( + n, + m, + pixidx, + polang, + hwp_angle, + time_mask, + nobs_matrix, + ) + + return nobs_matrix + + +def make_h_maps( + observations: Observation | list[Observation], + nside: int, + pointings: np.ndarray | list[np.ndarray] | None = None, + n_m_couples: np.ndarray = np.array(np.meshgrid([0, 2, 4], [0])).T.reshape(-1, 2), + hwp: HWP | None = None, + output_coordinate_system: CoordinateSystem = CoordinateSystem.Galactic, + detector_split: str = "full", + time_split: str = "full", + pointings_dtype=np.float64, + save_to_file: bool = True, + output_directory: str = "./h_n_maps", +) -> HnMapResult: + """Generate complex harmonic maps :math:`h_{n,m}` from observations. + + The map for ``(n, m) = (0, 0)`` contains hit counts per pixel. + + :param observations: Observation object or list of observations used to build + the maps. + :type observations: Observation | list[Observation] + :param nside: HEALPix ``nside`` of the output maps. + :type nside: int + :param pointings: Optional external pointings. Expected shape is + ``(n_detectors, n_samples, 3)`` per observation. + :type pointings: np.ndarray | list[np.ndarray] | None + :param n_m_couples: Array of ``(n, m)`` pairs with shape ``(N, 2)``. + :type n_m_couples: np.ndarray + :param hwp: Half-wave plate model. + :type hwp: HWP | None + :param output_coordinate_system: Coordinate system of the output maps. + :type output_coordinate_system: CoordinateSystem + :param detector_split: Detector split to apply. + :type detector_split: str + :param time_split: Time split to apply. + :type time_split: str + :param pointings_dtype: Data type used when pointings are computed on the fly. + :type pointings_dtype: type + :param save_to_file: If ``True``, save maps to HDF5 files. + :type save_to_file: bool + :param output_directory: Output directory for generated HDF5 files. + :type output_directory: str + :returns: Result container with all computed maps and metadata. + :rtype: HnMapResult + """ + + assert ( + n_m_couples.shape[1] == 2 + ), """the n,m couples should be passed in this shape: array([[n1, m1], + [n2, m2], + [n2, m3]])""" + + h_maps = {} + obs_list, ptg_list = _normalize_observations_and_pointings( + observations=observations, pointings=pointings + ) + all_dets_list = np.concatenate([obs.name for obs in obs_list]) + + detector_mask_list = _build_mask_detector_split(detector_split, obs_list) + + time_mask_list = _build_mask_time_split(time_split, obs_list) + + h_maps = {det: {} for det in all_dets_list} + + for cur_obs, cur_ptg, cur_d_mask, cur_t_mask in zip( + obs_list, ptg_list, detector_mask_list, time_mask_list + ): + assert cur_obs.n_blocks_time == 1, ( + "The current implementation of make_h_maps only supports one time block per observation, you can set several detector blocks though." + ) + for idet in range(cur_obs.n_detectors): + if not cur_d_mask[idet]: + continue + log.info( + f" Computing pixel indices and angles for detector {all_dets_list[idet]}" + ) + + hpx = Healpix_Base(nside, "RING") + hwp_angle = _get_hwp_angle( + obs=cur_obs, hwp=hwp, pointing_dtype=pointings_dtype + ) + # print(np.shape(hwp_angle)) + # print(f"size of hwp array: {hwp_angle.nbytes}") + duration = cur_obs.end_time_global + sampling_rate = cur_obs.sampling_rate_hz + pixidx, polang = _compute_pixel_indices_single_detector( + hpx=hpx, + pointings=cur_ptg, + pol_angle_detector=cur_obs.pol_angle_rad, + num_of_samples=cur_obs.n_samples, + detector_index=idet, + hwp_angle=hwp_angle, + output_coordinate_system=output_coordinate_system, + pointings_dtype=pointings_dtype, + hmap_generation=True, + ) + log.info( + f"Pixel indices and angles for detector {all_dets_list[idet]} computed, now building nobs matrices" + ) + + for i in range(n_m_couples.shape[0]): + n = n_m_couples[i, 0] + m = n_m_couples[i, 1] + + nobs_matrix = _build_nobs_matrix( + n, + m, + nside=nside, + pixidx=pixidx, + polang=polang, + hwp_angle=hwp_angle, + time_mask=cur_t_mask, + ) + rhs = _extract_rhs(nobs_matrix) + _solve_binning(nobs_matrix, rhs) + h_maps[all_dets_list[idet]][n, m] = h_map_Re_and_Im( + real=rhs.T[0].copy(), + imag=rhs.T[1].copy(), + det_info=all_dets_list[idet], + n=n, + m=m, + ) + log.info( + f" h_map n={n} m={m} for detector {all_dets_list[idet]} computed." + ) + del rhs + del nobs_matrix + gc.collect() + del pixidx + del polang + gc.collect() + + result = HnMapResult( + h_maps=h_maps, + coordinate_system=output_coordinate_system, + duration_s=duration, + sampling_rate_Hz=sampling_rate, + detector_split=detector_split, + time_split=time_split, + ) + if save_to_file: + save_hn_maps(result, output_directory) + log.info(f"h_n maps saved to directory: {output_directory}") + return result + + +def save_hn_maps(result, output_directory: str) -> None: + """Save :math:`h_n` maps to HDF5 files. + + One file is produced per detector in ``output_directory``. The directory is + created if it does not exist. + + :param result: Container with maps and metadata to save. + :type result: HnMapResult + :param output_directory: Destination directory for output HDF5 files. + :type output_directory: str + :returns: ``None`` + :rtype: None + """ + + if not os.path.exists(output_directory): + os.makedirs(output_directory) + for det in result.h_maps.keys(): + with h5py.File( + os.path.join(output_directory, f"h_maps_det_{det}.h5"), "w" + ) as f: + f.attrs["coordinate_system"] = result.coordinate_system.name + f.attrs["det"] = str(det) + f.attrs["detector_split"] = result.detector_split + f.attrs["time_split"] = result.time_split + f.attrs["duration_s"] = result.duration_s + f.attrs["sampling_rate_Hz"] = result.sampling_rate_Hz + for hn_map in result.h_maps[det].values(): + grp = f.create_group(f"{hn_map.n},{hn_map.m}") + grp.create_dataset("Re", data=hn_map.real) + grp.create_dataset("Im", data=hn_map.imag) diff --git a/litebird_sim/simulations.py b/litebird_sim/simulations.py index 5ebada2d..f530cee9 100644 --- a/litebird_sim/simulations.py +++ b/litebird_sim/simulations.py @@ -44,6 +44,11 @@ from .imo.imo import Imo from .io import read_list_of_observations, write_list_of_observations from .mapmaking import ( + make_binned_map, + make_h_maps, + HnMapResult, + make_brahmap_gls_map, + check_valid_splits, BinnerResult, DestriperParameters, DestriperResult, @@ -2463,6 +2468,39 @@ def make_binned_map( pointings_dtype=pointings_dtype, ) + @_profile + def make_h_maps( + self, + nside: int, + n_list: list[int] = [0, 2, 4], + hwp: HWP | None = None, + m_list: list[int] = [0], + output_coordinate_system: CoordinateSystem = CoordinateSystem.Galactic, + detector_split: str = "full", + time_split: str = "full", + pointings_dtype=np.float64, + save_to_file: bool = True, + output_directory: str = "./h_n_maps", + ) -> HnMapResult: + """ + Computes the Hn maps from the pointings of `sim.observations`. + This is a wrapper around :func:`litebird_sim.mapmaking.make_h_maps`. + + """ + return make_h_maps( + observations=self.observations, + nside=nside, + n_list=n_list, + hwp=hwp, + m_list=m_list, + output_coordinate_system=output_coordinate_system, + detector_split=detector_split, + time_split=time_split, + pointings_dtype=pointings_dtype, + save_to_file=save_to_file, + output_directory=output_directory, + ) + def _impose_and_check_full_split(self, detector_splits, time_splits): """ Impose the full split if it is not present in the splits.