FrechetCellFilter is added for variable cell relaxation in Relaxer cl…

…ass (#275)

* 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

* FrechetCellFilter class for variable cell relaxation is added

src/matgl/ext/ase.py

@@ -17,6 +17,7 @@
 from ase import Atoms, units
 from ase.calculators.calculator import Calculator, all_changes
 from ase.constraints import ExpCellFilter
+from ase.filters import FrechetCellFilter
 from ase.md import Langevin
 from ase.md.andersen import Andersen
 from ase.md.npt import NPT
@@ -253,7 +254,8 @@ def relax(
         traj_file: str | None = None,
         interval: int = 1,
         verbose: bool = False,
-        params_expcellfilter: dict | None = None,
+        ase_cellfilter: Literal["Frechet", "Exp"] = "Frechet",
+        params_asecellfilter: dict | None = None,
         **kwargs,
     ):
         """
@@ -267,27 +269,34 @@ def relax(
             traj_file (str): the trajectory file for saving
             interval (int): the step interval for saving the trajectories
             verbose (bool): Whether to have verbose output.
-            params_expcellfilter (dict): Parameters to be passed to ExpCellFilter. Allows
+            ase_cellfilter (literal): which filter is used for variable cell relaxation. Default is Frechet.
+            params_asecellfilter (dict): Parameters to be passed to ExpCellFilter or FrechetCellFilter. Allows
                 setting of constant pressure or constant volume relaxations, for example. Refer to
-                https://wiki.fysik.dtu.dk/ase/ase/filters.html#the-expcellfilter-class for more information.
+                https://wiki.fysik.dtu.dk/ase/ase/filters.html#FrechetCellFilter for more information.
             **kwargs: Kwargs pass-through to optimizer.
         """
         if isinstance(atoms, (Structure, Molecule)):
             atoms = self.ase_adaptor.get_atoms(atoms)
         atoms.set_calculator(self.calculator)
         stream = sys.stdout if verbose else io.StringIO()
-        params_expcellfilter = params_expcellfilter or {}
+        params_asecellfilter = params_asecellfilter or {}
         with contextlib.redirect_stdout(stream):
             obs = TrajectoryObserver(atoms)
             if self.relax_cell:
-                atoms = ExpCellFilter(atoms, **params_expcellfilter)
+                atoms = (
+                    FrechetCellFilter(atoms, **params_asecellfilter)
+                    if ase_cellfilter == "Frechet"
+                    else ExpCellFilter(atoms, **params_asecellfilter)
+                )
+
             optimizer = self.optimizer(atoms, **kwargs)
             optimizer.attach(obs, interval=interval)
             optimizer.run(fmax=fmax, steps=steps)
             obs()
         if traj_file is not None:
             obs.save(traj_file)
-        if isinstance(atoms, ExpCellFilter):
+
+        if isinstance(atoms, (FrechetCellFilter, ExpCellFilter)):
             atoms = atoms.atoms
 
         return {