From dd96aca675359a51acc8d39412ac066aea3a17db Mon Sep 17 00:00:00 2001 From: maxbalrog Date: Thu, 5 Mar 2026 16:28:17 +0100 Subject: [PATCH 01/10] remove field allocation at each time step --- src/quvac/field/external_field.py | 21 ++++++++++++++++++++- src/quvac/integrator/vacuum_emission.py | 3 ++- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/src/quvac/field/external_field.py b/src/quvac/field/external_field.py index 8bcef64..554eb4d 100644 --- a/src/quvac/field/external_field.py +++ b/src/quvac/field/external_field.py @@ -43,6 +43,11 @@ def __init__(self, fields_params, grid, nthreads=None): self.nthreads = nthreads if nthreads else os.cpu_count() + # Reallocation of E_out and B_out fields at each time step was removed. + # If Maxwell fields are present in the setup, the zeroing out fields is + # automatic + self.need_to_zero_fields = True + maxwell_params = [ params for params in fields_params @@ -54,7 +59,12 @@ def __init__(self, fields_params, grid, nthreads=None): if not params["field_type"].endswith("maxwell") ] if maxwell_params: - new_params.append(maxwell_params) + # MaxwellField directly overwrites given E_out and B_out arrays. + # If it's used together with other fields (analytic profiles) then + # it needs to go first so total external field is correct. + new_params.insert(0, maxwell_params) + # MaxwellField takes care of zeroing out. + self.need_to_zero_fields = False _logger.info( f"{self.__class__.__name__}\n" @@ -97,10 +107,19 @@ def setup_field(self, field_params): self.fields.append(field) _logger.info(f"Base class: {field.__class__.__name__}") + def maybe_zero_out_field_buffers(self, E_out, B_out): + output_given = (E_out is not None) and (B_out is not None) + if output_given and self.need_to_zero_fields: + for idx in range(3): + E_out[idx][:] = 0. + B_out[idx][:] = 0. + return E_out, B_out + def calculate_field(self, t, E_out=None, B_out=None): """ Calculates the electric and magnetic fields at a given time step. """ + E_out, B_out = self.maybe_zero_out_field_buffers(E_out, B_out) for field in self.fields: E_out, B_out = field.calculate_field(t, E_out=E_out, B_out=B_out) return E_out, B_out diff --git a/src/quvac/integrator/vacuum_emission.py b/src/quvac/integrator/vacuum_emission.py index 5a94b36..1c87649 100644 --- a/src/quvac/integrator/vacuum_emission.py +++ b/src/quvac/integrator/vacuum_emission.py @@ -162,7 +162,7 @@ def calculate_one_time_step(self, t, weight=1): Calculate the field and U terms (integrand) for one time step. """ # Calculate fields - self._allocate_fields() + # self._allocate_fields() if not self.channels: self.field.calculate_field(t, E_out=self.E_out, B_out=self.B_out) @@ -244,6 +244,7 @@ def calculate_amplitudes( # Allocate resources self._allocate_result_arrays() self._allocate_fft() + self._allocate_fields() time_integral_start = time.perf_counter() self.calculate_time_integral(t_grid, integration_method) From 7ca8483140a8db6fc028ecf961bfb8ba21ec2601 Mon Sep 17 00:00:00 2001 From: maxbalrog Date: Fri, 6 Mar 2026 13:20:18 +0100 Subject: [PATCH 02/10] add optional FFTW_FLAG parameter --- src/quvac/config.py | 2 ++ src/quvac/field/abc.py | 2 +- src/quvac/field/maxwell.py | 6 ++--- src/quvac/integrator/vacuum_emission.py | 29 ++++++++++--------------- src/quvac/simulation.py | 20 ++++++++++++++++- 5 files changed, 37 insertions(+), 22 deletions(-) diff --git a/src/quvac/config.py b/src/quvac/config.py index e30b149..b2fc201 100644 --- a/src/quvac/config.py +++ b/src/quvac/config.py @@ -17,3 +17,5 @@ FDTYPE = "float64" CDTYPE = "complex128" + +FFTW_FLAG = "FFTW_MEASURE" diff --git a/src/quvac/field/abc.py b/src/quvac/field/abc.py index 3b0191c..ff394ab 100644 --- a/src/quvac/field/abc.py +++ b/src/quvac/field/abc.py @@ -241,7 +241,7 @@ def _allocate_fft(self): a, axes=(0, 1, 2), direction="FFTW_FORWARD", - flags=("FFTW_MEASURE",), + flags=(config.FFTW_FLAG,), threads=1, ) for a in self.Ef diff --git a/src/quvac/field/maxwell.py b/src/quvac/field/maxwell.py index 3e310b7..c0003db 100644 --- a/src/quvac/field/maxwell.py +++ b/src/quvac/field/maxwell.py @@ -91,7 +91,7 @@ def _allocate_tmp(self): self.tmp, axes=(0, 1, 2), direction="FFTW_BACKWARD", - flags=("FFTW_MEASURE",), + flags=(config.FFTW_FLAG,), threads=self.nthreads, ) @@ -122,9 +122,9 @@ def calculate_field(self, t, E_out=None, B_out=None): ne.evaluate(self.EB_expr[idx], local_dict=self.EB_dict, out=self.tmp) self.EB_fftw.execute() if idx < 3: - E_out[idx][:] = self.tmp.astype(config.CDTYPE) + np.copyto(E_out[idx], self.tmp) else: - B_out[idx-3][:] = self.tmp.astype(config.CDTYPE) + np.copyto(B_out[idx-3], self.tmp) return E_out, B_out diff --git a/src/quvac/integrator/vacuum_emission.py b/src/quvac/integrator/vacuum_emission.py index 1c87649..2651822 100644 --- a/src/quvac/integrator/vacuum_emission.py +++ b/src/quvac/integrator/vacuum_emission.py @@ -76,7 +76,7 @@ def __init__(self, field, grid, nthreads=None, channels=False): for j in range(1, 3) } for key, val in self.I_ij.items(): - self.__dict__[f"I_{key}_expr"] = val + setattr(self, f"I_{key}_expr", val) def _define_channel_variables(self): """ @@ -126,21 +126,16 @@ def _allocate_fft(self): """ Allocate memory for FFT calculations. """ - self.tmp = [ - pyfftw.zeros_aligned(self.grid_shape, dtype=config.CDTYPE) for _ in range(3) - ] - # Add number of threads - self.tmp_fftw = [ - pyfftw.FFTW( - a, - a, + self.tmp = pyfftw.zeros_aligned(self.grid_shape, dtype=config.CDTYPE) + self.tmp_fftw = pyfftw.FFTW( + self.tmp, + self.tmp, axes=(0, 1, 2), direction="FFTW_FORWARD", - flags=("FFTW_MEASURE",), + flags=(config.FFTW_FLAG,), threads=self.nthreads, - ) - for a in self.tmp - ] + ) + self.prefactor = np.zeros(self.grid_shape, dtype="complex128") self.prefactor_dict = { @@ -202,20 +197,20 @@ def calculate_one_time_step(self, t, weight=1): for U_key, U_expr in zip(["U1_acc", "U2_acc"], [self.U1, self.U2]): # noqa: B905 for i, expr in enumerate(U_expr): U_acc = getattr(self, U_key)[i] - ne.evaluate(expr, global_dict=self.U_dict, out=self.tmp[i]) - self.tmp_fftw[i].execute() + ne.evaluate(expr, global_dict=self.U_dict, out=self.tmp) + self.tmp_fftw.execute() U_res = ne.evaluate( "U_acc + U*prefactor*dt*dV", global_dict={ "U_acc": U_acc, - "U": self.tmp[i], + "U": self.tmp, "prefactor": self.prefactor, "dt": self.dt, "dV": self.dV, }, ) - U_acc[:] = U_res.astype(config.CDTYPE) + np.copyto(U_acc, U_res) def calculate_time_integral(self, t_grid, integration_method="trapezoid"): """ diff --git a/src/quvac/simulation.py b/src/quvac/simulation.py index ae62b0e..f4c3b8c 100755 --- a/src/quvac/simulation.py +++ b/src/quvac/simulation.py @@ -155,6 +155,19 @@ def set_precision(precision): config.CDTYPE = "complex128" +def set_pyfftw_flag(flag): + """ + Set global planning flag for pyfftw. + + Parameters + ---------- + flag : str + How much to plan the transform. One of `FFTW_ESTIMATE`, `FFTW_MEASURE`, + `FFTW_PATIENT` and `FFTW_EXHAUSTIVE`. + """ + config.FFTW_FLAG = flag + + def create_basic_logger(filename): logging.basicConfig( filename=filename, @@ -353,11 +366,16 @@ def quvac_simulation(ini_file, save_path=None, wisdom_file="wisdom/fftw-wisdom") timings['start'] = time.perf_counter() memory = {'maxrss_amplitudes': 0} - # Set up global precision for calculations perf_params = ini_config.get("performance", {}) + # Set up global precision for calculations + # Note: since we use numexpr, internally it casts to + # float64 and complex128 anyway! precision = perf_params.get("precision", "float64") set_precision(precision) _logger.info(f"Using {precision} precision") + # Set pyfftw planning flag + pyfftw_flag = perf_params.get("pyfftw_flag", "FFTW_MEASURE") + set_pyfftw_flag(pyfftw_flag) fields_params = ini_config["fields"] if isinstance(fields_params, dict): From d94e4b28d15cfef440e21dd208af404aa030b11c Mon Sep 17 00:00:00 2001 From: maxbalrog Date: Fri, 6 Mar 2026 16:10:55 +0100 Subject: [PATCH 03/10] start unifying pyfftw interfaces in the code --- src/quvac/field/abc.py | 38 +++++++++++------- src/quvac/field/external_field.py | 15 +++++--- src/quvac/field/maxwell.py | 43 ++++++++++++--------- src/quvac/integrator/vacuum_emission.py | 36 +++++++++-------- src/quvac/postprocess.py | 15 ++++++-- src/quvac/pyfftw_executor.py | 51 +++++++++++++++++++++++++ src/quvac/simulation.py | 12 +++++- 7 files changed, 153 insertions(+), 57 deletions(-) create mode 100644 src/quvac/pyfftw_executor.py diff --git a/src/quvac/field/abc.py b/src/quvac/field/abc.py index ff394ab..c2c1b4e 100644 --- a/src/quvac/field/abc.py +++ b/src/quvac/field/abc.py @@ -16,6 +16,7 @@ from quvac import config from quvac.field.utils import get_field_energy, get_field_energy_kspace +from quvac.pyfftw_executor import setup_fftw_executor ANGLE_KEYS = [ "theta", @@ -234,18 +235,23 @@ def _allocate_fft(self): self.Ef = [ pyfftw.zeros_aligned(self.grid_shape, dtype="complex128") for _ in range(3) ] + self.fft_executor = setup_fftw_executor(self.fft_executor, self.grid_shape) + # if self.fft_executor is None: + # self.fft_executor = FFTExecutor(self.grid_shape) + # self.fft_executor.allocate_fft() + # pyfftw scheme - self.Ef_fftw = [ - pyfftw.FFTW( - a, - a, - axes=(0, 1, 2), - direction="FFTW_FORWARD", - flags=(config.FFTW_FLAG,), - threads=1, - ) - for a in self.Ef - ] + # self.Ef_fftw = [ + # pyfftw.FFTW( + # a, + # a, + # axes=(0, 1, 2), + # direction="FFTW_FORWARD", + # flags=(config.FFTW_FLAG,), + # threads=1, + # ) + # for a in self.Ef + # ] def _check_energy_kspace(self): # Fix energy @@ -267,7 +273,7 @@ def _check_energy_kspace(self): _logger.info(f' Energy after "correction": {W_corrected:.3f} J') - def get_a12(self, t0=None): + def get_a12(self, t0=None, fft_executor=None): """ Calculates the a1 and a2 coefficients at a given time step. @@ -289,11 +295,15 @@ def get_a12(self, t0=None): is corrected to the desired value. """ t0 = t0 if t0 is not None else self.t0 + self.fft_executor = fft_executor self._allocate_fft() self.calculate_field(t0, E_out=self.Ef, mode="complex") for idx in range(3): - self.Ef_fftw[idx].execute() + np.copyto(self.fft_executor.tmp, self.Ef[idx]) + self.fft_executor.forward_fftw.execute() + np.copyto(self.Ef[idx], self.fft_executor.tmp) + # self.Ef_fftw[idx].execute() # Calculate a1, a2 coefficients Efx, Efy, Efz = self.Ef @@ -307,7 +317,7 @@ def get_a12(self, t0=None): self._check_energy_kspace() - del self.Ef, self.Ef_fftw + del self.Ef return self.a1, self.a2 diff --git a/src/quvac/field/external_field.py b/src/quvac/field/external_field.py index 554eb4d..e5400ab 100644 --- a/src/quvac/field/external_field.py +++ b/src/quvac/field/external_field.py @@ -37,11 +37,12 @@ class ExternalField(Field): Number of threads to use for calculations. """ - def __init__(self, fields_params, grid, nthreads=None): + def __init__(self, fields_params, grid, fft_executor=None, nthreads=None): self.fields = [] self.grid_xyz = grid self.nthreads = nthreads if nthreads else os.cpu_count() + self.fft_executor = fft_executor # Reallocation of E_out and B_out fields at each time step was removed. # If Maxwell fields are present in the setup, the zeroing out fields is @@ -97,7 +98,8 @@ def setup_field(self, field_params): field = ANALYTIC_FIELDS[field_type](field_params, self.grid_xyz) elif field_type == "maxwell": field = MaxwellMultiple( - field_params, self.grid_xyz, nthreads=self.nthreads + field_params, self.grid_xyz, self.fft_executor, + nthreads=self.nthreads ) else: raise NotImplementedError( @@ -161,7 +163,8 @@ class ProbePumpField(Field): ExternalField object for the pump fields. """ - def __init__(self, fields_params, grid, probe_pump_idx=None, nthreads=None): + def __init__(self, fields_params, grid, fft_executor=None, probe_pump_idx=None, + nthreads=None): if not probe_pump_idx: probe_pump_idx = {"probe": [0], "pump": [1]} self.probe_pump_idx = probe_pump_idx @@ -171,8 +174,10 @@ def __init__(self, fields_params, grid, probe_pump_idx=None, nthreads=None): probe_params = [fields_params[idx] for idx in probe_pump_idx["probe"]] pump_params = [fields_params[idx] for idx in probe_pump_idx["pump"]] - self.probe_field = ExternalField(probe_params, grid, nthreads=nthreads) - self.pump_field = ExternalField(pump_params, grid, nthreads=nthreads) + self.probe_field = ExternalField(probe_params, grid, fft_executor, + nthreads=nthreads) + self.pump_field = ExternalField(pump_params, grid, fft_executor, + nthreads=nthreads) def calculate_field(self, t, E_probe=None, B_probe=None, E_pump=None, B_pump=None): """ diff --git a/src/quvac/field/maxwell.py b/src/quvac/field/maxwell.py index c0003db..5995118 100644 --- a/src/quvac/field/maxwell.py +++ b/src/quvac/field/maxwell.py @@ -14,6 +14,7 @@ from quvac import config from quvac.field import SPATIAL_MODEL_FIELDS from quvac.field.abc import Field +from quvac.pyfftw_executor import setup_fftw_executor _logger = logging.getLogger("simulation") @@ -36,11 +37,12 @@ class MaxwellField(Field): the number of CPU cores. """ - def __init__(self, grid, nthreads=None): + def __init__(self, grid, fft_executor=None, nthreads=None): self.grid_xyz = grid self.__dict__.update(self.grid_xyz.__dict__) self.nthreads = nthreads if nthreads else os.cpu_count() + self.fft_executor = fft_executor self.c = c self.norm_ifft = self.dVk / (2.0 * pi) ** 3 @@ -85,15 +87,16 @@ def _allocate_tmp(self): """ Allocate temporary memory for FFT calculations. """ - self.tmp = pyfftw.zeros_aligned(self.grid_shape, dtype="complex128") - self.EB_fftw = pyfftw.FFTW( - self.tmp, - self.tmp, - axes=(0, 1, 2), - direction="FFTW_BACKWARD", - flags=(config.FFTW_FLAG,), - threads=self.nthreads, - ) + # self.tmp = pyfftw.zeros_aligned(self.grid_shape, dtype="complex128") + # self.EB_fftw = pyfftw.FFTW( + # self.tmp, + # self.tmp, + # axes=(0, 1, 2), + # direction="FFTW_BACKWARD", + # flags=(config.FFTW_FLAG,), + # threads=self.nthreads, + # ) + self.fft_executor = setup_fftw_executor(self.fft_executor, self.grid_shape) def calculate_field(self, t, E_out=None, B_out=None): """ @@ -119,12 +122,17 @@ def calculate_field(self, t, E_out=None, B_out=None): # Calculate fourier of fields at time t and transform back to # spatial domain for idx in range(6): - ne.evaluate(self.EB_expr[idx], local_dict=self.EB_dict, out=self.tmp) - self.EB_fftw.execute() + np.copyto( + self.fft_executor.tmp, + ne.evaluate(self.EB_expr[idx], local_dict=self.EB_dict) + ) + # ne.evaluate(self.EB_expr[idx], local_dict=self.EB_dict, out=self.tmp) + # self.EB_fftw.execute() + self.fft_executor.backward_fftw.execute() if idx < 3: - np.copyto(E_out[idx], self.tmp) + np.copyto(E_out[idx], self.fft_executor.tmp) else: - np.copyto(B_out[idx-3], self.tmp) + np.copyto(B_out[idx-3], self.fft_executor.tmp) return E_out, B_out @@ -155,8 +163,9 @@ class MaxwellMultiple(MaxwellField): Initial time for the field. """ - def __init__(self, fields, grid, nthreads=None): - super().__init__(grid, nthreads) + def __init__(self, fields, grid, fft_executor=None, nthreads=None): + self.fft_executor = fft_executor + super().__init__(grid, fft_executor, nthreads) self.a1, self.a2 = [ pyfftw.zeros_aligned(self.grid_shape, dtype=config.CDTYPE) for _ in range(2) @@ -195,7 +204,7 @@ def get_a12_from_field(self, field_params): _logger.info(f" {field_type}: {cls.__name__}") ini_field = cls(field_params, self.grid_xyz) self.t0 = ini_field.t0 - a1, a2 = ini_field.get_a12(ini_field.t0) + a1, a2 = ini_field.get_a12(ini_field.t0, self.fft_executor) else: raise NotImplementedError( f"`{field_type}` is not implemented, use" diff --git a/src/quvac/integrator/vacuum_emission.py b/src/quvac/integrator/vacuum_emission.py index 2651822..f28f79b 100644 --- a/src/quvac/integrator/vacuum_emission.py +++ b/src/quvac/integrator/vacuum_emission.py @@ -19,10 +19,10 @@ import numexpr as ne import numpy as np -import pyfftw from scipy.constants import alpha, c, e, hbar, m_e, pi from quvac import config +from quvac.pyfftw_executor import setup_fftw_executor BS = m_e**2 * c**2 / (hbar * e) # Schwinger magnetic field @@ -46,7 +46,7 @@ class VacuumEmission: """ - def __init__(self, field, grid, nthreads=None, channels=False): + def __init__(self, field, grid, fft_executor, nthreads=None, channels=False): self.field = field self.grid_xyz = grid # Update local dict with variables from GridXYZ class @@ -56,6 +56,8 @@ def __init__(self, field, grid, nthreads=None, channels=False): self.c = c self.nthreads = nthreads if nthreads else os.cpu_count() + self.fft_executor = fft_executor + # Define symbolic expressions to evaluate later self.F_expr = "(Bx**2 + By**2 + Bz**2 - Ex**2 - Ey**2 - Ez**2)/2" self.G_expr = "-(Ex*Bx + Ey*By + Ez*Bz)" @@ -126,15 +128,18 @@ def _allocate_fft(self): """ Allocate memory for FFT calculations. """ - self.tmp = pyfftw.zeros_aligned(self.grid_shape, dtype=config.CDTYPE) - self.tmp_fftw = pyfftw.FFTW( - self.tmp, - self.tmp, - axes=(0, 1, 2), - direction="FFTW_FORWARD", - flags=(config.FFTW_FLAG,), - threads=self.nthreads, - ) + # self.tmp = pyfftw.zeros_aligned(self.grid_shape, dtype=config.CDTYPE) + # self.tmp_fftw = pyfftw.FFTW( + # self.tmp, + # self.tmp, + # axes=(0, 1, 2), + # direction="FFTW_FORWARD", + # flags=(config.FFTW_FLAG,), + # threads=self.nthreads, + # ) + # self.fft_executor.allocate_fft() + # self.tmp = self.fft_executor.tmp + self.fft_executor = setup_fftw_executor(self.fft_executor, self.grid_shape) self.prefactor = np.zeros(self.grid_shape, dtype="complex128") @@ -150,7 +155,8 @@ def _free_resources(self): Free allocated resources. """ del self.E_out, self.B_out - del self.tmp, self.tmp_fftw + # del self.tmp, self.tmp_fftw + del self.fft_executor def calculate_one_time_step(self, t, weight=1): """ @@ -197,14 +203,14 @@ def calculate_one_time_step(self, t, weight=1): for U_key, U_expr in zip(["U1_acc", "U2_acc"], [self.U1, self.U2]): # noqa: B905 for i, expr in enumerate(U_expr): U_acc = getattr(self, U_key)[i] - ne.evaluate(expr, global_dict=self.U_dict, out=self.tmp) - self.tmp_fftw.execute() + ne.evaluate(expr, global_dict=self.U_dict, out=self.fft_executor.tmp) + self.fft_executor.forward_fftw.execute() U_res = ne.evaluate( "U_acc + U*prefactor*dt*dV", global_dict={ "U_acc": U_acc, - "U": self.tmp, + "U": self.fft_executor.tmp, "prefactor": self.prefactor, "dt": self.dt, "dV": self.dV, diff --git a/src/quvac/postprocess.py b/src/quvac/postprocess.py index 25329b3..3e98428 100644 --- a/src/quvac/postprocess.py +++ b/src/quvac/postprocess.py @@ -29,6 +29,7 @@ from quvac.field.maxwell import MaxwellMultiple from quvac.grid import GridXYZ, get_pol_basis, setup_grids from quvac.log import sph_interp_warn +from quvac.pyfftw_executor import FFTExecutor from quvac.utils import read_yaml _logger = logging.getLogger("simulation") @@ -452,6 +453,9 @@ def __init__(self, fields_params, data_path, save_path=None): # Update local dict with variables from GridXYZ class self.__dict__.update(self.grid_xyz.__dict__) + # Set up FFT executor + self.fft_executor = FFTExecutor(self.grid_xyz.grid_shape) + for ax in "xyz": kx = getattr(self, f"k{ax}") setattr(self, f"k{ax}", np.fft.fftshift(kx)) @@ -480,7 +484,7 @@ def get_polarization_from_field(self): Tuple containing the polarization components (efx, efy, efz). """ perp_field_params = self.fields_params[self.perp_field_idx] - field = MaxwellMultiple([perp_field_params], self.grid_xyz) + field = MaxwellMultiple([perp_field_params], self.grid_xyz, self.fft_executor) a1, a2 = field.a1, field.a2 Ex = ne.evaluate("e1x*a1 + e2x*a2", global_dict=self.__dict__) Ey = ne.evaluate("e1y*a1 + e2y*a2", global_dict=self.__dict__) @@ -662,9 +666,11 @@ def get_background_xyz(self, add_to_cls_dict=True, bgr_idx=None): Background field spectra. """ if bgr_idx is not None: - bgr_field = MaxwellMultiple(self.fields_params[bgr_idx], self.grid_xyz) + bgr_field = MaxwellMultiple(self.fields_params[bgr_idx], self.grid_xyz, + self.fft_executor) else: - bgr_field = MaxwellMultiple(self.fields_params, self.grid_xyz) + bgr_field = MaxwellMultiple(self.fields_params, self.grid_xyz, + self.fft_executor) bgr_N_xyz = get_photon_spectrum_from_a12( bgr_field.a1, bgr_field.a2, self.kabs ) @@ -745,7 +751,8 @@ def get_signal_a12(self, add_signal_bg=False): """ self.a1_sig, self.a2_sig = transform_S12_to_a12(self.S1, self.S2, self.kabs) if add_signal_bg: - bgr_field = MaxwellMultiple(self.fields_params, self.grid_xyz) + bgr_field = MaxwellMultiple(self.fields_params, self.grid_xyz, + self.fft_executor) a1_bgr, a2_bgr = bgr_field.a1, bgr_field.a2 self.a1_mix, self.a2_mix = self.a1_sig + a1_bgr, self.a2_sig + a2_bgr self.a1_mix, self.a2_mix = [np.fft.fftshift(a) for a diff --git a/src/quvac/pyfftw_executor.py b/src/quvac/pyfftw_executor.py new file mode 100644 index 0000000..bc14f33 --- /dev/null +++ b/src/quvac/pyfftw_executor.py @@ -0,0 +1,51 @@ +''' +Pyfftw executor that is used in all 3D FFTs. + +At the moment, batched FFT transforms of 3 3D arrays are performed. +''' +import os + +import pyfftw + +from quvac import config + + +class FFTExecutor: + def __init__(self, grid_shape, nthreads=None, fft_axes=(0,1,2)): + self.grid_shape = grid_shape + self.nthreads = nthreads if nthreads else os.cpu_count() + self.fft_axes = fft_axes + + def allocate_fft(self): + if getattr(self, "tmp_shape", None) is None: + self._allocate_fft() + + def _allocate_fft(self): + # self.tmp_shape = (3,) + self.grid_shape + self.tmp_shape = self.grid_shape + self.tmp = pyfftw.zeros_aligned(self.tmp_shape, dtype=config.CDTYPE) + + self.forward_fftw = pyfftw.FFTW( + self.tmp, + self.tmp, + axes=self.fft_axes, + direction="FFTW_FORWARD", + flags=(config.FFTW_FLAG,), + threads=self.nthreads, + ) + + self.backward_fftw = pyfftw.FFTW( + self.tmp, + self.tmp, + axes=self.fft_axes, + direction="FFTW_BACKWARD", + flags=(config.FFTW_FLAG,), + threads=self.nthreads, + ) + + +def setup_fftw_executor(fft_executor, grid_shape): + if fft_executor is None: + fft_executor = FFTExecutor(grid_shape) + fft_executor.allocate_fft() + return fft_executor diff --git a/src/quvac/simulation.py b/src/quvac/simulation.py index f4c3b8c..5120c9c 100755 --- a/src/quvac/simulation.py +++ b/src/quvac/simulation.py @@ -32,6 +32,7 @@ log_time, ) 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 _logger = logging.getLogger("simulation") @@ -233,22 +234,28 @@ def run_simulation(ini_config, fields_params, files, timings, memory): _logger.info(grid_print) _logger.info("MILESTONE: Grids are created\n") + # Set up FFT executor + fft_executor = FFTExecutor(grid_xyz.grid_shape, pyfftw_threads) + # Shorten time grid for the test run if test_run: expected_timesteps = len(grid_t) grid_t = grid_t[:test_timesteps] _logger.info(f"Performing test run for {test_timesteps} timesteps\n") + # Field setup _logger.info( "Field constructor:\n" "====================================================" ) if not channels: - field = ExternalField(fields_params, grid_xyz, nthreads=pyfftw_threads) + field = ExternalField(fields_params, grid_xyz, fft_executor, + nthreads=pyfftw_threads) else: field = ProbePumpField( fields_params, grid_xyz, + fft_executor, probe_pump_idx=probe_pump_idx, nthreads=pyfftw_threads, ) @@ -267,7 +274,8 @@ def run_simulation(ini_config, fields_params, files, timings, memory): f" Pump idx: {probe_pump['pump']}" ) _logger.info(log_message) - vacem = VacuumEmission(field, grid_xyz, nthreads=pyfftw_threads, channels=channels) + 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']) From 9876cb8276a322bb89500658d0293ea2f79ea4f0 Mon Sep 17 00:00:00 2001 From: maxbalrog Date: Mon, 9 Mar 2026 10:49:48 +0100 Subject: [PATCH 04/10] start unitying all components of vectors into arrays --- src/quvac/field/abc.py | 34 +++++--- src/quvac/field/maxwell.py | 75 ++++++++++++------ src/quvac/grid.py | 35 ++++++-- src/quvac/integrator/vacuum_emission.py | 101 ++++++++++++++++-------- src/quvac/postprocess.py | 2 +- src/quvac/pyfftw_executor.py | 8 +- src/quvac/simulation.py | 2 +- 7 files changed, 175 insertions(+), 82 deletions(-) diff --git a/src/quvac/field/abc.py b/src/quvac/field/abc.py index c2c1b4e..12ef657 100644 --- a/src/quvac/field/abc.py +++ b/src/quvac/field/abc.py @@ -130,9 +130,11 @@ def get_EB_out(self, E_out, B_out, mode): """ dtype = np.float64 if mode == "real" else np.complex128 if E_out is None: - E_out = [np.zeros(self.grid_shape, dtype=dtype) for _ in range(3)] + # E_out = [np.zeros(self.grid_shape, dtype=dtype) for _ in range(3)] + E_out = np.zeros(self.vector_shape, dtype=dtype) if B_out is None: - B_out = [np.zeros(self.grid_shape, dtype=dtype) for _ in range(3)] + # B_out = [np.zeros(self.grid_shape, dtype=dtype) for _ in range(3)] + B_out = np.zeros(self.vector_shape, dtype=dtype) return E_out, B_out def rotate_fields_back(self, E_out, B_out, mode): @@ -232,10 +234,11 @@ def _allocate_fft(self): """ Allocates memory for FFT calculations. """ - self.Ef = [ - pyfftw.zeros_aligned(self.grid_shape, dtype="complex128") for _ in range(3) - ] - self.fft_executor = setup_fftw_executor(self.fft_executor, self.grid_shape) + # self.Ef = [ + # pyfftw.zeros_aligned(self.grid_shape, dtype="complex128") for _ in range(3) + # ] + self.Ef = pyfftw.zeros_aligned(self.vector_shape, dtype="complex128") + self.fft_executor = setup_fftw_executor(self.fft_executor, self.vector_shape) # if self.fft_executor is None: # self.fft_executor = FFTExecutor(self.grid_shape) # self.fft_executor.allocate_fft() @@ -299,10 +302,14 @@ def get_a12(self, t0=None, fft_executor=None): self._allocate_fft() self.calculate_field(t0, E_out=self.Ef, mode="complex") - for idx in range(3): - np.copyto(self.fft_executor.tmp, self.Ef[idx]) - self.fft_executor.forward_fftw.execute() - np.copyto(self.Ef[idx], self.fft_executor.tmp) + np.copyto(self.fft_executor.tmp, self.Ef) + self.fft_executor.forward_fftw.execute() + np.copyto(self.Ef, self.fft_executor.tmp) + + # for idx in range(3): + # np.copyto(self.fft_executor.tmp, self.Ef[idx]) + # self.fft_executor.forward_fftw.execute() + # np.copyto(self.Ef[idx], self.fft_executor.tmp) # self.Ef_fftw[idx].execute() # Calculate a1, a2 coefficients @@ -315,6 +322,13 @@ def get_a12(self, t0=None, fft_executor=None): "dV * (e2x*Efx + e2y*Efy + e2z*Efz)", global_dict=self.__dict__ ) + # self.a1 = ne.evaluate( + # "dV * (e1*Ef)", global_dict=self.__dict__ + # ) + # self.a2 = ne.evaluate( + # "dV * (e2*Ef)", global_dict=self.__dict__ + # ) + self._check_energy_kspace() del self.Ef diff --git a/src/quvac/field/maxwell.py b/src/quvac/field/maxwell.py index 5995118..939323f 100644 --- a/src/quvac/field/maxwell.py +++ b/src/quvac/field/maxwell.py @@ -48,9 +48,11 @@ def __init__(self, grid, fft_executor=None, nthreads=None): self.norm_ifft = self.dVk / (2.0 * pi) ** 3 # 1st list for E, 2nd list for B - self.EB_expr = [f"(e1{ax}*a1t + e2{ax}*a2t)" for ax in "xyz"] + [ - f"(e2{ax}*a1t - e1{ax}*a2t)" for ax in "xyz" - ] + # self.EB_expr = [f"(e1{ax}*a1t + e2{ax}*a2t)" for ax in "xyz"] + [ + # f"(e2{ax}*a1t - e1{ax}*a2t)" for ax in "xyz" + # ] + self.E_expr = "(e1*a1t + e2*a2t)" + self.B_expr = "(e2*a1t - e1*a2t)" self._allocate_tmp() @@ -72,13 +74,19 @@ def _allocate_ifft(self): "a2": self.a2, } + # self.EB_dict = { + # "e1x": self.e1x, + # "e1y": self.e1y, + # "e1z": self.e1z, + # "e2x": self.e2x, + # "e2y": self.e2y, + # "e2z": self.e2z, + # "a1t": self.a1t, + # "a2t": self.a2t, + # } self.EB_dict = { - "e1x": self.e1x, - "e1y": self.e1y, - "e1z": self.e1z, - "e2x": self.e2x, - "e2y": self.e2y, - "e2z": self.e2z, + "e1": self.e1, + "e2": self.e2, "a1t": self.a1t, "a2t": self.a2t, } @@ -96,15 +104,17 @@ def _allocate_tmp(self): # flags=(config.FFTW_FLAG,), # threads=self.nthreads, # ) - self.fft_executor = setup_fftw_executor(self.fft_executor, self.grid_shape) + self.fft_executor = setup_fftw_executor(self.fft_executor, self.vector_shape) def calculate_field(self, t, E_out=None, B_out=None): """ Calculates the electric and magnetic fields at a given time step. """ if E_out is None: - E_out = [np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3)] - B_out = [np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3)] + # E_out = [np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3)] + # B_out = [np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3)] + E_out = np.zeros(self.vector_shape, dtype=config.CDTYPE) + B_out = np.zeros(self.vector_shape, dtype=config.CDTYPE) # Calculate a1,a2 at time t self.a_dict.update({"t": t}) @@ -121,18 +131,35 @@ def calculate_field(self, t, E_out=None, B_out=None): # Calculate fourier of fields at time t and transform back to # spatial domain - for idx in range(6): - np.copyto( - self.fft_executor.tmp, - ne.evaluate(self.EB_expr[idx], local_dict=self.EB_dict) - ) - # ne.evaluate(self.EB_expr[idx], local_dict=self.EB_dict, out=self.tmp) - # self.EB_fftw.execute() - self.fft_executor.backward_fftw.execute() - if idx < 3: - np.copyto(E_out[idx], self.fft_executor.tmp) - else: - np.copyto(B_out[idx-3], self.fft_executor.tmp) + + # Transform electric field + np.copyto( + self.fft_executor.tmp, + ne.evaluate(self.E_expr, local_dict=self.EB_dict) + ) + self.fft_executor.backward_fftw.execute() + np.copyto(E_out, self.fft_executor.tmp) + + # Transform magnetic field + np.copyto( + self.fft_executor.tmp, + ne.evaluate(self.B_expr, local_dict=self.EB_dict) + ) + self.fft_executor.backward_fftw.execute() + np.copyto(B_out, self.fft_executor.tmp) + + # for idx in range(6): + # np.copyto( + # self.fft_executor.tmp, + # ne.evaluate(self.EB_expr[idx], local_dict=self.EB_dict) + # ) + # # ne.evaluate(self.EB_expr[idx], local_dict=self.EB_dict, out=self.tmp) + # # self.EB_fftw.execute() + # self.fft_executor.backward_fftw.execute() + # if idx < 3: + # np.copyto(E_out[idx], self.fft_executor.tmp) + # else: + # np.copyto(B_out[idx-3], self.fft_executor.tmp) return E_out, B_out diff --git a/src/quvac/grid.py b/src/quvac/grid.py index e7f7c06..d68be38 100644 --- a/src/quvac/grid.py +++ b/src/quvac/grid.py @@ -78,6 +78,7 @@ def __init__(self, grid): ) self.dxyz = tuple(ax[1] - ax[0] for ax in grid) self.dV = np.prod(self.dxyz) + self.vector_shape = (3,) + self.grid_shape self.kgrid = None def get_k_grid(self): @@ -92,6 +93,8 @@ 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)] + for i, ax in enumerate("xyz"): Nx, dx = self.grid_shape[i], self.dxyz[i] k = (2 * pi * np.fft.fftfreq(Nx, dx)).astype(config.FDTYPE) @@ -112,16 +115,32 @@ def _calculate_k_grid(self): kperp = ne.evaluate("sqrt(kx**2 + ky**2)") # noqa: F841 # Polarization vectors - self.e1x = ne.evaluate( - "where((kx==0) & (ky==0), 2*(kz>0)-1, kx * kz / (kperp*kabs))" + ne.evaluate( + "where((kx==0) & (ky==0), 2*(kz>0)-1, kx * kz / (kperp*kabs))", + out=self.e1[0] ) - 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)") + ne.evaluate("where((kx==0) & (ky==0), 0, ky * kz / (kperp*kabs))", + out=self.e1[1]) + ne.evaluate("where((kx==0) & (ky==0), 0, -kperp / kabs)", + out=self.e1[2]) + self.e1x, self.e1y, self.e1z = self.e1 + + 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.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 f28f79b..11bd984 100644 --- a/src/quvac/integrator/vacuum_emission.py +++ b/src/quvac/integrator/vacuum_emission.py @@ -67,8 +67,10 @@ def __init__(self, field, grid, fft_executor, nthreads=None, channels=False): ] if not self.channels: - self.U1 = [f"(4*E{ax}*F + 7*B{ax}*G)" for ax in "xyz"] - self.U2 = [f"(4*B{ax}*F - 7*E{ax}*G)" for ax in "xyz"] + # self.U1 = [f"(4*E{ax}*F + 7*B{ax}*G)" for ax in "xyz"] + # self.U2 = [f"(4*B{ax}*F - 7*E{ax}*G)" for ax in "xyz"] + self.U1 = "(4*E*F + 7*B*G)" + self.U2 = "(4*B*F - 7*E*G)" else: self._define_channel_variables() @@ -92,35 +94,44 @@ def _define_channel_variables(self): np.zeros(self.grid_shape, dtype=config.FDTYPE) for _ in range(3) ] - self.U1 = [ - f"(4*(Ep{ax}*F + E{ax}*F_B_Bp) + 7*(Bp{ax}*G + B{ax}*(G_Ep_B + G_E_Bp)))" - for ax in "xyz" - ] - self.U2 = [ - f"(4*(Bp{ax}*F + B{ax}*F_B_Bp) - 7*(Ep{ax}*G + E{ax}*(G_Ep_B + G_E_Bp)))" - for ax in "xyz" - ] + # self.U1 = [ + # f"(4*(Ep{ax}*F + E{ax}*F_B_Bp) + 7*(Bp{ax}*G + B{ax}*(G_Ep_B + G_E_Bp)))" + # for ax in "xyz" + # ] + # self.U2 = [ + # f"(4*(Bp{ax}*F + B{ax}*F_B_Bp) - 7*(Ep{ax}*G + E{ax}*(G_Ep_B + G_E_Bp)))" + # for ax in "xyz" + # ] + 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)))" def _allocate_fields(self): """ Allocate memory for field calculations. """ - self.E_out = [np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3)] - self.B_out = [np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3)] + # self.E_out = [np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3)] + # self.B_out = [np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3)] + # if self.channels: + # self.E_probe = [ + # np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3) + # ] + # self.B_probe = [ + # np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3) + # ] + self.E_out = np.zeros(self.vector_shape, dtype=config.CDTYPE) + self.B_out = np.zeros(self.vector_shape, dtype=config.CDTYPE) if self.channels: - self.E_probe = [ - np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3) - ] - self.B_probe = [ - np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3) - ] + self.E_probe = np.zeros(self.vector_shape, dtype=config.CDTYPE) + self.B_probe = np.zeros(self.vector_shape, dtype=config.CDTYPE) def _allocate_result_arrays(self): """ Allocate memory for result arrays. """ - self.U1_acc = [np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3)] - self.U2_acc = [np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3)] + # self.U1_acc = [np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3)] + # self.U2_acc = [np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3)] + self.U1_acc = np.zeros(self.vector_shape, dtype=config.CDTYPE) + self.U2_acc = np.zeros(self.vector_shape, dtype=config.CDTYPE) self.U1_acc_x, self.U1_acc_y, self.U1_acc_z = self.U1_acc self.U2_acc_x, self.U2_acc_y, self.U2_acc_z = self.U2_acc @@ -139,7 +150,7 @@ def _allocate_fft(self): # ) # self.fft_executor.allocate_fft() # self.tmp = self.fft_executor.tmp - self.fft_executor = setup_fftw_executor(self.fft_executor, self.grid_shape) + self.fft_executor = setup_fftw_executor(self.fft_executor, self.vector_shape) self.prefactor = np.zeros(self.grid_shape, dtype="complex128") @@ -175,12 +186,17 @@ def calculate_one_time_step(self, t, weight=1): E_pump=self.E_out, B_pump=self.B_out, ) - Epx, Epy, Epz = [E.real for E in self.E_probe] - Bpx, Bpy, Bpz = [B.real for B in self.B_probe] - - Ex, Ey, Ez = [E.real for E in self.E_out] - Bx, By, Bz = [B.real for B in self.B_out] - self.U_dict.update({"Ex": Ex, "Ey": Ey, "Ez": Ez, "Bx": Bx, "By": By, "Bz": Bz}) + # Epx, Epy, Epz = [E.real for E in self.E_probe] + # Bpx, Bpy, Bpz = [B.real for B in self.B_probe] + Ep = Epx, Epy, Epz = np.real(self.E_probe) + Bp = Bpx, Bpy, Bpz = np.real(self.B_probe) + + # Ex, Ey, Ez = [E.real for E in self.E_out] + # Bx, By, Bz = [B.real for B in self.B_out] + E = Ex, Ey, Ez = np.real(self.E_out) + B = Bx, By, Bz = np.real(self.B_out) + self.U_dict.update({"Ex": Ex, "Ey": Ey, "Ez": Ez, "Bx": Bx, "By": By, "Bz": Bz, + "E": E, "B": B}) ne.evaluate(self.F_expr, out=self.F) ne.evaluate(self.G_expr, out=self.G) @@ -193,7 +209,8 @@ def calculate_one_time_step(self, t, weight=1): "Bpx": Bpx, "Bpy": Bpy, "Bpz": Bpz, "F_B_Bp": self.F_B_Bp, "G_Ep_B": self.G_Ep_B, - "G_E_Bp": self.G_E_Bp,}) + "G_E_Bp": self.G_E_Bp, + "Ep": Ep, "Bp": Bp,}) # Evaluate U1 and U2 expressions self.prefactor_dict.update({"t": t}) @@ -201,12 +218,11 @@ def calculate_one_time_step(self, t, weight=1): "exp(1j*kabs*c*t)", local_dict=self.prefactor_dict, out=self.prefactor ) for U_key, U_expr in zip(["U1_acc", "U2_acc"], [self.U1, self.U2]): # noqa: B905 - for i, expr in enumerate(U_expr): - U_acc = getattr(self, U_key)[i] - ne.evaluate(expr, global_dict=self.U_dict, out=self.fft_executor.tmp) - self.fft_executor.forward_fftw.execute() + U_acc = getattr(self, U_key) + ne.evaluate(U_expr, global_dict=self.U_dict, out=self.fft_executor.tmp) + self.fft_executor.forward_fftw.execute() - U_res = ne.evaluate( + U_res = ne.evaluate( "U_acc + U*prefactor*dt*dV", global_dict={ "U_acc": U_acc, @@ -216,7 +232,24 @@ def calculate_one_time_step(self, t, weight=1): "dV": self.dV, }, ) - np.copyto(U_acc, U_res) + np.copyto(U_acc, U_res) + + # for i, expr in enumerate(U_expr): + # U_acc = getattr(self, U_key)[i] + # ne.evaluate(expr, global_dict=self.U_dict, out=self.fft_executor.tmp) + # self.fft_executor.forward_fftw.execute() + + # U_res = ne.evaluate( + # "U_acc + U*prefactor*dt*dV", + # global_dict={ + # "U_acc": U_acc, + # "U": self.fft_executor.tmp, + # "prefactor": self.prefactor, + # "dt": self.dt, + # "dV": self.dV, + # }, + # ) + # np.copyto(U_acc, U_res) def calculate_time_integral(self, t_grid, integration_method="trapezoid"): """ diff --git a/src/quvac/postprocess.py b/src/quvac/postprocess.py index 3e98428..d8489b5 100644 --- a/src/quvac/postprocess.py +++ b/src/quvac/postprocess.py @@ -454,7 +454,7 @@ def __init__(self, fields_params, data_path, save_path=None): self.__dict__.update(self.grid_xyz.__dict__) # Set up FFT executor - self.fft_executor = FFTExecutor(self.grid_xyz.grid_shape) + self.fft_executor = FFTExecutor(self.grid_xyz.vector_shape) for ax in "xyz": kx = getattr(self, f"k{ax}") diff --git a/src/quvac/pyfftw_executor.py b/src/quvac/pyfftw_executor.py index bc14f33..b4dabbe 100644 --- a/src/quvac/pyfftw_executor.py +++ b/src/quvac/pyfftw_executor.py @@ -11,18 +11,18 @@ class FFTExecutor: - def __init__(self, grid_shape, nthreads=None, fft_axes=(0,1,2)): - self.grid_shape = grid_shape + def __init__(self, tmp_shape, nthreads=None, fft_axes=(1,2,3)): + self.tmp_shape = tmp_shape self.nthreads = nthreads if nthreads else os.cpu_count() self.fft_axes = fft_axes def allocate_fft(self): - if getattr(self, "tmp_shape", None) is None: + if getattr(self, "tmp", None) is None: self._allocate_fft() def _allocate_fft(self): # self.tmp_shape = (3,) + self.grid_shape - self.tmp_shape = self.grid_shape + # self.tmp_shape = self.grid_shape self.tmp = pyfftw.zeros_aligned(self.tmp_shape, dtype=config.CDTYPE) self.forward_fftw = pyfftw.FFTW( diff --git a/src/quvac/simulation.py b/src/quvac/simulation.py index 5120c9c..000b589 100755 --- a/src/quvac/simulation.py +++ b/src/quvac/simulation.py @@ -235,7 +235,7 @@ def run_simulation(ini_config, fields_params, files, timings, memory): _logger.info("MILESTONE: Grids are created\n") # Set up FFT executor - fft_executor = FFTExecutor(grid_xyz.grid_shape, pyfftw_threads) + fft_executor = FFTExecutor(grid_xyz.vector_shape, pyfftw_threads) # Shorten time grid for the test run if test_run: From fdd4d2878da2ba53e318777681f2bab3a2c8417b Mon Sep 17 00:00:00 2001 From: maxbalrog Date: Mon, 9 Mar 2026 13:41:01 +0100 Subject: [PATCH 05/10] remove comments from field scripts --- src/quvac/field/abc.py | 42 ++------------------ src/quvac/field/external_field.py | 9 ++--- src/quvac/field/maxwell.py | 66 +++++-------------------------- src/quvac/field/utils.py | 2 +- src/quvac/pyfftw_executor.py | 30 +++++++------- src/quvac/simulation.py | 2 +- 6 files changed, 32 insertions(+), 119 deletions(-) diff --git a/src/quvac/field/abc.py b/src/quvac/field/abc.py index 12ef657..28ce9d1 100644 --- a/src/quvac/field/abc.py +++ b/src/quvac/field/abc.py @@ -11,7 +11,6 @@ import numexpr as ne import numpy as np import pyfftw -from scipy.constants import pi from scipy.spatial.transform import Rotation from quvac import config @@ -60,7 +59,7 @@ def __init__(self, field_params, grid): # self. for key, val in field_params.items(): if key in ANGLE_KEYS: - val *= pi / 180.0 + val = np.radians(val) setattr(self, key, val) def _check_energy(self): @@ -70,8 +69,8 @@ def _check_energy(self): E, B = self.calculate_field(t=0) W = get_field_energy(E, B, self.dV) - self.modify_energy = "W" in self.__dict__.keys() and not np.isclose(W, self.W, - rtol=1e-5) + self.modify_energy = getattr(self, "W", None) and not np.isclose(W, self.W, + rtol=1e-5) if self.modify_energy: self.W_correction = np.sqrt(self.W / W) if W > 0 else 0 self.W_num = self.W @@ -130,10 +129,8 @@ def get_EB_out(self, E_out, B_out, mode): """ dtype = np.float64 if mode == "real" else np.complex128 if E_out is None: - # E_out = [np.zeros(self.grid_shape, dtype=dtype) for _ in range(3)] E_out = np.zeros(self.vector_shape, dtype=dtype) if B_out is None: - # B_out = [np.zeros(self.grid_shape, dtype=dtype) for _ in range(3)] B_out = np.zeros(self.vector_shape, dtype=dtype) return E_out, B_out @@ -174,7 +171,6 @@ def convert_fields_to_real(self): for field in "Ex Ey Ez Bx By Bz".split(): setattr(self, field, np.real(getattr(self, field))) - @abstractmethod def calculate_field(self, t, E_out=None, B_out=None, **kwargs): """ @@ -234,27 +230,8 @@ def _allocate_fft(self): """ Allocates memory for FFT calculations. """ - # self.Ef = [ - # pyfftw.zeros_aligned(self.grid_shape, dtype="complex128") for _ in range(3) - # ] self.Ef = pyfftw.zeros_aligned(self.vector_shape, dtype="complex128") self.fft_executor = setup_fftw_executor(self.fft_executor, self.vector_shape) - # if self.fft_executor is None: - # self.fft_executor = FFTExecutor(self.grid_shape) - # self.fft_executor.allocate_fft() - - # pyfftw scheme - # self.Ef_fftw = [ - # pyfftw.FFTW( - # a, - # a, - # axes=(0, 1, 2), - # direction="FFTW_FORWARD", - # flags=(config.FFTW_FLAG,), - # threads=1, - # ) - # for a in self.Ef - # ] def _check_energy_kspace(self): # Fix energy @@ -306,12 +283,6 @@ def get_a12(self, t0=None, fft_executor=None): self.fft_executor.forward_fftw.execute() np.copyto(self.Ef, self.fft_executor.tmp) - # for idx in range(3): - # np.copyto(self.fft_executor.tmp, self.Ef[idx]) - # self.fft_executor.forward_fftw.execute() - # np.copyto(self.Ef[idx], self.fft_executor.tmp) - # self.Ef_fftw[idx].execute() - # Calculate a1, a2 coefficients Efx, Efy, Efz = self.Ef @@ -322,13 +293,6 @@ def get_a12(self, t0=None, fft_executor=None): "dV * (e2x*Efx + e2y*Efy + e2z*Efz)", global_dict=self.__dict__ ) - # self.a1 = ne.evaluate( - # "dV * (e1*Ef)", global_dict=self.__dict__ - # ) - # self.a2 = ne.evaluate( - # "dV * (e2*Ef)", global_dict=self.__dict__ - # ) - self._check_energy_kspace() del self.Ef diff --git a/src/quvac/field/external_field.py b/src/quvac/field/external_field.py index e5400ab..b3cd51a 100644 --- a/src/quvac/field/external_field.py +++ b/src/quvac/field/external_field.py @@ -4,7 +4,6 @@ """ import logging -import os from quvac.field import ANALYTIC_FIELDS from quvac.field.abc import Field @@ -41,7 +40,7 @@ def __init__(self, fields_params, grid, fft_executor=None, nthreads=None): self.fields = [] self.grid_xyz = grid - self.nthreads = nthreads if nthreads else os.cpu_count() + self.nthreads = nthreads self.fft_executor = fft_executor # Reallocation of E_out and B_out fields at each time step was removed. @@ -163,14 +162,12 @@ class ProbePumpField(Field): ExternalField object for the pump fields. """ - def __init__(self, fields_params, grid, fft_executor=None, probe_pump_idx=None, - nthreads=None): + def __init__(self, fields_params, grid, probe_pump_idx=None, + fft_executor=None, nthreads=None): if not probe_pump_idx: probe_pump_idx = {"probe": [0], "pump": [1]} self.probe_pump_idx = probe_pump_idx - self.nthreads = nthreads if nthreads else os.cpu_count() - probe_params = [fields_params[idx] for idx in probe_pump_idx["probe"]] pump_params = [fields_params[idx] for idx in probe_pump_idx["pump"]] diff --git a/src/quvac/field/maxwell.py b/src/quvac/field/maxwell.py index 939323f..aa670f5 100644 --- a/src/quvac/field/maxwell.py +++ b/src/quvac/field/maxwell.py @@ -4,7 +4,6 @@ """ import logging -import os import numexpr as ne import numpy as np @@ -41,16 +40,12 @@ def __init__(self, grid, fft_executor=None, nthreads=None): self.grid_xyz = grid self.__dict__.update(self.grid_xyz.__dict__) - self.nthreads = nthreads if nthreads else os.cpu_count() + self.nthreads = nthreads self.fft_executor = fft_executor self.c = c self.norm_ifft = self.dVk / (2.0 * pi) ** 3 - # 1st list for E, 2nd list for B - # self.EB_expr = [f"(e1{ax}*a1t + e2{ax}*a2t)" for ax in "xyz"] + [ - # f"(e2{ax}*a1t - e1{ax}*a2t)" for ax in "xyz" - # ] self.E_expr = "(e1*a1t + e2*a2t)" self.B_expr = "(e2*a1t - e1*a2t)" @@ -74,16 +69,6 @@ def _allocate_ifft(self): "a2": self.a2, } - # self.EB_dict = { - # "e1x": self.e1x, - # "e1y": self.e1y, - # "e1z": self.e1z, - # "e2x": self.e2x, - # "e2y": self.e2y, - # "e2z": self.e2z, - # "a1t": self.a1t, - # "a2t": self.a2t, - # } self.EB_dict = { "e1": self.e1, "e2": self.e2, @@ -95,24 +80,14 @@ def _allocate_tmp(self): """ Allocate temporary memory for FFT calculations. """ - # self.tmp = pyfftw.zeros_aligned(self.grid_shape, dtype="complex128") - # self.EB_fftw = pyfftw.FFTW( - # self.tmp, - # self.tmp, - # axes=(0, 1, 2), - # direction="FFTW_BACKWARD", - # flags=(config.FFTW_FLAG,), - # threads=self.nthreads, - # ) - self.fft_executor = setup_fftw_executor(self.fft_executor, self.vector_shape) + self.fft_executor = setup_fftw_executor(self.fft_executor, self.vector_shape, + self.nthreads) def calculate_field(self, t, E_out=None, B_out=None): """ Calculates the electric and magnetic fields at a given time step. """ if E_out is None: - # E_out = [np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3)] - # B_out = [np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3)] E_out = np.zeros(self.vector_shape, dtype=config.CDTYPE) B_out = np.zeros(self.vector_shape, dtype=config.CDTYPE) @@ -131,35 +106,14 @@ def calculate_field(self, t, E_out=None, B_out=None): # Calculate fourier of fields at time t and transform back to # spatial domain - - # Transform electric field - np.copyto( - self.fft_executor.tmp, - ne.evaluate(self.E_expr, local_dict=self.EB_dict) - ) - self.fft_executor.backward_fftw.execute() - np.copyto(E_out, self.fft_executor.tmp) - - # Transform magnetic field - np.copyto( - self.fft_executor.tmp, - ne.evaluate(self.B_expr, local_dict=self.EB_dict) - ) - self.fft_executor.backward_fftw.execute() - np.copyto(B_out, self.fft_executor.tmp) + for expr,out_array in zip([self.E_expr, self.B_expr], [E_out, B_out]): # noqa: B905 + np.copyto( + self.fft_executor.tmp, + ne.evaluate(expr, local_dict=self.EB_dict) + ) + self.fft_executor.backward_fftw.execute() + np.copyto(out_array, self.fft_executor.tmp) - # for idx in range(6): - # np.copyto( - # self.fft_executor.tmp, - # ne.evaluate(self.EB_expr[idx], local_dict=self.EB_dict) - # ) - # # ne.evaluate(self.EB_expr[idx], local_dict=self.EB_dict, out=self.tmp) - # # self.EB_fftw.execute() - # self.fft_executor.backward_fftw.execute() - # if idx < 3: - # np.copyto(E_out[idx], self.fft_executor.tmp) - # else: - # np.copyto(B_out[idx-3], self.fft_executor.tmp) return E_out, B_out diff --git a/src/quvac/field/utils.py b/src/quvac/field/utils.py index 42fd3fb..72a6f5e 100644 --- a/src/quvac/field/utils.py +++ b/src/quvac/field/utils.py @@ -178,6 +178,6 @@ def get_intensity(field, t): The intensity is calculated as (abs(E)**2 + abs(B)**2)/2 """ E, B = field.calculate_field(t=t) - E, B = [np.real(Ex) for Ex in E], [np.real(Bx) for Bx in B] + E, B = np.real(E), np.real(B) intensity = (E[0]**2 + E[1]**2 + E[2]**2 + B[0]**2 + B[1]**2 + B[2]**2)/2 return intensity \ No newline at end of file diff --git a/src/quvac/pyfftw_executor.py b/src/quvac/pyfftw_executor.py index b4dabbe..3979617 100644 --- a/src/quvac/pyfftw_executor.py +++ b/src/quvac/pyfftw_executor.py @@ -21,31 +21,29 @@ def allocate_fft(self): self._allocate_fft() def _allocate_fft(self): - # self.tmp_shape = (3,) + self.grid_shape - # self.tmp_shape = self.grid_shape self.tmp = pyfftw.zeros_aligned(self.tmp_shape, dtype=config.CDTYPE) self.forward_fftw = pyfftw.FFTW( - self.tmp, - self.tmp, - axes=self.fft_axes, - direction="FFTW_FORWARD", - flags=(config.FFTW_FLAG,), - threads=self.nthreads, + self.tmp, + self.tmp, + axes=self.fft_axes, + direction="FFTW_FORWARD", + flags=(config.FFTW_FLAG,), + threads=self.nthreads, ) self.backward_fftw = pyfftw.FFTW( - self.tmp, - self.tmp, - axes=self.fft_axes, - direction="FFTW_BACKWARD", - flags=(config.FFTW_FLAG,), - threads=self.nthreads, + self.tmp, + self.tmp, + axes=self.fft_axes, + direction="FFTW_BACKWARD", + flags=(config.FFTW_FLAG,), + threads=self.nthreads, ) -def setup_fftw_executor(fft_executor, grid_shape): +def setup_fftw_executor(fft_executor, grid_shape, nthreads=None): if fft_executor is None: - fft_executor = FFTExecutor(grid_shape) + fft_executor = FFTExecutor(grid_shape, nthreads) fft_executor.allocate_fft() return fft_executor diff --git a/src/quvac/simulation.py b/src/quvac/simulation.py index 000b589..488f501 100755 --- a/src/quvac/simulation.py +++ b/src/quvac/simulation.py @@ -255,8 +255,8 @@ def run_simulation(ini_config, fields_params, files, timings, memory): field = ProbePumpField( fields_params, grid_xyz, - fft_executor, probe_pump_idx=probe_pump_idx, + fft_executor=fft_executor, nthreads=pyfftw_threads, ) timings['field_setup'] = time.perf_counter() From c4275b5aed210ab4cb888731889482e4c2d2f78e Mon Sep 17 00:00:00 2001 From: maxbalrog Date: Mon, 9 Mar 2026 13:50:55 +0100 Subject: [PATCH 06/10] remove commented code from amplitude calculation --- src/quvac/integrator/vacuum_emission.py | 76 ++++--------------------- 1 file changed, 11 insertions(+), 65 deletions(-) diff --git a/src/quvac/integrator/vacuum_emission.py b/src/quvac/integrator/vacuum_emission.py index 11bd984..a1c05a8 100644 --- a/src/quvac/integrator/vacuum_emission.py +++ b/src/quvac/integrator/vacuum_emission.py @@ -54,7 +54,7 @@ def __init__(self, field, grid, fft_executor, nthreads=None, channels=False): self.channels = channels self.c = c - self.nthreads = nthreads if nthreads else os.cpu_count() + self.nthreads = nthreads self.fft_executor = fft_executor @@ -67,8 +67,6 @@ def __init__(self, field, grid, fft_executor, nthreads=None, channels=False): ] if not self.channels: - # self.U1 = [f"(4*E{ax}*F + 7*B{ax}*G)" for ax in "xyz"] - # self.U2 = [f"(4*B{ax}*F - 7*E{ax}*G)" for ax in "xyz"] self.U1 = "(4*E*F + 7*B*G)" self.U2 = "(4*B*F - 7*E*G)" else: @@ -94,14 +92,6 @@ def _define_channel_variables(self): np.zeros(self.grid_shape, dtype=config.FDTYPE) for _ in range(3) ] - # self.U1 = [ - # f"(4*(Ep{ax}*F + E{ax}*F_B_Bp) + 7*(Bp{ax}*G + B{ax}*(G_Ep_B + G_E_Bp)))" - # for ax in "xyz" - # ] - # self.U2 = [ - # f"(4*(Bp{ax}*F + B{ax}*F_B_Bp) - 7*(Ep{ax}*G + E{ax}*(G_Ep_B + G_E_Bp)))" - # for ax in "xyz" - # ] 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)))" @@ -109,15 +99,6 @@ def _allocate_fields(self): """ Allocate memory for field calculations. """ - # self.E_out = [np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3)] - # self.B_out = [np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3)] - # if self.channels: - # self.E_probe = [ - # np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3) - # ] - # self.B_probe = [ - # np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3) - # ] self.E_out = np.zeros(self.vector_shape, dtype=config.CDTYPE) self.B_out = np.zeros(self.vector_shape, dtype=config.CDTYPE) if self.channels: @@ -128,8 +109,6 @@ def _allocate_result_arrays(self): """ Allocate memory for result arrays. """ - # self.U1_acc = [np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3)] - # self.U2_acc = [np.zeros(self.grid_shape, dtype=config.CDTYPE) for _ in range(3)] self.U1_acc = np.zeros(self.vector_shape, dtype=config.CDTYPE) self.U2_acc = np.zeros(self.vector_shape, dtype=config.CDTYPE) self.U1_acc_x, self.U1_acc_y, self.U1_acc_z = self.U1_acc @@ -139,18 +118,8 @@ def _allocate_fft(self): """ Allocate memory for FFT calculations. """ - # self.tmp = pyfftw.zeros_aligned(self.grid_shape, dtype=config.CDTYPE) - # self.tmp_fftw = pyfftw.FFTW( - # self.tmp, - # self.tmp, - # axes=(0, 1, 2), - # direction="FFTW_FORWARD", - # flags=(config.FFTW_FLAG,), - # threads=self.nthreads, - # ) - # self.fft_executor.allocate_fft() - # self.tmp = self.fft_executor.tmp - self.fft_executor = setup_fftw_executor(self.fft_executor, self.vector_shape) + self.fft_executor = setup_fftw_executor(self.fft_executor, self.vector_shape, + self.nthreads) self.prefactor = np.zeros(self.grid_shape, dtype="complex128") @@ -166,7 +135,6 @@ def _free_resources(self): Free allocated resources. """ del self.E_out, self.B_out - # del self.tmp, self.tmp_fftw del self.fft_executor def calculate_one_time_step(self, t, weight=1): @@ -174,8 +142,6 @@ def calculate_one_time_step(self, t, weight=1): Calculate the field and U terms (integrand) for one time step. """ # Calculate fields - # self._allocate_fields() - if not self.channels: self.field.calculate_field(t, E_out=self.E_out, B_out=self.B_out) else: @@ -186,17 +152,14 @@ def calculate_one_time_step(self, t, weight=1): E_pump=self.E_out, B_pump=self.B_out, ) - # Epx, Epy, Epz = [E.real for E in self.E_probe] - # Bpx, Bpy, Bpz = [B.real for B in self.B_probe] Ep = Epx, Epy, Epz = np.real(self.E_probe) Bp = Bpx, Bpy, Bpz = np.real(self.B_probe) - # Ex, Ey, Ez = [E.real for E in self.E_out] - # Bx, By, Bz = [B.real for B in self.B_out] E = Ex, Ey, Ez = np.real(self.E_out) B = Bx, By, Bz = np.real(self.B_out) - self.U_dict.update({"Ex": Ex, "Ey": Ey, "Ez": Ez, "Bx": Bx, "By": By, "Bz": Bz, - "E": E, "B": B}) + self.U_dict.update({"E": E, "B": B, + "Ex": Ex, "Ey": Ey, "Ez": Ez, + "Bx": Bx, "By": By, "Bz": Bz,}) ne.evaluate(self.F_expr, out=self.F) ne.evaluate(self.G_expr, out=self.G) @@ -205,20 +168,20 @@ def calculate_one_time_step(self, t, weight=1): ne.evaluate(self.F_B_Bp_expr, out=self.F_B_Bp) ne.evaluate(self.G_Ep_B_expr, out=self.G_Ep_B) ne.evaluate(self.G_E_Bp_expr, out=self.G_E_Bp) - self.U_dict.update({"Epx": Epx, "Epy": Epy, "Epz": Epz, + self.U_dict.update({"Ep": Ep, "Bp": Bp, + "Epx": Epx, "Epy": Epy, "Epz": Epz, "Bpx": Bpx, "Bpy": Bpy, "Bpz": Bpz, "F_B_Bp": self.F_B_Bp, "G_Ep_B": self.G_Ep_B, - "G_E_Bp": self.G_E_Bp, - "Ep": Ep, "Bp": Bp,}) + "G_E_Bp": self.G_E_Bp,}) # Evaluate U1 and U2 expressions self.prefactor_dict.update({"t": t}) ne.evaluate( "exp(1j*kabs*c*t)", local_dict=self.prefactor_dict, out=self.prefactor ) - for U_key, U_expr in zip(["U1_acc", "U2_acc"], [self.U1, self.U2]): # noqa: B905 - U_acc = getattr(self, U_key) + for U_acc, U_expr in zip([self.U1_acc, self.U2_acc], [self.U1, self.U2]): # noqa: B905 + # U_acc = getattr(self, U_key) ne.evaluate(U_expr, global_dict=self.U_dict, out=self.fft_executor.tmp) self.fft_executor.forward_fftw.execute() @@ -233,23 +196,6 @@ def calculate_one_time_step(self, t, weight=1): }, ) np.copyto(U_acc, U_res) - - # for i, expr in enumerate(U_expr): - # U_acc = getattr(self, U_key)[i] - # ne.evaluate(expr, global_dict=self.U_dict, out=self.fft_executor.tmp) - # self.fft_executor.forward_fftw.execute() - - # U_res = ne.evaluate( - # "U_acc + U*prefactor*dt*dV", - # global_dict={ - # "U_acc": U_acc, - # "U": self.fft_executor.tmp, - # "prefactor": self.prefactor, - # "dt": self.dt, - # "dV": self.dV, - # }, - # ) - # np.copyto(U_acc, U_res) def calculate_time_integral(self, t_grid, integration_method="trapezoid"): """ From 809c2183d4697b18147d95f83b4d409fa05f930e Mon Sep 17 00:00:00 2001 From: maxbalrog Date: Mon, 9 Mar 2026 14:46:58 +0100 Subject: [PATCH 07/10] move exponential prefactor calculation out of the loop --- src/quvac/integrator/vacuum_emission.py | 61 ++++++++++++++++--------- 1 file changed, 40 insertions(+), 21 deletions(-) diff --git a/src/quvac/integrator/vacuum_emission.py b/src/quvac/integrator/vacuum_emission.py index a1c05a8..c5d5a76 100644 --- a/src/quvac/integrator/vacuum_emission.py +++ b/src/quvac/integrator/vacuum_emission.py @@ -121,12 +121,8 @@ def _allocate_fft(self): self.fft_executor = setup_fftw_executor(self.fft_executor, self.vector_shape, self.nthreads) - self.prefactor = np.zeros(self.grid_shape, dtype="complex128") - - self.prefactor_dict = { - "kabs": self.kabs, - "c": c, - } + self.prefactor = np.ones(self.grid_shape, dtype="complex128") + self.prefactor_step = np.zeros(self.grid_shape, dtype="complex128") self.U_dict = {"F": self.F, "G": self.G} @@ -175,33 +171,53 @@ def calculate_one_time_step(self, t, weight=1): "G_Ep_B": self.G_Ep_B, "G_E_Bp": self.G_E_Bp,}) + # Update prefactor: this prescription is valid only for equidistant grids!!! + ne.evaluate("prefactor*prefactor_step", local_dict=self.__dict__, + out=self.prefactor) + # Evaluate U1 and U2 expressions - self.prefactor_dict.update({"t": t}) - ne.evaluate( - "exp(1j*kabs*c*t)", local_dict=self.prefactor_dict, out=self.prefactor - ) for U_acc, U_expr in zip([self.U1_acc, self.U2_acc], [self.U1, self.U2]): # noqa: B905 - # U_acc = getattr(self, U_key) ne.evaluate(U_expr, global_dict=self.U_dict, out=self.fft_executor.tmp) self.fft_executor.forward_fftw.execute() U_res = ne.evaluate( - "U_acc + U*prefactor*dt*dV", - global_dict={ - "U_acc": U_acc, - "U": self.fft_executor.tmp, - "prefactor": self.prefactor, - "dt": self.dt, - "dV": self.dV, - }, - ) + "U_acc + U*prefactor", + global_dict={ + "U_acc": U_acc, + "U": self.fft_executor.tmp, + "prefactor": self.prefactor, + }, + ) np.copyto(U_acc, U_res) + def multiply_integration_result(self, t_grid): + self._free_resources() + # prefactor related to time grid and discretization + self.prefactor_dict.update({"t": t_grid[0], "dt": self.dt, "dV": self.dV,}) + ne.evaluate( + "exp(1j*kabs*c*t)*dt*dV", + local_dict=self.prefactor_dict, + out=self.prefactor + ) + for acc in [self.U1_acc, self.U2_acc]: + ne.evaluate("acc*prefactor", global_dict={"prefactor": self.prefactor}, + out=acc) + def calculate_time_integral(self, t_grid, integration_method="trapezoid"): """ Calculate the time integral. """ self.dt = t_grid[1] - t_grid[0] + # prefactor-related variables + self.prefactor_dict = { + "kabs": self.kabs, + "c": c, + "dt": self.dt, + } + ne.evaluate( + "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): @@ -214,6 +230,9 @@ def calculate_time_integral(self, t_grid, integration_method="trapezoid"): f"passed {integration_method}" ) raise NotImplementedError(err_msg) + + # finish calculation + self.multiply_integration_result(t_grid) def calculate_amplitudes( self, t_grid, integration_method="trapezoid", save_path=None @@ -230,7 +249,7 @@ def calculate_amplitudes( self.calculate_time_integral(t_grid, integration_method) time_integral_end = time.perf_counter() time_integral = time_integral_end - time_integral_start - self._free_resources() + # self._free_resources() # Results should be in U1_acc and U2_acc dims = 1 / BS**3 * m_e**2 * c**3 / hbar**2 From 5d3c31452d5df6135389dad6ffa75ccac6b4bed9 Mon Sep 17 00:00:00 2001 From: maxbalrog Date: Mon, 9 Mar 2026 15:22:43 +0100 Subject: [PATCH 08/10] rework maxwell propagation calculation --- src/quvac/field/abc.py | 11 +++++-- src/quvac/field/maxwell.py | 40 +++++++++---------------- src/quvac/integrator/vacuum_emission.py | 22 ++++++++------ 3 files changed, 36 insertions(+), 37 deletions(-) diff --git a/src/quvac/field/abc.py b/src/quvac/field/abc.py index 28ce9d1..73c49f7 100644 --- a/src/quvac/field/abc.py +++ b/src/quvac/field/abc.py @@ -233,6 +233,13 @@ def _allocate_fft(self): self.Ef = pyfftw.zeros_aligned(self.vector_shape, dtype="complex128") self.fft_executor = setup_fftw_executor(self.fft_executor, self.vector_shape) + self.a_dict = { + "e1x": self.e1x, "e1y": self.e1y, "e1z": self.e1z, + "e2x": self.e2x, "e2y": self.e2y, "e2z": self.e2z, + "Efx": self.Ef[0], "Efy": self.Ef[1], "Efz": self.Ef[2], + "dV": self.dV, + } + def _check_energy_kspace(self): # Fix energy W_upd = get_field_energy_kspace( @@ -287,10 +294,10 @@ def get_a12(self, t0=None, fft_executor=None): Efx, Efy, Efz = self.Ef self.a1 = ne.evaluate( - "dV * (e1x*Efx + e1y*Efy + e1z*Efz)", global_dict=self.__dict__ + "dV * (e1x*Efx + e1y*Efy + e1z*Efz)", local_dict=self.a_dict ) self.a2 = ne.evaluate( - "dV * (e2x*Efx + e2y*Efy + e2z*Efz)", global_dict=self.__dict__ + "dV * (e2x*Efx + e2y*Efy + e2z*Efz)", global_dict=self.a_dict ) self._check_energy_kspace() diff --git a/src/quvac/field/maxwell.py b/src/quvac/field/maxwell.py index aa670f5..59db943 100644 --- a/src/quvac/field/maxwell.py +++ b/src/quvac/field/maxwell.py @@ -46,8 +46,10 @@ def __init__(self, grid, fft_executor=None, nthreads=None): self.c = c self.norm_ifft = self.dVk / (2.0 * pi) ** 3 - self.E_expr = "(e1*a1t + e2*a2t)" - self.B_expr = "(e2*a1t - e1*a2t)" + self.prefactor_expr = "exp(-1j*kabs*c*(t-t0)) * norm_ifft" + + self.E_expr = "prefactor * (e1*a1 + e2*a2)" + self.B_expr = "prefactor * (e2*a1 - e1*a2)" self._allocate_tmp() @@ -56,24 +58,21 @@ def _allocate_ifft(self): Allocate memory for inverse FFT calculations and define dictionaries for numexpr.evaluate(). """ - self.a1t, self.a2t = [ - np.zeros(self.grid_shape, dtype="complex128") for _ in range(2) - ] + self.prefactor = np.zeros(self.grid_shape, dtype=config.CDTYPE) - self.a_dict = { + self.prefactor_dict = { "kabs": self.kabs, "c": c, "t0": self.t0, "norm_ifft": self.norm_ifft, - "a1": self.a1, - "a2": self.a2, } self.EB_dict = { "e1": self.e1, "e2": self.e2, - "a1t": self.a1t, - "a2t": self.a2t, + "a1": self.a1, + "a2": self.a2, + "prefactor": self.prefactor, } def _allocate_tmp(self): @@ -91,26 +90,15 @@ def calculate_field(self, t, E_out=None, B_out=None): E_out = np.zeros(self.vector_shape, dtype=config.CDTYPE) B_out = np.zeros(self.vector_shape, dtype=config.CDTYPE) - # Calculate a1,a2 at time t - self.a_dict.update({"t": t}) - ne.evaluate( - "exp(-1j*kabs*c*(t-t0)) * a1 * norm_ifft", - local_dict=self.a_dict, - out=self.a1t, - ) - ne.evaluate( - "exp(-1j*kabs*c*(t-t0)) * a2 * norm_ifft", - local_dict=self.a_dict, - out=self.a2t, - ) + # Calculate prefactor at time t + self.prefactor_dict.update({"t": t}) + ne.evaluate(self.prefactor_expr, local_dict=self.prefactor_dict, + out=self.prefactor) # Calculate fourier of fields at time t and transform back to # spatial domain for expr,out_array in zip([self.E_expr, self.B_expr], [E_out, B_out]): # noqa: B905 - np.copyto( - self.fft_executor.tmp, - ne.evaluate(expr, local_dict=self.EB_dict) - ) + ne.evaluate(expr, local_dict=self.EB_dict, out=self.fft_executor.tmp) self.fft_executor.backward_fftw.execute() np.copyto(out_array, self.fft_executor.tmp) diff --git a/src/quvac/integrator/vacuum_emission.py b/src/quvac/integrator/vacuum_emission.py index c5d5a76..f35e3ce 100644 --- a/src/quvac/integrator/vacuum_emission.py +++ b/src/quvac/integrator/vacuum_emission.py @@ -125,6 +125,10 @@ def _allocate_fft(self): self.prefactor_step = np.zeros(self.grid_shape, dtype="complex128") self.U_dict = {"F": self.F, "G": self.G} + self.U_acc_dict = { + "U": self.fft_executor.tmp, + "prefactor": self.prefactor, + } def _free_resources(self): """ @@ -157,6 +161,7 @@ def calculate_one_time_step(self, t, weight=1): "Ex": Ex, "Ey": Ey, "Ez": Ez, "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) @@ -172,7 +177,7 @@ def calculate_one_time_step(self, t, weight=1): "G_E_Bp": self.G_E_Bp,}) # Update prefactor: this prescription is valid only for equidistant grids!!! - ne.evaluate("prefactor*prefactor_step", local_dict=self.__dict__, + ne.evaluate("prefactor*prefactor_step", local_dict=self.prefactor_dict, out=self.prefactor) # Evaluate U1 and U2 expressions @@ -180,15 +185,12 @@ def calculate_one_time_step(self, t, weight=1): ne.evaluate(U_expr, global_dict=self.U_dict, out=self.fft_executor.tmp) self.fft_executor.forward_fftw.execute() - U_res = ne.evaluate( + self.U_acc_dict.update({"U_acc": U_acc}) + ne.evaluate( "U_acc + U*prefactor", - global_dict={ - "U_acc": U_acc, - "U": self.fft_executor.tmp, - "prefactor": self.prefactor, - }, + local_dict=self.U_acc_dict, + out=U_acc, ) - np.copyto(U_acc, U_res) def multiply_integration_result(self, t_grid): self._free_resources() @@ -213,6 +215,8 @@ def calculate_time_integral(self, t_grid, integration_method="trapezoid"): "kabs": self.kabs, "c": c, "dt": self.dt, + "prefactor": self.prefactor, + "prefactor_step": self.prefactor_step, } ne.evaluate( "exp(1j*kabs*c*dt)", local_dict=self.prefactor_dict, out=self.prefactor_step @@ -249,7 +253,6 @@ def calculate_amplitudes( self.calculate_time_integral(t_grid, integration_method) time_integral_end = time.perf_counter() time_integral = time_integral_end - time_integral_start - # self._free_resources() # Results should be in U1_acc and U2_acc dims = 1 / BS**3 * m_e**2 * c**3 / hbar**2 @@ -263,6 +266,7 @@ def calculate_amplitudes( f"prefactor * ({self.I_12_expr} + {self.I_21_expr})", global_dict=self.__dict__, ).astype(config.CDTYPE) + # Save amplitudes if save_path: self.save_amplitudes(save_path) From 6d84061f052eba4b16397e9ea362a6ccc12dd9f3 Mon Sep 17 00:00:00 2001 From: maxbalrog Date: Mon, 9 Mar 2026 15:59:27 +0100 Subject: [PATCH 09/10] remove zip from time for loops --- src/quvac/field/maxwell.py | 19 +++++---- src/quvac/integrator/vacuum_emission.py | 55 +++++++++++++++---------- 2 files changed, 44 insertions(+), 30 deletions(-) diff --git a/src/quvac/field/maxwell.py b/src/quvac/field/maxwell.py index 59db943..cb5937c 100644 --- a/src/quvac/field/maxwell.py +++ b/src/quvac/field/maxwell.py @@ -51,13 +51,19 @@ def __init__(self, grid, fft_executor=None, nthreads=None): self.E_expr = "prefactor * (e1*a1 + e2*a2)" self.B_expr = "prefactor * (e2*a1 - e1*a2)" - self._allocate_tmp() + self.EB_pairs = [ + [self.E_expr, None], + [self.B_expr, None], + ] def _allocate_ifft(self): """ Allocate memory for inverse FFT calculations and define dictionaries for numexpr.evaluate(). """ + self.fft_executor = setup_fftw_executor(self.fft_executor, self.vector_shape, + self.nthreads) + self.prefactor = np.zeros(self.grid_shape, dtype=config.CDTYPE) self.prefactor_dict = { @@ -75,13 +81,6 @@ def _allocate_ifft(self): "prefactor": self.prefactor, } - def _allocate_tmp(self): - """ - Allocate temporary memory for FFT calculations. - """ - self.fft_executor = setup_fftw_executor(self.fft_executor, self.vector_shape, - self.nthreads) - def calculate_field(self, t, E_out=None, B_out=None): """ Calculates the electric and magnetic fields at a given time step. @@ -89,6 +88,8 @@ def calculate_field(self, t, E_out=None, B_out=None): if E_out is None: E_out = np.zeros(self.vector_shape, dtype=config.CDTYPE) B_out = np.zeros(self.vector_shape, dtype=config.CDTYPE) + self.EB_pairs[0][1] = E_out + self.EB_pairs[1][1] = B_out # Calculate prefactor at time t self.prefactor_dict.update({"t": t}) @@ -97,7 +98,7 @@ def calculate_field(self, t, E_out=None, B_out=None): # Calculate fourier of fields at time t and transform back to # spatial domain - for expr,out_array in zip([self.E_expr, self.B_expr], [E_out, B_out]): # noqa: B905 + for expr,out_array in self.EB_pairs: # noqa: B905 ne.evaluate(expr, local_dict=self.EB_dict, out=self.fft_executor.tmp) self.fft_executor.backward_fftw.execute() np.copyto(out_array, self.fft_executor.tmp) diff --git a/src/quvac/integrator/vacuum_emission.py b/src/quvac/integrator/vacuum_emission.py index f35e3ce..6f807b4 100644 --- a/src/quvac/integrator/vacuum_emission.py +++ b/src/quvac/integrator/vacuum_emission.py @@ -114,6 +114,16 @@ def _allocate_result_arrays(self): self.U1_acc_x, self.U1_acc_y, self.U1_acc_z = self.U1_acc 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.prefactor = np.ones(self.grid_shape, dtype="complex128") + self.prefactor_step = np.zeros(self.grid_shape, dtype="complex128") + + self.U_dict = {"F": self.F, "G": self.G} + def _allocate_fft(self): """ Allocate memory for FFT calculations. @@ -121,15 +131,19 @@ def _allocate_fft(self): self.fft_executor = setup_fftw_executor(self.fft_executor, self.vector_shape, self.nthreads) - self.prefactor = np.ones(self.grid_shape, dtype="complex128") - self.prefactor_step = np.zeros(self.grid_shape, dtype="complex128") - - self.U_dict = {"F": self.F, "G": self.G} self.U_acc_dict = { "U": self.fft_executor.tmp, "prefactor": self.prefactor, } + def _allocate_resources(self): + """ + Allocate arrays needed for calculation. + """ + self._allocate_result_arrays() + self._allocate_fft() + self._allocate_fields() + def _free_resources(self): """ Free allocated resources. @@ -181,7 +195,7 @@ def calculate_one_time_step(self, t, weight=1): out=self.prefactor) # Evaluate U1 and U2 expressions - for U_acc, U_expr in zip([self.U1_acc, self.U2_acc], [self.U1, self.U2]): # noqa: B905 + for U_acc, U_expr in self.U_pairs: # noqa: B905 ne.evaluate(U_expr, global_dict=self.U_dict, out=self.fft_executor.tmp) self.fft_executor.forward_fftw.execute() @@ -238,34 +252,33 @@ def calculate_time_integral(self, t_grid, integration_method="trapezoid"): # finish calculation self.multiply_integration_result(t_grid) + 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( + f"prefactor * ({self.I_11_expr} - {self.I_22_expr})", + global_dict=self.__dict__, + ).astype(config.CDTYPE) + self.S2 = ne.evaluate( + f"prefactor * ({self.I_12_expr} + {self.I_21_expr})", + global_dict=self.__dict__, + ).astype(config.CDTYPE) + def calculate_amplitudes( self, t_grid, integration_method="trapezoid", save_path=None ): """ Calculate the vacuum emission amplitudes and save the result. """ - # Allocate resources - self._allocate_result_arrays() - self._allocate_fft() - self._allocate_fields() + self._allocate_resources() time_integral_start = time.perf_counter() self.calculate_time_integral(t_grid, integration_method) time_integral_end = time.perf_counter() time_integral = time_integral_end - time_integral_start - # Results should be in U1_acc and U2_acc - 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 - # Next time need to be careful with f-strings and brackets - self.S1 = ne.evaluate( - f"prefactor * ({self.I_11_expr} - {self.I_22_expr})", - global_dict=self.__dict__, - ).astype(config.CDTYPE) - self.S2 = ne.evaluate( - f"prefactor * ({self.I_12_expr} + {self.I_21_expr})", - global_dict=self.__dict__, - ).astype(config.CDTYPE) + # Calculate + self._calculate_S1_S2() # Save amplitudes if save_path: From 82d1457b2e421d6765763bb3ffba47f1ae917957 Mon Sep 17 00:00:00 2001 From: maxbalrog Date: Mon, 9 Mar 2026 16:33:04 +0100 Subject: [PATCH 10/10] update docstrings --- src/quvac/field/external_field.py | 6 ++++ src/quvac/field/maxwell.py | 4 +++ src/quvac/integrator/vacuum_emission.py | 10 +++++- src/quvac/pyfftw_executor.py | 44 +++++++++++++++++++++++-- 4 files changed, 60 insertions(+), 4 deletions(-) diff --git a/src/quvac/field/external_field.py b/src/quvac/field/external_field.py index b3cd51a..6c2028e 100644 --- a/src/quvac/field/external_field.py +++ b/src/quvac/field/external_field.py @@ -22,6 +22,8 @@ class ExternalField(Field): List of dictionaries containing the parameters for each external field. grid : quvac.grid.GridXYZ Spatial and spectral grid. + fft_executor: quvac.pyfftw_executor.FFTExecutor, optional + Executor that performs FFTs. nthreads : int, optional Number of threads to use for calculations. If not provided, defaults to the number of CPU cores. @@ -32,6 +34,8 @@ class ExternalField(Field): List of field objects. grid_xyz : tuple of np.array The spatial grid. + fft_executor: quvac.pyfftw_executor.FFTExecutor + Executor that performs FFTs. nthreads : int Number of threads to use for calculations. """ @@ -146,6 +150,8 @@ class ProbePumpField(Field): If not provided, defaults to {"probe": [0], "pump": [1]}. + fft_executor: quvac.pyfftw_executor.FFTExecutor, optional + Executor that performs FFTs. nthreads : int, optional Number of threads to use for calculations. If not provided, defaults to the number of CPU cores. diff --git a/src/quvac/field/maxwell.py b/src/quvac/field/maxwell.py index cb5937c..d68387d 100644 --- a/src/quvac/field/maxwell.py +++ b/src/quvac/field/maxwell.py @@ -31,6 +31,8 @@ class MaxwellField(Field): ---------- grid : quvac.grid.GridXYZ Spatial and spectral grid. + fft_executor: quvac.pyfftw_executor.FFTExecutor, optional + Executor that performs FFTs. nthreads : int, optional Number of threads to use for calculations. If not provided, defaults to the number of CPU cores. @@ -117,6 +119,8 @@ class MaxwellMultiple(MaxwellField): Parameters of the fields. grid : quvac.grid.GridXYZ Spatial and spectral grid. + fft_executor: quvac.pyfftw_executor.FFTExecutor, optional + Executor that performs FFTs. nthreads : int, optional Number of threads to use for calculations. If not provided, defaults to the number of CPU cores. diff --git a/src/quvac/integrator/vacuum_emission.py b/src/quvac/integrator/vacuum_emission.py index 6f807b4..6372c21 100644 --- a/src/quvac/integrator/vacuum_emission.py +++ b/src/quvac/integrator/vacuum_emission.py @@ -37,6 +37,8 @@ class VacuumEmission: External fields. grid : quvac.grid.GridXYZ Spatial and spectral grid. + fft_executor: quvac.pyfftw_executor.FFTExecutor, optional + Executor that performs FFTs. nthreads : int, optional Number of threads to use for calculations. If not provided, defaults to the number of CPU cores. @@ -46,7 +48,7 @@ class VacuumEmission: """ - def __init__(self, field, grid, fft_executor, nthreads=None, channels=False): + def __init__(self, field, grid, fft_executor=None, nthreads=None, channels=False): self.field = field self.grid_xyz = grid # Update local dict with variables from GridXYZ class @@ -207,6 +209,9 @@ def calculate_one_time_step(self, t, weight=1): ) def multiply_integration_result(self, t_grid): + """ + Multiply the integral by common prefactors. + """ self._free_resources() # prefactor related to time grid and discretization self.prefactor_dict.update({"t": t_grid[0], "dt": self.dt, "dV": self.dV,}) @@ -253,6 +258,9 @@ def calculate_time_integral(self, t_grid, integration_method="trapezoid"): self.multiply_integration_result(t_grid) def _calculate_S1_S2(self): + """ + Calculate S1 and S2 amplitudes from accumulated values of integrals U1 and U2. + """ 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( diff --git a/src/quvac/pyfftw_executor.py b/src/quvac/pyfftw_executor.py index 3979617..f9f17e4 100644 --- a/src/quvac/pyfftw_executor.py +++ b/src/quvac/pyfftw_executor.py @@ -1,7 +1,8 @@ ''' Pyfftw executor that is used in all 3D FFTs. -At the moment, batched FFT transforms of 3 3D arrays are performed. +At the moment, batched FFT transforms of three 3D arrays are performed (corresponding +to a vector field on 3D grid). ''' import os @@ -11,12 +12,44 @@ class FFTExecutor: + """ + Class to unite the FFT calculation. + + Parameters + ---------- + tmp_shape : tuple or list + Size of the buffer array over which FFTs are performed. By default, it expects + (3, Nx, Ny, Nz). + nthreads : int, optional + Number of threads to use for calculations. If not provided, defaults to the + number of CPU cores. + fft_axes : tuple or list + Axes over which to perform FFT. By default, (1,2,3). + + Attributes + ---------- + tmp_shape : tuple + Size of the buffer array. + nthreads : int + Number of threads to use for FFTs. + fft_axes : tuple + Axes over which to perform FFT. + tmp : np.ndarray of size `tmp_shape` + Buffer array used for FFTs. + forward_fftw : pyfftw.FFTW + Forward FFT executor. + backward_fftw : pyfftw.FFTW + Backward FFT executor. + """ def __init__(self, tmp_shape, nthreads=None, fft_axes=(1,2,3)): - self.tmp_shape = tmp_shape + self.tmp_shape = tuple(tmp_shape) self.nthreads = nthreads if nthreads else os.cpu_count() - self.fft_axes = fft_axes + self.fft_axes = tuple(fft_axes) def allocate_fft(self): + """ + Allocate buffer array and set up FFT executors if they were not already. + """ if getattr(self, "tmp", None) is None: self._allocate_fft() @@ -43,6 +76,11 @@ 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. + - Allocate buffer arrays and FFT executors. + """ if fft_executor is None: fft_executor = FFTExecutor(grid_shape, nthreads) fft_executor.allocate_fft()