Ndtyjky

I'm trying to make an animation of the following wavefunction:

For some reason my code works for n=0 but then it doesn't for any other n. I checked the values of the function for various n, x and t and it seems to work fine but for some reason matplotlib isn't animating. Here's the code:

%matplotlib qt



import numpy as np

from sympy import hermite

from scipy.special import gamma

import matplotlib.pyplot as plt

from matplotlib import animation





def psi(n, x, t):

    A = (np.pi ** -0.25) * (1 / np.sqrt((2 ** n) * gamma(n + 1)))

    E = n + 0.5

    f = A * hermite(n, x) * np.exp(-(x ** 2) / 2) * np.cos(E * t)

    return f





def animar(f, x0=0, xf=1, dx=0.01, t0=0, tf=1, dt=0.01, ym=-2, yM=2):

    nf = int((xf - x0) / dx + 1)

    nt = int((tf - t0) / dt + 1)

    x = np.linspace(x0, xf, nf)

    t = np.linspace(t0, tf, nt)



    fig, ax = plt.subplots()



    ax.set_xlim((x0, xf))

    ax.set_ylim((ym, yM))



    line, = ax.plot(, , lw=2)



    def init():

        line.set_data(, )

        return line,



    def animate(i):

        y = f(x, i)

        line.set_data(x, y)

        return line,



    anim = animation.FuncAnimation(fig, animate, init_func=init,

                                   frames=5 * t, interval=20, blit=True)

    plt.show()

    return anim





F = lambda x,t: psi(1, x, t)

anim = animar(F, x0=-3, xf=3, tf=3, ym=-1, yM=1)

I don't think you want to evaluate the hermite polynomials symbolically. Instead, why not use numpy.polynomial.hermite.hermval? Also maybe use scipy.special.factorial to calculate the factorial.

import numpy as np

from numpy.polynomial.hermite import hermval

from scipy.special import factorial

import matplotlib.pyplot as plt

from matplotlib import animation





def psi(n, x, t):

    A = (np.pi ** -0.25) * (1 / np.sqrt((2 ** n) * factorial(n)))

    E = n + 0.5

    nar = np.zeros(n+1)

    nar[-1] = 1

    f = A * hermval(x,nar) * np.exp(-(x ** 2) / 2) * np.cos(E * t)

    return f





def animar(f, x0=0, xf=1, dx=0.01, t0=0, tf=1, dt=0.01, ym=-2, yM=2):

    nf = int((xf - x0) / dx + 1)

    nt = int((tf - t0) / dt + 1)

    x = np.linspace(x0, xf, nf)

    t = np.linspace(t0, tf, nt)



    fig, ax = plt.subplots()



    ax.set_xlim((x0, xf))

    ax.set_ylim((ym, yM))



    line, = ax.plot(, , lw=2)



    def init():

        line.set_data(, )

        return line,



    def animate(i):

        y = f(x, i)

        line.set_data(x, y)

        return line,



    anim = animation.FuncAnimation(fig, animate, init_func=init,

                                   frames=5 * t, interval=20, blit=True)

    plt.show()

    return anim





F = lambda x,t: psi(3, x, t)

anim = animar(F, x0=-3, xf=3, tf=3, ym=-1, yM=1)

Here's a fix. Redefine your psi function to be:

def psi(n, x, t):

    A = (np.pi ** -0.25) * (1 / np.sqrt((2 ** n) * gamma(n + 1)))

    E = n + 0.5

    herms = np.array([hermite(n, xelem) for xelem in x])

    f = A * herms * np.exp(-(x ** 2) / 2) * np.cos(E * t)

    return f

Basically, hermite seems to be a little weird/buggy. If x if a numpy array, then hermite(n=0, x) works just fine. However, if n>0 then hermite will complain bitterly if x is anything other than a single value. So we work around that by feeding in the values from x to hermite one at a time, then make the herms array from those results.

Quantum mechanics was many years ago now, so I can't tell you if there is a justifiable reason for the inconsistency of hermite vis-a-vis array inputs. My guess is that it's just down to a slightly buggy implementation in sympy. In any case, once you fix up those two lines in psi the animations seem to work fine.