Source code for wannierberri.w90files.amn

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()