JP-2152: Refactor resample_variance_array to use less memory (#6251)

* Refactor resample_variance_array to use less memory

* Resample variance, not inverse variance

* Ignore divide by zero

* Add unit test

(cherry picked from commit 446cc4f)

jwst/resample/resample.py

@@ -1,6 +1,5 @@
 import logging
 import re
-import warnings
 
 import numpy as np
 from drizzle import util
@@ -188,48 +187,42 @@ def resample_many_to_one(self):
     def resample_variance_array(self, name, output_model):
         """Resample variance arrays from self.input_models to the output_model
 
-        Loop through self.input_models and resample the `name` variance array
-        to the same name in output_model.  This modifies output_model in-place.
+        Resample the `name` variance array to the same name in output_model,
+        using a cummulative sum.
+
+        This modifies output_model in-place.
         """
-        outwht = np.zeros_like(output_model.data)
-        outcon = np.zeros_like(output_model.con)
         output_wcs = output_model.meta.wcs
-        output_arrays = []
+        inverse_variance_sum = np.zeros_like(output_model.data)
 
         log.info(f"Resampling {name}")
         for model in self.input_models:
-            # Unity weight but ignore pixels that are NON_SCIENCE or REFERENCE_PIXEL
+            variance = getattr(model, name)
+            # Make input weight map of unity where there is science data
             inwht = resample_utils.build_driz_weight(model, weight_type=None,
                                                      good_bits="~NON_SCIENCE+REFERENCE_PIXEL")
-            input_array = getattr(model, name)
-            output_array = np.zeros_like(output_model.data)
 
-            # Resample the variance array
-            self.drizzle_arrays(input_array, inwht, model.meta.wcs,
-                                output_wcs, output_array, outwht, outcon,
+            resampled_variance = np.zeros_like(output_model.data)
+            outwht = np.zeros_like(output_model.data)
+            outcon = np.zeros_like(output_model.con)
+
+            # Resample the variance array.  Use fillval=np.inf so that when we
+            # take the reciprocal for summing, it is zero where there is zero weight
+            self.drizzle_arrays(variance, inwht, model.meta.wcs,
+                                output_wcs, resampled_variance, outwht, outcon,
                                 pixfrac=self.pixfrac, kernel=self.kernel,
-                                fillval=np.nan)
-            output_arrays.append(output_array)
+                                fillval=np.inf)
 
-        # Stack and coadd as inverse variance if there any to stack
-        if len(output_arrays) > 1:
-            stacked = np.stack(output_arrays)
+            # Add the inverse of the resampled variance to a running sum
             with np.errstate(divide="ignore"):
-                inv_var = np.reciprocal(stacked)
-            inv_var[~np.isfinite(inv_var)] = np.nan
-
-            with warnings.catch_warnings():
-                warnings.filterwarnings(action="ignore",
-                                        message="Mean of empty slice",
-                                        category=RuntimeWarning)
-                inv_var_mean = np.nanmean(inv_var, axis=0)
-            with np.errstate(divide="ignore"):
-                output_err_array = np.reciprocal(inv_var_mean)
-        else:
-            output_err_array = output_arrays[0]
-        output_err_array[~np.isfinite(output_err_array)] = np.nan
+                inverse_variance_sum += np.reciprocal(resampled_variance)
 
-        setattr(output_model, name, output_err_array)
+        # We now have a sum of the inverse resampled variances.  We need the
+        # inverse of that to get back to units of variance.
+        with np.errstate(divide="ignore"):
+            output_variance = np.reciprocal(inverse_variance_sum)
+        output_variance[~np.isfinite(output_variance)] = np.nan
+        setattr(output_model, name, output_variance)
 
     def update_exposure_times(self, output_model):
         """Modify exposure time metadata in-place"""
@@ -249,7 +242,7 @@ def update_exposure_times(self, output_model):
     @staticmethod
     def drizzle_arrays(insci, inwht, input_wcs, output_wcs, outsci, outwht, outcon,
                        uniqid=1, xmin=None, xmax=None, ymin=None, ymax=None,
-                       pixfrac=1.0, kernel='square', fillval="INDEF"):
+                       pixfrac=1.0, kernel='square', fillval="INDEF", wtscale=1.0):
         """
         Low level routine for performing 'drizzle' operation on one image.
 
@@ -408,7 +401,7 @@ def drizzle_arrays(insci, inwht, input_wcs, output_wcs, outsci, outwht, outcon,
             kernel=kernel,
             in_units="cps",
             expscale=1.0,
-            wtscale=1.0,
+            wtscale=wtscale,
             fillstr=fillval
         )
 

jwst/resample/tests/test_resample_step.py

@@ -398,3 +398,27 @@ def test_wcs_keywords(nircam_rate):
     assert result.meta.wcsinfo.roll_ref is None
     assert result.meta.wcsinfo.v3yangle is None
     assert result.meta.wcsinfo.vparity is None
+
+
+@pytest.mark.parametrize("n_images", [1, 2, 3, 9])
+def test_resample_variance(nircam_rate, n_images):
+    """Test that resampled variance and error arrays are computed properly"""
+    err = 0.02429
+    var_rnoise = 0.00034
+    var_poisson = 0.00025
+    im = AssignWcsStep.call(nircam_rate)
+    im.var_rnoise += var_rnoise
+    im.var_poisson += var_poisson
+    im.err += err
+    im.meta.filename = "foo.fits"
+
+    c = ModelContainer()
+    for n in range(n_images):
+        c.append(im.copy())
+
+    result = ResampleStep.call(c, blendheaders=False)
+
+    # Verify that the combined uncertainty goes as 1 / sqrt(N)
+    assert_allclose(result.err[5:-5,5:-5].mean(), err / np.sqrt(n_images), atol=1e-5)
+    assert_allclose(result.var_rnoise[5:-5,5:-5].mean(), var_rnoise / n_images, atol=1e-7)
+    assert_allclose(result.var_poisson[5:-5,5:-5].mean(), var_poisson / n_images, atol=1e-7)