def initiate_wf_and_mm(self):
SIZE = 60
self.wf = Wavefunction.from_directory('.')
vol = self.wf.structure.volume
self.realspace_wf = self.wf.get_state_realspace(0,
0,
0,
dim=(SIZE, SIZE, SIZE),
remove_phase=True)
#self.realspace_chg = np.abs(self.realspace_wf)**2
self.realspace_chg = self.wf.get_state_realspace_density(0,
0,
0,
dim=(SIZE,
SIZE,
SIZE))
self.recipspace_wf = np.fft.fftn(
self.realspace_wf) / SIZE**3 * np.sqrt(vol)
self.recipspace_chg = np.fft.fftn(self.realspace_chg) / SIZE**3 * vol
self.mm_real = MomentumMatrix(self.wf, encut=3000)
self.mm_direct = MomentumMatrix(self.wf)
self.ncl_wf = NCLWavefunction.from_directory('noncollinear')
self.ncl_realspace_wf = self.ncl_wf.get_state_realspace(
0, 0, 0, dim=(SIZE, SIZE, SIZE), remove_phase=True)
self.ncl_recipspace_wf = (np.fft.fftn(self.ncl_realspace_wf[0]) / SIZE**3 * np.sqrt(vol),\
np.fft.fftn(self.ncl_realspace_wf[1]) / SIZE**3 * np.sqrt(vol))
def get_converged_encut(wavefunction,
bs,
iband=None,
max_encut=500,
n_samples=1000,
std_tol=0.02):
from pawpyseed.core.momentum import MomentumMatrix
nspins = wavefunction.nspin
nkpoints = wavefunction.kpts.shape[0]
if iband is None:
iband = {s: len(bs.bands[s]) for s in bs.spins}
sample_points = sample_random_kpoints(nspins, nkpoints, iband, n_samples)
origin = sample_points[0]
sample_points = sample_points[1:]
mm = MomentumMatrix(wavefunction, encut=max_encut)
true_overlaps = get_overlaps(mm, origin, sample_points)
# filter points to only include these with reasonable overlaps
mask = true_overlaps > 0.05
true_overlaps = true_overlaps[mask]
sample_points = sample_points[mask]
for encut in np.arange(100, max_encut, 50):
mm = MomentumMatrix(wavefunction, encut=encut)
fake_overlaps = get_overlaps(mm, origin, sample_points)
diff = (true_overlaps / fake_overlaps) - 1
if diff.std() < std_tol:
return encut
return max_encut
def get_wavefunction_coefficients(wavefunction,
bs,
iband=None,
ikpoints=None,
encut=600,
pbar=True):
from pawpyseed.core.momentum import MomentumMatrix
mm = MomentumMatrix(wavefunction, encut=encut)
if not iband:
iband = {}
if not ikpoints:
ikpoints = np.arange(mm.wf.kpts.shape[0])
coeffs = {}
for spin_idx in range(wavefunction.nspin):
spin = int_to_spin[spin_idx]
spin_iband = iband.get(spin, None)
coeffs[spin] = _get_spin_wavefunction_coefficients(mm,
bs,
spin,
iband=spin_iband,
ikpoints=ikpoints,
pbar=pbar)
return coeffs, mm.momentum_grid, wavefunction.kpts[ikpoints]
def get_wavefunction_coefficients(wavefunction,
bs,
iband=None,
encut=600,
pbar=True):
from pawpyseed.core.momentum import MomentumMatrix
mm = MomentumMatrix(wavefunction, encut=encut)
if not iband:
iband = {}
coeffs = {}
for spin_idx in range(wavefunction.nspin):
spin = int_to_spin[spin_idx]
spin_iband = iband.get(spin, None)
coeffs[spin] = _get_spin_wavefunction_coefficients(mm,
bs,
spin,
iband=spin_iband,
pbar=pbar)
return coeffs
class TestMomentumMatrix:
def setup(self):
self.currdir = os.getcwd()
os.chdir(os.path.join(MODULE_DIR, "../../../test_files"))
self.initialize_wf_and_mm()
def initialize_wf_and_mm(self):
SIZE = 60
self.wf = Wavefunction.from_directory(".")
vol = self.wf.structure.volume
self.realspace_wf = self.wf.get_state_realspace(0,
0,
0,
dim=(SIZE, SIZE, SIZE),
remove_phase=True)
# self.realspace_chg = np.abs(self.realspace_wf)**2
self.realspace_chg = self.wf.get_state_realspace_density(0,
0,
0,
dim=(SIZE,
SIZE,
SIZE))
self.recipspace_wf = np.fft.fftn(
self.realspace_wf) / SIZE**3 * np.sqrt(vol)
self.recipspace_chg = np.fft.fftn(self.realspace_chg) / SIZE**3 * vol
self.mm_real = MomentumMatrix(self.wf, encut=3000)
self.mm_direct = MomentumMatrix(self.wf)
self.mm_direct2 = MomentumMatrix(self.wf, encut=self.wf.encut)
self.ncl_wf = NCLWavefunction.from_directory("noncollinear")
self.ncl_realspace_wf = self.ncl_wf.get_state_realspace(
0, 0, 0, dim=(SIZE, SIZE, SIZE), remove_phase=True)
self.ncl_recipspace_wf = (
np.fft.fftn(self.ncl_realspace_wf[0]) / SIZE**3 * np.sqrt(vol),
np.fft.fftn(self.ncl_realspace_wf[1]) / SIZE**3 * np.sqrt(vol),
)
def teardown(self):
os.chdir(self.currdir)
def test_ncl_transform(self):
chg = np.sum(np.abs(self.ncl_recipspace_wf[0])**2) + np.sum(
np.abs(self.ncl_recipspace_wf[1])**2)
assert_array_almost_equal(chg, 1, 3)
def test_encut_insensitivity(self):
res = self.mm_direct.get_momentum_matrix_elems(0, 0, 0, 0, 0, 0)
res2 = self.mm_direct2.get_momentum_matrix_elems(0, 0, 0, 0, 0, 0)
assert_almost_equal(res[0], 1, 4)
assert_almost_equal(res2[0], 1, 4)
assert_almost_equal(res[:6], res2[:6], 7)
def test_get_momentum_matrix_elems(self):
res = self.mm_direct.get_momentum_matrix_elems(0, 0, 0, 0, 0, 0)
grid = self.mm_direct.momentum_grid
for i in range(grid.shape[0]):
if (np.abs(grid[i]) < 3).all():
# print(grid[i], res[i], self.recipspace_chg[grid[i][0],grid[i][1],grid[i][2]])
assert_almost_equal(
res[i], self.recipspace_chg[grid[i][0], grid[i][1],
grid[i][2]], 3)
with assert_raises(ValueError):
self.mm_direct.get_momentum_matrix_elems(0, 0, 0, 0, -1, 0)
def test_get_reciprocal_fullfw(self):
res = self.mm_real.get_reciprocal_fullfw(0, 0, 0)
print("check size", np.sum(np.abs(res)**2))
grid = self.mm_real.momentum_grid
for i in range(grid.shape[0]):
if (np.abs(grid[i]) < 2).all():
# print(grid[i], res[i], self.recipspace_wf[grid[i][0],grid[i][1],grid[i][2]])
assert_almost_equal(
res[i], self.recipspace_wf[grid[i][0], grid[i][1],
grid[i][2]], 3)
with assert_raises(ValueError):
self.mm_real.get_reciprocal_fullfw(50, 0, 0)
def test_g_from_wf(self):
grid = self.mm_real.momentum_grid
for i in range(grid.shape[0]):
if (np.abs(grid[i]) < 2).all():
# print(grid[i], self.mm_real.g_from_wf(0,0,0,0,0,0,grid[i]), self.recipspace_chg[grid[i][0],grid[i][1],grid[i][2]])
assert_almost_equal(
self.mm_real.g_from_wf(0, 0, 0, 0, 0, 0, grid[i]),
self.recipspace_chg[grid[i][0], grid[i][1], grid[i][2]],
3,
)
with assert_raises(ValueError):
self.mm_real.g_from_wf(100, 0, 0, 0, 0, 0, [0, 0, 0])