From 17e96c877ceee13426d4374a8cac964ace41cdeb Mon Sep 17 00:00:00 2001
From: William Jamieson <wjamieson@stsci.edu>
Date: Wed, 1 Nov 2023 11:11:45 -0400
Subject: [PATCH 1/4] Move ramp initialization to a local operation This
 significantly reduces memory usage ~150Gb -> 1Gb

---
 src/stcal/ramp_fitting/ols_cas22/_core.pxd    |   7 +-
 src/stcal/ramp_fitting/ols_cas22/_core.pyx    | 188 ++++++++++--------
 .../ramp_fitting/ols_cas22/_fit_ramps.pyx     |  16 +-
 src/stcal/ramp_fitting/ols_cas22/_pixel.pyx   |   8 +-
 src/stcal/ramp_fitting/ols_cas22_fit.py       |   3 +-
 5 files changed, 121 insertions(+), 101 deletions(-)

diff --git a/src/stcal/ramp_fitting/ols_cas22/_core.pxd b/src/stcal/ramp_fitting/ols_cas22/_core.pxd
index f7fe7877..63ee6323 100644
--- a/src/stcal/ramp_fitting/ols_cas22/_core.pxd
+++ b/src/stcal/ramp_fitting/ols_cas22/_core.pxd
@@ -1,6 +1,5 @@
 from libcpp.vector cimport vector
 from libcpp.stack cimport stack
-from libcpp.deque cimport deque
 
 
 cdef struct RampIndex:
@@ -15,8 +14,8 @@ cdef struct RampFit:
 
 
 cdef struct RampFits:
-    vector[RampFit] fits
-    vector[RampIndex] index
+    # vector[RampFit] fits
+    # vector[RampIndex] index
     vector[int] jumps
     RampFit average
 
@@ -54,5 +53,5 @@ cpdef enum RampJumpDQ:
 
 cpdef float threshold(Thresh thresh, float slope)
 cdef float get_power(float s)
-cdef deque[stack[RampIndex]] init_ramps(int[:, :] dq)
+cdef stack[RampIndex] init_ramps(int[:, :] dq, int n_resultants, int index_pixel)
 cpdef ReadPatternMetadata metadata_from_read_pattern(list[list[int]] read_pattern, float read_time)
diff --git a/src/stcal/ramp_fitting/ols_cas22/_core.pyx b/src/stcal/ramp_fitting/ols_cas22/_core.pyx
index 5c10cb4b..5726fea1 100644
--- a/src/stcal/ramp_fitting/ols_cas22/_core.pyx
+++ b/src/stcal/ramp_fitting/ols_cas22/_core.pyx
@@ -134,7 +134,7 @@ cpdef inline float threshold(Thresh thresh, float slope):
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-cdef inline deque[stack[RampIndex]] init_ramps(int[:, :] dq):
+cdef inline stack[RampIndex] init_ramps(int[:, :] dq, int n_resultants, int index_pixel):
     """
     Create the initial ramp stack for each pixel
         if dq[index_resultant, index_pixel] == 0, then the resultant is in a ramp
@@ -144,94 +144,118 @@ cdef inline deque[stack[RampIndex]] init_ramps(int[:, :] dq):
     ----------
     dq : int[n_resultants, n_pixel]
         DQ array
+    n_resultants : int
+        Number of resultants
+    index_pixel : int
+        The index of the pixel to find ramps for
 
     Returns
     -------
-    deque of stacks of RampIndex objects
-        - deque with entry for each pixel
-            Chosen to be deque because need element access to loop
+    stack of RampIndex objects
         - stack with entry for each ramp found (top of stack is last ramp found)
         - RampIndex with start and end indices of the ramp in the resultants
     """
-    cdef int n_pixel, n_resultants
-
-    n_resultants, n_pixel = np.array(dq).shape
-    cdef deque[stack[RampIndex]] pixel_ramps
-
-    cdef int index_resultant, index_pixel
-    cdef stack[RampIndex] ramps
-    cdef RampIndex ramp
-
-    for index_pixel in range(n_pixel):
-        ramps = stack[RampIndex]()
-
-        # Note: if start/end are -1, then no value has been assigned
-        # ramp.start == -1 means we have not started a ramp
-        # dq[index_resultant, index_pixel] == 0 means resultant is in ramp
-        ramp = RampIndex(-1, -1)
-        for index_resultant in range(n_resultants):
-            if ramp.start == -1:
-                # Looking for the start of a ramp
-                if dq[index_resultant, index_pixel] == 0:
-                    # We have found the start of a ramp!
-                    ramp.start = index_resultant
-                else:
-                    # This is not the start of the ramp yet
-                    continue
+    cdef stack[RampIndex] ramps = stack[RampIndex]()
+
+    # Note: if start/end are -1, then no value has been assigned
+    # ramp.start == -1 means we have not started a ramp
+    # dq[index_resultant, index_pixel] == 0 means resultant is in ramp
+    cdef RampIndex ramp = RampIndex(-1, -1)
+    for index_resultant in range(n_resultants):
+        if ramp.start == -1:
+            # Looking for the start of a ramp
+            if dq[index_resultant, index_pixel] == 0:
+                # We have found the start of a ramp!
+                ramp.start = index_resultant
             else:
-                # Looking for the end of a ramp
-                if dq[index_resultant, index_pixel] == 0:
-                    # This pixel is in the ramp do nothing
-                    continue
-                else:
-                    # This pixel is not in the ramp
-                    # => index_resultant - 1 is the end of the ramp
-                    ramp.end = index_resultant - 1
-
-                    # Add completed ramp to stack and reset ramp
-                    ramps.push(ramp)
-                    ramp = RampIndex(-1, -1)
-
-        # Handle case where last resultant is in ramp (so no end has been set)
-        if ramp.start != -1 and ramp.end == -1:
-            # Last resultant is end of the ramp => set then add to stack
-            ramp.end = n_resultants - 1
-            ramps.push(ramp)
-
-        # Add ramp stack for pixel to list
-        pixel_ramps.push_back(ramps)
-
-    return pixel_ramps
-
-
-def _init_ramps_list(np.ndarray[int, ndim=2] dq):
-    """
-    This is a wrapper for init_ramps so that it can be fully inspected from pure
-    python. A cpdef cannot be used in that case becase a stack has no direct python
-    analog. Instead this function turns that stack into a list ordered in the same
-    order as the stack; meaning that, the first element of the list is the top of
-    the stack.
-        Note this function is for testing purposes only and so is marked as private
-        within this private module
-    """
-    cdef deque[stack[RampIndex]] raw = init_ramps(dq)
-
-    # Have to turn deque and stack into python compatible objects
-    cdef RampIndex index
-    cdef stack[RampIndex] ramp
-    cdef list out = []
-    cdef list stack_out
-    for ramp in raw:
-        stack_out = []
-        while not ramp.empty():
-            index = ramp.top()
-            ramp.pop()
-            # So top of stack is first item of list
-            stack_out = [index] + stack_out
-
-        out.append(stack_out)
-
-    return out
+                # This is not the start of the ramp yet
+                continue
+        else:
+            # Looking for the end of a ramp
+            if dq[index_resultant, index_pixel] == 0:
+                # This pixel is in the ramp do nothing
+                continue
+            else:
+                # This pixel is not in the ramp
+                # => index_resultant - 1 is the end of the ramp
+                ramp.end = index_resultant - 1
+
+                # Add completed ramp to stack and reset ramp
+                ramps.push(ramp)
+                ramp = RampIndex(-1, -1)
+
+    # Handle case where last resultant is in ramp (so no end has been set)
+    if ramp.start != -1 and ramp.end == -1:
+        # Last resultant is end of the ramp => set then add to stack
+        ramp.end = n_resultants - 1
+        ramps.push(ramp)
+
+    return ramps
+
+
+# @cython.boundscheck(False)
+# @cython.wraparound(False)
+# cdef inline deque[stack[RampIndex]] init_ramps(int[:, :] dq):
+#     """
+#     Create the initial ramp stack for each pixel
+#         if dq[index_resultant, index_pixel] == 0, then the resultant is in a ramp
+#         otherwise, the resultant is not in a ramp
+
+#     Parameters
+#     ----------
+#     dq : int[n_resultants, n_pixel]
+#         DQ array
+
+#     Returns
+#     -------
+#     deque of stacks of RampIndex objects
+#         - deque with entry for each pixel
+#             Chosen to be deque because need element access to loop
+#         - stack with entry for each ramp found (top of stack is last ramp found)
+#         - RampIndex with start and end indices of the ramp in the resultants
+#     """
+#     cdef int n_pixel, n_resultants
+
+#     n_resultants, n_pixel = np.array(dq).shape
+#     cdef deque[stack[RampIndex]] pixel_ramps
+
+#     cdef int index_pixel
+
+#     for index_pixel in range(n_pixel):
+#         # Add ramp stack for pixel to list
+#         pixel_ramps.push_back(_init_ramps_pixel(dq, n_resultants, index_pixel))
+
+#     return pixel_ramps
+
+
+# def _init_ramps_list(np.ndarray[int, ndim=2] dq):
+#     """
+#     This is a wrapper for init_ramps so that it can be fully inspected from pure
+#     python. A cpdef cannot be used in that case becase a stack has no direct python
+#     analog. Instead this function turns that stack into a list ordered in the same
+#     order as the stack; meaning that, the first element of the list is the top of
+#     the stack.
+#         Note this function is for testing purposes only and so is marked as private
+#         within this private module
+#     """
+#     cdef deque[stack[RampIndex]] raw = init_ramps(dq)
+
+#     # Have to turn deque and stack into python compatible objects
+#     cdef RampIndex index
+#     cdef stack[RampIndex] ramp
+#     cdef list out = []
+#     cdef list stack_out
+#     for ramp in raw:
+#         stack_out = []
+#         while not ramp.empty():
+#             index = ramp.top()
+#             ramp.pop()
+#             # So top of stack is first item of list
+#             stack_out = [index] + stack_out
+
+#         out.append(stack_out)
+
+#     return out
 
 
 @cython.boundscheck(False)
diff --git a/src/stcal/ramp_fitting/ols_cas22/_fit_ramps.pyx b/src/stcal/ramp_fitting/ols_cas22/_fit_ramps.pyx
index ad19fb28..78c77672 100644
--- a/src/stcal/ramp_fitting/ols_cas22/_fit_ramps.pyx
+++ b/src/stcal/ramp_fitting/ols_cas22/_fit_ramps.pyx
@@ -1,9 +1,7 @@
 import numpy as np
 cimport numpy as np
 from libcpp cimport bool
-from libcpp.stack cimport stack
 from libcpp.list cimport list as cpp_list
-from libcpp.deque cimport deque
 cimport cython
 
 from stcal.ramp_fitting.ols_cas22._core cimport (RampFits, RampIndex, Thresh,
@@ -42,7 +40,7 @@ class RampFitOutputs(NamedTuple):
             the dq array, with additional flags set for jumps detected by the
             jump detection algorithm.
     """
-    fits: list
+    # fits: list
     parameters: np.ndarray
     variances: np.ndarray
     dq: np.ndarray
@@ -102,13 +100,10 @@ def fit_ramps(np.ndarray[float, ndim=2] resultants,
                                                         Thresh(intercept, constant),
                                                         use_jump)
 
-    # Compute all the initial sets of ramps
-    cdef deque[stack[RampIndex]] pixel_ramps = init_ramps(dq)
-
     # Use list because this might grow very large which would require constant
     #    reallocation. We don't need random access, and this gets cast to a python
     #    list in the end.
-    cdef cpp_list[RampFits] ramp_fits
+    # cdef cpp_list[RampFits] ramp_fits
 
     cdef np.ndarray[float, ndim=2] parameters = np.zeros((n_pixels, 2), dtype=np.float32)
     cdef np.ndarray[float, ndim=2] variances = np.zeros((n_pixels, 3), dtype=np.float32)
@@ -119,7 +114,7 @@ def fit_ramps(np.ndarray[float, ndim=2] resultants,
     for index in range(n_pixels):
         # Fit all the ramps for the given pixel
         fit = make_pixel(fixed, read_noise[index],
-                         resultants[:, index]).fit_ramps(pixel_ramps[index])
+                         resultants[:, index]).fit_ramps(init_ramps(dq, n_resultants, index))
 
         parameters[index, Parameter.slope] = fit.average.slope
 
@@ -130,6 +125,7 @@ def fit_ramps(np.ndarray[float, ndim=2] resultants,
         for jump in fit.jumps:
             dq[jump, index] = RampJumpDQ.JUMP_DET
 
-        ramp_fits.push_back(fit)
+        # ramp_fits.push_back(fit)
 
-    return RampFitOutputs(ramp_fits, parameters, variances, dq)
+    # return RampFitOutputs(ramp_fits, parameters, variances, dq)
+    return RampFitOutputs(parameters, variances, dq)
diff --git a/src/stcal/ramp_fitting/ols_cas22/_pixel.pyx b/src/stcal/ramp_fitting/ols_cas22/_pixel.pyx
index 88544243..5ddfe256 100644
--- a/src/stcal/ramp_fitting/ols_cas22/_pixel.pyx
+++ b/src/stcal/ramp_fitting/ols_cas22/_pixel.pyx
@@ -446,8 +446,8 @@ cdef class Pixel:
             #    than threshold
             # Note that ramps are computed backward in time meaning we need to
             #  reverse the order of the fits at the end
-            ramp_fits.fits.push_back(ramp_fit)
-            ramp_fits.index.push_back(ramp)
+            # ramp_fits.fits.push_back(ramp_fit)
+            # ramp_fits.index.push_back(ramp)
 
             # Start computing the averages
             #    Note we do not do anything in the NaN case for degenerate ramps
@@ -462,8 +462,8 @@ cdef class Pixel:
                 ramp_fits.average.poisson_var += weight**2 * ramp_fit.poisson_var
 
         # Reverse to order in time
-        ramp_fits.fits = ramp_fits.fits[::-1]
-        ramp_fits.index = ramp_fits.index[::-1]
+        # ramp_fits.fits = ramp_fits.fits[::-1]
+        # ramp_fits.index = ramp_fits.index[::-1]
 
         # Finish computing averages
         ramp_fits.average.slope /= total_weight if total_weight != 0 else 1
diff --git a/src/stcal/ramp_fitting/ols_cas22_fit.py b/src/stcal/ramp_fitting/ols_cas22_fit.py
index 9584970e..8a8b6a3e 100644
--- a/src/stcal/ramp_fitting/ols_cas22_fit.py
+++ b/src/stcal/ramp_fitting/ols_cas22_fit.py
@@ -131,4 +131,5 @@ def fit_ramps_casertano(
     if resultants_unit is not None:
         parameters = parameters * resultants_unit
 
-    return ols_cas22.RampFitOutputs(output.fits, parameters, variances, dq)
+    # return ols_cas22.RampFitOutputs(output.fits, parameters, variances, dq)
+    return ols_cas22.RampFitOutputs(parameters, variances, dq)

From 5770243c9677fff000b5583b4a7ca68dd23adebd Mon Sep 17 00:00:00 2001
From: William Jamieson <wjamieson@stsci.edu>
Date: Wed, 1 Nov 2023 12:23:34 -0400
Subject: [PATCH 2/4] Fix tests

---
 src/stcal/ramp_fitting/ols_cas22/_core.pxd    |  6 +-
 src/stcal/ramp_fitting/ols_cas22/_core.pyx    | 87 +++++--------------
 .../ramp_fitting/ols_cas22/_fit_ramps.pyx     | 24 ++---
 src/stcal/ramp_fitting/ols_cas22/_fixed.pyx   |  1 +
 src/stcal/ramp_fitting/ols_cas22/_pixel.pxd   |  3 +-
 src/stcal/ramp_fitting/ols_cas22/_pixel.pyx   | 18 ++--
 src/stcal/ramp_fitting/ols_cas22_fit.py       | 18 ++--
 tests/test_jump_cas22.py                      | 12 +--
 8 files changed, 75 insertions(+), 94 deletions(-)

diff --git a/src/stcal/ramp_fitting/ols_cas22/_core.pxd b/src/stcal/ramp_fitting/ols_cas22/_core.pxd
index 63ee6323..694fcc2a 100644
--- a/src/stcal/ramp_fitting/ols_cas22/_core.pxd
+++ b/src/stcal/ramp_fitting/ols_cas22/_core.pxd
@@ -14,10 +14,10 @@ cdef struct RampFit:
 
 
 cdef struct RampFits:
-    # vector[RampFit] fits
-    # vector[RampIndex] index
-    vector[int] jumps
     RampFit average
+    vector[int] jumps
+    vector[RampFit] fits
+    vector[RampIndex] index
 
 
 cdef struct ReadPatternMetadata:
diff --git a/src/stcal/ramp_fitting/ols_cas22/_core.pyx b/src/stcal/ramp_fitting/ols_cas22/_core.pyx
index 5726fea1..3e336fe3 100644
--- a/src/stcal/ramp_fitting/ols_cas22/_core.pyx
+++ b/src/stcal/ramp_fitting/ols_cas22/_core.pyx
@@ -193,69 +193,30 @@ cdef inline stack[RampIndex] init_ramps(int[:, :] dq, int n_resultants, int inde
     return ramps
 
 
-# @cython.boundscheck(False)
-# @cython.wraparound(False)
-# cdef inline deque[stack[RampIndex]] init_ramps(int[:, :] dq):
-#     """
-#     Create the initial ramp stack for each pixel
-#         if dq[index_resultant, index_pixel] == 0, then the resultant is in a ramp
-#         otherwise, the resultant is not in a ramp
-
-#     Parameters
-#     ----------
-#     dq : int[n_resultants, n_pixel]
-#         DQ array
-
-#     Returns
-#     -------
-#     deque of stacks of RampIndex objects
-#         - deque with entry for each pixel
-#             Chosen to be deque because need element access to loop
-#         - stack with entry for each ramp found (top of stack is last ramp found)
-#         - RampIndex with start and end indices of the ramp in the resultants
-#     """
-#     cdef int n_pixel, n_resultants
-
-#     n_resultants, n_pixel = np.array(dq).shape
-#     cdef deque[stack[RampIndex]] pixel_ramps
-
-#     cdef int index_pixel
-
-#     for index_pixel in range(n_pixel):
-#         # Add ramp stack for pixel to list
-#         pixel_ramps.push_back(_init_ramps_pixel(dq, n_resultants, index_pixel))
-
-#     return pixel_ramps
-
-
-# def _init_ramps_list(np.ndarray[int, ndim=2] dq):
-#     """
-#     This is a wrapper for init_ramps so that it can be fully inspected from pure
-#     python. A cpdef cannot be used in that case becase a stack has no direct python
-#     analog. Instead this function turns that stack into a list ordered in the same
-#     order as the stack; meaning that, the first element of the list is the top of
-#     the stack.
-#         Note this function is for testing purposes only and so is marked as private
-#         within this private module
-#     """
-#     cdef deque[stack[RampIndex]] raw = init_ramps(dq)
-
-#     # Have to turn deque and stack into python compatible objects
-#     cdef RampIndex index
-#     cdef stack[RampIndex] ramp
-#     cdef list out = []
-#     cdef list stack_out
-#     for ramp in raw:
-#         stack_out = []
-#         while not ramp.empty():
-#             index = ramp.top()
-#             ramp.pop()
-#             # So top of stack is first item of list
-#             stack_out = [index] + stack_out
-
-#         out.append(stack_out)
-
-#     return out
+def _init_ramps_list(np.ndarray[int, ndim=2] dq, int n_resultants, int index_pixel):
+    """
+    This is a wrapper for init_ramps so that it can be fully inspected from pure
+    python. A cpdef cannot be used in that case becase a stack has no direct python
+    analog. Instead this function turns that stack into a list ordered in the same
+    order as the stack; meaning that, the first element of the list is the top of
+    the stack.
+        Note this function is for testing purposes only and so is marked as private
+        within this private module
+    """
+    cdef stack[RampIndex] ramp = init_ramps(dq, n_resultants, index_pixel)
+
+    # Have to turn deque and stack into python compatible objects
+    cdef RampIndex index
+    cdef list out = []
+
+    out = []
+    while not ramp.empty():
+        index = ramp.top()
+        ramp.pop()
+        # So top of stack is first item of list
+        out = [index] + out
+
+    return out
 
 
 @cython.boundscheck(False)
diff --git a/src/stcal/ramp_fitting/ols_cas22/_fit_ramps.pyx b/src/stcal/ramp_fitting/ols_cas22/_fit_ramps.pyx
index 78c77672..4b2e39e3 100644
--- a/src/stcal/ramp_fitting/ols_cas22/_fit_ramps.pyx
+++ b/src/stcal/ramp_fitting/ols_cas22/_fit_ramps.pyx
@@ -10,7 +10,7 @@ from stcal.ramp_fitting.ols_cas22._core cimport (RampFits, RampIndex, Thresh,
 from stcal.ramp_fitting.ols_cas22._fixed cimport fixed_values_from_metadata, FixedValues
 from stcal.ramp_fitting.ols_cas22._pixel cimport make_pixel
 
-from typing import NamedTuple
+from typing import NamedTuple, Optional
 
 
 # Fix the default Threshold values at compile time these values cannot be overridden
@@ -26,9 +26,6 @@ class RampFitOutputs(NamedTuple):
 
     Attributes
     ----------
-        fits: list of RampFits
-            the raw ramp fit outputs, these are all structs which will get mapped to
-            python dictionaries.
         parameters: np.ndarray[n_pixel, 2]
             the slope and intercept for each pixel's ramp fit. see Parameter enum
             for indexing indicating slope/intercept in the second dimension.
@@ -39,11 +36,14 @@ class RampFitOutputs(NamedTuple):
         dq: np.ndarray[n_resultants, n_pixel]
             the dq array, with additional flags set for jumps detected by the
             jump detection algorithm.
+        fits: list of RampFits
+            the raw ramp fit outputs, these are all structs which will get mapped to
+            python dictionaries.
     """
-    # fits: list
     parameters: np.ndarray
     variances: np.ndarray
     dq: np.ndarray
+    fits: Optional[list] = None
 
 
 @cython.boundscheck(False)
@@ -55,7 +55,8 @@ def fit_ramps(np.ndarray[float, ndim=2] resultants,
               list[list[int]] read_pattern,
               bool use_jump=False,
               float intercept=DefaultIntercept,
-              float constant=DefaultConstant):
+              float constant=DefaultConstant,
+              bool include_diagnostic=False):
     """Fit ramps using the Casertano+22 algorithm.
         This implementation uses the Cas22 algorithm to fit ramps, where
         ramps are fit between bad resultants marked by dq flags for each pixel
@@ -82,6 +83,8 @@ def fit_ramps(np.ndarray[float, ndim=2] resultants,
         The intercept value for the threshold function. Default=5.5
     constant : float
         The constant value for the threshold function. Default=1/3.0
+    include_diagnostic : bool
+        If True, include the raw ramp fits in the output. Default=False
 
     Returns
     -------
@@ -103,7 +106,7 @@ def fit_ramps(np.ndarray[float, ndim=2] resultants,
     # Use list because this might grow very large which would require constant
     #    reallocation. We don't need random access, and this gets cast to a python
     #    list in the end.
-    # cdef cpp_list[RampFits] ramp_fits
+    cdef cpp_list[RampFits] ramp_fits
 
     cdef np.ndarray[float, ndim=2] parameters = np.zeros((n_pixels, 2), dtype=np.float32)
     cdef np.ndarray[float, ndim=2] variances = np.zeros((n_pixels, 3), dtype=np.float32)
@@ -114,7 +117,7 @@ def fit_ramps(np.ndarray[float, ndim=2] resultants,
     for index in range(n_pixels):
         # Fit all the ramps for the given pixel
         fit = make_pixel(fixed, read_noise[index],
-                         resultants[:, index]).fit_ramps(init_ramps(dq, n_resultants, index))
+                         resultants[:, index]).fit_ramps(init_ramps(dq, n_resultants, index), include_diagnostic)
 
         parameters[index, Parameter.slope] = fit.average.slope
 
@@ -125,7 +128,8 @@ def fit_ramps(np.ndarray[float, ndim=2] resultants,
         for jump in fit.jumps:
             dq[jump, index] = RampJumpDQ.JUMP_DET
 
-        # ramp_fits.push_back(fit)
+        if include_diagnostic:
+            ramp_fits.push_back(fit)
 
     # return RampFitOutputs(ramp_fits, parameters, variances, dq)
-    return RampFitOutputs(parameters, variances, dq)
+    return RampFitOutputs(parameters, variances, dq, ramp_fits if include_diagnostic else None)
diff --git a/src/stcal/ramp_fitting/ols_cas22/_fixed.pyx b/src/stcal/ramp_fitting/ols_cas22/_fixed.pyx
index 6bd72b07..9a6b4071 100644
--- a/src/stcal/ramp_fitting/ols_cas22/_fixed.pyx
+++ b/src/stcal/ramp_fitting/ols_cas22/_fixed.pyx
@@ -18,6 +18,7 @@ Functions
 import numpy as np
 cimport numpy as np
 cimport cython
+from libcpp cimport bool
 
 from stcal.ramp_fitting.ols_cas22._core cimport Thresh, ReadPatternMetadata, Diff
 from stcal.ramp_fitting.ols_cas22._fixed cimport FixedValues
diff --git a/src/stcal/ramp_fitting/ols_cas22/_pixel.pxd b/src/stcal/ramp_fitting/ols_cas22/_pixel.pxd
index bf390419..e24272ef 100644
--- a/src/stcal/ramp_fitting/ols_cas22/_pixel.pxd
+++ b/src/stcal/ramp_fitting/ols_cas22/_pixel.pxd
@@ -1,3 +1,4 @@
+from libcpp cimport bool
 from libcpp.stack cimport stack
 
 from stcal.ramp_fitting.ols_cas22._core cimport RampFit, RampFits, RampIndex
@@ -17,7 +18,7 @@ cdef class Pixel:
     cdef float correction(Pixel self, RampIndex ramp, float slope)
     cdef float stat(Pixel self, float slope, RampIndex ramp, int index, int diff)
     cdef float[:] stats(Pixel self, float slope, RampIndex ramp)
-    cdef RampFits fit_ramps(Pixel self, stack[RampIndex] ramps)
+    cdef RampFits fit_ramps(Pixel self, stack[RampIndex] ramps, bool include_diagnostic)
 
 
 cpdef Pixel make_pixel(FixedValues fixed, float read_noise, float [:] resultants)
diff --git a/src/stcal/ramp_fitting/ols_cas22/_pixel.pyx b/src/stcal/ramp_fitting/ols_cas22/_pixel.pyx
index 5ddfe256..9ece0f8c 100644
--- a/src/stcal/ramp_fitting/ols_cas22/_pixel.pyx
+++ b/src/stcal/ramp_fitting/ols_cas22/_pixel.pyx
@@ -14,6 +14,7 @@ Functions
             - cpdef gives a python wrapper, but the python version of this method
               is considered private, only to be used for testing
 """
+from libcpp cimport bool
 from libc.math cimport sqrt, fabs
 from libcpp.vector cimport vector
 from libcpp.stack cimport stack
@@ -330,7 +331,7 @@ cdef class Pixel:
     @cython.boundscheck(False)
     @cython.wraparound(False)
     @cython.cdivision(True)
-    cdef inline RampFits fit_ramps(Pixel self, stack[RampIndex] ramps):
+    cdef inline RampFits fit_ramps(Pixel self, stack[RampIndex] ramps, bool include_diagnostic):
         """
         Compute all the ramps for a single pixel using the Casertano+22 algorithm
             with jump detection.
@@ -398,8 +399,9 @@ cdef class Pixel:
                     #     consideration.
                     jump0 = np.argmax(stats) + ramp.start
                     jump1 = jump0 + 1
-                    ramp_fits.jumps.push_back(jump0)
-                    ramp_fits.jumps.push_back(jump1)
+                    if include_diagnostic:
+                        ramp_fits.jumps.push_back(jump0)
+                        ramp_fits.jumps.push_back(jump1)
 
                     # The two resultant indicies need to be skipped, therefore
                     # the two
@@ -446,8 +448,9 @@ cdef class Pixel:
             #    than threshold
             # Note that ramps are computed backward in time meaning we need to
             #  reverse the order of the fits at the end
-            # ramp_fits.fits.push_back(ramp_fit)
-            # ramp_fits.index.push_back(ramp)
+            if include_diagnostic:
+                ramp_fits.fits.push_back(ramp_fit)
+                ramp_fits.index.push_back(ramp)
 
             # Start computing the averages
             #    Note we do not do anything in the NaN case for degenerate ramps
@@ -462,8 +465,9 @@ cdef class Pixel:
                 ramp_fits.average.poisson_var += weight**2 * ramp_fit.poisson_var
 
         # Reverse to order in time
-        # ramp_fits.fits = ramp_fits.fits[::-1]
-        # ramp_fits.index = ramp_fits.index[::-1]
+        if include_diagnostic:
+            ramp_fits.fits = ramp_fits.fits[::-1]
+            ramp_fits.index = ramp_fits.index[::-1]
 
         # Finish computing averages
         ramp_fits.average.slope /= total_weight if total_weight != 0 else 1
diff --git a/src/stcal/ramp_fitting/ols_cas22_fit.py b/src/stcal/ramp_fitting/ols_cas22_fit.py
index 8a8b6a3e..663ede7d 100644
--- a/src/stcal/ramp_fitting/ols_cas22_fit.py
+++ b/src/stcal/ramp_fitting/ols_cas22_fit.py
@@ -78,11 +78,19 @@ def fit_ramps_casertano(
 
     Returns
     -------
-    par : np.ndarray[..., 2] (float)
-        the best fit pedestal and slope for each pixel
-    var : np.ndarray[..., 3, 2, 2] (float)
-        the covariance matrix of par, for each of three noise terms:
-        the read noise, Poisson source noise, and total noise.
+    RampFitOutputs
+        parameters: np.ndarray[n_pixel, 2]
+            the slope and intercept for each pixel's ramp fit. see Parameter enum
+            for indexing indicating slope/intercept in the second dimension.
+        variances: np.ndarray[n_pixel, 3]
+            the read, poisson, and total variances for each pixel's ramp fit.
+            see Variance enum for indexing indicating read/poisson/total in the
+            second dimension.
+        dq: np.ndarray[n_resultants, n_pixel]
+            the dq array, with additional flags set for jumps detected by the
+            jump detection algorithm.
+        fits: always None, this is a hold over which can contain the diagnostic
+            fit information from the jump detection algorithm.
     """
 
     # Trickery to avoid having to specify the defaults for the threshold
diff --git a/tests/test_jump_cas22.py b/tests/test_jump_cas22.py
index 5d49597f..6fc945fb 100644
--- a/tests/test_jump_cas22.py
+++ b/tests/test_jump_cas22.py
@@ -89,7 +89,9 @@ def test_init_ramps():
                    [0, 0, 0, 1, 0, 0, 1, 0, 1, 0, 1, 1, 0, 1, 1, 1],
                    [0, 0, 0, 0, 1, 0, 0, 1, 0, 1, 1, 0, 1, 1, 1, 1]], dtype=np.int32)
 
-    ramps = _init_ramps_list(dq)
+    n_resultants, n_pixels = dq.shape
+    ramps = [_init_ramps_list(dq, n_resultants, index_pixel) for index_pixel in range(n_pixels)]
+
     assert len(ramps) == dq.shape[1] == 16
 
     # Check that the ramps are correct
@@ -418,7 +420,7 @@ def test_fit_ramps(detector_data, use_jump, use_dq):
     if not use_dq:
         assert okay.all()
 
-    output = fit_ramps(resultants, dq, read_noise, READ_TIME, read_pattern, use_jump=use_jump)
+    output = fit_ramps(resultants, dq, read_noise, READ_TIME, read_pattern, use_jump=use_jump, include_diagnostic=True)
     assert len(output.fits) == N_PIXELS  # sanity check that a fit is output for each pixel
 
     chi2 = 0
@@ -456,7 +458,7 @@ def test_fit_ramps_array_outputs(detector_data, use_jump):
     resultants, read_noise, read_pattern = detector_data
     dq = np.zeros(resultants.shape, dtype=np.int32)
 
-    output = fit_ramps(resultants, dq, read_noise, READ_TIME, read_pattern, use_jump=use_jump)
+    output = fit_ramps(resultants, dq, read_noise, READ_TIME, read_pattern, use_jump=use_jump, include_diagnostic=True)
 
     for fit, par, var in zip(output.fits, output.parameters, output.variances):
         assert par[Parameter.intercept] == 0
@@ -528,7 +530,7 @@ def test_find_jumps(jump_data):
     resultants, read_noise, read_pattern, jump_reads, jump_resultants = jump_data
     dq = np.zeros(resultants.shape, dtype=np.int32)
 
-    output = fit_ramps(resultants, dq, read_noise, READ_TIME, read_pattern, use_jump=True)
+    output = fit_ramps(resultants, dq, read_noise, READ_TIME, read_pattern, use_jump=True, include_diagnostic=True)
     assert len(output.fits) == len(jump_reads)  # sanity check that a fit/jump is set for every pixel
 
     chi2 = 0
@@ -603,7 +605,7 @@ def test_jump_dq_set(jump_data):
     resultants, read_noise, read_pattern, jump_reads, jump_resultants = jump_data
     dq = np.zeros(resultants.shape, dtype=np.int32)
 
-    output = fit_ramps(resultants, dq, read_noise, READ_TIME, read_pattern, use_jump=True)
+    output = fit_ramps(resultants, dq, read_noise, READ_TIME, read_pattern, use_jump=True, include_diagnostic=True)
 
     for fit, pixel_dq in zip(output.fits, output.dq.transpose()):
         # Check that all jumps found get marked

From bbd70000ca5e4f24676ba8ff1d7cc9820af6e388 Mon Sep 17 00:00:00 2001
From: William Jamieson <wjamieson@stsci.edu>
Date: Wed, 1 Nov 2023 12:54:00 -0400
Subject: [PATCH 3/4] Use vector instead of stack

---
 src/stcal/ramp_fitting/ols_cas22/_core.pxd  |  3 +-
 src/stcal/ramp_fitting/ols_cas22/_core.pyx  | 41 ++++-----------------
 src/stcal/ramp_fitting/ols_cas22/_pixel.pxd |  4 +-
 src/stcal/ramp_fitting/ols_cas22/_pixel.pyx | 15 ++++----
 tests/test_jump_cas22.py                    |  6 +--
 5 files changed, 20 insertions(+), 49 deletions(-)

diff --git a/src/stcal/ramp_fitting/ols_cas22/_core.pxd b/src/stcal/ramp_fitting/ols_cas22/_core.pxd
index 694fcc2a..641c206e 100644
--- a/src/stcal/ramp_fitting/ols_cas22/_core.pxd
+++ b/src/stcal/ramp_fitting/ols_cas22/_core.pxd
@@ -1,5 +1,4 @@
 from libcpp.vector cimport vector
-from libcpp.stack cimport stack
 
 
 cdef struct RampIndex:
@@ -53,5 +52,5 @@ cpdef enum RampJumpDQ:
 
 cpdef float threshold(Thresh thresh, float slope)
 cdef float get_power(float s)
-cdef stack[RampIndex] init_ramps(int[:, :] dq, int n_resultants, int index_pixel)
+cpdef vector[RampIndex] init_ramps(int[:, :] dq, int n_resultants, int index_pixel)
 cpdef ReadPatternMetadata metadata_from_read_pattern(list[list[int]] read_pattern, float read_time)
diff --git a/src/stcal/ramp_fitting/ols_cas22/_core.pyx b/src/stcal/ramp_fitting/ols_cas22/_core.pyx
index 3e336fe3..b1864cca 100644
--- a/src/stcal/ramp_fitting/ols_cas22/_core.pyx
+++ b/src/stcal/ramp_fitting/ols_cas22/_core.pyx
@@ -74,8 +74,7 @@ Functions
         - cpdef gives a python wrapper, but the python version of this method
           is considered private, only to be used for testing
 """
-from libcpp.stack cimport stack
-from libcpp.deque cimport deque
+from libcpp.vector cimport vector
 from libc.math cimport log10
 
 import numpy as np
@@ -134,7 +133,7 @@ cpdef inline float threshold(Thresh thresh, float slope):
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-cdef inline stack[RampIndex] init_ramps(int[:, :] dq, int n_resultants, int index_pixel):
+cpdef inline vector[RampIndex] init_ramps(int[:, :] dq, int n_resultants, int index_pixel):
     """
     Create the initial ramp stack for each pixel
         if dq[index_resultant, index_pixel] == 0, then the resultant is in a ramp
@@ -151,11 +150,11 @@ cdef inline stack[RampIndex] init_ramps(int[:, :] dq, int n_resultants, int inde
 
     Returns
     -------
-    stack of RampIndex objects
-        - stack with entry for each ramp found (top of stack is last ramp found)
+    vector of RampIndex objects
+        - vector with entry for each ramp found (last entry is last ramp found)
         - RampIndex with start and end indices of the ramp in the resultants
     """
-    cdef stack[RampIndex] ramps = stack[RampIndex]()
+    cdef vector[RampIndex] ramps = vector[RampIndex]()
 
     # Note: if start/end are -1, then no value has been assigned
     # ramp.start == -1 means we have not started a ramp
@@ -181,44 +180,18 @@ cdef inline stack[RampIndex] init_ramps(int[:, :] dq, int n_resultants, int inde
                 ramp.end = index_resultant - 1
 
                 # Add completed ramp to stack and reset ramp
-                ramps.push(ramp)
+                ramps.push_back(ramp)
                 ramp = RampIndex(-1, -1)
 
     # Handle case where last resultant is in ramp (so no end has been set)
     if ramp.start != -1 and ramp.end == -1:
         # Last resultant is end of the ramp => set then add to stack
         ramp.end = n_resultants - 1
-        ramps.push(ramp)
+        ramps.push_back(ramp)
 
     return ramps
 
 
-def _init_ramps_list(np.ndarray[int, ndim=2] dq, int n_resultants, int index_pixel):
-    """
-    This is a wrapper for init_ramps so that it can be fully inspected from pure
-    python. A cpdef cannot be used in that case becase a stack has no direct python
-    analog. Instead this function turns that stack into a list ordered in the same
-    order as the stack; meaning that, the first element of the list is the top of
-    the stack.
-        Note this function is for testing purposes only and so is marked as private
-        within this private module
-    """
-    cdef stack[RampIndex] ramp = init_ramps(dq, n_resultants, index_pixel)
-
-    # Have to turn deque and stack into python compatible objects
-    cdef RampIndex index
-    cdef list out = []
-
-    out = []
-    while not ramp.empty():
-        index = ramp.top()
-        ramp.pop()
-        # So top of stack is first item of list
-        out = [index] + out
-
-    return out
-
-
 @cython.boundscheck(False)
 @cython.wraparound(False)
 cpdef ReadPatternMetadata metadata_from_read_pattern(list[list[int]] read_pattern, float read_time):
diff --git a/src/stcal/ramp_fitting/ols_cas22/_pixel.pxd b/src/stcal/ramp_fitting/ols_cas22/_pixel.pxd
index e24272ef..1539724b 100644
--- a/src/stcal/ramp_fitting/ols_cas22/_pixel.pxd
+++ b/src/stcal/ramp_fitting/ols_cas22/_pixel.pxd
@@ -1,5 +1,5 @@
 from libcpp cimport bool
-from libcpp.stack cimport stack
+from libcpp.vector cimport vector
 
 from stcal.ramp_fitting.ols_cas22._core cimport RampFit, RampFits, RampIndex
 from stcal.ramp_fitting.ols_cas22._fixed cimport FixedValues
@@ -18,7 +18,7 @@ cdef class Pixel:
     cdef float correction(Pixel self, RampIndex ramp, float slope)
     cdef float stat(Pixel self, float slope, RampIndex ramp, int index, int diff)
     cdef float[:] stats(Pixel self, float slope, RampIndex ramp)
-    cdef RampFits fit_ramps(Pixel self, stack[RampIndex] ramps, bool include_diagnostic)
+    cdef RampFits fit_ramps(Pixel self, vector[RampIndex] ramps, bool include_diagnostic)
 
 
 cpdef Pixel make_pixel(FixedValues fixed, float read_noise, float [:] resultants)
diff --git a/src/stcal/ramp_fitting/ols_cas22/_pixel.pyx b/src/stcal/ramp_fitting/ols_cas22/_pixel.pyx
index 9ece0f8c..6daa33ca 100644
--- a/src/stcal/ramp_fitting/ols_cas22/_pixel.pyx
+++ b/src/stcal/ramp_fitting/ols_cas22/_pixel.pyx
@@ -17,7 +17,6 @@ Functions
 from libcpp cimport bool
 from libc.math cimport sqrt, fabs
 from libcpp.vector cimport vector
-from libcpp.stack cimport stack
 
 import numpy as np
 cimport numpy as np
@@ -331,15 +330,15 @@ cdef class Pixel:
     @cython.boundscheck(False)
     @cython.wraparound(False)
     @cython.cdivision(True)
-    cdef inline RampFits fit_ramps(Pixel self, stack[RampIndex] ramps, bool include_diagnostic):
+    cdef inline RampFits fit_ramps(Pixel self, vector[RampIndex] ramps, bool include_diagnostic):
         """
         Compute all the ramps for a single pixel using the Casertano+22 algorithm
             with jump detection.
 
         Parameters
         ----------
-        ramps : stack[RampIndex]
-            Stack of initial ramps to fit for a single pixel
+        ramps : vector[RampIndex]
+            Vector of initial ramps to fit for a single pixel
             multiple ramps are possible due to dq flags
 
         Returns
@@ -362,8 +361,8 @@ cdef class Pixel:
         # Run while the stack is non-empty
         while not ramps.empty():
             # Remove top ramp of the stack to use
-            ramp = ramps.top()
-            ramps.pop()
+            ramp = ramps.back()
+            ramps.pop_back()
 
             # Compute fit
             ramp_fit = self.fit_ramp(ramp)
@@ -426,7 +425,7 @@ cdef class Pixel:
                         #    important that we exclude it.
                         # Note that jump0 < ramp.start is not possible because
                         # the argmax is always >= 0
-                        ramps.push(RampIndex(ramp.start, jump0 - 1))
+                        ramps.push_back(RampIndex(ramp.start, jump0 - 1))
 
                     if jump1 < ramp.end:
                         # Note that if jump1 == ramp.end, we have detected a
@@ -440,7 +439,7 @@ cdef class Pixel:
                         # resultants which are not considered part of the ramp
                         # under consideration. Therefore, we have to exlude all
                         # of those values.
-                        ramps.push(RampIndex(jump1 + 1, ramp.end))
+                        ramps.push_back(RampIndex(jump1 + 1, ramp.end))
 
                     continue
 
diff --git a/tests/test_jump_cas22.py b/tests/test_jump_cas22.py
index 6fc945fb..867bfadd 100644
--- a/tests/test_jump_cas22.py
+++ b/tests/test_jump_cas22.py
@@ -2,7 +2,7 @@
 import pytest
 from numpy.testing import assert_allclose
 
-from stcal.ramp_fitting.ols_cas22._core import metadata_from_read_pattern, threshold
+from stcal.ramp_fitting.ols_cas22._core import init_ramps, metadata_from_read_pattern, threshold
 from stcal.ramp_fitting.ols_cas22._fixed import fixed_values_from_metadata
 from stcal.ramp_fitting.ols_cas22._pixel import make_pixel
 
@@ -82,7 +82,7 @@ def test_init_ramps():
         a direct python equivalent, we call the wrapper for `init_ramps` which
         converts that stack into a list ordered in the same fashion as the stack
     """
-    from stcal.ramp_fitting.ols_cas22._core import _init_ramps_list
+    # from stcal.ramp_fitting.ols_cas22._core import _init_ramps_list
 
     dq = np.array([[0, 1, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 1, 0, 1],
                    [0, 0, 1, 0, 0, 1, 0, 0, 1, 1, 0, 1, 1, 0, 1, 1],
@@ -90,7 +90,7 @@ def test_init_ramps():
                    [0, 0, 0, 0, 1, 0, 0, 1, 0, 1, 1, 0, 1, 1, 1, 1]], dtype=np.int32)
 
     n_resultants, n_pixels = dq.shape
-    ramps = [_init_ramps_list(dq, n_resultants, index_pixel) for index_pixel in range(n_pixels)]
+    ramps = [init_ramps(dq, n_resultants, index_pixel) for index_pixel in range(n_pixels)]
 
     assert len(ramps) == dq.shape[1] == 16
 

From c6f90bece6485ab348afca976b042d533d181aba Mon Sep 17 00:00:00 2001
From: William Jamieson <wjamieson@stsci.edu>
Date: Wed, 1 Nov 2023 12:57:06 -0400
Subject: [PATCH 4/4] Update changes

---
 CHANGES.rst              |  1 +
 tests/test_jump_cas22.py | 12 ++++++++----
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/CHANGES.rst b/CHANGES.rst
index dd186cf8..ff8324ba 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -8,6 +8,7 @@ ramp_fitting
 
 - Refactor Casertano, et.al, 2022 uneven ramp fitting and incorporate the matching
   jump detection algorithm into it. [#215]
+- Fix memory issue with uneven ramp fitting [#226]
 
 Changes to API
 --------------
diff --git a/tests/test_jump_cas22.py b/tests/test_jump_cas22.py
index 867bfadd..d5c40eee 100644
--- a/tests/test_jump_cas22.py
+++ b/tests/test_jump_cas22.py
@@ -420,7 +420,8 @@ def test_fit_ramps(detector_data, use_jump, use_dq):
     if not use_dq:
         assert okay.all()
 
-    output = fit_ramps(resultants, dq, read_noise, READ_TIME, read_pattern, use_jump=use_jump, include_diagnostic=True)
+    output = fit_ramps(resultants, dq, read_noise, READ_TIME, read_pattern, use_jump=use_jump,
+                       include_diagnostic=True)
     assert len(output.fits) == N_PIXELS  # sanity check that a fit is output for each pixel
 
     chi2 = 0
@@ -458,7 +459,8 @@ def test_fit_ramps_array_outputs(detector_data, use_jump):
     resultants, read_noise, read_pattern = detector_data
     dq = np.zeros(resultants.shape, dtype=np.int32)
 
-    output = fit_ramps(resultants, dq, read_noise, READ_TIME, read_pattern, use_jump=use_jump, include_diagnostic=True)
+    output = fit_ramps(resultants, dq, read_noise, READ_TIME, read_pattern, use_jump=use_jump,
+                       include_diagnostic=True)
 
     for fit, par, var in zip(output.fits, output.parameters, output.variances):
         assert par[Parameter.intercept] == 0
@@ -530,7 +532,8 @@ def test_find_jumps(jump_data):
     resultants, read_noise, read_pattern, jump_reads, jump_resultants = jump_data
     dq = np.zeros(resultants.shape, dtype=np.int32)
 
-    output = fit_ramps(resultants, dq, read_noise, READ_TIME, read_pattern, use_jump=True, include_diagnostic=True)
+    output = fit_ramps(resultants, dq, read_noise, READ_TIME, read_pattern, use_jump=True,
+                       include_diagnostic=True)
     assert len(output.fits) == len(jump_reads)  # sanity check that a fit/jump is set for every pixel
 
     chi2 = 0
@@ -605,7 +608,8 @@ def test_jump_dq_set(jump_data):
     resultants, read_noise, read_pattern, jump_reads, jump_resultants = jump_data
     dq = np.zeros(resultants.shape, dtype=np.int32)
 
-    output = fit_ramps(resultants, dq, read_noise, READ_TIME, read_pattern, use_jump=True, include_diagnostic=True)
+    output = fit_ramps(resultants, dq, read_noise, READ_TIME, read_pattern, use_jump=True,
+                       include_diagnostic=True)
 
     for fit, pixel_dq in zip(output.fits, output.dq.transpose()):
         # Check that all jumps found get marked