def notestPropagatePlaneRealistic(self): label = join(test_folder, 'siesta') eigs, kpts, EF = readInfoForSTMSimulation(label=label) rho, variables, dh_cell = readCubeDensity(label=label) dh = np.array(map(vectorLength, dh_cell)) band = 30 kn = 4 kappa = 1.0 rho = rho.sum(axis=0) c, variables, domain = surfaceIntegrator(rho=rho, variables=variables) wf, variables1, dh_cell = readCubeWavefunction(label=label, band=band, kn=kn) kpt, wk = kpts[kn-1] eikr = phaseFactor(kpt, variables1) u1 = wf/eikr u2, variables2 = propagateWaveFunction( c=c, kpt=kpt, u=u1, variables=variables, kappa=kappa, ) grad_u = np.array(np.gradient(u2, *dh)) index = 5 grad_uz = grad_u[2, :, :, index] u2_plane = u2[:,:, index] plane_z = variables2[2,0,0, index] r_plane = variables2[:2, :,:, index] u3, variables3 = propagateWaveFunctionPlane( kpt=kpt, u=u2_plane, grad_uz=grad_uz, variables=r_plane, kappa=kappa, plane_z=plane_z, ) from matplotlib import pylab as plt fig = plt.figure(figsize=(8, 8)) ax = plt.gca() z = variables[2,0,0,:] v = abs(u1).sum(axis=(0,1)) ax.semilogy(z, v) z = variables2[2,0,0,:] v = abs(u2).sum(axis=(0,1)) ax.semilogy(z, v) z = variables3[2,0,0,:] v = abs(u3).sum(axis=(0,1)) #ax.semilogy(z, v) print abs(u2).sum(axis=(0,1)) print variables2[2,0,0,:] plt.show() eeeeee
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