-
Notifications
You must be signed in to change notification settings - Fork 0
/
PyStructureFactor.py
344 lines (302 loc) · 12.9 KB
/
PyStructureFactor.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
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
import numpy
import math
import multiprocessing as mp
from multiprocessing.dummy import Pool as ThreadPool
from pyscf import df, scf, mcscf
from scipy import special
from pyscf import gto, lib, dft
from pyscf.dft import radi
from wigner import wigner_dl
def convert_sph2cart(xyz):
"""
Transforms the spherical coordinates to Cartesian coordinates.
Parameters:
xyz: 2D float array, shape (N,3)
Cartesian coordinates [x,y,z]
Returns:
2D float array, shape (N,3)
Spherical coordinates [r,theta,phi] corresponding to [x,y,z]
"""
ptsnew = numpy.zeros(xyz.shape)
xy2 = xyz[:, 0] ** 2 + xyz[:, 1] ** 2
ptsnew[:, 0] = numpy.sqrt(xy2 + xyz[:, 2] ** 2)
ptsnew[:, 1] = numpy.arctan2(numpy.sqrt(xy2), xyz[:, 2])
ptsnew[:, 2] = numpy.arctan2(xyz[:, 1], xyz[:, 0])
return ptsnew
def get_homo_index(np1):
"""
To find the index of HOMO orbit according to the orbital occupation. Returns i-1, denoting i the first orbit whose occupancy number is zero.
Parameters:
np1: Orbital occupancy number, array.
Returns:
The index of HOMO orbit starting from 0.
"""
for i in range(np1.size):
if np1[i] == 0:
return i - 1
def vc_orbit(mol, z, coeff, coords, dm, hf_method='RHF'):
"""To calculate the core potential times the ionizing orbital wave function, that is, V_c ψ_0.
Parameters:
mol: Mole
The molecule object.
z: int
The charge of the parent ion.
coeff: 2D array, shape (:,index)
Molecular orbital coefficient of the ionizing orbit.
coords: 2D array, shape (N,3)
Cartesian grid points(x,y,z)
dm: ndarray
Density matrix of molecule.
hf_method : str
Indicates whether 'RHF' or 'UHF' should be used in molecular HF calculation. Default is 'RHF'.
Returns:
The core potential oprator times the orbital wave function.
"""
ngrids = coords.shape[0]
vnuc_r = z / numpy.einsum('xi,xi->x', coords, coords) ** .5
vnuc = 0
for i in range(mol.natm):
r = mol.atom_coord(i)
z = mol.atom_charge(i)
rp = r - coords
vnuc += z / numpy.einsum('xi,xi->x', rp, rp) ** .5
if hf_method == 'UHF':
dm_vex = dm[0]
dm_vele = dm[1] + dm[0]
elif hf_method == 'RHF':
dm_vex = dm / 2
dm_vele = dm
else:
raise Exception("error: method must be 'RHF' or 'UHF'")
blksize = min(8000, ngrids)
ao = numpy.zeros((ngrids, mol.nao))
for ip0, ip1 in lib.prange(0, ngrids, blksize):
ao[ip0:ip1, :] = mol.eval_gto('GTOval', coords[ip0:ip1])
ion_orbit = numpy.dot(ao, coeff.T)
vex_ionorbit = numpy.zeros(shape=ngrids)
mep = numpy.zeros(shape=ngrids)
lsrange = list(lib.prange(0, ngrids, 600))
def part(p):
p0 = p[0]
p1 = p[1]
fakemol = gto.fakemol_for_charges(coords[p0:p1])
ints = df.incore.aux_e2(mol, fakemol)
del fakemol
fg = numpy.dot(coeff, ints)
fg = fg.reshape(mol.nao, p1 - p0)
mg = numpy.dot(ao[p0:p1], dm_vex)
vex_ionorbit[p0:p1] = -numpy.einsum('br,br->r', fg, mg.T)
vele = numpy.einsum('ijp,ij->p', ints, dm_vele)
mep[p0:p1] = -vnuc[p0:p1] + vele
del ints
return 0
num_cores = int(mp.cpu_count()) - 2
pool = ThreadPool(num_cores)
pool.map(part, lsrange)
pool.close()
pool.join()
vex_ionorbit = numpy.expand_dims(vex_ionorbit, axis=1)
mep_vnuc = numpy.expand_dims((mep + vnuc_r), axis=1)
vc_orbit = vex_ionorbit + mep_vnuc * ion_orbit
return vc_orbit
def omega_wvl(m, n_xi, kappa, z, l):
"""
The normalization coefficient of Omega.
Parameters
m: int
Magnetic quantum number, should satisfy `l >= m >= -l`.
l: int
Angular quantum number, should be a non-negative integer.
n_ep: int
Parabolic quantum number, should be a non-negative integer.
z: int
The charge of the parent ion.
kappa: float
`sqrt(-2*E0)`, E0 is orbital energy.
Returns:
The normalization coefficient of Omega for the given parameter.
"""
part1 = ((-1) ** (l + (abs(m) - m) / 2 + 1)) * (2 ** (l + 3 / 2)) * (
kappa ** (z / kappa - (abs(m) + 1) / 2 - n_xi))
part2 = math.sqrt(
(2 * l + 1) * math.factorial(l + m) * math.factorial(l - m) * math.factorial(
abs(m) + n_xi) * math.factorial(
n_xi)) * math.factorial(l) / math.factorial(2 * l + 1)
range_len = min(n_xi, l - abs(m))
part3 = 0
for k in range(range_len + 1):
part3 += special.gamma(l + 1 - z / kappa + n_xi - k) / (
math.factorial(k) * math.factorial(l - k) * math.factorial(abs(m) + k) * math.factorial(
l - abs(m) - k) * math.factorial(n_xi - k))
omega_norm = part1 * part2 * part3
return omega_norm
def omega_r(kappa, z, l, coords):
"""
Main part of Omega excluding the normalization coefficient. Will calculate all m' in [-l,l].
# Parameters:
l : int
Angular quantum number, should be a non-negative integer.
z : int
The charge of the parent ion.
kappa : float
`sqrt(-2*E0)`, E0 is orbital energy.
Returns:
An numpy array of shape(2l+1, len(coords)).
"""
coords_sph = convert_sph2cart(coords)
m1 = [[i] for i in range(-l, l + 1)]
r = coords_sph[:, 0]
theta = coords_sph[:, 1]
phi = coords_sph[:, 2]
ylm = (special.sph_harm(m1, l, phi, theta)).conjugate()
hypf = special.hyp1f1(l + 1 - z / kappa, 2 * l + 2, 2 * kappa * r)
rvl = hypf * (kappa * r) ** l * numpy.exp(-kappa * r)
rst = rvl * ylm
omega = numpy.array(rst)
return omega
def orbital_dip(coeff,mol):
with mol.with_common_orig((0, 0, 0)):
ao_dip1 = mol.intor_symmetric('int1e_r', comp=3)
dm_01 = coeff.T * coeff
u_lab = -numpy.einsum('xij,ji->x', ao_dip1, dm_01).real
return u_lab
def get_structure_factor(mol,
orbital_index = 0,
channel = (0,0),
lmax = 10,
hf_method = 'RHF',
casscf_conf = None,
atom_grid_level = 3,
orient_grid_size= (90,1),
move_dip_zero = True,
rmax = 40):
"""
Calculates the molecular structure factor $G$ according to the given parameters.
# Parameters:
mol : Mole
The molecule object. Initialized by invoking `pyscf.M` or `pyscf.gto.M`.
orbital_index : int
Index of the ionizing orbital relative to the HOMO. Default is 0. e.g., HOMO -> 0, LUMO -> +1, HOMO-1 -> -1, ...
channel : tuple
Parabolic channel ν=(nξ, m). Default is (0,0).
lmax : int
The cut-off limit of the angular quantum number (larger l would be cut off) used in the summation. Default is 10.
hf_method : str
Indicates whether 'RHF' or 'UHF' should be used in molecular HF calculation. Default is 'RHF'. [!] Note: Must use 'UHF' for open-shell molecules.
casscf_conf : tuple
Configuration of CASSCF calculation consisting of (n_active_orb, n_active_elec). Specifying None (by default) indicates employing HF instead of CASSCF.
atom_grid_level : int
Level of fineness of the grid used in integration (see also pyscf.dft.Grid), which controls the number of radial and angular grids around each atom in the evaluation of the integration, ranging from 0 to 9. Default is 3.
orient_grid_size : tuple
Indicates the size of the output (β,γ) grid which defines the orientation of the molecule with respect to the polarization direction of the laser field. Default is (90,1). The grid is uniform, with β ranging from [0,π) and γ ranging from [0,2π).
move_dip_zero : bool
Indicates whether to move the molecule so that the dipole of the parent ion vanishes. Default true.
rmax : float
Indicates the cut off limit of the radial grid points, points of radius>rmax would not be accounted in calculation. Default is 40.
# Returns:
A numpy array containing the structure factors $G$ on the (β,γ) orientation grid. Shape = orient_grid_size.
"""
Z = mol.charge + 1
n_xi = channel[0]
m = channel[1]
n_beta = orient_grid_size[0]
n_gamma = orient_grid_size[1]
if hf_method == 'RHF':
if mol.spin != 0:
print("[!] Warning: Spin != 0 but 'RHF' method is used. You should use 'UHF' method for non-zero spin.")
print("[i] Info: Running HF ...")
hftask = scf.RHF(mol).run(verbose=0)
if not (casscf_conf is None):
print("[i] Info: Running CASSCF ...")
casscftask = mcscf.CASSCF(hftask, casscf_conf[0], casscf_conf[1]) .run(verbose=0)
mo_occ = hftask.mo_occ
orbit_energy = hftask.mo_energy
if not (casscf_conf is None):
coeff = casscftask.mo_coeff
else:
coeff = hftask.mo_coeff
elif hf_method == 'UHF':
print("[i] Info: Running HF ...")
hftask = scf.UHF(mol).run(verbose=0)
if not (casscf_conf is None):
print("[i] Info: Running CASSCF ...")
casscftask = mcscf.UCASSCF(hftask, casscf_conf[0], casscf_conf[1]) .run(verbose=0)
mo_occ = hftask.mo_occ[0]
orbit_energy = hftask.mo_energy[0]
if not (casscf_conf is None):
coeff = casscftask.mo_coeff[0]
else:
coeff = hftask.mo_coeff[0]
else:
raise Exception("error: hf_method must be either 'RHF' or 'UHF'")
if not hftask.converged:
raise Exception("error: The Hartree-Fock calculation failed to converge, check your parameters.")
if not (casscf_conf is None):
if not casscftask.converged:
raise Exception("error: The CASSCF calculation failed to converge, check your parameters.")
print("[i] Info: Calculating structure factors ...")
index = get_homo_index(mo_occ) + orbital_index
energy_index = orbit_energy[index]
coeff = numpy.expand_dims(coeff[:, index], axis=0)
kappa = math.sqrt(-2 * energy_index)
dm = hftask.make_rdm1()
if move_dip_zero:
import copy
atom = copy.deepcopy(mol._atom)
tDict = {}
for k in gto.Mole.build.__code__.co_varnames:
if k in mol.__dict__ and not k.startswith("_"):
tDict[k] = mol.__dict__[k]
tDict["atom"] = atom
tDict["unit"] = 'Bohr'
mol = gto.M(**tDict)
u = orbital_dip(coeff,mol)
dip_moment = hftask.dip_moment(unit="A.U.",verbose=0)
D = (u - dip_moment)
move_distance = D/Z
atom = copy.deepcopy(mol._atom)
for i in range(len(atom)):
for j in range(len(atom[i][1])):
atom[i][1][j] += move_distance[j]
tDict = {}
for k in gto.Mole.build.__code__.co_varnames:
if k in mol.__dict__ and not k.startswith("_"):
tDict[k] = mol.__dict__[k]
tDict["atom"] = atom
tDict["unit"] = 'Bohr'
mol = gto.M(**tDict)
g = dft.Grids(mol)
g.level = atom_grid_level
g.radii_adjust = radi.becke_atomic_radii_adjust
g.atomic_radii = radi.COVALENT_RADII
g.radi_method = radi.gauss_chebyshev
g.prune = None
g.build()
weights = g.weights
coords_ori = g.coords
weights_ori = numpy.expand_dims(weights, axis=1)
weights_coords_ori = numpy.hstack((weights_ori, coords_ori))
weights_coords = weights_coords_ori[(weights_coords_ori[:,1])**2 + (weights_coords_ori[:,2])**2 + (weights_coords_ori[:,3])**2 < rmax**2]
weights = weights_coords[:,0]
coords = weights_coords[:,1:4]
u_lab = orbital_dip(coeff, mol)
vc_ionorbit = vc_orbit(mol, Z, coeff, coords, dm, hf_method=hf_method)
I = [None]*(lmax+1) # I[l][l+m'] stores the integrals I_{lm'}^{\nu}
for l in range(abs(m), lmax + 1):
wvl = omega_wvl(m, n_xi, kappa, Z, l)
omega = omega_r(kappa, Z, l, coords) * wvl * weights
I[l] = numpy.dot(omega, vc_ionorbit)
def factor_G(beta,gamma):
uz = (-u_lab[0] * numpy.cos(gamma) + u_lab[1] * numpy.sin(gamma))*numpy.sin(beta) + u_lab[2] * numpy.cos(beta)
Gsum = numpy.complex128(0.0)
for l in range(abs(m), lmax + 1): # l in |m|:lmax
for mp in range(-l,l+1): # m' in -l:l
Gsum += I[l][l+mp] * wigner_dl(l,l,m,mp,beta)[0] * numpy.exp(-1j*mp*gamma)
return Gsum * numpy.exp(-kappa * uz)
beta_grid = numpy.linspace(0, math.pi, n_beta)
gamma_grid = numpy.linspace(0, 2*math.pi, n_gamma)
G_grid = numpy.zeros(shape=(n_beta,n_gamma), dtype=numpy.complex128)
for ibeta in range(n_beta):
for igamma in range(n_gamma):
G_grid[ibeta,igamma] = factor_G(beta_grid[ibeta],gamma_grid[igamma])
return G_grid