diff --git a/litebird_sim/observations.py b/litebird_sim/observations.py index 508ba126..f99f6993 100644 --- a/litebird_sim/observations.py +++ b/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,7 @@ def get_pointings( pointing_buffer: npt.NDArray | None = None, hwp_buffer: npt.NDArray | None = None, pointings_dtype=np.float64, + nside_centering: int | None = None, ) -> tuple[npt.NDArray, npt.NDArray | None]: """Compute the pointings for one or more detectors in this observation @@ -1035,6 +1037,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. @@ -1077,6 +1081,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 @@ -1125,6 +1130,14 @@ 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): + 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 diff --git a/litebird_sim/pointings_in_obs.py b/litebird_sim/pointings_in_obs.py index a4142e85..ad2ccbdc 100644 --- a/litebird_sim/pointings_in_obs.py +++ b/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], diff --git a/litebird_sim/simulations.py b/litebird_sim/simulations.py index df2a3cb1..1e0f4c50 100644 --- a/litebird_sim/simulations.py +++ b/litebird_sim/simulations.py @@ -33,15 +33,17 @@ add_convolved_sky_to_observations, ) from .beam_synthesis import generate_gauss_beam_alms +from .constants import NUMBA_NUM_THREADS_ENVVAR from .coordinates import CoordinateSystem -from .detectors import DetectorInfo, FreqChannelInfo, InstrumentInfo, UUID +from .detectors import UUID, DetectorInfo, FreqChannelInfo, InstrumentInfo from .dipole import DipoleType, add_dipole_to_observations from .distribute import distribute_evenly, distribute_optimally -from .gaindrifts import GainDriftType, GainDriftParams, apply_gaindrift_to_observations +from .gaindrifts import GainDriftParams, GainDriftType, apply_gaindrift_to_observations from .healpix import write_healpix_map_to_file from .hwp import HWP from .hwp_diff_emiss import add_2f_to_observations from .imo.imo import Imo +from .input_sky import SkyGenerationParams, SkyGenerator from .io import read_list_of_observations, write_list_of_observations from .mapmaking import ( BinnerResult, @@ -54,24 +56,27 @@ make_destriped_map, save_destriper_results, ) -from .input_sky import SkyGenerator, SkyGenerationParams -from .mpi import MPI_ENABLED, MPI_COMM_WORLD, MPI_COMM_GRID +from .maps_and_harmonics import HealpixMap, SphericalHarmonics +from .mpi import MPI_COMM_GRID, MPI_COMM_WORLD, MPI_ENABLED 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, +) from .profiler import TimeProfiler, profile_list_to_speedscope from .scan_map import scan_map_in_observations from .scanning import ScanningStrategy, SpinningScanningStrategy from .seeding import RNGHierarchy from .spacecraft import SpacecraftOrbit, spacecraft_pos_and_vel -from .maps_and_harmonics import SphericalHarmonics, HealpixMap from .units import Units from .version import ( - __version__ as litebird_sim_version, __author__ as litebird_sim_author, ) -from .constants import NUMBA_NUM_THREADS_ENVVAR +from .version import ( + __version__ as litebird_sim_version, +) DEFAULT_BASE_IMO_URL = "https://litebirdimo.ssdc.asi.it" @@ -1664,7 +1669,8 @@ 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, ) @_profile diff --git a/test/test_simulations.py b/test/test_simulations.py index f50b4cfd..eba033ca 100644 --- a/test/test_simulations.py +++ b/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 @@ -473,7 +473,9 @@ def _configure_simulation_for_pointings( sim.prepare_pointings() if store_full_pointings: - sim.precompute_pointings(pointings_dtype=dtype) + sim.precompute_pointings( + pointings_dtype=dtype, + ) return sim @@ -887,3 +889,30 @@ 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=False, + dtype=dtype, + ) + + for cur_obs in sim.observations: + pointing_matrix, _ = cur_obs.get_pointings( + 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, + )