""" From "COMPUTATIONAL PHYSICS" & "COMPUTER PROBLEMS in PHYSICS"
    by RH Landau, MJ Paez, and CC Bordeianu (deceased)
    Copyright R Landau, Oregon State Unv, MJ Paez, Univ Antioquia, 
    C Bordeianu, Univ Bucharest, 2018. 
    Please respect copyright & acknowledge our work."""

# HOchargeMat.py charge in HO plus E field wi Matplotlib

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

dx = 0.06
dx2 = dx * dx
k0 = 5.5 * np.pi
dt = dx2 / 8.
xmax = 6.
xs = np.arange(-xmax, xmax + dx / 2, dx)             # x array
psr = np.exp(-0.5 * (xs / 0.5)**2) * np.cos(k0 * xs)       # Re Psi
psi = np.exp(-0.5 * (xs / 0.5)**2) * np.sin(k0 * xs)       # Im Psi
E = 70                                        # E field
v = 25.0 * xs**2 - E * xs                         # V HO + E
fig = plt.figure()
ax = fig.add_subplot(111, autoscale_on=False,
                     xlim=(-xmax, xmax), ylim=(0, 1.5))
ax.grid()                                    # Plot grid
plt.title("Charged Harmonic Oscillator in E Field")
line, = ax.plot(xs, psr * psr + psi * psi, lw=2)


def animate(dum):
    psr[1:-1] = psr[1:-1] - (dt / dx2) * (psi[2:] + psi[:-2]
                                          - 2 * psi[1:-1]) + dt * v[1:-1] * psi[1:-1]
    psi[1:-1] = psi[1:-1] + (dt / dx2) * (psr[2:] + psr[:-2]
                                          - 2 * psr[1:-1]) - dt * v[1:-1] * psr[1:-1]
    line.set_data(xs, psr**2 + psi**2)
    return line,


ani = animation.FuncAnimation(fig, animate, 1, blit=True)
plt.show()
