CHGNet-matgl implementation

pretrained_models/CHGNet-MPtrj-2023.12.1-PES-2.7M/README.md

@@ -0,0 +1,33 @@
+# Description
+
+This model is a CHGNet universal potential trained from the Materials Project trajectory (MPtrj) dataset
+that contains over 1.5 million structures with 89 elements.
+This Matgl implementation has slight modification from original pytorch implementation by adding directed edge updates.
+
+# Training dataset
+
+MPtrj-2022.9: Materials Project trajectory dataset that contains GGA and GGA+U static and relaxation calculations.
+- Train-Val-Test splitting with mp-id: 0.9 - 0.5 - 0.5
+- Train set size: 1419861
+- Validation set size: 79719
+- Test set size: 79182
+
+# Performance metrics
+## Training and validation errors
+
+| partition  | Energy (meV/atom) | Force (meV/A) | stress (GPa) | magmom (muB) |
+| ---------- | ----------------- | ------------- | ------------ | ------------ |
+| Train      | 26.45             | 49            | 0.173        | 0.036        |
+| Validation | 30.31             | 70            | 0.297        | 0.037        |
+| Test       | 30.80             | 66            | 0.296        | 0.038        |
+
+
+# References
+
+```txt
+Deng, B. et al. CHGNet as a pretrained universal neural network potential for charge-informed atomistic modelling.
+Nat. Mach. Intell. 1–11 (2023) doi:10.1038/s42256-023-00716-3.
+  ```
+
+####  Date: 2023.12.1
+####  Author: Bowen Deng

pretrained_models/CHGNet-MPtrj-2023.12.1-PES-2.7M/model.json

@@ -0,0 +1,66 @@
+{
+    "@class": "Potential",
+    "@module": "matgl.apps.pes",
+    "@model_version": 1,
+    "metadata": null,
+    "kwargs": {
+        "model": {
+            "@class": "CHGNet",
+            "@module": "matgl.models._chgnet",
+            "@model_version": 1,
+            "init_args": {
+                "element_types": null,
+                "dim_state_feats": null,
+                "non_linear_bond_embedding": false,
+                "non_linear_angle_embedding": false,
+                "cutoff": 6.0,
+                "threebody_cutoff": 3.0,
+                "cutoff_exponent": 5,
+                "max_f": 32,
+                "learn_basis": false,
+                "num_blocks": 5,
+                "shared_bond_weights": "both",
+                "final_mlp_type": "mlp",
+                "final_hidden_dims": [
+                    128,
+                    128,
+                    128
+                ],
+                "final_dropout": 0.0,
+                "pooling_operation": "sum",
+                "readout_field": "atom_feat",
+                "activation_type": "swish",
+                "is_intensive": false,
+                "num_targets": 1,
+                "num_site_targets": 1,
+                "task_type": "regression",
+                "angle_update_hidden_dims": [],
+                "atom_conv_hidden_dims": [
+                    128
+                ],
+                "bond_conv_hidden_dims": [
+                    128
+                ],
+                "bond_update_hidden_dims": [
+                    128
+                ],
+                "conv_dropout": 0.0,
+                "dim_angle_embedding": 128,
+                "dim_atom_embedding": 128,
+                "dim_bond_embedding": 128,
+                "dim_state_embedding": null,
+                "layer_bond_weights": null,
+                "max_n": 63,
+                "normalization": "layer",
+                "normalize_hidden": false
+            }
+        },
+        "data_mean": "tensor(0.)",
+        "data_std": "tensor(1.)",
+        "element_refs": "tensor([ -3.4524,  -0.2535,  -3.1356,  -3.5818,  -7.5282,  -8.2669,  -7.7537,\n         -8.3183,  -5.6419,  -0.0301,  -1.9928,  -1.5805,  -4.3933,  -6.2148,\n         -6.3137,  -5.6612,  -3.6236,  -0.0632,  -1.7023,  -3.7368,  -6.8803,\n         -9.4099,  -9.5156,  -9.5164,  -9.0957,  -7.9901,  -6.4274,  -5.5935,\n         -3.3122,  -0.8411,  -3.2124,  -4.8460,  -4.6307,  -4.8599,  -3.1515,\n          0.8167,  -1.5970,  -3.4679,  -7.7371,  -9.5454, -10.5613,  -9.9911,\n         -6.7590,  -8.2281,  -7.0194,  -5.0765,  -1.8264,  -0.3508,  -2.5801,\n         -3.9090,  -4.0772,  -3.8814,  -2.4899,   3.0526,  -2.2544,  -3.9207,\n         -7.2230,  -7.4592,  -6.5778,  -6.7207,  -5.1486,  -6.7854, -11.7604,\n        -16.2078,  -6.4504,  -6.4809,  -6.3964,  -6.3873,  -6.3777,  -2.8369,\n         -6.4433, -10.6572, -12.3228, -11.8578, -10.5498,  -9.1838,  -8.1969,\n         -6.0177,  -2.7653,   0.6777,  -1.6524,  -3.1770,  -3.3326,   0.0000,\n          0.0000,   0.0000,   0.0000,   0.0000,  -4.2516,  -9.0753, -10.3286,\n        -12.5580, -12.7219, -14.3059])",
+        "calc_forces": true,
+        "calc_stresses": true,
+        "calc_hessian": false,
+        "calc_site_wise": true
+    }
+}

pretrained_models/CHGNet-MPtrj-2024.2.13-PES-11M/README.md

@@ -0,0 +1,32 @@
+# Description
+
+This model is a CHGNet universal potential trained from the Materials Project trajectory (MPtrj) dataset
+that contains over 1.5 million structures with 89 elements.
+This Matgl implementation has slight modification from original pytorch implementation by adding directed edge updates.
+
+# Training dataset
+
+MPtrj-2022.9: Materials Project trajectory dataset that contains GGA and GGA+U static and relaxation calculations.
+- Train-Val-Test splitting with mp-id: 0.95 - 0.5
+- Train set size: 1499043
+- Validation set size: 79719
+- Test set size: 0
+
+# Performance metrics
+## Training and validation errors
+
+| partition  | Energy (meV/atom) | Force (meV/A) | stress (GPa) | magmom (muB) |
+| ---------- | ----------------- | ------------- | ------------ | ------------ |
+| Train      | 25.6              | 47.6          | 0.177        | 0.017        |
+| Validation | 27.7              | 62.5          | 0.288        | 0.017        |
+
+
+# References
+
+```txt
+Deng, B. et al. CHGNet as a pretrained universal neural network potential for charge-informed atomistic modelling.
+Nat. Mach. Intell. 1–11 (2023) doi:10.1038/s42256-023-00716-3.
+  ```
+
+####  Date: 2024.2.13
+####  Author: Bowen Deng

pretrained_models/CHGNet-MPtrj-2024.2.13-PES-11M/model.json

@@ -0,0 +1,67 @@
+{
+    "@class": "Potential",
+    "@module": "matgl.apps.pes",
+    "@model_version": 2,
+    "metadata": null,
+    "kwargs": {
+        "model": {
+            "@class": "CHGNet",
+            "@module": "matgl.models._chgnet",
+            "@model_version": 1,
+            "init_args": {
+                "element_types": null,
+                "dim_state_feats": null,
+                "non_linear_bond_embedding": false,
+                "non_linear_angle_embedding": false,
+                "cutoff": 6.0,
+                "threebody_cutoff": 3.0,
+                "cutoff_exponent": 5,
+                "max_f": 32,
+                "learn_basis": false,
+                "num_blocks": 5,
+                "shared_bond_weights": "both",
+                "final_mlp_type": "mlp",
+                "final_hidden_dims": [
+                    256,
+                    256,
+                    256
+                ],
+                "final_dropout": 0.0,
+                "pooling_operation": "sum",
+                "readout_field": "atom_feat",
+                "activation_type": "swish",
+                "is_intensive": false,
+                "num_targets": 1,
+                "num_site_targets": 1,
+                "task_type": "regression",
+                "angle_update_hidden_dims": [],
+                "atom_conv_hidden_dims": [
+                    256
+                ],
+                "bond_conv_hidden_dims": [
+                    256
+                ],
+                "bond_update_hidden_dims": [
+                    256
+                ],
+                "conv_dropout": 0.0,
+                "dim_angle_embedding": 256,
+                "dim_atom_embedding": 256,
+                "dim_bond_embedding": 256,
+                "dim_state_embedding": null,
+                "layer_bond_weights": null,
+                "max_n": 63,
+                "normalization": "layer",
+                "normalize_hidden": false
+            }
+        },
+        "data_mean": "tensor(0.)",
+        "data_std": "tensor(1.)",
+        "element_refs": "tensor([ -3.4524,  -0.2535,  -3.1356,  -3.5818,  -7.5282,  -8.2669,  -7.7537,\n         -8.3183,  -5.6419,  -0.0301,  -1.9928,  -1.5805,  -4.3933,  -6.2148,\n         -6.3137,  -5.6612,  -3.6236,  -0.0632,  -1.7023,  -3.7368,  -6.8803,\n         -9.4099,  -9.5156,  -9.5164,  -9.0957,  -7.9901,  -6.4274,  -5.5935,\n         -3.3122,  -0.8411,  -3.2124,  -4.8460,  -4.6307,  -4.8599,  -3.1515,\n          0.8167,  -1.5970,  -3.4679,  -7.7371,  -9.5454, -10.5613,  -9.9911,\n         -6.7590,  -8.2281,  -7.0194,  -5.0765,  -1.8264,  -0.3508,  -2.5801,\n         -3.9090,  -4.0772,  -3.8814,  -2.4899,   3.0526,  -2.2544,  -3.9207,\n         -7.2230,  -7.4592,  -6.5778,  -6.7207,  -5.1486,  -6.7854, -11.7604,\n        -16.2078,  -6.4504,  -6.4809,  -6.3964,  -6.3873,  -6.3777,  -2.8369,\n         -6.4433, -10.6572, -12.3228, -11.8578, -10.5498,  -9.1838,  -8.1969,\n         -6.0177,  -2.7653,   0.6777,  -1.6524,  -3.1770,  -3.3326,   0.0000,\n          0.0000,   0.0000,   0.0000,   0.0000,  -4.2516,  -9.0753, -10.3286,\n        -12.5580, -12.7219, -14.3059])",
+        "calc_forces": true,
+        "calc_stresses": true,
+        "calc_hessian": false,
+        "calc_site_wise": true,
+        "debug_mode": false
+    }
+}

pyproject.toml

@@ -15,6 +15,8 @@ authors = [
     { name = "Ji Qi", email = "[email protected]" },
     { name = "Santiago Miret", email = "[email protected]" },
     { name = "Eliott Liu", email = "[email protected]" },
+    { name = "Bowen Deng", email = "[email protected]" },
+    { name = "Luis Barroso-Luque", email = "[email protected]" },
     { name = "Shyue Ping Ong", email = "[email protected]" },
 ]
 description = "MatGL is a framework for graph deep learning for materials science."

src/matgl/apps/pes.py

@@ -95,7 +95,7 @@ def forward(
         st = lat.new_zeros([g.batch_size, 3, 3])
         if self.calc_stresses:
             st.requires_grad_(True)
-        lattice = lat @ (torch.eye(3).to(st.device) + st)
+        lattice = lat @ (torch.eye(3, device=lat.device) + st)
         g.edata["lattice"] = torch.repeat_interleave(lattice, g.batch_num_edges(), dim=0)
         g.edata["pbc_offshift"] = (g.edata["pbc_offset"].unsqueeze(dim=-1) * g.edata["lattice"]).sum(dim=1)
         g.ndata["pos"] = (

src/matgl/ext/ase.py

@@ -122,7 +122,7 @@ def get_graph(self, atoms: Atoms) -> tuple[dgl.DGLGraph, torch.Tensor, list]:
 class PESCalculator(Calculator):
     """Potential calculator for ASE."""
 
-    implemented_properties = ("energy", "free_energy", "forces", "stress", "hessian")
+    implemented_properties = ("energy", "free_energy", "forces", "stress", "hessian", "magmoms")
 
     def __init__(
         self,
@@ -145,6 +145,7 @@ def __init__(
         self.potential = potential
         self.compute_stress = potential.calc_stresses
         self.compute_hessian = potential.calc_hessian
+        self.compute_magmoms = potential.calc_site_wise
         self.stress_weight = stress_weight
         self.state_attr = state_attr
         self.element_types = potential.model.element_types  # type: ignore
@@ -172,18 +173,20 @@ def calculate(
         graph, lattice, state_attr_default = Atoms2Graph(self.element_types, self.cutoff).get_graph(atoms)
         # type: ignore
         if self.state_attr is not None:
-            energies, forces, stresses, hessians = self.potential(graph, lattice, self.state_attr)
+            calc_result = self.potential(graph, lattice, self.state_attr)
         else:
-            energies, forces, stresses, hessians = self.potential(graph, lattice, state_attr_default)
+            calc_result = self.potential(graph, lattice, state_attr_default)
         self.results.update(
-            energy=energies.detach().cpu().numpy().item(),
-            free_energy=energies.detach().cpu().numpy().item(),
-            forces=forces.detach().cpu().numpy(),
+            energy=calc_result[0].detach().cpu().numpy().item(),
+            free_energy=calc_result[0].detach().cpu().numpy(),
+            forces=calc_result[1].detach().cpu().numpy(),
         )
         if self.compute_stress:
-            self.results.update(stress=stresses.detach().cpu().numpy() * self.stress_weight)
+            self.results.update(stress=calc_result[2].detach().cpu().numpy() * self.stress_weight)
         if self.compute_hessian:
-            self.results.update(hessian=hessians.detach().cpu().numpy())
+            self.results.update(hessian=calc_result[3].detach().cpu().numpy())
+        if self.compute_magmoms:
+            self.results.update(magmoms=calc_result[4].detach().cpu().numpy())
 
 
 # for backward compatibility

src/matgl/graph/compute.py

@@ -78,7 +78,8 @@ def create_line_graph(g: dgl.DGLGraph, threebody_cutoff: float, directed: bool =
     Args:
         g: DGL graph
         threebody_cutoff (float): cutoff for three-body interactions
-        directed (bool): Whether to create a directed line graph, or an m3gnet 3body line graph (default: False, m3gnet)
+        directed (bool): Whether to create a directed line graph, or an M3gnet 3body line graph
+            Default = False (M3Gnet)
 
     Returns:
         l_g: DGL graph containing three body information from graph
@@ -226,10 +227,10 @@ def _create_directed_line_graph(graph: dgl.DGLGraph, threebody_cutoff: float) ->
         threebody_cutoff: cutoff for three-body interactions
 
     Returns:
-        line_graph: DGL graph line graph of pruned graph to three body cutoff
+        line_graph: DGL line graph of pruned graph to three body cutoff
     """
     with torch.no_grad():
-        pg = prune_edges_by_features(graph, feat_name="bond_dist", condition=lambda x: x > threebody_cutoff)
+        pg = prune_edges_by_features(graph, feat_name="bond_dist", condition=lambda x: torch.gt(x, threebody_cutoff))
         src_indices, dst_indices = pg.edges()
         images = pg.edata["pbc_offset"]
         all_indices = torch.arange(pg.number_of_nodes(), device=graph.device).unsqueeze(dim=0)