diff --git a/.github/workflows/sphinx.yml b/.github/workflows/sphinx.yml index c52301e..fcec085 100644 --- a/.github/workflows/sphinx.yml +++ b/.github/workflows/sphinx.yml @@ -10,6 +10,7 @@ jobs: cleanup: runs-on: ubuntu-latest permissions: write-all + if: github.event_name == 'push' && github.ref == 'refs/heads/main' steps: - name: Delete deployment uses: strumwolf/delete-deployment-environment@v2 @@ -21,6 +22,7 @@ jobs: build: runs-on: ubuntu-latest needs: cleanup + if: always() permissions: contents: write steps: diff --git a/docs/source/implementation.rst b/docs/source/implementation.rst index 21c72ee..06c2087 100644 --- a/docs/source/implementation.rst +++ b/docs/source/implementation.rst @@ -148,7 +148,7 @@ References ^^^^^^^^^^ .. [1] F. Karbstein. "Probing vacuum polarization effects with high-intensity lasers." - Particles 3.1 (2020): 39-61. + Particles 3.1 (2020): 39-61 `(article) `_. .. [2] A. Blinne, et al. "All-optical signatures of quantum vacuum nonlinearities - in generic laser fields." PRD 99.1 (2019): 016006. \ No newline at end of file + in generic laser fields." PRD 99.1 (2019): 016006 `(article) `_. \ No newline at end of file diff --git a/docs/source/input_file.rst b/docs/source/input_file.rst index e2a5320..77ed4e2 100644 --- a/docs/source/input_file.rst +++ b/docs/source/input_file.rst @@ -120,7 +120,7 @@ Keys for ``dynamic`` mode: ``integrator`` (optional) ---------------------- +------------------------- Keys: - ``type``: str ``vacuum_emission`` (calculate the total vacuum emission amplitude) or ``vacuum_emission_channels`` (calculate the amplitude linearized in the probe field) @@ -130,9 +130,13 @@ Keys: Indices of the probe field, by default [0]. - ``pump``: list of int Indices of the pump field, by default [1]. + - ``integration_method``: ``'trapezoid'`` or ``'simpson'`` + Quadrature rule for discretized time integral. + - ``load_integration_weights``: bool + Whether to load integration weights from a separate file (usefule for parallel simulation). ``performance`` (optional) ----------------------- +-------------------------- Keys: - ``precision``: str Numerical precision for calculations: ``float32`` or (by default) ``float64``. @@ -146,10 +150,13 @@ Keys: Number of timesteps for a test run, by default 5. - ``use_wisdom``: bool Whether to use existing wisdom file for ``pyfftw`` planning. + - ``pyfftw_flag``: str + How much to plan the optimal execution of FFT with ``pyfftw``. One of + ``'FFTW_ESTIMATE'``, ``'FFTW_MEASURE'``, ``'FFTW_PATIENT'`` and ``'FFTW_EXHAUSTIVE'``. ``postprocessing`` (optional) -------------------------- +----------------------------- This section is relevant only when ``mode`` is ``postprocess`` or ``simulation_postprocess``. Relevant keys for the polarization-insensitive signals: - ``calculate_xyz_background`` : bool, optional Whether to calculate the background spectra on Cartesian grid, @@ -180,7 +187,7 @@ Relevant keys for the polarization-sensitive signals: Whether to calculate Stokes parameters, by default False. ``cluster_params`` (for ``quvac-simulation-parallel``) --------------------------------------------------- +------------------------------------------------------ Keys: - ``cluster_type``: str, Where perform calculations, ``local`` or ``slurm``. @@ -194,7 +201,7 @@ Keys: ``quvac.config.DEFAULT_SLURM_PARAMS``. ``gridscan`` (for ``quvac-gridscan``) ----------------------------------- +-------------------------------------- Keys: - ``create_grids``: bool Flag to create grids given [start, end, n_steps]. @@ -207,9 +214,11 @@ Keys: Maximal number of submitted jobs in parallel. - ``sbatch_params``: dict Submission parameters for a single job. + - ``estimate_memory_usage``: bool + Whether to estimame memory usage for each job separately. ``optimization`` (for ``quvac-optimization``) ------------------------------------------ +--------------------------------------------- Keys: - ``name``: str Optimization name. diff --git a/pyproject.toml b/pyproject.toml index 17565df..c1c20f6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,7 @@ quvac-optimization = "quvac.cluster.optimization:main_optimization" [project.optional-dependencies] optimization = [ - "ax-platform", + "ax-platform>=1.2.3", ] plot = [ "matplotlib>=3.10.1", @@ -48,6 +48,9 @@ plot = [ test = [ "pytest>=8.3.5", "pytest-cov>=6.1.0", + "scalene>=2.1.4", + "memray>=1.19", + "pytest-memray", ] docs = [ "astroid==3.3.8", diff --git a/src/quvac/grid.py b/src/quvac/grid.py index d68be38..f4f663c 100644 --- a/src/quvac/grid.py +++ b/src/quvac/grid.py @@ -93,7 +93,7 @@ def get_k_grid(self): self._calculate_k_grid() def _calculate_k_grid(self): - self.e1, self.e2 = [np.zeros((3,) + self.grid_shape) for _ in range(2)] + self.e1, self.e2 = [np.zeros(self.vector_shape) for _ in range(2)] for i, ax in enumerate("xyz"): Nx, dx = self.grid_shape[i], self.dxyz[i] @@ -127,21 +127,8 @@ def _calculate_k_grid(self): ne.evaluate("where((kx==0) & (ky==0), 0, -ky / kperp)", out=self.e2[0]) ne.evaluate("where((kx==0) & (ky==0), 1, kx / kperp)", out=self.e2[1]) - # self.e2y = ne.evaluate("where((kx==0) & (ky==0), 2*(kz>0)-1, kx / kperp)") - # self.e2z = 0 self.e2x, self.e2y, self.e2z = self.e2 - # self.e1x = ne.evaluate( - # "where((kx==0) & (ky==0), 2*(kz>0)-1, kx * kz / (kperp*kabs))" - # ) - # self.e1y = ne.evaluate("where((kx==0) & (ky==0), 0, ky * kz / (kperp*kabs))") - # self.e1z = ne.evaluate("where((kx==0) & (ky==0), 0, -kperp / kabs)") - - # self.e2x = ne.evaluate("where((kx==0) & (ky==0), 0, -ky / kperp)") - # self.e2y = ne.evaluate("where((kx==0) & (ky==0), 1, kx / kperp)") - # # self.e2y = ne.evaluate("where((kx==0) & (ky==0), 2*(kz>0)-1, kx / kperp)") - # self.e2z = 0 - def get_ek(theta, phi): """ diff --git a/src/quvac/integrator/vacuum_emission.py b/src/quvac/integrator/vacuum_emission.py index 6372c21..1711324 100644 --- a/src/quvac/integrator/vacuum_emission.py +++ b/src/quvac/integrator/vacuum_emission.py @@ -23,10 +23,48 @@ from quvac import config from quvac.pyfftw_executor import setup_fftw_executor +from quvac.utils import free_memory BS = m_e**2 * c**2 / (hbar * e) # Schwinger magnetic field +def determine_integration_scheme(Nt, integration_method): + """ + Determine integration weights based on the integration scheme. + + Parameters + ---------- + Nt : int + Number of time points in a discretized time integral. + integration_method: 'trapezoid' or 'simpson' + Quadrature rule to use. + + Returns + ------- + np.ndarray of length Nt + Integration weights for a given time interval. + """ + integration_weights = np.ones(Nt) + match integration_method: + case "trapezoid": + integration_weights[0] = 0.5 + integration_weights[-1] = 0.5 + case "simpson": + number_of_intervals = Nt - 1 + integration_weights /= 3. + idx_even = [2*i for i in range(1, number_of_intervals//2-1)] + idx_odd = [2*i-1 for i in range(1, number_of_intervals//2)] + integration_weights[idx_even] *= 2 + integration_weights[idx_odd] *= 4 + case _: + err_msg = ( + "integration_method should be one of ['trapezoid', 'simpson'] but you " + f"passed {integration_method}" + ) + raise NotImplementedError(err_msg) + return integration_weights + + class VacuumEmission: """ Calculator of Vacuum Emission amplitude from given fields @@ -45,7 +83,6 @@ class VacuumEmission: channels : bool, optional Whether to calculate a particular channel in vacuum emission amplitude. Default is False. - """ def __init__(self, field, grid, fft_executor=None, nthreads=None, channels=False): @@ -69,8 +106,8 @@ def __init__(self, field, grid, fft_executor=None, nthreads=None, channels=False ] if not self.channels: - self.U1 = "(4*E*F + 7*B*G)" - self.U2 = "(4*B*F - 7*E*G)" + self.U1_expr = "(4*E*F + 7*B*G)" + self.U2_expr = "(4*B*F - 7*E*G)" else: self._define_channel_variables() @@ -94,8 +131,8 @@ def _define_channel_variables(self): np.zeros(self.grid_shape, dtype=config.FDTYPE) for _ in range(3) ] - self.U1 = "(4*(Ep*F + E*F_B_Bp) + 7*(Bp*G + B*(G_Ep_B + G_E_Bp)))" - self.U2 = "(4*(Bp*F + B*F_B_Bp) - 7*(Ep*G + E*(G_Ep_B + G_E_Bp)))" + self.U1_expr = "(4*(Ep*F + E*F_B_Bp) + 7*(Bp*G + B*(G_Ep_B + G_E_Bp)))" + self.U2_expr = "(4*(Bp*F + B*F_B_Bp) - 7*(Ep*G + E*(G_Ep_B + G_E_Bp)))" def _allocate_fields(self): """ @@ -117,12 +154,16 @@ def _allocate_result_arrays(self): self.U2_acc_x, self.U2_acc_y, self.U2_acc_z = self.U2_acc self.U_pairs = [ - (self.U1_acc, self.U1), - (self.U2_acc, self.U2), + (self.U1_acc, self.U1_expr), + (self.U2_acc, self.U2_expr), ] - self.prefactor = np.ones(self.grid_shape, dtype="complex128") - self.prefactor_step = np.zeros(self.grid_shape, dtype="complex128") + # During the main loop calculation these arrays serve as buffers for calculation + # of prefactor and prefactor_step. When the calculation is finished, this space + # is used to store the transition amplitudes. It seems that Python garbage + # collector doesn't manage to quickly clean the arrays that are not used. + self.prefactor = self.S1 = np.ones(self.grid_shape, dtype=config.CDTYPE) + self.prefactor_step = self.S2 = np.zeros(self.grid_shape, dtype=config.CDTYPE) self.U_dict = {"F": self.F, "G": self.G} @@ -152,6 +193,11 @@ def _free_resources(self): """ del self.E_out, self.B_out del self.fft_executor + del self.F, self.G + if self.channels: + del self.E_probe, self.B_probe + del self.F_B_Bp, self.G_Ep_B, self.G_E_Bp + free_memory() def calculate_one_time_step(self, t, weight=1): """ @@ -178,8 +224,8 @@ def calculate_one_time_step(self, t, weight=1): "Bx": Bx, "By": By, "Bz": Bz,}) # Evaluate F and G - ne.evaluate(self.F_expr, out=self.F) - ne.evaluate(self.G_expr, out=self.G) + ne.evaluate(self.F_expr, local_dict=self.U_dict, out=self.F) + ne.evaluate(self.G_expr, local_dict=self.U_dict, out=self.G) if self.channels: ne.evaluate(self.F_B_Bp_expr, out=self.F_B_Bp) @@ -197,13 +243,13 @@ def calculate_one_time_step(self, t, weight=1): out=self.prefactor) # Evaluate U1 and U2 expressions - for U_acc, U_expr in self.U_pairs: # noqa: B905 + for U_acc, U_expr in self.U_pairs: ne.evaluate(U_expr, global_dict=self.U_dict, out=self.fft_executor.tmp) self.fft_executor.forward_fftw.execute() - self.U_acc_dict.update({"U_acc": U_acc}) + self.U_acc_dict.update({"U_acc": U_acc, "weight": weight}) ne.evaluate( - "U_acc + U*prefactor", + "U_acc + U*prefactor*weight", local_dict=self.U_acc_dict, out=U_acc, ) @@ -224,7 +270,7 @@ def multiply_integration_result(self, t_grid): ne.evaluate("acc*prefactor", global_dict={"prefactor": self.prefactor}, out=acc) - def calculate_time_integral(self, t_grid, integration_method="trapezoid"): + def calculate_time_integral(self, t_grid, integration_weights): """ Calculate the time integral. """ @@ -241,18 +287,8 @@ def calculate_time_integral(self, t_grid, integration_method="trapezoid"): "exp(1j*kabs*c*dt)", local_dict=self.prefactor_dict, out=self.prefactor_step ) - if integration_method == "trapezoid": - # end_pts = (0, len(t_grid) - 1) - for _, t in enumerate(t_grid): - # weight = 0.5 if i in end_pts else 1. - weight = 1 - self.calculate_one_time_step(t, weight=weight) - else: - err_msg = ( - "integration_method should be one of ['trapezoid'] but you " - f"passed {integration_method}" - ) - raise NotImplementedError(err_msg) + for t,weight in zip(t_grid, integration_weights, strict=True): + self.calculate_one_time_step(t, weight=weight) # finish calculation self.multiply_integration_result(t_grid) @@ -263,25 +299,31 @@ def _calculate_S1_S2(self): """ dims = 1 / BS**3 * m_e**2 * c**3 / hbar**2 prefactor = -1j * np.sqrt(alpha * self.kabs) / (2 * pi) ** 1.5 / 45 * dims # noqa: F841 - self.S1 = ne.evaluate( + ne.evaluate( f"prefactor * ({self.I_11_expr} - {self.I_22_expr})", - global_dict=self.__dict__, - ).astype(config.CDTYPE) - self.S2 = ne.evaluate( + global_dict=self.__dict__, out=self.S1 + ) + ne.evaluate( f"prefactor * ({self.I_12_expr} + {self.I_21_expr})", - global_dict=self.__dict__, - ).astype(config.CDTYPE) + global_dict=self.__dict__, out=self.S2 + ) def calculate_amplitudes( - self, t_grid, integration_method="trapezoid", save_path=None + self, t_grid, integration_method="trapezoid", integration_weights=None, + save_path=None ): """ Calculate the vacuum emission amplitudes and save the result. """ self._allocate_resources() + if integration_weights is None: + integration_weights = determine_integration_scheme( + len(t_grid), integration_method, + ) + time_integral_start = time.perf_counter() - self.calculate_time_integral(t_grid, integration_method) + self.calculate_time_integral(t_grid, integration_weights) time_integral_end = time.perf_counter() time_integral = time_integral_end - time_integral_start diff --git a/src/quvac/parallel.py b/src/quvac/parallel.py index 56748da..c6c1798 100644 --- a/src/quvac/parallel.py +++ b/src/quvac/parallel.py @@ -8,6 +8,7 @@ from quvac.config import DEFAULT_SLURM_PARAMS from quvac.simulation import quvac_simulation +from quvac.utils import estimate_memory_usage _logger = logging.getLogger("simulation") @@ -118,8 +119,34 @@ def setup_job_executor_from_params(cluster_params, save_path, return executor +def submit_jobs_with_memory_estimation(executor, ini_files): + """ + Submit jobs for a list of initialization files estimating memory usage + for each of them. + + Parameters + ---------- + executor : submitit.AutoExecutor + Executor for running jobs. + ini_files : list of str + List of paths to the initialization files for each job. + + Returns + ------- + list of submitit jobs + Submitted jobs. + """ + jobs = [] + for ini_file in ini_files: + memory = estimate_memory_usage(ini_file) + executor.update_parameters(slurm_mem=memory) + job = executor.submit(quvac_simulation, ini_file) + jobs.append(job) + return jobs + + def run_simulations_with_job_executor(ini_files, cluster_params, save_path, - max_parallel_jobs_default=2): + max_parallel_jobs_default=2,): """ Run simulations for a list of ini files with job executor. @@ -138,8 +165,12 @@ def run_simulations_with_job_executor(ini_files, cluster_params, save_path, max_parallel_jobs_default) # Submit jobs + estimate_memory_usage = cluster_params.get("estimate_memory_usage", False) _logger.info("MILESTONE: Submitting jobs...") - jobs = executor.map_array(quvac_simulation, ini_files) + if estimate_memory_usage: + jobs = submit_jobs_with_memory_estimation(executor, ini_files) + else: + jobs = executor.map_array(quvac_simulation, ini_files) _logger.info("MILESTONE: Jobs submitted, waiting for results...") # Wait till all jobs end diff --git a/src/quvac/pyfftw_executor.py b/src/quvac/pyfftw_executor.py index f9f17e4..f695246 100644 --- a/src/quvac/pyfftw_executor.py +++ b/src/quvac/pyfftw_executor.py @@ -78,7 +78,7 @@ def _allocate_fft(self): def setup_fftw_executor(fft_executor, grid_shape, nthreads=None): """ Unified function to: - - (Optional )Set up FFTExecutor if it was not already. + - (Optional) Set up FFTExecutor if it was not already. - Allocate buffer arrays and FFT executors. """ if fft_executor is None: diff --git a/src/quvac/simulation.py b/src/quvac/simulation.py index 488f501..4373724 100755 --- a/src/quvac/simulation.py +++ b/src/quvac/simulation.py @@ -17,6 +17,7 @@ import time import numexpr as ne +import numpy as np import pyfftw from quvac import config @@ -33,7 +34,14 @@ ) from quvac.postprocess import VacuumEmissionAnalyzer from quvac.pyfftw_executor import FFTExecutor -from quvac.utils import get_maxrss, load_wisdom, read_yaml, save_wisdom, write_yaml +from quvac.utils import ( + free_memory, + get_maxrss, + load_wisdom, + read_yaml, + save_wisdom, + write_yaml, +) _logger = logging.getLogger("simulation") @@ -105,14 +113,16 @@ def get_filenames(ini_file, save_path, wisdom_file, mode=None): ini_config = read_yaml(ini_file) if mode is None: mode = ini_config.get('mode', 'simulation_postprocess') - files = {} - files['save_path'] = save_path - files['ini'] = ini_file - files['wisdom'] = wisdom_file - files['amplitudes'] = os.path.join(save_path, "amplitudes.npz") - files['spectra'] = os.path.join(save_path, "spectra.npz") - files['performance'] = os.path.join(save_path, f"{mode}_performance.yml") - files['logger'] = os.path.join(save_path, f"{mode}.log") + files = { + "save_path": save_path, + "ini": ini_file, + "wisdom": wisdom_file, + "amplitudes": os.path.join(save_path, "amplitudes.npz"), + "spectra": os.path.join(save_path, "spectra.npz"), + "performance": os.path.join(save_path, f"{mode}_performance.yml"), + "logger": os.path.join(save_path, f"{mode}.log"), + "integration_weights": os.path.join(save_path, "integration_weights.npy"), + } return files @@ -170,6 +180,9 @@ def set_pyfftw_flag(flag): def create_basic_logger(filename): + """ + Set up basic logger configuration. + """ logging.basicConfig( filename=filename, filemode="w", @@ -213,6 +226,14 @@ def run_simulation(ini_config, fields_params, files, timings, memory): if channels: probe_pump_idx = integrator_params.get("probe_pump_idx", None) + # Determine time quadrature rule + integration_method = integrator_params.get("integration_method", "trapezoid") + load_integration_weights = integrator_params.get("load_integration_weights", False) + if load_integration_weights: + integration_weights = np.load(files["integration_weights"]) + else: + integration_weights = None + # Set up number of threads nthreads = perf_params.get("nthreads", os.cpu_count()) ne.set_num_threads(nthreads) @@ -243,7 +264,6 @@ def run_simulation(ini_config, fields_params, files, timings, memory): grid_t = grid_t[:test_timesteps] _logger.info(f"Performing test run for {test_timesteps} timesteps\n") - # Field setup _logger.info( "Field constructor:\n" "====================================================" @@ -277,8 +297,14 @@ def run_simulation(ini_config, fields_params, files, timings, memory): vacem = VacuumEmission(field, grid_xyz, fft_executor, nthreads=pyfftw_threads, channels=channels) timings['vacem_setup'] = time.perf_counter() - timings['integral'] = vacem.calculate_amplitudes(grid_t, - save_path=files['amplitudes']) + + timings['integral'] = vacem.calculate_amplitudes( + grid_t, + integration_method=integration_method, + integration_weights=integration_weights, + save_path=files['amplitudes'] + ) + timings['amplitudes'] = time.perf_counter() memory['maxrss_amplitudes'] = get_maxrss() _logger.info("MILESTONE: Amplitudes are calculated") @@ -392,6 +418,7 @@ def quvac_simulation(ini_file, save_path=None, wisdom_file="wisdom/fftw-wisdom") if do_simulation: timings, memory = run_simulation(ini_config, fields_params, files, timings, memory) + free_memory() # Calculate spectra if do_postprocess: postprocess_simulation(ini_config, files, fields_params) diff --git a/src/quvac/simulation_parallel.py b/src/quvac/simulation_parallel.py index 0f7569b..8874e72 100755 --- a/src/quvac/simulation_parallel.py +++ b/src/quvac/simulation_parallel.py @@ -20,6 +20,7 @@ import numpy as np from quvac.grid import setup_grids +from quvac.integrator.vacuum_emission import determine_integration_scheme from quvac.log import get_parallel_performance_stats, log_time from quvac.parallel import run_simulations_with_job_executor from quvac.simulation import ( @@ -33,7 +34,15 @@ _logger = logging.getLogger("simulation") -def create_ini_files_for_parallel(ini_config, grid_xyz, grid_t, n_jobs, save_path): +def _setup_integration_scheme(ini_config): + integrator_params = ini_config.get("integrator", {}) + integrator_params["load_integration_weights"] = True + ini_config["integrator"] = integrator_params + return ini_config + + +def create_ini_files_for_parallel(ini_config, grid_xyz, grid_t, n_jobs, save_path, + integration_weights): """ Create initialization files for parallel jobs. @@ -49,6 +58,8 @@ def create_ini_files_for_parallel(ini_config, grid_xyz, grid_t, n_jobs, save_pat Number of parallel jobs. save_path : str Path to save the initialization files. + integration_weights : np.ndarray + Integration weights for a given temporal grid. Returns ------- @@ -61,6 +72,7 @@ def create_ini_files_for_parallel(ini_config, grid_xyz, grid_t, n_jobs, save_pat ini_job = deepcopy(ini_config) ini_job["mode"] = "simulation" ini_job["postprocess"] = {} + ini_job = _setup_integration_scheme(ini_job) box_xyz = [float(-ax[0] * 2) for ax in grid_xyz.grid] Nxyz = [int(N) for N in grid_xyz.grid_shape] @@ -80,9 +92,16 @@ def create_ini_files_for_parallel(ini_config, grid_xyz, grid_t, n_jobs, save_pat "Nt": Nt, } ini_job["grid"].update(grid_params_job) - ini_path_job = os.path.join(save_path, f"job_{str(idx).zfill(2)}", "ini.yml") + job_folder = os.path.join(save_path, f"job_{str(idx).zfill(2)}") + ini_path_job = os.path.join(job_folder, "ini.yml") Path(os.path.dirname(ini_path_job)).mkdir(parents=True, exist_ok=True) write_yaml(ini_path_job, ini_job) + + # write integration weights to a file + weights_for_job = integration_weights[idx_start:idx_end+1] + weights_path = os.path.join(job_folder, "integration_weights.npy") + np.save(weights_path, weights_for_job) + ini_files.append(ini_path_job) return ini_files @@ -113,6 +132,13 @@ def collect_results(ini_files, amplitudes_file): np.savez(amplitudes_file, **amplitude_total) +def configure_integration_weights(ini_config, Nt): + integrator_params = ini_config.get("integrator", {}) + integration_method = integrator_params.get("integration_method", "trapezoid") + integration_weights = determine_integration_scheme(Nt, integration_method) + return integration_weights + + def quvac_simulation_parallel( ini_file, save_path=None, wisdom_file="wisdom/fftw-wisdom" ): @@ -163,8 +189,12 @@ def quvac_simulation_parallel( # Get grids grid_xyz, grid_t = setup_grids(fields_params, grid_params) + # Configure integration weights + integration_weights = configure_integration_weights(ini_config, len(grid_t)) + ini_files = create_ini_files_for_parallel( - ini_config, grid_xyz, grid_t, number_of_time_intervals, files['save_path'] + ini_config, grid_xyz, grid_t, number_of_time_intervals, files['save_path'], + integration_weights, ) run_simulations_with_job_executor( diff --git a/src/quvac/utils.py b/src/quvac/utils.py index 652e456..163355f 100644 --- a/src/quvac/utils.py +++ b/src/quvac/utils.py @@ -2,19 +2,24 @@ Useful generic utilities. """ +import gc import importlib import inspect +import math import os from pathlib import Path import pkgutil import platform import resource import shutil +import sys import numpy as np import pyfftw import yaml +from quvac.grid import setup_grids + def read_yaml(yaml_file): """ @@ -215,4 +220,72 @@ def find_classes_in_package(package_name): def round_to_n(x, n): + """ + Round up to n significant digits. + + Parameters + ---------- + x : int or float + Number to round up. + n : int + Number of digits to round up to. + + Returns + ------- + int or float: + Rounded number. + """ return round(x, -int(np.floor(np.sign(x) * np.log10(abs(x)))) + n) + + +def size_to_Gb(size): + """ + Convert the size of float64 array to GBs. + """ + return size*8 / 1024**3 + + +def estimate_max_required_memory(size): + """ + 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 + 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)) + + +def estimate_memory_usage(ini_file): + """ + Estimate potential memory usage for a given ini file. + + Parameters + ---------- + ini_file: str + Path to the initialization file + + Returns + ------- + str + Required memory in format 'GB'. + """ + ini_config = read_yaml(ini_file) + grid_xyz, _ = setup_grids( + ini_config.get("fields", None), + ini_config.get("grid", None), + ) + + required_memory = estimate_max_required_memory(grid_xyz.grid_shape) + return f"{required_memory}GB" + + +def free_memory(): + gc.collect() + if sys.platform == "linux": + import ctypes + ctypes.CDLL("libc.so.6").malloc_trim(0) diff --git a/tests/bench_profiler.py b/tests/bench_profiler.py new file mode 100644 index 0000000..e8aaa82 --- /dev/null +++ b/tests/bench_profiler.py @@ -0,0 +1,58 @@ +""" +Test simulation is run and profiler reports are automatically generated. + +Currently two profilers are used: +1. `scalene` to study the time performance with line-by-line timings. +2. `memray` to study memory consumption. +""" + +import os +from pathlib import Path + +import pytest + +from quvac.utils import read_yaml, write_yaml +from tests.config_for_tests import PROFILER_CONFIG_PATH, SIMULATION_SCRIPT + + +@pytest.mark.benchmark +def run_scalene_and_memray(path, ini_data): + current_dir = os.path.abspath(os.getcwd()) + + Path(path).mkdir(parents=True, exist_ok=True) + + ini_file = os.path.join(path, "ini.yml") + write_yaml(ini_file, ini_data) + + profile_folder = os.path.join(path, "profiler_info") + Path(profile_folder).mkdir(parents=True, exist_ok=True) + scalene_file = "report-scalene.json" + scalene_path = os.path.join(profile_folder, scalene_file) + memray_file = os.path.join(profile_folder, "report-memray.bin") + + if os.path.isfile(scalene_file): + os.remove(scalene_file) + if os.path.isfile(memray_file): + os.remove(memray_file) + + # Run scalene profiler + status = os.system(f"scalene run -o {scalene_path} --cpu-only " + f"{SIMULATION_SCRIPT} --input {ini_file}") + assert status == 0, "Scalene execution did not finish successfully" + os.system(f"cd {profile_folder} && scalene view --standalone {scalene_file}") + os.system(f"cd {current_dir}") + + # Run memray profiler + status = os.system(f"memray run -o {memray_file} " + f"{SIMULATION_SCRIPT} --input {ini_file}") + assert status == 0, "Memray execution did not finish successfully" + os.system(f"memray flamegraph {memray_file}") + + +@pytest.mark.benchmark +def test_scalene_and_memray(): + # Load default simulation parameters + ini_data = read_yaml(PROFILER_CONFIG_PATH) + + path = "data/profiler" + run_scalene_and_memray(path, ini_data) \ No newline at end of file diff --git a/tests/config_for_tests.py b/tests/config_for_tests.py index c6168f3..adb9444 100644 --- a/tests/config_for_tests.py +++ b/tests/config_for_tests.py @@ -1,5 +1,6 @@ DEFAULT_CONFIG_PATH = "tests/default_config.yml" BENCHMARK_CONFIG_PATH = "tests/benchmark_config.yml" +PROFILER_CONFIG_PATH = "tests/profiler_config.yml" SIMULATION_SCRIPT = "src/quvac/simulation.py" GRIDSCAN_SCRIPT = "src/quvac/cluster/gridscan.py" OPTIMIZATION_SCRIPT = "src/quvac/cluster/optimization.py" diff --git a/tests/test_simulation.py b/tests/test_simulation.py index fb3b9b4..2e06cf2 100644 --- a/tests/test_simulation.py +++ b/tests/test_simulation.py @@ -71,6 +71,8 @@ def test_precision(tmp_path): quvac_simulation(ini_file) + + ########################################################################################## # POSTPROCESS ########################################################################################## @@ -164,6 +166,24 @@ def test_mix_bg_and_signal(tmp_path): quvac_simulation(ini_file) +def test_integration_methods(tmp_path): + ini_data = read_yaml(DEFAULT_CONFIG_PATH) + integration_schemes = ["trapezoid", "simpson"] + + N_totals = [] + for integration_scheme in integration_schemes: + ini_data["integrator"]["integration_scheme"] = integration_scheme + ini_file = save_ini(tmp_path, ini_data) + quvac_simulation(ini_file) + + folder = os.path.dirname(ini_file) + data = np.load(os.path.join(folder, "spectra_total.npz")) + N_totals.append(data["N_total"]) + + err_msg = "Different integration schemes do not give the same result" + assert np.isclose(N_totals[0], N_totals[1]), err_msg + + ########################################################################################## # SIMULATION PARALLEL ########################################################################################## @@ -185,7 +205,29 @@ def test_parallel_simulation(tmp_path): data = np.load(os.path.join(folder, "spectra_total.npz")) N_total_parallel = data["N_total"] - assert np.isclose(N_total, N_total_parallel), "Sequential and parallel results " - "should be the same." + err_msg = "Sequential and parallel results should be the same." + assert np.isclose(N_total, N_total_parallel), err_msg + + +def test_parallel_integration_schemes(tmp_path): + integration_scheme = "simpson" + # run usual simulation + ini_data = read_yaml(DEFAULT_CONFIG_PATH) + ini_data["integrator"]["integration_scheme"] = integration_scheme + ini_file = save_ini(tmp_path, ini_data) + quvac_simulation(ini_file) + + folder = os.path.dirname(ini_file) + data = np.load(os.path.join(folder, "spectra_total.npz")) + N_total = data["N_total"] + + ini_data["performance"]["nthreads"] = 2 + ini_file = save_ini(tmp_path, ini_data) + quvac_simulation_parallel(ini_file) + data = np.load(os.path.join(folder, "spectra_total.npz")) + N_total_parallel = data["N_total"] + err_msg = (f"Sequential and parallel results for {integration_scheme} " + "should be the same.") + assert np.isclose(N_total, N_total_parallel), err_msg