From 1e17901598afb55f5f0fb11e080118e6854a2939 Mon Sep 17 00:00:00 2001
From: Seth Siegel <ssiegel@cedar5.cedar.computecanada.ca>
Date: Fri, 28 Aug 2020 11:31:34 -0700
Subject: [PATCH 1/5] feat(sidereal,containers): sample variance over sidereal
 days

sidereal.SiderealStacker:
- Add the capabililty to use uniform or inverse variance weighting.
- Calculate the sample variance over sidereal days.
- Switch to an algorithm that calculates the current mean at each
iteration and updates the sum of square differences from the
current mean, which is less prone to numerical issues.

containers.SampleVarianceContainer
- Acts as super class for both SiderealStream and TrackBeam.
+ Sets default component axis on instantiation.
+ Defines properties to access common datasets.
+ Defines properties for rotating sample_variance to common bases.
+ Defines sample_weight property that can be compared to the weight
dataset.
- Add sample_variance and nsample datasets to SiderealStream
and TrackBeam.  By default these are not initialized.
---
 draco/analysis/sidereal.py | 128 +++++++++++++++++++----
 draco/core/containers.py   | 201 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 303 insertions(+), 26 deletions(-)

diff --git a/draco/analysis/sidereal.py b/draco/analysis/sidereal.py
index b6ee91f7b..93ccb9bd0 100644
--- a/draco/analysis/sidereal.py
+++ b/draco/analysis/sidereal.py
@@ -249,11 +249,20 @@ def process(self, data):
 class SiderealStacker(task.SingleTask):
     """Take in a set of sidereal days, and stack them up.
 
-    This will apply relative calibration.
+    Also computes the variance over sideral days using an
+    algorithm that updates the sum of square differences from
+    the current mean, which is less prone to numerical issues.
+
+    Attributes
+    ----------
+    weight: str (default: "inverse_variance")
+        The weighting to use in the stack.
+        Either `uniform` or `inverse_variance`.
     """
 
     stack = None
-    lsd_list = None
+
+    weight = config.enum(["uniform", "inverse_variance"], default="inverse_variance")
 
     def process(self, sdata):
         """Stack up sidereal days.
@@ -261,14 +270,12 @@ def process(self, sdata):
         Parameters
         ----------
         sdata : containers.SiderealStream
-            Individual sidereal day to stack up.
+            Individual sidereal day to add to stack.
         """
-
         sdata.redistribute("freq")
 
-        # Get the LSD label out of the data (resort to using a CSD if it's
-        # present). If there's no label just use a place holder and stack
-        # anyway.
+        # Get the LSD (or CSD) label out of the input's attributes.
+        # If there is no label, use a placeholder.
         if "lsd" in sdata.attrs:
             input_lsd = sdata.attrs["lsd"]
         elif "csd" in sdata.attrs:
@@ -278,43 +285,122 @@ def process(self, sdata):
 
         input_lsd = _ensure_list(input_lsd)
 
+        # If this is our first sidereal day, then initialize the
+        # container that will hold the stack.
         if self.stack is None:
 
             self.stack = containers.empty_like(sdata)
-            self.stack.redistribute("freq")
 
-            self.stack.vis[:] = sdata.vis[:] * sdata.weight[:]
-            self.stack.weight[:] = sdata.weight[:]
+            # Add stack-specific datasets
+            if "sample_variance" not in self.stack.datasets:
+                self.stack.add_dataset("sample_variance")
 
-            self.lsd_list = input_lsd
+            if "nsample" not in self.stack.datasets:
+                self.stack.add_dataset("nsample")
 
-            self.log.info("Starting stack with LSD:%i", sdata.attrs["lsd"])
+            self.stack.redistribute("freq")
 
-            return
+            # Initialize all datasets to zero.
+            for data in self.stack.datasets.values():
+                data[:] = 0
 
-        self.log.info("Adding LSD:%i to stack", sdata.attrs["lsd"])
+            self.lsd_list = []
 
-        # note: Eventually we should fix up gains
+            # Keep track of the sum of squared weights
+            # to perform Bessel's correction at the end.
+            self.sum_coeff_sq = np.zeros_like(self.stack.weight[:].view(np.ndarray))
 
-        # Combine stacks with inverse `noise' weighting
-        self.stack.vis[:] += sdata.vis[:] * sdata.weight[:]
-        self.stack.weight[:] += sdata.weight[:]
+        # Accumulate
+        self.log.info("Adding to stack LSD(s): %s" % input_lsd)
 
         self.lsd_list += input_lsd
 
+        if "nsample" in sdata.datasets:
+            # The input sidereal stream is already a stack
+            # over multiple sidereal days. Use the nsample
+            # dataset as the weight for the uniform case.
+            count = sdata.nsample[:]
+        else:
+            # The input sidereal stream contains a single
+            # sidereal day.  Use a boolean array that
+            # indicates a non-zero weight dataset as
+            # the weight for the uniform case.
+            dtype = self.stack.nsample.dtype
+            count = (sdata.weight[:] > 0.0).astype(dtype)
+
+        # Determine the weights to be used in the average.
+        if self.weight == "uniform":
+            coeff = count.astype(np.float32)
+            # Accumulate the variances in the stack.weight dataset.
+            self.stack.weight[:] += (coeff ** 2) * tools.invert_no_zero(sdata.weight[:])
+        else:
+            coeff = sdata.weight[:]
+            # We are using inverse variance weights.  In this case, the
+            # expression above reduces such that we are accumulating
+            # the sum of inverse variances in the stack.weight dataset.
+            # Do that directly to prevent an unneccessary division.
+            self.stack.weight[:] += sdata.weight[:]
+
+        # Accumulate the total number of samples and the sum of squared coefficients.
+        self.stack.nsample[:] += count
+        self.sum_coeff_sq += coeff ** 2
+
+        # Below we will need to normalize by the current sum of coefficients.
+        # Can be found in the stack.nsample dataset for uniform case or
+        # the stack.weight dataset for inverse variance case.
+        if self.weight == "uniform":
+            sum_coeff = self.stack.nsample[:].astype(np.float32)
+        else:
+            sum_coeff = self.stack.weight[:]
+
+        # Calculate weighted difference between the new data and the current mean.
+        delta_before = coeff * (sdata.vis[:] - self.stack.vis[:])
+
+        # Update the mean value
+        self.stack.vis[:] += delta_before * tools.invert_no_zero(sum_coeff)
+
+        # Calculate the difference between the new data and the updated mean.
+        delta_after = sdata.vis[:] - self.stack.vis[:]
+
+        # Update the sample variance.
+        self.stack.sample_variance[0] += delta_before.real * delta_after.real
+        self.stack.sample_variance[1] += delta_before.real * delta_after.imag
+        self.stack.sample_variance[2] += delta_before.imag * delta_after.imag
+
     def process_finish(self):
-        """Construct and emit sidereal stack.
+        """Normalize the stack and return the result.
 
         Returns
         -------
         stack : containers.SiderealStream
             Stack of sidereal days.
         """
-
         self.stack.attrs["tag"] = "stack"
         self.stack.attrs["lsd"] = np.array(self.lsd_list)
 
-        self.stack.vis[:] *= tools.invert_no_zero(self.stack.weight[:])
+        nsample = self.stack.nsample[:]
+
+        # We need to normalize the sample variance by the sum of coefficients.
+        # Can be found in the stack.nsample dataset for uniform case
+        # or the stack.weight dataset for inverse variance case.
+        if self.weight == "uniform":
+            norm = nsample.astype(np.float32)
+            # Normalize the accumulated variances by the total number of samples
+            # and then invert to finalize stack.weight.
+            self.stack.weight[:] = (
+                tools.invert_no_zero(self.stack.weight[:]) * norm ** 2
+            )
+        else:
+            norm = self.stack.weight[:]
+
+        # Perform Bessel's correction.  In the case of
+        # uniform  weighting, norm will be equal to nsample - 1.
+        norm = norm - self.sum_coeff_sq * tools.invert_no_zero(norm)
+
+        # Normalize the sample variance.
+        self.stack.sample_variance[:] *= np.where(
+            nsample > 1, tools.invert_no_zero(norm), 0.0
+        )[np.newaxis, :]
 
         return self.stack
 
diff --git a/draco/core/containers.py b/draco/core/containers.py
index 9a9fd18e7..91c45976f 100644
--- a/draco/core/containers.py
+++ b/draco/core/containers.py
@@ -8,11 +8,26 @@
 
     TimeStream
     SiderealStream
-    GainData
-    StaticGainData
     Map
+    GridBeam
+    TrackBeam
     MModes
-    RingMap
+    SVDModes
+    KLModes
+    CommonModeGainData
+    CommonModeSiderealGainData
+    GainData
+    SiderealGainData
+    StaticGainData
+    DelaySpectrum
+    Powerspectrum2D
+    SVDSpectrum
+    FrequencyStack
+    SourceCatalog
+    SpectroscopicCatalog
+    FormedBeam
+    FormedBeamHA
+
 
 Container Base Classes
 ----------------------
@@ -23,6 +38,7 @@
     ContainerBase
     TODContainer
     VisContainer
+    SampleVarianceContainer
 
 Helper Routines
 ---------------
@@ -50,6 +66,8 @@
 
 from caput import memh5, tod
 
+from ..util import tools
+
 # Try to import bitshuffle to set the default compression options
 try:
     import bitshuffle.h5
@@ -726,6 +744,131 @@ def is_stacked(self):
         return len(self.stack) != len(self.prod)
 
 
+class SampleVarianceContainer(ContainerBase):
+    """Base container for holding the sample variance over observations.
+
+    This works like :class:`ContainerBase` but provides additional capabilities
+    for containers that may be used to hold the sample mean and variance over
+    complex-valued observations.  These capabilities include automatic definition
+    of the component axis, properties for accessing standard datasets, properties
+    that rotate the sample variance into common bases, and a `sample_weight` property
+    that provides an equivalent to the `weight` dataset that is determined from the
+    sample variance over observations.
+
+    Subclasses must include a `sample_variance` and `nsample` dataset
+    in there `_dataset_spec` dictionary.  They must also specify a
+    `_mean` property that returns the dataset containing the mean over observations.
+    """
+
+    _axes = ("component",)
+
+    def __init__(self, *args, **kwargs):
+
+        # Set component axis to default real-imaginary basis if not already provided
+        if "component" not in kwargs:
+            kwargs["component"] = np.array(
+                [("real", "real"), ("real", "imag"), ("imag", "imag")],
+                dtype=[("component_a", "<U8"), ("component_b", "<U8")],
+            )
+
+        super(SampleVarianceContainer, self).__init__(*args, **kwargs)
+
+    @property
+    def component(self):
+        return self.index_map["component"]
+
+    @property
+    def sample_variance(self):
+        """Convience access to the sample variance dataset.
+
+        Returns
+        -------
+        C: np.ndarray[ncomponent, ...]
+            The variance over the dimension that was stacked
+            (e.g., sidereal days, holographic observations)
+            in the default real-imaginary basis. The array is packed
+            into upper-triangle format such that the component axis
+            contains [('real', 'real'), ('real', 'imag'), ('imag', 'imag')].
+        """
+        if "sample_variance" in self.datasets:
+            return self.datasets["sample_variance"]
+        else:
+            raise KeyError("Dataset 'sample_variance' not initialised.")
+
+    @property
+    def sample_variance_iq(self):
+        """Rotate the sample variance to the in-phase/quadrature basis.
+
+        Returns
+        -------
+        C: np.ndarray[ncomponent, ...]
+            The `sample_variance` dataset in the in-phase/quadrature basis,
+            packed into upper triangle format such that the component axis
+            contains [('I', 'I'), ('I', 'Q'), ('Q', 'Q')].
+        """
+        C = self.sample_variance[:].view(np.ndarray)
+
+        # Construct rotation coefficients from average vis angle
+        phi = np.angle(self._mean[:].view(np.ndarray))
+        cc = np.cos(phi) ** 2
+        cs = np.cos(phi) * np.sin(phi)
+        ss = np.sin(phi) ** 2
+
+        # Rotate the covariance matrix from real-imag to in-phase/quadrature
+        Cphi = np.zeros_like(C)
+        Cphi[0] = cc * C[0] + 2 * cs * C[1] + ss * C[2]
+        Cphi[1] = -cs * C[0] + (cc - ss) * C[1] + cs * C[2]
+        Cphi[2] = ss * C[0] - 2 * cs * C[1] + cc * C[2]
+
+        return Cphi
+
+    @property
+    def sample_variance_amp_phase(self):
+        """Calculate the amplitude/phase covariance.
+
+        This interpretation is only valid if the fractional
+        variations in the amplitude and phase are small.
+
+        Returns
+        -------
+        C: np.ndarray[ncomponent, ...]
+            The observed amplitude/phase covariance matrix, packed
+            into upper triangle format such that the component axis
+            contains [('amp', 'amp'), ('amp', 'phase'), ('phase', 'phase')].
+        """
+        # Rotate to in-phase/quadrature basis and then
+        # normalize by squared amplitude to convert to
+        # fractional units (amplitude) and radians (phase).
+        return self.sample_variance_iq * tools.invert_no_zero(
+            np.abs(self._mean[:][np.newaxis, ...]) ** 2
+        )
+
+    @property
+    def nsample(self):
+        if "nsample" in self.datasets:
+            return self.datasets["nsample"]
+        else:
+            raise KeyError("Dataset 'nsample' not initialised.")
+
+    @property
+    def sample_weight(self):
+        """Calculate a weight from the sample variance.
+
+        Returns
+        -------
+        weight: np.ndarray[...]
+            The trace of the `sample_variance` dataset is used
+            as an estimate of the total variance and divided by the
+            `nsample` dataset to yield the uncertainty on the mean.
+            The inverse of this quantity is returned, and can be compared
+            directly to the `weight` dataset.
+        """
+        C = self.sample_variance[:].view(np.ndarray)
+        nsample = self.nsample[:].view(np.ndarray)
+
+        return nsample * tools.invert_no_zero(C[0] + C[2])
+
+
 class Map(ContainerBase):
     """Container for holding multifrequency sky maps.
 
@@ -774,7 +917,7 @@ def map(self):
         return self.datasets["map"]
 
 
-class SiderealStream(VisContainer):
+class SiderealStream(VisContainer, SampleVarianceContainer):
     """A container for holding a visibility dataset in sidereal time.
 
     Parameters
@@ -819,6 +962,26 @@ class SiderealStream(VisContainer):
             "distributed": True,
             "distributed_axis": "freq",
         },
+        "sample_variance": {
+            "axes": ["component", "freq", "stack", "ra"],
+            "dtype": np.float32,
+            "initialise": False,
+            "distributed": True,
+            "distributed_axis": "freq",
+            "compression": COMPRESSION,
+            "compression_opts": COMPRESSION_OPTS,
+            "chunks": (3, 64, 256, 128),
+        },
+        "nsample": {
+            "axes": ["freq", "stack", "ra"],
+            "dtype": np.uint16,
+            "initialise": False,
+            "distributed": True,
+            "distributed_axis": "freq",
+            "compression": COMPRESSION,
+            "compression_opts": COMPRESSION_OPTS,
+            "chunks": (64, 256, 128),
+        },
     }
 
     def __init__(self, ra=None, *args, **kwargs):
@@ -843,6 +1006,10 @@ def input_flags(self):
     def ra(self):
         return self.index_map["ra"]
 
+    @property
+    def _mean(self):
+        return self.datasets["vis"]
+
 
 class SystemSensitivity(TODContainer):
     """A container for holding the total system sensitivity.
@@ -1066,7 +1233,7 @@ def phi(self):
         return self.index_map["phi"]
 
 
-class TrackBeam(ContainerBase):
+class TrackBeam(SampleVarianceContainer):
     """ Container for a sequence of beam samples at arbitrary locations
         on the sphere. The axis of the beam samples is 'pix', defined by
         the numpy.dtype [('theta', np.float32), ('phi', np.float32)].
@@ -1095,6 +1262,26 @@ class TrackBeam(ContainerBase):
             "compression_opts": COMPRESSION_OPTS,
             "chunks": (128, 2, 128, 128),
         },
+        "sample_variance": {
+            "axes": ["component", "freq", "pol", "input", "pix"],
+            "dtype": np.float32,
+            "initialise": False,
+            "distributed": True,
+            "distributed_axis": "freq",
+            "compression": COMPRESSION,
+            "compression_opts": COMPRESSION_OPTS,
+            "chunks": (3, 128, 2, 128, 128),
+        },
+        "nsample": {
+            "axes": ["freq", "pol", "input", "pix"],
+            "dtype": np.uint8,
+            "initialise": False,
+            "distributed": True,
+            "distributed_axis": "freq",
+            "compression": COMPRESSION,
+            "compression_opts": COMPRESSION_OPTS,
+            "chunks": (128, 2, 128, 128),
+        },
     }
 
     def __init__(
@@ -1164,6 +1351,10 @@ def input(self):
     def pix(self):
         return self.index_map["pix"]
 
+    @property
+    def _mean(self):
+        return self.datasets["beam"]
+
 
 class MModes(VisContainer):
     """Parallel container for holding m-mode data.

From deb8fc41316886d6b35d5e77e87bc1b4562cb713 Mon Sep 17 00:00:00 2001
From: Seth Siegel <ssiegel@cedar5.cedar.computecanada.ca>
Date: Mon, 31 Aug 2020 20:16:19 -0700
Subject: [PATCH 2/5] fix(sidereal): clarify comments

---
 draco/analysis/sidereal.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/draco/analysis/sidereal.py b/draco/analysis/sidereal.py
index 93ccb9bd0..54724a7d6 100644
--- a/draco/analysis/sidereal.py
+++ b/draco/analysis/sidereal.py
@@ -335,10 +335,10 @@ def process(self, sdata):
             self.stack.weight[:] += (coeff ** 2) * tools.invert_no_zero(sdata.weight[:])
         else:
             coeff = sdata.weight[:]
-            # We are using inverse variance weights.  In this case, the
-            # expression above reduces such that we are accumulating
-            # the sum of inverse variances in the stack.weight dataset.
-            # Do that directly to prevent an unneccessary division.
+            # We are using inverse variance weights.  In this case,
+            # we accumulate the inverse variances in the stack.weight
+            # dataset.  Do that directly to avoid an unneccessary
+            # division in the more general expression above.
             self.stack.weight[:] += sdata.weight[:]
 
         # Accumulate the total number of samples and the sum of squared coefficients.
@@ -356,7 +356,7 @@ def process(self, sdata):
         # Calculate weighted difference between the new data and the current mean.
         delta_before = coeff * (sdata.vis[:] - self.stack.vis[:])
 
-        # Update the mean value
+        # Update the mean.
         self.stack.vis[:] += delta_before * tools.invert_no_zero(sum_coeff)
 
         # Calculate the difference between the new data and the updated mean.

From 57b290feeef47ff208cd81b9962573b602335475 Mon Sep 17 00:00:00 2001
From: ssiegelx <sethrsiegel@gmail.com>
Date: Mon, 31 Aug 2020 23:30:31 -0400
Subject: [PATCH 3/5] style: run black v20.8b1

---
 draco/analysis/beamform.py    | 37 ++++++++------------
 draco/analysis/fgfilter.py    |  3 +-
 draco/analysis/flagging.py    |  3 +-
 draco/analysis/sourcestack.py |  2 +-
 draco/core/containers.py      | 45 +++++++++---------------
 draco/core/io.py              |  8 ++---
 draco/core/misc.py            |  2 +-
 draco/core/task.py            |  9 ++---
 draco/synthesis/stream.py     |  6 ++--
 draco/util/tools.py           | 66 +++++++++++++++++------------------
 10 files changed, 77 insertions(+), 104 deletions(-)

diff --git a/draco/analysis/beamform.py b/draco/analysis/beamform.py
index 65492ce11..756eeaf07 100644
--- a/draco/analysis/beamform.py
+++ b/draco/analysis/beamform.py
@@ -24,7 +24,7 @@
 
 
 class BeamFormBase(task.SingleTask):
-    """ Base class for beam forming tasks.
+    """Base class for beam forming tasks.
 
     Defines a few useful methods. Not to be used directly
     but as parent class for BeamForm and BeamFormCat.
@@ -60,7 +60,7 @@ class BeamFormBase(task.SingleTask):
     freqside = config.Property(proptype=int, default=None)
 
     def setup(self, manager):
-        """ Generic setup method.
+        """Generic setup method.
 
         To be complemented by specific
         setup methods in daughter tasks.
@@ -97,7 +97,7 @@ def setup(self, manager):
         self.latitude = np.deg2rad(self.telescope.latitude)
 
     def process(self):
-        """ Generic process method.
+        """Generic process method.
 
         Performs all the beamforming,
         but not the data parsing. To be complemented by specific
@@ -331,7 +331,7 @@ def process(self):
         return formed_beam
 
     def _ha_side(self, data, timetrack=900.0):
-        """ Number of RA/time bins to track the source at each side of transit.
+        """Number of RA/time bins to track the source at each side of transit.
 
         Parameters
         ----------
@@ -358,7 +358,7 @@ def _ha_side(self, data, timetrack=900.0):
         return int(timetrack / approx_time_perbin)
 
     def _ha_array(self, ra, source_ra_index, source_ra, ha_side, is_sstream=True):
-        """ Hour angle for each RA/time bin to be processed.
+        """Hour angle for each RA/time bin to be processed.
 
         Also return the indices of these bins in the full RA/time axis.
 
@@ -416,7 +416,7 @@ def _ha_array(self, ra, source_ra_index, source_ra, ha_side, is_sstream=True):
 
     # TODO: This is very CHIME specific. Should probably be moved somewhere else.
     def _beamfunc(self, ha, pol, freq, dec, zenith=0.70999994):
-        """ Simple and fast beam model to be used as beamforming weights.
+        """Simple and fast beam model to be used as beamforming weights.
 
         Parameters
         ----------
@@ -447,14 +447,12 @@ def _beamfunc(self, ha, pol, freq, dec, zenith=0.70999994):
             pol = pollist.index(pol)
 
         def _sig(pp, freq, dec):
-            """
-            """
+            """"""
             sig_amps = [14.87857614, 9.95746878]
             return sig_amps[pp] / freq / np.cos(dec)
 
         def _amp(pp, dec, zenith):
-            """
-            """
+            """"""
 
             def _flat_top_gauss6(x, A, sig, x0):
                 """Flat-top gaussian. Power of 6."""
@@ -488,8 +486,7 @@ def _flat_top_gauss3(x, A, sig, x0):
             ) ** 0.5
 
     def _process_data(self, data):
-        """ Store code for parsing and formating data prior to beamforming.
-        """
+        """Store code for parsing and formating data prior to beamforming."""
         # Easy access to communicator
         self.comm_ = data.comm
 
@@ -627,12 +624,10 @@ def _process_catalog(self, catalog):
 
 
 class BeamForm(BeamFormBase):
-    """ BeamForm for a single source catalog and multiple visibility datasets.
-
-    """
+    """BeamForm for a single source catalog and multiple visibility datasets."""
 
     def setup(self, manager, source_cat):
-        """ Parse the source catalog and performs the generic setup.
+        """Parse the source catalog and performs the generic setup.
 
         Parameters
         ----------
@@ -647,7 +642,7 @@ def setup(self, manager, source_cat):
         self.catalog = source_cat
 
     def process(self, data):
-        """ Parse the visibility data and beamforms all sources.
+        """Parse the visibility data and beamforms all sources.
 
         Parameters
         ----------
@@ -668,12 +663,10 @@ def process(self, data):
 
 
 class BeamFormCat(BeamFormBase):
-    """ BeamForm for multiple source catalogs and a single visibility dataset.
-
-    """
+    """BeamForm for multiple source catalogs and a single visibility dataset."""
 
     def setup(self, manager, data):
-        """ Parse the visibility data and performs the generic setup.
+        """Parse the visibility data and performs the generic setup.
 
         Parameters
         ----------
@@ -690,7 +683,7 @@ def setup(self, manager, data):
         self._process_data(data)
 
     def process(self, source_cat):
-        """ Parse the source catalog and beamforms all sources.
+        """Parse the source catalog and beamforms all sources.
 
         Parameters
         ----------
diff --git a/draco/analysis/fgfilter.py b/draco/analysis/fgfilter.py
index dd298fdc0..a412fe374 100644
--- a/draco/analysis/fgfilter.py
+++ b/draco/analysis/fgfilter.py
@@ -169,8 +169,7 @@ class KLModeProject(_ProjectFilterBase):
     klname = config.Property(proptype=str)
 
     def setup(self, manager):
-        """Set the product manager that holds the saved KL modes.
-        """
+        """Set the product manager that holds the saved KL modes."""
         self.product_manager = manager
 
     def _forward(self, svdmodes):
diff --git a/draco/analysis/flagging.py b/draco/analysis/flagging.py
index 6fded3d39..eed3a6e45 100644
--- a/draco/analysis/flagging.py
+++ b/draco/analysis/flagging.py
@@ -674,8 +674,7 @@ def _apply_sir(self, mask, baseflag, eta=0.2):
         return flagsir
 
     def _mad_tv_mask(self, data, start_flag, freq):
-        """Use the specific scattered TV channel flagging.
-        """
+        """Use the specific scattered TV channel flagging."""
         # Make copy of data
         data = np.copy(data)
 
diff --git a/draco/analysis/sourcestack.py b/draco/analysis/sourcestack.py
index 8cae1c1e1..3c577128d 100644
--- a/draco/analysis/sourcestack.py
+++ b/draco/analysis/sourcestack.py
@@ -36,7 +36,7 @@ class SourceStack(task.SingleTask):
     freqside = config.Property(proptype=int, default=50)
 
     def process(self, formed_beam):
-        """ Receives a formed beam object and stack across sources.
+        """Receives a formed beam object and stack across sources.
 
         Parameters
         ----------
diff --git a/draco/core/containers.py b/draco/core/containers.py
index 91c45976f..efb7dbd8c 100644
--- a/draco/core/containers.py
+++ b/draco/core/containers.py
@@ -330,8 +330,7 @@ def dataset_spec(self):
 
     @classmethod
     def _class_axes(cls):
-        """Return the set of axes for this container defined by this class and the base classes.
-        """
+        """Return the set of axes for this container defined by this class and the base classes."""
         axes = set()
 
         # Iterate over the reversed MRO and look for _table_spec attributes
@@ -350,8 +349,7 @@ def _class_axes(cls):
 
     @property
     def axes(self):
-        """The set of axes for this container including any defined on the instance.
-        """
+        """The set of axes for this container including any defined on the instance."""
         axes = set(self._class_axes())
 
         # Add in any axes found on the instance (this is needed to support the table classes where
@@ -1161,8 +1159,8 @@ def input_flags(self):
 
 
 class GridBeam(ContainerBase):
-    """ Generic container for representing the 2-d beam in spherical
-        coordinates on a rectangular grid.
+    """Generic container for representing the 2-d beam in spherical
+    coordinates on a rectangular grid.
     """
 
     _axes = ("freq", "pol", "input", "theta", "phi")
@@ -1234,9 +1232,9 @@ def phi(self):
 
 
 class TrackBeam(SampleVarianceContainer):
-    """ Container for a sequence of beam samples at arbitrary locations
-        on the sphere. The axis of the beam samples is 'pix', defined by
-        the numpy.dtype [('theta', np.float32), ('phi', np.float32)].
+    """Container for a sequence of beam samples at arbitrary locations
+    on the sphere. The axis of the beam samples is 'pix', defined by
+    the numpy.dtype [('theta', np.float32), ('phi', np.float32)].
     """
 
     _axes = ("freq", "pol", "input", "pix")
@@ -1486,8 +1484,7 @@ class KLModes(SVDModes):
 
 
 class CommonModeGainData(TODContainer):
-    """Parallel container for holding gain data common to all inputs.
-    """
+    """Parallel container for holding gain data common to all inputs."""
 
     _axes = ("freq",)
 
@@ -1525,8 +1522,7 @@ def freq(self):
 
 
 class CommonModeSiderealGainData(ContainerBase):
-    """Parallel container for holding sidereal gain data common to all inputs.
-    """
+    """Parallel container for holding sidereal gain data common to all inputs."""
 
     _axes = ("freq", "ra")
 
@@ -1568,8 +1564,7 @@ def ra(self):
 
 
 class GainData(TODContainer):
-    """Parallel container for holding gain data.
-    """
+    """Parallel container for holding gain data."""
 
     _axes = ("freq", "input")
 
@@ -1611,8 +1606,7 @@ def input(self):
 
 
 class SiderealGainData(ContainerBase):
-    """Parallel container for holding sidereal gain data.
-    """
+    """Parallel container for holding sidereal gain data."""
 
     _axes = ("freq", "input", "ra")
 
@@ -1658,8 +1652,7 @@ def ra(self):
 
 
 class StaticGainData(ContainerBase):
-    """Parallel container for holding static gain data (i.e. non time varying).
-    """
+    """Parallel container for holding static gain data (i.e. non time varying)."""
 
     _axes = ("freq", "input")
 
@@ -1698,8 +1691,7 @@ def input(self):
 
 
 class DelaySpectrum(ContainerBase):
-    """Container for a delay spectrum.
-    """
+    """Container for a delay spectrum."""
 
     _axes = ("baseline", "delay")
 
@@ -1791,8 +1783,7 @@ def C_inv(self):
 
 
 class SVDSpectrum(ContainerBase):
-    """Container for an m-mode SVD spectrum.
-    """
+    """Container for an m-mode SVD spectrum."""
 
     _axes = ("m", "singularvalue")
 
@@ -1866,8 +1857,7 @@ class SourceCatalog(TableBase):
 
 
 class SpectroscopicCatalog(SourceCatalog):
-    """A container for spectroscopic catalogs.
-    """
+    """A container for spectroscopic catalogs."""
 
     _table_spec = {
         "redshift": {
@@ -1878,8 +1868,7 @@ class SpectroscopicCatalog(SourceCatalog):
 
 
 class FormedBeam(ContainerBase):
-    """Container for formed beams.
-    """
+    """Container for formed beams."""
 
     _axes = ("object_id", "pol", "freq")
 
@@ -1939,7 +1928,7 @@ def pol(self):
 
 class FormedBeamHA(FormedBeam):
     """Container for formed beams.
-       These have not been collapsed in the hour angle (HA) axis
+    These have not been collapsed in the hour angle (HA) axis
     """
 
     _axes = ("object_id", "pol", "freq", "ha")
diff --git a/draco/core/io.py b/draco/core/io.py
index 1e35b7ce2..745b1541d 100644
--- a/draco/core/io.py
+++ b/draco/core/io.py
@@ -272,7 +272,7 @@ def process(self):
 
 
 class FindFiles(pipeline.TaskBase):
-    """ Take a glob or list of files specified as a parameter in the
+    """Take a glob or list of files specified as a parameter in the
     configuration file and pass on to other tasks.
 
     Parameters
@@ -283,8 +283,7 @@ class FindFiles(pipeline.TaskBase):
     files = config.Property(proptype=_list_or_glob)
 
     def setup(self):
-        """ Return list of files specified in the parameters.
-        """
+        """Return list of files specified in the parameters."""
         if not isinstance(self.files, (list, tuple)):
             raise RuntimeError("Argument must be list of files.")
 
@@ -353,8 +352,7 @@ def next(self, data):
 
 
 class Print(pipeline.TaskBase):
-    """Stupid module which just prints whatever it gets. Good for debugging.
-    """
+    """Stupid module which just prints whatever it gets. Good for debugging."""
 
     def next(self, input_):
 
diff --git a/draco/core/misc.py b/draco/core/misc.py
index 0fd849908..b6e5bb12f 100644
--- a/draco/core/misc.py
+++ b/draco/core/misc.py
@@ -47,7 +47,7 @@ def process(self, tstream, gain):
         ----------
         tstream : TimeStream like or SiderealStream
             Time stream to apply gains to. The gains are applied in place.
-        gain : StaticGainData, GainData, SiderealGainData, CommonModeGainData 
+        gain : StaticGainData, GainData, SiderealGainData, CommonModeGainData
             or CommonModeSiderealGainData. Gains to apply.
 
         Returns
diff --git a/draco/core/task.py b/draco/core/task.py
index 2aa8a345d..2ef2d9d29 100644
--- a/draco/core/task.py
+++ b/draco/core/task.py
@@ -146,8 +146,7 @@ def __init__(self):
 
 
 class LoggedTask(pipeline.TaskBase):
-    """A task with logger support.
-    """
+    """A task with logger support."""
 
     log_level = config.Property(proptype=_log_level, default=None)
 
@@ -164,8 +163,7 @@ def __init__(self):
 
     @property
     def log(self):
-        """The logger object for this task.
-        """
+        """The logger object for this task."""
         return self._log
 
 
@@ -207,8 +205,7 @@ def process(self, msg, kwargs):
 
 
 class MPILoggedTask(MPITask, LoggedTask):
-    """A task base that has MPI aware logging.
-    """
+    """A task base that has MPI aware logging."""
 
     def __init__(self):
 
diff --git a/draco/synthesis/stream.py b/draco/synthesis/stream.py
index 9e43047c5..e01fadf2a 100644
--- a/draco/synthesis/stream.py
+++ b/draco/synthesis/stream.py
@@ -31,8 +31,7 @@
 
 
 class SimulateSidereal(task.SingleTask):
-    """Create a simulated sidereal dataset from an input map.
-    """
+    """Create a simulated sidereal dataset from an input map."""
 
     done = False
 
@@ -178,8 +177,7 @@ def process(self, map_):
 
 
 class ExpandProducts(task.SingleTask):
-    """Un-wrap collated products to full triangle.
-    """
+    """Un-wrap collated products to full triangle."""
 
     def setup(self, telescope):
         """Get a reference to the telescope class.
diff --git a/draco/util/tools.py b/draco/util/tools.py
index 8472d32cd..080dc30d8 100644
--- a/draco/util/tools.py
+++ b/draco/util/tools.py
@@ -397,24 +397,24 @@ def redefine_stack_index_map(telescope, inputs, prod, stack, reverse_stack):
 
 
 def polarization_map(index_map, telescope, exclude_autos=True):
-    """ Map the visibilities corresponding to entries in
-        pol = ['XX', 'XY', 'YX', 'YY'].
-
-        Parameters
-        ----------
-        index_map : h5py.group or dict
-            Index map to map into polarizations. Must contain a `stack`
-            entry and an `input` entry.
-        telescope : :class: `drift.core.telescope`
-            Telescope object containing feed information.
-        exclude_autos: bool
-            If True (default), auto-correlations are set to -1.
-
-        Returns
-        -------
-        polmap : array of int
-            Array of size `nstack`. Each entry is the index to the
-            corresponding polarization in pol = ['XX', 'XY', 'YX', 'YY']
+    """Map the visibilities corresponding to entries in
+    pol = ['XX', 'XY', 'YX', 'YY'].
+
+    Parameters
+    ----------
+    index_map : h5py.group or dict
+        Index map to map into polarizations. Must contain a `stack`
+        entry and an `input` entry.
+    telescope : :class: `drift.core.telescope`
+        Telescope object containing feed information.
+    exclude_autos: bool
+        If True (default), auto-correlations are set to -1.
+
+    Returns
+    -------
+    polmap : array of int
+        Array of size `nstack`. Each entry is the index to the
+        corresponding polarization in pol = ['XX', 'XY', 'YX', 'YY']
 
     """
     # Old versions of telescope object don't have the `stack_type`
@@ -485,21 +485,21 @@ def polarization_map(index_map, telescope, exclude_autos=True):
 
 
 def baseline_vector(index_map, telescope):
-    """ Baseline vectors in meters.
-
-        Parameters
-        ----------
-        index_map : h5py.group or dict
-            Index map to map into polarizations. Must contain a `stack`
-            entry and an `input` entry.
-        telescope : :class: `drift.core.telescope`
-            Telescope object containing feed information.
-
-        Returns
-        -------
-        bvec_m : array
-            Array of shape (2, nstack). The 2D baseline vector
-            (in meters) for each visibility in index_map['stack']
+    """Baseline vectors in meters.
+
+    Parameters
+    ----------
+    index_map : h5py.group or dict
+        Index map to map into polarizations. Must contain a `stack`
+        entry and an `input` entry.
+    telescope : :class: `drift.core.telescope`
+        Telescope object containing feed information.
+
+    Returns
+    -------
+    bvec_m : array
+        Array of shape (2, nstack). The 2D baseline vector
+        (in meters) for each visibility in index_map['stack']
     """
     nstack = len(index_map["stack"])
     # Baseline vectors in meters.

From d80f831482c78f05ed8ea7e1447b1720c8b060d8 Mon Sep 17 00:00:00 2001
From: Seth Siegel <seth@Seths-Mac-mini.local>
Date: Fri, 14 May 2021 19:18:16 -0400
Subject: [PATCH 4/5] feat(sidereal): add with_sample_variance config property

Now SiderealStacker only calculates the sample variance if requested.
---
 draco/analysis/sidereal.py | 80 ++++++++++++++++++++++----------------
 1 file changed, 47 insertions(+), 33 deletions(-)

diff --git a/draco/analysis/sidereal.py b/draco/analysis/sidereal.py
index d088574e0..3c0d45cc8 100644
--- a/draco/analysis/sidereal.py
+++ b/draco/analysis/sidereal.py
@@ -494,21 +494,26 @@ class SiderealStacker(task.SingleTask):
     Also computes the variance over sideral days using an
     algorithm that updates the sum of square differences from
     the current mean, which is less prone to numerical issues.
+    See West, D.H.D. (1979). https://doi.org/10.1145/359146.359153.
 
     Attributes
     ----------
+    tag : str (default: "stack")
+        The tag to give the stack.
     weight: str (default: "inverse_variance")
         The weighting to use in the stack.
         Either `uniform` or `inverse_variance`.
-    tag : str
-        The tag to give the stack.
+    with_sample_variance : bool (default: False)
+        Add a dataset containing the sample variance
+        of the visibilities over sidereal days to the
+        sidereal stack.
     """
 
     stack = None
 
-    weight = config.enum(["uniform", "inverse_variance"], default="inverse_variance")
-
     tag = config.Property(proptype=str, default="stack")
+    weight = config.enum(["uniform", "inverse_variance"], default="inverse_variance")
+    with_sample_variance = config.Property(proptype=bool, default=False)
 
     def process(self, sdata):
         """Stack up sidereal days.
@@ -538,12 +543,14 @@ def process(self, sdata):
             self.stack = containers.empty_like(sdata)
 
             # Add stack-specific datasets
-            if "sample_variance" not in self.stack.datasets:
-                self.stack.add_dataset("sample_variance")
-
             if "nsample" not in self.stack.datasets:
                 self.stack.add_dataset("nsample")
 
+            if self.with_sample_variance and (
+                "sample_variance" not in self.stack.datasets
+            ):
+                self.stack.add_dataset("sample_variance")
+
             self.stack.redistribute("freq")
 
             # Initialize all datasets to zero.
@@ -554,7 +561,8 @@ def process(self, sdata):
 
             # Keep track of the sum of squared weights
             # to perform Bessel's correction at the end.
-            self.sum_coeff_sq = np.zeros_like(self.stack.weight[:].view(np.ndarray))
+            if self.with_sample_variance:
+                self.sum_coeff_sq = np.zeros_like(self.stack.weight[:].view(np.ndarray))
 
         # Accumulate
         self.log.info("Adding to stack LSD(s): %s" % input_lsd)
@@ -587,9 +595,8 @@ def process(self, sdata):
             # division in the more general expression above.
             self.stack.weight[:] += sdata.weight[:]
 
-        # Accumulate the total number of samples and the sum of squared coefficients.
+        # Accumulate the total number of samples.
         self.stack.nsample[:] += count
-        self.sum_coeff_sq += coeff ** 2
 
         # Below we will need to normalize by the current sum of coefficients.
         # Can be found in the stack.nsample dataset for uniform case or
@@ -605,13 +612,19 @@ def process(self, sdata):
         # Update the mean.
         self.stack.vis[:] += delta_before * tools.invert_no_zero(sum_coeff)
 
-        # Calculate the difference between the new data and the updated mean.
-        delta_after = sdata.vis[:] - self.stack.vis[:]
+        # The calculations below are only required if the sample variance was requested
+        if self.with_sample_variance:
+
+            # Accumulate the sum of squared coefficients.
+            self.sum_coeff_sq += coeff ** 2
+
+            # Calculate the difference between the new data and the updated mean.
+            delta_after = sdata.vis[:] - self.stack.vis[:]
 
-        # Update the sample variance.
-        self.stack.sample_variance[0] += delta_before.real * delta_after.real
-        self.stack.sample_variance[1] += delta_before.real * delta_after.imag
-        self.stack.sample_variance[2] += delta_before.imag * delta_after.imag
+            # Update the sample variance.
+            self.stack.sample_variance[0] += delta_before.real * delta_after.real
+            self.stack.sample_variance[1] += delta_before.real * delta_after.imag
+            self.stack.sample_variance[2] += delta_before.imag * delta_after.imag
 
     def process_finish(self):
         """Normalize the stack and return the result.
@@ -624,29 +637,30 @@ def process_finish(self):
         self.stack.attrs["tag"] = self.tag
         self.stack.attrs["lsd"] = np.array(self.lsd_list)
 
-        nsample = self.stack.nsample[:]
-
-        # We need to normalize the sample variance by the sum of coefficients.
-        # Can be found in the stack.nsample dataset for uniform case
-        # or the stack.weight dataset for inverse variance case.
+        # For uniform weighting, normalize the accumulated variances by the total
+        # number of samples squared and then invert to finalize stack.weight.
         if self.weight == "uniform":
-            norm = nsample.astype(np.float32)
-            # Normalize the accumulated variances by the total number of samples
-            # and then invert to finalize stack.weight.
+            norm = self.stack.nsample[:].astype(np.float32)
             self.stack.weight[:] = (
                 tools.invert_no_zero(self.stack.weight[:]) * norm ** 2
             )
-        else:
-            norm = self.stack.weight[:]
 
-        # Perform Bessel's correction.  In the case of
-        # uniform  weighting, norm will be equal to nsample - 1.
-        norm = norm - self.sum_coeff_sq * tools.invert_no_zero(norm)
+        # We need to normalize the sample variance by the sum of coefficients.
+        # Can be found in the stack.nsample dataset for uniform case
+        # or the stack.weight dataset for inverse variance case.
+        if self.with_sample_variance:
+
+            if self.weight != "uniform":
+                norm = self.stack.weight[:]
+
+            # Perform Bessel's correction.  In the case of
+            # uniform  weighting, norm will be equal to nsample - 1.
+            norm = norm - self.sum_coeff_sq * tools.invert_no_zero(norm)
 
-        # Normalize the sample variance.
-        self.stack.sample_variance[:] *= np.where(
-            nsample > 1, tools.invert_no_zero(norm), 0.0
-        )[np.newaxis, :]
+            # Normalize the sample variance.
+            self.stack.sample_variance[:] *= np.where(
+                self.stack.nsample[:] > 1, tools.invert_no_zero(norm), 0.0
+            )[np.newaxis, :]
 
         return self.stack
 

From 4c567abd70e5b3e4b43793d1b3aa249b1798da5f Mon Sep 17 00:00:00 2001
From: Seth Siegel <seth@Seths-Mac-mini.local>
Date: Wed, 26 May 2021 12:29:18 -0400
Subject: [PATCH 5/5] style(containers): run black

---
 draco/core/containers.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/draco/core/containers.py b/draco/core/containers.py
index 23ac4608a..e8d3819bf 100644
--- a/draco/core/containers.py
+++ b/draco/core/containers.py
@@ -1016,7 +1016,9 @@ def map(self):
         return self.datasets["map"]
 
 
-class SiderealStream(FreqContainer, VisContainer, SiderealContainer, SampleVarianceContainer):
+class SiderealStream(
+    FreqContainer, VisContainer, SiderealContainer, SampleVarianceContainer
+):
     """A container for holding a visibility dataset in sidereal time.
 
     Parameters