add the package

setup.py

@@ -0,0 +1,28 @@
+from __future__ import absolute_import, division, print_function
+
+from setuptools import setup, find_packages
+
+
+setup(
+    name='wannier90-utils',
+    version='0.1.0',
+    description='Wannier90 utility library',
+    author='Jamal I. Mustafa',
+    author_email='jimustafa@gmail.com',
+    license='BSD',
+    classifiers=[
+        'Intended Audience :: Science/Research',
+        'License :: OSI Approved :: BSD License',
+        'Natural Language :: English',
+        'Operating System :: POSIX :: Linux',
+        'Programming Language :: Python :: 2.7',
+        'Topic :: Scientific/Engineering :: Physics',
+        'Topic :: Software Development :: Libraries :: Python Modules',
+    ],
+    packages=find_packages(exclude=['docs', 'tests']),
+    install_requires=[
+        'numpy',
+        'scipy',
+    ],
+    tests_require=['pytest'],
+)

w90utils/__init__.py

@@ -0,0 +1,8 @@
+"""Wannier90 utility library"""
+from __future__ import absolute_import, division, print_function
+
+from . import io
+from . import sprd
+from ._amn import expand_amn
+from ._mmn import rotate_mmn
+from ._utils import unitarize

w90utils/_amn.py

@@ -0,0 +1,47 @@
+from __future__ import absolute_import, division, print_function
+
+import numpy as np
+
+
+def expand_amn(a, kpoints, idx, Rvectors, nproj_atom=None):
+    """
+    Expand the projections matrix by translations of the orbitals
+
+    Parameters
+    ----------
+    a : ndarray, shape (nkpts, nbnds, nproj)
+    kpoints : ndarray, shape (nkpts, 3)
+    idx : ndarray
+        indices of translated orbitals
+    Rvectors: ndarray
+        translation vectors for the orbitals
+    nproj_atom: ndarray, optional
+        number of projections on each atom, with idx and Rvectors now describing
+        atoms instead of orbitals
+
+    """
+    assert len(Rvectors) == len(idx)
+
+    if nproj_atom is not None:
+        assert len(nproj_atom) == len(idx)
+        idx_new = []
+        Rvectors_new = []
+        for iatom, i in enumerate(idx):
+            offset = np.sum(nproj_atom[:i])
+            for j in range(nproj_atom[i]):
+                idx_new.append(offset+j)
+                Rvectors_new.append(Rvectors[iatom])
+
+        idx = idx_new
+        Rvectors = Rvectors_new
+
+    nkpts, nbnds, nproj = a.shape
+
+    a1 = np.zeros((nkpts, nbnds, len(idx)), dtype=complex)
+
+    k_dot_R = np.einsum('ki,ri->kr', kpoints, Rvectors)
+    exp_factors = np.exp(-1j * 2*np.pi * k_dot_R)
+
+    a1 = a[:, :, idx] * exp_factors[:, np.newaxis, :]
+
+    return a1

w90utils/_mmn.py

@@ -0,0 +1,65 @@
+from __future__ import division, print_function
+
+import numpy as np
+
+
+def rotate_mmn(mmn, umn, kpb_kidx, window=None):
+    """
+    Rotate the overlap matrices according to
+    :math:`U^{(\mathbf{k})\dagger}M^{(\mathbf{k},\mathbf{b})}U^{(\mathbf{k}+\mathbf{b})}`
+
+    """
+    (nkpts, nntot, nbnds, nbnds) = mmn.shape
+    nproj = umn[0].shape[1]
+
+    mmn_rotated = np.empty((nkpts, nntot, nproj, nproj), dtype=complex)
+
+    if window is not None:
+        for ikpt in range(nkpts):
+            for inn in range(nntot):
+                ikpb = kpb_kidx[ikpt][inn]
+                mmn_rotated[ikpt][inn] = (
+                    np.dot(
+                        np.dot(
+                            umn[ikpt].conj().T, mmn[ikpt][inn][window[ikpt]][:, window[ikpb]],
+                        ),
+                        umn[ikpb]
+                    )
+                )
+    else:
+        for ikpt in range(nkpts):
+            for inn in range(nntot):
+                ikpb = kpb_kidx[ikpt][inn]
+                mmn_rotated[ikpt][inn] = np.dot(np.dot(umn[ikpt].conj().T, mmn[ikpt][inn]), umn[ikpb])
+
+    return mmn_rotated
+
+
+# def change_gauge_k(m, u, setup_file):
+#     (nkpts, nntot, nbnds, nbnds) = m.shape
+#     nproj = u[0].shape[1]
+
+#     m_rotated = np.zeros((nkpts, nntot, nproj, nbnds), dtype=complex)
+
+#     for ikpt in range(nkpts):
+#         for inn in range(nntot):
+#             m_rotated[ikpt][inn] = np.dot(u[ikpt].conj().T, m[ikpt][inn])
+
+#     return m_rotated
+
+
+# def change_gauge_kpb(m, u, setup_file, kpb_kidx=None):
+#     (nkpts, nntot, nbnds, nbnds) = m.shape
+#     nproj = u[0].shape[1]
+
+#     if kpb_kidx is None:
+#         kpb_kidx = setup_file.kpb_kidx
+
+#     m_rotated = np.zeros((nkpts, nntot, nbnds, nproj), dtype=complex)
+
+#     for ikpt in range(nkpts):
+#         for inn in range(nntot):
+#             ikpb = kpb_kidx[ikpt][inn]
+#             m_rotated[ikpt][inn] = np.dot(m[ikpt][inn], u[ikpb])
+
+#     return m_rotated

w90utils/_utils.py

@@ -0,0 +1,8 @@
+from __future__ import absolute_import, division, print_function
+
+import numpy as np
+
+
+def unitarize(a):
+    u, _, v = np.linalg.svd(a, full_matrices=False)
+    return np.einsum('...ik,...kj->...ij', u, v)

w90utils/io/__init__.py

@@ -0,0 +1,74 @@
+"""Wannier90 I/O library"""
+from __future__ import absolute_import, division, print_function
+import collections
+
+from ._amn import *
+from ._bands import *
+from ._chk import *
+from ._eig import *
+from ._hr import *
+from ._mmn import *
+from ._unk import *
+from . import nnkp
+from . import postw90
+from . import utils
+from . import win
+from . import wout
+from . import _utils
+
+
+Wannier90Data = collections.namedtuple(
+    'Wannier90Data',
+    [
+        'dlv', 'rlv',
+        'amn', 'mmn', 'eig',
+        'kpoints', 'kpb_kidx', 'kpb_g',
+        'bv', 'bw',
+        'length_unit', 'energy_unit'
+    ])
+
+
+def read_data(seedname='wannier', **kwargs):
+    """
+    Read all Wannier90 input data files from the current directory.
+
+    Parameters
+    ----------
+    seedname : str, optional
+        seedname for the Wannier90 files, the default is "wannier"
+
+    """
+    dlv = kwargs.get('dlv', nnkp.read_dlv(seedname+'.nnkp', units='angstrom'))
+    rlv = kwargs.get('rlv', nnkp.read_rlv(seedname+'.nnkp', units='angstrom'))
+    try:
+        amn = kwargs['amn']
+    except KeyError:
+        amn = read_amn(seedname+'.amn')
+    try:
+        mmn = kwargs['mmn']
+    except KeyError:
+        mmn = read_mmn(seedname+'.mmn')
+    try:
+        eig = kwargs['eig']
+    except KeyError:
+        eig = read_eig(seedname+'.eig')
+    kpoints = kwargs.get('kpoints', nnkp.read_kpoints(seedname+'.nnkp'))
+    kpb_kidx = kwargs.get('kpb_kidx', nnkp.read_nnkpts(seedname+'.nnkp')[0])
+    kpb_g = kwargs.get('kpb_kidx', nnkp.read_nnkpts(seedname+'.nnkp')[1])
+    bv = kwargs.get('bv', nnkp.read_bvectors(seedname+'.nnkp', units='angstrom'))
+    bw = kwargs.get('bw', _utils.bweights(bv))
+
+    return Wannier90Data(
+        dlv=dlv,
+        rlv=rlv,
+        amn=amn,
+        mmn=mmn,
+        eig=eig,
+        kpoints=kpoints,
+        kpb_kidx=kpb_kidx,
+        kpb_g=kpb_g,
+        bv=bv,
+        bw=bw,
+        length_unit='angstrom',
+        energy_unit='eV',
+        )

w90utils/io/_amn.py

@@ -0,0 +1,52 @@
+from __future__ import absolute_import, division, print_function
+
+import numpy as np
+
+
+__all__ = ['read_amn', 'write_amn']
+
+
+def read_amn(fname):
+    """
+    Read AMN file.
+
+    Parameters
+    ----------
+    fname : str
+
+    Returns
+    -------
+    amn : ndarray, shape (nkpts, nbnds, nproj)
+
+    """
+    with open(fname, 'r') as f:
+        f.readline()    # header
+        [nbnds, nkpts, nproj] = map(int, f.readline().split())
+        data_str = f.read()
+
+    raw_data = np.fromstring(data_str, sep='\n').reshape((nkpts*nbnds*nproj, 5))
+    amn = raw_data[:, 3] + 1j*raw_data[:, 4]
+    amn = np.copy(np.transpose(amn.reshape((nbnds, nproj, nkpts), order='F'), axes=(2, 0, 1)), order='C')
+
+    return amn
+
+
+def write_amn(fname, amn, header='HEADER'):
+    r"""
+    Write :math:`A^{(\mathbf{k})}_{mn}` to AMN file.
+
+    Parameters
+    ----------
+    fname : str
+    amn : ndarray, shape (nkpts, nbnds, nproj)
+    header : str
+
+    """
+    (nkpts, nbnds, nproj) = amn.shape
+    indices = np.mgrid[:nbnds, :nproj, :nkpts].reshape((3, -1), order='F') + 1
+    amn = np.transpose(amn, axes=(1, 2, 0)).flatten(order='F').view(float).reshape((-1, 2))
+    data_out = np.column_stack((indices.transpose(), amn))
+    with open(fname, 'w') as f:
+        print(header, file=f)
+        print('%13d%13d%13d' % (nbnds, nkpts, nproj), file=f)
+        np.savetxt(f, data_out, fmt='%5d%5d%5d%18.12f%18.12f')

w90utils/io/_bands.py

@@ -0,0 +1,34 @@
+from __future__ import absolute_import, division, print_function
+
+import numpy as np
+
+
+__all__ = ['read_kpoints', 'read_bands']
+
+
+def read_kpoints(fname):
+    raw_data = np.loadtxt(fname, skiprows=1)
+    kpoints = raw_data[:, (0, 1, 2)]
+    kweights = raw_data[:, 3]
+
+    return kpoints
+
+
+def read_bands(fname):
+    nkpts = None
+    with open(fname, 'r') as f:
+        for iln, line in enumerate(f):
+            if len(line.strip()) == 0:
+                nkpts = iln
+                break
+
+    if nkpts is None:
+        raise Exception
+
+    print(nkpts)
+
+    raw_data = np.loadtxt(fname)
+    nbnds = len(raw_data) // nkpts
+    bands = raw_data[:, 1].reshape((nbnds, nkpts)).transpose()
+
+    return bands

w90utils/io/_chk.py

@@ -0,0 +1,53 @@
+from __future__ import absolute_import, division, print_function
+import contextlib
+import cStringIO as StringIO
+
+import numpy as np
+from scipy.io import FortranFile
+
+
+__all__ = ['CheckpointIO']
+
+
+class CheckpointIO(object):
+    def __init__(self, fname=None, auto_read=True):
+        if fname and auto_read:
+            self.from_file(fname)
+
+    def from_file(self, fname):
+        with FortranFile(fname, 'r') as f:
+            self.header = ''.join(f.read_record('c'))
+            self.nbnds = f.read_ints()[0]
+            self.nbnds_excl = f.read_ints()[0]
+            self.bands_excl = f.read_ints()
+            self.dlv = f.read_reals().reshape((3, 3), order='F')
+            self.rlv = f.read_reals().reshape((3, 3), order='F')
+            self.nkpts = f.read_ints()[0]
+            self.grid_dims = f.read_ints()
+            self.kpoints = f.read_reals().reshape((-1, 3))
+            self.nntot = f.read_ints()[0]
+            self.nwann = f.read_ints()[0]
+            self.chkpt = f.read_record('c')
+            self.disentanglement = bool(f.read_ints()[0])
+
+            if self.disentanglement:
+                self.omega_invariant = f.read_reals()[0]
+                self.windows = f.read_ints().reshape((self.nbnds, self.nkpts), order='F').transpose
+                f.read_ints()
+                self.umat_opt = np.transpose(f.read_reals().view(complex).reshape((self.nbnds, self.nwann, self.nkpts), order='F'), axes=(2, 0, 1))
+
+            self.umat = np.transpose(f.read_reals().view(complex).reshape((self.nwann, self.nwann, self.nkpts), order='F'), axes=(2, 0, 1))
+            self.mmat = np.transpose(f.read_reals().view(complex).reshape((self.nwann, self.nwann, self.nntot, self.nkpts), order='F'), axes=(3, 2, 0, 1))
+            self.wannier_centers = f.read_reals().reshape((-1, 3))
+            self.wannier_spreads = f.read_reals()
+
+    def __str__(self):
+        with contextlib.closing(StringIO.StringIO()) as sio:
+            print(self.header, file=sio)
+            print(self.nbnds, file=sio)
+            print(self.nbnds_excl, file=sio)
+            print(self.dlv, file=sio)
+            print(self.rlv, file=sio)
+            print(self.nkpts, file=sio)
+            s = sio.getvalue()
+        return s