Untitled

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

h_bar = 1.0545718e-34  # Reduced Planck constant (J.s)
m = 9.10938356e-31  # Electron mass (kg)
L = 0.1e-9  # Length of the box (m)

def energy_level(n):
    return (n**2 * np.pi**2 * h_bar**2)/(2 * m * L**2)

def wavefunction(n, x):
    return np.sqrt(2/L) * np.sin(n * np.pi * x / L)

def plot_superposition_at_time(weight1, weight2, n1, n2, t):
    x = np.linspace(0, L, 1000)
    psi1 = wavefunction(n1, x)
    psi2 = wavefunction(n2, x)
    magnitude = (weight1 ** 2 + weight2 ** 2) ** 0.5
    weight1 = weight1 / magnitude
    weight2 = weight2 / magnitude
    prob = ((weight1**2) * (psi1**2))  + ((weight2**2) * (psi2**2)) + 2 * weight1 * weight2 * psi1 * psi2 * np.cos(((energy_level(n2) - energy_level(n1)) * t) / h_bar)

    fig, ax = plt.subplots(figsize=(8, 6))
    ax.plot(x, prob)
    ax.set_xlabel('x')
    ax.set_ylabel('Probability')
    ax.set_title('Probability Density')

    plt.show()

plot_superposition_at_time(np.sqrt(5/10),np.sqrt(5/10),1,2,3)
''''