From cc19551b84f244f099e533c53b9bfaac903b0871 Mon Sep 17 00:00:00 2001 From: Daniel Zuegner Date: Thu, 27 Aug 2026 14:39:09 +0000 Subject: [PATCH 1/2] Package runtime resources for installed CLI Move sampling configurations into the Python package and lazily retrieve reference datasets from the MatterGen Hugging Face repository when they are not available in a source checkout.\n\nThe reference datasets must be uploaded to the Hugging Face repository before release.\n\nCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>\nCopilot-Session: 6d151412-121a-41f5-a0ab-bd6a3614c96e --- README.md | 9 ++--- mattergen/common/utils/globals.py | 2 +- mattergen/evaluation/reference/presets.py | 33 ++++++++++++++++--- .../sampling_conf}/csp.yaml | 8 ++--- .../sampling_conf}/default.yaml | 0 mattergen/tests/test_packaging_resources.py | 28 ++++++++++++++++ 6 files changed, 65 insertions(+), 15 deletions(-) rename {sampling_conf => mattergen/sampling_conf}/csp.yaml (97%) rename {sampling_conf => mattergen/sampling_conf}/default.yaml (100%) create mode 100644 mattergen/tests/test_packaging_resources.py diff --git a/README.md b/README.md index 1eb6eaa0..704d81c2 100644 --- a/README.md +++ b/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: diff --git a/mattergen/common/utils/globals.py b/mattergen/common/utils/globals.py index e35ec6a0..238e37a6 100644 --- a/mattergen/common/utils/globals.py +++ b/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", diff --git a/mattergen/evaluation/reference/presets.py b/mattergen/evaluation/reference/presets.py index f3f79f9c..cbe1bbc7 100644 --- a/mattergen/evaluation/reference/presets.py +++ b/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 diff --git a/sampling_conf/csp.yaml b/mattergen/sampling_conf/csp.yaml similarity index 97% rename from sampling_conf/csp.yaml rename to mattergen/sampling_conf/csp.yaml index b81a15ab..ff6c9401 100644 --- a/sampling_conf/csp.yaml +++ b/mattergen/sampling_conf/csp.yaml @@ -2,7 +2,7 @@ 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: @@ -10,7 +10,7 @@ sampler_partial: keep_conditioning_fn: _target_: mattergen.property_embeddings.SetConditionalEmbeddingType predictor_partials: - pos: + pos: _target_: mattergen.diffusion.wrapped.wrapped_predictors_correctors.WrappedAncestralSamplingPredictor _partial_: true cell: @@ -18,12 +18,12 @@ sampler_partial: _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 diff --git a/sampling_conf/default.yaml b/mattergen/sampling_conf/default.yaml similarity index 100% rename from sampling_conf/default.yaml rename to mattergen/sampling_conf/default.yaml diff --git a/mattergen/tests/test_packaging_resources.py b/mattergen/tests/test_packaging_resources.py new file mode 100644 index 00000000..f31fe777 --- /dev/null +++ b/mattergen/tests/test_packaging_resources.py @@ -0,0 +1,28 @@ +from pathlib import Path +from unittest.mock import patch + +from mattergen.common.utils.globals import DEFAULT_SAMPLING_CONFIG_PATH +from mattergen.evaluation.reference.presets import get_reference_dataset_path + + +def test_default_sampling_configs_are_packaged() -> None: + assert (DEFAULT_SAMPLING_CONFIG_PATH / "default.yaml").is_file() + assert (DEFAULT_SAMPLING_CONFIG_PATH / "csp.yaml").is_file() + + +def test_reference_dataset_is_downloaded_when_not_in_source_checkout(tmp_path: Path) -> None: + downloaded_path = tmp_path / "reference.gz" + with ( + patch("mattergen.evaluation.reference.presets._is_gzip_file", return_value=False), + patch( + "mattergen.evaluation.reference.presets.hf_hub_download", + return_value=str(downloaded_path), + ) as download, + ): + result = get_reference_dataset_path("reference_MP2020correction.gz") + + assert result == downloaded_path + download.assert_called_once_with( + repo_id="microsoft/mattergen", + filename="data-release/alex-mp/reference_MP2020correction.gz", + ) From 4425ceb0039473f619e7fc9efcdb76e3114d068f Mon Sep 17 00:00:00 2001 From: Daniel Zuegner Date: Thu, 27 Aug 2026 15:40:50 +0000 Subject: [PATCH 2/2] Remove packaging resource tests Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>\nCopilot-Session: 6d151412-121a-41f5-a0ab-bd6a3614c96e --- mattergen/tests/test_packaging_resources.py | 28 --------------------- 1 file changed, 28 deletions(-) delete mode 100644 mattergen/tests/test_packaging_resources.py diff --git a/mattergen/tests/test_packaging_resources.py b/mattergen/tests/test_packaging_resources.py deleted file mode 100644 index f31fe777..00000000 --- a/mattergen/tests/test_packaging_resources.py +++ /dev/null @@ -1,28 +0,0 @@ -from pathlib import Path -from unittest.mock import patch - -from mattergen.common.utils.globals import DEFAULT_SAMPLING_CONFIG_PATH -from mattergen.evaluation.reference.presets import get_reference_dataset_path - - -def test_default_sampling_configs_are_packaged() -> None: - assert (DEFAULT_SAMPLING_CONFIG_PATH / "default.yaml").is_file() - assert (DEFAULT_SAMPLING_CONFIG_PATH / "csp.yaml").is_file() - - -def test_reference_dataset_is_downloaded_when_not_in_source_checkout(tmp_path: Path) -> None: - downloaded_path = tmp_path / "reference.gz" - with ( - patch("mattergen.evaluation.reference.presets._is_gzip_file", return_value=False), - patch( - "mattergen.evaluation.reference.presets.hf_hub_download", - return_value=str(downloaded_path), - ) as download, - ): - result = get_reference_dataset_path("reference_MP2020correction.gz") - - assert result == downloaded_path - download.assert_called_once_with( - repo_id="microsoft/mattergen", - filename="data-release/alex-mp/reference_MP2020correction.gz", - )