diff --git a/changes/278.apichange.rst b/changes/278.apichange.rst new file mode 100644 index 00000000..e847fc7b --- /dev/null +++ b/changes/278.apichange.rst @@ -0,0 +1 @@ +[ramp_fitting] Add the likelihood algorithm to ramp fitting. diff --git a/docs/stcal/ramp_fitting/description.rst b/docs/stcal/ramp_fitting/description.rst index 0fd8279c..49f0dd0b 100644 --- a/docs/stcal/ramp_fitting/description.rst +++ b/docs/stcal/ramp_fitting/description.rst @@ -2,15 +2,14 @@ Description =========== This step determines the mean count rate, in units of counts per second, for -each pixel by performing a linear fit to the data in the input file. The fit -is done using the "ordinary least squares" method. -The fit is performed independently for each pixel. +each pixel by performing a linear fit to the data in the input file. The default +method uses "ordinary least squares" based on the Fixsen fitting algorithm +described by +`Fixsen et al. (2011) `_. The count rate for each pixel is determined by a linear fit to the cosmic-ray-free and saturation-free ramp intervals for each pixel; hereafter -this interval will be referred to as a "segment." The fitting algorithm uses an -'optimal' weighting scheme, as described by -`Fixsen et al. (2011) `_. +this interval will be referred to as a "segment." Segments are determined using the 4-D GROUPDQ array of the input data set, under the assumption that the jump @@ -19,6 +18,11 @@ saturation flags are found. Pixels are processed simultaneously in blocks using the array-based functionality of numpy. The size of the block depends on the image size and the number of groups. + +There is also a new algorithm available for testing, the likelihood +algorithm, implementing an algorithm based on the group differences +of a ramp. See :ref:`likelihood algorithm `. + .. _ramp_output_products: Output Products @@ -317,3 +321,39 @@ that pixel will be flagged as JUMP_DET in the corresponding integration in the "rateints" product. That pixel will also be flagged as JUMP_DET in the "rate" product. +.. _likelihood_algo: + +Likelihood Algorithm Details +---------------------------- +As an alternative to the OLS algorithm, a likelihood algorithm can be selected +with the step argument ``--ramp_fitting.algorithm=LIKELY``. This algorithm has +its own algorithm for jump detection. The normal jump detection algorithm runs +by default, but can be skipped when selecting the LIKELY algorithm. This +algorithm removes jump detection flags and sets jump detection flags. This jump +detection algorithm requires a minimum of four (4) NGROUPS. If the LIKELY +algorithm is selected for data with NGROUPS less than four, the ramp fitting +algorithm is changed to OLS_C. + +Each pixel is independently processed, but rather than operate on each +group/resultant directly, the likelihood algorithm is based on differences of +the groups/resultants :math:`d_i = r_i - r_{i-1}`. The model used to determine +the slope/countrate, :math:`a`, is: + +.. math:: + \chi^2 = ({\bf d} - a \cdot {\bf 1})^T C ({\bf d} - a \cdot {\bf 1}) \,, + +Differentiating, setting to zero, then solving for :math:`a` results in + +.. math:: + a = ({\bf 1}^T C {\bf d})({\bf 1}^T C {\bf 1})^T \,, + +The covariance matrix :math:`C` is a tridiagonal matrix, due to the nature of the +differences. Because the covariance matrix is tridiagonal, the computational +complexity reduces from :math:`O(n^3)` to :math:`O(n)`. To see the detailed +derivation and computations implemented, refer to the links above. +The Poisson and read noise computations are based on equations (27) and (28), in +`Brandt (2024) `__ + +For more details, especially for the jump detection portion in the liklihood +algorithm, see +`Brandt (2024) `__. diff --git a/src/stcal/ramp_fitting/likely_algo_classes.py b/src/stcal/ramp_fitting/likely_algo_classes.py new file mode 100644 index 00000000..fc71fd9f --- /dev/null +++ b/src/stcal/ramp_fitting/likely_algo_classes.py @@ -0,0 +1,319 @@ +import numpy as np +from scipy import special + + +class IntegInfo: + """ + Storage for the integration information for ramp fitting computations. + """ + def __init__(self, nints, nrows, ncols): + """ + Initialize output arrays. + + Parameters + ---------- + nints : int + The number of integrations in the data. + + nrows : int + The number of rows in the data. + + ncols : int + The number of columns in the data. + """ + dims = (nints, nrows, ncols) + self.data = np.zeros(shape=dims, dtype=np.float32) + + self.dq = np.zeros(shape=dims, dtype=np.uint32) + + self.var_poisson = np.zeros(shape=dims, dtype=np.float32) + self.var_rnoise = np.zeros(shape=dims, dtype=np.float32) + + self.err = np.zeros(shape=dims, dtype=np.float32) + + def prepare_info(self): + """ + Arrange output arrays as a tuple, which the ramp fit step expects. + """ + return (self.data, self.dq, self.var_poisson, self.var_rnoise, self.err) + + def get_results(self, result, integ, row): + """ + Capture the ramp fitting computation. + + Parameters + ---------- + result : RampResult + Holds computed ramp fitting information. + + integ : int + The current integration being operated on. + + row : int + The current row being operated on. + """ + self.data[integ, row, :] = result.countrate + self.err[integ, row, :] = result.uncert + self.var_poisson[integ, row, :] = result.var_poisson + self.var_rnoise[integ, row, :] = result.var_rnoise + + +class RampResult: + def __init__(self): + """ + Contains the ramp fitting results. + """ + self.countrate = None + self.chisq = None + self.uncert = None + self.var_poisson = None + self.var_rnoise = None + self.weights = None + + self.countrate_two_omit = None + self.chisq_two_omit = None + self.uncert_two_omit = None + + self.countrate_one_omit = None + self.jumpval_one_omit = None + self.jumpsig_one_omit = None + self.chisq_one_omit = None + self.uncert_one_omit = None + + def __repr__(self): + """ + Return string of information about the class. + """ + ostring = f"countrate = \n{self.countrate}" + ostring += f"\nchisq = \n{self.chisq}" + ostring += f"\nucert = \n{self.uncert}" + ''' + ostring += f"\nweights = \n{self.weights}" + + ostring += f"\ncountrate_two_omit = \n{self.countrate_two_omit}" + ostring += f"\nchisq_two_omit = \n{self.chisq_two_omit}" + ostring += f"\nuncert_two_omit = \n{self.uncert_two_omit}" + + ostring += f"\ncountrate_one_omit = \n{self.countrate_one_omit}" + ostring += f"\njumpval_one_omit = \n{self.jumpval_one_omit}" + ostring += f"\njumpsig_one_omit = \n{self.jumpsig_one_omit}" + ostring += f"\nchisq_one_omit = \n{self.chisq_one_omit}" + ostring += f"\nuncert_one_omit = \n{self.uncert_one_omit}" + ''' + + return ostring + + def fill_masked_reads(self, diffs2use): + """ + Mask groups to use for ramp fitting. + + Replace countrates, uncertainties, and chi squared values that + are NaN because resultant differences were doubly omitted. + For these cases, revert to the corresponding values in with + fewer omitted resultant differences to get the correct values + without double-counting omissions. + + This function replaces the relevant entries of + self.countrate_two_omit, self.chisq_two_omit, + self.uncert_two_omit, self.countrate_one_omit, and + self.chisq_one_omit in place. It does not return a value. + + Parameters + ---------- + diffs2use : ndarray + A 2D array matching self.countrate_one_omit in shape with zero + for resultant differences that were masked and one for + differences that were not masked. + """ + # replace entries that would be nan (from trying to + # doubly exclude read differences) with the global fits. + omit = diffs2use == 0 + ones = np.ones(diffs2use.shape) + + self.countrate_one_omit[omit] = (self.countrate * ones)[omit] + self.chisq_one_omit[omit] = (self.chisq * ones)[omit] + self.uncert_one_omit[omit] = (self.uncert * ones)[omit] + + omit = diffs2use[1:] == 0 + + self.countrate_two_omit[omit] = (self.countrate_one_omit[:-1])[omit] + self.chisq_two_omit[omit] = (self.chisq_one_omit[:-1])[omit] + self.uncert_two_omit[omit] = (self.uncert_one_omit[:-1])[omit] + + omit = diffs2use[:-1] == 0 + + self.countrate_two_omit[omit] = (self.countrate_one_omit[1:])[omit] + self.chisq_two_omit[omit] = (self.chisq_one_omit[1:])[omit] + self.uncert_two_omit[omit] = (self.uncert_one_omit[1:])[omit] + + +class Covar: + """ + class Covar holding read and photon noise components of alpha and + beta and the time intervals between the resultant midpoints + """ + def __init__(self, readtimes): + """ + Compute alpha and beta, the diagonal and off-diagonal elements of + the covariance matrix of the resultant differences, and the time + intervals between the resultant midpoints. + + Parameters + ---------- + readtimes : list + List of values or lists for the times of reads. If a list of + lists, times for reads that are averaged together to produce + a resultant. + """ + # Equations (4) and (11) in paper 1. + mean_t, tau, N, delta_t = self._compute_means_and_taus(readtimes) + + self.delta_t = delta_t + self.mean_t = mean_t + self.tau = tau + self.Nreads = N + + # Equations (28) and (29) in paper 1. + self._compute_alphas_and_betas(mean_t, tau, N, delta_t) + + def _compute_means_and_taus(self, readtimes): + """ + Computes the means and taus of defined in EQNs 4 and 11 in paper 1. + + Parameters + ---------- + readtimes : list + List of values or lists for the times of reads. If a list of + lists, times for reads that are averaged together to produce + a resultant. + """ + mean_t = [] # mean time of the resultant as defined in the paper + tau = [] # variance-weighted mean time of the resultant + N = [] # Number of reads per resultant + + for times in readtimes: + mean_t += [np.mean(times)] + + if hasattr(times, "__len__"): + # eqn 11 + N += [len(times)] + k = np.arange(1, N[-1] + 1) + if False: + tau += [ + 1 + / N[-1] ** 2 + * np.sum((2 * N[-1] + 1 - 2 * k) * np.array(times)) + ] + # tau += [(np.sum((2*N[-1] + 1 - 2*k)*np.array(times))) / N[-1]**2] + else: + length = N[-1] + tmp0 = (2 * length + 1) - (2 * k) + sm = np.sum(tmp0 * np.array(times)) + tmp = sm / length**2 + tau.append(tmp) + else: + tau += [times] + N += [1] + + # readtimes is a list of lists, so mean_t is the list of each + # mean of each list. + mean_t = np.array(mean_t) + tau = np.array(tau) + N = np.array(N) + delta_t = mean_t[1:] - mean_t[:-1] + + return mean_t, tau, N, delta_t + + def _compute_alphas_and_betas(self, mean_t, tau, N, delta_t): + """ + Computes the means and taus defined in EQNs 28 and 29 in paper 1. + + Parameters + ---------- + mean_t : ndarray + The means of the reads for each group. + + tau : ndarray + Intermediate computation. + + N : ndarray + The number of reads in each group. + + delta_t : ndarray + The group differences of integration ramps. + """ + self.alpha_readnoise = (1 / N[:-1] + 1 / N[1:]) / delta_t**2 + self.beta_readnoise = -1 / (N[1:-1] * delta_t[1:] * delta_t[:-1]) + + self.alpha_phnoise = (tau[:-1] + tau[1:] - 2 * mean_t[:-1]) / delta_t**2 + self.beta_phnoise = (mean_t[1:-1] - tau[1:-1]) / (delta_t[1:] * delta_t[:-1]) + + def calc_bias(self, countrates, sig, cvec, da=1e-7): + """ + Calculate the bias in the best-fit count rate from estimating the covariance matrix. + + Section 5 of paper 1. + + Arguments: + Parameters + ---------- + countrates : ndarray + Array of count rates at which the bias is desired. + + sig : float + Single read noise. + + cvec : ndarray + Weight vector on resultant differences for initial estimation + of count rate for the covariance matrix. Will be renormalized + inside this function. + + da : float + Fraction of the count rate plus sig**2 to use for finite difference + estimate of the derivative. Optional parameter. Default 1e-7. + + Returns + ------- + bias : ndarray + Bias of the best-fit count rate from using cvec plus the observed + resultants to estimate the covariance matrix. + """ + alpha = countrates[np.newaxis, :] * self.alpha_phnoise[:, np.newaxis] + alpha += sig**2 * self.alpha_readnoise[:, np.newaxis] + beta = countrates[np.newaxis, :] * self.beta_phnoise[:, np.newaxis] + beta += sig**2 * self.beta_readnoise[:, np.newaxis] + + # we only want the weights; it doesn't matter what the count rates are. + n = alpha.shape[0] + z = np.zeros((len(cvec), len(countrates))) + result_low_a = fit_ramps(z, self, sig, countrateguess=countrates) + + # try to avoid problems with roundoff error + da_incr = da * (countrates[np.newaxis, :] + sig**2) + + dalpha = da_incr * self.alpha_phnoise[:, np.newaxis] + dbeta = da_incr * self.beta_phnoise[:, np.newaxis] + result_high_a = fit_ramps(z, self, sig, countrateguess=countrates + da_incr) + # finite difference approximation to dw/da + + dw_da = (result_high_a.weights - result_low_a.weights) / da_incr + + bias = np.zeros(len(countrates)) + c = cvec / np.sum(cvec) + + for i in range(len(countrates)): + + C = np.zeros((n, n)) + for j in range(n): + C[j, j] = alpha[j, i] + for j in range(n - 1): + C[j + 1, j] = C[j, j + 1] = beta[j, i] + + bias[i] = np.linalg.multi_dot([c[np.newaxis, :], C, dw_da[:, i]]) + + sig_a = np.sqrt( + np.linalg.multi_dot([c[np.newaxis, :], C, c[:, np.newaxis]]) + ) + bias[i] *= 0.5 * (1 + special.erf(countrates[i] / sig_a / 2**0.5)) + + return bias diff --git a/src/stcal/ramp_fitting/likely_fit.py b/src/stcal/ramp_fitting/likely_fit.py new file mode 100644 index 00000000..7a17fa97 --- /dev/null +++ b/src/stcal/ramp_fitting/likely_fit.py @@ -0,0 +1,1255 @@ +#! /usr/bin/env python + +import logging +import multiprocessing +import time +import scipy +import sys +import warnings + +from multiprocessing import cpu_count +from pprint import pprint + +import numpy as np + +from . import ramp_fit_class, utils +from .likely_algo_classes import IntegInfo, RampResult, Covar + + +DELIM = "=" * 80 +SQRT2 = np.sqrt(2) +LIKELY_MIN_NGROUPS = 4 + +log = logging.getLogger(__name__) +log.setLevel(logging.DEBUG) + +log = logging.getLogger(__name__) +log.setLevel(logging.DEBUG) + + +def likely_ramp_fit(ramp_data, readnoise_2d, gain_2d): + """ + Invoke ramp fitting using the likelihood algorithm. + + Parameters + ---------- + ramp_data : RampData + Input data necessary for computing ramp fitting. + + readnoise_2d : ndarray + readnoise for all pixels + + gain_2d : ndarray + gain for all pixels + + 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. + """ + image_info, integ_info, opt_info = None, None, None + + nints, ngroups, nrows, ncols = ramp_data.data.shape + + if ngroups < LIKELY_MIN_NGROUPS: + raise ValueError("Likelihood fit requires at least 4 groups.") + + + remove_jump_detection_flags(ramp_data) + + readtimes = get_readtimes(ramp_data) + + covar = Covar(readtimes) + integ_class = IntegInfo(nints, nrows, ncols) + + readnoise_2d = readnoise_2d / SQRT2 + + for integ in range(nints): + data = ramp_data.data[integ, :, :, :] + gdq = ramp_data.groupdq[integ, :, :, :].copy() + pdq = ramp_data.pixeldq[:, :].copy() + + # Eqn (5) + diff = (data[1:] - data[:-1]) / covar.delta_t[:, np.newaxis, np.newaxis] + alldiffs2use = np.ones(diff.shape, np.uint8) + + for row in range(nrows): + d2use = determine_diffs2use(ramp_data, integ, row, diff) + d2use_copy = d2use.copy() # Use to flag jumps + if ramp_data.rejection_threshold is not None: + threshold_one_omit = ramp_data.rejection_threshold**2 + pval = scipy.special.erfc(ramp_data.rejection_threshold/SQRT2) + threshold_two_omit = scipy.stats.chi2.isf(pval, 2) + if np.isinf(threshold_two_omit): + threshold_two_omit = threshold_one_omit + 10 + d2use, countrates = mask_jumps( + diff[:, row], covar, readnoise_2d[row], gain_2d[row], + threshold_one_omit=threshold_one_omit, + threshold_two_omit=threshold_two_omit, + diffs2use=d2use + ) + else: + d2use, countrates = mask_jumps( + diff[:, row], covar, readnoise_2d[row], gain_2d[row], + diffs2use=d2use + ) + + # Set jump detection flags + jump_locs = d2use_copy ^ d2use + jump_locs[jump_locs > 0] = ramp_data.flags_jump_det + gdq[1:, row, :] |= jump_locs + + alldiffs2use[:, row] = d2use + + rateguess = countrates * (countrates > 0) + ramp_data.average_dark_current[row, :] + result = fit_ramps( + diff[:, row], + covar, + gain_2d[row], + readnoise_2d[row], + diffs2use=d2use, + count_rate_guess=rateguess, + ) + integ_class.get_results(result, integ, row) + + pdq = utils.dq_compress_sect(ramp_data, integ, gdq, pdq) + integ_class.dq[integ, :, :] = pdq + + del gdq + + integ_info = integ_class.prepare_info() + image_info = compute_image_info(integ_class, ramp_data) + + return image_info, integ_info, opt_info + + +def mask_jumps( + diffs, + covar, + rnoise, + gain, + threshold_one_omit=20.25, + threshold_two_omit=23.8, + diffs2use=None, +): + + """ + Implements a likelihood-based, iterative jump detection algorithm. + + Parameters + ---------- + diffs : ndarray + The group differences of the data array for a given integration and row + (ngroups-1, ncols). + + covar : Covar + The class instance that computes and contains the covariance matrix info. + + rnoise : ndarray + The read noise (ncols,) + + gain : ndarray + The gain (ncols,) + + threshold_one_omit : float + Minimum chisq improvement to exclude a single resultant difference. + Default: 20.25. + + threshold_two_omit : float + Minimum chisq improvement to exclude two sequential resultant differences. + Default 23.8. + + diffs2use : ndarray + A boolean array definined the segmented ramps for each pixel in a row. + (ngroups-1, ncols) + + Returns + ------- + dffs2use : ndarray + A boolean array definined the segmented ramps for each pixel in a row. + (ngroups-1, ncols) + + countrates : ndarray + Count rate estimates used to estimate the covariance matrix. + Optional, default is None. + + """ + # Force a copy of the input array for more efficient memory access. + # loc_diff = diffs * 1 + loc_diff = diffs + + # We can use one-omit searches only where the reads immediately + # preceding and following have just one read. If a readout + # pattern has more than one read per resultant but significant + # gaps between resultants then a one-omit search might still be a + # good idea even with multiple-read resultants. + one_omit_ok = covar.Nreads[1:] * covar.Nreads[:-1] >= 1 + one_omit_ok[0] = one_omit_ok[-1] = True + + # Other than that, we need to omit two. If a resultant has more + # than two reads, we need to omit both differences containing it + # (one pair of omissions in the differences). + two_omit_ok = covar.Nreads[1:-1] > 1 + + # This is the array to return: one for resultant differences to + # use, zero for resultant differences to ignore. + if diffs2use is None: + diffs2use = np.ones(loc_diff.shape, np.uint8) + + # We need to estimate the covariance matrix. I'll use the median + # here for now to limit problems with the count rate in reads with + # jumps (which is what we are looking for) since we'll be using + # likelihoods and chi squared; getting the covariance matrix + # reasonably close to correct is important. + count_rate_guess = np.median(loc_diff, axis=0)[np.newaxis, :] + + count_rate_guess *= count_rate_guess > 0 + + # boolean arrays to be used later + recheck = np.ones(loc_diff.shape[1]) == 1 + dropped = np.ones(loc_diff.shape[1]) == 0 + + for j in range(loc_diff.shape[0]): + # No need for indexing on the first pass. + if j == 0: + result = fit_ramps( + loc_diff, + covar, + gain, + rnoise, + count_rate_guess=count_rate_guess, + diffs2use=diffs2use, + detect_jumps=True, + ) + # Also save the count rates so that we can use them later + # for debiasing. + countrate = result.countrate * 1.0 + else: + result = fit_ramps( + loc_diff[:, recheck], + covar, + gain[recheck], + rnoise[recheck], + count_rate_guess=count_rate_guess[:, recheck], + diffs2use=diffs2use[:, recheck], + detect_jumps=True, + ) + + # Chi squared improvements + dchisq_two = result.chisq - result.chisq_two_omit + dchisq_one = result.chisq - result.chisq_one_omit + + # We want the largest chi squared difference + best_dchisq_one = np.amax(dchisq_one * one_omit_ok[:, np.newaxis], axis=0) + best_dchisq_two = np.amax( + dchisq_two * two_omit_ok[:, np.newaxis], axis=0 + ) + + # Is the best improvement from dropping one resultant + # difference or two? Two drops will always offer more + # improvement than one so penalize them by the respective + # thresholds. Then find the chi squared improvement + # corresponding to dropping either one or two reads, whichever + # is better, if either exceeded the threshold. + onedropbetter = ( + best_dchisq_one - threshold_one_omit > best_dchisq_two - threshold_two_omit + ) + + best_dchisq = ( + best_dchisq_one * (best_dchisq_one > threshold_one_omit) * onedropbetter + ) + best_dchisq += ( + best_dchisq_two * (best_dchisq_two > threshold_two_omit) * (~onedropbetter) + ) + + # If nothing exceeded the threshold set the improvement to + # NaN so that dchisq==best_dchisq is guaranteed to be False. + best_dchisq[best_dchisq == 0] = np.nan + + # Now make the masks for which resultant difference(s) to + # drop, count the number of ramps affected, and drop them. + # If no ramps were affected break the loop. + dropone = dchisq_one == best_dchisq + droptwo = dchisq_two == best_dchisq + + drop = np.any([np.sum(dropone, axis=0), np.sum(droptwo, axis=0)], axis=0) + + if np.sum(drop) == 0: + break + + # Store the updated counts with omitted reads + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + new_cts = np.zeros(np.sum(recheck)) + i_d1 = np.sum(dropone, axis=0) > 0 + new_cts[i_d1] = np.sum(result.countrate_one_omit * dropone, axis=0)[i_d1] + i_d2 = np.sum(droptwo, axis=0) > 0 + new_cts[i_d2] = np.sum(result.countrate_two_omit * droptwo, axis=0)[i_d2] + + # zero out count rates with drops and add their new values back in + countrate[recheck] *= drop == 0 + countrate[recheck] += new_cts + + # Drop the read (set diffs2use=0) if the boolean array is True. + diffs2use[:, recheck] *= ~dropone + diffs2use[:-1, recheck] *= ~droptwo + diffs2use[1:, recheck] *= ~droptwo + + # No need to repeat this on the entire ramp, only re-search + # ramps that had a resultant difference dropped this time. + dropped[:] = False + dropped[recheck] = drop + recheck[:] = dropped + + # Do not try to search for bad resultants if we have already + # given up on all but one, two, or three resultant differences + # in the ramp. If there are only two left we have no way of + # choosing which one is "good". If there are three left we + # run into trouble in case we need to discard two. + recheck[np.sum(diffs2use, axis=0) <= 3] = False + + return diffs2use, countrate + + +def get_readtimes(ramp_data): + """ + Get the read times needed to compute the covariance matrices. + + If there is already a read_pattern in the ramp_data class, then just get it. + If not, then one needs to be constructed. If one needs to be constructed it + is assumed the groups are evenly spaced in time, as are the frames that make + up the group. If each group has only one frame and no group gap, then a list + of the group times is returned. If nframes > 0, then a list of lists of each + frame time in each group is returned with the assumption: + group_time = (nframes + groupgap) * frame_time + + Parameters + ---------- + ramp_data : RampData + Input data necessary for computing ramp fitting. + + Returns + ------- + readtimes : list + A list of frame times for each frame used in the computation of the ramp. + """ + if ramp_data.read_pattern is not None: + return ramp_data.read_pattern + + ngroups = ramp_data.data.shape[1] + tot_frames = ramp_data.nframes + ramp_data.groupgap + tot_nreads = np.arange(1, ramp_data.nframes + 1) + rtimes = [ + (tot_nreads + k * tot_frames) * ramp_data.frame_time for k in range(ngroups) + ] + + return rtimes + + +def compute_image_info(integ_class, ramp_data): + """ + Combine all integrations into a single image of rates, + variances, and DQ flags. + + Parameters + ---------- + integ_class : IntegInfo + Contains the rateints product calculations. + + ramp_data : RampData + Input data necessary for computing ramp fitting. + + Returns + ------- + image_info : tuple + The list of arrays for the rate product. + """ + if integ_class.data.shape[0] == 1: + data = integ_class.data[0, :, :] + dq = integ_class.dq[0, :, :] + var_p = integ_class.var_poisson[0, :, :] + var_r = integ_class.var_rnoise[0, :, :] + var_e = integ_class.err[0, :, :] + return (data, dq, var_p, var_r, var_e) + + dq = utils.dq_compress_final(integ_class.dq, ramp_data) + + slope = np.nanmedian(integ_class.data, axis=0) + for _ in range(2): + rate_scale = slope[np.newaxis, :] / integ_class.data + rate_scale[(~np.isfinite(rate_scale)) | (rate_scale < 0)] = 0 + all_var_p = integ_class.var_poisson * rate_scale + weight = 1/(all_var_p + integ_class.var_rnoise) + weight /= np.nansum(weight, axis=0)[np.newaxis, :] + tmp_slope = integ_class.data * weight + all_nan = np.all(np.isnan(tmp_slope), axis=0) + slope = np.nansum(tmp_slope, axis=0) + slope[all_nan] = np.nan + + tmp_v = all_var_p * weight**2 + all_nan = np.all(np.isnan(tmp_v), axis=0) + var_p = np.nansum(tmp_v, axis=0) + var_p[all_nan] = np.nan + + tmp_v = integ_class.var_rnoise * weight**2 + all_nan = np.all(np.isnan(tmp_v), axis=0) + var_r = np.nansum(tmp_v, axis=0) + var_r[all_nan] = np.nan + + err = np.sqrt(var_p + var_r) + + return (slope, dq, var_p, var_r, err) + + +def remove_jump_detection_flags(ramp_data): + """ + Remove the JUMP_DET flag from the group DQ array + + Parameters + ---------- + ramp_data : RampData + Input data necessary for computing ramp fitting. + """ + jump = ramp_data.flags_jump_det + gdq = ramp_data.groupdq + wh_jump = np.where(np.bitwise_and(gdq.astype(np.uint32), jump)) + gdq[wh_jump] -= jump + ramp_data.groupdq = gdq + + +def determine_diffs2use(ramp_data, integ, row, diffs): + """ + Compute the diffs2use mask based on DQ flags of a row. + + Parameters + ---------- + ramp_data : RampData + Input data necessary for computing ramp fitting. + + integ : int + The current integration being processed. + + row : int + The current row being processed. + + diffs : ndarray + The group differences of the data array for a given integration and row + (ngroups-1, ncols). + + Returns + ------- + d2use : ndarray + A boolean array definined the segmented ramps for each pixel in a row. + (ngroups-1, ncols) + """ + # import ipdb; ipdb.set_trace() + _, ngroups, _, ncols = ramp_data.data.shape + dq = np.zeros(shape=(ngroups, ncols), dtype=np.uint8) + dq[:, :] = ramp_data.groupdq[integ, :, row, :] + d2use_tmp = np.ones(shape=diffs.shape, dtype=np.uint8) + d2use = d2use_tmp[:, row] + + d2use[dq[1:, :] != 0] = 0 + d2use[dq[:-1, :] != 0] = 0 + + return d2use + + +def initial_count_rate_guess(diffs, diffs2use): + """ + Compute the initial count rate. + + Parameters + ---------- + diffs : ndarray + The group differences of the data (ngroups-1, nrows, ncols). + + diffs2use : ndarray + Boolean mask determining with group differences to use (ngroups-1, ncols). + + Returns + ------- + count_rate_guess : ndarray + The initial count rate. + """ + # initial guess for count rate is the average of the unmasked + # group differences unless otherwise specified. + num = np.sum((diffs * diffs2use), axis=0) + den = np.sum(diffs2use, axis=0) + + count_rate_guess = num / den + count_rate_guess *= count_rate_guess > 0 + + return count_rate_guess + + +# RAMP FITTING BEGIN +def fit_ramps( + diffs, + covar, + gain, + rnoise, # Referred to as 'sig' in fitramp repo + count_rate_guess=None, + diffs2use=None, + detect_jumps=False, + rescale=True, + dn_scale=10.0, +): + """ + Fit ramps on a row of data. + + Fits ramps to read differences using the covariance matrix for + the read differences as given by the diagonal elements and the + off-diagonal elements. + + Parameters + ---------- + diffs : ndarray + The group differences of the data (ngroups-1, nrows, ncols). + + covar : Covar + The class instance that computes and contains the covariance matrix info. + + gain : ndarray + The gain (ncols,) + + rnoise : ndarray + The read noise (ncols,) + + count_rate_guess : ndarray + Count rate estimates used to estimate the covariance matrix. + Optional, default is None. + + diffs2use : ndarray + Boolean mask determining with group differences to use (ngroups-1, ncols). + Optional, default is None, which results in a mask of all 1's. + + detect_jumps : boolean + Run jump detection. + Optional, default is False. + + rescale : boolean + Scale the covariance matrix internally to avoid possible + overflow/underflow problems for long ramps. + Optional, default is True. + + dn_scale : float + Scaling factor. + + Returns + ------- + result : RampResult + Holds computed ramp fitting information. XXX - rename + """ + if diffs2use is None: + # Use all diffs + diffs2use = np.ones(diffs.shape, np.uint8) + + # diffs is (ngroups, ncols) of the current row + if count_rate_guess is None: + count_rate_guess = initial_count_rate_guess(covar, diffs, diffs2use) + + alpha_tuple, beta_tuple, scale = compute_alphas_betas( + count_rate_guess, gain, rnoise, covar, rescale, diffs, dn_scale + ) + alpha, alpha_phnoise, alpha_readnoise = alpha_tuple + beta, beta_phnoise, beta_readnoise = beta_tuple + + ndiffs, npix = diffs.shape + + # Mask group differences that should be ignored. This is half + # of what we need to do to mask these group differences; the + # rest comes later. + diff_mask = diffs * diffs2use + beta = beta * diffs2use[1:] * diffs2use[:-1] + + # All definitions and formulas here are in the paper. + # --- Till line 237: Paper 1 section 4 + theta = compute_thetas(ndiffs, npix, alpha, beta) # EQNs 38-40 + phi = compute_phis(ndiffs, npix, alpha, beta) # EQNs 41-43 + + sgn = np.ones((ndiffs, npix)) + sgn[::2] = -1 + + Phi = compute_Phis(ndiffs, npix, beta, phi, sgn) # EQN 46 + PhiD = compute_PhiDs(ndiffs, npix, beta, phi, sgn, diff_mask) # EQN ?? + Theta = compute_Thetas(ndiffs, npix, beta, theta, sgn) # EQN 47 + ThetaD = compute_ThetaDs(ndiffs, npix, beta, theta, sgn, diff_mask) # EQN 48 + + dB, dC, A, B, C = matrix_computations( + ndiffs, + npix, + sgn, + diff_mask, + diffs2use, + beta, + phi, + Phi, + PhiD, + theta, + Theta, + ThetaD, + ) + + result = get_ramp_result( + dC, + A, + B, + C, + scale, + alpha_phnoise, + alpha_readnoise, + beta_phnoise, + beta_readnoise, + ) + + # --- Beginning at line 250: Paper 1 section 4 + + if detect_jumps: + # result = compute_jump_detects( + # result, ndiffs, diffs2use, dC, dB, A, B, C, scale, beta, phi, theta, covar + # ) + result = compute_jump_detects( + result, ndiffs, diffs2use, dC, dB, A, B, C, scale, beta, phi, theta + ) + + return result + + +# RAMP FITTING END + + +def compute_jump_detects( + result, ndiffs, diffs2use, dC, dB, A, B, C, scale, beta, phi, theta +): + """ + Detect jumps in ramps. + + The code below computes the best chi squared, best-fit slope, + and its uncertainty leaving out each group difference in + turn. There are ndiffs possible differences that can be + omitted. + + Then do it omitting two consecutive reads. There are ndiffs-1 + possible pairs of adjacent reads that can be omitted. + + Paper II, sections 3.1 and 3.2 + + Parameters + ---------- + result : RampResult + The results of the ramp fitting for a given row of pixels in an integration. + + ndiffs : int + Number of differences. + + diffs2use : ndarray + Boolean mask determining with group differences to use (ngroups-1, ncols). + + dC : ndarray + Intermediate computation. + + dB : ndarray + Intermediate computation. + + A : ndarray + Intermediate computation. + + B : ndarray + Intermediate computation. + + C : ndarray + Intermediate computation. + + scale : ndarray or integer + Overflow/underflow prevention scale. + + beta : ndarray + Off diagonal of covariance matrix. + + phi : ndarray + Intermediate computation. + + theta : ndarray + Intermediate computation. + + covar : Covar + The class instance that computes and contains the covariance matrix info. + + + Returns + ------- + result : RampResult + The results of the ramp fitting for a given row of pixels in an integration. + """ + # Diagonal elements of the inverse covariance matrix + Cinv_diag = theta[:-1] * phi[1:] / theta[ndiffs] + Cinv_diag *= diffs2use + + # Off-diagonal elements of the inverse covariance matrix + # one spot above and below for the case of two adjacent + # differences to be masked + Cinv_offdiag = -beta * theta[:-2] * phi[2:] / theta[ndiffs] + + # Equations in the paper: best-fit a, b + # + # Catch warnings in case there are masked group + # differences, since these will be overwritten later. No need + # to warn about division by zero here. + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + a = (Cinv_diag * B - dB * dC) / (C * Cinv_diag - dC**2) + b = (dB - a * dC) / Cinv_diag + + result.countrate_one_omit = a + result.jumpval_one_omit = b + + # Use the best-fit a, b to get chi squared + result.chisq_one_omit = ( + A + + a**2 * C + - 2 * a * B + + b**2 * Cinv_diag + - 2 * b * dB + + 2 * a * b * dC + ) + + # invert the covariance matrix of a, b to get the uncertainty on a + result.uncert_one_omit = np.sqrt(Cinv_diag / (C * Cinv_diag - dC**2)) + result.jumpsig_one_omit = np.sqrt(C / (C * Cinv_diag - dC**2)) + + result.chisq_one_omit /= scale + result.uncert_one_omit *= np.sqrt(scale) + result.jumpsig_one_omit *= np.sqrt(scale) + + # Now for two omissions in a row. This is more work. Again, + # all equations are in the paper. I first define three + # factors that will be used more than once to save a bit of + # computational effort. + cpj_fac = dC[:-1] ** 2 - C * Cinv_diag[:-1] + cjck_fac = dC[:-1] * dC[1:] - C * Cinv_offdiag + bcpj_fac = B * dC[:-1] - dB[:-1] * C + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + # best-fit a, b, c + c = bcpj_fac / cpj_fac - (B * dC[1:] - dB[1:] * C) / cjck_fac + c /= cjck_fac / cpj_fac - (dC[1:] ** 2 - C * Cinv_diag[1:]) / cjck_fac + b = (bcpj_fac - c * cjck_fac) / cpj_fac + a = (B - b * dC[:-1] - c * dC[1:]) / C + result.countrate_two_omit = a + + # best-fit chi squared + result.chisq_two_omit = ( + A + a**2 * C + b**2 * Cinv_diag[:-1] + c**2 * Cinv_diag[1:] + ) + result.chisq_two_omit -= 2 * a * B + 2 * b * dB[:-1] + 2 * c * dB[1:] + result.chisq_two_omit += ( + 2 * a * b * dC[:-1] + 2 * a * c * dC[1:] + 2 * b * c * Cinv_offdiag + ) + result.chisq_two_omit /= scale + + # uncertainty on the slope from inverting the (a, b, c) + # covariance matrix + fac = Cinv_diag[1:] * Cinv_diag[:-1] - Cinv_offdiag**2 + term2 = dC[:-1] * (dC[:-1] * Cinv_diag[1:] - Cinv_offdiag * dC[1:]) + term3 = dC[1:] * (dC[:-1] * Cinv_offdiag - Cinv_diag[:-1] * dC[1:]) + result.uncert_two_omit = np.sqrt(fac / (C * fac - term2 + term3)) + result.uncert_two_omit *= np.sqrt(scale) + + result.fill_masked_reads(diffs2use) + + return result + + +def compute_alphas_betas(count_rate_guess, gain, rnoise, covar, rescale, diffs, dn_scale): + """ + Compute alpha, beta, and scale needed for ramp fit. + + Elements of the covariance matrix. + + Parameters + ---------- + count_rate_guess : ndarray + Initial guess (ncols,) + + gain : ndarray + Gain (ncols,) + + rnoise : ndarray + Read noise (ncols,) + + covar : Covar + The class instance that computes and contains the covariance matrix info. + + rescale : bool + Determination to rescale covariance matrix. + + diffs : ndarray + The group differences of the data (ngroups-1, nrows, ncols). + + dn_scale : float + Scaling factor. + + Returns + ------- + alpha : ndarray + Diagonal of covariance matrix. + + beta : ndarray + Off diagonal of covariance matrix. + + scale : ndarray or integer + Overflow/underflow prevention scale. + + """ + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + + alpha_phnoise = count_rate_guess / gain * covar.alpha_phnoise[:, np.newaxis] + alpha_readnoise = rnoise**2 * covar.alpha_readnoise[:, np.newaxis] + alpha = alpha_phnoise + alpha_readnoise + + beta_phnoise = count_rate_guess / gain * covar.beta_phnoise[:, np.newaxis] + beta_readnoise = rnoise**2 * covar.beta_readnoise[:, np.newaxis] + beta = beta_phnoise + beta_readnoise + + warnings.resetwarnings() + + ndiffs, npix = diffs.shape + + # Rescale the covariance matrix to a determinant of 1 to + # avoid possible overflow/underflow. The uncertainty and chi + # squared value will need to be scaled back later. Note that + # theta[-1] is the determinant of the covariance matrix. + # + # The method below uses the fact that if all alpha and beta + # are multiplied by f, theta[i] is multiplied by f**i. Keep + # a running track of these factors to construct the scale at + # the end, and keep scaling throughout so that we never risk + # overflow or underflow. + + if rescale: + # scale = np.exp(np.mean(np.log(alpha), axis=0)) + theta = np.ones((ndiffs + 1, npix)) + theta[1] = alpha[0] + + scale = theta[0] * 1 + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + + for i in range(2, ndiffs + 1): + theta[i] = ( + alpha[i - 1] / scale * theta[i - 1] + - beta[i - 2] ** 2 / scale**2 * theta[i - 2] + ) + + # Scaling every ten steps in safe for alpha up to 1e20 + # or so and incurs a negligible computational cost for + # the fractional power. + + if i % int(dn_scale) == 0 or i == ndiffs: + f = theta[i] ** (1 / i) + scale *= f + tmp = theta[i] / f + theta[i - 1] /= tmp + theta[i - 2] /= tmp / f + theta[i] = 1 + + warnings.resetwarnings() + else: + scale = 1 + + alpha /= scale + beta /= scale + + alpha_tuple = (alpha, alpha_phnoise, alpha_readnoise) + beta_tuple = (beta, beta_phnoise, beta_readnoise) + + return alpha_tuple, beta_tuple, scale + + +def compute_thetas(ndiffs, npix, alpha, beta): + """ + Computes intermediate theta values for ramp fitting. + + EQNs 38-40 + + Parameters + ---------- + ndiffs : int + Number of differences. + + npix : int + Number of columns in a row. + + alpha : ndarray + Diagonal of covariance matrix. + + beta : ndarray + Off diagonal of covariance matrix. + + Returns + ------- + theta : ndarray + """ + theta = np.ones((ndiffs + 1, npix)) + theta[1] = alpha[0] + for i in range(2, ndiffs + 1): + theta[i] = alpha[i - 1] * theta[i - 1] - beta[i - 2] ** 2 * theta[i - 2] + return theta + + +def compute_phis(ndiffs, npix, alpha, beta): + """ + Computes intermediate phi values for ramp fitting. + + EQNs 41-43 + + Parameters + ---------- + ndiffs : int + Number of differences. + + npix : int + Number of columns in a row. + + alpha : ndarray + Diagonal of covariance matrix. + + beta : ndarray + Off diagonal of covariance matrix. + + Returns + ------- + phi : ndarray + """ + phi = np.ones((ndiffs + 1, npix)) + phi[ndiffs - 1] = alpha[ndiffs - 1] + for i in range(ndiffs - 2, -1, -1): + phi[i] = alpha[i] * phi[i + 1] - beta[i] ** 2 * phi[i + 2] + return phi + + +def compute_Phis(ndiffs, npix, beta, phi, sgn): + """ + Computes intermediate Phi values for ramp fitting. + + EQN 46 + + Parameters + ---------- + ndiffs : int + Number of differences. + + npix : int + Number of columns in a row. + + alpha : ndarray + Diagonal of covariance matrix. + + beta : ndarray + Off diagonal of covariance matrix. + + sgn : ndarray + Oscillating 1, -1 sequence. + + Returns + ------- + Phi : ndarray + """ + Phi = np.zeros((ndiffs, npix)) + for i in range(ndiffs - 2, -1, -1): + Phi[i] = Phi[i + 1] * beta[i] + sgn[i + 1] * beta[i] * phi[i + 2] + return Phi + + +def compute_PhiDs(ndiffs, npix, beta, phi, sgn, diff_mask): + """ + EQN 4, Paper II + This one is defined later in the paper and is used for jump detection. + + Parameters + ---------- + ndiffs : int + Number of differences. + + npix : int + Number of columns in a row. + + beta : ndarray + Off diagonal of covariance matrix. + + phi : ndarray + Intermediate computation. + + sgn : ndarray + Oscillating 1, -1 sequence. + + diff_mask : ndarray + Mask of differences used. + + Returns + ------- + PhiD: ndarray + """ + PhiD = np.zeros((ndiffs, npix)) + for i in range(ndiffs - 2, -1, -1): + PhiD[i] = (PhiD[i + 1] + sgn[i + 1] * diff_mask[i + 1] * phi[i + 2]) * beta[i] + return PhiD + + +def compute_Thetas(ndiffs, npix, beta, theta, sgn): + """ + EQN 47 + + Parameters + ---------- + ndiffs : int + Number of differences. + + npix : int + Number of columns in a row. + + beta : ndarray + Off diagonal of covariance matrix. + + theta : ndarray + Intermediate computation. + + sgn : ndarray + Oscillating 1, -1 sequence. + + Returns + ------- + Theta : ndarray + """ + Theta = np.zeros((ndiffs, npix)) + Theta[0] = -theta[0] + for i in range(1, ndiffs): + Theta[i] = Theta[i - 1] * beta[i - 1] + sgn[i] * theta[i] + return Theta + + +def compute_ThetaDs(ndiffs, npix, beta, theta, sgn, diff_mask): + """ + EQN 48 + + Parameters + ---------- + ndiffs : int + Number of differences. + + npix : int + Number of columns in a row. + + beta : ndarray + Off diagonal of covariance matrix. + + theta : ndarray + Intermediate computation. + + sgn : ndarray + Oscillating 1, -1 sequence. + + diff_mask : ndarray + Mask of differences used. + + Returns + ------- + ThetaD : ndarray + """ + ThetaD = np.zeros((ndiffs + 1, npix)) + ThetaD[1] = -diff_mask[0] * theta[0] + for i in range(1, ndiffs): + ThetaD[i + 1] = beta[i - 1] * ThetaD[i] + sgn[i] * diff_mask[i] * theta[i] + return ThetaD + + +def matrix_computations( + ndiffs, npix, sgn, diff_mask, diffs2use, beta, phi, Phi, PhiD, theta, Theta, ThetaD +): + """ + Compute matrix computations needed for ramp fitting. + + EQNs 61-63, 71, 75 + + Parameters + ---------- + ndiffs : int + Number of differences. + + npix : int + Number of columns in a row. + + sgn : ndarray + Oscillating 1, -1 sequence. + + diff_mask : ndarray + Mask of differences used. + + diff2use : ndarray + Masked differences. + + beta : ndarray + Off diagonal of covariance matrix. + + phi : ndarray + Intermediate computation. + + Phi : ndarray + Intermediate computation. + + PhiD : ndarray + Intermediate computation. + + theta : ndarray + Intermediate computation. + + Theta : ndarray + Intermediate computation. + + ThetaD : ndarray + Intermediate computation. + + Returns + ------- + dB : ndarray + Intermediate computation. + + dC : ndarray + Intermediate computation. + + A : ndarray + Intermediate computation. + + B : ndarray + Intermediate computation. + + C : ndarray + Intermediate computation. + """ + beta_extended = np.ones((ndiffs, npix)) + beta_extended[1:] = beta + + # C' and B' in the paper + + dC = sgn / theta[ndiffs] * (phi[1:] * Theta + theta[:-1] * Phi) + dC *= diffs2use # EQN 71 + + dB = sgn / theta[ndiffs] * (phi[1:] * ThetaD[1:] + theta[:-1] * PhiD) # EQN 75 + + # {\cal A}, {\cal B}, {\cal C} in the paper + + # EQNs 61-63 + A = 2 * np.sum( + diff_mask * sgn / theta[-1] * beta_extended * phi[1:] * ThetaD[:-1], axis=0 + ) + A += np.sum(diff_mask**2 * theta[:-1] * phi[1:] / theta[ndiffs], axis=0) + + B = np.sum(diff_mask * dC, axis=0) + C = np.sum(dC, axis=0) + + return dB, dC, A, B, C + + +def get_ramp_result( + dC, A, B, C, scale, alpha_phnoise, alpha_readnoise, + beta_phnoise, beta_readnoise, +): + """ + Use intermediate computations to fit the ramp and save the results. + + Parameters + ---------- + dC : ndarray + Intermediate computation. + + A : ndarray + Intermediate computation. + + B : ndarray + Intermediate computation. + + C : ndarray + Intermediate computation. + + rescale : float + Factor applied to each element of the covariance matrix to + normalize its determinant in order to avoid possible + overflow/underflow problems for long ramps. + + alpha_phnoise : ndarray + The photon noise contribution to the alphas. + + alpha_readnoise : ndarray + The read noise contribution to the alphas. + + beta_phnoise : ndarray + The photon noise contribution to the betas. + + beta_readnoise : ndarray + The read noise contribution to the betas. + + Returns + ------- + result : RampResult + The results of the ramp fitting for a given row of pixels in an integration. + """ + result = RampResult() + + # Finally, save the best-fit count rate, chi squared, uncertainty + # in the count rate, and the weights used to combine the + # groups. + + warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning) + warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning) + invC = 1 / C + result.countrate = B * invC + result.chisq = (A - B**2 / C) / scale + + result.uncert = np.sqrt(scale / C) + result.weights = dC / C + + result.var_poisson = np.sum(result.weights**2 * alpha_phnoise, axis=0) + result.var_poisson += 2 * np.sum( + result.weights[1:] * result.weights[:-1] * beta_phnoise, axis=0 + ) + + result.var_rnoise = np.sum(result.weights**2 * alpha_readnoise, axis=0) + result.var_rnoise += 2 * np.sum( + result.weights[1:] * result.weights[:-1] * beta_readnoise, axis=0 + ) + + warnings.resetwarnings() + + return result + + +################################################################################ +################################## DEBUG ####################################### + + +def dbg_print_info(group_time, readtimes, data, diff): + print(DELIM) + print(f"group_time = {group_time}") + print(DELIM) + print(f"readtimes = {array_string(np.array(readtimes))}") + print(DELIM) + print(f"data = {array_string(data[:, 0, 0])}") + print(DELIM) + data_gt = data / group_time + print(f"data / gt = {array_string(data_gt[:, 0, 0])}") + print(DELIM) + print(f"diff = {array_string(diff[:, 0, 0])}") + print(DELIM) + + +def array_string(arr, prec=4): + return np.array2string(arr, precision=prec, max_line_width=np.nan, separator=", ") diff --git a/src/stcal/ramp_fitting/ramp_fit.py b/src/stcal/ramp_fitting/ramp_fit.py index 799c8fb9..b5981c07 100755 --- a/src/stcal/ramp_fitting/ramp_fit.py +++ b/src/stcal/ramp_fitting/ramp_fit.py @@ -19,8 +19,9 @@ from astropy import units as u from . import ( - gls_fit, # used only if algorithm is "GLS" - ols_fit, # used only if algorithm is "OLS" + gls_fit, # used only if algorithm is "GLS" + likely_fit, # used only if algorithm is "LIKELY" + ols_fit, # used only if algorithm is "OLS" ramp_fit_class, ) @@ -92,6 +93,10 @@ def create_ramp_fit_class(model, algorithm, dqflags=None, suppress_one_group=Fal ramp_data.zeroframe = model.zeroframe ramp_data.algorithm = algorithm + + if hasattr(model.meta.exposure, "read_pattern"): + ramp_data.read_pattern = [list(reads) for reads in model.meta.exposure.read_pattern] + ramp_data.set_dqflags(dqflags) ramp_data.start_row = 0 ramp_data.num_rows = ramp_data.data.shape[2] @@ -141,6 +146,7 @@ def ramp_fit( algorithm : str 'OLS' specifies that ordinary least squares should be used; 'GLS' specifies that generalized least squares should be used. + 'LIKELY' specifies that maximum likelihood should be used. weighting : str 'optimal' specifies that optimal weighting should be used; @@ -186,9 +192,6 @@ def ramp_fit( # data models. ramp_data = create_ramp_fit_class(model, algorithm, dqflags, suppress_one_group) - if algorithm.upper() == "OLS_C": - ramp_data.run_c_code = True - return ramp_fit_data( ramp_data, buffsize, save_opt, readnoise_2d, gain_2d, algorithm, weighting, max_cores, dqflags ) @@ -222,6 +225,7 @@ def ramp_fit_data( algorithm : str 'OLS' specifies that ordinary least squares should be used; 'GLS' specifies that generalized least squares should be used. + 'LIKELY' specifies that maximum likelihood should be used. weighting : str 'optimal' specifies that optimal weighting should be used; @@ -253,11 +257,29 @@ def ramp_fit_data( Object containing optional GLS-specific ramp fitting data for the exposure """ + # For the LIKELY algorithm, due to the jump detection portion of the code + # a minimum of a four group ramp is needed. + ngroups = ramp_data.data.shape[1] + if algorithm.upper() == "LIKELY" and ngroups < likely_fit.LIKELY_MIN_NGROUPS: + log.info(f"When selecting the LIKELY ramp fitting algorithm the" + " ngroups needs to be a minimum of {likely_fit.LIKELY_MIN_NGROUPS}," + " but ngroups = {ngroups}. Due to this, the ramp fitting algorithm" + " is being changed to OLS_C") + algorithm = "OLS_C" + + if algorithm.upper() == "OLS_C": + ramp_data.run_c_code = True + if algorithm.upper() == "GLS": image_info, integ_info, gls_opt_info = gls_fit.gls_ramp_fit( ramp_data, buffsize, save_opt, readnoise_2d, gain_2d, max_cores ) opt_info = None + elif algorithm.upper() == "LIKELY" and ngroups >= likely_fit.LIKELY_MIN_NGROUPS: + image_info, integ_info, opt_info = likely_fit.likely_ramp_fit( + ramp_data, readnoise_2d, gain_2d + ) + gls_opt_info = None else: # Default to OLS. # Get readnoise array for calculation of variance of noiseless ramps, and diff --git a/src/stcal/ramp_fitting/ramp_fit_class.py b/src/stcal/ramp_fitting/ramp_fit_class.py index 05243e3e..d48e5a86 100644 --- a/src/stcal/ramp_fitting/ramp_fit_class.py +++ b/src/stcal/ramp_fitting/ramp_fit_class.py @@ -16,6 +16,8 @@ def __init__(self): # Meta information self.instrument_name = None + self.read_pattern = None + self.rejection_threshold = None self.frame_time = None self.group_time = None diff --git a/tests/test_ramp_fitting_likely_fit.py b/tests/test_ramp_fitting_likely_fit.py new file mode 100644 index 00000000..09ac883b --- /dev/null +++ b/tests/test_ramp_fitting_likely_fit.py @@ -0,0 +1,688 @@ +import numpy as np +import pytest + +from stcal.ramp_fitting.ramp_fit import ramp_fit_class, ramp_fit_data +from stcal.ramp_fitting.ramp_fit_class import RampData +from stcal.ramp_fitting.likely_fit import likely_ramp_fit + + +test_dq_flags = { + "GOOD": 0, + "DO_NOT_USE": 1, + "SATURATED": 2, + "JUMP_DET": 4, + "NO_GAIN_VALUE": 8, + "UNRELIABLE_SLOPE": 16, +} + +GOOD = test_dq_flags["GOOD"] +DNU = test_dq_flags["DO_NOT_USE"] +JMP = test_dq_flags["JUMP_DET"] +SAT = test_dq_flags["SATURATED"] +NGV = test_dq_flags["NO_GAIN_VALUE"] +USLOPE = test_dq_flags["UNRELIABLE_SLOPE"] + +DELIM = "-" * 70 + + +def create_blank_ramp_data(dims, var, tm): + """ + Create empty RampData classes, as well as gain and read noise arrays, + based on dimensional, variance, and timing input. + """ + nints, ngroups, nrows, ncols = dims + rnval, gval = var + frame_time, nframes, groupgap = tm + group_time = (nframes + groupgap) * frame_time + + data = np.zeros(shape=(nints, ngroups, nrows, ncols), dtype=np.float32) + err = np.ones(shape=(nints, ngroups, nrows, ncols), dtype=np.float32) + pixdq = np.zeros(shape=(nrows, ncols), dtype=np.uint32) + gdq = np.zeros(shape=(nints, ngroups, nrows, ncols), dtype=np.uint8) + dark_current = np.zeros(shape=(nrows, ncols), dtype = np.float32) + + ramp_data = RampData() + ramp_data.set_arrays(data=data, err=err, groupdq=gdq, pixeldq=pixdq, average_dark_current=dark_current) + ramp_data.set_meta( + name="NIRSpec", + frame_time=frame_time, + group_time=group_time, + groupgap=groupgap, + nframes=nframes, + drop_frames1=None, + ) + ramp_data.set_dqflags(test_dq_flags) + + gain = np.ones(shape=(nrows, ncols), dtype=np.float32) * gval + rnoise = np.ones(shape=(nrows, ncols), dtype=np.float32) * rnval + + return ramp_data, gain, rnoise + + +# ----------------------------------------------------------------------------- + + +def test_basic_ramp(): + """ + Test a basic ramp with a linear progression up the ramp. Compare the + integration results from the LIKELY algorithm to the OLS algorithm. + """ + nints, ngroups, nrows, ncols = 1, 10, 1, 1 + rnval, gval = 10.0, 5.0 + frame_time, nframes, groupgap = 10.736, 4, 1 + + dims = nints, ngroups, nrows, ncols + var = rnval, gval + tm = frame_time, nframes, groupgap + + ramp_data, gain2d, rnoise2d = create_blank_ramp_data(dims, var, tm) + + # Create a simple linear ramp. + ramp = np.array(list(range(ngroups))) * 20 + 10 + ramp_data.data[0, :, 0, 0] = ramp + + save_opt, algo, ncores = False, "LIKELY", "none" + slopes, cube, ols_opt, gls_opt = ramp_fit_data( + ramp_data, 512, save_opt, rnoise2d, gain2d, algo, "optimal", ncores, test_dq_flags + ) + + data = cube[0][0, 0, 0] + ddiff = (ramp_data.data[0, ngroups-1, 0, 0] - ramp_data.data[0, 0, 0, 0]) + check = ddiff / float(ngroups-1) + check = check / ramp_data.group_time + tol = 1.e-5 + diff = abs(data - check) + assert diff < tol + slope = slopes[0][0, 0] + assert abs(slope - data) < tol + + + # Check against OLS. + ramp_data1, gain2d1, rnoise2d1 = create_blank_ramp_data(dims, var, tm) + + ramp = np.array(list(range(ngroups))) * 20 + 10 + ramp_data1.data[0, :, 0, 0] = ramp + + save_opt, algo, ncores = False, "OLS", "none" + slopes1, cube1, ols_opt1, gls_opt1 = ramp_fit_data( + ramp_data1, 512, save_opt, rnoise2d1, gain2d1, algo, "optimal", ncores, test_dq_flags + ) + + data1 = cube1[0][0, 0, 0] + diff = abs(data - data1) + assert diff < tol + + +def test_basic_ramp_multi_pixel(): + """ + Test a basic ramp with a linear progression up the ramp. Compare the + integration results from the LIKELY algorithm to the OLS algorithm. + """ + nints, ngroups, nrows, ncols = 1, 10, 2, 2 + rnval, gval = 10.0, 5.0 + frame_time, nframes, groupgap = 10.736, 4, 1 + + dims = nints, ngroups, nrows, ncols + var = rnval, gval + tm = frame_time, nframes, groupgap + + ramp_data, gain2d, rnoise2d = create_blank_ramp_data(dims, var, tm) + + # Create a simple linear ramp. + ramp = np.array(list(range(ngroups))) * 20 + 10 + ramp_data.data[0, :, 0, 0] = ramp + ramp_data.data[0, :, 0, 1] = ramp + ramp_data.data[0, :, 1, 0] = ramp + ramp_data.data[0, :, 1, 1] = ramp + + save_opt, algo, ncores = False, "LIKELY", "none" + slopes, cube, ols_opt, gls_opt = ramp_fit_data( + ramp_data, 512, save_opt, rnoise2d, gain2d, algo, "optimal", ncores, test_dq_flags + ) + + ramp_data1, gain2d1, rnoise2d1 = create_blank_ramp_data(dims, var, tm) + + # Create a simple linear ramp. + ramp = np.array(list(range(ngroups))) * 20 + 10 + ramp_data1.data[0, :, 0, 0] = ramp + ramp_data1.data[0, :, 0, 1] = ramp + ramp_data1.data[0, :, 1, 0] = ramp + ramp_data1.data[0, :, 1, 1] = ramp + + save_opt, algo, ncores = False, "OLS", "none" + slopes1, cube1, ols_opt1, gls_opt1 = ramp_fit_data( + ramp_data1, 512, save_opt, rnoise2d1, gain2d1, algo, "optimal", ncores, test_dq_flags + ) + + data, dq, vp, vr, err = slopes + data1, dq1, vp1, vr1, err1 = slopes1 + + tol = 1.e-4 + np.testing.assert_allclose(data, data1, tol) + + +def test_basic_ramp_2integ(): + """ + Test a basic ramp with a linear progression up the ramp. Compare the + integration results from the LIKELY algorithm to the OLS algorithm. + """ + nints, ngroups, nrows, ncols = 2, 10, 1, 1 + rnval, gval = 10.0, 5.0 + frame_time, nframes, groupgap = 10.736, 4, 1 + + dims = nints, ngroups, nrows, ncols + var = rnval, gval + tm = frame_time, nframes, groupgap + + ramp_data, gain2d, rnoise2d = create_blank_ramp_data(dims, var, tm) + + # Create a simple linear ramp. + ramp = np.array(list(range(ngroups))) * 20 + 10 + ramp_data.data[0, :, 0, 0] = ramp + ramp_data.data[1, :, 0, 0] = ramp + + save_opt, algo, ncores = False, "LIKELY", "none" + slopes, cube, ols_opt, gls_opt = ramp_fit_data( + ramp_data, 512, save_opt, rnoise2d, gain2d, algo, "optimal", ncores, test_dq_flags + ) + + # Check against OLS. + ramp_data1, gain2d1, rnoise2d1 = create_blank_ramp_data(dims, var, tm) + + ramp = np.array(list(range(ngroups))) * 20 + 10 + ramp_data1.data[0, :, 0, 0] = ramp + ramp_data1.data[1, :, 0, 0] = ramp + + save_opt, algo, ncores = False, "OLS", "none" + slopes1, cube1, ols_opt1, gls_opt1 = ramp_fit_data( + ramp_data1, 512, save_opt, rnoise2d1, gain2d1, algo, "optimal", ncores, test_dq_flags + ) + + tol = 1.e-5 + data = cube[0][0, 0, 0] + data1 = cube1[0][0, 0, 0] + diff = abs(data - data1) + assert diff < tol + + dbg_print_slope_slope1(slopes, slopes1, (0, 0)) + + +def flagged_ramp_data(): + nints, ngroups, nrows, ncols = 1, 20, 1, 1 + rnval, gval = 10.0, 5.0 + frame_time, nframes, groupgap = 10.736, 4, 1 + + dims = nints, ngroups, nrows, ncols + var = rnval, gval + tm = frame_time, nframes, groupgap + + ramp_data, gain2d, rnoise2d = create_blank_ramp_data(dims, var, tm) + + ramp = np.array(list(range(ngroups))) * 20 + 10 + ramp_data.data[0, :, 0, 0] = ramp + ramp_data.data[0, 10:, 0, 0] += 150. # Add a jump. + + # Create segments in the ramp, including a jump and saturation at the end. + dq = np.array([GOOD] * ngroups) + dq[2] = DNU + dq[17:] = SAT + dq[10] = JMP + ramp_data.groupdq[0, :, 0, 0] = dq + + return ramp_data, gain2d, rnoise2d + + +def test_flagged_ramp(): + """ + Test flagged ramp. The flags will cause segments, as well as ramp + truncation. Compare the integration results from the LIKELY algorithm + to the OLS algorithm. + """ + ramp_data, gain2d, rnoise2d = flagged_ramp_data() + + save_opt, algo, ncores = False, "LIKELY", "none" + slopes, cube, ols_opt, gls_opt = ramp_fit_data( + ramp_data, 512, save_opt, rnoise2d, gain2d, algo, "optimal", ncores, test_dq_flags + ) + + data = cube[0][0, 0, 0] + dq = cube[1][0, 0, 0] + + # Check against OLS. + ramp_data, gain2d, rnoise2d = flagged_ramp_data() + + save_opt, algo, ncores = False, "OLS", "none" + slopes1, cube1, ols_opt, gls_opt = ramp_fit_data( + ramp_data, 512, save_opt, rnoise2d, gain2d, algo, "optimal", ncores, test_dq_flags + ) + + data_ols = cube1[0][0, 0, 0] + dq_ols = cube1[1][0, 0, 0] + + tol = 1.e-5 + diff = abs(data - data_ols) + assert diff < tol + assert dq == dq_ols + + dbg_print_slope_slope1(slopes, slopes1, (0, 0)) + + +def test_random_ramp(): + """ + Created a slope with a base slope of 150., with random Poisson + noise with lambda 5.0. At group 4 is a jump of 1100.0. + """ + nints, ngroups, nrows, ncols = 1, 10, 1, 1 + rnval, gval = 10.0, 5.0 + frame_time, nframes, groupgap = 10.736, 5, 2 + + dims = nints, ngroups, nrows, ncols + var = rnval, gval + tm = frame_time, nframes, groupgap + + ramp_data, gain2d, rnoise2d = create_blank_ramp_data(dims, var, tm) + + # A randomly generated ramp by setting up a ramp that has a slope of 150. + # with some randomly added Poisson values, with lambda=5., and a jump + # at group 4. + ramp = np.array([153., 307., 457., 604., 1853., 2002., 2159., 2308., 2459., 2601.]) + ramp_data.data[0, :, 0, 0] = ramp + + # Create a jump, but don't mark it to make sure it gets detected. + dq = np.array([GOOD] * ngroups) + ramp_data.groupdq[0, :, 0, 0] = dq + + save_opt, algo, ncores = False, "LIKELY", "none" + slopes, cube, ols_opt, gls_opt = ramp_fit_data( + ramp_data, 512, save_opt, rnoise2d, gain2d, algo, "optimal", ncores, test_dq_flags + ) + + data, dq, vp, vr, err = slopes + tol = 1.e-4 + + assert abs(data[0, 0] - 1.9960526) < tol + assert dq[0, 0] == JMP + assert abs(vp[0, 0] - 0.00064461) < tol + assert abs(vr[0, 0] - 0.00018037) < tol + + +def test_long_ramp(): + """ + Test a long ramp with hundreds of groups. + """ + nints, ngroups, nrows, ncols = 1, 200, 1, 1 + rnval, gval = 10.0, 5.0 + frame_time, nframes, groupgap = 10.736, 4, 1 + + dims = nints, ngroups, nrows, ncols + var = rnval, gval + tm = frame_time, nframes, groupgap + + ramp_data, gain2d, rnoise2d = create_blank_ramp_data(dims, var, tm) + + # Create a simple linear ramp. + ramp = np.array(list(range(ngroups))) * 20 + 10 + ramp_data.data[0, :, 0, 0] = ramp + + save_opt, algo, ncores = False, "LIKELY", "none" + slopes, cube, ols_opt, gls_opt = ramp_fit_data( + ramp_data, 512, save_opt, rnoise2d, gain2d, algo, "optimal", ncores, test_dq_flags + ) + + data = cube[0][0, 0, 0] + ddiff = (ramp_data.data[0, ngroups-1, 0, 0] - ramp_data.data[0, 0, 0, 0]) + check = ddiff / float(ngroups-1) + check = check / ramp_data.group_time + tol = 1.e-5 + diff = abs(data - check) + assert diff < tol + + # Check against OLS. + ramp_data1, gain2d1, rnoise2d1 = create_blank_ramp_data(dims, var, tm) + + ramp = np.array(list(range(ngroups))) * 20 + 10 + ramp_data1.data[0, :, 0, 0] = ramp + + save_opt, algo, ncores = False, "OLS", "none" + slopes1, cube1, ols_opt1, gls_opt1 = ramp_fit_data( + ramp_data1, 512, save_opt, rnoise2d1, gain2d1, algo, "optimal", ncores, test_dq_flags + ) + + data1 = cube1[0][0, 0, 0] + diff = abs(data - data1) + assert diff < tol + dbg_print_slope_slope1(slopes, slopes1, (0, 0)) + + +@pytest.mark.parametrize("ngroups", [1, 2]) +def test_too_few_group_ramp(ngroups): + """ + It's supposed to fail. The likelihood algorithm needs at least two + groups to work. + """ + nints, nrows, ncols = 1, 1, 1 + rnval, gval = 10.0, 5.0 + frame_time, nframes, groupgap = 10.736, 1, 0 + + dims = nints, ngroups, nrows, ncols + var = rnval, gval + tm = frame_time, nframes, groupgap + + ramp_data, gain2d, rnoise2d = create_blank_ramp_data(dims, var, tm) + + # Create a simple linear ramp. + ramp = np.array(list(range(ngroups))) * 20 + 10 + ramp_data.data[0, :, 0, 0] = ramp + + with pytest.raises(ValueError): + image_info, integ_info, opt_info = likely_ramp_fit( + ramp_data, rnoise2d, gain2d + ) + + +@pytest.mark.parametrize("nframes", [1, 2, 4, 8]) +def test_short_group_ramp(nframes): + """ + Test short ramps with various nframes. + """ + nints, ngroups, nrows, ncols = 1, 4, 1, 1 + rnval, gval = 10.0, 5.0 + # frame_time, nframes, groupgap = 10.736, 4, 1 + frame_time, groupgap = 10.736, 1 + + dims = nints, ngroups, nrows, ncols + var = rnval, gval + tm = frame_time, nframes, groupgap + + ramp_data, gain2d, rnoise2d = create_blank_ramp_data(dims, var, tm) + + # Create a simple linear ramp. + ramp = np.array(list(range(ngroups))) * 20 + 10 + ramp_data.data[0, :, 0, 0] = ramp + + save_opt, algo, ncores = False, "LIKELY", "none" + slopes, cube, ols_opt, gls_opt = ramp_fit_data( + ramp_data, 512, save_opt, rnoise2d, gain2d, algo, "optimal", ncores, test_dq_flags + ) + + data = cube[0][0, 0, 0] + ddiff = (ramp_data.data[0, ngroups-1, 0, 0] - ramp_data.data[0, 0, 0, 0]) + check = ddiff / float(ngroups-1) + check = check / ramp_data.group_time + tol = 1.e-5 + diff = abs(data - check) + assert diff < tol + + # Check against OLS. + ramp_data1, gain2d1, rnoise2d1 = create_blank_ramp_data(dims, var, tm) + + ramp = np.array(list(range(ngroups))) * 20 + 10 + ramp_data1.data[0, :, 0, 0] = ramp + + save_opt, algo, ncores = False, "OLS", "none" + slopes1, cube1, ols_opt1, gls_opt1 = ramp_fit_data( + ramp_data1, 512, save_opt, rnoise2d1, gain2d1, algo, "optimal", ncores, test_dq_flags + ) + + data1 = cube1[0][0, 0, 0] + diff = abs(data - data1) + assert diff < tol + # dbg_print_slope_slope1(slopes, slopes1, (0, 0)) + + +def data_small_good_groups(): + nints, ngroups, nrows, ncols = 1, 10, 1, 1 + rnval, gval = 10.0, 5.0 + frame_time, nframes, groupgap = 10.736, 4, 1 + + dims = nints, ngroups, nrows, ncols + var = rnval, gval + tm = frame_time, nframes, groupgap + + ramp_data, gain2d, rnoise2d = create_blank_ramp_data(dims, var, tm) + + # Create a simple linear ramp. + ramp = np.array(list(range(ngroups))) * 20 + 10 + ramp_data.data[0, :, 0, 0] = ramp + + dq = np.array([SAT] * ngroups, dtype=np.uint8) + ramp_data.groupdq[0, :, 0, 0] = dq + + return ramp_data, gain2d, rnoise2d + + +# XXX One good group in a ramp may not be any good +# @pytest.mark.parametrize("ngood", [1, 2]) +@pytest.mark.parametrize("ngood", [2]) +def test_small_good_groups(ngood): + """ + Test ramps with only one or two good groups. + """ + ramp_data, gain2d, rnoise2d = data_small_good_groups() + ramp_data.groupdq[0, :ngood, 0, 0] = GOOD + + save_opt, algo, ncores = False, "LIKELY", "none" + slopes, cube, ols_opt, gls_opt = ramp_fit_data( + ramp_data, 512, save_opt, rnoise2d, gain2d, algo, "optimal", ncores, test_dq_flags + ) + + lik_slope = slopes[0][0, 0] + + + # Check against OLS. + ramp_data1, gain2d1, rnoise2d1 = data_small_good_groups() + ramp_data1.groupdq[0, :ngood, 0, 0] = GOOD + + save_opt, algo, ncores = False, "OLS", "none" + slopes1, cube1, ols_opt1, gls_opt1 = ramp_fit_data( + ramp_data1, 512, save_opt, rnoise2d1, gain2d1, algo, "optimal", ncores, test_dq_flags + ) + ols_slope = slopes1[0][0, 0] + + tol = 1.e-4 + diff = abs(ols_slope - lik_slope) + if ngood==2: + assert diff < tol + else: + print(f"ols_slope = {ols_slope}") + print(f"lik_slope = {lik_slope}") + print(f"diff = {diff}") + + +def test_jump_detect(): + """ + Create a simple ramp with a (2, 2) image that has a jump in two + different ramps and the computed slopes are still close. + """ + nints, ngroups, nrows, ncols = 1, 10, 2, 2 + rnval, gval = 10.0, 5.0 + frame_time, nframes, groupgap = 10.736, 5, 2 + # frame_time, nframes, groupgap = 1., 1, 0 + + dims = nints, ngroups, nrows, ncols + var = rnval, gval + tm = frame_time, nframes, groupgap + + ramp_data, gain2d, rnoise2d = create_blank_ramp_data(dims, var, tm) + + # Create a ramp with a jump to see if it gets detected. + base, cr, jump_loc = 15., 1000., 6 + ramp = np.array([(k+1) * base for k in range(ngroups)]) + ramp_data.data[0, :, 0, 1] = ramp + if nrows > 1: + ramp_data.data[0, :, 1, 0] = ramp + ramp[jump_loc:] += cr + ramp_data.data[0, :, 0, 0] = ramp + ramp[jump_loc-1] += cr + if nrows > 1: + ramp_data.data[0, :, 1, 1] = ramp + + save_opt, algo, ncores = False, "LIKELY", "none" + slopes, cube, ols_opt, gls_opt = ramp_fit_data( + ramp_data, 512, save_opt, rnoise2d, gain2d, algo, "optimal", ncores, test_dq_flags + ) + + data, dq, vp, vr, err = slopes + slope_est = base / ramp_data.group_time + + tol = 1.e-4 + assert abs(data[0, 0] - slope_est) < tol + assert abs(data[0, 1] - slope_est) < tol + assert abs(data[1, 0] - slope_est) < tol + assert abs(data[1, 1] - slope_est) < tol + assert dq[0, 0] == JMP + assert dq[0, 1] == GOOD + assert dq[1, 0] == GOOD + assert dq[1, 1] == JMP + + +def test_too_few_groups(caplog): + """ + Ensure + """ + nints, ngroups, nrows, ncols = 1, 3, 1, 1 + rnval, gval = 10.0, 5.0 + frame_time, nframes, groupgap = 10.736, 4, 1 + + dims = nints, ngroups, nrows, ncols + var = rnval, gval + tm = frame_time, nframes, groupgap + + ramp_data, gain2d, rnoise2d = create_blank_ramp_data(dims, var, tm) + + # Create a simple linear ramp. + ramp = np.array(list(range(ngroups))) * 20 + 10 + ramp_data.data[0, :, 0, 0] = ramp + + save_opt, algo, ncores = False, "LIKELY", "none" + slopes, cube, ols_opt, gls_opt = ramp_fit_data( + ramp_data, 512, save_opt, rnoise2d, gain2d, algo, "optimal", ncores, test_dq_flags + ) + + expected_log = "ramp fitting algorithm is being changed to OLS_C" + assert expected_log in caplog.text + + +# ----------------------------------------------------------------- +# DEBUG +# ----------------------------------------------------------------- +def dbg_print_basic_ramp(ramp_data, pix=(0, 0)): + row, col = pix + nints = ramp_data.data.shape[0] + data = ramp_data.data[:, :, row, col] + dq = ramp_data.groupdq[:, :, row, col] + + print(" ") + print(DELIM) + print(f"Data Shape: {ramp_data.data.shape}") + print(DELIM) + print("Data:") + for integ in range(nints): + arr_str = np.array2string(data[integ, :], max_line_width=np.nan, separator=", ") + print(f"[{integ}] {arr_str}") + print(DELIM) + + print("DQ:") + for integ in range(nints): + arr_str = np.array2string(dq[integ, :], max_line_width=np.nan, separator=", ") + print(f"[{integ}] {arr_str}") + print(DELIM) + + +def dbg_print_slopes(slope, pix=(0, 0), label=None): + data, dq, vp, vr, err = slope + row, col = pix + + print(" ") + print(DELIM) + if label is not None: + print("Slope Information: ({label})") + else: + print("Slope Information:") + print(f" Pixel = ({row}, {col})") + + print(f"data = {data[row, col]}") + print(f"dq = {dq[row, col]}") + print(f"vp = {vp[row, col]}") + print(f"vr = {vr[row, col]}\n") + + print(DELIM) + + +def dbg_print_cube(cube, pix=(0, 0), label=None): + data, dq, vp, vr, err = cube + data1, dq1, vp1, vr1, err1 = cube1 + row, col = pix + nints = data1.shape[0] + + print(" ") + print(DELIM) + if label is not None: + print("Cube Information: ({label})") + else: + print("Cube Information:") + print(f" Pixel = ({row}, {col})") + print(f" Number of Integrations = {nints}") + + print(f"data = {data[:, row, col]}") + print(f"dq = {dq[:, row, col]}") + print(f"vp = {vp[:, row, col]}") + print(f"vr = {vr[:, row, col]}") + + print(DELIM) + + +def dbg_print_slope_slope1(slopes, slopes1, pix): + data, dq, vp, vr, err = slopes + data1, dq1, vp1, vr1, err1 = slopes1 + row, col = pix + + print(" ") + print(DELIM) + print("Slope Information:") + print(f" Pixel = ({row}, {col})") + + print(f"data LIK = {data[row, col]:.12f}") + print(f"data OLS = {data1[row, col]:.12f}\n") + + # print(f"dq LIK = {dq[row, col]}") + # print(f"dq OLS = {dq1[row, col]}\n") + + print(f"vp LIK = {vp[row, col]:.12f}") + print(f"vp OLS = {vp1[row, col]:.12f}\n") + + print(f"vr LIK = {vr[row, col]:.12f}") + print(f"vr OLS = {vr1[row, col]:.12f}\n") + + print(DELIM) + + +def dbg_print_cube_cube1(cube, cube1, pix): + data, dq, vp, vr, err = cube + data1, dq1, vp1, vr1, err1 = cube1 + row, col = pix + nints = data1.shape[0] + + print(" ") + print(DELIM) + print("Cube Information:") + print(f" Pixel = ({row}, {col})") + print(f" Number of Integrations = {nints}") + + print(f"data LIK = {data[:, row, col]}") + print(f"data OLS = {data1[:, row, col]}\n") + + # print(f"dq LIK = {dq[:, row, col]}") + # print(f"dq OLS = {dq1[:, row, col]}\n") + + print(f"vp LIK = {vp[:, row, col]}") + print(f"vp OLS = {vp1[:, row, col]}\n") + + print(f"vr LIK = {vr[:, row, col]}") + print(f"vr OLS = {vr1[:, row, col]}\n") + + print(DELIM) + + +def array_string(arr, prec=4): + return np.array2string(arr, precision=prec, max_line_width=np.nan, separator=", ")