def notestSaveReadWaveFunctionsTest(self):
band = 21
kn = 4
ws = 2
potential = Topographic(voltage=-2)
rho, r, dh_cell = readCubeDensity(label=test_label)
dh = np.array(map(vectorLength, dh_cell))
rho = rho.sum(axis=0)
c, variables, domain = surfaceIntegrator(rho=rho, variables=r, rho0=1e-3, DS=1)
thick_surface = vectorLength(c, axis=0) > 1e-9
take_z = thick_surface.sum(axis=0)
take_z = take_z.sum(axis=0) > 0
wf, r, dh_cell = readCubeWavefunction(label=test_label, band=band, kn=kn)
EF, eigs= readEigenvalues(label=test_label)
kpts = readKpoints(label=test_label)
energy = eigs[kn-1, band-1, 0]
we = STMWeighting(energies=[energy], potential=potential, EF=EF)
we = we[0]
kpt, wk = kpts[kn-1]
kappa = getVacuumDecayRate(energy)
eikr = phaseFactor(kpt, r)
u = wf/eikr
grad_u = np.array(np.gradient(u, *dh))
A = (c*grad_u).sum(axis=0)
B = u*c
A=A[:,:,take_z]
B=B[:,:,:,take_z]
r=r[:,:,:,take_z]
u, variables = propagateWaveFunction(
A=A,
B=B,
kpt=kpt,
r=r,
kappa=kappa,
height=10,
)
weight = we*ws*wk
Itest = weight*abs(u)**2
self.assertFalse(os.path.exists(testfile))
saveWavefunctionsToPlane(filename=testfile, label=test_label)
self.assertTrue(os.path.exists(testfile))
f = netcdf.netcdf_file(testfile, 'r')
self.assertTrue(len(f.variables.keys()) >0)
f.close()
I, r = STMFromFile(
label=join(test_folder, 'test'),
potentials=[potential],
cells=None,
)
self.assertArraysEqual(np.log(Itest[0,0]), np.log(I[0,0,0]), 5)
def saveWavefunctionsToPlane(filename, label='./siesta', density_range=(-3.5, -2.5)):
rho, r, dh_cell = readCubeDensity(label=label)
x, y, z = r
dh = np.array([dh_cell[0,0], dh_cell[1,1], dh_cell[2,2]])
eigs, kpts, EF = readInfoForSTMSimulation(label=label)
eigs, rho, w_s = handleSpinPolarization(eigs, rho)
is_spin_polarized = w_s == 1
if is_spin_polarized:
rho = rho.sum(axis=0)
energies = np.array(eigs)[:,0]
i_o = (energies < EF).sum() - 1
i_u = i_o + 1
i_t = i_o - 2
rho0 = 10**(np.array(density_range).mean())
DS = density_range[1] - density_range[0]
c, c_r, domain = surfaceIntegrator(rho=rho, variables=r, rho0=rho0, DS=DS)
thick_surface = vectorLength(c, axis=0) > 1e-9
take_z = thick_surface.sum(axis=0)
take_z = take_z.sum(axis=0) > 0
thick_surface = thick_surface[:,:,take_z]
x, y, z = c_r
wfs = []
EF, energies = readEigenvalues(label=label)
folder = os.path.dirname(label)
for wf_filename in os.listdir(folder):
wf_filename = wf_filename.split('.')
if wf_filename[0] == 'siesta' and wf_filename[-2] == 'REAL' and wf_filename[-1] == 'cube':
kn = int(wf_filename[1][1:])
band = int(wf_filename[2][2:])
spin = wf_filename[3]
if is_spin_polarized:
spin = {'UP':0, 'DOWN':1}[spin]
elif spin=='DOWN':
continue
else:
spin = 0
wfs.append((energies[kn-1, band-1, spin], kn-1, band-1, spin))
wfs.sort()
nc_defined = False
f = netcdf.netcdf_file(filename, 'w')
for i, info in enumerate(wfs):
energy, kn, band, spin = info
kpt, w_k = kpts[kn]
if energy > 0:
continue
kappa = getVacuumDecayRate(energy)
wf, r, dh_cell = readCubeWavefunction(band + 1, kn=kn+1, spin=spin+1, folder=folder)
eikr = phaseFactor(kpt, r)
u = wf/eikr
grad_u = np.array(np.gradient(u, *dh))
A = (c*grad_u).sum(axis=0)
r = r[:,:,:, take_z]
u = u[:,:,take_z]
A = A[:,:,take_z]
ct = c[:, :, :, take_z]
if not nc_defined:
nwf = f.createDimension('nwf', len(wfs))
xd = f.createDimension('x', u.shape[0])
yd = f.createDimension('y', u.shape[1])
zd = f.createDimension('z', u.shape[2])
ran = f.createDimension('range', 2)
points = f.createDimension('points', thick_surface.sum())
vector3d = f.createDimension('vector3d', 3)
comp = f.createDimension('complex', 2)
d_var = f.createVariable('thick_surface', dimensions=('x', 'y', 'z'), type=np.dtype('b'))
ran_var = f.createVariable('grid_cell', dimensions=('vector3d', 'vector3d'), type=np.dtype('f'))
ran_var[:] = dh_cell
origin_var = f.createVariable('origin', dimensions=('vector3d', ), type=np.dtype('f'))
origin_var[:] = np.array([r[0].min(), r[1].min(), r[2].min()])
d_var[:] = thick_surface
c_var = f.createVariable('c', dimensions=('vector3d', 'points'), type=np.dtype('f4'))
for ci in range(3):
c_var[ci,:] = ct[ci, thick_surface]
spin_weight = f.createVariable('spin_weight', dimensions=tuple(), type=np.dtype('i'))
spin_weight.assignValue(w_s)
fermi_energy = f.createVariable('EF', dimensions=tuple(), type=np.dtype('f'))
fermi_energy.assignValue(EF)
kpts_var = f.createVariable('kpts', type=np.dtype(float), dimensions=('nwf', 'vector3d'))
wk_var = f.createVariable('wk', type=np.dtype(float), dimensions=('nwf',))
energy_var = f.createVariable('energy', type=np.dtype(float), dimensions=('nwf',))
spin_var = f.createVariable('spin', type=np.dtype('i'), dimensions=('nwf',))
u_var = f.createVariable('u', type=np.dtype('f4'), dimensions=('nwf', 'points', 'complex'))
A_var = f.createVariable('A', type=np.dtype('f4'), dimensions=('nwf', 'points', 'complex'))
nc_defined = True
wk_var[i] = w_k
kpts_var[i] = kpt
energy_var[i] = energy
spin_var[i] = spin
u_var[i,:, 0] = u[thick_surface].real
u_var[i,:, 1] = u[thick_surface].imag
A_var[i,:, 0] = A[thick_surface].real
A_var[i,:, 1] = A[thick_surface].imag
f.close()
return True