Myujth

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.