JP-3706: Set outlier detection buffer size based on total available memory

CHANGES.rst

@@ -496,6 +496,9 @@ outlier_detection
 
 - Fix errors in documentation describing arguments. [#8603]
 
+- Set buffer size for ``create_median`` based on the amount of available memory
+  instead of hard-coding a value, and make it respect ``allowed_memory`` parameter. [#8756]
+
 pathloss
 --------
 

jwst/outlier_detection/imaging.py

@@ -100,7 +100,7 @@ def detect_outliers(
                 input_models.shelve(model, modify=True)
 
     # Perform median combination on set of drizzled mosaics
-    median_data = create_median(drizzled_models, maskpt, on_disk=not in_memory)
+    median_data = create_median(drizzled_models, maskpt, on_disk=not in_memory, allowed_memory=allowed_memory)
 
     if save_intermediate_results:
         # make a median model

jwst/outlier_detection/utils.py

@@ -2,6 +2,7 @@
 The ever-present utils sub-module. A home for all...
 """
 import warnings
+import psutil
 
 import numpy as np
 
@@ -32,7 +33,7 @@
     return median
 
 
-def create_median(resampled_models, maskpt, on_disk=True, buffer_size=10.0):
+def create_median(resampled_models, maskpt, on_disk=True, allowed_memory=None):
     """Create a median image from the singly resampled images.
 
     Parameters
@@ -46,8 +47,9 @@
     on_disk : bool
         If True, the input models are on disk and will be read in chunks.
 
-    buffer_size : float
-        The size of chunk in MB, per input model, that will be read into memory.
+    allowed_memory : float
+        The fraction of available memory to be used by create_median.
+        If None, use 50% of the machine's available memory.
         This parameter has no effect if on_disk is False.
 
     Returns
@@ -75,6 +77,16 @@
         return np.nanmedian(np.array(model_list), axis=0)
     else:
         # set up buffered access to all input models
+        log.info("Computing median in chunks to save memory")
+        if allowed_memory is None:
+            allowed_memory = 0.5
+        machine_available_memory = psutil.virtual_memory().available + psutil.swap_memory().total
+        available_memory = machine_available_memory * allowed_memory
+
+        # buffer_size sets size of chunk in MB that will be read into memory per model
+        buffer_size = available_memory / len(resampled_models.group_names) / _ONE_MB
+        log.info(f"Chunk size {buffer_size} MB per model for {len(resampled_models)} models (totaling {allowed_memory * 100}% of available memory)")
+
         with resampled_models:
             example_model = resampled_models.borrow(0)
             shp = example_model.data.shape
@@ -83,6 +95,11 @@
             resampled_models.shelve(example_model, modify=False)
             del example_model
 
+        if nsections == 1:
+            log.info("Chunk size is large enough to read entire model at once")
+        else:
+            log.info(f"Data will be read in {nsections} sections of {section_nrows} rows each")
+
         # get spatial sections of library and compute timewise median, one by one
         resampled_sections = _get_sections(resampled_models, nsections, section_nrows, shp[0])
         median_image_empty = np.empty(shp, dtype) * np.nan