model version for Potential class is added

dev/refactor.py

@@ -2,7 +2,6 @@
 from __future__ import annotations
 
 import torch
-
 from matgl.models import MEGNet
 
 # model_path = "pretrained_models/MEGNet-MP-2019.4.1-BandGap-mfi"

pretrained_models/M3GNet-MP-2021.2.8-DIRECT-PES/model.json

@@ -1,7 +1,7 @@
 {
     "@class": "Potential",
     "@module": "matgl.apps.pes",
-    "@model_version": 1,
+    "@model_version": 2,
     "metadata": null,
     "kwargs": {
         "model": {

pretrained_models/M3GNet-MP-2021.2.8-PES/model.json

@@ -1,7 +1,7 @@
 {
     "@class": "Potential",
     "@module": "matgl.apps.pes",
-    "@model_version": 1,
+    "@model_version": 2,
     "metadata": null,
     "kwargs": {
         "model": {

src/matgl/apps/pes.py

@@ -19,7 +19,7 @@
 class Potential(nn.Module, IOMixIn):
     """A class representing an interatomic potential."""
 
-    __version__ = 1
+    __version__ = 2
 
     def __init__(
         self,

tests/apps/test_pes.py

@@ -1,14 +1,13 @@
 from __future__ import annotations
 
+import matgl
 import numpy as np
 import pytest
 import torch
-from pymatgen.core import Lattice, Structure
-
-import matgl
 from matgl.apps.pes import Potential
 from matgl.ext.pymatgen import Structure2Graph, get_element_list
 from matgl.models._m3gnet import M3GNet
+from pymatgen.core import Lattice, Structure
 
 
 @pytest.fixture()

tests/conftest.py

@@ -10,16 +10,15 @@
 """
 from __future__ import annotations
 
+import matgl
 import pytest
 import torch
-from pymatgen.core import Lattice, Molecule, Structure
-from pymatgen.util.testing import PymatgenTest
-
-import matgl
 from matgl.ext.pymatgen import Molecule2Graph, Structure2Graph, get_element_list
 from matgl.graph.compute import (
     compute_pair_vector_and_distance,
 )
+from pymatgen.core import Lattice, Molecule, Structure
+from pymatgen.util.testing import PymatgenTest
 
 matgl.clear_cache(confirm=False)
 

tests/data/test_transformer.py

@@ -2,7 +2,6 @@
 
 import pytest
 import torch
-
 from matgl.data.transformer import LogTransformer, Normalizer
 
 

tests/ext/test_ase.py

@@ -6,10 +6,9 @@
 import pytest
 import torch
 from ase.build import molecule
-from pymatgen.io.ase import AseAtomsAdaptor
-
 from matgl import load_model
 from matgl.ext.ase import Atoms2Graph, M3GNetCalculator, MolecularDynamics, Relaxer
+from pymatgen.io.ase import AseAtomsAdaptor
 
 
 def test_M3GNetCalculator(MoS):

tests/ext/test_pymatgen.py

@@ -4,9 +4,8 @@
 
 import numpy as np
 import torch
-from pymatgen.core import Lattice, Structure
-
 from matgl.ext.pymatgen import Structure2Graph, get_element_list
+from pymatgen.core import Lattice, Structure
 
 module_dir = os.path.dirname(os.path.abspath(__file__))
 

tests/graph/test_compute.py

@@ -2,13 +2,11 @@
 
 from functools import partial
 
+import matgl
 import numpy as np
 import pytest
 import torch
 import torch.testing as tt
-from pymatgen.core import Lattice, Structure
-
-import matgl
 from matgl.ext.pymatgen import Structure2Graph, get_element_list
 from matgl.graph.compute import (
     compute_pair_vector_and_distance,
@@ -18,6 +16,7 @@
     ensure_line_graph_compatibility,
     prune_edges_by_features,
 )
+from pymatgen.core import Lattice, Structure
 
 
 def _loop_indices(bond_atom_indices, pair_dist, cutoff=4.0):

tests/graph/test_data.py

@@ -7,10 +7,9 @@
 
 import numpy as np
 from dgl.data.utils import split_dataset
-from pymatgen.core import Molecule
-
 from matgl.ext.pymatgen import Molecule2Graph, Structure2Graph, get_element_list
 from matgl.graph.data import M3GNetDataset, MEGNetDataset, MGLDataLoader, collate_fn, collate_fn_efs
+from pymatgen.core import Molecule
 
 module_dir = os.path.dirname(os.path.abspath(__file__))
 

tests/layers/test_activations.py

@@ -3,7 +3,6 @@
 import numpy as np
 import pytest
 import torch
-
 from matgl.layers._activations import SoftExponential, SoftPlus2
 
 

tests/layers/test_atom_ref.py

@@ -3,7 +3,6 @@
 import dgl
 import numpy as np
 import torch
-
 from matgl.layers._atom_ref import AtomRef
 
 

tests/layers/test_basis.py

@@ -3,8 +3,6 @@
 import numpy as np
 import pytest
 import torch
-from torch.testing import assert_close
-
 from matgl.graph.compute import (
     compute_theta_and_phi,
     create_line_graph,
@@ -19,6 +17,7 @@
     spherical_bessel_smooth,
 )
 from matgl.layers._three_body import combine_sbf_shf
+from torch.testing import assert_close
 
 
 def test_gaussian():

tests/layers/test_bond.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import pytest
-
 from matgl.layers import BondExpansion
 
 

tests/layers/test_core_and_embedding.py

@@ -2,10 +2,9 @@
 
 import pytest
 import torch
-from torch import nn
-
 from matgl.layers import BondExpansion, EmbeddingBlock
 from matgl.layers._core import MLP, GatedMLP
+from torch import nn
 
 
 @pytest.fixture()

tests/layers/test_graph_conv.py

@@ -4,8 +4,6 @@
 
 import dgl
 import torch
-from torch import nn
-
 from matgl.layers import BondExpansion, EmbeddingBlock
 from matgl.layers._graph_convolution import (
     MLP,
@@ -15,6 +13,7 @@
     MEGNetGraphConv,
 )
 from matgl.utils.cutoff import polynomial_cutoff
+from torch import nn
 
 
 class Graph(NamedTuple):

tests/layers/test_readout.py

@@ -2,15 +2,14 @@
 
 import pytest
 import torch
-from torch import nn
-
 from matgl.layers import BondExpansion, EmbeddingBlock
 from matgl.layers._readout import (
     ReduceReadOut,
     Set2SetReadOut,
     WeightedReadOut,
     WeightedReadOutPair,
 )
+from torch import nn
 
 
 class TestReadOut:

tests/layers/test_three_body.py

@@ -1,8 +1,6 @@
 from __future__ import annotations
 
 import torch
-from torch import nn
-
 from matgl.graph.compute import (
     compute_theta_and_phi,
     create_line_graph,
@@ -11,6 +9,7 @@
 from matgl.layers._core import MLP, GatedMLP
 from matgl.layers._three_body import ThreeBodyInteractions
 from matgl.utils.cutoff import polynomial_cutoff
+from torch import nn
 
 
 def test_three_body_interactions(graph_MoS):

tests/models/test_m3gnet.py

@@ -2,11 +2,10 @@
 
 import os
 
+import matgl
 import numpy as np
 import pytest
 import torch
-
-import matgl
 from matgl.models import M3GNet
 
 

tests/models/test_megnet.py

@@ -2,15 +2,14 @@
 
 import os
 
+import matgl
 import numpy as np
 import pytest
 import torch as th
-from pymatgen.core import Lattice, Structure
-
-import matgl
 from matgl.graph.compute import compute_pair_vector_and_distance
 from matgl.layers import BondExpansion
 from matgl.models import MEGNet
+from pymatgen.core import Lattice, Structure
 
 
 class TestMEGNet:

tests/models/test_wrapper.py

@@ -2,7 +2,6 @@
 
 import pytest
 import torch.nn
-
 from matgl.data.transformer import Normalizer
 from matgl.models._wrappers import TransformedTargetModel
 

tests/test_integration.py

@@ -1,9 +1,8 @@
 """This is an integration test file that checks on pre-trained models to ensure they still work."""
 from __future__ import annotations
 
-import pytest
-
 import matgl
+import pytest
 
 
 def test_form_e(LiFePO4):

tests/utils/test_cutoff.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 
 import torch
-from torch.testing import assert_close
-
 from matgl.layers._basis import SphericalBesselFunction
 from matgl.utils.cutoff import cosine_cutoff, polynomial_cutoff
+from torch.testing import assert_close
 
 
 def test_cosine():

tests/utils/test_io.py

@@ -7,7 +7,6 @@
 import pytest
 import requests
 import torch
-
 from matgl.utils.io import IOMixIn, RemoteFile, get_available_pretrained_models, load_model
 
 this_dir = Path(os.path.abspath(os.path.dirname(__file__)))

tests/utils/test_maths.py

@@ -4,7 +4,6 @@
 import numpy as np
 import pytest
 import torch
-
 from matgl.utils.maths import (
     SPHERICAL_BESSEL_ROOTS,
     broadcast_states_to_atoms,

tests/utils/test_training.py

@@ -11,12 +11,11 @@
 import pytorch_lightning as pl
 import torch.backends.mps
 from dgl.data.utils import split_dataset
-from pymatgen.core import Lattice, Structure
-
 from matgl.ext.pymatgen import Structure2Graph, get_element_list
 from matgl.graph.data import M3GNetDataset, MEGNetDataset, MGLDataLoader, collate_fn, collate_fn_efs
 from matgl.models import M3GNet, MEGNet
 from matgl.utils.training import ModelLightningModule, PotentialLightningModule, xavier_init
+from pymatgen.core import Lattice, Structure
 
 module_dir = os.path.dirname(os.path.abspath(__file__))