diff --git a/src/toast/ops/CMakeLists.txt b/src/toast/ops/CMakeLists.txt index 2624f3fcc..2a807aa43 100644 --- a/src/toast/ops/CMakeLists.txt +++ b/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 ) diff --git a/src/toast/ops/__init__.py b/src/toast/ops/__init__.py index b0021d3a1..7cd2381eb 100644 --- a/src/toast/ops/__init__.py +++ b/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 diff --git a/src/toast/ops/variable_noise_model.py b/src/toast/ops/variable_noise_model.py new file mode 100644 index 000000000..44305033d --- /dev/null +++ b/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