1#!/usr/local/bin/python 2 3# Atomistic Simulation Environment (ASE) calculator interface for JDFTx 4# See http://jdftx.org for JDFTx and https://wiki.fysik.dtu.dk/ase/ for ASE 5# Authors: Deniz Gunceler, Ravishankar Sundararaman 6 7from __future__ import print_function #For Python2 compatibility 8 9import os, scipy, subprocess, tempfile, re 10from ase.calculators.interface import Calculator 11from ase.units import Bohr, Hartree 12 13#Run shell command and return output as a string: 14def shell(cmd): 15 return subprocess.check_output(cmd, shell=True) 16 17#Return var, replacing it with environment variable varName if it is None 18def replaceVariable(var, varName): 19 if (var is None) and (varName in os.environ): 20 return os.environ[varName] 21 else: 22 return var 23 24class JDFTx(Calculator): 25 26 def __init__(self, executable=None, pseudoDir=None, pseudoSet='GBRV', commands=None, ignoreStress=False): 27 #Valid pseudopotential sets (mapping to path and suffix): 28 pseudoSetMap = { 29 'SG15' : 'SG15/$ID_ONCV_PBE.upf', 30 'GBRV' : 'GBRV/$ID_pbe.uspp', 31 'GBRV-pbe' : 'GBRV/$ID_pbe.uspp', 32 'GBRV-lda' : 'GBRV/$ID_lda.uspp', 33 'GBRV-pbesol' : 'GBRV/$ID_pbesol.uspp' 34 } 35 36 #Get default values from environment: 37 self.executable = replaceVariable(executable, 'JDFTx') #Path to the jdftx executable (cpu or gpu) 38 self.pseudoDir = replaceVariable(pseudoDir, 'JDFTx_pseudo') #Path to the pseudopotentials folder 39 self.ignoreStress = ignoreStress 40 41 if (self.executable is None): 42 raise Exception('Specify path to jdftx in argument \'executable\' or in environment variable \'JDFTx\'.') 43 if (self.pseudoDir is None) and (not (pseudoSet in pseudoSetMap)): 44 raise Exception('Specify pseudopotential path in argument \'pseudoDir\' or in environment variable \'JDFTx_pseudo\', or specify a valid \'pseudoSet\'.') 45 46 if pseudoSet in pseudoSetMap: 47 self.pseudoSetCmd = 'ion-species ' + pseudoSetMap[pseudoSet] 48 else: 49 self.pseudoSetCmd = '' 50 51 # Gets the input file template 52 self.acceptableCommands = set(['electronic-SCF']) 53 template = str(shell('%s -t' % (self.executable))) 54 for match in re.findall(r"# (\S+) ", template): 55 self.acceptableCommands.add(match) 56 57 self.input = [ 58 ('dump-name', '$VAR'), 59 ('initial-state', '$VAR') 60 ] 61 62 # Parse commands, which can be a dict or a list of 2-tuples. 63 if isinstance(commands, dict): 64 commands = commands.items() 65 elif commands is None: 66 commands = [] 67 for cmd, v in commands: 68 self.addCommand(cmd, v) 69 70 # Accepted pseudopotential formats 71 self.pseudopotentials = ['fhi', 'uspp', 'upf'] 72 73 # Current results 74 self.E = None 75 self.Forces = None 76 77 # History 78 self.lastAtoms = None 79 self.lastInput = None 80 81 # k-points 82 self.kpoints = None 83 84 # Dumps 85 self.dumps = [] 86 self.addDump("End", "State") 87 self.addDump("End", "Forces") 88 self.addDump("End", "Ecomponents") 89 90 #Run directory: 91 self.runDir = tempfile.mkdtemp() 92 print('Set up JDFTx calculator with run files in \'' + self.runDir + '\'') 93 94 ########### Interface Functions ########### 95 96 def addCommand(self, cmd, v): 97 if(not self.validCommand(cmd)): 98 raise IOError('%s is not a valid JDFTx command!\nLook at the input file template (jdftx -t) for a list of commands.' % (cmd)) 99 self.input.append((cmd, v)) 100 101 def addDump(self, when, what): 102 self.dumps.append((when, what)) 103 104 def addKPoint(self, pt, w=1): 105 b1, b2, b3 = pt 106 if self.kpoints is None: 107 self.kpoints = [] 108 self.kpoints.append((b1, b2, b3, w)) 109 110 def clean(self): 111 shell('rm -rf ' + self.runDir) 112 113 def calculation_required(self, atoms, quantities): 114 if((self.E is None) or (self.Forces is None)): 115 return True 116 if((self.lastAtoms != atoms) or (self.input != self.lastInput)): 117 return True 118 return False 119 120 def get_forces(self, atoms): 121 if(self.calculation_required(atoms, None)): 122 self.update(atoms) 123 return self.Forces 124 125 def get_potential_energy(self, atoms, force_consistent=False): 126 if(self.calculation_required(atoms, None)): 127 self.update(atoms) 128 return self.E 129 130 def get_stress(self, atoms): 131 if self.ignoreStress: 132 return scipy.zeros((3,3)) 133 else: 134 raise NotImplementedError('Stress calculation not implemented in JDFTx interface: set ignoreStress=True to ignore.') 135 136 ################### I/O ################### 137 138 def __readEnergy(self, filename): 139 Efinal = None 140 for line in open(filename): 141 tokens = line.split() 142 if len(tokens)==3: 143 Efinal = float(tokens[2]) 144 if Efinal is None: 145 raise IOError('Error: Energy not found.') 146 return Efinal * Hartree #Return energy from final line (Etot, F or G) 147 148 def __readForces(self, filename): 149 idxMap = {} 150 symbolList = self.lastAtoms.get_chemical_symbols() 151 for i, symbol in enumerate(symbolList): 152 if symbol not in idxMap: 153 idxMap[symbol] = [] 154 idxMap[symbol].append(i) 155 forces = [0]*len(symbolList) 156 for line in open(filename): 157 if line.startswith('force '): 158 tokens = line.split() 159 idx = idxMap[tokens[1]].pop(0) # tokens[1] is chemical symbol 160 forces[idx] = [float(word) for word in tokens[2:5]] # tokens[2:5]: force components 161 162 if(len(forces) == 0): 163 raise IOError('Error: Forces not found.') 164 return (Hartree / Bohr) * scipy.array(forces) 165 166 ############## Running JDFTx ############## 167 168 def update(self, atoms): 169 self.runJDFTx(self.constructInput(atoms)) 170 171 def runJDFTx(self, inputfile): 172 """ Runs a JDFTx calculation """ 173 #Write input file: 174 fp = open(self.runDir+'/in', 'w') 175 fp.write(inputfile) 176 fp.close() 177 #Run jdftx: 178 shell('cd %s && %s -i in -o out' % (self.runDir, self.executable)) 179 self.E = self.__readEnergy('%s/Ecomponents' % (self.runDir)) 180 self.Forces = self.__readForces('%s/force' % (self.runDir)) 181 182 def constructInput(self, atoms): 183 """ Constructs a JDFTx input string using the input atoms and the input file arguments (kwargs) in self.input """ 184 inputfile = '' 185 186 # Add lattice info 187 R = atoms.get_cell() / Bohr 188 inputfile += 'lattice \\\n' 189 for i in range(3): 190 for j in range(3): 191 inputfile += '%f ' % (R[j, i]) 192 if(i != 2): 193 inputfile += '\\' 194 inputfile += '\n' 195 196 # Construct most of the input file 197 inputfile += '\n' 198 for cmd, v in self.input: 199 inputfile += '%s %s\n' % (cmd, str(v)) 200 201 # Add ion info 202 atomPos = [x / Bohr for x in list(atoms.get_positions())] # Also convert to bohr 203 atomNames = atoms.get_chemical_symbols() # Get element names in a list 204 inputfile += '\ncoords-type cartesian\n' 205 for i in range(len(atomPos)): 206 inputfile += 'ion %s %f %f %f \t 1\n' % (atomNames[i], atomPos[i][0], atomPos[i][1], atomPos[i][2]) 207 del i 208 209 # Add k-points 210 if self.kpoints: 211 for pt in self.kpoints: 212 inputfile += 'kpoint %.8f %.8f %.8f %.14f\n' % pt 213 214 #Add pseudopotentials 215 inputfile += '\n' 216 if not (self.pseudoDir is None): 217 added = [] # List of pseudopotential that have already been added 218 for atom in atomNames: 219 if(sum([x == atom for x in added]) == 0.): # Add ion-species command if not already added 220 for filetype in self.pseudopotentials: 221 try: 222 shell('ls %s | grep %s.%s' % (self.pseudoDir, atom, filetype)) 223 inputfile += 'ion-species %s/%s.%s\n' % (self.pseudoDir, atom, filetype) 224 added.append(atom) 225 break 226 except: 227 pass 228 inputfile += self.pseudoSetCmd + '\n' #Pseudopotential sets 229 230 # Add truncation info (periodic vs isolated) 231 inputfile += '\ncoulomb-interaction ' 232 pbc = list(atoms.get_pbc()) 233 if(sum(pbc) == 3): 234 inputfile += 'periodic\n' 235 elif(sum(pbc) == 0): 236 inputfile += 'isolated\n' 237 elif(sum(pbc) == 1): 238 inputfile += 'wire %i%i%i\n' % (pbc[0], pbc[1], pbc[2]) 239 elif(sum(pbc) == 2): 240 inputfile += 'slab %i%i%i\n' % (not pbc[0], not pbc[1], not pbc[2]) 241 #--- add truncation center: 242 if(sum(pbc) < 3): 243 center = scipy.mean(scipy.array(atomPos), axis=0) 244 inputfile += 'coulomb-truncation-embed %g %g %g\n' % tuple(center.tolist()) 245 246 #Add dump commands 247 inputfile += "".join(["dump %s %s\n" % (when, what) for when, what in self.dumps]) 248 249 # Cache this calculation to history 250 self.lastAtoms = atoms.copy() 251 self.lastInput = list(self.input) 252 return inputfile 253 254 ############## JDFTx command structure ############## 255 256 def validCommand(self, command): 257 """ Checks whether the input string is a valid jdftx command \nby comparing to the input template (jdft -t)""" 258 if(type(command) != str): 259 raise IOError('Please enter a string as the name of the command!\n') 260 return command in self.acceptableCommands 261 262 def help(self, command=None): 263 """ Use this function to get help about a JDFTx command """ 264 if(command is None): 265 print('This is the help function for JDFTx-ASE interface. \ 266 \nPlease use the command variable to get information on a specific command. \ 267 \nVisit jdftx.sourceforge.net for more info.') 268 elif(self.validCommand(command)): 269 raise NotImplementedError('Template help is not yet implemented') 270 else: 271 raise IOError('%s is not a valid command' % (command)) 272