#!/usr/bin/env python
"""
Created on Tir 29 Oct 2013.

Script uses the forward euler method to solve the Schrodinger equation for a
potential well.

Animation is done using the matplotlib.pyplot library.

Usage:
    python ForwardEulerPotentialWell.py
or equivalent:
    ./ForwardEulerPotentialWell.py
No commandline arguments are needed.

Notice how low the time step needs to be for the method to be stable. Try
raising the timestep to 1e6 and see what happens. Also, compare with the needed
time step for the Crank Nicolson method.

Another thing to note is that when using this method for integration it is
necessary to normalize the wave function for each time step. This is because
the forward euler scheme does not contain the normalization.

@author Kristoffer Braekken
"""
# Tools for sparse matrices
import scipy.sparse as sparse
import scipy.sparse.linalg

# Integration tool needed for normalization
import scipy.integrate

# Numerical tools
from numpy import *

# Plotting library
from matplotlib.pyplot import *

"""Physical constants"""
_E0p = 938.27        # Rest energy for a proton [MeV]
_hbarc = 0.1973      # [MeV pm]
_c = 3.0e2           # Spees of light [pm / as]

def Psi0( x ):
    '''
    Initial state for a travelling gaussian wave packet.
    '''
    x0 = 0 # [pm]
    a = 0.0050 # [pm]

    A = ( 1. / ( 2 * pi * a**2 ) )**0.25
    K1 = exp( - ( x - x0 )**2 / ( 4. * a**2 ) )

    return A * K1

def potentialWell( x, depth=35, width=0.15 ):
    """
    Gives the potential for a potential well of depth depth and width width.

    @param depth Gives the depth of the potential well. Given as the magnitude
    (positive integer / double / float).
    @param width Gives the width of the potential well. Given as positive
    integer definig the fraction of the x spectrum to contain the well. For
    example, 1 will mean that the well covers the whole spectrum and 0.5 that
    it covers half of it.
    """
    # Declare new empty array with same length as x
    potential = zeros( len( x ) )

    potential[ (0.5-0.5*width)*len(potential) : (0.5+0.5*width)*len(potential) ] = -depth

    return potential

if __name__ == '__main__':
    nx = 1001 # Number of points in x direction
    dx = 0.001 # Distance between x points [pm]

    # Use zero as center, same amount of points each side
    a = - 0.5 * nx * dx
    b = 0.5 * nx * dx
    x = linspace( a, b, nx )

    # Time parameters
    T = 0.2 # How long to run simulation [as]
    dt = 1e-7 # The time step [as]
    t = 0
    time_steps = int( T / dt ) # Number of time steps

    # Fetch potential
    V = potentialWell(x)

    # Constants - save time by calculating outside of loop
    k1 = - ( 1j * _hbarc * _c) / (2. * _E0p )
    k2 = ( 1j * _c * V ) / _hbarc

    # Create initial state Psi
    Psi = Psi0(x)

    # Initialize empty array for storing second derivative. `complex128`
    # argument necessary for calculations with complex numbers.
    D2Psi = zeros(nx).astype(complex128)

    # Precalculate indexes for finding second order derivative
    IND = arange(1,nx-1)
    INDP = arange(2,nx)
    INDM = arange(0,nx-2)

    # Put mmatplotlib in interactive mode for animation
    ion()

    # Setup the figure before starting animation
    fig = figure() # Create window
    ax = fig.add_subplot(111) # Add axes
    line, = ax.plot( x, abs(Psi)**2, label='$|\Psi(x,t)|^2$' ) # Fetch the line object

    # Also draw a green line illustrating the potential
    ax.plot( x, V, label='$V(x)$' )

    # Add other properties to the plot to make it elegant
    fig.suptitle("Solution of Schrodinger's equation with potential well") # Title of plot
    ax.grid('on') # Square grid lines in plot
    ax.set_xlabel('$x$ [pm]') # X label of axes
    ax.set_ylabel('$|\Psi(x, t)|^2$ [1/pm] and $V(x)$ [MeV]') # Y label of axes
    ax.legend(loc='best') # Adds labels of the lines to the window
    draw() # Draws first window

    # Dont plot all frames. Creates better "framerate" when dealing with
    # methods requiring low time step.
    PLOT_EVERY = 100
    counter = 0 # Counter for drawing frames

    # Time loop
    while t < T:
        # Second order derivative
        D2Psi[IND] = (Psi[INDP]-2.*Psi[IND]+Psi[INDM])/dx**2

        # Wavefunction after timestep
        Psi = Psi + dt*(k1*D2Psi+k2*Psi)

        # Normalize the new wave since this method doesnt keep normalization
        Psi /= scipy.integrate.simps(abs(Psi)**2, dx=dx)

        # Update time
    	t += dt

        counter += 1

        if counter == PLOT_EVERY:
            # Plot this new state
            line.set_ydata( abs(Psi)**2 ) # Update the y values of the Psi line
            draw() # Update the plot
            counter = 0

    # Turn off interactive mode
    ioff()

    # Add show so that windows do not automatically close
    show()