Released 0.3.0 pretrained weights

chgnet/data/dataset.py

@@ -48,7 +48,8 @@ def __init__(
             magmoms (list[list[float]], optional): [data_size, n_atoms, 1]
             structure_ids (list[str], optional): a list of ids to track the structures
             graph_converter (CrystalGraphConverter, optional): Converts the structures
-                to graphs. If None, it will be set to CHGNet default converter.
+                to graphs. If None, it will be set to CHGNet 0.3.0 converter
+                with AtomGraph cutoff = 6A.
 
         Raises:
             RuntimeError: if the length of structures and labels (energies, forces,
@@ -74,7 +75,7 @@ def __init__(
         random.shuffle(self.keys)
         print(f"{len(structures)} structures imported")
         self.graph_converter = graph_converter or CrystalGraphConverter(
-            atom_graph_cutoff=5, bond_graph_cutoff=3
+            atom_graph_cutoff=6, bond_graph_cutoff=3
         )
         self.failed_idx: list[int] = []
         self.failed_graph_id: dict[str, str] = {}
@@ -157,9 +158,9 @@ def __init__(
             labels (str, dict): the path or dictionary of labels
             targets ("ef" | "efs" | "efm" | "efsm"): The training targets.
                 Default = "efsm"
-            graph_converter (CrystalGraphConverter, optional):
-                a CrystalGraphConverter to convert the structures,
-                if None, it will be set to CHGNet default converter
+            graph_converter (CrystalGraphConverter, optional): Converts the structures
+                to graphs. If None, it will be set to CHGNet 0.3.0 converter
+                with AtomGraph cutoff = 6A.
             energy_key (str, optional): the key of energy in the labels.
                 Default = "energy_per_atom".
             force_key (str, optional): the key of force in the labels.
@@ -175,7 +176,7 @@ def __init__(
         random.shuffle(self.cif_ids)
         print(f"{cif_path}: {len(self.cif_ids):,} structures imported")
         self.graph_converter = graph_converter or CrystalGraphConverter(
-            atom_graph_cutoff=5, bond_graph_cutoff=3
+            atom_graph_cutoff=6, bond_graph_cutoff=3
         )
 
         self.energy_key = energy_key

chgnet/model/model.py

@@ -650,15 +650,26 @@ def from_file(cls, path, **kwargs):
         return CHGNet.from_dict(state["model"], **kwargs)
 
     @classmethod
-    def load(cls, model_name="MPtrj-efsm"):
+    def load(cls, model_name="0.3.0"):
         """Load pretrained CHGNet."""
         current_dir = os.path.dirname(os.path.abspath(__file__))
-        if model_name == "MPtrj-efsm":
+        if model_name == "0.3.0":
             return cls.from_file(
-                os.path.join(current_dir, "../pretrained/e30f77s348m32.pth.tar"),
+                os.path.join(
+                    current_dir,
+                    "../pretrained/0.3.0/chgnet_0.3.0_e29f68s314m37.pth.tar",
+                )
+            )
+        elif model_name == "0.2.0":  # noqa: RET505
+            return cls.from_file(
+                os.path.join(
+                    current_dir,
+                    "../pretrained/0.2.0/chgnet_0.2.0_e30f77s348m32.pth.tar",
+                ),
                 mlp_out_bias=True,
             )
-        raise ValueError(f"Unknown {model_name=}")
+        else:
+            raise ValueError(f"Unknown {model_name=}")
 
 
 @dataclass

chgnet/pretrained/0.2.0/README.md

@@ -0,0 +1,74 @@
+## Model 0.2.0
+
+This is the pretrained weights published with CHGNet Nature Machine Intelligence paper.
+All the experiments and results shown in the paper were performed with this version of weights.
+
+Date: 2/24/2023
+
+Author: Bowen Deng
+
+## Model Parameters
+
+```python
+model = CHGNet(
+    atom_fea_dim=64,
+    bond_fea_dim=64,
+    angle_fea_dim=64,
+    composition_model="MPtrj",
+    num_radial=9,
+    num_angular=9,
+    n_conv=4,
+    atom_conv_hidden_dim=64,
+    update_bond=True,
+    bond_conv_hidden_dim=64,
+    update_angle=True,
+    angle_layer_hidden_dim=0,
+    conv_dropout=0,
+    read_out="ave",
+    mlp_hidden_dims=[64, 64],
+    mlp_first=True,
+    is_intensive=True,
+    non_linearity="silu",
+    atom_graph_cutoff=5,
+    bond_graph_cutoff=3,
+    graph_converter_algorithm="fast",
+    cutoff_coeff=5,
+    learnable_rbf=True,
+    mlp_out_bias=True,
+)
+```
+
+## Dataset Used
+
+MPtrj dataset with 8-1-1 train-val-test splitting
+
+## Trainer
+
+```python
+trainer = Trainer(
+    model=model,
+    targets='efsm',
+    energy_loss_ratio=1,
+    force_loss_ratio=1,
+    stress_loss_ratio=0.1,
+    mag_loss_ratio=0.1,
+    optimizer='Adam',
+    weight_decay=0,
+    scheduler='CosLR',
+    criterion='Huber',
+    delta=0.1,
+    epochs=20,
+    starting_epoch=0,
+    learning_rate=1e-3,
+    use_device='cuda',
+    print_freq=1000
+)
+```
+
+## Mean Absolute Error (MAE) logs
+
+| partition  | Energy (meV/atom) | Force (meV/A) | stress (GPa) | magmom (muB) |
+| ---------- | ----------------- | ------------- | ------------ | ------------ |
+| Train      | 22                | 59            | 0.246        | 0.030        |
+| Validation | 20                | 75            | 0.350        | 0.033        |
+| Test       | 30                | 77            | 0.348        | 0.032        |

chgnet/pretrained/0.3.0/README.md

@@ -0,0 +1,80 @@
+## Model 0.3.0
+
+Major changes:
+
+1. Increased AtomGraph cutoff to 6A
+2. Resolved discontinuity issue when no BondGraph presents
+3. Added some normalization layers
+4. Slight improvements on energy, force, stress accuracies
+
+Date: 10/22/2023
+
+Author: Bowen Deng
+
+## Model Parameters
+
+```python
+model = CHGNet(
+    atom_fea_dim=64,
+    bond_fea_dim=64,
+    angle_fea_dim=64,
+    composition_model="MPtrj",
+    num_radial=31,
+    num_angular=31,
+    n_conv=4,
+    atom_conv_hidden_dim=64,
+    update_bond=True,
+    bond_conv_hidden_dim=64,
+    update_angle=True,
+    angle_layer_hidden_dim=0,
+    conv_dropout=0,
+    read_out="ave",
+    gMLP_norm='layer',
+    readout_norm='layer',
+    mlp_hidden_dims=[64, 64, 64],
+    mlp_first=True,
+    is_intensive=True,
+    non_linearity="silu",
+    atom_graph_cutoff=6,
+    bond_graph_cutoff=3,
+    graph_converter_algorithm="fast",
+    cutoff_coeff=8,
+    learnable_rbf=True,
+)
+```
+
+## Dataset Used
+
+MPtrj dataset with 9-0.5-0.5 train-val-test splitting
+
+## Trainer
+
+```python
+trainer = Trainer(
+    model=model,
+    targets='efsm',
+    energy_loss_ratio=1,
+    force_loss_ratio=1,
+    stress_loss_ratio=0.1,
+    mag_loss_ratio=0.1,
+    optimizer='Adam',
+    weight_decay=0,
+    scheduler='CosLR',
+    scheduler_params={'decay_fraction': 0.5e-2},
+    criterion='Huber',
+    delta=0.1,
+    epochs=30,
+    starting_epoch=0,
+    learning_rate=5e-3,
+    use_device='cuda',
+    print_freq=1000
+)
+```
+
+## Mean Absolute Error (MAE) logs
+
+| partition  | Energy (meV/atom) | Force (meV/A) | stress (GPa) | magmom (muB) |
+| ---------- | ----------------- | ------------- | ------------ | ------------ |
+| Train      | 26                | 60            | 0.266        | 0.037        |
+| Validation | 29                | 70            | 0.308        | 0.037        |
+| Test       | 29                | 68            | 0.314        | 0.037        |