Variable noise model

src/toast/ops/CMakeLists.txt

@@ -66,6 +66,7 @@ install(FILES
     yield_cut.py
     azimuth_intervals.py
     calibrate.py
+    variable_noise_model.py
     DESTINATION ${PYTHON_SITE}/toast/ops
 )
 

src/toast/ops/__init__.py

@@ -71,5 +71,6 @@
 from .stokes_weights import StokesWeights
 from .time_constant import TimeConstant
 from .totalconvolve import SimTotalconvolve
+from .variable_noise_model import VariableNoiseModel
 from .weather_model import WeatherModel
 from .yield_cut import YieldCut

src/toast/ops/variable_noise_model.py

@@ -0,0 +1,103 @@
+import traitlets
+
+from .. import rng
+from ..noise_sim import AnalyticNoise
+from ..timing import function_timer
+from ..traits import Float, Int, Unicode
+from ..utils import Logger
+from .operator import Operator
+
+
+class VariableNoiseModel(Operator):
+    """Create a noise model that varies from detector to detector."""
+
+    # Class traits
+    API = Int(0, help="Internal interface version for this operator")
+    noise_model = Unicode(
+        "var_noise_model", help="The observation key for storing the noise model"
+    )
+    scatter = Float(0.1, help="Fractional scatter in the noise parameters")
+    realization = Int(0, help="The model realization index")
+
+    @traitlets.validate("realization")
+    def _check_realization(self, proposal):
+        check = proposal["value"]
+        if check < 0:
+            raise traitlets.TraitError("realization index must be positive")
+        return check
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    @function_timer
+    def _exec(self, data, detectors=None, **kwargs):
+        log = Logger.get()
+
+        noise_keys = set(["psd_fmin", "psd_fknee", "psd_alpha", "psd_net"])
+
+        for ob in data.obs:
+            sindx = ob.session.uid
+            telescope = ob.telescope.uid
+            fp_data = ob.telescope.focalplane.detector_data
+            has_parameters = False
+            for key in noise_keys:
+                if key not in fp_data.colnames:
+                    break
+            else:
+                has_parameters = True
+            if not has_parameters:
+                msg = f"Observation {ob.name} does not have a focalplane with "
+                msg += "noise parameters.  Skipping."
+                log.warning(msg)
+                ob[self.noise_model] = None
+                continue
+
+            local_dets = set(ob.local_detectors)
+
+            dets = []
+            fmin = {}
+            fknee = {}
+            alpha = {}
+            NET = {}
+            rates = {}
+            indices = {}
+
+            key1 = (
+                int(self.realization) * int(4294967296)
+                + int(telescope) * int(65536)
+                + int(sindx)
+            )
+
+            for row in fp_data:
+                name = row["name"]
+                detindx = row["uid"]
+                if name not in local_dets:
+                    continue
+                dets.append(name)
+                rates[name] = ob.telescope.focalplane.sample_rate
+                rngdata = rng.random(4, sampler="gaussian", key=(key1, detindx))
+                fmin[name] = row["psd_fmin"] * (1 + self.scatter * rngdata[0])
+                fknee[name] = row["psd_fknee"] * (1 + self.scatter * rngdata[1])
+                alpha[name] = row["psd_alpha"] * (1 + self.scatter * rngdata[2])
+                NET[name] = row["psd_net"] * (1 + self.scatter * rngdata[3])
+                indices[name] = detindx
+
+            ob[self.noise_model] = AnalyticNoise(
+                rate=rates,
+                fmin=fmin,
+                detectors=dets,
+                fknee=fknee,
+                alpha=alpha,
+                NET=NET,
+                indices=indices,
+            )
+
+    def _finalize(self, data, **kwargs):
+        return
+
+    def _requires(self):
+        return dict()
+
+    def _provides(self):
+        prov = {"meta": [self.noise_model]}
+        return prov