Add performance-related flexibility

.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:

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) <https://arxiv.org/abs/1912.11698>`_.
 
 .. [2] A. Blinne, et al. "All-optical signatures of quantum vacuum nonlinearities 
-    in generic laser fields." PRD 99.1 (2019): 016006.
+    in generic laser fields." PRD 99.1 (2019): 016006 `(article) <https://arxiv.org/abs/1811.08895>`_.

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.

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",

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):
     """

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