diff --git a/CHANGES.rst b/CHANGES.rst index b465fe6ab..2a89cb16d 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -1,5 +1,12 @@ +0.2.0 (2021/05/18) +================== + +- Added ramp fitting code [#6] + 0.1.0 (2021/03/19) ================== - Added code to manipulate bitmasks. + + diff --git a/setup.cfg b/setup.cfg index 1eb9aded5..fc11ffe32 100644 --- a/setup.cfg +++ b/setup.cfg @@ -21,7 +21,7 @@ setup_requires = setuptools_scm install_requires = - numpy>=1.19 + numpy>=1.17 astropy package_dir = diff --git a/src/stcal/__init__.py b/src/stcal/__init__.py index e69de29bb..b869c0932 100644 --- a/src/stcal/__init__.py +++ b/src/stcal/__init__.py @@ -0,0 +1,4 @@ +from ._version import version as __version__ + + +__all__ = ['__version__'] diff --git a/src/stcal/ramp_fitting/__init__.py b/src/stcal/ramp_fitting/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/stcal/ramp_fitting/constants.py b/src/stcal/ramp_fitting/constants.py new file mode 100644 index 000000000..47b1b4506 --- /dev/null +++ b/src/stcal/ramp_fitting/constants.py @@ -0,0 +1,5 @@ +DO_NOT_USE = 2**0 # Bad pixel. Do not use. +SATURATED = 2**1 # Pixel saturated during exposure. +JUMP_DET = 2**2 # Jump detected during exposure +NO_GAIN_VALUE = 2**19 # Gain cannot be measured +UNRELIABLE_SLOPE = 2**24 # Slope variance large (i.e., noisy pixel) diff --git a/src/stcal/ramp_fitting/gls_fit.py b/src/stcal/ramp_fitting/gls_fit.py new file mode 100644 index 000000000..e0297081e --- /dev/null +++ b/src/stcal/ramp_fitting/gls_fit.py @@ -0,0 +1,1184 @@ +#! /usr/bin/env python + +# !!!!!!!!!!!!!!!!!!! NOTE !!!!!!!!!!!!!!!!!!! +# Needs work. +# Also, this code makes reference to `nreads` as a the second dimension +# of the 4-D data set, while `ngroups` makes reference to the NGROUPS +# key word in the exposure metadata. This should be changed, removing +# reference to the NGROUPS key word and using ngroups as the second +# dimension of the 4-D data set. +# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! + + +''' +import logging +from multiprocessing.pool import Pool as Pool +import numpy as np +import numpy.linalg as la +import time + +from .. import datamodels +from ..datamodels import dqflags + +from . import utils + +log = logging.getLogger(__name__) +log.setLevel(logging.DEBUG) + +DO_NOT_USE = dqflags.group['DO_NOT_USE'] +JUMP_DET = dqflags.group['JUMP_DET'] +SATURATED = dqflags.group['SATURATED'] +UNRELIABLE_SLOPE = dqflags.pixel['UNRELIABLE_SLOPE'] + +BUFSIZE = 1024 * 300000 # 300Mb cache size for data section + + +log = logging.getLogger(__name__) +log.setLevel(logging.DEBUG) + +# This is the number of iterations that will be done with use_extra_terms +# set to False. If this is zero, use_extra_terms will be set to True even +# for the first iteration. +# NUM_ITER_NO_EXTRA_TERMS = 1 +NUM_ITER_NO_EXTRA_TERMS = 0 +# These are the lower and upper limits of the number of iterations that +# will be done by determine_slope. +# MIN_ITER = NUM_ITER_NO_EXTRA_TERMS + 1 +# MAX_ITER = 3 +MIN_ITER = 1 +MAX_ITER = 1 + +# This is a term to add for saturated pixels to give them low weight. +HUGE_FOR_LOW_WEIGHT = 1.e20 + +# This is a value to replace zero or negative values in a fit, to make +# all values of the fit positive and to give low weight where the fit was +# zero or negative. +FIT_MUST_BE_POSITIVE = 1.e10 + + +def gls_ramp_fit(input_model, buffsize, save_opt, readnoise_model, gain_model, max_cores): + """Fit a ramp using generalized least squares. + + Extended Summary + ---------------- + Calculate the count rate for each pixel in the data ramp, for every + integration. Generalized least squares is used for fitting the ramp + in order to take into account the correlation between reads. If the + input file contains multiple integrations, a second output file will + be written, containing per-integration count rates. + + One additional file can optionally be written (if save_opt is True), + containing per-integration data. + + Parameters + ---------- + model : data model + Input data model, assumed to be of type RampModel. + + buffsize : int + Size of data section (buffer) in bytes. + + save_opt : boolean + Calculate optional fitting results. + + readnoise_model : instance of data Model + Readnoise for all pixels. + + gain_model : instance of gain model + Gain for all pixels. + + max_cores : string + Number of cores to use for multiprocessing. If set to 'none' (the default), + then no multiprocessing will be done. The other allowable values are 'quarter', + 'half', and 'all'. This is the fraction of cores to use for multi-proc. The + total number of cores includes the SMT cores (Hyper Threading for Intel). + + Returns + ------- + new_model : Data Model object + DM object containing a rate image averaged over all integrations in + the exposure. + + int_model : Data Model object or None + DM object containing rate images for each integration in the exposure, + or None if there is only one integration. + + gls_opt_model : GLS_RampFitModel object or None + Object containing optional GLS-specific ramp fitting data for the + exposure; this will be None if save_opt is False. + """ + number_slices = utils.compute_slices(max_cores) + + # Get needed sizes and shapes + nreads, npix, imshape, cubeshape, n_int, instrume, frame_time, ngroups, \ + group_time = utils.get_dataset_info(input_model) + + (group_time, frames_per_group, saturated_flag, jump_flag) = \ + utils.get_more_info(input_model) + # Get readnoise array for calculation of variance of noiseless ramps, and + # gain array in case optimal weighting is to be done + # KDG - not sure what this means and no optimal weigting in GLS + readnoise_2d, gain_2d = utils.get_ref_subs(input_model, readnoise_model, + gain_model, frames_per_group) + # Flag any bad pixels in the gain + pixeldq = utils.reset_bad_gain(input_model.pixeldq, gain_2d) + log.info("number of processes being used is %d" % number_slices) + + total_rows = input_model.data.shape[2] + + tstart = time.time() + + # Determine the maximum number of cosmic ray hits for any pixel. + max_num_cr = -1 # invalid initial value + for num_int in range(n_int): + i_max_num_cr = utils.get_max_num_cr(input_model.groupdq[num_int, :, :, :], jump_flag) + max_num_cr = max(max_num_cr, i_max_num_cr) + + # Calculate effective integration time (once EFFINTIM has been populated + # and accessible, will use that instead), and other keywords that will + # needed if the pedestal calculation is requested. Note 'nframes' + # is the number of given by the NFRAMES keyword, and is the number of + # frames averaged on-board for a group, i.e., it does not include the + # groupgap. + effintim, nframes, groupgap, dropframes1 = utils.get_efftim_ped(input_model) + + if number_slices == 1: + rows_per_slice = total_rows + slopes, slope_int, slope_err_int, pixeldq_sect, dq_int, sum_weight, \ + intercept_int, intercept_err_int, pedestal_int, ampl_int, ampl_err_int = \ + gls_fit_all_integrations(frame_time, gain_2d, input_model.groupdq, + group_time, jump_flag, max_num_cr, input_model.data, + input_model.err, nframes, pixeldq, readnoise_2d, + saturated_flag, save_opt) + else: + rows_per_slice = round(total_rows / number_slices) + pool = Pool(processes=number_slices) + slices = [] + slopes = np.zeros(imshape, dtype=np.float32) + sum_weight = np.zeros(imshape, dtype=np.float32) + + # For multiple-integration datasets, will output integration-specific + # results to separate file named + '_rateints.fits'. + # Even if there's only one integration, the output results will be + # saved in these arrays. + slope_int = np.zeros((n_int,) + imshape, dtype=np.float32) + slope_err_int = np.zeros((n_int,) + imshape, dtype=np.float32) + dq_int = np.zeros((n_int,) + imshape, dtype=np.uint32) + out_pixeldq = np.zeros(imshape, dtype=np.uint32) + if save_opt: + # Create arrays for the fitted values of zero-point intercept and + # cosmic-ray amplitudes, and their errors. + intercept_int = np.zeros((n_int,) + imshape, dtype=np.float32) + intercept_err_int = np.zeros((n_int,) + imshape, dtype=np.float32) + # The pedestal is the extrapolation of the first group back to zero + # time, for each integration. + pedestal_int = np.zeros((n_int,) + imshape, dtype=np.float32) + # If there are no cosmic rays, set the last axis length to 1. + shape_ampl = (n_int, imshape[0], imshape[1], max(1, max_num_cr)) + ampl_int = np.zeros(shape_ampl, dtype=np.float32) + ampl_err_int = np.zeros(shape_ampl, dtype=np.float32) + +# Loop over number of processes + for i in range(number_slices - 1): + start_row = i * rows_per_slice + stop_row = (i + 1) * rows_per_slice + readnoise_slice = readnoise_2d[start_row: stop_row, :] + gain_slice = gain_2d[start_row: stop_row, :] + data_slice = input_model.data[:, :, start_row:stop_row, :].copy() + err_slice = input_model.err[:, :, start_row: stop_row, :].copy() + groupdq_slice = input_model.groupdq[:, :, start_row: stop_row, :].copy() + pixeldq_slice = pixeldq[start_row: stop_row, :].copy() + slices.insert(i, + (frame_time, gain_slice, groupdq_slice, group_time, + jump_flag, max_num_cr, data_slice, err_slice, frames_per_group, pixeldq_slice, + readnoise_slice, saturated_flag, save_opt)) + # The last slice takes the remainder of the rows + start_row = (number_slices - 1) * rows_per_slice + readnoise_slice = readnoise_2d[start_row: total_rows, :] + gain_slice = gain_2d[start_row: total_rows, :] + data_slice = input_model.data[:, :, start_row: total_rows, :].copy() + err_slice = input_model.err[:, :, start_row: total_rows, :].copy() + groupdq_slice = input_model.groupdq[:, :, start_row: total_rows, :].copy() + pixeldq_slice = input_model.pixeldq[start_row: total_rows, :].copy() + slices.insert(number_slices - 1, + (frame_time, gain_slice, groupdq_slice, group_time, + jump_flag, max_num_cr, data_slice, err_slice, frames_per_group, pixeldq_slice, + readnoise_slice, saturated_flag, save_opt)) + + log.debug("Creating %d processes for ramp fitting " % number_slices) + real_results = pool.starmap(gls_fit_all_integrations, slices) + pool.close() + pool.join() + k = 0 + log.debug("All processes complete") + for resultslice in real_results: + start_row = k * rows_per_slice + if len(real_results) == k + 1: # last result + slopes[start_row:total_rows, :] = resultslice[0] + slope_int[:, start_row:total_rows, :] = resultslice[1] + slope_err_int[:, start_row:total_rows, :] = resultslice[2] + out_pixeldq[start_row:total_rows, :] = resultslice[3] + if resultslice[4] is not None: + dq_int[:, start_row:total_rows, :] = resultslice[4] # nint > 1 + sum_weight[start_row:total_rows, :] = resultslice[5] # nint > 1 + if resultslice[6] is not None: + intercept_int[:, start_row: total_rows, :] = resultslice[6] # optional + intercept_err_int[:, start_row:total_rows, :] = resultslice[7] # optional + pedestal_int[:, start_row: total_rows, :] = resultslice[8] # optional + ampl_int[:, start_row:total_rows, :] = resultslice[9] # optional + ampl_err_int[:, start_row: total_rows, :] = resultslice[10] # optional + else: + stop_row = (k + 1) * rows_per_slice + slopes[start_row:stop_row, :] = resultslice[0] + slope_int[:, start_row:stop_row, :] = resultslice[1] + slope_err_int[:, start_row:stop_row, :] = resultslice[2] + out_pixeldq[start_row:stop_row, :] = resultslice[3] + if resultslice[4] is not None: + dq_int[:, start_row:stop_row, :] = resultslice[4] # nint > 1 + sum_weight[start_row:stop_row, :] = resultslice[5] # nint > 1 + if resultslice[6] is not None: + intercept_int[:, start_row: stop_row, :] = resultslice[6] # optional + intercept_err_int[:, start_row:stop_row, :] = resultslice[7] # optional + pedestal_int[:, start_row: stop_row, :] = resultslice[8] # optional + ampl_int[:, start_row:stop_row, :] = resultslice[9] # optional + ampl_err_int[:, start_row: stop_row, :] = resultslice[10] # optional + k = k + 1 + # Average the slopes over all integrations. + if n_int > 1: + sum_weight = np.where(sum_weight <= 0., 1., sum_weight) + recip_sum_weight = 1. / sum_weight + slopes *= recip_sum_weight + gls_err = np.sqrt(recip_sum_weight) + + # Convert back from electrons to DN. + slope_int /= gain_2d + slope_err_int /= gain_2d + if n_int > 1: + slopes /= gain_2d + gls_err /= gain_2d + if save_opt: + intercept_int /= gain_2d + intercept_err_int /= gain_2d + pedestal_int /= gain_2d + gain_shape = gain_2d.shape + gain_4d = gain_2d.reshape((1, gain_shape[0], gain_shape[1], 1)) + ampl_int /= gain_4d + ampl_err_int /= gain_4d + del gain_4d + del gain_2d + + # Compress all integration's dq arrays to create 2D PIXELDDQ array for + # primary output + final_pixeldq = utils.dq_compress_final(dq_int, n_int) + + int_model = utils.gls_output_integ(input_model, slope_int, slope_err_int, dq_int) + + if save_opt: # collect optional results for output + # Get the zero-point intercepts and the cosmic-ray amplitudes for + # each integration (even if there's only one integration). + gls_opt_model = utils.gls_output_optional( + input_model, intercept_int, intercept_err_int, pedestal_int, ampl_int, ampl_err_int) + else: + gls_opt_model = None + + tstop = time.time() + + if n_int > 1: + utils.log_stats(slopes) + else: + utils.log_stats(slope_int[0]) + + log.debug('Instrument: %s' % instrume) + log.debug('Number of pixels in 2D array: %d' % npix) + log.debug('Shape of 2D image: (%d, %d)' % imshape) + log.debug('Shape of data cube: (%d, %d, %d)' % cubeshape) + log.debug('Buffer size (bytes): %d' % buffsize) + log.debug('Number of rows per slice: %d' % rows_per_slice) + log.info('Number of groups per integration: %d' % nreads) + log.info('Number of integrations: %d' % n_int) + log.debug('The execution time in seconds: %f' % (tstop - tstart,)) + + # Create new model... + if n_int > 1: + new_model = datamodels.ImageModel( + data=slopes.astype(np.float32), dq=final_pixeldq, err=gls_err.astype(np.float32)) + else: + new_model = datamodels.ImageModel( + data=slope_int[0], dq=final_pixeldq, err=slope_err_int[0]) + + new_model.update(input_model) # ... and add all keys from input + + return new_model, int_model, gls_opt_model + + +def gls_fit_all_integrations( + frame_time, gain_2d, gdq_cube, group_time, jump_flag, max_num_cr, data_sect, + input_var_sect, nframes_used, pixeldq, readnoise_2d, saturated_flag, save_opt): + """ + This method will fit the rate for all pixels and all integrations using the Generalized Least + Squares (GLS) method. + Parameters + ---------- + frame_time : float32 + The time to read one frame + gain_2d : 2D float32 + The gain in electrons per DN for each pixel + gdq_cube : 4-D DQ Flags + The group dq flag values for all groups in the exposure + group_time : float32 + The time to read one group + jump_flag : DQ flag + The DQ value to mark a jump + max_num_cr : int + The largest number of cosmic rays found in any integration + data_sect : 4-D float32 + The input ramp cube with the sample values for each group of each integration for each pixel + input_var_sect: 4-D float32 + The input variance for each group of each integration for each pixel + nframes_used : int + The number of frames used to form each group average + pixel_dq : 2-D DQ flags + The pixel DQ flags for all pixels + readnoise_2d : 2-D float32 + The read noise for each pixel + saturated_flag : DQ flag + The DQ flag value to mark saturation + save_opt : boolean + Set to true to return the optional output model + Returns + -------- + slopes : 2-D float32 + The output rate for each pixel + slope_int : 2-D float32 + The output y-intercept for each pixel + slope_var_sect : 2-D float32 + The variance of the rate for each pixel + pixeldq_sect : 2-D DQ flag + The pixel dq for each pixel + dq_int : 3-D DQ flag + The pixel dq for each integration for each pixel + sum_weight : 2-D float32 + The sum of the weights for each pixel + intercept_int : 3-D float32 + The y-intercept for each integration for each pixel + intercept_err_int : 3-D float32 + The uncertainty of the y-intercept for each pixel of each integration + pedestal_int : 3-D float32 + The pedestal value for each integration for each pixel + ampl_int : 3-D float32 + The amplitude of each cosmic ray for each pixel + ampl_err_int : + The variance of the amplitude of each cosmic ray for each pixel + """ + number_ints = data_sect.shape[0] + number_rows = data_sect.shape[2] + number_cols = data_sect.shape[3] + imshape = (data_sect.shape[2], data_sect.shape[3]) + slope_int = np.zeros((number_ints, number_rows, number_cols), dtype=np.float32) + slope_err_int = np.zeros((number_ints, number_rows, number_cols), dtype=np.float32) + dq_int = np.zeros((number_ints, number_rows, number_cols), dtype=np.uint32) + temp_dq = np.zeros((number_rows, number_cols), dtype=np.uint32) + slopes = np.zeros((number_rows, number_cols), dtype=np.float32) + sum_weight = np.zeros((number_rows, number_cols), dtype=np.float32) + if save_opt: + # Create arrays for the fitted values of zero-point intercept and + # cosmic-ray amplitudes, and their errors. + intercept_int = np.zeros((number_ints,) + imshape, dtype=np.float32) + intercept_err_int = np.zeros((number_ints,) + imshape, dtype=np.float32) + # The pedestal is the extrapolation of the first group back to zero + # time, for each integration. + pedestal_int = np.zeros((number_ints,) + imshape, dtype=np.float32) + # The first group, for calculating the pedestal. (This only needs + # to be nrows high, but we don't have nrows yet. xxx) + first_group = np.zeros(imshape, dtype=np.float32) + # If there are no cosmic rays, set the last axis length to 1. + shape_ampl = (number_ints, imshape[0], imshape[1], max(1, max_num_cr)) + ampl_int = np.zeros(shape_ampl, dtype=np.float32) + ampl_err_int = np.zeros(shape_ampl, dtype=np.float32) + else: + intercept_int = None + intercept_err_int = None + pedestal_int = None + first_group = None + shape_ampl = None + ampl_int = None + ampl_err_int = None + # loop over data integrations + for num_int in range(number_ints): + if save_opt: + first_group[:, :] = 0. # re-use this for each integration + + # We'll propagate error estimates from previous steps to the + # current step by using the variance. + input_var_sect = input_var_sect ** 2 + + # Convert the data section from DN to electrons. + data_sect *= gain_2d + if save_opt: + first_group[:, :] = data_sect[num_int, 0, :, :].copy() + + intercept_sect, intercept_var_sect, slope_sect, slope_var_sect, cr_sect, cr_var_sect = \ + determine_slope(data_sect[num_int, :, :, :], input_var_sect[num_int, :, :, :], + gdq_cube[num_int, :, :, :], readnoise_2d, gain_2d, frame_time, group_time, + nframes_used, max_num_cr, saturated_flag, jump_flag) + + slope_int[num_int, :, :] = slope_sect.copy() + v_mask = (slope_var_sect <= 0.) + if v_mask.any(): + # Replace negative or zero variances with a large value. + slope_var_sect[v_mask] = utils.LARGE_VARIANCE + # Also set a flag in the pixel dq array. + temp_dq[:, :][v_mask] = UNRELIABLE_SLOPE + del v_mask + # If a pixel was flagged (by an earlier step) as saturated in + # the first group, flag the pixel as bad. + # Note: save s_mask until after the call to utils.gls_pedestal. + s_mask = (gdq_cube[0] == saturated_flag) + if s_mask.any(): + temp_dq[:, :][s_mask] = UNRELIABLE_SLOPE + slope_err_int[num_int, :, :] = np.sqrt(slope_var_sect) + + # We need to take a weighted average if (and only if) number_ints > 1. + # Accumulate sum of slopes and sum of weights. + if number_ints > 1: + weight = 1. / slope_var_sect + slopes[:, :] += (slope_sect * weight) + sum_weight[:, :] += weight + + if save_opt: + # Save the intercepts and cosmic-ray amplitudes for the + # current integration. + intercept_int[num_int, :, :] = intercept_sect.copy() + intercept_err_int[num_int, :, :] = np.sqrt(np.abs(intercept_var_sect)) + pedestal_int[num_int, :, :] = utils.gls_pedestal(first_group[:, :], + slope_int[num_int, :, :], + s_mask, + frame_time, nframes_used) + ampl_int[num_int, :, :, :] = cr_sect.copy() + ampl_err_int[num_int, :, :, :] = np.sqrt(np.abs(cr_var_sect)) + + # Compress 4D->2D dq arrays for saturated and jump-detected + # pixels + pixeldq_sect = pixeldq[:, :].copy() + dq_int[num_int, :, :] = \ + utils.dq_compress_sect(gdq_cube[num_int, :, :, :], pixeldq_sect).copy() + + dq_int[num_int, :, :] |= temp_dq + temp_dq[:, :] = 0 # initialize for next integration + + return slopes, slope_int, slope_var_sect, pixeldq_sect, dq_int, sum_weight, \ + intercept_int, intercept_err_int, pedestal_int, ampl_int, ampl_err_int + + +def determine_slope(data_sect, input_var_sect, + gdq_sect, readnoise_sect, gain_sect, + frame_time, group_time, nframes_used, + max_num_cr, saturated_flag, jump_flag): + """Iteratively fit a slope, intercept, and cosmic rays to a ramp. + + This function fits a ramp, possibly with discontinuities (cosmic-ray + hits), to a 3-D data "cube" with shape (number of groups, number of + pixels high, number of pixels wide). The fit will be done multiple + times, with the previous fit being used to assign weights (via the + covariance matrix) for the current fit. The iterations stop either + when the maximum number of iterations has been reached or when the + maximum difference between the previous fit and the current fit is + below a cutoff. This function calls compute_slope and evaluate_fit. + + compute_slope creates arrays for the slope, intercept, and cosmic-ray + amplitudes (the arrays that will be returned by determine_slope). Then + it loops over the number of cosmic rays, from 0 to max_num_cr + inclusive. Within this loop, compute_slope copies to temporary arrays + the ramp data for all the pixels that have the current number of cosmic + ray hits, calls gls_fit to compute the fit, then copies the results + of the fit (slope, etc.) to the output arrays for just those pixels. + + The input to gls_fit is ramp data for a subset of pixels (nz in number) + that all have the same number of cosmic-ray hits. gls_fit solves + matrix equations (one for each of the nz pixels) of the form: + + y = x * p + + where y is a column vector containing the observed data values in + electrons for each group (the length of y is ngroups, the number of + groups); x is a matrix with ngroups rows and 2 + num_cr columns, where + num_cr is the number of cosmic rays being included in the fit; and p + is the solution, a column vector containing the intercept, slope, and + the amplitude of each of the num_cr cosmic rays. The first column of + x is all ones, for fitting to the intercept. The second column of x + is the time (seconds) at the beginning of each group. The remaining + num_cr columns (if num_cr > 0) are Heaviside functions, 0 for the + first rows and 1 for all rows at and following the group containing a + cosmic-ray hit (each row corresponds to a group). There will be one + such column for each cosmic ray, so that the cosmic rays will be fit + independently of each other. Whether a cosmic ray hit the detector + during a particular group was determined by a previous step, and the + affected groups are flagged in the group data quality array. In order + to account for the variance of each observed value and the covariance + between them (since they're measurements along a ramp), the solution + is computed in this form (the @ sign represents matrix multiplication): + + (xT @ C^-1 @ x)^-1 @ [xT @ C^-1 @ y] + + where C is the ngroups by ngroups covariance matrix, ^-1 means matrix + inverse, and xT is the transpose of x. + + Summary of the notation: + + data_sect is 3-D, (ngroups, ny, nx); this is the ramp of science data. + cr_flagged is 3-D, (ngroups, ny, nx); 1 indicates a cosmic ray, e.g.: + cr_flagged = np.where(np.bitwise_and(gdq_sect, jump_flag), 1, 0) + cr_flagged_2d is 2-D, (ngroups, nz); this gives the location within + the ramp of each cosmic ray, for the subset of pixels (nz of them) + that have a total of num_cr cosmic ray hits at each pixel. This + is passed to gls_fit(), which fits a slope to the ramp. + + ramp_data has shape (ngroups, nz); this will be a ramp with a 1-D + array of pixels copied out of data_sect. The pixels will be those + that have a particular number of cosmic-ray hits, somewhere within + the ramp. + + Sum cr_flagged over groups to get an (ny, nx) image of the number of + cosmic rays (i.e. accumulated over the ramp) in each pixel. + sum_flagged = cr_flagged.sum(axis=0, ...) + sum_flagged is used to extract the nz pixels from (ny, nx) that have a + specified number of cosmic ray hits, e.g.: + for num_cr in range(max_num_cr + 1): + ncr_mask = (sum_flagged == num_cr) + nz = ncr_mask.sum(dtype=np.int32) + for k in range(ngroups): + ramp_data[k] = data_sect[k][ncr_mask] + cr_flagged_2d[k] = cr_flagged[k][ncr_mask] + + gls_fit is called for the subset of pixels (nz of them) that have + num_cr cosmic ray hits within the ramp, the same number for every + pixel. + + Parameters + ---------- + data_sect: 3-D ndarray, shape (ngroups, ny, nx) + The ramp data for one integration. This may be a subarray in + detector coordinates, but covering all groups. ngroups is the + number of groups; ny is the number of pixels in the Y direction; + nx is the number of pixels in the X (more rapidly varying) + direction. The units should be electrons. + + input_var_sect: 3-D ndarray, shape (ngroups, ny, nx) + The square of the input ERR array, matching data_sect. + + gdq_sect: 3-D ndarray, shape (ngroups, ny, nx) + The group data quality array. This may be a subarray, matching + data_sect. + + readnoise_sect: 2-D ndarray, shape (ny, nx) + The read noise in electrons at each detector pixel (i.e. not a + ramp). This may be a subarray, similar to data_sect. + + gain_sect: 2-D ndarray, or None, shape (ny, nx) + The gain in electrons per DN at each detector pixel (i.e. not a + ramp). This may be a subarray, matching readnoise_sect. If + gain_sect is None, a value of 1 will be assumed. + + frame_time: float + The time to read one frame, in seconds (e.g. 10.6 s). + + group_time: float + Time increment between groups, in seconds. + + nframes_used: int + Number of frames that were averaged together to make a group. + Note that this value does not include the number (if any) of + skipped frames. + + max_num_cr: non-negative int + The maximum number of cosmic rays that should be handled. This + must be specified by the caller, because determine_slope may be + called for different sections of the input data, and those sections + may have differing maximum numbers of cosmic rays. + + saturated_flag: int + dqflags.group['SATURATED'] + + jump_flag: int + dqflags.group['JUMP_DET'] + + Returns + ------- + tuple: (intercept_sect, int_var_sect, slope_sect, slope_var_sect, + cr_sect, cr_var_sect) + intercept_sect: 2-D ndarray, shape (ny, nx) + The intercept of the ramp at each pixel. + int_var_sect: 2-D ndarray, shape (ny, nx) + The variance of the intercept at each pixel. + slope_sect: 2-D ndarray, shape (ny, nx) + The ramp slope at each pixel of data_sect. + slope_var_sect: 2-D ndarray, shape (ny, nx) + The variance of the slope at each pixel. + cr_sect: 3-D ndarray, shape (ny, nx, cr_dimen) + The amplitude of each cosmic ray at each pixel. cr_dimen is + max_num_cr or 1, whichever is larger. + cr_var_sect: 3-D ndarray, shape (ny, nx, cr_dimen) + The variance of each cosmic-ray amplitude. + """ + + slope_diff_cutoff = 1.e-5 + + # These will be updated in the loop. + prev_slope_sect = (data_sect[1, :, :] - data_sect[0, :, :]) / group_time + prev_fit = data_sect.copy() + + use_extra_terms = True + + iter = 0 + done = False + if NUM_ITER_NO_EXTRA_TERMS <= 0: + # Even the first iteration uses the extra terms. + temp_use_extra_terms = True + else: + temp_use_extra_terms = False + while not done: + (intercept_sect, int_var_sect, slope_sect, slope_var_sect, + cr_sect, cr_var_sect) = \ + compute_slope(data_sect, input_var_sect, + gdq_sect, readnoise_sect, gain_sect, + prev_fit, prev_slope_sect, + frame_time, group_time, nframes_used, + max_num_cr, saturated_flag, jump_flag, + temp_use_extra_terms) + iter += 1 + if iter == NUM_ITER_NO_EXTRA_TERMS: + temp_use_extra_terms = use_extra_terms + if iter >= MAX_ITER: + done = True + else: + # If there are pixels with zero or negative variance, ignore + # them when taking the difference between the slopes computed + # in the current and previous iterations. + slope_diff = np.where(slope_var_sect > 0., + prev_slope_sect - slope_sect, 0.) + max_slope_diff = np.abs(slope_diff).max() + if iter >= MIN_ITER and max_slope_diff < slope_diff_cutoff: + done = True + current_fit = evaluate_fit(intercept_sect, slope_sect, cr_sect, + frame_time, group_time, + gdq_sect, jump_flag) + prev_fit = positive_fit(current_fit) # use for next iteration + del current_fit + prev_slope_sect = slope_sect.copy() + + return (intercept_sect, int_var_sect, slope_sect, slope_var_sect, + cr_sect, cr_var_sect) + + +def evaluate_fit(intercept_sect, slope_sect, cr_sect, + frame_time, group_time, + gdq_sect, jump_flag): + """Evaluate the fit (intercept, slope, cosmic-ray amplitudes). + + Parameters + ---------- + intercept_sect: 2-D ndarray + The intercept of the ramp at each pixel of data_sect (one of the + arguments to determine_slope). + + slope_sect: 2-D ndarray + The ramp slope at each pixel of data_sect. + + cr_sect: 3-D ndarray + The amplitude of each cosmic ray at each pixel of data_sect. + + frame_time: float + The time to read one frame, in seconds (e.g. 10.6 s). + + group_time: float + Time increment between groups, in seconds. + + gdq_sect: 3-D ndarray; indices: group, y, x + The group data quality array. This may be a subarray, matching + data_sect. + + jump_flag: int + dqflags.group['JUMP_DET'] + + Returns + ------- + fit_model: 3-D ndarray, shape (ngroups, ny, nx) + This is the same shape as data_sect, and if the fit is good, + fit_model and data_sect should not differ by much. + """ + + shape_3d = gdq_sect.shape # the ramp, (ngroups, ny, nx) + ngroups = gdq_sect.shape[0] + # This array is also created in function compute_slope. + cr_flagged = np.empty(shape_3d, dtype=np.uint8) + cr_flagged[:] = np.where(np.bitwise_and(gdq_sect, jump_flag), 1, 0) + + sum_flagged = cr_flagged.sum(axis=0, dtype=np.int32) + # local_max_num_cr is local to this function. It may be smaller than + # the max_num_cr that's an argument to determine_slope, and it can even + # be zero. + local_max_num_cr = sum_flagged.max() + del sum_flagged + + # The independent variable, in seconds at each image pixel. + ind_var = np.zeros(shape_3d, dtype=np.float64) + M = round(group_time / frame_time) + iv = np.arange(ngroups, dtype=np.float64) * group_time + \ + frame_time * (M + 1.) / 2. + iv = iv.reshape((ngroups, 1, 1)) + ind_var += iv + + # No cosmic rays yet; these will be accounted for below. + # ind_var has a different shape (ngroups, ny, nx) from slope_sect and + # intercept_sect, but their last dimensions are the same. + fit_model = ind_var * slope_sect + intercept_sect + + # heaviside and cr_flagged have shape (ngroups, ny, nx). + heaviside = np.zeros(shape_3d, dtype=np.float64) + cr_cumsum = cr_flagged.cumsum(axis=0, dtype=np.int16) + + # Add an offset for each cosmic ray. + for n in range(local_max_num_cr): + heaviside[:] = np.where(cr_cumsum > n, 1., 0.) + fit_model += (heaviside * cr_sect[:, :, n]) + + return fit_model + + +def positive_fit(current_fit): + """Replace zero and negative values with a positive number. + + Ramp data should be positive, since they are based on counts. The + fit to a ramp can go negative, however, due e.g. to extrapolation + beyond where the data are saturated. To avoid negative elements in + the covariance matrix (populated in part with the fit to the ramp), + this function replaces zero or negative values in the fit with a + positive number. + + Parameters + ---------- + current_fit: 3-D ndarray, shape (ngroups, ny, nx) + The fit returned by evaluate_fit. + + Returns + ------- + current_fit: 3-D ndarray, shape (ngroups, ny, nx) + This is the same as the input current_fit, except that zero and + negative values will have been replaced by a positive value. + """ + + return np.where(current_fit <= 0., FIT_MUST_BE_POSITIVE, current_fit) + + +def compute_slope(data_sect, input_var_sect, + gdq_sect, readnoise_sect, gain_sect, + prev_fit, prev_slope_sect, + frame_time, group_time, nframes_used, + max_num_cr, saturated_flag, jump_flag, + use_extra_terms): + """Set up the call to fit a slope to ramp data. + + This loops over the number of cosmic rays (jumps). That is, all the + ramps with no cosmic rays are processed first, then all the ramps with + one cosmic ray, then with two, etc. + + Parameters + ---------- + data_sect: 3-D ndarray; shape (ngroups, ny, nx) + The ramp data for one of the integrations in an exposure. This + may be a subarray in detector coordinates, but covering all groups. + + input_var_sect: 3-D ndarray, shape (ngroups, ny, nx) + The square of the input ERR array, matching data_sect. + + gdq_sect: 3-D ndarray; shape (ngroups, ny, nx) + The group data quality array. This may be a subarray, matching + data_sect. + + readnoise_sect: 2-D ndarray; shape (ny, nx) + The read noise in electrons at each detector pixel (i.e. not a + ramp). This may be a subarray, similar to data_sect. + + gain_sect: 2-D ndarray, or None; shape (ny, nx) + The gain in electrons per DN at each detector pixel (i.e. not a + ramp). This may be a subarray, matching readnoise_sect. If + gain_sect is None, a value of 1 will be assumed. + + prev_fit: 3-D ndarray; shape (ngroups, ny, nx) + The previous fit (intercept, slope, cosmic-ray amplitudes) + evaluated for each pixel in the subarray. data_sect itself may be + used for the first iteration. + + prev_slope_sect: 2-D ndarray; shape (ny, nx) + An estimate (e.g. from a previous iteration) of the slope at each + pixel, in electrons per second. This may be a subarray, similar to + data_sect. + + frame_time: float + The time to read one frame, in seconds (e.g. 10.6 s). + + group_time: float + Time increment between groups, in seconds. + + nframes_used: int + Number of frames that were averaged together to make a group. + This value does not include the number (if any) of skipped frames. + + max_num_cr: non-negative int + The maximum number of cosmic rays that should be handled. + + saturated_flag: int + dqflags.group['SATURATED'] + + jump_flag: int + dqflags.group['JUMP_DET'] + + use_extra_terms: bool + True if we should include Massimo Robberto's terms in the + covariance matrix. + See JWST-STScI-003193.pdf + + Returns + ------- + tuple: (intercept_sect, int_var_sect, slope_sect, slope_var_sect, + cr_sect, cr_var_sect) + intercept_sect is a 2-D ndarray, the intercept of the ramp at each + pixel of data_sect. + int_var_sect is a 2-D ndarray, the variance of the intercept at + each pixel of data_sect. + slope_sect is a 2-D ndarray, the ramp slope at each pixel of + data_sect. + slope_var_sect is a 2-D ndarray, the variance of the slope at each + pixel of data_sect. + cr_sect is a 3-D ndarray, shape (ny, nx, cr_dimen), the amplitude + of each cosmic ray at each pixel of data_sect. cr_dimen is + max_num_cr or 1, whichever is larger. + cr_var_sect is a 3-D ndarray, the variance of each cosmic ray + amplitude. + """ + + cr_flagged = np.empty(data_sect.shape, dtype=np.uint8) + cr_flagged[:] = np.where(np.bitwise_and(gdq_sect, jump_flag), 1, 0) + + # If a pixel is flagged as a jump in the first group, we can't fit to + # the ramp, because a matrix that we need to invert would be singular. + # If there's only one group, we can't fit a ramp to it anyway, so + # at this point we wouldn't need to be concerned about a jump. If + # there is more than one group, just ignore any jump the first group. + if data_sect.shape[0] > 1: + cr_flagged[0, :, :] = 0 + + # Sum over groups to get an (ny, nx) image of the number of cosmic + # rays in each pixel, accumulated over the ramp. + sum_flagged = cr_flagged.sum(axis=0, dtype=np.int32) + + # If a pixel is flagged as saturated in the first or second group, we + # don't want to even attempt to fit a slope to the ramp for that pixel. + # Handle this case by setting the corresponding pixel in sum_flagged to + # a negative number. The test `ncr_mask = (sum_flagged == num_cr)` + # will therefore never match, since num_cr is zero or larger, and the + # pixel will not be included in any ncr_mask. + mask1 = (gdq_sect[0, :, :] == saturated_flag) + sum_flagged[mask1] = -1 + # one_group_mask flags pixels that are not saturated in the first + # group but are saturated in the second group (if there is a second + # group). For these pixels, we will assign a value to the slope + # image by just dividing the value in the first group by group_time. + if len(gdq_sect) > 1: + mask2 = (gdq_sect[1, :, :] == saturated_flag) + sum_flagged[mask2] = -1 + one_group_mask = np.bitwise_and(mask2, np.bitwise_not(mask1)) + del mask2 + else: + one_group_mask = np.bitwise_not(mask1) + del mask1 + + # Set elements of this array to a huge value if the corresponding + # pixels are saturated. This is not a flag, it's a value to be + # added to the diagonal of the covariance matrix. + saturated = np.empty(data_sect.shape, dtype=np.float64) + saturated[:] = np.where(np.bitwise_and(gdq_sect, saturated_flag), + HUGE_FOR_LOW_WEIGHT, 0.) + + # Create arrays to be populated and then returned. + shape = data_sect.shape + # Lower limit of one, in case there are no cosmic rays at all. + cr_dimen = max(1, max_num_cr) + intercept_sect = np.zeros((shape[1], shape[2]), dtype=data_sect.dtype) + slope_sect = np.zeros((shape[1], shape[2]), dtype=data_sect.dtype) + cr_sect = np.zeros((shape[1], shape[2], cr_dimen), + dtype=data_sect.dtype) + int_var_sect = np.zeros((shape[1], shape[2]), dtype=data_sect.dtype) + slope_var_sect = np.zeros((shape[1], shape[2]), dtype=data_sect.dtype) + cr_var_sect = np.zeros((shape[1], shape[2], cr_dimen), + dtype=data_sect.dtype) + + # This takes care of the case that there's only one group, as well as + # pixels that are saturated in the second but not the first group of a + # multi-group file + if one_group_mask.any(): + slope_sect[one_group_mask] = data_sect[0, one_group_mask] / group_time + del one_group_mask + + # Fit slopes for all pixels that have no cosmic ray hits anywhere in + # the ramp, then fit slopes with one CR hit, then with two, etc. + for num_cr in range(max_num_cr + 1): + ngroups = len(data_sect) + ncr_mask = (sum_flagged == num_cr) + # Number of detector pixels flagged with num_cr CRs within the ramp. + nz = ncr_mask.sum(dtype=np.int32) + if nz <= 0: + continue + + # ramp_data will be a ramp with a 1-D array of pixels copied out + # of data_sect. + ramp_data = np.empty((ngroups, nz), dtype=data_sect.dtype) + input_var_data = np.empty((ngroups, nz), dtype=data_sect.dtype) + prev_fit_data = np.empty((ngroups, nz), dtype=prev_fit.dtype) + prev_slope_data = np.empty(nz, dtype=prev_slope_sect.dtype) + prev_slope_data[:] = prev_slope_sect[ncr_mask] + readnoise = np.empty(nz, dtype=readnoise_sect.dtype) + readnoise[:] = readnoise_sect[ncr_mask] + if gain_sect is None: + gain = None + else: + gain = np.empty(nz, dtype=gain_sect.dtype) + gain[:] = gain_sect[ncr_mask] + cr_flagged_2d = np.empty((ngroups, nz), dtype=cr_flagged.dtype) + saturated_data = np.empty((ngroups, nz), dtype=prev_fit.dtype) + for k in range(ngroups): + ramp_data[k] = data_sect[k][ncr_mask] + input_var_data[k] = input_var_sect[k][ncr_mask] + prev_fit_data[k] = prev_fit[k][ncr_mask] + cr_flagged_2d[k] = cr_flagged[k][ncr_mask] + # This is for clobbering saturated pixels. + saturated_data[k] = saturated[k][ncr_mask] + + (result, variances) = \ + gls_fit(ramp_data, + prev_fit_data, prev_slope_data, + readnoise, gain, + frame_time, group_time, nframes_used, + num_cr, cr_flagged_2d, saturated_data) + # Copy the intercept, slope, and cosmic-ray amplitudes and their + # variances to the arrays to be returned. + # ncr_mask is a mask array that is True for each pixel that has the + # current number (num_cr) of cosmic rays. Thus, the output arrays + # are being populated here in sets, a different set of pixels with + # each iteration of this loop. + intercept_sect[ncr_mask] = result[:, 0].copy() + int_var_sect[ncr_mask] = variances[:, 0].copy() + slope_sect[ncr_mask] = result[:, 1].copy() + slope_var_sect[ncr_mask] = variances[:, 1].copy() + # In this loop, i is just an index. cr_sect is populated for + # number of cosmic rays = 1 to num_cr, inclusive. + for i in range(num_cr): + cr_sect[ncr_mask, i] = result[:, 2 + i].copy() + cr_var_sect[ncr_mask, i] = variances[:, 2 + i].copy() + + return (intercept_sect, int_var_sect, slope_sect, slope_var_sect, + cr_sect, cr_var_sect) + + +def gls_fit(ramp_data, + prev_fit_data, prev_slope_data, + readnoise, gain, + frame_time, group_time, nframes_used, + num_cr, cr_flagged_2d, saturated_data): + """Generalized least squares linear fit. + + It is assumed that every input pixel has num_cr cosmic-ray hits + somewhere within the ramp. This function should be called separately + for different values of num_cr. + + Notes + ----- + Curently the noise model is assumed to be a combination of + read and photon noise alone. + Same technique could be used with more complex noise models, but then + the ramp covariance matrix should be input. + + Parameters + ---------- + ramp_data: 2-D ndarray; indices: group, pixel number + The ramp data for one of the integrations in an exposure. This + may be a subset in detector coordinates, but covering all groups. + The shape is (ngroups, nz), where ngroups is the length of the + ramp, and nz is the number of pixels in the current subset. + + prev_fit_data: 2-D ndarray, shape (ngroups, nz) + The fit to ramp_data, based on applying the values of intercept, + slope, and cosmic-ray amplitudes that were determined in a previous + call to gls_fit. This array is only used for setting up the + covariance matrix. + + prev_slope_data: 1-D ndarray, length nz. + An estimate (e.g. from a previous iteration) of the slope at each + pixel, in electrons per second. + + readnoise: 1-D ndarray, length nz. + The read noise in electrons at each detector pixel. + + gain: 1-D ndarray, shape (nz,) + The analog-to-digital gain (electrons per dn) at each detector + pixel. + + frame_time: float + The time to read one frame, in seconds (e.g. 10.6 s). + + group_time: float + Time increment between groups, in seconds. + + nframes_used: int + Number of frames that were averaged together to make a group. + Note that this value does not include the number (if any) of + skipped frames. + + num_cr: int + The number of cosmic rays that will be handled. All pixels in the + current set (ramp_data) are assumed to have this many cosmic ray + hits somewhere within the ramp. + + cr_flagged_2d: 2-D ndarray, shape (ngroups, nz) + The values should be 0 or 1; 1 indicates that a cosmic ray was + detected (by another step) at that point. + + saturated_data: 2-D ndarray, shape (ngroups, nz) + Normal values are zero; the value will be a huge number for + saturated pixels. This will be added to the main diagonal of the + inverse of the weight matrix to greatly reduce the weight for + saturated pixels. + + Returns + ------- + tuple: (result2d, variances) + result2d is a 2-D ndarray; shape (nz, 2 + num_cr) + The computed values of intercept, slope, and cosmic-ray amplitudes + (there will be num_cr cosmic-ray amplitudes) for each of the nz + pixels. + + variances is a 2-D ndarray; shape (nz, 2 + num_cr) + The variance for the intercept, slope, and for the amplitude of + each cosmic ray that was detected. + """ + + M = float(nframes_used) + + ngroups = ramp_data.shape[0] + nz = ramp_data.shape[1] + num_cr = int(num_cr) + + # x is an array (length nz) of matrices, each of which is the + # independent variable of a linear equation. Each such matrix + # has ngroups rows and 2 + num_cr columns. The first column is set + # to 1, for finding the intercept. The second column is the time at + # each group, for finding the slope. The remaining columns (if any), + # are 0 for all rows prior to a certain point, then 1 for all + # subsequent rows (i.e. the Heaviside function). The transition from + # 0 to 1 is the location of a cosmic ray hit; the first 1 in a column + # corresponds to the value in cr_flagged_2d being 1. + x = np.zeros((nz, ngroups, 2 + num_cr), dtype=np.float64) + x[:, :, 0] = 1. + x[:, :, 1] = np.arange(ngroups, dtype=np.float64) * group_time + \ + frame_time * (M + 1.) / 2. + + if num_cr > 0: + sum_crs = cr_flagged_2d.cumsum(axis=0) + for k in range(ngroups): + s = slice(k, ngroups) + for n in range(1, num_cr + 1): + temp = np.where(np.logical_and(cr_flagged_2d[k] == 1, + sum_crs[k] == n)) + if len(temp[0]) > 0: + index = (temp[0], s, n + 1) + x[index] = 1 + del temp, index + + y = np.transpose(ramp_data, (1, 0)).reshape((nz, ngroups, 1)) + + # ramp_cov is an array of nz matrices, each ngroups x ngroups. + # each matrix gives the covariance of that pixel's ramp data + ramp_cov = np.ones((nz, ngroups, ngroups), dtype=np.float64) + + # Use the previous fit to the data to populate the covariance matrix, + # for each of the nz pixels. prev_fit_data has shape (ngroups, nz), + # similar to the ramp data, but we want the nz axis to be the first + # (we're constructing an array of nz matrix equations), so transpose + # prev_fit_data. + prev_fit_T = np.transpose(prev_fit_data, (1, 0)) + for k in range(ngroups): + # Populate the upper right, row by row. + ramp_cov[:, k, k:ngroups] = prev_fit_T[:, k:k + 1] + # Populate the lower left, column by column. + ramp_cov[:, k:ngroups, k] = prev_fit_T[:, k:k + 1] + # Give saturated pixels a very high high variance (hence a low weight) + ramp_cov[:, k, k] += saturated_data[k, :] + del prev_fit_T + + # iden is 2-D, but it can broadcast to 4-D. This is used to add terms to + # the diagonal of the covariance matrix. + iden = np.identity(ngroups) + + rn3d = readnoise.reshape((nz, 1, 1)) + ramp_cov += (iden * rn3d**2) + + # prev_slope_data must be non-negative. + flags = prev_slope_data < 0. + prev_slope_data[flags] = 1. + + # The resulting fit parameters are + # (xT @ ramp_cov^-1 @ x)^-1 @ [xT @ ramp_cov^-1 @ y] + # = [y-intercept, slope, cr_amplitude_1, cr_amplitude_2, ...] + # where @ means matrix multiplication. + + # shape of xT is (nz, 2 + num_cr, ngroups) + xT = np.transpose(x, (0, 2, 1)) + + # shape of `ramp_invcov` is (nz, ngroups, ngroups) + iden = iden.reshape((1, ngroups, ngroups)) + ramp_invcov = la.solve(ramp_cov, iden) + + del iden + + # temp1 = xT @ ramp_invcov + # np.einsum use is equivalent to matrix multiplication + # shape of temp1 is (nz, 2 + num_cr, ngroups) + temp1 = np.einsum('...ij,...jk->...ik', xT, ramp_invcov) + + # temp_var = xT @ ramp_invcov @ x + # shape of temp_var is (nz, 2 + num_cr, 2 + num_cr) + temp_var = np.einsum('...ij,...jk->...ik', temp1, x) + + # `fitparam_cov` is an array of nz covariance matrices. + # fitparam_cov = (xT @ ramp_invcov @ x)^-1 + # shape of fitparam_covar is (nz, 2 + num_cr, 2 + num_cr) + I_2 = np.eye(2 + num_cr).reshape((1, 2 + num_cr, 2 + num_cr)) + try: + # inverse of temp_var + fitparam_cov = la.solve(temp_var, I_2) + except la.LinAlgError: + # find the pixel with the singular matrix + for z in range(nz): + try: + la.solve(temp_var[z], I_2) + except la.LinAlgError as msg2: + log.warning("singular matrix, z = %d" % z) + raise la.LinAlgError(msg2) + del I_2 + + # [xT @ ramp_invcov @ y] + # shape of temp2 is (nz, 2 + num_cr, 1) + temp2 = np.einsum('...ij,...jk->...ik', temp1, y) + + # shape of fitparam is (nz, 2 + num_cr, 1) + fitparam = np.einsum('...ij,...jk->...ik', fitparam_cov, temp2) + r_shape = fitparam.shape + fitparam2d = fitparam.reshape((r_shape[0], r_shape[1])) + del fitparam + + # shape of both result2d and variances is (nz, 2 + num_cr) + fitparam_uncs = fitparam_cov.diagonal(axis1=1, axis2=2).copy() + + return (fitparam2d, fitparam_uncs) +''' diff --git a/src/stcal/ramp_fitting/ols_fit.py b/src/stcal/ramp_fitting/ols_fit.py new file mode 100644 index 000000000..16957229a --- /dev/null +++ b/src/stcal/ramp_fitting/ols_fit.py @@ -0,0 +1,3485 @@ +#! /usr/bin/env python + + +import logging +# from multiprocessing.pool import Pool as Pool +import numpy as np +import time + +import warnings + +from . import constants +from . import utils + +log = logging.getLogger(__name__) +log.setLevel(logging.DEBUG) + +# TODO Should figure out a better way to do this +DO_NOT_USE = constants.DO_NOT_USE +SATURATED = constants.SATURATED +JUMP_DET = constants.JUMP_DET +UNRELIABLE_SLOPE = constants.UNRELIABLE_SLOPE + +BUFSIZE = 1024 * 300000 # 300Mb cache size for data section + + +log = logging.getLogger(__name__) +log.setLevel(logging.DEBUG) + + +def ols_ramp_fit_multi( + input_model, buffsize, save_opt, readnoise_2d, gain_2d, weighting, max_cores): + """ + Setup the inputs to ols_ramp_fit with and without multiprocessing. The + inputs will be sliced into the number of cores that are being used for + multiprocessing. Because the data models cannot be pickled, only numpy + arrays are passed and returned as parameters to ols_ramp_fit. + + Parameters + ---------- + input_model : data model + input data model, assumed to be of type RampModel + + buffsize : int + size of data section (buffer) in bytes (not used) + + save_opt : bool + calculate optional fitting results + + readnoise_2d : ndarray + readnoise for all pixels + + gain_2d : ndarray + gain for all pixels + + algorithm : str + 'OLS' specifies that ordinary least squares should be used; + 'GLS' specifies that generalized least squares should be used. + + weighting : str + 'optimal' specifies that optimal weighting should be used; + currently the only weighting supported. + + max_cores : str + Number of cores to use for multiprocessing. If set to 'none' (the default), + then no multiprocessing will be done. The other allowable values are 'quarter', + 'half', and 'all'. This is the fraction of cores to use for multi-proc. The + total number of cores includes the SMT cores (Hyper Threading for Intel). + + Returns + ------- + image_info: tuple + The tuple of computed ramp fitting arrays. + + integ_info: tuple + The tuple of computed integration fitting arrays. + + opt_info: tuple + The tuple of computed optional results arrays for fitting. + + gls_opt_model : GLS_RampFitModel object or None + Object containing optional GLS-specific ramp fitting data for the + exposure + """ + + # Determine number of slices to use for multi-processor computations + number_slices = utils.compute_slices(max_cores) + + # Copy the int_times table for TSO data + int_times = input_model.int_times + + # total_rows = input_model.data.shape[2] + # total_cols = input_model.data.shape[3] + number_of_integrations = input_model.data.shape[0] + + # For MIRI datasets having >1 group, if all pixels in the final group are + # flagged as DO_NOT_USE, resize the input model arrays to exclude the + # final group. Similarly, if leading groups 1 though N have all pixels + # flagged as DO_NOT_USE, those groups will be ignored by ramp fitting, and + # the input model arrays will be resized appropriately. If all pixels in + # all groups are flagged, return None for the models. + if input_model.meta.instrument.name == 'MIRI' and input_model.data.shape[1] > 1: + miri_ans = discard_miri_groups(input_model) + # The function returns False if the removed groups leaves no data to be + # processed. If this is the case, return None for all expected variables + # returned by ramp_fit + if miri_ans is not True: + return [None] * 3 + + # Call ramp fitting for the single processor (1 data slice) case + if number_slices == 1: + max_segments, max_CRs = calc_num_seg(input_model.groupdq, number_of_integrations) + log.debug(f"Max segments={max_segments}") + + # Single threaded computation + image_info, integ_info, opt_info = ols_ramp_fit_single( + input_model, int_times, buffsize, save_opt, readnoise_2d, gain_2d, weighting) + if image_info is None: + return None, None, None + + return image_info, integ_info, opt_info + + # Call ramp fitting for multi-processor (multiple data slices) case + else: + return None, None, None + + +''' +# Remove BEGIN multiprocessing + log.debug(f'number of processes being used is {number_slices}') + rows_per_slice = round(total_rows / number_slices) + pool = Pool(processes=number_slices) + slices = [] + + # Populate the first n-1 slices + for i in range(number_slices - 1): + start_row = i * rows_per_slice + stop_row = (i + 1) * rows_per_slice + readnoise_slice = readnoise_2d[start_row: stop_row, :] + gain_slice = gain_2d[start_row: stop_row, :] + data_slice = input_model.data[:, :, start_row: stop_row, :].copy() + err_slice = input_model.err[:, :, start_row: stop_row, :].copy() + groupdq_slice = input_model.groupdq[:, :, start_row:stop_row, :].copy() + pixeldq_slice = input_model.pixeldq[start_row:stop_row, :].copy() + + slices.insert( + i, + (data_slice, err_slice, groupdq_slice, pixeldq_slice, buffsize, save_opt, + readnoise_slice, gain_slice, weighting, + input_model.meta.instrument.name, input_model.meta.exposure.frame_time, + input_model.meta.exposure.ngroups, input_model.meta.exposure.group_time, + input_model.meta.exposure.groupgap, input_model.meta.exposure.nframes, + input_model.meta.exposure.drop_frames1, int_times)) + + # last slice gets the rest + start_row = (number_slices - 1) * rows_per_slice + readnoise_slice = readnoise_2d[start_row: total_rows, :] + gain_slice = gain_2d[start_row: total_rows, :] + data_slice = input_model.data[:, :, start_row: total_rows, :].copy() + err_slice = input_model.err[:, :, start_row: total_rows, :].copy() + groupdq_slice = input_model.groupdq[:, :, start_row: total_rows, :].copy() + pixeldq_slice = input_model.pixeldq[start_row: total_rows, :].copy() + slices.insert(number_slices - 1, + (data_slice, err_slice, groupdq_slice, pixeldq_slice, buffsize, save_opt, + readnoise_slice, gain_slice, weighting, + input_model.meta.instrument.name, input_model.meta.exposure.frame_time, + input_model.meta.exposure.ngroups, input_model.meta.exposure.group_time, + input_model.meta.exposure.groupgap, input_model.meta.exposure.nframes, + input_model.meta.exposure.drop_frames1, int_times)) + + # Start up the processes for each slice + log.debug("Creating %d processes for ramp fitting " % number_slices) + + # Use starmap because input is iterable as well. + real_results = pool.starmap(ols_ramp_fit_sliced, slices) + pool.close() + pool.join() + k = 0 + log.debug("All processes complete") + + # Check that all slices got processed properly + for resultslice in real_results: + if resultslice[0] is None: + return None, None, None + + # Create new model for the primary output. + actual_segments = real_results[0][20] + actual_CRs = real_results[0][21] + int_model, opt_model, out_model = \ + create_output_models(input_model, number_of_integrations, save_opt, + total_cols, total_rows, actual_segments, actual_CRs) + int_model.int_times = int_times + + # iterate over the number of slices and place the results into the output models + for resultslice in real_results: + start_row = k * rows_per_slice + if len(real_results) == k + 1: # last result + out_model.data[start_row: total_rows, :] = resultslice[0] + out_model.dq[start_row:total_rows, :] = resultslice[1] + out_model.var_poisson[start_row:total_rows, :] = resultslice[2] + out_model.var_rnoise[start_row:total_rows, :] = resultslice[3] + out_model.err[start_row:total_rows, :] = resultslice[4] + if resultslice[5] is not None: # Integration results exist + int_model.data[:, start_row:total_rows, :] = resultslice[5] + int_model.dq[:, start_row:total_rows, :] = resultslice[6] + int_model.var_poisson[:, start_row:total_rows, :] = resultslice[7] + int_model.var_rnoise[:, start_row:total_rows, :] = resultslice[8] + int_model.err[:, start_row:total_rows, :] = resultslice[9] + if resultslice[11] is not None: # Optional results exist + opt_model.slope[:, :, start_row:total_rows, :] = resultslice[11] + opt_model.sigslope[:, :, start_row:total_rows, :] = resultslice[12] + opt_model.var_poisson[:, :, start_row:total_rows, :] = resultslice[13] + opt_model.var_rnoise[:, :, start_row:total_rows, :] = resultslice[14] + opt_model.yint[:, :, start_row:total_rows, :] = resultslice[15] + opt_model.sigyint[:, :, start_row:total_rows, :] = resultslice[16] + opt_model.pedestal[:, start_row:total_rows, :] = resultslice[17] + opt_model.weights[:, :, start_row:total_rows, :] = resultslice[18] + opt_model.crmag[:, :, start_row:total_rows, :] = resultslice[19] + else: # all but last slice + stop_row = (k + 1) * rows_per_slice + out_model.data[start_row: stop_row, :] = resultslice[0] + out_model.dq[start_row: stop_row, :] = resultslice[1] + out_model.var_poisson[start_row: stop_row, :] = resultslice[2] + out_model.var_rnoise[start_row: stop_row, :] = resultslice[3] + out_model.err[start_row: stop_row, :] = resultslice[4] + if resultslice[5] is not None: # Multiple integration results exist + int_model.data[:, start_row: stop_row, :] = resultslice[5] + int_model.dq[:, start_row: stop_row, :] = resultslice[6] + int_model.var_poisson[:, start_row: stop_row, :] = resultslice[7] + int_model.var_rnoise[:, start_row: stop_row, :] = resultslice[8] + int_model.err[:, start_row: stop_row, :] = resultslice[9] + if resultslice[11] is not None: # Optional Results exist + opt_model.slope[:, :, start_row: (k + 1) * rows_per_slice, :] = resultslice[11] + opt_model.sigslope[:, :, start_row: (k + 1) * rows_per_slice, :] = resultslice[12] + opt_model.var_poisson[:, :, start_row: (k + 1) * rows_per_slice, :] = resultslice[13] + opt_model.var_rnoise[:, :, start_row: (k + 1) * rows_per_slice, :] = resultslice[14] + opt_model.yint[:, :, start_row: (k + 1) * rows_per_slice, :] = resultslice[15] + opt_model.sigyint[:, :, start_row: (k + 1) * rows_per_slice, :] = resultslice[16] + opt_model.pedestal[:, start_row: (k + 1) * rows_per_slice, :] = resultslice[17] + opt_model.weights[:, :, start_row: (k + 1) * rows_per_slice, :] = resultslice[18] + opt_model.crmag[:, :, start_row: (k + 1) * rows_per_slice, :] = resultslice[19] + k = k + 1 + + return out_model, int_model, opt_model + + +def create_output_models(input_model, number_of_integrations, save_opt, + total_cols, total_rows, actual_segments, actual_CRs): + """ + Create_output_models is used to make blank output models to hold the results from the OLS + ramp fitting. + + Parameters + ---------- + input_model : DataModel + The input ramp model + + number_of_integrations : int + The number of integration in the input model + + save_opt : bool + Whether to save the optional outputs + + total_cols : int + The number of columns in the input image + + total_rows : int + The number of rows in the input image + + actual_segments : int + The largest number of segments in the integration resulting from cosmic rays + + actual_CRs : int + The largest number of cosmic rays jumps found in any integration + + Returns + ------------ + int_model : DataModel + The per integration output model + + opt_model : DataModel + The optional output model + + out_model : RampFitOutputModel + The standard rate output model + + """ + # TODO Remove function + imshape = (total_rows, total_cols) + out_model = datamodels.ImageModel(data=np.zeros(imshape, dtype=np.float32), + dq=np.zeros(imshape, dtype=np.uint32), + var_poisson=np.zeros(imshape, dtype=np.float32), + var_rnoise=np.zeros(imshape, dtype=np.float32), + err=np.zeros(imshape, dtype=np.float32)) + # ... and add all keys from input + out_model.update(input_model) + + # create per integrations model + # TODO Remove function + int_model = datamodels.CubeModel( + data=np.zeros((number_of_integrations,) + imshape, dtype=np.float32), + dq=np.zeros((number_of_integrations,) + imshape, dtype=np.uint32), + var_poisson=np.zeros((number_of_integrations,) + imshape, dtype=np.float32), + var_rnoise=np.zeros((number_of_integrations,) + imshape, dtype=np.float32), + err=np.zeros((number_of_integrations,) + imshape, dtype=np.float32)) + int_model.int_times = None + int_model.update(input_model) # ... and add all keys from input + + # Create model for the optional output + # TODO Remove function + if save_opt: + opt_model = datamodels.RampFitOutputModel( + slope=np.zeros((number_of_integrations,) + (actual_segments,) + imshape, dtype=np.float32), + yint=np.zeros((number_of_integrations,) + (actual_segments,) + imshape, dtype=np.float32), + sigyint=np.zeros((number_of_integrations,) + (actual_segments,) + imshape, dtype=np.float32), + sigslope=np.zeros((number_of_integrations,) + (actual_segments,) + imshape, dtype=np.float32), + weights=np.zeros((number_of_integrations,) + (actual_segments,) + imshape, dtype=np.float32), + firstf_int=np.zeros((number_of_integrations,) + imshape, dtype=np.float32), + pedestal=np.zeros((number_of_integrations,) + imshape, dtype=np.float32), + crmag=np.zeros((number_of_integrations,) + (actual_CRs,) + imshape, dtype=np.float32), + var_poisson=np.zeros((number_of_integrations,) + (actual_segments,) + imshape, dtype=np.float32), + var_rnoise=np.zeros((number_of_integrations,) + (actual_segments,) + imshape, dtype=np.float32), + ) + + opt_model.meta.filename = input_model.meta.filename + opt_model.update(input_model) # ... and add all keys from input + else: + opt_model = None + + return int_model, opt_model, out_model + + +def ols_ramp_fit_sliced( + data, err, groupdq, inpixeldq, buffsize, save_opt, readnoise_2d, gain_2d, + weighting, instrume, frame_time, ngroups, group_time, groupgap, nframes, + dropframes1, int_times): + + """ + Fit a ramp using ordinary least squares. Calculate the count rate for each + pixel in all data cube sections and all integrations, equal to the weighted + slope for all sections (intervals between cosmic rays) of the pixel's ramp + divided by the effective integration time. This function wraps the single + threaded ols_ramp_fit_single function and is used in order to properly handle + the slicing of the data for multiprocessing. + + Parameters + ---------- + data : The input 4-D array with ramp data (num_integrations, num_groups, num_rows, num_cols) + The input ramp data + + err : ndarray + The input 4-D error that matches the ramp data + + groupdq : ndarray + The input 4-D group DQ flags + + inpixeldq : ndarray + The input 2-D pixel DQ flags + + buffsize : int + The working buffer size + + save_opt : bool + Whether to return the optional output model + + readnoise_2d : ndarray + The read noise of each pixel + + gain_2d : ndarray + The gain of each pixel + + weighting : str + 'optimal' is the only valid value + + instrume : str + Instrument name + + frame_time : float32 + The time to read one frame. + + ngroups : int + The number of groups in each integration + + group_time : float32 + The time to read one group. + + groupgap : int + The number of frames that are not included in the group average + + nframes : int + The number of frames that are included in the group average + + dropframes1 : + The number of frames dropped at the beginning of every integration + + int_times : None + Not used + + Returns + ------- + new_model.data : ndarray + The output final rate of each pixel, 2-D float + + new_model.dq : ndarray + The output pixel dq for each pixel, 2-D flag + + new_model.var_poisson : ndarray + The variance in each pixel due to Poisson noise, 2-D float32 + + new_model.var_rnoise : ndarray + The variance in each pixel due to read noise, 2-D float32 + + new_model.err : ndarray + The output total variance for each pixel, 2-D float32 + + int_data : ndarray + The rate for each pixel in each integration, 3-D float32 + + int_dq : ndarray + The pixel dq flag for each integration, 3-D float32 + + int_var_poisson : ndarray + The variance of the rate for each integration due to Poisson noise, 3-D float32 + + int_var_rnoise : ndarray + The variance of the rate for each integration due to read noise, 3-D float32 + + int_err : ndarray + The total variance of the rate for each integration, 3-D float32 + + int_int_times : 3-D + The total time for each integration + + opt_slope : ndarray + The rate of each segment in each integration, 4-D float32 + + opt_sigslope : ndarray + The total variance of the rate for each pixel in each segment of each + integration, 4-D float32 + + opt_var_poisson : ndarray + The Poisson variance of the rate for each pixel in each segment of each + integration, 4-D float32 + + opt_var_rnoise : ndarray + The read noise variance of the rate for each pixel in each segment of + each integration, 4-D float32 + + opt_yint : ndarray + The y-intercept for each pixel in each segment of each integration, 4-D float32 + + opt_sigyint : ndarray + The variance for each pixel in each segment of each integration, 4-D float32 + + opt_pedestal : ndarray + The zero point for each pixel in each segment of each integration, 4-D float32 + + opt_weights : ndarray + The weight of each pixel to use in combining the segments, 4-D float32 + + opt_crmag : ndarray + The magnitude of each CR in each integration, 4-D float32 + + actual_segments : int + The actual maximum number of segments in any integration + + actual_CRs : int + The actual maximum number of CRs in any integration + """ + # Package image data into a RampModel + input_model = RampModel() + + input_model.data = data + input_model.err = err + input_model.groupdq = groupdq + input_model.pixeldq = inpixeldq + + # Capture exposure and instrument data into the RampModel + input_model.meta.instrument.name = instrume + + input_model.meta.exposure.frame_time = frame_time + input_model.meta.exposure.ngroups = ngroups + input_model.meta.exposure.group_time = group_time + input_model.meta.exposure.groupgap = groupgap + input_model.meta.exposure.nframes = nframes + input_model.meta.exposure.drop_frames1 = dropframes1 + + # Compute ramp fitting + new_model, int_model, opt_res = ols_ramp_fit_single( + input_model, int_times, buffsize, save_opt, readnoise_2d, gain_2d, weighting) + + if new_model is None: + return [None] * 22 + + # Package computed data for return + if int_model is not None: + int_data = int_model.data.copy() + int_dq = int_model.dq.copy() + int_var_poisson = int_model.var_poisson.copy() + int_var_rnoise = int_model.var_rnoise.copy() + int_err = int_model.err.copy() + int_int_times = int_model.int_times.copy() + else: + int_data = None + int_dq = None + int_var_poisson = None + int_var_rnoise = None + int_err = None + int_int_times = None + + if opt_res is not None: + opt_slope = opt_res.slope.copy() + opt_sigslope = opt_res.sigslope.copy() + opt_var_poisson = opt_res.var_poisson.copy() + opt_var_rnoise = opt_res.var_rnoise.copy() + opt_yint = opt_res.yint.copy() + opt_sigyint = opt_res.sigyint.copy() + opt_pedestal = opt_res.pedestal.copy() + opt_weights = opt_res.weights.copy() + opt_crmag = opt_res.crmag.copy() + actual_segments = opt_slope.shape[1] + actual_CRs = opt_crmag.shape[1] + else: + opt_slope = None + opt_sigslope = None + opt_var_poisson = None + opt_var_rnoise = None + opt_yint = None + opt_sigyint = None + opt_pedestal = None + opt_weights = None + opt_crmag = None + actual_segments = 0 + actual_CRs = 0 + + return new_model.data, new_model.dq, new_model.var_poisson, \ + new_model.var_rnoise, new_model.err, \ + int_data, int_dq, int_var_poisson, int_var_rnoise, int_err, int_int_times, \ + opt_slope, opt_sigslope, opt_var_poisson, opt_var_rnoise, opt_yint, opt_sigyint, \ + opt_pedestal, opt_weights, opt_crmag, actual_segments, actual_CRs +# END Multiprocessing +''' + + +def ols_ramp_fit_single( + input_model, int_times, buffsize, save_opt, readnoise_2d, gain_2d, weighting): + """ + Fit a ramp using ordinary least squares. Calculate the count rate for each + pixel in all data cube sections and all integrations, equal to the weighted + slope for all sections (intervals between cosmic rays) of the pixel's ramp + divided by the effective integration time. + + Parameters + ---------- + input_model: RampModel + + int_times : None + Not used + + buffsize : int + The working buffer size + + save_opt : bool + Whether to return the optional output model + + readnoise_2d : ndarray + The read noise of each pixel + + gain_2d : ndarray + The gain of each pixel + + weighting : str + 'optimal' is the only valid value + + Return + ------ + image_info: tuple + The tuple of computed ramp fitting arrays. + + integ_info: tuple + The tuple of computed integration fitting arrays. + + opt_info: tuple + The tuple of computed optional results arrays for fitting. + """ + tstart = time.time() + + # Save original shapes for writing to log file, as these may change for MIRI + n_int, ngroups, nrows, ncols = input_model.data.shape + orig_ngroups = ngroups + orig_cubeshape = (ngroups, nrows, ncols) + + if ngroups == 1: + log.warning('Dataset has NGROUPS=1, so count rates for each integration ') + log.warning('will be calculated as the value of that 1 group divided by ') + log.warning('the group exposure time.') + + # In this 'First Pass' over the data, loop over integrations and data + # sections to calculate the estimated median slopes, which will be used + # to calculate the variances. This is the same method to estimate slopes + # as is done in the jump detection step, except here CR-affected and + # saturated groups have already been flagged. The actual, fit, slopes for + # each segment are also calculated here. + fit_slopes_ans = ramp_fit_slopes( + input_model, gain_2d, readnoise_2d, save_opt, weighting) + if fit_slopes_ans[0] == "saturated": + return fit_slopes_ans[1:] + + # In this 'Second Pass' over the data, loop over integrations and data + # sections to calculate the variances of the slope using the estimated + # median slopes from the 'First Pass'. These variances are due to Poisson + # noise only, read noise only, and the combination of Poisson noise and + # read noise. The integration-specific variances are 3D arrays, and the + # segment-specific variances are 4D arrays. + variances_ans = \ + ramp_fit_compute_variances(input_model, gain_2d, readnoise_2d, fit_slopes_ans) + + # Now that the segment-specific and integration-specific variances have + # been calculated, the segment-specific, integration-specific, and + # overall slopes will be calculated. The integration-specific slope is + # calculated as a weighted average of the segments in the integration: + # slope_int = sum_over_segs(slope_seg/var_seg)/ sum_over_segs(1/var_seg) + # The overall slope is calculated as a weighted average of the segments in + # all integrations: + # slope = sum_over_integs_and_segs(slope_seg/var_seg)/ + # sum_over_integs_and_segs(1/var_seg) + image_info, integ_info, opt_info = ramp_fit_overall( + input_model, orig_cubeshape, orig_ngroups, buffsize, fit_slopes_ans, + variances_ans, save_opt, int_times, tstart) + + return image_info, integ_info, opt_info + + +def discard_miri_groups(input_model): + """ + For MIRI datasets having >1 group, if all pixels in the final group are + flagged as DO_NOT_USE, resize the input model arrays to exclude the + final group. Similarly, if leading groups 1 though N have all pixels + flagged as DO_NOT_USE, those groups will be ignored by ramp fitting, and + the input model arrays will be resized appropriately. If all pixels in + all groups are flagged, return None for the models. + + Parameters + ---------- + input_model: RampModel + The input model containing the image data. + + Returns + ------- + bool: + False if no data to process after discarding unusable data. + True if useable data available for further processing. + """ + data = input_model.data + err = input_model.err + groupdq = input_model.groupdq + + n_int, ngroups, nrows, ncols = data.shape + + num_bad_slices = 0 # number of initial groups that are all DO_NOT_USE + + while np.all(np.bitwise_and(groupdq[:, 0, :, :], DO_NOT_USE)): + num_bad_slices += 1 + ngroups -= 1 + + # Check if there are remaining groups before accessing data + if ngroups < 1: # no usable data + log.error('1. All groups have all pixels flagged as DO_NOT_USE,') + log.error(' so will not process this dataset.') + return False + + groupdq = groupdq[:, 1:, :, :] + + # Where the initial group of the just-truncated data is a cosmic ray, + # remove the JUMP_DET flag from the group dq for those pixels so + # that those groups will be included in the fit. + wh_cr = np.where(np.bitwise_and(groupdq[:, 0, :, :], JUMP_DET)) + num_cr_1st = len(wh_cr[0]) + + for ii in range(num_cr_1st): + groupdq[wh_cr[0][ii], 0, wh_cr[1][ii], wh_cr[2][ii]] -= JUMP_DET + + if num_bad_slices > 0: + data = data[:, num_bad_slices:, :, :] + err = err[:, num_bad_slices:, :, :] + + log.info('Number of leading groups that are flagged as DO_NOT_USE: %s', num_bad_slices) + + # If all groups were flagged, the final group would have been picked up + # in the while loop above, ngroups would have been set to 0, and Nones + # would have been returned. If execution has gotten here, there must + # be at least 1 remaining group that is not all flagged. + if np.all(np.bitwise_and(groupdq[:, -1, :, :], DO_NOT_USE)): + ngroups -= 1 + + # Check if there are remaining groups before accessing data + if ngroups < 1: # no usable data + log.error('2. All groups have all pixels flagged as DO_NOT_USE,') + log.error(' so will not process this dataset.') + return False + + data = data[:, :-1, :, :] + err = err[:, :-1, :, :] + groupdq = groupdq[:, :-1, :, :] + + log.info('MIRI dataset has all pixels in the final group flagged as DO_NOT_USE.') + + # Next block is to satisfy github issue 1681: + # "MIRI FirstFrame and LastFrame minimum number of groups" + if ngroups < 2: + log.warning('MIRI datasets require at least 2 groups/integration') + log.warning('(NGROUPS), so will not process this dataset.') + return False + + input_model.data = data + input_model.err = err + input_model.groupdq = groupdq + return True + + +def ramp_fit_slopes(input_model, gain_2d, readnoise_2d, save_opt, weighting): + """ + Calculate effective integration time (once EFFINTIM has been populated accessible, will + use that instead), and other keywords that will needed if the pedestal calculation is + requested. Note 'nframes' is the number of given by the NFRAMES keyword, and is the + number of frames averaged on-board for a group, i.e., it does not include the groupgap. + + Parameters + ---------- + input_model: RampModel + The input model containing the image data. + + gain_2d : ndarrays + gain for all pixels + + readnoise_2d : ndarrays + readnoise for all pixels + + save_opt : bool + calculate optional fitting results + + weighting : str + 'optimal' specifies that optimal weighting should be used; + currently the only weighting supported. + + Return + ------ + max_seg : int + Maximum possible number of segments over all groups and segments + + gdq_cube_shape : ndarray + Group DQ dimensions + + effintim : float + effective integration time for a single group + + f_max_seg : int + Actual maximum number of segments over all groups and segments + + dq_int : ndarray + The pixel dq for each integration for each pixel + + sum_weight : ndarray + The sum of the weights for each pixel + + num_seg_per_int : ndarray + Cube of numbers of segments for all integrations and pixels, 3-D int + + sat_0th_group_int : ndarray + Integration-specific slice whose value for a pixel is 1 if the initial + group of the ramp is saturated, 3-D uint8 + + opt_res : OptRes + Object to hold optional results for all good pixels. + + pixeldq : ndarray + The input 2-D pixel DQ flags + + inv_var : ndarray + values of 1/variance for good pixels, 1-D float + + med_rates : ndarray + Rate array + """ + + # Get image data information + data = input_model.data + err = input_model.err + groupdq = input_model.groupdq + inpixeldq = input_model.pixeldq + + # Get instrument and exposure data + frame_time = input_model.meta.exposure.frame_time + group_time = input_model.meta.exposure.group_time + groupgap = input_model.meta.exposure.groupgap + nframes = input_model.meta.exposure.nframes + + # Get needed sizes and shapes + n_int, ngroups, nrows, ncols = data.shape + imshape = (nrows, ncols) + cubeshape = (ngroups,) + imshape + + # If all the pixels have their initial groups flagged as saturated, the DQ + # in the primary and integration-specific output products are updated, + # the other arrays in all output products are populated with zeros, and + # the output products are returned to ramp_fit(). If the initial group of + # a ramp is saturated, it is assumed that all groups are saturated. + first_gdq = groupdq[:, 0, :, :] + if np.all(np.bitwise_and(first_gdq, SATURATED)): + image_info, integ_info, opt_info = utils.do_all_sat( + inpixeldq, groupdq, imshape, n_int, save_opt) + + return "saturated", image_info, integ_info, opt_info + + # Calculate effective integration time (once EFFINTIM has been populated + # and accessible, will use that instead), and other keywords that will + # needed if the pedestal calculation is requested. Note 'nframes' + # is the number of given by the NFRAMES keyword, and is the number of + # frames averaged on-board for a group, i.e., it does not include the + # groupgap. + effintim = (nframes + groupgap) * frame_time + + # Get GROUP DQ and ERR arrays from input file + gdq_cube = groupdq + gdq_cube_shape = gdq_cube.shape + + # Get max number of segments fit in all integrations + # TODO: this computation is already done in ramp_fit_multi + max_seg, num_CRs = calc_num_seg(gdq_cube, n_int) + del gdq_cube + + f_max_seg = 0 # final number to use, usually overwritten by actual value + + dq_int, median_diffs_2d, num_seg_per_int, sat_0th_group_int =\ + utils.alloc_arrays_1(n_int, imshape) + + opt_res = utils.OptRes(n_int, imshape, max_seg, ngroups, save_opt) + + # Get Pixel DQ array from input file. The incoming RampModel has uint32 + # PIXELDQ, but ramp fitting will update this array here by flagging + # the 2D PIXELDQ locations where the ramp data has been previously + # flagged as jump-detected or saturated. These additional bit values + # require this local variable to be uint16, and it will be used as the + # (uint16) PIXELDQ in the outgoing ImageModel. + pixeldq = inpixeldq.copy() + pixeldq = utils.reset_bad_gain(pixeldq, gain_2d) # Flag bad pixels in gain + + # In this 'First Pass' over the data, loop over integrations and data + # sections to calculate the estimated median slopes, which will be used + # to calculate the variances. This is the same method to estimate slopes + # as is done in the jump detection step, except here CR-affected and + # saturated groups have already been flagged. The actual, fit, slopes for + # each segment are also calculated here. + + # Loop over data integrations: + for num_int in range(0, n_int): + # Loop over data sections + for rlo in range(0, cubeshape[1], nrows): + rhi = rlo + nrows + + if rhi > cubeshape[1]: + rhi = cubeshape[1] + + # Skip data section if it is all NaNs + # data_sect = np.float32(data[num_int, :, :, :]) + data_sect = data[num_int, :, :, :] + if np.all(np.isnan(data_sect)): + log.error('Current data section is all nans, so not processing the section.') + continue + + # first frame section for 1st group of current integration + ff_sect = data[num_int, 0, rlo:rhi, :] + + # Get appropriate sections + gdq_sect = groupdq[num_int, :, :, :] + rn_sect = readnoise_2d[rlo:rhi, :] + gain_sect = gain_2d[rlo:rhi, :] + + # Reset all saturated groups in the input data array to NaN + where_sat = np.where(np.bitwise_and(gdq_sect, SATURATED)) + + data_sect[where_sat] = np.NaN + del where_sat + + # Compute the first differences of all groups + first_diffs_sect = np.diff(data_sect, axis=0) + + # If the dataset has only 1 group/integ, assume the 'previous group' + # is all zeros, so just use data as the difference + if first_diffs_sect.shape[0] == 0: + first_diffs_sect = data_sect.copy() + else: + # Similarly, for datasets having >1 group/integ and having + # single-group segments, just use the data as the difference + wh_nan = np.where(np.isnan(first_diffs_sect[0, :, :])) + + if len(wh_nan[0]) > 0: + first_diffs_sect[0, :, :][wh_nan] = data_sect[0, :, :][wh_nan] + + del wh_nan + + # Mask all the first differences that are affected by a CR, + # starting at group 1. The purpose of starting at index 1 is + # to shift all the indices down by 1, so they line up with the + # indices in first_diffs. + i_group, i_yy, i_xx, = np.where(np.bitwise_and(gdq_sect[1:, :, :], JUMP_DET)) + first_diffs_sect[i_group - 1, i_yy, i_xx] = np.NaN + + del i_group, i_yy, i_xx + + # Check for pixels in which there is good data in 0th group, but + # all first_diffs for this ramp are NaN because there are too + # few good groups past the 0th. Due to the shortage of good + # data, the first_diffs will be set here equal to the data in + # the 0th group. + wh_min = np.where(np.logical_and( + np.isnan(first_diffs_sect).all(axis=0), np.isfinite(data_sect[0, :, :]))) + if len(wh_min[0] > 0): + first_diffs_sect[0, :, :][wh_min] = data_sect[0, :, :][wh_min] + + del wh_min + + # All first differences affected by saturation and CRs have been set + # to NaN, so compute the median of all non-NaN first differences. + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", "All-NaN.*", RuntimeWarning) + nan_med = np.nanmedian(first_diffs_sect, axis=0) + nan_med[np.isnan(nan_med)] = 0. # if all first_diffs_sect are nans + median_diffs_2d[rlo:rhi, :] += nan_med + + # Calculate the slope of each segment + # note that the name "opt_res", which stands for "optional results", + # is deceiving; this in fact contains all the per-integration and + # per-segment results that will eventually be used to compute the + # final slopes, sigmas, etc. for the main (non-optional) products + t_dq_cube, inv_var, opt_res, f_max_seg, num_seg = \ + calc_slope(data_sect, gdq_sect, frame_time, opt_res, save_opt, rn_sect, + gain_sect, max_seg, ngroups, weighting, f_max_seg) + + del gain_sect + + # Populate 3D num_seg { integ, y, x } with 2D num_seg for this data + # section (y,x) and integration (num_int) + sect_shape = data_sect.shape[-2:] + num_seg_per_int[num_int, rlo:rhi, :] = num_seg.reshape(sect_shape) + + # Populate integ-spec slice which is set if 0th group has SAT + wh_sat0 = np.where(np.bitwise_and(gdq_sect[0, :, :], SATURATED)) + if len(wh_sat0[0]) > 0: + sat_0th_group_int[num_int, rlo:rhi, :][wh_sat0] = 1 + + del wh_sat0 + + pixeldq_sect = pixeldq[rlo:rhi, :].copy() + dq_int[num_int, rlo:rhi, :] = utils.dq_compress_sect(t_dq_cube, pixeldq_sect).copy() + + del t_dq_cube + + # Loop over the segments and copy the reshaped 2D segment-specific + # results for the current data section to the 4D output arrays. + opt_res.reshape_res(num_int, rlo, rhi, sect_shape, ff_sect, save_opt) + + if save_opt: + # Calculate difference between each slice and the previous slice + # as approximation to cosmic ray amplitude for those pixels + # having their DQ set for cosmic rays + data_diff = data_sect - utils.shift_z(data_sect, -1) + dq_cr = np.bitwise_and(JUMP_DET, gdq_sect) + + opt_res.cr_mag_seg[num_int, :, rlo:rhi, :] = data_diff * (dq_cr != 0) + + del data_diff + + del data_sect + del ff_sect + del gdq_sect + + if pixeldq_sect is not None: + del pixeldq_sect + + # Compute the final 2D array of differences; create rate array + median_diffs_2d /= n_int + med_rates = median_diffs_2d / group_time + + del median_diffs_2d + del first_diffs_sect + + input_model.data = data + input_model.err = err + input_model.groupdq = groupdq + input_model.pixeldq = inpixeldq + + return max_seg, gdq_cube_shape, effintim, f_max_seg, dq_int, num_seg_per_int,\ + sat_0th_group_int, opt_res, pixeldq, inv_var, med_rates + + +def ramp_fit_compute_variances(input_model, gain_2d, readnoise_2d, fit_slopes_ans): + """ + In this 'Second Pass' over the data, loop over integrations and data + sections to calculate the variances of the slope using the estimated + median slopes from the 'First Pass'. These variances are due to Poisson + noise only, read noise only, and the combination of Poisson noise and + read noise. The integration-specific variances are 3D arrays, and the + segment-specific variances are 4D arrays. + + The naming convention for the arrays: + 'var': a variance + 'p3': intermediate 3D array for variance due to Poisson noise + 'r4': intermediate 4D array for variance due to read noise + 'both4': intermediate 4D array for combined variance due to both Poisson and read noise + 'inv_': intermediate array = 1/ + 's_inv_': intermediate array = 1/, summed over integrations + + + Parameters + ---------- + input_model: RampModel + The input model containing the image data. + + gain_2d : ndarray + gain for all pixels + + readnoise_2d : ndarray + The read noise for each pixel + + fit_slopes_ans : tuple + Contains intermediate values computed in the first pass over the data. + + Return + ------ + var_p3 : ndarray + 3-D variance based on Poisson noise + + var_r3 : ndarray + 3-D variance based on read noise + + var_p4 : ndarray + 4-D variance based on Poisson noise + + var_r4 : ndarray + 4-D variance based on read noise + + var_both4 : ndarray + 4-D array for combined variance due to both Poisson and read noise + + var_both3 : ndarray + 3-D array for combined variance due to both Poisson and read noise + + inv_var_both4 : ndarray + 1 / var_both4 + + s_inv_var_p3 : ndarray + 1 / var_p3, summed over integrations + + s_inv_var_r3 : ndarray + 1 / var_r3, summed over integrations + + s_inv_var_both3 : ndarray + 1 / var_both3, summed over integrations + """ + + # Get image data information + data = input_model.data + err = input_model.err + groupdq = input_model.groupdq + inpixeldq = input_model.pixeldq + + # Get instrument and exposure data + group_time = input_model.meta.exposure.group_time + + # Get needed sizes and shapes + n_int, ngroups, nrows, ncols = data.shape + imshape = (nrows, ncols) + cubeshape = (ngroups,) + imshape + + max_seg = fit_slopes_ans[0] + num_seg_per_int = fit_slopes_ans[5] + med_rates = fit_slopes_ans[10] + + var_p3, var_r3, var_p4, var_r4, var_both4, var_both3, \ + inv_var_both4, s_inv_var_p3, s_inv_var_r3, s_inv_var_both3, segs_4 = \ + utils.alloc_arrays_2(n_int, imshape, max_seg) + + # Loop over data integrations + for num_int in range(n_int): + + # Loop over data sections + for rlo in range(0, cubeshape[1], nrows): + rhi = rlo + nrows + + if rhi > cubeshape[1]: + rhi = cubeshape[1] + + gdq_sect = groupdq[num_int, :, rlo:rhi, :] + rn_sect = readnoise_2d[rlo:rhi, :] + gain_sect = gain_2d[rlo:rhi, :] + + # Calculate results needed to compute the variance arrays + den_r3, den_p3, num_r3, segs_beg_3 = \ + utils.calc_slope_vars(rn_sect, gain_sect, gdq_sect, group_time, max_seg) + + segs_4[num_int, :, rlo:rhi, :] = segs_beg_3 + + # Suppress harmless arithmetic warnings for now + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + var_p4[num_int, :, rlo:rhi, :] = den_p3 * med_rates[rlo:rhi, :] + + # Find the segment variance due to read noise and convert back to DN + var_r4[num_int, :, rlo:rhi, :] = num_r3 * den_r3 / gain_sect**2 + + # Reset the warnings filter to its original state + warnings.resetwarnings() + + del den_r3, den_p3, num_r3, segs_beg_3 + del gain_sect + del gdq_sect + + # The next 4 statements zero out entries for non-existing segments, and + # set the variances for segments having negative slopes (the segment + # variance is proportional to the median estimated slope) to + # outrageously large values so that they will have negligible + # contributions. + var_p4[num_int, :, :, :] *= (segs_4[num_int, :, :, :] > 0) + + # Suppress, then re-enable harmless arithmetic warnings + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + var_p4[var_p4 <= 0.] = utils.LARGE_VARIANCE + + var_r4[num_int, :, :, :] *= (segs_4[num_int, :, :, :] > 0) + var_r4[var_r4 <= 0.] = utils.LARGE_VARIANCE + + # The sums of inverses of the variances are needed for later + # variance calculations. + s_inv_var_p3[num_int, :, :] = (1. / var_p4[num_int, :, :, :]).sum(axis=0) + var_p3[num_int, :, :] = 1. / s_inv_var_p3[num_int, :, :] + s_inv_var_r3[num_int, :, :] = (1. / var_r4[num_int, :, :, :]).sum(axis=0) + var_r3[num_int, :, :] = 1. / s_inv_var_r3[num_int, :, :] + + # Huge variances correspond to non-existing segments, so are reset to 0 + # to nullify their contribution. + var_p3[var_p3 > 0.1 * utils.LARGE_VARIANCE] = 0. + warnings.resetwarnings() + + var_both4[num_int, :, :, :] = var_r4[num_int, :, :, :] + var_p4[num_int, :, :, :] + inv_var_both4[num_int, :, :, :] = 1. / var_both4[num_int, :, :, :] + + # Want to retain values in the 4D arrays only for the segments that each + # pixel has, so will zero out values for the higher indices. Creating + # and manipulating intermediate arrays (views, such as var_p4_int + # will zero out the appropriate indices in var_p4 and var_r4.) + # Extract the slice of 4D arrays for the current integration + var_p4_int = var_p4[num_int, :, :, :] # [ segment, y, x ] + inv_var_both4_int = inv_var_both4[num_int, :, :, :] + + # Zero out non-existing segments + var_p4_int *= (segs_4[num_int, :, :, :] > 0) + inv_var_both4_int *= (segs_4[num_int, :, :, :] > 0) + + # reshape these arrays to simplify masking [ segment, 1D pixel ] + var_p4_int2 = var_p4_int.reshape( + (var_p4_int.shape[0], var_p4_int.shape[1] * var_p4_int.shape[2])) + + s_inv_var_both3[num_int, :, :] = (inv_var_both4[num_int, :, :, :]).sum(axis=0) + + # Suppress, then re-enable harmless arithmetic warnings + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + var_both3[num_int, :, :] = 1. / s_inv_var_both3[num_int, :, :] + warnings.resetwarnings() + + del var_p4_int + del var_p4_int2 + + del gain_2d + + var_p4 *= (segs_4[:, :, :, :] > 0) # Zero out non-existing segments + var_r4 *= (segs_4[:, :, :, :] > 0) + + # Delete lots of arrays no longer needed + if inv_var_both4_int is not None: + del inv_var_both4_int + + if med_rates is not None: + del med_rates + + if num_seg_per_int is not None: + del num_seg_per_int + + if readnoise_2d is not None: + del readnoise_2d + + if rn_sect is not None: + del rn_sect + + if segs_4 is not None: + del segs_4 + + input_model.data = data + input_model.err = err + input_model.groupdq = groupdq + input_model.pixeldq = inpixeldq + + return var_p3, var_r3, var_p4, var_r4, var_both4, var_both3, inv_var_both4, \ + s_inv_var_p3, s_inv_var_r3, s_inv_var_both3 + + +def ramp_fit_overall( + input_model, orig_cubeshape, orig_ngroups, buffsize, fit_slopes_ans, + variances_ans, save_opt, int_times, tstart): + """ + Parameters + ---------- + input_model : data model + input data model, assumed to be of type RampModel + + orig_cubeshape : tuple + Original shape cube of input dataset + + orig_ngroups: int + Original number of groups + + buffsize : int + Size of data section (buffer) in bytes + + fit_slopes_ans : tuple + Return values from ramp_fit_slopes + + variances_ans : tuple + Return values from ramp_fit_compute_variances + + save_opt : bool + Calculate optional fitting results. + + int_times : bintable, or None + The INT_TIMES table, if it exists in the input, else None + + tstart : float + Start time. + + Return + ------ + image_info: tuple + The tuple of computed ramp fitting arrays. + + integ_info: tuple + The tuple of computed integration fitting arrays. + + opt_info: tuple + The tuple of computed optional results arrays for fitting. + """ + # Get image data information + data = input_model.data + groupdq = input_model.groupdq + + # Get instrument and exposure data + instrume = input_model.meta.instrument.name + + groupgap = input_model.meta.exposure.groupgap + nframes = input_model.meta.exposure.nframes + dropframes1 = input_model.meta.exposure.drop_frames1 + if dropframes1 is None: # set to default if missing + dropframes1 = 0 + log.debug('Missing keyword DRPFRMS1, so setting to default value of 0') + + # Get needed sizes and shapes + n_int, ngroups, nrows, ncols = data.shape + imshape = (nrows, ncols) + + # Unpack intermediate computations from preious steps + max_seg, gdq_cube_shape, effintim, f_max_seg, dq_int, num_seg_per_int = fit_slopes_ans[:6] + sat_0th_group_int, opt_res, pixeldq, inv_var, med_rates = fit_slopes_ans[6:] + + var_p3, var_r3, var_p4, var_r4, var_both4, var_both3 = variances_ans[:6] + inv_var_both4, s_inv_var_p3, s_inv_var_r3, s_inv_var_both3 = variances_ans[6:] + + slope_by_var4 = opt_res.slope_seg.copy() / var_both4 + + del var_both4 + + s_slope_by_var3 = slope_by_var4.sum(axis=1) # sum over segments (not integs) + s_slope_by_var2 = s_slope_by_var3.sum(axis=0) # sum over integrations + s_inv_var_both2 = s_inv_var_both3.sum(axis=0) + + # Compute the 'dataset-averaged' slope + # Suppress, then re-enable harmless arithmetic warnings + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + slope_dataset2 = s_slope_by_var2 / s_inv_var_both2 + warnings.resetwarnings() + + del s_slope_by_var2, s_slope_by_var3, slope_by_var4 + del s_inv_var_both2, s_inv_var_both3 + + # Replace nans in slope_dataset2 with 0 (for non-existing segments) + slope_dataset2[np.isnan(slope_dataset2)] = 0. + + # Compute the integration-specific slope + the_num = (opt_res.slope_seg * inv_var_both4).sum(axis=1) + + the_den = (inv_var_both4).sum(axis=1) + + # Suppress, then re-enable harmless arithmetic warnings + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + slope_int = the_num / the_den + warnings.resetwarnings() + + del the_num, the_den + + # Clean up ramps that are SAT on their initial groups; set ramp parameters + # for variances and slope so they will not contribute + var_p3, var_both3, slope_int, dq_int = utils.fix_sat_ramps( + sat_0th_group_int, var_p3, var_both3, slope_int, dq_int) + + if sat_0th_group_int is not None: + del sat_0th_group_int + + # Loop over data integrations to calculate integration-specific pedestal + if save_opt: + dq_slice = np.zeros((gdq_cube_shape[2], gdq_cube_shape[3]), dtype=np.uint32) + + for num_int in range(0, n_int): + dq_slice = groupdq[num_int, 0, :, :] + opt_res.ped_int[num_int, :, :] = \ + utils.calc_pedestal(num_int, slope_int, opt_res.firstf_int, + dq_slice, nframes, groupgap, dropframes1) + + del dq_slice + + # Collect optional results for output + if save_opt: + gdq_cube = groupdq + opt_res.shrink_crmag(n_int, gdq_cube, imshape, ngroups) + del gdq_cube + + # Some contributions to these vars may be NaN as they are from ramps + # having PIXELDQ=DO_NOT_USE + var_p4[np.isnan(var_p4)] = 0. + var_r4[np.isnan(var_r4)] = 0. + + # Truncate results at the maximum number of segments found + opt_res.slope_seg = opt_res.slope_seg[:, :f_max_seg, :, :] + opt_res.sigslope_seg = opt_res.sigslope_seg[:, :f_max_seg, :, :] + opt_res.yint_seg = opt_res.yint_seg[:, :f_max_seg, :, :] + opt_res.sigyint_seg = opt_res.sigyint_seg[:, :f_max_seg, :, :] + opt_res.weights = (inv_var_both4[:, :f_max_seg, :, :])**2. + opt_res.var_p_seg = var_p4[:, :f_max_seg, :, :] + opt_res.var_r_seg = var_r4[:, :f_max_seg, :, :] + + opt_info = opt_res.output_optional(effintim) + else: + opt_info = None + + if inv_var_both4 is not None: + del inv_var_both4 + + if var_p4 is not None: + del var_p4 + + if var_r4 is not None: + del var_r4 + + if inv_var is not None: + del inv_var + + if pixeldq is not None: + del pixeldq + + # Output integration-specific results to separate file + integ_info = utils.output_integ( + slope_int, dq_int, effintim, var_p3, var_r3, var_both3, int_times) + + if opt_res is not None: + del opt_res + + if slope_int is not None: + del slope_int + del var_p3 + del var_r3 + del var_both3 + if int_times is not None: + del int_times + + # Divide slopes by total (summed over all integrations) effective + # integration time to give count rates. + c_rates = slope_dataset2 / effintim + + # Compress all integration's dq arrays to create 2D PIXELDDQ array for + # primary output + final_pixeldq = utils.dq_compress_final(dq_int, n_int) + + if dq_int is not None: + del dq_int + + tstop = time.time() + + utils.log_stats(c_rates) + + log.debug('Instrument: %s', instrume) + log.debug('Number of pixels in 2D array: %d', nrows * ncols) + log.debug('Shape of 2D image: (%d, %d)' % (imshape)) + log.debug('Shape of data cube: (%d, %d, %d)' % (orig_cubeshape)) + log.debug('Buffer size (bytes): %d', buffsize) + log.debug('Number of rows per buffer: %d', nrows) + log.info('Number of groups per integration: %d', orig_ngroups) + log.info('Number of integrations: %d', n_int) + log.debug('The execution time in seconds: %f', tstop - tstart) + + # Compute the 2D variances due to Poisson and read noise + var_p2 = 1 / (s_inv_var_p3.sum(axis=0)) + var_r2 = 1 / (s_inv_var_r3.sum(axis=0)) + + # Huge variances correspond to non-existing segments, so are reset to 0 + # to nullify their contribution. + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", "invalid value.*", RuntimeWarning) + var_p2[var_p2 > 0.1 * utils.LARGE_VARIANCE] = 0. + var_r2[var_r2 > 0.1 * utils.LARGE_VARIANCE] = 0. + + # Some contributions to these vars may be NaN as they are from ramps + # having PIXELDQ=DO_NOT_USE + var_p2[np.isnan(var_p2)] = 0. + var_r2[np.isnan(var_r2)] = 0. + + # Suppress, then re-enable, harmless arithmetic warning + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + err_tot = np.sqrt(var_p2 + var_r2) + warnings.resetwarnings() + + del s_inv_var_p3 + del s_inv_var_r3 + + # Create new model for the primary output. + data = c_rates.astype(np.float32) + dq = final_pixeldq.astype(np.uint32) + var_poisson = var_p2.astype(np.float32) + var_rnoise = var_r2.astype(np.float32) + err = err_tot.astype(np.float32) + image_info = (data, dq, var_poisson, var_rnoise, err) + + return image_info, integ_info, opt_info + + +def calc_power(snr): + """ + Using the given SNR, calculate the weighting exponent, which is from + `Fixsen, D.J., Offenberg, J.D., Hanisch, R.J., Mather, J.C, Nieto, + Santisteban, M.A., Sengupta, R., & Stockman, H.S., 2000, PASP, 112, 1350`. + + Parameters + ---------- + snr : ndarray + signal-to-noise for the ramp segments, 1-D float + + Returns + ------- + pow_wt : ndarray + weighting exponent, 1-D float + """ + pow_wt = snr.copy() * 0.0 + pow_wt[np.where(snr > 5.)] = 0.4 + pow_wt[np.where(snr > 10.)] = 1.0 + pow_wt[np.where(snr > 20.)] = 3.0 + pow_wt[np.where(snr > 50.)] = 6.0 + pow_wt[np.where(snr > 100.)] = 10.0 + + return pow_wt.ravel() + + +def interpolate_power(snr): + """ + Using the given SNR, interpolate the weighting exponent, which is from + `Fixsen, D.J., Offenberg, J.D., Hanisch, R.J., Mather, J.C, Nieto, + Santisteban, M.A., Sengupta, R., & Stockman, H.S., 2000, PASP, 112, 1350`. + + Parameters + ---------- + snr : ndarray + signal-to-noise for the ramp segments, 1-D float + + Returns + ------- + pow_wt : ndarray + weighting exponent, 1-D float + """ + pow_wt = snr.copy() * 0.0 + pow_wt[np.where(snr > 5.)] = ((snr[snr > 5] - 5) / (10 - 5)) * 0.6 + 0.4 + pow_wt[np.where(snr > 10.)] = ((snr[snr > 10] - 10) / (20 - 10)) * 2.0 + 1.0 + pow_wt[np.where(snr > 20.)] = ((snr[snr > 20] - 20)) / (50 - 20) * 3.0 + 3.0 + pow_wt[np.where(snr > 50.)] = ((snr[snr > 50] - 50)) / (100 - 50) * 4.0 + 6.0 + pow_wt[np.where(snr > 100.)] = 10.0 + + return pow_wt.ravel() + + +def calc_slope(data_sect, gdq_sect, frame_time, opt_res, save_opt, rn_sect, + gain_sect, i_max_seg, ngroups, weighting, f_max_seg): + """ + Compute the slope of each segment for each pixel in the data cube section + for the current integration. Each segment has its slope fit in fit_lines(); + that slope and other quantities from the fit are added to the 'optional + result' object by append_arr() from the appropriate 'CASE' (type of segment) + in fit_next_segment(). + + Parameters + ---------- + data_sect : ndarray + section of input data cube array, 3-D float + + gdq_sect : ndarray + section of GROUPDQ data quality array, 3-D int + + frame_time : float + integration time + + opt_res : OptRes object + Contains all quantities derived from fitting all segments in all + pixels in all integrations, which will eventually be used to compute + per-integration and per-exposure quantities for all pixels. It's + also used to populate the optional product, when requested. + + save_opt : bool + save optional fitting results + + rn_sect : ndarray + read noise values for all pixels in data section + + gain_sect : ndarray + gain values for all pixels in data section + + i_max_seg : int + used for size of initial allocation of arrays for optional results; + maximum possible number of segments within the ramp, based on the + number of CR flags + + ngroups : int + number of groups per integration + + weighting : str + 'optimal' specifies that optimal weighting should be used; currently + the only weighting supported. + + f_max_seg : int + actual maximum number of segments within a ramp, based on the fitting + of all ramps; later used when truncating arrays before output. + + Returns + ------- + gdq_sect : ndarray + data quality flags for pixels in section, 3-D int + + inv_var : ndarray + values of 1/variance for good pixels, 1-D float + + opt_res : OptRes object + contains all quantities related to fitting for use in computing final + slopes, variances, etc. and is used to populate the optional output + + f_max_seg : int + actual maximum number of segments within a ramp, updated here based on + fitting ramps in the current data section; later used when truncating + arrays before output. + + num_seg : ndarray + numbers of segments for good pixels, 1-D int + """ + ngroups, nrows, ncols = data_sect.shape + npix = nrows * ncols # number of pixels in section of 2D array + + all_pix = np.arange(npix) + arange_ngroups_col = np.arange(ngroups)[:, np.newaxis] + start = np.zeros(npix, dtype=np.int32) # lowest channel in fit + + # Highest channel in fit initialized to last read + end = np.zeros(npix, dtype=np.int32) + (ngroups - 1) + + pixel_done = (end < 0) # False until processing is done + + inv_var = np.zeros(npix, dtype=np.float32) # inverse of fit variance + num_seg = np.zeros(npix, dtype=np.int32) # number of segments per pixel + + # End stack array - endpoints for each pixel + # initialize with ngroups for each pixel; set 1st channel to 0 + end_st = np.zeros((ngroups + 1, npix), dtype=np.int32) + end_st[0, :] = ngroups - 1 + + # end_heads is initially a tuple populated with every pixel that is + # either saturated or contains a cosmic ray based on the input DQ + # array, so is sized to accomodate the maximum possible number of + # pixels flagged. It is later compressed to be an array denoting + # the number of endpoints per pixel. + end_heads = np.ones(npix * ngroups, dtype=np.int32) + + # Create nominal 2D ERR array, which is 1st slice of + # avged_data_cube * readtime + err_2d_array = data_sect[0, :, :] * frame_time + + # Suppress, then re-enable, harmless arithmetic warnings + ''' + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + ''' + err_2d_array[err_2d_array < 0] = 0 + warnings.resetwarnings() + + # Frames >= start and <= end will be masked. However, the first channel + # to be included in fit will be the read in which a cosmic ray has + # been flagged + mask_2d = ((arange_ngroups_col >= start[np.newaxis, :]) & + (arange_ngroups_col <= end[np.newaxis, :])) + + end = 0 # array no longer needed + + # Section of GROUPDQ dq section, excluding bad dq values in mask + gdq_sect_r = np.reshape(gdq_sect, (ngroups, npix)) + mask_2d[gdq_sect_r != 0] = False # saturated or CR-affected + mask_2d_init = mask_2d.copy() # initial flags for entire ramp + + wh_f = np.where(np.logical_not(mask_2d)) + + these_p = wh_f[1] # coordinates of pixels flagged as False + these_r = wh_f[0] # reads of pixels flagged as False + + del wh_f + + # Populate end_st to contain the set of end points for each pixel. + # Populate end_heads to initially include every pixel that is either + # saturated or contains a cosmic ray. Skips the duplicated final group + # for saturated pixels. Saturated pixels resulting in a contiguous set + # of intervals of length 1 will later be flagged as too short + # to fit well. + for ii, val in enumerate(these_p): + if these_r[ii] != (ngroups - 1): + end_st[end_heads[these_p[ii]], these_p[ii]] = these_r[ii] + end_heads[these_p[ii]] += 1 + + # Sort and reverse array to handle the order that saturated pixels + # were added + end_st.sort(axis=0) + end_st = end_st[::-1] + + # Reformat to designate the number of endpoints per pixel; compress + # to specify number of groups per pixel + end_heads = (end_st > 0).sum(axis=0) + + # Create object to hold optional results + opt_res.init_2d(npix, i_max_seg, save_opt) + + # LS fit until 'ngroups' iterations or all pixels in + # section have been processed + for iter_num in range(ngroups): + if pixel_done.all(): + break + + # frames >= start and <= end_st will be included in fit + mask_2d = \ + ((arange_ngroups_col >= start) + & (arange_ngroups_col < (end_st[end_heads[all_pix] - 1, all_pix] + 1))) + + mask_2d[gdq_sect_r != 0] = False # RE-exclude bad group dq values + + # for all pixels, update arrays, summing slope and variance + f_max_seg, num_seg = \ + fit_next_segment(start, end_st, end_heads, pixel_done, data_sect, mask_2d, + mask_2d_init, inv_var, num_seg, opt_res, save_opt, rn_sect, + gain_sect, ngroups, weighting, f_max_seg) + + if f_max_seg is None: + f_max_seg = 1 + + arange_ngroups_col = 0 + all_pix = 0 + + return gdq_sect, inv_var, opt_res, f_max_seg, num_seg + + +def fit_next_segment(start, end_st, end_heads, pixel_done, data_sect, mask_2d, + mask_2d_init, inv_var, num_seg, opt_res, save_opt, rn_sect, + gain_sect, ngroups, weighting, f_max_seg): + """ + Call routine to LS fit masked data for a single segment for all pixels in + data section. Then categorize each pixel's fitting interval based on + interval length, and whether the interval is at the end of the array. + Update the start array, the end stack array, the end_heads array which + contains the number of endpoints. For pixels in which the fitting intervals + are long enough, the resulting slope and variance are added to the + appropriate stack arrays. The first channel to fit in a segment is either + the first group in the ramp, or a group in which a cosmic ray has been + flagged. + + Parameters + ---------- + start : ndarray + lowest channel in fit, 1-D int + + end_st : ndarray + stack array of endpoints, 2-D int + + end_heads : ndarray + number of endpoints for each pixel, 1-D int + + pixel_done : ndarray + whether each pixel's calculations are completed, 1-D bool + + data_sect : ndarray + data cube section, 3-D float + + mask_2d : ndarray + delineates which channels to fit for each pixel, 2-D bool + + mask_2d_init : ndarray + copy of intial mask_2d, 2-D bool + + inv_var : ndarray + values of 1/variance for good pixels, 1-D float + + num_seg : ndarray + numbers of segments for good pixels, 1-D int + + opt_res : OptRes object + all fitting quantities, used to compute final results + and to populate optional output product + + save_opt : bool + save optional fitting results + + rn_sect : ndarray + read noise values for all pixels in data section + + gain_sect : ndarray + gain values for all pixels in data section + + ngroups : int + number of groups per integration + + weighting : str + 'optimal' specifies that optimal weighting should be used; currently + the only weighting supported. + + f_max_seg : int + actual maximum number of segments within a ramp, updated here based on + fitting ramps in the current data section; later used when truncating + arrays before output. + + Returns + ------- + f_max_seg : int + actual maximum number of segments within a ramp, updated here based on + fitting ramps in the current data section; later used when truncating + arrays before output. + + num_seg : ndarray + numbers of segments for good pixels, 1-D int + """ + ngroups, nrows, ncols = data_sect.shape + all_pix = np.arange(nrows * ncols) + + ramp_mask_sum = mask_2d_init.sum(axis=0) + + # Compute fit quantities for the next segment of all pixels + # Each returned array below is 1D, for all npix pixels for current segment + slope, intercept, variance, sig_intercept, sig_slope = \ + fit_lines(data_sect, mask_2d, rn_sect, gain_sect, ngroups, weighting) + + end_locs = end_st[end_heads[all_pix] - 1, all_pix] + + # Set the fitting interval length; for a segment having >1 groups, this is + # the number of groups-1 + l_interval = end_locs - start + + wh_done = (start == -1) # done pixels + l_interval[wh_done] = 0 # set interval lengths for done pixels to 0 + + # Create array to set when each good pixel is classified for the current + # semiramp (to enable unclassified pixels to have their arrays updated) + got_case = np.zeros((ncols * nrows), dtype=bool) + + # Special case fit with NGROUPS being 1 or 2. + if ngroups == 1 or ngroups == 2: + return fit_short_ngroups( + ngroups, start, end_st, end_heads, pixel_done, all_pix, + inv_var, num_seg, slope, intercept, variance, sig_intercept, sig_slope, + opt_res, save_opt, mask_2d_init, ramp_mask_sum) + + # CASE: Long enough (semiramp has >2 groups), at end of ramp + wh_check = np.where((l_interval > 1) & (end_locs == ngroups - 1) & (~pixel_done)) + if len(wh_check[0]) > 0: + f_max_seg = fit_next_segment_long_end_of_ramp( + wh_check, start, end_st, end_heads, pixel_done, got_case, f_max_seg, + inv_var, num_seg, slope, intercept, variance, sig_intercept, sig_slope, + opt_res, save_opt) + + # CASE: Long enough (semiramp has >2 groups ), not at array end (meaning + # final group for this semiramp is not final group of the whole ramp) + wh_check = np.where((l_interval > 2) & (end_locs != ngroups - 1) & ~pixel_done) + if len(wh_check[0]) > 0: + f_max_seg = fit_next_segment_long_not_end_of_ramp( + wh_check, start, end_st, end_heads, pixel_done, got_case, f_max_seg, + inv_var, num_seg, slope, intercept, variance, sig_intercept, sig_slope, + opt_res, save_opt, mask_2d_init, end_locs, ngroups) + + # CASE: interval too short to fit normally (only 2 good groups) + # At end of array, NGROUPS>1, but exclude NGROUPS==2 datasets + # as they are covered in `fit_short_ngroups`. + wh_check = np.where((l_interval == 1) & (end_locs == ngroups - 1) + & (ngroups > 2) & (~pixel_done)) + + if len(wh_check[0]) > 0: + f_max_seg = fit_next_segment_short_seg_at_end( + wh_check, start, end_st, end_heads, pixel_done, got_case, f_max_seg, + inv_var, num_seg, slope, intercept, variance, sig_intercept, sig_slope, + opt_res, save_opt, mask_2d_init) + + # CASE: full-length ramp has 2 good groups not at array end + wh_check = np.where((l_interval == 2) & (ngroups > 2) + & (end_locs != ngroups - 1) & ~pixel_done) + + if len(wh_check[0]) > 0: + f_max_seg = fit_next_segment_short_seg_not_at_end( + wh_check, start, end_st, end_heads, pixel_done, got_case, f_max_seg, + inv_var, num_seg, slope, intercept, variance, sig_intercept, sig_slope, + opt_res, save_opt, mask_2d_init, end_locs, ngroups) + + # CASE: full-length ramp has a good group on 0th group of the entire ramp, + # and no later good groups. Will use single good group data as the slope. + wh_check = np.where( + mask_2d_init[0, :] & ~mask_2d_init[1, :] & (ramp_mask_sum == 1) & ~pixel_done) + + if len(wh_check[0]) > 0: + f_max_seg = fit_next_segment_only_good_0th_group( + wh_check, start, end_st, end_heads, pixel_done, got_case, f_max_seg, + inv_var, num_seg, slope, intercept, variance, sig_intercept, sig_slope, + opt_res, save_opt, mask_2d_init) + + # CASE: the segment has a good 0th group and a bad 1st group. + wh_check = np.where(mask_2d_init[0, :] & ~mask_2d_init[1, :] & ~pixel_done + & (end_locs == 1) & (start == 0)) + + if len(wh_check[0]) > 0: + fit_next_segment_good_0th_bad_1st( + wh_check, start, end_st, end_heads, got_case, ngroups) + + # CASE OTHER: all other types of segments not covered earlier. No segments + # handled here have adequate data, but the stack arrays are updated. + wh_check = np.asarray(np.where(~pixel_done & ~got_case)) + if len(wh_check[0]) > 0: + fit_next_segment_all_other(wh_check, start, end_st, end_heads, ngroups) + + return f_max_seg, num_seg + + +def fit_next_segment_all_other(wh_check, start, end_st, end_heads, ngroups): + """ + Catch all other types of segments not covered earlier. No segments + handled here have adequate data, but the stack arrays are updated. + - increment start array + - remove current end from end stack + - decrement number of ends + + Parameters + ---------- + wh_check: ndarray + pixels for current segment processing and updating, 1-D + + start : ndarray + lowest channel in fit, 1-D int + + end_st : ndarray + stack array of endpoints, 2-D int + + end_heads : ndarray + number of endpoints for each pixel, 1-D int + + ngroups: int + number of groups in exposure + """ + these_pix = wh_check[0] + start[these_pix] += 1 + start[start > ngroups - 1] = ngroups - 1 # to keep at max level + end_st[end_heads[these_pix] - 1, these_pix] = 0 + end_heads[these_pix] -= 1 + end_heads[end_heads < 0.] = 0. + + +def fit_next_segment_good_0th_bad_1st( + wh_check, start, end_st, end_heads, got_case, ngroups): + """ + The segment has a good 0th group and a bad 1st group. For the + data from the 0th good group of this segment to possibly be used as a + slope, that group must necessarily be the 0th group of the entire ramp. + It is possible to have a single 'good' group segment after the 0th group + of the ramp; in that case the 0th group and the 1st group would both have + to be CRs, and the data of the 0th group would not be included as a slope. + For a good 0th group in a ramp followed by a bad 1st group there must be + good groups later in the segment because if there were not, the segment + would be done in `fit_next_segment_only_good_0th_group`. In this situation, + since here are later good groups in the segment, those later good groups + will be used in the slope computation, and the 0th good group will not be. + As a result, for all instances of these types of segments, the data in the + initial good group will not be used in the slope calculation, but the + arrays for the indices for the ramp (end_st, etc) are appropriately + adjusted. + - increment start array + - remove current end from end stack + - decrement number of ends + + Parameters + ---------- + wh_check: ndarray + pixels for current segment processing and updating, 1-D + + start : ndarray + lowest channel in fit, 1-D int + + end_st : ndarray + stack array of endpoints, 2-D int + + end_heads : ndarray + number of endpoints for each pixel, 1-D int + + got_case: ndarray + classification of pixel for current semiramp, 1-D + + ngroups: int + number of groups in exposure + """ + these_pix = wh_check[0] + got_case[these_pix] = True + start[these_pix] += 1 + start[start > ngroups - 1] = ngroups - 1 # to keep at max level + end_st[end_heads[these_pix] - 1, these_pix] = 0 + end_heads[these_pix] -= 1 + end_heads[end_heads < 0.] = 0. + + +def fit_next_segment_only_good_0th_group( + wh_check, start, end_st, end_heads, pixel_done, got_case, f_max_seg, + inv_var, num_seg, slope, intercept, variance, sig_intercept, sig_slope, + opt_res, save_opt, mask_2d_init): + """ + Full-length ramp has a good group on 0th group of the entire ramp, + and no later good groups. Will use single good group data as the slope. + - set start to -1 to designate all fitting done + - remove current end from end stack + - set number of end to 0 + - add slopes and variances to running sums + - set pixel_done to True to designate all fitting done + + Parameters + ---------- + wh_check: ndarray + pixels for current segment processing and updating, 1-D + + start : ndarray + lowest channel in fit, 1-D int + + end_st : ndarray + stack array of endpoints, 2-D int + + end_heads : ndarray + number of endpoints for each pixel, 1-D int + + pixel_done : ndarray + whether each pixel's calculations are completed, 1-D bool + + got_case: ndarray + classification of pixel for current semiramp, 1-D + + f_max_seg : int + actual maximum number of segments within a ramp, updated here based on + fitting ramps in the current data section; later used when truncating + arrays before output. + + inv_var : ndarray + values of 1/variance for good pixels, 1-D float + + num_seg : ndarray + numbers of segments for good pixels, 1-D int + + slope: ndarray + weighted slope for current iteration's pixels for data section, 1-D + float + + intercept: ndarray + y-intercepts from fit for data section, 1-D float + + variance: ndarray + variance of residuals for fit for data section, 1-D float + + sig_intercept: ndarray + sigma of y-intercepts from fit for data section, 1-D float + + sig_slope: ndarray + sigma of slopes from fit for data section (for a single segment), 1-D + float + + opt_res : OptRes object + all fitting quantities, used to compute final results + and to populate optional output product + + save_opt : bool + save optional fitting results + + mask_2d_init : ndarray + copy of intial mask_2d, 2-D bool + + Returns + ------- + f_max_seg : int + actual maximum number of segments within a ramp, updated here based on + fitting ramps in the current data section; later used when truncating + arrays before output. + """ + these_pix = wh_check[0] + got_case[these_pix] = True + + start[these_pix] = -1 + end_st[end_heads[these_pix] - 1, these_pix] = 0 + end_heads[these_pix] = 0 + pixel_done[these_pix] = True # all processing for pixel is completed + inv_var[these_pix] += 1.0 / variance[these_pix] + + # Append results to arrays + opt_res.append_arr(num_seg, these_pix, intercept, slope, + sig_intercept, sig_slope, inv_var, save_opt) + + num_seg[these_pix] += 1 + f_max_seg = max(f_max_seg, num_seg.max()) + + return f_max_seg + + +def fit_next_segment_short_seg_not_at_end( + wh_check, start, end_st, end_heads, pixel_done, got_case, f_max_seg, + inv_var, num_seg, slope, intercept, variance, sig_intercept, sig_slope, + opt_res, save_opt, mask_2d_init, end_locs, ngroups): + """ + Special case + Full-length ramp has 2 good groups not at array end + - use the 2 good reads to get the slope + - set start to -1 to designate all fitting done + - remove current end from end stack + - set number of end to 0 + - add slopes and variances to running sums + - set pixel_done to True to designate all fitting done + For segments of this type, the final good group in the segment is + followed by a group that is flagged as a CR and/or SAT and is not the + final group in the ramp, and the variable `l_interval` used below is + equal to 2, which is the number of the segment's groups. + + Parameters + ---------- + wh_check: ndarray + pixels for current segment processing and updating, 1-D + + start : ndarray + lowest channel in fit, 1-D int + + end_st : ndarray + stack array of endpoints, 2-D int + + end_heads : ndarray + number of endpoints for each pixel, 1-D int + + pixel_done : ndarray + whether each pixel's calculations are completed, 1-D bool + + got_case: ndarray + classification of pixel for current semiramp, 1-D + + f_max_seg : int + actual maximum number of segments within a ramp, updated here based on + fitting ramps in the current data section; later used when truncating + arrays before output. + + inv_var : ndarray + values of 1/variance for good pixels, 1-D float + + num_seg : ndarray + numbers of segments for good pixels, 1-D int + + slope: ndarray + weighted slope for current iteration's pixels for data section, 1-D + float + + intercept: ndarray + y-intercepts from fit for data section, 1-D float + + variance: ndarray + variance of residuals for fit for data section, 1-D float + + sig_intercept: ndarray + sigma of y-intercepts from fit for data section, 1-D float + + sig_slope: ndarray + sigma of slopes from fit for data section (for a single segment), 1-D + float + + opt_res : OptRes object + all fitting quantities, used to compute final results + and to populate optional output product + + save_opt : bool + save optional fitting results + + mask_2d_init : ndarray + copy of intial mask_2d, 2-D bool + + end_locs: ndarray + end locations, 1-D + + ngroups: int + number of groups in exposure + + Returns + ------- + f_max_seg : int + actual maximum number of segments within a ramp, updated here based on + fitting ramps in the current data section; later used when truncating + arrays before output. + """ + # Copy mask, as will modify when calculating the number of later good groups + c_mask_2d_init = mask_2d_init.copy() + + these_pix = wh_check[0] + got_case[these_pix] = True + + # Suppress, then re-enable, harmless arithmetic warnings + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + inv_var[these_pix] += 1.0 / variance[these_pix] + warnings.resetwarnings() + + # create array: 0...ngroups-1 in a column for each pixel + arr_ind_all = np.array( + [np.arange(ngroups), ] * c_mask_2d_init.shape[1]).transpose() + wh_c_start_all = np.zeros(mask_2d_init.shape[1], dtype=np.uint8) + wh_c_start_all[these_pix] = start[these_pix] + + # set to False all groups before start group + c_mask_2d_init[arr_ind_all < wh_c_start_all] = 0 + tot_good_groups = c_mask_2d_init.sum(axis=0) + + # Select pixels having at least 2 later good groups (these later good + # groups are a segment whose slope will be calculated) + wh_more = np.where(tot_good_groups[these_pix] > 1) + pix_more = these_pix[wh_more] + start[pix_more] = end_locs[pix_more] + end_st[end_heads[pix_more] - 1, pix_more] = 0 + end_heads[pix_more] -= 1 + + # Select pixels having less than 2 later good groups (these later good + # groups will not be used) + wh_only = np.where(tot_good_groups[these_pix] <= 1) + pix_only = these_pix[wh_only] + start[pix_only] = -1 + end_st[end_heads[pix_only] - 1, pix_only] = 0 + end_heads[pix_only] = 0 + pixel_done[pix_only] = True # all processing for pixel is completed + end_heads[(end_heads < 0.)] = 0. + + # Append results to arrays + opt_res.append_arr(num_seg, these_pix, intercept, slope, + sig_intercept, sig_slope, inv_var, save_opt) + + num_seg[these_pix] += 1 + f_max_seg = max(f_max_seg, num_seg.max()) + + return f_max_seg + + +def fit_next_segment_short_seg_at_end( + wh_check, start, end_st, end_heads, pixel_done, got_case, f_max_seg, + inv_var, num_seg, slope, intercept, variance, sig_intercept, sig_slope, + opt_res, save_opt, mask_2d_init): + """ + Interval too short to fit normally (only 2 good groups) + At end of array, NGROUPS>1, but exclude NGROUPS==2 datasets + as they are covered in `fit_short_groups`. + - set start to -1 to designate all fitting done + - remove current end from end stack + - set number of ends to 0 + - add slopes and variances to running sums + - set pixel_done to True to designate all fitting done + For segments of this type, the final good group is the final group in the + ramp, and the variable `l_interval` used below = 1, and the number of + groups in the segment = 2 + + Parameters + ---------- + wh_check: ndarray + pixels for current segment processing and updating, 1-D + + start : ndarray + lowest channel in fit, 1-D int + + end_st : ndarray + stack array of endpoints, 2-D int + + end_heads : ndarray + number of endpoints for each pixel, 1-D int + + pixel_done : ndarray + whether each pixel's calculations are completed, 1-D bool + + got_case: ndarray + classification of pixel for current semiramp, 1-D + + f_max_seg : int + actual maximum number of segments within a ramp, updated here based on + fitting ramps in the current data section; later used when truncating + arrays before output. + + inv_var : ndarray + values of 1/variance for good pixels, 1-D float + + num_seg : ndarray + numbers of segments for good pixels, 1-D int + + slope: ndarray + weighted slope for current iteration's pixels for data section, 1-D + float + + intercept: ndarray + y-intercepts from fit for data section, 1-D float + + variance: ndarray + variance of residuals for fit for data section, 1-D float + + sig_intercept: ndarray + sigma of y-intercepts from fit for data section, 1-D float + + sig_slope: ndarray + sigma of slopes from fit for data section (for a single segment), 1-D + float + + opt_res : OptRes object + all fitting quantities, used to compute final results + and to populate optional output product + + save_opt : bool + save optional fitting results + + mask_2d_init : ndarray + copy of intial mask_2d, 2-D bool + + Returns + ------- + f_max_seg : int + actual maximum number of segments within a ramp, updated here based on + fitting ramps in the current data section; later used when truncating + arrays before output. + """ + # Require that pixels to be processed here have at least 1 good group out + # of the final 2 groups (these ramps have 2 groups and are at the end of + # the array). + wh_list = [] + + num_wh = len(wh_check[0]) + for ii in range(num_wh): # locate pixels with at least 1 good group + this_pix = wh_check[0][ii] + sum_final_2 = mask_2d_init[start[this_pix]:, this_pix].sum() + + if sum_final_2 > 0: + wh_list.append(wh_check[0][ii]) # add to list to be fit + + if len(wh_list) > 0: + these_pix = np.asarray(wh_list) + got_case[these_pix] = True + + start[these_pix] = -1 + end_st[end_heads[these_pix] - 1, these_pix] = 0 + end_heads[these_pix] = 0 + pixel_done[these_pix] = True + + g_pix = these_pix[variance[these_pix] > 0.] # good pixels + + if len(g_pix) > 0: + inv_var[g_pix] += 1.0 / variance[g_pix] + + # Append results to arrays + opt_res.append_arr(num_seg, g_pix, intercept, slope, + sig_intercept, sig_slope, inv_var, save_opt) + + num_seg[g_pix] += 1 + f_max_seg = max(f_max_seg, num_seg.max()) + + return f_max_seg + + +def fit_next_segment_long_not_end_of_ramp( + wh_check, start, end_st, end_heads, pixel_done, got_case, f_max_seg, + inv_var, num_seg, slope, intercept, variance, sig_intercept, sig_slope, + opt_res, save_opt, mask_2d_init, end_locs, ngroups): + """ + Special case fitting long segment at the end of ramp. + Long enough (semiramp has >2 groups ), not at array end (meaning + final group for this semiramp is not final group of the whole ramp) + - remove current end from end stack + - decrement number of ends + - add slopes and variances to running sums + For segments of this type, the final good group in the segment is a CR + and/or SAT and is not the final group in the ramp, and the variable + `l_interval` used below is equal to the number of the segment's groups. + + Parameters + ---------- + wh_check: ndarray + pixels for current segment processing and updating, 1-D + + start : ndarray + lowest channel in fit, 1-D int + + end_st : ndarray + stack array of endpoints, 2-D int + + end_heads : ndarray + number of endpoints for each pixel, 1-D int + + pixel_done : ndarray + whether each pixel's calculations are completed, 1-D bool + + got_case: ndarray + classification of pixel for current semiramp, 1-D + + f_max_seg : int + actual maximum number of segments within a ramp, updated here based on + fitting ramps in the current data section; later used when truncating + arrays before output. + + inv_var : ndarray + values of 1/variance for good pixels, 1-D float + + num_seg : ndarray + numbers of segments for good pixels, 1-D int + + slope: ndarray + weighted slope for current iteration's pixels for data section, 1-D + float + + intercept: ndarray + y-intercepts from fit for data section, 1-D float + + variance: ndarray + variance of residuals for fit for data section, 1-D float + + sig_intercept: ndarray + sigma of y-intercepts from fit for data section, 1-D float + + sig_slope: ndarray + sigma of slopes from fit for data section (for a single segment), 1-D + float + + opt_res : OptRes object + all fitting quantities, used to compute final results + and to populate optional output product + + save_opt : bool + save optional fitting results + + end_locs: ndarray + end locations, 1-D + + mask_2d_init : ndarray + copy of intial mask_2d, 2-D bool + + ngroups: int + number of groups in exposure + + Returns + ------- + f_max_seg : int + actual maximum number of segments within a ramp, updated here based on + fitting ramps in the current data section; later used when truncating + arrays before output. + """ + these_pix = wh_check[0] + got_case[these_pix] = True + + start[these_pix] = end_locs[these_pix] + end_st[end_heads[these_pix] - 1, these_pix] = 0 + end_heads[these_pix] -= 1 + end_heads[end_heads < 0.] = 0. + + g_pix = these_pix[variance[these_pix] > 0.] # good pixels + + if len(g_pix) > 0: + inv_var[g_pix] += 1.0 / variance[g_pix] + + # Append results to arrays + opt_res.append_arr(num_seg, g_pix, intercept, slope, sig_intercept, + sig_slope, inv_var, save_opt) + + num_seg[g_pix] += 1 + f_max_seg = max(f_max_seg, num_seg.max()) + + # If there are pixels with no later good groups, update stack + # arrays accordingly + c_mask_2d_init = mask_2d_init.copy() + + # create array: 0...ngroups-1 in a column for each pixel + arr_ind_all = np.array( + [np.arange(ngroups), ] * c_mask_2d_init.shape[1]).transpose() + + wh_c_start_all = np.zeros(c_mask_2d_init.shape[1], dtype=np.uint8) + wh_c_start_all[g_pix] = start[g_pix] + + # set to False all groups before start group + c_mask_2d_init[arr_ind_all < wh_c_start_all] = False + + # select pixels having all groups False from start to ramp end + wh_rest_false = np.where(c_mask_2d_init.sum(axis=0) == 0) + if len(wh_rest_false[0]) > 0: + pix_rest_false = wh_rest_false[0] + start[pix_rest_false] = -1 + end_st[end_heads[pix_rest_false] - 1, pix_rest_false] = 0 + end_heads[pix_rest_false] = 0 + pixel_done[pix_rest_false] = True # all processing is complete + + return f_max_seg + + +def fit_next_segment_long_end_of_ramp( + wh_check, start, end_st, end_heads, pixel_done, got_case, f_max_seg, + inv_var, num_seg, slope, intercept, variance, sig_intercept, sig_slope, + opt_res, save_opt): + """ + Long enough (semiramp has >2 groups), at end of ramp + - set start to -1 to designate all fitting done + - remove current end from end stack + - set number of ends to 0 + - add slopes and variances to running sums + For segments of this type, the final good group is the final group in the + ramp, and the variable `l_interval` used below is equal to the number of + the segment's groups minus 1. + + Parameters + ---------- + wh_check: ndarray + pixels for current segment processing and updating, 1-D + + start : ndarray + lowest channel in fit, 1-D int + + end_st : ndarray + stack array of endpoints, 2-D int + + end_heads : ndarray + number of endpoints for each pixel, 1-D int + + pixel_done : ndarray + whether each pixel's calculations are completed, 1-D bool + + got_case: ndarray + classification of pixel for current semiramp, 1-D + + f_max_seg : int + actual maximum number of segments within a ramp, updated here based on + fitting ramps in the current data section; later used when truncating + arrays before output. + + inv_var : ndarray + values of 1/variance for good pixels, 1-D float + + num_seg : ndarray + numbers of segments for good pixels, 1-D int + + slope: ndarray + weighted slope for current iteration's pixels for data section, 1-D + float + + intercept: ndarray + y-intercepts from fit for data section, 1-D float + + variance: ndarray + variance of residuals for fit for data section, 1-D float + + sig_intercept: ndarray + sigma of y-intercepts from fit for data section, 1-D float + + sig_slope: ndarray + sigma of slopes from fit for data section (for a single segment), 1-D + float + + opt_res : OptRes object + all fitting quantities, used to compute final results + and to populate optional output product + + save_opt : bool + save optional fitting results + + Returns + ------- + f_max_seg : int + actual maximum number of segments within a ramp, updated here based on + fitting ramps in the current data section; later used when truncating + arrays before output. + """ + these_pix = wh_check[0] + start[these_pix] = -1 # all processing for this pixel is completed + end_st[end_heads[these_pix] - 1, these_pix] = 0 + end_heads[these_pix] = 0 + pixel_done[these_pix] = True # all processing for pixel is completed + got_case[these_pix] = True + + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", "invalid value.*", RuntimeWarning) + g_pix = these_pix[variance[these_pix] > 0.] # good pixels + if len(g_pix) > 0: + inv_var[g_pix] += 1.0 / variance[g_pix] + + # Append results to arrays + opt_res.append_arr(num_seg, g_pix, intercept, slope, sig_intercept, + sig_slope, inv_var, save_opt) + + num_seg[g_pix] += 1 + f_max_seg = max(f_max_seg, num_seg.max()) + return f_max_seg + + +def fit_short_ngroups( + ngroups, start, end_st, end_heads, pixel_done, all_pix, + inv_var, num_seg, slope, intercept, variance, sig_intercept, sig_slope, + opt_res, save_opt, mask_2d_init, ramp_mask_sum): + """ + Special case fitting for short ngroups fit. + + Parameters + ---------- + ngroups: int + number of groups in exposure + + start : ndarray + lowest channel in fit, 1-D int + + end_st : ndarray + stack array of endpoints, 2-D int + + end_heads : ndarray + number of endpoints for each pixel, 1-D int + + pixel_done : ndarray + whether each pixel's calculations are completed, 1-D bool + + all_pix: ndarray + all pixels in image, 1-D + + inv_var : ndarray + values of 1/variance for good pixels, 1-D float + + num_seg : ndarray + numbers of segments for good pixels, 1-D int + + slope: ndarray + weighted slope for current iteration's pixels for data section, 1-D + float + + intercept: ndarray + y-intercepts from fit for data section, 1-D float + + variance: float, ndarray + variance of residuals for fit for data section, 1-D + + sig_intercept: ndarray + sigma of y-intercepts from fit for data section, 1-D float + + sig_slope: ndarray + sigma of slopes from fit for data section (for a single segment), 1-D + float + + opt_res : OptRes object + all fitting quantities, used to compute final results + and to populate optional output product + + save_opt : bool + save optional fitting results + + mask_2d_init : ndarray + copy of intial mask_2d, 2-D bool + + ramp_mask_sum: ndarray + number of channels to fit for each pixel, 1-D int + + Returns + ------- + f_max_seg : int + actual maximum number of segments within a ramp, updated here based on + fitting ramps in the current data section; later used when truncating + arrays before output. + + num_seg : ndarray + numbers of segments for good pixels, 1-D int + """ + + # Dataset has NGROUPS=2, so special fitting is done for all pixels. + # All segments are at the end of the array. + # - set start to -1 to designate all fitting done + # - remove current end from end stack + # - set number of ends to 0 + # - add slopes and variances to running sums + # - set pixel_done to True to designate all fitting done + if ngroups == 2: + start[all_pix] = -1 + end_st[end_heads[all_pix] - 1, all_pix] = 0 + end_heads[all_pix] = 0 + pixel_done[all_pix] = True + + g_pix = all_pix[variance[all_pix] > 0.] + if len(g_pix) > 0: + inv_var[g_pix] += 1.0 / variance[g_pix] + + opt_res.append_arr(num_seg, g_pix, intercept, slope, sig_intercept, + sig_slope, inv_var, save_opt) + + num_seg[g_pix] = 1 + + return 1, num_seg + + # Dataset has NGROUPS=1 ; so special fitting is done for all pixels + # and all intervals are at the end of the array. + # - set start to -1 to designate all fitting done + # - remove current end from end stack + # - set number of ends to 0 + # - add slopes and variances to running sums + # - set pixel_done to True to designate all fitting done + start[all_pix] = -1 + end_st[end_heads[all_pix] - 1, all_pix] = 0 + end_heads[all_pix] = 0 + pixel_done[all_pix] = True + + wh_check = np.where(mask_2d_init[0, :] & (ramp_mask_sum == 1)) + if len(wh_check[0]) > 0: + g_pix = wh_check[0] + + # Ignore all pixels having no good groups (so the single group is bad) + if len(g_pix) > 0: + inv_var[g_pix] += 1.0 / variance[g_pix] + + # Append results to arrays + opt_res.append_arr(num_seg, g_pix, intercept, slope, sig_intercept, + sig_slope, inv_var, save_opt) + + num_seg[g_pix] = 1 + + return 1, num_seg + + +def fit_lines(data, mask_2d, rn_sect, gain_sect, ngroups, weighting): + """ + Do linear least squares fit to data cube in this integration for a single + segment for all pixels. In addition to applying the mask due to identified + cosmic rays, the data is also masked to exclude intervals that are too short + to fit well. The first channel to fit in a segment is either the first group + in the ramp, or a group in which a cosmic ray has been flagged. + + Parameters + ---------- + data : ndarray + array of values for current data section, 3-D float + + mask_2d : ndarray + delineates which channels to fit for each pixel, 2-D bool + + rn_sect : ndarray + read noise values for all pixels in data section + + gain_sect : ndarray + gain values for all pixels in data section + + ngroups : int + number of groups per integration + + weighting : str + 'optimal' specifies that optimal weighting should be used; currently + the only weighting supported. + + Returns + ------- + Note - all of these pertain to a single segment (hence '_s') + + slope_s : ndarray + 1-D weighted slope for current iteration's pixels for data section + + intercept_s : ndarray + 1-D y-intercepts from fit for data section + + variance_s : ndarray + 1-D variance of residuals for fit for data section + + sig_intercept_s : ndarray + 1-D sigma of y-intercepts from fit for data section + + sig_slope_s : ndarray + 1-D sigma of slopes from fit for data section (for a single segment) + + """ + # To ensure that the first channel to be fit is the cosmic-ray-affected + # group, the channel previous to each channel masked as good is + # also masked as good. This is only for the local purpose of setting + # the first channel, and will not propagate beyond this current function + # call. + c_mask_2d = mask_2d.copy() + wh_mask_2d = np.where(c_mask_2d) + c_mask_2d[np.maximum(wh_mask_2d[0] - 1, 0), wh_mask_2d[1]] = True + + del wh_mask_2d + + # num of reads/pixel unmasked + nreads_1d = c_mask_2d.astype(np.int16).sum(axis=0) + npix = c_mask_2d.shape[1] + + slope_s = np.zeros(npix, dtype=np.float32) + variance_s = np.zeros(npix, dtype=np.float32) + intercept_s = np.zeros(npix, dtype=np.float32) + sig_intercept_s = np.zeros(npix, dtype=np.float32) + sig_slope_s = np.zeros(npix, dtype=np.float32) + + # Calculate slopes etc. for datasets having either 1 or 2 groups per + # integration, and return + if ngroups == 1: # process all pixels in 1 group/integration dataset + slope_s, intercept_s, variance_s, sig_intercept_s, sig_slope_s = \ + fit_1_group(slope_s, intercept_s, variance_s, sig_intercept_s, + sig_slope_s, npix, data, c_mask_2d) + + return slope_s, intercept_s, variance_s, sig_intercept_s, sig_slope_s + + if ngroups == 2: # process all pixels in 2 group/integration dataset + rn_sect_1d = rn_sect.reshape(npix) + slope_s, intercept_s, variance_s, sig_intercept_s, sig_slope_s = \ + fit_2_group(slope_s, intercept_s, variance_s, sig_intercept_s, + sig_slope_s, npix, data, c_mask_2d, rn_sect_1d) + + return slope_s, intercept_s, variance_s, sig_intercept_s, sig_slope_s + + # reshape data_masked + data_masked = data * np.reshape(c_mask_2d, data.shape) + data_masked = np.reshape(data_masked, (data_masked.shape[0], npix)) + + # For datasets having >2 groups/integration, for any semiramp in which the + # 0th group is good and the 1st group is bad, determine whether or not to + # use the 0th group. + wh_pix_1r = np.where(c_mask_2d[0, :] & (np.logical_not(c_mask_2d[1, :]))) + + if len(wh_pix_1r[0]) > 0: + slope_s, intercept_s, variance_s, sig_intercept_s, \ + sig_slope_s = fit_single_read(slope_s, intercept_s, variance_s, + sig_intercept_s, sig_slope_s, npix, + data, wh_pix_1r) + + del wh_pix_1r + + # For datasets having >2 groups/integrations, for any semiramp in which only + # the 0th and 1st group are good, set slope, etc + wh_pix_2r = np.where(c_mask_2d.sum(axis=0) == 2) # ramps with 2 good groups + slope_s, intercept_s, variance_s, sig_slope_s, sig_intercept_s = \ + fit_double_read(c_mask_2d, wh_pix_2r, data_masked, slope_s, intercept_s, + variance_s, sig_slope_s, sig_intercept_s, rn_sect) + + del wh_pix_2r + + # Select ramps having >2 good groups + wh_pix_to_use = np.where(c_mask_2d.sum(axis=0) > 2) + + good_pix = wh_pix_to_use[0] # Ramps with >2 good groups + data_masked = data_masked[:, good_pix] + + del wh_pix_to_use + + xvalues = np.arange(data_masked.shape[0])[:, np.newaxis] * c_mask_2d + xvalues = xvalues[:, good_pix] # set to those pixels to be used + + c_mask_2d = c_mask_2d[:, good_pix] + nreads_1d = nreads_1d[good_pix] + + if weighting.lower() == 'optimal': # fit using optimal weighting + # get sums from optimal weighting + sumx, sumxx, sumxy, sumy, nreads_wtd, xvalues = \ + calc_opt_sums(rn_sect, gain_sect, data_masked, c_mask_2d, xvalues, good_pix) + + slope, intercept, sig_slope, sig_intercept = \ + calc_opt_fit(nreads_wtd, sumxx, sumx, sumxy, sumy) + + variance = sig_slope**2. # variance due to fit values + + elif weighting.lower() == 'unweighted': # fit using unweighted weighting + # get sums from unweighted weighting + sumx, sumxx, sumxy, sumy = calc_unwtd_sums(data_masked, xvalues) + + slope, intercept, sig_slope, sig_intercept, line_fit =\ + calc_unwtd_fit(xvalues, nreads_1d, sumxx, sumx, sumxy, sumy) + + denominator = nreads_1d * sumxx - sumx**2 + + # In case this branch is ever used again, disable, and then re-enable + # harmless arithmetic warrnings + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + variance = nreads_1d / denominator + warnings.resetwarnings() + + denominator = 0 + + else: # unsupported weighting type specified + log.error('FATAL ERROR: unsupported weighting type specified.') + + slope_s[good_pix] = slope + variance_s[good_pix] = variance + intercept_s[good_pix] = intercept + sig_intercept_s[good_pix] = sig_intercept + sig_slope_s[good_pix] = sig_slope + + return slope_s, intercept_s, variance_s, sig_intercept_s, sig_slope_s + + +def fit_single_read(slope_s, intercept_s, variance_s, sig_intercept_s, + sig_slope_s, npix, data, wh_pix_1r): + """ + For datasets having >2 groups/integrations, for any semiramp in which the + 0th group is good and the 1st group is either SAT or CR, set slope, etc. + + Parameters + ---------- + slope_s : ndarray + 1-D weighted slope for current iteration's pixels for data section + + intercept_s : ndarray + 1-D y-intercepts from fit for data section + + variance_s : ndarray + 1-D variance of residuals for fit for data section + + sig_intercept_s : ndarray + 1-D sigma of y-intercepts from fit for data section + + sig_slope_s : ndarray + 1-D sigma of slopes from fit for data section + + npix : int + number of pixels in 2D array + + data : float + array of values for current data section + + wh_pix_1r : tuple + locations of pixels whose only good group is the 0th group + + Returns + ------- + slope_s : ndarray + 1-D weighted slope for current iteration's pixels for data section + + intercept_s : ndarray + 1-D y-intercepts from fit for data section + + variance_s : ndarray + 1-D variance of residuals for fit for data section + + sig_slope_s : ndarray + 1-D sigma of slopes from fit for data section + + sig_intercept_s : ndarray + 1-D sigma of y-intercepts from fit for data section + """ + data0_slice = data[0, :, :].reshape(npix) + slope_s[wh_pix_1r] = data0_slice[wh_pix_1r] + + # The following arrays will have values correctly calculated later; for + # now they are just place-holders + variance_s[wh_pix_1r] = utils.LARGE_VARIANCE + sig_slope_s[wh_pix_1r] = 0. + intercept_s[wh_pix_1r] = 0. + sig_intercept_s[wh_pix_1r] = 0. + + return slope_s, intercept_s, variance_s, sig_slope_s, sig_intercept_s + + +def fit_double_read(mask_2d, wh_pix_2r, data_masked, slope_s, intercept_s, + variance_s, sig_slope_s, sig_intercept_s, rn_sect): + """ + Process all semi-ramps having exactly 2 good groups. May need to optimize + later to remove loop over pixels. + + Parameters + ---------- + mask_2d : ndarray + 2-D bool delineates which channels to fit for each pixel + + wh_pix_2r : tuple + locations of pixels whose only good groups are the 0th and the 1st + + data_masked : ndarray + 2-D masked values for all pixels in data section + + slope_s : ndarray + 1-D weighted slope for current iteration's pixels for data section + + intercept_s : ndarray + 1-D y-intercepts from fit for data section + + variance_s : ndarray + 1-D variance of residuals for fit for data section + + sig_slope_s : ndarray + 1-D sigma of slopes from fit for data section + + sig_intercept_s : ndarray + 1-D sigma of y-intercepts from fit for data section + + rn_sect : ndarray + 2-D read noise values for all pixels in data section + + Returns + ------- + slope_s : ndarray + 1-D weighted slope for current iteration's pixels for data section + + intercept_s : ndarray + 1-D y-intercepts from fit for data section + + variance_s : ndarray + 1-D variance of residuals for fit for data section + + sig_slope_s : ndarray + 1-D sigma of slopes from fit for data section + + sig_intercept_s : ndarray + 1-D sigma of y-intercepts from fit for data section + """ + rn_sect_flattened = rn_sect.flatten() + + for ff in range(len(wh_pix_2r[0])): # loop over the pixels + pixel_ff = wh_pix_2r[0][ff] # pixel index (1d) + + rn = rn_sect_flattened[pixel_ff] # read noise for this pixel + + read_nums = np.where(mask_2d[:, pixel_ff]) + second_read = read_nums[0][1] + data_ramp = data_masked[:, pixel_ff] * mask_2d[:, pixel_ff] + data_semi = data_ramp[mask_2d[:, pixel_ff]] # picks only the 2 + diff_data = data_semi[1] - data_semi[0] + + slope_s[pixel_ff] = diff_data + intercept_s[pixel_ff] = \ + data_semi[1] * (1. - second_read) + data_semi[0] * second_read # by geometry + variance_s[pixel_ff] = 2.0 * rn * rn + sig_slope_s[pixel_ff] = np.sqrt(2) * rn + sig_intercept_s[pixel_ff] = np.sqrt(2) * rn + + return slope_s, intercept_s, variance_s, sig_slope_s, sig_intercept_s + + +def calc_unwtd_fit(xvalues, nreads_1d, sumxx, sumx, sumxy, sumy): + """ + Do linear least squares fit to data cube in this integration, using + unweighted fits to the segments. Currently not supported. + + Parameters + ---------- + xvalues : ndarray + 1-D int indices of valid pixel values for all groups + + nreads_1d : ndarray + 1-D int number of reads in an integration + + sumxx : float + sum of squares of xvalues + + sumx : float + sum of xvalues + + sumxy : float + sum of product of xvalues and data + + sumy : float + sum of data + + Returns + ------- + slope : ndarray + 1-D weighted slope for current iteration's pixels for data section + + intercept : ndarray + 1-D y-intercepts from fit for data section + + sig_slope : ndarray + 1-D sigma of slopes from fit for data section + + sig_intercept : ndarray + 1-D sigma of y-intercepts from fit for data section + + line_fit : ndarray + 1-D values of fit using slope and intercept + """ + + denominator = nreads_1d * sumxx - sumx**2 + + # In case this branch is ever used again, suppress, and then re-enable + # harmless arithmetic warnings + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + slope = (nreads_1d * sumxy - sumx * sumy) / denominator + intercept = (sumxx * sumy - sumx * sumxy) / denominator + sig_intercept = (sumxx / denominator)**0.5 + sig_slope = (nreads_1d / denominator)**0.5 + warnings.resetwarnings() + + line_fit = (slope * xvalues) + intercept + + return slope, intercept, sig_slope, sig_intercept, line_fit + + +def calc_opt_fit(nreads_wtd, sumxx, sumx, sumxy, sumy): + """ + Do linear least squares fit to data cube in this integration for a single + semi-ramp for all pixels, using optimally weighted fits to the semi_ramps. + The weighting uses the formulation by Fixsen (Fixsen et al, PASP, 112, 1350). + Note - these weights, sigmas, and variances pertain only to the fitting, and + the variances are *NOT* the variances of the slope due to noise. + + Parameters + ---------- + nreads_wtd : ndarray + sum of product of data and optimal weight, 1-D float + + sumxx : ndarray + sum of squares of xvalues, 1-D float + + sumx : ndarray + sum of xvalues, 1-D float + + sumxy : ndarray + sum of product of xvalues and data, 1-D float + + sumy : ndarray + sum of data, 1-D float + + Returns + ------- + slope : ndarray + weighted slope for current iteration's pixels for data section, 1-D + float + + intercept : ndarray + y-intercepts from fit for data section, 1-D float + + sig_slope : ndarray + sigma of slopes from fit for data section, 1-D float + + sig_intercept : ndarray + sigma of y-intercepts from fit for data section, 1-D float + """ + denominator = nreads_wtd * sumxx - sumx**2 + + # Suppress, and then re-enable harmless arithmetic warnings + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + + slope = (nreads_wtd * sumxy - sumx * sumy) / denominator + intercept = (sumxx * sumy - sumx * sumxy) / denominator + sig_intercept = (sumxx / denominator)**0.5 + sig_slope = (nreads_wtd / denominator)**0.5 # STD of the slope's fit + + warnings.resetwarnings() + + return slope, intercept, sig_slope, sig_intercept + + +def fit_1_group(slope_s, intercept_s, variance_s, sig_intercept_s, + sig_slope_s, npix, data, mask_2d): + """ + This function sets the fitting arrays for datasets having only 1 group + per integration. + + Parameters + ---------- + slope_s : ndarray + weighted slope for current iteration's pixels for data section, 1-D float + + intercept_s : ndarray + y-intercepts from fit for data section, 1-D float + + variance_s : ndarray + variance of residuals for fit for data section, 1-D float + + sig_intercept_s : ndarray + sigma of y-intercepts from fit for data section, 1-D float + + sig_slope_s : ndarray + sigma of slopes from fit for data section, 1-D float + + npix : int + number of pixels in 2d array + + data : float + array of values for current data section + + mask_2d : ndarray + delineates which channels to fit for each pixel, 2-D bool + + Returns + ------- + slope_s : ndarray + weighted slope for current iteration's pixels for data section, 1-D float + + intercept_s : ndarray + y-intercepts from fit for data section, 1-D float + + variance_s : ndarray + variance of residuals for fit for data section, 1-D float + + sig_intercept_s : ndarray + sigma of y-intercepts from fit for data section, 1-D float + + sig_slope_s : ndarray + sigma of slopes from fit for data section, 1-D float + """ + # For pixels not saturated, recalculate the slope as the value of the SCI + # data in that group, which will later be divided by the group exposure + # time to give the count rate. Recalculate other fit quantities to be + # benign. + slope_s = data[0, :, :].reshape(npix) + + # The following arrays will have values correctly calculated later; for + # now they are just place-holders + variance_s = np.zeros(npix, dtype=np.float32) + utils.LARGE_VARIANCE + sig_slope_s = slope_s * 0. + intercept_s = slope_s * 0. + sig_intercept_s = slope_s * 0. + + # For saturated pixels, overwrite slope with benign values. + wh_sat0 = np.where(np.logical_not(mask_2d[0, :])) + + if len(wh_sat0[0]) > 0: + sat_pix = wh_sat0[0] + slope_s[sat_pix] = 0. + + return slope_s, intercept_s, variance_s, sig_intercept_s, sig_slope_s + + +def fit_2_group(slope_s, intercept_s, variance_s, sig_intercept_s, + sig_slope_s, npix, data, mask_2d, rn_sect_1d): + """ + This function sets the fitting arrays for datasets having only 2 groups + per integration. + + Parameters + ---------- + slope_s : ndarray + weighted slope for current iteration's pixels for data section, 1-D + float + + intercept_s : ndarray + y-intercepts from fit for data section, 1-D float + + variance_s : ndarray + variance of residuals for fit for data section, 1-D float + + sig_intercept_s : ndarray + sigma of y-intercepts from fit for data section, 1-D float + + sig_slope_s : ndarray + sigma of slopes from fit for data section, 1-D float + + npix : int + number of pixels in 2d array + + data : float + array of values for current data section + + mask_2d : ndarray + delineates which channels to fit for each pixel, 2-D bool + + rn_sect_1d : ndarray + read noise values for all pixels in data section, 1-D float + + Returns + ------- + slope_s : ndarray + weighted slope for current iteration's pixels for data section, 1-D float + + intercept_s : ndarray + y-intercepts from fit for data section, 1-D float + + variance_s : ndarray + variance of residuals for fit for data section, 1-D float + + sig_intercept_s : ndarray + sigma of y-intercepts from fit for data section, 1-D float + + sig_slope_s : ndarray + sigma of slopes from fit for data section, 1-D float + """ + # For pixels saturated on the first group, overwrite fit values with + # benign values to be recalculated later. + wh_sat0 = np.where(np.logical_not(mask_2d[0, :])) + + if len(wh_sat0[0]) > 0: + sat_pix = wh_sat0[0] + slope_s[sat_pix] = 0. + variance_s[sat_pix] = 0. + sig_slope_s[sat_pix] = 0. + intercept_s[sat_pix] = 0. + sig_intercept_s[sat_pix] = 0. + del wh_sat0 + + # For pixels saturated on the second group, recalculate the slope as + # the value of the SCI data in the first group, which will later be + # divided by the group exposure time to give the count rate, and + # recalculate the other fit quantities to be benign. Note: these pixels + # will already have been handled earlier (for intervals of arbitrary + # length) in this function, but are being included here to explicitly + # cover all possibilities for pixels in datasets with ngroups=2. Will + # later consider refactoring. + wh_sat1 = np.where((mask_2d[:, :].sum(axis=0) == 1) & mask_2d[0, :]) + + if len(wh_sat1[0]) > 0: + data0_slice = data[0, :, :].reshape(npix) + slope_s[wh_sat1] = data0_slice[wh_sat1] + # set variance non-zero because calling function uses variance=0 to + # throw out bad results; this is not bad + variance_s[wh_sat1] = 1. + sig_slope_s[wh_sat1] = 0. + intercept_s[wh_sat1] = 0. + sig_intercept_s[wh_sat1] = 0. + del wh_sat1 + + # For pixels with no saturated values, recalculate the slope as the + # difference between the values of the second and first groups (1-based), + # which will later be divided by the group exposure time to give the count + # rate, and recalculate other fit quantities to be benign. + wh_sat_no = np.where(mask_2d[:, :].sum(axis=0) == 2) + + if len(wh_sat_no[0]) > 0: + data0_slice = data[0, :, :].reshape(npix) + data1_slice = data[1, :, :].reshape(npix) + slope_s[wh_sat_no] = data1_slice[wh_sat_no] - data0_slice[wh_sat_no] + sig_slope_s[wh_sat_no] = np.sqrt(2) * rn_sect_1d[wh_sat_no] + intercept_s[wh_sat_no] = data0_slice[wh_sat_no] -\ + data1_slice[wh_sat_no] # by geometry + sig_intercept_s[wh_sat_no] = np.sqrt(2) * rn_sect_1d[wh_sat_no] + variance_s[wh_sat_no] = np.sqrt(2) * rn_sect_1d[wh_sat_no] + + del wh_sat_no + + return slope_s, intercept_s, variance_s, sig_intercept_s, sig_slope_s + + +def calc_num_seg(gdq, n_int): + """ + Calculate the maximum number of segments that will be fit within an + integration, calculated over all pixels and all integrations. This value + is based on the locations of cosmic ray-affected pixels in all of the ramps, + and will be used to allocate arrays used for the optional output product. + + Parameters + ---------- + gdq : ndarray + cube of GROUPDQ array for a data, 3-D flag + + n_int : int (unused) + total number of integrations in data set + + Return: + ------- + max_num_seg: int + The maximum number of segements within an integration + max_cr: int + The maximum number of cosmic rays within an integration + """ + max_cr = 0 # max number of CRS for all integrations + + # For all 2d pixels, get max number of CRs or DO_NOT_USE flags along their + # ramps, to use as a surrogate for the number of segments along the ramps + # Note that we only care about flags that are NOT in the first or last groups, + # because exclusion of a first or last group won't result in an additional segment. + max_cr = np.count_nonzero(np.bitwise_and(gdq[:, 1:-1], JUMP_DET | DO_NOT_USE), axis=1).max() + + # Do not want to return a value > the number of groups, which can occur if + # this is a MIRI dataset in which the first or last group was flagged as + # DO_NOT_USE and also flagged as a jump. + max_num_seg = int(max_cr) + 1 # n CRS implies n+1 segments + if max_num_seg > gdq.shape[1]: + max_num_seg = gdq.shape[1] + + return max_num_seg, max_cr + + +def calc_unwtd_sums(data_masked, xvalues): + """ + Calculate the sums needed to determine the slope and intercept (and sigma + of each) using an unweighted fit. Unweighted fitting currently not + supported. + + Parameters + ---------- + data_masked : ndarray + masked values for all pixels in data section, 2-D float + + xvalues : ndarray + indices of valid pixel values for all groups, 1-D int + + Return: + ------- + sumx : float + sum of xvalues + + sumxx : float + sum of squares of xvalues + + sumxy : float + sum of product of xvalues and data + + sumy : float + sum of data + + """ + sumx = xvalues.sum(axis=0) + sumxx = (xvalues**2).sum(axis=0) + sumy = (np.reshape(data_masked.sum(axis=0), sumx.shape)) + sumxy = (xvalues * np.reshape(data_masked, xvalues.shape)).sum(axis=0) + + return sumx, sumxx, sumxy, sumy + + +def calc_opt_sums(rn_sect, gain_sect, data_masked, mask_2d, xvalues, good_pix): + """ + Calculate the sums needed to determine the slope and intercept (and sigma of + each) using the optimal weights. For each good pixel's segment, from the + initial and final indices and the corresponding number of counts, calculate + the SNR. From the SNR, calculate the weighting exponent using the formulation + by Fixsen (Fixsen et al, PASP, 112, 1350). Using this exponent and the gain + and the readnoise, the weights are calculated from which the sums are + calculated. + + Parameters + ---------- + rn_sect : ndarray + read noise values for all pixels in data section, 2-D float + + gain_sect : ndarray + gain values for all pixels in data section, 2-D float + + data_masked : ndarray + masked values for all pixels in data section, 2-D float + + mask_2d : ndarray + delineates which channels to fit for each pixel, 2-D bool + + xvalues : ndarray + indices of valid pixel values for all groups, 2-D int + + good_pix : ndarray + indices of pixels having valid data for all groups, 1-D int + + Return: + ------- + sumx : float + sum of xvalues + + sumxx : float + sum of squares of xvalues + + sumxy : float + sum of product of xvalues and data + + sumy : float + sum of data + + nreads_wtd : ndarray + sum of optimal weights, 1-D float + + xvalues : ndarray + rolled up indices of valid pixel values for all groups, 2-D int + """ + c_mask_2d = mask_2d.copy() # copy the mask to prevent propagation + rn_sect = np.float32(rn_sect) + + # Return 'empty' sums if there is no more data to fit + if data_masked.size == 0: + return np.array([]), np.array([]), np.array([]), np.array([]),\ + np.array([]), np.array([]) + + # get initial group for each good pixel for this semiramp + fnz = np.argmax(c_mask_2d, axis=0) + + # For those pixels that are all False, set to sentinel value of -1 + fnz[c_mask_2d.sum(axis=0) == 0] = -1 + + mask_2d_sum = c_mask_2d.sum(axis=0) # number of valid groups/pixel + + # get final valid group for each pixel for this semiramp + ind_lastnz = fnz + mask_2d_sum - 1 + + # get SCI value of initial good group for semiramp + data_zero = data_masked[fnz, range(data_masked.shape[1])] + + # get SCI value of final good group for semiramp + data_final = data_masked[(ind_lastnz), range(data_masked.shape[1])] + data_diff = data_final - data_zero # correctly does *NOT* have nans + + ind_lastnz = 0 + + # Use the readnoise and gain for good pixels only + rn_sect_rav = rn_sect.flatten()[good_pix] + rn_2_r = rn_sect_rav * rn_sect_rav + + gain_sect_r = gain_sect.flatten()[good_pix] + + # Calculate the sigma for nonzero gain values + sigma_ir = data_final.copy() * 0.0 + numer_ir = data_final.copy() * 0.0 + + # Calculate the SNR for pixels from the readnoise, the gain, and the + # difference between the last and first reads for pixels where this results + # in a positive SNR. Otherwise set the SNR to 0. + sqrt_arg = rn_2_r + data_diff * gain_sect_r + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", "invalid value.*", RuntimeWarning) + wh_pos = np.where((sqrt_arg >= 0.) & (gain_sect_r != 0.)) + numer_ir[wh_pos] = \ + np.sqrt(rn_2_r[wh_pos] + data_diff[wh_pos] * gain_sect_r[wh_pos]) + sigma_ir[wh_pos] = numer_ir[wh_pos] / gain_sect_r[wh_pos] + snr = data_diff * 0. + snr[wh_pos] = data_diff[wh_pos] / sigma_ir[wh_pos] + snr[np.isnan(snr)] = 0.0 + snr[snr < 0.] = 0.0 + + del wh_pos + + gain_sect_r = 0 + numer_ir = 0 + data_diff = 0 + sigma_ir = 0 + + power_wt_r = calc_power(snr) # Get the interpolated power for this SNR + # Make array of number of good groups, and exponents for each pixel + num_nz = (data_masked != 0.).sum(0) # number of nonzero groups per pixel + nrd_data_a = num_nz.copy() + num_nz = 0 + + nrd_prime = (nrd_data_a - 1) / 2. + nrd_data_a = 0 + + # Calculate inverse read noise^2 for use in weights + # Suppress, then re-enable, harmless arithmetic warning + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + invrdns2_r = 1. / rn_2_r + warnings.resetwarnings() + + rn_sect = 0 + fnz = 0 + + # Set optimal weights for each group of each pixel; + # for all pixels at once, loop over the groups + wt_h = np.zeros(data_masked.shape, dtype=np.float32) + + for jj_rd in range(data_masked.shape[0]): + wt_h[jj_rd, :] = \ + abs((abs(jj_rd - nrd_prime) / nrd_prime) ** power_wt_r) * invrdns2_r + + wt_h[np.isnan(wt_h)] = 0. + wt_h[np.isinf(wt_h)] = 0. + + # For all pixels, 'roll' up the leading zeros such that the 0th group of + # each pixel is the lowest nonzero group for that pixel + wh_m2d_f = np.logical_not(c_mask_2d[0, :]) # ramps with initial group False + while wh_m2d_f.sum() > 0: + data_masked[:, wh_m2d_f] = np.roll(data_masked[:, wh_m2d_f], -1, axis=0) + c_mask_2d[:, wh_m2d_f] = np.roll(c_mask_2d[:, wh_m2d_f], -1, axis=0) + xvalues[:, wh_m2d_f] = np.roll(xvalues[:, wh_m2d_f], -1, axis=0) + wh_m2d_f = np.logical_not(c_mask_2d[0, :]) + + # Create weighted sums for Poisson noise and read noise + nreads_wtd = (wt_h * c_mask_2d).sum(axis=0) # using optimal weights + + sumx = (xvalues * wt_h).sum(axis=0) + sumxx = (xvalues**2 * wt_h).sum(axis=0) + + c_data_masked = data_masked.copy() + c_data_masked[np.isnan(c_data_masked)] = 0. + sumy = (np.reshape((c_data_masked * wt_h).sum(axis=0), sumx.shape)) + sumxy = (xvalues * wt_h * np.reshape(c_data_masked, xvalues.shape)).sum(axis=0) + + return sumx, sumxx, sumxy, sumy, nreads_wtd, xvalues diff --git a/src/stcal/ramp_fitting/ramp_fit.py b/src/stcal/ramp_fitting/ramp_fit.py new file mode 100755 index 000000000..b26e94680 --- /dev/null +++ b/src/stcal/ramp_fitting/ramp_fit.py @@ -0,0 +1,104 @@ +#! /usr/bin/env python +# +# ramp_fit.py - calculate weighted mean of slope, based on Massimo +# Robberto's "On the Optimal Strategy to fit MULTIACCUM +# ramps in the presence of cosmic rays." +# (JWST-STScI-0001490,SM-12; 07/25/08). The derivation +# is a generalization for >1 cosmic rays, calculating +# the slope and variance of the slope for each section +# of the ramp (in between cosmic rays). The intervals are +# determined from the input data quality arrays. +# +# Note: +# In this module, comments on the 'first group','second group', etc are +# 1-based, unless noted otherwise. + +import numpy as np +import logging + +# from . import gls_fit # used only if algorithm is "GLS" +from . import ols_fit # used only if algorithm is "OLS" + +log = logging.getLogger(__name__) +log.setLevel(logging.DEBUG) + +BUFSIZE = 1024 * 300000 # 300Mb cache size for data section + + +def ramp_fit(model, buffsize, save_opt, readnoise_2d, gain_2d, + algorithm, weighting, max_cores): + """ + Calculate the count rate for each pixel in all data cube sections and all + integrations, equal to the slope for all sections (intervals between + cosmic rays) of the pixel's ramp divided by the effective integration time. + The weighting parameter must currently be set to 'optim', to use the optimal + weighting (paper by Fixsen, ref. TBA) will be used in the fitting; this is + currently the only supported weighting scheme. + + Parameters + ---------- + model : data model + input data model, assumed to be of type RampModel + + buffsize : int + size of data section (buffer) in bytes + + save_opt : bool + calculate optional fitting results + + readnoise_2d: ndarray + 2-D array readnoise for all pixels + + gain_2d: ndarray + 2-D array gain for all pixels + + algorithm : str + 'OLS' specifies that ordinary least squares should be used; + 'GLS' specifies that generalized least squares should be used. + + weighting : str + 'optimal' specifies that optimal weighting should be used; + currently the only weighting supported. + + max_cores : str + Number of cores to use for multiprocessing. If set to 'none' (the + default), then no multiprocessing will be done. The other allowable + values are 'quarter', 'half', and 'all'. This is the fraction of cores + to use for multi-proc. The total number of cores includes the SMT cores + (Hyper Threading for Intel). + + Returns + ------- + image_info: tuple + The tuple of computed ramp fitting arrays. + + integ_info: tuple + The tuple of computed integration fitting arrays. + + opt_info: tuple + The tuple of computed optional results arrays for fitting. + + gls_opt_model : GLS_RampFitModel object or None (Unused for now) + Object containing optional GLS-specific ramp fitting data for the + exposure + """ + if algorithm.upper() == "GLS": + # !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! + # !!!!! Reference to ReadModel and GainModel changed to simple ndarrays !!!!! + # !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! + # new_model, int_model, gls_opt_model = gls_fit.gls_ramp_fit( + # model, buffsize, save_opt, readnoise_model, gain_model, max_cores) + image_info, integ_info, gls_opt_model = None, None, None + opt_info = None + else: + # Get readnoise array for calculation of variance of noiseless ramps, and + # gain array in case optimal weighting is to be done + nframes = model.meta.exposure.nframes + readnoise_2d *= gain_2d / np.sqrt(2. * nframes) + + # Compute ramp fitting using ordinary least squares. + image_info, integ_info, opt_info = ols_fit.ols_ramp_fit_multi( + model, buffsize, save_opt, readnoise_2d, gain_2d, weighting, max_cores) + gls_opt_model = None + + return image_info, integ_info, opt_info, gls_opt_model diff --git a/src/stcal/ramp_fitting/utils.py b/src/stcal/ramp_fitting/utils.py new file mode 100644 index 000000000..80f278748 --- /dev/null +++ b/src/stcal/ramp_fitting/utils.py @@ -0,0 +1,1455 @@ +#! /usr/bin/env python +# +# utils.py: utility functions +import logging +import multiprocessing +import numpy as np +import warnings + +from . import constants + +log = logging.getLogger(__name__) +log.setLevel(logging.DEBUG) + +# Replace zero or negative variances with this: +LARGE_VARIANCE = 1.e8 + +# TODO Should figure out a better way to do this +DO_NOT_USE = constants.DO_NOT_USE +SATURATED = constants.SATURATED +JUMP_DET = constants.JUMP_DET +NO_GAIN_VALUE = constants.NO_GAIN_VALUE +UNRELIABLE_SLOPE = constants.UNRELIABLE_SLOPE + + +class OptRes: + """ + Object to hold optional results for all good pixels for + y-intercept, slope, uncertainty for y-intercept, uncertainty for + slope, inverse variance, first frame (for pedestal image), and + cosmic ray magnitude. + """ + + def __init__(self, n_int, imshape, max_seg, nreads, save_opt): + """ + Initialize the optional attributes. These are 4D arrays for the + segment-specific values of the y-intercept, the slope, the uncertainty + associated with both, the weights, the approximate cosmic ray + magnitudes, and the inverse variance. These are 3D arrays for the + integration-specific first frame and pedestal values. + + Parameters + ---------- + n_int : int + number of integrations in data set + + imshape : tuple + shape of 2D image + + max_seg : int + maximum number of segments fit + + nreads : int + number of reads in an integration + + save_opt : bool + save optional fitting results + """ + self.slope_seg = np.zeros((n_int,) + (max_seg,) + imshape, dtype=np.float32) + if save_opt: + self.yint_seg = np.zeros((n_int,) + (max_seg,) + imshape, dtype=np.float32) + self.sigyint_seg = np.zeros((n_int,) + (max_seg,) + imshape, dtype=np.float32) + self.sigslope_seg = np.zeros((n_int,) + (max_seg,) + imshape, dtype=np.float32) + self.inv_var_seg = np.zeros((n_int,) + (max_seg,) + imshape, dtype=np.float32) + self.firstf_int = np.zeros((n_int,) + imshape, dtype=np.float32) + self.ped_int = np.zeros((n_int,) + imshape, dtype=np.float32) + self.cr_mag_seg = np.zeros((n_int,) + (nreads,) + imshape, dtype=np.float32) + self.var_p_seg = np.zeros((n_int,) + (max_seg,) + imshape, dtype=np.float32) + self.var_r_seg = np.zeros((n_int,) + (max_seg,) + imshape, dtype=np.float32) + + def init_2d(self, npix, max_seg, save_opt): + """ + Initialize the 2D segment-specific attributes for the current data + section. + + Parameters + ---------- + npix : integer + number of pixels in section of 2D array + + max_seg : integer + maximum number of segments that will be fit within an + integration, calculated over all pixels and all integrations + + save_opt : bool + save optional fitting results + + Returns + ------- + None + """ + self.slope_2d = np.zeros((max_seg, npix), dtype=np.float32) + if save_opt: + self.interc_2d = np.zeros((max_seg, npix), dtype=np.float32) + self.siginterc_2d = np.zeros((max_seg, npix), dtype=np.float32) + self.sigslope_2d = np.zeros((max_seg, npix), dtype=np.float32) + self.inv_var_2d = np.zeros((max_seg, npix), dtype=np.float32) + self.firstf_2d = np.zeros((max_seg, npix), dtype=np.float32) + self.var_s_2d = np.zeros((max_seg, npix), dtype=np.float32) + self.var_r_2d = np.zeros((max_seg, npix), dtype=np.float32) + + def reshape_res(self, num_int, rlo, rhi, sect_shape, ff_sect, save_opt): + """ + Loop over the segments and copy the reshaped 2D segment-specific + results for the current data section to the 4D output arrays. + + Parameters + ---------- + num_int : int + integration number + + rlo : int + first column of section + + rhi : int + last column of section + + sect_sect : tuple + shape of section image + + ff_sect : ndarray + first frame data, 2-D float + + save_opt : bool + save optional fitting results + + Returns + ------- + """ + for ii_seg in range(0, self.slope_seg.shape[1]): + self.slope_seg[num_int, ii_seg, rlo:rhi, :] = \ + self.slope_2d[ii_seg, :].reshape(sect_shape) + + if save_opt: + self.yint_seg[num_int, ii_seg, rlo:rhi, :] = \ + self.interc_2d[ii_seg, :].reshape(sect_shape) + self.slope_seg[num_int, ii_seg, rlo:rhi, :] = \ + self.slope_2d[ii_seg, :].reshape(sect_shape) + self.sigyint_seg[num_int, ii_seg, rlo:rhi, :] = \ + self.siginterc_2d[ii_seg, :].reshape(sect_shape) + self.sigslope_seg[num_int, ii_seg, rlo:rhi, :] = \ + self.sigslope_2d[ii_seg, :].reshape(sect_shape) + self.inv_var_seg[num_int, ii_seg, rlo:rhi, :] = \ + self.inv_var_2d[ii_seg, :].reshape(sect_shape) + self.firstf_int[num_int, rlo:rhi, :] = ff_sect + + def append_arr(self, num_seg, g_pix, intercept, slope, sig_intercept, + sig_slope, inv_var, save_opt): + """ + Add the fitting results for the current segment to the 2d arrays. + + Parameters + ---------- + num_seg : ndarray + counter for segment number within the section, 1-D int + + g_pix : ndarray + pixels having fitting results in current section, 1-D int + + intercept : ndarray + intercepts for pixels in current segment and section, 1-D float + + slope : ndarray + slopes for pixels in current segment and section, 1-D float + + sig_intercept : ndarray + uncertainties of intercepts for pixels in current segment + and section, 1-D float + + sig_slope : ndarray + uncertainties of slopes for pixels in current segment and + section, 1-D float + + inv_var : ndarray + reciprocals of variances for fits of pixels in current + segment and section, 1-D float + + save_opt : bool + save optional fitting results + + Returns + ------- + None + """ + self.slope_2d[num_seg[g_pix], g_pix] = slope[g_pix] + + if save_opt: + self.interc_2d[num_seg[g_pix], g_pix] = intercept[g_pix] + self.siginterc_2d[num_seg[g_pix], g_pix] = sig_intercept[g_pix] + self.sigslope_2d[num_seg[g_pix], g_pix] = sig_slope[g_pix] + self.inv_var_2d[num_seg[g_pix], g_pix] = inv_var[g_pix] + + def shrink_crmag(self, n_int, dq_cube, imshape, nreads): + """ + Compress the 4D cosmic ray magnitude array for the current + integration, removing all groups whose cr magnitude is 0 for + pixels having at least one group with a non-zero magnitude. For + every integration, the depth of the array is equal to the + maximum number of cosmic rays flagged in all pixels in all + integrations. This routine currently involves a loop over all + pixels having at least 1 group flagged; if this algorithm takes + too long for datasets having an overabundance of cosmic rays, + this routine will require further optimization. + + Parameters + ---------- + n_int : int + number of integrations in dataset + + dq_cube : ndarray + input data quality array, 4-D flag + + imshape : tuple + shape of a single input image + + nreads : int + number of reads in an integration + + Returns + ---------- + None + + """ + # Loop over data integrations to find max num of crs flagged per pixel + # (this could exceed the maximum number of segments fit) + max_cr = 0 + for ii_int in range(0, n_int): + dq_int = dq_cube[ii_int, :, :, :] + dq_cr = np.bitwise_and(JUMP_DET, dq_int) + max_cr_int = (dq_cr > 0.).sum(axis=0).max() + max_cr = max(max_cr, max_cr_int) + + # Allocate compressed array based on max number of crs + cr_com = np.zeros((n_int,) + (max_cr,) + imshape, dtype=np.float32) + + # Loop over integrations and groups: for those pix having a cr, add + # the magnitude to the compressed array + for ii_int in range(0, n_int): + cr_mag_int = self.cr_mag_seg[ii_int, :, :, :] + cr_int_has_cr = np.where(cr_mag_int.sum(axis=0) != 0) + + # Initialize number of crs for each image pixel for this integration + end_cr = np.zeros(imshape, dtype=np.int16) + + for k_rd in range(nreads): + # loop over pixels having a CR + for nn in range(len(cr_int_has_cr[0])): + y, x = cr_int_has_cr[0][nn], cr_int_has_cr[1][nn] + + if (cr_mag_int[k_rd, y, x] > 0.): + cr_com[ii_int, end_cr[y, x], y, x] = cr_mag_int[k_rd, y, x] + end_cr[y, x] += 1 + + max_num_crs = end_cr.max() + if max_num_crs == 0: + max_num_crs = 1 + self.cr_mag_seg = np.zeros(shape=(n_int, 1, imshape[0], imshape[1])) + else: + self.cr_mag_seg = cr_com[:, :max_num_crs, :, :] + + def output_optional(self, effintim): + """ + These results are the cosmic ray magnitudes in the + segment-specific results for the count rates, y-intercept, + uncertainty in the slope, uncertainty in the y-intercept, + pedestal image, fitting weights, and the uncertainties in + the slope due to poisson noise only and read noise only, and + the integration-specific results for the pedestal image. The + slopes are divided by the effective integration time here to + yield the count rates. Any variance values that are a large fraction + of the default value LARGE_VARIANCE correspond to non-existent segments, + so will be set to 0 here before output. + + Parameters + ---------- + effintim : float + effective integration time for a single group + + Returns + ------- + opt_info: tuple + The tuple of computed optional results arrays for fitting. + """ + self.var_p_seg[self.var_p_seg > 0.4 * LARGE_VARIANCE] = 0. + self.var_r_seg[self.var_r_seg > 0.4 * LARGE_VARIANCE] = 0. + + # Suppress, then re-enable, arithmetic warnings + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + + # Tiny 'weights' values correspond to non-existent segments, so set to 0. + self.weights[1. / self.weights > 0.4 * LARGE_VARIANCE] = 0. + warnings.resetwarnings() + + self.slope_seg /= effintim + + opt_info = (self.slope_seg, self.sigslope_seg, self.var_p_seg, + self.var_r_seg, self.yint_seg, self.sigyint_seg, + self.ped_int, self.weights, self.cr_mag_seg) + + return opt_info + + def print_full(self): # pragma: no cover + """ + Diagnostic function for printing optional output arrays; most + useful for tiny datasets + + Parameters + ---------- + None + + Returns + ------- + None + """ + print('Will now print all optional output arrays - ') + print(' yint_seg: ') + print((self.yint_seg)) + print(' ') + print(' slope_seg: ') + print(self.slope_seg) + print(' ') + print(' sigyint_seg: ') + print(self.sigyint_seg) + print(' ') + print(' sigslope_seg: ') + print(self.sigslope_seg) + print(' ') + print(' inv_var_2d: ') + print((self.inv_var_2d)) + print(' ') + print(' firstf_int: ') + print((self.firstf_int)) + print(' ') + print(' ped_int: ') + print((self.ped_int)) + print(' ') + print(' cr_mag_seg: ') + print((self.cr_mag_seg)) + + +def alloc_arrays_1(n_int, imshape): + """ + Allocate arrays for integration-specific results and segment-specific + results and variances. + + Parameters + ---------- + n_int : int + number of integrations + + imshape : tuple + shape of a single image + + Returns + ------- + dq_int : ndarray + Cube of integration-specific group data quality values, 3-D flag + + median_diffs_2d : ndarray + Estimated median slopes, 2-D float + + num_seg_per_int : ndarray + Cube of numbers of segments for all integrations and pixels, 3-D int + + sat_0th_group_int : ndarray + Integration-specific slice whose value for a pixel is 1 if the initial + group of the ramp is saturated, 3-D uint8 + """ + dq_int = np.zeros((n_int,) + imshape, dtype=np.uint32) + num_seg_per_int = np.zeros((n_int,) + imshape, dtype=np.uint8) + + # for estimated median slopes + median_diffs_2d = np.zeros(imshape, dtype=np.float32) + sat_0th_group_int = np.zeros((n_int,) + imshape, dtype=np.uint8) + + return (dq_int, median_diffs_2d, num_seg_per_int, sat_0th_group_int) + + +def alloc_arrays_2(n_int, imshape, max_seg): + """ + Allocate arrays for integration-specific results and segment-specific + results and variances. + + Parameters + ---------- + n_int : int + number of integrations + + imshape : tuple + shape of a single image + + max_seg : int + maximum number of segments fit + + Returns + ------- + var_p3 : ndarray + Cube of integration-specific values for the slope variance due to + Poisson noise only, 3-D float + + var_r3 : ndarray + Cube of integration-specific values for the slope variance due to + readnoise only, 3-D float + + var_p4 : ndarray + Hypercube of segment- and integration-specific values for the slope + variance due to Poisson noise only, 4-D float + + var_r4 : ndarray + Hypercube of segment- and integration-specific values for the slope + variance due to read noise only, 4-D float + + var_both4 : ndarray + Hypercube of segment- and integration-specific values for the slope + variance due to combined Poisson noise and read noise, 4-D float + + var_both3 : ndarray + Cube of segment- and integration-specific values for the slope + variance due to combined Poisson noise and read noise, 3-D float + + inv_var_both4 : ndarray + Hypercube of reciprocals of segment- and integration-specific values for + the slope variance due to combined Poisson noise and read noise, 4-D float + + s_inv_var_p3 : ndarray + Cube of reciprocals of segment- and integration-specific values for + the slope variance due to Poisson noise only, summed over segments, 3-D float + + s_inv_var_r3 : ndarray + Cube of reciprocals of segment- and integration-specific values for + the slope variance due to read noise only, summed over segments, 3-D float + + s_inv_var_both3 : ndarray + Cube of reciprocals of segment- and integration-specific values for + the slope variance due to combined Poisson noise and read noise, summed + over segments, 3-D float + + segs_4 : ndarray + Hypercube of lengths of segments for all integrations and pixels, 4-D + int + """ + # Initialize variances so that non-existing ramps and segments will have + # negligible contributions + # Integration-specific: + var_p3 = np.zeros((n_int,) + imshape, dtype=np.float32) + LARGE_VARIANCE + var_r3 = var_p3.copy() + var_both3 = var_p3.copy() + s_inv_var_p3 = np.zeros_like(var_p3) + s_inv_var_r3 = np.zeros_like(var_p3) + s_inv_var_both3 = np.zeros_like(var_p3) + + # Segment-specific: + var_p4 = np.zeros((n_int,) + (max_seg,) + imshape, dtype=np.float32) + LARGE_VARIANCE + var_r4 = var_p4.copy() + var_both4 = var_p4.copy() + inv_var_both4 = np.zeros_like(var_p4) + + # number of segments + segs_4 = np.zeros((n_int,) + (max_seg,) + imshape, dtype=np.uint8) + + return (var_p3, var_r3, var_p4, var_r4, var_both4, var_both3, + inv_var_both4, s_inv_var_p3, s_inv_var_r3, + s_inv_var_both3, segs_4) + + +def calc_slope_vars(rn_sect, gain_sect, gdq_sect, group_time, max_seg): + """ + Calculate the segment-specific variance arrays for the given + integration. + + Parameters + ---------- + rn_sect : ndarray + read noise values for all pixels in data section, 2-D float + + gain_sect : ndarray + gain values for all pixels in data section, 2-D float + + gdq_sect : ndarray + data quality flags for pixels in section, 3-D int + + group_time : float + Time increment between groups, in seconds. + + max_seg : int + maximum number of segments fit + + Returns + ------- + den_r3 : ndarray + for a given integration, the reciprocal of the denominator of the + segment-specific variance of the segment's slope due to read noise, 3-D float + + den_p3 : ndarray + for a given integration, the reciprocal of the denominator of the + segment-specific variance of the segment's slope due to Poisson noise, 3-D float + + num_r3 : ndarray + numerator of the segment-specific variance of the segment's slope + due to read noise, 3-D float + + segs_beg_3 : ndarray + lengths of segments for all pixels in the given data section and + integration, 3-D int + """ + (nreads, asize2, asize1) = gdq_sect.shape + npix = asize1 * asize2 + imshape = (asize2, asize1) + + # Create integration-specific sections of input arrays for determination + # of the variances. + gdq_2d = gdq_sect[:, :, :].reshape((nreads, npix)) + gain_1d = gain_sect.reshape(npix) + gdq_2d_nan = gdq_2d.copy() # group dq with SATS will be replaced by nans + gdq_2d_nan = gdq_2d_nan.astype(np.float32) + + wh_sat = np.where(np.bitwise_and(gdq_2d, SATURATED)) + if len(wh_sat[0]) > 0: + gdq_2d_nan[wh_sat] = np.nan # set all SAT groups to nan + + del wh_sat + + # Get lengths of semiramps for all pix [number_of_semiramps, number_of_pix] + segs = np.zeros_like(gdq_2d) + + # Counter of semiramp for each pixel + sr_index = np.zeros(npix, dtype=np.uint8) + pix_not_done = np.ones(npix, dtype=bool) # initialize to True + + i_read = 0 + # Loop over reads for all pixels to get segments (segments per pixel) + while (i_read < nreads and np.any(pix_not_done)): + gdq_1d = gdq_2d_nan[i_read, :] + wh_good = np.where(gdq_1d == 0) # good groups + + # if this group is good, increment those pixels' segments' lengths + if len(wh_good[0]) > 0: + segs[sr_index[wh_good], wh_good] += 1 + del wh_good + + # Locate any CRs that appear before the first SAT group... + wh_cr = np.where(gdq_2d_nan[i_read, :].astype(np.int32) & JUMP_DET > 0) + + # ... but not on final read: + if (len(wh_cr[0]) > 0 and (i_read < nreads - 1)): + sr_index[wh_cr[0]] += 1 + segs[sr_index[wh_cr], wh_cr] += 1 + + del wh_cr + + # If current group is a NaN, this pixel is done (pix_not_done is False) + wh_nan = np.where(np.isnan(gdq_2d_nan[i_read, :])) + if len(wh_nan[0]) > 0: + pix_not_done[wh_nan[0]] = False + + del wh_nan + + i_read += 1 + + segs = segs.astype(np.uint8) + segs_beg = segs[:max_seg, :] # the leading nonzero lengths + + # Create reshaped version [ segs, y, x ] to simplify computation + segs_beg_3 = segs_beg.reshape(max_seg, imshape[0], imshape[1]) + segs_beg_3 = remove_bad_singles(segs_beg_3) + + # Create a version 1 less for later calculations for the variance due to + # Poisson, with a floor=1 to handle single-group segments + wh_pos_3 = np.where(segs_beg_3 > 1) + segs_beg_3_m1 = segs_beg_3.copy() + segs_beg_3_m1[wh_pos_3] -= 1 + segs_beg_3_m1[segs_beg_3_m1 < 1] = 1 + + # For a segment, the variance due to Poisson noise + # = slope/(tgroup * gain * (ngroups-1)), + # where slope is the estimated median slope, tgroup is the group time, + # and ngroups is the number of groups in the segment. + # Here the denominator of this quantity will be computed, which will be + # later multiplied by the estimated median slope. + + # Suppress, then re-enable, harmless arithmetic warnings, as NaN will be + # checked for and handled later + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + den_p3 = 1. / (group_time * gain_1d.reshape(imshape) * segs_beg_3_m1) + warnings.resetwarnings() + + # For a segment, the variance due to readnoise noise + # = 12 * readnoise**2 /(ngroups_seg**3. - ngroups_seg)/( tgroup **2.) + num_r3 = 12. * (rn_sect / group_time)**2. # always >0 + + # Reshape for every group, every pixel in section + num_r3 = np.dstack([num_r3] * max_seg) + num_r3 = np.transpose(num_r3, (2, 0, 1)) + + # Denominator den_r3 = 1./(segs_beg_3 **3.-segs_beg_3). The minimum number + # of allowed groups is 2, which will apply if there is actually only 1 + # group; in this case den_r3 = 1/6. This covers the case in which there is + # only one good group at the beginning of the integration, so it will be + # be compared to the plane of (near) zeros resulting from the reset. For + # longer segments, this value is overwritten below. + den_r3 = num_r3.copy() * 0. + 1. / 6 + wh_seg_pos = np.where(segs_beg_3 > 1) + + # Suppress, then, re-enable harmless arithmetic warnings, as NaN will be + # checked for and handled later + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + den_r3[wh_seg_pos] = 1. / (segs_beg_3[wh_seg_pos] ** 3. - + segs_beg_3[wh_seg_pos]) # overwrite where segs>1 + warnings.resetwarnings() + + return (den_r3, den_p3, num_r3, segs_beg_3) + + +def calc_pedestal(num_int, slope_int, firstf_int, dq_first, nframes, groupgap, + dropframes1): + """ + The pedestal is calculated by extrapolating the final slope for each pixel + from its value at the first sample in the integration to an exposure time + of zero; this calculation accounts for the values of nframes and groupgap. + Any pixel that is saturated on the 1st group is given a pedestal value of 0. + + Parameters + ---------- + num_int : int + integration number + + slope_int : ndarray + cube of integration-specific slopes, 3-D float + + firstf_int : ndarray + integration-specific first frame array, 3-D float + + dq_first : ndarray + DQ of the initial group for all ramps in the given integration, 2-D + flag + + nframes : int + number of frames averaged per group; from the NFRAMES keyword. Does + not contain the groupgap. + + groupgap : int + number of frames dropped between groups, from the GROUPGAP keyword. + + dropframes1 : int + number of frames dropped at the beginning of every integration, from + the DRPFRMS1 keyword. + + Returns + ------- + ped : ndarray + pedestal image, 2-D float + """ + ff_all = firstf_int[num_int, :, :].astype(np.float32) + ped = ff_all - slope_int[num_int, ::] * \ + (((nframes + 1.) / 2. + dropframes1) / (nframes + groupgap)) + + ped[np.bitwise_and(dq_first, SATURATED) == SATURATED] = 0 + ped[np.isnan(ped)] = 0. + + return ped + + +def output_integ(slope_int, dq_int, effintim, var_p3, var_r3, var_both3, + int_times): + """ + For the OLS algorithm, construct the output integration-specific results. + Any variance values that are a large fraction of the default value + LARGE_VARIANCE correspond to non-existent segments, so will be set to 0 + here before output. + + Parameters + ---------- + model : instance of Data Model + DM object for input + + slope_int : ndarray + Data cube of weighted slopes for each integration, 3-D float + + dq_int : ndarray + Data cube of DQ arrays for each integration, 3-D int + + effintim : float + Effective integration time per integration + + var_p3 : ndarray + Cube of integration-specific values for the slope variance due to + Poisson noise only, 3-D float + + var_r3 : ndarray + Cube of integration-specific values for the slope variance due to + read noise only, 3-D float + + var_both3 : ndarray + Cube of integration-specific values for the slope variance due to + read noise and Poisson noise, 3-D float + + int_times : bintable, or None + The INT_TIMES table, if it exists in the input, else None + + Returns + ------- + integ_info: tuple + The tuple of computed integration:t fitting arrays. + + """ + # Suppress harmless arithmetic warnings for now + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + + var_p3[var_p3 > 0.4 * LARGE_VARIANCE] = 0. + var_r3[var_r3 > 0.4 * LARGE_VARIANCE] = 0. + var_both3[var_both3 > 0.4 * LARGE_VARIANCE] = 0. + + data = slope_int / effintim + err = np.sqrt(var_both3) + dq = dq_int + var_poisson = var_p3 + var_rnoise = var_r3 + int_times = int_times + integ_info = (data, dq, var_poisson, var_rnoise, int_times, err) + + # Reset the warnings filter to its original state + warnings.resetwarnings() + + return integ_info + + +''' +# BEGIN remove GLS +def gls_output_integ(model, slope_int, slope_err_int, dq_int): + """ + For the GLS algorithm, construct the output integration-specific results. + Parameters + ---------- + model : instance of Data Model + DM object for input + + slope_int : ndarray + Data cube of weighted slopes for each integration, 3-D float + + slope_err_int : ndarray + Data cube of slope errors for each integration, 3-D float + + dq_int : ndarray + Data cube of DQ arrays for each integration, 3-D flag + + Returns + ------- + cubemod : Data Model object + """ + # Suppress harmless arithmetic warnings for now + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + + cubemod = datamodels.CubeModel() + cubemod.data = slope_int + cubemod.err = slope_err_int + cubemod.dq = dq_int + + # Reset the warnings filter to its original state + warnings.resetwarnings() + + cubemod.update(model) # keys from input needed for photom step + + return cubemod + + +def gls_output_optional(model, intercept_int, intercept_err_int, + pedestal_int, + ampl_int, ampl_err_int): # pragma: no cover + """Construct the optional results for the GLS algorithm. + + Extended Summary + ---------------- + Construct the GLS-specific optional output data. + These results are the Y-intercepts, uncertainties in the intercepts, + pedestal (first group extrapolated back to zero time), + cosmic ray magnitudes, and uncertainties in the CR magnitudes. + + Parameters + ---------- + model : instance of Data Model + Data model object for input; this is used only for the file name. + + intercept_int : 3-D ndarray, float32, shape (n_int, ny, nx) + Y-intercept for each integration, at each pixel. + + intercept_err_int : 3-D ndarray, float32, shape (n_int, ny, nx) + Uncertainties for Y-intercept for each integration, at each pixel. + + pedestal_int : 3-D ndarray, float32, shape (n_int, ny, nx) + The pedestal, for each integration and each pixel. + + ampl_int : 4-D ndarray, float32, shape (n_int, ny, nx, max_num_cr) + Cosmic-ray amplitudes for each integration, at each pixel, and for + each CR hit in the ramp. max_num_cr will be the maximum number of + CRs within the ramp for any pixel, or it will be one if there were + no CRs at all. + + ampl_err_int : 4-D ndarray, float32, shape (n_int, ny, nx, max_num_cr) + Uncertainties for cosmic-ray amplitudes for each integration, at + each pixel, and for each CR in the ramp. + + Returns + ------- + gls_ramp_model : GLS_RampFitModel object + GLS-specific ramp fit data for the exposure. + """ + + gls_ramp_model = datamodels.GLS_RampFitModel() + + gls_ramp_model.yint = intercept_int + gls_ramp_model.sigyint = intercept_err_int + gls_ramp_model.pedestal = pedestal_int + gls_ramp_model.crmag = ampl_int + gls_ramp_model.sigcrmag = ampl_err_int + + return gls_ramp_model + + +def gls_pedestal(first_group, slope_int, s_mask, + frame_time, nframes_used): # pragma: no cover + """Calculate the pedestal for the GLS case. + + The pedestal is the first group, but extrapolated back to zero time + using the slope obtained by the fit to the whole ramp. The time of + the first group is the frame time multiplied by (M + 1) / 2, where M + is the number of frames per group, not including the number (if any) + of skipped frames. + + The input arrays and output pedestal are slices of the full arrays. + They are just the relevant data for the current integration (assuming + that this function is called within a loop over integrations), and + they may include only a subset of image lines. For example, this + function might be called with slope_int argument given as: + `slope_int[num_int, rlo:rhi, :]`. The input and output parameters + are in electrons. + + Parameters + ---------- + first_group : ndarray + A slice of the first group in the ramp, 2-D float + + slope_int : ndarray + The slope obtained by GLS ramp fitting. This is a slice for the + current integration and a subset of the image lines, 2-D float + + s_mask : ndarray + True for ramps that were saturated in the first group, 2-D bool + + frame_time : float + The time to read one frame, in seconds. + + nframes_used : int + Number of frames that were averaged together to make a group. + Exludes the groupgap. + + Returns + ------- + pedestal : ndarray + This is a slice of the full pedestal array, and it's for the + current integration, 2-D float + """ + + M = float(nframes_used) + pedestal = first_group - slope_int * frame_time * (M + 1.) / 2. + if s_mask.any(): + pedestal[s_mask] = 0. + + return pedestal +# END remove GLS +''' + + +def shift_z(a, off): + """ + Shift input 3D array by requested offset in z-direction, padding shifted + array (of the same size) by leading or trailing zeros as needed. + + Parameters + ---------- + a : ndarray + input array, 3-D float + + off : int + offset in z-direction + + Returns + ------- + b : ndarray + shifted array, 3-D float + + """ + # set initial and final indices along z-direction for original and + # shifted 3D arrays + ai_z = int((abs(off) + off) / 2) + af_z = a.shape[0] + int((-abs(off) + off) / 2) + + bi_z = a.shape[0] - af_z + bf_z = a.shape[0] - ai_z + + b = a * 0 + b[bi_z:bf_z, :, :] = a[ai_z:af_z, :, :] + + return b + + +def get_efftim_ped(model): + """ + Calculate the effective integration time for a single group, and return the + number of frames per group, and the number of frames dropped between groups. + + Parameters + ---------- + model : instance of Data Model + DM object for input + + Returns + ------- + effintim : float + effective integration time for a single group + + nframes : int + number of frames averaged per group; from the NFRAMES keyword. + + groupgap : int + number of frames dropped between groups; from the GROUPGAP keyword. + + dropframes1 : int + number of frames dropped at the beginning of every integration; from + the DRPFRMS1 keyword, or 0 if the keyword is missing + """ + groupgap = model.meta.exposure.groupgap + nframes = model.meta.exposure.nframes + frame_time = model.meta.exposure.frame_time + dropframes1 = model.meta.exposure.drop_frames1 + + if (dropframes1 is None): # set to default if missing + dropframes1 = 0 + log.debug('Missing keyword DRPFRMS1, so setting to default value of 0') + + try: + effintim = (nframes + groupgap) * frame_time + except TypeError: + log.error('Can not retrieve values needed to calculate integ. time') + + log.debug('Calculating effective integration time for a single group using:') + log.debug(' groupgap: %s' % (groupgap)) + log.debug(' nframes: %s' % (nframes)) + log.debug(' frame_time: %s' % (frame_time)) + log.debug(' dropframes1: %s' % (dropframes1)) + log.info('Effective integration time per group: %s' % (effintim)) + + return effintim, nframes, groupgap, dropframes1 + + +def get_dataset_info(model): + """ + Extract values for the number of groups, the number of pixels, dataset + shapes, the number of integrations, the instrument name, the frame time, + and the observation time. + + Parameters + ---------- + model : instance of Data Model + DM object for input + + Returns + ------- + nreads : int + number of reads in input dataset + + npix : int + number of pixels in 2D array + + imshape : tuple + shape of 2D image + + cubeshape : tuple + shape of input dataset + + n_int : int + number of integrations + + instrume : str + instrument + + frame_time : float + integration time from TGROUP + + ngroups : int + number of groups per integration + + group_time : float + Time increment between groups, in seconds. + """ + instrume = model.meta.instrument.name + frame_time = model.meta.exposure.frame_time + ngroups = model.meta.exposure.ngroups + group_time = model.meta.exposure.group_time + + n_int = model.data.shape[0] + nreads = model.data.shape[1] + asize2 = model.data.shape[2] + asize1 = model.data.shape[3] + + # If nreads and ngroups are not the same, override the value of ngroups + # with nreads, which is more likely to be correct, since it's based on + # the image shape. + if nreads != ngroups: + log.warning('The value from the key NGROUPS does not (but should) match') + log.warning(' the value of nreads from the data; will use value of') + log.warning(' nreads: %s' % (nreads)) + ngroups = nreads + + npix = asize2 * asize1 # number of pixels in 2D array + imshape = (asize2, asize1) + cubeshape = (nreads,) + imshape + + return nreads, npix, imshape, cubeshape, n_int, instrume, frame_time, \ + ngroups, group_time + + +def get_more_info(model): # pragma: no cover + """Get information used by GLS algorithm. + + Parameters + ---------- + model : instance of Data Model + DM object for input + + Returns + ------- + group_time : float + Time increment between groups, in seconds. + + nframes_used : int + Number of frames that were averaged together to make a group, + i.e. excluding skipped frames. + + saturated_flag : int + Group data quality flag that indicates a saturated pixel. + + jump_flag : int + Group data quality flag that indicates a cosmic ray hit. + """ + + group_time = model.meta.exposure.group_time + nframes_used = model.meta.exposure.nframes + saturated_flag = SATURATED + jump_flag = JUMP_DET + + return (group_time, nframes_used, saturated_flag, jump_flag) + + +def get_max_num_cr(gdq_cube, jump_flag): # pragma: no cover + """ + Find the maximum number of cosmic-ray hits in any one pixel. + + Parameters + ---------- + gdq_cube : ndarray + The group data quality array, 3-D flag + + jump_flag : int + The data quality flag indicating a cosmic-ray hit. + + Returns + ------- + max_num_cr : int + The maximum number of cosmic-ray hits for any pixel. + """ + cr_flagged = np.empty(gdq_cube.shape, dtype=np.uint8) + cr_flagged[:] = np.where(np.bitwise_and(gdq_cube, jump_flag), 1, 0) + max_num_cr = cr_flagged.sum(axis=0, dtype=np.int32).max() + del cr_flagged + + return max_num_cr + + +def reset_bad_gain(pdq, gain): + """ + For pixels in the gain array that are either non-positive or NaN, reset the + the corresponding pixels in the pixel DQ array to NO_GAIN_VALUE and + DO_NOT_USE so that they will be ignored. + + Parameters + ---------- + pdq : ndarray + pixel dq array of input model, 2-D int + + gain : ndarray + gain array from reference file, 2-D float + + Returns + ------- + pdq : ndarray + pixleldq array of input model, reset to NO_GAIN_VALUE and DO_NOT_USE + for pixels in the gain array that are either non-positive or NaN., 2-D + flag + """ + ''' + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", "invalid value.*", RuntimeWarning) + wh_g = np.where(gain <= 0.) + ''' + wh_g = np.where(gain <= 0.) + if len(wh_g[0]) > 0: + pdq[wh_g] = np.bitwise_or(pdq[wh_g], NO_GAIN_VALUE) + pdq[wh_g] = np.bitwise_or(pdq[wh_g], DO_NOT_USE) + + wh_g = np.where(np.isnan(gain)) + if len(wh_g[0]) > 0: + pdq[wh_g] = np.bitwise_or(pdq[wh_g], NO_GAIN_VALUE) + pdq[wh_g] = np.bitwise_or(pdq[wh_g], DO_NOT_USE) + + return pdq + + +def remove_bad_singles(segs_beg_3): + """ + For the current integration and data section, remove all segments having only + a single group if there are other segments in the ramp. This method allows + for the possibility that a ramp can have multiple (necessarily consecutive) + 1-group segments, which in principle could occur if there are consecutive + cosmic rays. + + Parameters + ---------- + segs_beg_3 : ndarray + lengths of all segments for all ramps in the given data section and + integration; some of these ramps may contain segments having a single + group, and another segment, 3-D int + + Returns + ------- + segs_beg_3 : ndarray + lengths of all segments for all ramps in the given data section and + integration; segments having a single group, and another segment + will be removed, 3-D int + """ + max_seg = segs_beg_3.shape[0] + + # get initial number of ramps having single-group segments + tot_num_single_grp_ramps = len(np.where((segs_beg_3 == 1) & + (segs_beg_3.sum(axis=0) > 1))[0]) + + while(tot_num_single_grp_ramps > 0): + # until there are no more single-group segments + for ii_0 in range(max_seg): + slice_0 = segs_beg_3[ii_0, :, :] + + for ii_1 in range(max_seg): # correctly includes EARLIER segments + if (ii_0 == ii_1): # don't compare with itself + continue + + slice_1 = segs_beg_3[ii_1, :, :] + + # Find ramps of a single-group segment and another segment + # either earlier or later + wh_y, wh_x = np.where((slice_0 == 1) & (slice_1 > 0)) + + if (len(wh_y) == 0): + # Are none, so go to next pair of segments to check + continue + + # Remove the 1-group segment + segs_beg_3[ii_0:-1, wh_y, wh_x] = segs_beg_3[ii_0 + 1:, wh_y, wh_x] + + # Zero the last segment entry for the ramp, which would otherwise + # remain non-zero due to the shift + segs_beg_3[-1, wh_y, wh_x] = 0 + + del wh_y, wh_x + + tot_num_single_grp_ramps = len(np.where((segs_beg_3 == 1) & + (segs_beg_3.sum(axis=0) > 1))[0]) + + return segs_beg_3 + + +def fix_sat_ramps(sat_0th_group_int, var_p3, var_both3, slope_int, dq_int): + """ + For ramps within an integration that are saturated on the initial group, + reset the integration-specific variances and slope so they will have no + contribution. + + Parameters + ---------- + sat_0th_group_int : ndarray + Integration-specific slice whose value for a pixel is 1 if the initial + group of the ramp is saturated, 3-D uint8 + + var_p3 : ndarray + Cube of integration-specific values for the slope variance due to + Poisson noise only; some ramps may be saturated in the initial group, + 3-D float + + var_both3 : ndarray + Cube of segment- and integration-specific values for the slope + variance due to combined Poisson noise and read noise; some ramps may + be saturated in the initial group, 3-D float + + slope_int : ndarray + Cube of integration-specific slopes. Some ramps may be saturated in the + initial group, 3-D float + + dq_int : ndarray + Cube of integration-specific DQ flags, 3-D flag + + Returns + ------- + var_p3 : ndarray + Cube of integration-specific values for the slope variance due to + Poisson noise only; for ramps that are saturated in the initial group, + this variance has been reset to a huge value to minimize the ramps + contribution, 3-D float + + var_both3 : ndarray + Cube of segment- and integration-specific values for the slope + variance due to combined Poisson noise and read noise; for ramps that + are saturated in the initial group, this variance has been reset to a + huge value to minimize the ramps contribution, 3-D float + + slope_int : ndarray + Cube of integration-specific slopes; for ramps that are saturated in + the initial group, this variance has been reset to a huge value to + minimize the ramps contribution, 3-D float + + dq_int : ndarray + Cube of integration-specific DQ flags. For ramps that are saturated in + the initial group, the flag 'DO_NOT_USE' is added, 3-D flag + """ + var_p3[sat_0th_group_int > 0] = LARGE_VARIANCE + var_both3[sat_0th_group_int > 0] = LARGE_VARIANCE + slope_int[sat_0th_group_int > 0] = 0. + dq_int[sat_0th_group_int > 0] = np.bitwise_or( + dq_int[sat_0th_group_int > 0], DO_NOT_USE) + + return var_p3, var_both3, slope_int, dq_int + + +def do_all_sat(pixeldq, groupdq, imshape, n_int, save_opt): + """ + For an input exposure where all groups in all integrations are saturated, + the DQ in the primary and integration-specific output products are updated, + and the other arrays in all output products are populated with zeros. + + Parameters + ---------- + model : instance of Data Model + DM object for input + + imshape : (int, int) tuple + shape of 2D image + + n_int : int + number of integrations + + save_opt : bool + save optional fitting results + + Returns + ------- + image_info: tuple + The tuple of computed ramp fitting arrays. + + integ_info: tuple + The tuple of computed integration fitting arrays. + + opt_info: tuple + The tuple of computed optional results arrays for fitting. + """ + # Create model for the primary output. Flag all pixels in the pixiel DQ + # extension as SATURATED and DO_NOT_USE. + pixeldq = np.bitwise_or(pixeldq, SATURATED) + pixeldq = np.bitwise_or(pixeldq, DO_NOT_USE) + + data = np.zeros(imshape, dtype=np.float32) + dq = pixeldq + var_poisson = np.zeros(imshape, dtype=np.float32) + var_rnoise = np.zeros(imshape, dtype=np.float32) + err = np.zeros(imshape, dtype=np.float32) + image_info = (data, dq, var_poisson, var_rnoise, err) + + # Create model for the integration-specific output. The 3D group DQ created + # is based on the 4D group DQ of the model, and all pixels in all + # integrations will be flagged here as DO_NOT_USE (they are already flagged + # as SATURATED). The INT_TIMES extension will be left as None. + if n_int > 1: + m_sh = groupdq.shape # (integ, grps/integ, y, x ) + groupdq_3d = np.zeros((m_sh[0], m_sh[2], m_sh[3]), dtype=np.uint32) + + for ii in range(n_int): # add SAT flag to existing groupdq in each slice + groupdq_3d[ii, :, :] = np.bitwise_or.reduce(groupdq[ii, :, :, :], + axis=0) + + groupdq_3d = np.bitwise_or(groupdq_3d, DO_NOT_USE) + + data = np.zeros((n_int,) + imshape, dtype=np.float32) + dq = groupdq_3d + var_poisson = np.zeros((n_int,) + imshape, dtype=np.float32) + var_rnoise = np.zeros((n_int,) + imshape, dtype=np.float32) + int_times = None + err = np.zeros((n_int,) + imshape, dtype=np.float32) + + integ_info = (data, dq, var_poisson, var_rnoise, int_times, err) + else: + integ_info = None + + # Create model for the optional output + if save_opt: + new_arr = np.zeros((n_int,) + (1,) + imshape, dtype=np.float32) + + slope = new_arr + sigslope = new_arr + var_poisson = new_arr + var_rnoise = new_arr + yint = new_arr + sigyint = new_arr + pedestal = np.zeros((n_int,) + imshape, dtype=np.float32) + weights = new_arr + crmag = new_arr + + opt_info = (slope, sigslope, var_poisson, var_rnoise, + yint, sigyint, pedestal, weights, crmag) + + else: + opt_info = None + + log.info('All groups of all integrations are saturated.') + + return image_info, integ_info, opt_info + + +def log_stats(c_rates): + """ + Optionally log statistics of detected cosmic rays + + Parameters + ---------- + c_rates : ndarray + weighted count rate, 2-D float + + Returns + ------- + None + """ + wh_c_0 = np.where(c_rates == 0.) # insuff data or no signal + + log.debug('The number of pixels having insufficient data') + log.debug('due to excessive CRs or saturation %d:', len(wh_c_0[0])) + log.debug('Count rates - min, mean, max, std: %f, %f, %f, %f' + % (c_rates.min(), c_rates.mean(), c_rates.max(), c_rates.std())) + + +def compute_slices(max_cores): + """ + Computes the number of slices to be created for multiprocessing. + + Parameters + ---------- + max_cores : str + Number of cores to use for multiprocessing. If set to 'none' (the default), + then no multiprocessing will be done. The other allowable values are 'quarter', + 'half', and 'all'. This is the fraction of cores to use for multi-proc. The + total number of cores includes the SMT cores (Hyper Threading for Intel). + + Returns + ------- + number_slices : int + The number of slices for multiprocessing. + """ + if max_cores == 'none': + number_slices = 1 + else: + num_cores = multiprocessing.cpu_count() + log.debug(f'Found {num_cores} possible cores to use for ramp fitting') + if max_cores == 'quarter': + number_slices = num_cores // 4 or 1 + elif max_cores == 'half': + number_slices = num_cores // 2 or 1 + elif max_cores == 'all': + number_slices = num_cores + else: + number_slices = 1 + return number_slices + + +def dq_compress_final(dq_int, n_int): + """ + Combine the integration-specific dq arrays (which have already been + compressed and combined with the PIXELDQ array) to create the dq array + of the primary output product. + + Parameters + ---------- + dq_int : ndarray + cube of combined dq arrays for all data sections in a single + integration, 3-D flag + + n_int : int + total number of integrations in data set + + Returns + ------- + f_dq : ndarray + combination of all integration's pixeldq arrays, 2-D flag + """ + f_dq = dq_int[0, :, :] + + for jj in range(1, n_int): + f_dq = np.bitwise_or(f_dq, dq_int[jj, :, :]) + + return f_dq + + +def dq_compress_sect(gdq_sect, pixeldq_sect): + """ + Get ramp locations where the data has been flagged as saturated in the 4D + GROUPDQ array for the current data section, find the corresponding image + locations, and set the SATURATED flag in those locations in the PIXELDQ + array. Similarly, get the ramp locations where the data has been flagged as + a jump detection in the 4D GROUPDQ array, find the corresponding image + locations, and set the COSMIC_BEFORE flag in those locations in the PIXELDQ + array. These modifications to the section of the PIXELDQ array are not used + to flag groups for any computations; they are used only in the integration- + specific output. + + Parameters + ---------- + gdq_sect : ndarray + cube of GROUPDQ array for a data section, 3-D flag + + pixeldq_sect : ndarray + dq array of data section of input model, 2-D flag + + Returns + ------- + pixeldq_sect : ndarray + dq array of data section updated with saturated and jump-detected + flags, 2-D flag + + """ + sat_loc_r = np.bitwise_and(gdq_sect, SATURATED) + sat_loc_im = np.where(sat_loc_r.sum(axis=0) > 0) + pixeldq_sect[sat_loc_im] = np.bitwise_or(pixeldq_sect[sat_loc_im], SATURATED) + + cr_loc_r = np.bitwise_and(gdq_sect, JUMP_DET) + cr_loc_im = np.where(cr_loc_r.sum(axis=0) > 0) + pixeldq_sect[cr_loc_im] = np.bitwise_or(pixeldq_sect[cr_loc_im], JUMP_DET) + + return pixeldq_sect