Implement (adaptable) memory limits for potentially memory-intense functions

deerlab/bootstrap_analysis.py

@@ -10,7 +10,7 @@
 from deerlab.classes import UQResult
 from deerlab.utils import isnumeric
 
-def bootstrap_analysis(fcn,Vexp,Vfit, samples=1000, resampling='gaussian', verbose = False, cores=1):
+def bootstrap_analysis(fcn,Vexp,Vfit, samples=1000, resampling='gaussian', verbose = False, cores=1, memorylimit=8):
     r""" 
     Bootstrap analysis for uncertainty quantification
 
@@ -50,6 +50,9 @@ def bootstrap_analysis(fcn,Vexp,Vfit, samples=1000, resampling='gaussian', verbo
         Specifies whether to print the progress of the bootstrap analysis on the 
         command window, the default is false.
 
+    memorylimit : 
+        Memory limit to be allocated for the full boostrap analysis. If the requested analysis exceeds this limit, the 
+        execution will be stopped. The default is 12GB.  
 
     Returns
     -------
@@ -122,6 +125,11 @@ def sample():
         if not all(isnumeric(x) for x in var):
             raise ValueError('Non-numeric output arguments by the analyzed function are not accepted.')
 
+    # Check that the full bootstrap analysis will not exceed the memory limits
+    memory_requirements = np.sum([nSamples*np.atleast_1d(var).size*np.atleast_1d(var).itemsize])/1e9 # GB
+    if memory_requirements > memorylimit: 
+        raise MemoryError(f'The requested bootstrap analysis requires {memory_requirements:.2f}GB, exceeding the current memory limit {memorylimit:.2f}GB.')
+
     # Get ndarray shapes of all outputs
     shapeout = []
     evals = []

deerlab/dipolarkernel.py

@@ -26,7 +26,7 @@ def ω0(g):
     return (μ0/2)*μB**2*g[0]*g[1]/h*1e21 # Hz m^3 -> MHz nm^3 -> rad μs^-1 nm^3
 
 def dipolarkernel(t, r, *, pathways=None, mod=None, bg=None, method='fresnel', excbandwidth=inf, orisel=None, g=[ge,ge], 
-                  integralop=True, nKnots=5001, clearcache=False):
+                  integralop=True, nKnots=5001, clearcache=False, memorylimit=8):
 #===================================================================================================
     r"""Compute the dipolar kernel operator which enables the linear transformation from
     distance-domain to time-domain data. 
@@ -86,6 +86,10 @@ def dipolarkernel(t, r, *, pathways=None, mod=None, bg=None, method='fresnel', e
     clearcache : boolean, optional
         Clear the cached dipolar kernels at the beginning of the function. Disabled by default.
 
+    memorylimit : 
+        Memory limit to be allocated for the dipolar kernel. If the requested kernel exceeds this limit, the 
+        execution will be stopped. The default is 12GB.  
+
     Returns
     --------
     K : ndarray
@@ -152,6 +156,10 @@ def dipolarkernel(t, r, *, pathways=None, mod=None, bg=None, method='fresnel', e
     # Ensure that inputs are Numpy arrays
     r,t,g = np.atleast_1d(r,t,g)
 
+    memory_requirement = 8*len(t)*len(r)/1e9 # GB
+    if memory_requirement > memorylimit: 
+        raise MemoryError(f'The requested kernel requires {memory_requirement:.2f}GB, exceeding the current memory limit {memorylimit:.2f}GB.')
+
     # Check validity of some inputs
     if len(g) == 1:
         g = [g, g]

test/test_bootstrap_analysis.py

@@ -170,3 +170,16 @@ def bootfcn(V):
 
     assert all(abs(paruq.mean - fit.nonlin) < 1.5e-2)
 # ======================================================================
+
+import pytest
+# ======================================================================
+def test_memory_limit():
+    "Check that the memory limit works for too large analyses"
+    # Request one million samples, approx. 80GB
+    yfit = np.linspace(0,1,int(1e4))    
+    yexp = yfit + whitegaussnoise(yfit,0.01)
+    def bootfcn(y):
+        return y
+    with pytest.raises(MemoryError):
+        paruq = bootstrap_analysis(bootfcn,yexp,yfit,1e6)
+# ======================================================================

test/test_dipolarkernel.py

@@ -472,3 +472,15 @@ def test_docstring():
     "Check that the docstring includes all variables and keywords."
     assert_docstring(dipolarkernel)
 # ======================================================================
+
+import pytest 
+
+def test_memory_limit():
+# ======================================================================
+    "Check that the memory limiter works"
+    # Construct 12GB kernel (should fit in RAM)
+    t = np.linspace(0,5,int(3e4))
+    r = np.linspace(1,6,int(5e4))
+    with pytest.raises(MemoryError):
+        K = dipolarkernel(t,r)
+# ======================================================================