From 1031c4580928c660f9e105f9f2ea613ebdd5cef9 Mon Sep 17 00:00:00 2001 From: Daniel Zuegner Date: Thu, 27 Aug 2026 14:50:28 +0000 Subject: [PATCH 1/2] Make GemNet scale-file configuration portable Use a package-relative default, honor explicit custom paths, preserve compatibility with legacy checkpoint configs, and remove internal paths from shipped configurations.\n\nThe corresponding Hugging Face checkpoint configs must be refreshed before release.\n\nCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>\nCopilot-Session: 6d151412-121a-41f5-a0ab-bd6a3614c96e --- checkpoints/chemical_system/config.yaml | 2 -- .../config.yaml | 2 -- checkpoints/dft_band_gap/config.yaml | 2 -- checkpoints/dft_mag_density/config.yaml | 2 -- .../dft_mag_density_hhi_score/config.yaml | 2 -- checkpoints/mattergen_base/config.yaml | 1 - checkpoints/ml_bulk_modulus/config.yaml | 2 -- checkpoints/mp_20_base/config.yaml | 1 - checkpoints/space_group/config.yaml | 2 -- mattergen/common/gemnet/gemnet.py | 23 ++++++++++++++++-- mattergen/common/tests/gemnet_test.py | 2 -- .../diffusion_module/model/mattergen.yaml | 1 - mattergen/tests/test_gemnet_config.py | 24 +++++++++++++++++++ 13 files changed, 45 insertions(+), 21 deletions(-) create mode 100644 mattergen/tests/test_gemnet_config.py diff --git a/checkpoints/chemical_system/config.yaml b/checkpoints/chemical_system/config.yaml index 671d0abb..b4578a58 100644 --- a/checkpoints/chemical_system/config.yaml +++ b/checkpoints/chemical_system/config.yaml @@ -21,7 +21,6 @@ adapter: num_targets: 1 otf_graph: true regress_stress: true - scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json hidden_dim: 512 property_embeddings: {} property_embeddings_adapt: @@ -150,7 +149,6 @@ lightning_module: num_targets: 1 otf_graph: true regress_stress: true - scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json hidden_dim: 512 property_embeddings: {} property_embeddings_adapt: diff --git a/checkpoints/chemical_system_energy_above_hull/config.yaml b/checkpoints/chemical_system_energy_above_hull/config.yaml index e0c95365..c87c61ba 100644 --- a/checkpoints/chemical_system_energy_above_hull/config.yaml +++ b/checkpoints/chemical_system_energy_above_hull/config.yaml @@ -22,7 +22,6 @@ adapter: num_targets: 1 otf_graph: true regress_stress: true - scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json hidden_dim: 512 property_embeddings: {} property_embeddings_adapt: @@ -166,7 +165,6 @@ lightning_module: num_targets: 1 otf_graph: true regress_stress: true - scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json hidden_dim: 512 property_embeddings: {} property_embeddings_adapt: diff --git a/checkpoints/dft_band_gap/config.yaml b/checkpoints/dft_band_gap/config.yaml index 1eebe5d5..d2853ecd 100644 --- a/checkpoints/dft_band_gap/config.yaml +++ b/checkpoints/dft_band_gap/config.yaml @@ -21,7 +21,6 @@ adapter: num_targets: 1 otf_graph: true regress_stress: true - scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json hidden_dim: 512 property_embeddings: {} property_embeddings_adapt: @@ -150,7 +149,6 @@ lightning_module: num_targets: 1 otf_graph: true regress_stress: true - scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json hidden_dim: 512 property_embeddings: {} property_embeddings_adapt: diff --git a/checkpoints/dft_mag_density/config.yaml b/checkpoints/dft_mag_density/config.yaml index bcb67a9b..5c9e5be8 100644 --- a/checkpoints/dft_mag_density/config.yaml +++ b/checkpoints/dft_mag_density/config.yaml @@ -21,7 +21,6 @@ adapter: num_targets: 1 otf_graph: true regress_stress: true - scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json hidden_dim: 512 property_embeddings: {} property_embeddings_adapt: @@ -150,7 +149,6 @@ lightning_module: num_targets: 1 otf_graph: true regress_stress: true - scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json hidden_dim: 512 property_embeddings: {} property_embeddings_adapt: diff --git a/checkpoints/dft_mag_density_hhi_score/config.yaml b/checkpoints/dft_mag_density_hhi_score/config.yaml index f86c90fd..3b1d2344 100644 --- a/checkpoints/dft_mag_density_hhi_score/config.yaml +++ b/checkpoints/dft_mag_density_hhi_score/config.yaml @@ -22,7 +22,6 @@ adapter: num_targets: 1 otf_graph: true regress_stress: true - scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json hidden_dim: 512 property_embeddings: {} property_embeddings_adapt: @@ -166,7 +165,6 @@ lightning_module: num_targets: 1 otf_graph: true regress_stress: true - scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json hidden_dim: 512 property_embeddings: {} property_embeddings_adapt: diff --git a/checkpoints/mattergen_base/config.yaml b/checkpoints/mattergen_base/config.yaml index 04012f3a..1aa63e72 100644 --- a/checkpoints/mattergen_base/config.yaml +++ b/checkpoints/mattergen_base/config.yaml @@ -118,7 +118,6 @@ lightning_module: num_targets: 1 otf_graph: true regress_stress: true - scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json hidden_dim: 512 property_embeddings: {} property_embeddings_adapt: {} diff --git a/checkpoints/ml_bulk_modulus/config.yaml b/checkpoints/ml_bulk_modulus/config.yaml index 413ffcf1..ad19ad6b 100644 --- a/checkpoints/ml_bulk_modulus/config.yaml +++ b/checkpoints/ml_bulk_modulus/config.yaml @@ -21,7 +21,6 @@ adapter: num_targets: 1 otf_graph: true regress_stress: true - scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json hidden_dim: 512 property_embeddings: {} property_embeddings_adapt: @@ -152,7 +151,6 @@ lightning_module: num_targets: 1 otf_graph: true regress_stress: true - scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json hidden_dim: 512 property_embeddings: {} property_embeddings_adapt: diff --git a/checkpoints/mp_20_base/config.yaml b/checkpoints/mp_20_base/config.yaml index 814a3656..16e80444 100644 --- a/checkpoints/mp_20_base/config.yaml +++ b/checkpoints/mp_20_base/config.yaml @@ -120,7 +120,6 @@ lightning_module: num_targets: 1 otf_graph: true regress_stress: true - scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json hidden_dim: 512 property_embeddings: {} property_embeddings_adapt: {} diff --git a/checkpoints/space_group/config.yaml b/checkpoints/space_group/config.yaml index 2861dcee..4dc5627e 100644 --- a/checkpoints/space_group/config.yaml +++ b/checkpoints/space_group/config.yaml @@ -21,7 +21,6 @@ adapter: num_targets: 1 otf_graph: true regress_stress: true - scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json hidden_dim: 512 property_embeddings: {} property_embeddings_adapt: @@ -150,7 +149,6 @@ lightning_module: num_targets: 1 otf_graph: true regress_stress: true - scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json hidden_dim: 512 property_embeddings: {} property_embeddings_adapt: diff --git a/mattergen/common/gemnet/gemnet.py b/mattergen/common/gemnet/gemnet.py index 0e500ad9..51dbe649 100644 --- a/mattergen/common/gemnet/gemnet.py +++ b/mattergen/common/gemnet/gemnet.py @@ -5,7 +5,9 @@ Adapted from https://github.com/FAIR-Chem/fairchem/blob/main/src/fairchem/core/models/gemnet/gemnet.py. """ +import warnings from dataclasses import dataclass +from pathlib import Path from typing import Optional, Tuple # import numpy as np @@ -37,6 +39,23 @@ from mattergen.common.utils.globals import MODELS_PROJECT_ROOT, get_device, get_pyg_device from mattergen.common.utils.lattice_score import edge_score_to_lattice_score_frac_symmetric +DEFAULT_SCALE_FILE = MODELS_PROJECT_ROOT / "common/gemnet/gemnet-dT.json" +LEGACY_SCALE_FILE = Path("/scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json") + + +def resolve_scale_file(scale_file: str | Path | None) -> str: + if scale_file is None: + return str(DEFAULT_SCALE_FILE) + if Path(scale_file) == LEGACY_SCALE_FILE: + warnings.warn( + "The checkpoint contains a legacy internal scale_file path; using the packaged " + "gemnet-dT.json instead.", + DeprecationWarning, + stacklevel=2, + ) + return str(DEFAULT_SCALE_FILE) + return str(scale_file) + @dataclass(frozen=True) class ModelOutput: @@ -206,13 +225,13 @@ def __init__( otf_graph: bool = False, output_init: str = "HeOrthogonal", activation: str = "swish", + scale_file: str | Path | None = None, max_cell_images_per_dim: int = 5, encoder_mode: bool = False, # **kwargs, ): super().__init__() - scale_file = f"{MODELS_PROJECT_ROOT}/common/gemnet/gemnet-dT.json" - assert scale_file is not None, "`scale_file` is required." + scale_file = resolve_scale_file(scale_file) self.encoder_mode = encoder_mode self.num_targets = num_targets diff --git a/mattergen/common/tests/gemnet_test.py b/mattergen/common/tests/gemnet_test.py index ee8ce64d..7c48bb56 100644 --- a/mattergen/common/tests/gemnet_test.py +++ b/mattergen/common/tests/gemnet_test.py @@ -20,7 +20,6 @@ lattice_params_to_matrix_torch, ) from mattergen.common.utils.eval_utils import make_structure -from mattergen.common.utils.globals import MODELS_PROJECT_ROOT ### UTILS ### @@ -37,7 +36,6 @@ def get_model(**kwargs) -> GemNetT: emb_size_trip=4, emb_size_bil_trip=4, otf_graph=True, - scale_file=f"{MODELS_PROJECT_ROOT}/common/gemnet/gemnet-dT.json", **kwargs, ) diff --git a/mattergen/conf/lightning_module/diffusion_module/model/mattergen.yaml b/mattergen/conf/lightning_module/diffusion_module/model/mattergen.yaml index b3f190b3..db8e8d1d 100644 --- a/mattergen/conf/lightning_module/diffusion_module/model/mattergen.yaml +++ b/mattergen/conf/lightning_module/diffusion_module/model/mattergen.yaml @@ -16,7 +16,6 @@ gemnet: num_blocks: 4 regress_stress: true otf_graph: true - scale_file: ${oc.env:PROJECT_ROOT}/common/gemnet/gemnet-dT.json denoise_atom_types: true atom_type_diffusion: mask property_embeddings_adapt: {} diff --git a/mattergen/tests/test_gemnet_config.py b/mattergen/tests/test_gemnet_config.py new file mode 100644 index 00000000..456225da --- /dev/null +++ b/mattergen/tests/test_gemnet_config.py @@ -0,0 +1,24 @@ +from pathlib import Path + +import pytest + +from mattergen.common.gemnet.gemnet import DEFAULT_SCALE_FILE, resolve_scale_file + + +def test_scale_file_defaults_to_packaged_resource() -> None: + assert resolve_scale_file(None) == str(DEFAULT_SCALE_FILE) + + +def test_custom_scale_file_is_honored(tmp_path: Path) -> None: + scale_file = tmp_path / "scales.json" + + assert resolve_scale_file(scale_file) == str(scale_file) + + +def test_legacy_internal_scale_file_uses_packaged_resource() -> None: + with pytest.warns(DeprecationWarning, match="legacy internal scale_file"): + result = resolve_scale_file( + "/scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json" + ) + + assert result == str(DEFAULT_SCALE_FILE) From ea5007f9b57ecc5f9ca373ca0d3b0cf1a7b96d09 Mon Sep 17 00:00:00 2001 From: Daniel Zuegner Date: Thu, 27 Aug 2026 15:17:36 +0000 Subject: [PATCH 2/2] Remove GemNet scale-file tests Address review feedback.\n\nCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>\nCopilot-Session: 6d151412-121a-41f5-a0ab-bd6a3614c96e --- mattergen/tests/test_gemnet_config.py | 24 ------------------------ 1 file changed, 24 deletions(-) delete mode 100644 mattergen/tests/test_gemnet_config.py diff --git a/mattergen/tests/test_gemnet_config.py b/mattergen/tests/test_gemnet_config.py deleted file mode 100644 index 456225da..00000000 --- a/mattergen/tests/test_gemnet_config.py +++ /dev/null @@ -1,24 +0,0 @@ -from pathlib import Path - -import pytest - -from mattergen.common.gemnet.gemnet import DEFAULT_SCALE_FILE, resolve_scale_file - - -def test_scale_file_defaults_to_packaged_resource() -> None: - assert resolve_scale_file(None) == str(DEFAULT_SCALE_FILE) - - -def test_custom_scale_file_is_honored(tmp_path: Path) -> None: - scale_file = tmp_path / "scales.json" - - assert resolve_scale_file(scale_file) == str(scale_file) - - -def test_legacy_internal_scale_file_uses_packaged_resource() -> None: - with pytest.warns(DeprecationWarning, match="legacy internal scale_file"): - result = resolve_scale_file( - "/scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json" - ) - - assert result == str(DEFAULT_SCALE_FILE)