From c88f093479890eae448646d281051e86f05f08d3 Mon Sep 17 00:00:00 2001 From: Clare Shanahan Date: Thu, 7 Oct 2021 09:53:43 -0400 Subject: [PATCH 1/2] common code for linearity correction --- CHANGES.rst | 8 ++ src/stcal/linearity/__init__.py | 0 src/stcal/linearity/linearity.py | 199 +++++++++++++++++++++++++++++++ tests/test_linearity.py | 88 ++++++++++++++ 4 files changed, 295 insertions(+) create mode 100644 src/stcal/linearity/__init__.py create mode 100644 src/stcal/linearity/linearity.py create mode 100644 tests/test_linearity.py diff --git a/CHANGES.rst b/CHANGES.rst index 13998b787..41bb5d5ee 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -1,3 +1,11 @@ +0.4.0(unreleased) +================== + +linearity +--------- + +- Adds common code for linearity correction [#55] + 0.3.0 (2021-09-28) ================== diff --git a/src/stcal/linearity/__init__.py b/src/stcal/linearity/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/stcal/linearity/linearity.py b/src/stcal/linearity/linearity.py new file mode 100644 index 000000000..dff00b4cb --- /dev/null +++ b/src/stcal/linearity/linearity.py @@ -0,0 +1,199 @@ +import numpy as np + + +def linearity_correction(data, gdq, pdq, lin_coeffs, lin_dq, dqflags): + """ + Apply linearity correction to individual groups in `data` to pixels that + haven't already been flagged as saturated. + + Pixels having at least one correction coefficient equal to NaN will not + have the linearity correction applied and the DQ flag `NO_LIN_CORR` is + added to the science exposure PIXELDQ array. Pixels that have the + `NO_LIN_CORR` flag set in the DQ array of the linearity reference file will + not have the correction applied and the “NO_LIN_CORR” flag is added to the + science exposure `PIXELDQ` array. Pixel values that have the `SATURATED` + flag set in a particular group of the science exposure GROUPDQ array will + not have the linearity correction applied to that group. Any groups for + that pixel that are not flagged as saturated will be corrected. + + The corrected data and updated `PIXELDQ` arrays are returned. + Reference: + http://www.stsci.edu/hst/wfc3/documents/handbooks/currentDHB/wfc3_Ch606.html + WFC3 Data Handbook, Chapter 6, WFC3-IR Error Sources, 6.5 Detector + Nonlinearity issues, 2.1 May 2011 + + Parameters + ---------- + data: `np.array` + Array of correction coefficients in reference file. + + gdq: `np.array` + Group dq array. + + pdq: `np.array` + Pixel dq array. + + lin_coeffs: `np.array` + Array containing pixel-by-pixel linearity coefficient values + or each term in the polynomial fit. + + dq_flags: dict + A dictionary with at least the following keywords: + SATURATED, NO_LIN_CORR. + + Returns + ------- + data: 3D array + Updated array of correction coefficients in reference file. + + new_pdq: `np.array` + Updated pixel data quality flags. + + """ + + # Retrieve the ramp data cube characteristics + nints, ngroups, nrows, ncols = data.shape + + # Number of coeffs is equal to the number of planes in coeff cube + ncoeffs = lin_coeffs.shape[0] + + # Combine the DQ arrays using bitwise_or + new_pdq = np.bitwise_or(pdq, lin_dq) + + # Check for NO_LIN_CORR flags in the DQ extension of the ref file + lin_coeffs = correct_for_flag(lin_coeffs, lin_dq, dqflags) + + # Check for NaNs in the COEFFS extension of the ref file + lin_coeffs, new_pdq = correct_for_NaN(lin_coeffs, new_pdq, dqflags) + + # Apply the linearity correction one integration at a time. + for ints in range(nints): + + # Apply the linearity correction one group at a time + for plane in range(ngroups): + + # Accumulate the polynomial terms into the corrected counts + scorr = lin_coeffs[ncoeffs - 1] * data[ints, plane] + for j in range(ncoeffs - 2, 0, -1): + scorr = (scorr + lin_coeffs[j]) * data[ints, plane] + scorr = lin_coeffs[0] + scorr + + # Only use the corrected signal where the original signal value + # has not been flagged by the saturation step. + # Otherwise use the original signal. + data[ints, plane, :, :] = \ + np.where(np.bitwise_and(gdq[ints, plane, :, :], dqflags['SATURATED']), + data[ints, plane, :, :], scorr) + + return data, new_pdq + + +def correct_for_NaN(lin_coeffs, pixeldq, dqflags): + """ + Check for NaNs in the COEFFS extension of the ref file in case there are + pixels that should have been (but were not) flagged there as NO_LIN_CORR + (linearity correction not determined for pixel). + + For such pixels, update the + coefficients so that there is effectively no correction, and flag their + pixeldq values in place as NO_LIN_CORR in the step output. + + Parameters + ---------- + lin_coeffs: 3D array + array of correction coefficients in reference file + + input: data model object + science data model to be corrected in place + + Returns + ------- + lin_coeffs: 3D array + updated array of correction coefficients in reference file + """ + + wh_nan = np.where(np.isnan(lin_coeffs)) + znan, ynan, xnan = wh_nan[0], wh_nan[1], wh_nan[2] + num_nan = 0 + + nan_array = np.zeros((lin_coeffs.shape[1], lin_coeffs.shape[2]), + dtype=np.uint32) + + # If there are NaNs as the correction coefficients, update those + # coefficients so that those SCI values will be unchanged. + if len(znan) > 0: + ben_cor = ben_coeffs(lin_coeffs) # get benign coefficients + num_nan = len(znan) + + for ii in range(num_nan): + lin_coeffs[:, ynan[ii], xnan[ii]] = ben_cor + nan_array[ynan[ii], xnan[ii]] = dqflags['NO_LIN_CORR'] + + # Include these pixels in the output pixeldq + pixeldq = np.bitwise_or(pixeldq, nan_array) + + return lin_coeffs, pixeldq + + +def correct_for_flag(lin_coeffs, lin_dq, dqflags): + """ + Short Summary + ------------- + Check for pixels that are flagged as NO_LIN_CORR + ('No linearity correction available') in the DQ extension of the ref data. + For such pixels, update the coefficients so that there is effectively no + correction. Because these are already flagged in the ref file, they will + also be flagged in the output dq. + + Parameters + ---------- + lin_coeffs: 3D array + array of correction coefficients in reference file + + lin_dq: 2D array + array of data quality flags in reference file + + Returns + ------- + lin_coeffs: 3D array + updated array of correction coefficients in reference file + """ + + wh_flag = np.bitwise_and(lin_dq, dqflags['NO_LIN_CORR']) + num_flag = len(np.where(wh_flag > 0)[0]) + + wh_lin = np.where(wh_flag == dqflags['NO_LIN_CORR']) + yf, xf = wh_lin[0], wh_lin[1] + + # If there are pixels flagged as 'NO_LIN_CORR', update the corresponding + # coefficients so that those SCI values will be unchanged. + if (num_flag > 0): + ben_cor = ben_coeffs(lin_coeffs) # get benign coefficients + + for ii in range(num_flag): + lin_coeffs[:, yf[ii], xf[ii]] = ben_cor + + return lin_coeffs + + +def ben_coeffs(lin_coeffs): + """ + Short Summary + ------------- + For pixels having at least 1 NaN coefficient, reset the coefficients to be + benign, which will effectively leave the SCI values unaffected. + + Parameters + ---------- + lin_coeffs: 3D array + array of correction coefficients in reference file + + Returns + ------- + ben_cor: 1D array + benign coefficients - all ben_cor[:] = 0.0 except ben_cor[1] = 1.0 + """ + ben_cor = np.zeros(lin_coeffs.shape[0]) + ben_cor[1] = 1.0 + + return ben_cor diff --git a/tests/test_linearity.py b/tests/test_linearity.py new file mode 100644 index 000000000..5937c54c0 --- /dev/null +++ b/tests/test_linearity.py @@ -0,0 +1,88 @@ +""" + +Unit tests for linearity correction + +""" + +import numpy as np +import pytest + +from stcal.linearity.linearity import linearity_correction + +DQFLAGS = {'GOOD': 0, 'DO_NOT_USE': 1, 'SATURATED': 2, 'DEAD': 1024, 'HOT': 2048, 'NO_LIN_CORR': 1048576} + + +def test_coeff_dq(): + """Test linearity algorithm with random data ramp (does algorithm match expected algorithm) + also test a variety of dq flags and expected output """ + + # size of integration + nints = 1 + ngroups = 160 + xsize = 103 + ysize = 102 + + # Create data array and group/pixel dq arrays + data = np.ones((nints, ngroups, ysize, xsize)).astype(np.float32) + pdq = np.zeros((ysize, xsize)).astype(np.uint32) + gdq = np.zeros((nints, ngroups, ysize, xsize)).astype(np.uint32) + + # Create reference file data and dq arrays + numcoeffs = 5 + lin_coeffs = np.zeros((numcoeffs, ysize, xsize)) + + # Set coefficient values in reference file to check the algorithm + # Equation is DNcorr = L0 + L1*DN(i) + L2*DN(i)^2 + L3*DN(i)^3 + L4*DN(i)^4 + # DN(i) = signal in pixel, Ln = coefficient from ref file + # L0 = 0 for all pixels for CDP6 + L0 = 0.0e+00 + L1 = 0.85 + L2 = 4.62E-6 + L3 = -6.16E-11 + L4 = 7.23E-16 + + coeffs = np.asfarray([L0, L1, L2, L3, L4]) + lin_coeffs[:, 30, 50] = coeffs + lin_dq = np.zeros((ysize, xsize), dtype=np.uint32) + + # check behavior with NaN coefficients: should not alter pixel values + coeffs2 = np.asfarray([L0, np.nan, L2, L3, L4]) + lin_coeffs[:, 20, 50] = coeffs2 + data[0, 50, 20, 50] = 500.0 + + tgroup = 2.775 + + # set pixel values (DN) for specific pixels up the ramp + data[0, :, 30, 50] = np.arange(ngroups) * 100 * tgroup + + scival = 40000.0 + data[0, 45, 30, 50] = scival # to check linearity multiplication is done correctly + data[0, 30, 35, 36] = 35 # pixel to check that dq=2 meant no correction was applied + + # check if dq flags in pixeldq are correctly populated in output + pdq[50, 40] = DQFLAGS['DO_NOT_USE'] + pdq[50, 41] = DQFLAGS['SATURATED'] + pdq[50, 42] = DQFLAGS['DEAD'] + pdq[50, 43] = DQFLAGS['HOT'] + + # set dq flags in DQ of reference file + lin_dq[35, 35] = DQFLAGS['DO_NOT_USE'] + lin_dq[35, 36] = DQFLAGS['NO_LIN_CORR'] + lin_dq[30, 50] = DQFLAGS['GOOD'] + + np.bitwise_or(pdq, lin_dq) + + # run linearity correction + output_data, output_pdq = linearity_correction(data, gdq, pdq, lin_coeffs, lin_dq, DQFLAGS) + + # check that multiplication of polynomial was done correctly for specified pixel + outval = L0 + (L1 * scival) + (L2 * scival**2) + (L3 * scival**3) + (L4 * scival**4) + assert(np.isclose(output_data[0, 45, 30, 50], outval, rtol=0.00001)) + + # check that dq value was handled correctly + assert output_pdq[35, 35] == DQFLAGS['DO_NOT_USE'] + assert output_pdq[35, 36] == DQFLAGS['NO_LIN_CORR'] + # NO_LIN_CORR, sci value should not change + assert output_data[0, 30, 35, 36] == 35 + # NaN coefficient should not change data value + assert output_data[0, 50, 20, 50] == 500.0 From 2d72add8db0b2cd98e17b2926524fe753ab47727 Mon Sep 17 00:00:00 2001 From: Clare Shanahan Date: Tue, 12 Oct 2021 10:09:01 -0400 Subject: [PATCH 2/2] removed unused import --- tests/test_linearity.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_linearity.py b/tests/test_linearity.py index 5937c54c0..e054830e6 100644 --- a/tests/test_linearity.py +++ b/tests/test_linearity.py @@ -5,7 +5,6 @@ """ import numpy as np -import pytest from stcal.linearity.linearity import linearity_correction