Better npixels estimate for optimal extraction; fix background fit test

jwst/extract_1d/extract1d.py

@@ -401,9 +401,12 @@ def extract1d(image, profiles_2d, variance_rn, variance_phnoise, variance_flat,
 
         coefs = np.nansum(pixwgt * image.T[:, :, np.newaxis], axis=1)
 
-        # Number of contributing pixels at each wavelength for each source.
-
-        npixels = np.sum(pixwgt[..., -nobjects:] != 0, axis=1).T
+        # Effective number of contributing pixels at each wavelength for each source.
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=RuntimeWarning, message="invalid value")
+            wgt_src_pix = [profiles_2d[i] * weights / np.sum(profiles_2d[i] ** 2, axis=0)
+                           for i in range(nobjects)]
+        npixels = np.sum(wgt_src_pix, axis=1)
 
         if order > -1:
             bkg_2d = np.sum(coefs[:, np.newaxis, :order + 1] * coefmatrix[..., :order + 1], axis=-1).T

jwst/extract_1d/tests/test_extract_src_flux.py

@@ -124,7 +124,7 @@ def test_extract_optimal(inputs_with_source):
     # Total flux should be well modeled
     assert np.allclose(total_flux[0], 20)
     assert np.allclose(bkg_flux[0], 0)
-    assert np.allclose(npixels[0], 3)
+    assert np.allclose(npixels[0], 2.66667)
 
     # set a NaN value in a column of interest
     image[4, 2] = np.nan
@@ -137,7 +137,7 @@ def test_extract_optimal(inputs_with_source):
 
     # Total flux is still well modeled from 2 pixels
     assert np.allclose(total_flux[0], 20)
-    assert npixels[0, 2] == 2
+    assert np.isclose(npixels[0, 2], 1.33333)
 
     # set the whole column to NaN
     image[:, 2] = np.nan

jwst/extract_1d/tests/test_fit_background_model.py

@@ -32,9 +32,9 @@ def inputs_constant():
 def inputs_with_source():
     shape = (9, 5)
     image = np.full(shape, 1.0)
-    image[3] = 5.0
-    image[4] = 10.0
-    image[5] = 5.0
+    image[3] += 5.0
+    image[4] += 10.0
+    image[5] += 5.0
 
     var_rnoise = image * 0.05
     var_poisson = image * 0.05
@@ -217,15 +217,8 @@ def test_fit_background_optimal(inputs_with_source, bkg_order_val):
         bkg_fit_type='poly', bkg_order=bkg_order_val,
         extraction_type='optimal')
 
-    names = ('total_flux', 'f_var_rnoise', 'f_var_poisson', 'f_var_flat',
-             'bkg_flux', 'b_var_rnoise', 'b_var_poisson', 'b_var_flat',
-             'npixels', 'model')
-    for name, data in zip(names, result):
-        print(name, data)
-
     flux = result[0][0]
     background = result[4][0]
 
-    # this should be exact, not sure why background fit is off
-    assert np.allclose(flux, 20.0 - 1.0 * 3, atol=1.0)
-    assert np.allclose(background, 3.0, atol=1.0)
+    assert np.allclose(flux, 20.0)
+    assert np.allclose(background, 2.66667)