Consistent naming of obs / ob. Port memory counter operator and unit tests. Add clear() method to Observation class.

Author: tskisner

src/toast/future_ops/CMakeLists.txt

@@ -4,6 +4,7 @@
 install(FILES
     __init__.py
     pipeline.py
+    memory_counter.py
     sim_hwp.py
     sim_tod_noise.py
     sim_satellite.py

src/toast/future_ops/__init__.py

@@ -4,6 +4,8 @@
 
 # import functions into our public API
 
+from .memory_counter import MemoryCounter
+
 from .pipeline import Pipeline
 
 from .sim_satellite import SimSatellite

src/toast/future_ops/noise_model.py

@@ -43,15 +43,15 @@ def _exec(self, data, detectors=None, **kwargs):
 
         comm = data.comm
 
-        for obs in data.obs:
+        for ob in data.obs:
             # Get the detectors we are using for this observation
-            dets = obs.select_local_detectors(detectors)
+            dets = ob.select_local_detectors(detectors)
             if len(dets) == 0:
                 # Nothing to do for this observation
                 continue
 
             # The focalplane for this observation
-            focalplane = obs.telescope.focalplane
+            focalplane = ob.telescope.focalplane
 
             # Every process has a copy of the focalplane, and every process may want
             # the noise model for all detectors (not just our local detectors).
@@ -73,7 +73,7 @@ def _exec(self, data, detectors=None, **kwargs):
                 rate=rates, fmin=fmin, detectors=dets, fknee=fknee, alpha=alpha, NET=NET
             )
 
-            obs[self.noise_model] = noise
+            ob[self.noise_model] = noise
 
         return
 

src/toast/future_ops/pointing_healpix.py

@@ -162,61 +162,61 @@ def _exec(self, data, detectors=None, **kwargs):
         # overhead from temporary arrays.
         tod_buffer_length = env.tod_buffer_length()
 
-        for obs in data.obs:
+        for ob in data.obs:
             # Get the detectors we are using for this observation
-            dets = obs.select_local_detectors(detectors)
+            dets = ob.select_local_detectors(detectors)
             if len(dets) == 0:
                 # Nothing to do for this observation
                 continue
 
             # Get the flags if needed
             flags = None
             if self.flags is not None:
-                flags = np.array(obs.shared[self.flags])
+                flags = np.array(ob.shared[self.flags])
                 flags &= self.flag_mask
 
             # HWP angle if needed
             hwp_angle = None
             if self.hwp_angle is not None:
-                hwp_angle = obs.shared[self.hwp_angle]
+                hwp_angle = ob.shared[self.hwp_angle]
 
             # Boresight pointing quaternions
-            boresight = obs.shared[self.boresight]
+            boresight = ob.shared[self.boresight]
 
             # Focalplane for this observation
-            focalplane = obs.telescope.focalplane
+            focalplane = ob.telescope.focalplane
 
             # Optional calibration
             cal = None
             if self.cal is not None:
-                cal = obs[self.cal]
+                cal = ob[self.cal]
 
             # Create output data for the pixels, weights and optionally the
             # detector quaternions.
 
             if self.single_precision:
-                obs.detdata.create(
+                ob.detdata.create(
                     self.pixels, detshape=(), dtype=np.int32, detectors=dets
                 )
-                obs.detdata.create(
+                ob.detdata.create(
                     self.weights,
                     detshape=(self._nnz,),
                     dtype=np.float32,
                     detectors=dets,
                 )
             else:
-                obs.detdata.create(
+                ob.detdata.create(
                     self.pixels, detshape=(), dtype=np.int64, detectors=dets
                 )
-                obs.detdata.create(
+                ob.detdata.create(
                     self.weights,
                     detshape=(self._nnz,),
                     dtype=np.float64,
                     detectors=dets,
                 )
 
             if self.quats is not None:
-                obs.detdata.create(
+                ob.detdata.create(
                     self.quats,
                     detshape=(4,),
                     dtype=np.float64,
@@ -237,7 +237,7 @@ def _exec(self, data, detectors=None, **kwargs):
                 # Timestream of detector quaternions
                 quats = qa.mult(boresight, detquat)
                 if self.quats is not None:
-                    obs.detdata[self.quats][det, :] = quats
+                    ob.detdata[self.quats][det, :] = quats
 
                 # Cal for this detector
                 dcal = 1.0
@@ -247,9 +247,9 @@ def _exec(self, data, detectors=None, **kwargs):
                 # Buffered pointing calculation
                 buf_off = 0
                 buf_n = tod_buffer_length
-                while buf_off < obs.n_local_samples:
-                    if buf_off + buf_n > obs.n_local_samples:
-                        buf_n = obs.n_local_samples - buf_off
+                while buf_off < ob.n_local_samples:
+                    if buf_off + buf_n > ob.n_local_samples:
+                        buf_n = ob.n_local_samples - buf_off
                     bslice = slice(buf_off, buf_off + buf_n)
 
                     # This buffer of detector quaternions
@@ -266,8 +266,8 @@ def _exec(self, data, detectors=None, **kwargs):
                         fslice = flags[bslice].reshape(-1)
 
                     # Pixel and weight buffers
-                    pxslice = obs.detdata[self.pixels][det, bslice].reshape(-1)
-                    wtslice = obs.detdata[self.weights][det, bslice].reshape(-1)
+                    pxslice = ob.detdata[self.pixels][det, bslice].reshape(-1)
+                    wtslice = ob.detdata[self.weights][det, bslice].reshape(-1)
 
                     pbuf = pxslice
                     wbuf = wtslice
@@ -296,7 +296,7 @@ def _exec(self, data, detectors=None, **kwargs):
 
                 if self.create_dist is not None:
                     self._local_submaps[
-                        obs.detdata["pixels"][det] // self._n_pix_submap
+                        ob.detdata["pixels"][det] // self._n_pix_submap
                     ] = True
         return
 

src/toast/future_ops/sim_hwp.py

@@ -4,95 +4,115 @@
 
 import numpy as np
 
+from astropy import units as u
+
 from ..timing import function_timer, Timer
 
 
 @function_timer
-def simulate_hwp_angle(
-    obs, obs_key, hwp_start_s, hwp_rpm, hwp_step_deg, hwp_step_time_m
+def simulate_hwp_response(
+    ob,
+    ob_angle_key=None,
+    ob_mueller_key=None,
+    hwp_start=None,
+    hwp_rpm=None,
+    hwp_step=None,
+    hwp_step_time=None,
 ):
     """Simulate and store the HWP angle for one observation.
 
     Args:
-        obs (Observation): The observation to populate.
-        obs_key (str): The observation key for the HWP angle.
-        hwp_start_s (float): The mission starting time in seconds of the HWP rotation.
+        ob (Observation):  The observation to populate.
+        ob_time_key (str):  The observation shared key for timestamps.
+        ob_angle_key (str):  (optional) The output observation key for the HWP angle.
+        ob_mueller_key (str):  (optional) The output observation key for the full
+            Mueller matrix.
+        hwp_start (Quantity): The mission starting time of the HWP rotation.
         hwp_rpm (float): The HWP rotation rate in Revolutions Per Minute.
-        hwp_step_deg (float): The HWP step size in degrees.
-        hwp_step_time_m (float): The time in minutes between steps.
+        hwp_step (Quantity): The HWP step size.
+        hwp_step_time (Quantity): The time between steps.
 
     Returns:
         None
 
     """
-    if hwp_rpm is None and hwp_step_deg is None:
+    if ob_mueller_key is not None:
+        raise NotImplementedError("Mueller matrix not yet implemented")
+
+    if hwp_rpm is None and hwp_step is None:
         # Nothing to do!
         return
 
-    if (hwp_rpm is not None) and (hwp_step_deg is not None):
+    if (hwp_rpm is not None) and (hwp_step is not None):
         raise RuntimeError("choose either continuously rotating or stepped HWP")
 
-    if hwp_step_deg is not None and hwp_step_time_m is None:
+    if hwp_step is not None and hwp_step_time is None:
         raise RuntimeError("for a stepped HWP, you must specify the time between steps")
 
+    hwp_start_s = hwp_start.to_value(u.second)
+
     # compute effective sample rate
-    times = obs.times
+    times = ob.shared[ob_time_key]
     dt = np.mean(times[1:-1] - times[0:-2])
     rate = 1.0 / dt
 
     hwp_rate = None
-    hwp_step = None
-    hwp_step_time = None
+    hwp_step_rad = None
+    hwp_step_time_s = None
 
     if hwp_rpm is not None:
         # convert to radians / second
         hwp_rate = hwp_rpm * 2.0 * np.pi / 60.0
 
     if hwp_step_deg is not None:
         # convert to radians and seconds
-        hwp_step = hwp_step_deg * np.pi / 180.0
-        hwp_step_time = hwp_step_time_m * 60.0
-
-    first_samp, n_samp = obs.local_samples
-
-    obs.shared.create(
-        obs_key, shape=(n_samp,), dtype=np.float64, comm=obs.grid_comm_col
-    )
+        hwp_step_rad = hwp_set.to_value(u.radian)
+        hwp_step_time_s = hwp_step_time.to_value(u.second)
 
     # Only the first process in each grid column simulates the common HWP angle
 
     start_sample = int(hwp_start_s * rate)
+    first_sample = ob.local_index_offset
+    n_sample = ob.n_local_samples
+
     hwp_angle = None
+    hwp_mueller = None
 
-    if obs.grid_comm_col is None or obs.grid_comm_col.rank == 0:
+    if ob.grid_comm_col is None or ob.grid_comm_col.rank == 0:
         if hwp_rate is not None:
             # continuous HWP
             # HWP increment per sample is:
             # (hwprate / samplerate)
             hwpincr = hwp_rate / rate
-            startang = np.fmod((start_sample + first_samp) * hwpincr, 2 * np.pi)
-            hwp_angle = hwpincr * np.arange(n_samp, dtype=np.float64)
+            startang = np.fmod((start_sample + first_sample) * hwpincr, 2 * np.pi)
+            hwp_angle = hwpincr * np.arange(n_sample, dtype=np.float64)
             hwp_angle += startang
         elif hwp_step is not None:
             # stepped HWP
-            hwp_angle = np.ones(n_samp, dtype=np.float64)
-            stepsamples = int(hwp_step_time * rate)
-            wholesteps = int((start_sample + first_samp) / stepsamples)
-            remsamples = (start_sample + first_samp) - wholesteps * stepsamples
-            curang = np.fmod(wholesteps * hwp_step, 2 * np.pi)
+            hwp_angle = np.ones(n_sample, dtype=np.float64)
+            stepsamples = int(hwp_step_time_s * rate)
+            wholesteps = int((start_sample + first_sample) / stepsamples)
+            remsamples = (start_sample + first_sample) - wholesteps * stepsamples
+            curang = np.fmod(wholesteps * hwp_step_rad, 2 * np.pi)
             curoff = 0
             fill = remsamples
-            while curoff < n_samp:
-                if curoff + fill > n_samp:
-                    fill = n_samp - curoff
+            while curoff < n_sample:
+                if curoff + fill > n_sample:
+                    fill = n_sample - curoff
                 hwp_angle[curoff:fill] *= curang
                 curang += hwp_step
                 curoff += fill
                 fill = stepsamples
-        if hwp_angle is not None:
-            # Choose the HWP angle between [0, 2*pi)
-            hwp_angle %= 2 * np.pi
+        # Choose the HWP angle between [0, 2*pi)
+        hwp_angle %= 2 * np.pi
+
+        # Create a Mueller matrix if we will be writing that...
 
-    obs.shared[obs_key].set(hwp_angle, offset=(0,), fromrank=0)
+    # Store the angle and / or the Mueller matrix
+    if ob_angle_key is not None:
+        ob.shared.create(
+            ob_angle_key, shape=(n_sample,), dtype=np.float64, comm=ob.grid_comm_col
+        )
+        ob.shared[ob_angle_key].set(hwp_angle, offset=(0,), fromrank=0)
 
     return