parallelised spin-wave calculation with a ProcessPoolExecutor

Author: alexhroom

pyspinw/calculations/spinwave.py

@@ -1,5 +1,6 @@
 """Spinwave Calculations"""
 
+from concurrent.futures import as_completed, ProcessPoolExecutor
 from dataclasses import dataclass
 from enum import Enum
 
@@ -58,8 +59,8 @@ def spinwave_calculation(rotations: np.ndarray,
     root_mags = np.sqrt(0.5*magnitudes) # S_i / sqrt(2)
     spin_coefficients = root_mags.reshape(-1, 1) * root_mags.reshape(1, -1)
 
-    energies = []
-    methods = []
+    energies = {}
+    methods = {}
 
     # calculate the C matrix for h(q), which is q-independent
     C = np.zeros((n_sites, n_sites), dtype=complex)
@@ -70,24 +71,37 @@ def spinwave_calculation(rotations: np.ndarray,
 
     C *= 2
 
-    for q in q_vectors:
-        eigenvalues, method = spinwave_calculation_single_q(q, C, n_sites, z, z_conj, spin_coefficients, couplings)
-        energies.append(eigenvalues)  # These are currently the square energies
-        methods.append(method)
+    # we index the single-q calculations so that we can create a dictionary of which order the energies are in
+    # and can sort them later: we use q_index as the key because the q-vectors themselves are not hashable
+    with ProcessPoolExecutor() as executor:
+        q_calculations = {executor.submit(_calc_single_q, q, C, n_sites, z, z_conj, spin_coefficients, couplings)
+                          : q_index for q_index, q in enumerate(q_vectors)}
+    for future in as_completed(q_calculations):
+        q_index = q_calculations[future]
+        try:
+            eigenvalues, method = future.result()
+        except Exception as err:
+            raise err
+        else:
+            energies[q_index] = eigenvalues
+            methods[q_index] = method
+
+    sorted_energies = [energies[i] for i in range(q_vectors.shape[0])]
+    sorted_methods = [methods[i] for i in range(q_vectors.shape[0])]
 
     return SpinwaveResult(
                 q_vectors=q_vectors,
-                raw_energies=energies,
-                method=methods)
+                raw_energies=sorted_energies,
+                method=sorted_methods)
 
 
-def spinwave_calculation_single_q(q: float,
-                                  C: np.ndarray,
-                                  n_sites: int,
-                                  z: np.ndarray,
-                                  z_conj: np.ndarray,
-                                  spin_coefficients: np.ndarray,
-                                  couplings: list[Coupling]):
+def _calc_single_q(q: float,
+                   C: np.ndarray,
+                   n_sites: int,
+                   z: np.ndarray,
+                   z_conj: np.ndarray,
+                   spin_coefficients: np.ndarray,
+                   couplings: list[Coupling]):
     """Calculate the energies for the 'Hamiltonian' h(q) for a single q-value."""
     A = np.zeros((n_sites, n_sites), dtype=complex)
     B = np.zeros((n_sites, n_sites), dtype=complex)