Rework the core calculation scripts (long awaited...)

src/quvac/config.py

@@ -17,3 +17,5 @@
 
 FDTYPE = "float64"
 CDTYPE = "complex128"
+
+FFTW_FLAG = "FFTW_MEASURE"

src/quvac/field/abc.py

@@ -11,11 +11,11 @@
 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
 from quvac.field.utils import get_field_energy, get_field_energy_kspace
+from quvac.pyfftw_executor import setup_fftw_executor
 
 ANGLE_KEYS = [
     "theta",
@@ -59,7 +59,7 @@ def __init__(self, field_params, grid):
         # self.<variable_name>
         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):
@@ -69,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
@@ -129,9 +129,9 @@ 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
 
     def rotate_fields_back(self, E_out, B_out, mode):
@@ -171,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):
         """
@@ -231,21 +230,15 @@ def _allocate_fft(self):
         """
         Allocates memory for FFT calculations.
         """
-        self.Ef = [
-            pyfftw.zeros_aligned(self.grid_shape, dtype="complex128") for _ in range(3)
-        ]
-        # pyfftw scheme
-        self.Ef_fftw = [
-            pyfftw.FFTW(
-                a,
-                a,
-                axes=(0, 1, 2),
-                direction="FFTW_FORWARD",
-                flags=("FFTW_MEASURE",),
-                threads=1,
-            )
-            for a in self.Ef
-        ]
+        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
@@ -267,7 +260,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,25 +282,27 @@ 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)
+        self.fft_executor.forward_fftw.execute()
+        np.copyto(self.Ef, self.fft_executor.tmp)
 
         # Calculate a1, a2 coefficients
         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()
 
-        del self.Ef, self.Ef_fftw
+        del self.Ef
         return self.a1, self.a2
 
 

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
@@ -23,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.
@@ -33,15 +34,23 @@ 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.
     """
 
-    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.nthreads = nthreads
+        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 
+        # automatic
+        self.need_to_zero_fields = True
 
         maxwell_params = [
             params
@@ -54,7 +63,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"
@@ -87,7 +101,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(
@@ -97,10 +112,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
@@ -126,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.
@@ -142,18 +168,19 @@ 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, 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"]]
 
-        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):
         """