Implementation of h_n maps

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 <https://github.com/simonsobs/smarties/tree/main>`_ 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 <https://arxiv.org/abs/2109.05038>`_), extended to the HWP case in Takase et al (2024) (`arXiv:2408.03040 <https://arxiv.org/abs/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
+

docs/source/part4.rst

@@ -8,3 +8,4 @@ Systematic effects
    non_linearity.rst
    hwp_sys.rst
    pointing_sys.rst
+   h_maps.rst

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,
 )

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

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

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:

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