Fix add_masked_image to properly handle weight arrays.

lsst · Jul 8, 2024 · 3873452 · 3873452
1 parent 4e26c3a
commit 3873452
Showing 1 changed file with 19 additions and 7 deletions.
diff --git a/python/lsst/meas/algorithms/accumulator_mean_stack.py b/python/lsst/meas/algorithms/accumulator_mean_stack.py
@@ -27,7 +27,7 @@
 __all__ = ['AccumulatorMeanStack']
 
 
-class AccumulatorMeanStack(object):
+class AccumulatorMeanStack:
     """Stack masked images.
 
     Parameters
@@ -106,25 +106,37 @@ def add_masked_image(self, masked_image, weight=1.0):
         good_pixels = np.where(((masked_image.mask.array & self.bit_mask_value) == 0)
                                & np.isfinite(masked_image.mask.array))
 
-        self.sum_weight[good_pixels] += weight
-        self.sum_wdata[good_pixels] += weight*masked_image.image.array[good_pixels]
+        weight_array = False
+        if isinstance(weight, np.ndarray):
+            if weight.shape != masked_image.mask.array.shape:
+                raise RuntimeError("Weight image must be the same shape as masked_image.mask.array")
+            _weight = weight[good_pixels]
+            weight_array = True
+        else:
+            _weight = weight
+
+        self.sum_weight[good_pixels] += _weight
+        self.sum_wdata[good_pixels] += _weight*masked_image.image.array[good_pixels]
 
         if self.compute_n_image:
             self.n_image[good_pixels] += 1
 
         if self.calc_error_from_input_variance:
-            self.sum_w2var[good_pixels] += (weight**2.)*masked_image.variance.array[good_pixels]
+            self.sum_w2var[good_pixels] += (_weight**2.)*masked_image.variance.array[good_pixels]
         else:
-            self.sum_weight2[good_pixels] += weight**2.
-            self.sum_wdata2[good_pixels] += weight*(masked_image.image.array[good_pixels]**2.)
+            self.sum_weight2[good_pixels] += _weight**2.
+            self.sum_wdata2[good_pixels] += _weight*(masked_image.image.array[good_pixels]**2.)
 
         # Mask bits are propagated for good pixels
         self.or_mask[good_pixels] |= masked_image.mask.array[good_pixels]
 
         # Bad pixels are only tracked if they cross a threshold
         for bit in self.mask_threshold_dict:
             bad_pixels = ((masked_image.mask.array & 2**bit) > 0)
-            self.rejected_weights_by_bit[bit][bad_pixels] += weight
+            if weight_array:
+                self.rejected_weights_by_bit[bit][bad_pixels] += weight[bad_pixels]
+            else:
+                self.rejected_weights_by_bit[bit][bad_pixels] += weight
             self.masked_pixels_mask[bad_pixels] |= 2**bit
 
     def fill_stacked_masked_image(self, stacked_masked_image):