feat(sidereal): fixes to gradient correction

draco/analysis/sidereal.py

@@ -9,8 +9,6 @@
 :class:`SiderealStacker` if you want to combine the different days.
 """
 
-from typing import Union
-
 import numpy as np
 import scipy.linalg as la
 from caput import config, mpiarray, tod
@@ -583,7 +581,7 @@ class BinCentreGradientCorrection(task.SingleTask):
     create the gradient.
     """
 
-    def setup(self, sstream_ref: Union[containers.SiderealStream, None] = None) -> None:
+    def setup(self, sstream_ref: containers.SiderealStream) -> None:
         """Provide the dataset to use in the gradient calculation.
 
         This dataset must have complete sidereal coverage.
@@ -595,9 +593,7 @@ def setup(self, sstream_ref: Union[containers.SiderealStream, None] = None) -> N
         """
         self.sstream_ref = sstream_ref
 
-    def process(
-        self, sstream: containers.SiderealStreamRebin
-    ) -> containers.SiderealStreamRebin:
+    def process(self, sstream: containers.SiderealStream) -> containers.SiderealStream:
         """Apply the gradient correction to the input dataset.
 
         Parameters
@@ -613,26 +609,22 @@ def process(
         # Allows a normal sidereal stream to pass through this task.
         # Helpful for creating generic configs
         try:
-            sera = sstream.datasets["effective_ra"][:].local_array
+            sera = sstream.effective_ra[:].local_array
         except KeyError:
             self.log.info(
                 f"Dataset of type ({type(sstream)}) does not have an effective "
                 "ra dataset. No correction will be applied."
             )
             return sstream
 
-        # If the reference stream was never set, use the data stream
-        if self.sstream_ref is None:
-            self.sstream_ref = sstream
-
         self.sstream_ref.redistribute("freq")
         sstream.redistribute("freq")
 
         try:
             # If the reference dataset has an effective ra dataset, use this
             # when calculating the gradient. This could be true if the reference
             # and target datasets are the same
-            ref_ra = self.sstream_ref.datasets["effective_ra"][:].local_array
+            ref_ra = self.sstream_ref.effective_ra[:].local_array
         except KeyError:
             # Use fixed ra, which should be regularly sampled
             ref_ra = self.sstream_ref.ra
@@ -643,37 +635,32 @@ def process(
         ref_vis = self.sstream_ref.vis[:].local_array
         ref_weight = self.sstream_ref.weight[:].local_array
 
+        # Iterate over frequencies and baselines for memory
         for fi in range(vis.shape[0]):
             # Skip if entirely masked already
             if not np.any(weight[fi]):
                 continue
 
-            # Depends on whether using effective or grid ra in reference
-            rra = ref_ra[fi] if ref_ra.ndim > 1 else ref_ra
-
             for vi in range(vis.shape[1]):
                 # Skip if entire baseline is masked
                 if not np.any(weight[fi, vi]):
                     continue
 
                 # Depends on whether the effective ra has baseline dependence
-                rra_i = rra[vi] if rra.ndim > 1 else rra
+                rra = ref_ra[fi, vi] if ref_ra.ndim > 1 else ref_ra
 
                 # Calculate the weighted gradient for each baseline, sampled
                 # at the reference RA points
                 norm = tools.invert_no_zero(ref_weight[fi, vi])
-                grad = norm * np.gradient(ref_vis[fi, vi] * ref_weight[fi, vi], rra_i)
+                grad = norm * np.gradient(ref_vis[fi, vi] * ref_weight[fi, vi], rra)
 
                 # Apply the correction to estimate the sample value at the
                 # RA bin centre
                 vis[fi, vi] -= grad * (sera[fi, vi] - sstream.ra)
 
-        # Since the correction has been applied, remove the `effective_ra`
-        # dataset to save memory. Also, remove the reference stream object
-        # to avoid having up to three copies of the dataset in memory
-        del self.sstream_ref
-
-        return containers.SiderealStream(copy_from=sstream)
+        # Return a shallow copy of the container without
+        # the `effective_ra` dataset
+        return sstream.copy(shallow=True, skipped=["effective_ra"])
 
 
 class SiderealStacker(task.SingleTask):