NVE ensemble is added (#266)

* model version for Potential class is added

* model version for Potential class is modified

* Enable the smooth version of Spherical Bessel function in TensorNet

* max_n, max_l for SphericalBessel radial basis functions are included in TensorNet class

* adding united tests for improving the coverage score

* little clean up in _so3.py and so3.py

* remove unnecessary data storage in dgl graphs

* update pymatgen version to fix the bug

* refractor all include_states into include_state for consistency

* change include_states into include_state in test_graph_conv.py

* Ensure the state attr from molecule graph is consistent with matgl.float_th and including linear layer in TensorNet to match the original implementations

* Fix the jupyter-notebook for M3GNet training

* included more united tests to improve code coverage

* update pyproject.toml and requirements.txt to adapt dgl-2.x

* downgrade dgl to 2.1.0

* Improve more code coverage in _chgnet.py

* update pytorch-lightning

* NVE ensemble is added

README.md

@@ -185,12 +185,12 @@ model = torch.hub.load("materialsvirtuallab/matgl", 'm3gnet_universal_potential'
 ```
 ## Model Training
 
-In the PES training, the units of energies, forces and stresses (optional) in the training, validation and test sets are extremely important to be consistent with the units used in MatGL.
+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

src/matgl/ext/ase.py

@@ -22,6 +22,7 @@
 from ase.md.npt import NPT
 from ase.md.nptberendsen import Inhomogeneous_NPTBerendsen, NPTBerendsen
 from ase.md.nvtberendsen import NVTBerendsen
+from ase.md.verlet import VelocityVerlet
 from pymatgen.core.structure import Molecule, Structure
 from pymatgen.io.ase import AseAtomsAdaptor
 from pymatgen.optimization.neighbors import find_points_in_spheres
@@ -367,7 +368,9 @@ def __init__(
         potential: Potential,
         state_attr: torch.Tensor | None = None,
         stress_weight: float = 1 / 160.21766208,
-        ensemble: Literal["nvt", "nvt_langevin", "nvt_andersen", "npt", "npt_berendsen", "npt_nose_hoover"] = "nvt",
+        ensemble: Literal[
+            "nve", "nvt", "nvt_langevin", "nvt_andersen", "npt", "npt_berendsen", "npt_nose_hoover"
+        ] = "nvt",
         temperature: int = 300,
         timestep: float = 1.0,
         pressure: float = 1.01325 * units.bar,
@@ -394,8 +397,8 @@ def __init__(
             stress of the atoms
             state_attr (torch.Tensor): State attr.
             stress_weight (float): conversion factor from GPa to eV/A^3
-            ensemble (str): choose from 'nvt' or 'npt'. NPT is not tested,
-            use with extra caution
+            ensemble (str): choose from "nve", "nvt", "nvt_langevin", "nvt_andersen", "npt", "npt_berendsen",
+            "npt_nose_hoover"
             temperature (float): temperature for MD simulation, in K
             timestep (float): time step in fs
             pressure (float): pressure in eV/A^3
@@ -443,6 +446,16 @@ def __init__(
                 append_trajectory=append_trajectory,
             )
 
+        elif ensemble.lower() == "nve":
+            self.dyn = VelocityVerlet(
+                self.atoms,
+                timestep * units.fs,
+                trajectory=trajectory,
+                logfile=logfile,
+                loginterval=loginterval,
+                append_trajectory=append_trajectory,
+            )
+
         elif ensemble.lower() == "nvt_langevin":
             self.dyn = Langevin(
                 self.atoms,

tests/ext/test_ase.py

@@ -117,7 +117,7 @@ def test_get_graph_from_atoms_mol():
 
 def test_molecular_dynamics(MoS2):
     pot = load_model("pretrained_models/M3GNet-MP-2021.2.8-PES/")
-    for ensemble in ["nvt", "nvt_langevin", "nvt_andersen", "npt", "npt_berendsen", "npt_nose_hoover"]:
+    for ensemble in ["nvt", "nve", "nvt_langevin", "nvt_andersen", "npt", "npt_berendsen", "npt_nose_hoover"]:
         md = MolecularDynamics(MoS2, potential=pot, ensemble=ensemble, taut=0.1, taup=0.1, compressibility_au=10)
         md.run(10)
         assert md.dyn is not None