from itertools import islice
import multiprocessing
from time import time
import numpy as np
from .utility import convert
from .w90file import W90_file
[docs]
class MMN(W90_file):
"""
class to store overlaps between Bloch functions at neighbouring k-points
the MMN file of wannier90
MMN.data[ik, ib, m, n] = <u_{m,k}|u_{n,k+b}>
Parameters
----------
seedname : str
the prefix of the file (including relative/absolute path, but not including the extension `.mmn`)
npar : int
the number of parallel processes to be used for reading the file
Attributes
----------
data : np.ndarray(shape=(NK, NNB, NB, NB), dtype=complex)
the overlap matrix elements between the Wavefunctions at neighbouring k-points
neighbours : np.ndarray(shape=(NK, NNB), dtype=int)
the indices of the neighbouring k-points
G : np.ndarray(shape=(NK, NNB, 3), dtype=int)
the reciprocal lattice vectors connecting the k-points
"""
@property
def n_neighb(self):
"""
number of nearest neighbours indices
"""
return 1
def __init__(self, seedname="wannier90", npar=multiprocessing.cpu_count(), **kwargs):
self.npz_tags = ["data", "neighbours", "G"]
super().__init__(seedname, ext="mmn", npar=npar, **kwargs)
[docs]
def from_w90_file(self, seedname, npar):
t0 = time()
f_mmn_in = open(seedname + ".mmn", "r")
f_mmn_in.readline()
NB, NK, NNB = np.array(f_mmn_in.readline().split(), dtype=int)
block = 1 + NB * NB
data = []
headstring = []
mult = 4
# FIXME: npar = 0 does not work
if npar > 0:
pool = multiprocessing.Pool(npar)
for j in range(0, NNB * NK, npar * mult):
x = list(islice(f_mmn_in, int(block * npar * mult)))
if len(x) == 0:
break
headstring += x[::block]
y = [x[i * block + 1:(i + 1) * block] for i in range(npar * mult) if (i + 1) * block <= len(x)]
if npar > 0:
data += pool.map(convert, y)
else:
data += [convert(z) for z in y]
if npar > 0:
pool.close()
pool.join()
f_mmn_in.close()
t1 = time()
data = [d[:, 0] + 1j * d[:, 1] for d in data]
self.data = np.array(data).reshape(NK, NNB, NB, NB).transpose((0, 1, 3, 2))
headstring = np.array([s.split() for s in headstring], dtype=int).reshape(NK, NNB, 5)
assert np.all(headstring[:, :, 0] - 1 == np.arange(NK)[:, None])
self.neighbours = headstring[:, :, 1] - 1
self.G = headstring[:, :, 2:]
t2 = time()
print(f"Time for MMN.__init__() : {t2 - t0} , read : {t1 - t0} , headstring {t2 - t1}")
def to_w90_file(self, seedname):
f_mmn_out = open(seedname + ".mmn", "w")
f_mmn_out.write("MMN file\n")
f_mmn_out.write(f"{self.NB} {self.NK} {self.NNB}\n")
for ik in range(self.NK):
for ib in range(self.NNB):
f_mmn_out.write(f"{ik + 1} {self.neighbours[ik, ib] + 1} {' '.join(map(str, self.G[ik, ib]))}\n")
for m in range(self.NB):
for n in range(self.NB):
f_mmn_out.write(f"{self.data[ik, ib, n, m].real} {self.data[ik, ib, n, m].imag}\n")
f_mmn_out.close()
def apply_window(self, selected_bands):
if selected_bands is not None:
self.data = self.data[:, :, selected_bands, :][:, :, :, selected_bands]
[docs]
def get_disentangled(self, v_left, v_right):
"""
Reduce number of bands
Parameters
----------
v_matrix : np.ndarray(NB,num_wann)
the matrix of column vectors defining the Wannier gauge
v_matrix_dagger : np.ndarray(num_wann,NB)
the Hermitian conjugate of `v_matrix`
Returns
-------
`~wannierberri.system.w90_files.MMN`
the disentangled MMN object
"""
print(f"v shape {v_left.shape} {v_right.shape}")
data = np.zeros((self.NK, self.NNB, v_right.shape[2], v_right.shape[2]), dtype=complex)
print(f"shape of data {data.shape} , old {self.data.shape}")
for ik in range(self.NK):
for ib, iknb in enumerate(self.neighbours[ik]):
data[ik, ib] = v_left[ik] @ self.data[ik, ib] @ v_right[iknb]
return self.__class__(data=data)
def set_bk(self, kpt_latt, mp_grid, recip_lattice, kmesh_tol=1e-7, bk_complete_tol=1e-5):
try:
self.bk_cart
self.wk
self.wk_unique
self.bk_latt_unique
self.bk_cart_unique
self.ib_unique_map
self.ib_unique_map_inverse
self.neighbours_unique
return
except AttributeError:
bk_latt = np.array(
np.round(
[
(kpt_latt[nbrs] - kpt_latt + G) * mp_grid[None, :]
for nbrs, G in zip(self.neighbours.T, self.G.transpose(1, 0, 2))
]).transpose(1, 0, 2),
dtype=int)
bk_latt_unique = np.array([b for b in set(tuple(bk) for bk in bk_latt.reshape(-1, 3))], dtype=int)
assert len(bk_latt_unique) == self.NNB
bk_cart_unique = bk_latt_unique.dot(recip_lattice / mp_grid[:, None])
bk_cart_unique_length = np.linalg.norm(bk_cart_unique, axis=1)
srt = np.argsort(bk_cart_unique_length)
bk_latt_unique = bk_latt_unique[srt]
bk_cart_unique = bk_cart_unique[srt]
bk_cart_unique_length = bk_cart_unique_length[srt]
brd = [
0,
] + list(np.where(bk_cart_unique_length[1:] - bk_cart_unique_length[:-1] > kmesh_tol)[0] + 1) + [
self.NNB,
]
shell_mat = np.array([bk_cart_unique[b1:b2].T.dot(bk_cart_unique[b1:b2]) for b1, b2 in zip(brd, brd[1:])])
shell_mat_line = shell_mat.reshape(-1, 9)
u, s, v = np.linalg.svd(shell_mat_line, full_matrices=False)
s = 1. / s
weight_shell = np.eye(3).reshape(1, -1).dot(v.T.dot(np.diag(s)).dot(u.T)).reshape(-1)
check_eye = sum(w * m for w, m in zip(weight_shell, shell_mat))
tol = np.linalg.norm(check_eye - np.eye(3))
if tol > bk_complete_tol:
raise RuntimeError(
f"Error while determining shell weights. the following matrix :\n {check_eye} \n"
f"failed to be identity by an error of {tol}. Further debug information : \n"
f"bk_latt_unique={bk_latt_unique} \n bk_cart_unique={bk_cart_unique} \n"
f"bk_cart_unique_length={bk_cart_unique_length}\n shell_mat={shell_mat}\n"
f"weight_shell={weight_shell}\n")
weight = np.array([w for w, b1, b2 in zip(weight_shell, brd, brd[1:]) for i in range(b1, b2)])
weight_dict = {tuple(bk): w for bk, w in zip(bk_latt_unique, weight)}
bk_cart_dict = {tuple(bk): bkcart for bk, bkcart in zip(bk_latt_unique, bk_cart_unique)}
self.bk_cart = np.array([[bk_cart_dict[tuple(bkl)] for bkl in bklk] for bklk in bk_latt])
self.wk = np.array([[weight_dict[tuple(bkl)] for bkl in bklk] for bklk in bk_latt])
#############
### Oscar ###
###################################################################
# Wannier90 provides a list of nearest-neighbor vectors b for every
# q point. For Jae-Mo's finite-difference scheme it is convenient
# to evaluate the Fourier transform of the matrix elements in the
# original ab-initio mesh before performing the sum over
# nearest-neighbor vectors. This requires defining a mapping from
# any pair {q,b} to a unique list of b vectors that is independent
# of q.
bk_latt = np.rint((self.bk_cart @ np.linalg.inv(recip_lattice)) * mp_grid[None, None, :]).astype(int)
bk_latt_unique = np.unique(bk_latt.reshape(-1, 3), axis=0)
bk_cart_unique = (bk_latt_unique / mp_grid[None, :]) @ recip_lattice
assert bk_latt_unique.shape == (self.NNB, 3)
ib_unique_map = np.zeros((self.NK, self.NNB), dtype=int)
ib_unique_map_inverse = np.zeros((self.NK, self.NNB), dtype=int)
bk_latt_unique_tuples = [tuple(b) for b in bk_latt_unique]
for ik in range(self.NK):
for ib in range(self.NNB):
b_latt = np.rint((self.bk_cart[ik, ib, :] @ np.linalg.inv(recip_lattice)) * mp_grid).astype(int)
ib_unique = bk_latt_unique_tuples.index(tuple(b_latt))
assert np.allclose(bk_cart_unique[ib_unique, :], self.bk_cart[ik, ib, :])
ib_unique_map[ik, ib] = ib_unique
ib_unique_map_inverse[ik, ib_unique] = ib
self.bk_latt_unique = bk_latt_unique
self.bk_cart_unique = bk_cart_unique
self.ib_unique_map = ib_unique_map
###################################################################
self.ib_unique_map_inverse = ib_unique_map_inverse
self.wk_unique = self.wk[0, ib_unique_map_inverse[0]]
self.neighbours_unique = np.array([neigh[order] for neigh, order in
zip(self.neighbours, self.ib_unique_map_inverse)])
for ik in range(0, self.NK):
order = ib_unique_map_inverse[ik]
assert np.allclose(self.wk[ik, order], self.wk_unique)
assert np.allclose(self.neighbours[ik, order], self.neighbours_unique[ik])
self.bk_dot_bk = self.bk_cart_unique @ self.bk_cart_unique.T
def set_bk_chk(self, chk, **argv):
self.set_bk(chk.kpt_latt, chk.mp_grid, chk.recip_lattice, **argv)