Nonlinearity module implementation (#331)

* differential emission from hwp

* add nonlinearity

* example notebook

* fix typo

* documentation

* add docstrings

* Add non-linearity documentation

* Add non-linearity documentation

* Add non-linearity documentation

* Add reference for nonlinearity

* Add test for nonlinearity

* small reworking of the modules

* incorrect use of Typing fixed

* documentation and tests fixed

* CHANGELOG updated

---------

Co-authored-by: Luca Pagano <[email protected]>
Co-authored-by: Maurizio Tomasi <[email protected]>

CHANGELOG.md

@@ -1,6 +1,9 @@
 # HEAD
 
+-   Module for including nonlinearity in the simulations [#331](https://github.com/litebird/litebird_sim/pull/331)
+
 -   Improve the documentation of the binner and the destriper [#333](https://github.com/litebird/litebird_sim/pull/333)
+
 -   Make the code compatible with Python 3.12 [#332](https://github.com/litebird/litebird_sim/pull/332)
 
 # Version 0.13.0

docs/source/index.rst

@@ -27,6 +27,7 @@ Welcome to litebird_sim's documentation!
    mpi
    reports
    gaindrifts
+   non_linearity
    external_modules
    profiling
    integrating

docs/source/non_linearity.rst

@@ -0,0 +1,173 @@
+Non-linearity injection
+=======================
+
+Non-linearity is the effect of a non-ideal TES detectors' response. This means that the responsivity :math:`S` is not constant as the optical power varies. 
+The LiteBIRD Simulation Framework provides a non-linearity simulation module to simulate the effect of non-linearity on TODs.
+
+The framework provides the simplest case, which is a quadratic non-linearity. 
+This case is described in `Micheli+2024 <https://arxiv.org/pdf/2407.15294>`_, where the effect of non-linearity is propagated to the estimation of the tensor-to-scalar ratio. 
+
+Considering a first order correction of the usual linear gain, a TOD :math:`d(t)` is modified according to:
+
+.. math::
+    d(t) = [1+g_1 d(t)] d(t)
+
+where :math:`g_1` is the detector non-linearity factor in units of :math:`K^{-1}`.
+
+To simulate a quadratic non-linearity, one can use the method of :class:`.Simulation` class :meth:`.Simulation.apply_quadratic_nonlin()`, 
+or any of the low level functions: :func:`.apply_quadratic_nonlin_to_observations()`, :func:`.apply_quadratic_nonlin_for_one_detector()`. 
+The following example shows the typical usage of the method and low level functions:
+
+
+.. code-block:: python
+
+    import numpy as np
+    import litebird_sim as lbs
+    from astropy.time import Time
+
+    start_time = Time("2025-02-02T00:00:00")
+    mission_time_days = 1
+    sampling_hz = 19
+
+    dets = [
+        lbs.DetectorInfo(
+            name="det_A", sampling_rate_hz=sampling_hz
+        ),
+        lbs.DetectorInfo(
+            name="det_B", sampling_rate_hz=sampling_hz
+        ),
+    ]
+
+    sim = lbs.Simulation(
+        base_path="nonlin_example",
+        start_time=start_time,
+        duration_s=mission_time_days * 24 * 3600.0,
+        random_seed=12345,
+    )
+
+    sim.create_observations(
+        detectors=dets,
+        split_list_over_processes=False,
+    )
+    
+    # Creating fiducial TODs
+    sim.observations[0].nl_2_self = np.ones_like(sim.observations[0].tod)
+    sim.observations[0].nl_2_obs = np.ones_like(sim.observations[0].tod)
+    sim.observations[0].nl_2_det = np.ones_like(sim.observations[0].tod)
+
+If the non-linearity parameter is not read from the IMo, one has to specify :math:`g_1` using the ``g_one_over_k`` argument as in the following example:
+
+.. code-block:: python
+
+    # Define non-linear parameters for the detectors. We choose the same value for both detectors in this example, but it is not necessary.
+    sim.observations[0].g_one_over_k = np.ones(len(dets)) * 1e-3
+    
+    # Applying non-linearity using the `Simulation` class method
+    sim.apply_quadratic_nonlin(component = "nl_2_self",)
+
+    # Applying non-linearity on the given TOD component of an `Observation` object
+    lbs.non_linearity.apply_quadratic_nonlin_to_observations(
+        observations=sim.observations,
+        component="nl_2_obs",
+    )
+
+    # Applying non-linearity on the TOD arrays of the individual detectors.
+    for idx, tod in enumerate(sim.observations[0].nl_2_det):
+        lbs.non_linearity.apply_quadratic_nonlin_for_one_detector(
+            tod_det=tod,
+            g_one_over_k=sim.observations[0].g_one_over_k[idx],
+        )
+
+In particular, the effect of detector non-linearity needs to be included to assess its impact when coupled to other systematic effects. As described in `Micheli+2024 <https://arxiv.org/pdf/2407.15294>`_, a typical case is the coupling with HWP synchronous signal (HWPSS) appearing at twice the rotation frequency of the HWP. 
+This kind of signal can be produced by non-idealities of the HWP, such as its differential transmission and emission.
+
+In that case, the usual TOD :math:`d(t)` will contain an additional term, and can be written as:
+
+.. math::
+    d(t) = d(t) = I + \mathrm{Re}[\epsilon_{\mathrm{pol}}e^{4i\chi}(Q+iU)]+A_2 \cos(2 \omega_{HWP} t) + N
+    
+where :math:`A_2` is the amplitude of the HWPSS and :math:`\omega_{HWP}` is the rotation speed of the HWP. In presence of detector non-linearity, the 2f signal is up-modulated to 4f, affecting the science band.
+
+The framework provides an independent module to introduce this signal in the simulation, adding it to the TODs. To simulate the 2f signal from a rotating, non-ideal  HWP, one can use the method of :class:`.Simulation` class :meth:`.Simulation.add_2f()`, 
+or any of the low level functions: :func:`.add_2f_to_observations()`, :func:`.add_2f_for_one_detector()`. 
+
+If the 2f amplitude is not read from the IMo, one has to specify :math:`A_2` using the ``amplitude_2f_k`` argument. The argument ``optical_power_k`` allows to include the integrated nominal optical power expected for each channel as in the following example:
+
+.. code-block:: python
+
+    import numpy as np
+    import litebird_sim as lbs
+    from astropy.time import Time
+    from litebird_sim.pointings import get_hwp_angle
+
+    
+    telescope = "LFT"
+    channel = "L4-140"
+    detlist = ["000_001_017_QB_140_T", "000_001_017_QB_140_B"]
+    imo_version = "vPTEP"
+    start_time = Time("2025-02-02T00:00:00")
+    mission_time_days = 1
+
+    imo = lbs.Imo(flatfile_location=lbs.PTEP_IMO_LOCATION)
+
+
+    sim = lbs.Simulation(
+        base_path="nonlin_example",
+        start_time=start_time,
+        imo=imo,
+        duration_s=mission_time_days * 24 * 3600.0,
+        random_seed=12345,
+    )
+    
+    # Load the definition of the instrument 
+    sim.set_instrument(
+        lbs.InstrumentInfo.from_imo(
+            imo,
+            f"/releases/{imo_version}/satellite/{telescope}/instrument_info",
+        )
+    )
+
+    dets = [] 
+    for n_det in detlist:
+        det = lbs.DetectorInfo.from_imo(
+            url=f"/releases/{imo_version}/satellite/{telescope}/{channel}/{n_det}/detector_info",
+            imo=imo,)
+        det.sampling_rate_hz = 1
+        dets.append(det)
+
+    sim.create_observations(
+        detectors=dets,
+        split_list_over_processes=False,
+    )
+    
+    sim.set_scanning_strategy(imo_url=f"/releases/{imo_version}/satellite/scanning_parameters/")
+    
+    sim.set_hwp(
+        lbs.IdealHWP(sim.instrument.hwp_rpm * 2 * np.pi / 60,),
+    )
+    
+    sim.prepare_pointings()
+    sim.precompute_pointings()
+
+    # Creating fiducial TODs
+    sim.observations[0].tod_2f = np.zeros_like(sim.observations[0].tod)
+
+    # Define differential emission parameters for the detectors.
+    sim.observations[0].amplitude_2f_k = np.array([0.1, 0.1])
+    sim.observations[0].optical_power_k = np.array([1.0, 1.0])
+
+    # Adding 2f signal from HWP differential emission using the `Simulation` class method
+    sim.add_2f(component="tod_2f")
+
+
+Refer to the :ref:`gd-api-reference` for the full list of non-linearity simulation parameters.
+
+.. _gd-api-reference:
+
+API reference
+-------------
+
+.. automodule:: litebird_sim.non_linearity
+    :members:
+    :undoc-members:
+    :show-inheritance:

litebird_sim/__init__.py

@@ -55,6 +55,7 @@
     IdealHWP,
     read_hwp_from_hdf5,
 )
+from .hwp_diff_emiss import add_2f, add_2f_to_observations
 from .hwp_sys.hwp_sys import (
     HwpSys,
 )
@@ -325,4 +326,7 @@ def destripe_with_toast2(*args, **kwargs):
     "FocalplaneCoord",
     "SpacecraftCoord",
     "PointingSys",
+    # hwp_diff_emiss.py
+    "add_2f",
+    "add_2f_to_observations",
 ]

litebird_sim/hwp_diff_emiss.py

@@ -0,0 +1,118 @@
+# -*- encoding: utf-8 -*-
+
+import numpy as np
+from numba import njit, prange
+
+from typing import Union, List, Optional
+from numbers import Number
+
+from .observations import Observation
+from .hwp import HWP
+from .pointings import get_hwp_angle
+
+
+"""def convert_pW_to_K(power_pW, NET, NEP):
+    temperature_K = power_pW * NET*1e-6/NEP*1e-18/sqrt(2) 
+    return(temperature_K)"""  # we could use this if we prefer having IMo quantities in pW
+
+
+# We calculate the additive signal coming from hwp harmonics.
+# here we calculate 2f directly.
+@njit
+def compute_2f_for_one_sample(angle_rad, amplitude_k, monopole_k):
+    return amplitude_k * np.cos(angle_rad) + monopole_k
+
+
+@njit(parallel=True)
+def add_2f_for_one_detector(tod_det, angle_det_rad, amplitude_k, monopole_k):
+    for i in prange(len(tod_det)):
+        tod_det[i] += compute_2f_for_one_sample(
+            angle_rad=angle_det_rad[i], amplitude_k=amplitude_k, monopole_k=monopole_k
+        )
+
+
+def add_2f(
+    tod,
+    hwp_angle,
+    amplitude_2f_k: float,
+    optical_power_k: float,
+):
+    """Add the HWP differential emission to some time-ordered data
+
+    This functions modifies the values in `tod` by adding the contribution of the HWP
+    synchronous signal coming from differential emission. The `amplitude_2f_k` argument must be
+    a N_dets array containing the amplitude of the HWPSS. The `optical_power_k` argument must have
+    the same size and contain the value of the nominal optical power for the considered frequency channel."""
+
+    assert len(tod.shape) == 2
+    num_of_dets = tod.shape[0]
+
+    if isinstance(amplitude_2f_k, Number):
+        amplitude_2f_k = np.array([amplitude_2f_k] * num_of_dets)
+
+    if isinstance(optical_power_k, Number):
+        optical_power_k = np.array([optical_power_k] * num_of_dets)
+
+    assert len(amplitude_2f_k) == num_of_dets
+    assert len(optical_power_k) == num_of_dets
+
+    for detector_idx in range(tod.shape[0]):
+        add_2f_for_one_detector(
+            tod_det=tod[detector_idx],
+            angle_det_rad=hwp_angle,
+            amplitude_k=amplitude_2f_k[detector_idx],
+            monopole_k=optical_power_k[detector_idx],
+        )
+
+
+def add_2f_to_observations(
+    observations: Union[Observation, List[Observation]],
+    hwp: Optional[HWP] = None,
+    component: str = "tod",
+    amplitude_2f_k: Union[float, None] = None,
+    optical_power_k: Union[float, None] = None,
+):
+    """Add the HWP differential emission to some time-ordered data
+
+    This is a wrapper around the :func:`.add_2f` function that applies to the TOD
+    stored in `observations`, which can either be one :class:`.Observation` instance
+    or a list of observations.
+
+    By default, the TOD is added to ``Observation.tod``. If you want to add it to some
+    other field of the :class:`.Observation` class, use `component`::
+
+    for cur_obs in sim.observations:
+        # Allocate a new TOD for the 2f alone
+        cur_obs.2f_tod = np.zeros_like(cur_obs.tod)
+
+        # Ask `add_2f_to_observations` to store the 2f
+        # in `observations.2f_tod`
+        add_2f_to_observations(sim.observations, component="2f_tod")
+    """
+    if isinstance(observations, Observation):
+        obs_list = [observations]
+    else:
+        obs_list = observations
+
+    # iterate through each observation
+    for cur_obs in obs_list:
+        if amplitude_2f_k is None:
+            amplitude_2f_k = cur_obs.amplitude_2f_k
+
+        if optical_power_k is None:
+            optical_power_k = cur_obs.optical_power_k
+
+        if hwp is None:
+            if hasattr(cur_obs, "hwp_angle"):
+                hwp_angle = cur_obs.hwp_angle
+            else:
+                hwp_angle = None
+        else:
+            hwp_angle = get_hwp_angle(cur_obs, hwp)
+
+        add_2f(
+            tod=getattr(cur_obs, component),
+            hwp_angle=hwp_angle,
+            amplitude_2f_k=amplitude_2f_k,
+            optical_power_k=optical_power_k,
+        )

litebird_sim/noise.py

@@ -235,7 +235,7 @@ def add_noise_to_observations(
         obs_list = observations
 
     # iterate through each observation
-    for i, cur_obs in enumerate(obs_list):
+    for cur_obs in obs_list:
         add_noise(
             tod=getattr(cur_obs, component),
             noise_type=noise_type,