Source code for wannierberri.w90files.spn

import numpy as np
from .w90file import W90_file
from ..io import FortranFileR


[docs] class SPN(W90_file): """ SPN.data[ik, m, n, ipol] = <u_{m,k}|S_ipol|u_{n,k}> """ def __init__(self, seedname="wannier90", **kwargs): self.npz_tags = ["data"] super().__init__(seedname=seedname, ext="spn", **kwargs)
[docs] def from_w90_file(self, seedname='wannier90', formatted=False): print("----------\n SPN \n---------\n") if formatted: f_spn_in = open(seedname + ".spn", 'r') SPNheader = f_spn_in.readline().strip() nbnd, NK = (int(x) for x in f_spn_in.readline().split()) else: f_spn_in = FortranFileR(seedname + ".spn") SPNheader = (f_spn_in.read_record(dtype='c')) nbnd, NK = f_spn_in.read_record(dtype=np.int32) SPNheader = "".join(a.decode('ascii') for a in SPNheader) print(f"reading {seedname}.spn : {SPNheader}") indm, indn = np.tril_indices(nbnd) self.data = np.zeros((NK, nbnd, nbnd, 3), dtype=complex) for ik in range(NK): A = np.zeros((3, nbnd, nbnd), dtype=complex) if formatted: tmp = np.array([f_spn_in.readline().split() for i in range(3 * nbnd * (nbnd + 1) // 2)], dtype=float) tmp = tmp[:, 0] + 1.j * tmp[:, 1] else: tmp = f_spn_in.read_record(dtype=np.complex128) A[:, indn, indm] = tmp.reshape(3, nbnd * (nbnd + 1) // 2, order='F') check = np.einsum('ijj->', np.abs(A.imag)) A[:, indm, indn] = A[:, indn, indm].conj() if check > 1e-10: raise RuntimeError(f"REAL DIAG CHECK FAILED : {check}") self.data[ik] = A.transpose(1, 2, 0) print("----------\n SPN OK \n---------\n")
[docs] def select_bands(self, selected_bands): if selected_bands is not None: self.data = self.data[:, selected_bands, :, :][:, :, selected_bands, :]