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] def queue_wrapper(result_queue, wid, func, args):
def partition_matrix(n_sample, size, n_cpus):
def partition_matrix_custom(n_sample: int, start1, end1, start2, end2, n_cpus):
def partition_vector(n_sample, size, n_cpus):
def partition_force_energy_block(n_sample: int, size1: int, size2: int, n_cpus: int):
def partition_update(n_sample, size, old_size, n_cpus):
def obtain_noise_len(hyps, hyps_mask):
def parallel_matrix_construction(
    pack_function,
    hyps,
    name,
    kernel,
    cutoffs,
    hyps_mask,
    block_id,
    nbatch,
    size1,
    size2,
    m1,
    m2,
    symm=True,
):
def parallel_vector_construction(
    pack_function,
    name,
    x,
    kernel,
    hyps,
    cutoffs,
    hyps_mask,
    block_id,
    nbatch,
    size,
    mult,
    d_1=None,
):
def multiple_array_construction(
    pack_function,
    name,
    x,
    kernel,
    hyps,
    cutoffs,
    hyps_mask,
    block_id,
    nbatch,
    size,
    mult,
    array_sizes,
):
def get_distance_mat_pack(
    hyps: np.ndarray,
    name: str,
    s1: int,
    e1: int,
    s2: int,
    e2: int,
    same: bool,
    kernel,
    cutoffs,
    hyps_mask,
):
def kernel_distance_mat(
    hyps: np.ndarray,
    name: str,
    kernel,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
def get_force_block_pack(
    hyps: np.ndarray,
    name: str,
    s1: int,
    e1: int,
    s2: int,
    e2: int,
    same: bool,
    kernel,
    cutoffs,
    hyps_mask,
):
def get_energy_block_pack(
    hyps: np.ndarray,
    name: str,
    s1: int,
    e1: int,
    s2: int,
    e2: int,
    same: bool,
    kernel,
    cutoffs,
    hyps_mask,
):
def get_force_energy_block_pack(
    hyps: np.ndarray,
    name: str,
    s1: int,
    e1: int,
    s2: int,
    e2: int,
    same: bool,
    kernel,
    cutoffs,
    hyps_mask,
):
def get_force_block(
    hyps: np.ndarray,
    name: str,
    kernel,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
def get_energy_block(
    hyps: np.ndarray,
    name: str,
    kernel,
    energy_noise,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
def get_force_energy_block(
    hyps: np.ndarray,
    name: str,
    kernel,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
def get_Ky_mat(
    hyps: np.ndarray,
    name: str,
    force_kernel: Callable,
    energy_kernel: Callable,
    force_energy_kernel: Callable,
    energy_noise,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
def update_force_block(
    ky_mat_old: np.ndarray,
    n_envs_prev: int,
    hyps: np.ndarray,
    name: str,
    kernel,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
def update_energy_block(
    ky_mat_old: np.ndarray,
    n_envs_prev: int,
    hyps: np.ndarray,
    name: str,
    kernel,
    energy_noise: float,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
def update_force_energy_block(
    ky_mat_old: np.ndarray,
    n_envs_prev: int,
    hyps: np.ndarray,
    name: str,
    kernel,
    energy_noise: float,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
def get_ky_mat_update(
    ky_mat_old: np.ndarray,
    n_envs_prev: int,
    hyps: np.ndarray,
    name: str,
    force_kernel: Callable,
    energy_kernel: Callable,
    force_energy_kernel: Callable,
    energy_noise: float,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
def energy_energy_vector_unit(
    name, s, e, x, kernel, hyps, cutoffs=None, hyps_mask=None, d_1=None
):
def energy_force_vector_unit(
    name, s, e, x, kernel, hyps, cutoffs=None, hyps_mask=None, d_1=None
):
def force_energy_vector_unit(name, s, e, x, kernel, hyps, cutoffs, hyps_mask, d_1):
def force_force_vector_unit(name, s, e, x, kernel, hyps, cutoffs, hyps_mask, d_1):
def efs_force_vector_unit(name, s, e, x, efs_force_kernel, hyps, cutoffs, hyps_mask):
def efs_energy_vector_unit(name, s, e, x, efs_energy_kernel, hyps, cutoffs, hyps_mask):
def energy_energy_vector(
    name, kernel, x, hyps, cutoffs=None, hyps_mask=None, n_cpus=1, n_sample=100
):
def energy_force_vector(
    name, kernel, x, hyps, cutoffs=None, hyps_mask=None, n_cpus=1, n_sample=100
):
def force_energy_vector(
    name, kernel, x, d_1, hyps, cutoffs=None, hyps_mask=None, n_cpus=1, n_sample=100
):
def force_force_vector(
    name, kernel, x, d_1, hyps, cutoffs=None, hyps_mask=None, n_cpus=1, n_sample=100
):
def efs_force_vector(
    name,
    efs_force_kernel,
    x,
    hyps,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
def efs_energy_vector(
    name,
    efs_energy_kernel,
    x,
    hyps,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
def get_kernel_vector(
    name,
    force_force_kernel: Callable,
    force_energy_kernel: Callable,
    x,
    d_1,
    hyps,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
def en_kern_vec(
    name,
    energy_force_kernel,
    energy_energy_kernel,
    x,
    hyps,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
def efs_kern_vec(
    name,
    efs_force_kernel,
    efs_energy_kernel,
    x,
    hyps,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
def get_ky_and_hyp_pack(
    name,
    s1,
    e1,
    s2,
    e2,
    same: bool,
    hyps: np.ndarray,
    kernel_grad,
    cutoffs=None,
    hyps_mask=None,
):
def get_ky_and_hyp(
    hyps: np.ndarray,
    name,
    kernel_grad,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
def get_like_from_mats(ky_mat, l_mat, alpha, name):
def get_neg_like(
    hyps: np.ndarray,
    name: str,
    force_kernel,
    logger_name=None,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
def get_neg_like_grad(
    hyps: np.ndarray,
    name: str,
    kernel_grad,
    logger_name: str = None,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
def get_like_grad_from_mats(ky_mat, hyp_mat, name): def str_to_kernel_set(
    kernels: list = ["twobody", "threebody"],
    component: str = "mc",
    hyps_mask: dict = None,
):
    """
    return kernels and kernel gradient function base on a string.
    If it contains 'sc', it will use the kernel in sc module;
    otherwise, it uses the kernel in mc_simple;
    if sc is not included and multihyps is True,
    it will use the kernel in mc_sephyps module
    otherwise, it will use the kernel in the sc module

    Args:

    name (str): name for kernels. example: "2+3mc"
    multihyps (bool, optional): True for using multiple hyperparameter groups

    :return: kernel function, kernel gradient, energy kernel,
             energy_and_force kernel

    """

    #  kernel name should be replace with kernel array
    if component == "sc":
        stk = sc._str_to_kernel
    else:
        multihyps = True
        if hyps_mask is None:
            multihyps = False
        elif hyps_mask["nspecie"] == 1:
            multihyps = False
        if multihyps:
            stk = mc_sephyps._str_to_kernel
        else:
            stk = mc_simple._str_to_kernel

    # b2 = Two body in use, b3 = Three body in use
    str_terms = {
        "2": ["2", "two", "twobody"],
        "3": ["3", "three", "threebody"],
        "many": ["mb", "manybody", "many"],
    }

    if isinstance(kernels, str):
        kernels = [kernels]

    prefix = ""
    for term in str_terms:
        add = False
        for s in str_terms[term]:
            for k in kernels:
                if s in k.lower():
                    add = True
        if add:
            if len(prefix) > 0:
                prefix += "+"
            prefix += term

    if len(prefix) == 0:
        raise RuntimeError(f"the name has to include at least one number {kernels}")

    for suffix in [
        "",
        "_grad",
        "_en",
        "_force_en",
        "_efs_energy",
        "_efs_force",
        "_efs_self",
    ]:
        if prefix + suffix not in stk:
            raise RuntimeError(f"cannot find kernel function of {prefix}{suffix}")

    return (
        stk[prefix],
        stk[prefix + "_grad"],
        stk[prefix + "_en"],
        stk[prefix + "_force_en"],
        stk[prefix + "_efs_energy"],
        stk[prefix + "_efs_force"],
        stk[prefix + "_efs_self"],
    ) def from_mask_to_args(hyps, cutoffs, hyps_mask=None):
    """Return the tuple of arguments needed for kernel function.
    The order of the tuple has to be exactly the same as the one taken by
        the kernel function.

    :param hyps: list of hyperparmeter values
    :type hyps: nd.array
    :param hyps_mask: all the remaining parameters needed
    :type hyps_mask: dictionary
    :param cutoffs: cutoffs used

    :return: args
    """

    # no special setting
    multihyps = True
    if hyps_mask is None:
        multihyps = False
    elif hyps_mask["nspecie"] == 1:
        multihyps = False

    if not multihyps:

        cutoffs_array = [0, 0, 0]
        cutoffs_array[0] = cutoffs.get("twobody", 0)
        cutoffs_array[1] = cutoffs.get("threebody", 0)
        cutoffs_array[2] = cutoffs.get("manybody", 0)
        return (hyps, cutoffs_array)

    # setting for mc_sephyps
    nspecie = hyps_mask["nspecie"]
    n2b = hyps_mask.get("ntwobody", 0)

    n3b = hyps_mask.get("nthreebody", 0)
    nmanybody = hyps_mask.get("nmanybody", 0)
    ncut3b = hyps_mask.get("ncut3b", 0)

    twobody_mask = hyps_mask.get("twobody_mask", None)
    threebody_mask = hyps_mask.get("threebody_mask", None)
    manybody_mask = hyps_mask.get("manybody_mask", None)
    cut3b_mask = hyps_mask.get("cut3b_mask", None)

    # TO DO , should instead use the non-sephyps kernel
    if n2b == 1:
        twobody_mask = np.zeros(nspecie ** 2, dtype=int)
    if n3b == 1:
        threebody_mask = np.zeros(nspecie ** 3, dtype=int)
    if nmanybody == 1:
        manybody_mask = np.zeros(nspecie ** 2, dtype=int)

    cutoff_2b = cutoffs.get("twobody", 0)
    cutoff_3b = cutoffs.get("threebody", 0)
    cutoff_mb = cutoffs.get("manybody", 0)

    if "bond_cutoff_list" in hyps_mask:
        cutoff_2b = hyps_mask["bond_cutoff_list"]
    else:
        cutoff_2b = np.ones(nspecie ** 2, dtype=float) * cutoff_2b

    if "threebody_cutoff_list" in hyps_mask:
        cutoff_3b = hyps_mask["threebody_cutoff_list"]
    if "manybody_cutoff_list" in hyps_mask:
        cutoff_mb = hyps_mask["manybody_cutoff_list"]

    (sig2, ls2) = Parameters.get_component_hyps(hyps_mask, "twobody", hyps=hyps)
    (sig3, ls3) = Parameters.get_component_hyps(hyps_mask, "threebody", hyps=hyps)
    (sigm, lsm) = Parameters.get_component_hyps(hyps_mask, "manybody", hyps=hyps)

    return (
        cutoff_2b,
        cutoff_3b,
        cutoff_mb,
        nspecie,
        np.array(hyps_mask["species_mask"]),
        n2b,
        twobody_mask,
        n3b,
        threebody_mask,
        ncut3b,
        cut3b_mask,
        nmanybody,
        manybody_mask,
        sig2,
        ls2,
        sig3,
        ls3,
        sigm,
        lsm,
    ) def kernel_str_to_array(kernel_name: str):
    """
    Args:

    name (str): name for kernels. example: "2+3mc"

    :return: kernel function, kernel gradient, energy kernel,
             energy_and_force kernel
    """

    #  kernel name should be replace with kernel array
    str_terms = {
        "twobody": ["2", "two", "twobody"],
        "threebody": ["3", "three", "threebody"],
        "manybody": ["mb", "manybody", "many"],
    }

    array = []
    for term in str_terms:
        add = False
        for s in str_terms[term]:
            if s in kernel_name.lower():
                add = True
        if add:
            array += [term]
    return array 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 set_logger(
    name: str, stream: bool, fileout_name: str = None, verbose: str = "info"
):
    """
    set up a logger with handlers

    :param name: unique name of the logger in logging module
    :type name: str
    :param stream: if True, set up a screen output
    :type stream: bool
    :param fileout_name: name for log file
    :type fileout_name: str
    :param verbose: verbose level
    :type verbose: str
    """
    logger = logging.getLogger(name)
    logger.propagate = False
    logger.handlers = []
    logger.setLevel(getattr(logging, verbose.upper()))
    if stream:
        add_stream(logger, verbose)
    if fileout_name is not None:
        add_file(logger, fileout_name, verbose)
    return logger class Parameters:
    """"""

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

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

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

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

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

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

        if isinstance(cutoffs, dict):
            return cutoffs

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

    @staticmethod
    def backward(kernels, param_dict):

        if param_dict is None:
            param_dict = {}

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

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

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

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

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

        return param_dict

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

        :return:
        """

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

        param_dict["cutoffs"] = cutoffs

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

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

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

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

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

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

            if n > 1:

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

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

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

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

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

        if "map" in param_dict:

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

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

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

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

        hyps = Parameters.get_hyps(param_dict, hyps)

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

        return param_dict

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

        Args:

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

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

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

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

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

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

        if constraint:
            return newhyps, newopt
        else:
            return newhyps

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

        Args:

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

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

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

            new_dict[kernel_name + "_start"] = 0

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

            if kernel_name in Parameters.cutoff_types_values:

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

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

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

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

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

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

        Args:

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

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

        Args:

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

        """

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

        if f"{kernel_name}_cutoff_list" in param_dict:

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

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

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

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

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

        Args:

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

        """

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

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

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

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

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

        if dict1 is None:
            return True

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

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

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

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

        return True class 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 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 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]
import json
import logging
import multiprocessing as mp
import pickle
import time
import numpy as np
from collections import Counter
from copy import deepcopy
from typing import List, Union, Tuple
from flare.env import AtomicEnvironment
from flare.gp_algebra import (
    get_like_from_mats,
    get_neg_like_grad,
    get_ky_mat_update,
    _global_training_data,
    _global_training_labels,
    _global_training_structures,
    _global_energy_labels,
    get_Ky_mat,
    get_kernel_vector,
    en_kern_vec,
    efs_kern_vec,
)
from flare.kernels.utils import (
    str_to_kernel_set,
    from_mask_to_args,
    kernel_str_to_array,
)
from flare.output import Output, set_logger
from flare.parameters import Parameters
from flare.struc import Structure
from flare.utils.element_coder import NumpyEncoder, Z_to_element
from numpy.random import random
from scipy.linalg import solve_triangular
from scipy.optimize import minimize
            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