// packet -- This demonstration computes and plots the evolution of a free
//           wave packet.  The wave frequency is given by the ``omega''
//           function at the end of the comments.  This may be changed to
//           suit.
//
// Function arguments:
//           nx -- number of points on the x-axis
//           dx -- distance between x-axis points
//           k0 -- central wavenumber of wave packet
//           width -- half-width (on x-axis) of wave packet
//           dt -- time step between successive plots
//           nt -- number of plots
//
// Returns:  This function returns a list.  The first member of the
//           list is the column vector of x positions.  The second
//           member is the row vector of times for which the wave
//           function is presented.  The third argument is a matrix
//           composed of columns.  Each column is the (complex) wave
//           packet as a function of x.
//
// Trial values of arguments: Try ``a = packet(101,1,.5,10,2.5,21);''
// See if you can determine what the phase and group velocities
// are for the default dispersion relation.  Also try the dispersion
// relations 1) omega = 1/k, 2) omega = 1, and 3) omega = k.
// Note that you may have to use different values of dt to get the
// wave function to change at the right rate -- trial and error here.
// Compare with the theory of phase and group velocity.  Once the
// program is run, you can obtain a more compact plot of the ups
// and downs of the wave function by doing ``pgplot2(real(a.psi))''.
// Alternatively, try ``pgplot3(real(a.psi))''.

//******************************************************************
// dispersion relation -- Change this as you please, but if you do,
//                        also change the ``dispreln'' string in the
//                        function ``packet'' to reflect what this
//                        function does.  This will get the label
//                        right on the output graph.
omega = function(k)
{
  return(k*k);
//  return(1/k);
//  return(k);
//  return(1);
}

//******************************************************************
// packet -- This is the main function
packet = function(nx, dx, k0, width, dt, nt)
{
  if (nargs != 6) {
    error("Usage: packet(nx,dx,k0,width,dt,nt)");
  }

// set dispersion relation string -- change this if the function
// ``omega'' is changed
  dispreln = "omega = k*k";

// initialize the wave packet
  x = [0:nx - 1]'*dx;
  xc = dx*nx/2;
  psi0 = exp(-(x - xc).^2/width^2 + (0 + 1i)*k0.*x);

// get the fourier transform
  dk = 6.28318/(nx*dx);
  k = ([1:nx] - nx/2)*dk;
  q = exp((0 + 1i)*x*k)/sqrt(nx);
  psik = q'*psi0;

// initial matrix whose columns are wavefunction at different times
  psi = zeros(nx,nt);
  for (ix in 1:nx) {
    psi[ix;1] = psi0[ix];
  }
  phase = zeros(nx,1);
  for (ik in 1:nx) {
    phase[ik] = exp((0 - 1i)*omega(k[ik])*dt);
  }

// set up graphics
  pgopen();
  pgclear();
  pgwindow(1,0,0,15,10);
  pgaxes(1,x[1],x[nx],6,"x",-1,1,5,"real(psi)");
  pgline(1,5,x,real(psi0));
  pglabel(1,5,.8,"Initial time");
  pglabel(1,5,.7,dispreln);
  pgpause();

// loop on time
  time = ([1:nt] - 1)*dt;
  for (it in 2:nt) {
    psik = psik.*phase;
    psi0 = q*psik;
    for (ix in 1:nx) {
      psi[ix;it] = psi0[ix];
    }
    pgclear();
    pgwindow(1,0,0,15,10);
    pgaxes(1,x[1],x[nx],6,"x",-1,1,5,"real(psi)");
    pgline(1,5,x,real(psi0));
    sprintf(timetext,"Time = %f",time[it]);
    pglabel(1,5,.8,timetext);
    pglabel(1,5,.7,dispreln);
    pgpause();
  }
  pgclose();

// return a list containing information
  return(<<x = x; time = time; psi = psi>>);
}