Center pointings

CHANGELOG.md

@@ -1,5 +1,7 @@
 # HEAD
 
+-   Allow center pointings in get_pointings, remove _get_centered_pointings function and move its centering logic inside _get_pointings_array [#506](https://github.com/litebird/litebird_sim/pull/506)
+
 -   Add complete HWP Jones formalism, including band integration [#499](https://github.com/litebird/litebird_sim/pull/499)
 
 -   Fix missing parallelizations in multiple places [#521](https://github.com/litebird/litebird_sim/pull/521)

litebird_sim/beam_convolution.py

@@ -1,5 +1,4 @@
 import logging
-import os
 from dataclasses import dataclass
 
 import numpy as np
@@ -13,11 +12,10 @@
 from .pointings_in_obs import (
     _get_hwp_angle,
     _normalize_observations_and_pointings,
-    _get_centered_pointings,
     _get_pointings_array,
 )
 from .maps_and_harmonics import SphericalHarmonics
-from .constants import NUM_THREADS_ENVVAR
+from .utilities import resolve_nthreads
 
 
 @dataclass
@@ -53,7 +51,6 @@ def add_convolved_sky_to_one_detector(
     mueller_matrix,
     hwp_angle,
     convolution_params: BeamConvolutionParameters | None = None,
-    nside_centering: int | None = None,
     nthreads: int = 0,
 ):
     """
@@ -77,9 +74,6 @@ def add_convolved_sky_to_one_detector(
     convolution_params : BeamConvolutionParameters, optional
         Parameters controlling the convolution, such as resolution and precision. If None,
         reasonable defaults are chosen based on the sky and beam properties.
-    nside_centering : int, default=None
-        If set, shifts the detector pointings to the centers of the corresponding HEALPix pixels
-        at the given NSIDE resolution. If None, no centering is applied.
     nthreads : int, default=0
         Number of threads to use for convolution. If set to 0, all available CPU cores
         will be used.
@@ -189,13 +183,7 @@ def add_convolved_sky_to_one_detector(
                     pointings_det.dtype, real_type
                 )
             )
-        tod_det += inter.interpol(
-            pointings_det.astype(real_type, copy=False)
-            if nside_centering is None
-            else _get_centered_pointings(
-                pointings_det.astype(real_type, copy=False), nside_centering, nthreads
-            )
-        )[0]
+        tod_det += inter.interpol(pointings_det.astype(real_type, copy=False))[0]
     else:
         fullconv = MuellerConvolver(
             slm=_slm,
@@ -208,11 +196,7 @@ def add_convolved_sky_to_one_detector(
             nthreads=nthreads,
         )
         tod_det += fullconv.signal(
-            ptg=(
-                pointings_det
-                if nside_centering is None
-                else _get_centered_pointings(pointings_det, nside_centering, nthreads)
-            ),
+            ptg=pointings_det,
             alpha=hwp_angle,
             strict_typing=convolution_params.strict_typing,
         )
@@ -270,8 +254,8 @@ def add_convolved_sky(
         If set, shifts the detector pointings to the centers of the corresponding HEALPix pixels
         at the given NSIDE resolution. If None, no centering is applied.
     nthreads : int, default=0
-        Number of threads to use for convolution. If set to 0, all available CPU cores
-        will be used.
+        Number of threads to use for convolution and in case for HEALPix operations.
+        If set to 0, all available CPU cores will be used.
 
     Raises
     ------
@@ -336,6 +320,7 @@ def add_convolved_sky(
             pointings=pointings,
             hwp_angle=hwp_angle,
             output_coordinate_system=coordinates,
+            nside_centering=nside_centering,
             pointings_dtype=pointings_dtype,
         )
 
@@ -350,7 +335,6 @@ def add_convolved_sky(
             mueller_matrix=mueller_matrix,
             hwp_angle=hwp_angle,
             convolution_params=convolution_params,
-            nside_centering=nside_centering,
             nthreads=nthreads,
         )
 
@@ -478,8 +462,7 @@ def add_convolved_sky_to_observations(
         hwp_angle = _get_hwp_angle(obs=cur_obs, hwp=hwp, pointing_dtype=pointings_dtype)
 
         # Set number of threads
-        if nthreads is None:
-            nthreads = int(os.environ.get(NUM_THREADS_ENVVAR, 0))
+        nthreads = resolve_nthreads(nthreads)
 
         # Perform convolution
         add_convolved_sky(

litebird_sim/grasp2alm.py

@@ -13,7 +13,6 @@
 import typing
 import warnings
 from dataclasses import dataclass
-import os
 
 import ducc0.healpix
 import ducc0.sht
@@ -23,7 +22,7 @@
 
 from .maps_and_harmonics import HealpixMap, SphericalHarmonics
 from .units import Units
-from .constants import NUM_THREADS_ENVVAR
+from .utilities import resolve_nthreads
 
 REASON_DESCRIPTION = {
     1: "Approximate solution found",
@@ -352,8 +351,7 @@ def to_map(
             beam_polar = self
 
         # Build the Stokes maps
-        if nthreads is None:
-            nthreads = int(os.environ.get(NUM_THREADS_ENVVAR, 0))
+        nthreads = resolve_nthreads(nthreads)
         pixel_indexes = base.ang2pix(self.theta_phi_values_rad, nthreads=nthreads)
         beam_map = np.empty((nstokes, npix), dtype=float)
         hit_map = np.empty(npix, dtype=int)

litebird_sim/mapmaking/binner.py

@@ -10,7 +10,6 @@
 from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Any
-import os
 
 import healpy as hp
 import numpy as np
@@ -28,7 +27,7 @@
 )
 from litebird_sim.maps_and_harmonics import HealpixMap
 
-from ..constants import NUM_THREADS_ENVVAR
+from ..utilities import resolve_nthreads
 
 from .common import (
     COND_THRESHOLD,
@@ -342,8 +341,7 @@ def make_binned_map(
     time_mask_list = _build_mask_time_split(time_split, obs_list)
 
     # Set number of threads
-    if nthreads is None:
-        nthreads = int(os.environ.get(NUM_THREADS_ENVVAR, 0))
+    nthreads = resolve_nthreads(nthreads)
 
     nobs_matrix = _build_nobs_matrix(
         nside=nside,

litebird_sim/mapmaking/destriper.py

@@ -5,7 +5,6 @@
 from pathlib import Path
 from typing import Any, cast
 from collections.abc import Callable
-import os
 
 import healpy as hp
 import numpy as np
@@ -22,7 +21,7 @@
     _normalize_observations_and_pointings,
 )
 
-from ..constants import NUM_THREADS_ENVVAR
+from ..utilities import resolve_nthreads
 
 from .common import (
     _compute_pixel_indices,
@@ -462,8 +461,7 @@ def _store_pixel_idx_and_pol_angle_in_obs(
         hwp_angle = _get_hwp_angle(obs=cur_obs, hwp=hwp, pointing_dtype=pointings_dtype)
 
         # Set number of threads
-        if nthreads is None:
-            nthreads = int(os.environ.get(NUM_THREADS_ENVVAR, 0))
+        nthreads = resolve_nthreads(nthreads)
 
         (
             cur_obs.destriper_pixel_idx,
@@ -1596,8 +1594,7 @@ def my_gui_callback(
     hpx = Healpix_Base(nside=nside, scheme="RING")
 
     # Set number of threads
-    if nthreads is None:
-        nthreads = int(os.environ.get(NUM_THREADS_ENVVAR, 0))
+    nthreads = resolve_nthreads(nthreads)
 
     # Convert pointings and ψ angles according to the coordinate system,
     # convert them into Healpix indices and save the result into

litebird_sim/mapmaking/pair_differencing.py

@@ -10,7 +10,6 @@
 from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Any
-import os
 
 import healpy as hp
 import numpy as np
@@ -28,7 +27,7 @@
 )
 from litebird_sim.maps_and_harmonics import HealpixMap
 
-from ..constants import NUM_THREADS_ENVVAR
+from ..utilities import resolve_nthreads
 
 from .common import (
     COND_THRESHOLD,
@@ -434,8 +433,7 @@ def make_pair_differenced_map(
     time_mask_list = _build_mask_time_split(time_split, obs_list)
 
     # Set number of threads
-    if nthreads is None:
-        nthreads = int(os.environ.get(NUM_THREADS_ENVVAR, 0))
+    nthreads = resolve_nthreads(nthreads)
 
     nobs_matrix, rhs = _build_nobs_matrix(
         nside=nside,

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
@@ -18,6 +19,7 @@
 from .pointings import DEFAULT_INTERNAL_BUFFER_SIZE_FOR_POINTINGS_MB, PointingProvider
 from .scanning import RotQuaternion
 from .units import Units
+from .utilities import resolve_nthreads
 
 
 @dataclass
@@ -996,6 +998,8 @@ def get_pointings(
         pointing_buffer: npt.NDArray | None = None,
         hwp_buffer: npt.NDArray | None = None,
         pointings_dtype=np.float64,
+        nside_centering: int | None = None,
+        nthreads: int | None = None,
     ) -> tuple[npt.NDArray, npt.NDArray | None]:
         """Compute the pointings for one or more detectors in this observation
 
@@ -1037,6 +1041,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 setting
+        nside_centering to the nside of an HEALPix map.
         The HWP angle is always a vector with shape ``(N_samples,)``, as it does
         not depend on the list of detectors.
 
@@ -1079,6 +1085,7 @@ def get_pointings(
                 pointing_buffer=pointing_buffer,
                 hwp_buffer=hwp_buffer,
                 pointings_dtype=pointings_dtype,
+                nside_centering=nside_centering,
             )
 
         # Most complex case: an explicit list (or NumPy array) of detectors
@@ -1127,6 +1134,17 @@ def get_pointings(
                 abs_det_idx,
                 pointing_buffer=pointing_buffer[rel_det_idx, :, :],
                 hwp_buffer=hwp_buffer,
+                nside_centering=nside_centering,
+            )
+
+        if isinstance(nside_centering, int):
+            nthreads = resolve_nthreads(nthreads)
+
+            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], nthreads=nthreads),
+                nthreads=nthreads,
             )
 
         return pointing_buffer, hwp_buffer