class AtomicEnvironment:
    """Contains information about the local environment of an atom,
    including arrays of pair and triplet distances and the chemical
    species of atoms in the environment.

    :param structure: Structure of atoms.
    :type structure: struc.Structure
    :param atom: Index of the atom in the structure.
    :type atom: int
    :param cutoffs: 2- and 3-body cutoff radii. 2-body if one cutoff is
    given, 2+3-body if two are passed.
    :type cutoffs: np.ndarray
    :param cutoffs_mask: a dictionary to store multiple cutoffs if neede
                         it should be exactly the same as the hyps mask
    :type cutoffs_mask: dict

    The cutoffs_mask allows the user to define multiple cutoffs for different
    bonds, triples, and many body interaction. This dictionary should be
    consistent with the hyps_mask used in the GuassianProcess object.

    * species_mask: 118-long integer array descirbing which elements belong to
        like groups for determining which bond hyperparameters to use.
        For instance, [0,0,1,1,0 ...] assigns H to group 0, He and
        Li to group 1, and Be to group 0 (the 0th register is ignored).
    * nspecie: Integer, number of different species groups (equal to number of
        unique values in species_mask).
    * ntwobody: Integer, number of different hyperparameter/cutoff sets to
        associate with different 2-body pairings of atoms in groups defined in
        species_mask.
    * twobody_mask: Array of length nspecie^2, which describes the cutoff to
        associate with different pairings of species types. For example, if
        there are atoms of type 0 and 1, then twobody_mask defines which cutoff
        to use for parings [0-0, 0-1, 1-0, 1-1]: if we wanted cutoff0 for
        0-0 parings and set 1 for 0-1 and 1-1 pairings, then we would make
        twobody_mask [0, 1, 1, 1].
    * twobody_cutoff_list: Array of length ntwobody, which stores the cutoff
        used for different types of bonds defined in twobody_mask
    * ncut3b: Integer, number of different cutoffs sets to associate
        with different 3-body pariings of atoms in groups defined in
        species_mask.
    * cut3b_mask: Array of length nspecie^2, which describes the cutoff to
        associate with different bond types in triplets. For example, in a
        triplet (C, O, H) , there are three cutoffs. Cutoffs for CH bond, CO
        bond and OH bond. If C and O are associate with atom group 1 in
        species_mask and H are associate with group 0 in species_mask, the
        cut3b_mask[1*nspecie+0] determines the C/O-H bond cutoff, and
        cut3b_mask[1*nspecie+1] determines the C-O bond cutoff. If we want the
        former one to use the 1st cutoff in threebody_cutoff_list and the later
        to use the 2nd cutoff in threebody_cutoff_list, the cut3b_mask should
        be [0, 0, 0, 1].
    * threebody_cutoff_list: Array of length ncut3b, which stores the cutoff
        used for different types of bonds in triplets.
    * nmanybody: Integer, number of different cutoffs set to associate with
        different coordination numbers.
    * manybody_mask: Similar to twobody_mask and cut3b_mask.
    * manybody_cutoff_list: Array of length nmanybody, stores the cutoff used
        for different many body terms

    Examples can be found at the end of in tests/test_env.py

    """

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

    def __init__(self, structure: Structure, atom: int, cutoffs, cutoffs_mask=None):

        self.structure = structure
        self.positions = structure.wrapped_positions
        self.cell = np.array(structure.cell)
        self.species = structure.coded_species

        # backward compatability
        if not isinstance(cutoffs, dict):
            newcutoffs = {"twobody": cutoffs[0]}
            if len(cutoffs) > 1:
                newcutoffs["threebody"] = cutoffs[1]
            if len(cutoffs) > 2:
                newcutoffs["manybody"] = cutoffs[2]
            cutoffs = newcutoffs

        if cutoffs_mask is None:
            cutoffs_mask = {"cutoffs": cutoffs}
        elif cutoffs is not None:
            cutoffs_mask["cutoffs"] = deepcopy(cutoffs)

        # Set the sweep array based on the max cutoff.
        sweep_val = ceil(np.max(list(cutoffs.values())) / structure.max_cutoff)
        self.sweep_val = sweep_val
        self.sweep_array = np.arange(-sweep_val, sweep_val + 1, 1)

        self.atom = atom
        self.ctype = structure.coded_species[atom]

        self.twobody_cutoff = 0
        self.threebody_cutoff = 0
        self.manybody_cutoff = 0

        self.ntwobody = 1
        self.ncut3b = 0
        self.nmanybody = 0

        self.nspecie = 1
        self.species_mask = None
        self.twobody_mask = None
        self.threebody_mask = None
        self.manybody_mask = None
        self.twobody_cutoff_list = None
        self.threebody_cutoff_list = None
        self.manybody_cutoff_list = None

        self.setup_mask(cutoffs_mask)

        assert (
            self.threebody_cutoff <= self.twobody_cutoff
        ), "2b cutoff has to be larger than 3b cutoff"
        # # TO DO, once the mb function is updated to use the bond_array_2
        # # this block should be activated.
        # assert self.manybody_cutoff <= self.twobody_cutoff, \
        #         "mb cutoff has to be larger than mb cutoff"

        self.bond_array_2 = None
        self.etypes = None
        self.bond_inds = None
        self.bond_array_3 = None
        self.cross_bond_inds = None
        self.cross_bond_dists = None
        self.triplet_counts = None
        self.q_array = None
        self.q_neigh_array = None
        self.q_grads = None
        self.q_neigh_grads = None
        self.unique_species = None
        self.etypes_mb = None

        self.compute_env()

    def setup_mask(self, cutoffs_mask):

        self.cutoffs_mask = deepcopy(cutoffs_mask)
        self.cutoffs = cutoffs_mask["cutoffs"]

        for kernel in AtomicEnvironment.all_kernel_types:
            ndim = AtomicEnvironment.ndim[kernel]
            if kernel in self.cutoffs:
                setattr(self, kernel + "_cutoff", self.cutoffs[kernel])

        if self.twobody_cutoff == 0:
            self.twobody_cutoff = np.max([self.threebody_cutoff, self.manybody_cutoff])

            self.cutoffs["twobody"] = self.twobody_cutoff

        self.nspecie = cutoffs_mask.get("nspecie", 1)
        if "species_mask" in cutoffs_mask:
            self.species_mask = np.array(cutoffs_mask["species_mask"], dtype=int)

        for kernel in AtomicEnvironment.all_kernel_types:
            ndim = AtomicEnvironment.ndim[kernel]
            if kernel in self.cutoffs:
                setattr(self, kernel + "_cutoff", self.cutoffs[kernel])
                setattr(self, "n" + kernel, 1)
                if kernel != "threebody":
                    name_list = [
                        kernel + "_cutoff_list",
                        "n" + kernel,
                        kernel + "_mask",
                    ]
                    for name in name_list:
                        if name in cutoffs_mask:
                            setattr(self, name, cutoffs_mask[name])
                else:
                    self.ncut3b = cutoffs_mask.get("ncut3b", 1)
                    self.cut3b_mask = cutoffs_mask.get("cut3b_mask", None)
                    if "threebody_cutoff_list" in cutoffs_mask:
                        self.threebody_cutoff_list = np.array(
                            cutoffs_mask["threebody_cutoff_list"], dtype=np.float
                        )

    def compute_env(self):

        # get 2-body arrays
        if self.ntwobody >= 1:
            bond_array_2, bond_positions_2, etypes, bond_inds = get_2_body_arrays(
                self.positions,
                self.atom,
                self.cell,
                self.twobody_cutoff,
                self.twobody_cutoff_list,
                self.species,
                self.sweep_array,
                self.nspecie,
                self.species_mask,
                self.twobody_mask,
            )

            self.bond_array_2 = bond_array_2
            self.etypes = etypes
            self.bond_inds = bond_inds

        # if 2 cutoffs are given, create 3-body arrays
        if self.ncut3b > 0:
            (
                bond_array_3,
                cross_bond_inds,
                cross_bond_dists,
                triplet_counts,
            ) = get_3_body_arrays(
                bond_array_2,
                bond_positions_2,
                self.species[self.atom],
                etypes,
                self.threebody_cutoff,
                self.threebody_cutoff_list,
                self.nspecie,
                self.species_mask,
                self.cut3b_mask,
            )

            self.bond_array_3 = bond_array_3
            self.cross_bond_inds = cross_bond_inds
            self.cross_bond_dists = cross_bond_dists
            self.triplet_counts = triplet_counts

        # if 3 cutoffs are given, create many-body arrays
        if self.nmanybody > 0:
            (
                self.q_array,
                self.q_neigh_array,
                self.q_grads,
                self.q_neigh_grads,
                self.unique_species,
                self.etypes_mb,
            ) = get_m2_body_arrays(
                self.positions,
                self.atom,
                self.cell,
                self.manybody_cutoff,
                self.manybody_cutoff_list,
                self.species,
                self.sweep_array,
                self.nspecie,
                self.species_mask,
                self.manybody_mask,
                cf.quadratic_cutoff,
            )

    def as_str(self) -> str:
        """
        Returns string dictionary serialization cast as string.

        :return: output of as_dict method cast as string
        :rtype: str
        """
        return dumps(self.as_dict(), cls=NumpyEncoder)

    def as_dict(self, include_structure: bool = False):
        """
        Returns Atomic Environment object as a dictionary for serialization
        purposes. Optional to not include the structure to avoid redundant
        information.
        :return:
        """

        dictionary = dict(vars(self))
        dictionary["object"] = "AtomicEnvironment"
        dictionary["forces"] = self.structure.forces

        # Backward compatibility for older models: Cutoffs mask.
        # Can be deleted one day if support is dropped for older (Pre June
        # 2020) pickled environment objects.
        cutoffs_mask = getattr(self, "cutoffs_mask", {"cutoffs": self.cutoffs})
        if not hasattr(self, "cutoffs_mask"):
            self.cutoffs_mask = cutoffs_mask
        dictionary["cutoffs_mask"] = cutoffs_mask

        if not include_structure:
            del dictionary["structure"]
        else:
            dictionary["structure"] = dictionary["structure"].as_dict()

        return dictionary

    @staticmethod
    def from_dict(dictionary):
        """
        Loads in atomic environment object from a dictionary which was
        serialized by the to_dict method.

        :param dictionary: Dictionary describing atomic environment.
        """
        # TODO Instead of re-computing 2 and 3 body environment,
        # directly load in, this would be much more efficient

        struc = Structure(
            cell=np.array(dictionary["cell"]),
            positions=dictionary["positions"],
            species=dictionary["species"],
        )
        index = dictionary["atom"]

        if dictionary.get("cutoffs") is not None:
            cutoffs = dictionary["cutoffs"]
        else:
            cutoffs = {}

        cutoffs_mask = dictionary.get("cutoffs_mask", None)

        return AtomicEnvironment(struc, index, cutoffs, cutoffs_mask=cutoffs_mask)

    @staticmethod
    def from_file(
        file_name: str, format: str = ""
    ) -> Union["flare.env.AtomicEnvironment", List["flare.env.AtomicEnvironment"]]:
        """
        Load an atomic environment from a file or a series of atomic environments
        :param file_name:
        :param format:
        :return:
        """

        # Ensure the file specified exists.
        with open(file_name, "r") as _:
            pass

        if "pickle" in file_name or (
            "pickle" in format.lower() or "binary" in format.lower()
        ):
            with open(file_name, "rb") as f:
                return pickle.load(f)

        if "json" in format.lower() or ".json" in file_name:
            # Assumed format is one atomic environment per line,
            # or one line with many atomic environments
            with open(file_name, "r") as f:
                thelines = f.readlines()

                non_empty_lines = [loads(line) for line in thelines if len(line) > 2]

            envs = [
                AtomicEnvironment.from_dict(env_dict) for env_dict in non_empty_lines
            ]

            if len(envs) == 1:
                return envs[0]
            else:
                return envs

    def __str__(self):
        atom_type = self.ctype
        neighbor_types = self.etypes
        n_neighbors = len(self.bond_array_2)
        string = "Atomic Env. of Type {} surrounded by {} atoms of Types {}".format(
            atom_type, n_neighbors, sorted(list(set(neighbor_types)))
        )

        return string

    @property
    def force(self) -> "np.ndarray":
        return self.structure.forces[self.atom] class GaussianProcess:
    """Gaussian process force field. Implementation is based on Algorithm 2.1
    (pg. 19) of "Gaussian Processes for Machine Learning" by Rasmussen and
    Williams.

    Methods within GaussianProcess allow you to make predictions on
    AtomicEnvironment objects (see env.py) generated from
    FLARE Structures (see struc.py), and after data points are added,
    optimize hyperparameters based on available training data (train method).

    Args:
        kernels (list, optional): Determine the type of kernels. Example:
            ['twbody', 'threebody'], ['2', '3', 'mb'], ['2']. Defaults to [
            'twboody', 'threebody']
        component (str, optional): Determine single- ("sc") or multi-
            component ("mc") kernel to use. Defaults to "mc"
        hyps (np.ndarray, optional): Hyperparameters of the GP.
        cutoffs (Dict, optional): Cutoffs of the GP kernel. For simple hyper-
            parameter setups, formatted like {"twobody":7, "threebody":4.5},
            etc.
        hyp_labels (List, optional): List of hyperparameter labels. Defaults
            to None.
        opt_algorithm (str, optional): Hyperparameter optimization algorithm.
            Defaults to 'L-BFGS-B'.
        maxiter (int, optional): Maximum number of iterations of the
            hyperparameter optimization algorithm. Defaults to 10.
        parallel (bool, optional): If True, the covariance matrix K of the GP is
            computed in parallel. Defaults to False.
        n_cpus (int, optional): Number of cpus used for parallel
            calculations. Defaults to 1 (serial)
        n_sample (int, optional): Size of submatrix to use when parallelizing
            predictions.
        output (Output, optional): Output object used to dump hyperparameters
            during optimization. Defaults to None.
        hyps_mask (dict, optional): hyps_mask can set up which hyper parameter
            is used for what interaction. Details see kernels/mc_sephyps.py
        name (str, optional): Name for the GP instance which dictates global
            memory access.
    """

    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,
    ):
        """Initialize GP parameters and training data."""

        # load arguments into attributes
        self.name = name

        self.output = output
        self.opt_algorithm = opt_algorithm

        self.per_atom_par = per_atom_par
        self.maxiter = maxiter

        # set up parallelization
        self.n_cpus = n_cpus
        self.n_sample = n_sample
        self.parallel = parallel

        self.component = component
        self.kernels = (
            ["twobody", "threebody"]
            if kernels is None
            else kernel_str_to_array("".join(kernels))
        )
        self.cutoffs = {} if cutoffs is None else cutoffs
        self.hyp_labels = hyp_labels
        self.hyps_mask = {} if hyps_mask is None else hyps_mask
        self.hyps = hyps

        GaussianProcess.backward_arguments(kwargs, self.__dict__)
        GaussianProcess.backward_attributes(self.__dict__)

        # ------------  "computed" attributes ------------

        if self.output is None:
            self.logger_name = self.name + "GaussianProcess"
            set_logger(self.logger_name, stream=True, fileout_name=None, verbose="info")
        else:
            self.logger_name = self.output.basename + "log"

        if self.hyps is None:
            # If no hyperparameters are passed in, assume 2 hyps for each
            # kernel, plus one noise hyperparameter, and use a guess value
            self.hyps = np.array([0.1] * (1 + 2 * len(self.kernels)))
        else:
            self.hyps = np.array(self.hyps, dtype=np.float64)

        kernel, grad, ek, efk, efs_e, efs_f, efs_self = str_to_kernel_set(
            self.kernels, self.component, self.hyps_mask
        )
        self.kernel = kernel
        self.kernel_grad = grad
        self.energy_force_kernel = efk
        self.energy_kernel = ek
        self.efs_energy_kernel = efs_e
        self.efs_force_kernel = efs_f
        self.efs_self_kernel = efs_self
        self.kernels = kernel_str_to_array(kernel.__name__)

        # parallelization
        if self.parallel:
            if self.n_cpus is None:
                self.n_cpus = mp.cpu_count()
            else:
                self.n_cpus = n_cpus
        else:
            self.n_cpus = 1

        self.training_data = []  # Atomic environments
        self.training_labels = []  # Forces acting on central atoms
        self.training_labels_np = np.empty(
            0,
        )
        self.n_envs_prev = len(self.training_data)

        # Attributes to accomodate energy labels:
        self.training_structures = []  # Environments of each structure
        self.energy_labels = []  # Energies of training structures
        self.energy_labels_np = np.empty(
            0,
        )
        self.energy_noise = energy_noise
        self.all_labels = np.empty(
            0,
        )

        # Parameters set during training
        self.ky_mat = None
        self.force_block = None
        self.energy_block = None
        self.force_energy_block = None
        self.l_mat = None
        self.l_mat_inv = None
        self.alpha = None
        self.ky_mat_inv = None
        self.likelihood = None
        self.likelihood_gradient = None
        self.bounds = None

        # File used for reading / writing model if model is large
        self.ky_mat_file = None
        # Flag if too-big warning has been printed for this model
        self.large_warning = False

        if self.logger_name is None:
            if self.output is None:
                self.logger_name = self.name + "GaussianProcess"
                set_logger(
                    self.logger_name, stream=True, fileout_name=None, verbose="info"
                )
            else:
                self.logger_name = self.output.basename + "log"
        logger = logging.getLogger(self.logger_name)

        if self.cutoffs == {}:
            # If no cutoffs are passed in, assume 7 A for 2 body, 3.5 for
            # 3-body.
            cutoffs = {}
            if "twobody" in self.kernels:
                cutoffs["twobody"] = 7
            if "threebody" in self.kernels:
                cutoffs["threebody"] = 3.5
            if "manybody" in self.kernels:
                raise ValueError(
                    "No cutoff was set for the manybody kernel."
                    "A default value will not be set by default."
                )

            self.cutoffs = cutoffs
            logger.warning(
                "Warning: No cutoffs were set for your GP."
                "Default values have been assigned but you "
                "should think carefully about which are "
                "appropriate for your use case."
            )

        self.check_instantiation()

    @property
    def force_noise(self):
        return Parameters.get_noise(self.hyps_mask, self.hyps, constraint=False)

    @property
    def hyps_and_labels(self):
        return Parameters.get_hyps(
            self.hyps_mask, self.hyps, constraint=False, label=True
        )

    def check_instantiation(self):
        """
        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:
        """
        logger = logging.getLogger(self.logger_name)

        # check whether it's be loaded before
        loaded = False
        if self.name in _global_training_labels:
            if (
                _global_training_labels.get(self.name, None)
                is not self.training_labels_np
            ):
                loaded = True
        if self.name in _global_energy_labels:
            if _global_energy_labels.get(self.name, None) is not self.energy_labels_np:
                loaded = True

        if loaded:

            base = f"{self.name}"
            count = 2
            while self.name in _global_training_labels and count < 100:
                time.sleep(random())
                self.name = f"{base}_{count}"
                logger.debug(
                    "Specified GP name is present in global memory; "
                    "Attempting to rename the "
                    f"GP instance to {self.name}"
                )
                count += 1
            if self.name in _global_training_labels:
                milliseconds = int(round(time.time() * 1000) % 10000000)
                self.name = f"{base}_{milliseconds}"
                logger.debug(
                    "Specified GP name still present in global memory: "
                    f"renaming the gp instance to {self.name}"
                )
            logger.debug(f"Final name of the gp instance is {self.name}")

        self.sync_data()

        self.hyps_mask = Parameters.check_instantiation(
            hyps=self.hyps,
            cutoffs=self.cutoffs,
            kernels=self.kernels,
            param_dict=self.hyps_mask,
        )

        self.bounds = deepcopy(self.hyps_mask.get("bounds", None))

    def update_kernel(
        self,
        kernels: List[str],
        component: str = "mc",
        hyps=None,
        cutoffs: dict = None,
        hyps_mask: dict = None,
    ):
        kernel, grad, ek, efk, _, _, _ = str_to_kernel_set(
            kernels, component, hyps_mask
        )
        self.kernel = kernel
        self.kernel_grad = grad
        self.energy_force_kernel = efk
        self.energy_kernel = ek
        self.kernels = kernel_str_to_array(kernel.__name__)

        if hyps_mask is not None:
            self.hyps_mask = hyps_mask
        # Cutoffs argument will override hyps mask's cutoffs key, if present
        if isinstance(hyps_mask, dict) and cutoffs is None:
            cutoffs = hyps_mask.get("cutoffs", None)

        if cutoffs is not None:
            if self.cutoffs != cutoffs:
                self.adjust_cutoffs(cutoffs, train=False, new_hyps_mask=hyps_mask)
            self.cutoffs = cutoffs

        if isinstance(hyps_mask, dict) and hyps is None:
            hyps = hyps_mask.get("hyps", None)

        if hyps is not None:
            self.hyps = hyps

    def update_db(
        self,
        struc: Structure,
        forces: "ndarray" = None,
        custom_range: List[int] = (),
        energy: float = None,
        stress: "ndarray" = None,
    ):
        """Given a structure and forces, add local environments from the
        structure to the training set of the GP. If energy is given, add the
        entire structure to the training set.

        Args:
            struc (Structure): Input structure. Local environments of atoms
                in this structure will be added to the training set of the GP.

            forces (np.ndarray): Forces on atoms in the structure.

            custom_range (List[int]): Indices of atoms whose local
                environments will be added to the training set of the GP.

            energy (float): Energy of the structure.

            stress (np.ndarray): Stress tensor of the structure. The stress
                tensor components should be given in the following order:
                xx, xy, xz, yy, yz, zz.
        """

        # By default, use all atoms in the structure
        noa = len(struc.positions)
        update_indices = custom_range or list(range(noa))

        # If forces are given, update the environment list.
        if forces is not None:
            for atom in update_indices:
                env_curr = AtomicEnvironment(
                    struc, atom, self.cutoffs, cutoffs_mask=self.hyps_mask
                )
                forces_curr = np.array(forces[atom])

                self.training_data.append(env_curr)
                self.training_labels.append(forces_curr)

            # create numpy array of training labels
            self.training_labels_np = np.hstack(self.training_labels)

        # If an energy is given, update the structure list.
        if energy is not None:
            structure_list = []  # Populate with all environments of the struc
            for atom in range(noa):
                env_curr = AtomicEnvironment(
                    struc, atom, self.cutoffs, cutoffs_mask=self.hyps_mask
                )
                structure_list.append(env_curr)

            self.energy_labels.append(energy)
            self.training_structures.append(structure_list)
            self.energy_labels_np = np.array(self.energy_labels)

        if forces is None and energy is None and stress is None:
            logger = logging.getLogger(self.logger_name)
            logger.warn(
                "Update DB method called with data but no labels!"
                "The GP has not been updated with data!"
            )
        # update list of all labels
        self.all_labels = np.concatenate(
            (self.training_labels_np, self.energy_labels_np)
        )
        self.sync_data()

    def add_one_env(
        self,
        env: AtomicEnvironment,
        force: "np.ndarray" = None,
        train: bool = False,
        **kwargs,
    ):
        """Add a single local environment to the training set of the GP.

        Args:
            env (AtomicEnvironment): Local environment to be added to the
                training set of the GP.
            force (np.ndarray): Force on the central atom of the local
                environment in the form of a 3-component Numpy array
                containing the x, y, and z components.
            train (bool): If True, the GP is trained after the local
                environment is added.
        """
        self.training_data.append(env)
        if force is None:
            self.training_labels.append(env.force)
        else:
            self.training_labels.append(force)
        self.training_labels_np = np.hstack(self.training_labels)
        self.sync_data()

        # update list of all labels
        self.all_labels = np.concatenate(
            (self.training_labels_np, self.energy_labels_np)
        )

        if train:
            self.train(**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,
    ):
        """Train Gaussian Process model on training data. Tunes the
        hyperparameters to maximize the likelihood, then computes L and alpha
        (related to the covariance matrix of the training set).

        Args:
            logger (logging.logger): logger object specifying where to write the
                progress of the optimization.
            custom_bounds (np.ndarray): Custom bounds on the hyperparameters.
            grad_tol (float): Tolerance of the hyperparameter gradient that
                determines when hyperparameter optimization is terminated.
            x_tol (float): Tolerance on the x values used to decide when
                Nelder-Mead hyperparameter optimization is terminated.
            line_steps (int): Maximum number of line steps for L-BFGS
                hyperparameter optimization.
                :param logger_name:
                :param print_progress:
        """

        verbose = "warning"
        if print_progress:
            verbose = "info"
        if logger_name is None:
            set_logger(
                "gp_algebra",
                stream=True,
                fileout_name="log.gp_algebra",
                verbose=verbose,
            )
            logger_name = "gp_algebra"

        disp = print_progress

        if (
            max(len(self.training_data), len(self.training_labels)) > 5000
            and not self.large_warning
        ):

            self.large_warning = True
            warning_message = (
                "WARNING! Your GP is very large (>5000 atomic "
                "environments). The hyperparameter optimization process "
                "does not scale favorably with increasing atomic "
                "environments"
                " (roughly N^2)"
                "and so your GP may take a very long time to train."
                "Consider finding a way to reduce the number of atomic "
                "environments in your model if you want to optimize the "
                "hyperparameters or optimize them by a different route."
            )

            logger = logging.getLogger(self.logger_name)
            logger.warning(warning_message)

        if len(self.training_data) == 0 or len(self.training_labels) == 0:
            raise Warning(
                "You are attempting to train a GP with no "
                "training data. Add environments and forces "
                "to the GP and try again."
            )

        x_0 = self.hyps

        args = (
            self.name,
            self.kernel_grad,
            logger_name,
            self.cutoffs,
            self.hyps_mask,
            self.n_cpus,
            self.n_sample,
        )

        res = None

        if self.opt_algorithm == "L-BFGS-B":

            # bound signal noise below to avoid overfitting
            if self.bounds is None:
                bounds = np.array([(1e-6, np.inf)] * len(x_0))
                bounds[-1, 0] = 1e-3
            else:
                bounds = self.bounds

            # Catch linear algebra errors and switch to BFGS if necessary
            try:
                res = minimize(
                    get_neg_like_grad,
                    x_0,
                    args,
                    method="L-BFGS-B",
                    jac=True,
                    bounds=bounds,
                    options={
                        "disp": disp,
                        "gtol": grad_tol,
                        "maxls": line_steps,
                        "maxiter": self.maxiter,
                    },
                )
            except np.linalg.LinAlgError:
                logger = logging.getLogger(self.logger_name)
                logger.warning(
                    "Algorithm for L-BFGS-B failed. Changing to "
                    "BFGS for remainder of run."
                )
                self.opt_algorithm = "BFGS"

        if custom_bounds is not None:
            res = minimize(
                get_neg_like_grad,
                x_0,
                args,
                method="L-BFGS-B",
                jac=True,
                bounds=custom_bounds,
                options={
                    "disp": disp,
                    "gtol": grad_tol,
                    "maxls": line_steps,
                    "maxiter": self.maxiter,
                },
            )

        elif self.opt_algorithm == "BFGS":
            res = minimize(
                get_neg_like_grad,
                x_0,
                args,
                method="BFGS",
                jac=True,
                options={"disp": disp, "gtol": grad_tol, "maxiter": self.maxiter},
            )

        if res is None:
            raise RuntimeError("Optimization failed for some reason.")
        self.hyps = res.x
        self.set_L_alpha()
        self.likelihood = -res.fun
        self.likelihood_gradient = -res.jac

        return res

    def check_L_alpha(self):
        """
        Check that the alpha vector is up to date with the training set. If
        not, update_L_alpha is called.
        """

        # Check that alpha is up to date with training set
        size3 = len(self.training_data) * 3 + len(self.training_structures)

        # If model is empty, then just return
        if size3 == 0:
            return

        if self.alpha is None:
            self.update_L_alpha()
        elif size3 > self.alpha.shape[0]:
            self.update_L_alpha()
        elif size3 != self.alpha.shape[0]:
            self.set_L_alpha()

    def predict(self, x_t: AtomicEnvironment, d: int) -> [float, float]:
        """
        Predict a force component of the central atom of a local environment.

        Args:
            x_t (AtomicEnvironment): Input local environment.
            d (int): Force component to be predicted (1 is x, 2 is y, and
                3 is z).

        Return:
            (float, float): Mean and epistemic variance of the prediction.
        """

        if d not in [1, 2, 3]:
            raise ValueError("d should be 1, 2, or 3")

        # Kernel vector allows for evaluation of atomic environments.
        if self.parallel and not self.per_atom_par:
            n_cpus = self.n_cpus
        else:
            n_cpus = 1

        self.sync_data()

        k_v = get_kernel_vector(
            self.name,
            self.kernel,
            self.energy_force_kernel,
            x_t,
            d,
            self.hyps,
            cutoffs=self.cutoffs,
            hyps_mask=self.hyps_mask,
            n_cpus=n_cpus,
            n_sample=self.n_sample,
        )

        # Guarantee that alpha is up to date with training set
        self.check_L_alpha()

        # get predictive mean
        pred_mean = np.matmul(k_v, self.alpha)

        # get predictive variance without cholesky (possibly faster)
        # pass args to kernel based on if mult. hyperparameters in use
        args = from_mask_to_args(self.hyps, self.cutoffs, self.hyps_mask)

        self_kern = self.kernel(x_t, x_t, d, d, *args)
        pred_var = self_kern - np.matmul(np.matmul(k_v, self.ky_mat_inv), k_v)

        return pred_mean, pred_var

    def predict_force_xyz(self, x_t: AtomicEnvironment) -> ("np.ndarray", "np.ndarray"):
        """
        Simple wrapper to predict all three components of a force in one go.
        :param x_t:
        :return:
        """
        forces = []
        stds = []
        for d in (1, 2, 3):
            force, std = self.predict(x_t, d)
            forces.append(force)
            stds.append(std)
        return np.array(forces), np.array(stds)

    def predict_local_energy(self, x_t: AtomicEnvironment) -> float:
        """Predict the local energy of a local environment.

        Args:
            x_t (AtomicEnvironment): Input local environment.

        Return:
            float: Local energy predicted by the GP.
        """

        if self.parallel and not self.per_atom_par:
            n_cpus = self.n_cpus
        else:
            n_cpus = 1

        self.sync_data()

        k_v = en_kern_vec(
            self.name,
            self.energy_force_kernel,
            self.energy_kernel,
            x_t,
            self.hyps,
            cutoffs=self.cutoffs,
            hyps_mask=self.hyps_mask,
            n_cpus=n_cpus,
            n_sample=self.n_sample,
        )

        pred_mean = np.matmul(k_v, self.alpha)

        return pred_mean

    def predict_local_energy_and_var(self, x_t: AtomicEnvironment):
        """Predict the local energy of a local environment and its
        uncertainty.

        Args:
            x_t (AtomicEnvironment): Input local environment.

        Return:
            (float, float): Mean and predictive variance predicted by the GP.
        """

        if self.parallel and not self.per_atom_par:
            n_cpus = self.n_cpus
        else:
            n_cpus = 1

        self.sync_data()

        # get kernel vector
        k_v = en_kern_vec(
            self.name,
            self.energy_force_kernel,
            self.energy_kernel,
            x_t,
            self.hyps,
            cutoffs=self.cutoffs,
            hyps_mask=self.hyps_mask,
            n_cpus=n_cpus,
            n_sample=self.n_sample,
        )

        # get predictive mean
        pred_mean = np.matmul(k_v, self.alpha)

        # get predictive variance
        v_vec = solve_triangular(self.l_mat, k_v, lower=True)
        args = from_mask_to_args(self.hyps, self.cutoffs, self.hyps_mask)

        self_kern = self.energy_kernel(x_t, x_t, *args)

        pred_var = self_kern - np.matmul(v_vec, v_vec)

        return pred_mean, pred_var

    def predict_efs(self, x_t: AtomicEnvironment):
        """Predict the local energy, forces, and partial stresses of an
        atomic environment and their predictive variances."""

        # Kernel vector allows for evaluation of atomic environments.
        if self.parallel and not self.per_atom_par:
            n_cpus = self.n_cpus
        else:
            n_cpus = 1

        self.sync_data()

        energy_vector, force_array, stress_array = efs_kern_vec(
            self.name,
            self.efs_force_kernel,
            self.efs_energy_kernel,
            x_t,
            self.hyps,
            cutoffs=self.cutoffs,
            hyps_mask=self.hyps_mask,
            n_cpus=n_cpus,
            n_sample=self.n_sample,
        )

        # Check that alpha is up to date with training set.
        self.check_L_alpha()

        # Compute mean predictions.
        en_pred = np.matmul(energy_vector, self.alpha)
        force_pred = np.matmul(force_array, self.alpha)
        stress_pred = np.matmul(stress_array, self.alpha)

        # Compute uncertainties.
        args = from_mask_to_args(self.hyps, self.cutoffs, self.hyps_mask)
        self_en, self_force, self_stress = self.efs_self_kernel(x_t, *args)

        en_var = self_en - np.matmul(
            np.matmul(energy_vector, self.ky_mat_inv), energy_vector
        )
        force_var = self_force - np.diag(
            np.matmul(np.matmul(force_array, self.ky_mat_inv), force_array.transpose())
        )
        stress_var = self_stress - np.diag(
            np.matmul(
                np.matmul(stress_array, self.ky_mat_inv), stress_array.transpose()
            )
        )

        return en_pred, force_pred, stress_pred, en_var, force_var, stress_var

    def set_L_alpha(self):
        """
        Invert the covariance matrix, setting L (a lower triangular
        matrix s.t. L L^T = (K + sig_n^2 I)) and alpha, the inverse
        covariance matrix multiplied by the vector of training labels.
        The forces and variances are later obtained using alpha.
        """

        self.sync_data()

        ky_mat = get_Ky_mat(
            self.hyps,
            self.name,
            self.kernel,
            self.energy_kernel,
            self.energy_force_kernel,
            self.energy_noise,
            cutoffs=self.cutoffs,
            hyps_mask=self.hyps_mask,
            n_cpus=self.n_cpus,
            n_sample=self.n_sample,
        )

        l_mat = np.linalg.cholesky(ky_mat)
        l_mat_inv = np.linalg.inv(l_mat)
        ky_mat_inv = l_mat_inv.T @ l_mat_inv
        alpha = np.matmul(ky_mat_inv, self.all_labels)

        self.ky_mat = ky_mat
        self.l_mat = l_mat
        self.alpha = alpha
        self.ky_mat_inv = ky_mat_inv

        self.likelihood = get_like_from_mats(ky_mat, l_mat, alpha, self.name)
        self.n_envs_prev = len(self.training_data)

    def update_L_alpha(self):
        """
        Update the GP's L matrix and alpha vector without recalculating
        the entire covariance matrix K.
        """

        # Set L matrix and alpha if set_L_alpha has not been called yet
        if self.l_mat is None or np.array(self.ky_mat) is np.array(None):
            self.set_L_alpha()
            return

        # Reset global variables.
        self.sync_data()

        ky_mat = get_ky_mat_update(
            self.ky_mat,
            self.n_envs_prev,
            self.hyps,
            self.name,
            self.kernel,
            self.energy_kernel,
            self.energy_force_kernel,
            self.energy_noise,
            cutoffs=self.cutoffs,
            hyps_mask=self.hyps_mask,
            n_cpus=self.n_cpus,
            n_sample=self.n_sample,
        )

        l_mat = np.linalg.cholesky(ky_mat)
        l_mat_inv = np.linalg.inv(l_mat)
        ky_mat_inv = l_mat_inv.T @ l_mat_inv
        alpha = np.matmul(ky_mat_inv, self.all_labels)

        self.ky_mat = ky_mat
        self.l_mat = l_mat
        self.alpha = alpha
        self.ky_mat_inv = ky_mat_inv
        self.n_envs_prev = len(self.training_data)

    def __str__(self):
        """String representation of the GP model."""

        thestr = ""
        thestr += f"Number of cpu cores: {self.n_cpus}\n"
        thestr += f"Kernel: {self.kernels}\n"
        thestr += f"Training points: {len(self.training_data)}\n"
        thestr += f"Cutoffs: {self.cutoffs}\n"

        thestr += f"Number of hyperparameters: {len(self.hyps)}\n"
        thestr += f"Hyperparameter array: {str(self.hyps)}\n"

        if self.hyp_labels is None:
            # Put unlabeled hyperparameters on one line
            thestr = thestr[:-1]
            thestr += str(self.hyps) + "\n"
        else:
            for hyp, label in zip(self.hyps, self.hyp_labels):
                thestr += f"{label}: {hyp} \n"

        return thestr

    def as_dict(self):
        """Dictionary representation of the GP model."""

        self.check_L_alpha()

        out_dict = dict(vars(self))

        out_dict["training_data"] = [env.as_dict() for env in self.training_data]

        # Write training structures (which are just list of environments)
        out_dict["training_structures"] = []
        for n, env_list in enumerate(self.training_structures):
            out_dict["training_structures"].append([])
            for env_curr in env_list:
                out_dict["training_structures"][n].append(env_curr.as_dict())

        # Remove the callables
        for key in [
            "kernel",
            "kernel_grad",
            "energy_kernel",
            "energy_force_kernel",
            "efs_energy_kernel",
            "efs_force_kernel",
            "efs_self_kernel",
            "output",
        ]:
            out_dict.pop(key)

        return out_dict

    def sync_data(self):
        _global_training_data[self.name] = self.training_data
        _global_training_labels[self.name] = self.training_labels_np
        _global_training_structures[self.name] = self.training_structures
        _global_energy_labels[self.name] = self.energy_labels_np

    @staticmethod
    def from_dict(dictionary):
        """Create GP object from dictionary representation."""

        GaussianProcess.backward_arguments(dictionary, dictionary)
        GaussianProcess.backward_attributes(dictionary)

        new_gp = GaussianProcess(**dictionary)

        # Save time by attempting to load in computed attributes
        if "training_data" in dictionary:
            new_gp.training_data = [
                AtomicEnvironment.from_dict(env) for env in dictionary["training_data"]
            ]
            new_gp.training_labels = deepcopy(dictionary["training_labels"])
            new_gp.training_labels_np = deepcopy(dictionary["training_labels_np"])
            new_gp.sync_data()

        # Reconstruct training structures.
        if "training_structures" in dictionary:
            new_gp.training_structures = []
            for n, env_list in enumerate(dictionary["training_structures"]):
                new_gp.training_structures.append([])
                for env_curr in env_list:
                    new_gp.training_structures[n].append(
                        AtomicEnvironment.from_dict(env_curr)
                    )
            new_gp.energy_labels = deepcopy(dictionary["energy_labels"])
            new_gp.energy_labels_np = deepcopy(dictionary["energy_labels_np"])
            new_gp.sync_data()

        new_gp.all_labels = np.concatenate(
            (new_gp.training_labels_np, new_gp.energy_labels_np)
        )

        new_gp.likelihood = dictionary.get("likelihood", None)
        new_gp.likelihood_gradient = dictionary.get("likelihood_gradient", None)

        new_gp.n_envs_prev = len(new_gp.training_data)

        # Save time by attempting to load in computed attributes
        if dictionary.get("ky_mat_file"):
            try:
                new_gp.ky_mat = np.load(dictionary["ky_mat_file"])
                new_gp.compute_matrices()
                new_gp.ky_mat_file = None

            except FileNotFoundError:
                new_gp.ky_mat = None
                new_gp.l_mat = None
                new_gp.alpha = None
                new_gp.ky_mat_inv = None
                filename = dictionary.get("ky_mat_file")
                logger = logging.getLogger(new_gp.logger_name)
                logger.warning(
                    "the covariance matrices are not loaded"
                    f"because {filename} cannot be found"
                )
        else:
            new_gp.ky_mat = (
                np.array(dictionary["ky_mat"])
                if dictionary.get("ky_mat") is not None
                else None
            )
            new_gp.ky_mat_inv = (
                np.array(dictionary["ky_mat_inv"])
                if dictionary.get("ky_mat_inv") is not None
                else None
            )
            new_gp.ky_mat = (
                np.array(dictionary["ky_mat"])
                if dictionary.get("ky_mat") is not None
                else None
            )
            new_gp.l_mat = (
                np.array(dictionary["l_mat"])
                if dictionary.get("l_mat") is not None
                else None
            )
            new_gp.alpha = (
                np.array(dictionary["alpha"])
                if dictionary.get("alpha") is not None
                else None
            )

        return new_gp

    def compute_matrices(self):
        """
        When covariance matrix is known, reconstruct other matrices.
        Used in re-loading large GPs.
        :return:
        """
        ky_mat = self.ky_mat

        if ky_mat is None or (isinstance(ky_mat, np.ndarray) and not np.any(ky_mat)):
            Warning(
                "Warning: Covariance matrix was not loaded but "
                "compute_matrices was called. Computing covariance "
                "matrix and proceeding..."
            )
            self.set_L_alpha()

        else:
            self.l_mat = np.linalg.cholesky(ky_mat)
            self.l_mat_inv = np.linalg.inv(self.l_mat)
            self.ky_mat_inv = self.l_mat_inv.T @ self.l_mat_inv
            self.alpha = np.matmul(self.ky_mat_inv, self.all_labels)

    def adjust_cutoffs(
        self,
        new_cutoffs: Union[list, tuple, "np.ndarray"] = None,
        reset_L_alpha=True,
        train=True,
        new_hyps_mask=None,
    ):
        """
        Loop through atomic environment objects stored in the training data,
        and re-compute cutoffs for each. Useful if you want to gauge the
        impact of cutoffs given a certain training set! Unless you know
        *exactly* what you are doing for some development or test purpose,
        it is **highly** suggested that you call set_L_alpha and
        re-optimize your hyperparameters afterwards as is default here.

        A helpful way to update the cutoffs and kernel for an extant
        GP is to perform the following commands:
        >> hyps_mask = pm.as_dict()
        >> hyps = hyps_mask['hyps']
        >> cutoffs = hyps_mask['cutoffs']
        >> kernels = hyps_mask['kernels']
        >> gp_model.update_kernel(kernels, 'mc', hyps, cutoffs, hyps_mask)

        :param reset_L_alpha:
        :param train:
        :param new_hyps_mask:
        :param new_cutoffs:
        :return:
        """

        if new_hyps_mask is not None:
            hm = new_hyps_mask
            self.hyps_mask = new_hyps_mask
        else:
            hm = self.hyps_mask
        if new_cutoffs is None:
            try:
                new_cutoffs = hm["cutoffs"]
            except KeyError:
                raise KeyError(
                    "New cutoffs not found in the hyps_mask"
                    "dictionary via call to 'cutoffs' key."
                )

        # update environment
        nenv = len(self.training_data)
        for i in range(nenv):
            self.training_data[i].cutoffs = new_cutoffs
            self.training_data[i].cutoffs_mask = hm
            self.training_data[i].setup_mask(hm)
            self.training_data[i].compute_env()

        # Ensure that training data and labels are still consistent
        self.sync_data()

        self.cutoffs = new_cutoffs

        if reset_L_alpha:
            del self.l_mat
            del self.ky_mat
            self.set_L_alpha()

        if train:
            self.train()

    def remove_force_data(
        self, indexes: Union[int, List[int]], update_matrices: bool = True
    ) -> Tuple[List[Structure], List["ndarray"]]:
        """
        Remove force components from the model. Convenience function which
        deletes individual data points.

        Matrices should *always* be updated if you intend to use the GP to make
        predictions afterwards. This might be time consuming for large GPs,
        so, it is provided as an option, but, only do so with extreme caution.
        (Undefined behavior may result if you try to make predictions and/or
        add to the training set afterwards).

        Returns training data which was removed akin to a pop method, in order
        of lowest to highest index passed in.

        :param indexes: Indexes of envs in training data to remove.
        :param update_matrices: If false, will not update the GP's matrices
            afterwards (which can be time consuming for large models).
            This should essentially always be true except for niche development
            applications.
        :return:
        """

        # Listify input even if one integer
        if isinstance(indexes, int):
            indexes = [indexes]

        if max(indexes) > len(self.training_data):
            raise ValueError("Index out of range of data")

        if len(indexes) == 0:
            return [], []

        # Get in reverse order so that modifying higher indexes doesn't affect
        # lower indexes
        indexes.sort(reverse=True)
        removed_data = []
        removed_labels = []
        for i in indexes:
            removed_data.append(self.training_data.pop(i))
            removed_labels.append(self.training_labels.pop(i))

        self.training_labels_np = np.hstack(self.training_labels)
        self.all_labels = np.concatenate(
            (self.training_labels_np, self.energy_labels_np)
        )
        self.sync_data()

        if update_matrices:
            self.set_L_alpha()

        # Put removed data in order of lowest to highest index
        removed_data.reverse()
        removed_labels.reverse()

        return removed_data, removed_labels

    def write_model(
        self, name: str, format: str = None, split_matrix_size_cutoff: int = 5000
    ):
        """
        Write model in a variety of formats to a file for later re-use.
        JSON files are open to visual inspection and are easier to use
        across different versions of FLARE or GP implementations. However,
        they are larger and loading them in takes longer (by setting up a
        new GP from the specifications). Pickled files can be faster to
        read & write, and they take up less memory.

        Args:
            name (str): Output name.
            format (str): Output format.
            split_matrix_size_cutoff (int): If there are more than this
            number of training points in the set, save the matrices seperately.
        """

        if len(self.training_data) > split_matrix_size_cutoff:
            np.save(f"{name}_ky_mat.npy", self.ky_mat)
            self.ky_mat_file = f"{name}_ky_mat.npy"

            temp_ky_mat = self.ky_mat
            temp_l_mat = self.l_mat
            temp_alpha = self.alpha
            temp_ky_mat_inv = self.ky_mat_inv

            self.ky_mat = None
            self.l_mat = None
            self.alpha = None
            self.ky_mat_inv = None

        # Automatically detect output format from name variable

        for detect in ["json", "pickle", "binary"]:
            if detect in name.lower():
                format = detect
                break

        if format is None:
            format = "json"

        supported_formats = ["json", "pickle", "binary"]

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

        elif format.lower() == "pickle" or format.lower() == "binary":
            if ".pickle" != name[-7:]:
                name += ".pickle"
            with open(name, "wb") as f:
                pickle.dump(self, f)

        else:
            raise ValueError(
                "Output format not supported: try from {}".format(supported_formats)
            )

        if len(self.training_data) > split_matrix_size_cutoff:
            self.ky_mat = temp_ky_mat
            self.l_mat = temp_l_mat
            self.alpha = temp_alpha
            self.ky_mat_inv = temp_ky_mat_inv

    @staticmethod
    def from_file(filename: str, format: str = ""):
        """
        One-line convenience method to load a GP from a file stored using
        write_file

        Args:
            filename (str): path to GP model
            format (str): json or pickle if format is not in filename
        :return:
        """

        if ".json" in filename or "json" in format:
            with open(filename, "r") as f:
                gp_model = GaussianProcess.from_dict(json.loads(f.readline()))

        elif ".pickle" in filename or "pickle" in format:
            with open(filename, "rb") as f:

                gp_model = pickle.load(f)

                GaussianProcess.backward_arguments(gp_model.__dict__, gp_model.__dict__)

                GaussianProcess.backward_attributes(gp_model.__dict__)

                if hasattr(gp_model, "ky_mat_file") and gp_model.ky_mat_file:
                    try:
                        gp_model.ky_mat = np.load(
                            gp_model.ky_mat_file, allow_pickle=True
                        )
                        gp_model.compute_matrices()
                    except FileNotFoundError:
                        gp_model.ky_mat = None
                        gp_model.l_mat = None
                        gp_model.alpha = None
                        gp_model.ky_mat_inv = None
                        Warning(
                            f"the covariance matrices are not loaded, "
                            f"this can take a long time to recompute"
                        )

        else:
            raise ValueError(
                "Warning: Format unspecieified or file is not .json or .pickle format."
            )

        gp_model.check_instantiation()
        return gp_model

    def __len__(self):
        return len(self.training_data)

    @property
    def training_statistics(self) -> dict:
        """
        Return a dictionary with statistics about the current training data.
        Useful for quickly summarizing info about the GP.
        :return:
        """

        data = dict()

        data["N"] = len(self.training_data)

        # Count all of the present species in the atomic env. data
        present_species = []
        for env, _ in zip(self.training_data, self.training_labels):
            present_species.append(Z_to_element(env.structure.coded_species[env.atom]))

        # Summarize the relevant information
        data["species"] = list(set(present_species))
        data["envs_by_species"] = dict(Counter(present_species))

        return data

    @property
    def par(self):
        """
        Backwards compability attribute
        :return:
        """
        return self.parallel

    def __deepcopy__(self, memo):
        # this way can also deepcopy the training data in _global_training dicts
        return GaussianProcess.from_dict(self.as_dict())

    def __del__(self):
        if self is None:
            return
        if self.name in _global_training_labels:
            return (
                _global_training_data.pop(self.name, None),
                _global_training_labels.pop(self.name, None),
                _global_training_structures.pop(self.name, None),
                _global_energy_labels.pop(self.name, None),
            )

    @staticmethod
    def backward_arguments(kwargs, new_args={}):
        """
        update the initialize arguments that were renamed
        """

        if "kernel_name" in kwargs:
            DeprecationWarning("kernel_name is being replaced with kernels")
            new_args["kernels"] = kernel_str_to_array(kwargs["kernel_name"])
            kwargs.pop("kernel_name")
        if "nsample" in kwargs:
            DeprecationWarning("nsample is being replaced with n_sample")
            new_args["n_sample"] = kwargs["nsample"]
            kwargs.pop("nsample")
        if "par" in kwargs:
            DeprecationWarning("par is being replaced with parallel")
            new_args["parallel"] = kwargs["par"]
            kwargs.pop("par")
        if "no_cpus" in kwargs:
            DeprecationWarning("no_cpus is being replaced with n_cpu")
            new_args["n_cpus"] = kwargs["no_cpus"]
            kwargs.pop("no_cpus")
        if "multihyps" in kwargs:
            DeprecationWarning("multihyps is removed")
            kwargs.pop("multihyps")

        return new_args

    @staticmethod
    def backward_attributes(dictionary):
        """
        add new attributes to old instance
        or update attribute types
        """

        if "name" not in dictionary:
            dictionary["name"] = "default_gp"
        if "per_atom_par" not in dictionary:
            dictionary["per_atom_par"] = True
        if "opt_algorithm" not in dictionary:
            dictionary["opt_algorithm"] = "L-BFGS-B"
        if "hyps_mask" not in dictionary:
            dictionary["hyps_mask"] = None
        if "parallel" not in dictionary:
            dictionary["parallel"] = False
        if "component" not in dictionary:
            dictionary["component"] = "mc"

        if "training_structures" not in dictionary:
            # Environments of each structure
            dictionary["training_structures"] = []
            dictionary["energy_labels"] = []  # Energies of training structures
            dictionary["energy_labels_np"] = np.empty(
                0,
            )

        if "training_labels" not in dictionary:
            dictionary["training_labels"] = []
            dictionary["training_labels_np"] = np.empty(
                0,
            )

        if "energy_noise" not in dictionary:
            dictionary["energy_noise"] = 0.01

        if not isinstance(dictionary["cutoffs"], dict):
            dictionary["cutoffs"] = Parameters.cutoff_array_to_dict(
                dictionary["cutoffs"]
            )

        dictionary["hyps_mask"] = Parameters.backward(
            dictionary["kernels"], deepcopy(dictionary["hyps_mask"])
        )

        if "logger_name" not in dictionary:
            dictionary["logger_name"] = None class Output:
    """
    This is an I/O class that hosts the log files for OTF and Trajectories
    class. It is also used in get_neg_like_grad and get_neg_likelihood in
    gp_algebra to print intermediate results.

    It opens and print files with the basename prefix and different
    suffixes corresponding to different kinds of output data.

    :param basename: Base output file name, suffixes will be added
    :type basename: str, optional
    :param verbose: print level. The same as logging level. It can be
                    CRITICAL, ERROR, WARNING, INFO, DEBUG, NOTSET
    :type verbose: str, optional
    :param always_flush: Always write to file instantly
    :type always_flus: bool, optional
    """

    def __init__(
        self,
        basename: str = "otf_run",
        verbose: str = "INFO",
        print_as_xyz: bool = False,
        always_flush: bool = False,
    ):
        """
        Construction. Open files.
        """
        self.basename = f"{basename}"
        self.print_as_xyz = print_as_xyz
        self.always_flush = always_flush

        filesuffix = {"log": ".out", "hyps": "-hyps.dat"}
        if print_as_xyz:
            filesuffix["xyz"] = ".xyz"

        self.logger = []
        for filetype in filesuffix:
            self.open_new_log(filetype, filesuffix[filetype], verbose)

    def conclude_run(self, extra_strings: List[str] = None):
        """
        destruction function that closes all files
        """

        logger = logging.getLogger(self.basename + "log")
        logger.info("-" * 20)
        logger.info("Run complete.")
        if extra_strings is not None:
            for string in extra_strings:
                logger.info(string)
        logging.shutdown()
        self.logger = []

    def open_new_log(self, filetype: str, suffix: str, verbose="info"):
        """
        Open files.  If files with the same
        name are exist, they are backed up with a suffix "-bak".

        :param filetype: the key name for logging
        :param suffix: the suffix of the file to be opened
        :param verbose: the verbose level for the logger
        """

        if filetype not in self.logger:
            set_logger(
                self.basename + filetype,
                stream=False,
                fileout_name=self.basename + suffix,
                verbose=verbose,
            )
            self.logger += [filetype]

    def write_to_log(self, logstring: str, name: str = "log", flush: bool = False):
        """
        Write any string to logfile

        :param logstring: the string to write
        :param name: the key name of the file to logger named 'log'
        :param flush: whether it should be flushed
        """
        logger = logging.getLogger(self.basename + name)
        logger.info(logstring)

        if flush or self.always_flush:
            logger.handlers[0].flush()

    def write_header(
        self,
        gp_str: str,
        dt: float = None,
        Nsteps: int = None,
        structure: Structure = None,
        std_tolerance: Union[float, int] = None,
        optional: dict = None,
    ):
        """
        TO DO: this should be replace by the string method of GP and OTF, GPFA

        Write header to the log function. Designed for Trajectory Trainer and
        OTF runs and can take flexible input for both.

        :param gp_str: string representation of the GP
        :param dt: timestep for OTF MD
        :param Nsteps: total number of steps for OTF MD
        :param structure: initial structure
        :param std_tolerance: tolarence for active learning
        :param optional: a dictionary of all the other parameters
        """

        f = logging.getLogger(self.basename + "log")
        f.info(f"{datetime.datetime.now()}")

        # Write package version numbers
        f.info(f"python {sys.version}")
        f.info(f"numpy {np.__version__}")
        try:
            import ase

            f.info(f"ase {ase.__version__}")
        except:
            pass

        try:
            import pymatgen

            f.info(f"pymatgen {pymatgen.__version__}")
        except:
            pass

        try:
            f.info(f"flare {flare.__version__}")
        except AttributeError:  # using customized package not from pip install
            dirs = flare.__path__[0].split("/")
            dirs_rep = copy.copy(dirs)
            dirs_rep[-1] = "setup.py"
            setup_path = "/".join(dirs_rep)
            try:
                setup_py = open(setup_path)

                for line in setup_py.readlines():
                    line = line.strip()
                    if "version=" in line:
                        f.info(f"flare {line[9:len(line)-2]}")
                        break
                setup_py.close()
            # Catch case where the version can't be found.
            # (This happens in the Action test of the flare++ tutorial.)
            except FileNotFoundError:
                pass

        try:
            import flare_pp

            f.info(f"flare_pp {flare_pp.__version__}")
        except AttributeError:
            import flare_pp

            dirs = flare_pp.__path__[0].split("/")
            dirs_rep = copy.copy(dirs)
            dirs_rep[-1] = "setup.py"
            setup_path = "/".join(dirs_rep)

            try:
                setup_py = open(setup_path)

                for line in setup_py.readlines():
                    line = line.strip()
                    if "version=" in line:
                        f.info(f"flare_pp {line[9:len(line)-2]}")
                        break
                setup_py.close()
            except FileNotFoundError:
                pass
        except ModuleNotFoundError:
            pass

        # Write uncertainty tolerance
        if isinstance(std_tolerance, tuple):
            std_string = (
                "Relative uncertainty tolerance: "
                f"{std_tolerance[0]} times noise hyperparameter \n"
            )
            std_string += f"Absolute uncertainty tolerance: {std_tolerance[1]} eV/A\n"
        elif std_tolerance < 0:
            std_string = f"Uncertainty tolerance: {np.abs(std_tolerance)} eV/A\n"
        elif std_tolerance > 0:
            std_string = (
                f"Uncertainty tolerance: {np.abs(std_tolerance)} "
                "times noise hyperparameter \n"
            )
        else:
            std_string = ""

        headerstring = "\n"
        headerstring += gp_str
        headerstring += "\n"
        headerstring += std_string
        if dt is not None:
            headerstring += f"Timestep (ps): {dt}\n"
        headerstring += f"Number of frames: {Nsteps}\n"
        if structure is not None:
            headerstring += f"Number of atoms: {structure.nat}\n"
            headerstring += f"System species: {set(structure.species_labels)}\n"
            headerstring += "Periodic cell (A): \n"
            headerstring += str(np.array(structure.cell)) + "\n"

        if optional:
            for key, value in optional.items():
                headerstring += f"{key}: {value} \n"

        # report previous positions
        if structure is not None:
            headerstring += "\nPrevious positions (A):\n"
            for i in range(len(structure.positions)):
                headerstring += f"{structure.species_labels[i]:5}"
                for j in range(3):
                    headerstring += f"{structure.prev_positions[i][j]:10.4f}"
                headerstring += "\n"
        headerstring += "-" * 80

        f.info(headerstring)

        if self.always_flush:
            f.handlers[0].flush()

    def write_md_header(self, dt, curr_step, dft_step):
        string = ""
        # Mark if a frame had DFT forces with an asterisk
        if not dft_step:
            string += "-" * 80 + "\n"
            string += f"-Frame: {curr_step} "
        else:
            string += f"\n*-Frame: {curr_step} "

        string += f"\nSimulation Time: {(dt * curr_step):10.3f} ps \n"
        return string

    def write_md_config(
        self,
        dt,
        curr_step,
        structure,
        temperature,
        KE,
        start_time,
        dft_step,
        velocities,
    ):
        """write md configuration in log file

        :param dt: timestemp of OTF MD
        :param curr_step: current timestep of OTF MD
        :param structure: atomic structure
        :param temperature: current temperature
        :param KE: current total kinetic energy
        :param local_energies: local energy
        :param start_time: starting time for time profiling
        :param dft_step: # of DFT calls
        :param velocities: list of velocities

        :return:
        """

        string = self.write_md_header(dt, curr_step, dft_step)

        # Construct Header line
        n_space = 30
        string += str.ljust("El", 5)
        string += str.center("Position (A)", n_space)
        string += " " * 4
        if not dft_step:
            string += str.center("GP Force (ev/A)", n_space)
            string += " " * 4
        else:
            string += str.center("DFT Force (ev/A)", n_space)
            string += " " * 4
        string += str.center("Std. Dev (ev/A)", n_space) + " " * 4
        string += str.center("Velocities (A/ps)", n_space) + "\n"

        # Construct atom-by-atom description
        for i in range(len(structure.positions)):
            string += f"{structure.species_labels[i]:5}"
            # string += '\t'
            for j in range(3):
                string += f"{structure.positions[i][j]:10.4f}"
            string += " " * 4
            for j in range(3):
                string += f"{structure.forces[i][j]:10.4f}"
            string += " " * 4
            for j in range(3):
                string += f"{structure.stds[i][j]:10.4f}"
            string += " " * 4
            for j in range(3):
                string += f"{velocities[i][j]:10.4f}"
            string += "\n"

        string += "\n"

        # Report cell if stress attribute is present.
        if structure.stress is not None:
            string += "Periodic cell (A): \n"
            string += str(np.array(structure.cell)) + "\n\n"

        # Report stress tensor.
        pressure = None
        if structure.stress is not None:
            stress_tensor = structure.stress * eva_to_gpa  # Convert to GPa
            s8 = " " * 8
            string += "Stress tensor (GPa):\n"
            string += (
                " " * 7
                + "xx"
                + s8
                + "yy"
                + s8
                + "zz"
                + s8
                + "yz"
                + s8
                + "xz"
                + s8
                + "xy\n"
            )
            for p in range(6):
                string += f"{stress_tensor[p]:10.3f}"
            string += "\n\n"
            pressure = (stress_tensor[0] + stress_tensor[1] + stress_tensor[2]) / 3

        # Report stress tensor uncertainties.
        if structure.stress_stds is not None:
            stress_stds = structure.stress_stds * eva_to_gpa  # Convert to GPa
            string += "Stress tensor uncertainties (GPa):\n"
            for p in range(6):
                string += f"{stress_stds[p]:10.3f}"
            string += "\n\n"

        # Report pressure.
        if pressure is not None:
            string += f"Pressure (GPa): {pressure:.6f} \n"

        string += f"Temperature (K): {temperature:.2f} \n"
        string += f"Kinetic energy (eV): {KE:.6f} \n"

        # Report potential energy.
        if structure.potential_energy is not None:
            string += f"Potential energy (eV): {structure.potential_energy:.6f} \n"

        # Report potential energy uncertainty.
        if structure.local_energy_stds is not None:
            pot_en_std = np.sqrt(np.sum(structure.local_energy_stds ** 2))
            string += f"Uncertainty (eV): {pot_en_std:.6f} \n"

        # Report total energy.
        if structure.potential_energy is not None:
            tot_en = KE + structure.potential_energy
            string += f"Total energy (eV): {tot_en:.6f} \n"

        logger = logging.getLogger(self.basename + "log")
        logger.info(string)
        self.write_wall_time(start_time)

        if self.always_flush:
            logger.handlers[0].flush()

    def write_xyz_config(
        self,
        curr_step,
        structure,
        forces: np.array = None,
        stds: np.array = None,
        dft_forces: np.array = None,
        dft_energy=0,
        predict_energy=float("nan"),
        target_atoms=None,
    ):
        """write atomic configuration in xyz file

        :param curr_step: Int, number of frames to note in the comment line
        :param structure: Structure, contain positions and forces
        :param forces: Optional list of forces to xyz file
        :param stds: Optional list of uncertanties to xyz file
        :param dft_forces: Optional second list of forces (e.g. DFT forces)

        :return:
        """

        xyz_str = structure.to_xyz(
            extended_xyz=True,
            print_stds=True,
            print_forces=True,
            print_max_stds=False,
            print_energies=True,
            predict_energy=predict_energy,
            dft_forces=dft_forces,
            dft_energy=dft_energy,
            timestep=curr_step,
            write_file="",
            append=False,
            target_atoms=target_atoms,
        )

        logger = logging.getLogger(self.basename + "xyz")
        logger.info(xyz_str)
        if self.always_flush:
            logger.handlers[0].flush()

    def write_hyps(
        self, hyp_labels, hyps, start_time, like, like_grad, name="log", hyps_mask=None
    ):
        """write hyperparameters to logfile

        :param name:
        :param hyp_labels: labels for hyper-parameters. can be None
        :param hyps: list of hyper-parameters
        :param start_time: start time for time profiling
        :param like: likelihood
        :param like_grad: gradient of likelihood

        :return:
        """
        f = logging.getLogger(self.basename + name)

        f.info("\nGP hyperparameters: ")

        if hyp_labels is not None:
            for i, label in enumerate(hyp_labels):
                f.info(f"Hyp{i} : {label:30s} = {hyps[i]:.4f}")
        else:
            for i, hyp in enumerate(hyps):
                f.info(f"Hyp{i} : {hyp:.4f}")
        if isinstance(like, list):
            f.info(f"Likelihood: {np.round(like,4)}")
        else:
            f.info(f"Likelihood: {like:.4f}")
        f.info(f"Likelihood gradient: {like_grad}")

        if start_time:
            self.write_wall_time(start_time)

        if self.always_flush:
            f.handlers[0].flush()

    def write_wall_time(self, start_time):
        time_curr = time.time() - start_time
        f = logging.getLogger(self.basename + "log")
        f.info(f"Wall time from start: {time_curr:.2f} s")

    def conclude_dft(self, dft_count, start_time):
        f = logging.getLogger(self.basename + "log")
        f.info("DFT run complete.")
        f.info(f"Number of DFT calls: {dft_count}")
        self.write_wall_time(start_time)

    def add_atom_info(self, train_atoms, stds):
        f = logging.getLogger(self.basename + "log")
        f.info(f"Adding atom {train_atoms} to the training set.")
        f.info(f"Uncertainty: {stds[train_atoms[0]]}")

    def write_gp_dft_comparison(
        self,
        curr_step,
        frame,
        start_time,
        dft_forces,
        dft_energy,
        error,
        local_energies=None,
        KE=None,
        mgp=False,
        cell=None,
        stress=None,
    ):
        """Write the comparison to logfile.

        :param curr_step: current timestep
        :param frame: Structure object that contains the current GP calculation
            results.
        :param start_time: start time for time profiling
        :param dft_forces: list of forces computed by DFT
        :param dft_energy: total energy computed by DFT
        :param error: list of force differences between DFT and GP prediction
        :param local_energies: local atomic energy
        :param KE: total kinetic energy
        :param cell: print the unit cell of the structure
        :param stress: print the stress acting on the cell

        :return:
        """

        string = ""

        # Mark if a frame had DFT forces with an asterisk
        string += f"\n*-Frame: {curr_step}"

        # Construct Header line
        string += "\nEl  Position (A) \t\t\t\t "
        if mgp:
            string += "M"
        string += "GP Force (ev/A)  \t\t\t\t"
        string += "Std. Dev (ev/A) \t\t\t\t"
        string += "DFT Force (ev/A)  \t\t\t\t \n"

        # Construct atom-by-atom description
        for i in range(len(frame.positions)):
            string += f"{frame.species_labels[i]} "
            for j in range(3):
                string += f"{frame.positions[i][j]:10.5} "
            string += "\t"
            for j in range(3):
                string += f"{frame.forces[i][j]:10.5} "
            string += "\t"
            for j in range(3):
                string += f"{frame.stds[i][j]:10.5} "
            string += "\t"
            for j in range(3):
                string += f"{dft_forces[i][j]:10.5} "
            string += "\n"

        string += "\n"

        # Print stress & cell related parameters
        if cell is not None:
            rounded_cell = np.round(cell, 4)
            string += f"cell: {[list(vec) for vec in rounded_cell]} \n"
        if stress:
            raise NotImplementedError

        # Compute errors and errors by species
        mae = np.nanmean(error) * 1000
        mac = np.mean(np.abs(dft_forces)) * 1000
        string += f"mean absolute error: {mae:.2f} meV/A\n"
        string += f"mean absolute dft component: {mac:.2f} meV/A\n"
        stat = f"{curr_step} {mae:.2} {mac:.2}"

        mae_per_species = {}
        count_per_species = {}
        species = [Z_to_element(Z) for Z in set(frame.coded_species)]
        for ele in species:
            mae_per_species[ele] = 0
            count_per_species[ele] = 0

        for atom in range(frame.nat):
            Z = frame.coded_species[atom]
            ele = Z_to_element(Z)
            if np.isnan(np.sum(error[atom, :])):
                continue
            mae_per_species[ele] += np.sum(error[atom, :])
            count_per_species[ele] += 1

        string += "mae per species\n"
        for ele in species:
            if count_per_species[ele] > 0:
                mae_per_species[ele] /= count_per_species[ele] * 3
                mae_per_species[ele] *= 1000  # Put in meV/A
                string += f"type {ele} mae: {mae_per_species[ele]:.2f} meV/A\n"
            stat += f" {mae_per_species[ele]:.2f}"

        # calculate potential and total energy
        if local_energies is not None:
            pot_en = 0
            pot_en = np.sum(local_energies)
            tot_en = KE + pot_en
            string += f"potential energy: {pot_en:10.6} eV (DFT: {dft_energy} eV\n"
            string += f"total energy: {tot_en:10.6} eV \n"
            stat += f" {pot_en:10.6} {tot_en:10.6}"
        else:
            pot_en = float("nan")

        if self.print_as_xyz:
            self.write_xyz_config(
                curr_step,
                frame,
                forces=frame.forces,
                stds=frame.stds,
                dft_forces=dft_forces,
                dft_energy=dft_energy,
                predict_energy=pot_en,
            )

        f = logging.getLogger(self.basename + "log")
        f.info(string)
        self.write_wall_time(start_time)

        # stat += f' {dt}\n'
        # logging.getLogger('stat').write(stat)

        if self.always_flush:
            f.handlers[0].flush() def predict_on_structure_par(
    structure: Structure,
    gp: GaussianProcess,
    n_cpus: int = None,
    write_to_structure: bool = True,
    selective_atoms: List[int] = None,
    skipped_atom_value=0,
) -> ("np.ndarray", "np.ndarray"):

    """
    Return the forces/std. dev. uncertainty associated with each
    individual atom in a structure. Forces are stored directly to the
    structure and are also returned.

    :param structure: FLARE structure to obtain forces for, with N atoms
    :param gp: Gaussian Process model
    :param n_cpus: Number of cores to parallelize over
    :param write_to_structure: Write results to structure's forces,
                            std attributes
    :param selective_atoms: Only predict on these atoms; e.g. [0,1,2] will
                                only predict and return for those atoms
    :param skipped_atom_value: What value to use for atoms that are skipped.
            Defaults to 0 but other options could be e.g. NaN. Will NOT
            write this to the structure if write_to_structure is True.
    :return: N x 3 array of forces, N x 3 array of uncertainties
    :rtype: (np.ndarray, np.ndarray)
    """
    # Just work in serial in the number of cpus is 1
    # or the gp is not parallelized by atoms
    if n_cpus == 1 or not gp.per_atom_par:
        return predict_on_structure(
            structure=structure,
            gp=gp,
            n_cpus=n_cpus,
            write_to_structure=write_to_structure,
            selective_atoms=selective_atoms,
            skipped_atom_value=skipped_atom_value,
        )

    forces = np.zeros(shape=(structure.nat, 3))
    stds = np.zeros(shape=(structure.nat, 3))

    if selective_atoms:
        forces.fill(skipped_atom_value)
        stds.fill(skipped_atom_value)
    else:
        selective_atoms = []

    # Automatically detect number of cpus available.
    if n_cpus is None:
        pool = mp.Pool(processes=mp.cpu_count())
    else:
        pool = mp.Pool(processes=n_cpus)

    # Parallelize over atoms in structure.
    results = []
    for atom in range(structure.nat):
        # If selective atoms is on, skip ones that was skipped.
        if atom not in selective_atoms and selective_atoms:
            # Keep length of results equal to nat.
            results.append(None)
            continue
        results.append(pool.apply_async(predict_on_atom, args=[(structure, atom, gp)]))
    pool.close()
    pool.join()

    for i in range(structure.nat):
        if i not in selective_atoms and selective_atoms:
            continue
        r = results[i].get()
        forces[i] = r[0]
        stds[i] = r[1]
        if write_to_structure:
            structure.forces[i] = r[0]
            structure.stds[i] = r[1]

    return forces, stds def predict_on_structure_par_en(
    structure: Structure,
    gp: GaussianProcess,
    n_cpus: int = None,
    write_to_structure: bool = True,
    selective_atoms: List[int] = None,
    skipped_atom_value=0,
) -> ("np.ndarray", "np.ndarray", "np.ndarray"):
    """
    Return the forces/std. dev. uncertainty / local energy associated with each
    individual atom in a structure, parallelized over atoms. Forces are
    stored directly to the structure and are also returned.

    :param structure: FLARE structure to obtain forces for, with N atoms
    :param gp: Gaussian Process model
    :param n_cpus: Number of cores to parallelize over
    :return: N x 3 array of forces, N x 3 array of uncertainties,
        N-length array of energies
    :rtype: (np.ndarray, np.ndarray, np.ndarray)
    """
    # Work in serial if the number of cpus is 1
    # or the gp is not parallelized by atoms
    if (n_cpus == 1) or (not gp.per_atom_par):
        predict_on_structure_en(
            structure,
            gp,
            n_cpus,
            write_to_structure,
            selective_atoms,
            skipped_atom_value,
        )

    forces = np.zeros((structure.nat, 3))
    stds = np.zeros((structure.nat, 3))
    local_energies = np.zeros(structure.nat)

    if selective_atoms:
        forces.fill(skipped_atom_value)
        stds.fill(skipped_atom_value)
        local_energies.fill(skipped_atom_value)
    else:
        selective_atoms = []

    if n_cpus is None:
        pool = mp.Pool(processes=mp.cpu_count())
    else:
        pool = mp.Pool(processes=n_cpus)

    # Parallelize over atoms in structure
    results = []
    for atom_i in range(structure.nat):

        if atom_i not in selective_atoms and selective_atoms:
            results.append(None)
            continue

        results.append(
            pool.apply_async(predict_on_atom_en, args=[(structure, atom_i, gp)])
        )
    pool.close()
    pool.join()

    # Compile results
    for i in range(structure.nat):

        if i not in selective_atoms and selective_atoms:
            continue

        r = results[i].get()
        forces[i][:] = r[0]
        stds[i][:] = r[1]
        local_energies[i] = r[2]

        if write_to_structure:
            structure.forces[i] = forces[i]
            structure.stds[i] = stds[i]

    return forces, stds, local_energies def predict_on_structure_mgp(
    structure: Structure,
    mgp: MappedGaussianProcess,
    output=None,
    output_name=None,
    n_cpus: int = None,
    write_to_structure: bool = True,
    selective_atoms: List[int] = None,
    skipped_atom_value: Union[float, int] = 0,
    energy: bool = False,
) -> Union[Tuple["np.ndarray", "np.ndarray", float], Tuple["np.ndarray", "np.ndarray"]]:
    """
    Assign forces to structure based on an mgp
    """
    if output and output_name:
        output.write_to_output("\npredict with mapping:\n", output_name)

    forces = np.zeros(shape=(structure.nat, 3))
    stds = np.zeros(shape=(structure.nat, 3))
    local_energy = np.zeros(shape=(structure.nat))

    if selective_atoms:
        forces.fill(skipped_atom_value)
        stds.fill(skipped_atom_value)
    else:
        selective_atoms = []

    for n in range(structure.nat):

        if n not in selective_atoms and selective_atoms:
            continue

        forces[n, :], stds[n, :], local_energy[n] = predict_on_atom_mgp(
            n, structure, mgp, write_to_structure
        )

    if energy:
        return forces, stds, local_energy
    else:
        return forces, stds class Structure:
    """
    Contains information about a periodic structure of atoms, including the
    periodic cell boundaries, atomic species, and coordinates.

    *Note that input positions are assumed to be Cartesian.*

    :param cell: 3x3 array whose rows are the Bravais lattice vectors of the
        cell.
    :type cell: np.ndarray
    :param species: List of atomic species, which are represented either as
        integers or chemical symbols.
    :type species: List
    :param positions: Nx3 array of atomic coordinates.
    :type positions: np.ndarray
    :param mass_dict: Dictionary of atomic masses used in MD simulations.
    :type mass_dict: dict
    :param prev_positions: Nx3 array of previous atomic coordinates used in
        MD simulations.
    :type prev_positions: np.ndarray
    :param species_labels: List of chemical symbols. Used in the output file
        of on-the-fly runs.
    :type species_labels: List[str]
    :param stds: Uncertainty associated with forces
    :type stds: np.ndarray
    """

    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,
    ):

        # Define cell (each row is a Bravais lattice vector).
        self.cell = np.array(cell)

        # Compute the max cutoff compatible with a 3x3x3 supercell of the
        # structure.
        self.max_cutoff = get_max_cutoff(self.cell)

        # Set positions.
        self.positions = np.array(positions)

        # If species are strings, convert species to integers by atomic number
        if species_labels is None:
            self.species_labels = species
        else:
            self.species_labels = species_labels
        self.coded_species = np.array([element_to_Z(spec) for spec in species])
        self.nat = len(species)

        # Default: atoms have no velocity
        if prev_positions is None:
            self.prev_positions = np.copy(self.positions)
        else:
            assert len(positions) == len(
                prev_positions
            ), "Previous positions and positions are not same length"
            self.prev_positions = prev_positions

        # Set forces, energies, and stresses and their uncertainties.
        if forces is not None:
            self.forces = np.array(forces)
        else:
            self.forces = np.zeros((len(positions), 3))

        if stds is not None:
            self.stds = np.array(stds)
        else:
            self.stds = np.zeros((len(positions), 3))

        self.energy = energy

        self.local_energies = None
        self.local_energy_stds = None
        self.partial_stresses = None
        self.partial_stress_stds = None
        self.stress = None
        self.stress_stds = None

        # Potential energy attribute needed to mirror ASE atoms object.
        self.potential_energy = None

        self.mass_dict = mass_dict

        # Convert from elements to atomic numbers in mass dict
        if mass_dict is not None:
            keys = list(mass_dict.keys())
            for elt in keys:
                if isinstance(elt, str):
                    mass_dict[element_to_Z(elt)] = mass_dict[elt]
                    if elt.isnumeric():
                        mass_dict[int(elt)] = mass_dict[elt]

    @property
    def positions(self):
        return self._positions

    @property
    def wrapped_positions(self):
        return self._wrapped_positions

    @positions.setter
    def positions(self, position_array):
        self._positions = position_array
        self._wrapped_positions = self.wrap_positions()

    @property
    def cell(self):
        return self._cell

    @property
    def vec1(self):
        return self._vec1

    @property
    def vec2(self):
        return self._vec2

    @property
    def vec3(self):
        return self._vec3

    @property
    def cell_transpose(self):
        return self._cell_transpose

    @property
    def cell_transpose_inverse(self):
        return self._cell_transpose_inverse

    @property
    def cell_dot(self):
        return self._cell_dot

    @property
    def cell_dot_inverse(self):
        return self._cell_dot_inverse

    @cell.setter
    def cell(self, cell_array):
        """Set the cell and related properties."""
        self._cell = cell_array
        self._vec1 = cell_array[0, :]
        self._vec2 = cell_array[1, :]
        self._vec3 = cell_array[2, :]
        self._cell_transpose = cell_array.transpose()
        self._cell_transpose_inverse = np.linalg.inv(self._cell_transpose)
        self._cell_dot = self.get_cell_dot(cell_array)
        self._cell_dot_inverse = np.linalg.inv(self._cell_dot)

    @staticmethod
    def get_cell_dot(cell_array):
        """
        Compute 3x3 array of dot products of cell vectors used to
        fold atoms back to the unit cell.

        :return: 3x3 array of cell vector dot products.
        :rtype: np.ndarray
        """

        cell_dot = np.zeros((3, 3))

        for m in range(3):
            for n in range(3):
                cell_dot[m, n] = np.dot(cell_array[m], cell_array[n])

        return cell_dot

    @staticmethod
    def raw_to_relative(
        positions: "ndarray", cell_transpose: "ndarray", cell_dot_inverse: "ndarray"
    ) -> "ndarray":
        """Convert Cartesian coordinates to relative (fractional) coordinates,
        expressed in terms of the cell vectors set in self.cell.

        :param positions: Cartesian coordinates.
        :type positions: np.ndarray
        :param cell_transpose: Transpose of the cell array.
        :type cell_transpose: np.ndarray
        :param cell_dot_inverse: Inverse of the array of dot products of
            cell vectors.
        :type cell_dot_inverse: np.ndarray
        :return: Relative positions.
        :rtype: np.ndarray
        """

        relative_positions = np.matmul(
            np.matmul(positions, cell_transpose), cell_dot_inverse
        )

        return relative_positions

    @staticmethod
    def relative_to_raw(
        relative_positions: "ndarray",
        cell_transpose_inverse: "ndarray",
        cell_dot: "ndarray",
    ) -> "ndarray":
        """Convert fractional coordinates to raw (Cartesian) coordinates.

        :param relative_positions: fractional coordinates.
        :type relative_positions: np.ndarray
        :param cell_transpose_inverse: Transpose of the cell array.
        :type cell_transpose_inverse: np.ndarray
        :param cell_dot: Dot products of cell vectors
        :type cell_dot: np.ndarray
        :return: Cartesian positions.
        :rtype: np.ndarray
        """

        return np.matmul(
            np.matmul(relative_positions, cell_dot), cell_transpose_inverse
        )

    def wrap_positions(self) -> "ndarray":
        """
        Convenience function which folds atoms outside of the unit cell back
        into the unit cell. in_place flag controls if the wrapped positions
        are set in the class.

        :return: Cartesian coordinates of positions all in unit cell
        :rtype: np.ndarray
        """
        rel_pos = self.raw_to_relative(
            self.positions, self.cell_transpose, self.cell_dot_inverse
        )

        rel_wrap = rel_pos - np.floor(rel_pos)

        pos_wrap = self.relative_to_raw(
            rel_wrap, self.cell_transpose_inverse, self.cell_dot
        )

        return pos_wrap

    def indices_of_specie(self, specie: Union[int, str]) -> List[int]:
        """
        Return the indices of a given species within atoms of the structure.

        :param specie: Element to target, can be string or integer
        :return: The indices in the structure at which this element occurs
        :rtype: List[str]
        """
        return [i for i, spec in enumerate(self.coded_species) if spec == specie]

    # TODO make more descriptive
    def __str__(self) -> str:
        """
        Simple descriptive string of structure.

        :return: One-line descriptor of number of atoms and species present.
        :rtype: str
        """

        return "Structure with {} atoms of types {}".format(
            self.nat, set(self.species_labels)
        )

    def __len__(self) -> int:
        """
        Returns number of atoms in structure.

        :return: number of atoms in structure.
        :rtype: int
        """
        return self.nat

    def as_dict(self) -> dict:
        """
        Returns structure as a dictionary; useful for serialization purposes.

        :return: Dictionary version of current structure
        :rtype: dict
        """
        return dict(vars(self))

    def as_str(self) -> str:
        """
        Returns string dictionary serialization cast as string.

        :return: output of as_dict method cast as string
        :rtype: str
        """
        return dumps(self.as_dict(), cls=NumpyEncoder)

    @staticmethod
    def from_dict(dictionary: dict) -> "flare.struc.Structure":
        """
        Assembles a Structure object from a dictionary parameterizing one.

        :param dictionary: dict describing structure parameters.
        :return: FLARE structure assembled from dictionary
        """
        struc = Structure(
            cell=np.array(dictionary.get("_cell", dictionary.get("cell"))),
            species=dictionary["coded_species"],
            positions=np.array(
                dictionary.get("_positions", dictionary.get("positions"))
            ),
            mass_dict=dictionary.get("mass_dict"),
            prev_positions=dictionary.get("prev_positions", None),
            species_labels=dictionary.get("species_labels"),
            forces=np.array(dictionary.get("forces")),
            stds=np.array(dictionary.get("stds")),
            energy=dictionary.get("energy", None),
        )

        struc.stress = dictionary.get("stress", None)

        return struc

    @staticmethod
    def from_ase_atoms(atoms: "ase.Atoms", cell=None) -> "flare.struc.Structure":
        """
        From an ASE Atoms object, return a FLARE structure

        :param atoms: ASE Atoms object
        :type atoms: ASE Atoms object
        :return: A FLARE structure from an ASE atoms object
        """

        if cell is None:
            cell = np.array(atoms.cell)

        try:
            forces = atoms.get_forces()
        except:
            forces = None

        try:
            stds = atoms.get_uncertainties()
        except:
            stds = None

        try:
            energy = atoms.get_potential_energy()
        except:
            energy = None

        try:
            stress = atoms.get_stress()  # TODO: what stress order should we use?
        except:
            stress = None

        struc = Structure(
            cell=cell,
            positions=atoms.positions,
            species=atoms.get_chemical_symbols(),
            forces=forces,
            stds=stds,
            energy=energy,
        )
        struc.stress = stress
        return struc

    def to_ase_atoms(self) -> "ase.Atoms":
        from ase import Atoms
        from ase.calculators.singlepoint import SinglePointCalculator

        atoms = Atoms(
            self.species_labels, positions=self.positions, cell=self.cell, pbc=True
        )

        results = {}
        properties = ["forces", "energy", "stress"]
        for p in properties:
            results[p] = getattr(self, p)
        atoms.calc = SinglePointCalculator(atoms, **results)

        return atoms

    def to_pmg_structure(self):
        """
        Returns FLARE structure as a pymatgen structure.

        :return: Pymatgen structure corresponding to current FLARE structure
        """

        if not _pmg_present:
            raise ModuleNotFoundError(
                "Pymatgen is not present. Please install Pymatgen and try again"
            )

        if self.forces is None:
            forces_temp = np.zeros((len(self.positions), 3))
            site_properties = {"force:": forces_temp, "std": self.stds}
        else:
            site_properties = {"force:": self.forces, "std": self.stds}

        return pmgstruc.Structure(
            lattice=self.cell,
            species=self.species_labels,
            coords=self.positions,
            coords_are_cartesian=True,
            site_properties=site_properties,
        )

    @staticmethod
    def from_pmg_structure(structure: "pymatgen Structure") -> "flare Structure":
        """
        Returns Pymatgen structure as FLARE structure.

        :param structure: Pymatgen Structure
        :type structure: Pymatgen Structure
        :return: FLARE Structure
        """

        cell = structure.lattice.matrix.copy()
        species = [str(spec) for spec in structure.species]
        positions = structure.cart_coords.copy()

        new_struc = Structure(cell=cell, species=species, positions=positions)

        site_props = structure.site_properties

        if "force" in site_props.keys():
            forces = site_props["force"]
            new_struc.forces = [np.array(force) for force in forces]

        if "std" in site_props.keys():
            stds = site_props["std"]
            new_struc.stds = [np.array(std) for std in stds]

        return new_struc

    def is_valid(self, tolerance: float = 0.5) -> bool:
        """
        Plugin to pymatgen's is_valid method to gauge if a structure
        has atoms packed too closely together, which is likely
        to cause unphysically large forces / energies or present convergence
        issues in DFT.
        :return:
        """
        pmg_structure = self.to_pmg_structure()
        return pmg_structure.is_valid(tol=tolerance)

    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:
        """
        Convenience function which turns a structure into an extended .xyz
        file; useful for further input into visualization programs like VESTA
        or Ovito. Can be saved to an output file via write_file.

        :param print_stds: Print the stds associated with the structure.
        :param print_forces:
        :param extended_xyz:
        :param print_max_stds:
        :param write_file:
        :return:
        """
        species_list = [Z_to_element(x) for x in self.coded_species]
        xyz_str = ""
        xyz_str += f"{len(self.coded_species)} \n"

        # Add header line with info about lattice and properties if extended
        #  xyz option is called.
        if extended_xyz:
            cell = self.cell

            xyz_str += f'Lattice="{cell[0,0]} {cell[0,1]} {cell[0,2]}'
            xyz_str += f" {cell[1,0]} {cell[1,1]} {cell[1,2]}"
            xyz_str += f' {cell[2,0]} {cell[2,1]} {cell[2,2]}"'
            if timestep > 0:
                xyz_str += f" Timestep={timestep}"
            if predict_energy:
                xyz_str += f" PE={predict_energy}"
            if dft_energy is not None:
                xyz_str += f" DFT_PE={dft_energy}"
            if target_atoms is not None:
                indices = " ".join([str(t) for t in target_atoms])
                xyz_str += f' target_atoms="{indices}"'
            xyz_str += f' Proprties="species:S:1:pos:R:3'

            if print_stds:
                xyz_str += ":stds:R:3"
                stds = self.stds
            if print_forces:
                xyz_str += ":forces:R:3"
                forces = self.forces
            if print_max_stds:
                xyz_str += ":max_std:R:1"
                stds = self.stds
            if print_energies:
                if self.local_energies is None:
                    print_energies = False
                else:
                    xyz_str += ":local_energy:R:1"
                    local_energies = self.local_energies
            if dft_forces is not None:
                xyz_str += ":dft_forces:R:3"
            xyz_str += "\n"
        else:
            xyz_str += "\n"

        for i, pos in enumerate(self.positions):
            # Write positions
            xyz_str += f"{species_list[i]} {pos[0]} {pos[1]} {pos[2]}"

            # If extended XYZ: Add in extra information
            if print_stds and extended_xyz:
                xyz_str += f" {stds[i,0]} {stds[i,1]} {stds[i,2]}"
            if print_forces and extended_xyz:
                xyz_str += f" {forces[i,0]} {forces[i,1]} {forces[i,2]}"
            if print_energies and extended_xyz:
                xyz_str += f" {local_energies[i]}"
            if print_max_stds and extended_xyz:
                xyz_str += f" {np.max(stds[i,:])} "
            if dft_forces is not None:
                xyz_str += f" {dft_forces[i, 0]} {dft_forces[i,1]} {dft_forces[i, 2]}"
            if i < (len(self.positions) - 1):
                xyz_str += "\n"

        # Write to file, optionally
        if write_file:
            if append:
                fmt = "a"
            else:
                fmt = "w"
            with open(write_file, fmt) as f:
                f.write(xyz_str)
                f.write("\n")

        return xyz_str

    @staticmethod
    def from_file(
        file_name: str, format: str = "", as_trajectory: bool = False
    ) -> Union["flare.struc.Structure", List["flare.struc.Structure"]]:
        """
        Load a FLARE structure from a file or a series of FLARE structures
        :param file_name:
        :param format:
        :return:
        """

        # Ensure the file specified exists.
        with open(file_name, "r") as _:
            pass

        if "pickle" in file_name or (
            "pickle" in format.lower() or "binary" in format.lower()
        ):
            with open(file_name, "rb") as f:
                return pickle.load(f)

        if "xyz" in file_name or "xyz" in format.lower():
            raise NotImplementedError

        if "json" in format.lower() or ".json" in file_name:
            # Assumed format is one FLARE structure per line,
            # or one line with many FLARE structures
            with open(file_name, "r") as f:
                thelines = f.readlines()

                non_empty_lines = [loads(line) for line in thelines if len(line) > 2]

            structures = [
                Structure.from_dict(struc_dict) for struc_dict in non_empty_lines
            ]

            if len(structures) == 1:
                return structures[0]
            else:
                return structures

        is_poscar = (
            "POSCAR" in file_name or "CONTCAR" in file_name or "vasp" in format.lower()
        )
        if is_poscar and _pmg_present:
            pmg_structure = pmgvaspio.Poscar.from_file(file_name).structure
            return Structure.from_pmg_structure(pmg_structure)
        elif is_poscar and not _pmg_present:
            raise ImportError("Pymatgen not imported; functionality requires pymatgen.")

    def __eq__(self, other):
        """
        Very crude equality measure that will check to see if positions,
        cell, and species are identical. Will not detect e.g. symmetry
        equivalence, translation, if one structure is a supercell of another,
        etc. Requires that species are in the same order.
        :param other:
        :return:
        """
        if not isinstance(other, Structure):
            return False

        if len(other) != len(self):
            return False
        if not all(other.coded_species == self.coded_species):
            return False
        if not np.allclose(other.positions, self.positions):
            return False
        if not np.allclose(other.cell, self.cell):
            return False

        return True class Trajectory(StructureSource, ForceSource):
    def __init__(
        self, frames: List[Structure] = None, iterate_strategy: Union[int, str] = 1
    ):
        if frames is None:
            frames = []
        self.frames = frames
        self.cur_idx = 0
        self.iterate_strategy = iterate_strategy
        self.seen_before = []

        if self.iterate_strategy == "shuffle":
            self.frames = np.random.shuffle(frames)
        if isinstance(iterate_strategy, int):
            self.frames = frames[::iterate_strategy]

    def get_next_structure(self) -> Union[Structure, None]:

        if self.cur_idx == len(self):
            self.cur_idx = 0
            return None

        cur_frame = self.frames[self.cur_idx]
        self.cur_idx += 1
        return cur_frame

    @property
    def cur_frame(self):
        return self.frames[self.cur_idx]

    @property
    def cur_forces(self):
        return self.frames[self.cur_idx].forces

    def get_next_force(self, index: int = -1) -> "np.ndarray":
        """
        Return the forces associated with a current structure,
        and if an index is passed, that atom.
        :param index:
        :return:
        """
        if index != 1:
            return self.cur_frame.forces[index]
        return self.cur_frame.forces

    def __getitem__(self, item):
        return self.frames[item]

    def __len__(self):
        return len(self.frames)

    def __next__(self):
        struc = self.get_next_structure()
        if struc is None:
            raise StopIteration()
        return struc

    def __iter__(self):
        self.cur_idx = 0
        return self

    def append(self, frame: Structure):
        self.frames.append(frame) def element_to_Z(element: str) -> int:
    """
    Returns the atomic number Z associated with an elements 1-2 letter name.
    Returns the same integer if an integer is passed in.

    :param element:
    :return:
    """

    if element in _user_element_to_Z:
        return _user_element_to_Z[element]

    # If already integer, do nothing
    if isinstance(element, (int, np.integer)):
        return element
    if type(element).__module__ == "numpy" and np.issubdtype(type(element), np.integer):
        return element

    # If a string-casted integer, do nothing
    if isinstance(element, str) and element.isnumeric():
        return int(element)

    # Check that a valid element was passed in then return
    if _element_to_Z.get(element, None) is None:
        warn(
            f"Element as specified not found in list of element-Z mappings. "
            f"If you would like to specify a custom element, use an integer"
            f" of your choosing instead. Setting element {{element}} to intege"
            f"r 0"
        )
    return _element_to_Z.get(element, 0) def Z_to_element(Z: int) -> str:
    """
    Maps atomic numbers Z to element name, e.g. 1->"H".

    :param Z: Atomic number corresponding to element.
    :return: One or two-letter name of element.
    """

    if Z in _user_Z_to_element:
        return _user_Z_to_element[Z]

    # Check proper formatting
    if isinstance(Z, str):
        if Z.isnumeric():
            Z = int(Z)
        else:
            raise ValueError("Input Z is not a number. It should be an integer")

    return _Z_to_element[Z] class NumpyEncoder(JSONEncoder):
    """
    Special json encoder for numpy types for serialization
    use as

    json.loads(... cls = NumpyEncoder)

    or:

    json.dumps(... cls = NumpyEncoder)

    Thanks to StackOverflow users karlB and fnunnari, who contributed this from:
    `https://stackoverflow.com/a/47626762`
    """

    def default(self, obj):
        """"""
        if isinstance(
            obj,
            (
                np.int_,
                np.intc,
                np.intp,
                np.int8,
                np.int16,
                np.int32,
                np.int64,
                np.uint8,
                np.uint16,
                np.uint32,
                np.uint64,
            ),
        ):
            return int(obj)
        elif isinstance(obj, (np.float_, np.float16, np.float32, np.float64)):
            return float(obj)
        elif isinstance(obj, (np.ndarray,)):
            return obj.tolist()
        return JSONEncoder.default(self, obj) def subset_of_frame_by_element(
    frame: "flare.Structure", predict_atoms_per_element: dict
) -> List[int]:
    """
    Given a structure and a dictionary formatted as {"Symbol":int,
    ..} describing a number of atoms per element, return a sorted list of
    indices corresponding to a random subset of atoms by species
    :param frame:
    :param predict_atoms_by_species:
    :return:
    """

    # Null case: No dictionary or empty dict passed in; just return all indices
    if not predict_atoms_per_element:
        return list(range(len(frame)))

    # Keep track of atoms which were considered (dictionary may only cover a
    #  subset of species of the whole frame)
    all_atoms = set(range(len(frame)))
    return_atoms = []
    considered_atoms = set([])

    species = frame.species_labels

    # Main loop: Obtain the number of relevant atoms for each element
    for elt, n in predict_atoms_per_element.items():

        matching_atoms = [i for i in all_atoms if species[i] == elt]
        considered_atoms.update(matching_atoms)

        if len(matching_atoms) == 0:
            continue
        # Choose the atoms to add
        to_add_atoms = np.random.choice(
            matching_atoms, replace=False, size=min(n, len(matching_atoms))
        )
        return_atoms += list(to_add_atoms)

    return_atoms += list(all_atoms - considered_atoms)

    return_atoms.sort()

    return return_atoms def is_std_in_bound_per_species(
    rel_std_tolerance: float,
    abs_std_tolerance: float,
    noise: float,
    structure: "flare.struc.Structure",
    max_atoms_added: int = inf,
    max_by_species: dict = {},
) -> (bool, List[int]):
    """
    Checks the stds of GP prediction assigned to the structure, returns a
    list of atoms which either meet an absolute threshold or a relative
    threshold defined by rel_std_tolerance * noise. Can limit the
    total number of target atoms via max_atoms_added, and limit per species
    by max_by_species.

    The max_atoms_added argument will 'overrule' the
    max by species; e.g. if max_atoms_added is 2 and max_by_species is {"H":3},
    then at most two atoms will be added.

    :param rel_std_tolerance: Multiplied by noise to get a lower
        bound for the uncertainty threshold defined relative to the model.
    :param abs_std_tolerance: Used as an absolute lower bound for the
        uncertainty threshold.
    :param noise: Noise hyperparameter for model, used to define relative
        uncertainty cutoff.
    :param structure: FLARE structure decorated with
        uncertainties in structure.stds.
    :param max_atoms_added: Maximum number of atoms to return from structure.
    :param max_by_species: Dictionary describing maximum number of atoms to
        return by species (e.g. {'H':1,'He':2} will return at most 1 H and 2 He
        atoms.)
    :return: Bool indicating if any atoms exceeded the uncertainty
        threshold, and a list of indices of atoms which did, sorted by their
        uncertainty.
    """

    # Always returns true; use this when you want to test model performance
    # without updating the training set.
    if rel_std_tolerance == 0 and abs_std_tolerance == 0:
        return True, [-1]

    # set uncertainty threshold based on if only one or the other is passed in,
    # and use the lower of the two.

    if rel_std_tolerance is None or rel_std_tolerance == 0:
        threshold = abs_std_tolerance
    elif abs_std_tolerance is None or abs_std_tolerance == 0:
        threshold = rel_std_tolerance * np.abs(noise)
    else:
        threshold = min(rel_std_tolerance * np.abs(noise), abs_std_tolerance)

    # Determine if any std component will trigger the threshold
    # before looking through individual species.
    max_std_components = [np.nanmax(std) for std in structure.stds]
    if np.nanmax(max_std_components) < threshold:
        return True, [-1]

    target_atoms = []

    # Sort from greatest to smallest max. std component
    std_arg_sorted = np.flip(np.argsort(max_std_components))

    present_species = {spec: 0 for spec in set(structure.species_labels)}

    # Loop through atoms and add until cutoffs are met.
    for i in std_arg_sorted:

        # If max atoms added reached or stds of atoms considered are now below
        # threshold, conclude
        if len(target_atoms) == max_atoms_added or (
            max_std_components[i] < threshold and max_std_components[i] != np.nan
        ):
            break

        if np.isnan(max_std_components[i]):
            continue

        # Only add up to species allowance, if it exists
        cur_spec = structure.species_labels[i]
        if present_species[cur_spec] < max_by_species.get(cur_spec, inf):
            target_atoms.append(i)
            present_species[cur_spec] += 1

    # Check in case that nothing was added, e.g. due to species limitations
    if len(target_atoms):
        return False, target_atoms

    return True, [-1] def is_force_in_bound_per_species(
    abs_force_tolerance: float,
    predicted_forces: "ndarray",
    label_forces: "ndarray",
    structure,
    max_atoms_added: int = inf,
    max_by_species: dict = {},
    max_force_error: float = inf,
) -> (bool, List[int]):
    """
    Checks the forces of GP prediction assigned to the structure against a
    DFT calculation, and return a list of atoms which meet an absolute
    threshold abs_force_tolerance.

    Can limit the total number of target atoms via max_atoms_added, and limit
    per species by max_by_species.

    The max_atoms_added argument will 'overrule' the
    max by species; e.g. if max_atoms_added is 2 and max_by_species is {"H":3},
    then at most two atoms total will be added.

    Because adding atoms which are in configurations which are far outside
    of the potential energy surface may not always be
    desirable, a maximum force error can be passed in; atoms with

    :param abs_force_tolerance: If error exceeds this value, then return
        atom index
    :param predicted_forces: Force predictions made by GP model
    :param label_forces: "True" forces computed by DFT
    :param structure: FLARE Structure
    :param max_atoms_added: Maximum atoms to return
    :param max_by_species: Limit to a maximum number of atoms by species
    :param max_force_error: In order to avoid counting in highly unlikely
        configurations, if the error exceeds this, do not add atom
    :return: Bool indicating if any atoms exceeded the error
        threshold, and a list of indices of atoms which did sorted by their
        error.
    """

    # Always returns true; use this when you want to test model performance
    # without updating the training set.
    if abs_force_tolerance == 0:
        return True, [-1]

    errors = np.abs(predicted_forces - label_forces)

    # Determine if any force component will trigger the threshold
    max_error_components = np.amax(errors, axis=1)
    if np.nanmax(max_error_components) < abs_force_tolerance:
        return True, [-1]

    target_atoms = []

    # Sort from greatest to smallest error
    force_arg_sorted = np.flip(np.argsort(max_error_components))

    present_species = {spec: 0 for spec in set(structure.species_labels)}

    # Only add atoms up to the bound
    for i in force_arg_sorted:

        # If max atoms added reached or force errors are now below threshold,
        # conclude
        if len(target_atoms) == max_atoms_added or (
            max_error_components[i] < abs_force_tolerance
            and max_error_components[i] != np.nan
        ):
            break

        cur_spec = structure.species_labels[i]

        # Only add up to species allowance, if it exists
        if (
            present_species[cur_spec] < max_by_species.get(cur_spec, inf)
            and max_error_components[i] < max_force_error
        ):
            target_atoms.append(i)
            present_species[cur_spec] += 1

    # Check in case that nothing was added e.g. due to species or error
    # limitations
    if len(target_atoms):
        return False, target_atoms
    else:
        return True, [-1] def evaluate_training_atoms(
    pred_forces: "np.ndarray" = None,
    dft_forces: "np.ndarray" = None,
    rel_std_tolerance: float = 4,
    abs_std_tolerance: float = 0,
    noise: float = 0,
    abs_force_tolerance: float = 0.15,
    max_force_error: float = inf,
    structure: "flare.struc.Structure" = None,
    max_atoms_from_frame: int = None,
    max_elts_per_frame: dict = None,
    max_model_elts: dict = None,
    training_statistics: dict = None,
):
    # Set max elements per frame based on model size.
    # E.g. if model will have at most 100 Carbon atoms,
    # and 5 carbon atoms per frame are allowed,
    # and the GP currently has 96,
    # set the next max Carbon atoms to 4.
    max_atoms_by_elt = {}
    if max_model_elts and training_statistics:
        for key, val in max_model_elts.items():
            max_atoms_by_elt[key] = val - training_statistics["envs_by_species"].get(
                key, 0
            )
            max_atoms_by_elt[key] = max(max_atoms_by_elt[key], 0)
    if max_elts_per_frame:
        for key, val in max_elts_per_frame.items():
            max_atoms_by_elt[key] = min(max_atoms_by_elt.get(key, inf), val)
    if not max_atoms_by_elt:
        for spec in structure.species_labels:
            max_atoms_by_elt[spec] = inf

    std_in_bound, std_train_atoms = is_std_in_bound_per_species(
        rel_std_tolerance=rel_std_tolerance,
        abs_std_tolerance=abs_std_tolerance,
        noise=noise,
        structure=structure,
        max_atoms_added=max_atoms_from_frame,
        max_by_species=max_atoms_by_elt,
    )

    # Get max force error atoms
    if not dft_forces is None:
        force_in_bound, force_train_atoms = is_force_in_bound_per_species(
            abs_force_tolerance=abs_force_tolerance,
            predicted_forces=pred_forces,
            label_forces=dft_forces,
            structure=structure,
            max_atoms_added=max_atoms_from_frame,
            max_by_species=max_atoms_by_elt,
            max_force_error=max_force_error,
        )
    else:
        force_in_bound = True
        force_train_atoms = {-1}

    in_bound = std_in_bound and force_in_bound

    train_atoms = list(set(force_train_atoms).union(std_train_atoms) - {-1})

    return in_bound, train_atoms 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"] 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
import json as json
import logging
import numpy as np
import time
import warnings
from copy import deepcopy
from math import inf
from typing import List, Tuple, Union, Dict
from random import sample
from flare.env import AtomicEnvironment
from flare.gp import GaussianProcess
from flare.output import Output
from flare.predict import (
    predict_on_structure_par,
    predict_on_structure_par_en,
    predict_on_structure_mgp,
)
from flare.struc import Structure, Trajectory
from flare.utils.element_coder import element_to_Z, Z_to_element, NumpyEncoder
from flare.utils.learner import (
    subset_of_frame_by_element,
    is_std_in_bound_per_species,
    is_force_in_bound_per_species,
    evaluate_training_atoms,
)
from flare.mgp import MappedGaussianProcess
from flare.parameters import Parameters
            self.output.write_gp_dft_comparison(
                curr_step=i,
                frame=dummy_frame,
                start_time=frame_start_time,
                dft_forces=dft_forces,
                dft_energy=dft_energy,
                error=error,
                local_energies=local_energies,
                KE=0,
                cell=cur_frame.cell,
            )

            logger.debug(
                f"Single frame calculation time {time.time()-frame_start_time}"
            )

            if i < train_frame:
                # Noise hyperparameter & relative std tolerance is not for gp_is_mapped.
                if self.gp_is_mapped:
                    noise = 0
                else:
                    noise = Parameters.get_noise(
                        self.gp.hyps_mask, self.gp.hyps, constraint=False
                    )
                std_in_bound, std_train_atoms = is_std_in_bound_per_species(
                    rel_std_tolerance=self.rel_std_tolerance,
                    abs_std_tolerance=self.abs_std_tolerance,
                    noise=noise,
                    structure=dummy_frame,
                    max_atoms_added=self.max_atoms_from_frame,
                    max_by_species=self.train_env_per_species,
                )

                # Get max force error atoms
                force_in_bound, force_train_atoms = is_force_in_bound_per_species(
                    abs_force_tolerance=self.abs_force_tolerance,
                    predicted_forces=pred_forces,
                    label_forces=dft_forces,
                    structure=dummy_frame,
                    max_atoms_added=self.max_atoms_from_frame,
                    max_by_species=self.train_env_per_species,
                    max_force_error=self.max_force_error,
                )

                if not std_in_bound or not force_in_bound:

                    # -1 is returned from the is_in_bound methods,
                    # so filter that out and the use sets to remove repeats
                    train_atoms = list(
                        set(std_train_atoms).union(force_train_atoms) - {-1}
                    )

                    # Record frame and training atoms, uncertainty, error
                    force_errors = list(np.abs(pred_forces - dft_forces))
                    uncertainties = list(dummy_frame.stds)
                    training_plan[int(i)] = [
                        (int(a), uncertainties[a], force_errors[a]) for a in train_atoms
                    ]

                    # Compute mae and write to output;
                    # Add max uncertainty atoms to training set
                    self.update_gp_and_print(
                        cur_frame,
                        train_atoms=train_atoms,
                        uncertainties=pred_stds[train_atoms],
                        train=False,
                    )
                    cur_atoms_added_train += len(train_atoms)
                    cur_atoms_added_write += len(train_atoms)
                    # Re-train if number of sampled atoms is high enough

                    if (
                        cur_atoms_added_train >= self.min_atoms_per_train
                        or (i + 1) == train_frame
                    ):
                        if self.train_count < self.max_trains:
                            self.train_gp()
                            cur_trains_done_write += 1
                        else:
                            self.gp.update_L_alpha()
                        cur_atoms_added_train = 0
                    else:
                        self.gp.update_L_alpha()

                    # Loop to decide of a model should be written this
                    # iteration
                    will_write = False

                    if (
                        self.train_checkpoint_interval
                        and cur_trains_done_write
                        and self.train_checkpoint_interval <= cur_trains_done_write
                    ):
                        will_write = True
                        cur_trains_done_write = 0

                    if (
                        self.atom_checkpoint_interval
                        and cur_atoms_added_write
                        and self.atom_checkpoint_interval <= cur_atoms_added_write
                    ):
                        will_write = True
                        cur_atoms_added_write = 0

                    if self.model_format and will_write:
                        self.gp.write_model(
                            f"{self.output_name}_checkpt", self.model_format
                        )

                if (i + 1) == train_frame and not self.gp_is_mapped:
                    self.gp.check_L_alpha()

        # Print training statistics for GP model used
        conclusion_strings = [
            "Final GP statistics:" + json.dumps(self.gp.training_statistics)
        ]
        self.output.conclude_run(conclusion_strings)

        if self.print_training_plan:
            with open(f"{self.output_name}_training_plan.json", "w") as f:

                f.write(json.dumps(training_plan, cls=NumpyEncoder))