dft_tools/python/sumk_dft_tools.py

##########################################################################
#
# TRIQS: a Toolbox for Research in Interacting Quantum Systems
#
# Copyright (C) 2011 by M. Aichhorn, L. Pourovskii, V. Vildosola
#
# 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/>.
#
##########################################################################
import sys
from types import *
import numpy
from pytriqs.gf import *
import pytriqs.utility.mpi as mpi
from symmetry import *
from sumk_dft import SumkDFT
from scipy.integrate import *
from scipy.interpolate import *

if not hasattr(numpy, 'full'):
    # polyfill full for older numpy:
    numpy.full = lambda a, f: numpy.zeros(a) + f

class SumkDFTTools(SumkDFT):
    """
    Extends the SumkDFT class with some tools for analysing the data.
    """

    def __init__(self, hdf_file, h_field=0.0, use_dft_blocks=False, dft_data='dft_input', symmcorr_data='dft_symmcorr_input',
                 parproj_data='dft_parproj_input', symmpar_data='dft_symmpar_input', bands_data='dft_bands_input',
                 transp_data='dft_transp_input', misc_data='dft_misc_input'):
        """
        Initialisation of the class. Parameters are exactly as for SumKDFT.
        """

        SumkDFT.__init__(self, hdf_file=hdf_file, h_field=h_field, use_dft_blocks=use_dft_blocks,
                         dft_data=dft_data, symmcorr_data=symmcorr_data, parproj_data=parproj_data,
                         symmpar_data=symmpar_data, bands_data=bands_data, transp_data=transp_data,
                         misc_data=misc_data)

    # Uses .data of only GfReFreq objects.
    def dos_wannier_basis(self, mu=None, broadening=None, mesh=None, with_Sigma=True, with_dc=True, save_to_file=True):
        """
        Calculates the density of states in the basis of the Wannier functions.

        Parameters
        ----------
        mu : double, optional
             Chemical potential, overrides the one stored in the hdf5 archive.
        broadening : double, optional
                     Lorentzian broadening of the spectra. If not given, standard value of lattice_gf is used.
        mesh : real frequency MeshType, optional
               Omega mesh for the real-frequency Green's function. Given as parameter to lattice_gf.
        with_Sigma : boolean, optional
                     If True, the self energy is used for the calculation. If false, the DOS is calculated without self energy.
        with_dc : boolean, optional
                  If True the double counting correction is used.
        save_to_file : boolean, optional
                       If True, text files with the calculated data will be created.

        Returns
        -------
        DOS : Dict of numpy arrays
              Contains the full density of states.
        DOSproj :  Dict of numpy arrays
                   DOS projected to atoms.
        DOSproj_orb : Dict of numpy arrays
                      DOS projected to atoms and resolved into orbital contributions.
        """

        if (mesh is None) and (not with_Sigma):
            raise ValueError, "lattice_gf: Give the mesh=(om_min,om_max,n_points) for the lattice GfReFreq."
        if mesh is None:
            om_mesh = [x.real for x in self.Sigma_imp_w[0].mesh]
            om_min = om_mesh[0]
            om_max = om_mesh[-1]
            n_om = len(om_mesh)
            mesh = (om_min, om_max, n_om)
        else:
            om_min, om_max, n_om = mesh
            om_mesh = numpy.linspace(om_min, om_max, n_om)

        G_loc = []
        for icrsh in range(self.n_corr_shells):
            spn = self.spin_block_names[self.corr_shells[icrsh]['SO']]
            glist = [GfReFreq(indices=inner, window=(om_min, om_max), n_points=n_om)
                     for block, inner in self.gf_struct_sumk[icrsh]]
            G_loc.append(
                BlockGf(name_list=spn, block_list=glist, make_copies=False))
        for icrsh in range(self.n_corr_shells):
            G_loc[icrsh].zero()

        DOS = {sp: numpy.zeros([n_om], numpy.float_)
               for sp in self.spin_block_names[self.SO]}
        DOSproj = [{} for ish in range(self.n_inequiv_shells)]
        DOSproj_orb = [{} for ish in range(self.n_inequiv_shells)]
        for ish in range(self.n_inequiv_shells):
            for sp in self.spin_block_names[self.corr_shells[self.inequiv_to_corr[ish]]['SO']]:
                dim = self.corr_shells[self.inequiv_to_corr[ish]]['dim']
                DOSproj[ish][sp] = numpy.zeros([n_om], numpy.float_)
                DOSproj_orb[ish][sp] = numpy.zeros(
                    [n_om, dim, dim], numpy.complex_)

        ikarray = numpy.array(range(self.n_k))
        for ik in mpi.slice_array(ikarray):

            G_latt_w = self.lattice_gf(
                ik=ik, mu=mu, iw_or_w="w", broadening=broadening, mesh=mesh, with_Sigma=with_Sigma, with_dc=with_dc)
            G_latt_w *= self.bz_weights[ik]

            # Non-projected DOS
            for iom in range(n_om):
                for bname, gf in G_latt_w:
                    DOS[bname][iom] -= gf.data[iom, :, :].imag.trace() / \
                        numpy.pi

            # Projected DOS:
            for icrsh in range(self.n_corr_shells):
                tmp = G_loc[icrsh].copy()
                for bname, gf in tmp:
                    tmp[bname] << self.downfold(ik, icrsh, bname, G_latt_w[
                                                bname], gf)  # downfolding G
                G_loc[icrsh] += tmp

        # Collect data from mpi:
        for bname in DOS:
            DOS[bname] = mpi.all_reduce(
                mpi.world, DOS[bname], lambda x, y: x + y)
        for icrsh in range(self.n_corr_shells):
            G_loc[icrsh] << mpi.all_reduce(
                mpi.world, G_loc[icrsh], lambda x, y: x + y)
        mpi.barrier()

        # Symmetrize and rotate to local coord. system if needed:
        if self.symm_op != 0:
            G_loc = self.symmcorr.symmetrize(G_loc)
        if self.use_rotations:
            for icrsh in range(self.n_corr_shells):
                for bname, gf in G_loc[icrsh]:
                    G_loc[icrsh][bname] << self.rotloc(
                        icrsh, gf, direction='toLocal')

        # G_loc can now also be used to look at orbitally-resolved quantities
        for ish in range(self.n_inequiv_shells):
            for bname, gf in G_loc[self.inequiv_to_corr[ish]]:  # loop over spins
                for iom in range(n_om):
                    DOSproj[ish][bname][iom] -= gf.data[iom,
                                                        :, :].imag.trace() / numpy.pi
                DOSproj_orb[ish][bname][
                    :, :, :] += (1.0j*(gf-gf.conjugate().transpose())/2.0/numpy.pi).data[:,:,:]

        # Write to files
        if save_to_file and mpi.is_master_node():
            for sp in self.spin_block_names[self.SO]:
                f = open('DOS_wann_%s.dat' % sp, 'w')
                for iom in range(n_om):
                    f.write("%s    %s\n" % (om_mesh[iom], DOS[sp][iom]))
                f.close()

                # Partial
                for ish in range(self.n_inequiv_shells):
                    f = open('DOS_wann_%s_proj%s.dat' % (sp, ish), 'w')
                    for iom in range(n_om):
                        f.write("%s    %s\n" %
                                (om_mesh[iom], DOSproj[ish][sp][iom]))
                    f.close()

                    # Orbitally-resolved
                    for i in range(self.corr_shells[self.inequiv_to_corr[ish]]['dim']):
                        for j in range(i, self.corr_shells[self.inequiv_to_corr[ish]]['dim']):
                            f = open('DOS_wann_' + sp + '_proj' + str(ish) +
                                     '_' + str(i) + '_' + str(j) + '.dat', 'w')
                            for iom in range(n_om):
                                f.write("%s    %s    %s\n" % (
                                    om_mesh[iom], DOSproj_orb[ish][sp][iom, i, j].real,DOSproj_orb[ish][sp][iom, i, j].imag))
                            f.close()

        return DOS, DOSproj, DOSproj_orb

    # Uses .data of only GfReFreq objects.
    def dos_parproj_basis(self, mu=None, broadening=None, mesh=None, with_Sigma=True, with_dc=True, save_to_file=True):
        """
        Calculates the orbitally-resolved DOS.
        Different to dos_Wannier_basis is that here we calculate projections also to non-Wannier projectors, in the
        flavour of Wien2k QTL calculatuions.

        Parameters
        ----------
        mu : double, optional
             Chemical potential, overrides the one stored in the hdf5 archive.
        broadening : double, optional
                     Lorentzian broadening of the spectra. If not given, standard value of lattice_gf is used.
        mesh : real frequency MeshType, optional
               Omega mesh for the real-frequency Green's function. Given as parameter to lattice_gf.
        with_Sigma : boolean, optional
                     If True, the self energy is used for the calculation. If false, the DOS is calculated without self energy.
        with_dc : boolean, optional
                  If True the double counting correction is used.
        save_to_file : boolean, optional
                       If True, text files with the calculated data will be created.

        Returns
        -------
        DOS : Dict of numpy arrays
              Contains the full density of states.
        DOSproj :  Dict of numpy arrays
                   DOS projected to atoms.
        DOSproj_orb : Dict of numpy arrays
                      DOS projected to atoms and resolved into orbital contributions.
        """

        things_to_read = ['n_parproj', 'proj_mat_all',
                          'rot_mat_all', 'rot_mat_all_time_inv']
        value_read = self.read_input_from_hdf(
            subgrp=self.parproj_data, things_to_read=things_to_read)
        if not value_read:
            return value_read
        if self.symm_op:
            self.symmpar = Symmetry(self.hdf_file, subgroup=self.symmpar_data)

        if (mesh is None) and (not with_Sigma):
            raise ValueError, "lattice_gf: Give the mesh=(om_min,om_max,n_points) for the lattice GfReFreq."
        if mesh is None:
            om_mesh = [x.real for x in self.Sigma_imp_w[0].mesh]
            om_min = om_mesh[0]
            om_max = om_mesh[-1]
            n_om = len(om_mesh)
            mesh = (om_min, om_max, n_om)
        else:
            om_min, om_max, n_om = mesh
            om_mesh = numpy.linspace(om_min, om_max, n_om)

        G_loc = []
        spn = self.spin_block_names[self.SO]
        gf_struct_parproj = [[(sp, range(self.shells[ish]['dim'])) for sp in spn]
                             for ish in range(self.n_shells)]
        for ish in range(self.n_shells):
            glist = [GfReFreq(indices=inner, window=(om_min, om_max), n_points=n_om)
                     for block, inner in gf_struct_parproj[ish]]
            G_loc.append(
                BlockGf(name_list=spn, block_list=glist, make_copies=False))
        for ish in range(self.n_shells):
            G_loc[ish].zero()

        DOS = {sp: numpy.zeros([n_om], numpy.float_)
               for sp in self.spin_block_names[self.SO]}
        DOSproj = [{} for ish in range(self.n_shells)]
        DOSproj_orb = [{} for ish in range(self.n_shells)]
        for ish in range(self.n_shells):
            for sp in self.spin_block_names[self.SO]:
                dim = self.shells[ish]['dim']
                DOSproj[ish][sp] = numpy.zeros([n_om], numpy.float_)
                DOSproj_orb[ish][sp] = numpy.zeros(
                    [n_om, dim, dim], numpy.complex_)

        ikarray = numpy.array(range(self.n_k))
        for ik in mpi.slice_array(ikarray):

            G_latt_w = self.lattice_gf(
                ik=ik, mu=mu, iw_or_w="w", broadening=broadening, mesh=mesh, with_Sigma=with_Sigma, with_dc=with_dc)
            G_latt_w *= self.bz_weights[ik]

            # Non-projected DOS
            for iom in range(n_om):
                for bname, gf in G_latt_w:
                    DOS[bname][iom] -= gf.data[iom, :, :].imag.trace() / \
                        numpy.pi

            # Projected DOS:
            for ish in range(self.n_shells):
                tmp = G_loc[ish].copy()
                for ir in range(self.n_parproj[ish]):
                    for bname, gf in tmp:
                        tmp[bname] << self.downfold(ik, ish, bname, G_latt_w[
                                                    bname], gf, shells='all', ir=ir)
                    G_loc[ish] += tmp

        # Collect data from mpi:
        for bname in DOS:
            DOS[bname] = mpi.all_reduce(
                mpi.world, DOS[bname], lambda x, y: x + y)
        for ish in range(self.n_shells):
            G_loc[ish] << mpi.all_reduce(
                mpi.world, G_loc[ish], lambda x, y: x + y)
        mpi.barrier()

        # Symmetrize and rotate to local coord. system if needed:
        if self.symm_op != 0:
            G_loc = self.symmpar.symmetrize(G_loc)
        if self.use_rotations:
            for ish in range(self.n_shells):
                for bname, gf in G_loc[ish]:
                    G_loc[ish][bname] << self.rotloc(
                        ish, gf, direction='toLocal', shells='all')

        # G_loc can now also be used to look at orbitally-resolved quantities
        for ish in range(self.n_shells):
            for bname, gf in G_loc[ish]:
                for iom in range(n_om):
                    DOSproj[ish][bname][iom] -= gf.data[iom,
                                                        :, :].imag.trace() / numpy.pi
                DOSproj_orb[ish][bname][
                    :, :, :] += (1.0j*(gf-gf.conjugate().transpose())/2.0/numpy.pi).data[:,:,:]

        # Write to files
        if save_to_file and mpi.is_master_node():
            for sp in self.spin_block_names[self.SO]:
                f = open('DOS_parproj_%s.dat' % sp, 'w')
                for iom in range(n_om):
                    f.write("%s    %s\n" % (om_mesh[iom], DOS[sp][iom]))
                f.close()

                # Partial
                for ish in range(self.n_shells):
                    f = open('DOS_parproj_%s_proj%s.dat' % (sp, ish), 'w')
                    for iom in range(n_om):
                        f.write("%s    %s\n" %
                                (om_mesh[iom], DOSproj[ish][sp][iom]))
                    f.close()

                    # Orbitally-resolved
                    for i in range(self.shells[ish]['dim']):
                        for j in range(i, self.shells[ish]['dim']):
                            f = open('DOS_parproj_' + sp + '_proj' + str(ish) +
                                     '_' + str(i) + '_' + str(j) + '.dat', 'w')
                            for iom in range(n_om):
                                f.write("%s    %s    %s\n" % (
                                    om_mesh[iom], DOSproj_orb[ish][sp][iom, i, j].real,DOSproj_orb[ish][sp][iom, i, j].imag))
                            f.close()

        return DOS, DOSproj, DOSproj_orb

    # Uses .data of only GfReFreq objects.
    def spaghettis(self, broadening=None, plot_shift=0.0, plot_range=None, ishell=None, mu=None, save_to_file='Akw_'):
        """
        Calculates the correlated band structure using a real-frequency self energy.

        Parameters
        ----------
        mu : double, optional
             Chemical potential, overrides the one stored in the hdf5 archive.
        broadening : double, optional
                     Lorentzian broadening of the spectra. If not given, standard value of lattice_gf is used.
        plot_shift : double, optional
                     Offset for each A(k,w) for stacked plotting of spectra.
        plot_range : list of double, optional
                     Sets the energy window for plotting to (plot_range[0],plot_range[1]). If not provided, the energy mesh of the self energy is used.
        ishell : integer, optional
                 Contains the index of the shell on which the spectral function is projected. If ishell=None, the total spectrum without projection is calculated.
        save_to_file : string, optional
                       Filename where the spectra are stored.

        Returns
        -------
        Akw : Dict of numpy arrays
              Data as it is also written to the files.
        """

        assert hasattr(
            self, "Sigma_imp_w"), "spaghettis: Set Sigma_imp_w first."
        things_to_read = ['n_k', 'n_orbitals', 'proj_mat',
                          'hopping', 'n_parproj', 'proj_mat_all']
        value_read = self.read_input_from_hdf(
            subgrp=self.bands_data, things_to_read=things_to_read)
        if not value_read:
            return value_read
        if ishell is not None:
            things_to_read = ['rot_mat_all', 'rot_mat_all_time_inv']
            value_read = self.read_input_from_hdf(
                subgrp=self.parproj_data, things_to_read=things_to_read)
            if not value_read:
                return value_read

        if mu is None:
            mu = self.chemical_potential
        spn = self.spin_block_names[self.SO]
        mesh = [x.real for x in self.Sigma_imp_w[0].mesh]
        n_om = len(mesh)

        if plot_range is None:
            om_minplot = mesh[0] - 0.001
            om_maxplot = mesh[n_om - 1] + 0.001
        else:
            om_minplot = plot_range[0]
            om_maxplot = plot_range[1]

        if ishell is None:
            Akw = {sp: numpy.zeros([self.n_k, n_om], numpy.float_)
                   for sp in spn}
        else:
            Akw = {sp: numpy.zeros(
                [self.shells[ishell]['dim'], self.n_k, n_om], numpy.float_) for sp in spn}

        if not ishell is None:
            gf_struct_parproj = [
                (sp, range(self.shells[ishell]['dim'])) for sp in spn]
            G_loc = BlockGf(name_block_generator=[(block, GfReFreq(indices=inner, mesh=self.Sigma_imp_w[0].mesh))
                                                  for block, inner in gf_struct_parproj], make_copies=False)
            G_loc.zero()

        ikarray = numpy.array(range(self.n_k))
        for ik in mpi.slice_array(ikarray):

            G_latt_w = self.lattice_gf(
                ik=ik, mu=mu, iw_or_w="w", broadening=broadening)

            if ishell is None:
                # Non-projected A(k,w)
                for iom in range(n_om):
                    if (mesh[iom] > om_minplot) and (mesh[iom] < om_maxplot):
                        for bname, gf in G_latt_w:
                            Akw[bname][ik, iom] += gf.data[iom, :,
                                                           :].imag.trace() / (-1.0 * numpy.pi)
                        # shift Akw for plotting stacked k-resolved eps(k)
                        # curves
                        Akw[bname][ik, iom] += ik * plot_shift

            else:  # ishell not None
                # Projected A(k,w):
                G_loc.zero()
                tmp = G_loc.copy()
                for ir in range(self.n_parproj[ishell]):
                    for bname, gf in tmp:
                        tmp[bname] << self.downfold(ik, ishell, bname, G_latt_w[
                                                    bname], gf, shells='all', ir=ir)
                    G_loc += tmp

                # Rotate to local frame
                if self.use_rotations:
                    for bname, gf in G_loc:
                        G_loc[bname] << self.rotloc(
                            ishell, gf, direction='toLocal', shells='all')

                for iom in range(n_om):
                    if (mesh[iom] > om_minplot) and (mesh[iom] < om_maxplot):
                        for ish in range(self.shells[ishell]['dim']):
                            for sp in spn:
                                Akw[sp][ish, ik, iom] = G_loc[sp].data[
                                    iom, ish, ish].imag / (-1.0 * numpy.pi)

        # Collect data from mpi
        for sp in spn:
            Akw[sp] = mpi.all_reduce(mpi.world, Akw[sp], lambda x, y: x + y)
        mpi.barrier()

        if save_to_file and mpi.is_master_node():
            if ishell is None:
                for sp in spn:  # loop over GF blocs:
                    # Open file for storage:
                    f = open(save_to_file + sp + '.dat', 'w')
                    for ik in range(self.n_k):
                        for iom in range(n_om):
                            if (mesh[iom] > om_minplot) and (mesh[iom] < om_maxplot):
                                if plot_shift > 0.0001:
                                    f.write('%s      %s\n' %
                                            (mesh[iom], Akw[sp][ik, iom]))
                                else:
                                    f.write('%s     %s      %s\n' %
                                            (ik, mesh[iom], Akw[sp][ik, iom]))
                        f.write('\n')
                    f.close()

            else:  # ishell is not None
                for sp in spn:
                    for ish in range(self.shells[ishell]['dim']):
                        # Open file for storage:
                        f = open(save_to_file + str(ishell) + '_' +
                                 sp + '_proj' + str(ish) + '.dat', 'w')
                        for ik in range(self.n_k):
                            for iom in range(n_om):
                                if (mesh[iom] > om_minplot) and (mesh[iom] < om_maxplot):
                                    if plot_shift > 0.0001:
                                        f.write('%s      %s\n' % (
                                            mesh[iom], Akw[sp][ish, ik, iom]))
                                    else:
                                        f.write('%s     %s      %s\n' % (
                                            ik, mesh[iom], Akw[sp][ish, ik, iom]))
                            f.write('\n')
                        f.close()

        return Akw

    def partial_charges(self, beta=40, mu=None, with_Sigma=True, with_dc=True):
        """
        Calculates the orbitally-resolved density matrix for all the orbitals considered in the input, consistent with
        the definition of Wien2k. Hence, (possibly non-orthonormal) projectors have to be provided in the partial projectors subgroup of
        the hdf5 archive.

        Parameters
        ----------

        with_Sigma : boolean, optional
                     If True, the self energy is used for the calculation. If false, partial charges are calculated without self-energy correction.
        beta : double, optional
               In case the self-energy correction is not used, the inverse temperature where the calculation should be done has to be given here.
        mu : double, optional
             Chemical potential, overrides the one stored in the hdf5 archive.
        with_dc : boolean, optional
                  If True the double counting correction is used.

        Returns
        -------
        dens_mat : list of numpy array
                   A list of density matrices projected to all shells provided in the input.
        """

        things_to_read = ['dens_mat_below', 'n_parproj',
                          'proj_mat_all', 'rot_mat_all', 'rot_mat_all_time_inv']
        value_read = self.read_input_from_hdf(
            subgrp=self.parproj_data, things_to_read=things_to_read)
        if not value_read:
            return value_read
        if self.symm_op:
            self.symmpar = Symmetry(self.hdf_file, subgroup=self.symmpar_data)

        spn = self.spin_block_names[self.SO]
        ntoi = self.spin_names_to_ind[self.SO]
        # Density matrix in the window
        self.dens_mat_window = [[numpy.zeros([self.shells[ish]['dim'], self.shells[ish]['dim']], numpy.complex_)
                                 for ish in range(self.n_shells)]
                                for isp in range(len(spn))]
        # Set up G_loc
        gf_struct_parproj = [[(sp, range(self.shells[ish]['dim'])) for sp in spn]
                             for ish in range(self.n_shells)]
        if with_Sigma:
            G_loc = [BlockGf(name_block_generator=[(block, GfImFreq(indices=inner, mesh=self.Sigma_imp_iw[0].mesh))
                                                   for block, inner in gf_struct_parproj[ish]], make_copies=False)
                     for ish in range(self.n_shells)]
            beta = self.Sigma_imp_iw[0].mesh.beta
        else:
            G_loc = [BlockGf(name_block_generator=[(block, GfImFreq(indices=inner, beta=beta))
                                                   for block, inner in gf_struct_parproj[ish]], make_copies=False)
                     for ish in range(self.n_shells)]
        for ish in range(self.n_shells):
            G_loc[ish].zero()

        ikarray = numpy.array(range(self.n_k))
        for ik in mpi.slice_array(ikarray):

            G_latt_iw = self.lattice_gf(
                ik=ik, mu=mu, iw_or_w="iw", beta=beta, with_Sigma=with_Sigma, with_dc=with_dc)
            G_latt_iw *= self.bz_weights[ik]
            for ish in range(self.n_shells):
                tmp = G_loc[ish].copy()
                for ir in range(self.n_parproj[ish]):
                    for bname, gf in tmp:
                        tmp[bname] << self.downfold(ik, ish, bname, G_latt_iw[
                                                    bname], gf, shells='all', ir=ir)
                    G_loc[ish] += tmp

        # Collect data from mpi:
        for ish in range(self.n_shells):
            G_loc[ish] << mpi.all_reduce(
                mpi.world, G_loc[ish], lambda x, y: x + y)
        mpi.barrier()

        # Symmetrize and rotate to local coord. system if needed:
        if self.symm_op != 0:
            G_loc = self.symmpar.symmetrize(G_loc)
        if self.use_rotations:
            for ish in range(self.n_shells):
                for bname, gf in G_loc[ish]:
                    G_loc[ish][bname] << self.rotloc(
                        ish, gf, direction='toLocal', shells='all')

        for ish in range(self.n_shells):
            isp = 0
            for bname, gf in G_loc[ish]:
                self.dens_mat_window[isp][ish] = G_loc[ish].density()[bname]
                isp += 1

        # Add density matrices to get the total:
        dens_mat = [[self.dens_mat_below[ntoi[spn[isp]]][ish] + self.dens_mat_window[isp][ish]
                     for ish in range(self.n_shells)]
                    for isp in range(len(spn))]

        return dens_mat

    def print_hamiltonian(self):
        """
        Prints the Kohn-Sham Hamiltonian to the text files hamup.dat and hamdn.dat (no spin orbit-coupling), or to ham.dat (with spin-orbit coupling).
        """

        if self.SP == 1 and self.SO == 0:
            f1 = open('hamup.dat', 'w')
            f2 = open('hamdn.dat', 'w')
            for ik in range(self.n_k):
                for i in range(self.n_orbitals[ik, 0]):
                    f1.write('%s    %s\n' %
                             (ik, self.hopping[ik, 0, i, i].real))
                for i in range(self.n_orbitals[ik, 1]):
                    f2.write('%s    %s\n' %
                             (ik, self.hopping[ik, 1, i, i].real))
                f1.write('\n')
                f2.write('\n')
            f1.close()
            f2.close()
        else:
            f = open('ham.dat', 'w')
            for ik in range(self.n_k):
                for i in range(self.n_orbitals[ik, 0]):
                    f.write('%s    %s\n' %
                            (ik, self.hopping[ik, 0, i, i].real))
                f.write('\n')
            f.close()


# ----------------- transport -----------------------

    def read_transport_input_from_hdf(self):
        r"""
        Reads the data for transport calculations from the hdf5 archive.
        """