Update containers to allow for concurrency

lisainstrument/containers.py

@@ -9,6 +9,7 @@ Authors:
 
 import abc
 import logging
+from concurrent.futures import ThreadPoolExecutor
 import h5py
 import numpy as np
 import matplotlib.pyplot as plt
@@ -19,16 +20,23 @@ logger = logging.getLogger(__name__)
 class ForEachObject(abc.ABC):
     """Abstract class which represents a dictionary holding a value for each object."""
 
-    def __init__(self, values):
+    def __init__(self, values, concurrent=False):
         """Initialize an object with a value or a function of the mosa index.
 
         Args:
             values: a value, a dictionary of values, or a function (mosa -> value)
+            concurrent (bool): whether to parallelize using a thread pool
         """
         if isinstance(values, dict):
             self.dict = {mosa: values[mosa] for mosa in self.indices()}
         elif callable(values):
-            self.dict = {mosa: values(mosa) for mosa in self.indices()}
+            if concurrent:
+                with ThreadPoolExecutor() as executor:
+                    indices = self.indices()
+                    computed_values = executor.map(values, indices)
+                    self.dict = dict(zip(indices, computed_values))
+            else:
+                self.dict = {mosa: values(mosa) for mosa in self.indices()}
         elif isinstance(values, h5py.Dataset):
             self.dict = {mosa: values[mosa] for mosa in self.indices()}
         else:
@@ -40,13 +48,14 @@ class ForEachObject(abc.ABC):
         """Return list of object indices."""
         raise NotImplementedError
 
-    def transformed(self, transformation):
+    def transformed(self, transformation, concurrent=False):
         """Return a new dictionary from transformed objects.
 
         Args:
             transformation: function (mosa, value -> new_value)
+            concurrent (bool): whether to parallelize using a thread pool
         """
-        return self.__class__(lambda mosa: transformation(mosa, self[mosa]))
+        return self.__class__(lambda mosa: transformation(mosa, self[mosa]), concurrent)
 
     def collapsed(self):
         """Turn a numpy arrays containing identical elements into a scalar.