From 25b2d2e984feaa36c34a4df90f0ef326fbbf3795 Mon Sep 17 00:00:00 2001 From: Melanie Clarke Date: Fri, 22 Nov 2024 09:50:03 -0500 Subject: [PATCH] Better npixels estimate for optimal extraction; fix background fit test --- jwst/extract_1d/extract1d.py | 9 ++++++--- jwst/extract_1d/tests/test_extract_src_flux.py | 4 ++-- .../tests/test_fit_background_model.py | 17 +++++------------ 3 files changed, 13 insertions(+), 17 deletions(-) diff --git a/jwst/extract_1d/extract1d.py b/jwst/extract_1d/extract1d.py index 73e2a690a6..2cb3da87b5 100644 --- a/jwst/extract_1d/extract1d.py +++ b/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 diff --git a/jwst/extract_1d/tests/test_extract_src_flux.py b/jwst/extract_1d/tests/test_extract_src_flux.py index 4c50636ec0..041026bbd3 100644 --- a/jwst/extract_1d/tests/test_extract_src_flux.py +++ b/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 diff --git a/jwst/extract_1d/tests/test_fit_background_model.py b/jwst/extract_1d/tests/test_fit_background_model.py index dd2d1cf53b..7079363bb2 100644 --- a/jwst/extract_1d/tests/test_fit_background_model.py +++ b/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)