add memory buffer keyword

src/quvac/parallel.py

@@ -119,7 +119,7 @@ def setup_job_executor_from_params(cluster_params, save_path,
     return executor
 
 
-def submit_jobs_with_memory_estimation(executor, ini_files):
+def submit_jobs_with_memory_estimation(executor, ini_files, memory_buffer):
     """
     Submit jobs for a list of initialization files estimating memory usage
     for each of them.
@@ -138,7 +138,7 @@ def submit_jobs_with_memory_estimation(executor, ini_files):
     """
     jobs = []
     for ini_file in ini_files:
-        memory = estimate_memory_usage(ini_file)
+        memory = estimate_memory_usage(ini_file, memory_buffer)
         executor.update_parameters(slurm_mem=memory)
         job = executor.submit(quvac_simulation, ini_file)
         jobs.append(job)
@@ -168,12 +168,14 @@ def run_simulations_with_job_executor(ini_files, cluster_params, save_path,
     estimate_memory_usage = cluster_params.get("estimate_memory_usage", False)
     _logger.info("MILESTONE: Submitting jobs...")
     if estimate_memory_usage:
-        jobs = submit_jobs_with_memory_estimation(executor, ini_files)
+        memory_buffer = cluster_params.get("memory_buffer", 10)
+        jobs = submit_jobs_with_memory_estimation(executor, ini_files, memory_buffer)
     else:
         jobs = executor.map_array(quvac_simulation, ini_files)
     _logger.info("MILESTONE: Jobs submitted, waiting for results...")
 
     # Wait till all jobs end
-    _ = [job.result() for job in jobs]
+    for job in submitit.helpers.as_completed(jobs):
+        job.result()
     _logger.info("MILESTONE: Jobs are finished")
 

src/quvac/postprocess.py

@@ -388,29 +388,29 @@ def get_spectra_from_Stokes(P0, P1, P2, P3, basis="linear"):
 
 
 def get_photon_spectrum_from_a12(a1, a2, k):
-        """
-        Calculates the photon spectrum from spectral coefficients a1 and a2.
+    """
+    Calculates the photon spectrum from spectral coefficients a1 and a2.
 
-        Parameters
-        ----------
-        a1 : numpy.ndarray
-            Spectral coefficient a1.
-        a2 : numpy.ndarray
-            Spectral coefficient a2.
-        k : numpy.ndarray
-            Wavevector
+    Parameters
+    ----------
+    a1 : numpy.ndarray
+        Spectral coefficient a1.
+    a2 : numpy.ndarray
+        Spectral coefficient a2.
+    k : numpy.ndarray
+        Wavevector
 
-        Returns
-        -------
-        numpy.ndarray
-            Photon spectrum.
-        """
-        prefactor = 0.5 * epsilon_0 * c / hbar
-        N_xyz = ne.evaluate(
-            "prefactor * (a1.real**2 + a1.imag**2 + a2.real**2 + a2.imag**2) / k"
-        )
-        N_xyz[0, 0, 0] = 0.0
-        return np.fft.fftshift(N_xyz / (2*pi)**3)
+    Returns
+    -------
+    numpy.ndarray
+        Photon spectrum.
+    """
+    prefactor = 0.5 * epsilon_0 * c / hbar
+    N_xyz = ne.evaluate(
+        "prefactor * (a1.real**2 + a1.imag**2 + a2.real**2 + a2.imag**2) / k"
+    )
+    N_xyz[0, 0, 0] = 0.0
+    return np.fft.fftshift(N_xyz / (2*pi)**3)
 
 
 class VacuumEmissionAnalyzer:
@@ -703,7 +703,7 @@ def get_background(self, discernibility="angular", bgr_idx=None,
         _, bgr_N_sph = cartesian_to_spherical_array(
             bgr_N_xyz,
             self.grid_xyz,
-            spherical_grid=self.__dict__.get('spherical_grid', None),
+            spherical_grid=getattr(self, 'spherical_grid', None),
             **interp_kwargs,
         )
 

src/quvac/utils.py

@@ -245,22 +245,20 @@ def size_to_Gb(size):
     return size*8 / 1024**3
 
 
-def estimate_max_required_memory(size):
+def estimate_max_required_memory(size, memory_buffer=10):
     """
     Estimate max requred memory based on the grid size.
     """
     # this value is estimated by running the simulation with different grid sizes,
     # looking at the max used memory and fitting a line to the dependency
     # max memory vs grid size
-    MEMORY_SCALING = 56
+    MEMORY_SCALING = 52
     estimated_mem = size_to_Gb(math.prod(size))*MEMORY_SCALING
 
-    # memory buffer just in case
-    SAFE_BUFFER = 10
-    return int(np.ceil(estimated_mem + SAFE_BUFFER))
+    return int(np.ceil(estimated_mem + memory_buffer))
 
 
-def estimate_memory_usage(ini_file):
+def estimate_memory_usage(ini_file, memory_buffer=10):
     """
     Estimate potential memory usage for a given ini file.
 
@@ -280,7 +278,7 @@ def estimate_memory_usage(ini_file):
         ini_config.get("grid", None),
     )
 
-    required_memory = estimate_max_required_memory(grid_xyz.grid_shape)
+    required_memory = estimate_max_required_memory(grid_xyz.grid_shape, memory_buffer)
     return f"{required_memory}GB"