+#!/usr/bin/env python3
+################################################################################
+#
+# TRIQS: a Toolbox for Research in Interacting Quantum Systems
+#
+# Copyright (C) 2016-2018, N. Wentzell
+# Copyright (C) 2018-2019, Simons Foundation
+# author: N. Wentzell
+#
+# TRIQS 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.
+#
+# TRIQS 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
+# TRIQS. If not, see <http://www.gnu.org/licenses/>.
+#
+################################################################################
+
+"""
+Reads DMFT_ouput observables such as real-frequency Sigma and a Wannier90
+TB Hamiltonian to compute spectral properties. It runs in two modes,
+either calculating the bandstructure or Fermi slice.
+
+Written by Sophie Beck, 2021-2022
+
+TODO:
+- extend to multi impurity systems
+- make proper use of rot_mat from DFT_Tools (atm it assumed that wannier_hr and Sigma are written in the same basis)
+"""
+
+import matplotlib.pyplot as plt
+from matplotlib.ticker import MaxNLocator
+from matplotlib.colors import Normalize
+from matplotlib import cm
+from scipy.optimize import brentq
+from scipy.interpolate import interp1d
+from scipy.signal import argrelextrema
+import numpy as np
+import itertools
+import skimage.measure
+from warnings import warn
+
+from h5 import HDFArchive
+from triqs.gf import BlockGf, MeshReFreq, Gf
+from triqs.lattice.utils import TB_from_wannier90, k_space_path
+from triqs_dft_tools.sumk_dft import SumkDFT
+
+
+def lambda_matrix_w90_t2g(add_lambda):
+
+ lambda_x, lambda_y, lambda_z = add_lambda
+
+ lambda_matrix = np.zeros((6, 6), dtype=complex)
+ lambda_matrix[0, 1] = +1j*lambda_z/2.0
+ lambda_matrix[0, 5] = -1j*lambda_x/2.0
+ lambda_matrix[1, 5] = -lambda_y/2.0
+ lambda_matrix[2, 3] = +1j*lambda_x/2.0
+ lambda_matrix[2, 4] = +lambda_y/2.0
+ lambda_matrix[3, 4] = -1j*lambda_z/2.0
+ lambda_matrix += np.transpose(np.conjugate(lambda_matrix))
+
+ return lambda_matrix
+
+
+def change_basis(n_orb, orbital_order_to, orbital_order_from):
+
+ change_of_basis = np.eye(n_orb)
+ for ct, orb in enumerate(orbital_order_to):
+ orb_idx = orbital_order_from.index(orb)
+ change_of_basis[orb_idx, :] = np.roll(np.eye(n_orb, 1), ct)[:, 0]
+
+ return change_of_basis
+
+
+def print_matrix(matrix, n_orb, text):
+
+ print('{}:'.format(text))
+
+ if np.any(matrix.imag > 1e-4):
+ fmt = '{:16.4f}' * n_orb
+ else:
+ fmt = '{:8.4f}' * n_orb
+ matrix = matrix.real
+
+ for row in matrix:
+ print((' '*4 + fmt).format(*row))
+
+
+def _sigma_from_dmft(n_orb, orbital_order, with_sigma, spin, orbital_order_dmft=None, **specs):
+
+ if orbital_order_dmft is None:
+ orbital_order_dmft = orbital_order
+
+ if with_sigma == 'calc':
+ print('Setting Sigma from {}'.format(specs['dmft_path']))
+
+ sigma_imp_list = []
+ dc_imp_list = []
+ with HDFArchive(specs['dmft_path'], 'r') as ar:
+ for icrsh in range(ar['dft_input']['n_inequiv_shells']):
+ try:
+ sigma = ar['DMFT_results'][specs['it']][f'Sigma_freq_{icrsh}']
+ assert isinstance(sigma.mesh, MeshReFreq), 'Imported Greens function must be real frequency'
+ except(KeyError, AssertionError):
+ try:
+ sigma = ar['DMFT_results'][specs['it']][f'Sigma_maxent_{icrsh}']
+ except KeyError:
+ try:
+ sigma = ar['DMFT_results'][specs['it']][f'Sigma_Refreq_{icrsh}']
+ except KeyError:
+ raise KeyError('Provide either "Sigma_freq_0" in real frequency, "Sigma_Refreq_0" or "Sigma_maxent_0".')
+ sigma_imp_list.append(sigma)
+
+ for ish in range(ar['dft_input']['n_corr_shells']):
+ dc_imp_list.append(ar['DMFT_results'][specs['it']]['DC_pot'][ish])
+
+ mu_dmft = ar['DMFT_results'][specs['it']]['chemical_potential_post']
+
+ sum_k = SumkDFT(specs['dmft_path'], mesh=sigma.mesh)
+ sum_k.block_structure = ar['DMFT_input/block_structure']
+ sum_k.deg_shells = ar['DMFT_input/deg_shells']
+ sum_k.set_mu(mu_dmft)
+ # set Sigma and DC into sum_k
+ sum_k.dc_imp = dc_imp_list
+ sum_k.set_Sigma(sigma_imp_list)
+
+ # use add_dc function to rotate to sumk block structure and subtract the DC
+ sigma_sumk = sum_k.add_dc()
+
+ assert np.allclose(sum_k.proj_mat[0], sum_k.proj_mat[-1]), 'upfolding works only when proj_mat is the same for all kpoints (wannier mode)'
+
+ # now upfold with proj_mat to band basis, this only works for the
+ # case where proj_mat is equal for all k points (wannier mode)
+ sigma = Gf(mesh=sigma.mesh, target_shape=[n_orb, n_orb])
+ for ish in range(ar['dft_input']['n_corr_shells']):
+ sigma += sum_k.upfold(ik=0, ish=ish,
+ bname=spin, gf_to_upfold=sigma_sumk[ish][spin],
+ gf_inp=sigma)
+
+ # already subtracted
+ dc = 0.0
+
+ else:
+ print('Setting Sigma from memory')
+
+ sigma = with_sigma[spin]
+ dc = specs['dc'][0][spin][0, 0]
+ mu_dmft = specs['mu_dmft']
+
+ SOC = (spin == 'ud')
+ w_mesh_dmft = np.linspace(sigma.mesh.w_min, sigma.mesh.w_max, len(sigma.mesh))
+ assert sigma.target_shape[0] == n_orb, f'Number of Wannier orbitals: {n_orb} and self-energy target_shape {sigma.target_shape} does not match'
+
+ sigma_mat = sigma.data.real - np.eye(n_orb) * dc + 1j * sigma.data.imag
+
+ # rotate sigma from orbital_order_dmft to orbital_order
+ change_of_basis = change_basis(n_orb, orbital_order, orbital_order_dmft)
+ sigma_mat = np.einsum('ij, kjl -> kil', np.linalg.inv(change_of_basis), np.einsum('ijk, kl -> ijl', sigma_mat, change_of_basis))
+
+ # set up mesh
+ if 'w_mesh' in specs:
+ freq_dict = specs['w_mesh']
+ w_mesh = np.linspace(*freq_dict['window'], freq_dict['n_w'])
+ freq_dict.update({'w_mesh': w_mesh})
+ else:
+ w_mesh = w_mesh_dmft
+ freq_dict = {'w_mesh': w_mesh_dmft, 'n_w': len(sigma.mesh), 'window': [sigma.mesh.w_min, sigma.mesh.w_max]}
+
+ sigma_interpolated = np.zeros((n_orb, n_orb, freq_dict['n_w']), dtype=complex)
+
+ # interpolate sigma
+ def interpolate_sigma(w_mesh, w_mesh_dmft, orb1, orb2): return np.interp(w_mesh, w_mesh_dmft, sigma_mat[:, orb1, orb2])
+
+ for ct1, ct2 in itertools.product(range(n_orb), range(n_orb)):
+ sigma_interpolated[ct1, ct2] = interpolate_sigma(w_mesh, w_mesh_dmft, ct1, ct2)
+
+ return sigma_interpolated, mu_dmft, freq_dict
+
+
+def sigma_FL(n_orb, orbital_order, Sigma_0, Sigma_Z, freq_dict, eta=0.0, mu_dmft=None):
+
+ print('Setting Re[Sigma] with Fermi liquid approximation')
+
+ if np.any(Sigma_0) and mu_dmft == None:
+ raise ValueError('Sigma_0 does not preserve electron count. Please provide "mu_dmft".')
+ elif not np.any(Sigma_0) and mu_dmft == None:
+ mu_dmft = 0.
+
+ eta = eta * 1j
+
+ # set up mesh
+ w_mesh = np.linspace(*freq_dict['window'], freq_dict['n_w'])
+ freq_dict.update({'w_mesh': w_mesh})
+
+ # setting up sigma
+ sigma_array = np.zeros((n_orb, n_orb, freq_dict['n_w']), dtype=complex)
+ def approximate_sigma(orb): return (1-1/Sigma_Z[orb]) * freq_dict['w_mesh'] + Sigma_0[orb] - mu_dmft
+ for ct, orb in enumerate(orbital_order):
+ sigma_array[ct, ct] = approximate_sigma(ct) + 1j * eta
+
+ return sigma_array, freq_dict
+
+
+def _calc_alatt(n_orb, mu, eta, e_mat, sigma, qp_bands=False, e_vecs=None,
+ proj_nuk=None, trace=True, **freq_dict):
+ '''
+ calculate slice of lattice spectral function for given TB dispersion / e_mat and self-energy
+
+ Parameters
+ ----------
+ n_orb : int
+ number of Wannier orbitals
+ proj_nuk : optinal, 2D numpy array (n_orb, n_k)
+ projections to be applied on A(k,w) in band basis. Only works when band_basis=True
+
+ Returns
+ -------
+ alatt_k_w : numpy array, either (n_k, n_w) or if trace=False (n_k, n_w, n_orb)
+ Lattice Green's function on specified k-path / mesh
+
+ '''
+
+ # adjust to system size
+ def upscale(quantity, n_orb): return quantity * np.identity(n_orb)
+ mu = upscale(mu, n_orb)
+ eta = upscale(eta, n_orb)
+ if isinstance(e_vecs, np.ndarray):
+ sigma_rot = np.zeros(sigma.shape, dtype=complex)
+
+ w_vec = np.array([upscale(freq_dict['w_mesh'][w], n_orb) for w in range(freq_dict['n_w'])])
+ n_k = e_mat.shape[2]
+
+ if not qp_bands:
+ if trace:
+ alatt_k_w = np.zeros((n_k, freq_dict['n_w']))
+ else:
+ alatt_k_w = np.zeros((n_k, freq_dict['n_w'], n_orb))
+
+ def invert_and_trace(w, eta, mu, e_mat, sigma, trace, proj=None):
+ # inversion is automatically vectorized over first axis of 3D array (omega first index now)
+ Glatt = np.linalg.inv(w + eta[None, ...] + mu[None, ...] - e_mat[None, ...] - sigma.transpose(2, 0, 1))
+ A_w_nu = -1.0/(2.0 * np.pi)* np.diagonal(Glatt - Glatt.transpose(0,2,1).conj(), axis1=1, axis2=2).imag
+ if isinstance(proj, np.ndarray):
+ A_w_nu = A_w_nu * proj[None, :]
+ if trace:
+ return np.sum(A_w_nu, axis=1)
+ else:
+ return A_w_nu
+
+ for ik in range(n_k):
+ # if evecs are given transform sigma into band basis
+ if isinstance(e_vecs, np.ndarray):
+ sigma_rot = np.einsum('ij,jkw->ikw', e_vecs[:, :, ik].conjugate().transpose(), np.einsum('ijw,jk->ikw', sigma, e_vecs[:, :, ik]))
+ if isinstance(proj_nuk, np.ndarray):
+ alatt_k_w[ik, :] = invert_and_trace(w_vec, eta, mu, e_mat[:, :, ik], sigma_rot, trace, proj_nuk[:, ik])
+ else:
+ alatt_k_w[ik, :] = invert_and_trace(w_vec, eta, mu, e_mat[:, :, ik], sigma_rot, trace)
+ else:
+ alatt_k_w[ik, :] = invert_and_trace(w_vec, eta, mu, e_mat[:, :, ik], sigma, trace)
+
+ else:
+ alatt_k_w = np.zeros((n_k, n_orb))
+ kslice = np.zeros((freq_dict['n_w'], n_orb))
+ def kslice_interp(orb): return interp1d(freq_dict['w_mesh'], kslice[:, orb])
+
+ for ik in range(n_k):
+ for iw, w in enumerate(freq_dict['w_mesh']):
+ np.fill_diagonal(sigma[:, :, iw], np.diag(sigma[:, :, iw]).real)
+ #sigma[:,:,iw] = sigma[:,:,iw].real
+ kslice[iw], _ = np.linalg.eigh(upscale(w, n_orb) + eta + mu - e_mat[:, :, ik] - sigma[:, :, iw])
+
+ for orb in range(n_orb):
+ w_min, w_max = freq_dict['window']
+ try:
+ x0 = brentq(kslice_interp(orb), w_min, w_max)
+ w_bin = int((x0 - w_min) / ((w_max - w_min) / freq_dict['n_w']))
+ alatt_k_w[ik, orb] = freq_dict['w_mesh'][w_bin]
+ except ValueError:
+ pass
+
+ return alatt_k_w
+
+
+def _calc_kslice(n_orb, mu, eta, e_mat, sigma, qp_bands, e_vecs=None, proj_nuk=None, **freq_dict):
+ '''
+ calculate lattice spectral function for given TB dispersion / e_mat and self-energy
+
+ Parameters
+ ----------
+ n_orb : int
+ number of Wannier orbitals
+ proj_nuk : optinal, 2D numpy array (n_orb, n_k)
+ projections to be applied on A(k,w) in band basis. Only works when band_basis=True
+
+ Returns
+ -------
+ alatt_k_w : numpy array, either (n_k, n_w) or if trace=False (n_k, n_w, n_orb)
+ Lattice Green's function on specified k-path / mesh
+
+ '''
+
+ # adjust to system size
+ def upscale(quantity, n_orb): return quantity * np.identity(n_orb)
+ mu = upscale(mu, n_orb)
+ eta = upscale(eta, n_orb)
+
+ iw0 = np.where(np.sign(freq_dict['w_mesh']) == True)[0][0]-1
+ print_matrix(sigma[:, :, iw0], n_orb, 'Zero-frequency Sigma')
+
+ if isinstance(e_vecs, np.ndarray):
+ sigma_rot = np.zeros(sigma.shape, dtype=complex)
+
+ n_kx, n_ky = e_mat.shape[2:4]
+
+ if not qp_bands:
+ alatt_k_w = np.zeros((n_kx, n_ky))
+
+ def invert_and_trace(w, eta, mu, e_mat, sigma, proj=None):
+ # inversion is automatically vectorized over first axis of 3D array (omega first index now)
+ Glatt = np.linalg.inv(w + eta + mu - e_mat - sigma)
+ A_nu = -1.0/(2.0 * np.pi)* np.diagonal(Glatt - Glatt.transpose().conj()).imag
+ if isinstance(proj, np.ndarray):
+ A_nu = A_nu * proj
+ return np.sum(A_nu)
+
+ for ikx, iky in itertools.product(range(n_kx), range(n_ky)):
+ if isinstance(e_vecs, np.ndarray):
+ sigma_rot = np.einsum('ij,jk->ik',
+ e_vecs[:, :, ikx, iky].conjugate().transpose(),
+ np.einsum('ij,jk->ik', sigma[:, :, iw0], e_vecs[:, :, ikx, iky]))
+ else:
+ sigma_rot = sigma[:, :, iw0]
+
+ if isinstance(proj_nuk, np.ndarray):
+ alatt_k_w[ikx, iky] = invert_and_trace(upscale(freq_dict['w_mesh'][iw0], n_orb), eta, mu,
+ e_mat[:, :, ikx, iky], sigma_rot, proj_nuk[:, ikx, iky])
+ else:
+ alatt_k_w[ikx, iky] = invert_and_trace(upscale(freq_dict['w_mesh'][iw0], n_orb), eta, mu, e_mat[:, :, ikx, iky], sigma_rot)
+
+ else:
+ assert n_kx == n_ky, 'Not implemented for N_kx != N_ky'
+
+ def search_for_extrema(data):
+ # return None for no extrema, [] if ends of interval are the only extrema,
+ # list of indices if local extrema are present
+ answer = np.all(data > 0) or np.all(data < 0)
+ if answer:
+ return
+ else:
+ roots = []
+ roots.append(list(argrelextrema(data, np.greater)[0]))
+ roots.append(list(argrelextrema(data, np.less)[0]))
+ roots = sorted([item for sublist in roots for item in sublist])
+ return roots
+
+ alatt_k_w = np.zeros((n_kx, n_ky, n_orb))
+ # go through grid horizontally, then vertically
+ for it in range(2):
+ kslice = np.zeros((n_kx, n_ky, n_orb))
+
+ for ik1 in range(n_kx):
+ e_temp = e_mat[:, :, :, ik1] if it == 0 else e_mat[:, :, ik1, :]
+ for ik2 in range(n_kx):
+ e_val, _ = np.linalg.eigh(eta + mu - e_temp[:, :, ik2] - sigma[:, :, iw0])
+ k1, k2 = [ik2, ik1] if it == 0 else [ik1, ik2]
+ kslice[k1, k2] = e_val
+
+ for orb in range(n_orb):
+ temp_kslice = kslice[:,ik1,orb] if it == 0 else kslice[ik1,:,orb]
+ roots = search_for_extrema(temp_kslice)
+ # iterate through sections between extrema
+ if roots is not None:
+ idx_1 = 0
+ for root_ct in range(len(roots) + 1):
+ idx_2 = roots[root_ct] if root_ct < len(roots) else n_kx
+ root_section = temp_kslice[idx_1:idx_2+1]
+ try:
+ x0 = brentq(interp1d(np.linspace(idx_1, idx_2, len(root_section)), root_section), idx_1, idx_2)
+ k1, k2 = [int(np.floor(x0)), ik1] if it == 0 else [ik1, int(np.floor(x0))]
+ alatt_k_w[k1, k2, orb] += 1
+ except(ValueError):
+ pass
+ idx_1 = idx_2
+
+ alatt_k_w[np.where(alatt_k_w > 1)] = 1
+
+ return alatt_k_w
+
+
+[docs]def get_tb_bands(e_mat, proj_on_orb=[None], **specs):
+
'''
+
calculate eigenvalues and eigenvectors for given list of e_mat on kmesh
+
+
Parameters
+
----------
+
e_mat : numpy array of shape (n_orb, n_orb, nk) or (n_orb, n_orb, nk, nk)
+
+
Returns
+
-------
+
e_val : numpy array of shape (n_orb, n_orb, nk) or (n_orb, n_orb, nk, nk)
+
eigenvalues as matrix
+
e_vec : numpy array of shape (n_orb, n_orb, nk) or (n_orb, n_orb, nk, nk)
+
eigenvectors as matrix
+
'''
+
+
e_val = np.zeros((e_mat.shape), dtype=complex)
+
e_vec = np.zeros((e_mat.shape), dtype=complex)
+
n_orb = e_mat.shape[0]
+
+
for ikx in range(e_mat.shape[2]):
+
# if we have a 2d kmesh e_mat is dim=4
+
if len(e_mat.shape) == 4:
+
for iky in range(e_mat.shape[3]):
+
e_val[range(n_orb), range(n_orb), ikx, iky], e_vec[:, :, ikx, iky] = np.linalg.eigh(e_mat[:, :, ikx, iky])
+
else:
+
e_val[range(n_orb), range(n_orb), ikx], e_vec[:, :, ikx] = np.linalg.eigh(e_mat[:, :, ikx])
+
+
if proj_on_orb[0] is not None:
+
print(f'calculating projection on orbitals {proj_on_orb}')
+
total_proj = np.zeros(np.shape(e_vec[0]))
+
for band in range(n_orb):
+
for orb in proj_on_orb:
+
total_proj[band] += np.real(e_vec[orb, band] * e_vec[orb, band].conjugate())
+
else:
+
total_proj = None
+
+
return e_val, e_vec, total_proj
+
+
+def get_tb_kslice(tb, mu_tb, **specs):
+
+ w90_paths = list(map(lambda section: (np.array(specs[section[0]]), np.array(specs[section[1]])), specs['bands_path']))
+ upper_left = w90_paths[0][0]
+ lower_right = w90_paths[1][-1]
+ origin = w90_paths[0][1]
+ assert np.allclose(origin, w90_paths[1][0]), '"bands_path" coordinates for origin of Fermi surface needs to be consistent'
+ if 'kz' in specs and specs['kz'] != 0.:
+ assert 'Z' in specs, 'Please provide Z point coordinate in tb_data_dict as input coordinate'
+ Z = np.array(specs['Z'])
+ kz = specs['kz']
+ else:
+ kz = 0.
+ Z = np.zeros((3))
+
+ # calculate FS at the mu_tb value
+ FS_kx_ky, band_char = get_kx_ky_FS(lower_right, upper_left, origin, Z, tb, N_kxy=specs['n_k'], kz=kz, fermi=mu_tb)
+
+ return FS_kx_ky, band_char
+
+
+def _fract_ind_to_val(x, ind):
+ ind[ind == len(x)-1] = len(x)-1-1e-6
+ int_ind = [int(indi) for indi in ind]
+ int_ind_p1 = [int(indi)+1 for indi in ind]
+ return x[int_ind] + (x[int_ind_p1] - x[int_ind])*(np.array(ind)-np.array(int_ind))
+
+
+def get_kx_ky_FS(lower_right, upper_left, origin, Z, tb, select=None, N_kxy=10, kz=0.0, fermi=0.0):
+
+ # create mesh
+ kx = np.linspace(0, 0.5, N_kxy)
+ ky = np.linspace(0, 0.5, N_kxy)
+
+ if select is None:
+ select = np.array(range(tb.n_orbitals))
+
+ # go in horizontal arrays from bottom to top
+ E_FS = np.zeros((tb.n_orbitals, N_kxy, N_kxy))
+ for kyi in range(N_kxy):
+ path_FS = [((upper_left - origin)/(N_kxy-1)*kyi+kz*Z+origin, origin+(lower_right-origin)+(upper_left-origin)/(N_kxy-1)*kyi+kz*Z)]
+ k_vec, dst, tks = k_space_path(path_FS, num=N_kxy)
+ E_FS[:, :, kyi] = tb.dispersion(k_vec).transpose() - fermi
+
+ contours = {}
+ FS_kx_ky = {}
+ FS_kx_ky_prim = {}
+ band_char = {}
+ # contour for each sheet
+ for sheet in range(tb.n_orbitals):
+ contours[sheet] = skimage.measure.find_contours(E_FS[sheet, :, :], 0.0)
+
+ sheet_ct = 0
+ for sheet in contours.keys():
+ for sec_per_sheet in range(len(contours[sheet])):
+ # once on 2D cubic mesh
+ FS_kx_ky[sheet_ct] = np.vstack([_fract_ind_to_val(kx, contours[sheet][sec_per_sheet][:, 0]),
+ _fract_ind_to_val(ky, contours[sheet][sec_per_sheet][:, 1]),
+ kz*np.ones(len(contours[sheet][sec_per_sheet][:, 0]))]).T.reshape(-1, 3)
+ # repeat on actual mesh for computing the weights
+ ks_skimage = contours[sheet][sec_per_sheet]/(N_kxy-1)
+ FS_kx_ky_prim[sheet_ct] = (+ np.einsum('i,j->ij', ks_skimage[:, 0], lower_right)
+ + np.einsum('i,j->ij', ks_skimage[:, 1], upper_left)
+ + np.einsum('i,j->ij', kz * np.ones(ks_skimage.shape[0]), Z))
+ band_char[sheet_ct] = {}
+ # compute the weight aka band character
+ for ct_k, k_on_sheet in enumerate(FS_kx_ky_prim[sheet_ct]):
+ E_mat = tb.fourier(k_on_sheet)
+ e_val, e_vec = np.linalg.eigh(E_mat[select[:, np.newaxis], select])
+ orb_on_FS = np.argmin(np.abs(e_val))
+
+ band_char[sheet_ct][ct_k] = [np.round(np.real(e_vec[orb, orb_on_FS]*np.conjugate(e_vec[orb, orb_on_FS])), 4) for orb in range(len(select))]
+ sheet_ct += 1
+
+ return FS_kx_ky, band_char
+
+
+def _setup_plot_bands(ax, special_k, k_points_labels, freq_dict):
+
+ ax.axhline(y=0, c='gray', ls='--', lw=0.8, zorder=0)
+ ax.set_ylabel(r'$\epsilon - \mu$ (eV)')
+# ax.set_ylim(*freq_dict['window'])
+ for ik in special_k:
+ ax.axvline(x=ik, linewidth=0.7, color='k', zorder=0.5)
+ ax.set_xticks(special_k)
+ ax.set_xlim(special_k[0], special_k[-1])
+ k_points_labels = [r'$\Gamma$' if k == 'G' else k for k in k_points_labels]
+ ax.set_xticklabels(k_points_labels)
+
+
+def setup_plot_kslice(ax):
+
+ ax.set_aspect(1)
+ # ax.set_xlim(0,1)
+ # ax.set_ylim(0,1)
+ ax.xaxis.set_major_locator(MaxNLocator(integer=True))
+ ax.yaxis.set_major_locator(MaxNLocator(integer=True))
+ ax.set_xlabel(r'$k_x\pi/a$')
+ ax.set_ylabel(r'$k_y\pi/b$')
+
+
+def plot_bands(fig, ax, alatt_k_w, tb_data, freq_dict, n_orb, tb=True, alatt=False, qp_bands=False, **plot_dict):
+
+ assert tb_data['special_k'] is not None, 'a regular k point mesh has been used, please call plot_dos'
+
+ proj_on_orb = tb_data['proj_on_orb']
+ total_proj = tb_data['proj_nuk']
+
+ if alatt:
+ if alatt_k_w is None:
+ raise ValueError('A(k,w) unknown. Specify "with_sigma = True"')
+ if qp_bands:
+ for orb in range(n_orb):
+ ax.scatter(tb_data['k_mesh'], alatt_k_w[:, orb].T, c=np.array([eval('cm.'+plot_dict['colorscheme_qpbands'])(1.0)]), zorder=2., s=1.)
+ else:
+ kw_x, kw_y = np.meshgrid(tb_data['k_mesh'], freq_dict['w_mesh'])
+
+ vmax = plot_dict['vmax'] if 'vmax' in plot_dict else np.max(alatt_k_w)
+ vmin = plot_dict['vmin'] if 'vmin' in plot_dict else 0.0
+
+ graph = ax.pcolormesh(kw_x, kw_y, alatt_k_w.T, cmap=plot_dict['colorscheme_alatt'],
+ norm=Normalize(vmin=vmin, vmax=vmax), shading='gouraud')
+
+ if 'colorbar' not in plot_dict or plot_dict['colorbar']:
+ colorbar = plt.colorbar(graph)
+ colorbar.set_label(r'$A(k, \omega)$')
+
+ if tb:
+ # if projection is requested, _get_tb_bands() ran already
+ if proj_on_orb[0] is not None:
+ eps_nuk = tb_data['e_mat']
+ evec_nuk = tb_data['e_vecs']
+ else:
+ eps_nuk, evec_nuk, _ = get_tb_bands(**tb_data)
+ for band in range(n_orb):
+ if not proj_on_orb[0] is not None:
+ if isinstance(plot_dict['colorscheme_bands'], str):
+ color = eval('cm.'+plot_dict['colorscheme_bands'])(1.0)
+ else:
+ color = plot_dict['colorscheme_bands']
+ ax.plot(tb_data['k_mesh'], eps_nuk[band, band].real - tb_data['mu_tb'], c=color, label=r'tight-binding', zorder=1., lw=1)
+ else:
+ color = eval('cm.'+plot_dict['colorscheme_bands'])(total_proj[band])
+ ax.scatter(tb_data['k_mesh'], eps_nuk[band, band].real - tb_data['mu_tb'], c=color, s=1, label=r'tight-binding', zorder=1.)
+
+ _setup_plot_bands(ax, tb_data['special_k'], tb_data['k_points_labels'], freq_dict)
+
+
+def plot_dos(fig, ax, alatt_k_w, tb_data, freq_dict, tb=False, alatt=True, label=None, color=None):
+
+ assert tb == False, 'plotting TB DOS is not supported yet.'
+
+ assert len(alatt_k_w.shape) == 2, 'input Akw should only have a k and omega index'
+
+ if not label:
+ label = ''
+
+ if not color:
+ ax.plot(freq_dict['w_mesh'], np.sum(alatt_k_w, axis=0)/alatt_k_w.shape[0], label=label)
+ else:
+ ax.plot(freq_dict['w_mesh'], np.sum(alatt_k_w, axis=0) /
+ alatt_k_w.shape[0], label=label, color=color)
+
+ ax.axvline(x=0, c='gray', ls='--', zorder=0)
+ ax.set_xlabel(r'$\epsilon - \mu$ (eV)')
+ ax.set_ylabel(r'A($\omega$)')
+
+ ax.set_xlim(*freq_dict['window'])
+
+ return
+
+def plot_kslice(fig, ax, alatt_k_w, tb_data, freq_dict, n_orb, tb_dict, tb=True, alatt=False, quarter=0, **plot_dict):
+
+ proj_on_orb = tb_data['proj_on_orb']
+ if quarter:
+ assert isinstance(quarter, int) or all(isinstance(x, int) for x in quarter), 'quarter should be'\
+ f'an integer or list of integers, but is {type(quarter)}.'
+
+ if isinstance(quarter, int):
+ quarter = [quarter]
+
+ sign = [1, -1]
+ quarters = np.array([sign, sign])
+ four_quarters = list(itertools.product(*quarters))
+ used_quarters = [four_quarters[x] for x in quarter]
+
+ vmax = plot_dict['vmax'] if 'vmax' in plot_dict else np.max(alatt_k_w)
+ vmin = plot_dict['vmin'] if 'vmin' in plot_dict else 0.0
+ assert vmax > vmin, 'vmax needs to be larger than vmin'
+
+ if alatt:
+ if alatt_k_w is None:
+ raise ValueError('A(k,w) unknown. Specify "with_sigma = True"')
+ n_kx, n_ky = tb_data['e_mat'].shape[2:4]
+ kx, ky = np.meshgrid(range(n_kx), range(n_ky))
+ draw_colorbar = True
+ for (qx, qy) in used_quarters:
+ if len(alatt_k_w.shape) > 2:
+ for orb in range(n_orb):
+ ax.contour(qx * kx/(n_kx-1), qy * ky/(n_ky-1), alatt_k_w[:, :, orb].T,
+ colors=np.array([eval('cm.'+plot_dict['colorscheme_qpbands'])(0.7)]), levels=1, zorder=2)
+ else:
+ graph = ax.pcolormesh(qx * kx/(n_kx-1), qy * ky/(n_ky-1), alatt_k_w.T,
+ cmap=plot_dict['colorscheme_kslice'],
+ norm=Normalize(vmin=vmin, vmax=vmax),
+ shading='gouraud')
+
+ if draw_colorbar and ('colorbar' not in plot_dict or plot_dict['colorbar']):
+ colorbar = plt.colorbar(graph)
+ colorbar.set_label(r'$A(k, 0$)')
+ draw_colorbar = False
+
+ if tb:
+ FS_kx_ky, band_char = get_tb_kslice(tb_data['tb'], tb_data['mu_tb'], **tb_dict)
+ for sheet in FS_kx_ky.keys():
+ for k_on_sheet in range(FS_kx_ky[sheet].shape[0]):
+ if not proj_on_orb[0] is not None:
+ if isinstance(plot_dict['colorscheme_bands'], str):
+ color = eval('cm.'+plot_dict['colorscheme_bands'])(1.0)
+ else:
+ color = plot_dict['colorscheme_bands']
+ else:
+ total_proj = 0
+ for orb in proj_on_orb:
+ total_proj += band_char[sheet][k_on_sheet][orb]
+ color = eval('cm.'+plot_dict['colorscheme_bands'])(total_proj)
+ for (qx, qy) in used_quarters:
+ ax.plot(2*qx * FS_kx_ky[sheet][k_on_sheet:k_on_sheet+2, 0], 2*qy * FS_kx_ky[sheet][k_on_sheet:k_on_sheet+2, 1], '-',
+ solid_capstyle='round', c=color, zorder=1., label=plot_dict['label'] if 'label' in plot_dict else '')
+
+ setup_plot_kslice(ax)
+
+ return ax
+
+
+[docs]def get_dmft_bands(n_orb, mu_tb, w90_path=None, w90_seed=None, TB_obj=None, add_spin=False, add_lambda=None, add_local=None,
+
with_sigma=None, fermi_slice=False, qp_bands=False, orbital_order_to=None,
+
add_mu_tb=False, band_basis=False, proj_on_orb=None, trace=True, eta=0.0,
+
mu_shift=0.0, proj_nuk=None, **specs):
+
'''
+
Extract tight-binding from given w90 seed_hr.dat and seed.wout files or alternatively given TB_obj, and then extract from
+
given solid_dmft calculation the self-energy and construct the spectral function A(k,w) on
+
given k-path.
+
+
Parameters
+
----------
+
n_orb : int
+
Number of Wannier orbitals in seed_hr.dat
+
mu_tb : float
+
Chemical potential of tight-binding calculation
+
w90_path : string
+
Path to w90 files
+
w90_seed : string
+
Seed of wannier90 calculation, i.e. seed_hr.dat and seed.wout
+
TB_obj : TB object
+
Tight-binding object from TB_from_wannier90
+
add_spin : bool, default=False
+
Extend w90 Hamiltonian by spin indices
+
add_lambda : float, default=None
+
Add SOC term with strength add_lambda (works only for t2g shells)
+
add_local : numpy array, default=None
+
Add local term of dimension (n_orb x n_orb)
+
with_sigma : str, or BlockGf, default=None
+
Add self-energy to spectral function? Can be either directly take
+
a triqs BlockGf object or can be either 'calc' or 'model'
+
'calc' reads results from h5 archive (solid_dmft)
+
in case 'calc' or 'model' are specified a extra kwargs dict has
+
to be given sigma_dict containing information about the self-energy
+
add_mu_tb : bool, default=False
+
Add the TB specified chemical potential to the lattice Green function
+
set to True if DMFT calculation was performed with DFT fermi subtracted.
+
proj_on_orb : int or list of int, default=None
+
orbital projections to be made for the spectral function and TB bands
+
the integer refer to the orbitals read
+
trace : bool, default=True
+
Return trace over orbitals for spectral function. For special
+
post-processing purposes this can be set to False giving the returned
+
alatt_k_w an extra dimension n_orb
+
eta : float, default=0.0
+
Broadening of spectral function, finitie shift on imaginary axis
+
if with_sigma=None it has to be provided !=0.0
+
mu_shift : float, default=0.0
+
Manual extra shift when calculating the spectral function
+
proj_nuk : numpy array, default [None]
+
Extra projections to be applied to the final spectral function
+
per orbital and k-point. Has to match shape of final lattice Green
+
function. Will be applied together with proj_on_orb if specified.
+
+
Returns
+
-------
+
tb_data : dict
+
tight binding dict containing the kpoint mesh, dispersion / emat, and eigenvectors
+
+
alatt_k_w : numpy array (float) of dim n_k x n_w ( x n_orb if trace=False)
+
lattice spectral function data on the kpoint mesh defined in tb_data and frequency
+
mesh defined in freq_dict
+
+
freq_dict : dict
+
frequency mesh information on which alatt_k_w is evaluated
+
'''
+
+
# set default ordering
+
if 'orbital_order_w90' in specs:
+
orbital_order_w90 = specs['orbital_order_w90']
+
else:
+
orbital_order_w90 = list(range(n_orb))
+
+
if orbital_order_to is None:
+
orbital_order_to = orbital_order_w90
+
+
# checks
+
assert len(set(orbital_order_to)) == len(orbital_order_to), 'Please provide a unique identifier for each orbital.'
+
+
assert set(orbital_order_w90) == set(orbital_order_to), f'Identifiers of orbital_order_to and orbital_order_w90'\
+
f'do not match! orbital_order_to is {orbital_order_to}, but orbital_order_w90 is {orbital_order_w90}.'
+
+
assert with_sigma or eta != 0.0, 'if no Sigma is provided eta has to be different from 0.0'
+
+
# proj_on_orb
+
assert isinstance(proj_on_orb, (int, type(None))) or all(isinstance(x, (int, type(None))) for x in proj_on_orb), 'proj_on_orb should be '\
+
f'an integer or list of integers, but is {type(specs["proj_on_orb"])}.'
+
+
if isinstance(proj_on_orb, (int, type(None))):
+
proj_on_orb = [proj_on_orb]
+
else:
+
proj_on_orb = proj_on_orb
+
+
# if projection is requested we have to use band_basis
+
if proj_on_orb[0] is not None:
+
band_basis = True
+
+
# if proj_nuk is given we need to use the band_basis
+
if isinstance(proj_nuk, np.ndarray) and not band_basis:
+
band_basis = True
+
+
if TB_obj is None:
+
assert w90_path is not None and w90_seed is not None, 'Please provide either a TB object or a path to the wannier90 files'
+
# set up Wannier Hamiltonian
+
n_orb = 2 * n_orb if add_spin else n_orb
+
change_of_basis = change_basis(n_orb, orbital_order_to, orbital_order_w90)
+
H_add_loc = np.zeros((n_orb, n_orb), dtype=complex)
+
if not isinstance(add_local, type(None)):
+
assert np.shape(add_local) == (n_orb, n_orb), 'add_local must have dimension (n_orb, n_orb), but has '\
+
f'dimension {np.shape(add_local)}'
+
H_add_loc += add_local
+
if add_spin and add_lambda:
+
H_add_loc += lambda_matrix_w90_t2g(add_lambda)
+
+
tb = TB_from_wannier90(path=w90_path, seed=w90_seed, extend_to_spin=add_spin, add_local=H_add_loc)
+
else:
+
assert not add_spin, 'add_spin is only valid when reading from wannier90 files'
+
change_of_basis = change_basis(n_orb, orbital_order_to, orbital_order_w90)
+
tb = TB_obj
+
+
eta = eta * 1j
+
+
# print local H(R)
+
h_of_r = np.einsum('ij, jk -> ik', np.linalg.inv(change_of_basis), np.einsum('ij, jk -> ik', tb.hoppings[(0, 0, 0)], change_of_basis))
+
if n_orb <= 12:
+
print_matrix(h_of_r, n_orb, 'H(R=0)')
+
+
# kmesh prep
+
if ('bands_path' in specs and 'kmesh' in specs) or ('bands_path' not in specs and 'kmesh' not in specs):
+
raise ValueError('choose either a bands_path or kmesh!')
+
elif 'bands_path' in specs:
+
w90_paths = list(map(lambda section: (
+
np.array(specs[section[0]]), np.array(specs[section[1]])), specs['bands_path']))
+
k_points_labels = [k[0] for k in specs['bands_path']] + [specs['bands_path'][-1][1]]
+
n_k = specs['n_k']
+
k_vec, k_1d, special_k = k_space_path(w90_paths, bz=tb.bz, num=n_k)
+
elif 'kmesh' in specs:
+
assert 'reg' in specs['kmesh'], 'only regular kmesh is implemented'
+
+
special_k = k_points_labels = None
+
+
# read kmesh size
+
if 'n_k' in specs:
+
k_dim = specs['n_k']
+
elif 'k_dim' in specs:
+
k_dim = specs['k_dim']
+
else:
+
raise ValueError('please specify either n_k or k_dim')
+
+
# create regular kmesh
+
if isinstance(k_dim, int):
+
k_spacing = np.linspace(0, 1, k_dim, endpoint=False)
+
k_vec = np.array(np.meshgrid(k_spacing, k_spacing, k_spacing)).T.reshape(-1, 3)
+
n_k = k_dim**3
+
k_1d = (k_dim, k_dim, k_dim)
+
elif all(isinstance(x, int) for x in k_dim) and len(k_dim) == 3:
+
k_x = np.linspace(0, 1, k_dim[0], endpoint=False)
+
k_y = np.linspace(0, 1, k_dim[1], endpoint=False)
+
k_z = np.linspace(0, 1, k_dim[2], endpoint=False)
+
k_vec = np.array(np.meshgrid(k_x, k_y, k_z)).T.reshape(-1, 3)
+
n_k = k_dim[0]*k_dim[1]*k_dim[2]
+
k_1d = k_dim
+
else:
+
raise ValueError(
+
'k_dim / n_k needs to be either an int or a list / tuple of int length 3')
+
+
# calculate tight-binding eigenvalues for non slices
+
if not fermi_slice:
+
# Fourier trafo on input grid / path
+
e_mat = tb.fourier(k_vec).transpose(1, 2, 0)
+
e_mat = np.einsum('ij, jkl -> ikl', np.linalg.inv(change_of_basis),
+
np.einsum('ijk, jm -> imk', e_mat, change_of_basis))
+
else:
+
if 'kz' in specs and specs['kz'] != 0.:
+
assert 'Z' in specs, 'Please provide Z point coordinate in tb_data_dict as input coordinate'
+
Z = np.array(specs['Z'])
+
kz = specs['kz']
+
else:
+
kz = 0.
+
Z = np.zeros((3))
+
+
k_vec = np.zeros((n_k*n_k, 3))
+
e_mat = np.zeros((n_orb, n_orb, n_k, n_k), dtype=complex)
+
+
upper_left = w90_paths[0][0]
+
lower_right = w90_paths[1][-1]
+
origin = w90_paths[0][1]
+
for ik_y in range(n_k):
+
path_along_x = [((upper_left - origin)/(n_k-1)*ik_y+kz*Z+origin, origin+(lower_right-origin)+(upper_left-origin)/(n_k-1)*ik_y+kz*Z)]
+
k_vec[ik_y*n_k:ik_y*n_k+n_k, :], k_1d, special_k = k_space_path(path_along_x, bz=tb.bz, num=n_k)
+
e_mat[:, :, :, ik_y] = tb.fourier(k_vec[ik_y*n_k:ik_y*n_k+n_k, :]).transpose(1, 2, 0)
+
#if add_spin:
+
# e_mat = e_mat[2:5, 2:5]
+
e_mat = np.einsum('ij, jklm -> iklm', np.linalg.inv(change_of_basis), np.einsum('ijkl, jm -> imkl', e_mat, change_of_basis))
+
+
if band_basis:
+
e_mat, e_vecs, orb_proj = get_tb_bands(e_mat, proj_on_orb)
+
else:
+
e_vecs = total_proj = orb_proj = None
+
+
# now we merge proj_nuk and orb_proj (has reverse shape)
+
if isinstance(proj_nuk, np.ndarray) and isinstance(orb_proj, np.ndarray):
+
proj_nuk = proj_nuk * orb_proj
+
elif not isinstance(proj_nuk, np.ndarray) and isinstance(orb_proj, np.ndarray):
+
proj_nuk = orb_proj
+
+
# dmft output
+
if with_sigma:
+
sigma_types = ['calc', 'model']
+
if isinstance(with_sigma, str):
+
if with_sigma not in sigma_types:
+
raise ValueError('Invalid sigma type. Expected one of: {}'.format(sigma_types))
+
elif not isinstance(with_sigma, BlockGf):
+
raise ValueError('Invalid sigma type. Expected BlockGf.')
+
+
# get sigma
+
if with_sigma == 'model':
+
mu_dmft = None if 'mu_dmft' not in specs else specs['mu_dmft']
+
delta_sigma, freq_dict = sigma_FL(n_orb, orbital_order_to, specs['Sigma_0'], specs['Sigma_Z'], specs['w_mesh'], eta=eta, mu_dmft=mu_dmft)
+
mu = mu_tb + mu_shift
+
# else is from dmft or memory:
+
else:
+
delta_sigma, mu_dmft, freq_dict = _sigma_from_dmft(n_orb, orbital_order_to, with_sigma, **specs)
+
mu = mu_dmft + mu_shift
+
+
freq_dict['sigma_upfolded'] = delta_sigma
+
if add_mu_tb:
+
print('Adding mu_tb to DMFT μ; assuming DMFT was run with subtracted dft μ.')
+
mu += mu_tb
+
+
print('μ={:2.4f} eV set for calculating A(k,ω)'.format(mu))
+
+
assert n_orb == delta_sigma.shape[0] and n_orb == delta_sigma.shape[
+
1], f'Number of orbitals n_orb={n_orb} and shape of sigma: {delta_sigma.shape} does not match'
+
if isinstance(proj_nuk, np.ndarray):
+
assert n_orb == proj_nuk.shape[0], f'Number of orbitals n_orb={n_orb} does not match shape of proj_nuk: {proj_nuk.shape[0]}'
+
if not fermi_slice:
+
assert proj_nuk.shape[-1] == e_vecs.shape[
+
2], f'Number of kpoints in proj_nuk : {proj_nuk.shape[-1]} does not match number of kpoints in e_vecs: {e_vecs.shape[2]}'
+
else:
+
assert proj_nuk.shape == tuple([n_orb, e_vecs.shape[2], e_vecs.shape[3]]
+
), f'shape of projectors {proj_nuk.shape} does not match expected shape of [{n_orb},{e_vecs.shape[2]},{e_vecs.shape[3]}]'
+
+
# calculate alatt
+
if not fermi_slice:
+
alatt_k_w = _calc_alatt(n_orb, mu, eta, e_mat, delta_sigma, qp_bands, e_vecs=e_vecs,
+
trace=trace, proj_nuk=proj_nuk, **freq_dict)
+
else:
+
alatt_k_w = _calc_kslice(n_orb, mu, eta, e_mat, delta_sigma, qp_bands, e_vecs=e_vecs,
+
proj_nuk=proj_nuk, **freq_dict)
+
else:
+
freq_dict = {}
+
freq_dict['w_mesh'] = None
+
freq_dict['window'] = None
+
alatt_k_w = None
+
+
tb_data = {'k_mesh': k_1d, 'special_k': special_k, 'k_points': k_vec,
+
'k_points_labels': k_points_labels, 'e_mat': e_mat,
+
'e_vecs': e_vecs, 'tb': tb, 'mu_tb': mu_tb,
+
'proj_on_orb': proj_on_orb, 'proj_nuk': proj_nuk}
+
+
return tb_data, alatt_k_w, freq_dict
+