system.py (1885B)
1 #!/usr/bin/python3.9 2 3 import numpy as np 4 from scipy import sparse 5 from scipy.sparse.linalg import spsolve 6 from scipy.sparse.linalg import inv 7 8 import time 9 10 def indicator(V0, x, y, x0, x1, y0, y1): 11 n = len(x) 12 h = np.zeros(n*n) 13 for i in range(n): 14 for j in range(n): 15 if x[i] >= x0 and x[i] <= x1: 16 if y[j] >= y0 and y[j] <= y1: 17 h[i + j*n] = 1 18 19 return V0*h 20 21 def hamiltonian(n, V, dx): 22 dim, n = n, n*n 23 24 d0 = -2*(1/dx**2 + 1/dx**2) * np.ones(n) + V 25 d1 = np.array([1/dx**2 if i%dim != 0 else 0 for i in range(1, n)]) 26 dn = 1/dx**2 * np.ones(n-dim) 27 28 return d0, d1, dn 29 30 def left_side(n, V, dx, dt): 31 dim, n = n, n*n 32 33 d0, d1, dn = hamiltonian(dim, V, dx) 34 35 d_0 = -dt/2 * d0 + 1j 36 d_1 = -dt/2 * d1 37 d_n = -dt/2 * dn 38 39 return sparse.diags([d_n, d_1, d_0, d_1, d_n], [-dim, -1, 0, 1, dim], format='csc') 40 41 42 def right_side(n, V, dx, dt): 43 dim, n = n, n*n 44 45 d0, d1, dn = hamiltonian(dim, V, dx) 46 47 d_0 = dt/2 * d0 + 1j 48 d_1 = dt/2 * d1 49 d_n = dt/2 * dn 50 51 return sparse.diags([d_n, d_1, d_0, d_1, d_n], [-dim, -1, 0, 1, dim], format='csc') 52 53 54 def wave_packet(x, y, dx, x0, y0, k0): 55 return 1/(2*dx**2*np.pi)**(1/2) *\ 56 np.exp(-((x-x0)/(2*dx)) ** 2) *\ 57 np.exp(-((y-y0)/(2*dx)) ** 2) *\ 58 np.exp(1.j * (k0*x)) 59 60 def mysolver(A, M, u0, dt, steps): 61 U = np.array([u0] + [np.zeros(u0.size) for i in range(steps)]) 62 for i in range(steps): 63 U[i+1] = spsolve(A, M.dot(U[i])) 64 return U 65 66 def norm(u, x): 67 n = len(x) 68 u = np.sqrt(np.real(u)**2 + np.imag(u)**2)**2 69 u = u.reshape(n, n) 70 return np.trapz(np.trapz(u, x), x) 71 72 def convert(U, x, y): 73 """normalization of each solution""" 74 n = len(x) 75 Unew = [] 76 for u in U: 77 unew = abs(u)**2/norm(u, x) 78 Unew.append(unew.reshape(n, n)) 79 return Unew 80