class Structure:
class StructureSource(object):
class ForceSource(object):
class Trajectory(StructureSource, ForceSource):
    def __init__(
        self,
        cell: "ndarray",
        species: Union[List[str], List[int]],
        positions: "ndarray",
        mass_dict: dict = None,
        prev_positions: "ndarray" = None,
        species_labels: List[str] = None,
        forces=None,
        stds=None,
        energy: float = None,
    ):
    def positions(self):
    def wrapped_positions(self):
    def positions(self, position_array):
    def cell(self):
    def vec1(self):
    def vec2(self):
    def vec3(self):
    def cell_transpose(self):
    def cell_transpose_inverse(self):
    def cell_dot(self):
    def cell_dot_inverse(self):
    def cell(self, cell_array):
    def get_cell_dot(cell_array):
    def raw_to_relative(
        positions: "ndarray", cell_transpose: "ndarray", cell_dot_inverse: "ndarray"
    ) -> "ndarray":
    def relative_to_raw(
        relative_positions: "ndarray",
        cell_transpose_inverse: "ndarray",
        cell_dot: "ndarray",
    ) -> "ndarray":
    def wrap_positions(self) -> "ndarray":
    def indices_of_specie(self, specie: Union[int, str]) -> List[int]:
    def __str__(self) -> str:
    def __len__(self) -> int:
    def as_dict(self) -> dict:
    def as_str(self) -> str:
    def from_dict(dictionary: dict) -> "flare.struc.Structure":
    def from_ase_atoms(atoms: "ase.Atoms", cell=None) -> "flare.struc.Structure":
    def to_ase_atoms(self) -> "ase.Atoms":
    def to_pmg_structure(self):
    def from_pmg_structure(structure: "pymatgen Structure") -> "flare Structure":
    def is_valid(self, tolerance: float = 0.5) -> bool:
    def to_xyz(
        self,
        extended_xyz: bool = True,
        print_stds: bool = False,
        print_forces: bool = False,
        print_max_stds: bool = False,
        print_energies: bool = False,
        predict_energy=None,
        dft_forces=None,
        dft_energy=None,
        target_atoms=None,
        timestep=-1,
        write_file: str = "",
        append: bool = False,
    ) -> str:
    def from_file(
        file_name: str, format: str = "", as_trajectory: bool = False
    ) -> Union["flare.struc.Structure", List["flare.struc.Structure"]]:
    def __eq__(self, other):
    def __init__(self):
    def get_next_structure(self) -> Structure:
    def write_file(self):
    def __init__(self):
    def get_next_force(self, *args, **kwargs) -> "np.ndarray":
    def pre_force(self, *args, **kwargs):
    def __init__(
        self, frames: List[Structure] = None, iterate_strategy: Union[int, str] = 1
    ):
    def get_next_structure(self) -> Union[Structure, None]:
    def cur_frame(self):
    def cur_forces(self):
    def get_next_force(self, index: int = -1) -> "np.ndarray":
    def __getitem__(self, item):
    def __len__(self):
    def __next__(self):
    def __iter__(self):
    def append(self, frame: Structure):
def get_unique_species(species: List[Any]) -> (List, List[int]): class AtomicEnvironment:
    def __init__(self, structure: Structure, atom: int, cutoffs, cutoffs_mask=None):
    def setup_mask(self, cutoffs_mask):
    def compute_env(self):
    def as_str(self) -> str:
    def as_dict(self, include_structure: bool = False):
    def from_dict(dictionary):
    def from_file(
        file_name: str, format: str = ""
    ) -> Union["flare.env.AtomicEnvironment", List["flare.env.AtomicEnvironment"]]:
    def __str__(self):
    def force(self) -> "np.ndarray": class GaussianProcess:
    def __init__(
        self,
        kernels: List[str] = None,
        component: str = "mc",
        hyps: "ndarray" = None,
        cutoffs: dict = None,
        hyps_mask: dict = None,
        hyp_labels: List = None,
        opt_algorithm: str = "L-BFGS-B",
        maxiter: int = 10,
        parallel: bool = False,
        per_atom_par: bool = True,
        n_cpus: int = 1,
        n_sample: int = 100,
        output: Output = None,
        name="default_gp",
        energy_noise: float = 0.01,
        **kwargs,
    ):
    def force_noise(self):
    def hyps_and_labels(self):
    def check_instantiation(self):
    def update_kernel(
        self,
        kernels: List[str],
        component: str = "mc",
        hyps=None,
        cutoffs: dict = None,
        hyps_mask: dict = None,
    ):
    def update_db(
        self,
        struc: Structure,
        forces: "ndarray" = None,
        custom_range: List[int] = (),
        energy: float = None,
        stress: "ndarray" = None,
    ):
    def add_one_env(
        self,
        env: AtomicEnvironment,
        force: "np.ndarray" = None,
        train: bool = False,
        **kwargs,
    ):
    def train(
        self,
        logger_name: str = None,
        custom_bounds=None,
        grad_tol: float = 1e-4,
        x_tol: float = 1e-5,
        line_steps: int = 20,
        print_progress: bool = False,
    ):
    def check_L_alpha(self):
    def predict(self, x_t: AtomicEnvironment, d: int) -> [float, float]:
    def predict_force_xyz(self, x_t: AtomicEnvironment) -> ("np.ndarray", "np.ndarray"):
    def predict_local_energy(self, x_t: AtomicEnvironment) -> float:
    def predict_local_energy_and_var(self, x_t: AtomicEnvironment):
    def predict_efs(self, x_t: AtomicEnvironment):
    def set_L_alpha(self):
    def update_L_alpha(self):
    def __str__(self):
    def as_dict(self):
    def sync_data(self):
    def from_dict(dictionary):
    def compute_matrices(self):
    def adjust_cutoffs(
        self,
        new_cutoffs: Union[list, tuple, "np.ndarray"] = None,
        reset_L_alpha=True,
        train=True,
        new_hyps_mask=None,
    ):
    def remove_force_data(
        self, indexes: Union[int, List[int]], update_matrices: bool = True
    ) -> Tuple[List[Structure], List["ndarray"]]:
    def write_model(
        self, name: str, format: str = None, split_matrix_size_cutoff: int = 5000
    ):
    def from_file(filename: str, format: str = ""):
    def __len__(self):
    def training_statistics(self) -> dict:
    def par(self):
    def __deepcopy__(self, memo):
    def __del__(self):
    def backward_arguments(kwargs, new_args={}):
    def backward_attributes(dictionary): class Parameters:
    """"""

    all_kernel_types = ["twobody", "threebody", "manybody"]
    cutoff_types = {"cut3b": "threebody"}
    ndim = {"twobody": 2, "threebody": 3, "manybody": 2, "cut3b": 2}
    n_kernel_parameters = {"twobody": 2, "threebody": 2, "manybody": 2, "cut3b": 0}

    cutoff_types_keys = list(cutoff_types.keys())
    cutoff_types_values = list(cutoff_types.values())

    logger = set_logger("Parameters", stream=True, fileout_name=None, verbose="info")

    def __init__(self):
        """
        Enumerate all keys and  their default values that hyps_mask should store
        """

        self.param_dict = {
            "nspecie": 1,
            "ntwobody": 0,
            "nthreebody": 0,
            "ncut3b": 0,
            "nmanybody": 0,
            "species_mask": None,
            "twobody_mask": None,
            "threebody_mask": None,
            "cut3b_mask": None,
            "manybody_mask": None,
            "twobody_cutoff_list": None,
            "threebody_cutoff_list": None,
            "manybody_cutoff_list": None,
            "train_noise": True,
            "energy_noise": 0,
            "map": None,
            "original_hyps": [],
            "original_labels": [],
        }
        self.hyps = None
        self.hyp_labels = None
        self.cutoffs = {}
        self.kernels = []

    @staticmethod
    def cutoff_array_to_dict(cutoffs):
        """
        Convert old cutoffs array to the new dictionary format
        """

        if isinstance(cutoffs, dict):
            return cutoffs

        if (cutoffs is not None) and not isinstance(cutoffs, dict):
            DeprecationWarning("cutoffs is replace by dictionary")
            newcutoffs = {"twobody": cutoffs[0]}
            if len(cutoffs) > 1:
                newcutoffs["threebody"] = cutoffs[1]
            if len(cutoffs) > 2:
                newcutoffs["manybody"] = cutoffs[2]
            Parameters.logger.debug("Convert cutoffs array to cutoffs dict")
            Parameters.logger.debug("Original", cutoffs)
            Parameters.logger.debug("Now", newcutoffs)
            return newcutoffs
        else:
            raise TypeError("cannot handle cutoffs with {type(cutoffs)} type")

    @staticmethod
    def backward(kernels, param_dict):

        if param_dict is None:
            param_dict = {}

        # update old keys to new keys. for example nspec to nspecies
        replace_list = {
            "spec": "specie",
            "bond": "twobody",
            "triplet": "threebody",
            "mb": "manybody",
        }
        keys = list(param_dict.keys())
        for key in keys:
            for original in replace_list:
                if original in key and replace_list[original] not in key:
                    newkey = key.replace(original, replace_list[original])
                    param_dict[newkey] = param_dict[key]
                    DeprecationWarning("{key} is being replaced with {newkey}")

        # add a couple new keys that was not there previously
        if "train_noise" not in param_dict:
            param_dict["train_noise"] = True
            DeprecationWarning("train_noise has to be in hyps_mask, set to True")
        if "nspecie" not in param_dict:
            param_dict["nspecie"] = 1

        # sort the kernels dictionary again. but this can result in
        # wrong results...
        if set(kernels) != set(param_dict.get("kernels", [])):

            start = 0
            for k in Parameters.all_kernel_types:
                if k in kernels:
                    if k + "_start" not in param_dict:
                        param_dict[k + "_start"] = start
                    if "n" + k not in param_dict:
                        Parameters.logger.debug(
                            "add in hyper parameter separatorsfor", k
                        )
                        param_dict["n" + k] = 1
                        start += Parameters.n_kernel_parameters[k]
                    else:
                        start += param_dict["n" + k] * Parameters.n_kernel_parameters[k]
                else:
                    Warning(
                        "inconsistency between input kernel and kernel list"
                        "stored in hyps_mask"
                    )

            Parameters.logger.debug("Replace kernel array in param_dict")
            param_dict["kernels"] = deepcopy(kernels)

        return param_dict

    @staticmethod
    def check_instantiation(hyps, cutoffs, kernels, param_dict):
        """
        Runs a series of checks to ensure that the user has not supplied
        contradictory arguments which will result in undefined behavior
        with multiple hyperparameters.

        :return:
        """

        assert isinstance(param_dict, dict)
        assert isinstance(cutoffs, dict)
        assert isinstance(kernels, list)

        param_dict["cutoffs"] = cutoffs

        # double check nspecie is there
        nspecie = param_dict["nspecie"]
        if nspecie > 1:
            assert (
                "species_mask" in param_dict
            ), "species_mask key missing in param_dict dictionary"
            param_dict["species_mask"] = nparray(param_dict["species_mask"], dtype=int)

        # for each kernel, check whether it is defined
        # and the length of corresponding hyper-parameters
        hyps_length = 0
        used_parameters = np.zeros_like(hyps, dtype=bool)
        for kernel in kernels + list(Parameters.cutoff_types.keys()):

            n = param_dict.get(f"n{kernel}", 0)
            assert isinstance(n, int)

            if kernel not in list(Parameters.cutoff_types.keys()):
                hyps_length += Parameters.n_kernel_parameters[kernel] * n
                assert n > 0, f"{kernel} has 0 hyperparameters defined"

                # check all corresponding keys exist
                assert kernel in cutoffs.keys()
                assert kernel + "_start" in param_dict

                # check the partition of hyperparameters are not used
                start = param_dict[kernel + "_start"]
                length = Parameters.n_kernel_parameters[kernel] * n
                assert not used_parameters[start : start + length].any()
                used_parameters[start : start + length] = True

            if n > 1:

                assert (
                    f"{kernel}_mask" in param_dict
                ), f"{kernel}_mask key missing in param_dict dictionary"

                # check mask has the right dimension and values
                mask = param_dict[f"{kernel}_mask"]
                param_dict[f"{kernel}_mask"] = nparray(mask, dtype=np.int)

                assert npmax(mask) < n
                dim = Parameters.ndim[kernel]
                assert len(mask) == nspecie ** dim, (
                    f"wrong dimension of {kernel}_mask: "
                    f" {len(mask)} != nspec ^ {dim} {nspecie**dim}"
                )

                # check whether the mask array is symmetrical
                # enumerate all possible combinations
                all_comb = list(combinations_with_replacement(np.arange(nspecie), dim))
                for comb in all_comb:
                    mask_value = None
                    perm = list(permutations(comb))
                    for ele_list in perm:
                        mask_id = 0
                        for ele in ele_list:
                            mask_id += ele
                            mask_id *= nspecie
                        mask_id = mask_id // nspecie
                        if mask_value == None:
                            mask_value = mask[mask_id]
                        else:
                            assert (
                                mask[mask_id] == mask_value
                            ), f"{kernel}_mask has to be symmetrical"

                if kernel not in list(Parameters.cutoff_types.keys()):
                    if kernel + "_cutoff_list" in param_dict:
                        cutoff_list = param_dict[kernel + "_cutoff_list"]
                        assert (
                            len(cutoff_list) == n
                        ), f"number of cutoffs should be the same as n {n}"
                        assert npmax(cutoff_list) <= cutoffs[kernel]
            else:
                assert (
                    f"{kernel}_mask" not in param_dict
                ), f"{kernel}_mask should not be in param_dict"
                assert (
                    f"{kernel}_cutoff_list" not in param_dict
                ), f"{kernel}_cutoff_list should not be in param_dict"

        if "map" in param_dict:

            assert (
                "original_hyps" in param_dict
            ), "original hyper parameters have to be defined"

            # Ensure typed correctly as numpy array
            param_dict["original_hyps"] = nparray(
                param_dict["original_hyps"], dtype=np.float
            )
            if (len(param_dict["original_hyps"]) - 1) not in param_dict["map"]:
                assert (
                    param_dict["train_noise"] is False
                ), "train_noise should be False when noise is not in hyps"

            assert len(param_dict["map"]) == len(
                hyps
            ), "the hyperparmeter length is inconsistent with the mask"
            assert npmax(param_dict["map"]) < len(param_dict["original_hyps"])

        else:
            assert (
                param_dict["train_noise"] is True
            ), "train_noise should be True when map is not used"

        hyps = Parameters.get_hyps(param_dict, hyps)

        hyps_length += 1
        assert hyps_length == len(
            hyps
        ), "the hyperparmeter length is inconsistent with the mask"

        return param_dict

    @staticmethod
    def get_component_hyps(
        param_dict, kernel_name, hyps=None, constraint=False, noise=False
    ):
        """
        return the hyper-parameters correspond to the kernel specified by kernel_name

        Args:

        param_dict (dict): the hyps_mask dictionary used/stored in GaussianProcess
        kernel_name (str): the name of the kernel.
        hyps (np.array): if hyps is None, use the one stored in param_dict
        constraint (bool): if True, return one additional list that shows whether the
                           hyper-parmaeters can be trained
        noise (bool): if True, the last element of returned hyper-parameters is
                      the noise variance.

        return: hyper-parameters, and whether they can be optimized
        """

        if kernel_name not in param_dict["kernels"]:
            if constraint:
                if noise:
                    return [None, None, None], [None, None]
                else:
                    return [None, None], [None, None]
            else:
                if noise:
                    return [None, None, None]
                else:
                    return [None, None]

        hyps, opt = Parameters.get_hyps(param_dict, hyps=hyps, constraint=True)
        s = param_dict[kernel_name + "_start"]
        n = param_dict[f"n{kernel_name}"]

        newhyps = [hyps[s : s + n], hyps[s + n : s + 2 * n]]
        newopt = [opt[s : s + n], opt[s + n : s + 2 * n]]

        if noise:
            newhyps += [hyps[-1]]

        if constraint:
            return newhyps, newopt
        else:
            return newhyps

    @staticmethod
    def get_component_mask(param_dict, kernel_name, hyps=None):
        """
        return the hyper-parameter masking correspond to the kernel specified by kernel_name

        Args:

        param_dict (dict): the hyps_mask dictionary used/stored in GaussianProcess
        kernel_name (str): the name of the kernel.
        hyps (np.array): if hyps is None, use the one stored in param_dict

        return: hyper-parameters, cutoffs, and new hyps_mask
        """

        if kernel_name in param_dict["kernels"]:
            new_dict = {}
            new_dict["kernels"] = [kernel_name]

            new_dict[kernel_name + "_start"] = 0

            name_list = [
                "nspecie",
                "species_mask",
                "n" + kernel_name,
                kernel_name + "_mask",
                kernel_name + "_cutoff_list",
            ]

            if kernel_name in Parameters.cutoff_types_values:

                key_ind = Parameters.cutoff_types_values.index(kernel_name)
                cutoff_key = Parameters.cutoff_types_keys[key_ind]
                name_list += ["n" + cutoff_key, cutoff_key + "_mask"]

            for name in name_list:
                if name in param_dict:
                    new_dict[name] = deepcopy(param_dict[name])

            hyps = np.hstack(
                Parameters.get_component_hyps(
                    param_dict, kernel_name, hyps=hyps, noise=True
                )
            )

            cutoffs = {}
            if "twobody" in param_dict["cutoffs"]:
                cutoffs["twobody"] = param_dict["cutoffs"]["twobody"]
            cutoffs[kernel_name] = param_dict["cutoffs"][kernel_name]

            return hyps, cutoffs, new_dict
        else:
            return [], {}, {}

    @staticmethod
    def get_noise(param_dict, hyps=None, constraint=False):
        """
        get the noise parameters

        Args:

        constraint (bool): if True, return one additional list that shows whether the
                           hyper-parmaeters can be trained
        noise (bool): if True, the last element of returned hyper-parameters is
                      the noise variance.
        """
        hyps = Parameters.get_hyps(param_dict, hyps=hyps)
        if constraint:
            return hyps[-1], param_dict["train_noise"]
        else:
            return hyps[-1]

    @staticmethod
    def get_cutoff(kernel_name, coded_species, param_dict):
        """
        get the cutoff

        Args:

        kernel_name (str): name of the kernel
        coded_species (list): list of element names
        param_dict (dict): hyps_mask

        """

        cutoffs = param_dict["cutoffs"]
        universal_cutoff = cutoffs[kernel_name]

        if f"{kernel_name}_cutoff_list" in param_dict:

            species_mask = param_dict["species_mask"]
            cutoff_list = param_dict[f"{kernel_name}_cutoff_list"]
            nspecie = param_dict["nspecie"]

            if kernel_name not in Parameters.cutoff_types_values:
                mask_id = 0
                for ele in coded_species:
                    mask_id += species_mask[ele]
                    mask_id *= nspecie
                mask_id = mask_id // nspecie
                mask_id = param_dict[kernel_name + "_mask"][mask_id]
                return cutoff_list[mask_id]
            else:

                key_ind = Parameters.cutoff_types_values.index(kernel_name)
                cutoff_key = Parameters.cutoff_types_keys[cutoff_key]

                cut_mask = param_dict[cutoff_key + "_mask"]
                ele1 = species_mask[coded_species[0]]
                ele2 = species_mask[coded_species[1]]
                ele3 = species_mask[coded_species[2]]
                twobody1 = cut_mask[param_dict["nspecie"] * ele1 + ele2]
                twobody2 = cut_mask[param_dict["nspecie"] * ele1 + ele3]
                twobody12 = cut_mask[param_dict["nspecie"] * ele2 + ele3]
                return np.array(
                    [
                        cutoff_list[twobody1],
                        cutoff_list[twobody2],
                        cutoff_list[twobody12],
                    ]
                )
        else:
            if kernel_name != "threebody":
                return [universal_cutoff]
            else:
                return [universal_cutoff] * 3

    @staticmethod
    def get_hyps(param_dict, hyps=None, constraint=False, label=False):
        """
        get the cutoff

        Args:

        kernel_name (str): name of the kernel
        coded_species (list): list of element names
        param_dict (dict): hyps_mask

        """

        if hyps is None:
            hyps = param_dict["hyps"]

        if "map" in param_dict:
            label_key = "original_labels"
            newhyps = np.copy(param_dict["original_hyps"])
            opt = np.zeros_like(newhyps, dtype=bool)
            for i, ori in enumerate(param_dict["map"]):
                newhyps[ori] = hyps[i]
                opt[ori] = True
        else:
            label_key = "hyp_labels"
            newhyps = np.copy(hyps)
            opt = np.zeros_like(hyps, dtype=bool)

        if constraint:
            if label:
                return newhyps, opt, param_dict.get(label_key, None)
            else:
                return newhyps, opt
        else:
            if label:
                return newhyps, param_dict.get(label_key, None)
            else:
                return newhyps

    @staticmethod
    def compare_dict(dict1, dict2):
        """
        compare whether two hyps_masks are the same
        """

        if type(dict1) != type(dict2):
            return False

        if dict1 is None:
            return True

        list_of_names = ["nspecie", "species_mask", "map", "original_hyps"]
        for k in Parameters.all_kernel_types:
            list_of_names += ["n" + k]
            list_of_names += [k + "_mask"]
            list_of_names += ["cutoff_" + k]
            list_of_names += [k + "_cutoff_list"]
        for k in Parameters.cutoff_types:
            list_of_names += ["n" + k]
            list_of_names += [k + "_mask"]

        for k in list_of_names:
            if (k in dict1) != (k in dict2):
                return False
            elif k in dict1:
                if not (np.isclose(dict1[k], dict2[k]).all()):
                    return False

        for k in ["hyp_labels", "original_labels"]:
            if (k in dict1) != (k in dict2):
                return False
            elif k in dict1:
                if not (dict1[k] == dict2[k]):
                    return False

        for k in ["train_noise"]:
            if (k in dict1) != (k in dict2):
                return False
            elif k in dict1:
                if dict1[k] != dict2[k]:
                    return False

        return True class MappedGaussianProcess:
    """
    Build Mapped Gaussian Process (MGP)
    and automatically save coefficients for LAMMPS pair style.

    Args:
        grid_params (dict): Parameters for the mapping itself, such as
            grid size of spline fit, etc. As described below.
        unique_species (dict): List of all the (unique) species included during
            the training that need to be mapped
        GP (GaussianProcess): None or a GaussianProcess object. If a GP is input,
            and container_only is False, automatically build a mapping corresponding
            to the GaussianProcess.
        var_map (str): if None: only build mapping for mean (force). If 'pca', then
            use PCA to map the variance, based on `grid_params['xxbody']['svd_rank']`.
            If 'simple', then only map the diagonal of covariance, and predict the
            upper bound of variance. The 'pca' mode is much heavier in terms of
            memory, but its prediction is much closer to GP variance.
        container_only (bool): if True: only build splines container
            (with no coefficients); if False: Attempt to build map immediately
        lmp_file_name (str): LAMMPS coefficient file name
        n_cpus (int): Default None. Set to the number of cores needed for
            parallelization. Used in the construction of the map.
        n_sample (int): Default 10. The batch size for building map. Not used now.

    Examples:

    >>> # build 2 + 3 body map
    >>> grid_params = {'twobody': {'grid_num': [64]},
    ...                'threebody': {'grid_num': [64, 64, 64]}}

    For `grid_params`, the following keys and values are allowed

    Args:
        'twobody' (dict, optional): if 2-body is present, set as a dictionary
            of parameters for 2-body mapping. Parameters see below.
        'threebody' (dict, optional): if 3-body is present, set as a dictionary
            of parameters for 3-body mapping. Parameters see below.
        'load_grid' (str, optional): Default None. the path to the directory
            where the previously generated grids (``grid_*.npy``) are stored.
            If no path is specified, MGP will construct grids from scratch.
        'lower_bound_relax' (float, optional): Default 0.1. if 'lower_bound' is
            set to 'auto' this value will be used as a relaxation of lower
            bound. (see below the description of 'lower_bound')

    For two/three body parameter dictionary, the following keys and values are allowed

    Args:
        'grid_num' (list): a list of integers, the number of grid points for
            interpolation. The larger the number, the better the approximation
            of MGP is compared with GP.
        'lower_bound' (str or list, optional): Default 'auto', the lower bound
            of the spline interpolation will be searched. First, search the
            training set of GP and find the minimal interatomic distance r_min.
            Then, the ``lower_bound = r_min - lower_bound_relax``. The user
            can set their own lower_bound, of the same shape as 'grid_num'.
            E.g. for threebody, the customized lower bound can be set as
            [1.2, 1.2, 1.2].
        'upper_bound' (str or list, optional): Default 'auto', the upper bound
            of the spline interpolation will be the cutoffs of GP. The user
            can set their own upper_bound, of the same shape as 'grid_num'.
            E.g. for threebody, the customized lower bound can be set as
            [3.5, 3.5, 3.5].
        'svd_rank' (int, optional): Default 'auto'. If the variance mapping is
            needed, it is set as the rank of the mapping. 'auto' uses full
            rank, which is the smaller one between the total number of grid
            points and training set size. i.e.
            ``full_rank = min(np.prod(grid_num), 3 * N_train)``
    """

    def __init__(
        self,
        grid_params: dict,
        unique_species: list = [],
        GP: GaussianProcess = None,
        var_map: str = None,
        container_only: bool = True,
        lmp_file_name: str = "lmp",
        n_cpus: int = None,
        n_sample: int = 10,
    ):

        # load all arguments as attributes
        self.var_map = var_map
        self.lmp_file_name = lmp_file_name
        self.n_cpus = n_cpus
        self.n_sample = n_sample
        self.grid_params = grid_params
        self.species_labels = []
        self.coded_species = []

        self.hyps_mask = None
        self.cutoffs = None
        self.training_statistics = None

        species_labels = []
        coded_species = []
        for i, ele in enumerate(unique_species):
            if isinstance(ele, str):
                species_labels.append(ele)
                coded_species.append(element_to_Z(ele))
            elif isinstance(ele, int):
                coded_species.append(ele)
                species_labels.append(Z_to_element(ele))
            else:
                print("element type not accepted", ele, type(ele))
        sort_id = np.argsort(coded_species)
        for i in sort_id:
            self.coded_species.append(coded_species[i])
            self.species_labels.append(species_labels[i])

        self.load_grid = grid_params.get("load_grid", None)
        self.update = grid_params.get("update", False)
        self.lower_bound_relax = grid_params.get("lower_bound_relax", 0.1)

        self.maps = {}

        optional_xb_params = ["lower_bound", "upper_bound", "svd_rank"]
        for key in grid_params:
            if "body" in key:
                if "twobody" == key:
                    mapxbody = Map2body
                elif "threebody" == key:
                    mapxbody = Map3body
                else:
                    raise KeyError("Only 'twobody' & 'threebody' are allowed")

                xb_dict = grid_params[key]

                # set to 'auto' if the param is not given
                args = {}
                for oxp in optional_xb_params:
                    args[oxp] = xb_dict.get(oxp, "auto")
                args["grid_num"] = xb_dict.get("grid_num", None)

                for k in xb_dict:
                    args[k] = xb_dict[k]

                xb_maps = mapxbody(**args, **self.__dict__)
                self.maps[key] = xb_maps

    def build_map(self, GP: GaussianProcess):
        self.hyps_mask = GP.hyps_mask
        self.cutoffs = GP.cutoffs
        self.training_statistics = GP.training_statistics

        for xb in self.maps:
            self.maps[xb].build_map(GP)

        # write to lammps pair style coefficient file
        self.write_lmp_file(self.lmp_file_name + ".mgp", write_var=False)
        if self.var_map == "simple":
            self.write_lmp_file(self.lmp_file_name + ".std", write_var=True)

    def predict(
        self, atom_env: AtomicEnvironment
    ) -> ("ndarray", "ndarray", "ndarray", float):
        """
        predict force, variance, stress and local energy for given
        atomic environment

        Args:
            atom_env: atomic environment (with a center atom and its neighbors)

        Return:
            force: 3d array of atomic force
            variance: 3d array of the predictive variance
            stress: 6d array of the virial stress
            energy: the local energy (atomic energy)
        """

        force = virial = kern = v = energy = 0
        rebuild_dict = {}
        newbound_dict = {}
        for xb in self.maps:
            try:
                pred = self.maps[xb].predict(atom_env)
                force += pred[0]
                virial += pred[1]
                kern += pred[2]
                v += pred[3]
                energy += pred[4]
            # record maps when test envs have smaller distance than lower bound
            except ValueError as err_msg:
                rebuild_dict[xb] = err_msg.args[0]
                newbound_dict[xb] = err_msg.args[1]

        if len(rebuild_dict) > 0:
            err_str = "\n"
            for xb in rebuild_dict:
                nb = len(rebuild_dict[xb])
                err_str += f"{nb} {xb} maps need re-construction\n"
            raise ValueError(rebuild_dict, newbound_dict, err_str)

        if self.var_map == "simple":
            variance = v ** 2
        else:
            variance = kern - np.sum(v ** 2, axis=0)

        return force, variance, virial, energy

    def write_lmp_file(self, lammps_name: str, write_var: bool = False):
        """
        write the coefficients to a file that can be used by lammps pair style
        """

        f = open(lammps_name, "w")

        # write header
        header_comment = """# #2bodyarray #3bodyarray\n# elem1 elem2 a b order\n\n"""
        f.write(header_comment)
        header = ""
        xbodies = ["twobody", "threebody"]
        for xb in xbodies:
            if xb in self.maps:
                num = self.maps[xb].num_lmp_maps  # len(self.maps[xb].maps)
            else:
                num = 0
            header += f"{num} "
        f.write(header + "\n")

        # write coefficients
        for xb in self.maps:
            self.maps[xb].write(f, write_var)

        f.close()

    def as_dict(self) -> dict:
        """
        Dictionary representation of the MGP model.
        """

        out_dict = deepcopy(dict(vars(self)))
        out_dict.pop("maps")

        # Uncertainty mappings currently not serializable;
        if self.var_map == "pca":
            warnings.warn(
                "Uncertainty mappings cannot be serialized, "
                "and so the MGP dict outputted will not have "
                "them.",
                Warning,
            )
            out_dict["var_map"] = None

        # only save the coefficients
        maps_dict = {}
        for m in self.maps:
            maps_dict[m] = self.maps[m].as_dict()
        out_dict["maps"] = maps_dict

        return out_dict

    @staticmethod
    def from_dict(dictionary: dict) -> "MappedGaussianProcess":
        """
        Create MGP object from dictionary representation.
        """

        # Set GP
        if dictionary.get("GP"):
            GP = GaussianProcess.from_dict(dictionary.get("GP"))
        else:
            dictionary["GP"] = None

        dictionary["unique_species"] = list(set(dictionary["species_labels"]))
        if "container_only" not in dictionary:
            dictionary["container_only"] = True

        init_arg_name = [
            "grid_params",
            "unique_species",
            "GP",
            "var_map",
            "container_only",
            "lmp_file_name",
            "n_cpus",
            "n_sample",
        ]
        kwargs = {key: dictionary[key] for key in init_arg_name}
        new_mgp = MappedGaussianProcess(**kwargs)

        # Fill up the model with the saved coeffs
        if "twobody" in new_mgp.maps:
            new_mgp.maps["twobody"] = Map2body.from_dict(
                dictionary["maps"]["twobody"], Map2body
            )
        if "threebody" in new_mgp.maps:
            new_mgp.maps["threebody"] = Map3body.from_dict(
                dictionary["maps"]["threebody"], Map3body
            )

        return new_mgp

    def write_model(self, name: str, format: str = "json"):
        """
        Write everything necessary to re-load and re-use the model
        :param model_name:
        :return:
        """
        if "json" in format.lower() or "json" in name:
            with open(f"{name}.json", "w") as f:
                json.dump(self.as_dict(), f, cls=NumpyEncoder)

        elif (
            "pickle" in format.lower()
            or "binary" in format.lower()
            or "pickle" in name
            or "binary" in name
        ):
            with open(f"{name}.pickle", "wb") as f:
                pickle.dump(self, f)

        else:
            raise ValueError("Requested format not found.")

    @staticmethod
    def from_file(filename: str, format: str = ""):
        if ".json" in filename or format.lower() == "json":
            with open(filename, "r") as f:
                model = MappedGaussianProcess.from_dict(json.loads(f.readline()))
            return model

        elif "pickle" in filename or format.lower() in ["binary", "pickle"]:
            with open(filename, "rb") as f:
                return pickle.load(f)
        else:
            raise NotImplementedError

    def __len__(self):
        return self.training_statistics["N"] def run_lammps(lammps_executable, input_file, output_file):
def lammps_parser(dump_file, std=False):
def lammps_dat(structure, atom_types, atom_masses, species):
def lammps_dat_charged(structure, atom_types, atom_charges, atom_masses, species):
def lammps_cell_text(structure):
def lammps_pos_text(structure, species):
def lammps_pos_text_charged(structure, charges, species):
def write_text(file, text):
def generic_lammps_input(
    dat_file,
    style_string,
    coeff_string,
    dump_file,
    newton=False,
    std_string="",
    std_style=None,
):
def ewald_input(dat_file, short_cut, kspace_accuracy, dump_file, newton=True): def inject_user_definition(element: str, Z: int):
def element_to_Z(element: str) -> int:
    def default(self, obj):
def Z_to_element(Z: int) -> str:
class NumpyEncoder(JSONEncoder):
            Z = int(Z) class FLARE_Calculator(Calculator):
    """
    Build FLARE as an ASE Calculator, which is compatible with ASE Atoms and
    Molecular Dynamics.
    Args:
        gp_model (GaussianProcess): FLARE's Gaussian process object
        mgp_model (MappedGaussianProcess): FLARE's Mapped Gaussian Process
            object. `None` by default. MGP will only be used if `use_mapping`
            is set to True.
        par (Bool): set to `True` if parallelize the prediction. `False` by
            default.
        use_mapping (Bool): set to `True` if use MGP for prediction. `False`
            by default.
    """

    implemented_properties = ["energy", "forces", "stress", "stds"]

    def __init__(
        self, gp_model, mgp_model=None, par=False, use_mapping=False, **kwargs
    ):
        super().__init__()  # all set to default values, TODO: change
        self.mgp_model = mgp_model
        self.gp_model = gp_model
        self.use_mapping = use_mapping
        self.par = par
        self.results = {}

    def get_uncertainties(self, atoms):
        return self.get_property("stds", atoms)

    def calculate(self, atoms=None, properties=None, system_changes=all_changes):
        """
        Calculate properties including: energy, local energies, forces,
            stress, uncertainties.

        Args:
            atoms (FLARE_Atoms): FLARE_Atoms object
        """

        super().calculate(
            atoms=atoms, properties=properties, system_changes=system_changes
        )

        if properties is None:
            properties = self.implemented_properties

        if self.use_mapping:
            self.calculate_mgp(atoms)
        else:
            self.calculate_gp(atoms)

        # get global properties
        volume = atoms.get_volume()
        total_stress = np.sum(self.results["partial_stresses"], axis=0)
        self.results["stress"] = total_stress / volume
        self.results["energy"] = np.sum(self.results["local_energies"])

    def calculate_gp(self, atoms):
        # Compute energy, forces, and stresses and their uncertainties
        if self.par:
            res = predict_on_structure_efs_par(
                atoms, self.gp_model, write_to_structure=False
            )
        else:
            res = predict_on_structure_efs(
                atoms, self.gp_model, write_to_structure=False
            )

        # Set the energy, force, and stress attributes of the calculator.
        self.results["local_energies"] = res[0]
        self.results["forces"] = res[1]
        self.results["partial_stresses"] = -res[2][:, [0, 3, 5, 4, 2, 1]]
        self.results["local_energy_stds"] = res[3]
        self.results["stds"] = res[4]
        self.results["partial_stress_stds"] = res[5][:, [0, 3, 5, 4, 2, 1]]

    def calculate_mgp(self, atoms):
        nat = len(atoms)

        self.results["forces"] = np.zeros((nat, 3))
        self.results["partial_stresses"] = np.zeros((nat, 6))
        self.results["stds"] = np.zeros((nat, 3))
        self.results["local_energies"] = np.zeros(nat)

        rebuild_dict = {}
        newbound_dict = {}
        repredict_atoms = []
        for n in range(nat):
            chemenv = AtomicEnvironment(
                atoms, n, self.gp_model.cutoffs, cutoffs_mask=self.mgp_model.hyps_mask
            )

            # TODO: Check that stress is being calculated correctly.
            try:
                f, v, vir, e = self.mgp_model.predict(chemenv)
                self.results["forces"][n] = f
                self.results["partial_stresses"][n] = vir
                self.results["stds"][n][0] = np.sqrt(np.absolute(v))
                self.results["local_energies"][n] = e

            except ValueError as err_msg:  # if lower_bound error is raised
                get_rebuild_from_err(err_msg, rebuild_dict, newbound_dict)
                repredict_atoms.append((n, chemenv))

        if len(rebuild_dict) > 0:
            # rebuild map for those problematic species
            for xb in rebuild_dict:
                for s_ind, spc in enumerate(rebuild_dict[xb]):
                    map_ind = self.mgp_model.maps[xb].find_map_index(spc)
                    rebuild_map = self.mgp_model.maps[xb].maps[map_ind]
                    rebuild_map.set_bounds(
                        newbound_dict[xb][s_ind], rebuild_map.bounds[1]
                    )
                    rebuild_map.build_map_container()
                    rebuild_map.build_map(self.gp_model)

            # re-predict forces, energies, etc. for those problematic atoms
            for ra in repredict_atoms:
                n, chemenv = ra
                f, v, vir, e = self.mgp_model.predict(chemenv)
                self.results["forces"][n] = f
                self.results["partial_stresses"][n] = -vir[[0, 3, 5, 4, 2, 1]]
                self.results["stds"][n][0] = np.sqrt(np.absolute(v))
                self.results["local_energies"][n] = e

    def calculation_required(self, atoms, quantities):
        return True

    def as_dict(self):
        outdict = {}

        outdict["class"] = self.__class__.__name__

        gp_dict = self.gp_model.as_dict()
        outdict["gp_model"] = gp_dict

        outdict["use_mapping"] = self.use_mapping
        if self.use_mapping:
            mgp_dict = self.mgp_model.as_dict()
            outdict["mgp_model"] = mgp_dict
        else:
            outdict["mgp_model"] = None

        outdict["par"] = self.par
        outdict["results"] = self.results
        return outdict

    @staticmethod
    def from_dict(dct):
        dct["gp_model"] = GaussianProcess.from_dict(dct["gp_model"])
        if dct["use_mapping"]:
            dct["mgp_model"] = MappedGaussianProcess.from_dict(dct["mgp_model"])

        calc = FLARE_Calculator(**dct)
        res = dct["results"]
        for key in res:
            if isinstance(res[key], float):
                calc.results[key] = res[key]
            if isinstance(res[key], list):
                calc.results[key] = np.array(res[key])

        if dct["use_mapping"]:
            for xb in calc.mgp_model.maps:
                xb_map = calc.mgp_model.maps[xb]
                xb_map.hyps_mask = calc.gp_model.hyps_mask

        return calc

    def write_model(self, name):
        if ".json" != name[-5:]:
            name += ".json"
        with open(name, "w") as f:
            json.dump(self.as_dict(), f, cls=NumpyEncoder)

    @staticmethod
    def from_file(name):
        with open(name, "r") as f:
            calc = FLARE_Calculator.from_dict(json.loads(f.readline()))

        return calc

    def build_map(self):
        self.mgp_model.build_map(self.gp_model) class FLARE_Atoms(Atoms):
    """
    The `FLARE_Atoms` class is a child class of ASE `Atoms`,
    which has completely the same usage as the primitive ASE `Atoms`, and
    in the meanwhile mimic `Structure` class. It is used in the `OTF` module
    with ASE engine (by `OTF_ASE` module). It enables attributes to be
    obtained by both the name from ASE `Atoms` and `Structure`.

    The input arguments are the same as ASE `Atoms`.
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.prev_positions = np.zeros_like(self.positions)

    @staticmethod
    def from_ase_atoms(atoms):
        """
        Args:
            atoms (ASE Atoms): the ase atoms to build from
        """
        new_atoms = deepcopy(atoms)
        new_atoms.__class__ = FLARE_Atoms
        new_atoms.prev_positions = np.zeros_like(new_atoms.positions)
        return new_atoms

    @property
    def nat(self):
        return len(self)

    @property
    def species_labels(self):
        return self.symbols

    @property
    def coded_species(self):
        return self.numbers

    @property
    def forces(self):
        return self.get_forces()

    @forces.setter
    def forces(self, forces_array):
        pass

    @property
    def potential_energy(self):
        return self.get_potential_energy()

    @property
    def stress(self):
        return self.get_stress()

    @property
    def stress_stds(self):
        return None  # TODO: to implement

    @property
    def local_energy_stds(self):
        return None  # TODO: to implement

    @property
    def stds(self):
        try:  # when self.calc is not FLARE, there's no get_uncertainties()
            stds = self.calc.results["stds"]
        except:
            stds = np.zeros_like(self.positions)
        return stds

    def wrap_positions(self):
        return self.get_positions(wrap=True)

    @property
    def wrapped_positions(self):
        return self.get_positions(wrap=True)

    @property
    def max_cutoff(self):
        return get_max_cutoff(self.cell)

    def as_dict(self):
        return self.todict()

    @staticmethod
    def from_dict(dct):
        atoms = Atoms.fromdict(dct)
        return FLARE_Atoms.from_ase_atoms(atoms) def get_gp(
    bodies,
    kernel_type="mc",
    multihyps=True,
    cellabc=[1, 1, 1.5],
    force_only=False,
    noa=5,
) -> GaussianProcess:
    """Returns a GP instance with a two-body numba-based kernel"""
    print("\nSetting up...\n")

    # params
    cell = np.diag(cellabc)
    unique_species = [1, 2]

    ntwobody = 0
    nthreebody = 0
    prefix = bodies
    if "2" in bodies or "two" in bodies:
        ntwobody = 1
    if "3" in bodies or "three" in bodies:
        nthreebody = 1

    hyps, hm, _ = generate_hm(ntwobody, nthreebody, nmanybody=0, multihyps=multihyps)
    cutoffs = hm["cutoffs"]
    kernels = hm["kernels"]
    hl = hm["hyp_labels"]

    # create test structure
    test_structure, forces = get_random_structure(cell, unique_species, noa)
    energy = 3.14

    # test update_db
    gaussian = GaussianProcess(
        kernels=kernels,
        component=kernel_type,
        hyps=hyps,
        hyp_labels=hl,
        cutoffs=cutoffs,
        hyps_mask=hm,
        parallel=False,
        n_cpus=1,
    )

    if force_only:
        gaussian.update_db(test_structure, forces)
    else:
        gaussian.update_db(test_structure, forces, energy=energy)
    gaussian.check_L_alpha()

    # print(gaussian.alpha)

    return gaussian def get_random_structure(cell, unique_species, noa, set_seed: int = None):
    """Create a random test structure"""
    if set_seed:
        np.random.seed(set_seed)

    forces = (np.random.random([noa, 3]) - 0.5) * 2
    positions = np.random.random([noa, 3])
    species = [
        unique_species[np.random.randint(0, len(unique_species))] for i in range(noa)
    ]

    test_structure = Structure(cell, species, positions)

    return test_structure, forces def clean(prefix="tmp"):
    for f in os.listdir("./"):
        if re.search("mgp_grids", f):
            shutil.rmtree(f)
        if re.search("kv3", f):
            os.rmdir(f)
        if prefix in f:
            if os.path.isdir(f):
                shutil.rmtree(f)
            else:
                os.remove(f) def compare_triplet(mgp_model, gp_model, atom_env):
    spcs, comp_r, comp_xyz = mgp_model.get_arrays(atom_env)
    for i, spc in enumerate(spcs):
        lengths = np.array(comp_r[i])
        xyzs = np.array(comp_xyz[i])

        print("compare triplet spc, lengths, xyz", spc)

        gp_f = []
        gp_e = []
        grid_env = get_grid_env(gp_model, spc, 3)
        for l in range(lengths.shape[0]):
            r1, r2, r12 = lengths[l, :]
            grid_env = get_triplet_env(r1, r2, r12, grid_env)
            gp_pred = np.array([gp_model.predict(grid_env, d + 1) for d in range(3)]).T
            gp_en, _ = gp_model.predict_local_energy_and_var(grid_env)
            gp_f.append(gp_pred[0])
            gp_e.append(gp_en)
        gp_force = np.sum(gp_f, axis=0)
        gp_energy = np.sum(gp_e, axis=0)
        print("gp_e", gp_e)
        print("gp_f")
        print(gp_f)

        map_ind = mgp_model.find_map_index(spc)
        xyzs = np.zeros_like(xyzs)
        xyzs[:, 0] = np.ones_like(xyzs[:, 0])
        f, _, _, e = mgp_model.maps[map_ind].predict(lengths, xyzs)

        assert np.allclose(gp_force, f, rtol=1e-2)
        assert np.allclose(gp_energy, e, rtol=1e-2) def predict_atom_diag_var(atom_env, gp_model, kernel_name, force_kernel=False):
    print("predict diag var")
    if kernel_name == "twobody":
        return predict_atom_diag_var_2b(atom_env, gp_model, force_kernel)
    elif kernel_name == "threebody":
        return predict_atom_diag_var_3b(atom_env, gp_model, force_kernel)
import numpy as np
import os
import pickle
import pytest
import re
import time
import shutil
from copy import deepcopy
from numpy import allclose, isclose
from flare import struc, env, gp
from flare.parameters import Parameters
from flare.mgp import MappedGaussianProcess
from flare.lammps import lammps_calculator
from flare.utils.element_coder import _Z_to_mass, _Z_to_element, _element_to_Z
from flare.ase.calculator import FLARE_Calculator
from flare.ase.atoms import FLARE_Atoms
from ase.calculators.lammpsrun import LAMMPS
from .fake_gp import get_gp, get_random_structure
from .mgp_test import clean, compare_triplet, predict_atom_diag_var
        struc_test, 0, cutoffs, cutoffs_mask=gp_model.hyps_mask
    )

    if "2" in bodies:
        kernel_name = "twobody"
    elif "3" in bodies:
        kernel_name = "threebody"
        # compare_triplet(mgp_model.maps['threebody'], gp_model, test_envi)

    mgp_f, mgp_e_var, mgp_s, mgp_e = mgp_model.predict(test_envi)

    assert Parameters.compare_dict(
        gp_model.hyps_mask, mgp_model.maps[kernel_name].hyps_mask
    )

    if multihyps:
        gp_e, gp_e_var = gp_model.predict_local_energy_and_var(test_envi)
        gp_f, gp_f_var = gp_model.predict_force_xyz(test_envi)
    else:
        gp_e, gp_f, gp_s, gp_e_var, _, _ = gp_model.predict_efs(test_envi)
        gp_s = -gp_s[[0, 3, 5, 4, 2, 1]]

        # check stress
        assert np.allclose(mgp_s, gp_s, rtol=1e-2)

    # check mgp is within 2 meV/A of the gp
    print("mgp_en, gp_en", mgp_e, gp_e)
    assert np.allclose(mgp_e, gp_e, rtol=2e-3), (
        f"{bodies} body" f" energy mapping is wrong"
    )

    # check forces
    print("isclose?", mgp_f - gp_f, gp_f)
    assert np.allclose(mgp_f, gp_f, atol=1e-3), f"{bodies} body force mapping is wrong"

    if mgp_model.var_map == "simple":
        print(bodies, multihyps)
        for i in range(struc_test.nat):
            test_envi = env.AtomicEnvironment(
                struc_test, i, cutoffs, cutoffs_mask=gp_model.hyps_mask
            )
            mgp_pred = mgp_model.predict(test_envi)
            mgp_var = mgp_pred[1]
            gp_var = predict_atom_diag_var(test_envi, gp_model, kernel_name)
            print("mgp_var, gp_var", mgp_var, gp_var)
            assert np.allclose(mgp_var, gp_var, rtol=1e-2)

    print("struc_test positions", struc_test.positions, struc_test.species_labels)


@pytest.mark.skipif(
    not os.environ.get("lmp", False),
    reason=(
        "lmp not found "
        "in environment: Please install LAMMPS "
        "and set the $lmp env. "
        "variable to point to the executatble."
    ),
)
@pytest.mark.parametrize("bodies", body_list)
@pytest.mark.parametrize("multihyps", multi_list)
def test_lmp_predict(all_lmp, all_gp, all_mgp, bodies, multihyps):
    """
    test the lammps implementation
    """

    # pytest.skip()

    prefix = f"{bodies}{multihyps}"

    mgp_model = all_mgp[prefix]
    gp_model = all_gp[prefix]
    lmp_calculator = all_lmp[prefix]
    ase_calculator = FLARE_Calculator(gp_model, mgp_model, par=False, use_mapping=True)

    # create test structure
    np.random.seed(1)
    cell = np.diag(np.array([1, 1, 1])) * 4
    nenv = 10
    unique_species = gp_model.training_statistics["species"]
    cutoffs = gp_model.cutoffs
    struc_test, f = get_random_structure(cell, unique_species, nenv)

    # build ase atom from struc
    ase_atoms_flare = struc_test.to_ase_atoms()
    ase_atoms_flare = FLARE_Atoms.from_ase_atoms(ase_atoms_flare)
    ase_atoms_flare.calc = ase_calculator

    ase_atoms_lmp = deepcopy(struc_test).to_ase_atoms()
    ase_atoms_lmp.calc = lmp_calculator

    try:
        lmp_en = ase_atoms_lmp.get_potential_energy()
        flare_en = ase_atoms_flare.get_potential_energy()

        lmp_stress = ase_atoms_lmp.get_stress()
        flare_stress = ase_atoms_flare.get_stress()

        lmp_forces = ase_atoms_lmp.get_forces()
        flare_forces = ase_atoms_flare.get_forces()
    except Exception as e:
        os.chdir(curr_path)
        print(e)
        raise e

    os.chdir(curr_path)

    # check that lammps agrees with mgp to within 1 meV/A
    print("energy", lmp_en - flare_en, flare_en)
    assert np.isclose(lmp_en, flare_en, atol=1e-3)
    print("force", lmp_forces - flare_forces, flare_forces)
    assert np.isclose(lmp_forces, flare_forces, atol=1e-3).all()
    print("stress", lmp_stress - flare_stress, flare_stress)
    assert np.isclose(lmp_stress, flare_stress, atol=1e-3).all()

    # check the lmp var
    # mgp_std = np.sqrt(mgp_pred[1])
    # print("isclose? diff:", lammps_stds[atom_num]-mgp_std, "mgp value", mgp_std)
    # assert np.isclose(lammps_stds[atom_num], mgp_std, rtol=1e-2)


    clean(prefix=prefix)