feat(containers): add elevation-dependent weights to the HybridVisStream

draco/analysis/transform.py

@@ -490,6 +490,37 @@ def _set_freq_selection(self, **kwargs):
             setattr(self, key, value)
 
 
+class ElevationDependentHybridVisWeight(task.SingleTask):
+    """Add elevation dependence to hybrid visibility weights."""
+
+    def process(self, data: containers.HybridVisStream):
+        """Remove the weights dataset and broadcast to elevation weights.
+
+        Parameters
+        ----------
+        data
+            Hybrid visibilities with elevation-independent weights.
+
+        Returns
+        -------
+        data
+            Input container with different weights dataset
+        """
+        weights = data.weight[:].local_array
+
+        # Remove the reference to the vis_weight dataset
+        del data["vis_weight"]
+
+        # Add the new elevation-dependent weight dataset
+        data.add_dataset("elevation_vis_weight")
+
+        # Write the weights into the new dataset, broadcasting over
+        # the elevation axis
+        data.weight[:].local_array[:] = weights[..., np.newaxis, :]
+
+        return data
+
+
 class MModeTransform(task.SingleTask):
     """Transform a sidereal stream to m-modes.
 

draco/core/containers.py

@@ -1961,6 +1961,16 @@ class HybridVisStream(FreqContainer, SiderealContainer, VisBase):
             "compression": COMPRESSION,
             "compression_opts": COMPRESSION_OPTS,
         },
+        "elevation_vis_weight": {
+            "axes": ["pol", "freq", "ew", "el", "ra"],
+            "dtype": np.float32,
+            "initialise": False,
+            "distributed": True,
+            "distributed_axis": "freq",
+            "chunks": (1, 32, 4, 512, 2048),
+            "compression": COMPRESSION,
+            "compression_opts": COMPRESSION_OPTS,
+        },
         "effective_ra": {
             "axes": ["pol", "freq", "ew", "ra"],
             "dtype": np.float32,
@@ -2029,6 +2039,19 @@ def add_dataset(self, name):
                 "Requesting creation of complex-valued filter but "
                 "real filter already exists."
             )
+        if name == "vis_weight" and "elevation_vis_weight" in self.datasets:
+            raise RuntimeError(
+                "Requesting creation of elevation-independent weights but "
+                "elevation-dependent weights already exist."
+            )
+        if name == "elevation_vis_weight":
+            if "vis_weight" in self.datasets:
+                raise RuntimeError(
+                    "Requesting creation of elevation-dependent weights but "
+                    "elevation-independent weights already exist."
+                )
+            # Make this the default weight dataset
+            self._weight_dset_name = "elevation_vis_weight"
         return super().add_dataset(name)
 
     @property