HWP simulation function now uses units. Add extra helper function to run a simple satellite sim for use in unit tests. Restore unit tests for healpix pointing.

Author: tskisner

src/toast/future_ops/sim_hwp.py

@@ -12,6 +12,7 @@
 @function_timer
 def simulate_hwp_response(
     ob,
+    ob_time_key=None,
     ob_angle_key=None,
     ob_mueller_key=None,
     hwp_start=None,
@@ -64,9 +65,9 @@ def simulate_hwp_response(
         # convert to radians / second
         hwp_rate = hwp_rpm * 2.0 * np.pi / 60.0
 
-    if hwp_step_deg is not None:
+    if hwp_step is not None:
         # convert to radians and seconds
-        hwp_step_rad = hwp_set.to_value(u.radian)
+        hwp_step_rad = hwp_step.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
@@ -78,7 +79,7 @@ def simulate_hwp_response(
     hwp_angle = None
     hwp_mueller = None
 
-    if ob.grid_comm_col is None or ob.grid_comm_col.rank == 0:
+    if ob.comm_col is None or ob.comm_col.rank == 0:
         if hwp_rate is not None:
             # continuous HWP
             # HWP increment per sample is:
@@ -111,7 +112,7 @@ def simulate_hwp_response(
     # 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_angle_key, shape=(n_sample,), dtype=np.float64, comm=ob.comm_col
         )
         ob.shared[ob_angle_key].set(hwp_angle, offset=(0,), fromrank=0)
 

src/toast/future_ops/sim_satellite.py

@@ -510,6 +510,7 @@ def _exec(self, data, detectors=None, **kwargs):
 
             simulate_hwp_response(
                 ob,
+                ob_time_key=self.times,
                 ob_angle_key=self.hwp_angle,
                 ob_mueller_key=None,
                 hwp_start=obsrange[obindx].start * u.second,

src/toast/mpi.py

@@ -40,6 +40,7 @@
 
 import sys
 import itertools
+from contextlib import contextmanager
 
 import numpy as np
 
@@ -234,8 +235,7 @@ def comm_rank(self):
 
     @property
     def cuda(self):
-        """The CUDA device properties for this process.
-        """
+        """The CUDA device properties for this process."""
         return self._cuda
 
     def __repr__(self):
@@ -253,3 +253,34 @@ def __repr__(self):
         else:
             lines.append("  Using CUDA device {}".format(self._cuda.device_index))
         return "<toast.Comm\n{}\n>".format("\n".join(lines))
+
+
+@contextmanager
+def exception_guard(comm=None):
+    """Ensure that if one MPI process raises an un-caught exception, all of them do.
+
+    Args:
+        comm (mpi4py.MPI.Comm): The MPI communicator or None.
+
+    """
+    log = Logger.get()
+    failed = 0
+    try:
+        yield
+    except:
+        msg = "Exception on process {}:\n".format(comm.rank)
+        exc_type, exc_value, exc_traceback = sys.exc_info()
+        lines = traceback.format_exception(exc_type, exc_value, exc_traceback)
+        msg += "\n".join(lines)
+        log.error(msg)
+        failed = 1
+
+    failcount = None
+    if comm is None:
+        failcount = failed
+    else:
+        failcount = comm.allreduce(failed, op=MPI.SUM)
+    if failcount > 0:
+        raise RuntimeError("One or more MPI processes raised an exception")
+
+    return

src/toast/tests/CMakeLists.txt

@@ -21,7 +21,7 @@ install(FILES
     ops_applygain.py
     ops_sim_tod_noise.py
     covariance.py
-    ops_pmat.py
+    ops_pointing_healpix.py
     ops_dipole.py
     ops_groundfilter.py
     sim_focalplane.py

src/toast/tests/_helpers.py

@@ -6,7 +6,7 @@
 
 import numpy as np
 
-# from contextlib import contextmanager
+from astropy import units as u
 
 from ..mpi import Comm
 
@@ -20,8 +20,10 @@
 
 from ..observation import DetectorData, Observation
 
+from .. import future_ops as ops
 
-ZAXIS = np.array([0, 0, 1.0])
+
+ZAXIS = np.array([0.0, 0.0, 1.0])
 
 
 # These are helper routines for common operations used in the unit tests.
@@ -78,7 +80,7 @@ def create_comm(mpicomm):
     return toastcomm
 
 
-def create_telescope(group_size):
+def create_telescope(group_size, sample_rate=10.0 * u.Hz):
     """Create a fake telescope with at least one detector per process."""
     npix = 1
     ring = 1
@@ -90,10 +92,10 @@ def create_telescope(group_size):
 
 
 def create_distdata(mpicomm, obs_per_group=1, samples=10):
-    """Create a toast communicator and distributed data object.
+    """Create a toast communicator and (empty) distributed data object.
 
     Use the specified MPI communicator to attempt to create 2 process groups,
-    each with some observations.
+    each with some empty observations.
 
     Args:
         mpicomm (MPI.Comm): the MPI communicator (or None).
@@ -119,6 +121,42 @@ def create_distdata(mpicomm, obs_per_group=1, samples=10):
     return data
 
 
+def create_satellite_data(
+    mpicomm, obs_per_group=1, sample_rate=10.0 * u.Hz, obs_time=5.0 * u.minute
+):
+    """Create a data object with a simple satellite sim.
+
+    Use the specified MPI communicator to attempt to create 2 process groups.  Create
+    a fake telescope and run the satellite sim to make some observations for each
+    group.  This is useful for testing many operators that need some pre-existing
+    observations with boresight pointing.
+
+    Args:
+        mpicomm (MPI.Comm): the MPI communicator (or None).
+        obs_per_group (int): the number of observations assigned to each group.
+        samples (int): number of samples per observation.
+
+    Returns:
+        toast.Data: the distributed data with named observations.
+
+    """
+    toastcomm = create_comm(mpicomm)
+    data = Data(toastcomm)
+
+    tele = create_telescope(toastcomm.group_size, sample_rate=sample_rate)
+
+    sim_sat = ops.SimSatellite(
+        name="sim_sat",
+        n_observation=(toastcomm.ngroups * obs_per_group),
+        telescope=tele,
+        hwp_rpm=10.0,
+        observation_time=obs_time,
+    )
+    sim_sat.apply(data)
+
+    return data
+
+
 def uniform_chunks(samples, nchunk=100):
     """Divide some number of samples into chunks.
 
@@ -141,37 +179,3 @@ def uniform_chunks(samples, nchunk=100):
     for r in range(remain):
         chunks[r] += 1
     return chunks
-
-
-#
-# @contextmanager
-# def mpi_guard(comm=MPI.COMM_WORLD):
-#     """Ensure that if one MPI process raises an exception, all of them do.
-#
-#     Args:
-#         comm (mpi4py.MPI.Comm): The MPI communicator.
-#
-#     """
-#     failed = 0
-#     print(comm.rank, ": guard: enter", flush=True)
-#     try:
-#         print(comm.rank, ": guard: yield", flush=True)
-#         yield
-#     except:
-#         print(comm.rank, ": guard: except", flush=True)
-#         msg = "Exception on process {}:\n".format(comm.rank)
-#         exc_type, exc_value, exc_traceback = sys.exc_info()
-#         lines = traceback.format_exception(exc_type, exc_value,
-#             exc_traceback)
-#         msg += "\n".join(lines)
-#         print(msg, flush=True)
-#         failed = 1
-#         print(comm.rank, ": guard: except done", flush=True)
-#
-#     print(comm.rank, ": guard: failcount reduce", flush=True)
-#     failcount = comm.allreduce(failed, op=MPI.SUM)
-#     if failcount > 0:
-#         raise RuntimeError("One or more MPI processes raised an exception")
-#     print(comm.rank, ": guard: done", flush=True)
-#
-#     return

src/toast/tests/ops_memory_counter.py

@@ -11,47 +11,30 @@
 
 from .. import future_ops as ops
 
-from .. import config as tc
-
-from ._helpers import create_outdir, create_distdata, create_telescope
+from ._helpers import create_outdir, create_satellite_data
 
 
 class MemoryCounterTest(MPITestCase):
     def setUp(self):
         fixture_name = os.path.splitext(os.path.basename(__file__))[0]
         self.outdir = create_outdir(self.comm, fixture_name)
 
-        # One observation per group
-        self.data = create_distdata(self.comm, obs_per_group=1)
-
-        # Make a fake telescope for every observation
-
-        tele = create_telescope(self.data.comm.group_size)
+    def test_counter(self):
+        # Create a fake satellite data set for testing
+        data = create_satellite_data(self.comm)
 
         # Set up a pipeline that generates some data
-
         pipe_ops = [
-            ops.SimSatellite(
-                name="sim_satellite",
-                telescope=tele,
-                n_observation=self.data.comm.ngroups,
-            ),
-            ops.DefaultNoiseModel(name="noise_model"),
-            ops.SimNoise(name="sim_noise"),
+            ops.DefaultNoiseModel(),
+            ops.SimNoise(),
         ]
 
-        self.pipe = ops.Pipeline(name="sim_pipe")
-        self.pipe.operators = pipe_ops
-
-    def test_counter(self):
-        # Start with empty data
-        self.data.clear()
-
-        # Run a standard pipeline to simulate some data
-        self.pipe.apply(self.data)
+        pipe = ops.Pipeline()
+        pipe.operators = pipe_ops
+        pipe.apply(data)
 
         # Get the memory used
         mcount = ops.MemoryCounter()
-        bytes = mcount.apply(self.data)
+        bytes = mcount.apply(data)
 
         return

src/toast/tests/ops_pmat.py renamed to src/toast/tests/ops_pointing_healpix.py

@@ -10,54 +10,21 @@
 import numpy as np
 
 from .._libtoast import pointing_matrix_healpix
-from ..healpix import HealpixPixels
-from ..todmap import TODHpixSpiral, OpPointingHpix
+
 from .. import qarray as qa
 
-from ._helpers import create_outdir, create_distdata, boresight_focalplane
+from ..healpix import HealpixPixels
+
+from .. import future_ops as ops
+
+from ._helpers import create_outdir, create_satellite_data
 
 
-class OpPointingHpixTest(MPITestCase):
+class PointingHealpixTest(MPITestCase):
     def setUp(self):
         fixture_name = os.path.splitext(os.path.basename(__file__))[0]
         self.outdir = create_outdir(self.comm, fixture_name)
 
-        # Create one observation per group, and each observation will have
-        # one detector per process and a single chunk.  Data within an
-        # observation is distributed by detector.
-
-        self.data = create_distdata(self.comm, obs_per_group=1)
-        self.ndet = self.data.comm.group_size
-
-        # Create detectors with default properties
-        (
-            dnames,
-            dquat,
-            depsilon,
-            drate,
-            dnet,
-            dfmin,
-            dfknee,
-            dalpha,
-        ) = boresight_focalplane(self.ndet)
-
-        # A small number of samples
-        self.totsamp = 10
-
-        # Populate the observations (one per group)
-
-        tod = TODHpixSpiral(
-            self.data.comm.comm_group,
-            dquat,
-            self.totsamp,
-            detranks=self.data.comm.group_size,
-        )
-
-        self.data.obs[0]["tod"] = tod
-
-    def tearDown(self):
-        del self.data
-
     def test_pointing_matrix_healpix2(self):
         nside = 64
         npix = 12 * nside ** 2
@@ -255,31 +222,37 @@ def test_pointing_matrix_healpix_hwp(self):
         return
 
     def test_hpix_simple(self):
+        # Create a fake satellite data set for testing
+        data = create_satellite_data(self.comm)
+
+        pointing = ops.PointingHealpix(nside=64, mode="IQU", hwp_angle="hwp_angle")
+        pointing.apply(data)
+
         rank = 0
         if self.comm is not None:
             rank = self.comm.rank
-        op = OpPointingHpix()
-        op.exec(self.data)
 
         handle = None
         if rank == 0:
             handle = open(os.path.join(self.outdir, "out_test_hpix_simple_info"), "w")
-        self.data.info(handle=handle)
+        data.info(handle=handle)
         if rank == 0:
             handle.close()
-        return
 
     def test_hpix_hwpnull(self):
+        # Create a fake satellite data set for testing
+        data = create_satellite_data(self.comm)
+
+        pointing = ops.PointingHealpix(nside=64, mode="IQU")
+        pointing.apply(data)
+
         rank = 0
         if self.comm is not None:
             rank = self.comm.rank
-        op = OpPointingHpix(mode="IQU")
-        op.exec(self.data)
 
         handle = None
         if rank == 0:
-            handle = open(os.path.join(self.outdir, "out_test_hpix_hwpnull_info"), "w")
-        self.data.info(handle=handle)
+            handle = open(os.path.join(self.outdir, "out_test_hpix_hwpnull"), "w")
+        data.info(handle=handle)
         if rank == 0:
             handle.close()
-        return