#!/usr/bin/python
#
"""
Solve Schroedinger equation in Morse potential using a 
finite-difference scheme. This script / function assumes
as input parameters
x:  discrete position grid,
    assumed equidistant for the moment
D, a, r0, V0: Morse parameters
Et: target energy eigenvalue,
    sometimes we need eigenfunctions above the ground level,
    centered around a laser frequency, for example
"""

tsize = 18
xsize = 6.5
ysize = xsize*0.5*(sqrt(5)-1)
plt.rcParams['figure.figsize'] = (xsize, ysize)
plt.rcParams['figure.autolayout'] = True
plt.rcParams['font.size'] = tsize-2
plt.rcParams['axes.labelsize'] = tsize
plt.rcParams['xtick.labelsize'] = tsize-4
plt.rcParams['ytick.labelsize'] = tsize-4

from pylab import *
from scipy.sparse import diags
from scipy.sparse.linalg import eigsh
from scipy.sparse import csc_matrix

# convergence checks that are possible: 
# -- is the grid large enough? (are the highest states still bound?)
# -- is the grid fine enough? (accuracy of eigenvalues)


if True:
    # the following relatively arbitrary parameters 
    # are the default values
    nb_levels = 3;
    D = 20.; a = 5.; r0 = 2.; V0 = 5.; Et = 10.; # and grid x
    #
    nb_x = 256 + 128 + 128
    x0, x1 = r0 - 1.2*a, r0 + 4.5*a
    x_g, dx = linspace(x0, x1, nb_x, retstep = True)

    # order of the finite-difference scheme for the second derivative
    # = 2 (= 4 ) is a second (fourth) order scheme. Computation time is
    # roughly the same, but the fourth order scheme is superior in accuracy,
    # compared to exact solution.
#    stencil = 2
    stencil = 4

def solve_Morse(D = 20., a = 5., r0 = 2., V0 = 5., Et = 10., \
                nb_levels = 3, x = x_g, Q_plot = False):
    """
    Solve Schroedinger equation for the Morse potential
    D(1 - exp(-(z - r0)/a))**2 + V0
    levels, states = solve_Morse()
    levels: list (length nb_levels) with eigenvalues
    states: array (size(x), nb_levels) with wave functions

    Meaning of parameters:
    D: dissociation energy
    V0: minimum energy
    a: length parameter
    r0: equilibrium bond length
    Et: targeted energy eigenvalue
    nb_levels: number of eigenvalues
    x: spatial grid (must be equidistant)
    Q_plot = False: plot the solution

    The script makes a few simple checks whether the grid is
    wide enough. Boundary conditions assumed:
    Dirichlet on inner side
    Neumann on outer side
    Can be adjusted by changing in the script the definition
    of type_BC = ['D', 'N'] (inner, outer).
    """
    # type of boundary conditions at left and right end
    type_BC = ['D', 'N'] 
    # D: Dirichlet
    # N: Neumann
    #
    nb_x = shape(x)[0]
    x0, x1 = x[0], x[-1]
    dx = min(abs(diff(x)))
    if abs(dx - max(diff(x))) > 1e-5*dx:
        print('... non-equidistant grid: sorry to return nonsense')
    #
    def morse(z):
        return D*(1. - exp(-(z - r0)/a))**2 + V0
#        return D*((z - r0)/a)**2 + V0
#         return V0*ones(size(z))
    # make a few tests of grid size, turning points at target energy Et
    if Et < V0:
        print('Error: increase Et above V0 = %4.4g' %(V0))
    if (morse(x[0]) < Et) | (morse(x[-1]) < Et):
        print('Error: turning points not in grid')
    #
    # set up Schroedinger matrix
    #
    h = 0.5; # hbar^2 / (2 m) in front of second derivative
    hb_omega = sqrt(2.*h*2.0*D/a**2); # harmonic approximation
    x_e = hb_omega / (4.*D); # anharmonicity parameter
    # checked that the Morse spectrum is indeed
    # E_k = V0 + hb_omega*(k + 1/2) - hb_omega*x_e*(k + 1/2)**2
    #

    # try to understand convergence towards exact eigenvalues
    # it is indeed a question of step size! It seems that 
    # with a smaller step (up to 1700 points were no problem numerically)
    # the agreement with the exact spectrum becomes better and better.
    # The width of the grid must be large enough, too, of course.
    #
    # Another idea: try to use a 'five-point stencil' for the second
    # derivative: accuracy improves from O(dx^2) to O(dx^4) -- works
    # very well, and I did not see any increase in computing time ...
    #
    # From https://en.wikipedia.org/wiki/Finite_difference_coefficient
    # replace band matrix with diagonals (1, -2, 1) [/dx**2]
    # by band matrix with entries (-1./12, 4./3, -2.5, 4./3, -1./12) [/dx**2]
    # guess that Dirichlet boundary condition is handled as before,
    # while Neumann BC (at the left end, e.g.) boils down to replacing
    # by the four-point stencil (., 5./4, -2.5, 4./3, -1./12 [/dx**2])
    # or by the three-point stencil (., ., -1.25, 4./3, -1./12 [/dx**2])
    #
    if stencil == 2: # (1, -2, 1) [/dx**2]
        diag_0 =  h*2./dx**2 + morse(x)
        diag_1 = -h*1./dx**2
        for b in [0,1]:
            if type_BC[b] == 'N':
                # adjust diagonal values at the end for Neumann b.condition
                diag_0[-b] = morse(x[-b]) + h*1./dx**2
        # sparse (banded) matrix from diagonals, secondary diags are -1 and +1
        Schroed = diags([diag_1, diag_0, diag_1], [-1, 0, 1], shape = (nb_x, nb_x))
    #
    elif stencil == 4: # (-1./12, 4./3, -2.5, 4./3, -1./12) [/dx**2]
        diag_0 =  h*2.5/dx**2 + morse(x)        # need the arrays
        diag_1 = -h*(4./3)/dx**2 * ones(nb_x-1) # here for the BC
        diag_2 =  h*(1./12)/dx**2
        for b in [0,1]:
            if type_BC[b] == 'N': 
                # 1st line (., ., -1.25, 4./3, -1./12 [/dx**2])
                # 2nd line (., 5./4, -2.5, 4./3, -1./12 [/dx**2])
                # adjust diagonal (-1,0,1) values at the end for Neumann b.condition
                # (just a guess, not obvious!)
                diag_0[-b] = morse(x[-b]) + h*1.25/dx**2
                diag_1[-b] = -h*1.25/dx**2
        # sparse (banded) matrix from diagonals, secondary diags are -1 and +1
        Schroed = diags([diag_2, diag_1, diag_0, diag_1, diag_2], [-2,-1, 0, 1, 2], \
                        shape = (nb_x, nb_x))
    #
    #              eigsh is the (sparse) solver for symmetric/hermitean problems
    # transform to 'column-indexed sparse coordinates'
    Schroed = csc_matrix(Schroed)
    level, state = eigsh(Schroed, k = nb_levels, sigma = Et, \
                         mode = 'normal', which = 'LM')
    # sigma: find eigenvalues around Et
    # mode: related to re-scaling of eigenvalues
    # which: L(argest) M(agnitude) for 1/(E - Et)
    if Q_plot:
        # make plot of found wave functions
        scale_psi = 0.5/sqrt(dx)
        offset_E = 0.03*scale_psi
        L = x1 - x0
        mysize = 18; tsize = mysize - 4

        plot( x, morse(x), 'k', lw = 2. )
        for k in range(nb_levels):
            scaled_Ek = (level[k] - V0) / hb_omega # re-scaled to convenient unit
            plot( x, level[k] + scale_psi*state[:,k], 'b', linewidth = 1.5 )
#            plot( [x[0], x[-1]], [level[k], level[k]], c = 'gray', linewidth = 1.0 )
            text( x1 - 0.3*L, level[k] + offset_E, \
#                  '$E_{%i}\,=\,%4.4g\,\hbar\omega$' %(k, scaled_Ek), \
                  '$%4.4g\,\hbar\omega$' %(scaled_Ek), \
                  fontsize = tsize, clip_on=True)
        # add numbers for exact Morse eigenvalues
        for k in range(int(2*D/hb_omega)):   
            expected_k = V0 + hb_omega*(k + 1./2) - hb_omega*x_e*(k + 1./2)**2
            if (expected_k - min(level) > -0.45*hb_omega) & \
               (expected_k - max(level) < 0.45*hb_omega):
                scaled_k = (expected_k - V0) / hb_omega
                plot( [x1 - 0.12*L, x1], [expected_k, expected_k], \
                      c = 'magenta', linewidth = 1.0 )
#                text( x1 - 0.12*L, expected_k + offset_E, '$%4.4g$' %(scaled_k), \
#                      fontsize = tsize, clip_on=True)

        xticks(size = tsize), yticks(size = tsize)
        ylim(min(level) - 1.2*hb_omega, max(level) + 1.2*hb_omega)    
        xlabel( '$\mathsf{position}\, x$', fontsize = mysize )
        ylabel( '$\mathsf{wave\,functions}\, \psi_n(x)$', fontsize = mysize )
        show(block=False)

    return level, state

# show(block=False)


Low-energy spectrum of the Morse potential. Magenta lines on right: exact energy eigenvalues.