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. OK, I Understand
 
Top