def traj(self):
"""Get a trajectory.
This reads Atoms objects from vasprun.xml. By default returns
all images. If index is an integer, return that image.
Technically, this is just a list of atoms with a
SinglePointCalculator attached to them.
This is usually only relevant if you have done a
relaxation. If the calculation is an NEB, the images are
returned.
"""
from ase.calculators.singlepoint import SinglePointCalculator as SPC
self.update()
if self.neb:
images, energies = self.get_neb()
tatoms = [x.copy() for x in images]
for i, x in enumerate(tatoms):
x.set_calculator(SPC(x, energy=energies[i]))
return tatoms
LOA = []
for atoms in read(os.path.join(self.directory, 'vasprun.xml'), ':'):
catoms = atoms.copy()
catoms = catoms[self.resort]
catoms.set_calculator(
SPC(catoms,
energy=atoms.get_potential_energy(),
forces=atoms.get_forces()[self.resort]))
LOA += [catoms]
return LOA
def encode_to_atoms(encode, out_file='input.traj'):
""" Dump the encoding to a local traj file."""
# First, decode the trajectory
data = json.loads(encode, encoding='utf-8')
# Construct the initial atoms object
atoms = Atoms(
data['numbers'],
data['trajectory']['0']['positions'],
cell=data['trajectory']['0']['cell'],
pbc=data['pbc'])
atoms.info = data['calculator_parameters']
atoms.set_constraint([dict2constraint(_) for _ in data['constraints']])
if data.get('initial_magmoms') is not None:
atoms.set_initial_magnetic_moments(data['initial_magmoms'])
# Attach the calculator
calc = SPC(
atoms=atoms,
energy=data['trajectory']['0'].get('energy'),
forces=data['trajectory']['0'].get('forces'),
stress=data['trajectory']['0'].get('stress'))
atoms.set_calculator(calc)
# Collect the rest of the trajectory information
images = [atoms]
for i in range(len(data['trajectory']))[1:]:
atoms = atoms.copy()
if data['trajectory'][str(i)]['cell']:
atoms.set_cell(data['trajectory'][str(i)]['cell'])
if data['trajectory'][str(i)]['positions']:
atoms.set_positions(data['trajectory'][str(i)]['positions'])
calc = SPC(
atoms=atoms,
energy=data['trajectory'][str(i)].get('energy'),
forces=data['trajectory'][str(i)].get('forces'),
stress=data['trajectory'][str(i)].get('stress'))
atoms.set_calculator(calc)
images += [atoms]
# Write the traj file
if out_file:
write(out_file, images)
return images
def encode_to_atoms(encode, out_file='input.traj'):
"""Dump the encoding to a local traj file."""
# First, decode the trajectory
data = json.loads(encode, encoding='utf-8')
# Construct the initial atoms object
atoms = Atoms(
data['numbers'],
data['trajectory']['0']['positions'],
cell=data['trajectory']['0']['cell'],
pbc=data['pbc'])
atoms.info = data['calculator_parameters']
atoms.set_constraint([dict2constraint(_) for _ in data['constraints']])
initial_magmoms = data.get('initial_magmoms')
if initial_magmoms:
atoms.set_initial_magnetic_moments(initial_magmoms)
# Attach the calculator
results = {'atoms': atoms}
for prop in supported_properties:
results.update({prop: data['trajectory']['0'].get(prop)})
calc = SPC(**results)
atoms.set_calculator(calc)
# Collect the rest of the trajectory information
images = [atoms]
for i in range(len(data['trajectory']))[1:]:
atoms = atoms.copy()
if data['trajectory'][str(i)]['cell']:
atoms.set_cell(data['trajectory'][str(i)]['cell'])
if data['trajectory'][str(i)]['positions']:
atoms.set_positions(data['trajectory'][str(i)]['positions'])
results = {'atoms': atoms}
for prop in supported_properties:
results.update({prop: data['trajectory'][str(i)].get(prop)})
calc = SPC(**results)
atoms.set_calculator(calc)
images += [atoms]
# Write the traj file
if out_file:
write(out_file, images)
return images
def _relax_with_vasp(atoms, vasp_flags):
'''
Perform a DFT relaxation with VASP and then write the trajectory to the
'all.traj' file and save the log file.
Args:
atoms `ase.Atoms` object of the structure we want to relax
vasp_flags A dictionary of settings we want to pass to the `Vasp2`
calculator
Returns:
atoms The relaxed `ase.Atoms` structure
'''
# Run the calculation
calc = Vasp2(**vasp_flags)
atoms.set_calculator(calc)
atoms.get_potential_energy()
# Read the trajectory from the output file
images = []
for atoms in ase.io.read('vasprun.xml', ':'):
image = atoms.copy()
image = image[calc.resort]
image.set_calculator(SPC(image,
energy=atoms.get_potential_energy(),
forces=atoms.get_forces()[calc.resort]))
images += [image]
# Write the trajectory
with TrajectoryWriter('all.traj', 'a') as tj:
for atoms in images:
tj.write(atoms)
return images[-1]
def attach_results(f, atoms, write_file=True):
""" Return the TS corrected energy for a scf instance
in a log file and attach them to the given atoms obejct.
Will also attach the forces and stress if applicable.
"""
energy, forces, stress = None, None, None
line = f.readline()
while '! total energy' not in line:
line = f.readline()
energy = float(line.split()[-2]) * Ry
# Correct for non-zero temperature smearing
for i in range(20):
line = f.readline()
if ' smearing contrib.' in line:
energy -= 0.5 * float(line.split()[-2]) * Ry
# Collect the forces on the atoms
if 'Forces acting on atoms (Ry/au):' in line:
for _ in range(4):
line = f.readline()
if 'atom' in line:
break
forces = []
for _ in range(len(atoms)):
forces += [line.split()[-3:]]
line = f.readline()
forces = np.array(forces, dtype=float) / Ry * aul
# If forces were located, attempt to find stress
for i in range(10):
line = f.readline()
if 'total stress' in line:
stress = []
for _ in range(3):
line = f.readline()
stress += [line.split()[-3:]]
stress = np.array(stress, dtype=float) / Ry * Bohr ** 3
break
# attach the calculator
calc = SPC(
atoms=atoms,
energy=energy,
forces=forces,
stress=stress)
atoms.set_calculator(calc)
return atoms
def vasptraj2db(db, directory):
# traj = get_traj(calc=calc)
dir_name = directory.split('/')[-1]
vasprun = '{0}/vasprun.xml'.format(directory)
for i, atoms in enumerate(read(vasprun, ':')):
atoms.set_constraint(None)
atoms.set_calculator(
SPC(atoms,
energy=atoms.get_potential_energy(),
forces=atoms.get_forces()))
db.write(atoms, image_no=int(i), path=directory)
return
def xml_to_traj(xml='vasprun.xml'):
'''
Converts an XML file into both a trajectory file while also returning the
trajectory as a list of `ase.Atoms` objects
Args:
xml String indicating the XML file to read from
Returns:
traj A list of `ase.Atoms` objects
'''
traj = ase.io.read(xml, ':')
for atoms in traj:
atoms.set_calculator(
SPC(atoms,
energy=atoms.get_potential_energy(),
forces=atoms.get_forces()))
return traj
def __collect_configs(self):
"""
Collect final configurations after relaxation
"""
init_db = connect(os.path.join(self.traj_path, 'initial.db'))
init_cnstrt = [
entry.toatoms().constraints for entry in init_db.select()
]
config_steps = [[] for _ in range(init_db.count())]
config_convgs = [False] * init_db.count()
for i in range(self.n_step):
tmp_db = connect(
os.path.join(self.traj_path, f'to-cal-step{i}-gt.db'))
tmp_configs = [entry.toatoms() for entry in tmp_db.select()]
tmp_frs = [atoms.get_forces() for atoms in tmp_configs]
tmp_nrgs = [atoms.get_potential_energy() for atoms in tmp_configs]
j = 0
for k, convg in enumerate(config_convgs):
if not convg:
f = tmp_frs[j]
for c in init_cnstrt[k]:
c.adjust_forces(None, f)
mf = np.linalg.norm(f, axis=1).max()
if mf <= self.fmax:
config_convgs[k] = True
dummy = tmp_configs[j].copy()
dummy.set_constraint(init_cnstrt[k])
dummy.set_calculator(
SPC(dummy, energy=tmp_nrgs[j], forces=tmp_frs[j]))
config_steps[k].append(dummy)
j += 1
final_db = connect(os.path.join(self.traj_path, 'final.db'))
for traj in config_steps:
final_db.write(traj[-1])
return final_db
def __get_ground_truth(self):
"""
Calculate the groud truth energetics for uncertain configurations.
"""
print(f'Step {self.n_step}: get groud truth data')
self.log_file.write(f'Step {self.n_step}: get groud truth data \n')
to_cal_db_path = os.path.join(self.traj_path,
f'to-cal-step{self.n_step}.db')
caled_db_path = os.path.join(self.traj_path,
f'to-cal-step{self.n_step}-gt.db')
to_cal_db = connect(to_cal_db_path)
caled_db = connect(caled_db_path)
for entry in to_cal_db.select():
atoms = entry.toatoms()
atoms.set_calculator(self.gt_calculator)
nrg = atoms.get_potential_energy()
frs = atoms.get_forces()
atoms.set_calculator(SPC(atoms, energy=nrg, forces=frs))
caled_db.write(atoms)
print(f'Step {self.n_step}: groud truth data calculation done')
self.log_file.write(
f'Step {self.n_step}: groud truth data calculation done \n')
return None
def runVasp(fname_in, fname_out, vaspflags, npar=4):
'''
This function is meant to be sent to each cluster and then used to run our rockets.
As such, it has algorithms to run differently depending on the cluster that is trying
to use this function.
Inputs:
fname_in
fname_out
vaspflags
npar
'''
fname_in = str(fname_in)
fname_out = str(fname_out)
# read the input atoms object
atoms = read(str(fname_in))
# Check that the unit vectors obey the right-hand rule, (X x Y points in Z) and if not
# Flip the order of X and Y to enforce this so that VASP is happy.
if np.dot(np.cross(atoms.cell[0], atoms.cell[1]), atoms.cell[2]) < 0:
atoms.set_cell(atoms.cell[[1, 0, 2], :])
vasp_cmd = 'vasp_std'
os.environ['PBS_SERVER'] = 'gilgamesh.cheme.cmu.edu'
if 'PBS_NODEFILE' in os.environ:
NPROCS = NPROCS = len(open(os.environ['PBS_NODEFILE']).readlines())
elif 'SLURM_CLUSTER_NAME' in os.environ:
if 'SLURM_NPROCS' in os.environ:
# We're on cori haswell
NPROCS = int(os.environ['SLURM_NPROCS'])
else:
# we're on cori KNL, just one processor
NPROCS = 1
# If we're on Gilgamesh...
if 'PBS_NODEFILE' in os.environ and os.environ['PBS_SERVER'] == 'gilgamesh.cheme.cmu.edu':
vaspflags['npar'] = 4
vasp_cmd = '/home-research/zhongnanxu/opt/vasp-5.3.5/bin/vasp-vtst-beef-parallel'
NPROCS = NPROCS = len(open(os.environ['PBS_NODEFILE']).readlines())
mpicall = lambda x, y: 'mpirun -np %i %s' % (x, y) # noqa: E731
# If we're on Arjuna...
elif 'SLURM_CLUSTER_NAME' in os.environ and os.environ['SLURM_CLUSTER_NAME'] == 'arjuna':
# If this is a GPU job...
if os.environ['CUDA_VISIBLE_DEVICES'] != 'NoDevFiles':
vaspflags['ncore'] = 1
vaspflags['kpar'] = 16
vaspflags['nsim'] = 8
vaspflags['lreal'] = 'Auto'
vasp_cmd = 'vasp_gpu'
mpicall = lambda x, y: 'mpirun -np %i %s' % (x, y) # noqa: E731
# If this is a CPU job...
else:
if NPROCS > 16:
vaspflags['ncore'] = 4
vaspflags['kpar'] = 4
else:
vaspflags['kpar'] = 1
vaspflags['ncore'] = 4
mpicall = lambda x, y: 'mpirun -np %i %s' % (x, y) # noqa: E731
# If we're on Cori, use SLURM. Note that we decrease the priority by 1000
# in order to prioritize other things higher, such as modeling and prediction
# in GASpy_regression
elif 'SLURM_CLUSTER_NAME' in os.environ and os.environ['SLURM_CLUSTER_NAME'] == 'cori':
# If we're on a Haswell node...
if os.environ['CRAY_CPU_TARGET'] == 'haswell' and 'knl' not in os.environ['PATH']:
NNODES = int(os.environ['SLURM_NNODES'])
vaspflags['kpar'] = NNODES
mpicall = lambda x, y: 'srun -n %d %s' % (x, y) # noqa: E731
# If we're on a KNL node...
elif 'knl' in os.environ['PATH']:
mpicall = lambda x, y: 'srun -n %d -c8 --cpu_bind=cores %s' % (x*32, y) # noqa: E731
vaspflags['ncore'] = 1
# Set the pseudopotential type by setting 'xc' in Vasp()
if vaspflags['pp'].lower() == 'lda':
vaspflags['xc'] = 'lda'
elif vaspflags['pp'].lower() == 'pbe':
vaspflags['xc'] = 'PBE'
pseudopotential = vaspflags['pp_version']
os.environ['VASP_PP_PATH'] = os.environ['VASP_PP_BASE'] + '/' + str(pseudopotential) + '/'
del vaspflags['pp_version']
os.environ['VASP_COMMAND'] = mpicall(NPROCS, vasp_cmd)
# Detect whether or not there are constraints that cannot be handled by VASP
allowable_constraints = ['FixAtoms']
constraint_not_allowable = [constraint.todict()['name']
not in allowable_constraints
for constraint in atoms.constraints]
vasp_incompatible_constraints = np.any(constraint_not_allowable)
# If there are incompatible constraints, we need to switch to an ASE-based optimizer
if vasp_incompatible_constraints:
vaspflags['ibrion'] = 2
vaspflags['nsw'] = 0
calc = Vasp2(**vaspflags)
atoms.set_calculator(calc)
qn = BFGS(atoms, logfile='relax.log', trajectory='all.traj')
qn.run(fmax=vaspflags['ediffg'] if 'ediffg' in vaspflags else 0.05)
finalimage = atoms
else:
# set up the calculation and run
calc = Vasp2(**vaspflags)
atoms.set_calculator(calc)
# Trigger the calculation
atoms.get_potential_energy()
atomslist = []
for atoms in read('vasprun.xml', ':'):
catoms = atoms.copy()
catoms = catoms[calc.resort]
catoms.set_calculator(SPC(catoms,
energy=atoms.get_potential_energy(),
forces=atoms.get_forces()[calc.resort]))
atomslist += [catoms]
# Get the final trajectory
finalimage = atoms
# Write a traj file for the optimization
tj = TrajectoryWriter('all.traj', 'a')
for atoms in atomslist:
print('writing trajectory file!')
print(atoms)
tj.write(atoms)
tj.close()
# Write the final structure
finalimage.write(fname_out)
# Write a text file with the energy
with open('energy.out', 'w') as fhandle:
fhandle.write(str(finalimage.get_potential_energy()))
try:
os.remove('CHGCAR')
except OSError:
pass
try:
os.remove('WAVECAR')
except OSError:
pass
try:
os.remove('CHG')
except OSError:
pass
return str(atoms), open('all.traj', 'r').read().encode('hex'), finalimage.get_potential_energy()
def __update_training_data(self):
"""
Add calcualted ground truth data into training file
Check the convergence of atoms configurations
"""
gt_path = os.path.join(self.traj_path,
f'to-cal-step{self.n_step}-gt.db')
gt_atoms = [entry.toatoms() for entry in connect(gt_path).select()]
to_relax_path = os.path.join(self.traj_path,
f'to-relax-step{self.n_step}.db')
to_relax_db = connect(to_relax_path)
max_frs = []
to_added_atoms = []
if self.n_step == 0:
previous_data = []
else:
train_db_path = os.path.join(self.traj_path,
f'train-set-step{self.n_step-1}.db')
previous_db = connect(train_db_path)
previous_data = [entry.toatoms() for entry in previous_db.select()]
for i, atoms in enumerate(gt_atoms):
if not self.relaxed[i]:
to_added_atoms.append(atoms) # for training
dummy = atoms.copy() # for relaxation
dummy.set_constraint(self.constraints[i])
dummy.set_calculator(
SPC(dummy,
energy=atoms.get_potential_energy(),
forces=atoms.get_forces()))
tmp_frs = atoms.get_forces()
for c in self.constraints[i]:
c.adjust_forces(None, tmp_frs)
max_fr = np.linalg.norm(tmp_frs, axis=1).max()
if max_fr <= self.fmax:
self.relaxed[i] = True
else:
to_relax_db.write(dummy)
max_frs.append(round(max_fr, 3))
all_data = previous_data + to_added_atoms
new_train_path = os.path.join(self.traj_path,
f'train-set-step{self.n_step}.db')
new_train_db = connect(new_train_path)
for atoms in all_data[-self.train_size:]:
new_train_db.write(atoms)
self.constraints = [
self.constraints[i] for i in range(len(self.relaxed))
if not self.relaxed[i]
]
self.relaxed = [entry for entry in self.relaxed if not entry]
print('max force for each configuration: ')
print(max_frs)
self.log_file.write('max force for each configuration: \n')
self.log_file.write(f'{max_frs} \n')
self.log_file.flush()
if not self.relaxed:
return True
else:
return False
print('=================================================================================================='.center(120))
print('ASE GUI plot of specified atoms.'.center(120))
print('=================================================================================================='.center(120))
print('')
## Argparse
args = argparse()
atom_i = args.atom_ind
cut_r = args.cutoff_radius
from ase.io import read
alist = read(args.alist_file, args.image_slice)
if not isinstance(alist, list):
alist = [alist]
new_alist = []
for i in range(len(alist)):
d = alist[i].get_distances(
atom_i,
np.arange(len(alist[i])).tolist(),
mic=True,
)
m = d < cut_r
new_alist.append(alist[i][m])
from ase.calculators.singlepoint import SinglePointCalculator as SPC
new_alist[-1].calc = SPC(new_alist[-1], forces=alist[i].get_forces()[m])
print('\nimg #{} - {}'.format(i, np.arange(len(alist[i]))[m].tolist()))
from ase.visualize import view
view(new_alist)
def runVasp(fname_in, fname_out, vaspflags, npar=4):
'''
This function is meant to be sent to each cluster and then used to run our rockets.
As such, it has algorithms to run differently depending on the cluster that is trying
to use this function.
Inputs:
fname_in
fname_out
vaspflags
npar
'''
fname_in = str(fname_in)
fname_out = str(fname_out)
# read the input atoms object
atoms = read(str(fname_in))
# Check that the unit vectors obey the right-hand rule, (X x Y points in Z) and if not
# Flip the order of X and Y to enforce this so that VASP is happy.
if np.dot(np.cross(atoms.cell[0], atoms.cell[1]), atoms.cell[2]) < 0:
atoms.set_cell(atoms.cell[[1, 0, 2], :])
# If we're on UON rcg
NPROCS = int(os.environ['NCPUS'])
queue = os.environ['PBS_QUEUE']
vaspflags['npar'] = 4
if 'xeon5' in str(queue):
vasp_cmd = 'vasp_std'
else:
vasp_cmd = '/home/ajp/bin/vasp_std_5.4.intelmpi'
mpicall = lambda x, y: 'mpirun -np %i %s' % (x, y) # noqa: E731
# Set the pseudopotential type by setting 'xc' in Vasp()
if vaspflags['pp'].lower() == 'lda':
vaspflags['xc'] = 'lda'
elif vaspflags['pp'].lower() == 'pbe':
vaspflags['xc'] = 'PBE'
# pseudopotential = vaspflags['pp_version']
# s.environ['VASP_PP_PATH'] = os.environ['VASP_PP_BASE'] + '/' + str(pseudopotential) + '/'
del vaspflags['pp_version']
os.environ['VASP_COMMAND'] = mpicall(NPROCS, vasp_cmd)
# Detect whether or not there are constraints that cannot be handled by VASP
allowable_constraints = ['FixAtoms']
constraint_not_allowable = [
constraint.todict()['name'] not in allowable_constraints
for constraint in atoms.constraints
]
vasp_incompatible_constraints = np.any(constraint_not_allowable)
# If there are incompatible constraints, we need to switch to an ASE-based optimizer
if vasp_incompatible_constraints:
vaspflags['ibrion'] = 2
vaspflags['nsw'] = 0
calc = Vasp2(**vaspflags)
atoms.set_calculator(calc)
qn = BFGS(atoms, logfile='relax.log', trajectory='all.traj')
qn.run(fmax=vaspflags['ediffg'] if 'ediffg' in vaspflags else 0.05)
finalimage = atoms
else:
# set up the calculation and run
calc = Vasp2(**vaspflags)
atoms.set_calculator(calc)
# Trigger the calculation
atoms.get_potential_energy()
atomslist = []
for atoms in read('vasprun.xml', ':'):
catoms = atoms.copy()
catoms = catoms[calc.resort]
catoms.set_calculator(
SPC(catoms,
energy=atoms.get_potential_energy(),
forces=atoms.get_forces()[calc.resort]))
atomslist += [catoms]
# Get the final trajectory
finalimage = atoms
# Write a traj file for the optimization
tj = TrajectoryWriter('all.traj', 'a')
for atoms in atomslist:
print('writing trajectory file!')
print(atoms)
tj.write(atoms)
tj.close()
# Write the final structure
finalimage.write(fname_out)
# Write a text file with the energy
with open('energy.out', 'w') as fhandle:
fhandle.write(str(finalimage.get_potential_energy()))
try:
os.remove('CHGCAR')
except OSError:
pass
try:
os.remove('WAVECAR')
except OSError:
pass
try:
os.remove('CHG')
except OSError:
pass
return str(atoms), open(
'all.traj', 'rb').read().hex(), finalimage.get_potential_energy()
