Fix indexing errors, add tests

jwst/background/background_sub.py

@@ -91,24 +91,24 @@ def get_subset_array(self, other):
         if jmax < jmin:
             jmax = jmin
 
-        # To ensure that we mix and match subarray obs, take
+        # To ensure that we can mix and match subarray obs, take
         # the x/y shape from self.data. To ensure we can work
         # with mismatched NINTS, if we have that third dimension
         # use the value from other
-        data_shape = copy.deepcopy(self.data.shape)
+        data_shape = list(self.data.shape)
         if self.im_dim == 3:
             data_shape[0] = other.data.shape[0]
 
         # Set up arrays, NaN out data/err for sigma clipping, keep DQ as 0 for bitwise_or
-        data_overlap = np.ones(data_shape) * np.nan
-        err_overlap = np.ones(data_shape) * np.nan
+        data_overlap = np.full(data_shape, np.nan, dtype=other.data.dtype)
+        err_overlap = np.full(data_shape, np.nan, dtype=other.data.dtype)
         dq_overlap = np.zeros(data_shape, dtype=np.uint32)
 
         if self.im_dim == 2:
             idx = (slice(jmin - other.jmin, jmax - other.jmin),
                    slice(imin - other.imin, imax - other.imin),
                    )
-        if self.im_dim == 3:
+        else:
             idx = (slice(None, None),
                    slice(jmin - other.jmin, jmax - other.jmin),
                    slice(imin - other.imin, imax - other.imin),
@@ -120,9 +120,9 @@ def get_subset_array(self, other):
         dq_cutout = other.dq[idx]
 
         # Put them into the right place in the original image array shape
-        data_overlap[:data_cutout.shape[0], :data_cutout.shape[1]] = copy.deepcopy(data_cutout)
-        err_overlap[:data_cutout.shape[0], :data_cutout.shape[1]] = copy.deepcopy(err_cutout)
-        dq_overlap[:data_cutout.shape[0], :data_cutout.shape[1]] = copy.deepcopy(dq_cutout)
+        data_overlap[..., :data_cutout.shape[-2], :data_cutout.shape[-1]] = copy.deepcopy(data_cutout)
+        err_overlap[..., :data_cutout.shape[-2], :data_cutout.shape[-1]] = copy.deepcopy(err_cutout)
+        dq_overlap[..., :data_cutout.shape[-2], :data_cutout.shape[-1]] = copy.deepcopy(dq_cutout)
 
         return data_overlap, err_overlap, dq_overlap
 

jwst/background/tests/test_background.py

@@ -80,6 +80,47 @@ def science_image():
     return image
 
 
+def miri_rate_model(data_shape, value=1.0):
+    """
+    Generate a MIRI image subarray rate or rateints image.
+
+    Parameters
+    ----------
+    data_shape : tuple of int
+        Shape of the rate data. 2 values for rate, 3 for rateints.
+    value : float, optional
+        Value to set in the data array.
+
+    Returns
+    -------
+    image : DataModel
+        An open datamodel containing MIRI subarray rate or rateints
+        data.
+    """
+
+    if len(data_shape) == 2:
+        image = datamodels.ImageModel(data_shape)
+    else:
+        image = datamodels.CubeModel(data_shape)
+
+    image.data[:, :] = value
+    image.meta.instrument.name = 'MIRI'
+    image.meta.instrument.detector = 'MIRIMAGE'
+    image.meta.instrument.filter = 'F2100W'
+    image.meta.exposure.type = 'MIR_IMAGE'
+    image.meta.observation.date = '2019-02-27'
+    image.meta.observation.time = '13:37:18.548'
+    image.meta.date = '2019-02-27T13:37:18.548'
+
+    image.meta.subarray.xstart = 1
+    image.meta.subarray.ystart = 1
+
+    image.meta.subarray.xsize = image.data.shape[-1]
+    image.meta.subarray.ysize = image.data.shape[-2]
+
+    return image
+
+
 def test_nirspec_gwa(tmp_cwd, background, science_image):
     """Verify NIRSPEC GWA logic for in the science and background"""
 
@@ -281,6 +322,55 @@ def test_nis_wfss_background(filters, pupils, make_wfss_datamodel):
         assert np.isclose([pipeline_reference_mean], [test_reference_mean], rtol=1e-1)
 
 
+@pytest.mark.parametrize('data_shape,background_shape',
+                         [((10, 10), (10, 10)),
+                          ((10, 10), (20, 20)),
+                          ((2, 10, 10), (2, 10, 10)),
+                          ((2, 10, 10), (2, 20, 20)),
+                          ((2, 10, 10), (3, 10, 10)),
+                          ((2, 10, 10), (3, 20, 20)),
+                          ((3, 10, 10), (2, 10, 10)),
+                          ((3, 10, 10), (2, 20, 20))])
+def test_miri_subarray_full_overlap(data_shape, background_shape):
+    image_value = 10.0
+    background_value = 1.0
+    image = miri_rate_model(data_shape, value=image_value)
+    background = miri_rate_model(background_shape, value=background_value)
+
+    result = BackgroundStep.call(image, [background])
+
+    assert_allclose(result.data, image_value - background_value)
+    assert type(result) is type(image)
+    assert result.meta.cal_step.back_sub == 'COMPLETE'
+
+    image.close()
+    background.close()
+
+
+@pytest.mark.parametrize('data_shape,background_shape',
+                         [((20, 20), (10, 10)),
+                          ((2, 20, 20), (2, 10, 10),),
+                          ((3, 20, 20), (2, 10, 10),),
+                          ((2, 20, 20), (3, 10, 10),)])
+def test_miri_subarray_partial_overlap(data_shape, background_shape):
+    image_value = 10.0
+    background_value = 1.0
+    image = miri_rate_model(data_shape, value=image_value)
+    background = miri_rate_model(background_shape, value=background_value)
+
+    result = BackgroundStep.call(image, [background])
+
+    assert_allclose(result.data[..., :background_shape[-2], :background_shape[-1]],
+                    image_value - background_value)
+    assert_allclose(result.data[..., background_shape[-2]:, :], image_value)
+    assert_allclose(result.data[..., :, background_shape[-1]:], image_value)
+    assert type(result) is type(image)
+    assert result.meta.cal_step.back_sub == 'COMPLETE'
+
+    image.close()
+    background.close()
+
+
 def test_robust_mean():
     """Test robust mean calculation"""
     data = np.random.rand(2048, 2048)
@@ -290,7 +380,7 @@ def test_robust_mean():
     assert np.isclose([test], [result], rtol=1e-3)
 
 
-def test_no_Nan():
+def test_no_nan():
     """Make sure that nan values are filled with fill value"""
     # Make data model
     model = datamodels.ImageModel()