[wannier] few minor tidying changes while reading through

python/converters/wannier90_converter.py

@@ -1,5 +1,5 @@
 
-################################################################################
+##########################################################################
 #
 # TRIQS: a Toolbox for Research in Interacting Quantum Systems
 #
@@ -18,7 +18,7 @@
 # You should have received a copy of the GNU General Public License along with
 # TRIQS. If not, see <http://www.gnu.org/licenses/>.
 #
-################################################################################
+##########################################################################
 
 ###
 #  Wannier90 to HDF5 converter for the SumkDFT class of dfttools/TRIQS;
@@ -49,8 +49,10 @@ import numpy
 import math
 from pytriqs.archive import *
 from converter_tools import *
+from itertools import product
 import os.path
 
+
 class Wannier90Converter(ConverterTools):
     """
     Conversion from Wannier90 output to an hdf5 file that can be used as input for the SumkDFT class.
@@ -77,10 +79,13 @@ class Wannier90Converter(ConverterTools):
         """
 
         self._name = "Wannier90Converter"
-        assert type(seedname)==StringType, self._name + ": Please provide the DFT files' base name as a string."
+        assert type(seedname) == StringType, self._name + \
-        if hdf_filename is None: hdf_filename = seedname+'.h5'
+            ": Please provide the DFT files' base name as a string."
+        if hdf_filename is None:
+            hdf_filename = seedname + '.h5'
         self.hdf_file = hdf_filename
-        # if the w90 output is seedname_hr.dat, the input file for the converter must be called seedname.inp
+        # if the w90 output is seedname_hr.dat, the input file for the
+        # converter must be called seedname.inp
         self.inp_file = seedname + '.inp'
         self.w90_seed = seedname
         self.dft_subgrp = dft_subgrp
@@ -93,7 +98,6 @@ class Wannier90Converter(ConverterTools):
         if (os.path.exists(self.hdf_file) and repacking):
             ConverterTools.repack(self)
 
-
     def convert_dft_input(self):
         """
         Reads the appropriate files and stores the data for the
@@ -106,41 +110,50 @@ class Wannier90Converter(ConverterTools):
         """
 
         # Read and write only on the master node
-        if not (mpi.is_master_node()): return
+        if not (mpi.is_master_node()):
+            return
         mpi.report("Reading input from %s..." % self.inp_file)
 
         # R is a generator : each R.Next() will return the next number in the file
-        R = ConverterTools.read_fortran_file(self,self.inp_file,self.fortran_to_replace)
+        R = ConverterTools.read_fortran_file(
+            self, self.inp_file, self.fortran_to_replace)
         shell_entries = ['atom', 'sort', 'l', 'dim']
         corr_shell_entries = ['atom', 'sort', 'l', 'dim', 'SO', 'irep']
         # First, let's read the input file with the parameters needed for the conversion
         try:
-            kmesh_mode = int(R.next())           # read k-point mesh generation option
+            # read k - point mesh generation option
+            kmesh_mode = int(R.next())
             if kmesh_mode >= 0:
                 # read k-point mesh size from input
                 nki = [int(R.next()) for idir in range(3)]
             else:
                 # some default grid, if everything else fails...
                 nki = [8, 8, 8]
-            density_required = float(R.next())   # read the total number of electrons per cell
+            # read the total number of electrons per cell
+            density_required = float(R.next())
             # we do not read shells, because we have no additional shells beyond correlated ones,
             # and the data will be copied from corr_shells into shells (see below)
-            n_corr_shells = int(R.next())                 # number of corr. shells (e.g. Fe d, Ce f) in the unit cell, 
+            # number of corr. shells (e.g. Fe d, Ce f) in the unit cell,
+            n_corr_shells = int(R.next())
             # now read the information about the correlated shells (atom, sort, l, dim, SO flag, irep):
-            corr_shells = [ {name: int(val) for name, val in zip(corr_shell_entries, R)} for icrsh in range(n_corr_shells) ]
+            corr_shells = [{name: int(val) for name, val in zip(
+                corr_shell_entries, R)} for icrsh in range(n_corr_shells)]
         except StopIteration:  # a more explicit error if the file is corrupted.
-            mpi.report(self._name + ": reading input file %s failed!"%self.inp_file)
+            mpi.report(self._name + ": reading input file %s failed!" %
+                       self.inp_file)
         # close the input file
         R.close()
 
         # Set or derive some quantities
-        symm_op = 0                     # Wannier90 does not use symmetries to reduce the k-points
+        # Wannier90 does not use symmetries to reduce the k-points
         # the following might change in future versions
-        ### copy corr_shells into shells (see above)
+        symm_op = 0                     
+        # copy corr_shells into shells (see above)
         n_shells = n_corr_shells
         shells = []
         for ish in range(n_shells):
-            shells.append({key: corr_shells[ish].get(key,None) for key in shell_entries})
+            shells.append({key: corr_shells[ish].get(
+                key, None) for key in shell_entries})
         ###
         SP = 0                          # NO spin-polarised calculations for now
         SO = 0                          # NO spin-orbit calculation for now
@@ -150,13 +163,16 @@ class Wannier90Converter(ConverterTools):
         # this is more general
         n_spin = SP + 1 - SO
         dim_corr_shells = sum([sh['dim'] for sh in corr_shells])
-        mpi.report("Total number of WFs expected in the correlated shells: %d"%dim_corr_shells)
+        mpi.report(
+            "Total number of WFs expected in the correlated shells: %d" % dim_corr_shells)
 
         # determine the number of inequivalent correlated shells and maps, needed for further processing
-        n_inequiv_shells, corr_to_inequiv, inequiv_to_corr = ConverterTools.det_shell_equivalence(self,corr_shells)
+        n_inequiv_shells, corr_to_inequiv, inequiv_to_corr = ConverterTools.det_shell_equivalence(
+            self, corr_shells)
         mpi.report("Number of inequivalent shells: %d" % n_inequiv_shells)
         mpi.report("Shell representatives: " + format(inequiv_to_corr))
-        shells_map = [inequiv_to_corr[corr_to_inequiv[ish]] for ish in range(n_corr_shells)]
+        shells_map = [inequiv_to_corr[corr_to_inequiv[ish]]
+                      for ish in range(n_corr_shells)]
         mpi.report("Mapping: " + format(shells_map))
 
         # build the k-point mesh, if its size was given on input (kmesh_mode >= 0),
@@ -175,26 +191,26 @@ class Wannier90Converter(ConverterTools):
             lmax = ll * (corr_shells[inequiv_to_corr[ish]]['SO'] + 1)
             T.append(numpy.zeros([lmax, lmax], numpy.complex_))
 
-
         spin_w90name = ['_up', '_down']
         hamr_full = []
 
         # TODO: generalise to SP=1 (only partially done)
         rot_mat_time_inv = [0 for i in range(n_corr_shells)]
 
-
         # Second, let's read the file containing the Hamiltonian in WF basis produced by Wannier90
         for isp in range(n_spin):
-        ### begin loop on isp
+            # begin loop on isp
 
             # build filename according to wannier90 conventions
             if SP == 1:
-                mpi.report("Reading information for spin component n. %d"%isp)
+                mpi.report(
+                    "Reading information for spin component n. %d" % isp)
                 hr_file = self.w90_seed + spin_w90name[isp] + '_hr.dat'
             else:
                 hr_file = self.w90_seed + '_hr.dat'
             # now grab the data from the H(R) file
-            mpi.report("The Hamiltonian in MLWF basis is extracted from %s ..."%hr_file)
+            mpi.report(
+                "The Hamiltonian in MLWF basis is extracted from %s ..." % hr_file)
             nr, rvec, rdeg, nw, hamr = self.read_wannier90hr(hr_file)
             # number of R vectors, their indices, their degeneracy, number of WFs, H(R)
             mpi.report("... done: %d R vectors, %d WFs found" % (nr, nw))
@@ -230,15 +246,18 @@ class Wannier90Converter(ConverterTools):
                     mpi.report("Number of WFs equal to number of correlated orbitals")
 
                 # we assume spin up and spin down always have same total number of WFs
-                n_orbitals = numpy.ones([self.n_k,n_spin],numpy.int)*self.nwfs
+                n_orbitals = numpy.ones(
+                    [self.n_k, n_spin], numpy.int) * self.nwfs
 
             else:
                 # consistency check between the _up and _down file contents
-                if nr != self.nrpt: mpi.report("Different number of R vectors for spin-up/spin-down!")
+                if nr != self.nrpt:
-                if nw != self.nwfs: mpi.report("Different number of WFs for spin-up/spin-down!")
+                    mpi.report("Different number of R vectors for spin-up/spin-down!")
+                if nw != self.nwfs:
+                    mpi.report("Different number of WFs for spin-up/spin-down!")
 
             hamr_full.append(hamr)
-            ##FIXME: when do we actually need deepcopy()?
+            # FIXME: when do we actually need deepcopy()?
             # hamr_full.append(deepcopy(hamr))
 
             for ir in range(nr):
@@ -265,18 +284,18 @@ class Wannier90Converter(ConverterTools):
                 for icrsh in range(n_corr_shells):
                     if not numpy.allclose(rot_mat_[icrsh], rot_mat[icrsh], atol=self._w90zero, rtol=1.e-15):
                         mpi.report("Rotations for spin component n. %d do not match!" % isp)
-        ### end loop on isp
+        # end loop on isp
-
 
         mpi.report("The k-point grid has dimensions: %d, %d, %d" % tuple(nki))
         # if calculations are spin-polarized, then renormalize k-point weights
-        if SP == 1: bz_weights = 0.5 * bz_weights
+        if SP == 1:
+            bz_weights = 0.5 * bz_weights
 
         # Third, compute the hoppings in reciprocal space
         hopping = numpy.zeros([self.n_k, n_spin, numpy.max(n_orbitals), numpy.max(n_orbitals)], numpy.complex_)
         for isp in range(n_spin):
             # make Fourier transform H(R) -> H(k) : it can be done one spin at a time
-            hamk = self.fourierham(self.nwfs, hamr_full[isp])
+            hamk = self.fourier_ham(self.nwfs, hamr_full[isp])
             # copy the H(k) in the right place of hoppings... is there a better way to do this??
             for ik in range(self.n_k):
                 #hopping[ik,isp,:,:] = deepcopy(hamk[ik][:,:])*energy_unit
@@ -284,7 +303,8 @@ class Wannier90Converter(ConverterTools):
 
         # Then, initialise the projectors
         k_dep_projection = 0   # we always have the same number of WFs at each k-point
-        proj_mat = numpy.zeros([self.n_k,n_spin,n_corr_shells,max([crsh['dim'] for crsh in corr_shells]),numpy.max(n_orbitals)],numpy.complex_)
+        proj_mat = numpy.zeros([self.n_k, n_spin, n_corr_shells, max(
+            [crsh['dim'] for crsh in corr_shells]), numpy.max(n_orbitals)], numpy.complex_)
         iorb = 0
         # Projectors simply consist in identity matrix blocks selecting those MLWFs that
         # correspond to the specific correlated shell indexed by icrsh.
@@ -292,22 +312,23 @@ class Wannier90Converter(ConverterTools):
         # file and that the ordering of MLWFs matches the corr_shell info from the input.
         for icrsh in range(n_corr_shells):
             norb = corr_shells[icrsh]['dim']
-            proj_mat[:,:,icrsh,0:norb,iorb:iorb+norb] = numpy.identity(norb,numpy.complex_)
+            proj_mat[:, :, icrsh, 0:norb, iorb:iorb +
+                     norb] = numpy.identity(norb, numpy.complex_)
             iorb += norb
 
         # Finally, save all required data into the HDF archive:
         ar = HDFArchive(self.hdf_file, 'a')
-        if not (self.dft_subgrp in ar): ar.create_group(self.dft_subgrp) 
+        if not (self.dft_subgrp in ar):
+            ar.create_group(self.dft_subgrp)
         # The subgroup containing the data. If it does not exist, it is created. If it exists, the data is overwritten!
         things_to_save = ['energy_unit', 'n_k', 'k_dep_projection', 'SP', 'SO', 'charge_below', 'density_required',
                           'symm_op', 'n_shells', 'shells', 'n_corr_shells', 'corr_shells', 'use_rotations', 'rot_mat',
                           'rot_mat_time_inv', 'n_reps', 'dim_reps', 'T', 'n_orbitals', 'proj_mat', 'bz_weights', 'hopping',
                           'n_inequiv_shells', 'corr_to_inequiv', 'inequiv_to_corr']
-        for it in things_to_save: ar[self.dft_subgrp][it] = locals()[it]
+        for it in things_to_save:
+            ar[self.dft_subgrp][it] = locals()[it]
         del ar
 
-
-
     def read_wannier90hr(self, hr_filename="wannier_hr.dat"):
         """
         Method for reading the seedname_hr.dat file produced by Wannier90 (http://wannier.org)
@@ -333,7 +354,8 @@ class Wannier90Converter(ConverterTools):
         """
 
         # Read only from the master node
-        if not (mpi.is_master_node()): return
+        if not (mpi.is_master_node()):
+            return
 
         try:
             with open(hr_filename, "r") as hr_filedesc:
@@ -345,7 +367,8 @@ class Wannier90Converter(ConverterTools):
         mpi.report("Reading %s..." % hr_filename + hr_data[0])
 
         try:
-            num_wf = int(hr_data[1])  # reads number of Wannier functions per spin
+            # reads number of Wannier functions per spin
+            num_wf = int(hr_data[1])
             nrpt = int(hr_data[2])
         except ValueError:
             mpi.report("Could not read number of WFs or R vectors")
@@ -353,7 +376,8 @@ class Wannier90Converter(ConverterTools):
         # allocate arrays to save the R vector indexes and degeneracies and the Hamiltonian
         rvec_idx = numpy.zeros((nrpt, 3), dtype=int)
         rvec_deg = numpy.zeros(nrpt, dtype=int)
-        h_of_r = [numpy.zeros((num_wf, num_wf), dtype=numpy.complex_) for n in range(nrpt)]
+        h_of_r = [numpy.zeros((num_wf, num_wf), dtype=numpy.complex_)
+                  for n in range(nrpt)]
 
         # variable currpos points to the current line in the file
         currpos = 2
@@ -367,25 +391,26 @@ class Wannier90Converter(ConverterTools):
                         raise IndexError("wrong number of R vectors??")
                     rvec_deg[ir] = int(x)
                     ir += 1
-            # for each direct lattice vector R
+            # for each direct lattice vector R read the block of the
-            for ir in range(nrpt):
+            # Hamiltonian H(R)
-                # read the block of the Hamiltonian H(R)
+            for ir, jj, ii in product(range(nrpt), range(num_wf), range(num_wf)):
-                for jj in range(num_wf):
-                    for ii in range(num_wf):
                 # advance one line, split the line into tokens
                 currpos += 1
                 cline = hr_data[currpos].split()
                 # check if the orbital indexes in the file make sense
                 if int(cline[3]) != ii + 1 or int(cline[4]) != jj + 1:
-                            mpi.report("Inconsistent indices at %s%s of R n. %s"%(ii,jj,ir))
+                    mpi.report(
-                        rcurr = numpy.array([int(cline[0]), int(cline[1]), int(cline[2])])
+                        "Inconsistent indices at %s%s of R n. %s" % (ii, jj, ir))
+                rcurr = numpy.array(
+                    [int(cline[0]), int(cline[1]), int(cline[2])])
                 if ii == 0 and jj == 0:
                     rvec_idx[ir] = rcurr
                     rprec = rcurr
                 else:
                     # check if the vector indices are consistent
                     if not numpy.array_equal(rcurr, rprec):
-                                mpi.report("Inconsistent indices for R vector n. %s"%ir)
+                        mpi.report(
+                            "Inconsistent indices for R vector n. %s" % ir)
 
                 # fill h_of_r with the matrix elements of the Hamiltonian
                 h_of_r[ir][ii, jj] = complex(float(cline[5]), float(cline[6]))
@@ -396,8 +421,6 @@ class Wannier90Converter(ConverterTools):
         # return the data into variables
         return nrpt, rvec_idx, rvec_deg, num_wf, h_of_r
 
-
-
     def find_rot_mat(self, n_sh, sh_lst, sh_map, ham0):
         """
         Method for finding the matrices that bring from local to global coordinate systems
@@ -424,12 +447,14 @@ class Wannier90Converter(ConverterTools):
         """
 
         # initialize the rotation matrices to identities
-        rot_mat = [numpy.identity(sh_lst[ish]['dim'], dtype=complex) for ish in range(n_sh)]
+        rot_mat = [numpy.identity(sh_lst[ish]['dim'], dtype=complex)
+                   for ish in range(n_sh)]
         istatus = 0
 
         hs = ham0.shape
         if hs[0] != hs[1] or hs[0] != sum([sh['dim'] for sh in sh_lst]):
-            mpi.report("find_rot_mat: wrong block structure of input Hamiltonian!")
+            mpi.report(
+                "find_rot_mat: wrong block structure of input Hamiltonian!")
             istatus = 0
             # this error will lead into troubles later... early return
             return istatus, rot_mat
@@ -443,7 +468,8 @@ class Wannier90Converter(ConverterTools):
             # nw = number of orbitals in this shell
             nw = sh_lst[ish]["dim"]
             # diagonalize the sub-block of H(0) corresponding to this shell
-            eigval, eigvec = numpy.linalg.eigh(ham0[iwf:iwf+nw, iwf:iwf+nw])
+            eigval, eigvec = numpy.linalg.eigh(
+                ham0[iwf:iwf + nw, iwf:iwf + nw])
             # find the indices sorting the eigenvalues in ascending order
             eigsrt = eigval[0:nw].argsort()
             # order eigenvalues and eigenvectors and save in a list
@@ -460,15 +486,20 @@ class Wannier90Converter(ConverterTools):
 
         for ish in range(n_sh):
             try:
-                # build rotation matrices by combining the unitary transformations that diagonalize H(0)
+                # build rotation matrices by combining the unitary
-                rot_mat[ish] = numpy.dot(eigvec_lst[ish],eigvec_lst[sh_map[ish]].conjugate().transpose())
+                # transformations that diagonalize H(0)
+                rot_mat[ish] = numpy.dot(eigvec_lst[ish], eigvec_lst[
+                                         sh_map[ish]].conjugate().transpose())
             except ValueError:
-                mpi.report("Global-to-local rotation matrices cannot be constructed!")
+                mpi.report(
+                    "Global-to-local rotation matrices cannot be constructed!")
 
             istatus = 1
-            # check that eigenvalues are the same (within accuracy) for equivalent shells
+            # check that eigenvalues are the same (within accuracy) for
+            # equivalent shells
             if not numpy.allclose(eigval_lst[ish], eigval_lst[sh_map[ish]], atol=self._w90zero, rtol=1.e-15):
-                mpi.report("ERROR: eigenvalue mismatch between equivalent shells! %d"%ish)
+                mpi.report(
+                    "ERROR: eigenvalue mismatch between equivalent shells! %d" % ish)
                 eigval_diff = eigval_lst[ish] - eigval_lst[sh_map[ish]]
                 mpi.report("Eigenvalue difference: " + format(eigval_diff))
                 istatus = 0
@@ -477,8 +508,6 @@ class Wannier90Converter(ConverterTools):
 
         return istatus, rot_mat
 
-
-
     def kmesh_build(self, msize=None, mmode=0):
         """
         Method for the generation of the k-point mesh.
@@ -502,27 +531,24 @@ class Wannier90Converter(ConverterTools):
 
         """
 
-        if mmode == 0:
+        if mmode != 0:
+            raise ValueError("Mesh generation mode not supported: %s" % mmode)
+
         # a regular mesh including Gamma point
-            nkpt = msize[0] * msize[1] * msize[2]      # total number of k-points
+        # total number of k-points
+        nkpt = msize[0] * msize[1] * msize[2]
         kmesh = numpy.zeros((nkpt, 3), dtype=float)
         ii = 0
-            for ix in range(msize[0]):
+        for ix, iy, iz in product(range(msize[0]), range(msize[1]), range(msize[2])):
-                for iy in range(msize[1]):
-                    for iz in range(msize[2]):
             kmesh[ii, :] = [float(ix) / msize[0], float(iy) / msize[1], float(iz) / msize[2]]
             ii += 1
         # weight is equal for all k-points because wannier90 uses uniform grid on whole BZ
         # (normalization is always 1 and takes into account spin degeneracy)
         wk = numpy.ones([nkpt], dtype=float) / float(nkpt)
-        else:
-            raise ValueError("Mesh generation mode not supported: %s"%mmode)
 
         return nkpt, kmesh, wk
 
-
+    def fourier_ham(self, norb, h_of_r):
-
-    def fourierham(self, norb, h_of_r):
         """
         Method for obtaining H(k) from H(R) via Fourier transform
         The R vectors and k-point mesh are read from global module variables
@@ -541,16 +567,12 @@ class Wannier90Converter(ConverterTools):
 
         """
 
-        imag = 1j
         twopi = 2 * numpy.pi
-
         h_of_k = [numpy.zeros((norb, norb), dtype=numpy.complex_) for ik in range(self.n_k)]
-        for ik in range(self.n_k):
         ridx = numpy.array(range(self.nrpt))
-            for ir in ridx:
+        for ik, ir in product(range(self.n_k), ridx):
             rdotk = twopi * numpy.dot(self.k_mesh[ik], self.rvec[ir])
-                factor = (math.cos(rdotk) + imag * math.sin(rdotk)) / float(self.rdeg[ir])
+            factor = (math.cos(rdotk) + 1j * math.sin(rdotk)) / float(self.rdeg[ir])
             h_of_k[ik][:, :] += factor * h_of_r[ir][:, :]
 
         return h_of_k
-