Add NEBCalculator with M3GNet and CHGNet example notebook

examples/NEB_data/LiFePO4_supercell.cif

@@ -0,0 +1,138 @@
+# generated using pymatgen
+data_LiFePO4
+_symmetry_space_group_name_H-M   'P 1'
+_cell_length_a   10.23619605
+_cell_length_b   11.94151020
+_cell_length_c   9.30983438
+_cell_angle_alpha   90.00002115
+_cell_angle_beta   90.00000000
+_cell_angle_gamma   90.00000000
+_symmetry_Int_Tables_number   1
+_chemical_formula_structural   LiFePO4
+_chemical_formula_sum   'Li16 Fe16 P16 O64'
+_cell_volume   1137.99355945
+_cell_formula_units_Z   16
+loop_
+ _symmetry_equiv_pos_site_id
+ _symmetry_equiv_pos_as_xyz
+  1  'x, y, z'
+loop_
+ _atom_site_type_symbol
+ _atom_site_label
+ _atom_site_symmetry_multiplicity
+ _atom_site_fract_x
+ _atom_site_fract_y
+ _atom_site_fract_z
+ _atom_site_occupancy
+  Li  Li0  1  0.00000000  0.00000000  0.00000000  1
+  Li  Li1  1  0.00000000  0.00000000  0.50000000  1
+  Li  Li2  1  0.00000000  0.50000000  0.00000000  1
+  Li  Li3  1  0.00000000  0.50000000  0.50000000  1
+  Li  Li4  1  0.50000000  0.25000000  0.25000000  1
+  Li  Li5  1  0.50000000  0.25000000  0.75000000  1
+  Li  Li6  1  0.50000000  0.75000000  0.25000000  1
+  Li  Li7  1  0.50000000  0.75000000  0.75000000  1
+  Li  Li8  1  0.50000000  0.00000000  0.25000000  1
+  Li  Li9  1  0.50000000  0.00000000  0.75000000  1
+  Li  Li10  1  0.50000000  0.50000000  0.25000000  1
+  Li  Li11  1  0.50000000  0.50000000  0.75000000  1
+  Li  Li12  1  0.00000000  0.25000000  0.00000000  1
+  Li  Li13  1  0.00000000  0.25000000  0.50000000  1
+  Li  Li14  1  0.00000000  0.75000000  0.00000000  1
+  Li  Li15  1  0.00000000  0.75000000  0.50000000  1
+  Fe  Fe16  1  0.78115127  0.12500000  0.26493287  1
+  Fe  Fe17  1  0.78115127  0.12500000  0.76493287  1
+  Fe  Fe18  1  0.78115127  0.62500000  0.26493287  1
+  Fe  Fe19  1  0.78115127  0.62500000  0.76493287  1
+  Fe  Fe20  1  0.71884873  0.37500000  0.01493386  1
+  Fe  Fe21  1  0.71884873  0.37500000  0.51493387  1
+  Fe  Fe22  1  0.71884873  0.87500000  0.01493387  1
+  Fe  Fe23  1  0.71884873  0.87500000  0.51493386  1
+  Fe  Fe24  1  0.28115127  0.12500000  0.48506664  1
+  Fe  Fe25  1  0.28115127  0.12500000  0.98506664  1
+  Fe  Fe26  1  0.28115127  0.62500000  0.48506663  1
+  Fe  Fe27  1  0.28115127  0.62500000  0.98506664  1
+  Fe  Fe28  1  0.21884873  0.37500000  0.23506663  1
+  Fe  Fe29  1  0.21884873  0.37500000  0.73506664  1
+  Fe  Fe30  1  0.21884873  0.87500000  0.23506663  1
+  Fe  Fe31  1  0.21884873  0.87500000  0.73506664  1
+  P  P32  1  0.09386630  0.12500000  0.20931129  1
+  P  P33  1  0.09386630  0.12500000  0.70931129  1
+  P  P34  1  0.09386630  0.62500000  0.20931129  1
+  P  P35  1  0.09386630  0.62500000  0.70931128  1
+  P  P36  1  0.40613370  0.37500000  0.45931179  1
+  P  P37  1  0.40613370  0.37500000  0.95931179  1
+  P  P38  1  0.40613370  0.87500000  0.45931178  1
+  P  P39  1  0.40613370  0.87500000  0.95931179  1
+  P  P40  1  0.59386630  0.12500000  0.04068771  1
+  P  P41  1  0.59386630  0.12500000  0.54068772  1
+  P  P42  1  0.59386630  0.62500000  0.04068771  1
+  P  P43  1  0.59386630  0.62500000  0.54068772  1
+  P  P44  1  0.90613370  0.37500000  0.29068822  1
+  P  P45  1  0.90613370  0.37500000  0.79068822  1
+  P  P46  1  0.90613370  0.87500000  0.29068822  1
+  P  P47  1  0.90613370  0.87500000  0.79068821  1
+  O  O48  1  0.09423067  0.12500000  0.37239328  1
+  O  O49  1  0.09423067  0.12500000  0.87239328  1
+  O  O50  1  0.09423067  0.62500000  0.37239328  1
+  O  O51  1  0.09423067  0.62500000  0.87239328  1
+  O  O52  1  0.83415452  0.27277902  0.35686767  1
+  O  O53  1  0.83415452  0.27277902  0.85686767  1
+  O  O54  1  0.83415452  0.77277902  0.35686767  1
+  O  O55  1  0.83415452  0.77277902  0.85686767  1
+  O  O56  1  0.83415452  0.47722149  0.35686767  1
+  O  O57  1  0.83415452  0.47722149  0.85686767  1
+  O  O58  1  0.83415452  0.97722149  0.35686767  1
+  O  O59  1  0.83415452  0.97722149  0.85686767  1
+  O  O60  1  0.90577033  0.37500000  0.12760722  1
+  O  O61  1  0.90577033  0.37500000  0.62760722  1
+  O  O62  1  0.90577033  0.87500000  0.12760722  1
+  O  O63  1  0.90577033  0.87500000  0.62760722  1
+  O  O64  1  0.04430856  0.37500000  0.35493223  1
+  O  O65  1  0.04430856  0.37500000  0.85493222  1
+  O  O66  1  0.04430856  0.87500000  0.35493223  1
+  O  O67  1  0.04430856  0.87500000  0.85493223  1
+  O  O68  1  0.40576933  0.37500000  0.12239278  1
+  O  O69  1  0.40576933  0.37500000  0.62239278  1
+  O  O70  1  0.40576933  0.87500000  0.12239278  1
+  O  O71  1  0.40576933  0.87500000  0.62239278  1
+  O  O72  1  0.16584548  0.02277901  0.14313283  1
+  O  O73  1  0.16584548  0.02277902  0.64313283  1
+  O  O74  1  0.16584548  0.52277902  0.14313283  1
+  O  O75  1  0.16584548  0.52277902  0.64313283  1
+  O  O76  1  0.66584548  0.22722098  0.10686717  1
+  O  O77  1  0.66584548  0.22722098  0.60686717  1
+  O  O78  1  0.66584548  0.72722098  0.10686717  1
+  O  O79  1  0.66584548  0.72722098  0.60686717  1
+  O  O80  1  0.33415452  0.27277902  0.39313283  1
+  O  O81  1  0.33415452  0.27277902  0.89313283  1
+  O  O82  1  0.33415452  0.77277902  0.39313283  1
+  O  O83  1  0.33415452  0.77277902  0.89313283  1
+  O  O84  1  0.33415452  0.47722149  0.39313283  1
+  O  O85  1  0.33415452  0.47722149  0.89313283  1
+  O  O86  1  0.33415452  0.97722149  0.39313283  1
+  O  O87  1  0.33415452  0.97722149  0.89313283  1
+  O  O88  1  0.95569144  0.12500000  0.14506827  1
+  O  O89  1  0.95569144  0.12500000  0.64506827  1
+  O  O90  1  0.95569144  0.62500000  0.14506827  1
+  O  O91  1  0.95569144  0.62500000  0.64506828  1
+  O  O92  1  0.45569144  0.12500000  0.10493173  1
+  O  O93  1  0.45569144  0.12500000  0.60493173  1
+  O  O94  1  0.45569144  0.62500000  0.10493173  1
+  O  O95  1  0.45569144  0.62500000  0.60493173  1
+  O  O96  1  0.54430856  0.37500000  0.39506877  1
+  O  O97  1  0.54430856  0.37500000  0.89506878  1
+  O  O98  1  0.54430856  0.87500000  0.39506877  1
+  O  O99  1  0.54430856  0.87500000  0.89506878  1
+  O  O100  1  0.59423067  0.12500000  0.37760772  1
+  O  O101  1  0.59423067  0.12500000  0.87760772  1
+  O  O102  1  0.59423067  0.62500000  0.37760772  1
+  O  O103  1  0.59423067  0.62500000  0.87760772  1
+  O  O104  1  0.66584548  0.02277901  0.10686717  1
+  O  O105  1  0.66584548  0.02277901  0.60686717  1
+  O  O106  1  0.66584548  0.52277902  0.10686717  1
+  O  O107  1  0.66584548  0.52277902  0.60686717  1
+  O  O108  1  0.16584548  0.22722098  0.14313283  1
+  O  O109  1  0.16584548  0.22722099  0.64313283  1
+  O  O110  1  0.16584548  0.72722098  0.14313283  1
+  O  O111  1  0.16584548  0.72722098  0.64313283  1

matcalc/neb.py

@@ -0,0 +1,137 @@
+"""NEB calculations."""
+from __future__ import annotations
+
+import os
+from inspect import isclass
+from typing import TYPE_CHECKING
+
+from ase import optimize
+from ase.io import Trajectory
+from ase.neb import NEB, NEBTools
+from ase.optimize.optimize import Optimizer
+from pymatgen.core import Structure
+from pymatgen.io.ase import AseAtomsAdaptor
+
+from matcalc.base import PropCalc
+from matcalc.util import get_universal_calculator
+
+if TYPE_CHECKING:
+    from ase.calculators.calculator import Calculator
+
+
+class NEBCalc(PropCalc):
+    """Nudged Elastic Band calculator."""
+
+    def __init__(
+        self,
+        images: list[Structure],
+        calculator: str | Calculator = "M3GNet-MP-2021.2.8-DIRECT-PES",
+        optimizer: str | Optimizer = "BFGS",
+        traj_folder: str | None = None,
+        interval: int = 1,
+        climb: bool = True,
+        **kwargs,
+    ):
+        """
+        Args:
+            images(list): A list of pymatgen structures as NEB image structures.
+            calculator(str|Calculator): ASE Calculator to use. Default to M3GNet-MP-2021.2.8-DIRECT-PES.
+            optimizer(str|Optimizer): The optimization algorithm. Defaults to "BEGS".
+            traj_folder(str|None): The folder address to store NEB trajectories. Default to None.
+            interval(int): The step interval for saving the trajectories. Defaults to 1.
+            climb(bool): Whether to enable climb image NEB. Default to True.
+            kwargs: Other arguments passed to ASE NEB object.
+        """
+        self.images = images
+        self.calculator = calculator
+
+        # check str is valid optimizer key
+        def is_ase_optimizer(key):
+            return isclass(obj := getattr(optimize, key)) and issubclass(obj, Optimizer)
+
+        valid_keys = [key for key in dir(optimize) if is_ase_optimizer(key)]
+        if isinstance(optimizer, str) and optimizer not in valid_keys:
+            raise ValueError(f"Unknown {optimizer=}, must be one of {valid_keys}")
+
+        self.optimizer: Optimizer = getattr(optimize, optimizer) if isinstance(optimizer, str) else optimizer
+        self.traj_folder = traj_folder
+        self.interval = interval
+        self.climb = climb
+
+        self.images = []
+        for image in images:
+            if isinstance(image, Structure):
+                atoms = AseAtomsAdaptor().get_atoms(image)
+            atoms.calc = get_universal_calculator(self.calculator)
+            self.images.append(atoms)
+
+        self.neb = NEB(
+            self.images,
+            climb=self.climb,
+            allow_shared_calculator=True,
+            **kwargs,
+        )
+        self.optimizer = self.optimizer(self.neb)
+
+    @classmethod
+    def from_end_images(
+        cls,
+        start_struct: Structure,
+        end_struct: Structure,
+        calculator: str | Calculator = "M3GNet-MP-2021.2.8-DIRECT-PES",
+        nimages: int = 7,
+        interpolate_lattices: bool = False,
+        autosort_tol: float = 0.5,
+        **kwargs,
+    ):
+        """
+        Initialize a NEBCalc from end images.
+
+        Args:
+            start_struct(Structure): The starting image as a pymatgen Structure.
+            end_struct(Structure): The ending image as a pymatgen Structure.
+            calculator(str|Calculator): ASE Calculator to use. Default to M3GNet-MP-2021.2.8-DIRECT-PES.
+            nimages(int): The number of intermediate image structures to create.
+            interpolate_lattices(bool): Whether to interpolate the lattices when creating NEB
+                path with Structure.interpolate() in pymatgen. Default to False.
+            autosort_tol(float): A distance tolerance in angstrom in which to automatically
+                sort end_struct to match to the closest points in start_struct. This
+                argument is required for Structure.interpolate() in pymatgen.
+                Default to 0.5.
+            kwargs: Other arguments passed to construct NEBCalc.
+        """
+        images = start_struct.interpolate(
+            end_struct,
+            nimages=nimages + 1,
+            interpolate_lattices=interpolate_lattices,
+            pbc=False,
+            autosort_tol=autosort_tol,
+        )
+        return cls(images=images, calculator=calculator, **kwargs)
+
+    def calc(
+        self,
+        fmax: float = 0.1,
+        max_steps: int = 1000,
+    ):
+        """
+        Perform NEB calculation.
+
+        Args:
+            fmax (float): Convergence criteria for NEB calculations defined by Max forces.
+                Default to 0.1 eV/Angstrom.
+            max_steps (int): Maximum number of steps in NEB calculations. Default to 1000.
+
+        Returns:
+            NEB barrier.
+        """
+        if self.traj_folder is not None:
+            os.makedirs(self.traj_folder, exist_ok=True)
+            for i in range(len(self.images)):
+                self.optimizer.attach(
+                    Trajectory(f"{self.traj_folder}/image_{i}.traj", "w", self.images[i]),
+                    interval=self.interval,
+                )
+        self.optimizer.run(fmax=fmax, steps=max_steps)
+        neb_tool = NEBTools(self.neb.images)
+        return neb_tool.get_barrier()

matcalc/util.py

@@ -42,7 +42,7 @@ def get_universal_calculator(name: str | Calculator, **kwargs) -> Calculator:
         from matgl.ext.ase import M3GNetCalculator
 
         # M3GNet is shorthand for latest M3GNet based on DIRECT sampling.
-        name = {"m3gnet": "M3GNet-MP-2021.2.8-PES"}.get(name.lower(), name)
+        name = {"m3gnet": "M3GNet-MP-2021.2.8-DIRECT-PES"}.get(name.lower(), name)
         model = matgl.load_model(name)
         kwargs.setdefault("stress_weight", 1.0 / 160.21766208)
         return M3GNetCalculator(potential=model, **kwargs)

tests/test_neb.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+import pytest
+
+from matcalc.neb import NEBCalc
+
+
+def test_neb_calc(LiFePO4, M3GNetCalc):
+    """Tests for NEBCalc class"""
+    image_start = LiFePO4.copy()
+    image_start.remove_sites([2])
+    image_end = LiFePO4.copy()
+    image_end.remove_sites([3])
+    NEBcalc = NEBCalc.from_end_images(image_start, image_end, M3GNetCalc, nimages=5)
+    barriers = NEBcalc.calc(fmax=0.5)
+    print(barriers)
+    assert len(NEBcalc.neb.images) == 7
+    assert barriers[0] == pytest.approx(0.0184783935546875, rel=0.002)
+    assert barriers[1] == pytest.approx(0.0018920898, rel=0.002)