Modulated hwpss (#951)

* Modulate atmosphere-driven HWPSS according to supplied atmospheric signal

* Add workflow tools support for atmosphere-modulated HWPSS

* Full support for modulated HWPSS

* Don't remove the offset, it is important

sotodlib/toast/ops/sim_hwpss.py

@@ -58,6 +58,13 @@ class SimHWPSS(Operator):
         help="Observation detdata key for simulated signal",
     )
 
+    atmo_data = Unicode(
+        None,
+        allow_none=True,
+        help="Observation detdata key for simulated atmosphere "
+        "(modulates part of the HWPSS)",
+    )
+
     stokes_weights = Instance(
         klass=Operator,
         allow_none=True,
@@ -186,6 +193,10 @@ def _exec(self, data, detectors=None, **kwargs):
             chi = obs.shared[self.hwp_angle].data
             for det in dets:
                 signal = obs.detdata[self.det_data][det]
+                if self.atmo_data is None:
+                    atmo = None
+                else:
+                    atmo = obs.detdata[self.atmo_data][det]
                 band = focalplane[det]["band"]
                 freq = {
                     "SAT_f030" : "027",
@@ -238,42 +249,77 @@ def _exec(self, data, detectors=None, **kwargs):
                 theta_high = self.thetas[itheta_high]
                 r = (theta_deg - theta_low) / (theta_high - theta_low)
 
-                transmission = (
-                    (1 - r) * self.all_stokes[freq]["transmission"][itheta_low]
-                    + r * self.all_stokes[freq]["transmission"][itheta_high]
+                # HWPSS not from atmosphere
+
+                transmission_wo_atmo = (
+                    (1 - r) * self.all_stokes[freq]["transmission_wo_atmo"][itheta_low]
+                    + r * self.all_stokes[freq]["transmission_wo_atmo"][itheta_high]
                 )
-                reflection = (
-                    (1 - r) * self.all_stokes[freq]["reflection"][itheta_low]
-                    + r * self.all_stokes[freq]["reflection"][itheta_high]
+                reflection_wo_atmo = (
+                    (1 - r) * self.all_stokes[freq]["reflection_wo_atmo"][itheta_low]
+                    + r * self.all_stokes[freq]["reflection_wo_atmo"][itheta_high]
                 )
+                # Thermal emission from the HWP is not driven by the atmosphere
                 emission = (
-                    (1 - r) * self.all_stokes[freq]["emission"][itheta_low]
-                    + r * self.all_stokes[freq]["emission"][itheta_high]
+                    (1 - r) * self.all_stokes[freq]["emission_wo_atmo"][itheta_low]
+                    + r * self.all_stokes[freq]["emission_wo_atmo"][itheta_high]
                 )
 
+                # HWPSS from atmosphere
+
+                transmission_atmo = (
+                    (1 - r) * self.all_stokes[freq]["transmission_atmo"][itheta_low]
+                    + r * self.all_stokes[freq]["transmission_atmo"][itheta_high]
+                )
+                reflection_atmo = (
+                    (1 - r) * self.all_stokes[freq]["reflection_atmo"][itheta_low]
+                    + r * self.all_stokes[freq]["reflection_atmo"][itheta_high]
+                )
+
+                if atmo is None:
+                    transmission = transmission_wo_atmo + transmission_atmo
+                    reflection = reflection_wo_atmo + reflection_atmo
+                else:
+                    transmission = transmission_wo_atmo
+                    reflection = reflection_wo_atmo
+
                 # Scale HWPSS for observing elevation
 
                 el_ref = np.radians(50)
                 scale = np.sin(el_ref) / np.sin(el)
 
                 # Observe HWPSS with the detector
 
-                iquv = (transmission + reflection).T
+                iquv = (transmission + reflection)
+                iquv = iquv.T.copy()
                 iquss = (
                     iweights * splev(chi, splrep(self.chis, iquv[0], k=5)) +
                     qweights * splev(chi, splrep(self.chis, iquv[1], k=5)) +
                     uweights * splev(chi, splrep(self.chis, iquv[2], k=5))
                 ) * scale
 
-                iquv = emission.T
+                if atmo is not None:
+                    # Atmospheric HWPSS is modulated by the relative
+                    # atmospheric fluctuation
+                    modulation = atmo / np.median(atmo)
+                    iquv = (transmission_atmo + reflection_atmo)
+                    iquv = iquv.T.copy()
+                    # Replace the generic T offset with the simulated atmosphere
+                    # offset
+                    iquv[0] += np.median(atmo) / det_scale - np.mean(iquv[0])
+                    iquss += (
+                        iweights * splev(chi, splrep(self.chis, iquv[0], k=5)) +
+                        qweights * splev(chi, splrep(self.chis, iquv[1], k=5)) +
+                        uweights * splev(chi, splrep(self.chis, iquv[2], k=5))
+                    ) * scale * modulation
+
+                iquv = (emission).T.copy()
                 iquss += (
                     iweights * splev(chi, splrep(self.chis, iquv[0], k=5)) +
                     qweights * splev(chi, splrep(self.chis, iquv[1], k=5)) +
                     uweights * splev(chi, splrep(self.chis, iquv[2], k=5))
                 )
 
-                iquss -= np.median(iquss)
-
                 if self.hwpss_random_drift:
                     # Apply detector couplings to HWPSS random drift common mode
                     key1 = obs.telescope.uid
@@ -295,6 +341,10 @@ def _exec(self, data, detectors=None, **kwargs):
 
                 # Co-add with the cached signal
 
+                if atmo is not None:
+                    # Avoid double-counting the atmosphere.
+                    # HWPSS has a copy of it
+                    signal -= atmo
                 signal += det_scale * iquss * (1 + drift * coupling)
 
         return

sotodlib/toast/workflows/sim_atm.py

@@ -126,6 +126,17 @@ def simulate_atmosphere_signal(job, otherargs, runargs, data):
     # Configured operators for this job
     job_ops = job.operators
 
+    # Check if we need an extra copy of the atmospheric signal
+    final_signal = job_ops.sim_atmosphere.det_data
+    if (
+            hasattr(job_ops, "sim_hwpss")
+            and job_ops.sim_hwpss.enabled
+            and job_ops.sim_hwpss.atmo_data is not None
+        ):
+        temp_signal = job_ops.sim_hwpss.atmo_data
+    else:
+        temp_signal = None
+
     if otherargs.realization is not None:
         job_ops.sim_atmosphere_coarse.realization = 1000000 + otherargs.realization
         job_ops.sim_atmosphere.realization = otherargs.realization
@@ -136,8 +147,29 @@ def simulate_atmosphere_signal(job, otherargs, runargs, data):
         sim_atm.detector_pointing = job_ops.det_pointing_azel_sim
         if sim_atm.polarization_fraction != 0:
             sim_atm.detector_weights = job_ops.weights_azel
+        if temp_signal is not None:
+            # Write the simulated atmosphere to a temporary array
+            sim_atm.det_data = temp_signal
         log.info_rank(f"  Running {sim_atm.name}...", comm=data.comm.comm_world)
         sim_atm.apply(data)
         log.info_rank(
             f"  Applied {sim_atm.name} in", comm=data.comm.comm_world, timer=timer
         )
+        if temp_signal is not None:
+            # Restore original configuration
+            sim_atm.det_data = final_signal
+
+    if temp_signal is not None:
+        # Add the atmospheric signal to the final target but also keep the
+        # separate copy
+        combine = toast.ops.Combine(
+            op="add",
+            first=final_signal,
+            second=temp_signal,
+            result=final_signal,
+        ).apply(data)
+        log.info_rank(
+            f"  Added {temp_signal} to {final_signal} in",
+            comm=data.comm.comm_world,
+            timer=timer,
+        )