• Sign Up
• Login
• API
• FAQ
• Tools
• Archive
daily pastebin goal
65%
SHARE
TWEET

# Untitled

a guest Mar 26th, 2019 68 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
1. ___author___ = "Endre Jacobsen, Wilhelm A. Mangersnes"
2. import numpy as np
3. import matplotlib
4. from scipy import integrate
5. matplotlib.use('TkAgg')
6. import matplotlib.pyplot as plt
7. """
8. CONSTANTS
9. """
10. H_BAR = 1       # Planck's constant (Atomic units)
11. MASS = 1        # Mass of an electron (Atomic units)
12. K_0 = 20        # Wave number
13. L = 20          # Length of box
14. OMEGA = (H_BAR * K_0 ** 2)/(2*MASS)     # Wave speed
15. E = OMEGA * H_BAR       # Energy
16. V_G = H_BAR*K_0/MASS        # Group velocity
17. T = L/(2*V_G)       # Time it will take for the propagation
18. N_X = 2000       # Steps in x-space
19. N_T = 10000   # Number of iterations in time
20. DELTA_X = L/(N_X-1)     # x-step
21.
22. '''For problem 2. Change these values.'''
23. X_S = 5  # Start position.
24.
25. SIGMA_S = 1.3       # Sigma values to test: 0.5, 1.0, 1.3, 1.5, 2.0
26.
27.
28. def main():
29.     delta_t = (T/N_T)    # 1 time-step
30.     """
31.    Arrays
32.    """
33.     x_list = np.linspace(0, L, N_X)      # List containing all N_X steps in x-space
34.     pot_list = np.zeros(N_X)        # Empty array for storing potential values
35.     psi = np.zeros(N_X, dtype=complex)     # Complex array
36.
37.     '''PROBLEMS. COMMENT IN/OUT AS NEEDED'''
38.     problem1(delta_t, x_list, psi)
39.     # problem2(delta_t, x_list, pot_list, psi)
40.     # problem3(delta_t, x_list, pot_list, psi)
41.     # problem4(delta_t)
42.     # problem5(delta_t)
43.
44.
45. def problem1(delta_t, x_list, psi):
46.     """
47.    Problem 1: Visualize the initial values of both the complex and real parts of Psi
48.    :param delta_t: Time-step
49.    :param x_list: List of length N_X with step-length DELTA_X
50.    :param psi: Empty complex array for Psi values
51.    :return: Nothing
52.    """
53.     init_psi(psi, delta_t)      # Initialize Psi
54.     prob_dens = calc_prob_density(psi, 0, N_X)      # Calculate the probability density
55.
56.     '''PLOT IMAGINARY AND REAL PARTS'''
57.     real_plot, = plt.plot(x_list, np.real(psi), label='Real')
58.     im_plot, = plt.plot(x_list, np.imag(psi), label='Imaginary')
59.     prob_plot, = plt.plot(x_list, prob_dens, label='Prob. density')
60.     plt.xlabel('x')
61.     plt.ylabel('Psi')
62.     plt.title('Initial wave-packet')
63.     plt.legend(handles=[real_plot, im_plot, prob_plot])
64.     plt.show()
65.
66.
67. def problem2(delta_t, x_list, pot_list, psi):
68.     """
69.    Problem 2: Propagate the wave a distance L/2
70.    :param delta_t: Timestep
71.    :param x_list: List of length N_X with step-length DELTA_X
72.    :param pot_list: Empty array for potential values
73.    :param psi: Empty complex array for Psi values
74.    :return: Noting
75.    """
76.     '''Initialize and plot Psi values (Problem 1)'''
77.     init_psi(psi, delta_t)      # Initialize psi
78.     prob_dens = calc_prob_density(psi, 0, N_X)      # Calculate probability density of initial wave-packet
79.     in_psi, = plt.plot(x_list, psi, label='Initialized wave')       # Plot wave-packet
80.     in_prob, = plt.plot(x_list, prob_dens, label='Prob. density for initialized wave')     # Plot prob. density
81.     '''Propagate the wave packet'''
82.     psi = draw_psi(psi, pot_list, delta_t)      # Calls draw_psi() to propagate the wave
83.     prob_dens = calc_prob_density(psi, 0, N_X)      # Calculate the prob. density of the new wave-packet
84.     new_psi, = plt.plot(x_list, psi, label='New wave')       # Plot new wave-packet
85.     new_prob, = plt.plot(x_list, prob_dens, label='Prob. density for new wave')     # Plot prob. density
86.     plt.xlabel('x')
87.     plt.ylabel('Psi')
88.     plt.title('Propagating a wave-packet, no barrier, sigma_x = 1.3')
89.     plt.legend(handles=[in_psi, in_prob, new_psi, new_prob])
90.     plt.show()
91.
92.
93. def problem3(delta_t, x_list, pot_list, psi):
94.     """
95.    Problem 3: Introduce a potential barrier at L/2 with height 0.5E and calculate the prob. density of the reflected
96.    and the transmitted wave-packet.
97.    :param delta_t: Timestep
98.    :param x_list: List of length N_X with step-length DELTA_X
99.    :param pot_list: Empty array for potential values
100.    :param psi: Empty complex array for Psi values
101.    :return: Nothing
102.    """
103.     '''Create barrier and propagate the wave past it. Plot results'''
104.     pot_list = make_barrier(pot_list, 0.5*E, L/50)       # Make a barrier at L/2
105.     init_psi(psi, delta_t)       # Initialize Psi
106.     psi = draw_psi(psi, pot_list, delta_t)      # Propagate the wave a distance L/2
107.     prob_dens = calc_prob_density(psi, 0, N_X)      # Calculate the prob. density of the resulting wave
108.     psi_plot, = plt.plot(x_list, psi, label='Wave-packet')       # Plot the resulting wave
109.     prob_plot, = plt.plot(x_list, prob_dens, label='Prob. density')     # Plot the prob. density
110.     v_plot, = plt.plot(x_list, pot_list / (2*E), label='Potential barrier')        # Plot the barrier. Scaled down with a factor 2*E
111.     plt.xlabel('x')
112.     plt.ylabel('Psi')
113.     plt.title('Wave propagated through a square barrier')
114.     plt.legend(handles=[psi_plot, prob_plot, v_plot])
115.     plt.show()
116.
117.     '''Calculate the probability of reflection and transmission'''
118.     prob_t_dens_ = calc_prob_density(psi, int(np.ceil((L/2)/DELTA_X)), N_X)     # Prob. density of transmitted wave
119.     prob_t = ratio(prob_t_dens_) * 100      # Probability of transmission
120.     print('There is a', prob_t, '% chance of transmission.')
121.
122.     prob_r_dens = calc_prob_density(psi, 0, int(np.floor((L / 2) / DELTA_X)))   # Prob. density of reflected wave
123.     prob_r = ratio(prob_r_dens) * 100       # Probability of reflection
124.     print('There is a', prob_r, '% chance of reflection.')
125.
126.     print('prob_t + prob_r =', prob_t + prob_r, '% . This should be close to 100%.')
127.
128.
129. def problem4(delta_t):
130.     """
131.    Problem 4: Calculate the probabilities of reflection and transmission with 50 different barrier heights
132.    :param delta_t: Time-step
133.    :return: Nothing
134.    """
135.     v_max = (3 / 2) * E     # Maximum value of the
136.     total_barriers = 50     # Number of barriers to be tested
137.     barrier_heights = np.linspace(0, v_max, total_barriers)         # Create an array with all the barrier heights up to v_max
138.
139.     prob_t_list = np.zeros(total_barriers)      # Empty array to hold the probability of transmission at each barrier height
140.     prob_r_list = np.zeros(total_barriers)      # Empty array to hold the probability of reflection at each barrier height
141.
142.     start_psi = init_psi(np.zeros(N_X, dtype=complex), delta_t)  # Initialize Psi
143.
144.     for i in range(len(barrier_heights)):
145.         psi = np.copy(start_psi)        # Reset psi for every new barrier
146.         pot_list = make_barrier(np.zeros(N_X), barrier_heights[i], L/50)        # Make the barrier
147.         psi = draw_psi(psi, pot_list, delta_t)      # Propagate the wave
148.         prob_dens_trans = calc_prob_density(psi, int(np.ceil((L / 2) / DELTA_X)), N_X)      # Calculate the prob. density for transmission
149.         prob_t_list[i] = ratio(prob_dens_trans)     # Calculate the probability for transmission
150.         prob_r_list[i] = 1 - prob_t_list[i]         # Calculate the probability for reflection
151.
152.     '''Plot the probabilities'''
153.     r_plot, = plt.plot(barrier_heights, prob_r_list, label='Probability for reflection', color="blue")
154.     t_plot, = plt.plot(barrier_heights, prob_t_list, label='Probability for transmission', color="gold")
155.     plt.xlabel('Potential')
156.     plt.ylabel('Probability')
157.     plt.title('P(Reflected) vs. P(Transmitted)')
158.     plt.legend(handles=[r_plot, t_plot])
159.     plt.show()
160.
161.
162. def problem5(delta_t):
163.     """
164.    Problem 5: Calculate the probabilities of reflection and transmission with 50 different barrier widths
165.    :param delta_t: Time-step
166.    :return: Nothing
167.    """
168.     v_max = (9/10)*E        # Height of the barrier
169.     max_width = L/20        # Maximum width of the barrier
170.     total_barriers = 50     # Number of barriers to test
171.
172.     barrier_widths = np.linspace(0, max_width, total_barriers)      # Create an array of barrier widths
173.
174.     prob_t_list = np.zeros(total_barriers)      # Empty array to hold the probability of transmission at each barrier height
175.     prob_r_list = np.zeros(total_barriers)      # Empty array to hold the probability of reflection at each barrier height
176.
177.     start_psi = init_psi(np.zeros(N_X, dtype=complex), delta_t)  # Initialize Psi
178.
179.     for i in range(len(barrier_widths)):
180.         psi = np.copy(start_psi)        # Reset psi for every new barrier width
181.         pot_list = make_barrier(np.zeros(N_X), v_max, barrier_widths[i])        # Create a new barrier
182.         psi = draw_psi(psi, pot_list, delta_t)      # Propagate the wave-packet
183.         prob_t_dens = calc_prob_density(psi, int(np.ceil((L/2)/DELTA_X)), N_X)      # Calculate the prob. density for transmission
184.         prob_t_list[i] = ratio(prob_t_dens)     # Calculate the probability for transmission
185.         prob_r_list[i] = 1 - prob_t_list[i]     # Calculate the probability for reflection
186.
187.     '''Plot the probabilities'''
188.     r_plot, = plt.plot(barrier_widths, prob_r_list, label='Probability for reflection', color="blue")
189.     t_plot, = plt.plot(barrier_widths, prob_t_list, label='Probability for transmission', color ="gold")
190.     plt.xlabel('Barrier widths')
191.     plt.ylabel('Probability')
192.     plt.title('P(Reflected) vs. P(Transmitted)')
193.     plt.legend(handles=[r_plot, t_plot])
194.     plt.show()
195.
196.
197. def draw_psi(psi, pot_list, delta_t):
198.     """
199.    Function for propagating the wave-packet contained in psi
200.    :param psi: A complex array containing the wave-packet
201.    :param pot_list: Potential
202.    :param delta_t: Time-step
203.    :return: The propagated wave-packet psi
204.    """
205.     for t in range(1, N_T):
206.         psi = calc_psi_i(psi, pot_list, delta_t)        # Calculate Im(Psi(x,t))
207.         psi = calc_psi_r(psi, pot_list, delta_t)        # Calculate Re(Psi(x,t+delta_t/2
208.     return psi
209.
210.
211. def init_psi(psi, delta_t):
212.     """
213.    Initializes Psi at time zero.
214.    :param psi:     Wave equation
215.    :param delta_t:     Time-step
216.    :return:        Initialized array of wave equation values
217.    """
218.     c = calc_c()        # Calculate normalisation constant
219.     for x in range(1, N_X-1):
220.         '''
221.        Equation (8)
222.        t = 0 for Psi_I and t = delta_t/2 for Psi_R
223.        '''
224.         psi[x] = complex(c[0] * np.exp(calc_alpha(x*DELTA_X)) * np.cos(K_0 * x*DELTA_X - OMEGA*(delta_t/2)),
225.                          c[0] * np.exp(calc_alpha(x*DELTA_X)) * np.sin(K_0 * x*DELTA_X))
226.     return psi
227.
228.
229. def calc_alpha(x):
230.     """
231.    Calculates alpha (the first exponent) from equation (8).
232.    :param x:   Position
233.    :return:    Alpha
234.    """
235.     alpha = -((x-X_S) ** 2)/(2*(SIGMA_S ** 2))
236.     return alpha
237.
238.
239. def calc_prob_density(psi, start, stop):
240.     """
241.    Calculate the probability density for a complex array
242.    :param psi: Complex array
243.    :param start: Where to start the calculation
244.    :param stop: Where to stop the calculation
245.    :return: Prob list, a list of probability densities
246.    """
247.     prob_list = np.zeros(stop-start)        # Create empty array containing number of elements to be calculated
248.     for i in range(stop-start):
249.         real = np.real(psi[i+start])        # Get the real value of the complex number
250.         imaginary = np.imag(psi[i+start])   # Get the imaginary value of the complex number
251.         prob_list[i] = real**2 + imaginary ** 2     # Calculate the probability density
252.     return prob_list
253.
254.
255. def calc_psi_i(psi, pot_list, delta_t):
256.     """ Calculate Psi at position x and time t.
257.    :param psi:    The list of wave equation values
258.    :param pot_list:    The list containing the potential values
259.    :param delta_t:     Time-step
260.    :return:    New value of psi(x,t)
261.    """
262.     '''Equation (7a)'''
263.     val_i = np.imag(psi[1:N_X-1]) - delta_t * ((pot_list[1:N_X-1]/H_BAR) * np.real(psi[1:N_X-1]) - H_BAR/(2*MASS) *
264.                                                (np.real(psi[2:N_X]) - 2*np.real(psi[1:N_X-1]) +
265.                                                 np.real(psi[0:N_X-2]))/(DELTA_X ** 2))
266.
267.     psi[1:N_X-1] = np.real(psi[1:N_X-1]) + val_i*1j
268.     return psi
269.
270.
271. def calc_psi_r(psi, pot_list, delta_t):
272.     """ Calculate Psi at position x and time t. The imaginary part is calculated at t + delta_t/2
273.    :param psi: The list of wave equation values
274.    :param pot_list:    Potential values
275.    :param delta_t:     Time-step
276.    :return: The new value of psi
277.    """
278.     '''Equation (7b)'''
279.     val_r = np.real(psi[1:N_X-1]) + delta_t *((pot_list[1:N_X-1]/H_BAR) * np.imag(psi[1:N_X-1]) - H_BAR/(2*MASS) *
280.                                               (np.imag(psi[2:N_X]) - 2*np.imag(psi[1:N_X-1]) +
281.                                                np.imag(psi[0:N_X-2]))/(DELTA_X ** 2))
282.
283.     psi[1:N_X-1] = val_r + np.imag(psi[1:N_X-1])*1j
284.     return psi
285.
286.
287. def calc_c():
288.     """
289.    Calculate the normalisation constant for equation (8)
290.    :return: Normalisation constant
291.    """
292.     def normalize_function():
293.         return lambda x: np.exp(-((x-X_S)**2)/SIGMA_S**2)
294.     expr = normalize_function()
295.     ans = integrate.quadrature(expr, 0, L, args=(), tol=1.49e-08)
296.     c = 1/np.sqrt(ans)
297.     return c
298.
299.
300. def make_barrier(pot_list, height, width):
301.     """
302.    Create a barrier of potential value 'height'
303.    :param pot_list: The empty potential values
304.    :param height: The height that the barrier is to be set to
305.    :param width: The width of the potential barrier
306.    :return: The potential (barrier)
307.    """
308.     pot_list[int(np.floor((L/2 - width/2)/DELTA_X)):int(np.ceil((L/2 + width/2)/DELTA_X))] = height
309.     return pot_list
310.
311.
312. def ratio(prob_dens):
313.     """
314.    Calculate the probability for a given prob. density
315.    :param prob_dens: Array of probability density
316.    :return: The probability for the given prob. density
317.    """
318.     return np.sum(prob_dens) * DELTA_X      # With DELTA_X small enough this is approximately the same as integrating.
319.
320.
321. main()
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy.

Top