// scatter1_demo.r -- Demonstrate scattering in 1-D from a
//                    potential.  The equations are scaled and
//                    nondimensionalized.  A potential of square
//                    or Gaussian shape is placed at the center of
//                    the domain.  Alternatively, a step function
//                    potential with the edge at the domain center is used.
//                    An outward radiation condition is applied to
//                    the left end so that the incident wave comes
//                    from the right and scatters off the potential,
//                    resulting in a transmitted wave to the left and
//                    a reflected wave to the right.                        

// The scaled governing equations are
//        d psi/dx = chi
//        d chi/dx = 2*[U(x) - energy]*psi

// input arguments -- n -- number of grid points
//                    dx -- size of grid cells
//                    umax -- maximum value of potential
//                    width -- half width of potential
//                    shape -- integer specifying shape of potential --
//                             1 = step function
//                             2 = square
//                             3 = Gaussian
//                    energy -- total energy of the particle
//                    plotflag -- if > 0, do plot; if <= 0, don't do plot

// suggested to start with:
//      a = scatter1(2001,.01,1.5,1,2,1,1);

// graphics -- The program plots the absolute square of the wave function
//             psi in the upper graph and the potential function in the
//             lower graph.

// returns -- a two element list containing the transmission and
//            reflection probabilities is returned

static(umax1,width1,shape1,epsi);

scatter1 = function(n,dx,umax,width,shape,energy,plotflag)
{
  global(umax1,width1,shape1,scatter1_rhs,epsi);

// check arguments
  if (nargs != 7) {
    error("Usage: scatter1_demo(n,dx,umax,width,shape,energy,plotflag)");
  }
  umax1 = umax;
  width1 = width;
  shape1 = shape;
  epsi = energy;

// construct x values and potential and total energy
  x = ([1:n]' - 0.5*(1 + n))*dx;
  xstart = x[1];
  xend = x[n];
  u = [1:n]';
  e = [1:n]';
  for (i in 1:n) {
    u[i] = potential(x[i]);
    e[i] = epsi;
  }

// set up starting conditions
  vec0 = [1; -1i*sqrt(2*(epsi - u[1]))];

// do the integration
  vec = myode(scatter1_rhs,x[1],x[n],vec0,dx);

// compute the normalized probability
  apsi = abs(vec[1;]');
  psi = vec[1;]';
  prob = real(conj(psi) .* psi);
  rprob = prob[.75*n:n];
  maxval = sqrt(max(rprob));
  minval = sqrt(min(rprob));
  probnorm = (.5*(minval + maxval))^2;
  prob = prob/probnorm;
  eta = minval/maxval;
  reflprob = ((1 - eta)/(1 + eta))^2;
  tranprob = 1 - reflprob;

// do the plot if requested
  if (plotflag > 0) {
    pgopen();
    pgclear();
    pgwindow(1,0,5,15,10);
    pgwindow(2,0,0,15,5);
    pgaxes(1,xstart,xend,11,"x",0,4,5,"|psi|^2");
    pgaxes(2,xstart,xend,11,"x",-2,2,5,"energy");
    pgline(1,5,x,prob);
    pgline(2,9,x,u);
    pgline(2,13,x,e);
    sprintf(title1," Transmission probability = %g",tranprob);
    pglabel(1,xstart,3.5,title1);
    sprintf(title2," Reflection probability = %g",reflprob);
    pglabel(1,xstart,3,title2);
    pglabel(2,xstart,1.5,"Potential and Total Energies");
    pgpause();
    pgclose();
  }

// return the stuff
  return <<tranprob = tranprob; reflprob = reflprob>>;
}

// scatter1_rhs -- evaluates the right hand side for myode
scatter1_rhs = function(x,vec)
{
  global(potential,epsi);
  u = potential(x);
  return [vec[2]; 2.*(u - epsi)*vec[1]];
}

// myode -- My ode solver that solves systems of ordinary differential
//          equations.  Works on complex as well as real functions.
//          Uses Euler half step-full step algorithm.

// inputs -- rhs -- function(x,y) yielding the tendency vector given
//                  the value of the independent variable, x, and the
//                  dependent vector, y(x).
//           xstart -- starting value of independent variable
//           xend -- ending value of independent variable
//           ystart -- starting value of dependent vector (should be
//                     a column vector
//           dx -- increment value of independent variable

// returns -- a matrix with each row representing one of the dependent
//            variables as a function of the independent variable
myode = function(rhs,xstart,xend,ystart,dx)
{

// do initialization
  if (nargs != 5) {
    return "Usage: myode(rhs,xstart,xend,ystart,dx)";
  }
  sflag = size(ystart);
  if (sflag[2] != 1) {
    return "myode: ystart has wrong shape";
  }
  m = sflag[1];
  if (m < 1) {
    return "myode: number of equations must be > 0";
  }
  n = int((xend - xstart)/dx + 1.5);
  if (n < 1) {
    return "myode: inconsistent dx, xstart, xend";
  }
  dhalf = dx/2;
  ymid = [1:m]';

// allocate space for results and initialize
  y = zeros(m,n);
  y[;1] = ystart;

// loop on steps
  for (i in 2:n) {

// half step
    ytendency = rhs(xstart + (i - 2)*dx,y[;i - 1]);
    ymid = y[;i - 1] + ytendency*dhalf;

// full step
    ytendency = rhs(xstart + (i - 1.5)*dx,ymid);
    y[;i] = y[;i - 1] + ytendency*dx;
  }

// done
  return y;
}

// compute the potential as a function of x
potential = function(x)
{
  global(umax1,width1,shape1);

  if (shape1 == 3) {
    return umax1*exp(-x^2/width1^2);  // this gives a Gaussian potential
  }

  if (shape1 == 2) {
    if (abs(x) > abs(width1)) {       // this gives a square potential
      return 0;
    else
      return umax1;
    }
  }

  if (x < 0) {                        // this gives a step potential
    return umax1;
  else
    return 0;
  }
}