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: {}