thread or process pool execution does not help with performance

znh5md/utils.py

@@ -1,7 +1,7 @@
+
 import ase
 import numpy as np
 from ase.calculators.singlepoint import SinglePointCalculator
-import concurrent.futures
 
 
 def concatenate_varying_shape_arrays(arrays: list[np.ndarray]) -> np.ndarray:
@@ -125,9 +125,17 @@ def fill_dataset(dataset, new_data):
     dataset[old_shape[0] :] = padded_new_data
 
 
-def build_atoms(
-    atomic_numbers, positions, velocities, cell, pbc, calc_data, info_data, arrays_data, index
-) -> tuple[ase.Atoms, int]:
+def build_atoms(args) -> ase.Atoms:
+    (
+        atomic_numbers,
+        positions,
+        velocities,
+        cell,
+        pbc,
+        calc_data,
+        info_data,
+        arrays_data,
+    ) = args
     atomic_numbers = remove_nan_rows(atomic_numbers)
     if positions is not None:
         positions = remove_nan_rows(positions)
@@ -136,7 +144,9 @@ def build_atoms(
     if calc_data is not None:
         calc_data = {key: remove_nan_rows(value) for key, value in calc_data.items()}
     if arrays_data is not None:
-        arrays_data = {key: remove_nan_rows(value) for key, value in arrays_data.items()}
+        arrays_data = {
+            key: remove_nan_rows(value) for key, value in arrays_data.items()
+        }
     if info_data is not None:
         info_data = {key: remove_nan_rows(value) for key, value in info_data.items()}
 
@@ -154,7 +164,7 @@ def build_atoms(
         atoms.calc = SinglePointCalculator(atoms=atoms)
         atoms.calc.results = calc_data
 
-    return atoms, index
+    return atoms
 
 
 def build_structures(
@@ -170,19 +180,16 @@ def build_structures(
     velocities = arrays_data.pop("velocity", None)
     atomic_numbers = arrays_data.pop("species")
     if atomic_numbers is not None:
-        # with concurrent.futures.ProcessPoolExecutor() as executor:
         for idx in range(len(atomic_numbers)):
-            # ruff thinks, this is less complex than doing it in place ... ??
-            atoms, _ = build_atoms(
-                atomic_numbers=atomic_numbers[idx],
-                positions=positions[idx] if positions is not None else None,
-                velocities=velocities[idx] if velocities is not None else None,
-                cell=cell[idx] if cell is not None else None,
-                pbc=pbc[idx] if isinstance(pbc[0], np.ndarray) else pbc,
-                arrays_data={key: value[idx] for key, value in arrays_data.items()},
-                calc_data={key: value[idx] for key, value in calc_data.items()},
-                info_data={key: value[idx] for key, value in info_data.items()},
-                index=idx,
+            args = (
+                atomic_numbers[idx],
+                positions[idx] if positions is not None else None,
+                velocities[idx] if velocities is not None else None,
+                cell[idx] if cell is not None else None,
+                pbc[idx] if isinstance(pbc[0], np.ndarray) else pbc,
+                {key: value[idx] for key, value in calc_data.items()},
+                {key: value[idx] for key, value in info_data.items()},
+                {key: value[idx] for key, value in arrays_data.items()},
             )
-            structures.append(atoms)
+            structures.append(build_atoms(args))
     return structures