Center pointings

litebird_sim/observations.py

@@ -6,6 +6,7 @@
 import astropy.time
 import numpy as np
 import numpy.typing as npt
+from ducc0.healpix import Healpix_Base
 
 from .coordinates import DEFAULT_TIME_SCALE
 from .detectors import DetectorInfo, InstrumentInfo
@@ -14,7 +15,7 @@
 from .input_sky import SkyGenerationParams
 from .maps_and_harmonics import HealpixMap, SphericalHarmonics
 from .mpi import MPI_COMM_GRID, _SerialMpiCommunicator
-from .pointings import PointingProvider, DEFAULT_INTERNAL_BUFFER_SIZE_FOR_POINTINGS_MB
+from .pointings import DEFAULT_INTERNAL_BUFFER_SIZE_FOR_POINTINGS_MB, PointingProvider
 from .scanning import RotQuaternion
 from .units import Units
 
@@ -994,6 +995,8 @@ def get_pointings(
         pointing_buffer: npt.NDArray | None = None,
         hwp_buffer: npt.NDArray | None = None,
         pointings_dtype=np.float64,
+        center: bool = False,
+        nside_centering: int | None = None,
     ) -> tuple[npt.NDArray, npt.NDArray | None]:
         """Compute the pointings for one or more detectors in this observation
 
@@ -1035,6 +1038,8 @@ def get_pointings(
         and ψ (orientation angle, in radians). *Important*: if you ask for just *one*
         detector passing the index of the detector, the shape of the pointing matrix
         will always be ``(N_samples, 3)``.
+        The pointings can be aligned to the center of the HEALPix pixel they belong to
+        by setting center=True and setting nside_centering to the nside of the map.
         The HWP angle is always a vector with shape ``(N_samples,)``, as it does
         not depend on the list of detectors.
 
@@ -1047,6 +1052,11 @@ def get_pointings(
             "You must initialize pointing_provider; use Simulation.prepare_pointings()"
         )
 
+        if center:
+            assert nside_centering is not None, (
+                "You must set the parameter nside_centering when center=True"
+            )
+
         # Simplest case: we need just one detector
         if isinstance(detector_idx, int):
             assert (detector_idx >= 0) and (detector_idx < self.n_detectors), (
@@ -1077,6 +1087,8 @@ def get_pointings(
                 pointing_buffer=pointing_buffer,
                 hwp_buffer=hwp_buffer,
                 pointings_dtype=pointings_dtype,
+                center=center,
+                nside_centering=nside_centering,
             )
 
         # Most complex case: an explicit list (or NumPy array) of detectors
@@ -1125,6 +1137,15 @@ def get_pointings(
                 abs_det_idx,
                 pointing_buffer=pointing_buffer[rel_det_idx, :, :],
                 hwp_buffer=hwp_buffer,
+                center=center,
+                nside_centering=nside_centering,
+            )
+
+        if center:
+            hpx = Healpix_Base(nside_centering, "RING")
+            # Apply centering on the first two columns (θ, φ)
+            pointing_buffer[:, :, 0:2] = hpx.pix2ang(
+                hpx.ang2pix(pointing_buffer[:, :, 0:2])
             )
 
         return pointing_buffer, hwp_buffer

litebird_sim/pointings_in_obs.py

@@ -1,20 +1,17 @@
 from collections.abc import Callable
 
+import astropy.time
 import numpy as np
 import numpy.typing as npt
-import astropy.time
-
 from deprecated import deprecated
-
 from ducc0.healpix import Healpix_Base
 
+from .coordinates import CoordinateSystem, rotate_coordinates_e2g
 from .detectors import InstrumentInfo
 from .hwp import HWP
 from .observations import Observation
 from .scanning import RotQuaternion
 
-from .coordinates import CoordinateSystem, rotate_coordinates_e2g
-
 
 def prepare_pointings(
     observations: Observation | list[Observation],
@@ -71,6 +68,8 @@ def prepare_pointings(
 def precompute_pointings(
     observations: Observation | list[Observation],
     pointings_dtype=np.float64,
+    center: bool = False,
+    nside_centering: int | None = None,
 ) -> None:
     """Precompute pointing angles and HWP angles for a set of observations.
 
@@ -90,6 +89,15 @@ def precompute_pointings(
             Data type to use when allocating the pointing and HWP arrays.
             Defaults to `np.float64`.
 
+        center : bool
+            Flag to perform alignment of the pointings with the center of the HEALPix
+            pixel they belong to.
+            Defaults to False.
+
+        nside_centering : int
+            HEALPix NSIDE parameter used to determine the pixel centers.
+            Defaults to None.
+
     Returns:
         None
 
@@ -110,6 +118,14 @@ def precompute_pointings(
 
     for cur_obs in obs_list:
         cur_obs.precompute_pointings(pointings_dtype=pointings_dtype)
+        if center:
+            assert isinstance(nside_centering, int), (
+                "nside_centering is not a valid nside value"
+            )
+            for det_idx in range(len(cur_obs.pointing_matrix)):
+                cur_obs.pointing_matrix[det_idx] = _get_centered_pointings(
+                    cur_obs.pointing_matrix[det_idx], nside_centering=nside_centering
+                )
 
 
 @deprecated(

litebird_sim/simulations.py

@@ -59,7 +59,11 @@
 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 precompute_pointings, prepare_pointings
+from .pointings_in_obs import (
+    precompute_pointings,
+    prepare_pointings,
+    _get_centered_pointings,
+)
 from .profiler import TimeProfiler, profile_list_to_speedscope
 from .scan_map import scan_map_in_observations
 from .scanning import ScanningStrategy, SpinningScanningStrategy
@@ -1652,7 +1656,9 @@ def prepare_pointings(
                 memory_occupation=int(memory_occupation),
             )
 
-    def precompute_pointings(self, pointings_dtype=np.float64) -> None:
+    def precompute_pointings(
+        self, pointings_dtype=np.float64, center=False, nside_centering=None
+    ) -> None:
         """Compute all the pointings for all observations and save them
 
         Save the pointing matrix of each :class:`.Observation` object in this simulation
@@ -1664,9 +1670,24 @@ def precompute_pointings(self, pointings_dtype=np.float64) -> None:
         execution if you plan to access the pointings repeatedly during a simulation.
         """
         precompute_pointings(
-            observations=self.observations, pointings_dtype=pointings_dtype
+            observations=self.observations,
+            pointings_dtype=pointings_dtype,
+            center=center,
+            nside_centering=nside_centering,
         )
 
+    def center_pointings(self, nside) -> None:
+        """Force the pointings to the center of pixels for a given nside.
+        Changes the values in the field ``pointing matrix``
+        """
+        for obs_idx, cur_obs in enumerate(self.observations):
+            for det_idx in range(len(cur_obs.pointing_matrix)):
+                self.observations[obs_idx].pointing_matrix[det_idx] = (
+                    _get_centered_pointings(
+                        cur_obs.pointing_matrix[det_idx], nside_centering=nside
+                    )
+                )
+
     @_profile
     def compute_pos_and_vel(
         self,

test/test_simulations.py

@@ -3,15 +3,15 @@
 import os
 import pathlib
 from pathlib import Path
-from tempfile import TemporaryDirectory, NamedTemporaryFile
+from tempfile import NamedTemporaryFile, TemporaryDirectory
+from unittest.mock import MagicMock, patch
 from uuid import UUID
 
 import astropy
+import litebird_sim as lbs
 import numpy as np
 import pytest
-from unittest.mock import MagicMock, patch
-
-import litebird_sim as lbs
+from ducc0.healpix import Healpix_Base
 from litebird_sim.detectors import DetectorInfo, FreqChannelInfo
 
 
@@ -21,7 +21,7 @@ def savefig(*args, **kwargs):
 
 
 def test_healpix_map_write(tmp_path):
-    sim = lbs.Simulation(base_path=tmp_path / "simulation_dir", random_seed=12345)
+    sim = lbs.Simulation(base_path="./simulation_dir", random_seed=12345)
     output_file = sim.write_healpix_map(filename="test.fits.gz", pixels=np.zeros(12))
 
     assert isinstance(output_file, pathlib.Path)
@@ -40,7 +40,7 @@ def test_healpix_map_write(tmp_path):
 
 def test_markdown_report(tmp_path):
     sim = lbs.Simulation(
-        base_path=tmp_path / "simulation_dir",
+        base_path="./simulation_dir",
         name="My simulation",
         description="Lorem ipsum",
         start_time=1.0,
@@ -98,7 +98,7 @@ def test_imo_in_report(tmp_path):
     imo = lbs.Imo(flatfile_location=curpath / "test_imo")
 
     sim = lbs.Simulation(
-        base_path=tmp_path / "simulation_dir",
+        base_path="./simulation_dir",
         name="My simulation",
         description="Lorem ipsum",
         imo=imo,
@@ -414,6 +414,8 @@ def _configure_simulation_for_pointings(
     tmp_path: Path,
     include_hwp: bool,
     store_full_pointings: bool,
+    center_pointings: bool = False,
+    nside_centering: int | None = None,
     num_of_detectors: int = 1,
     dtype=np.float32,
 ) -> lbs.Simulation:
@@ -428,7 +430,7 @@ def _configure_simulation_for_pointings(
     )
 
     sim = lbs.Simulation(
-        base_path=tmp_path / "simulation_dir",
+        base_path="./simulation_dir",
         start_time=0.0,
         duration_s=61.0,
         random_seed=12345,
@@ -473,7 +475,11 @@ def _configure_simulation_for_pointings(
     sim.prepare_pointings()
 
     if store_full_pointings:
-        sim.precompute_pointings(pointings_dtype=dtype)
+        sim.precompute_pointings(
+            pointings_dtype=dtype,
+            center=center_pointings,
+            nside_centering=nside_centering,
+        )
 
     return sim
 
@@ -887,3 +893,64 @@ def det_side_effect(imo, url):
         # Asking for 5 detectors when only 2 exist in the channel
         with pytest.raises(ValueError, match="Expected 5 detectors, but got 2"):
             sim.set_detectors(channels=["url_ch_a"], detectors=5)
+
+
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def test_center_pointings(tmp_path, dtype):
+    sim = _configure_simulation_for_pointings(
+        tmp_path,
+        include_hwp=False,
+        store_full_pointings=True,
+        center_pointings=True,
+        nside_centering=16,
+        dtype=dtype,
+    )
+
+    for cur_obs in sim.observations:
+        assert "pointing_matrix" in dir(cur_obs)
+        assert cur_obs.pointing_matrix.dtype == dtype
+        assert cur_obs.pointing_matrix.shape == (
+            cur_obs.n_detectors,
+            cur_obs.n_samples,
+            3,
+        )
+
+        # confirming that the pointings are centered
+        # by confirming they are exactly the same after
+        # doing ang2pix and pix2ang in sequence
+        hpx = Healpix_Base(16, "RING")
+        aux_pointings = cur_obs.pointing_matrix.copy()
+        aux_pointings[:, :, 0:2] = hpx.pix2ang(
+            hpx.ang2pix(cur_obs.pointing_matrix[:, :, 0:2])
+        )
+
+        np.testing.assert_allclose(
+            cur_obs.pointing_matrix,
+            aux_pointings,
+        )
+
+    ###
+
+    sim_wo_precompute = _configure_simulation_for_pointings(
+        tmp_path,
+        include_hwp=False,
+        store_full_pointings=False,
+        dtype=dtype,
+    )
+
+    for cur_obs in sim_wo_precompute.observations:
+        pointing_matrix, _ = cur_obs.get_pointings(
+            center=True, nside_centering=16, pointings_dtype=dtype
+        )
+
+        # confirming that the pointings are centered
+        # by confirming they are exactly the same after
+        # doing ang2pix and pix2ang in sequence
+        hpx = Healpix_Base(16, "RING")
+        aux_pointings = pointing_matrix.copy()
+        aux_pointings[:, :, 0:2] = hpx.pix2ang(hpx.ang2pix(pointing_matrix[:, :, 0:2]))
+
+        np.testing.assert_allclose(
+            pointing_matrix,
+            aux_pointings,
+        )