# -*- coding: utf-8 -*-
################################################################################
#
# TPRF: Two-Particle Response Function (TPRF) Toolbox for TRIQS
#
# Copyright (C) 2018 by The Simons Foundation
# Copyright (C) 2020, S. Käser
# Copyright (C) 2023, Hugo U. R. Strand
# Authors: H. U.R. Strand, S. Käser
#
# TPRF is free software: you can redistribute it and/or modify it under the
# terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# TPRF is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# TPRF. If not, see <http://www.gnu.org/licenses/>.
#
################################################################################
import itertools
import multiprocessing
import numpy as np
# ----------------------------------------------------------------------
import triqs.utility.mpi as mpi
from triqs.gf import Gf
from triqs.gf import Idx
from triqs.gf import MeshImFreq
from triqs.gf import MeshDLRImFreq
from triqs.gf import MeshProduct
from triqs.gf import MeshBrZone
from triqs.gf import make_gf_dlr
from triqs.lattice import BrillouinZone
# ----------------------------------------------------------------------
from triqs_tprf.logo import tprf_banner
from triqs_tprf.lattice import lattice_dyson_g0_wk
from triqs_tprf.lattice import lattice_dyson_g_wk
from triqs_tprf.lattice import fourier_wk_to_wr
from triqs_tprf.lattice import fourier_wr_to_tr
from triqs_tprf.lattice import chi0_tr_from_grt_PH
from triqs_tprf.lattice import chi0_wr_from_grt_PH
from triqs_tprf.lattice import chi0_w0r_from_grt_PH
from triqs_tprf.lattice import chi_wr_from_chi_tr
from triqs_tprf.lattice import chi_w0r_from_chi_tr
from triqs_tprf.lattice import chi_wk_from_chi_wr
from triqs_tprf.lattice import dlr_on_imfreq
# ----------------------------------------------------------------------
def add_fake_bosonic_mesh(gf, beta=None):
""" Put a one value bosonic mesh as the first mesh argument of a
Green's function object.
Parameters
----------
gf : Gf,
Green's function on some arbitrary mesh. If 'beta' is not given
one mesh needs to be a 'MeshImFreq' to obtain a beta'
beta : float, optional
The inverse temperature used for the fake bosonic mesh.
Returns
-------
gf_w : Gf,
Green's function with an additional one value bosonic mesh
on its first position.
"""
mesh = gf.mesh
if isinstance(mesh, MeshProduct):
meshes = mesh.components
else:
meshes = (mesh,)
# If beta is not given access it from a 'MeshImFreq' of the 'Gf'
if not beta:
betas = [mesh.beta for mesh in meshes if hasattr(mesh, "beta")]
if len(betas) == 0:
raise ValueError(
"No 'beta' was given and the Green's function does not contain"
" a 'MeshImFreq'")
beta = betas[0]
wmesh = MeshImFreq(beta, 'Boson', 1)
mesh = (wmesh,) + meshes
mesh = MeshProduct(*mesh)
gf_w = Gf(mesh=mesh, target_shape=gf.target_shape)
gf_w.data[0,...] = gf.data
return gf_w
# ----------------------------------------------------------------------
def put_gf_on_mesh(g_in, wmesh):
assert( len(wmesh) <= len(g_in.mesh) )
g_out = Gf(mesh=wmesh, target_shape=g_in.target_shape)
for w in wmesh:
index = w.data_index + wmesh.first_index() # absolute index
g_out[w] = g_in[Idx(index)]
return g_out
# ----------------------------------------------------------------------
def strip_sigma(nw, beta, sigma_in, debug=False):
np.testing.assert_almost_equal(beta, sigma_in.mesh.beta)
wmesh = MeshImFreq(beta, 'Fermion', n_max=nw)
sigma = Gf(mesh=wmesh, target_shape=sigma_in.target_shape)
for w in wmesh:
index = w.data_index + wmesh.first_index() # absolute index
sigma[w] = sigma_in[Idx(index)]
if debug:
from triqs.plot.mpl_interface import oplot, plt
oplot(p.Sigmalatt_iw)
oplot(sigma,'x')
plt.show(); exit()
return sigma
# ----------------------------------------------------------------------
def bubble_setup(beta, mu, tb_lattice, nk, nw, sigma_w=None):
print((tprf_banner(), "\n"))
print(('beta =', beta))
print(('mu =', mu))
print(('sigma =', (not (sigma == None))))
norb = tb_lattice.NOrbitalsInUnitCell
print(('nk =', nk))
print(('nw =', nw))
print(('norb =', norb))
print()
ntau = 4 * nw
ntot = np.prod(nk) * norb**4 + np.prod(nk) * (nw + ntau) * norb**2
nbytes = ntot * np.complex128().nbytes
ngb = nbytes / 1024.**3
print(('Approx. Memory Utilization: %2.2f GB\n' % ngb))
bz = BrillouinZone(tb_lattice.bl)
bzmesh = MeshBrZone(bz, nk)
print('--> ek')
e_k = ek_tb_dispersion_on_bzmesh(tb_lattice, bzmesh, bz)
if sigma is None:
print('--> g0k')
wmesh = MeshImFreq(beta=beta, S='Fermion', n_max=nw)
g_wk = lattice_dyson_g0_wk(mu=mu, e_k=e_k, mesh=wmesh)
else:
print('--> gk')
sigma_w = strip_sigma(nw, beta, sigma)
g_wk = lattice_dyson_g_wk(mu=mu, e_k=e_k, sigma_w=sigma_w)
print('--> gr_from_gk (k->r)')
g_wr = fourier_wk_to_wr(g_wk)
del g_wk
print('--> grt_from_grw (w->tau)')
g_tr = fourier_wr_to_tr(g_wr)
del g_wr
if sigma is None:
return g_tr
else:
return g_tr, sigma_w
# ----------------------------------------------------------------------
[docs]
def imtime_bubble_chi0_wk(g_wk, nw=1, save_memory=False, verbose=True):
ncores = multiprocessing.cpu_count()
wmesh, kmesh = g_wk.mesh.components
norb = g_wk.target_shape[0]
beta = wmesh.beta
nw_g = len(wmesh)
nk = len(kmesh)
ntau = 2 * nw_g
is_dlr_mesh = type(wmesh) == MeshDLRImFreq
if is_dlr_mesh:
iw_values = np.fromiter(wmesh.values(), dtype=complex)
is_symmetrized = np.allclose(iw_values, -iw_values[::-1])
# -- Memory Approximation
ng_tr = ntau * np.prod(nk) * norb**2 # storing G(tau, r)
ng_wr = nw_g * np.prod(nk) * norb**2 # storing G(w, r)
ng_t = ntau * norb**2 # storing G(tau)
nchi_tr = ntau * np.prod(nk) * norb**4 # storing \chi(tau, r)
nchi_wr = nw * np.prod(nk) * norb**4 # storing \chi(w, r)
nchi_t = ntau * norb**4 # storing \chi(tau)
nchi_w = nw * norb**4 # storing \chi(w)
nchi_r = np.prod(nk) * norb**4 # storing \chi(r)
if nw == 1:
ntot_case_1 = ng_tr + ng_wr
ntot_case_2 = ng_tr + nchi_wr + ncores*(nchi_t + 2*ng_t)
ntot_case_3 = 4 * nchi_wr
ntot = max(ntot_case_1, ntot_case_2, ntot_case_3)
else:
ntot_case_1 = ng_tr + nchi_tr + ncores*(nchi_t + 2*ng_t)
ntot_case_2 = nchi_tr + nchi_wr + ncores*(nchi_w + nchi_t)
ntot = max(ntot_case_1, ntot_case_2)
nbytes = ntot * np.complex128().nbytes
ngb = nbytes / 1024.**3
if verbose and mpi.is_master_node():
print(tprf_banner(), "\n")
print('beta =', beta)
print('nk =', nk)
print('nw =', nw_g)
print('norb =', norb)
print()
print('Approx. Memory Utilization: %2.2f GB\n' % ngb)
if verbose: mpi.report('--> fourier_wk_to_wr')
g_wr = fourier_wk_to_wr(g_wk)
del g_wk
if verbose: mpi.report('--> fourier_wr_to_tr')
g_tr = fourier_wr_to_tr(g_wr)
del g_wr
if nw == 1:
if verbose: mpi.report('--> chi0_w0r_from_grt_PH (bubble in tau & r)')
chi0_wr = chi0_w0r_from_grt_PH(g_tr)
del g_tr
else:
if not save_memory:
if verbose: mpi.report('--> chi0_tr_from_grt_PH (bubble in tau & r)')
if is_dlr_mesh:
chi0_tr = chi0_tr_from_grt_PH(g_tr, is_symmetrized)
else:
chi0_tr = chi0_tr_from_grt_PH(g_tr)
del g_tr
if verbose: mpi.report('--> chi_wr_from_chi_tr')
chi0_wr = chi_wr_from_chi_tr(chi0_tr, nw=nw)
del chi0_tr
elif save_memory:
chi0_wr = chi0_wr_from_grt_PH(g_tr, nw=nw)
if verbose: mpi.report('--> chi_wk_from_chi_wr (r->k)')
chi0_wk = chi_wk_from_chi_wr(chi0_wr)
del chi0_wr
return chi0_wk
# ----------------------------------------------------------------------
[docs]
def chi_contraction(chi, op1, op2):
"""Contract a susceptibility with two operators
Parameters
----------
chi : Gf,
Susceptibility :math:`\chi(i\omega_n, \mathbf{k})`. The mesh attribute of
the Gf must be a MeshProduct with the components (MeshImFreq, MeshBrZone)
and its target_rank 4.
op1, op2 : np.ndarray,
Operators in matrix representation.
Returns
-------
Gf,
Susceptibility :math:`\chi(i\omega_n, \mathbf{k})`. With a target_rank of 0.
"""
if chi.target_shape[:2] != op1.shape or chi.target_shape[2:] != op2.shape:
raise ValueError('The shape of the operators %s and %s'%(op1.shape, op2.shape) +
' must fit the shape of chi %s.'%(chi.target_shape,))
chi_op1op2 = chi[0, 0, 0, 0].copy()
chi_op1op2.data[:] = np.einsum('...abcd,ab,cd->...', chi.data, op1, op2)
return chi_op1op2
# ----------------------------------------------------------------------
def chi0_w0k_tau_bubble(beta, mu, tb_lattice, nk, nw, sigma_w=None):
if sigma_w is None:
g_tr = bubble_setup(beta, mu, tb_lattice, nk, nw)
else:
g_tr, sigma_w_cut = bubble_setup(beta, mu, tb_lattice, nk, nw, sigma_w=sigma_w)
print('--> chi0_w0r_from_grt_PH (bubble in tau & r)')
chi0_wr = chi0_w0r_from_grt_PH(g_tr)
del grt
print('--> chi_wk_from_chi_wr (r->k)')
chi0_wk = chi_wk_from_chi_wr(chi0_wr)
del chi0_wr
if sigma is None:
return chi0_wk
else:
return chi0_wk, sigma_w_cut
# ----------------------------------------------------------------------
def ek_tb_dispersion_on_bzmesh(tb_lattice, bzmesh, bz):
""" Evaluate dispersion on bzmesh from tight binding model. """
n_orb = tb_lattice.NOrbitalsInUnitCell
ek = Gf(mesh=bzmesh, target_shape=[n_orb]*2)
k_vec = np.array([k.value for k in bzmesh])
k_vec_rel = get_relative_k_from_absolute(k_vec, bz.units)
ek.data[:] = tb_lattice.hopping(k_vec_rel.T).transpose(2, 0, 1)
return ek
# ----------------------------------------------------------------------
def get_relative_k_from_absolute(k_vec, units):
k_vec_rel = np.dot(np.linalg.inv(units).T, k_vec.T).T
return k_vec_rel
# ----------------------------------------------------------------------
def get_kidx_from_k_vec_relative(k_vec_rel, nk):
kidx = np.array(np.round(k_vec_rel * np.array(nk)[None, :]), dtype=np.int64)
return kidx
# ----------------------------------------------------------------------
def get_k_components_from_k_vec(k_vec, nk):
dim = 3
shape_4 = [dim] + list(nk)
shape_2 = [dim] + [np.prod(nk)]
k_vec = k_vec.swapaxes(0, -1)
k_vec = k_vec.reshape(shape_4)
# -- cut out values for each axis
k_out = []
for axis in range(dim):
cut = [0]*dim
cut[axis] = slice(None)
cut = [axis] + cut
k_out.append(k_vec[tuple(cut)])
return tuple(k_out)
# ----------------------------------------------------------------------
def cluster_mesh_fourier_interpolation(k, chiwr):
assert( len(k.shape) == 2 )
assert( k.shape[1] == 3 )
rmesh = chiwr.mesh.components[1]
r = np.array([r.value for r in rmesh])
exp_fact = np.exp(np.sum(-1.j * k[:, None, :] * r[None, :, :], axis=-1))
chi00wk_data = np.einsum('wrabcd,kr->wkabcd', chiwr.data, exp_fact)
return chi00wk_data
# ----------------------------------------------------------------------
def get_abs_k_chi_interpolator(values, bzmesh, bz, extend_bz=[0]):
k_mat = bz.units
k_vec = np.array([k.value for k in bzmesh])
# -- Extend with points beyond the first bz
k_vec_ext = []
values_ext = []
for k_shift in itertools.product(extend_bz, repeat=3):
k_shift = np.dot(k_mat.T, k_shift)
k_vec_ext.append( k_vec + k_shift[None, :] )
values_ext.append(values)
k_vec = np.vstack(k_vec_ext)
values = np.hstack(values_ext)
from scipy.interpolate import LinearNDInterpolator
interp = LinearNDInterpolator(k_vec, values, fill_value=float('nan'))
return interp
# ----------------------------------------------------------------------
def get_rel_k_chi_interpolator(values, bzmesh, bz, nk,
extend_boundary=True, interpolator='regular'):
k_mat = bz.units
k_vec = np.array([k.value for k in bzmesh])
k_vec_rel = get_relative_k_from_absolute(k_vec, bz.units)
k_idx = get_kidx_from_k_vec_relative(k_vec_rel, nk)
kx, ky, kz = get_k_components_from_k_vec(k_vec_rel, nk)
if extend_boundary:
values, k_vec_rel, (kx, ky, kz) = \
extend_data_on_boundary(values, nk)
else:
values = values.reshape(nk)
# -- select interpolator type
if interpolator == 'regular':
from scipy.interpolate import RegularGridInterpolator
interp = RegularGridInterpolator(
(kx, ky, kz), values, fill_value=float('nan'), bounds_error=False)
elif interpolator == 'nearest':
from scipy.interpolate import NearestNDInterpolator
interp = NearestNDInterpolator(k_vec_rel, values.flatten())
elif interpolator == 'linear':
from scipy.interpolate import LinearNDInterpolator
interp = LinearNDInterpolator(k_vec_rel, values.flatten(), fill_value=float('nan'))
else:
raise NotImplementedError
return interp
# ----------------------------------------------------------------------
def extend_data_on_boundary(values, nk):
# -- extended cube
nk = np.array(nk)
# -- this adds points on the boundary
# Add extra points in the positive index directions
nk_ext = nk + 1
coords = [ np.arange(0, n) for n in nk_ext ]
# This adds two extra points in both positive and negative directions
#nk_ext = nk + 4
#coords = [ np.arange(-2, n+2) for n in nk ]
Coords = np.meshgrid(*coords, indexing='ij')
Coords_mod = [ np.mod(x, n) for x, n in zip(Coords, nk) ]
values_ext = values.reshape(nk)[tuple([ X.flatten() for X in Coords_mod])]
values_ext = values_ext.reshape(nk_ext)
# -- compute kidx_ext
k_idx_ext = np.array([ X.flatten() for X in Coords ]).T
k_vec_rel_ext = np.array(k_idx_ext, dtype=np.float) / nk[None, :]
kxe, kye, kze = get_k_components_from_k_vec(k_vec_rel_ext, nk_ext)
return values_ext, k_vec_rel_ext, (kxe, kye, kze)
# ----------------------------------------------------------------------
[docs]
def k_space_path(paths, num=100, bz=None, relative_coordinates=False):
""" Generate an array of k-vectors along a path defined by a list of pairs of k-vectors
Parameters
----------
paths : list of pairs of three-vectors of floats
List of pairs of k-vectors in reciprocal units to create a path in-between.
num : int, default=100
Number of k-vectors along each segment of the overall path
bz : brillouin_zone, optional
When a Brillouin Zone is passed, calculate distance in absolute units
relative_coordinates : bool, optional
Return k-vectors in reciprocal units. (Default `False`)
Returns
-------
kvecs: numpy.ndarray [shape=(len(paths)*num,3)]
Two-dimensional numpy array containing the path vectors (in reciprocal units) as rows
dist: numpy.ndarray [shape=(kvecs.shape[0])]
One-dimensional numpy array containing, for each element in kvecs,
the distance travelled along the path. Useful for plotting.
If bz is provided, calculate the distance in absolute units.
The distances for the relevant k-points in paths can be obtained with dist[num::num].
ticks : numpy.ndarray [shape=(len(paths)+1)]
Array with tick points, i.e. distances at the interfaces of path segments.
Only returned if `return_ticks` is `True`.
"""
# Put this back once TPRF and TRIQS agree on default output units
#print("WARNING: triqs_tprf.lattice_utils.k_space_path has moved to triqs.lattice.utils.k_space_path")
from triqs.lattice.utils import k_space_path
return k_space_path(
paths, num=num, bz=bz, relative_coordinates=relative_coordinates)
# ----------------------------------------------------------------------
def gf_tensor_to_matrix(g):
N = g.target_shape[0] * g.target_shape[1]
M = g.target_shape[2] * g.target_shape[3]
g_mat = Gf(mesh=g.mesh, target_shape=[N, M])
g_mat.data[:] = np.reshape(
g.data, (g.data.shape[0], N, M))
return g_mat
# ----------------------------------------------------------------------
def gf_matrix_to_tensor(g_mat, target_shape):
g = Gf(mesh=g_mat.mesh, target_shape=target_shape)
shape = [g_mat.data.shape[0]] + list(target_shape)
g.data[:] = np.reshape(g_mat.data, shape)
return g
# ----------------------------------------------------------------------
def pade_analytical_continuation_wk(
g_wk, fmesh, n_points=32, freq_offset=0.05):
""" Perform Pade analytical continuation of a lattice Green's function
Parameters
----------
g_wk : triqs.gf.Gf with mesh (MeshImFreq or MeshDLRImFreq, MeshBrZone)
Lattice Green's function in Matsubara frequency and momentum space :math:`G(i\omega_n, \mathbf{k})`.
fmesh : MeshReFreq
Real frequency mesh to perform the analytical continuation to.
n_points : int
Number of Matsubara frequencies to use in the Pade fit.
freq_offset : float
Distance from the real axis used in the fit.
Returns
-------
g_fk : triqs.gf.Gf with mesh (MeshReFreq, MeshBrZone)
Analytically continued real-frequency lattice Green's function
"""
wmesh = g_wk.mesh[0]
kmesh = g_wk.mesh[1]
g_fk = Gf(mesh=MeshProduct(fmesh, kmesh), target_shape=g_wk.target_shape)
for k in kmesh:
g_f = g_fk[:, k]
# Make_gf_dlr is broken on gf views (issue #913),
# so make a copy of the data instead
if type(wmesh) == MeshDLRImFreq:
g_w = Gf(mesh=wmesh, target_shape=g_wk.target_shape)
g_w.data[:] = g_wk[:, k].data[:]
else:
g_w = g_wk[:, k]
if len(g_wk.target_shape) == 4:
g_w = gf_tensor_to_matrix(g_w)
g_f = gf_tensor_to_matrix(g_f)
if type(g_w.mesh) == MeshDLRImFreq:
g_c = make_gf_dlr(g_w)
small_mesh = MeshImFreq(g_w.mesh.beta, g_w.mesh.statistic, n_points)
g_w = dlr_on_imfreq(g_c, small_mesh)
g_f.set_from_pade(g_w, n_points=n_points, freq_offset=freq_offset)
if len(g_wk.target_shape) == 4:
g_f = gf_matrix_to_tensor(g_f, g_wk.target_shape)
g_fk[:, k] = g_f
return g_fk