Downgrade Pytorch and DGL version (#392)

* improve TensorNet model coverage

* Update pyproject.toml

Signed-off-by: Tsz Wai Ko <[email protected]>

* Improve the unit test for SO(3) equivarance in TensorNet class

* improve SO3Net model class coverage and simplify TensorNet implementations

* improve the coverage in MLP_norm class

* Improve the implementation of three-body interactions

* fixed black

* Optimize the speed of _compute_3body class

* type checking is added for scheduler

* update M3GNet Potential training notebook for the demonstration of obtaining and using element offsets

* Downgrade sympy to avoid crash of SO3 operations

* Smooth l1 loss function is added and united tests are improved

* merge the method predict_structure and featurize_structure into a function including both

* remove unnecessary else statement for training magmoms

* modify so3 operation implementation to make united tests pass due to the update of sympy

* skip test_load_all_models for MacOS pytest now

* Reference for CHGNet is added

* Update README.md and index.md for including CHGNet

Signed-off-by: Tsz Wai Ko <[email protected]>

* add more description for using CHGNet pretrained models in Relaxations and Simulations using the M3GNet Universal Potential.ipynb

* A command-line interface for performing ASE MD simulations is added

* added back py.typed

* ExpNormal Smearing for radial basis functions is added

* Changed deprecated torch.scalar_tensor into torch.Tensor

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Signed-off-by: Tsz Wai Ko <[email protected]>

* Converted the float number into tensor

Signed-off-by: Tsz Wai Ko <[email protected]>

* fix the united test in test_bond.py

* fix the error from the upgrade of boto3

* Downgrade DGL to 2.2.1

* Downgrade pytorch

* fix mypy by adding self.norm_layers is not None

---------

Signed-off-by: Tsz Wai Ko <[email protected]>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: Shyue Ping Ong <[email protected]>

pyproject.toml

@@ -50,10 +50,10 @@ classifiers = [
 ]
 dependencies = [
     "ase",
-    "dgl<=2.4.0",
+    "dgl<=2.2.1",
     "pymatgen",
-    "lightning<=2.4.0",
-    "torch<=2.2.1",
+    "lightning<=2.3.0",
+    "torch<=2.4.0",
     "pydantic",
     "torchdata<=0.7.1",
     "boto3",

requirements.txt

@@ -1,6 +1,6 @@
-dgl==2.4.0
+dgl==2.2.1
 torch==2.4.0
-lightning==2.4.0
+lightning==2.3.0
 pymatgen==2024.10.3
 ase==3.23.0
 pydantic==2.9.2

src/matgl/layers/_core.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Callable, Literal
+from typing import TYPE_CHECKING, Any, Callable, Literal, cast
 
 import dgl
 import torch
@@ -130,21 +130,23 @@ def __init__(
         self.activate_last = activate_last
         self.normalize_hidden = normalize_hidden
         norm_kwargs = norm_kwargs or {}
+        norm_kwargs = cast(dict, norm_kwargs)
 
         for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
             if i < self._depth - 1:
                 self.layers.append(Linear(in_dim, out_dim, bias=use_bias))
-                if normalize_hidden:
+                if normalize_hidden and self.norm_layers is not None:
                     if normalization == "graph":
                         self.norm_layers.append(GraphNorm(out_dim, **norm_kwargs))
                     elif normalization == "layer":
                         self.norm_layers.append(LayerNorm(out_dim, **norm_kwargs))
             else:
                 self.layers.append(Linear(in_dim, out_dim, bias=use_bias and bias_last))
-                if normalization == "graph":
-                    self.norm_layers.append(GraphNorm(out_dim, **norm_kwargs))
-                elif normalization == "layer":
-                    self.norm_layers.append(LayerNorm(out_dim, **norm_kwargs))
+                if self.norm_layers is not None:
+                    if normalization == "graph":
+                        self.norm_layers.append(GraphNorm(out_dim, **norm_kwargs))
+                    elif normalization == "layer":
+                        self.norm_layers.append(LayerNorm(out_dim, **norm_kwargs))
 
     def forward(self, inputs: torch.Tensor, g: dgl.Graph | None = None) -> torch.Tensor:
         """Applies all layers in turn.

tests/layers/test_bond.py

@@ -42,6 +42,9 @@ def test_exp_normal(self, graph_MoS, graph_CO):
         bond_basis = bond_expansion(g1.edata["bond_dist"])
         assert bond_basis.shape == (28, 9)
 
+        bond_basis = bond_expansion(g2.edata["bond_dist"])
+        assert bond_basis.shape == (2, 9)
+
     def test_exception(self):
         with pytest.raises(ValueError, match="Undefined rbf_type"):
             BondExpansion(rbf_type="nonsense")