Repository: kevin218/eureka
Issue: Running S3 with multiple threads seems to be slower than one thread

When testing the new jwst v1.3.3 I noticed that the background subtraction step in S3 ([here](https://github.com/kevin218/Eureka/blob/fc3c4c9cf972b9fc1896133af708349edc54c6b9/eureka/S3_data_reduction/background.py#L65-L73)) was requiring 6 seconds per iteration with 14 cpu threads, but when I tried to run it with only 1 cpu thread it only took 2 seconds per iteration. I can imagine the runtime not scaling linearly with the cpu threads, but running slower seems very unintuitive and is obviously a bad sign. We'll need to see if this is related to my other multithreading issue (#92) or if the there is a larger issue with multithreading here.

Changed files:
eureka/S3_data_reduction/nircam.py, eureka/S3_data_reduction/background.py, eureka/S3_data_reduction/wfc3.py, eureka/S3_data_reduction/miri.py, eureka/S3_data_reduction/nirspec.py

diff --git a/eureka/S3_data_reduction/background.py b/eureka/S3_data_reduction/background.py
index 7d2a7e88..78b2bcb2 100644
--- a/eureka/S3_data_reduction/background.py
+++ b/eureka/S3_data_reduction/background.py
@@ -65,6 +65,13 @@ def BGsubtraction(data, meta, log, isplots):
         data.subbg[n] = bg_data
         data.submask[n] = bg_mask
         return
+    def writeBG_WFC3(arg):
+        bg_data, bg_mask, datav0, datavariance, n = arg
+        data.subbg[n] = bg_data
+        data.submask[n] = bg_mask
+        data.subv0[n] = datav0
+        data.subvariance[n] = datavariance
+        return
 
     # Compute background for each integration
     log.writelog('  Performing background subtraction', mute=(not meta.verbose))
@@ -76,14 +83,27 @@ def BGsubtraction(data, meta, log, isplots):
             iterfn = tqdm(iterfn)
         for n in iterfn:
             # Fit sky background with out-of-spectra data
-            writeBG(inst.fit_bg(data, meta, n, isplots))
+            if meta.inst == 'niriss':
+                writeBG(inst.fit_bg(data, meta, n, isplots))
+            elif meta.inst == 'wfc3':
+                writeBG_WFC3(inst.fit_bg(data.subdata[n], data.submask[n], data.subv0[n], data.subvariance[n], n, meta, isplots))
+            else:
+                writeBG(inst.fit_bg(data.subdata[n], data.submask[n], n, meta, isplots))
     else:
         # Multiple CPUs
         pool = mp.Pool(meta.ncpu)
         args_list = []
-        for n in range(meta.int_start,meta.n_int):
-            args_list.append((data, meta, n, isplots))
-        jobs = [pool.apply_async(func=inst.fit_bg, args=(*args,), callback=writeBG) for args in args_list]
+
+        # Todo, convert NIRISS fit_bg to only accept individual frames (see nircam and below for example)
+        if meta.inst == 'niriss':
+            for n in range(meta.int_start,meta.n_int):
+                args_list.append((data, meta, n, isplots))
+            jobs = [pool.apply_async(func=inst.fit_bg, args=(*args,), callback=writeBG) for args in args_list]
+        elif meta.inst =='wfc3':
+            # The WFC3 background subtraction needs a few more inputs and outputs
+            jobs = [pool.apply_async(func=inst.fit_bg, args=(data.subdata[n], data.submask[n], data.subv0[n], data.subvariance[n], n, meta, isplots,), callback=writeBG_WFC3) for n in range(meta.int_start,meta.n_int)]
+        else: 
+            jobs = [pool.apply_async(func=inst.fit_bg, args=(data.subdata[n], data.submask[n], n, meta, isplots,), callback=writeBG) for n in range(meta.int_start,meta.n_int)]
         pool.close()
         iterfn = jobs
         if meta.verbose:
diff --git a/eureka/S3_data_reduction/miri.py b/eureka/S3_data_reduction/miri.py
index 8c54ba9f..80f50318 100644
--- a/eureka/S3_data_reduction/miri.py
+++ b/eureka/S3_data_reduction/miri.py
@@ -224,9 +224,10 @@ def flag_bg(data, meta):
     return nircam.flag_bg(data, meta)
 
 
-def fit_bg(data, meta, n, isplots=False):
+def fit_bg(dataim, datamask, n, meta, isplots=False):
     '''Fit for a non-uniform background.
 
-    Uses the code written for NIRCam and untested for MIRI, but likely to still work (as long as MIRI data gets rotated)
+    Uses the code written for NIRCam which works for MIRI as long as MIRI data gets rotated.
     '''
-    return nircam.fit_bg(data, meta, n, isplots=isplots)
+    return nircam.fit_bg(dataim, datamask, n, meta, isplots=isplots)
+
diff --git a/eureka/S3_data_reduction/nircam.py b/eureka/S3_data_reduction/nircam.py
index 84f69e32..29e5fe64 100644
--- a/eureka/S3_data_reduction/nircam.py
+++ b/eureka/S3_data_reduction/nircam.py
@@ -1,4 +1,3 @@
-
 # NIRCam specific rountines go here
 import numpy as np
 from astropy.io import fits
@@ -90,10 +89,10 @@ def flag_bg(data, meta):
     return data
 
 
-def fit_bg(data, meta, n, isplots=False):
+def fit_bg(dataim, datamask, n, meta, isplots=False):
     '''Fit for a non-uniform background.
     '''
 
-    bg, mask = background.fitbg(data.subdata[n], meta, data.submask[n], meta.bg_y1, meta.bg_y2, deg=meta.bg_deg,
+    bg, mask = background.fitbg(dataim, meta, datamask, meta.bg_y1, meta.bg_y2, deg=meta.bg_deg,
                                 threshold=meta.p3thresh, isrotate=2, isplots=isplots)
     return (bg, mask, n)
diff --git a/eureka/S3_data_reduction/nirspec.py b/eureka/S3_data_reduction/nirspec.py
index 3f46c1a6..1e2f5b15 100644
--- a/eureka/S3_data_reduction/nirspec.py
+++ b/eureka/S3_data_reduction/nirspec.py
@@ -99,9 +99,9 @@ def flag_bg(data, meta):
 
     return data
 
-def fit_bg(data, meta, n, isplots=False):
+def fit_bg(dataim, datamask, n, meta, isplots=False):
     '''Fit for a non-uniform background.
 
-    Uses the code written for NIRCam and untested for NIRSpec, but likely to still work
+    Uses the code written for NIRCam.
     '''
-    return nircam.fit_bg(data, meta, n, isplots=isplots)
+    return nircam.fit_bg(dataim, datamask, n, meta, isplots=isplots)
diff --git a/eureka/S3_data_reduction/wfc3.py b/eureka/S3_data_reduction/wfc3.py
index ab818b66..25fc6183 100644
--- a/eureka/S3_data_reduction/wfc3.py
+++ b/eureka/S3_data_reduction/wfc3.py
@@ -376,20 +376,20 @@ def flag_bg(data, meta):
     '''
     return nircam.flag_bg(data, meta)
 
-def fit_bg(data, meta, n, isplots=False):
+def fit_bg(dataim, datamask, datav0, datavariance, n, meta, isplots=False):
     '''Fit for a non-uniform background.
 
     Uses the code written for NIRCam, but adds on some extra steps 
     '''
-    bg, mask, n = nircam.fit_bg(data, meta, n, isplots=isplots)
+    bg, mask, n = nircam.fit_bg(dataim, datamask, n, meta, isplots=isplots)
 
     # Calculate variance assuming background dominated rather than read noise dominated
     bgerr       = np.std(bg[n], axis=0)/np.sqrt(np.sum(meta.subdiffmask[-1][n], axis=0))
     bgerr[np.where(np.logical_not(np.isfinite(bgerr)))] = 0.
-    data.subv0[n]      += np.mean(bgerr**2)
-    data.subvariance[n]    = abs(data.subdata[n]) / meta.gain + data.subv0[n]
+    datav0 += np.mean(bgerr**2)
+    datavariance = abs(dataim) / meta.gain + datav0
 
-    return (bg, mask, n)
+    return (dataim, datamask, datav0, datavariance, n)
 
 def correct_drift2D(data, meta, m):
     '''