import logging
import os
import random
import string
import warnings
import importlib.metadata as imd
import h5py
import numpy as np
from numba import jit
from pyscf.df import addons
from pyscf.pbc import gto, df, tools
from pyscf.pbc.lib import kpts as libkpts
from pyscf.lib import logger
from pyscf.pbc.lib.kpts_helper import unique
from pyscf.pbc.df.rsdf_builder import _RSGDFBuilder
from pyscf.pbc.df.gdf_builder import _CCGDFBuilder
import scipy.linalg as LA
from . import kpt_utils
# Linear dep threshold for J2C metric eigenvalues
J2C_LIN_DEP_THRESH = 1e-10
# AO block of X_k[ik] for ERI rotation. For X2C the spinor X is
# (2·nao, 2·nao) and we slice the upper-left nao×nao block — valid only
# because the CLI guards (pyscf_init / init_seet) restrict X2C
# orthogonalization to Löwdin variants, whose X inherits the
# block-diagonal spin structure of the spinor S. MO and natural-orbital
# rotations are refused upstream for X2C precisely because their X
# would not be block-diagonal. Rectangular X (rank-deficient case)
# is rejected here as it can't fit the (NQ, nao, nao) chunk shape.
def _x_block_for_eri(X_k, ik, nao):
Xi = np.asarray(X_k[ik])
if Xi.shape == (nao, nao):
return Xi
if Xi.shape == (2 * nao, 2 * nao):
return Xi[:nao, :nao]
raise ValueError(
f"_x_block_for_eri: X_k[{ik}] has shape {Xi.shape}; expected "
f"({nao}, {nao}) or ({2*nao}, {2*nao}). The three-center integral "
"storage path requires a square orthogonalizer; rank-deficient "
"transforms (lowdin_per_k with linear dependencies, or canonical "
"orthogonalization dropping modes) cannot be written into "
"df_hf_int / df_int."
)
# Lpq_ortho[Q] = X_i @ Lpq[Q] @ X_j.conj().T — the "X Z X†" convention used by
# the per-k ortho primitives; pairs with the dm stored by common_utils.orthogonalize
# so the Hartree contraction Σ_{ab} L[Q,a,b] * dm[b,a] inside mbpt's HF kernel
# reduces to its AO value.
def _rotate_Lpq(Lpq, X_i, X_j):
return np.einsum("Ia,Qab,Jb->QIJ", X_i, Lpq, X_j.conj(), optimize=True)
[docs]
def compute_kG(k, Gv, wrap_around, mesh, cell):
if abs(k).sum() > 1e-9:
kG = k + Gv
else:
kG = Gv
equal2boundary = np.zeros(Gv.shape[0], dtype=bool)
if wrap_around and abs(k).sum() > 1e-9:
# Here we 'wrap around' the high frequency k+G vectors into their lower
# frequency counterparts. Important if you want the gamma point and k-point
# answers to agree
b = cell.reciprocal_vectors()
box_edge = np.einsum('i,ij->ij', np.asarray(mesh)//2+0.5, b)
assert (all(np.linalg.solve(box_edge.T, k).round(9).astype(int)==0))
reduced_coords = np.linalg.solve(box_edge.T, kG.T).T.round(9)
on_edge = reduced_coords.astype(int)
if cell.dimension >= 1:
equal2boundary |= reduced_coords[:,0] == 1
equal2boundary |= reduced_coords[:,0] ==-1
kG[on_edge[:,0]== 1] -= 2 * box_edge[0]
kG[on_edge[:,0]==-1] += 2 * box_edge[0]
if cell.dimension >= 2:
equal2boundary |= reduced_coords[:,1] == 1
equal2boundary |= reduced_coords[:,1] ==-1
kG[on_edge[:,1]== 1] -= 2 * box_edge[1]
kG[on_edge[:,1]==-1] += 2 * box_edge[1]
if cell.dimension == 3:
equal2boundary |= reduced_coords[:,2] == 1
equal2boundary |= reduced_coords[:,2] ==-1
kG[on_edge[:,2]== 1] -= 2 * box_edge[2]
kG[on_edge[:,2]==-1] += 2 * box_edge[2]
return kG, equal2boundary
[docs]
def get_coarsegrained_coulG(lattice_kmesh, cell, k=np.zeros(3), exx=False, mf=None, mesh=None, Gv=None,
wrap_around=True, omega=None, **kwargs):
'''Calculate the coarse-grained Coulomb kernel for all G-vectors, handling G=0 and exchange.
This routine overrides get_coulG to perform interaction coarse-graining.
'''
exxdiv = exx
if isinstance(exx, str):
exxdiv = exx
elif exx and mf is not None:
exxdiv = mf.exxdiv
if mesh is None:
mesh = cell.mesh
if 'gs' in kwargs:
logging.warning('cell.gs is deprecated. It is replaced by cell.mesh,'
'the number of PWs (=2*gs+1) along each direction.')
mesh = [2*n+1 for n in kwargs['gs']]
if Gv is None:
Gv = cell.get_Gv(mesh)
absG2 = []
kG_0, equal2boundary_0 = compute_kG(k, Gv, wrap_around, mesh, cell)
absG2_0 = np.einsum('gi,gi->g', kG_0, kG_0)
for kp in lattice_kmesh:
kG, equal2boundary = compute_kG(k + kp, Gv, wrap_around, mesh, cell)
absG2.append(np.einsum('gi,gi->g', kG, kG))
absG2 = np.array(absG2)
if getattr(mf, 'kpts', None) is not None:
kpts = mf.kpts
else:
kpts = k.reshape(1,3)
Nk = len(kpts)
fullkpts = np.zeros([kpts.shape[0]*lattice_kmesh.shape[0],3])
for ikk, kk in enumerate(kpts):
for ikkp, kkp in enumerate(lattice_kmesh):
fullkpts[ikk*lattice_kmesh.shape[0] + ikkp] = kk + kkp
Nkk = Nk * lattice_kmesh.shape[0]
if exxdiv == 'vcut_sph': # PRB 77 193110
raise NotImplementedError
elif exxdiv == 'vcut_ws': # PRB 87, 165122
raise NotImplementedError
else:
# Ewald probe charge method to get the leading term of the finite size
# error in exchange integrals
G0_idx = np.where(absG2_0==0)[0]
if cell.dimension != 2 or cell.low_dim_ft_type == 'inf_vacuum':
coulG = np.zeros(absG2.shape)
for ikp in range(absG2.shape[0]):
#with np.errstate(divide='ignore'):
coulG[ikp] = 4*np.pi/absG2[ikp]
if np.sum(np.dot(k,k)) < 1e-9:
print(ikp, 4*np.pi/absG2[ikp])
if not np.isfinite(coulG[ikp, G0_idx]) :
coulG[ikp, G0_idx] = 0
coulG = np.sum(coulG, axis=0) / absG2.shape[0]
if np.sum(np.dot(k,k)) < 1e-9:
print("Sum:", coulG)
#coulG[G0_idx] = 0
elif cell.dimension == 2:
raise NotImplementedError
elif cell.dimension == 1:
logging.warning(cell, 'No method for PBC dimension 1, dim-type %s.'
' cell.low_dim_ft_type="inf_vacuum" should be set.',
cell.low_dim_ft_type)
raise NotImplementedError
# The divergent part of periodic summation of (ii|ii) integrals in
# Coulomb integrals were cancelled out by electron-nucleus
# interaction. The periodic part of (ii|ii) in exchange cannot be
# cancelled out by Coulomb integrals. Its leading term is calculated
# using Ewald probe charge (the function madelung below)
if cell.dimension > 0 and exxdiv == 'ewald' and len(G0_idx) > 0:
coulG[G0_idx] += Nkk*cell.vol*tools.madelung(cell, fullkpts)
coulG[equal2boundary_0] = 0
# Scale the coulG kernel for attenuated Coulomb integrals.
# * omega is used by RangeSeparatedJKBuilder which requires ewald probe charge
# being evaluated with regular Coulomb interaction (1/r12).
# * cell.omega, which affects the ewald probe charge, is often set by
# DFT-RSH functionals to build long-range HF-exchange for erf(omega*r12)/r12
if omega is not None:
if omega > 0:
# long range part
coulG *= np.exp(-.25/omega**2 * absG2)
elif omega < 0:
# short range part
coulG *= (1 - np.exp(-.25/omega**2 * absG2))
elif cell.omega > 0:
coulG *= np.exp(-.25/cell.omega**2 * absG2)
elif cell.omega < 0:
raise NotImplementedError
return coulG
[docs]
def weighted_coulG_ewald(mydf, kpt, exx, mesh, omega=None):
if omega is None:
return df.aft.weighted_coulG(mydf, kpt, "ewald", mesh)
return df.aft.weighted_coulG(mydf, kpt, "ewald", mesh, omega)
# a = lattice vectors / (2*pi)
[docs]
@jit(nopython=True)
def kpair_reduced_lists(kptis, kptjs, kptij_idx, kmesh, a):
nkpts = kmesh.shape[0]
print("nkpts = ", nkpts)
kptis = np.asarray(kptis)
kptjs = np.asarray(kptjs)
if kptis.shape[0] != kptjs.shape[0]:
raise ValueError("Error: Dimension of kptis and kptjs doesn't match.")
num_kpair = kptis.shape[0]
conj_list = np.arange(num_kpair, dtype=np.int64)
trans_list = np.arange(num_kpair, dtype=np.int64)
seen = np.zeros(num_kpair, dtype=np.int64)
for i in range(num_kpair):
if seen[i] != 0:
continue
k1 = kptis[i]
k2 = kptjs[i]
k1_idx, k2_idx = kptij_idx[i]
l = 0
# conj_list: (k1, k2) = (-kk1, -kk2)
for kk1_idx in range(k1_idx,nkpts):
kk1 = kmesh[kk1_idx]
kjdif = np.dot(a, k1 + kk1)
kjdif_int = np.rint(kjdif)
maskj = np.sum(np.abs(kjdif - kjdif_int)) < 1e-6
if maskj == True:
for kk2_idx in range(0,kk1_idx+1):
kk2 = kmesh[kk2_idx]
kidif = np.dot(a, k2 + kk2)
#kidif = np.einsum('wx,x->w', a, k2 + kk2)
kidif_int = np.rint(kidif)
maski = np.sum(np.abs(kidif - kidif_int)) < 1e-6
if maski == True:
j = int(kk1_idx * (kk1_idx + 1)/2 + kk2_idx)
conj_list[j] = i
seen[j] = 1
l = 1
break
if l == 1:
break
if l == 1:
seen[i] = 1
continue
# trans_list: (k1, k2) = (-kk2, -kk1)
for kk1_idx in range(k1_idx,nkpts):
kk1 = kmesh[kk1_idx]
kjdif = np.dot(a, k2 + kk1)
kjdif_int = np.rint(kjdif)
maskj = np.sum(np.abs(kjdif - kjdif_int)) < 1e-6
if maskj == True:
for kk2_idx in range(0,kk1_idx+1):
kk2 = kmesh[kk2_idx]
kidif = np.dot(a, k1 + kk2)
kidif_int = np.rint(kidif)
maski = np.sum(np.abs(kidif - kidif_int)) < 1e-6
if maski == True:
j = int(kk1_idx * (kk1_idx + 1)/2 + kk2_idx)
trans_list[j] = i
seen[j] = 1
l = 1
break
if l == 1:
break
seen[i] = 1
return conj_list, trans_list
[docs]
def make_gdf_kptij_lst_jk(kstruct):
'''
[WIP]
Build GDF k-point-pair list for get_jk
All combinations::
k_ibz != k_bz
k_bz == k_bz
'''
from pyscf.pbc.lib.kpts_helper import member
kptij_lst = [(kstruct.kpts[i], kstruct.kpts[i]) for i in range(kstruct.nkpts)]
kptij_idx_lst = [(i, i) for i in range(kstruct.nkpts)]
for i in range(kstruct.nkpts_ibz):
ki = kstruct.kpts_ibz[i]
where = member(ki, kstruct.kpts)
for j in range(kstruct.nkpts):
kj = kstruct.kpts[j]
if j not in where:
kptij_lst.extend([(ki,kj)])
kptij_idx_lst.extend([(i,j)])
kptij_lst = np.asarray(kptij_lst)
kptij_idx_lst = np.asarray(kptij_idx_lst)
# dummy lists for kij_conj and kij_trans
# kptij_conjlist =
return kptij_lst, kptij_idx_lst
[docs]
def integrals_grid(mycell, kmesh):
a_lattice = mycell.lattice_vectors() / (2*np.pi)
kptij_lst = [(ki, kmesh[j]) for i, ki in enumerate(kmesh) for j in range(i+1)]
kptij_idx = [(i, j) for i in range(kmesh.shape[0]) for j in range(i+1)]
kptij_lst = np.asarray(kptij_lst)
kptij_idx = np.asarray(kptij_idx)
kptis = kptij_lst[:,0]
kptjs = kptij_lst[:,1]
kij_conj, kij_trans = kpair_reduced_lists(kptis, kptjs, kptij_idx, kmesh, a_lattice)
kpair_irre_list = np.argwhere(kij_conj == kij_trans)[:,0]
num_kpair_stored = len(kpair_irre_list)
return kptij_idx, kij_conj, kij_trans, kpair_irre_list, num_kpair_stored, kptis, kptjs
[docs]
def compute_integrals(args, mycell, mydf, kmesh, nao, X_k=None, basename = "df_int", cderi_name="cderi.h5", keep=True, keep_after=False, cderi_name2="cderi_ewald.h5"):
kptij_idx, kij_conj, kij_trans, kpair_irre_list, num_kpair_stored, kptis, kptjs = integrals_grid(mycell, kmesh)
mydf.kpts = kmesh
filename = basename + "/meta.h5"
os.system("sync") # This is needed to syncronize the NFS between nodes
if os.path.exists(basename):
unsafe_rm = "rm " + basename + "/VQ*"
os.system(unsafe_rm)
unsafe_rm = "rm " + basename + "/meta*"
os.system(unsafe_rm)
# This is needed to ensure that files are removed both on the computational and head nodes:
os.system("sync")
os.system("mkdir -p " + basename) # Here "-p" is important and is needed if the if condition is trigerred
if os.path.exists(cderi_name) and keep :
mydf._cderi = cderi_name
else:
mydf._cderi_to_save = cderi_name
mydf.build()
correction_df = None
apply_correction = False
if os.path.exists(cderi_name2) :
import copy
apply_correction = True
correction_df = copy.copy(mydf)
correction_df._cderi = cderi_name2
auxcell = addons.make_auxmol(mycell, mydf.auxbasis)
NQ = auxcell.nao_nr()
print("NQ = ", NQ)
# compute partitioning
def compute_partitioning(tot_size, num_kpair_stored):
# We have a guess for each fitted density upper bound of 150M
memory = 700 if args.memory is None else args.memory
ubound = memory * 1024 * 1024
if tot_size > ubound :
mult = tot_size // ubound
chunks = num_kpair_stored // (mult+1)
if chunks == 0 :
print("\n\n\n Chunk size is bigger than upper memory bound per chunk you have \n\n\n")
chunks = 1
return chunks
return num_kpair_stored
single_rho_size = nao**2 * NQ * 16
full_rho_size = (num_kpair_stored * single_rho_size)
chunk_size = compute_partitioning(full_rho_size, num_kpair_stored)
print("The chunk size: ", chunk_size, " k-point pair")
# open file to write integrals in
if os.path.exists(filename) :
os.remove(filename)
data = h5py.File(filename, "w")
# Loop over k-point pair
# processed densities count
cnt = 0
# densities buffer
buffer = np.zeros((chunk_size, NQ, nao, nao), dtype=complex)
Lpq_mo = np.zeros((NQ, nao, nao), dtype=complex)
chunk_indices = []
# Rotate three-center integrals into the AO->ortho basis when
# args.orth != "none", so on-disk V matches S/F/T/dm in input.h5.
rotate = (X_k is not None and len(X_k) == kmesh.shape[0]
and args.orth != "none")
for i in kpair_irre_list:
k1 = kptis[i]
k2 = kptjs[i]
k1_idx, k2_idx = kptij_idx[i]
if rotate:
X_i = _x_block_for_eri(X_k, k1_idx, nao)
X_j = _x_block_for_eri(X_k, k2_idx, nao)
# auxiliary basis index
s1 = 0
for XXX in mydf.sr_loop((k1,k2), max_memory=4000, compact=False):
LpqR = XXX[0]
LpqI = XXX[1]
Lpq = (LpqR + LpqI*1j).reshape(LpqR.shape[0], nao, nao)
if rotate:
Lpq = _rotate_Lpq(Lpq, X_i, X_j)
buffer[cnt% chunk_size, s1:s1+Lpq.shape[0], :, :] = Lpq[0:Lpq.shape[0],:,:]
# s1 = NQ at maximum.
s1 += Lpq.shape[0]
if apply_correction and np.allclose(k1, k2) :
s1 = 0
for XXX in correction_df.sr_loop((k1,k1), max_memory=4000, compact=False):
LpqR = XXX[0]
LpqI = XXX[1]
Lpq = (LpqR + LpqI*1j).reshape(LpqR.shape[0], nao, nao)
if rotate:
Lpq = _rotate_Lpq(Lpq, X_i, X_j)
buffer[cnt% chunk_size, s1:s1+Lpq.shape[0], :, :] = Lpq[0:Lpq.shape[0],:,:]
# s1 = NQ at maximum.
s1 += Lpq.shape[0]
cnt += 1
# if reach chunk size: (cnt-chunk_size) equals to chunk id.
if cnt % chunk_size == 0:
chunk_name = basename + "/VQ_{}.h5".format(cnt - chunk_size)
if os.path.exists(chunk_name) :
os.remove(chunk_name)
VQ = h5py.File(chunk_name, "w")
VQ["{}".format(cnt - chunk_size)] = buffer.view(np.float64)
VQ.close()
chunk_indices.append(cnt - chunk_size)
buffer[:] = 0.0
# Deal the rest
if cnt % chunk_size != 0:
last_chunk = (num_kpair_stored // chunk_size) * chunk_size
chunk_name = basename + "/VQ_{}.h5".format(last_chunk)
if os.path.exists(chunk_name) :
os.remove(chunk_name)
VQ = h5py.File(chunk_name, "w")
VQ["{}".format(last_chunk)] = buffer.view(np.float64)
chunk_indices.append(last_chunk)
VQ.close()
buffer[:] = 0.0
data["chunk_size"] = chunk_size
data["chunk_indices"] = np.array(chunk_indices)
data.attrs["__green_version__"] = imd.version("green_mbtools")
data.close()
if not keep_after:
os.remove(cderi_name)
if apply_correction:
os.remove(cderi_name2)
os.system("sync")
print("Integrals have been computed and stored into {}".format(filename))
return kij_conj, kij_trans, kpair_irre_list, kptij_idx, num_kpair_stored
[docs]
def weighted_coulG_ewald_2nd(mydf, kpt, exx, mesh):
# this is a dirty hack
# PySCF needs to have a full k-grid to properly compute the madelung constant
# but we want to compute only the contribution for ki == kj to speedup this calculations
# TODO Double check where we define full_k_mesh?
if not hasattr(mydf.cell, 'full_k_mesh'):
raise RuntimeError("Using wrong DF object")
oldkpts = mydf.kpts
mydf.kpts = mydf.cell.full_k_mesh
coulG = df.aft.weighted_coulG(mydf, kpt, 'ewald', mesh)
mydf.kpts = oldkpts
return coulG
[docs]
def compute_ewald_correction(args, maindf, kmesh, nao, filename = "df_ewald.h5", X_k=None):
# global full_k_mesh
data = h5py.File(filename, "w")
EW = data.create_group("EW")
EW_bar = data.create_group("EW_bar")
# keep original method for computing Coulomb kernel
import importlib.util as iu
new_pyscf = iu.find_spec('pyscf.pbc.df.gdf_builder') is not None
if new_pyscf :
import pyscf.pbc.df.gdf_builder as gdf
weighted_coulG_old = gdf._CCGDFBuilder.weighted_coulG
else:
from pyscf.pbc import df as gdf
weighted_coulG_old = gdf.GDF.weighted_coulG
# density-fitting w/o ewald correction for fine grid
df2 = df.GDF(maindf.cell)
if hasattr(df2, "_prefer_ccdf"):
df2._prefer_ccdf = True # Disable RS-GDF switch for new pyscf versions
if maindf.auxbasis is not None:
df2.auxbasis = maindf.auxbasis
# Coulomb kernel mesh
df2.mesh = maindf.mesh
cderi_file_2 = ''.join(random.choices(string.ascii_uppercase + string.digits, k=10)) + ".h5"
df2._cderi_to_save = cderi_file_2
df2._cderi = cderi_file_2
df2.kpts = kmesh
df2.build()
# densities buffer
auxcell = addons.make_auxmol(maindf.cell, maindf.auxbasis)
NQ = auxcell.nao_nr()
buffer1 = np.zeros((NQ, nao, nao), dtype=np.complex128)
buffer2 = np.zeros((NQ, nao, nao), dtype=np.complex128)
Lpq_mo = np.zeros((NQ, nao, nao), dtype=np.complex128)
if new_pyscf :
gdf._CCGDFBuilder.weighted_coulG = weighted_coulG_ewald_2nd
else:
gdf.GDF.weighted_coulG = weighted_coulG_ewald_2nd
# df.aft.weighted_coulG = weighted_coulG_ewald_2nd
df1 = df.GDF(maindf.cell)
df1.cell.full_k_mesh = maindf.kpts
if hasattr(df1, "_prefer_ccdf"):
df1._prefer_ccdf = True # Disable RS-GDF switch for new pyscf versions
if maindf.auxbasis is not None:
df1.auxbasis = maindf.auxbasis
# Use Ewald for divergence treatment
df1.exxdiv = 'ewald'
# Coulomb kernel mesh
df1.mesh = maindf.mesh
cderi_file_1 = ''.join(random.choices(string.ascii_uppercase + string.digits, k=10)) + ".h5"
df1._cderi_to_save = cderi_file_1
df1._cderi = cderi_file_1
df1.kpts = kmesh
df1.build()
_, nk2, nk3 = args.nk
# Rotate Ewald-correction buffers to AO->ortho basis when
# args.orth != "none", to match the V integrals written by
# compute_integrals (avoids mixing V_ortho with EW_AO downstream).
rotate = (X_k is not None and len(X_k) == kmesh.shape[0]
and args.orth != "none")
# We know that G=0 contribution diverge only when q = 0
# so we loop over (k1,k1) pairs
for i, ki in enumerate(kmesh):
# Change the way to compute Coulomb kernel to include G=0 correction
# df.GDF.weighted_coulG = weighted_coulG_ewald_2nd
if new_pyscf :
gdf._CCGDFBuilder.weighted_coulG = weighted_coulG_ewald_2nd
else:
df.GDF.weighted_coulG = weighted_coulG_ewald_2nd
s1 = 0
# Compute three-point integrals with G=0 contribution included with Ewald correction
for XXX in df1.sr_loop((ki,ki), max_memory=4000, compact=False):
LpqR = XXX[0]
LpqI = XXX[1]
Lpq = (LpqR + LpqI*1.j).reshape(LpqR.shape[0], nao, nao)
for G in range(Lpq.shape[0]):
Lpq_mo[G] = Lpq[G]
buffer1[s1:s1+Lpq.shape[0], :, :] = Lpq_mo[0:Lpq.shape[0],:,:]
s1 += Lpq.shape[0]
# Restore the way to compute Coulomb kernel
# df.aft.weighted_coulG = weighted_coulG_old
if new_pyscf:
gdf._CCGDFBuilder.weighted_coulG = weighted_coulG_old
else:
df.GDF.weighted_coulG = weighted_coulG_old
s1 = 0
# Compute three-point integral without G=0 contribution included with Ewald correction
# and subtract it from the computed buffer to keep pure Ewald correction only
for XXX in df2.sr_loop((ki,ki), max_memory=4000, compact=False):
LpqR = XXX[0]
LpqI = XXX[1]
Lpq = (LpqR + LpqI*1.j).reshape(LpqR.shape[0], nao, nao)
for G in range(Lpq.shape[0]):
Lpq_mo[G] = Lpq[G]
buffer2[s1:s1+Lpq.shape[0], :, :] = Lpq_mo[0:Lpq.shape[0],:,:]
s1 += Lpq.shape[0]
if rotate:
X_i = _x_block_for_eri(X_k, i, nao)
ew_arr = _rotate_Lpq(buffer1 - buffer2, X_i, X_i)
ew_bar_arr = _rotate_Lpq(buffer2, X_i, X_i)
else:
ew_arr = buffer1 - buffer2
ew_bar_arr = buffer2
EW["{}".format(i)] = ew_arr.view(np.float64)
EW_bar["{}".format(i)] = ew_bar_arr.view(np.float64)
buffer1[:] = 0.0
buffer2[:] = 0.0
data.close()
# cleanup
# df.aft.weighted_coulG = weighted_coulG_old
if new_pyscf:
gdf._CCGDFBuilder.weighted_coulG = weighted_coulG_old
else:
gdf.GDF.weighted_coulG = weighted_coulG_old
os.remove(cderi_file_1)
os.remove(cderi_file_2)
print("Ewald correction has been computed and stored into {}".format(filename))
[docs]
class GreenGDF(df.GDF):
def __init__(self, cell, kpts=np.zeros((1,3))):
super().__init__(cell, kpts)
self.space_symm = True
self.tr_symm = True
self.x2c = 0
self.use_j2c_eig_decomposition = True
def _make_j3c(self, cell=None, auxcell=None, kptij_lst=None, cderi_file=None):
"""Customized ERI building to build and store j2c at the same time
Parameters
----------
cell : pyscf.pbc.gto.Cell, optional
periodic cell, by default None
auxcell : pyscf.pbc.gto.Cell, optional
auxiliary basis cell, by default None
kptij_lst : list, optional
list of k-point pairs, by default None
cderi_file : str, optional
output file to store ERI, by default None
"""
if cell is None: cell = self.cell
if auxcell is None: auxcell = self.auxcell
if cderi_file is None: cderi_file = self._cderi_to_save
# Logger
log = logger.new_logger(self)
# Remove duplicated k-points. Duplicated kpts may lead to a buffer
# located in incore.wrap_int3c larger than necessary. Integral code
# only fills necessary part of the buffer, leaving some space in the
# buffer unfilled.
if self.kpts_band is None:
kpts_union = self.kpts
else:
kpts_union = unique(np.vstack([self.kpts, self.kpts_band]))[0]
if self._prefer_ccdf or cell.omega > 0:
# For long-range integrals _CCGDFBuilder is the only option
dfbuilder = _CCGDFBuilder(cell, auxcell, kpts_union)
dfbuilder.eta = self.eta
else:
dfbuilder = _RSGDFBuilder(cell, auxcell, kpts_union)
# Keep configurable to support both legacy-reference compatibility and systematic eigenvalue workflow.
dfbuilder.j2c_eig_always = bool(getattr(self, 'use_j2c_eig_decomposition', True))
dfbuilder.mesh = self.mesh
dfbuilder.linear_dep_threshold = self.linear_dep_threshold
j_only = self._j_only or len(kpts_union) == 1
dfbuilder.make_j3c(cderi_file, j_only=j_only, dataname=self._dataname,
kptij_lst=kptij_lst)
# Build q = k1-k2 and reduce to a unique wrapped q-mesh, then construct a q-structure and use its functionalities
use_space_symm = bool(getattr(self, 'space_symm', True))
use_tr_symm = bool(getattr(self, 'tr_symm', True))
use_x2c = int(getattr(self, 'x2c', 0))
qstruct = kpt_utils.build_q_struct(cell, self.kpts, space_symm=(use_space_symm and use_x2c < 2), tr_symm=use_tr_symm)
uniq_qpts = qstruct.kpts
scaled_uniq_kpts = qstruct.kpts_scaled.round(5)
log.debug('Num uniq kpts %d', len(uniq_qpts))
log.debug2('scaled unique kpts %s', scaled_uniq_kpts)
# Group k-points into conjugate pairs and build a mapping from the original q-mesh to the unique q-mesh.
# qstruct.ibz2bz are the full-BZ indices of irreducible q-points.
# Store j2c strictly for q_IBZ representatives in this order.
q_ibz2bz = np.asarray(qstruct.ibz2bz, dtype=int)
j2c_uniq_kpts = uniq_qpts[q_ibz2bz]
feri = h5py.File(cderi_file, 'a')
all_j2c = list(dfbuilder.get_2c2e(j2c_uniq_kpts))
for i, j2c_idx in enumerate(q_ibz2bz):
j2c = all_j2c[i]
if j2c.dtype == np.complex128:
feri[f'j2c/{j2c_idx}'] = j2c
else:
feri[f'j2c/{j2c_idx}'] = j2c + 0.j
j2c = None
# Meta data
feri['j2c/qmesh'] = uniq_qpts
feri['j2c/scaled_qmesh'] = scaled_uniq_kpts
feri['j2c/j2c_uniq_kpts'] = j2c_uniq_kpts
# Compatibility metadata: pair each stored q_IBZ representative with itself.
feri['j2c/q_ibz2bz'] = q_ibz2bz
feri['j2c/q_bz2ibz'] = np.asarray(qstruct.bz2ibz, dtype=int)
feri['j2c'].attrs['j2c_decomposition'] = 'eigenvalue' if dfbuilder.j2c_eig_always else 'cholesky'
feri.close()
[docs]
def cholesky_decomposed_metric(j2c_k, cell, inv=False):
"""Calculate the Cholesky decomposition of the j2c metric for a given k-point, or its inverse if requested.
Parameters
----------
j2c_k : _type_
j2c metric for a specific k-point.
cell : _type_
information about the unit cell, used for error handling and potential adjustments based on dimensionality.
inv : bool, optional
Whether to return the inverse of the Cholesky decomposition, by default False
Returns
-------
j2c_sqrt_k or j2c_sqrt_inv_k : ndarray
Cholesky decomposition of the j2c metric (lower triangular matrix L such that j2c_k = L @ L†) or
its inverse (L⁻¹), depending on the value of `inv`.
j2c_negative : ndarray or None
If the j2c metric has negative eigenvalues (which can occur in 2D systems with certain Fourier
transform conventions), this will contain the corresponding eigenvectors. Otherwise, it will be None.
"""
# Cholesky: j2c = LL† (L lower triangular). PySCF computes B = L⁻¹ @ eri3c
# (by solving L†x = eri3c, i.e., x = L⁻¹ @ eri3c), so P0_tilde lives in the
# L-basis. Obar = L_bz⁻¹ @ mat_ao @ L_irre correctly maps P0_tilde between
# k-points. Using upper Cholesky (lower=False) gives the wrong convention.
nQ = j2c_k.shape[0]
try:
j2c_negative = None
j2c_sqrt_k = LA.cholesky(j2c_k, lower=True)
if inv:
# this is actually the inverse and not just the sqrt
j2c_sqrt_k = LA.solve_triangular( # = L^{-1}
j2c_sqrt_k, np.eye(nQ, dtype=np.complex128), lower=True
)
except LA.LinAlgError:
j2c_sqrt_k, j2c_negative = eigenvalue_decomposed_metric(j2c_k, cell, inv)
return j2c_sqrt_k, j2c_negative
[docs]
def eigenvalue_decomposed_metric(j2c_k, cell, inv=False):
j2c_negative = None
cond_num = np.linalg.cond(j2c_k)
if cond_num > 1e12:
warnings.warn(f"Condition number of j2c_k ({cond_num}) is too large. Consider using a smaller aux. basis.")
eigs, vecs = LA.eigh(j2c_k)
assert np.all(eigs > -J2C_LIN_DEP_THRESH), "j2c metric has non-positive eigenvalues"
eig_mask = eigs > np.max(eigs) * J2C_LIN_DEP_THRESH
eigs_trim = eigs[eig_mask]
if inv:
j2c_sqrt_k = vecs[:, eig_mask].conj().T / np.sqrt(eigs_trim).reshape(-1, 1)
else:
j2c_sqrt_k = vecs[:, eig_mask] * np.sqrt(eigs_trim).reshape(1, -1)
# negative eigenvalues can occur in 2D systems with certain Fourier transform conventions,
# but we can still use the corresponding eigenvectors to define a "negative" subspace for the J2C metric
if cell.dimension == 2 and cell.low_dim_ft_type != 'inf_vacuum':
idx = np.where(eigs < -J2C_LIN_DEP_THRESH)[0]
if len(idx) > 0:
j2c_negative = (vecs[:, idx] / np.sqrt(-eigs[idx])).conj().T
return j2c_sqrt_k, j2c_negative