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 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"],
    ) 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 def two_plus_three_body_mc(
    env1: AtomicEnvironment,
    env2: AtomicEnvironment,
    d1: int,
    d2: int,
    hyps: "ndarray",
    cutoffs: "ndarray",
    cutoff_func: Callable = cf.quadratic_cutoff,
) -> float:
    """2+3-body multi-element kernel between two force components.

    Args:
        env1 (AtomicEnvironment): First local environment.
        env2 (AtomicEnvironment): Second local environment.
        d1 (int): Force component of the first environment.
        d2 (int): Force component of the second environment.
        hyps (np.ndarray): Hyperparameters of the kernel function (sig1, ls1,
            sig2, ls2, sig_n).
        cutoffs (np.ndarray): Two-element array containing the 2- and 3-body
            cutoffs.
        cutoff_func (Callable): Cutoff function of the kernel.

    Return:
        float: Value of the 2+3-body kernel.
    """

    sig2 = hyps[0]
    ls2 = hyps[1]
    sig3 = hyps[2]
    ls3 = hyps[3]
    r_cut_2 = cutoffs[0]
    r_cut_3 = cutoffs[1]

    two_term = two_body_mc_jit(
        env1.bond_array_2,
        env1.ctype,
        env1.etypes,
        env2.bond_array_2,
        env2.ctype,
        env2.etypes,
        d1,
        d2,
        sig2,
        ls2,
        r_cut_2,
        cutoff_func,
    )

    three_term = three_body_mc_jit(
        env1.bond_array_3,
        env1.ctype,
        env1.etypes,
        env2.bond_array_3,
        env2.ctype,
        env2.etypes,
        env1.cross_bond_inds,
        env2.cross_bond_inds,
        env1.cross_bond_dists,
        env2.cross_bond_dists,
        env1.triplet_counts,
        env2.triplet_counts,
        d1,
        d2,
        sig3,
        ls3,
        r_cut_3,
        cutoff_func,
    )

    return two_term + three_term def two_plus_three_body_mc_grad(
    env1: AtomicEnvironment,
    env2: AtomicEnvironment,
    d1: int,
    d2: int,
    hyps: "ndarray",
    cutoffs: "ndarray",
    cutoff_func: Callable = cf.quadratic_cutoff,
) -> ("float", "ndarray"):
    """2+3-body multi-element kernel between two force components and its
    gradient with respect to the hyperparameters.

    Args:
        env1 (AtomicEnvironment): First local environment.
        env2 (AtomicEnvironment): Second local environment.
        d1 (int): Force component of the first environment.
        d2 (int): Force component of the second environment.
        hyps (np.ndarray): Hyperparameters of the kernel function (sig1, ls1,
            sig2, ls2, sig_n).
        cutoffs (np.ndarray): Two-element array containing the 2- and 3-body
            cutoffs.
        cutoff_func (Callable): Cutoff function of the kernel.

    Return:
        (float, np.ndarray):
            Value of the 2+3-body kernel and its gradient
            with respect to the hyperparameters.
    """

    sig2 = hyps[0]
    ls2 = hyps[1]
    sig3 = hyps[2]
    ls3 = hyps[3]
    r_cut_2 = cutoffs[0]
    r_cut_3 = cutoffs[1]

    kern2, grad2 = two_body_mc_grad_jit(
        env1.bond_array_2,
        env1.ctype,
        env1.etypes,
        env2.bond_array_2,
        env2.ctype,
        env2.etypes,
        d1,
        d2,
        sig2,
        ls2,
        r_cut_2,
        cutoff_func,
    )

    kern3, grad3 = three_body_mc_grad_jit(
        env1.bond_array_3,
        env1.ctype,
        env1.etypes,
        env2.bond_array_3,
        env2.ctype,
        env2.etypes,
        env1.cross_bond_inds,
        env2.cross_bond_inds,
        env1.cross_bond_dists,
        env2.cross_bond_dists,
        env1.triplet_counts,
        env2.triplet_counts,
        d1,
        d2,
        sig3,
        ls3,
        r_cut_3,
        cutoff_func,
    )

    return kern2 + kern3, np.array([grad2[0], grad2[1], grad3[0], grad3[1]]) def two_plus_three_mc_en(
    env1: AtomicEnvironment,
    env2: AtomicEnvironment,
    hyps: "ndarray",
    cutoffs: "ndarray",
    cutoff_func: Callable = cf.quadratic_cutoff,
) -> float:
    """2+3-body multi-element kernel between two local energies.

    Args:
        env1 (AtomicEnvironment): First local environment.
        env2 (AtomicEnvironment): Second local environment.
        hyps (np.ndarray): Hyperparameters of the kernel function (sig1, ls1,
            sig2, ls2).
        cutoffs (np.ndarray): Two-element array containing the 2- and 3-body
            cutoffs.
        cutoff_func (Callable): Cutoff function of the kernel.

    Return:
        float: Value of the 2+3-body energy/energy kernel.
    """

    sig2 = hyps[0]
    ls2 = hyps[1]
    sig3 = hyps[2]
    ls3 = hyps[3]
    r_cut_2 = cutoffs[0]
    r_cut_3 = cutoffs[1]

    # TODO: Move fractional factor to the njit function.
    two_term = (
        two_body_mc_en_jit(
            env1.bond_array_2,
            env1.ctype,
            env1.etypes,
            env2.bond_array_2,
            env2.ctype,
            env2.etypes,
            sig2,
            ls2,
            r_cut_2,
            cutoff_func,
        )
        / 4
    )

    # TODO: Move fractional factor to the njit function.
    three_term = (
        three_body_mc_en_jit(
            env1.bond_array_3,
            env1.ctype,
            env1.etypes,
            env2.bond_array_3,
            env2.ctype,
            env2.etypes,
            env1.cross_bond_inds,
            env2.cross_bond_inds,
            env1.cross_bond_dists,
            env2.cross_bond_dists,
            env1.triplet_counts,
            env2.triplet_counts,
            sig3,
            ls3,
            r_cut_3,
            cutoff_func,
        )
        / 9
    )

    return two_term + three_term def two_plus_three_mc_force_en(
    env1: AtomicEnvironment,
    env2: AtomicEnvironment,
    d1: int,
    hyps: "ndarray",
    cutoffs: "ndarray",
    cutoff_func: Callable = cf.quadratic_cutoff,
) -> float:
    """2+3-body multi-element kernel between a force component and a local
    energy.

    Args:
        env1 (AtomicEnvironment): Local environment associated with the
            force component.
        env2 (AtomicEnvironment): Local environment associated with the
            local energy.
        d1 (int): Force component of the first environment.
        hyps (np.ndarray): Hyperparameters of the kernel function (sig1, ls1,
            sig2, ls2).
        cutoffs (np.ndarray): Two-element array containing the 2- and 3-body
            cutoffs.
        cutoff_func (Callable): Cutoff function of the kernel.

    Return:
        float: Value of the 2+3-body force/energy kernel.
    """

    sig2 = hyps[0]
    ls2 = hyps[1]
    sig3 = hyps[2]
    ls3 = hyps[3]
    r_cut_2 = cutoffs[0]
    r_cut_3 = cutoffs[1]

    # TODO: Move fractional factor to the njit function.
    two_term = (
        two_body_mc_force_en_jit(
            env1.bond_array_2,
            env1.ctype,
            env1.etypes,
            env2.bond_array_2,
            env2.ctype,
            env2.etypes,
            d1,
            sig2,
            ls2,
            r_cut_2,
            cutoff_func,
        )
        / 2
    )

    # TODO: Move fractional factor to the njit function.
    three_term = (
        three_body_mc_force_en_jit(
            env1.bond_array_3,
            env1.ctype,
            env1.etypes,
            env2.bond_array_3,
            env2.ctype,
            env2.etypes,
            env1.cross_bond_inds,
            env2.cross_bond_inds,
            env1.cross_bond_dists,
            env2.cross_bond_dists,
            env1.triplet_counts,
            env2.triplet_counts,
            d1,
            sig3,
            ls3,
            r_cut_3,
            cutoff_func,
        )
        / 3
    )

    return two_term + three_term def two_plus_three_efs_energy(
    env1: AtomicEnvironment,
    env2: AtomicEnvironment,
    hyps: "ndarray",
    cutoffs: "ndarray",
    cutoff_func: Callable = cf.quadratic_cutoff,
):

    sig2 = hyps[0]
    ls2 = hyps[1]
    sig3 = hyps[2]
    ls3 = hyps[3]
    r_cut_2 = cutoffs[0]
    r_cut_3 = cutoffs[1]

    two_e, two_f, two_s = two_body_mc_simple.efs_energy(
        env1.bond_array_2,
        env1.ctype,
        env1.etypes,
        env2.bond_array_2,
        env2.ctype,
        env2.etypes,
        sig2,
        ls2,
        r_cut_2,
        cutoff_func,
    )

    three_e, three_f, three_s = three_body_mc_simple.efs_energy(
        env1.bond_array_3,
        env1.ctype,
        env1.etypes,
        env2.bond_array_3,
        env2.ctype,
        env2.etypes,
        env1.cross_bond_inds,
        env2.cross_bond_inds,
        env1.cross_bond_dists,
        env2.cross_bond_dists,
        env1.triplet_counts,
        env2.triplet_counts,
        sig3,
        ls3,
        r_cut_3,
        cutoff_func,
    )

    return two_e + three_e, two_f + three_f, two_s + three_s def two_plus_three_efs_force(
    env1: AtomicEnvironment,
    env2: AtomicEnvironment,
    hyps: "ndarray",
    cutoffs: "ndarray",
    cutoff_func: Callable = cf.quadratic_cutoff,
):

    sig2 = hyps[0]
    ls2 = hyps[1]
    sig3 = hyps[2]
    ls3 = hyps[3]
    r_cut_2 = cutoffs[0]
    r_cut_3 = cutoffs[1]

    two_e, two_f, two_s = two_body_mc_simple.efs_force(
        env1.bond_array_2,
        env1.ctype,
        env1.etypes,
        env2.bond_array_2,
        env2.ctype,
        env2.etypes,
        sig2,
        ls2,
        r_cut_2,
        cutoff_func,
    )

    three_e, three_f, three_s = three_body_mc_simple.efs_force(
        env1.bond_array_3,
        env1.ctype,
        env1.etypes,
        env2.bond_array_3,
        env2.ctype,
        env2.etypes,
        env1.cross_bond_inds,
        env2.cross_bond_inds,
        env1.cross_bond_dists,
        env2.cross_bond_dists,
        env1.triplet_counts,
        env2.triplet_counts,
        sig3,
        ls3,
        r_cut_3,
        cutoff_func,
    )

    return two_e + three_e, two_f + three_f, two_s + three_s def two_plus_three_body_mc(
    env1,
    env2,
    d1,
    d2,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
    """2+3-body multi-element kernel between two force components.

    Args:
        env1 (AtomicEnvironment): First local environment.
        env2 (AtomicEnvironment): Second local environment.
        d1 (int): Force component of the first environment.
        d2 (int): Force component of the second environment.
        cutoff_2b (float, np.ndarray): cutoff(s) for two-body interaction
        cutoff_3b (float, np.ndarray): cutoff(s) for three-body interaction
        cutoff_mb (float, np.ndarray): cutoff(s) for coordination-based manybody interaction
        nspec (int): number of different species groups
        spec_mask (np.ndarray): 118-long integer array that determines specie group
        nbond (int): number of different hyperparameter sets to associate with two-body pairings
        bond_mask (np.ndarray): nspec^2 long integer array
        ntriplet (int): number of different hyperparameter sets to associate with 3-body pairings
        triplet_mask (np.ndarray): nspec^3 long integer array
        ncut3b (int): number of different 3-body cutoff sets to associate with 3-body pairings
        cut3b_mask (np.ndarray): nspec^2 long integer array
        sig2 (np.ndarray): signal variances associates with two-body term
        ls2 (np.ndarray): length scales associates with two-body term
        sig3 (np.ndarray): signal variances associates with three-body term
        ls3 (np.ndarray): length scales associates with three-body term
        cutoff_func (Callable): Cutoff function of the kernel.

    Return:
        float: Value of the 2+3-body force/force kernel.
    """

    two_term = two_body_mc_jit(
        env1.bond_array_2,
        env1.ctype,
        env1.etypes,
        env2.bond_array_2,
        env2.ctype,
        env2.etypes,
        d1,
        d2,
        sig2,
        ls2,
        cutoff_2b,
        cutoff_func,
        nspec,
        spec_mask,
        bond_mask,
    )

    if ncut3b <= 1:
        tbmcj = three_body_mc_jit
    else:
        tbmcj = three_body_mc_sepcut_jit

    three_term = tbmcj(
        env1.bond_array_3,
        env1.ctype,
        env1.etypes,
        env2.bond_array_3,
        env2.ctype,
        env2.etypes,
        env1.cross_bond_inds,
        env2.cross_bond_inds,
        env1.cross_bond_dists,
        env2.cross_bond_dists,
        env1.triplet_counts,
        env2.triplet_counts,
        d1,
        d2,
        sig3,
        ls3,
        cutoff_3b,
        cutoff_func,
        nspec,
        spec_mask,
        triplet_mask,
        cut3b_mask,
    )

    return two_term + three_term def two_plus_three_body_mc_grad(
    env1,
    env2,
    d1,
    d2,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
    """2+3-body multi-element kernel between two force components and its
    gradient with respect to the hyperparameters.

    Args:
        env1 (AtomicEnvironment): First local environment.
        env2 (AtomicEnvironment): Second local environment.
        d1 (int): Force component of the first environment.
        d2 (int): Force component of the second environment.
        cutoff_2b (float, np.ndarray): cutoff(s) for two-body interaction
        cutoff_3b (float, np.ndarray): cutoff(s) for three-body interaction
        nspec (int): number of different species groups
        spec_mask (np.ndarray): 118-long integer array that determines specie group
        nbond (int): number of different hyperparameter sets to associate with two-body pairings
        bond_mask (np.ndarray): nspec^2 long integer array
        ntriplet (int): number of different hyperparameter sets to associate with 3-body pairings
        triplet_mask (np.ndarray): nspec^3 long integer array
        ncut3b (int): number of different 3-body cutoff sets to associate with 3-body pairings
        cut3b_mask (np.ndarray): nspec^2 long integer array
        sig2 (np.ndarray): signal variances associates with two-body term
        ls2 (np.ndarray): length scales associates with two-body term
        sig3 (np.ndarray): signal variances associates with three-body term
        ls3 (np.ndarray): length scales associates with three-body term
        cutoff_func (Callable): Cutoff function of the kernel.

    Return:
        (float, np.ndarray):
            Value of the 2+3-body kernel and its gradient
            with respect to the hyperparameters.
    """

    kern2, grad2 = two_body_mc_grad_jit(
        env1.bond_array_2,
        env1.ctype,
        env1.etypes,
        env2.bond_array_2,
        env2.ctype,
        env2.etypes,
        d1,
        d2,
        sig2,
        ls2,
        cutoff_2b,
        cutoff_func,
        nspec,
        spec_mask,
        nbond,
        bond_mask,
    )

    if ncut3b == 0:
        tbmcj = three_body_mc_grad_jit
    else:
        tbmcj = three_body_mc_grad_sepcut_jit

    kern3, grad3 = tbmcj(
        env1.bond_array_3,
        env1.ctype,
        env1.etypes,
        env2.bond_array_3,
        env2.ctype,
        env2.etypes,
        env1.cross_bond_inds,
        env2.cross_bond_inds,
        env1.cross_bond_dists,
        env2.cross_bond_dists,
        env1.triplet_counts,
        env2.triplet_counts,
        d1,
        d2,
        sig3,
        ls3,
        cutoff_3b,
        cutoff_func,
        nspec,
        spec_mask,
        ntriplet,
        triplet_mask,
        cut3b_mask,
    )

    g = np.hstack([grad2, grad3])

    return kern2 + kern3, g def two_three_many_body_mc(
    env1,
    env2,
    d1,
    d2,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def two_three_many_body_mc_grad(
    env1,
    env2,
    d1,
    d2,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def two_three_many_mc_force_en(
    env1,
    env2,
    d1,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def two_three_many_mc_en(
    env1,
    env2,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def two_plus_three_body_mc(
    env1,
    env2,
    d1,
    d2,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def two_plus_three_body_mc_grad(
    env1,
    env2,
    d1,
    d2,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def two_plus_three_mc_force_en(
    env1,
    env2,
    d1,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def two_plus_three_mc_en(
    env1,
    env2,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def three_body_mc(
    env1,
    env2,
    d1,
    d2,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def three_body_mc_grad(
    env1,
    env2,
    d1,
    d2,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def three_body_mc_force_en(
    env1,
    env2,
    d1,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def three_body_mc_en(
    env1,
    env2,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def two_body_mc(
    env1,
    env2,
    d1,
    d2,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def two_body_mc_grad(
    env1,
    env2,
    d1,
    d2,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def two_body_mc_force_en(
    env1,
    env2,
    d1,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def two_body_mc_en(
    env1,
    env2,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def three_body_mc_jit(
    bond_array_1,
    c1,
    etypes1,
    bond_array_2,
    c2,
    etypes2,
    cross_bond_inds_1,
    cross_bond_inds_2,
    cross_bond_dists_1,
    cross_bond_dists_2,
    triplets_1,
    triplets_2,
    d1,
    d2,
    sig,
    ls,
    r_cut,
    cutoff_func,
    nspec,
    spec_mask,
    triplet_mask,
    cut3b_mask,
):
def three_body_mc_grad_jit(
    bond_array_1,
    c1,
    etypes1,
    bond_array_2,
    c2,
    etypes2,
    cross_bond_inds_1,
    cross_bond_inds_2,
    cross_bond_dists_1,
    cross_bond_dists_2,
    triplets_1,
    triplets_2,
    d1,
    d2,
    sig,
    ls,
    r_cut,
    cutoff_func,
    nspec,
    spec_mask,
    ntriplet,
    triplet_mask,
    cut3b_mask,
):
def three_body_mc_force_en_jit(
    bond_array_1,
    c1,
    etypes1,
    bond_array_2,
    c2,
    etypes2,
    cross_bond_inds_1,
    cross_bond_inds_2,
    cross_bond_dists_1,
    cross_bond_dists_2,
    triplets_1,
    triplets_2,
    d1,
    sig,
    ls,
    r_cut,
    cutoff_func,
    nspec,
    spec_mask,
    triplet_mask,
    cut3b_mask,
):
def three_body_mc_en_jit(
    bond_array_1,
    c1,
    etypes1,
    bond_array_2,
    c2,
    etypes2,
    cross_bond_inds_1,
    cross_bond_inds_2,
    cross_bond_dists_1,
    cross_bond_dists_2,
    triplets_1,
    triplets_2,
    sig,
    ls,
    r_cut,
    cutoff_func,
    nspec,
    spec_mask,
    triplet_mask,
    cut3b_mask,
):
def two_body_mc_jit(
    bond_array_1,
    c1,
    etypes1,
    bond_array_2,
    c2,
    etypes2,
    d1,
    d2,
    sig,
    ls,
    r_cut,
    cutoff_func,
    nspec,
    spec_mask,
    bond_mask,
):
def two_body_mc_grad_jit(
    bond_array_1,
    c1,
    etypes1,
    bond_array_2,
    c2,
    etypes2,
    d1,
    d2,
    sig,
    ls,
    r_cut,
    cutoff_func,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
):
def two_body_mc_force_en_jit(
    bond_array_1,
    c1,
    etypes1,
    bond_array_2,
    c2,
    etypes2,
    d1,
    sig,
    ls,
    r_cut,
    cutoff_func,
    nspec,
    spec_mask,
    bond_mask,
):
def two_body_mc_en_jit(
    bond_array_1,
    c1,
    etypes1,
    bond_array_2,
    c2,
    etypes2,
    sig,
    ls,
    r_cut,
    cutoff_func,
    nspec,
    spec_mask,
    bond_mask,
):
def many_body_mc(
    env1,
    env2,
    d1,
    d2,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def many_body_mc_grad(
    env1,
    env2,
    d1,
    d2,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def many_body_mc_force_en(
    env1,
    env2,
    d1,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
def many_body_mc_en(
    env1,
    env2,
    cutoff_2b,
    cutoff_3b,
    cutoff_mb,
    nspec,
    spec_mask,
    nbond,
    bond_mask,
    ntriplet,
    triplet_mask,
    ncut3b,
    cut3b_mask,
    nmb,
    mb_mask,
    sig2,
    ls2,
    sig3,
    ls3,
    sigm,
    lsm,
    cutoff_func=cf.quadratic_cutoff,
):
                                        C1 = r11 * r11 + r22 * r22 + r33 * r33
                                        C3 = r12 * r12 + r21 * r21 + r33 * r33
                                        C5 = r13 * r13 + r21 * r21 + r32 * r32
                                        C2 = r11 * r11 + r23 * r23 + r32 * r32
                                        C6 = r13 * r13 + r22 * r22 + r31 * r31
                                        C4 = r12 * r12 + r23 * r23 + r31 * r31
                    A = ci * cj
                    B = r11 * ci
                    C = r11 * cj
                    D = r11 * r11
                    A = ci * cj
                    B = r11 * ci
                    C = r11 * cj
                    D = r11 * r11
                    B = r11 * ci
                    D = r11 * r11 class ParameterHelper:
    """
    A helper class to construct the hyps_mask dictionary for AtomicEnvironment
    , GaussianProcess and MappedGaussianProcess

    Args:
        hyps_mask (dict): Not implemented yet
        species (dict, list): Define specie groups
        kernels (dict, list): Define kernels and groups for the kernels
        cutoff_groups (dict): Define different cutoffs for different species
        parameters (dict): Define signal variance, length scales, and cutoffs
        constraints (dict): If listed as False, the cooresponding hyperparmeters
            will not be trained
        allseparate (bool): If True, define each type pair/triplet into a
            separate group.
        random (bool): If True, randomized all signal variances and lengthscales
        one (bool): If True, set all signal variances and lengthscales to one
        verbose (str): Level to print with "ERROR", "WARNING", "INFO", "DEBUG"

    * the ``species`` is an optional input. It can be left as None if the user only wants
      to set up one group of hyper-parameters for each kernel.
    * the ``kernels`` can be defined along with or without groups. But the later mode
      is not compatible with the ``allseparate`` flag.

      >>> kernels=['twobody', 'threebody'],

      or

      >>> kernels={'twobody':[['*', '*'], ['O','O']],
      ...          'threebody':[['*', '*', '*'],
      ...                       ['O','O', 'O']]},

      Current options for the kernels are twobody, threebody and manybody (based on coordination number).
    * See format of ``species``, ``kernels`` (dict), and ``cutoff_groups`` in ``list_groups()`` function.
    * See format of ``parameters`` and ``constraints`` in ``list_parameters()`` function.
    """

    # TO DO, sync it to kernel class
    #  need to be synced with kernel class

    # name of the kernels
    all_kernel_types = ["twobody", "threebody", "manybody"]
    cutoff_types = {"cut3b": "threebody"}
    cutoff_types_keys = list(cutoff_types.keys())
    cutoff_types_values = list(cutoff_types.values())
    additional_groups = []

    # dimension of the kernels
    ndim = {"twobody": 2, "threebody": 3, "manybody": 2, "cut3b": 2}
    n_kernel_parameters = {"twobody": 2, "threebody": 2, "manybody": 2, "cut3b": 0}

    def __init__(
        self,
        hyps_mask=None,
        species=None,
        kernels={},
        cutoff_groups={},
        parameters=None,
        constraints={},
        allseparate=False,
        random=False,
        ones=False,
        verbose="WARNING",
    ):

        self.logger = set_logger(
            "ParameterHelper", stream=True, fileout_name=None, verbose=verbose
        )

        self.all_group_types = (
            ParameterHelper.all_kernel_types
            + self.cutoff_types_keys
            + self.additional_groups
        )

        self.all_types = ["specie"] + self.all_group_types

        # number of groups {'twobody': 1, 'threebody': 2}
        self.n = {}
        # definition of groups {'specie': [['C', 'H'], ['O']],
        # 'twobody': [[['*', '*']], [[ele1, ele2]]]}
        self.groups = {}
        # joint values of the groups {'specie': ['C', 'H', 'O'],
        # 'twobody': [['*', '*'], [ele1, ele2]]}
        self.all_members = {}
        # names of each group {'specie': ['group1', 'group2'], 'twobody':
        # ['twobody0', 'twobody1']}
        self.all_group_names = {}
        # joint list of all the keys in self.all_group_names
        self.all_names = []

        # set up empty container
        for group_type in self.all_types:
            self.n[group_type] = 0
            self.groups[group_type] = []
            self.all_members[group_type] = []
            self.all_group_names[group_type] = []

        # store parameters, key should be the one used in
        # all_group_names or kernel_name
        self.sigma = {}
        self.ls = {}
        self.noise = 0.05
        self.energy_noise = 0.1
        self.opt = {"noise": True}

        # key should be sigma, lengthscale
        # cutoff_kernel_name
        self.universal = {}

        # key should be in all_group_names
        self.all_cutoff = {}

        # used for as_dict
        self.hyps_sig = {}
        self.hyps_ls = {}
        self.hyps_opt = {}
        self.cutoff_list = {}
        self.mask = {}

        self.hyps = None

        if isinstance(kernels, dict):
            self.kernel_dict = kernels
            self.kernels = list(kernels.keys())
            assert not allseparate
        elif isinstance(kernels, list):
            self.kernels = kernels
            # by default, there is only one group of hyperparameters
            # for each type of the kernel
            # unless allseparate is defined
            self.kernel_dict = {}
            for ktype in kernels:
                self.kernel_dict[ktype] = [["*"] * ParameterHelper.ndim[ktype]]

        if species is not None:
            self.list_groups("specie", species)

            # define groups
            if allseparate:
                for ktype in self.kernels:
                    self.all_separate_groups(ktype)
            else:
                for ktype in self.kernels:
                    self.list_groups(ktype, self.kernel_dict[ktype])

            # check for cut3b
            for group in cutoff_groups:
                self.list_groups(group, cutoff_groups[group])

            # define parameters
            if parameters is not None:
                self.list_parameters(parameters, constraints)

            if "lengthscale" in self.universal or "sigma" in self.universal:
                universal = True
            else:
                universal = False

            if (random + ones + universal) > 1:
                raise RuntimeError("random and ones cannot be simultaneously True")
            elif random or ones or universal:
                for ktype in self.kernels:
                    self.fill_in_parameters(
                        ktype, random=random, ones=ones, universal=universal
                    )

        elif len(self.kernels) > 0:
            self.list_groups("specie", ["*"])

            # define groups
            for ktype in self.kernels:
                self.list_groups(ktype, self.kernel_dict[ktype])

            # check for cut3b
            for group in cutoff_groups:
                self.list_groups(group, cutoff_groups[group])

            # define parameters
            if parameters is not None:
                self.list_parameters(parameters, constraints)

            if "lengthscale" in self.universal or "sigma" in self.universal:
                universal = True
            else:
                universal = False

            if (random + ones + universal) > 1:
                raise RuntimeError("random and ones cannot be simultaneously True")
            elif random or ones or universal:
                for ktype in self.kernels:
                    self.fill_in_parameters(
                        ktype, random=random, ones=ones, universal=universal
                    )

    def list_parameters(self, parameter_dict: dict, constraints: dict = {}):
        """Define many groups of parameters

        Args:
            parameter_dict (dict): dictionary of all parameters
            constraints (dict): dictionary of all constraints

        Example:

        >>> parameter_dict={"group_name":[sig, ls, cutoffs], ...}
        >>> constraints={"group_name":[True, False, False], ...}

        The name of parameters can be the group name previously defined in
        define_group or list_groups function. Aside from the group name,
        ``noise``, ``cutoff_twobody``, ``cutoff_threebody``, and
        ``cutoff_manybody`` are reserved for noise parmater
        and universal cutoffs, while ``sigma`` and ``lengthscale`` are
        reserved for universal signal variances and length scales.

        For non-reserved keys, the value should be a list of 2 to 3 elements,
        corresponding to the sigma, lengthscale and cutoff (if the third one
        is defined). For reserved keys, the value should be a float number.

        The parameter_dict and constraints should use the same set of keys.
        If a key in constraints is not used in parameter_dict, it will be ignored.

        The value in the constraints can be either a single bool, which apply
        to all parameters, or list of bools that apply to each parameter.
        """

        for name in parameter_dict:
            self.set_parameters(name, parameter_dict[name], constraints.get(name, True))
        for name in constraints:
            if name not in parameter_dict:
                self.set_constraints(name, constraints[name])

    def list_groups(self, group_type, definition_list):
        """define groups in batches.

        Args:
            group_type (str): "specie", "twobody", "threebody", "cut3b", "manybody"
            definition_list (list, dict): list of elements

        This function runs define_group in batch. Please first read
        the manual of define_group.

        If the definition_list is a list, it is equivalent to
        executing define_group through the definition_list.

        >>> for all terms in the list:
        >>>     define_group(group_type, group_type+'n', the nth term in the list)

        So the first twobody defined will be group twobody0, second one will be
        group twobody1. For specie, it will define all the listed elements as
        groups with only one element with their original name.

        If the definition_list is a dictionary, it is equivalent to

        >>> for k, v in the dict:
        >>>     define_group(group_type, k, v)

        It is not recommended to use the dictionary mode, especially when
        the group definitions are conflicting with each other. There is no
        guarantee that the priority order is the same as you want.

        Unlike ParameterHelper.define_group(), it can only be called once for each
        group_type, and not after any ParameterHelper.define_group() calls.

        """

        if group_type == "specie":
            if len(self.all_group_names["specie"]) > 0:
                raise RuntimeError(
                    "this function has to be run before any define_group"
                )
            if isinstance(definition_list, list):
                for ele in definition_list:
                    if isinstance(ele, list):
                        self.define_group("specie", ele, ele)
                    else:
                        self.define_group("specie", ele, [ele])
            elif isinstance(definition_list, dict):
                for ele in definition_list:
                    self.define_group("specie", ele, definition_list[ele])
            else:
                raise RuntimeError("type unknown")
        else:
            if self.n["specie"] == 0:
                raise RuntimeError(
                    "this function has to be run before any define_group"
                )
            if isinstance(definition_list, list):
                ngroup = len(definition_list)
                for idg in range(ngroup):
                    self.define_group(
                        group_type, f"{group_type}{idg}", definition_list[idg]
                    )
            elif isinstance(definition_list, dict):
                for name in definition_list:
                    if isinstance(definition_list[name][0], list):
                        for ele in definition_list[name]:
                            self.define_group(group_type, name, ele)
                    else:
                        self.define_group(group_type, name, definition_list[name])

    def all_separate_groups(self, group_type):
        """Separate all possible types of twobodys, threebodys, manybody.
        One type per group.

        Args:
            group_type (str): "specie", "twobody", "threebody", "cut3b", "manybody"

        """
        nspec = len(self.all_group_names["specie"])
        if nspec < 1:
            raise RuntimeError("the specie group has to be defined in advance")
        if group_type in self.all_group_types:
            # TO DO: the two blocks below can be replace by some upper triangle operation

            # generate all possible combination of group
            ele_grid = self.all_group_names["specie"]
            grid = np.meshgrid(*[ele_grid] * ParameterHelper.ndim[group_type])
            grid = np.array(grid).T.reshape(-1, ParameterHelper.ndim[group_type])

            # remove the redundant groups
            allgroup = []
            for group in grid:
                exist = False
                set_list_group = set(list(group))
                for prev_group in allgroup:
                    if set(prev_group) == set_list_group:
                        exist = True
                if not exist:
                    allgroup += [list(group)]

            # define the group
            tid = 0
            for group in allgroup:
                self.define_group(group_type, f"{group_type}{tid}", group)
                tid += 1
        else:
            self.logger.info(f"{group_type} will be ignored")

    def fill_in_parameters(self, group_type, random=False, ones=False, universal=False):
        """Separate all possible types of twobodys, threebodys, manybody.
        One type per group. And fill in either universal ls and sigma from
        pre-defined parameters from set_parameters("sigma", ..) and set_parameters("ls", ..)
        or random parameters if random is True.

        Args:
            group_type (str): "specie", "twobody", "threebody", "cut3b", "manybody"
            definition_list (list, dict): list of elements

        """
        nspec = len(self.all_group_names["specie"])
        if nspec < 1:
            raise RuntimeError("the specie group has to be defined in advance")
        if random:
            for group_name in self.all_group_names[group_type]:
                self.set_parameters(group_name, parameters=nprandom(2))
        elif ones:
            for group_name in self.all_group_names[group_type]:
                self.set_parameters(group_name, parameters=np.ones(2))
        elif universal:
            for group_name in self.all_group_names[group_type]:
                self.set_parameters(
                    group_name,
                    parameters=[self.universal["sigma"], self.universal["lengthscale"]],
                )

    def define_group(
        self, group_type, name, element_list, parameters=None, atomic_str=False
    ):
        """Define specie/twobody/threebody/3b cutoff/manybody group

        Args:
            group_type (str): "specie", "twobody", "threebody", "cut3b", "manybody"
            name (str): the name use for indexing. can be anything but "*"
            element_list (list): list of elements
            parameters (list): corresponding parameters for this group
            atomic_str (bool): whether the elements in element_list are
                specified by group names or periodic table element names.

        The function is helped to define different groups for specie/twobody/threebody
        /3b cutoff/manybody terms. This function can be used for many times.
        The later one always overrides the former one.

        The name of the group has to be unique string (but not "*"), that
        define a group of species or twobodys, etc. If the same name is used,
        in two function calls, the definitions of the group will be merged.
        Both calls will be effective.

        element_list has to be a list of atomic elements, or a list of
        specie group names (which should be defined in previous calls), or "*".
        "*" will loop the function over all previously defined species.
        It has to be two elements for twobody/3b cutoff/manybody term, or
        three elements for threebody. For specie group definition, it can be
        as many elements as you want.

        If multiple define_group calls have conflict with element, the later one
        has higher priority. For example, twobody 1-2 are defined as group1 in
        the first call, and as group2 in the second call. In the end, the twobody
        will be left as group2.

        Example 1:

        >>> define_group('specie', 'water', ['H', 'O'])
        >>> define_group('specie', 'salt', ['Cl', 'Na'])

        They define H and O to be group water, and Na and Cl to be group salt.

        Example 2.1:

        >>> define_group('twobody', 'in-water', ['H', 'H'], atomic_str=True)
        >>> define_group('twobody', 'in-water', ['H', 'O'], atomic_str=True)
        >>> define_group('twobody', 'in-water', ['O', 'O'], atomic_str=True)

        Example 2.2:

        >>> define_group('twobody', 'in-water', ['water', 'water'])

        The 2.1 is equivalent to 2.2.

        Example 3.1:

        >>> define_group('specie', '1', ['H'])
        >>> define_group('specie', '2', ['O'])
        >>> define_group('twobody', 'Hgroup', ['H', 'H'], atomic_str=True)
        >>> define_group('twobody', 'Hgroup', ['H', 'O'], atomic_str=True)
        >>> define_group('twobody', 'OO', ['O', 'O'], atomic_str=True)

        Example 3.2:

        >>> define_group('specie', '1', ['H'])
        >>> define_group('specie', '2', ['O'])
        >>> define_group('twobody', 'Hgroup', ['H', '*'], atomic_str=True)
        >>> define_group('twobody', 'OO', ['O', 'O'], atomic_str=True)

        Example 3.3:

        >>> list_groups('specie', ['H', 'O'])
        >>> define_group('twobody', 'Hgroup', ['H', '*'])
        >>> define_group('twobody', 'OO', ['O', 'O'])

        Example 3.4:

        >>> list_groups('specie', ['H', 'O'])
        >>> define_group('twobody', 'OO', ['*', '*'])
        >>> define_group('twobody', 'Hgroup', ['H', '*'])

        3.1 to 3.4 are all equivalent.
        """

        if name == "*" and group_type == "specie":
            name = "allspecie"
            element_list = ["H"]
        elif name == "*":
            raise ValueError(
                "* is reserved for substitution, cannot be used as a group name"
            )

        if group_type != "specie":

            assert len(element_list) == ParameterHelper.ndim[group_type]

            # Check all the other group_type to
            exclude_list = deepcopy(self.all_types)
            ide = exclude_list.index(group_type)
            exclude_list.pop(ide)

            for gt in exclude_list:
                if name in self.all_group_names[gt]:
                    raise ValueError(
                        "group name has to be unique across all types. "
                        f"{name} is found in type {gt}"
                    )

        if name in self.all_group_names[group_type]:
            groupid = self.all_group_names[group_type].index(name)
        else:
            groupid = self.n[group_type]
            self.all_group_names[group_type].append(name)
            self.groups[group_type].append([])
            self.n[group_type] += 1

        if group_type == "specie":
            for ele in element_list:
                assert (
                    ele not in self.all_members["specie"]
                ), "The element has already been defined"
                self.groups["specie"][groupid].append(ele)
                self.all_members["specie"].append(ele)
                self.logger.debug(f"Element {ele} will be defined as group {name}")
        else:
            if len(self.all_group_names["specie"]) == 0:
                raise RuntimeError("The atomic species have to bedefined in advance")

            # first translate element/group name to group name
            group_name_list = []
            if atomic_str:
                for ele_name in element_list:
                    if ele_name == "*":
                        gid += ["*"]
                    else:
                        for idx in range(self.n["specie"]):
                            group_name = self.all_group_names["specie"][idx]
                            if ele_name in self.groups["specie"][idx]:
                                group_name_list += [group_name]
                                self.logger.debug(
                                    f"Element {ele_name} is used for "
                                    "definition, but the whole group "
                                    f"{group_name} is affected"
                                )
            else:
                group_name_list = element_list

            if "*" not in group_name_list:

                gid = []
                for ele_name in group_name_list:
                    gid += [self.all_group_names["specie"].index(ele_name)]

                for ele in self.all_members[group_type]:
                    if set(gid) == set(ele):
                        self.logger.debug(
                            f"the definition of {group_type} {ele} will be overriden"
                        )
                self.groups[group_type][groupid].append(gid)
                self.all_members[group_type].append(gid)
                self.logger.debug(f"{group_type} {gid} will be defined as group {name}")
                if parameters is not None:
                    self.set_parameters(name, parameters)
            else:
                one_star_less = deepcopy(group_name_list)
                idstar = group_name_list.index("*")
                one_star_less.pop(idstar)
                for sub in self.all_group_names["specie"]:
                    self.logger.debug(f"{sub}, {one_star_less}")
                    self.define_group(
                        group_type,
                        name,
                        one_star_less + [sub],
                        parameters=parameters,
                        atomic_str=False,
                    )

    def find_group(self, group_type, element_list, atomic_str=False):
        """find the group that contains the input pair

        Args:
            group_type (str): species, twobody, threebody, cut3b, manybody
            element_list (list): list of elements for a pair/triplet/coordination-pair
            atomic_str (bool): whether the elements in element_list are
                specified by group names or periodic table element names.

        Return:
            name (str):
        """

        # remember the later command override the earlier ones
        if group_type == "specie":
            if not isinstance(element_list, str):
                self.logger.debug("for element, it has to be a string")
                return None
            name = None
            for igroup in range(self.n["specie"]):
                gname = self.all_group_names[group_type][igroup]
                allspec = self.groups[group_type][igroup]
                if element_list in allspec:
                    name = gname
            if name is None:
                self.logger.debug("cannot find the group")
            return name
        else:
            if "*" in element_list:
                self.logger.debug("* cannot be used for find")
                return None
            gid = []
            for ele_name in element_list:
                gid += [self.all_group_names["specie"].index(ele_name)]
            name = None
            for igroup in range(self.n[group_type]):
                gname = self.all_group_names[group_type][igroup]
                for ele in self.groups[group_type][igroup]:
                    if set(gid) == set(ele):
                        name = gname
            self.logger.debug(f"find the group {name}")
            return name

    def set_parameters(self, name, parameters, opt=True):
        """Set the parameters for certain group

        Args:
            name (str): name of the patermeters
            parameters (list): the sigma, lengthscale, and cutoff of each group.
            opt (bool, list): whether to optimize the parameter or not

        The name of parameters can be the group name previously defined in
        define_group or list_groups function. Aside from the group name,
        ``noise``, ``cutoff_twobody``, ``cutoff_threebody``, and
        ``cutoff_manybody`` are reserved for noise parmater
        and universal cutoffs, while ``sigma`` and ``lengthscale`` are
        reserved for universal signal variances and length scales.

        The parameter should be a list of 2-3 elements, for sigma,
        lengthscale (and cutoff if the third one is defined).

        The optimization flag can be a single bool, which apply to all
        parameters, or list of bools that apply to each parameter.
        """

        if name == "noise":
            self.noise = parameters
            self.opt["noise"] = opt
            self.logger.debug(f"noise will be set as {parameters} and opt {opt}")
            return
        elif name == "energy_noise":
            self.energy_noise = parameters
            self.opt["energy_noise"] = opt
            self.logger.debug(f"energy_noise will be set as {parameters} and opt {opt}")
            return
        elif "cutoff" in name:
            self.universal[name] = parameters
            self.logger.debug(f"universal cutoff {name} will be set as {parameters}")
            return
        elif name in ["sigma", "lengthscale"]:
            self.universal[name] = parameters
            self.opt[name] = opt
            self.logger.debug(
                f"universal {name} will be set as {parameters} and optimized {opt}"
            )
            return

        if isinstance(opt, bool):
            opt = [opt] * 2

        for cut_name in self.cutoff_types:
            if cut_name in name:
                self.all_cutoff[name] = float(parameters)
                self.logger.debug(
                    f"Cutoff for group {name} will be set as {parameters}"
                )
                return

        if name in self.sigma:
            self.logger.debug(f"the sig, ls of group {name} is overriden")
        self.sigma[name] = parameters[0]
        self.ls[name] = parameters[1]
        self.opt[name + "sig"] = opt[0]
        self.opt[name + "ls"] = opt[1]
        self.logger.debug(
            f"ParameterHelper for group {name} will be set as "
            f"sig={parameters[0]} ({opt[0]}) "
            f"ls={parameters[1]} ({opt[1]})"
        )
        if len(parameters) > 2:
            if name in self.all_cutoff:
                self.logger.debug(f"the cutoff of group {name} is overriden")
            self.all_cutoff[name] = parameters[2]
            self.logger.debug(f"Cutoff for group {name} will be set as {parameters[2]}")

    def set_constraints(self, name, opt):
        """Set the parameters for certain group

        Args:
            name (str): name of the patermeters
            opt (bool, list): whether to optimize the parameter or not

        The name of parameters can be the group name previously defined in
        define_group or list_groups function. Aside from the group name,
        ``noise``, ``cutoff_twobody``, ``cutoff_threebody``, and
        ``cutoff_manybody`` are reserved for noise parmater
        and universal cutoffs, while ``sigma`` and ``lengthscale`` are
        reserved for universal signal variances and length scales.

        The optimization flag can be a single bool, which apply to all
        parameters under that name, or list of bools that apply to each
        parameter.
        """

        if name == "noise":
            self.opt["noise"] = opt
            self.logger.debug(f"noise opt is set to{opt}")
            return

        if isinstance(opt, bool):
            opt = [opt, opt, opt]

        for cutname in self.cutoff_types:
            if cutname in name:
                return

        if name in self.sigma:
            self.logger.debug(f"the opt setting of group {name} is overriden")
        self.opt[name + "sig"] = opt[0]
        self.opt[name + "ls"] = opt[1]
        self.logger.debug(
            f"ParameterHelper for group {name} will be set as sig {opt[0]} ls {opt[1]}"
        )

    def summarize_group(self, group_type):
        """Sort and combine all the previous definition to internal varialbes

        Args:
            group_type (str): species, twobody, threebody, cut3b, manybody
        """

        aeg = self.all_group_names[group_type]
        nspecie = self.n["specie"]

        # specie need special sorting
        if group_type == "specie":

            self.nspecie = nspecie
            self.species_mask = np.ones(118, dtype=int) * (nspecie - 1)

            # mark the species_mask with atom type
            # default is nspecie-1
            for idt in range(self.nspecie):
                for ele in self.groups["specie"][idt]:
                    atom_n = element_to_Z(ele)
                    if atom_n >= len(self.species_mask):
                        new_mask = np.ones(atom_n, dtype=np.int) * (nspecie - 1)
                        new_mask[: len(self.species_mask)] = self.species_mask
                        self.species_mask = new_mask
                    self.species_mask[atom_n] = idt
                    self.logger.debug(
                        f"elemtn {ele} is defined as type {idt} with name {aeg[idt]}"
                    )
            self.logger.debug(f"All the remaining elements are left as type {idt}")

        elif group_type in self.all_group_types:

            if self.n[group_type] == 0:
                self.logger.debug(f"{group_type} is not defined. Skipped")
                return

            if (group_type not in self.kernels) and (
                group_type in ParameterHelper.all_kernel_types
            ):
                self.kernels.append(group_type)

            self.mask[group_type] = np.ones(
                nspecie ** ParameterHelper.ndim[group_type], dtype=int
            ) * (self.n[group_type] - 1)

            self.hyps_sig[group_type] = []
            self.hyps_ls[group_type] = []
            all_opt_sig = []
            all_opt_ls = []

            for idt in range(self.n[group_type]):
                name = aeg[idt]
                for ele_list in self.groups[group_type][idt]:
                    # generate all possible permutation
                    perms = list(permutations(ele_list))
                    for ele_list in perms:
                        mask_id = 0
                        for ele in ele_list:
                            mask_id += ele
                            mask_id *= nspecie
                        mask_id = mask_id // nspecie
                        self.mask[group_type][mask_id] = idt
                    def_str = "-".join(map(str, self.groups["specie"]))
                    self.logger.debug(
                        f"{group_type} {def_str} is defined as type {idt} "
                        f"with name {name}"
                    )

                if group_type not in self.cutoff_types:
                    sig = self.sigma.get(name, -1)
                    opt_sig = self.opt.get(name + "sig", True)
                    if sig == -1:
                        sig = self.sigma.get(group_type, -1)
                        opt_sig = self.opt.get(group_type + "sig", True)
                    if sig == -1:
                        sig = self.universal.get("sigma", -1)
                        opt_sig = self.opt.get("sigma", True)

                    ls = self.ls.get(name, -1)
                    opt_ls = self.opt.get(name + "ls", True)
                    if ls == -1:
                        ls = self.ls.get(group_type, -1)
                        opt_ls = self.opt.get(group_type + "ls", True)
                    if ls == -1:
                        ls = self.universal.get("lengthscale", -1)
                        opt_ls = self.opt.get("lengthscale", True)

                    if sig < 0 or ls < 0:
                        self.logger.error(
                            f"hyper parameters for group {name} is not defined"
                        )
                        raise RuntimeError
                    self.hyps_sig[group_type] += [sig]
                    self.hyps_ls[group_type] += [ls]
                    all_opt_sig += [opt_sig]
                    all_opt_ls += [opt_ls]
                    self.logger.debug(
                        f"   using hyper-parameters of {sig:6.2g} "
                        f"{ls:6.2g} {opt_sig} {opt_ls}"
                    )
            self.hyps_opt[group_type] = all_opt_sig + all_opt_ls
            self.logger.debug(f"All the remaining elements are left as type {idt}")

            # sort out the cutoffs
            if group_type in self.cutoff_types:
                universal_cutoff = self.universal.get(
                    "cutoff_" + self.cutoff_types[group_type], 0
                )
            else:
                universal_cutoff = self.universal.get("cutoff_" + group_type, 0)

            allcut = []
            alldefine = True
            for idt in range(self.n[group_type]):
                if aeg[idt] in self.all_cutoff:
                    allcut += [self.all_cutoff[aeg[idt]]]
                else:
                    alldefine = False
                    self.logger.info(
                        f"{aeg[idt]} cutoff is not defined. "
                        "it's going to use the universal cutoff."
                    )
            if group_type not in self.cutoff_types_values:

                if len(allcut) > 0:
                    if universal_cutoff <= 0:
                        universal_cutoff = np.max(allcut)
                        self.logger.info(
                            f"universal cutoff for {group_type} is defined as zero!"
                            f" reset it to {universal_cutoff}"
                        )

                    self.cutoff_list[group_type] = []
                    for idt in range(self.n[group_type]):
                        self.cutoff_list[group_type] += [
                            self.all_cutoff.get(aeg[idt], universal_cutoff)
                        ]
                    self.cutoff_list[group_type] = np.array(
                        self.cutoff_list[group_type], dtype=float
                    )

                    max_cutoff = np.max(self.cutoff_list[group_type])

                    # update the universal cutoff to make it higher than
                    if alldefine:
                        universal_cutoff = max_cutoff
                        self.logger.info(
                            f"universal cutoff is updated to {universal_cutoff}"
                        )
                    elif not np.any(self.cutoff_list[group_type] - max_cutoff):
                        # if not all the cutoffs are defined separately
                        # and they are all the same value
                        del self.cutoff_list[group_type]
                        universal_cutoff = max_cutoff
                        if group_type in self.cutoff_types:
                            self.n[group_type] = 0
                        self.logger.info(
                            f"universal cutoff is updated to {universal_cutoff}"
                        )

            else:
                if universal_cutoff <= 0 and len(allcut) > 0:
                    universal_cutoff = np.max(allcut)
                    self.logger.info(
                        "threebody universal cutoff is updated to"
                        f"{universal_cutoff}, but the separate definitions will"
                        "be ignored"
                    )

            if universal_cutoff > 0:
                if group_type in self.cutoff_types:
                    keyname = "cutoff_" + self.cutoff_types[group_type]
                else:
                    keyname = "cutoff_" + group_type
                self.universal[keyname] = universal_cutoff
            else:
                self.logger.error(f"cutoffs for {group_type} is undefined")
                raise RuntimeError

        else:
            pass

    def as_dict(self):
        """Dictionary representation of the mask. The output can be used for AtomicEnvironment
        or the GaussianProcess
        """

        # sort out all the definitions and resolve conflicts
        # cut3b has to be summarize before threebody
        # because the universal threebody cutoff is checked
        # at the end of threebody search

        self.summarize_group("specie")
        for ktype in self.cutoff_types:
            self.summarize_group(ktype)
        for ktype in ParameterHelper.additional_groups:
            self.summarize_group(ktype)
        for ktype in ParameterHelper.all_kernel_types:
            self.summarize_group(ktype)

        hyps_mask = {}
        cutoff_dict = {}

        hyps_mask["nspecie"] = self.n["specie"]
        if self.n["specie"] > 1:
            hyps_mask["species_mask"] = self.species_mask

        hyps = []
        hyp_labels = []
        opt = []
        for group in self.kernels:

            hyps_mask["n" + group] = self.n[group]
            hyps_mask[group + "_start"] = len(hyps)
            hyps += [self.hyps_sig[group]]
            hyps += [self.hyps_ls[group]]
            hyps = list(np.hstack(hyps))
            opt += [self.hyps_opt[group]]
            cutoff_dict[group] = self.universal["cutoff_" + group]

            if self.n[group] > 1:
                hyps_mask[group + "_mask"] = self.mask[group]
                # check parameters
                aeg = self.all_group_names[group]
                for idt in range(self.n[group]):
                    hyp_labels += ["Signal Std " + aeg[idt]]
                for idt in range(self.n[group]):
                    hyp_labels += ["Length Scale " + aeg[idt]]
            else:
                hyp_labels += ["Signal Std " + group]
                hyp_labels += ["Length Scale " + group]

            if group in self.cutoff_list:
                hyps_mask[group + "_cutoff_list"] = self.cutoff_list[group]

        for cutoff_name in self.cutoff_types:
            if self.n.get(cutoff_name, 0) >= 1:
                hyps_mask["n" + cutoff_name] = self.n[cutoff_name]
                hyps_mask[cutoff_name + "_mask"] = self.mask[cutoff_name]
                hyps_mask[
                    self.cutoff_types[cutoff_name] + "_cutoff_list"
                ] = self.cutoff_list[cutoff_name]

        hyps_mask["train_noise"] = self.opt["noise"]
        hyps_mask["energy_noise"] = self.energy_noise

        opt += [self.opt["noise"]]
        hyp_labels += ["Noise std"]
        hyps += [self.noise]
        hyps = np.hstack(hyps)
        opt = np.hstack(opt)

        # handle partial optimization if any constraints are defined
        if not opt.all():
            nhyps = len(hyps)
            hyps_mask["original_hyps"] = hyps
            hyps_mask["original_labels"] = hyp_labels
            mapping = []
            new_labels = []
            for i in range(nhyps):
                if opt[i]:
                    mapping += [i]
                    new_labels += [hyp_labels[i]]
            newhyps = hyps[mapping]
            hyps_mask["map"] = np.array(mapping, dtype=int)
        elif opt.any():
            newhyps = hyps
            new_labels = hyp_labels
        else:
            raise RuntimeError(
                "hyps has length zero."
                "at least one component of the hyper-parameters"
                "should be allowed to be optimized. \n"
            )

        if self.n["specie"] < 2:
            self.logger.debug(
                "only one type of elements was defined. Please use multihyps=False"
            )

        hyps_mask["kernels"] = self.kernels
        hyps_mask["kernel_name"] = "+".join(hyps_mask["kernels"])
        hyps_mask["cutoffs"] = cutoff_dict
        hyps_mask["hyps"] = newhyps
        hyps_mask["hyp_labels"] = new_labels

        logging.debug(str(hyps_mask))

        return hyps_mask

    @staticmethod
    def from_dict(hyps_mask, verbose=False, init_spec=[]):
        """convert dictionary mask to HM instance
        This function is not tested yet
        """

        Parameters.check_instantiation(
            hyps_mask["hyps"], hyps_mask["cutoffs"], hyps_mask["kernels"], hyps_mask
        )

        pm = ParameterHelper(verbose=verbose)

        nspecie = hyps_mask["nspecie"]
        if nspecie > 1:
            max_species = np.max(hyps_mask["species_mask"])
            species_mask = hyps_mask["species_mask"]
            for i in range(max_species + 1):
                elelist = np.where(species_mask == i)[0]
                if len(elelist) > 0:
                    for ele in elelist:
                        if ele != 0:
                            elename = Z_to_element(ele)
                            if len(init_spec) > 0:
                                if elename in init_spec:
                                    pm.define_group("specie", i, [elename])
                            else:
                                pm.define_group("specie", i, [elename])
        else:
            pm.define_group("specie", i, ["*"])

        for kernel in hyps_mask["kernels"] + ParameterHelper.cutoff_types_keys:
            n = hyps_mask.get("n" + kernel, 0)
            if n >= 0:
                if kernel not in ParameterHelper.cutoff_types:
                    chyps, copt = Parameters.get_component_hyps(
                        hyps_mask, kernel, constraint=True, noise=False
                    )
                    sig = chyps[0]
                    ls = chyps[1]
                    csig = copt[0]
                    cls = copt[1]
                    cutoff = hyps_mask["cutoffs"][kernel]
                    pm.set_parameters("cutoff_" + kernel, cutoff)
                    cutoff_list = hyps_mask.get(
                        f"{kernel}_cutoff_list", np.ones(len(sig)) * cutoff
                    )
                elif kernel in ParameterHelper.cutoff_types and n > 1:
                    cutoff_list = hyps_mask[
                        ParameterHelper.cutoff_types[kernel] + "_cutoff_list"
                    ]

                if n > 1:
                    all_specie = np.arange(nspecie)
                    all_comb = combinations_with_replacement(
                        all_specie, ParameterHelper.ndim[kernel]
                    )
                    for comb in all_comb:
                        mask_id = 0
                        for ele in comb:
                            mask_id += ele
                            mask_id *= nspecie
                        mask_id = mask_id // nspecie
                        ttype = hyps_mask[f"{kernel}_mask"][mask_id]
                        pm.define_group(f"{kernel}", f"{kernel}{ttype}", comb)

                        if (kernel not in ParameterHelper.cutoff_types) and (
                            kernel not in ParameterHelper.cutoff_types_values
                        ):
                            pm.set_parameters(
                                f"{kernel}{ttype}",
                                [sig[ttype], ls[ttype], cutoff_list[ttype]],
                                opt=[csig[ttype], cls[ttype]],
                            )
                        elif kernel in ParameterHelper.cutoff_types_values:
                            pm.set_parameters(
                                f"{kernel}{ttype}",
                                [sig[ttype], ls[ttype]],
                                opt=[csig[ttype], cls[ttype]],
                            )
                        else:
                            pm.set_parameters(f"{kernel}{ttype}", cutoff_list[ttype])
                else:
                    pm.define_group(
                        kernel, kernel, ["*"] * ParameterHelper.ndim[kernel]
                    )
                    if kernel not in ParameterHelper.cutoff_types_keys:
                        pm.set_parameters(
                            kernel, parameters=np.hstack([sig, ls, cutoff]), opt=copt
                        )
                    else:
                        pm.set_parameters(kernel, parameters=cutoff)

        hyps = Parameters.get_hyps(hyps_mask)
        pm.set_parameters("noise", hyps[-1])

        if "cutoffs" in hyps_mask:
            cutoffs = hyps_mask["cutoffs"]
            for k in cutoffs:
                pm.set_parameters(f"cutoff_{k}", cutoffs[k])

        return pm def get_like_grad_from_mats(ky_mat, hyp_mat, name):
    """compute the gradient of likelihood to hyper-parameters
    from covariance matrix and its gradient

    :param ky_mat: covariance matrix
    :type ky_mat: np.array
    :param hyp_mat: dky/d(hyper parameter) matrix
    :type hyp_mat: np.array
    :param name: name of the gp instance.

    :return: float, list. the likelihood and its gradients
    """

    number_of_hyps = hyp_mat.shape[0]

    # catch linear algebra errors
    try:
        ky_mat_inv = np.linalg.inv(ky_mat)
        l_mat = np.linalg.cholesky(ky_mat)
    except np.linalg.LinAlgError:
        return -1e8, np.zeros(number_of_hyps)

    labels = _global_training_labels[name]

    alpha = np.matmul(ky_mat_inv, labels)
    alpha_mat = np.matmul(alpha.reshape(-1, 1), alpha.reshape(1, -1))
    like_mat = alpha_mat - ky_mat_inv

    # calculate likelihood
    like = (
        -0.5 * np.matmul(labels, alpha)
        - np.sum(np.log(np.diagonal(l_mat)))
        - math.log(2 * np.pi) * ky_mat.shape[1] / 2
    )

    # calculate likelihood gradient
    like_grad = np.zeros(number_of_hyps)
    for n in range(number_of_hyps):
        like_grad[n] = 0.5 * np.trace(np.matmul(like_mat, hyp_mat[n, :, :]))

    return like, like_grad def get_force_block(
    hyps: np.ndarray,
    name: str,
    kernel,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
    """parallel version of get_ky_mat
    :param hyps: list of hyper-parameters
    :param name: name of the gp instance.
    :param kernel: function object of the kernel
    :param cutoffs: The cutoff values used for the atomic environments
    :type cutoffs: list of 2 float numbers
    :param hyps_mask: dictionary used for multi-group hyperparmeters

    :return: covariance matrix
    """

    training_data = _global_training_data[name]
    size = len(training_data)
    size3 = 3 * size

    if n_cpus is None:
        n_cpus = mp.cpu_count()
    if n_cpus == 1:
        k_mat = get_force_block_pack(
            hyps, name, 0, size, 0, size, True, kernel, cutoffs, hyps_mask
        )
    else:
        # initialize matrices
        block_id, nbatch = partition_matrix(n_sample, size, n_cpus)
        mult = 3

        k_mat = parallel_matrix_construction(
            get_force_block_pack,
            hyps,
            name,
            kernel,
            cutoffs,
            hyps_mask,
            block_id,
            nbatch,
            size3,
            size3,
            mult,
            mult,
        )

    sigma_n, _, __ = obtain_noise_len(hyps, hyps_mask)
    force_block = k_mat
    force_block += sigma_n ** 2 * np.eye(size3)

    return force_block def get_force_energy_block(
    hyps: np.ndarray,
    name: str,
    kernel,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
    training_data = _global_training_data[name]
    training_structures = _global_training_structures[name]
    size1 = len(training_data) * 3
    size2 = len(training_structures)
    size3 = len(training_data)

    if n_cpus is None:
        n_cpus = mp.cpu_count()
    if n_cpus == 1:
        force_energy_block = get_force_energy_block_pack(
            hyps, name, 0, size3, 0, size2, True, kernel, cutoffs, hyps_mask
        )
    else:
        # initialize matrices
        block_id, nbatch = partition_force_energy_block(n_sample, size3, size2, n_cpus)
        m1 = 3  # 3 force components per environment
        m2 = 1  # 1 energy per structure

        force_energy_block = parallel_matrix_construction(
            get_force_energy_block_pack,
            hyps,
            name,
            kernel,
            cutoffs,
            hyps_mask,
            block_id,
            nbatch,
            size1,
            size2,
            m1,
            m2,
            symm=False,
        )

    return force_energy_block def energy_energy_vector(
    name, kernel, x, hyps, cutoffs=None, hyps_mask=None, n_cpus=1, n_sample=100
):
    """
    Get a vector of covariances between the local energy of a test environment
    and the total energy labels in the training set.
    """

    size = len(_global_training_structures[name])

    if n_cpus is None:
        n_cpus = mp.cpu_count()
    if n_cpus == 1:
        return energy_energy_vector_unit(
            name, 0, size, x, kernel, hyps, cutoffs, hyps_mask
        )

    block_id, nbatch = partition_vector(n_sample, size, n_cpus)
    pack_function = energy_energy_vector_unit
    mult = 1

    force_energy_vector = parallel_vector_construction(
        pack_function,
        name,
        x,
        kernel,
        hyps,
        cutoffs,
        hyps_mask,
        block_id,
        nbatch,
        size,
        mult,
    )

    return force_energy_vector def energy_force_vector(
    name, kernel, x, hyps, cutoffs=None, hyps_mask=None, n_cpus=1, n_sample=100
):
    """
    Get the vector of covariances between the local energy of a test
    environment and the force labels in the training set.
    """

    training_data = _global_training_data[name]
    size = len(training_data)

    if n_cpus is None:
        n_cpus = mp.cpu_count()
    if n_cpus == 1:
        return energy_force_vector_unit(
            name, 0, size, x, kernel, hyps, cutoffs, hyps_mask
        )

    block_id, nbatch = partition_vector(n_sample, size, n_cpus)
    pack_function = energy_force_vector_unit
    mult = 3

    k12_v = parallel_vector_construction(
        pack_function,
        name,
        x,
        kernel,
        hyps,
        cutoffs,
        hyps_mask,
        block_id,
        nbatch,
        size,
        mult,
    )

    return k12_v def force_force_vector(
    name, kernel, x, d_1, hyps, cutoffs=None, hyps_mask=None, n_cpus=1, n_sample=100
):
    """
    Get a vector of covariances between a force component of a test environment
    and the force labels in the training set.
    """

    size = len(_global_training_data[name])

    if n_cpus is None:
        n_cpus = mp.cpu_count()
    if n_cpus == 1:
        return force_force_vector_unit(
            name, 0, size, x, kernel, hyps, cutoffs, hyps_mask, d_1
        )

    block_id, nbatch = partition_vector(n_sample, size, n_cpus)
    pack_function = force_force_vector_unit
    mult = 3

    k12_v = parallel_vector_construction(
        pack_function,
        name,
        x,
        kernel,
        hyps,
        cutoffs,
        hyps_mask,
        block_id,
        nbatch,
        size,
        mult,
        d_1,
    )

    return k12_v def force_energy_vector(
    name, kernel, x, d_1, hyps, cutoffs=None, hyps_mask=None, n_cpus=1, n_sample=100
):
    """
    Get a vector of covariances between a force component of a test environment
    and the total energy labels in the training set.
    """

    size = len(_global_training_structures[name])

    if n_cpus is None:
        n_cpus = mp.cpu_count()
    if n_cpus == 1:
        return force_energy_vector_unit(
            name, 0, size, x, kernel, hyps, cutoffs, hyps_mask, d_1
        )

    block_id, nbatch = partition_vector(n_sample, size, n_cpus)
    pack_function = force_energy_vector_unit
    mult = 1

    force_energy_vector = parallel_vector_construction(
        pack_function,
        name,
        x,
        kernel,
        hyps,
        cutoffs,
        hyps_mask,
        block_id,
        nbatch,
        size,
        mult,
        d_1,
    )

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

    size1 = len(_global_training_data[name])
    size2 = len(_global_training_structures[name])
    kernel_vector = np.zeros(size1 * 3 + size2)

    force_vector = force_force_vector(
        name, force_force_kernel, x, d_1, hyps, cutoffs, hyps_mask, n_cpus, n_sample
    )
    energy_vector = force_energy_vector(
        name, force_energy_kernel, x, d_1, hyps, cutoffs, hyps_mask, n_cpus, n_sample
    )

    kernel_vector[0 : size1 * 3] = force_vector
    kernel_vector[size1 * 3 :] = energy_vector

    return kernel_vector def en_kern_vec(
    name,
    energy_force_kernel,
    energy_energy_kernel,
    x,
    hyps,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):

    size1 = len(_global_training_data[name])
    size2 = len(_global_training_structures[name])
    kernel_vector = np.zeros(size1 * 3 + size2)

    force_vector = energy_force_vector(
        name, energy_force_kernel, x, hyps, cutoffs, hyps_mask, n_cpus, n_sample
    )
    energy_vector = energy_energy_vector(
        name, energy_energy_kernel, x, hyps, cutoffs, hyps_mask, n_cpus, n_sample
    )

    kernel_vector[0 : size1 * 3] = force_vector
    kernel_vector[size1 * 3 :] = energy_vector

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

    n_envs = len(_global_training_data[name])
    n_strucs = len(_global_training_structures[name])
    size1 = n_envs * 3
    size2 = n_strucs
    ky_mat = np.zeros((size1 + size2, size1 + size2))

    force_block = update_force_block(
        ky_mat_old,
        n_envs_prev,
        hyps,
        name,
        force_kernel,
        cutoffs,
        hyps_mask,
        n_cpus,
        n_sample,
    )

    energy_block = update_energy_block(
        ky_mat_old,
        n_envs_prev,
        hyps,
        name,
        energy_kernel,
        energy_noise,
        cutoffs,
        hyps_mask,
        n_cpus,
        n_sample,
    )

    force_energy_block = update_force_energy_block(
        ky_mat_old,
        n_envs_prev,
        hyps,
        name,
        force_energy_kernel,
        energy_noise,
        cutoffs,
        hyps_mask,
        n_cpus,
        n_sample,
    )

    ky_mat[0:size1, 0:size1] = force_block
    ky_mat[size1:, size1:] = energy_block
    ky_mat[0:size1, size1:] = force_energy_block
    ky_mat[size1:, 0:size1] = force_energy_block.transpose()

    return ky_mat def get_ky_and_hyp(
    hyps: np.ndarray,
    name,
    kernel_grad,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
    """
    parallel version of get_ky_and_hyp

    :param hyps: list of hyper-parameters
    :param name: name of the gp instance.
    :param kernel_grad: function object of the kernel gradient
    :param cutoffs: The cutoff values used for the atomic environments
    :type cutoffs: list of 2 float numbers
    :param hyps_mask: dictionary used for multi-group hyperparmeters
    :param n_cpus: number of cpus to use.
    :param n_sample: the size of block for matrix to compute

    :return: hyp_mat, ky_mat
    """

    training_data = _global_training_data[name]
    size = len(training_data)
    size3 = size * 3

    sigma_n, non_noise_hyps, train_noise = obtain_noise_len(hyps, hyps_mask)

    if n_cpus is None:
        n_cpus = mp.cpu_count()
    if n_cpus == 1:
        hyp_mat0, k_mat = get_ky_and_hyp_pack(
            name, 0, size, 0, size, True, hyps, kernel_grad, cutoffs, hyps_mask
        )
    else:

        block_id, nbatch = partition_matrix(n_sample, size, n_cpus)

        result_queue = mp.Queue()
        children = []
        for wid in range(nbatch):
            s1, e1, s2, e2 = block_id[wid]
            children.append(
                mp.Process(
                    target=queue_wrapper,
                    args=(
                        result_queue,
                        wid,
                        get_ky_and_hyp_pack,
                        (
                            name,
                            s1,
                            e1,
                            s2,
                            e2,
                            s1 == s2,
                            hyps,
                            kernel_grad,
                            cutoffs,
                            hyps_mask,
                        ),
                    ),
                )
            )

        # Run child processes.
        for c in children:
            c.start()

        k_mat = np.zeros([size3, size3])
        hyp_mat0 = np.zeros([non_noise_hyps, size3, size3])
        for _ in range(nbatch):
            wid, result_chunk = result_queue.get(block=True)
            s1, e1, s2, e2 = block_id[wid]
            h_mat_block, k_mat_block = result_chunk
            k_mat[s1 * 3 : e1 * 3, s2 * 3 : e2 * 3] = k_mat_block
            hyp_mat0[:, s1 * 3 : e1 * 3, s2 * 3 : e2 * 3] = h_mat_block
            if s1 != s2:
                k_mat[s2 * 3 : e2 * 3, s1 * 3 : e1 * 3] = k_mat_block.T
                for idx in range(hyp_mat0.shape[0]):
                    hyp_mat0[idx, s2 * 3 : e2 * 3, s1 * 3 : e1 * 3] = h_mat_block[idx].T

        # Join child processes (clean up zombies).
        for c in children:
            c.join()

    # add gradient of noise variance
    if train_noise:
        hyp_mat = np.zeros(
            [hyp_mat0.shape[0] + 1, hyp_mat0.shape[1], hyp_mat0.shape[2]]
        )
        hyp_mat[:-1, :, :] = hyp_mat0
        hyp_mat[-1, :, :] = np.eye(size3) * 2 * sigma_n
    else:
        hyp_mat = hyp_mat0

    # matrix manipulation
    ky_mat = k_mat + sigma_n ** 2 * np.eye(size3)

    return hyp_mat, ky_mat def get_energy_block(
    hyps: np.ndarray,
    name: str,
    kernel,
    energy_noise,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):
    training_structures = _global_training_structures[name]
    size = len(training_structures)

    if n_cpus is None:
        n_cpus = mp.cpu_count()
    if n_cpus == 1:
        k_mat = get_energy_block_pack(
            hyps, name, 0, size, 0, size, True, kernel, cutoffs, hyps_mask
        )
    else:
        block_id, nbatch = partition_matrix(n_sample, size, n_cpus)
        mult = 1

        k_mat = parallel_matrix_construction(
            get_energy_block_pack,
            hyps,
            name,
            kernel,
            cutoffs,
            hyps_mask,
            block_id,
            nbatch,
            size,
            size,
            mult,
            mult,
        )

    energy_block = k_mat
    energy_block += (energy_noise ** 2) * np.eye(size)

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

    training_data = _global_training_data[name]
    training_structures = _global_training_structures[name]
    size1 = len(training_data) * 3
    size2 = len(training_structures)

    # Initialize Ky.
    ky_mat = np.zeros((size1 + size2, size1 + size2))

    # Assemble the full covariance matrix block-by-block.
    force_block = get_force_block(
        hyps, name, force_kernel, cutoffs, hyps_mask, n_cpus, n_sample
    )

    energy_block = get_energy_block(
        hyps, name, energy_kernel, energy_noise, cutoffs, hyps_mask, n_cpus, n_sample
    )

    force_energy_block = get_force_energy_block(
        hyps, name, force_energy_kernel, cutoffs, hyps_mask, n_cpus, n_sample
    )

    ky_mat[0:size1, 0:size1] = force_block
    ky_mat[size1:, size1:] = energy_block
    ky_mat[0:size1, size1:] = force_energy_block
    ky_mat[size1:, 0:size1] = force_energy_block.transpose()

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

    old_size = n_envs_prev
    old_size3 = n_envs_prev * 3
    mult = 3
    size = len(_global_training_data[name])
    size3 = 3 * size

    if n_cpus is None:
        n_cpus = mp.cpu_count()

    # serial version
    if n_cpus == 1:
        force_block = np.zeros((size3, size3))
        new_mat = get_force_block_pack(
            hyps, name, 0, size, old_size, size, False, kernel, cutoffs, hyps_mask
        )
        force_block[:, old_size3:] = new_mat
        force_block[old_size3:, :] = new_mat.transpose()

    # parallel version
    else:
        block_id, nbatch = partition_update(n_sample, size, old_size, n_cpus)

        force_block = parallel_matrix_construction(
            get_force_block_pack,
            hyps,
            name,
            kernel,
            cutoffs,
            hyps_mask,
            block_id,
            nbatch,
            size3,
            size3,
            mult,
            mult,
        )

    # insert previous covariance matrix
    force_block[:old_size3, :old_size3] = ky_mat_old[:old_size3, :old_size3]

    # add the noise parameter
    sigma_n, _, _ = obtain_noise_len(hyps, hyps_mask)
    force_block[old_size3:, old_size3:] += sigma_n ** 2 * np.eye(size3 - old_size3)

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

    old_size = ky_mat_old.shape[0] - 3 * n_envs_prev
    mult = 1
    size = len(_global_training_structures[name])

    if n_cpus is None:
        n_cpus = mp.cpu_count()

    # serial version
    if n_cpus == 1:
        energy_block = np.zeros((size, size))
        new_mat = get_energy_block_pack(
            hyps, name, 0, size, old_size, size, False, kernel, cutoffs, hyps_mask
        )
        energy_block[:, old_size:] = new_mat
        energy_block[old_size:, :] = new_mat.transpose()

    # parallel version
    else:
        block_id, nbatch = partition_update(n_sample, size, old_size, n_cpus)

        energy_block = parallel_matrix_construction(
            get_energy_block_pack,
            hyps,
            name,
            kernel,
            cutoffs,
            hyps_mask,
            block_id,
            nbatch,
            size,
            size,
            mult,
            mult,
        )

    # insert previous covariance matrix (if it has nonzero size)
    if old_size > 0:
        energy_block[:old_size, :old_size] = ky_mat_old[-old_size:, -old_size:]

    # add the noise parameter
    energy_block[old_size:, old_size:] += (energy_noise ** 2) * np.eye(size - old_size)

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

    n_strucs_prev = ky_mat_old.shape[0] - 3 * n_envs_prev
    n_envs = len(_global_training_data[name])
    n_strucs = len(_global_training_structures[name])
    force_energy_block = np.zeros((n_envs * 3, n_strucs))

    if n_cpus is None:
        n_cpus = mp.cpu_count()

    # serial version
    if n_cpus == 1:
        mat_1 = get_force_energy_block_pack(
            hyps,
            name,
            0,
            n_envs,
            n_strucs_prev,
            n_strucs,
            False,
            kernel,
            cutoffs,
            hyps_mask,
        )

        mat_2 = get_force_energy_block_pack(
            hyps,
            name,
            n_envs_prev,
            n_envs,
            0,
            n_strucs_prev,
            False,
            kernel,
            cutoffs,
            hyps_mask,
        )

    # parallel version
    else:
        force_energy_block = np.zeros((n_envs * 3, n_strucs))

        block_id_1, nbatch_1 = partition_matrix_custom(
            n_sample, 0, n_envs, n_strucs_prev, n_strucs, n_cpus
        )
        block_id_2, nbatch_2 = partition_matrix_custom(
            n_sample, n_envs_prev, n_envs, 0, n_strucs_prev, n_cpus
        )

        size1 = n_envs * 3
        size2 = n_strucs
        m1 = 3
        m2 = 1

        mat_1 = parallel_matrix_construction(
            get_force_energy_block_pack,
            hyps,
            name,
            kernel,
            cutoffs,
            hyps_mask,
            block_id_1,
            nbatch_1,
            size1,
            size2,
            m1,
            m2,
            False,
        )
        mat_2 = parallel_matrix_construction(
            get_force_energy_block_pack,
            hyps,
            name,
            kernel,
            cutoffs,
            hyps_mask,
            block_id_2,
            nbatch_2,
            size1,
            size2,
            m1,
            m2,
            False,
        )

        # reduce the size of the matrices
        mat_1 = mat_1[: n_envs * 3, n_strucs_prev:n_strucs]
        mat_2 = mat_2[n_envs_prev * 3 : n_envs * 3, :n_strucs_prev]

    force_energy_block[: n_envs * 3, n_strucs_prev:n_strucs] = mat_1
    force_energy_block[n_envs_prev * 3 : n_envs * 3, :n_strucs_prev] = mat_2

    # insert previous covariance matrix
    force_energy_block[: n_envs_prev * 3, :n_strucs_prev] = ky_mat_old[
        : n_envs_prev * 3, n_envs_prev * 3 : n_envs_prev * 3 + n_strucs_prev
    ]

    return force_energy_block def efs_kern_vec(
    name,
    efs_force_kernel,
    efs_energy_kernel,
    x,
    hyps,
    cutoffs=None,
    hyps_mask=None,
    n_cpus=1,
    n_sample=100,
):

    size1 = len(_global_training_data[name])
    size2 = len(_global_training_structures[name])
    tot_size = size1 * 3 + size2

    # Initialize arrays.
    energy_vector = np.zeros(tot_size)
    force_array = np.zeros((3, tot_size))
    stress_array = np.zeros((6, tot_size))

    # Compute force and energy arrays.
    force_arrays = efs_force_vector(
        name, efs_force_kernel, x, hyps, cutoffs, hyps_mask, n_cpus, n_sample
    )

    energy_arrays = efs_energy_vector(
        name, efs_energy_kernel, x, hyps, cutoffs, hyps_mask, n_cpus, n_sample
    )

    # Populate arrays.
    energy_vector[0 : size1 * 3] = force_arrays[0]
    energy_vector[size1 * 3 :] = energy_arrays[0]
    force_array[:, 0 : size1 * 3] = force_arrays[1]
    force_array[:, size1 * 3 :] = energy_arrays[1]
    stress_array[:, 0 : size1 * 3] = force_arrays[2]
    stress_array[:, size1 * 3 :] = energy_arrays[2]

    return energy_vector, force_array, stress_array def get_tstp(hm=None) -> AtomicEnvironment:
    """Create test point for kernel to compare against"""
    # params
    cell = np.eye(3)
    unique_species = [2, 1]
    cutoffs = {"twobody": 0.8, "threebody": 0.8, "manybody": 0.8}
    noa = 10

    test_structure_2, _ = get_random_structure(cell, unique_species, noa)

    test_pt = AtomicEnvironment(test_structure_2, 0, cutoffs, cutoffs_mask=hm)
    return test_pt
import time
import pytest
import numpy as np
import flare.gp_algebra
    import cProfile
    import re
from flare.env import AtomicEnvironment
from flare.kernels.utils import str_to_kernel_set
from flare.struc import Structure
from flare.kernels.mc_simple import (
    two_plus_three_body_mc,
    two_plus_three_body_mc_grad,
    two_plus_three_mc_en,
    two_plus_three_mc_force_en,
    two_plus_three_efs_energy,
    two_plus_three_efs_force,
)
from flare.kernels.mc_sephyps import (
    two_plus_three_body_mc as two_plus_three_body_mc_multi,
)
from flare.kernels.mc_sephyps import (
    two_plus_three_body_mc_grad as two_plus_three_body_mc_grad_multi,
)
from flare.kernels import mc_sephyps
from flare.utils.parameter_helper import ParameterHelper
from flare.gp_algebra import (
    get_like_grad_from_mats,
    get_force_block,
    get_force_energy_block,
    energy_energy_vector,
    energy_force_vector,
    force_force_vector,
    force_energy_vector,
    get_kernel_vector,
    en_kern_vec,
    get_ky_mat_update,
    get_ky_and_hyp,
    get_energy_block,
    get_Ky_mat,
    update_force_block,
    update_energy_block,
    update_force_energy_block,
    efs_kern_vec,
)
from tests.fake_gp import get_tstp





@pytest.fixture(scope="module")
def params():

    parameters = get_random_training_set(10, 2)

    yield parameters
    del parameters


def get_random_training_set(nenv, nstruc):
    """Create a random training_set array with parameters
    And generate four different kinds of hyperparameter sets:
    * multi hypper parameters with two twobody type and two threebody type
    * constrained optimization, with noise parameter optimized
    * constrained optimization, without noise parameter optimized
    * simple hyper parameters without multihyps set up
    """

    np.random.seed(0)

    cutoffs = {"twobody": 0.8, "threebody": 0.8}
    parameters = {"cutoff_twobody": 0.8, "cutoff_threebody": 0.8, "noise": 0.05}

    # 9 different hyper-parameters
    pm = ParameterHelper(
        species=["H", "He"],
        kernels={
            "twobody": [["*", "*"], ["H", "H"]],
            "threebody": [["*", "*", "*"], ["H", "H", "H"]],
        },
        parameters=parameters,
        ones=True,
        random=False,
        verbose="DEBUG",
    )
    hyps_mask1 = pm.as_dict()

    # 9 different hyper-parameters, onlye train the 0, 2, 4, 6, 8
    pm.set_constraints("twobody0", [True, True])
    pm.set_constraints("twobody1", [False, False])
    pm.set_constraints("threebody0", [True, True])
    pm.set_constraints("threebody1", [False, False])
    hyps_mask2 = pm.as_dict()

    # 9 different hyper-parameters, only train the 0, 2, 4, 6
    pm.set_constraints("noise", False)
    hyps_mask3 = pm.as_dict()

    # 5 different hyper-parameters, equivalent to no multihyps
    pm = ParameterHelper(
        species=["H", "He"],
        kernels={"twobody": [["*", "*"]], "threebody": [["*", "*", "*"]]},
        parameters=parameters,
        ones=True,
        random=False,
        verbose="DEBUG",
    )
    hyps_mask4 = pm.as_dict()

    # 5 different hyper-parameters, no multihyps
    pm = ParameterHelper(
        kernels=["twobody", "threebody"],
        parameters=parameters,
        ones=True,
        random=False,
        verbose="DEBUG",
    )
    hyps_mask5 = pm.as_dict()

    hyps_mask_list = [hyps_mask1, hyps_mask2, hyps_mask3, hyps_mask4, hyps_mask5]

    # create training environments and forces
    cell = np.eye(3)
    unique_species = [0, 1]
    noa = 5
    training_data = []
    training_labels = []
    for _ in range(nenv):
        positions = np.random.uniform(-1, 1, [noa, 3])
        species = np.random.randint(0, len(unique_species), noa) + 1

        struc = Structure(cell, species, positions)