Updatae docs.

docs/changes.md

@@ -6,8 +6,16 @@ nav_order: 3
 
 # Change Log
 
-## 1.0.1
-- 
+## 1.1.2
+- Move AtomRef Fitting to numpy to avoid bug (@BowenD-UCB)
+- NVE ensemble added (@kenko911)
+- Migrate from pytorch_lightning to lightning.
+
+## 1.1.1
+- Pin dependencies to support latest DGL 2.x. @kenko911
+
+## 1.1.0
+- Implementation of CHGnet + pre-trained models. (@BowenD-UCB)
 
 ## 1.0.0
 - First 1.0.0 release to reflect the maturity of the matgl code! All changes below are the efforts of @kenko911.

docs/index.md

@@ -191,6 +191,15 @@ torch.hub.list("materialsvirtuallab/matgl", force_reload=True)
 # To load a model
 model = torch.hub.load("materialsvirtuallab/matgl", 'm3gnet_universal_potential')
 ```
+## Model Training
+
+In the PES training, the unit of energies, forces and stresses (optional) in the training, validation and test sets is extremely important to be consistent with the unit used in MatGL.
+
+- energies: a list of energies with unit eV.
+- forces: a list of nx3 force matrix with unit eV/Å, where n is the number of atom in each structure. n does not need to be the same for all structures.
+- stresses: a list of 3x3 stress matrices with unit GPa (optional)
+
+Note: For stresses, we use the convention that compressive stress gives negative values. Stresses obtained from VASP calculations (default unit is kBar) should be multiplied by -0.1 to work directly with the model.
 
 ## Tutorials
 

docs/tutorials/Benchmarking M3GNet Predictions of Cubic Lattice Parameters.md

@@ -105,7 +105,7 @@ In the next cell, we generate an initial structure for all the phases. The cubic
 predicted = []
 mp = []
 os.environ["MPRESTER_MUTE_PROGRESS_BARS"] = "true"
-mpr = MPRester("FwTXcju8unkI2VbInEgZDTN8coDB6S6U")
+mpr = MPRester("YOUR_API_KEY")
 
 # Load the pre-trained M3GNet Potential
 pot = matgl.load_model("M3GNet-MP-2021.2.8-PES")

docs/tutorials/Training a M3GNet Formation Energy Model with PyTorch Lightning.md

@@ -85,9 +85,13 @@ Here, we set up the dataset.
 elem_list = get_element_list(structures)
 # setup a graph converter
 converter = Structure2Graph(element_types=elem_list, cutoff=4.0)
-# convert the raw dataset into MEGNetDataset
+# convert the raw dataset into M3GNetDataset
 mp_dataset = MGLDataset(
-    threebody_cutoff=4.0, structures=structures, converter=converter, labels={"eform": eform_per_atom}
+    threebody_cutoff=4.0,
+    structures=structures,
+    converter=converter,
+    labels={"eform": eform_per_atom},
+    include_line_graph=True,
 )
 ```
 
@@ -114,18 +118,18 @@ train_loader, val_loader, test_loader = MGLDataLoader(
 
 # Model setup
 
-In the next step, we setup the model and the ModelLightningModule. Here, we have initialized a MEGNet model from scratch. Alternatively, you can also load one of the pre-trained models for transfer learning, which may speed up the training.
+In the next step, we setup the model and the ModelLightningModule. Here, we have initialized a M3GNet model from scratch. Alternatively, you can also load one of the pre-trained models for transfer learning, which may speed up the training.
 
 
 ```python
-# setup the architecture of MEGNet model
+# setup the architecture of M3GNet model
 model = M3GNet(
     element_types=elem_list,
     is_intensive=True,
     readout_type="set2set",
 )
-# setup the MEGNetTrainer
-lit_module = ModelLightningModule(model=model)
+# setup the M3GNetTrainer
+lit_module = ModelLightningModule(model=model, include_line_graph=True)
 ```
 
 # Training
@@ -143,7 +147,7 @@ trainer.fit(model=lit_module, train_dataloaders=train_loader, val_dataloaders=va
 
 # Visualizing the convergence
 
-Finally, we can plot the convergence plot for the loss metrics. You can see that the MAE is already going down nicely with 20 epochs. Obviously, this is nowhere state of the art performance for the formation energies, but a longer training time should lead to results consistent with what was reported in the original MEGNet work.
+Finally, we can plot the convergence plot for the loss metrics. You can see that the MAE is already going down nicely with 20 epochs. Obviously, this is nowhere state of the art performance for the formation energies, but a longer training time should lead to results consistent with what was reported in the original M3GNet work.
 
 
 ```python

docs/tutorials/Training a M3GNet Potential with PyTorch Lightning.md

@@ -18,15 +18,17 @@ import warnings
 
 import numpy as np
 import pytorch_lightning as pl
+from functools import partial
 from dgl.data.utils import split_dataset
 from mp_api.client import MPRester
 from pytorch_lightning.loggers import CSVLogger
 
 import matgl
 from matgl.ext.pymatgen import Structure2Graph, get_element_list
-from matgl.graph.data import MGLDataset, MGLDataLoader, collate_fn_efs
+from matgl.graph.data import MGLDataset, MGLDataLoader, collate_fn_pes
 from matgl.models import M3GNet
 from matgl.utils.training import PotentialLightningModule
+from matgl.config import DEFAULT_ELEMENTS
 
 # To suppress warnings for clearer output
 warnings.simplefilter("ignore")
@@ -37,8 +39,7 @@ For the purposes of demonstration, we will download all Si-O compounds in the Ma
 
 ```python
 # Obtain your API key here: https://next-gen.materialsproject.org/api
-# mpr = MPRester(api_key="YOUR_API_KEY")
-mpr = MPRester("FwTXcju8unkI2VbInEgZDTN8coDB6S6U")
+mpr = MPRester(api_key="YOUR_API_KEY")
 entries = mpr.get_entries_in_chemsys(["Si", "O"])
 structures = [e.structure for e in entries]
 energies = [e.energy for e in entries]
@@ -57,33 +58,31 @@ We will first setup the M3GNet model and the LightningModule.
 
 
 ```python
-element_types = get_element_list(structures)
+element_types = DEFAULT_ELEMENTS
 converter = Structure2Graph(element_types=element_types, cutoff=5.0)
 dataset = MGLDataset(
-    threebody_cutoff=4.0,
-    structures=structures,
-    converter=converter,
-    labels=labels,
+    threebody_cutoff=4.0, structures=structures, converter=converter, labels=labels, include_line_graph=True
 )
 train_data, val_data, test_data = split_dataset(
     dataset,
     frac_list=[0.8, 0.1, 0.1],
     shuffle=True,
     random_state=42,
 )
+my_collate_fn = partial(collate_fn_pes, include_line_graph=True)
 train_loader, val_loader, test_loader = MGLDataLoader(
     train_data=train_data,
     val_data=val_data,
     test_data=test_data,
-    collate_fn=collate_fn_efs,
+    collate_fn=my_collate_fn,
     batch_size=2,
     num_workers=0,
 )
 model = M3GNet(
     element_types=element_types,
     is_intensive=False,
 )
-lit_module = PotentialLightningModule(model=model)
+lit_module = PotentialLightningModule(model=model, include_line_graph=True)
 ```
 
 Finally, we will initialize the Pytorch Lightning trainer and run the fitting. Here, the max_epochs is set to 2 just for demonstration purposes. In a real fitting, this would be a much larger number. Also, the `accelerator="cpu"` was set just to ensure compatibility with M1 Macs. In a real world use case, please remove the kwarg or set it to cuda for GPU based training.
@@ -107,7 +106,7 @@ trainer.test(dataloaders=test_loader)
 ```python
 # save trained model
 model_export_path = "./trained_model/"
-model.save(model_export_path)
+lit_module.model.save(model_export_path)
 
 # load trained model
 model = matgl.load_model(path=model_export_path)
@@ -121,7 +120,7 @@ In the previous cells, we demonstrated the process of training an M3GNet from sc
 # download a pre-trained M3GNet
 m3gnet_nnp = matgl.load_model("M3GNet-MP-2021.2.8-PES")
 model_pretrained = m3gnet_nnp.model
-lit_module_finetune = PotentialLightningModule(model=model_pretrained, lr=1e-4)
+lit_module_finetune = PotentialLightningModule(model=model_pretrained, lr=1e-4, include_line_graph=True)
 ```
 
 
@@ -136,7 +135,7 @@ trainer.fit(model=lit_module_finetune, train_dataloaders=train_loader, val_datal
 ```python
 # save trained model
 model_save_path = "./finetuned_model/"
-model_pretrained.save(model_save_path)
+lit_module_finetune.model.save(model_save_path)
 # load trained model
 trained_model = matgl.load_model(path=model_save_path)
 ```

docs/tutorials/Using LAMMPS with MatGL.md

@@ -88,10 +88,10 @@ pair_style    m3gnet {HOME_DIR}/.local/share/lammps/potentials/M3GNET
 
 read_data     ./dat.lammps
 replicate     {x} 1 1
-pair_coeff    * *  M3GNet-MP-2021.2.8-DIRECT-PES  Zr O  # MatGL will be called
+pair_coeff    * *  M3GNet-MP-2021.2.8-DIRECT-PES  Mg O  # MatGL will be called
 
 dump          myDump all custom 10 xyz.lammpstrj id element x y z
-dump_modify   myDump sort id element Zr O
+dump_modify   myDump sort id element Mg O
 
 thermo_style  custom step time cpu pe ke etotal temp press vol density
 thermo        10