Skip to content

Commit

Permalink
Merge pull request #1 from mir-group/punchout
Browse files Browse the repository at this point in the history 
Add punchout-related files, methods, and functions
  • Loading branch information
@stevetorr
stevetorr authored Sep 10, 2018
2 parents 0ed5988 + 092b9f1 commit 153516f
Show file tree
Hide file tree
Showing 10 changed files with 805 additions and 216 deletions.
219 changes: 53 additions & 166 deletions src/otf.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand Down Expand Up @@ -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]
Expand All @@ -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' +
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Loading

0 comments on commit 153516f

Please sign in to comment.