Skip to content

Make GemNet scale-file configuration portable #262

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Daniel Zügner (danielzuegner) merged 2 commits into from
Aug 27, 2026
Merged

Make GemNet scale-file configuration portable #262

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1031c45
Make GemNet scale-file configuration portable
Aug 27, 2026
ea5007f
Remove GemNet scale-file tests
Aug 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 0 additions & 2 deletions checkpoints/chemical_system/config.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
2 changes: 0 additions & 2 deletions checkpoints/chemical_system_energy_above_hull/config.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
2 changes: 0 additions & 2 deletions checkpoints/dft_band_gap/config.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
2 changes: 0 additions & 2 deletions checkpoints/dft_mag_density/config.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
2 changes: 0 additions & 2 deletions checkpoints/dft_mag_density_hhi_score/config.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
1 change: 0 additions & 1 deletion checkpoints/mattergen_base/config.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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: {}
Expand Down
2 changes: 0 additions & 2 deletions checkpoints/ml_bulk_modulus/config.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
1 change: 0 additions & 1 deletion checkpoints/mp_20_base/config.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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: {}
Expand Down
2 changes: 0 additions & 2 deletions checkpoints/space_group/config.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
23 changes: 21 additions & 2 deletions mattergen/common/gemnet/gemnet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down
2 changes: 0 additions & 2 deletions mattergen/common/tests/gemnet_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 ###

Expand All @@ -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,
)

Expand Down
1 change: 0 additions & 1 deletion mattergen/conf/lightning_module/diffusion_module/model/mattergen.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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: {}
Expand Down
Loading