Adds a calc_many method that returns a generator using joblib.

matcalc/base.py

@@ -4,6 +4,8 @@
 import abc
 from typing import TYPE_CHECKING
 
+from joblib import Parallel, delayed
+
 if TYPE_CHECKING:
     from collections.abc import Generator, Sequence
 
@@ -25,16 +27,17 @@ def calc(self, structure: Structure) -> dict:
         Returns: {"prop name": value}
         """
 
-    def calc_many(self, structures: Sequence[Structure]) -> Generator[dict, None, None]:
+    def calc_many(self, structures: Sequence[Structure], n_jobs: None | int = None) -> Generator[dict, None, None]:
         """
         Performs calc on many structures. The return type is a generator given that the calc method can potentially be
         reasonably expensive. It is trivial to convert the generator to a list/tuple.
 
         Args:
-            structures: List or generator of Structures
+            structures: List or generator of Structures.
+            n_jobs: Number of processes to use for multiprocessing.Pool. Defaults to None, i.e., all CPUs.
 
         Returns:
             Generator of dicts.
         """
-        for structure in structures:
-            yield self.calc(structure)
+        parallel = Parallel(n_jobs=2, return_as="generator")
+        return parallel(delayed(self.calc)(s) for s in structures)

requirements.txt

@@ -1,2 +1,3 @@
 pymatgen==2023.7.11
-ase==3.22.1
+ase==3.22.1
+joblib==1.3.1

setup.py

@@ -31,12 +31,8 @@
         "deep learning",
     ],
     packages=find_packages(),
-    package_data={
-    },
-    install_requires=(
-        "ase",
-        "pymatgen",
-    ),
+    package_data={},
+    install_requires=("ase", "pymatgen", "joblib"),
     classifiers=[
         "Development Status :: 4 - Beta",
         "Intended Audience :: Science/Research",

tests/test_relaxation.py

@@ -10,8 +10,16 @@ def test_RelaxCalc(LiFePO4):
     potential = matgl.load_model("M3GNet-MP-2021.2.8-PES")
     calculator = M3GNetCalculator(potential=potential, stress_weight=0.01)
 
-    calc = RelaxCalc(calculator)
-    results = calc.calc(LiFePO4)
+    pcalc = RelaxCalc(calculator)
+    results = pcalc.calc(LiFePO4)
     assert results["a"] == pytest.approx(4.755711375217371)
     assert results["b"] == pytest.approx(6.131614236614623)
     assert results["c"] == pytest.approx(10.43859339794175)
+    assert results["alpha"] == pytest.approx(90, abs=1)
+    assert results["beta"] == pytest.approx(90, abs=1)
+    assert results["gamma"] == pytest.approx(90, abs=1)
+    assert results["volume"] == pytest.approx(results["a"] * results["b"] * results["c"], abs=0.1)
+
+    results = list(pcalc.calc_many([LiFePO4] * 2))
+    assert len(results) == 2
+    assert results[-1]["a"] == pytest.approx(4.755711375217371)