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()