chore(materials): Include callback to log progress when sampling (#214)

Enables custom logging of sampling progress.

Co-authored-by: Andrew Fowler <andrew31416@users.noreply.github.com>

Author: andrew31416

mattergen/common/utils/data_classes.py

@@ -6,7 +6,7 @@
 from dataclasses import asdict, dataclass, field
 from functools import cached_property
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, Protocol
 
 import numpy as np
 from huggingface_hub import hf_hub_download
@@ -154,3 +154,13 @@ def checkpoint_path(self) -> str:
             raise ValueError(f"Unrecognized load_epoch {self.load_epoch}")
         ckpt = ckpts[ckpt_ix]
         return ckpt
+
+
+class ProgressCallback(Protocol):
+    def __call__(self, progress: float):
+        """Callback which can be used to report progress on long-running inference.
+        
+        Args:
+            progress: Float between 0 and 1.
+        """
+        pass

mattergen/generator.py

@@ -32,6 +32,7 @@
 from mattergen.common.utils.globals import DEFAULT_SAMPLING_CONFIG_PATH, get_device
 from mattergen.diffusion.lightning_module import DiffusionLightningModule
 from mattergen.diffusion.sampling.pc_sampler import PredictorCorrector
+from mattergen.common.utils.data_classes import ProgressCallback
 
 
 def draw_samples_from_sampler(
@@ -41,6 +42,7 @@ def draw_samples_from_sampler(
     output_path: Path | None = None,
     cfg: DictConfig | None = None,
     record_trajectories: bool = True,
+    progress_callback: ProgressCallback | None = None,
 ) -> list[Structure]:
 
     # Dict
@@ -51,7 +53,9 @@ def draw_samples_from_sampler(
 
     all_samples_list = []
     all_trajs_list = []
-    for conditioning_data, mask in tqdm(condition_loader, desc="Generating samples"):
+    for batch_idx, (conditioning_data, mask) in enumerate(tqdm(condition_loader, desc="Generating samples")):
+        if progress_callback is not None:
+            progress_callback(progress=batch_idx / len(condition_loader))
 
         # generate samples
         if record_trajectories:
@@ -60,6 +64,11 @@ def draw_samples_from_sampler(
         else:
             sample, mean = sampler.sample(conditioning_data, mask)
         all_samples_list.extend(mean.to_data_list())
+
+    if progress_callback is not None:
+        # log 100% progress
+        progress_callback(progress=1.0)
+
     all_samples = collate(all_samples_list)
     assert isinstance(all_samples, ChemGraph)
     lengths, angles = lattice_matrix_to_params_torch(all_samples.cell)
@@ -199,6 +208,9 @@ class CrystalGenerator:
     _model: DiffusionLightningModule | None = None
     _cfg: DictConfig | None = None
 
+    # can be used to monitor progress of generation
+    progress_callback: ProgressCallback | None = None
+
     def __post_init__(self) -> None:
         assert self.num_atoms_distribution in NUM_ATOMS_DISTRIBUTIONS, (
             f"num_atoms_distribution must be one of {list(NUM_ATOMS_DISTRIBUTIONS.keys())}, "
@@ -374,6 +386,7 @@ def generate(
             output_path=Path(output_dir),
             properties_to_condition_on=self.properties_to_condition_on,
             record_trajectories=self.record_trajectories,
+            progress_callback=self.progress_callback,
         )
 
         return generated_structures

mattergen/scripts/generate.py

@@ -9,7 +9,7 @@
 from pymatgen.core.structure import Structure
 
 from mattergen.common.data.types import TargetProperty
-from mattergen.common.utils.data_classes import PRETRAINED_MODEL_NAME, MatterGenCheckpointInfo
+from mattergen.common.utils.data_classes import PRETRAINED_MODEL_NAME, MatterGenCheckpointInfo, ProgressCallback
 from mattergen.generator import CrystalGenerator
 
 
@@ -29,6 +29,7 @@ def main(
     diffusion_guidance_factor: float | None = None,
     strict_checkpoint_loading: bool = True,
     target_compositions: list[dict[str, int]] | None = None,
+    progress_callback: ProgressCallback | None = None,
 ) -> list[Structure]:
     """
     Evaluate diffusion model against molecular metrics.
@@ -46,7 +47,7 @@ def main(
         strict_checkpoint_loading: Whether to raise an exception when not all parameters from the checkpoint can be matched to the model.
         target_compositions: List of dictionaries with target compositions to condition on. Each dictionary should have the form `{element: number_of_atoms}`. If None, the target compositions are not conditioned on.
            Only supported for models trained for crystal structure prediction (CSP) (default: None)
-
+        progress_callback: Optional callback function that takes in a single float argument representing the progress of the generation process (between 0 and 1).
     NOTE: When specifying dictionary values via the CLI, make sure there is no whitespace between the key and value, e.g., `--properties_to_condition_on={key1:value1}`.
     """
     assert (
@@ -94,6 +95,7 @@ def main(
             diffusion_guidance_factor if diffusion_guidance_factor is not None else 0.0
         ),
         target_compositions_dict=target_compositions,
+        progress_callback=progress_callback,
     )
     return generator.generate(output_dir=Path(output_path))