Wave function calculation for a particle in a finite square-well potential

wavefunc.py

# This is a script that calculates the values for the wave functions
# of a particle in a finite box, in n = 1, 2 and 3 energy levels.
# The script also plots the graphs of the wave functions.
#
# Created by: Vicente K. Parmigiani (GRR20196908)

import math
# Initialize constants
PI = math.pi # PI
H = 6.62607015 * 10**(-34) # Planck's constant
HL = H / (2.0 * PI) # Reduced Planck's constant
M = 9.1093837 * 10**(-31) # Mass of the electron
L = 1.0 # Length of the box
V0 = 4.8197 * 10**(-37) # First value for V0
# V1 = 1.0 * 10**(-36) # Second value for V0 (not used)
X0 = -0.5 # Initial value for x
X1 = 1.5 # Final value for x
STEP = 0.1 # Step for x
V_DEF = [1.28, 2.54, 3.73] # Default values for v
# E_DEF = [3.9932E-38, 1.5687E-37, 3.3790E-37] # Default values for E (not used)
A_DEF = [72.4152E-11, 53.2617E-11, 23.6017E-11] # Default values for A 

# Function that calculates the energy of the particle in the box
def energy(n):
    return 2 * HL**2 * V_DEF[int(n) - 1]**2 / (M * L**2)

# Function that calculates the normalization constant of the wave function
def norm(e, v):

    # Simplify the expression
    OPERATOR = math.sqrt((2.0 * M) / (HL**2))

    # Calculate the values
    v1 = math.sqrt(2.0 * e) * 2 * M * (v - e)
    v2 = math.sqrt(2.0 * e + L * v * OPERATOR * math.sqrt(v - e))
    v3 = HL**2 * 2.0 * e + (L * v * OPERATOR * math.sqrt(v - e))

    # return (v1 * v2) / v3
    # If we use the calculated values for E here, the wave function is correct,
    # but the normalization constant is not. So we use the given values for A
    # instead.
    return A_DEF[int(e) - 1]

# Function that calculates the wave function
def wave_function(a, v, e, x, psi, n):

    M2HL2 = math.sqrt((2.0 * M) / (HL**2))

    if (psi == 1):
        return a * math.exp(M2HL2 * math.sqrt(v - e) * x)
    elif (psi == 2):
        return a * math.cos(M2HL2 * math.sqrt(e) * x) + math.sqrt(v - e) * (a/math.sqrt(e)) * math.sin(M2HL2 * math.sqrt(e) * x)
    elif (psi == 3):
        # Check if needs to be reversed
        reversed = 1
        if (n % 2 == 0):
            reversed = -1
        return reversed * a * math.exp(M2HL2 * math.sqrt(v - e) * L) * math.exp(-M2HL2 * math.sqrt(v - e) * x)
    else:
        return 0

# Main function
def main():
    # Calculate the values
    ENERGYLVL = [energy(1), energy(2), energy(3)]

    # Print Energy levels for debug
    print("Energia (n = 1) = " + str(ENERGYLVL[0]))

    # Calculate the values for the normalization constant
    normR = [[norm(ENERGYLVL[0], V0), norm(ENERGYLVL[1], V0), norm(ENERGYLVL[2], V0)], [norm(ENERGYLVL[0], V0), norm(ENERGYLVL[1], V0), norm(ENERGYLVL[2], V0)]]

    # Var to store the values of the wave functions
    wave_functions = [[[], [], []], [[], [], []]] # 2 wave functions, 3 individual wave functions, 16 values for x

    # Print the header
    print("Tarefa 2 - Funcao de onda de uma partícula em um poco finito retangular")
    print("n = 1, 2 e 3")
    print("")
    print("L = 1.0 | V0 = 5.8 * 10**(-37) | V0 = 1 * 10**(-36)")
    print("---------------------------------------------------------")
    print("")
    print("Energia (n = 1) = " + str(ENERGYLVL[0]))
    print("Energia (n = 2) = " + str(ENERGYLVL[1]))
    print("Energia (n = 3) = " + str(ENERGYLVL[2]))
    print("")
    
    # Print the values for the normalization constant
    for i in range(1, 3):
        for j in range(1, 4):
            print("A (n = " + str(j) + ", V = " + str(i) + ") = " + str(normR[i - 1][j - 1]))

    print("")
    print("---------------------------------------------------------")

    # Build all the values for x of the wave functions in a table

    # Range 1, 2 for only 1 value of V0, 1, 3 for 2 values of V0
    for i in range(1, 2):

        vvalue = V0

        # Used for more than 1 value of V0
        # if (i == 1):
        #     vvalue = V0
        # else:
        #     vvalue = V1

        # The 3 individual wave functions
        for j in range(1, 4):

            print("")
            print("Valores da funcao de onda #" + str(i) + " (n = " + str(j) + ", -0.5 <= x <= 1.5, passo 0.1)")
            print("")

            # Dynamic multiplier to avoid floating point errors based on STEP
            MULTIPLIER = 1 / STEP

            # The values for x
            for x in range(int(X0 * MULTIPLIER), int(X1 * MULTIPLIER) + 1, int(STEP * MULTIPLIER)):

                x = x / MULTIPLIER

                # Calculate the value of the wave function for 0.5 <= x <= 1.5 and
                # psi = 1 if x <= 0, psi = 2 if 0 < x < L and psi = 3 if x >= L
                if (x < 0):
                    result = wave_function(normR[i - 1][j - 1], vvalue, ENERGYLVL[j - 1], x, 1, j)
                    wave_functions[i - 1][j - 1] += [result]
                    print("Funcao " + str(i) + " (n = " + str(j) + ") | x = " + str(x) + " | " + str(result))
                elif (x <= L):
                    result = wave_function(normR[i - 1][j - 1], vvalue, ENERGYLVL[j - 1], x, 2, j)
                    wave_functions[i - 1][j - 1] += [result]
                    print("Funcao " + str(i) + " (n = " + str(j) + ") | x = " + str(x) + " | " + str(result))
                else:
                    result = wave_function(normR[i - 1][j - 1], vvalue, ENERGYLVL[j - 1], x, 3, j)
                    wave_functions[i - 1][j - 1] += [result]
                    print("Funcao " + str(i) + " (n = " + str(j) + ") | x = " + str(x) + " | " + str(result))
    
    print("")
    print("---------------------------------------------------------")

    # Wait for the user to press enter if the script should plot the graphs
    if (input("Deseja plotar os graficos? (s/N) ") == "s"):
        # Import the libraries
        import matplotlib.pyplot as plt
        import numpy as np

        # Plot the 3 graphs on the same window, 2 on the top and 1 on the bottom centered (2*2)
    
        plt.subplots_adjust(wspace=0.5, hspace=0.5)

        # Range 1, 2 for only 1 value of V0, 1, 3 for 2 values of V0
        for i in range(1, 2):
            for j in range(1, 4):
                # plt.subplot(2, 3, (i - 1) * 3 + j)
                plt.subplot(2, 2, j)
                plt.plot(np.arange(X0, X1 + STEP, STEP), wave_functions[i - 1][j - 1])
                plt.title("(n = " + str(j) + ")")
                plt.xlabel("x")
                plt.ylabel("ψ")
                plt.grid(True)

        # Show the graphs
        plt.show()

# Call the main function
main()