diff --git a/.ci_support/environment-lammps.yml b/.ci_support/environment-lammps.yml index b68d377dd..3d3abddcb 100644 --- a/.ci_support/environment-lammps.yml +++ b/.ci_support/environment-lammps.yml @@ -1,4 +1,5 @@ channels: - conda-forge dependencies: -- lammps =2024.02.07 \ No newline at end of file +- lammps =2024.02.07 +- fitsnap3 =3.1.0.2 \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index b6b5fb752..b9486cb1e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,7 @@ phonopy = [ "spglib==2.5.0", ] pyxtal = ["pyxtal==0.6.7"] +fitsnap3 = ["fitsnap3==3.1.0.2"] [tool.setuptools.packages.find] include = ["structuretoolkit*"] diff --git a/structuretoolkit/analyse/__init__.py b/structuretoolkit/analyse/__init__.py index e2af4e8f9..325c74da3 100644 --- a/structuretoolkit/analyse/__init__.py +++ b/structuretoolkit/analyse/__init__.py @@ -29,6 +29,19 @@ ) from structuretoolkit.analyse.strain import get_strain +try: + from structuretoolkit.analyse.fitsnap import ( + get_ace_descriptor_derivatives, + ) + from structuretoolkit.analyse.fitsnap import ( + get_ace_descriptor_derivatives as get_ace_descriptor_derivatives_fitsnap, + ) + from structuretoolkit.analyse.fitsnap import ( + get_snap_descriptor_derivatives as get_snap_descriptor_derivatives_fitsnap, + ) +except ImportError: + pass + def get_symmetry( structure, use_magmoms=False, use_elements=True, symprec=1e-5, angle_tolerance=-1.0 diff --git a/structuretoolkit/analyse/fitsnap.py b/structuretoolkit/analyse/fitsnap.py new file mode 100644 index 000000000..f9521fb06 --- /dev/null +++ b/structuretoolkit/analyse/fitsnap.py @@ -0,0 +1,251 @@ +import random +from typing import Optional, Union + +import numpy as np +from ase.atoms import Atoms +from fitsnap3lib.fitsnap import FitSnap + +from structuretoolkit.analyse.snap import _get_lammps_compatible_cell + + +def get_ace_descriptor_derivatives( + structure: Atoms, + atom_types: list[str], + ranks: list[int] = [1, 2, 3, 4], + lmax: list[int] = [0, 5, 2, 1], + nmax: list[int] = [22, 5, 3, 1], + mumax: int = 1, + nmaxbase: int = 22, + erefs: list[float] = [0.0], + rcutfac: float = 4.5, + rcinner: float = 1.2, + drcinner: float = 0.01, + RPI_heuristic: str = "root_SO3_span", + lambda_value: float = 1.275, + lmin: list[int] = [0, 0, 1, 1], + bzeroflag: bool = True, + cutoff: float = 10.0, +) -> np.ndarray: + """ + Calculate per atom ACE descriptors using FitSNAP https://fitsnap.github.io + + Args: + structure (ase.atoms.Atoms): atomistic structure as ASE atoms object + atom_types (list[str]): list of element types + ranks (list): + lmax (list): + nmax (list): + nmaxbase (int): + rcutfac (float): + lambda_value (float): + lmin (list): + cutoff (float): cutoff radius for the construction of the neighbor list + + Returns: + np.ndarray: Numpy array with the calculated descriptor derivatives + """ + settings = { + "ACE": { + "numTypes": len(atom_types), + "ranks": " ".join([str(r) for r in ranks]), + "lmax": " ".join([str(l) for l in lmax]), + "nmax": " ".join([str(n) for n in nmax]), + "mumax": mumax, + "nmaxbase": nmaxbase, + "rcutfac": rcutfac, + "erefs": " ".join([str(e) for e in erefs]), + "rcinner": rcinner, + "drcinner": drcinner, + "RPI_heuristic": RPI_heuristic, + "lambda": lambda_value, + "type": " ".join(atom_types), + "lmin": " ".join([str(l) for l in lmin]), + "bzeroflag": bzeroflag, + "bikflag": True, + "dgradflag": True, + }, + "CALCULATOR": { + "calculator": "LAMMPSPACE", + "energy": 1, + "force": 1, + "stress": 0, + }, + "REFERENCE": { + "units": "metal", + "atom_style": "atomic", + "pair_style": "zero " + str(cutoff), + "pair_coeff": "* *", + }, + } + fs = FitSnap(settings, comm=None, arglist=["--overwrite"]) + a, b, w = fs.calculator.process_single(_ase_scraper(data=[structure])[0]) + return a + + +def get_snap_descriptor_derivatives( + structure: Atoms, + atom_types: list[str], + twojmax: int = 6, + element_radius: list[int] = [4.0], + rcutfac: float = 1.0, + rfac0: float = 0.99363, + rmin0: float = 0.0, + bzeroflag: bool = False, + quadraticflag: bool = False, + weights: Optional[Union[list, np.ndarray]] = None, + cutoff: float = 10.0, +) -> np.ndarray: + """ + Calculate per atom SNAP descriptors using FitSNAP https://fitsnap.github.io + + Args: + structure (ase.atoms.Atoms): atomistic structure as ASE atoms object + atom_types (list[str]): list of element types + twojmax (int): band limit for bispectrum components (non-negative integer) + element_radius (list[int]): list of radii for the individual elements + rcutfac (float): scale factor applied to all cutoff radii (positive real) + rfac0 (float): parameter in distance to angle conversion (0 < rcutfac < 1) + rmin0 (float): parameter in distance to angle conversion (distance units) + bzeroflag (bool): subtract B0 + quadraticflag (bool): generate quadratic terms + weights (list/np.ndarry/None): list of neighbor weights, one for each type + cutoff (float): cutoff radius for the construction of the neighbor list + + Returns: + np.ndarray: Numpy array with the calculated descriptor derivatives + """ + if weights is None: + weights = [1.0] * len(atom_types) + settings = { + "BISPECTRUM": { + "numTypes": len(atom_types), + "twojmax": twojmax, + "rcutfac": rcutfac, + "rfac0": rfac0, + "rmin0": rmin0, + "wj": " ".join([str(w) for w in weights]), + "radelem": " ".join([str(r) for r in element_radius]), + "type": " ".join(atom_types), + "wselfallflag": 0, + "chemflag": 0, + "bzeroflag": bzeroflag, + "quadraticflag": quadraticflag, + }, + "CALCULATOR": { + "calculator": "LAMMPSSNAP", + "energy": 1, + "force": 1, + "stress": 0, + }, + "REFERENCE": { + "units": "metal", + "atom_style": "atomic", + "pair_style": "zero " + str(cutoff), + "pair_coeff": "* *", + }, + } + fs = FitSnap(settings, comm=None, arglist=["--overwrite"]) + a, b, w = fs.calculator.process_single(_ase_scraper(data=[structure])[0]) + return a + + +def _assign_validation(group_table): + """ + Given a dictionary of group info, add another key for test bools. + + Args: + group_table: Dictionary of group names. Must have keys "nconfigs" and "testing_size". + + Modifies the dictionary in place by adding another key "test_bools". + """ + + for name in group_table: + nconfigs = group_table[name]["nconfigs"] + assert "testing_size" in group_table[name] + assert group_table[name]["testing_size"] <= 1.0 + test_bools = [ + random.random() < group_table[name]["testing_size"] + for i in range(0, nconfigs) + ] + + group_table[name]["test_bools"] = test_bools + + +def _ase_scraper(data) -> list: + """ + Function to organize groups and allocate shared arrays used in Calculator. For now when using + ASE frames, we don't have groups. + + Args: + s: fitsnap instance. + data: List of ASE frames or dictionary group table containing frames. + + Returns a list of data dictionaries suitable for fitsnap descriptor calculator. + If running in parallel, this list will be distributed over procs, so that each proc will have a + portion of the list. + """ + + # Simply collate data from Atoms objects if we have a list of Atoms objects. + if type(data) == list: + # s.data = [collate_data(atoms) for atoms in data] + return [_collate_data(atoms) for atoms in data] + # If we have a dictionary, assume we are dealing with groups. + elif type(data) == dict: + _assign_validation(data) + # s.data = [] + ret = [] + for name in data: + frames = data[name]["frames"] + # Extend the fitsnap data list with this group. + # s.data.extend([collate_data(atoms, name, data[name]) for atoms in frames]) + ret.extend([_collate_data(atoms, name, data[name]) for atoms in frames]) + return ret + else: + raise Exception("Argument must be list or dictionary for ASE scraper.") + + +def _collate_data(atoms, name: str = None, group_dict: dict = None) -> dict: + """ + Function to organize fitting data for FitSNAP from ASE atoms objects. + + Args: + atoms: ASE atoms object for a single configuration of atoms. + name: Optional name of this configuration. + group_dict: Optional dictionary containing group information. + + Returns a data dictionary for a single configuration. + """ + + # Transform ASE cell to be appropriate for LAMMPS. + apre = _get_lammps_compatible_cell(cell=atoms.cell) + R = np.dot(np.linalg.inv(atoms.cell), apre) + positions = np.matmul(atoms.get_positions(), R) + cell = apre.T + + # Make a data dictionary for this config. + + data = {} + data["Group"] = name # 'ASE' # TODO: Make this customizable for ASE groups. + data["File"] = None + data["Stress"] = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]] + data["Positions"] = positions + data["Energy"] = 0.0 + data["AtomTypes"] = atoms.get_chemical_symbols() + data["NumAtoms"] = len(atoms) + data["Forces"] = np.array([0.0, 0.0, 0.0] * len(atoms)) + data["QMLattice"] = cell + data["test_bool"] = 0 + data["Lattice"] = cell + data["Rotation"] = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + data["Translation"] = np.zeros((len(atoms), 3)) + # Inject the weights. + if group_dict is not None: + data["eweight"] = group_dict["eweight"] if "eweight" in group_dict else 1.0 + data["fweight"] = group_dict["fweight"] if "fweight" in group_dict else 1.0 + data["vweight"] = group_dict["vweight"] if "vweight" in group_dict else 1.0 + else: + data["eweight"] = 1.0 + data["fweight"] = 1.0 + data["vweight"] = 1.0 + + return data diff --git a/tests/test_fitsnap.py b/tests/test_fitsnap.py new file mode 100644 index 000000000..fea3a709c --- /dev/null +++ b/tests/test_fitsnap.py @@ -0,0 +1,71 @@ +from ase.build.bulk import bulk +import structuretoolkit as stk +import unittest + + +try: + from structuretoolkit.analyse.fitsnap import ( + get_ace_descriptor_derivatives, + get_snap_descriptor_derivatives, + ) + + skip_snap_test = False +except ImportError: + skip_snap_test = True + + +@unittest.skipIf( + skip_snap_test, "LAMMPS is not installed, so the SNAP tests are skipped." +) +class TestSNAP(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.structure = bulk("Cu", cubic=True) + cls.numtypes = 1 + cls.twojmax = 6 + cls.rcutfac = 1.0 + cls.rfac0 = 0.99363 + cls.rmin0 = 0.0 + cls.bzeroflag = False + cls.quadraticflag = False + cls.radelem = [4.0] + cls.type = ["Cu"] + cls.wj = [1.0] + + def test_get_snap_descriptor_derivatives(self): + n_coeff = len(stk.analyse.get_snap_descriptor_names(twojmax=self.twojmax)) + mat_a = get_snap_descriptor_derivatives( + structure=self.structure, + atom_types=self.type, + twojmax=self.twojmax, + element_radius=self.radelem, + rcutfac=self.rcutfac, + rfac0=self.rfac0, + rmin0=self.rmin0, + bzeroflag=self.bzeroflag, + quadraticflag=self.quadraticflag, + weights=self.wj, + cutoff=10.0, + ) + self.assertEqual(mat_a.shape, (len(self.structure) * 3 + 7, n_coeff + 1)) + + def test_get_ace_descriptor_derivatives(self): + mat_a = get_ace_descriptor_derivatives( + structure=self.structure, + atom_types=self.type, + ranks=[1, 2, 3, 4], + lmax=[0, 5, 2, 1], + nmax=[22, 5, 3, 1], + mumax=1, + nmaxbase=22, + erefs=[0.0], + rcutfac=4.5, + rcinner=1.2, + drcinner=0.01, + RPI_heuristic="root_SO3_span", + lambda_value=1.275, + lmin=[0, 0, 1, 1], + bzeroflag=True, + cutoff=10.0, + ) + self.assertEqual(mat_a.shape, (16, 141))