Repository: hpc4cmb/toast (language: python)

Issue:
Implementing Cosmic ray operator
This operator aims at injecting cosmic ray glitches to TOD. This class can be of particular interest for space and baloon observations. 
For time being the operator relies on the user to provide the data encoding the glitch distribution as well as the amplitude of noise due to the low amplitude hits in the focalplane and the correlation among detectors. 

Several things need to be addressed before merging : 
- [x]  Random number generators, so far am using the toast random number generator for gaussian numbers, however in https://github.com/hpc4cmb/toast/blob/1ec1c2e925cbf36dbceee9a87c4576f06936ab34/src/toast/ops/sim_cosmic_rays.py#L231, https://github.com/hpc4cmb/toast/blob/1ec1c2e925cbf36dbceee9a87c4576f06936ab34/src/toast/ops/sim_cosmic_rays.py#L128  and in https://github.com/hpc4cmb/toast/blob/1ec1c2e925cbf36dbceee9a87c4576f06936ab34/src/toast/ops/sim_cosmic_rays.py#L236  using the `numpy` random numbers ,
- [ ] a traits  that can be very useful is the one encoding the conversion factor to convert cosmic ray tod data (which usually are in power units), to temperature units . 
- [x] to properly design a test for this operator can include a routine that generate on the fly a mock cosmic ray datafile   (see 05cd036)


Discussion / hints:
We should rebase this against the `toast3` branch to resolve the conflicts.

Fix patch:
diff --git a/src/toast/ops/CMakeLists.txt b/src/toast/ops/CMakeLists.txt
index 1a823ea37..328740ae9 100644
--- a/src/toast/ops/CMakeLists.txt
+++ b/src/toast/ops/CMakeLists.txt
@@ -41,5 +41,6 @@ install(FILES
     polyfilter.py
     groundfilter.py
     statistics.py
+    sim_cosmic_rays.py
     DESTINATION ${PYTHON_SITE}/toast/ops
 )
diff --git a/src/toast/ops/__init__.py b/src/toast/ops/__init__.py
index db3bcb911..c41c48549 100644
--- a/src/toast/ops/__init__.py
+++ b/src/toast/ops/__init__.py
@@ -66,6 +66,8 @@
 
 from .sim_gaindrifts import GainDrifter
 
+from .sim_cosmic_rays import InjectCosmicRays
+
 from .totalconvolve import SimTotalconvolve
 
 from .polyfilter import PolyFilter, PolyFilter2D, CommonModeFilter
diff --git a/src/toast/ops/sim_cosmic_rays.py b/src/toast/ops/sim_cosmic_rays.py
new file mode 100644
index 000000000..d9726f20c
--- /dev/null
+++ b/src/toast/ops/sim_cosmic_rays.py
@@ -0,0 +1,296 @@
+# Copyright (c) 2019-2021 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import traitlets
+
+import numpy as np
+from scipy import interpolate, signal
+
+from ..timing import function_timer
+
+from ..traits import trait_docs, Int, Float, Unicode, Bool, Quantity, Callable
+from astropy import units as u
+
+from .operator import Operator
+
+from ..utils import Environment, Logger
+from .. import rng
+
+
+class InjectCosmicRays(Operator):
+    """
+    Inject the cosmic rays  signal into the TOD. So far we inject two kinds of cosmic ray noise:
+
+    Wafer noise, due to ~400 impacts per second in the wafer   undistinguishable  individually.
+    For each observation and for each detector we inject low noise component as a _white noise signal_,
+    i.e. noraml distributed  random samples following the observed properties from simulations and read from disc.
+    This component is then coadded to the sky signal (as if it were a noise term) .
+
+    Common mode noise
+
+    A common mode  noise within each detector pair can be simulated given the properties of the wafer noise.
+    Will use the informations of correlations can be found in the file provided as an input to the simulations, if
+    present, otherwise 50% detector correlation is assumed.
+
+    Direct  hits (or Glitches)
+
+    Given the size of the detector we can derive the cosmic ray event rate and simulate the profile of a cosmic ray glitch.
+    We assume the glitch to be described as
+
+    .. math::
+    \gamma (t) = C_1 +C_2 e^{-t/\tau }
+
+    where :math:C_1 and :math:C_2 and the time constant :math:\tau  are drawn from a distribution of estimated values
+    from simulations.   For each observation and each detector, we estimate the number of hits expected
+    theroretically and draw a random integer, `N`,  with a Poissonian distribution given the expected number
+    of events, `Nexp`.  We then  select randomly `N` timestamps where   the hits will be injected into the tod simulated in TOAST.
+    Evaluate the function :math:\gamma at a higher sampling rate (~150 Hz), decimate it to the TOD sample rate and coadd  it.
+
+    Args:
+        crfile (string):  A `*.npz`  file encoding 4 attributes,
+            `low_noise` (mean and std. dev. of  the wafer noise)
+            `sampling_rate` sampling rate of the glitch simulations
+            `direct_hits` distribution of the glitch parameters
+            `correlation_matrix` correlation matrix for common mode
+            must have a tag {detector} that will be replaced with the detector index.
+        signal_name (string): the cache reference of the TOD data where the cosmic ray will be stored
+        realization (int): to run several    Monte-Carlo realizations of cosmic ray noise
+        eventrate (float) : the expected event rate of hits in a detector
+        inject_direct_hits (bool): will include also direct hits if set to True
+        conversion_factor (float): factor to convert the cosmic ray units to temperature units
+        common_mode (bool) :  will include also common mode per pixel pair  if set to True
+    """
+
+    API = Int(0, help="Internal interface version for this operator")
+
+    det_data = Unicode(
+        "signal", help="Observation detdata key to inject the gain drift"
+    )
+    crfile = Unicode(None, help="Path to the *.npz file encoding cosmic ray infos")
+    crdata_units = Quantity(1 * u.W, help="set the unities of the input amplitudes ")
+
+    realization = Int(0, help="integer to set a different random seed ")
+
+    eventrate = Float(
+        0.0015,
+        help="the expected event rate of hits in a detector",
+    )
+    inject_direct_hits = Bool(False, help="inject  direct hits as glitches in the TODs")
+    conversion_factor = Quantity(
+        1 * u.K / u.W,
+        help="factor to convert the cosmic ray signal (usually Watts) into temperature units",
+    )
+    include_common_mode = Bool(
+        False, help="will include also common mode per pixel pair  if set to True"
+    )
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def load_cosmic_ray_data(self, filename):
+        data_dic = np.load(filename)
+
+        return data_dic
+
+    def resample_cosmic_ray_statistics(self, arr, Nresamples, key, counter):
+        resampled = np.zeros((Nresamples, arr.shape[1]))
+
+        for ii in range(arr.shape[1]):
+            # Resample by considering the bulk of  the Parameter distribution ~2sigma central interval
+            bins = np.linspace(
+                np.quantile(
+                    arr[:, ii],
+                    0.025,
+                ),
+                np.quantile(arr[:, ii], 0.975),
+                30,
+            )
+
+            binned, edges = np.histogram(arr[:, ii], bins=bins)
+
+            xb = 0.5 * (edges[:-1] + edges[1:])
+            CDF = np.cumsum(binned) / binned.sum()
+
+            pinv = interpolate.interp1d(CDF, xb, fill_value="extrapolate")
+            # r = np.random.rand(Nresamples)
+            r = rng.random(
+                Nresamples,
+                sampler="uniform_01",
+                key=key,
+                counter=counter,
+            )
+
+            resampled[:, ii] = pinv(r)
+
+        return resampled
+
+    @function_timer
+    def _exec(self, data, detectors=None, **kwargs):
+        env = Environment.get()
+        log = Logger.get()
+        if self.crfile is None:
+            raise AttributeError(
+                "OpInjectCosmicRays cannot run if you don't provide cosmic ray data."
+            )
+        for ob in data.obs:
+            # Get the detectors we are using for this observation
+            dets = ob.select_local_detectors(detectors)
+            if len(dets) == 0:
+                # Nothing to do for this observation
+                continue
+            comm = ob.comm
+            rank = ob.comm_rank
+            # Make sure detector data output exists
+            ob.detdata.ensure(self.det_data, detectors=dets)
+            obsindx = ob.uid
+            telescope = ob.telescope.uid
+            focalplane = ob.telescope.focalplane
+            size = ob.detdata[self.det_data][dets[0]].size
+            samplerate = focalplane.sample_rate.to_value(u.Hz)
+
+            obstime_seconds = size / samplerate
+            n_events_expected = self.eventrate * obstime_seconds
+            key1 = self.realization * 4294967296 + telescope * 65536
+            counter2 = 0
+
+            for kk, det in enumerate(dets):
+                detindx = focalplane[det]["uid"]
+                key2 = obsindx
+                counter1 = detindx
+
+                rngdata = rng.random(
+                    size,
+                    sampler="gaussian",
+                    key=(key1, key2),
+                    counter=(counter1, counter2),
+                )
+                counter2 += size
+                filename = self.crfile.replace("detector", f"det{kk}")
+                data_dic = self.load_cosmic_ray_data(filename)
+                lownoise_params = data_dic["low_noise"]
+                var_tot = lownoise_params[1] ** 2
+                if not self.include_common_mode:
+                    lownoise_hits = lownoise_params[1] * rngdata + lownoise_params[0]
+                    tmparray = lownoise_hits
+                else:
+                    if kk % 2 != 0:  # if kk is odd
+                        detid_common = kk - 1
+                        kkcol = kk - 1
+                    else:  # kk even
+                        detid_common = kk
+                        kkcol = kk + 1
+
+                    filename_common = self.crfile.replace(
+                        "detector", f"det{detid_common}"
+                    )
+                    data_common = self.load_cosmic_ray_data(filename_common)
+                    try:
+                        corr_matr = data_common["correlation_matrix"]
+                        corr_frac = corr_matr[kk, kkcol]
+
+                    except KeyError:
+                        log.warning(
+                            "Correlation matrix not provided for common mode, assuming 50% correlation "
+                        )
+                        corr_frac = 0.5
+
+                    var_corr = corr_frac * data_common["low_noise"][1] ** 2
+                    var0 = var_tot - var_corr
+
+                    rngdata_common = rng.random(
+                        size,
+                        sampler="gaussian",
+                        key=(key1, key2),
+                        counter=(detid_common, counter2),
+                    )
+                    counter2 += size
+                    cr_common_mode = (
+                        np.sqrt(var_corr) * rngdata_common + data_common["low_noise"][0]
+                    )
+
+                    lownoise_hits = lownoise_params[1] * rngdata + lownoise_params[0]
+                    tmparray = lownoise_hits + cr_common_mode
+
+                if self.inject_direct_hits:
+                    glitches_param_distr = data_dic["direct_hits"]
+                    fsampl_sims = (data_dic["sampling_rate"][0] * u.Hz).value
+
+                    glitch_seconds = 0.15  # seconds, i.e. ~ 3samples at 19Hz
+                    # we approximate the number of samples to the closest integer
+                    nsamples_high = np.int_(np.around(glitch_seconds * fsampl_sims))
+                    nsamples_low = np.int_(np.around(glitch_seconds * samplerate))
+                    # import pdb; pdb.set_trace()
+                    # np.random.seed( obsindx//1e3  +detindx//1e3 )
+                    n_events = np.random.poisson(n_events_expected)
+                    params = self.resample_cosmic_ray_statistics(
+                        glitches_param_distr,
+                        Nresamples=n_events,
+                        key=(key1, key2),
+                        counter=(counter1, counter2),
+                    )
+                    counter2 += n_events
+                    # draw n_events uniformly from a continuous distribution
+                    # you want the events to happen during one observation
+                    # we also make sure that the glitch is injected at most
+                    # `glitch_seconds` before the end of the observation ,
+                    # otherwise we've problems in downsampling
+                    rngunif = rng.random(
+                        n_events,
+                        sampler="uniform_01",
+                        key=(key1, key2),
+                        counter=(counter1, counter2),
+                    )
+                    counter2 += n_events
+
+                    time_glitches = (obstime_seconds - glitch_seconds) * rngunif
+                    assert time_glitches.max() < obstime_seconds
+
+                    # otherwise we've problems in downsampling
+
+                    # estimate the timestamps rounding off the events in seconds
+                    time_stamp_glitches = np.int_(np.around(time_glitches * samplerate))
+                    # we measure the glitch and the bestfit timeconstant in millisec
+                    tglitch = np.linspace(0, glitch_seconds * 1e3, nsamples_high)
+                    glitch_func = lambda t, C1, C2, tau: C1 + (C2 * np.exp(-t / tau))
+                    for i in range(n_events):
+                        tmphit = glitch_func(tglitch, *params[i])
+                        tmparray[
+                            time_stamp_glitches[i] : time_stamp_glitches[i]
+                            + nsamples_low
+                        ] = signal.resample(tmphit, num=nsamples_low, t=tglitch)[0]
+                tmparray = tmparray * self.crdata_units
+                ob.detdata[self.det_data][det] += (
+                    self.conversion_factor * tmparray
+                ).value
+
+        return
+
+    def _finalize(self, data, **kwargs):
+        return
+
+    def _requires(self):
+        req = {
+            "meta": list(),
+            "shared": [
+                self.boresight,
+            ],
+            "detdata": list(),
+            "intervals": list(),
+        }
+        if self.view is not None:
+            req["intervals"].append(self.view)
+        return req
+
+    def _provides(self):
+        prov = {
+            "meta": list(),
+            "shared": list(),
+            "detdata": [
+                self.det_data,
+            ],
+        }
+        return prov
+
+    def _accelerators(self):
+        return list()


Test patch:
diff --git a/src/toast/tests/CMakeLists.txt b/src/toast/tests/CMakeLists.txt
index 595a1e6c5..e7d6cff1b 100644
--- a/src/toast/tests/CMakeLists.txt
+++ b/src/toast/tests/CMakeLists.txt
@@ -45,5 +45,6 @@ install(FILES
     ops_sim_gaindrifts.py
     ops_polyfilter.py
     ops_groundfilter.py
+    ops_sim_cosmic_rays.py
     DESTINATION ${PYTHON_SITE}/toast/tests
 )
diff --git a/src/toast/tests/ops_sim_cosmic_rays.py b/src/toast/tests/ops_sim_cosmic_rays.py
new file mode 100644
index 000000000..d78cb451a
--- /dev/null
+++ b/src/toast/tests/ops_sim_cosmic_rays.py
@@ -0,0 +1,128 @@
+# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import os
+
+import numpy as np
+import healpy as hp
+from .mpi import MPITestCase
+
+from ..vis import set_matplotlib_backend
+
+from .. import ops as ops
+from .. import rng
+
+from ..pixels import PixelDistribution, PixelData
+
+from ..pixels_io import write_healpix_fits
+
+from ..covariance import covariance_apply
+from ._helpers import (
+    create_outdir,
+    create_satellite_data,
+    create_satellite_data_big,
+    create_fake_sky,
+)
+
+
+class SimCosmicRayTest(MPITestCase):
+    def setUp(self):
+        fixture_name = os.path.splitext(os.path.basename(__file__))[0]
+        self.outdir = create_outdir(self.comm, fixture_name)
+
+    def make_mock_cosmic_ray_data(self, data, crfile):
+        # so far we hardcode typical values
+        low_noise = [-3e-17, 1e-17]
+        for obs in data.obs:
+            for kk, det in enumerate(obs.local_detectors):
+                direct_hits = np.zeros((400, 3))
+                # import pdb; pdb.set_trace()
+
+                direct_hits[:, 0] = 5e-17 * np.random.randn(400) + 3e-17  # C1, Watts
+                direct_hits[:, 1] = 1e-15 * np.random.randn(400) + 2e-15  # C2, watts
+                direct_hits[:, 2] = 2 * np.random.randn(400) + 4  # time constant, msec
+                filename = crfile.replace("detector", f"det{kk}")
+                np.savez(
+                    filename,
+                    low_noise=low_noise,
+                    direct_hits=direct_hits,
+                    sampling_rate=[156],
+                )
+        return
+
+    def test_cosmic_rays_wafer_noise(self):
+        # Create a fake satellite data set for testing
+        data = create_satellite_data(
+            self.comm,
+        )
+        crfile = f"{self.outdir}/cosmic_ray_glitches_detector.npz"
+        self.make_mock_cosmic_ray_data(data, crfile)
+        if self.comm is not None:
+            self.comm.Barrier()
+        # Simulate noise using this model
+        key = "my_signal"
+        sim_cosmic_rays = ops.InjectCosmicRays(
+            det_data=key,
+            crfile=crfile,
+        )
+        sim_cosmic_rays.apply(data)
+        for obs in data.obs:
+            telescope = obs.telescope.uid
+            focalplane = obs.telescope.focalplane
+            obsindx = obs.uid
+
+            for det in obs.local_detectors:
+                detindx = focalplane[det]["uid"]
+                assert obs.detdata[key][det].sum() != 0.0
+
+    def test_cosmic_rays_glitches(self):
+        # Create a fake satellite data set for testing
+        data = create_satellite_data(
+            self.comm,
+        )
+
+        crfile = f"{self.outdir}/cosmic_ray_glitches_detector.npz"
+        self.make_mock_cosmic_ray_data(data, crfile)
+        if self.comm is not None:
+            self.comm.Barrier()
+        # Simulate noise using this model
+        key = "my_signal"
+        sim_cosmic_rays = ops.InjectCosmicRays(
+            det_data=key,
+            crfile=crfile,
+            inject_direct_hits=True,
+            eventrate=0.1,  # we increase the eventrate artificially for testing
+        )
+        sim_cosmic_rays.apply(data)
+        for obs in data.obs:
+            telescope = obs.telescope.uid
+            focalplane = obs.telescope.focalplane
+            obsindx = obs.uid
+            for det in obs.local_detectors:
+                detindx = focalplane[det]["uid"]
+                assert obs.detdata[key][det].sum() != 0.0
+
+    def test_cosmic_rays_commonmode(self):
+        # Create a fake satellite data set for testing
+        data = create_satellite_data(
+            self.comm,
+        )
+        crfile = f"{self.outdir}/cosmic_ray_glitches_detector.npz"
+        self.make_mock_cosmic_ray_data(data, crfile)
+        if self.comm is not None:
+            self.comm.Barrier()
+        # Simulate noise using this model
+        key = "my_signal"
+        sim_cosmic_rays = ops.InjectCosmicRays(
+            det_data=key, crfile=crfile, include_common_mode=True
+        )
+        sim_cosmic_rays.apply(data)
+        for obs in data.obs:
+            telescope = obs.telescope.uid
+            focalplane = obs.telescope.focalplane
+            obsindx = obs.uid
+
+            for det in obs.local_detectors:
+                detindx = focalplane[det]["uid"]
+                assert obs.detdata[key][det].sum() != 0.0
diff --git a/src/toast/tests/runner.py b/src/toast/tests/runner.py
index 65d85c230..78e938bc7 100644
--- a/src/toast/tests/runner.py
+++ b/src/toast/tests/runner.py
@@ -57,7 +57,7 @@
 from . import ops_sim_gaindrifts as test_ops_sim_gaindrifts
 from . import ops_polyfilter as test_ops_polyfilter
 from . import ops_groundfilter as test_ops_groundfilter
-
+from . import ops_sim_cosmic_rays as test_ops_sim_cosmic_rays
 
 from . import covariance as test_covariance
 
@@ -178,6 +178,7 @@ def test(name=None, verbosity=2):
         suite.addTest(loader.loadTestsFromModule(test_ops_sim_gaindrifts))
         suite.addTest(loader.loadTestsFromModule(test_ops_polyfilter))
         suite.addTest(loader.loadTestsFromModule(test_ops_groundfilter))
+        suite.addTest(loader.loadTestsFromModule(test_ops_sim_cosmic_rays))
 
         suite.addTest(loader.loadTestsFromModule(test_covariance))