Package runtime resources for installed CLI

README.md

@@ -123,14 +123,13 @@ mattergen-generate $RESULTS_PATH --pretrained-name=$MODEL_NAME --batch_size=16 -
 
 Once you have generated a list of structures contained in `$RESULTS_PATH` (either using MatterGen or another method), you can relax the structures using the default MatterSim machine learning force field (see [repository](https://github.com/microsoft/mattersim)) and compute novelty, uniqueness, stability (using energy estimated by MatterSim), and other metrics via the following command:
 ```bash
-git lfs pull -I data-release/alex-mp/reference_MP2020correction.gz --exclude=""  # first download the MP2020 reference dataset from Git LFS
 mattergen-evaluate --structures_path=$RESULTS_PATH --relax=True --structure_matcher='disordered' --save_as="$RESULTS_PATH/metrics.json"
 ```
+The default MP2020 reference dataset is downloaded from Hugging Face on first use and reused from the local cache.
 
 If you want to use the reference dataset while applying the TRI2024 correction scheme (recommended), instead run the following:
 ```bash
-git lfs pull -I data-release/alex-mp/reference_TRI2024correction.gz --exclude=""  # ownload the TRI2024 reference datasets
-mattergen-evaluate --structures_path=$RESULTS_PATH --relax=True --structure_matcher='disordered' --save_as="$RESULTS_PATH/metrics.json" --reference_dataset_path="data-release/alex-mp/reference_TRI2024correction.gz"
+mattergen-evaluate --structures_path=$RESULTS_PATH --relax=True --structure_matcher='disordered' --save_as="$RESULTS_PATH/metrics.json" --energy_correction_scheme="TRI2024"
 ```
 
 This script will write `metrics.json` containing the metric results to `$RESULTS_PATH` and will print it to your console.
@@ -143,7 +142,6 @@ This script will write `metrics.json` containing the metric results to `$RESULTS
 
 If, instead, you have relaxed the structures and obtained the relaxed total energies via another mean (e.g., DFT), you can evaluate the metrics via:
 ```bash
-git lfs pull -I data-release/alex-mp/reference_MP2020correction.gz --exclude=""  # first download the reference dataset from Git LFS
 mattergen-evaluate --structures_path=$RESULTS_PATH --energies_path='energies.npy' --relax=False --structure_matcher='disordered' --save_as='metrics'
 ```
 This script will try to read structures from disk in the following precedence order:
@@ -156,8 +154,7 @@ Here, we expect `energies.npy` to be a numpy array with the entries being `float
 > [!IMPORTANT]
 > For any task beyond benchmarking against existing literature, we recommend using the TRI2024 correction scheme and reference dataset. To do so, run:
 ```bash
-git lfs pull -I data-release/alex-mp/reference_TRI2024correction.gz --exclude=""  # first download the reference dataset from Git LFS
-mattergen-evaluate --structures_path=$RESULTS_PATH --energies_path='energies.npy' --relax=False --structure_matcher='disordered' --save_as='metrics' --energy_correction_scheme="TRI2024" --reference_dataset_path="data-release/alex-mp/reference_TRI2024correction.gz" 
+mattergen-evaluate --structures_path=$RESULTS_PATH --energies_path='energies.npy' --relax=False --structure_matcher='disordered' --save_as='metrics' --energy_correction_scheme="TRI2024"
 ```
 
 If you want to save the relaxed structures, toghether with their energies, forces, and stresses, add `--structures_output_path=YOUR_PATH` to the script call, like so:

mattergen/common/utils/globals.py

@@ -38,7 +38,7 @@ def get_pyg_device() -> torch.device:
 # Set environment variable PROJECT_ROOT so that hydra / OmegaConf can access it.
 os.environ["PROJECT_ROOT"] = str(MODELS_PROJECT_ROOT)  # for hydra
 
-DEFAULT_SAMPLING_CONFIG_PATH = Path(__file__).resolve().parents[3] / "sampling_conf"
+DEFAULT_SAMPLING_CONFIG_PATH = MODELS_PROJECT_ROOT / "sampling_conf"
 PROPERTY_SOURCE_IDS = [
     "dft_mag_density",
     "dft_bulk_modulus",

mattergen/evaluation/reference/presets.py

@@ -4,9 +4,36 @@
 from functools import cached_property
 from pathlib import Path
 
+from huggingface_hub import hf_hub_download
+
 from mattergen.evaluation.reference.reference_dataset import ReferenceDataset
 from mattergen.evaluation.reference.reference_dataset_serializer import LMDBGZSerializer
 
+REFERENCE_REPOSITORY_NAME = "microsoft/mattergen"
+REFERENCE_DATASET_DIRECTORY = "data-release/alex-mp"
+
+
+def _is_gzip_file(path: Path) -> bool:
+    try:
+        with path.open("rb") as stream:
+            return stream.read(2) == b"\x1f\x8b"
+    except OSError:
+        return False
+
+
+def get_reference_dataset_path(filename: str) -> Path:
+    source_checkout_path = (
+        Path(__file__).resolve().parents[3] / REFERENCE_DATASET_DIRECTORY / filename
+    )
+    if _is_gzip_file(source_checkout_path):
+        return source_checkout_path
+    return Path(
+        hf_hub_download(
+            repo_id=REFERENCE_REPOSITORY_NAME,
+            filename=f"{REFERENCE_DATASET_DIRECTORY}/{filename}",
+        )
+    )
+
 
 class ReferenceMP2020Correction(ReferenceDataset):
     """Reference dataset using the MP2020 Energy Correction scheme.
@@ -20,9 +47,8 @@ def __init__(self):
 
     @classmethod
     def from_preset(cls) -> "ReferenceMP2020Correction":
-        current_dir = Path(__file__).parent
         return LMDBGZSerializer().deserialize(
-            f"{current_dir}/../../../data-release/alex-mp/reference_MP2020correction.gz"
+            get_reference_dataset_path("reference_MP2020correction.gz")
         )
 
     @cached_property
@@ -43,9 +69,8 @@ def __init__(self):
 
     @classmethod
     def from_preset(cls) -> "ReferenceTRI2024Correction":
-        current_dir = Path(__file__).parent
         return LMDBGZSerializer().deserialize(
-            f"{current_dir}/../../../data-release/alex-mp/reference_TRI2024correction.gz"
+            get_reference_dataset_path("reference_TRI2024correction.gz")
         )
 
     @cached_property

sampling_conf/csp.yaml → mattergen/sampling_conf/csp.yaml

@@ -2,28 +2,28 @@ sampler_partial:
   _target_: mattergen.diffusion.sampling.classifier_free_guidance.GuidedPredictorCorrector.from_pl_module
   N: 1000
   eps_t: ${eval:'1/${.N}'}
-  
+
   _partial_: true
   guidance_scale: 0.0
   remove_conditioning_fn:
     _target_: mattergen.property_embeddings.SetUnconditionalEmbeddingType
   keep_conditioning_fn:
     _target_: mattergen.property_embeddings.SetConditionalEmbeddingType
   predictor_partials:
-    pos: 
+    pos:
       _target_: mattergen.diffusion.wrapped.wrapped_predictors_correctors.WrappedAncestralSamplingPredictor
       _partial_: true
     cell:
       _target_: mattergen.common.diffusion.predictors_correctors.LatticeAncestralSamplingPredictor
       _partial_: true
 
   corrector_partials:
-    pos: 
+    pos:
       _target_: mattergen.diffusion.wrapped.wrapped_predictors_correctors.WrappedLangevinCorrector
       _partial_: true
       max_step_size: 1e6
       snr: 0.4
-    cell: 
+    cell:
       _target_: mattergen.common.diffusion.predictors_correctors.LatticeLangevinDiffCorrector
       _partial_: true
       max_step_size: 1e6