From b8ccf56a4e19d92dbec9718535fe70794c47950b Mon Sep 17 00:00:00 2001 From: Mark Dewing Date: Thu, 15 Sep 2016 22:06:41 -0500 Subject: [PATCH] Read and evaluate plane wave files. Reads the output of Quantum Espresso from pw2qmcpack.x. The code to perform the evaluation of the orbitals is derived from symbolic expressions. --- Wavefunctions/read_pw.py | 236 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 236 insertions(+) create mode 100644 Wavefunctions/read_pw.py diff --git a/Wavefunctions/read_pw.py b/Wavefunctions/read_pw.py new file mode 100644 index 0000000..e0338b4 --- /dev/null +++ b/Wavefunctions/read_pw.py @@ -0,0 +1,236 @@ +from __future__ import print_function + +import h5py +import sympy +from sympy.utilities.lambdify import lambdastr +import numpy as np +import math + +# Read and evaluate orbitals output from Quantum Espresso using pw2qmcpack.x + + +class PWOrbitalFile(object): + def __init__(self): + self.e, self.de, self.dde = create_eval_funcs() + + self.version = None + self.number_of_kpoints = 0 + self.number_of_electrons = 0 + self.atom_pos = None + self.primitive_vectors = None + self.kspace_basis = None + self.raw_kpoints = [] + self.kpoints = [] + self.raw_gvectors = None + self.gvectors = None + self.coeffs = [] + self.number_of_bands = 0 + + def evaluate(self, k_index, band_index, r): + coeff = self.coeffs[k_index][band_index] + k = self.kpoints[k_index] + n = self.gvectors.shape[0] + return self.e(n, coeff, self.gvectors, k, r) + + def evaluate_all(self, k_index, band_index, r): + coeff = self.coeffs[k_index][band_index] + n = self.gvectors.shape[0] + k = self.kpoints[k_index] + e = self.e(n, coeff, self.gvectors, k, r) + #fde = fake_de(n, coeff, self.gvectors, k, r) + de = self.de(n, coeff, self.gvectors, k, r) + dde = self.dde(n, coeff, self.gvectors, k, r) + return e, de, dde + + +# Reciprocal lattice vectors +def get_kspace_basis(basis): + # Volume factor for reciprocal lattice + a1, a2, a3 = basis + vol = a1.dot(np.cross(a2, a3)) + + pre = 2*math.pi + b1 = pre*np.cross(a2, a3)/vol + b2 = pre*np.cross(a3, a1)/vol + b3 = pre*np.cross(a1, a2)/vol + return [b1, b2, b3] + + +def read_pw_file(filename): + pw = PWOrbitalFile() + with h5py.File(filename) as f: + pw.version = f.get('application/version') + pw.number_of_kpoints = f.get('electrons/number_of_kpoints')[0] + pw.number_of_electrons = f.get('electrons/number_of_electrons') + pw.atom_pos = f.get('atoms/positions') + pw.primitive_vectors = f.get('supercell/primitive_vectors') + pw.kspace_basis = get_kspace_basis(pw.primitive_vectors) + + pw.raw_kpoints = [] + pw.kpoints = [] + for i in range(pw.number_of_kpoints): + raw_kpoint = np.array(f.get('electrons/kpoint_%d/reduced_k' % i)) + cell_kpoint = np.array(sum([pw.kspace_basis[j]*raw_kpoint[j] for j in range(3)])) + pw.raw_kpoints.append(raw_kpoint) + pw.kpoints.append(cell_kpoint) + + # We are assuming the G vectors are the same for each k-point (?) + pw.raw_gvectors = np.array(f.get('electrons/kpoint_0/gvectors')) + pw.gvectors = np.empty(pw.raw_gvectors.shape) + for i in range(pw.raw_gvectors.shape[0]): + conv_g = [pw.kspace_basis[j]*pw.raw_gvectors[i, j] for j in range(3)] + cell_g = np.array(sum(conv_g)) + pw.gvectors[i, :] = cell_g + + pw.coeffs = [] + pw.number_of_bands = f.get('electrons/kpoint_0/spin_0/number_of_states')[0] + for ik in range(pw.number_of_kpoints): + coeffs = [] + for ib in range(pw.number_of_bands): + tmp_c = f.get('electrons/kpoint_%d/spin_0/state_%d/psi_g' % (ik, ib)) + tmp_a = tmp_c[:, 0] + 1j*tmp_c[:, 1] + coeffs.append(tmp_a) + pw.coeffs.append(coeffs) + + return pw + + +# Use symbolic expressions for the orbital sum (Fourier transform) + +# For some reason, substitution doesn't work right for Sums. +# This is a workaround + +def sum_subs(s, old, new): + if isinstance(s, sympy.Sum): + new_func = s.function.subs(old, new) + return sympy.Sum(new_func, s.limits) + else: + return s.subs(old, new) + + +# before lambdify, to accept numpy arrays starting at 0 +def adjust_sum_limits(e): + if isinstance(e, sympy.Sum): + lim = e.limits[0] + return sympy.Sum(e.function, (lim[0], lim[1]-1, lim[2]-1)) + return e + + +def create_eval_funcs(): + n = sympy.Symbol('n', integer=True) + c = sympy.IndexedBase('c') + g = sympy.IndexedBase('g') + k = sympy.IndexedBase('k') + r = sympy.Symbol('r') + + # Later will change r from from Symbol to IndexedBase + real_r = sympy.IndexedBase('r') + + # Hack to use ':' in array index in lambdify + all_s = sympy.Symbol(':', integer=True, nonnegative=True) + + i = sympy.Symbol('i', integer=True, nonnegative=True) + j = sympy.Symbol('j', integer=True, nonnegative=True) + + # Fourier transform + e = sympy.Sum(c[i]*(sympy.cos((g[i]+k)*r) + sympy.I * sympy.sin((g[i]+k)*r)), (i, 1, n)) + #e = sympy.Sum(c[i]*(sympy.cos(g[i]*r) + 1j * sympy.sin(g[i]*r)), (i, 1, n)) + + # gradient + d_e = e.diff(r) + + # Laplacian + dd_e = e.diff(r, 2) + + # Mathematical notation often uses a variable 'level of detail' in its + # representation. We wish to imitate that here. Otherwise the full details + # of every symbol need to be carried around and represented from the very + # start, making some of these expressions visually complicated. + # Up until now a scalar symbol has been used for r and a 1-D array for g + # But r is really a length-3 vector, and so is each element of g. + # The whole of g can be stored as a 2-D array. + # This step replaces symbols g and r with the more detailed vector versions. + + dot = sympy.Sum((g[i, j] + k[j]) * real_r[j], (j, 0, 2)) + mag = sympy.Sum((g[i, j] + k[j]) * (g[i, j] + k[j]), (j, 0, 2)) + + e_vec = sum_subs(e, r*(g[i]+k), dot) + de_vec = sum_subs( + sum_subs(d_e, (g[i]+k)*r, dot), g[i], g[i, all_s]) + dde_vec = sum_subs( + sum_subs(dd_e, (g[i]+k)*r, dot), (g[i]+k)**2, mag) + + e_vec = adjust_sum_limits(e_vec) + de_vec = adjust_sum_limits(de_vec) + dde_vec = adjust_sum_limits(dde_vec) + + # Use the following to see what Python code is produced by lambdify + #e_str = lambdastr([n, c, g, k, r] ,e_vec) + #de_str = lambdastr([n, c, g, k, r], de_vec) + #dde_str = lambdastr([n, c, g, k, r], dde_vec) + #print('e lambdastr = ', e_str) + #print('de lambdastr = ', de_str) + #print('dde lambdastr = ', dde_str) + + e2 = sympy.lambdify([n, c, g, k, r], e_vec) + de2 = sympy.lambdify([n, c, g, k, r], de_vec) + dde2 = sympy.lambdify([n, c, g, k, r], dde_vec) + + # could apply CSE to make evaluating all of them faster + + return e2, de2, dde2 + + +# Expanded version of the output of lambdastr, for debugging +def fake_de(n, c, g, k, r): + s = 0.0 + for i in range(0, n): + tmp1 = sum(r[j]*(g[i, j]+k[j]) for j in range(0, 3)) + sin1 = -math.sin(tmp1) + #cos1 = sympy.I*math.cos(tmp1) + cos1 = 1j*math.cos(tmp1) + #print(i,'sin1',sin1,cos1) + #print('c',c[i]) + #tmp2 = sin1*(g[i,:] + k[:]) + #tmp3 = cos1*(g[i,:] + k[:]) + tmp4 = complex(cos1, sin1)*(g[i, :]+k[:]) + + #print(tmp2, tmp3) + s += (tmp4)*c[i] + return s + #sum( (-math.sin((sum(r[j]*g[i, j] for j in range(0, 2+1))))*g[i, :] + sympy.I*math.cos((sum(r[j]*g[i, j] for j in range(0, 2+1))))*g[i, :])*c[i] for i in range(0, n - 1+1)) + + +def test_eval_funcs(): + e, de, dde = create_eval_funcs() + + cc = np.array([complex(1, 0)]) + gg = np.array([[1., 0., 0.]]) + kk = np.array([0., 0., 0.]) + print('shape = ', cc.shape) + rr = np.array([1., 0., 0.]) + + out = e(cc.shape[0], cc, gg, kk, rr) + print(out) + d_out = de(cc.shape[0], cc, gg, kk, rr) + print('grad = ', d_out) + + dd_out = dde(cc.shape[0], cc, gg, kk, rr) + print('lap = ', dd_out) + +if __name__ == '__main__': + pw = read_pw_file('LiH-arb.pwscf.h5') + #print (pw.gvectors) + #print (pw.coeffs[0][0]) + + #test_eval_funcs() + + r = np.array([0.0, 0.0, 0.0]) + #v = pw.evaluate(0, 1, r) + #print(' v = ', v) + for i in range(36): + #x = 3.55*i/36 - 2*3.55/12 + x = 3.55*i/36 + r = np.array([x, 0.0, 0.0]) + v, dv, ddv = pw.evaluate_all(0, 0, r) + print (x, v, dv[0], ddv)