// shoot_demo -- This is a demonstration of solving the time-independent
// Schrodinger equation for a particle in a potential energy well by the
// shooting method -- that is, the equation is repeatedly integrated
// across the well from left to right with different energy values until
// the right-hand side boundary conditions are satisfied, i. e., psi
// goes to zero there.  The equation actually solved is scaled:
// d^2 psi/d x^2 + 2(eps - u(x))*psi = 0, where eps is the scaled total
// energy and u(x) is the scaled potential energy.  The scaling factor
// in each case is E_scale = hbar^2/(2m*L^2) where L is an assumed length
// scale.  Thus, eps = E/E_scale and u = U/E_scale.  The length scale
// comes from assuming that x --> L*x, where x after scaling is
// dimensionless.

// For the harmonic oscillator with spring constant k (the default
// potential), u = 0.5*k*m*L^2*L^2*x^2/hbar^2.  (The extra L^2 comes
// from non-dimensionalizing the x^2 from the potential energy function.)
//  We define L for maximum simplification, so that u = x^2/2.  This
// forces L to be given by L = (hbar^2/(k*m))^0.25.

// The single second order differential equation is equivalent to
// the pair of coupled equations d sig/d x + 2(eps - u(x))*psi = 0 and
// d psi/d x = sig.  These are solved by the Euler half step-full step
// method.  The program will repeatedly ask for values of eps to use.
// For each case the resulting psi(x) and u(x) will be plotted.  The
// value of eps is also plotted on the u(x) graph.

// Calling sequence: shoot(dx, nx):
//                   dx is the grid size in space
//                   nx is the number of grid elements

// Suggested starting calls for the harmonic oscillator (u(x) = x^2):
// Try shoot(0.05,201) with scaled energies near 0.5 and 1.5.  This
// creates a domain of integration from -5 to +5.  If you try different
// potentials, be sure to make the domain large enough that eps < u
// on both ends of the domain.  However, making this region of negative
// kinetic energy too large will cause numerical problems.  The plotting
// routine is set up to plot values of u from 0 to 5.  With different
// potential functions, you may have to change these limits.

// the potential energy function -- this may be changed if you want
// to try different potentials
potential = function(x)
{
  return x*x/2.;
}

// the main routine -- try to understand how this works
shoot = function(dx, nx)
{

// check to see if the number of arguments is correct
  if (nargs != 2) {
    error("Usage: shoot(dx,nx)");
  }

// allocate arrays
  x = ([0:nx - 1]' - (nx - 1)/2)*dx;
  psi = zeros(nx,1);
  sig = psi;
  ksq = sig;
  eps1 = sig;
  u = sig;
  for (ix in 1:nx) {
    u[ix] = potential(x[ix]);
  }

// loop on energy requests
  while (1) {
    fprintf("stdout","Type a guess for the energy (q to quit): ");
    inlist = getline("stdin");
    if (type(inlist.[1]) == "string") {
      break;
    }
    eps = inlist.[1];

// set up wavenumber squared
    for (ix in 1:nx) {
      eps1[ix] = eps;
      ksq[ix] = 2.*(eps - u[ix]);
    }

// initialize integration
    psi[1] = 0.;
    sig[1] = 1.;
    psimax = 0.;

// integration loop
    for (ix in 2:nx) {
      jx = ix - 1;

// half step for intermediate values
      sigstar = sig[jx] - ksq[jx]*psi[jx]*dx/2.;
      psistar = psi[jx] + sig[jx]*dx/2.;

// full step integration
      sig[ix] = sig[jx] - 0.5*(ksq[jx] + ksq[ix])*psistar*dx;
      psi[ix] = psi[jx] + sigstar*dx;
      if (ix < .7*nx && abs(psi[ix]) > psimax) {
        psimax = abs(psi[ix]);
      }
    }

// normalize the wavefunction
    psi = psi/psimax;

// plot it
    pgopen();
    pgclear();
    pgwindow(1,0,0,7.5,10);
    pgwindow(2,7.5,0,15,10);
    pgaxes(1,x[1],x[nx],5,"x",-1,1,11,"psi");
    pgaxes(2,x[1],x[nx],5,"x",0,5,11,"potential");
    pgline(1,5,x,psi);
    pgline(2,9,x,u);
    pgline(2,13,x,eps1);
    sprintf(seps,"%f",eps);
    pglabel(1,.95*x[1],-.85,"Energy = " + seps);
    pgpause();
    pgclose();
  }
}