diff --git a/draco/analysis/sidereal.py b/draco/analysis/sidereal.py index ec2672dfb..92ea066bb 100644 --- a/draco/analysis/sidereal.py +++ b/draco/analysis/sidereal.py @@ -492,18 +492,29 @@ def _regrid(self, vis, weight, lsd): 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. + See West, D.H.D. (1979). https://doi.org/10.1145/359146.359153. - Parameters + Attributes ---------- - tag : str + 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`. + 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 - lsd_list = None 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. @@ -511,14 +522,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: @@ -528,43 +537,131 @@ 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 "nsample" not in self.stack.datasets: + self.stack.add_dataset("nsample") - self.lsd_list = input_lsd + if self.with_sample_variance and ( + "sample_variance" not in self.stack.datasets + ): + self.stack.add_dataset("sample_variance") - 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. + if self.with_sample_variance: + 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, + # 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. + self.stack.nsample[:] += count + + # 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. + self.stack.vis[:] += delta_before * tools.invert_no_zero(sum_coeff) + + # 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 + 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"] = self.tag self.stack.attrs["lsd"] = np.array(self.lsd_list) - self.stack.vis[:] *= tools.invert_no_zero(self.stack.weight[:]) + # 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 = self.stack.nsample[:].astype(np.float32) + self.stack.weight[:] = ( + tools.invert_no_zero(self.stack.weight[:]) * norm ** 2 + ) + + # 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( + self.stack.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 66e9d8918..32878a493 100644 --- a/draco/core/containers.py +++ b/draco/core/containers.py @@ -38,6 +38,7 @@ - :py:class:`TableBase` - :py:class:`TODContainer` - :py:class:`VisContainer` +- :py:class:`SampleVarianceContainer` - :py:class:`FreqContainer` - :py:class:`SiderealContainer` - :py:class:`MContainer` @@ -56,6 +57,8 @@ import numpy as np from caput import memh5, tod +from ..util import tools + # Try to import bitshuffle to set the default compression options try: import bitshuffle.h5 @@ -734,6 +737,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", "