Implement adaptive batch size

pyproject.toml

@@ -135,6 +135,7 @@ types = [
     "types-decorator",
     "types-docutils",
     "types-tqdm",
+    "types-psutil"
 ]
 
 [tool.setuptools]

src/ert/analysis/_es_update.py

@@ -21,6 +21,7 @@
 
 import iterative_ensemble_smoother as ies
 import numpy as np
+import psutil
 import xarray as xr
 from iterative_ensemble_smoother.experimental import (
     AdaptiveESMDA,
@@ -570,6 +571,28 @@ def analysis_ES(
             )
             if module.localization:
                 num_params = temp_storage[param_group.name].shape[0]
+
+                # Calculate adaptive batch size.
+                # Adaptive Localization calculates the cross-covariance between
+                # parameters and responses.
+                # Cross-covariance is a matrix with shape num_params x num_obs
+                # which may be larger than memory.
+                # https://stackoverflow.com/a/21632554
+                process = psutil.Process()
+                available_memory_bytes = process.memory_info().rss
+                memory_safety_factor = 0.8
+                bytes_in_float64 = 8
+                batch_size = min(
+                    int(
+                        np.floor(
+                            available_memory_bytes
+                            * memory_safety_factor
+                            / (num_obs * bytes_in_float64)
+                        )
+                    ),
+                    num_params,
+                )
+
                 batches = _split_by_batchsize(np.arange(0, num_params), batch_size)
 
                 progress_callback(