Implement generator return type for Monte Carlo sims (#76)

tce/__init__.py

@@ -446,7 +446,7 @@ def callback(step_: int, num_steps_: int):
 
 """
 
-__version__ = "1.0.0"
+__version__ = "1.0.1"
 __authors__ = ["Jacob Jeffries"]
 
 __url__ = "https://github.com/MUEXLY/tce-lib"

tce/monte_carlo.py

@@ -4,7 +4,7 @@
 """
 
 
-from typing import Optional, Callable, TypeAlias, Sequence
+from typing import Generator, Iterator, Optional, Callable, TypeAlias, Sequence
 import logging
 import warnings
 
@@ -149,8 +149,9 @@ def monte_carlo(
     generator: Optional[np.random.Generator] = None,
     mc_step: Optional[Callable[[Atoms], Atoms]] = None,
     energy_modifier: Optional[Callable[[Atoms, Atoms], float]] = None,
-    callback: Optional[Callable[[int, int], None]] = None
-) -> list[Atoms]:
+    callback: Optional[Callable[[int, int], None]] = None,
+    return_generator: bool = False
+) -> list[Atoms] | Iterator[Atoms]:
 
     r"""
     New Monte Carlo simulation function that uses the `transform_model` function to transform the model to predict
@@ -191,6 +192,15 @@ def mc_step(atoms: Atoms) -> Atoms:
             MC, and will simply use the energy as the thermodynamic potential. See the grand canonical MC example 
             [here](https://github.com/MUEXLY/tce-lib/blob/main/examples/1-copper-nickel-mc2.py) for how to use this 
             argument to change the ensemble sampled.
+
+        callback (Callable[[int, int], None]):
+            Callback function to call at each step. If not specified, the function will log the current step 
+            number and total number of steps. 
+
+        return_generator (bool): 
+            If `True`, the function will return a generator so that simulation frames can be computed lazily within
+            an iteration loop. This is useful for the analysis of very long simulations, such as those executed on HPC clusters.
+            Defaults to `False` for backwards compatibility.
     """
 
     if not generator:
@@ -239,29 +249,34 @@ def energy_modifier(initial: Atoms, final: Atoms) -> float:
     )
     if isinstance(energy, np.ndarray):
         energy = energy.item()
+
+    def _generating_fn(initial_configuration: Atoms, energy: float) -> Generator[Atoms, None, None]: 
+        """Wrap the generator logic in a function so that we can return both output types."""
+
+        for step in range(num_steps):
+            callback(step, num_steps)
+
+            if not step % save_every:
+                to_save = initial_configuration.copy()
+                to_save.info["energy"] = energy
+                yield to_save
+                LOGGER.info(f"saved configuration at step {step:.0f}/{num_steps:.0f}")
+
+            new_configuration = mc_step(initial_configuration)
+            feature_diff = tce_calculator.get_feature_vector_difference(
+                initial_configuration, new_configuration
+            ).reshape(1, -1)
+            energy_diff = transformed_model.predict(feature_diff)
+            energy_diff += energy_modifier(initial_configuration, new_configuration)
+
+            if not isinstance(energy_diff, float):
+                energy_diff = energy_diff.item()
+            if np.exp(-beta_values[step] * energy_diff) > 1.0 - generator.random():
+                LOGGER.debug(f"move accepted with energy difference {energy_diff}")
+                initial_configuration = new_configuration
+                energy += energy_diff
+
+    if return_generator: 
+        return _generating_fn(initial_configuration, energy)
 
-    trajectory = []
-    for step in range(num_steps):
-        callback(step, num_steps)
-
-        if not step % save_every:
-            to_save = initial_configuration.copy()
-            to_save.info["energy"] = energy
-            trajectory.append(to_save)
-            LOGGER.info(f"saved configuration at step {step:.0f}/{num_steps:.0f}")
-
-        new_configuration = mc_step(initial_configuration)
-        feature_diff = tce_calculator.get_feature_vector_difference(
-            initial_configuration, new_configuration
-        ).reshape(1, -1)
-        energy_diff = transformed_model.predict(feature_diff)
-        energy_diff += energy_modifier(initial_configuration, new_configuration)
-
-        if not isinstance(energy_diff, float):
-            energy_diff = energy_diff.item()
-        if np.exp(-beta_values[step] * energy_diff) > 1.0 - generator.random():
-            LOGGER.debug(f"move accepted with energy difference {energy_diff}")
-            initial_configuration = new_configuration
-            energy += energy_diff
-
-    return trajectory
+    return list(_generating_fn(initial_configuration, energy))

test_lib.py

@@ -1,4 +1,5 @@
 from typing import Callable
+from types import GeneratorType
 from tempfile import TemporaryDirectory
 from pathlib import Path
 import pickle
@@ -591,3 +592,55 @@ def test_energy_diff_transform(model):
     feature_diff = second_feature_vector - first_feature_vector
     energy_diff_from_delta = calc.models["energy"].predict(feature_diff.reshape(1, -1)).squeeze()
     assert np.isclose(energy_diff, energy_diff_from_delta), f"{energy_diff}, {energy_diff_from_delta}"
+
+def test_monte_carlo_generator(): 
+    """Test if the list and generator return the same results"""
+
+    rng = np.random.default_rng(seed=0)
+    fe_cohesive_energy = 4.0
+    cr_cohesive_energy = 3.0
+
+    pure_fe = build.bulk("Fe", crystalstructure="bcc", a=3.0, cubic=True).repeat((5, 5, 5))
+    pure_fe.calc = SinglePointCalculator(pure_fe, energy=len(pure_fe) * -fe_cohesive_energy)
+    pure_cr = build.bulk("Cr", crystalstructure="bcc", a=3.0, cubic=True).repeat((5, 5, 5))
+    pure_cr.calc = SinglePointCalculator(pure_cr, energy=len(pure_cr) * -cr_cohesive_energy)
+    mixture = pure_fe.copy()
+    mixture.symbols = rng.choice(["Fe", "Cr"], size=len(mixture))
+    mixture.calc = SinglePointCalculator(
+        mixture,
+        energy=len(mixture) * -0.5 * (fe_cohesive_energy + cr_cohesive_energy)
+    )
+
+    calc = TCECalculator(
+        neighbor_cutoffs=CUTOFFS["bcc"][:2] * 3.0,
+        many_body_features=[(0, 0, 1)],
+        species=["Fe", "Cr"],
+        models={"energy": RidgeCV(fit_intercept=False)}
+    ).train([pure_fe, pure_cr, mixture])
+
+    new_mixture = pure_fe.copy()
+    new_mixture.symbols = rng.choice(["Fe", "Cr"], size=len(new_mixture))
+    new_mixture.calc = calc
+
+    num_steps = 10
+    beta = 11.1
+
+    list_results = monte_carlo(
+        initial_configuration=new_mixture,
+        tce_calculator=calc,
+        num_steps=num_steps,
+        beta=beta,
+        return_generator=False
+    )
+
+    generator_results = monte_carlo(
+        initial_configuration=new_mixture,
+        tce_calculator=calc,
+        num_steps=num_steps,
+        beta=beta,
+        return_generator=True
+    )
+
+    assert isinstance(list_results, list)
+    assert isinstance(generator_results, GeneratorType)  # check if generator
+    assert list_results == list(generator_results)