From 563a0f183a1919ac4fb783dd284f241d44075d85 Mon Sep 17 00:00:00 2001
From: Mike Regan <mregan@stsci.edu>
Date: Wed, 27 Dec 2023 15:46:56 -0500
Subject: [PATCH] Add Unit test

---
 src/stcal/ramp_fitting/ols_fit.py |  8 +++++---
 tests/test_ramp_fitting.py        | 23 +++++++++++++++++++++++
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/src/stcal/ramp_fitting/ols_fit.py b/src/stcal/ramp_fitting/ols_fit.py
index 3d121856..2b238428 100644
--- a/src/stcal/ramp_fitting/ols_fit.py
+++ b/src/stcal/ramp_fitting/ols_fit.py
@@ -1156,8 +1156,7 @@ def ramp_fit_compute_variances(ramp_data, gain_2d, readnoise_2d, fit_slopes_ans,
             # Suppress harmless arithmetic warnings for now
             warnings.filterwarnings("ignore", ".*invalid value.*", RuntimeWarning)
             warnings.filterwarnings("ignore", ".*divide by zero.*", RuntimeWarning)
-            var_p4[num_int, :, rlo:rhi, :] = den_p3 * (med_rates[rlo:rhi, :] +
-                                                       avg_dark_current)
+            var_p4[num_int, :, rlo:rhi, :] = (den_p3 * med_rates[rlo:rhi, :]) + avg_dark_current
 
             # Find the segment variance due to read noise and convert back to DN
             var_r4[num_int, :, rlo:rhi, :] = num_r3 * den_r3 / gain_sect**2
@@ -1197,7 +1196,10 @@ def ramp_fit_compute_variances(ramp_data, gain_2d, readnoise_2d, fit_slopes_ans,
         var_p3[:, med_rates <= 0.0] = 0.0
         warnings.resetwarnings()
 
-        var_p4[num_int, :, med_rates <= 0.0] = 0.0
+        # For pixels with zero or negative rates set the Poisson variance to the
+        # dark current rate.
+        var_p4[num_int, :, med_rates <= 0.0] = avg_dark_current
+
         var_both4[num_int, :, :, :] = var_r4[num_int, :, :, :] + var_p4[num_int, :, :, :]
         inv_var_both4[num_int, :, :, :] = 1.0 / var_both4[num_int, :, :, :]
 
diff --git a/tests/test_ramp_fitting.py b/tests/test_ramp_fitting.py
index fef42097..f566b0ba 100644
--- a/tests/test_ramp_fitting.py
+++ b/tests/test_ramp_fitting.py
@@ -3,6 +3,7 @@
 from stcal.ramp_fitting.ramp_fit import ramp_fit_data
 from stcal.ramp_fitting.ramp_fit_class import RampData
 from stcal.ramp_fitting.utils import compute_num_slices
+from stcal.ramp_fitting.ols_fit import ramp_fit_compute_variances
 
 DELIM = "=" * 70
 
@@ -1466,6 +1467,28 @@ def test_compute_num_slices():
     n_rows = 9
     assert compute_num_slices("21", n_rows, max_available_cores) == 9
 
+def test_ramp_fit_dark_current_variance():
+    nints, ngroups, nrows, ncols = 1, 5, 2, 2
+    rnval, gval = 10.0, 5.0
+    frame_time, nframes, groupgap = 10, 1, 0
+    var = rnval, gval
+    tm = frame_time, nframes, groupgap
+    dims = nints, ngroups, nrows, ncols
+    ramp, gain, rnoise = create_blank_ramp_data(dims, var, tm)
+    #Test for zero dark rate
+    avg_dark_current = 0.0
+    med_rates = np.zeros(shape=(2, 2))
+    fit_slopes_ans = (1, 0, 0, 0, 0, 1, 0, 0, 0, med_rates)
+    var_p3, var_r3, var_p4, var_r4, var_both4, var_both3, inv_var_both4, s_inv_var_p3, \
+    s_inv_var_r3, s_inv_var_both3 = ramp_fit_compute_variances(ramp, gain, rnoise, fit_slopes_ans, avg_dark_current)
+
+    assert(var_p4[0, 0, 0, 0] == 0.0)
+    
+    #Test for a dark rate of 1.5 e-/sec
+    avg_dark_current = 1.5
+    var_p3, var_r3, var_p4, var_r4, var_both4, var_both3, inv_var_both4, s_inv_var_p3, \
+    s_inv_var_r3, s_inv_var_both3 = ramp_fit_compute_variances(ramp, gain, rnoise, fit_slopes_ans, avg_dark_current)
+    assert(var_p4[0, 0, 0, 0] == 1.5)
 
 # -----------------------------------------------------------------------------
 #                           Set up functions