-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathsisl_jij_mpi.py
268 lines (229 loc) · 10.4 KB
/
sisl_jij_mpi.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
import numpy as np
import numpy.linalg as nl
import sisl
import tqdm
import sys
from mpi4py import MPI
import argparse
from itertools import permutations, product
from timeit import default_timer as timer
from scipy.special import roots_legendre
# define some useful functions
def hsk(dh,k=(0,0,0)):
'''
One way to speed up Hk and Sk generation
'''
k = np.asarray(k, np.float64) # this two conversion lines are from the sisl source
k.shape = (-1,) #
phases = np.exp(-1j * np.dot(np.dot(np.dot(dh.rcell, k), dh.cell), dh.sc.sc_off.T)) # this generates the list of phases
HKU = np.einsum('abc,c->ab',dh.hup,phases)
HKD = np.einsum('abc,c->ab',dh.hdo,phases)
SK = np.einsum('abc,c->ab',dh.sov,phases)
return HKU,HKD,SK
def make_contour(emin=-20,emax=0.0,enum=42,p=150):
'''
A more sophisticated contour generator
'''
x,wl = roots_legendre(enum)
R = (emax-emin)/2
z0 = (emax+emin)/2
y1 = -np.log(1+np.pi*p)
y2 = 0
y = (y2-y1)/2*x+(y2+y1)/2
phi = (np.exp(-y)-1)/p
ze = z0+R*np.exp(1j*phi)
we = -(y2-y1)/2*np.exp(-y)/p*1j*(ze-z0)*wl
class ccont:
#just an empty container class
pass
cont = ccont()
cont.R = R
cont.z0 = z0
cont.ze = ze
cont.we = we
cont.enum = enum
return cont
def make_kset(dirs='xyz',NUMK=20):
'''
Simple k-grid generator. Depending on the value of the dirs
argument k sampling in 1,2 or 3 dimensions is generated.
If dirs argument does not contain either of x,y or z
a kset of a single k-pont at the origin is returend.
'''
if not(sum([d in dirs for d in 'xyz'])):
return np.array([[0,0,0]])
kran=len(dirs)*[np.linspace(0,1,NUMK,endpoint=False)]
mg=np.meshgrid(*kran)
dirsdict=dict()
for d in enumerate(dirs):
dirsdict[d[1]]=mg[d[0]].flatten()
for d in 'xyz':
if not(d in dirs):
dirsdict[d]=0*dirsdict[dirs[0]]
kset = np.array([dirsdict[d] for d in 'xyz']).T
return kset
def make_atran(nauc,dirs='xyz',dist=1):
'''
Simple pair generator. Depending on the value of the dirs
argument sampling in 1,2 or 3 dimensions is generated.
If dirs argument does not contain either of x,y or z
a single pair is returend.
'''
if not(sum([d in dirs for d in 'xyz'])):
return (0,0,[1,0,0])
dran=len(dirs)*[np.arange(-dist,dist+1)]
mg=np.meshgrid(*dran)
dirsdict=dict()
for d in enumerate(dirs):
dirsdict[d[1]]=mg[d[0]].flatten()
for d in 'xyz':
if not(d in dirs):
dirsdict[d]=0*dirsdict[dirs[0]]
ucran = np.array([dirsdict[d] for d in 'xyz']).T
atran=[]
for i,j in list(product(range(nauc),repeat=2)):
for u in ucran:
if (abs(i-j)+sum(abs(u)))>0:
atran.append((i,j,list(u)))
return atran
#----------------------------------------------------------------------
start = timer()
# Some input parsing
parser = argparse.ArgumentParser()
parser.add_argument('--kset' , dest = 'kset' , default = 2 , type=int , help = 'k-space resolution of Jij calculation')
parser.add_argument('--kdirs' , dest = 'kdirs' , default = 'xyz' , help = 'Definition of k-space dimensionality')
parser.add_argument('--eset' , dest = 'eset' , default = 42 , type=int , help = 'Number of energy points on the contour')
parser.add_argument('--eset-p' , dest = 'esetp' , default = 10 , type=int , help = 'Parameter tuning the distribution on the contour')
parser.add_argument('--input' , dest = 'infile' , required = True , help = 'Input file name')
parser.add_argument('--output' , dest = 'outfile', required = True , help = 'Output file name')
parser.add_argument('--Ebot' , dest = 'Ebot' , default = -20.0 , type=float, help = 'Bottom energy of the contour')
parser.add_argument('--npairs' , dest = 'npairs' , default = 1 , type=int , help = 'Number of unitcell pairs in each direction for Jij calculation')
parser.add_argument('--adirs' , dest = 'adirs' , default = False , help = 'Definition of pair directions')
parser.add_argument('--use-tqdm', dest = 'usetqdm', default = 'not' , help = 'Use tqdm for progressbars or not')
parser.add_argument('--cutoff' , dest = 'cutoff' , default = 100.0 , type=float, help = 'Real space cutoff for pair generation in Angs')
args = parser.parse_args()
#----------------------------------------------------------------------
# MPI init
comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()
root_node = 0
if rank == root_node:
print('Number of nodes in the parallel cluster: ',size)
#----------------------------------------------------------------------
# importing the necessary structures from SIESTA output
dat = sisl.get_sile(args.infile)
dh = dat.read_hamiltonian()
# update datastructure of the hamiltonian
# this is needed for quick Hk building
dh.hup = dh.tocsr(0).toarray().reshape(dh.no,dh.n_s,dh.no).transpose(0,2,1).astype('complex128')
dh.hdo = dh.tocsr(1).toarray().reshape(dh.no,dh.n_s,dh.no).transpose(0,2,1).astype('complex128')
dh.sov = dh.tocsr(2).toarray().reshape(dh.no,dh.n_s,dh.no).transpose(0,2,1).astype('complex128')
#----------------------------------------------------------------------
# generate k space sampling
kset=make_kset(dirs=args.kdirs,NUMK=args.kset)
wk = 1/len(kset) # weight of a kpoint in BZ integral
kpcs = np.array_split(kset,size)
if 'k' in args.usetqdm:
kpcs[root_node] = tqdm.tqdm(kpcs[root_node],desc='k loop')
#----------------------------------------------------------------------
# generate pairs
args.adirs = args.adirs if args.adirs else args.kdirs # Take pair directions for k directions if adirs is not set
atran=make_atran(len(dh.atoms),args.adirs,dist=args.npairs) # definition of pairs in terms of integer coordinates refering to unicell distances and atomic positions
pairs=[]
for i,j,uc in atran:
if nl.norm(np.dot(uc,dh.cell)) < args.cutoff:
pairs.append(dict(
offset = uc, # lattice vector offset between the unitcells the two atoms are
aiij = [i,j], # indecies of the atoms in the unitcell
noij = [dh.atoms.orbitals[i],dh.atoms.orbitals[j]], # number of orbitals on the appropriate atoms
slij = [slice( *(lambda x:[min(x),max(x)+1])(dh.a2o(i,all=True)) ), # slices for
slice( *(lambda x:[min(x),max(x)+1])(dh.a2o(j,all=True)) )], # appropriate orbitals
rirj = [dh.axyz()[i],dh.axyz()[j]], # real space vectors of atoms in the unit cell
Rij = np.dot(uc,dh.cell), # real space distance vector between unit cells
rij = np.dot(uc,dh.cell)-dh.axyz()[i]+dh.axyz()[j], # real space vector between atoms
Jijz = [], # in this empty list are we going to gather the integrad of the energy integral
Jij = 0 # the final results of the calculation are going to be here on the root node
))
if rank == root_node:
print('Number of pairs beeing calculated: ',len(pairs))
#----------------------------------------------------------------------
# make energy contour
# we are working in eV now !
# and sisil shifts E_F to 0 !
cont = make_contour(emin=args.Ebot,enum=args.eset,p=args.esetp)
if (rank==root_node) and ('E' in args.usetqdm):
eran = tqdm.tqdm(cont.ze,desc='E loop')
else:
eran = cont.ze
#----------------------------------------------------------------------
# generating onsite matrix and overalp elements of all the atoms in the unitcell
# onsite of the origin supercell
orig_indx=np.arange(0,dh.no)+dh.sc_index([0,0,0])*dh.no
# spin up
uc_up = dh.tocsr(dh.UP )[:,orig_indx].toarray()
# spin down
uc_down = dh.tocsr(dh.DOWN )[:,orig_indx].toarray()
Hs=[]
# get number of atoms in the unit cell
for i in range(len(dh.atoms)):
at_indx=dh.a2o(i,all=True)
Hs.append(
uc_up[:,at_indx][at_indx,:]-
uc_down[:,at_indx][at_indx,:]
)
#----------------------------------------------------------------------
# sampling the integrand on the contour
for ze in eran:
# reset G-s
for pair in pairs:
noi , noj =pair['noij']
pair['Guij'] = np.zeros((noi,noj),dtype='complex128')
pair['Gdji'] = np.zeros((noj,noi),dtype='complex128')
pair['Guij_tmp'] = np.zeros((noi,noj),dtype='complex128')
pair['Gdji_tmp'] = np.zeros((noj,noi),dtype='complex128')
# doing parallel BZ integral
for k in kpcs[rank]:
k=np.array(k)
k.shape=(-1,)
HKU,HKD,SK = hsk(dh,k)
Gku = nl.inv((ze*SK-HKU))
Gkd = nl.inv((ze*SK-HKD))
for pair in pairs:
phase=np.exp(1j*np.dot(np.dot(k,dh.rcell),pair['Rij']))
si,sj=pair['slij']
# Fourier transform to real space
pair['Guij_tmp'] += Gku[si,sj]*phase*wk # ij gets exp(+i k R)
pair['Gdji_tmp'] += Gkd[sj,si]/phase*wk # ji gets exp(-i k R)
# summ reduce partial results of mpi nodes
for pair in pairs:
comm.Reduce(pair['Guij_tmp'],pair['Guij'],root=root_node)
comm.Reduce(pair['Gdji_tmp'],pair['Gdji'],root=root_node)
if rank==root_node:
# The Szunyogh-Lichtenstein formula
for pair in pairs:
i,j = pair['aiij']
pair['Jijz'].append(
np.trace(np.dot(
np.dot(Hs[i],pair['Guij']),
np.dot(Hs[j],pair['Gdji'])
)))
#----------------------------------------------------------------------
# evaluation of the contour integral on the root node
# and saveing output of the calculation
if rank==root_node:
for pair in pairs:
pair['Jijz'] = np.array(pair['Jijz'])
pair['Jij'] = np.trapz(np.imag(pair['Jijz']*cont.we)/(2*np.pi))
end = timer()
np.savetxt(args.outfile,
np.array([ [nl.norm(p['rij']),
p['Jij']*sisl.unit_convert('eV','Ry')*1000]+
p['aiij']+list(p['offset'])+list(p['rij'])
for p in pairs],
dtype=object),
header=str(args)+
'\nnumber of cores = '+str(size)+
'\ntime of calculation = '+str(end-start)+
'\nnorm(rij),Jij[mRy],aiij,offset,rij',
fmt="%s")