// forcing_demo.r -- demonstration of periodic forcing applied to
//                   any potential with bound states --
//                   the system is started out in a specified define
//                   state and a time-dependent potential energy
//                   of the form eps*x*cos(omega*t) is applied --
//                   the energies and eigenvectors for the bound
//                   state are obtained from a separate program
//                   such as ho_demo.r for the harmonic oscillator --
//                   generally only a subset of eigenstates of the
//                   full problem are used -- otherwise execution
//                   takes too long

// input arguments -- n0 -- number of initial state, with 1 being the
//                          ground state
//                    omega -- frequency of the forcing
//                    eps -- strength of the forcing
//                    a -- output list from ho_demo, anho_demo, etc.
//                    nstates -- max number of states to be included
//                               in the calculation
//                    dt -- time step -- be sure to make small enough
//                          so that the complex exponentials are not
//                          aliased
//                    nt -- number of time steps

// graphics -- a plot of the probability of being in each state as
//             as function of time is made -- the first curve is the
//             sum of all the probabilities and should be unity for
//             all times

// output -- a list containing a row vector of times, a column vector
//           of energies, and a matrix, the columns of which are the
//           state vectors (i. e., the amplitudes for each energy)
//           for each time

// trial command line arguments -- get the eigenvectors from (say)
// ho_demo.r as follows: a = ho_demo(200); -- use (say) the first
// 10 states
// run the program, starting in (say) state 1 with forcing frequency
// 1.2 and forcing strength 0.1:
//
// b = forcing_demo(1,1.2,0.1,a,10,0.1,501);
//
// a time step of 0.1 seems to work well in this case, and 501
// steps yields interesting results, but takes a few minutes to execute

static(vals1,omega1,dip,tip,nstates,march);

forcing_demo = function(n0,omega,eps,a,nstates1,dt,nt) {
  global(vals1,omega1,dip,tip,nstates);

// check arguments
  if (nargs != 7) {
    return "Usage: forcing_demo(n0,omega,eps,a,nstates,dt,nt)";
  }

// extract info from a
  x = a.x;
  vals = a.eigvals[1:nstates1];
  vecs = a.eigvecs[;1:nstates1];
  ncells = length(x);

// transfer arguments to global variables
  vals1 = vals;
  omega1 = omega;
  nstates = nstates1;

// compute the time grid
  time = [0:nt - 1]*dt;

// initialize state vector
  state0 = zeros(nstates,1)*(1 + 1i);
  state0[n0] = 1 + 0i;

// compute the x operator
  xop = zeros(ncells,ncells);
  for (i in 1:ncells) {
    xop[i;i] = x[i];
  }

// compute the dipole moments times eps/2 (assumed real)
  dip = zeros(nstates,nstates);
  for (i in 1:nstates) {
    psii = vecs[;i];
    for (f in 1:nstates) {
      psif = vecs[;f];
      dip[i;f] = 0.5*eps*psif'*xop*psii;
    }
  }

// allocate space for the time dependent part of the matrix
  tip = zeros(nstates,nstates);

// do the integration
  state = myode(march,0,dt*(nt - 1),state0,dt);

// plot the time series of probability and do a check sum
  probs = state .* conj(state);
  checksum = sum(probs);
  pgopen();
  pgclear();
  pgwindow(1,0,0,15,10);
  pgaxes(1,0,time[nt],11,"time",0,1.2,7,"probability");
  for (i in 1:nstates) {
    pgline(1,i,time',real(probs[i;]'));
  }
  pgline(1,1,time',real(checksum));
  sprintf(title,"Initial state: %d; Omega: %.3f; Eps: %.3f",n0,omega,eps);
  pglabel(1,0,1.1,title);
  pgpause();
  pgclose();

// returns a list with the energy eigenvalues, the time grid, and the
// time-marched state vector
  return <<time = time; ener = vals'; state = state>>;
}

// march -- right hand side function for myode

march = function(time,state) {
  global(vals1,omega1,dip,tip,nstates);

// compute the time dependent transformation matrix
  for (i in 1:nstates) {
    for (f in 1:nstates) {
      del1 = omega1 + vals1[f] - vals1[i];
      del2 = -omega1 + vals1[f] - vals1[i];
      tip[i;f] = dip[i;f]*(-1i*exp(1i*del1*time) - 1i*exp(1i*del2*time));
    }
  }

// return the result
  return tip*state;
}

// 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;
}