From 7f8b49890db0a1757d5205fcb3991a5bd4787771 Mon Sep 17 00:00:00 2001
From: Liam Gray <lgray@phas.ubc.ca>
Date: Thu, 2 May 2024 15:22:34 -0700
Subject: [PATCH] feat(sidereal): update SiderealStackerMatch to track
 effective ra bin centres

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

diff --git a/draco/analysis/sidereal.py b/draco/analysis/sidereal.py
index 858899da0..8af533f28 100644
--- a/draco/analysis/sidereal.py
+++ b/draco/analysis/sidereal.py
@@ -851,6 +851,7 @@ class SiderealStackerMatch(task.SingleTask):
 
     stack = None
     lsd_list = None
+    ra_correction = False
 
     tag = config.Property(proptype=str, default="stack")
 
@@ -864,6 +865,13 @@ def process(self, sdata):
         sdata : containers.SiderealStream
             Individual sidereal day to stack up.
         """
+        # Check that the input container is of the correct type
+        if (self.stack is not None) and not isinstance(sdata, type(self.stack)):
+            raise TypeError(
+                f"type(sdata) (={type(sdata)}) does not match "
+                f"type(stack) (={type(self.stack)})."
+            )
+
         sdata.redistribute("freq")
 
         if self.stack is None:
@@ -874,6 +882,9 @@ def process(self, sdata):
             self.stack.vis[:] = 0.0
             self.stack.weight[:] = 0.0
 
+            if "effective_ra" in self.stack.datasets:
+                self.ra_correction = True
+
             self.count = 0
             self.Ni_s = mpiarray.zeros(
                 (sdata.weight.shape[0], sdata.weight.shape[2]),
@@ -921,6 +932,20 @@ def process(self, sdata):
         # We need to keep the projection vector until the end
         self.Vm.append(v)
 
+        if self.ra_correction:
+            # Track the effective ra bin centres
+            rebin_weight = sdata.datasets["rebin_weight"][:].local_array
+            sum_rebin_weight = self.stack.datasets["rebin_weight"][:].local_array
+
+            effective_ra = sdata.datasets["effective_ra"][:].local_array
+            sum_effective_ra = self.stack.datasets["effective_ra"][:].local_array
+
+            # Accumulate the total rebin weight
+            sum_rebin_weight[:] += rebin_weight
+
+            delta = rebin_weight * (effective_ra - sum_effective_ra)
+            sum_effective_ra[:] += delta * tools.invert_no_zero(sum_rebin_weight)
+
         # 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.