From 2e28281b4767991ba088e73e763c06c5b28a6c52 Mon Sep 17 00:00:00 2001 From: William Jamieson Date: Wed, 4 Oct 2023 14:16:36 -0400 Subject: [PATCH] Move array building to cython --- .../ramp_fitting/ols_cas22/_fit_ramps.pyx | 18 +++++++++++--- src/stcal/ramp_fitting/ols_cas22_fit.py | 14 +---------- tests/test_jump_cas22.py | 24 ++++++++++++++++--- 3 files changed, 37 insertions(+), 19 deletions(-) diff --git a/src/stcal/ramp_fitting/ols_cas22/_fit_ramps.pyx b/src/stcal/ramp_fitting/ols_cas22/_fit_ramps.pyx index 1a5263527..46eab337a 100644 --- a/src/stcal/ramp_fitting/ols_cas22/_fit_ramps.pyx +++ b/src/stcal/ramp_fitting/ols_cas22/_fit_ramps.pyx @@ -69,11 +69,23 @@ def fit_ramps(np.ndarray[float, ndim=2] resultants, # list in the end. cdef cpp_list[RampFits] ramp_fits + cdef np.ndarray[float, ndim=2] parameters = np.zeros((n_pixels, 2), dtype=np.float32) + cdef np.ndarray[float, ndim=2] variances = np.zeros((n_pixels, 3), dtype=np.float32) + # Perform all of the fits + cdef RampFits fit cdef int index for index in range(n_pixels): # Fit all the ramps for the given pixel - ramp_fits.push_back(make_pixel(fixed, read_noise[index], - resultants[:, index]).fit_ramps(pixel_ramps[index])) + fit = make_pixel(fixed, read_noise[index], + resultants[:, index]).fit_ramps(pixel_ramps[index]) + + parameters[index, 1] = fit.average.slope + + variances[index, 0] = fit.average.read_var + variances[index, 1] = fit.average.poisson_var + variances[index, 2] = fit.average.read_var + fit.average.poisson_var + + ramp_fits.push_back(fit) - return ramp_fits + return ramp_fits, parameters, variances diff --git a/src/stcal/ramp_fitting/ols_cas22_fit.py b/src/stcal/ramp_fitting/ols_cas22_fit.py index 97c63c790..974c5d4e7 100644 --- a/src/stcal/ramp_fitting/ols_cas22_fit.py +++ b/src/stcal/ramp_fitting/ols_cas22_fit.py @@ -98,7 +98,7 @@ def fit_ramps_casertano(resultants, dq, read_noise, read_time, ma_table=None, re dq = dq.reshape(origshape + (1,)) read_noise = read_noise.reshape(origshape[1:] + (1,)) - ramp_fits = ols_cas22.fit_ramps( + ramp_fits, parameters, variances = ols_cas22.fit_ramps( resultants.reshape(resultants.shape[0], -1), dq.reshape(resultants.shape[0], -1), read_noise.reshape(-1), @@ -106,18 +106,6 @@ def fit_ramps_casertano(resultants, dq, read_noise, read_time, ma_table=None, re read_pattern, use_jump) - parameters = np.zeros((len(ramp_fits), 2), dtype=np.float32) - variances = np.zeros((len(ramp_fits), 3), dtype=np.float32) - - # Extract the data request from the ramp fits - for index, ramp_fit in enumerate(ramp_fits): - parameters[index, 1] = ramp_fit['average']['slope'] - - variances[index, 0] = ramp_fit['average']['read_var'] - variances[index, 1] = ramp_fit['average']['poisson_var'] - - variances[:, 2] = (variances[:, 0] + variances[:, 1]).astype(np.float32) - if resultants.shape != orig_shape: parameters = parameters[0] variances = variances[0] diff --git a/tests/test_jump_cas22.py b/tests/test_jump_cas22.py index 825137258..5849e2c32 100644 --- a/tests/test_jump_cas22.py +++ b/tests/test_jump_cas22.py @@ -249,6 +249,24 @@ def detector_data(ramp_data): return resultants, read_noise, read_pattern, N_PIXELS, FLUX +@pytest.mark.parametrize("use_jump", [True, False]) +def test_fit_ramps_array_outputs(detector_data, use_jump): + """ + Test that the array outputs line up with the dictionary output + """ + resultants, read_noise, read_pattern, n_pixels, flux = detector_data + dq = np.zeros(resultants.shape, dtype=np.int32) + + fits, parameters, variances = fit_ramps(resultants, dq, read_noise, ROMAN_READ_TIME, read_pattern, use_jump=use_jump) + + for fit, par, var in zip(fits, parameters, variances): + assert par[0] == 0 + assert par[1] == fit['average']['slope'] + + assert var[0] == fit['average']['read_var'] + assert var[1] == fit['average']['poisson_var'] + assert var[2] == np.float32(fit['average']['read_var'] + fit['average']['poisson_var']) + @pytest.mark.parametrize("use_jump", [True, False]) def test_fit_ramps_no_dq(detector_data, use_jump): @@ -260,7 +278,7 @@ def test_fit_ramps_no_dq(detector_data, use_jump): resultants, read_noise, read_pattern, n_pixels, flux = detector_data dq = np.zeros(resultants.shape, dtype=np.int32) - fits = fit_ramps(resultants, dq, read_noise, ROMAN_READ_TIME, read_pattern, use_jump=use_jump) + fits, _, _ = fit_ramps(resultants, dq, read_noise, ROMAN_READ_TIME, read_pattern, use_jump=use_jump) assert len(fits) == n_pixels # sanity check that a fit is output for each pixel # Check that the chi2 for the resulting fit relative to the assumed flux is ~1 @@ -291,7 +309,7 @@ def test_fit_ramps_dq(detector_data, use_jump): # i.e., we can make a measurement from them. okay = np.sum((dq[1:, :] == 0) & (dq[:-1, :] == 0), axis=0) != 0 - fits = fit_ramps(resultants, dq, read_noise, ROMAN_READ_TIME, read_pattern, use_jump=use_jump) + fits, _, _ = fit_ramps(resultants, dq, read_noise, ROMAN_READ_TIME, read_pattern, use_jump=use_jump) assert len(fits) == n_pixels # sanity check that a fit is output for each pixel chi2 = 0 @@ -353,7 +371,7 @@ def test_find_jumps(jump_data): resultants, read_noise, read_pattern, n_pixels, jumps = jump_data dq = np.zeros(resultants.shape, dtype=np.int32) - fits = fit_ramps(resultants, dq, read_noise, ROMAN_READ_TIME, read_pattern, use_jump=True) + fits, _, _ = fit_ramps(resultants, dq, read_noise, ROMAN_READ_TIME, read_pattern, use_jump=True) # Check that all the jumps have been located per the algorithm's constraints for index, (fit, jump) in enumerate(zip(fits, jumps)):