from datetime import datetime
import multiprocessing
import numpy as np
from .utility import str2arraymmn
from .w90file import W90_file
[docs]
class AMN(W90_file):
"""
Class to store the projection of the wavefunctions on the initial Wannier functions
AMN.data[ik, ib, iw] = <u_{i,k}|w_{i,w}>
Parameters
----------
seedname : str
the prefix of the file (including relative/absolute path, but not including the extension `.amn`)
npar : int
the number of parallel processes to be used for reading
Notes
-----
Attributes
----------
NB : int
number of bands
NW : int
number of Wannier functions
NK : int
number of k-points
data : numpy.ndarray( (NK, NB, NW), dtype=complex)
the data projections
"""
@property
def NB(self):
return self.data.shape[1]
def apply_window(self, selected_bands):
print(f"apply_window amn, selected_bands={selected_bands}")
if selected_bands is not None:
self.data = self.data[:, selected_bands, :]
@property
def NW(self):
return self.data.shape[2]
def __init__(self, seedname="wannier90", npar=multiprocessing.cpu_count(), **kwargs):
self.npz_tags = ["data"]
super().__init__(seedname, ext="amn", npar=npar, **kwargs)
[docs]
def from_w90_file(self, seedname, npar):
f_amn_in = open(seedname + ".amn", "r").readlines()
print(f"reading {seedname}.amn: " + f_amn_in[0].strip())
s = f_amn_in[1]
NB, NK, NW = np.array(s.split(), dtype=int)
block = NW * NB
allmmn = (f_amn_in[2 + j * block:2 + (j + 1) * block] for j in range(NK))
p = multiprocessing.Pool(npar)
self.data = np.array(p.map(str2arraymmn, allmmn)).reshape((NK, NW, NB)).transpose(0, 2, 1)
def to_w90_file(self, seedname):
f_amn_out = open(seedname + ".amn", "w")
f_amn_out.write(f"created by WannierBerri on {datetime.now()} \n")
print(f"writing {seedname}.amn: ")
f_amn_out.write(f" {self.NB:3d} {self.NK:3d} {self.NW:3d} \n")
for ik in range(self.NK):
for iw in range(self.NW):
for ib in range(self.NB):
f_amn_out.write(f"{ib+1:4d} {iw+1:4d} {ik+1:4d} {self.data[ik, ib, iw].real:17.12f} {self.data[ik, ib, iw].imag:17.12f}\n")
[docs]
def get_disentangled(self, v_left, v_right):
print(f"v shape {v_left.shape} {v_right.shape} , amn shape {self.data.shape} ")
data = np.einsum("klm,kmn->kln", v_left, self.data)
print(f"shape of data {data.shape} , old {self.data.shape}")
return self.__class__(data=data)
# def write(self, seedname, comment="written by WannierBerri"):
# comment = comment.strip()
# f_amn_out = open(seedname + ".amn", "w")
# print(f"writing {seedname}.amn: " + comment + "\n")
# f_amn_out.write(comment + "\n")
# f_amn_out.write(f" {self.NB:3d} {self.NK:3d} {self.NW:3d} \n")
# for ik in range(self.NK):
# f_amn_out.write("".join(" {:4d} {:4d} {:4d} {:17.12f} {:17.12f}\n".format(
# ib + 1, iw + 1, ik + 1, self.data[ik, ib, iw].real, self.data[ik, ib, iw].imag)
# for iw in range(self.NW) for ib in range(self.NB)))
# f_amn_out.close()