Coverage for gpaw/nlopt/shg.py: 74%
140 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-14 00:18 +0000
1import numpy as np
2from ase.parallel import parprint
3from ase.units import _e, Bohr, Ha, J
4from ase.utils.timing import Timer
6from gpaw.mpi import world
7from gpaw.nlopt.matrixel import get_derivative, get_rml
8from gpaw.utilities.progressbar import ProgressBar
11def get_shg(
12 nlodata,
13 freqs=[1.0],
14 eta=0.05,
15 pol='yyy',
16 eshift=0.0,
17 gauge='lg',
18 ftol=1e-4, Etol=1e-6,
19 band_n=None,
20 out_name='shg.npy'):
21 """
22 Calculate SHG spectrum within the independent particle approximation (IPA)
23 for nonmagnetic semiconductors.
25 Output: shg.npy file with numpy array containing the spectrum and
26 frequencies.
28 Parameters
29 ----------
30 nlodata
31 Data object of class NLOData. Contains energies, occupancies and
32 momentum matrix elements.
33 freqs
34 Excitation frequency array (a numpy array or list).
35 eta
36 Broadening, a number or an array (default 0.05 eV).
37 pol
38 Tensor element (default 'yyy').
39 gauge
40 Choose the gauge ('lg' or 'vg').
41 Etol, ftol
42 Tol. in energy and fermi to consider degeneracy.
43 band_n
44 List of bands in the sum (default 0 to nb).
45 out_name
46 Output filename (default 'shg.npy').
48 """
50 # Start a timer
51 timer = Timer()
52 parprint(f'Calculating SHG spectrum (in {world.size:d} cores).')
54 # Covert inputs in eV to Ha
55 freqs = np.array(freqs)
56 nw = len(freqs)
57 w_lc = (freqs + 1e-12 + 1j * eta) / Ha # Add small value to avoid 0
58 Etol /= Ha
59 eshift /= Ha
61 # Useful variables
62 pol_v = ['xyz'.index(ii) for ii in pol]
64 # Use the TRS to reduce calculation time
65 w_l = np.hstack((-w_lc[-1::-1], w_lc))
66 nw = 2 * nw
67 parprint(f'Calculation in the {gauge} gauge for element {pol}.')
69 # Load the required data
70 with timer('Load and distribute the data'):
71 k_info = nlodata.distribute()
72 if k_info:
73 tmp = list(k_info.values())[0]
74 nb = len(tmp[1])
75 nk = len(k_info) * world.size # Approximately
76 if band_n is None:
77 band_n = list(range(nb))
78 mem = 6 * 3 * nk * nb**2 * 16 / 2**20
79 parprint(f'At least {mem:.2f} MB of memory is required.')
81 # Initial call to print 0% progress
82 count = 0
83 ncount = len(k_info)
84 if world.rank == 0:
85 pb = ProgressBar()
87 # Initialize the outputs
88 sum2_l = np.zeros((nw), complex)
89 sum3_l = np.zeros((nw), complex)
91 # Do the calculations
92 for _, (we, f_n, E_n, p_vnn) in k_info.items():
93 # Which gauge
94 if gauge == 'vg':
95 with timer('Sum over bands'):
96 tmp = shg_velocity_gauge(
97 w_l, f_n, E_n, p_vnn, pol_v,
98 band_n, ftol, Etol, eshift)
99 elif gauge == 'lg':
100 with timer('Position matrix elements calculation'):
101 r_vnn, D_vnn = get_rml(E_n, p_vnn, pol_v, Etol=Etol)
103 with timer('Compute generalized derivative'):
104 rd_vvnn = get_derivative(E_n, r_vnn, D_vnn, pol_v, Etol=Etol)
106 with timer('Sum over bands'):
107 tmp = shg_length_gauge(
108 w_l, f_n, E_n, r_vnn, rd_vvnn, D_vnn, pol_v,
109 band_n, ftol, Etol, eshift)
110 else:
111 parprint('Gauge ' + gauge + ' not implemented.')
112 raise NotImplementedError
114 # Add it to previous with a weight
115 sum2_l += tmp[0] * we
116 sum3_l += tmp[1] * we
118 # Print the progress
119 if world.rank == 0:
120 pb.update(count / ncount)
121 count += 1
123 if world.rank == 0:
124 pb.finish()
125 with timer('Gather data from cores'):
126 world.sum(sum2_l)
127 world.sum(sum3_l)
129 # Make the output in SI unit
130 chi_l = make_output(gauge, sum2_l, sum3_l)
131 # A multi-col output
132 nw = len(freqs)
133 chi_l = chi_l[nw:] + chi_l[nw - 1::-1]
134 shg = np.vstack((freqs, chi_l))
136 # Save it to the file
137 if world.rank == 0:
138 np.save(out_name, shg)
140 # Print the timing
141 timer.write()
143 return shg
146def shg_velocity_gauge(
147 w_l, f_n, E_n, p_vnn, pol_v,
148 band_n=None, ftol=1e-4, Etol=1e-6, eshift=0):
149 """
150 Loop over bands for computing in velocity gauge
152 Input:
153 w_l Complex frequency array
154 f_n Fermi levels
155 E_n Energies
156 p_vnn Momentum matrix elements
157 pol_v Tensor element
158 band_n Band list
159 Etol, ftol Tol. in energy and fermi to consider degeneracy
160 eshift Bandgap correction
161 Output:
162 sum2_l, sum3_l Output 2 and 3 bands terms
163 """
165 # Initialize variables
166 nb = len(f_n)
167 if band_n is None:
168 band_n = list(range(nb))
169 sum2_l = np.zeros(w_l.size, complex)
170 sum3_l = np.zeros(w_l.size, complex)
172 # Loop over bands
173 for nni in band_n:
174 for mmi in band_n:
175 # Remove non important term using TRS
176 if mmi <= nni:
177 continue
179 # Useful variables
180 fnm = f_n[nni] - f_n[mmi]
181 Emn = E_n[mmi] - E_n[nni] + fnm * eshift
183 # Comute the 2-band term
184 if np.abs(Emn) > Etol and np.abs(fnm) > ftol:
185 pnml = (p_vnn[pol_v[0], nni, mmi]
186 * (p_vnn[pol_v[1], mmi, nni]
187 * (p_vnn[pol_v[2], mmi, mmi]
188 - p_vnn[pol_v[2], nni, nni])
189 + p_vnn[pol_v[2], mmi, nni]
190 * (p_vnn[pol_v[1], mmi, mmi]
191 - p_vnn[pol_v[1], nni, nni])) / 2)
192 sum2_l += 1j * fnm * np.imag(pnml) * \
193 (1 / (Emn**4 * (w_l - Emn)) -
194 16 / (Emn**4 * (2 * w_l - Emn)))
196 # Loop over the last band index
197 for lli in band_n:
198 fnl = f_n[nni] - f_n[lli]
199 fml = f_n[mmi] - f_n[lli]
201 # Do not do zero calculations
202 if np.abs(fnl) < ftol and np.abs(fml) < ftol:
203 continue
205 # Compute the susceptibility with 1/w form
206 Eln = E_n[lli] - E_n[nni] + fnl * eshift
207 Eml = E_n[mmi] - E_n[lli] - fml * eshift
208 pnml = (p_vnn[pol_v[0], nni, mmi]
209 * (p_vnn[pol_v[1], mmi, lli]
210 * p_vnn[pol_v[2], lli, nni]
211 + p_vnn[pol_v[2], mmi, lli]
212 * p_vnn[pol_v[1], lli, nni]))
213 pnml = 1j * np.imag(pnml) / 2
215 # Compute the divergence-free terms
216 if np.abs(Emn) > Etol and np.abs(
217 Eml) > Etol and np.abs(Eln) > Etol:
218 ftermD = (16 / (Emn**3 * (2 * w_l - Emn))
219 * (fnl / (Emn - 2 * Eln)
220 + fml / (Emn - 2 * Eml))) \
221 + fnl / (Eln**3 * (2 * Eln - Emn)
222 * (w_l - Eln)) \
223 + fml / (Eml**3 * (2 * Eml - Emn)
224 * (w_l - Eml))
225 sum3_l += pnml * ftermD
227 return sum2_l, sum3_l
230def shg_length_gauge(
231 w_l, f_n, E_n, r_vnn, rd_vvnn, D_vnn, pol_v,
232 band_n=None, ftol=1e-4, Etol=1e-6, eshift=0):
233 """
234 Loop over bands for computing in length gauge
236 Input:
237 w_l Complex frequency array
238 f_n Fermi levels
239 E_n Energies
240 r_vnn Momentum matrix elements
241 rd_vvnn Generalized derivative of position
242 D_vnn Velocity difference
243 pol_v Tensor element
244 band_n Band list
245 Etol, ftol Tol. in energy and fermi to consider degeneracy
246 eshift Band gap correction in eV
247 Output:
248 sum2_l, sum3_l Output 2 and 3 bands terms
249 """
251 # Initialize variables
252 nb = len(f_n)
253 if band_n is None:
254 band_n = list(range(nb))
255 sum2_l = np.zeros(w_l.size, complex)
256 sum3_l = np.zeros(w_l.size, complex)
258 # Loop over bands
259 for nni in band_n:
260 for mmi in band_n:
261 # Remove the non important term using TRS
262 if mmi <= nni:
263 continue
264 fnm = f_n[nni] - f_n[mmi]
265 Emn = E_n[mmi] - E_n[nni] + fnm * eshift
267 # Two band part
268 if np.abs(fnm) > ftol:
269 tmp = 2 * np.imag(
270 r_vnn[pol_v[0], nni, mmi]
271 * (rd_vvnn[pol_v[1], pol_v[2], mmi, nni]
272 + rd_vvnn[pol_v[2], pol_v[1], mmi, nni])) \
273 / (Emn * (2 * w_l - Emn))
274 tmp += np.imag(
275 r_vnn[pol_v[1], mmi, nni]
276 * rd_vvnn[pol_v[2], pol_v[0], nni, mmi]
277 + r_vnn[pol_v[2], mmi, nni]
278 * rd_vvnn[pol_v[1], pol_v[0], nni, mmi]) \
279 / (Emn * (w_l - Emn))
280 tmp += np.imag(
281 r_vnn[pol_v[0], nni, mmi]
282 * (r_vnn[pol_v[1], mmi, nni]
283 * D_vnn[pol_v[2], mmi, nni]
284 + r_vnn[pol_v[2], mmi, nni]
285 * D_vnn[pol_v[1], mmi, nni])) \
286 * (1 / (w_l - Emn)
287 - 4 / (2 * w_l - Emn)) / Emn**2
288 tmp -= np.imag(
289 r_vnn[pol_v[1], mmi, nni]
290 * rd_vvnn[pol_v[0], pol_v[2], nni, mmi]
291 + r_vnn[pol_v[2], mmi, nni]
292 * rd_vvnn[pol_v[0], pol_v[1], nni, mmi]) \
293 / (2 * Emn * (w_l - Emn))
294 sum2_l += 1j * fnm * tmp / 2 # 1j imag
296 # Three band term
297 for lli in band_n:
298 fnl = f_n[nni] - f_n[lli]
299 fml = f_n[mmi] - f_n[lli]
300 Eml = E_n[mmi] - E_n[lli] - fml * eshift
301 Eln = E_n[lli] - E_n[nni] + fnl * eshift
302 # Do not do zero calculations
303 if (np.abs(fnm) < ftol and np.abs(fnl) < ftol
304 and np.abs(fml) < ftol):
305 continue
306 if np.abs(Eln - Eml) < Etol:
307 continue
309 rnml = np.real(
310 r_vnn[pol_v[0], nni, mmi]
311 * (r_vnn[pol_v[1], mmi, lli]
312 * r_vnn[pol_v[2], lli, nni]
313 + r_vnn[pol_v[2], mmi, lli]
314 * r_vnn[pol_v[1], lli, nni])) / (2 * (Eln - Eml))
315 if np.abs(fnm) > ftol:
316 sum3_l += 2 * fnm / (2 * w_l - Emn) * rnml
317 if np.abs(fnl) > ftol:
318 sum3_l += -fnl / (w_l - Eln) * rnml
319 if np.abs(fml) > ftol:
320 sum3_l += fml / (w_l - Eml) * rnml
322 return sum2_l, sum3_l
325def make_output(gauge, sum2_l, sum3_l):
326 """
327 Multiply prefactors and return second-order chi in SI units [m / V]
329 Input:
330 gauge Chosen gauge
331 sum2_l 2-bands term
332 sum3_l 3-bands term
333 Output:
334 chi_l Output chi as an array
335 """
337 # 4 * pi from vacuum permittivty eps_0 in atomic units
338 prefactor = 4.0 * np.pi
339 # Pi factors from BZ volume
340 prefactor /= (2.0 * np.pi)**3
341 # atomic units [Bohr * elementary charge / Hartee] to SI units [m / V]
342 prefactor *= (Bohr * 1e-10 * _e) / (Ha / J)
344 if gauge == 'lg':
345 chi_l = prefactor * (1j * sum2_l + sum3_l)
346 elif gauge == 'vg':
347 chi_l = prefactor * 1j / 2 * (sum2_l + sum3_l)
348 else:
349 parprint('Gauge ' + gauge + ' not implemented.')
350 raise NotImplementedError
352 return chi_l