Feat: Copy the major codes of mattersim

src/mattersim/__init__.py

@@ -0,0 +1,2 @@
+# -*- coding: utf-8 -*-
+from .__version__ import __version__  # noqa: F401

src/mattersim/__version__.py

@@ -0,0 +1,5 @@
+# -*- coding: utf-8 -*-
+import pkg_resources
+
+# Get the version from setup.py
+__version__ = pkg_resources.get_distribution("mattersim").version

src/mattersim/applications/moldyn.py

@@ -0,0 +1,145 @@
+# -*- coding: utf-8 -*-
+from typing import Union
+
+from ase import Atoms, units
+from ase.io import Trajectory
+from ase.md.npt import NPT
+from ase.md.nvtberendsen import NVTBerendsen
+from ase.md.velocitydistribution import (  # noqa: E501
+    MaxwellBoltzmannDistribution,
+    Stationary,
+)
+
+
+class MolecularDynamics:
+    """
+    This class is used for Molecular Dynamics.
+    """
+
+    SUPPORTED_ENSEMBLE = ["NVT_BERENDSEN", "NVT_NOSE_HOOVER"]
+
+    def __init__(
+        self,
+        atoms: Atoms,
+        ensemble: str = "nvt_nose_hoover",
+        temperature: float = 300,
+        timestep: float = 1.0,
+        taut: Union[float, None] = None,
+        trajectory: Union[str, Trajectory, None] = None,
+        logfile: Union[str, None] = "-",
+        loginterval: int = 10,
+        append_trajectory: bool = False,
+    ):
+        """
+        Args:
+            atoms (Union[Atoms, Structure]): ASE atoms object contains
+                structure information and calculator.
+            ensemble (str, optional): Simulation ensemble choosen. Defaults
+                to nvt_nose_hoover'
+            temperature (float, optional): Simulation temperature, in Kelvin.
+                Defaults to 300 K.
+            timestep (float, optional): The simulation time step, in fs. Defaults
+                to 1 fs.
+            taut (float, optional): Characteristic timescale of the thermostat,
+                in fs. If is None, automatically set it to 1000 * timestep.
+            trajectory (Union[str, Trajectory], optional): Attach trajectory
+                object. If trajectory is a string a Trajectory will be constructed.
+                Defaults to None, which means for no trajectory.
+            logfile (str, optional): If logfile is a string, a file with that name
+                will be opened. Defaults to '-', which means output to stdout.
+            loginterval (int, optional): Only write a log line for every loginterval
+                time steps. Defaults to 10.
+            append_trajectory (bool, optional): If False the trajectory file to be
+                overwriten each time the dynamics is restarted from scratch. If True,
+                the new structures are appended to the trajectory file instead.
+
+        """
+        assert atoms.calc is not None, (
+            "Molecular Dynamics simulation only accept "
+            "ase atoms with an attached calculator"
+        )
+        if ensemble.upper() not in self.SUPPORTED_ENSEMBLE:
+            raise NotImplementedError(  # noqa: E501
+                f"Ensemble {ensemble} is not yet supported."
+            )
+
+        self.atoms = atoms
+
+        self.ensemble = ensemble.upper()
+        self._temperature = temperature
+        self.timestep = timestep
+
+        if taut is None:
+            taut = 1000 * timestep * units.fs
+        self.taut = taut
+
+        self._trajectory = trajectory
+        self.logfile = logfile
+        self.loginterval = loginterval
+        self.append_trajectory = append_trajectory
+
+        self._initialize_dynamics()
+
+    def _initialize_dynamics(self):
+        """
+        Initialize the Dynamic ensemble class.
+        """
+        MaxwellBoltzmannDistribution(
+            self.atoms, temperature_K=self._temperature, force_temp=True
+        )
+        Stationary(self.atoms)
+
+        if self.ensemble == "NVT_BERENDSEN":  # noqa: E501
+            self.dyn = NVTBerendsen(
+                self.atoms,
+                timestep=self.timestep * units.fs,
+                temperature_K=self._temperature,
+                taut=self.taut,
+                trajectory=self._trajectory,
+                logfile=self.logfile,
+                loginterval=self.loginterval,
+                append_trajectory=self.append_trajectory,
+            )
+        elif self.ensemble == "NVT_NOSE_HOOVER":
+            self.dyn = NPT(
+                self.atoms,
+                timestep=self.timestep * units.fs,
+                temperature_K=self._temperature,
+                ttime=self.taut,
+                pfactor=None,
+                trajectory=self._trajectory,
+                logfile=self.logfile,
+                loginterval=self.loginterval,
+                append_trajectory=self.append_trajectory,
+            )
+        else:
+            raise NotImplementedError(  # noqa: E501
+                f"Ensemble {self.ensemble} is not yet supported."
+            )
+
+    def run(self, n_steps: int = 1):
+        """
+        Run Molecular Dynamic simulation.
+
+        Args:
+            n_steps (int, optional): Number of steps to simulations. Defaults to 1.
+        """
+        self.dyn.run(n_steps)
+
+    @property
+    def temperature(self):
+        return self._temperature
+
+    @temperature.setter
+    def temperature(self, temperature: float):
+        self._temperature = temperature
+        self._initialize_dynamics()
+
+    @property
+    def trajectory(self):
+        return self._trajectory
+
+    @trajectory.setter
+    def trajectory(self, trajectory: Union[str, Trajectory, None]):
+        self._trajectory = trajectory
+        self._initialize_dynamics()

src/mattersim/applications/phonon.py

@@ -0,0 +1,228 @@
+# -*- coding: utf-8 -*-
+import datetime
+import os
+from typing import Iterable, Union
+
+import numpy as np
+from ase import Atoms
+from phonopy import Phonopy
+from tqdm import tqdm
+
+from mattersim.utils.phonon_utils import (
+    get_primitive_cell,
+    to_ase_atoms,
+    to_phonopy_atoms,
+)
+from mattersim.utils.supercell_utils import get_supercell_parameters
+
+
+class PhononWorkflow(object):
+    """
+    This class is used to calculate the phonon dispersion relationship of
+    material using phonopy
+    """
+
+    def __init__(
+        self,
+        atoms: Atoms,
+        find_prim: bool = False,
+        work_dir: str = None,
+        amplitude: float = 0.01,
+        supercell_matrix: np.ndarray = None,
+        qpoints_mesh: np.ndarray = None,
+        max_atoms: int = None,
+    ):
+        """_summary
+
+        Args:
+            atoms (Atoms): ASE atoms object contains structure information and
+                calculator.
+            find_prim (bool, optional): If find the primitive cell and use it
+                to calculate phonon. Default to False.
+            work_dir (str, optional): workplace path to contain phonon result.
+                Defaults to data + chemical_symbols + 'phonon'
+            amplitude (float, optional): Magnitude of the finite difference to
+                displace in force constant calculation, in Angstrom. Defaults
+                to 0.01 Angstrom.
+            supercell_matrix (nd.array, optional): Supercell matrix for constr
+                -uct supercell, priority over than max_atoms. Defaults to None.
+            qpoints_mesh (nd.array, optional): Qpoint mesh for IBZ integral,
+                priority over than max_atoms. Defaults to None.
+            max_atoms (int, optional): Maximum atoms number limitation for the
+                supercell generation. If not set, will automatic generate super
+                -cell based on symmetry. Defaults to None.
+        """
+        assert (
+            atoms.calc is not None
+        ), "PhononWorkflow only accepts ase atoms with an attached calculator"
+        if find_prim:
+            self.atoms = get_primitive_cell(atoms)
+            self.atoms.calc = atoms.calc
+        else:
+            self.atoms = atoms
+        if work_dir is not None:
+            self.work_dir = work_dir
+        else:
+            current_datetime = datetime.datetime.now()
+            formatted_datetime = current_datetime.strftime("%Y-%m-%d-%H-%M")
+            self.work_dir = (
+                f"{formatted_datetime}-{atoms.get_chemical_formula()}-phonon"
+            )
+        self.amplitude = amplitude
+        if supercell_matrix is not None:
+            if supercell_matrix.shape == (3, 3):
+                self.supercell_matrix = supercell_matrix
+            elif supercell_matrix.shape == (3,):
+                self.supercell_matrix = np.diag(supercell_matrix)
+            else:
+                assert (
+                    False
+                ), "Supercell matrix must be an array (3,1) or a matrix (3,3)."
+        else:
+            self.supercell_matrix = supercell_matrix
+
+        if qpoints_mesh is not None:
+            assert qpoints_mesh.shape == (3,), "Qpoints mesh must be an array (3,1)."
+            self.qpoints_mesh = qpoints_mesh
+        else:
+            self.qpoints_mesh = qpoints_mesh
+
+        self.max_atoms = max_atoms
+
+    def compute_force_constants(self, atoms: Atoms, nrep_second: np.ndarray):
+        """
+        Calculate force constants
+
+        Args:
+            atoms (Atoms): ASE atoms object to provide lattice informations.
+            nrep_second (np.ndarray): Supercell size used for 2nd force
+                constant calculations.
+        """
+        print(f"Supercell matrix for 2nd force constants : \n{nrep_second}")
+        # Generate phonopy object
+        phonon = Phonopy(
+            to_phonopy_atoms(atoms),
+            supercell_matrix=nrep_second,
+            primitive_matrix="auto",
+            log_level=2,
+        )
+
+        # Generate displacements
+        phonon.generate_displacements(distance=self.amplitude)
+
+        # Compute force constants
+        second_scs = phonon.supercells_with_displacements
+        second_force_sets = []
+        print("\n")
+        print("Inferring forces for displaced atoms and computing fcs ...")
+        for disp_second in tqdm(second_scs):
+            pa_second = to_ase_atoms(disp_second)
+            pa_second.calc = self.atoms.calc
+            second_force_sets.append(pa_second.get_forces())
+
+        phonon.forces = np.array(second_force_sets)
+        phonon.produce_force_constants()
+        phonon.symmetrize_force_constants()
+
+        return phonon
+
+    @staticmethod
+    def compute_phonon_spectrum_dos(
+        atoms: Atoms, phonon: Phonopy, k_point_mesh: Union[int, Iterable[int]]
+    ):
+        """
+        Calculate phonon spectrum and DOS based on force constant matrix in
+        phonon object
+
+        Args:
+            atoms (Atoms): ASE atoms object to provide lattice information
+            phonon (Phonopy): Phonopy object which contains force constants matrix
+            k_point_mesh (Union[int, Iterable[int]]): The qpoints number in First
+                Brillouin Zone in three directions for DOS calculation.
+        """
+        print(f"Qpoints mesh for Brillouin Zone integration : {k_point_mesh}")
+        phonon.run_mesh(k_point_mesh)
+        print(
+            "Dispersion relations using phonopy for "
+            + str(atoms.symbols)
+            + " ..."
+            + "\n"
+        )
+
+        # plot phonon spectrum
+        phonon.auto_band_structure(plot=True, write_yaml=True).savefig(
+            f"{str(atoms.symbols)}_phonon_band.png", dpi=300
+        )
+        phonon.auto_total_dos(plot=True, write_dat=True).savefig(
+            f"{str(atoms.symbols)}_phonon_dos.png", dpi=300
+        )
+
+        # Save additional files
+        phonon.save(settings={"force_constants": True})
+
+    @staticmethod
+    def check_imaginary_freq(phonon: Phonopy):
+        """
+        Check whether phonon has imaginary frequency.
+
+        Args:
+            phonon (Phonopy): Phonopy object which contains phonon spectrum frequency.
+        """
+        band_dict = phonon.get_band_structure_dict()
+        frequencies = np.concatenate(
+            [np.array(freq).flatten() for freq in band_dict["frequencies"]], axis=None
+        )
+        has_imaginary = False
+        if np.all(np.array(frequencies) >= -0.299):
+            pass
+        else:
+            print("Warning! Imaginary frequencies found!")
+            has_imaginary = True
+
+        return has_imaginary
+
+    def run(self):
+        """
+        The entrypoint to start the workflow.
+        """
+        current_path = os.path.abspath(".")
+        try:
+            # check folder exists
+            if not os.path.exists(self.work_dir):
+                os.makedirs(self.work_dir)
+
+            os.chdir(self.work_dir)
+
+            try:
+                # Generate supercell parameters based on optimized structures
+                nrep_second, k_point_mesh = get_supercell_parameters(
+                    self.atoms, self.supercell_matrix, self.qpoints_mesh, self.max_atoms
+                )
+            except Exception as e:
+                print("Error whille generating supercell parameters:", e)
+                raise
+
+            try:
+                # Calculate 2nd force constants
+                phonon = self.compute_force_constants(self.atoms, nrep_second)
+            except Exception as e:
+                print("Error while computing force constants:", e)
+                raise
+
+            try:
+                # Calculate phonon spectrum
+                self.compute_phonon_spectrum_dos(self.atoms, phonon, k_point_mesh)
+                # check whether has imaginary frequency
+                has_imaginary = self.check_imaginary_freq(phonon)
+            except Exception as e:
+                print("Error while computing phonon spectrum and dos:", e)
+                raise
+
+        except Exception as e:
+            print("An error occurred during the Phonon workflow:", e)
+            raise
+
+        finally:
+            os.chdir(current_path)
+
+        return has_imaginary, phonon