from __future__ import print_function, division
from copy import copy
from timeit import default_timer as timer
import numpy as np
from scipy.sparse import csr_matrix, coo_matrix
from scipy.linalg import blas
import scipy.sparse.linalg as splin
from pynao import mf
from pynao.m_tddft_iter_gpu import tddft_iter_gpu_c
from pynao.m_siesta_units import siesta_conv_coefficients
[docs]class chi0_matvec(mf):
"""
A class to organize the application of non-interacting response to a vector
"""
def __init__(self, **kw):
"""
The class handling the linear equation 2.74 in Ref 1 and setting up
the iterative solver.
Class Inputs:
krylov_solver (string): Type of Krylov solver to use to solve the
linear system of equation. The Krylov solvers are the ones offered
by Scipy library, see https://docs.scipy.org/doc/scipy/reference/sparse.linalg.html
Available solvers are
* lgmres
* bicgstab
* gmres
* gcrotmk
* qmr
* minres
* bicg
* cg
* cgs
Recommended solver is lgmres. bicgstab can provide better performances,
but can be unstable.
krylov_options (dict): the options for the iterative Krylov solver,
See the scipy solver specific page for the available options.
use_initial_guess_ite_solver (bool): Use initial guess for the Krylov
solver.
dealloc_hsx (bool): deallocate the hsx matrix in order to save memory
iter_broadening (float): the broadening in the iterative solver
See eps term in equation 2.61 of Ref. 1
default value 0.00367493 Ha
GPU (bool): Use GPU for the matvec operation
nfermi_tol (float): the tolerance for the determination of the Fermi
level from the orbital occupation. Default 1.0
telec (float): Temperature of the system. If not specified, it will be
the temperature of the DFT simulation. Unit is Ha??
fermi_energy (float): The fermi energy of the system. If not specified,
it will be the temperature of the DFT simulation. Unit is Ha??
References:
1. M. Barbry, “Plasmons in nanoparticles: Atomistic ab initio theory
for large systems, Ph.D. thesis, University of Basque Country,
Donostia-San Sebastián, Spain, 2018,
http://cfm.ehu.es/view/files/MArc_barbry_2-1.pdf
"""
from pynao.m_fermi_dirac import fermi_dirac_occupations
if "use_initial_guess_ite_solver" in kw:
self.use_initial_guess_ite_solver = kw["use_initial_guess_ite_solver"]
else:
self.use_initial_guess_ite_solver = False
krylov_solvers = {"lgmres": splin.lgmres,
"gmres": splin.gmres,
"gcrotmk": splin.gcrotmk,
"qmr": splin.qmr,
"minres": splin.minres,
"bicgstab": splin.bicgstab,
"bicg": splin.bicg,
"cg": splin.cg,
"cgs": splin.cgs}
if "krylov_solver" in kw:
self.krylov_solver = krylov_solvers[kw["krylov_solver"]]
else:
self.krylov_solver = krylov_solvers["lgmres"]
if "krylov_options" in kw:
self.krylov_options = kw["krylov_options"]
else:
self.krylov_options = {"tol": 1.0e-3,
"atol": 1.0e-3,
"maxiter": 1000}
mf.__init__(self, **kw)
if self.dtype == np.float32:
self.spmv = blas.sspmv
self.gemm = blas.sgemm
elif self.dtype == np.float64:
self.spmv = blas.dspmv
self.gemm = blas.dgemm
else:
raise ValueError("dtype can be only float32 or float64")
self.dealloc_hsx = kw['dealloc_hsx'] if 'dealloc_hsx' in kw else True
self.eps = kw['iter_broadening'] if 'iter_broadening' in kw else 0.00367493
self.GPU = GPU = kw['GPU'] if 'GPU' in kw else False
self.nfermi_tol = nfermi_tol = kw['nfermi_tol'] if 'nfermi_tol' in kw else 1e-5
self.telec = kw['telec'] if 'telec' in kw else self.telec
self.fermi_energy = kw['fermi_energy'] if 'fermi_energy' in kw else self.fermi_energy
if self.GPU:
try:
import cupy as cp
except ImportError as err:
raise ImportError("Could not import cupy:\n {}".format(err))
if not self.use_numba:
raise ValueError("GPU calculations require Numba")
if self.dtype == np.float32:
from pynao.m_div_eigenenergy_numba_gpu import div_eigenenergy_gpu_float32
self.div_numba = div_eigenenergy_gpu_float32
else:
from pynao.m_div_eigenenergy_numba_gpu import div_eigenenergy_gpu_float64
self.div_numba = div_eigenenergy_gpu_float64
elif self.use_numba:
from pynao.m_div_eigenenergy_numba import div_eigenenergy_numba
self.div_numba = div_eigenenergy_numba
else:
self.div_numba = None
# deallocate hsx
if hasattr(self, 'hsx') and self.dealloc_hsx:
self.hsx.deallocate()
# Just a pointer here. Is it ok?
self.ksn2e = self.mo_energy
if 'fermi_energy' in kw:
if self.verbosity > 0:
print(__name__, 'Fermi energy is specified => recompute occupations')
ksn2fd = fermi_dirac_occupations(self.telec, self.ksn2e, self.fermi_energy)
for s,n2fd in enumerate(ksn2fd[0]):
if not all(n2fd>self.nfermi_tol):
continue
print(self.telec, s, self.nfermi_tol, n2fd)
raise RuntimeError(__name__, 'telec is too high?')
ksn2f = self.ksn2f = (3-self.nspin)*ksn2fd
else:
ksn2f = self.ksn2f = self.mo_occ
self.nfermi = np.array([np.argmax(ksn2f[0,s] < self.nfermi_tol) \
for s in range(self.nspin)], dtype=int)
self.vstart = np.array([np.argmax(1.0-ksn2f[0,s]>=self.nfermi_tol) \
for s in range(self.nspin)], dtype=int)
self.xocc = [self.mo_coeff[0,s,:nfermi,:,0] for s,nfermi in enumerate(self.nfermi)]
self.xvrt = [self.mo_coeff[0,s,vstart:,:,0] for s,vstart in enumerate(self.vstart)]
if self.verbosity > 4:
print(__name__, '\t====> self.xocc[0].dtype ', self.xocc[0].dtype)
print(__name__, '\t====> self.xocc[0].shape ', self.xocc[0].shape)
print(__name__, '\t====> self.xvrt[0].dtype ', self.xvrt[0].dtype)
print(__name__, '\t====> self.xvrt[0].shape ', self.xvrt[0].shape)
print(__name__, '\t====> MO energies (ksn2e) (eV):\n{},\tType: {}'.format(
self.ksn2e*siesta_conv_coefficients["ha2ev"], self.ksn2e.dtype))
print(__name__, '\t====> Occupations (ksn2f):\n{},\tType: {}'.format(
self.ksn2f, self.ksn2f.dtype))
self.rf0_ncalls = 0
self.chi0_timing = np.zeros((9), dtype=np.float64)
self.moms0,self.moms1 = self.pb.comp_moments(dtype=self.dtype)
if self.GPU:
self.initialize_chi0_matvec_GPU()
[docs] def initialize_chi0_matvec_GPU(self):
"""
Initialize GPU variable.
Copy xocc, xvrt, ksn2e and ksn2f to the device
"""
import cupy as cp
block_sizes = np.array([32, 64, 128, 256, 512, 1024], dtype=np.int32)
self.block_size = [] # np.array([32, 32], dtype=np.int32) # threads by block
self.grid_size = [] #np.array([0, 0], dtype=np.int32) # number of blocks
self.xocc_gpu = []
self.xvrt_gpu = []
self.ksn2e_gpu = cp.asarray(self.ksn2e)
self.ksn2f_gpu = cp.asarray(self.ksn2f)
for spin in range(self.nspin):
dimensions = [self.nfermi[spin], self.nprod]
self.block_size.append([32, 32])
# Adjust the block size as function of the dimension of the array
for i in range(2):
ibs = 0
for bs in block_sizes:
if bs > dimensions[i]:
break
ibs += 1
if ibs > 0:
self.block_size[spin][i] = block_sizes[ibs-1]
else:
self.block_size[spin][i] = block_sizes[0]
self.grid_size.append([0, 0])
for i in range(2):
if dimensions[i] <= block_sizes[0]:
self.block_size[spin][i] = dimensions[i]
self.grid_size[spin][i] = 1
else:
self.grid_size[spin][i] = dimensions[i]//self.block_size[spin][i] + 1
print("spin {}: block_size: ({}, {}); grid_size: ({}, {})".format(
spin, self.block_size[spin][0], self.block_size[spin][1],
self.grid_size[spin][0], self.grid_size[spin][1]))
self.xocc_gpu.append(cp.asarray(self.xocc[spin]))
self.xvrt_gpu.append(cp.asarray(self.xvrt[spin]))
[docs] def apply_rf0(self, sp2v, comega, chi0_mv):
"""
This applies the non-interacting response function to a vector (a set of vectors?)
"""
#expect_shape=tuple([self.nspin*self.nprod])
# gw_iter does not have the same expected shape the tddft_iter ....
#assert np.all(sp2v.shape==expect_shape), "{} {}".format(sp2v.shape, expect_shape)
self.rf0_ncalls+=1
return chi0_mv(self, sp2v, comega)
[docs] def comp_polariz_nonin_xx_atom_split(self, comegas):
"""
Compute the non interacting polarizability along the xx direction
"""
aw2pxx = np.zeros((self.natoms, comegas.shape[0]), dtype=self.dtypeComplex)
vext = np.transpose(self.moms1)
for iw, comega in enumerate(comegas):
dn0 = self.apply_rf0(vext[0], comega, self.chi0_mv)
for ia in range(self.natoms):
dn0_atom = np.zeros(self.nprod, dtype=self.dtypeComplex)
st = self.pb.c2s[ia]
fn = self.pb.c2s[ia+1]
dn0_atom[st:fn] = dn0[st:fn]
aw2pxx[ia, iw] = dn0_atom.dot(vext[0])
self.write_chi0_mv_timing("tddft_iter_polariz_nonin_split_chi0_mv.txt")
return aw2pxx
[docs] def write_chi0_mv_timing(self, fname):
with open(fname, "w") as fl:
fl.write("# Total number call chi0_mv: {}\n".format(self.rf0_ncalls))
fl.write("# step time [s]\n")
for it, time in enumerate(self.chi0_timing):
fl.write("{}: {}\n".format(it, time))
[docs] def comp_dens_along_eext(self, comegas, eext=np.array([1.0, 0.0, 0.0]),
tmp_fname=None, inter=False):
"""
Compute a the average interacting polarizability along the eext direction
for the frequencies comegas.
Input Parameters:
comegas (1D array, complex): the real part contains the frequencies at which the polarizability
should be computed. The imaginary part id the width of the polarizability define as self.eps
eext (1D xyz array, real): direction of the external field
tmp_fname (string, optional): temporary file to where to store results
Particularly useful for large calculations
inter (bool, optional): Perform interactive calculations (or not)
Other Calculated quantity:
self.p_mat (complex array, dim: [3, 3, comega.size]): store the (3, 3) polarizability matrix
[[Pxx, Pxy, Pxz],
[Pyx, Pyy, Pyz],
[Pzx, Pzy, Pzz]] for each frequency.
self.dn (complex array, dim: [3, comegas.size, self.nprod]): store the density change
"""
if tmp_fname is not None:
if not isinstance(tmp_fname, str):
raise ValueError("tmp_fname must be a string")
else:
tmp_re = open(tmp_fname+".real", "w")
tmp_re.write("# All atomic units\n")
tmp_re.write("# w (Ha) Pxx Pxy Pxz Pyx Pyy Pyz Pzx Pzy Pzz\n")
tmp_im = open(tmp_fname+".imag", "w")
tmp_im.write("# All atomic units\n")
tmp_im.write("# w Pxx Pxy Pxz Pyx Pyy Pyz Pzx Pzy Pzz\n")
if isinstance(eext, list):
eext = np.array(eext)
assert eext.size == 3
nww = len(comegas)
p_mat = np.zeros((3, 3, nww), dtype=self.dtypeComplex)
dn = np.zeros((3, nww, self.nspin*self.nprod), dtype=self.dtypeComplex)
edir = eext/np.dot(eext, eext)
vext = np.zeros((self.nspin*self.nprod, 3), dtype=self.moms1.dtype)
#veff = np.zeros((self.nspin*self.nprod, 3), dtype=self.dtypeComplex)
for ixyz in range(3):
vext[:, ixyz] = np.concatenate([self.moms1[:, ixyz] for s in range(self.nspin)])
self.prev_sol = np.zeros((self.nspin*self.nprod, 3), dtype=self.dtypeComplex)
for iw, comega in enumerate(comegas):
if iw > 0:
dnprev = dn[:, iw-1, :]
dn[:, iw, :], p_mat[:, :, iw] = \
self.calc_dens_edir_omega(iw, nww, comega, edir, vext,
tmp_fname=tmp_fname, inter=inter)
if inter:
self.write_chi0_mv_timing("tddft_iter_inter_chi0_mv.txt")
else:
self.write_chi0_mv_timing("tddft_iter_nonin_chi0_mv.txt")
print("Total number of iterations: ", self.rf0_ncalls)
return dn, p_mat
[docs] def calc_dens_edir_omega(self, iw, nww, w, edir, vext, tmp_fname=None, inter=False):
"""
Calculate the density change and polarizability for a specific frequency
"""
Pmat = np.zeros((3, 3), dtype=self.dtypeComplex)
dn = np.zeros((3, self.nspin*self.nprod), dtype=self.dtypeComplex)
for xyz, Exyz in enumerate(edir):
if abs(Exyz) < 1.0e-12:
continue
chi0mv_ncalls_ite = self.rf0_ncalls
t1 = timer()
if inter:
veff = self.comp_veff(vext[:, xyz], w, prev_sol=self.prev_sol[:, xyz])
if self.use_initial_guess_ite_solver:
self.prev_sol[:, xyz] = veff
dn[xyz, :] = self.apply_rf0(veff, w, self.chi0_mv)
else:
dn[xyz, :] = self.apply_rf0(vext[:, xyz], w, self.chi0_mv)
t2 = timer()
chi0mv_ncalls_ite = self.rf0_ncalls - chi0mv_ncalls_ite
if self.verbosity > 1:
mess = "dir: {0}, iw: {1}/{2}; w: {3:.4f};".format(xyz, iw, nww,
w.real*siesta_conv_coefficients["ha2ev"])
mess += "nite: {0}; time: {1}".format(chi0mv_ncalls_ite, t2-t1)
print(mess)
for xyzp in range(edir.size):
Pmat[xyz, xyzp] = vext[:, xyzp].dot(dn[xyz, :])
if tmp_fname is not None:
tmp_re = open(tmp_fname + ".real", "a")
tmp_re.write("{0:.6e} ".format(w.real))
tmp_im = open(tmp_fname + ".imag", "a")
tmp_im.write("{0:.6e} ".format(w.real))
for i in range(3):
for j in range(3):
tmp_re.write("{0:.6e} ".format(Pmat[i, j].real))
tmp_im.write("{0:.6e} ".format(Pmat[i, j].imag))
tmp_re.write("\n")
tmp_im.write("\n")
tmp_re.close()
tmp_im.close()
# Need to open and close the file at every freq, otherwise
# tmp is written only at the end of the calculations, therefore,
# it is useless
return dn, Pmat
[docs]def write_occupation(ksn2f, fname="orbitals_occupation.txt"):
"""
Write the orbital occupation to the file orbitals_occupation.txt
"""
if ksn2f.shape[0] != 1:
raise ValueError("Periodic boundaries not supported")
nspin = ksn2f.shape[1]
with open(fname, "w") as fl:
fl.write("# Electronic occupation of orbitals from Fermi-Dirac statistic\n")
fl.write("# Spin")
for spin in range(nspin):
fl.write(" {}".format(spin))
fl.write("\n")
for iorb in range(ksn2f.shape[2]):
for spin in range(nspin):
fl.write(" {0:.6g}".format(ksn2f[0, spin, iorb]))
fl.write("\n")