main.py (1751B)
1 #!/usr/bin/python3.9 2 3 import numpy as np 4 from scipy import sparse 5 from scipy.sparse.linalg import spsolve 6 7 import matplotlib.pyplot as plt 8 import os 9 10 from system import * 11 from plotter import * 12 13 import time 14 15 """ 16 Solves the 2D time dependent Schrodinger Eq. on a square mesh. 17 Initial condition is a Gaussian wave packet. 18 The Potential is in form of the double/single slit 19 20 H * Psi = Psi_t 21 22 FD approximation of the Laplace Operator, 23 Crank Nicolson time discretization and 24 hbar = m = 1 makes the system equal to the following 25 26 (i - dt/2 * H) * Psi_{n+1} = (dt/2 * H + i) * Psi_n 27 """ 28 29 def main(): 30 dx = 0.7; dt = 0.2 31 n = 100; steps = 100 32 xmin = 0; xmax = 10 33 34 x = np.linspace(xmin, xmax, n) 35 X, Y = np.meshgrid(x, x) 36 37 x0 = 8; y0 = 5 38 k0 = -50; V0 = 100 39 40 # Initial Condition 41 psi0 = wave_packet(X, Y, dx, x0, y0, k0) 42 psi0 = (psi0/norm(psi0, x)).T.reshape(n*n) 43 44 # Potential 45 pos = 5; thickness = 0.07 46 xs1 = 4.6; xs2 = 5.4; 47 48 h0 = indicator(V0, x, x, xmin, xs1, pos, pos+thickness) 49 h1 = indicator(V0, x, x, xs2, xmax, pos, pos+thickness) 50 h2 = indicator(V0, x, x, 4.8, 5.2, pos, pos+thickness) 51 V = h0 + h1 + h2 52 53 #spalt = 0.2 54 #for i in range(1, int((xs2-xs1)/spalt)): 55 # if i%3 != 0: 56 # V += indicator(V0, x, x, xs1+spalt*i, xs1+spalt*(i+1), pos, pos+thickness) 57 58 59 # Assemble system 60 A = left_side(n, V, dx, dt) 61 M = right_side(n, V, dx, dt) 62 63 start = time.time() 64 Psi = mysolver(A, M, psi0, dt, steps) 65 timing = time.time() - start 66 67 print(f'Time for the calculation: {round(timing, 2)}') 68 69 U = convert(Psi, x, x) 70 71 plotter(U, V, dt, X, Y, 'single_slit') 72 73 if __name__ == "__main__": 74 main() 75