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