Merge pull request #259 from mir-group/bugfix/yu/uncertainty_print

Bugfix/yu/uncertainty print

flare/ase/atoms.py

@@ -31,33 +31,10 @@ def from_ase_atoms(atoms):
         Args:
             atoms (ASE Atoms): the ase atoms to build from
         """
-        kw = [
-            "symbols",  # use either "symbols" or "number"
-            "positions",
-            "tags",
-            "momenta",
-            "masses",
-            "magmoms",
-            "charges",
-            "scaled_positions",
-            "cell",
-            "pbc",
-            "celldisp",
-            "constraint",
-            "calculator",
-            "info",
-            "velocities",
-        ]
-        # The keywords are either the attr of atoms, or in dict atoms.arrays
-        kwargs = {}
-        for key in kw:
-            try:
-                kwargs[key] = getattr(atoms, key)
-            except:
-                if key in atoms.arrays:
-                    kwargs[key] = atoms.arrays[key]
-        kwargs["calculator"] = atoms.calc
-        return FLARE_Atoms(**kwargs)
+        new_atoms = deepcopy(atoms)
+        new_atoms.__class__ = FLARE_Atoms
+        new_atoms.prev_positions = np.zeros_like(new_atoms.positions)
+        return new_atoms
 
     @property
     def nat(self):

flare/ase/calculator.py

@@ -154,7 +154,7 @@ def calculate_mgp(self, atoms):
                 f, v, vir, e = self.mgp_model.predict(chemenv)
                 self.results["forces"][n] = f
                 self.results["partial_stresses"][n] = vir
-                self.results["stds"][n] = np.sqrt(np.absolute(v))
+                self.results["stds"][n][0] = np.sqrt(np.absolute(v))
                 self.results["local_energies"][n] = e
 
             except ValueError as err_msg:  # if lower_bound error is raised
@@ -179,7 +179,7 @@ def calculate_mgp(self, atoms):
                 f, v, vir, e = self.mgp_model.predict(chemenv)
                 self.results["forces"][n] = f
                 self.results["partial_stresses"][n] = vir
-                self.results["stds"][n] = np.sqrt(np.absolute(v))
+                self.results["stds"][n][0] = np.sqrt(np.absolute(v))
                 self.results["local_energies"][n] = e
 
     def calculation_required(self, atoms, quantities):

flare/ase/otf.py

@@ -230,7 +230,8 @@ def update_temperature(self):
         self.velocities = self.atoms.get_velocities() * units.fs * 1e3
 
     def update_gp(self, train_atoms, dft_frcs, dft_energy=None, dft_stress=None):
-        self.output.add_atom_info(train_atoms, self.structure.stds)
+        stds = self.flare_calc.results.get("stds", np.zeros_like(dft_frcs))
+        self.output.add_atom_info(train_atoms, stds)
 
         # Convert ASE stress (xx, yy, zz, yz, xz, xy) to FLARE stress
         # (xx, xy, xz, yy, yz, zz).

flare/env.py

@@ -40,15 +40,15 @@ class AtomicEnvironment:
     bonds, triples, and many body interaction. This dictionary should be
     consistent with the hyps_mask used in the GuassianProcess object.
 
-    * specie_mask: 118-long integer array descirbing which elements belong to
+    * 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 specie_mask).
+        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
-        specie_mask.
+        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
@@ -59,12 +59,12 @@ class AtomicEnvironment:
         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
-        specie_mask.
+        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
-        specie_mask and H are associate with group 0 in specie_mask, the
+        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
@@ -123,7 +123,7 @@ def __init__(self, structure: Structure, atom: int, cutoffs, cutoffs_mask=None):
         self.nmanybody = 0
 
         self.nspecie = 1
-        self.specie_mask = None
+        self.species_mask = None
         self.twobody_mask = None
         self.threebody_mask = None
         self.manybody_mask = None
@@ -173,8 +173,8 @@ def setup_mask(self, cutoffs_mask):
             self.cutoffs["twobody"] = self.twobody_cutoff
 
         self.nspecie = cutoffs_mask.get("nspecie", 1)
-        if "specie_mask" in cutoffs_mask:
-            self.specie_mask = np.array(cutoffs_mask["specie_mask"], dtype=np.int)
+        if "species_mask" in cutoffs_mask:
+            self.species_mask = np.array(cutoffs_mask["species_mask"], dtype=np.int)
 
         for kernel in AtomicEnvironment.all_kernel_types:
             ndim = AtomicEnvironment.ndim[kernel]
@@ -211,7 +211,7 @@ def compute_env(self):
                 self.species,
                 self.sweep_array,
                 self.nspecie,
-                self.specie_mask,
+                self.species_mask,
                 self.twobody_mask,
             )
 
@@ -234,7 +234,7 @@ def compute_env(self):
                 self.threebody_cutoff,
                 self.threebody_cutoff_list,
                 self.nspecie,
-                self.specie_mask,
+                self.species_mask,
                 self.cut3b_mask,
             )
 
@@ -261,7 +261,7 @@ def compute_env(self):
                 self.species,
                 self.sweep_array,
                 self.nspecie,
-                self.specie_mask,
+                self.species_mask,
                 self.manybody_mask,
                 cf.quadratic_cutoff,
             )

flare/kernels/utils.py

@@ -175,7 +175,7 @@ def from_mask_to_args(hyps, cutoffs, hyps_mask=None):
         cutoff_3b,
         cutoff_mb,
         nspecie,
-        np.array(hyps_mask["specie_mask"]),
+        np.array(hyps_mask["species_mask"]),
         n2b,
         twobody_mask,
         n3b,

flare/parameters.py

@@ -44,7 +44,7 @@ def __init__(self):
             "nthreebody": 0,
             "ncut3b": 0,
             "nmanybody": 0,
-            "specie_mask": None,
+            "species_mask": None,
             "twobody_mask": None,
             "threebody_mask": None,
             "cut3b_mask": None,
@@ -162,9 +162,9 @@ def check_instantiation(hyps, cutoffs, kernels, param_dict):
         nspecie = param_dict["nspecie"]
         if nspecie > 1:
             assert (
-                "specie_mask" in param_dict
-            ), "specie_mask key missing in param_dict dictionary"
-            param_dict["specie_mask"] = nparray(param_dict["specie_mask"], dtype=np.int)
+                "species_mask" in param_dict
+            ), "species_mask key missing in param_dict dictionary"
+            param_dict["species_mask"] = nparray(param_dict["species_mask"], dtype=np.int)
 
         # for each kernel, check whether it is defined
         # and the length of corresponding hyper-parameters
@@ -343,7 +343,7 @@ def get_component_mask(param_dict, kernel_name, hyps=None):
 
             name_list = [
                 "nspecie",
-                "specie_mask",
+                "species_mask",
                 "n" + kernel_name,
                 kernel_name + "_mask",
                 kernel_name + "_cutoff_list",
@@ -410,13 +410,14 @@ def get_cutoff(kernel_name, coded_species, param_dict):
 
         if f"{kernel_name}_cutoff_list" in param_dict:
 
-            specie_mask = param_dict["species_mask"]
+            species_mask = param_dict["species_mask"]
             cutoff_list = param_dict[f"{kernel_name}_cutoff_list"]
+            nspecie = param_dict["nspecie"]
 
             if kernel_name not in Parameters.cutoff_types_values:
                 mask_id = 0
-                for ele in coded_specie:
-                    mask_id += specie_mask[ele]
+                for ele in coded_species:
+                    mask_id += species_mask[ele]
                     mask_id *= nspecie
                 mask_id = mask_id // nspecie
                 mask_id = param_dict[kernel_name + "_mask"][mask_id]
@@ -497,7 +498,7 @@ def compare_dict(dict1, dict2):
         if dict1 is None:
             return True
 
-        list_of_names = ["nspecie", "specie_mask", "map", "original_hyps"]
+        list_of_names = ["nspecie", "species_mask", "map", "original_hyps"]
         for k in Parameters.all_kernel_types:
             list_of_names += ["n" + k]
             list_of_names += [k + "_mask"]

flare/utils/env_getarray.py

@@ -15,7 +15,7 @@ def get_2_body_arrays(
     species,
     sweep,
     nspecie,
-    specie_mask,
+    species_mask,
     twobody_mask,
 ):
     """Returns distances, coordinates, species of atoms, and indices of neighbors
@@ -35,7 +35,7 @@ def get_2_body_arrays(
     :type species: np.ndarray
     :param nspecie: number of atom types to define bonds
     :type: int
-    :param specie_mask: mapping from atomic number to atom types
+    :param species_mask: mapping from atomic number to atom types
     :type: np.ndarray
     :param twobody_mask: mapping from the types of end atoms to bond types
     :type: np.ndarray
@@ -72,15 +72,15 @@ def get_2_body_arrays(
     bcn = 0
     if nspecie > 1 and cutoff_2 is not None:
         sepcut = True
-        bc = specie_mask[species[atom]]
+        bc = species_mask[species[atom]]
         bcn = nspecie * bc
 
     # record distances and positions of images
     for n in range(noa):
         diff_curr = positions[n] - pos_atom
         im_count = 0
-        if sepcut and (specie_mask is not None) and (cutoff_2 is not None):
-            bn = specie_mask[species[n]]
+        if sepcut and (species_mask is not None) and (cutoff_2 is not None):
+            bn = species_mask[species[n]]
             r_cut = cutoff_2[twobody_mask[bn + bcn]]
 
         for s1 in sweep:
@@ -103,8 +103,8 @@ def get_2_body_arrays(
 
     for m in range(noa):
         spec_curr = species[m]
-        if sepcut and (specie_mask is not None) and (cutoff_2 is not None):
-            bm = specie_mask[species[m]]
+        if sepcut and (species_mask is not None) and (cutoff_2 is not None):
+            bm = species_mask[species[m]]
             r_cut = cutoff_2[twobody_mask[bm + bcn]]
         for im_count in range(super_count):
             dist_curr = dists[m, im_count]
@@ -136,7 +136,7 @@ def get_3_body_arrays(
     r_cut,
     cutoff_3,
     nspecie,
-    specie_mask,
+    species_mask,
     cut3b_mask,
 ):
     """Returns distances and coordinates of triplets of atoms in the
@@ -153,7 +153,7 @@ def get_3_body_arrays(
     :type cutoff_3: np.ndarray
     :param nspecie: number of atom types to define bonds
     :type: int
-    :param specie_mask: mapping from atomic number to atom types
+    :param species_mask: mapping from atomic number to atom types
     :type: np.ndarray
     :param cut3b_mask: mapping from the types of end atoms to bond types
     :type: np.ndarray
@@ -182,7 +182,7 @@ def get_3_body_arrays(
 
     sepcut = False
     if nspecie > 1 and cutoff_3 is not None:
-        bc = specie_mask[ctype]
+        bc = species_mask[ctype]
         bcn = nspecie * bc
         r_cut = np.max(cutoff_3)
         sepcut = True
@@ -212,12 +212,12 @@ def get_3_body_arrays(
 
         if (
             sepcut
-            and (specie_mask is not None)
+            and (species_mask is not None)
             and (cut3b_mask is not None)
             and (cutoff_3 is not None)
         ):
             # choose bond dependent bond
-            bm = specie_mask[etypes[m]]
+            bm = species_mask[etypes[m]]
             btype_m = cut3b_mask[bm + bcn]  # (m, c)
             cut_m = cutoff_3[btype_m]
             bmn = nspecie * bm  # for cross_dist usage
@@ -226,11 +226,11 @@ def get_3_body_arrays(
 
             if (
                 sepcut
-                and (specie_mask is not None)
+                and (species_mask is not None)
                 and (cut3b_mask is not None)
                 and (cutoff_3 is not None)
             ):
-                bn = specie_mask[etypes[n]]
+                bn = species_mask[etypes[n]]
                 btype_n = cut3b_mask[bn + bcn]  # (n, c)
                 cut_n = cutoff_3[btype_n]
 

flare/utils/parameter_helper.py

@@ -831,18 +831,18 @@ def summarize_group(self, group_type):
         if group_type == "specie":
 
             self.nspecie = nspecie
-            self.specie_mask = np.ones(118, dtype=np.int) * (nspecie - 1)
+            self.species_mask = np.ones(118, dtype=np.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.specie_mask):
+                    if atom_n >= len(self.species_mask):
                         new_mask = np.ones(atom_n, dtype=np.int) * (nspecie - 1)
-                        new_mask[: len(self.specie_mask)] = self.specie_mask
+                        new_mask[: len(self.species_mask)] = self.species_mask
                         self.species_mask = new_mask
-                    self.specie_mask[atom_n] = idt
+                    self.species_mask[atom_n] = idt
                     self.logger.debug(
                         f"elemtn {ele} is defined as type {idt} with name {aeg[idt]}"
                     )
@@ -1023,7 +1023,7 @@ def as_dict(self):
 
         hyps_mask["nspecie"] = self.n["specie"]
         if self.n["specie"] > 1:
-            hyps_mask["specie_mask"] = self.specie_mask
+            hyps_mask["species_mask"] = self.species_mask
 
         hyps = []
         hyp_labels = []
@@ -1122,10 +1122,10 @@ def from_dict(hyps_mask, verbose=False, init_spec=[]):
 
         nspecie = hyps_mask["nspecie"]
         if nspecie > 1:
-            max_species = np.max(hyps_mask["specie_mask"])
-            specie_mask = hyps_mask["specie_mask"]
+            max_species = np.max(hyps_mask["species_mask"])
+            species_mask = hyps_mask["species_mask"]
             for i in range(max_species + 1):
-                elelist = np.where(specie_mask == i)[0]
+                elelist = np.where(species_mask == i)[0]
                 if len(elelist) > 0:
                     for ele in elelist:
                         if ele != 0: