Merge pull request #128 from mir-group/feature/jon/gp-energy

Training on energy labels

flare/dft_interface/qe_util.py

@@ -87,6 +87,18 @@ def run_dft_en_par(dft_input, structure, dft_loc, n_cpus):
     return forces, energy
 
 
+def run_dft_en_npool(qe_input, structure, dft_loc, npool):
+    run_qe_path = qe_input
+    edit_dft_input_positions(run_qe_path, structure)
+    qe_command = \
+        'mpirun {0} -npool {1} < {2} > {3}'.format(dft_loc, npool, run_qe_path,
+                                                   'pwscf.out')
+    call(qe_command, shell=True)
+    forces, energy = parse_dft_forces_and_energy('pwscf.out')
+
+    return forces, energy
+
+
 def parse_dft_input(dft_input: str):
     """ parse the input to get information of atomic configuration
 

flare/env.py

@@ -8,25 +8,25 @@
 
 
 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.
+    """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.
+    given, 2+3-body if two are passed.
     :type cutoffs: np.ndarray
     """
 
-    def __init__(self, structure: Structure, atom: int, cutoffs):
+    def __init__(self, structure: Structure, atom: int, cutoffs, sweep=1):
         self.structure = structure
         self.positions = structure.wrapped_positions
         self.cell = structure.cell
         self.species = structure.coded_species
+        self.sweep_array = np.arange(-sweep, sweep+1, 1)
 
         self.atom = atom
         self.ctype = structure.coded_species[atom]
@@ -40,31 +40,33 @@ def compute_env(self):
         # get 2-body arrays
         bond_array_2, bond_positions_2, etypes = \
             get_2_body_arrays(self.positions, self.atom, self.cell,
-                              self.cutoffs[0], self.species)
+                              self.cutoffs[0], self.species, self.sweep_array)
         self.bond_array_2 = bond_array_2
         self.etypes = etypes
 
         # if 2 cutoffs are given, create 3-body arrays
         if len(self.cutoffs) > 1:
             bond_array_3, cross_bond_inds, cross_bond_dists, triplet_counts = \
-                get_3_body_arrays(bond_array_2, bond_positions_2, self.cutoffs[1])
+                get_3_body_arrays(bond_array_2, bond_positions_2,
+                                  self.cutoffs[1])
             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 len(self.cutoffs) > 2:
-            self.bond_array_mb, self.neigh_dists_mb, self.num_neighs_mb, self.etype_mb, \
-                    self.bond_array_mb_etypes = get_m_body_arrays(
-                self.positions, self.atom, self.cell, self.cutoffs[2], self.species)
+            self.bond_array_mb, self.neigh_dists_mb, self.num_neighs_mb, \
+                self.etype_mb, self.bond_array_mb_etypes = \
+                get_m_body_arrays(self.positions, self.atom,
+                                  self.cell, self.cutoffs[2],
+                                  self.species, self.sweep_array)
         else:
             self.bond_array_mb = None
             self.neigh_dists_mb = None
             self.num_neighs_mb = None
             self.etype_mb = None
 
-
     def as_dict(self):
         """
         Returns Atomic Environment object as a dictionary for serialization
@@ -123,7 +125,8 @@ def __str__(self):
 
 
 @njit
-def get_2_body_arrays(positions, atom: int, cell, cutoff_2: float, species):
+def get_2_body_arrays(positions: np.ndarray, atom: int, cell: np.ndarray,
+                      cutoff_2: float, species: np.ndarray, sweep: np.ndarray):
     """Returns distances, coordinates, and species of atoms in the 2-body
     local environment. This method is implemented outside the AtomicEnvironment
     class to allow for njit acceleration with Numba.
@@ -154,12 +157,13 @@ class to allow for njit acceleration with Numba.
      their atomic number.
     :rtype: np.ndarray, np.ndarray, np.ndarray
     """
+
     noa = len(positions)
     pos_atom = positions[atom]
-    coords = np.zeros((noa, 3, 27), dtype=np.float64)
-    dists = np.zeros((noa, 27), dtype=np.float64)
+    super_count = sweep.shape[0]**3
+    coords = np.zeros((noa, 3, super_count))
+    dists = np.zeros((noa, super_count))
     cutoff_count = 0
-    super_sweep = np.array([-1, 0, 1], dtype=np.float64)
 
     vec1 = cell[0]
     vec2 = cell[1]
@@ -169,9 +173,9 @@ class to allow for njit acceleration with Numba.
     for n in range(noa):
         diff_curr = positions[n] - pos_atom
         im_count = 0
-        for s1 in super_sweep:
-            for s2 in super_sweep:
-                for s3 in super_sweep:
+        for s1 in sweep:
+            for s2 in sweep:
+                for s3 in sweep:
                     im = diff_curr + s1 * vec1 + s2 * vec2 + s3 * vec3
                     dist = sqrt(im[0] * im[0] + im[1] * im[1] + im[2] * im[2])
                     if (dist < cutoff_2) and (dist != 0):
@@ -188,7 +192,7 @@ class to allow for njit acceleration with Numba.
 
     for m in range(noa):
         spec_curr = species[m]
-        for n in range(27):
+        for n in range(super_count):
             dist_curr = dists[m, n]
             if (dist_curr < cutoff_2) and (dist_curr != 0):
                 coord = coords[m, :, n]
@@ -371,7 +375,8 @@ def get_3_body_arrays(bond_array_2, bond_positions_2, cutoff_3: float):
 
 
 @njit
-def get_m_body_arrays(positions, atom: int, cell, cutoff_mb: float, species):
+def get_m_body_arrays(positions, atom: int, cell, cutoff_mb: float, species,
+                      sweep: np.ndarray):
     """Returns distances, and species of atoms in the many-body
     local environment, and returns distances and numbers of neighbours for atoms in the one
     many-body local environment. This method is implemented outside the AtomicEnvironment
@@ -421,23 +426,23 @@ class to allow for njit acceleration with Numba.
     neighbouring_etypes = []
     max_neighbours = 0
     for m in bond_inds:
-        neighbour_bond_array_2, ___, etypes_mb = get_2_body_arrays(positions, m, cell,
-                                                         cutoff_mb, species)
+        neighbour_bond_array_2, ___, etypes_mb = \
+            get_2_body_arrays(positions, m, cell, cutoff_mb, species, sweep)
         neighbouring_dists.append(neighbour_bond_array_2[:, 0])
         neighbouring_etypes.append(etypes_mb)
         if len(neighbour_bond_array_2[:, 0]) > max_neighbours:
             max_neighbours = len(neighbour_bond_array_2[:, 0])
 
     # Transform list of distances into Numpy array
-    neigh_dists_mb = np.zeros((len(bond_inds), max_neighbours), dtype=np.float64)
+    neigh_dists_mb = \
+        np.zeros((len(bond_inds), max_neighbours), dtype=np.float64)
     num_neighs_mb = np.zeros(len(bond_inds), dtype=np.int8)
     etypes_mb_array = np.zeros((len(bond_inds), max_neighbours), dtype=np.int8)
     for i in range(len(bond_inds)):
         num_neighs_mb[i] = len(neighbouring_dists[i])
         neigh_dists_mb[i, :num_neighs_mb[i]] = neighbouring_dists[i]
         etypes_mb_array[i, :num_neighs_mb[i]] = neighbouring_etypes[i]
 
-
     return bond_array_mb, neigh_dists_mb, num_neighs_mb, etypes_mb_array, etypes
 
 

flare/flare_io.py

@@ -19,5 +19,6 @@ def md_trajectory_from_file(filename: str):
     """
     with open(filename, 'r') as f:
         structure_list = load(f)
-        structures = [Structure.from_dict(dictionary) for dictionary in structure_list]
+        structures = \
+            [Structure.from_dict(dictionary) for dictionary in structure_list]
     return structures