Adding parsers for wavefunctions and U matrices

A parser for the _u.mat and _u_dis.mat files has been added.
Moreover, also a parser for wavefunctions (UNKNNNNN.N files)
has been added, but also in the (non-default) case where they
are printed in formatted format, for now.

Tests are added (skipping U matrices for Wannier 2.x where this
was not implemented) and CLI is expanded as well with two new commands.

docs/api.md

@@ -6,6 +6,8 @@
 ::: wannier90io.read_amn
 ::: wannier90io.write_amn
 ::: wannier90io.read_chk
+::: wannier90io.read_u
+::: wannier90io.read_unk_formatted
 ::: wannier90io.read_eig
 ::: wannier90io.write_eig
 ::: wannier90io.read_mmn

src/wannier90io/__init__.py

@@ -4,5 +4,7 @@
 from ._eig import *
 from ._mmn import *
 from ._nnkp import *
+from ._u import *
+from ._unk import *
 from ._win import *
 from ._wout import *

src/wannier90io/__main__.py

@@ -119,6 +119,26 @@ def info_chk(args):
     print(chk['have_disentangled'])
 
 
+def info_u(args):
+    with args.file:
+        kpoints, u_matrices = w90io.read_u(args.file)
+
+    print(f'Nk = {u_matrices.shape[0]}')
+    print(f'Nb = {u_matrices.shape[1]}')
+    print(f'Nw = {u_matrices.shape[2]}')
+
+
+def info_unk_formatted(args):
+    with args.file:
+        ik, wvfn = w90io.read_unk_formatted(args.file)
+
+    print(f'ik  = {ik}')
+    print(f'ngx = {wvfn.shape[0]}')
+    print(f'ngy = {wvfn.shape[1]}')
+    print(f'ngz = {wvfn.shape[2]}')
+    print(f'Nb  = {wvfn.shape[3]}')
+
+
 def main():
     parser = argparse.ArgumentParser()
     subparsers = parser.add_subparsers(dest='subparser', required=True)
@@ -157,6 +177,12 @@ def main():
     #
     parser_info_chk = subparsers.add_parser('info-chk', parents=[parser_common])
     parser_info_chk.set_defaults(func=info_chk)
+    #
+    parser_info_u = subparsers.add_parser('info-u', parents=[parser_common])
+    parser_info_u.set_defaults(func=info_u)
+    #
+    parser_info_unk_formatted = subparsers.add_parser('info-unk-formatted', parents=[parser_common])
+    parser_info_unk_formatted.set_defaults(func=info_unk_formatted)
 
     args = parser.parse_args()
     args.func(args)

src/wannier90io/_u.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+import typing
+
+import numpy as np
+
+# TODO: implement also read_u_dis
+__all__ = ['read_u']
+
+
+def read_u(stream: typing.TextIO) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Read unitary matrix file (seedname_u.mat) or the rectangular U_dis matrix
+    file (seedname_u_dis.mat).
+
+    Note: for the _u.mat file, num_bands == num_wann.
+
+    Arguments:
+        stream: a file-like stream
+
+    Returns:
+        kpoint coordinates in fractional coordinates (num_kpts, 3)
+        U matrix U(k) or U_dis(k) (num_kpts, num_bands, num_wann)
+
+
+    """
+    stream.readline()   # header
+
+    [nkpt, num_wann, num_bands] = np.fromstring(stream.readline(), sep=' ', dtype=int)
+    u_matrices = np.zeros((nkpt, num_bands, num_wann), dtype=complex)
+    kpoints = []
+
+    for ikpt in range(nkpt):
+        empty = stream.readline()   # header
+        assert not empty.strip(), f"Expected empty line but found instead: '{empty}'"
+
+        kpoint = np.fromstring(stream.readline(), sep=' ', dtype=float)
+        assert len(kpoint) == 3
+        kpoints.append(kpoint)
+        u_matrices[ikpt, :, :] = np.loadtxt(stream, max_rows=(num_wann * num_bands)).view(complex).reshape((num_bands, num_wann), order='F')
+
+    return np.array(kpoints), u_matrices

src/wannier90io/_unk.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+import typing
+
+import numpy as np
+
+# TODO: implement also read_unk_unformatted (that is the default)
+__all__ = ['read_unk_formatted']
+
+
+def read_unk_formatted(stream: typing.TextIO) -> tuple[int, np.ndarray]:
+    """
+    Read wavefunction files (UNKnnnnn.n files) in formatted format.
+
+    Note that the UNK files must have been created using the `wvfn_formatted`
+    option set to True in the interface code (e.g. pw2wannier90.x for the
+    Quantum ESPRESSO interface). Note that this is *not* the default, however
+    for reading into an external code, this is recommended for portability.
+
+    Note:
+       for now only works in the non-spinor case.
+       Spinor case still to be implemented.
+
+    Arguments:
+        stream: a file-like stream
+
+    Returns:
+        k-point index ik (integer)
+        complex wavefunction (ngx, ngy, ngz, Nb)
+
+
+    """
+    [ngx, ngy, ngz, ik, nbnd] = np.fromstring(stream.readline(), sep=' ', dtype=int)
+
+    wvfn = np.zeros((ngx, ngy, ngz, nbnd), dtype=complex)
+
+    for ibnd in range(nbnd):
+        wvfn[:, :, :, ibnd] = np.loadtxt(stream, max_rows=(ngx * ngy * ngz)).view(complex).reshape((ngx, ngy, ngz), order='F')
+
+    return (ik, wvfn)

tests/fixtures/fixtures.mk

@@ -20,6 +20,15 @@ else
 	WRITE_HR = "write_hr=true"
 endif
 
+# Only implemented in version 3.x
+ifeq ($(WANNIER90_VERSION), 2.0.1)
+	WRITE_U_MATRICES = ""
+else ($(WANNIER90_VERSION), 2.1)
+	WRITE_U_MATRICES = ""
+else
+	WRITE_U_MATRICES = "write_u_matrices=true"
+endif
+
 define modify_win
 	echo $(WRITE_HR) >> wannier.win
 	echo "write_xyz=true" >> wannier.win
@@ -78,6 +87,7 @@ run-example04:
 	$(W90) -pp wannier
 	$(call modify_win)
 	echo "geninterp_alsofirstder=true" >> wannier.win
+	echo $(WRITE_U_MATRICES) >> wannier.win
 	$(W90) wannier
 	$(W90CHK2CHK) -export wannier
 	echo "" >> wannier_geninterp.kpt

tests/test_u.py

@@ -0,0 +1,32 @@
+import pathlib
+
+import numpy as np
+import pytest
+
+import wannier90io as w90io
+
+
+@pytest.mark.parametrize('example', ['example04'])
+def test_read_u(wannier90, example):
+    with open(pathlib.Path(wannier90)/f'examples/{example}/wannier.win', 'r') as fh:
+        win = w90io.parse_win_raw(fh.read())
+    nkpt = len(win['kpoints']['kpoints'])
+    num_bands = win['parameters']['num_bands']
+    num_wann = win['parameters']['num_wann']
+
+    if 'write_u_matrices' not in win['parameters']:
+        # an old 2.x version, didn't have write_u_matrices, so skip this test
+        # The fixture-making script will already avoid to add the flag in W90 2.x
+        pytest.skip("U matrix printing not implemented in Wannier90 2.x")
+        return
+
+    with open(pathlib.Path(wannier90)/f'examples/{example}/wannier_u.mat', 'r') as fh:
+        (kpts, u_matrices) = w90io.read_u(fh)
+    assert np.allclose(kpts, np.asarray(win['kpoints']['kpoints'], dtype=float))
+    # Note that the U matrix is num_wann x num_wann, not num_bands x num_wann!
+    assert np.allclose(u_matrices.shape, [nkpt, num_wann, num_wann])
+
+    with open(pathlib.Path(wannier90)/f'examples/{example}/wannier_u_dis.mat', 'r') as fh:
+        (kpts, u_dis_matrices) = w90io.read_u(fh)
+    assert np.allclose(kpts, np.asarray(win['kpoints']['kpoints'], dtype=float))
+    assert np.allclose(u_dis_matrices.shape, [nkpt, num_bands, num_wann])

tests/test_unk_formatted.py

@@ -0,0 +1,17 @@
+import pathlib
+
+import numpy as np
+import pytest
+
+import wannier90io as w90io
+
+
+@pytest.mark.parametrize('example', ['example01'])
+def test_read_unk_formatted(wannier90, example):
+    for ik in range(1, 8+1):
+        with open(pathlib.Path(wannier90)/f'examples/{example}/UNK0000{ik}.1', 'r') as fh:
+            (ik_parsed, wvfn) = w90io.read_unk_formatted(fh)
+        assert ik == ik_parsed
+
+        ngx, ngy, ngz, nbnd = 20, 20, 20, 4
+        assert np.allclose(wvfn.shape, [ngx, ngy, ngz, nbnd])