// ho2_demo.r -- this simply calculates the energies for the 2-D
//               harmonic oscillator using the known formula and
//               orders them -- the Hamiltonian:
//
//               H = (p_x^2 + p_y^2 + x^2 + lambda^2 y^2 )/2
//
//               setting lambda != 1 means that the spring constants
//               in the x and y directions are different

// input arguments -- emax -- maximum energy
//                    lambda -- square root of ratio of y to x spring
//                              constants

// output -- this routine returns a matrix in which the first column
//           contains the magnitude-sorted energies, while the second
//           and third columns contain the x and y quantum numbers
//           associated with each total energy

ho2energies = function(emax,lambda) {
  if (nargs != 2) {
    return "Usage: ho2energies(emax,lambda)";
  }

// compute maximum x and y limits
  xlim = int(emax);
  ylim = int(emax/lambda);

// allocate space for energies and indices
  max = xlim*ylim;
  energies = zeros(max,1);
  m = [1:max]';
  n = m;
  sm = m;
  sn = n;

// loop over energy assignment
  for (x in 0:xlim - 1) {
    for (y in 0:ylim - 1) {
      i = y + x*ylim + 1;
      m[i] = x;
      n[i] = y;
      energies[i] = (x + lambda*y) + (1 + lambda)/2;
    }
  }

// sort energies
  sortout = sort(energies);

// get sorted energies and indices up to emax
  energies = sortout.val;
  indices = sortout.ind;
  lim = 0;
  for (i in 1:max) {
    sm[i] = m[indices[i]];
    sn[i] = n[indices[i]];
    if (energies[i] < emax) {
      lim = i;
    }
  }

// return a matrix with energy the first column, m second, and n third
  return [energies[1:lim], sm[1:lim], sn[1:lim]];
}