litebird_sim 0.12.0__py3-none-any.whl

litebird_sim/__init__.py

@@ -0,0 +1,287 @@
+# -*- encoding: utf-8 -*-
+import importlib
+from pathlib import Path
+import numba
+
+from .compress import (
+    rle_compress,
+    rle_decompress,
+)
+from .distribute import distribute_evenly, distribute_optimally
+from .detectors import (
+    DetectorInfo,
+    FreqChannelInfo,
+    InstrumentInfo,
+    detector_list_from_parameters,
+)
+from .bandpasses import BandPassInfo
+from .healpix import (
+    nside_to_npix,
+    npix_to_nside,
+    is_npix_ok,
+    map_type,
+    get_pixel_format,
+    write_healpix_map_to_hdu,
+    write_healpix_map_to_file,
+)
+from .hwp import (
+    HWP,
+    IdealHWP,
+)
+from .hwp_sys.hwp_sys import (
+    HwpSys,
+)
+from .imo import (
+    Imo,
+    FormatSpecification,
+    Entity,
+    Quantity,
+    Release,
+)
+from .madam import save_simulation_for_madam
+from .mbs.mbs import Mbs, MbsParameters, MbsSavedMapInfo
+from .mpi import MPI_COMM_WORLD, MPI_ENABLED, MPI_CONFIGURATION
+from .observations import Observation, TodDescription
+from .pointings import (
+    apply_hwp_to_obs,
+    get_pointing_buffer_shape,
+    get_pointings,
+    get_pointings_for_observations,
+)
+from .quaternions import (
+    quat_rotation_x,
+    quat_rotation_y,
+    quat_rotation_z,
+    quat_right_multiply,
+    quat_left_multiply,
+    rotate_vector,
+    rotate_x_vector,
+    rotate_y_vector,
+    rotate_z_vector,
+    all_rotate_vectors,
+    all_rotate_x_vectors,
+    all_rotate_y_vectors,
+    all_rotate_z_vectors,
+)
+from .scanning import (
+    compute_pointing_and_orientation,
+    all_compute_pointing_and_orientation,
+    spin_to_ecliptic,
+    all_spin_to_ecliptic,
+    calculate_sun_earth_angles_rad,
+    Spin2EclipticQuaternions,
+    ScanningStrategy,
+    SpinningScanningStrategy,
+    get_det2ecl_quaternions,
+    get_ecl2det_quaternions,
+)
+from litebird_sim.mapmaking import (
+    make_binned_map,
+    check_valid_splits,
+    BinnerResult,
+    make_destriped_map,
+    DestriperParameters,
+    DestriperResult,
+    ExternalDestriperParameters,
+)
+from .simulations import (
+    NUMBA_NUM_THREADS_ENVVAR,
+    Simulation,
+    MpiObservationDescr,
+    MpiProcessDescr,
+    MpiDistributionDescr,
+)
+from .noise import (
+    add_white_noise,
+    add_one_over_f_noise,
+    add_noise,
+    add_noise_to_observations,
+)
+from .scan_map import scan_map, scan_map_in_observations
+from .coordinates import (
+    DEFAULT_COORDINATE_SYSTEM,
+    DEFAULT_TIME_SCALE,
+    ECL_TO_GAL_ROT_MATRIX,
+    CoordinateSystem,
+    coord_sys_to_healpix_string,
+)
+
+from .spacecraft import (
+    compute_l2_pos_and_vel,
+    compute_lissajous_pos_and_vel,
+    spacecraft_pos_and_vel,
+    SpacecraftOrbit,
+    SpacecraftPositionAndVelocity,
+)
+from .dipole import add_dipole, add_dipole_to_observations, DipoleType
+from .io import (
+    write_list_of_observations,
+    write_observations,
+    read_list_of_observations,
+    read_observations,
+)
+
+from .gaindrifts import (
+    GainDriftType,
+    SamplingDist,
+    GainDriftParams,
+    apply_gaindrift_for_one_detector,
+    apply_gaindrift_to_tod,
+    apply_gaindrift_to_observations,
+)
+
+from .version import __author__, __version__
+
+
+# Check if the TOAST2 mapmaker is available
+TOAST_ENABLED = importlib.util.find_spec("OpMapMaker", "toast.todmap") is not None
+if not TOAST_ENABLED:
+
+    def destripe_with_toast2(*args, **kwargs):
+        raise ImportError(
+            "Install the toast package using `pip` to use destripe_with_toast2"
+        )
+
+    TOAST_ENABLED = False
+
+
+# Privilege TBB over OpenPM and the internal Numba implementation of a
+# work queue
+numba.config.THREADING_LAYER_CONFIG = ["tbb", "omp", "workqueue"]
+
+PTEP_IMO_LOCATION = Path(__file__).parent.parent / "default_imo"
+
+__all__ = [
+    "__author__",
+    "__version__",
+    "PTEP_IMO_LOCATION",
+    # compress.py
+    "rle_compress",
+    "rle_decompress",
+    # bandpasses.py
+    "BandPassInfo",
+    # healpix.py
+    "nside_to_npix",
+    "npix_to_nside",
+    "is_npix_ok",
+    "map_type",
+    "get_pixel_format",
+    "write_healpix_map_to_hdu",
+    "write_healpix_map_to_file",
+    # distribute.py
+    "distribute_evenly",
+    "distribute_optimally",
+    # imo.py
+    "Imo",
+    # imoobjects
+    "FormatSpecification",
+    "Entity",
+    "Quantity",
+    "Release",
+    # detectors.py
+    "DetectorInfo",
+    "FreqChannelInfo",
+    "InstrumentInfo",
+    "detector_list_from_parameters",
+    # hwp.py
+    "HWP",
+    "IdealHWP",
+    # hwp_sys/hwp_sys.py
+    "HwpSys",
+    # madam.py
+    "save_simulation_for_madam",
+    # mbs.py
+    "Mbs",
+    "MbsParameters",
+    "MbsSavedMapInfo",
+    # mpi.py
+    "MPI_COMM_WORLD",
+    "MPI_ENABLED",
+    "MPI_CONFIGURATION",
+    # observations.py
+    "Observation",
+    "TodDescription",
+    # quaternions.py
+    "quat_rotation_x",
+    "quat_rotation_y",
+    "quat_rotation_z",
+    "quat_right_multiply",
+    "quat_left_multiply",
+    "rotate_vector",
+    "rotate_x_vector",
+    "rotate_y_vector",
+    "rotate_z_vector",
+    "all_rotate_vectors",
+    "all_rotate_x_vectors",
+    "all_rotate_y_vectors",
+    "all_rotate_z_vectors",
+    # scanning.py
+    "compute_pointing_and_orientation",
+    "all_compute_pointing_and_orientation",
+    "spin_to_ecliptic",
+    "all_spin_to_ecliptic",
+    "calculate_sun_earth_angles_rad",
+    "Spin2EclipticQuaternions",
+    "ScanningStrategy",
+    "SpinningScanningStrategy",
+    "get_det2ecl_quaternions",
+    "get_ecl2det_quaternions",
+    # pointings.py
+    "apply_hwp_to_obs",
+    "get_pointing_buffer_shape",
+    "get_pointings",
+    "get_pointings_for_observations",
+    # mapmaking
+    "make_binned_map",
+    "check_valid_splits",
+    "BinnerResult",
+    "make_destriped_map",
+    "DestriperParameters",
+    "DestriperResult",
+    "ExternalDestriperParameters",
+    # toast_destriper.py
+    "TOAST_ENABLED",
+    "destripe_with_toast2",
+    # simulations.py
+    "NUMBA_NUM_THREADS_ENVVAR",
+    "Simulation",
+    "MpiObservationDescr",
+    "MpiProcessDescr",
+    "MpiDistributionDescr",
+    # noise.py
+    "add_white_noise",
+    "add_one_over_f_noise",
+    "add_noise",
+    "add_noise_to_observations",
+    # scan_map.py
+    "scan_map",
+    "scan_map_in_observations",
+    # dipole.py
+    "add_dipole",
+    "add_dipole_to_observations",
+    "DipoleType",
+    # coordinates.py
+    "DEFAULT_COORDINATE_SYSTEM",
+    "DEFAULT_TIME_SCALE",
+    "ECL_TO_GAL_ROT_MATRIX",
+    "CoordinateSystem",
+    "coord_sys_to_healpix_string",
+    # spacecraft.py
+    "compute_l2_pos_and_vel",
+    "compute_lissajous_pos_and_vel",
+    "spacecraft_pos_and_vel",
+    "SpacecraftOrbit",
+    "SpacecraftPositionAndVelocity",
+    # io
+    "write_list_of_observations",
+    "write_observations",
+    "read_list_of_observations",
+    "read_observations",
+    # gaindrifts.py
+    "GainDriftType",
+    "SamplingDist",
+    "GainDriftParams",
+    "apply_gaindrift_for_one_detector",
+    "apply_gaindrift_to_tod",
+    "apply_gaindrift_to_observations",
+]

litebird_sim/bandpasses.py

@@ -0,0 +1,344 @@
+# -*- encoding: utf-8 -*-
+
+from dataclasses import dataclass
+
+import numpy as np
+import scipy as sp
+import logging
+from typing import Optional, Union
+from uuid import UUID
+
+from .imo import Imo
+
+
+def _apod(x, a, b):
+    return (1 + np.cos((x - a) / (b - a) * np.pi)) / 2
+
+
+@dataclass
+class BandPassInfo:
+    """A class wrapping the basic information about a detector bandpass.
+
+    This class encodes the basic properties of a frequency band.
+
+    It can be initialized in three ways:
+
+    - The default constructor will generate top-hat bandpasses,
+      assuming that band centroid and width are provided.
+
+    - Through the class method :meth:`.from_imo`, which reads the
+      detector bandpass from the LiteBIRD Instrument Model (see
+      the :class:`.Imo` class).
+
+    Args:
+
+        - bandcenter_ghz (float):  center frequency in GHz
+
+        - bandwidth_ghz (float):  width of the band (default=0.)
+
+        - bandlow_ghz (float): lowest frequency to sample the band,
+          i.e. the first element of the frequency array considered
+          (default: bandcenter_ghz - bandwidth_ghz)
+
+        - bandhigh_ghz (float): highest frequency to sample the band,
+          i.e. the last element of frequency array considered
+          (default: bandcenter_ghz + bandwidth_ghz)
+
+        - nsamples_inband (int): number of elements to sample the band
+          (default=128)
+
+        - name (str) : ID of the band
+
+        - normalize(bool) : If set to true bandpass weights will be
+          normalized to 1
+
+        - bandtype (str): a string describing the band profile. It can be
+          one of the following:
+
+            - ``top-hat`` (default)
+
+            - ``top-hat-exp``: the edges of the band are apodized with an
+              exponential profile
+
+            - ``top-hat-cosine``: the edges of the band are apodized with
+              a cosine profile
+
+            - ``cheby``: the bandpass encodes a Chebyshev profile
+
+        - alpha_exp (float): out-of-band exponential decay index for
+          low freq edge.
+
+        - beta_exp (float) : out-of-band exponential decay index for
+          high freq edge
+
+        - cosine_apo_length (int): the numerical factor related to
+          the cosine apodization length
+
+        - cheby_poly_order (int): chebyshev filter order.
+
+        - cheby_ripple_dB (int): maximum ripple amplitude in decibels.
+    """
+
+    bandcenter_ghz: float = 0.0
+    bandwidth_ghz: float = 0.0
+    bandlow_ghz: Union[float, None] = None
+    bandhigh_ghz: Union[float, None] = None
+    nsamples_inband: int = 128
+    name: str = ""
+    bandtype: str = "top-hat"
+    normalize: bool = False
+    model: str = None
+
+    alpha_exp: float = 1
+    beta_exp: float = 1
+    cosine_apo_length: float = 5
+    cheby_poly_order: int = 3
+    cheby_ripple_dB: int = 3
+
+    def __post_init__(self):
+        """
+        Constructor of the class
+
+        """
+        if self.bandcenter_ghz <= 0.0:
+            raise ValueError(
+                "The band center must be strictly positive, \
+                assign a value to bandcenter_ghz > 0.0"
+            )
+        self.f0, self.f1 = self.get_edges()
+        # we extend the wings out-of-band of the top-hat bandpass
+        # before and after the edges
+        if not self.bandlow_ghz:
+            self.bandlow_ghz = self.bandcenter_ghz - self.bandwidth_ghz
+        if not self.bandhigh_ghz:
+            self.bandhigh_ghz = self.bandcenter_ghz + self.bandwidth_ghz
+        self.freqs_ghz = np.linspace(
+            self.bandlow_ghz, self.bandhigh_ghz, self.nsamples_inband
+        )
+        # checking that the bandpass edges lie within the freq range
+        if self.bandlow_ghz > self.f0 or self.bandhigh_ghz < self.f1 - 1:
+            raise ValueError("The bandpass is out of the frequency range")
+
+        self.isnormalized = False
+        if self.bandtype == "top-hat":
+            self._get_top_hat_bandpass(normalize=self.normalize)
+        elif self.bandtype == "top-hat-cosine":
+            self._get_top_hat_bandpass(apodization="cosine", normalize=self.normalize)
+        elif self.bandtype == "top-hat-exp":
+            self._get_top_hat_bandpass(apodization="exp", normalize=self.normalize)
+        elif self.bandtype == "cheby":
+            self._get_chebyshev_bandpass(normalize=self.normalize)
+        else:
+            logging.warning(f"{self.bandtype} profile not implemented. Assume top-hat")
+            self._get_top_hat_bandpass(normalize=self.normalize)
+
+    def get_edges(self):
+        """
+        get the edges of the tophat band
+        """
+        return (
+            self.bandcenter_ghz - self.bandwidth_ghz / 2,
+            self.bandcenter_ghz + self.bandwidth_ghz / 2,
+        )
+
+    def _get_top_hat_bandpass(self, normalize=False, apodization: Optional[str] = None):
+        """
+        Sample a top-hat bandpass, given the centroid and the bandwidth.
+        If the `normalize` flag is set to ``True``, the transmission
+        coefficients are
+        normalized so that its integral is 1 over the frequency band.
+        The parameter `apodization` must be either ``"cosine"``, ``"exp"``,
+        or ``None``: in the latter case, no apodization of the beamshape
+        is performed.
+        """
+
+        self.weights = np.zeros_like(self.freqs_ghz)
+        mask = np.ma.masked_inside(self.freqs_ghz, self.f0, self.f1).mask
+        self.weights[mask] = 1.0
+
+        if apodization == "cosine":
+            self._cosine_apodize_bandpass()
+        elif apodization == "exp":
+            self._exp_apodize_bandpass()
+
+        if normalize:
+            self.normalize_band()
+
+    def normalize_band(self):
+        """
+        Normalize the band transmission coefficients
+
+        """
+        A = np.trapz(self.weights, self.freqs_ghz)
+        self.weights /= A
+        self.isnormalized = True
+
+    def _exp_apodize_bandpass(self):
+        """Define a bandpass with exponential tails and unit transmission
+        in band freqs: frequency in GHz
+        """
+        mask_beta = np.ma.masked_greater(self.freqs_ghz, self.f1).mask
+        self.weights[mask_beta] = np.exp(
+            -self.beta_exp * (self.freqs_ghz[mask_beta] - self.f1)
+        )
+        mask_alpha = np.ma.masked_less(self.freqs_ghz, self.f0).mask
+        self.weights[mask_alpha] = np.exp(
+            self.alpha_exp * (self.freqs_ghz[mask_alpha] - self.f0)
+        )
+
+    def _cosine_apodize_bandpass(self):
+        """
+        Define a bandpass with cosine tails and unit transmission in band
+        """
+
+        apolength = self.bandwidth_ghz / self.cosine_apo_length
+        f_above = self.bandcenter_ghz * (1 + self.bandwidth_ghz / 2 + apolength)
+        f_below = self.bandcenter_ghz * (1 - self.bandwidth_ghz / 2 - apolength)
+        mask_above = np.ma.masked_inside(self.freqs_ghz, self.f1, f_above).mask
+
+        x_ab = np.linspace(self.f1, f_above, self.freqs_ghz[mask_above].size)
+
+        self.weights[mask_above] = _apod(x_ab, self.f1, f_above)
+
+        mask_below = np.ma.masked_inside(self.freqs_ghz, f_below, self.f0).mask
+        x_bel = np.linspace(f_below, self.f0, self.freqs_ghz[mask_below].size)
+        self.weights[mask_below] = _apod(x_bel, self.f0, f_below)
+
+    # Chebyshev profile bandpass
+    def _get_chebyshev_bandpass(self, normalize=False):
+        """
+        Define a bandpass with chebyshev prototype.
+        """
+        b, a = sp.signal.cheby1(
+            self.cheby_poly_order,
+            self.cheby_ripple_dB,
+            [2.0 * np.pi * self.f0 * 1e9, 2.0 * np.pi * self.f1 * 1e9],
+            "bandpass",
+            analog=True,
+        )
+        w, h = sp.signal.freqs(b, a, worN=self.freqs_ghz * 2 * np.pi * 1e9)
+
+        self.weights = abs(h)
+        if normalize:
+            A = self.get_normalization()
+            self.weights /= A
+            self.isnormalized = True
+
+    def get_normalization(self):
+        """
+        Estimate the integral over the frequency band
+        """
+        return np.trapz(self.weights, self.freqs_ghz)
+
+    # Find effective central frequency of a bandpass profile
+    def find_central_frequency(self):
+        """Find the effective central frequency of
+        a bandpass profile as defined in https://arxiv.org/abs/1303.5070
+        """
+        if self.isnormalized:
+            return np.trapz(self.freqs_ghz * self.weights, self.freqs_ghz)
+        else:
+            return (
+                np.trapz(self.freqs_ghz * self.weights, self.freqs_ghz)
+                / self.get_normalization()
+            )
+
+    @staticmethod
+    def from_imo(imo: Imo, url: Union[UUID, str]):
+        """Create a `BandPassInfo` object from a definition in the IMO
+        The `url` must either specify a UUID or a full URL to the
+        object.
+
+        """
+        obj = imo.query(url)
+        return BandPassInfo.from_dict(obj.metadata)
+
+    def _interpolate_band(self):
+        """
+        This function aims at building the sampler in order to
+        generate random samples
+        statistically equivalent to the model bandpass
+        """
+        # normalize band
+
+        if not self.isnormalized:
+            self.normalize_band()
+        # Interpolate the band
+        b = sp.interpolate.interp1d(x=self.freqs_ghz, y=self.weights)
+        # estimate the CDF
+        Pnu = np.array(
+            [
+                sp.integrate.quad(b, a=self.freqs_ghz.min(), b=inu)[0]
+                for inu in self.freqs_ghz[1:]
+            ]
+        )
+        # interpolate the inverse CDF
+        self.Sampler = sp.interpolate.interp1d(
+            Pnu,
+            self.freqs_ghz[:-1] + np.diff(self.freqs_ghz),
+            bounds_error=False,
+            fill_value="extrapolate",
+        )
+
+    def bandpass_resampling(self, bstrap_size=1000, nresample=54, model=None):
+        """
+        Resample a  bandpass with bootstrap resampling.
+
+        Note that the user can provide any sampler built with the
+        `interpolate_band` method; otherwise, an error will be raised!
+
+        This function should be used when the user wants to generate
+        many realizations of bandpasses, e.g. per detector bands.
+        There is no need to initialize many
+        instances of the class :class:`.BandPassInfo`, just rerun this
+        functions multiple times issuing the same bandpass model instance.
+
+        Args :
+
+        - `bstrap_size` (int): encodes the size of the random dataset
+           to be generated from the Sampler
+
+        - `nresample` (int): define how fine is the grid for the
+           resampled bandpass
+
+        - `model` (BandPassInfo.model): We can resample from a model previously
+           constructed with this function. The default value is set to ``None``:
+           in this case, it initializes the bandpass sampler with the model set
+           in the class instance (recommended use).
+        """
+
+        if model is not None:
+            logging.info(f"Bandpass sampler from {model.name }")
+            sampler = model.Sampler
+        else:
+            try:
+                sampler = self.Sampler
+            except AttributeError:
+                print(
+                    "Can't resample if no sampler is built and/or provided, "
+                    "initializing the sampler and interpolating the band"
+                )
+                self._interpolate_band()
+                sampler = self.Sampler
+            except AttributeError:
+                logging.warning(
+                    "Can't resample if no sampler is built and/or provided, "
+                    "initializing the sampler and interpolating the band"
+                )
+                self._interpolate_band()
+                sampler = self.Sampler
+
+        X = np.random.uniform(size=bstrap_size)
+        bins_nu = np.linspace(self.freqs_ghz.min(), self.freqs_ghz.max(), nresample)
+        h, xb = np.histogram(sampler(X), density=True, bins=bins_nu)
+
+        nu_b = xb[:-1] + np.diff(xb)
+        resampled_bpass = abs(
+            sp.interpolate.interp1d(
+                nu_b, h, kind="cubic", bounds_error=False, fill_value=0.0
+            )(self.freqs_ghz)
+        )
+        if self.isnormalized:
+            return resampled_bpass / np.trapz(resampled_bpass, self.freqs_ghz)
+        else:
+            return resampled_bpass

litebird_sim/compress.py

@@ -0,0 +1,41 @@
+# -*- encoding: utf-8 -*-
+
+import numpy as np
+from numba import njit
+
+
+def rle_compress(arr):
+    """Perform a Run-Length encoding of the input array
+
+    Returns a NumPy matrix of shape 2×N, where N is an integer; the first
+    row contains the *lengths* of each run, and the second row contains
+    the *values*."""
+
+    # This code was adapted from https://stackoverflow.com/a/32681075/3967151
+
+    ia = np.asarray(arr)  # force numpy
+    nsamples = len(ia)
+    if nsamples == 0:
+        return None, None
+    else:
+        y = ia[1:] != ia[:-1]
+        indexes = np.append(np.where(y), nsamples - 1)
+        runs = np.diff(np.append(-1, indexes))
+        return np.array([runs, ia[indexes]])
+
+
+@njit
+def _rle_decompress(compressed_arr, result):
+    out_idx = 0
+    for in_idx in range(compressed_arr.shape[1]):
+        length = compressed_arr[0, in_idx]
+        val = compressed_arr[1, in_idx]
+        for i in range(length):
+            result[out_idx] = val
+            out_idx += 1
+
+
+def rle_decompress(compressed_arr):
+    result = np.empty(np.sum(compressed_arr[0, :]))
+    _rle_decompress(compressed_arr, result)
+    return result

litebird_sim/constants.py

@@ -0,0 +1,14 @@
+# -*- encoding: utf-8 -*-
+from astropy.constants import c as c_light
+from astropy.constants import h, k_B
+
+C_LIGHT_KM_S = c_light.value / 1e3
+H_OVER_K_B = h.value / k_B.value
+
+T_CMB_K = 2.72548  # Fixsen 2009 http://arxiv.org/abs/0911.1955
+
+SOLAR_VELOCITY_KM_S = 369.8160
+SOLAR_VELOCITY_GAL_LAT_RAD = 0.842_173_724
+SOLAR_VELOCITY_GAL_LON_RAD = 4.608_035_744_4
+
+EARTH_L2_DISTANCE_KM = 1_496_509.30522