Add punchout-related files, methods, and functions

src/otf.py

@@ -8,17 +8,18 @@
 """
 
 import numpy as np
-import os
 import datetime
 from struc import Structure
 from math import sqrt
 from gp import GaussianProcess
 from env import ChemicalEnvironment
+from punchout import punchout
+from qe_util import run_espresso, parse_qe_input
 
 
 class OTF(object):
     def __init__(self, qe_input: str, dt: float, number_of_steps: int,
-                 kernel: str, cutoff: float):
+                 kernel: str, cutoff: float, punchout_d: float = None):
         """
         On-the-fly learning engine, containing methods to run OTF calculation
 
@@ -27,19 +28,24 @@ def __init__(self, qe_input: str, dt: float, number_of_steps: int,
         :param number_of_steps: int, Number of steps
         :param kernel: Type of kernel for GP regression model
         :param cutoff: Cutoff radius for kernel
+        :param punchout_d: Box distance around a high-uncertainty atom to
+                         punch out
         """
 
         self.qe_input = qe_input
         self.dt = dt
         self.Nsteps = number_of_steps
         self.gp = GaussianProcess(kernel)
+        self.cutoff = cutoff
+        self.punchout_d = punchout_d
 
         positions, species, cell, masses = parse_qe_input(self.qe_input)
         self.structure = Structure(lattice=cell, species=species,
                                    positions=positions, cutoff=cutoff,
                                    mass_dict=masses)
 
         self.curr_step = 0
+        self.train_structure = None
 
         # Create blank output file with time header and structure information
         with open('otf_run.out', 'w') as f:
@@ -88,15 +94,26 @@ def run_and_train(self):
 
         # Run espresso and write out results
         self.write_to_output('=' * 20 + '\n' + 'Calling QE... ')
-        forces = self.run_espresso()
+
+        # If not in punchout mode, run QE on the entire structure
+        if self.punchout_d is None:
+            self.train_structure = self.structure
+        # If first run, pick a random atom to punch out a structure around
+        elif self.train_structure is None:
+            target_atom = np.random.randint(0, len(self.structure.positions))
+            self.train_structure = punchout(self.structure, target_atom,
+                                            d=self.punchout_d)
+
+        forces = run_espresso(self.qe_input, self.train_structure)
+
         self.write_to_output('Done.\n')
 
         # Write input positions and force results
         qe_strings = 'Resultant Positions and Forces:\n'
-        for n in range(self.structure.nat):
-            qe_strings += self.structure.species[n] + ': '
+        for n in range(self.train_structure.nat):
+            qe_strings += self.train_structure.species[n] + ': '
             for i in range(3):
-                qe_strings += '%.3f  ' % self.structure.positions[n][i]
+                qe_strings += '%.3f  ' % self.train_structure.positions[n][i]
             qe_strings += '\t '
             for i in range(3):
                 qe_strings += '%.3f  ' % forces[n][i]
@@ -105,7 +122,7 @@ def run_and_train(self):
 
         # Update hyperparameters and write results
         self.write_to_output('Updating database hyperparameters...\n')
-        self.gp.update_db(self.structure, forces)
+        self.gp.update_db(self.train_structure, forces)
         self.gp.train()
         self.write_to_output('New GP Hyperparameters:\n' +
                              'Signal variance: \t' +
@@ -121,64 +138,22 @@ def update_positions(self):
         Apply a timestep to self.structure based on current structure's forces.
         """
 
-        # Maintain list of elemental masses in amu to calculate acceleration
+        # Precompute dt squared for efficiency
+        dtdt = self.dt ** 2
+
+        for i, pre_pos in enumerate(self.structure.prev_positions):
+            temp_pos = np.array(self.structure.positions[i])
+            mass = self.structure.mass_dict[self.structure.species[i]]
+            pos = self.structure.positions[i]
+            forces = self.structure.forces[i]
 
-        for i, prev_pos in enumerate(self.structure.prev_positions):
-            temp_pos = self.structure.positions[i]
-            self.structure.positions[i] = 2 * self.structure.positions[i] - \
-                prev_pos + self.dt ** 2 * self.structure.forces[i] / \
-                self.structure.mass_dict[self.structure.species[i]]
+            self.structure.positions[i] = 2 * pos - pre_pos + dtdt * forces / \
+                                                              mass
 
             self.structure.prev_positions[i] = np.array(temp_pos)
 
-    def run_espresso(self):
-        """
-        Calls quantum espresso from input located at self.qe_input
-
-        :return: List [nparray] List of forces
-        """
-
-        self.write_to_output('')
-        self.write_qe_input_positions()
-
-        pw_loc = os.environ.get('PWSCF_COMMAND', 'pwscf.x')
-
-        qe_command = '{0} < {1} > {2}'.format(pw_loc, self.qe_input,
-                                              'pwscf.out')
-
-        os.system(qe_command)
-
-        return parse_qe_forces('pwscf.out')
-
-    def write_qe_input_positions(self):
-        """
-        Write the current configuration of the OTF structure to the
-        qe input file
-        """
-
-        with open(self.qe_input, 'r') as f:
-            lines = f.readlines()
-
-        file_pos_index = None
-        for i, line in enumerate(lines):
-            if 'ATOMIC_POSITIONS' in line:
-                file_pos_index = int(i + 1)
-        assert file_pos_index is not None, 'Failed to find positions in input'
-
-        for pos_index, line_index in enumerate(
-                range(file_pos_index, len(lines))):
-            pos_string = ' '.join([self.structure.species[pos_index],
-                                   str(self.structure.positions[pos_index][0]),
-                                   str(self.structure.positions[pos_index][1]),
-                                   str(self.structure.positions[pos_index][
-                                           2])])
-            lines[line_index] = str(pos_string + '\n')
-
-        with open(self.qe_input, 'w') as f:
-            for line in lines:
-                f.write(line)
-
-    def write_to_output(self, string: str, output_file: str = 'otf_run.out'):
+    @staticmethod
+    def write_to_output(string: str, output_file: str = 'otf_run.out'):
         """
         Write a string or list of strings to the output file.
         :param string: String to print to output
@@ -218,8 +193,9 @@ def write_config(self):
                 string += str('%.6e' % self.structure.stds[i][j]) + ' '
             string += '\n'
 
-            self.write_to_output(string)
+        self.write_to_output(string)
 
+    # TODO change this to use the signal variance
     def is_std_in_bound(self):
         """
         Evaluate if the model error are within predefined criteria
@@ -229,110 +205,21 @@ def is_std_in_bound(self):
 
         # Some decision making criteria
 
-        return True
-
-
-def parse_qe_input(qe_input: str):
-    """
-    Reads the positions, species, and cell in from the qe input file
-
-    :param qe_input: str, Path to PWSCF input file
-    :return: List[nparray], List[str], nparray, Positions, species,
-                    3x3 Bravais cell
-    """
-    positions = []
-    species = []
-    cell = []
-
-    with open(qe_input) as f:
-        lines = f.readlines()
-
-    # Find the cell and positions in the output file
-    cell_index = None
-    positions_index = None
-    nat = None
-    species_index = None
-
-    for i, line in enumerate(lines):
-        if 'CELL_PARAMETERS' in line:
-            cell_index = int(i + 1)
-        if 'ATOMIC_POSITIONS' in line:
-            positions_index = int(i + 1)
-        if 'nat' in line:
-            nat = int(line.split('=')[1])
-        if 'ATOMIC_SPECIES' in line:
-            species_index = int(i + 1)
-
-    assert cell_index is not None, 'Failed to find cell in input'
-    assert positions_index is not None, 'Failed to find positions in input'
-    assert nat is not None, 'Failed to find number of atoms in input'
-    assert species_index is not None, 'Failed to find atomic species in input'
-
-    # Load cell
-    for i in range(cell_index, cell_index + 3):
-        cell_line = lines[i].strip()
-        cell.append(np.fromstring(cell_line, sep=' '))
-    cell = np.array(cell)
-
-    # Check cell IO
-    assert cell != [], 'Cell failed to load'
-    assert np.shape(cell) == (3, 3), 'Cell failed to load correctly'
-
-    # Load positions
-    for i in range(positions_index, positions_index + nat):
-        line_string = lines[i].strip().split()
-        species.append(line_string[0])
-
-        pos_string = ' '.join(line_string[1:4])
-
-        positions.append(np.fromstring(pos_string, sep=' '))
-    # Check position IO
-    assert positions != [], "Positions failed to load"
-
-    # Load masses
-    masses = {}
-    for i in range(species_index, species_index + len(set(species))):
-        # Expects lines of format like: H 1.0 H_pseudo_name.ext
-        line = lines[i].strip().split()
-        masses[line[0]] = float(line[1])
-
-    return positions, species, cell, masses
-
-
-def parse_qe_forces(outfile: str):
-    """
-    Get forces from a pwscf file in Ryd/bohr
-
-    :param outfile: str, Path to pwscf output file
-    :return: list[nparray] , List of forces acting on atoms
-    """
-    forces = []
-    total_energy = np.nan
-
-    with open(outfile, 'r') as outf:
-        for line in outf:
-            if line.lower().startswith('!    total energy'):
-                total_energy = float(line.split()[-2]) * 13.605698066
-
-            if line.find('force') != -1 and line.find('atom') != -1:
-                line = line.split('force =')[-1]
-                line = line.strip()
-                line = line.split(' ')
-                line = [x for x in line if x != '']
-                temp_forces = []
-                for x in line:
-                    temp_forces.append(float(x))
-                forces.append(np.array(list(temp_forces)))
-
-    assert total_energy != np.nan, "Quantum ESPRESSO parser failed to read " \
-                                   "the file {}. Run failed.".format(outfile)
+        # SKETCH OF EVENTUAL CODE:
+        """
+        high_error_atom = -1 
+
+        if self.punchout_d is not None:
+            self.train_structure= punchout(self.structure,atom= 
+            high_error_atom, d= self.punchout_d )
+            return False
+        """
 
-    return forces
+        return True
 
 
-# TODO Currently won't work: needs to be re-done in light of new output
-# formatting, but could wait until we finalize an output file format ,
-# will be useful for data analysis
+# TODO Currently won't work: needs to be re-done when we finalize our output
+#  formatting
 def parse_otf_output(outfile: str):
     """
     Parse the output of a otf run for analysis
@@ -378,8 +265,8 @@ def parse_otf_output(outfile: str):
 if __name__ == '__main__':
     # os.system('cp ../tests/test_files/qe_input_2.in pwscf.in')
 
-    # otf = OTF('pwscf.in', .1, 10, kernel='two_body',
-    #          cutoff=10)
-    # otf.run()
+    otf = OTF('pwscf.in', .1, 10, kernel='two_body',
+              cutoff=10, punchout_d=None)
+    otf.run()
     # parse_output('otf_run.out')
     pass