diff --git a/src/quvac/parallel.py b/src/quvac/parallel.py index c6c1798..e103a50 100644 --- a/src/quvac/parallel.py +++ b/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") diff --git a/src/quvac/postprocess.py b/src/quvac/postprocess.py index d8489b5..1207977 100644 --- a/src/quvac/postprocess.py +++ b/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, ) diff --git a/src/quvac/utils.py b/src/quvac/utils.py index 163355f..d14c608 100644 --- a/src/quvac/utils.py +++ b/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"