Multithread initial espes in optimization (JYU-IBA#193)

- add cleanup option for recoil files in calculate_espe

modules/element_simulation.py

@@ -833,15 +833,20 @@ def calculate_espe(
             ch: Optional[float] = None,
             fluence: Optional[float] = None,
             optimization_type: Optional[OptimizationType] = None,
-            write_to_file: bool = True) -> Tuple[List, Optional[Path]]:
+            write_to_file: bool = True,
+            remove_recoil_file: bool = False) -> Tuple[List, Optional[Path]]:
         """Calculate the energy spectrum from the MCERD result file.
+
         Args:
             recoil_element: Recoil element.
             verbose: In terminal (disabled by default).
             ch: Channel width to use.
             fluence: Fluence to use.
             optimization_type: Either recoil, fluence or None
             write_to_file: Whether spectrum is written to file
+            remove_recoil_file: Whether to remove temporary .recoil file
+                after getting the energy spectrum.
+
         Return:
             tuple consisting of spectrum data and espe file
         """
@@ -893,6 +898,10 @@ def calculate_espe(
             recoil_file=recoil_file,
             verbose=verbose
         )
+
+        if remove_recoil_file:
+            recoil_file.unlink()
+
         # TODO returning espe_file is a bit pointless if write_to_file is
         #   False
         return spectrum, output_file

modules/linear_optimization.py

@@ -1,3 +1,5 @@
+# Multithreaded espes
+
 """
 Created on 20.09.2021
 
@@ -26,6 +28,7 @@
 import copy
 import subprocess
 from collections import namedtuple
+from concurrent.futures import ThreadPoolExecutor
 from typing import Tuple, List, Optional, Callable, Union
 
 import numpy as np
@@ -150,13 +153,16 @@ def _generate_mev_to_nm_function(self) -> None:
 
         smoothing_width = round(self.measured_espe_x.shape[0] / 30)
 
+        solutions = [
+            get_solution6(nm, nm + self.sample_width, self.lower_limits, self.upper_limits)
+            for nm in nm_points]
+
+        espes = self._run_solutions(solutions)
+
         mevs = []
         prominences = []
         nms = []
-        for nm in nm_points:  # TODO: multi-thread this
-            solution = get_solution6(
-                nm, nm + self.sample_width, self.lower_limits, self.upper_limits)
-            espe = self._run_solution(solution)
+        for espe, nm in zip(espes, nm_points):
             if not espe:
                 raise ValueError(
                     "Ensure that there is simulated data for this recoil "
@@ -595,6 +601,49 @@ def initialize_solution(self):
 
         return solution
 
+    # TODO: This could be a more general solution (at least for NSGA-II too)
+    def _run_solutions(self, solutions: List["BaseSolution"]) -> List[Espe]:
+        """Form a recoil for each solution and return their espes.
+
+        This function is multithreaded.
+
+        Args:
+            solutions: solutions to run
+
+        Returns:
+            Energy spectra for each solution (in correct order)
+        """
+        if self.optimization_type is OptimizationType.RECOIL:
+            # TODO: Using a dummy value like this is questionable
+            self.element_simulation.optimization_recoils = [
+                self.form_recoil(copy.deepcopy(solutions[0]))]
+
+            recoil_elements = [self.form_recoil(sol, f"thread-{i}")
+                               for i, sol in enumerate(solutions)]
+
+            # TODO: Splitting the solutions over fewer threads (more than one
+            #  solution per thread) may be slightly faster than this
+            with ThreadPoolExecutor(max_workers=len(solutions)) as executor:
+                futures = [
+                    executor.submit(
+                        self.element_simulation.calculate_espe,
+                        rec,
+                        verbose=self.verbose,
+                        optimization_type=self.optimization_type,
+                        ch=self.channel_width,
+                        write_to_file=False,
+                        remove_recoil_file=True)
+                    for rec in recoil_elements]
+
+                espes = [future.result()[0] for future in futures]
+        elif self.optimization_type is OptimizationType.FLUENCE:
+            raise NotImplementedError
+        else:
+            raise ValueError(
+                f"Unknown optimization type {self.optimization_type}")
+
+        return espes
+
     def _run_solution(self, solution) -> Espe:
         """Form a recoil based on the given solution and return its espe."""
         if self.optimization_type is OptimizationType.RECOIL:
@@ -607,7 +656,8 @@ def _run_solution(self, solution) -> Espe:
                 verbose=self.verbose,
                 optimization_type=self.optimization_type,
                 ch=self.channel_width,
-                write_to_file=False)
+                write_to_file=False,
+                remove_recoil_file=True)
         elif self.optimization_type is OptimizationType.FLUENCE:
             raise NotImplementedError
         else: