% Find all the eigenvalues and eigenfunctions of an
% arbitrary 1D potential v1(x). The outputs are the eigenvalues EE,
% the array of wavefunctions psi, corresponding x-coordinate values x,
% and the uncertainty in the eigenvalues. Example of usage:
%     octave:1> [EE,psi, DeltaE] = eigenstates(@v1square, 2.5);
%
% v1 should be a function handle for a potential. The well has a _total_ width
% no greater than w, centered at x=0. In other words, the function v1 should
% return 0 for x outside this range, or else the returned eigenfunctions
% (which may "extend" beyond this range) will be goofed up. Fractional 
% tolerance on eigenvalues may be specified optionally (ftol, default 1e-12),
% and this tolerance also flows through to the integration procedure.
%
% The optional argument extend allows each plotted/returned eigenfunction
% to run over a wider domain (extend*w). This is a way to verify the
% repeatability of the integration (note instability in the V>=0 regime).
%
% CCK 2009-Apr-06
function [EE, psi, x, DeltaE] = ...
   eigenstates(v1, w, ftol, extend, oversample)
Npoints = 300; % Number of x points to use in integrating Schrodinger eqn.

if !exist('extend','var')
   extend = 1.25;  % default
   % Factor by which to extend w for plotting & normalization integral.
end
if !exist('ftol','var')
   ftol = 1e-12; % default value
end
if !exist('oversample','var')
   oversample = 10; % default value
end

% Set up lsode options
lsode_options('integration method','adams')
lsode_options('absolute tolerance', eps) % use machine precision here
lsode_options('relative tolerance', ftol)
lsode_options

% Automatically choose N_energies for the eigenvalue search (E_lowest < E < 0)
x1 = linspace(-w/2, w/2, Npoints);
V1arr = v1(x1);
predict_N_eigenvals = 1 + (1/pi)*sum( sqrt( -V1arr(V1arr < 0) ) ) * w/Npoints
N_energies = predict_N_eigenvals * oversample;
   % Always set N_energies to something larger than this.
E_lowest = min(V1arr);

E = linspace(E_lowest, ftol*E_lowest/N_energies/10, N_energies);  % We avoid E=0.
for i=1:N_energies
   err(i) = oneshot(E(i), v1, w, Npoints);
end

% This plot is a useful diagnostic to see that err(E) is well-sampled.
figure(1)
plot(E, err)
xlabel('{E}', 'fontsize', 18);
ylabel('{err}', 'fontsize', 18);
title('err(E): Eigenvalues are the roots.', 'fontsize', 20)

% Now use the tabulated err(E) to find the eigenvalues and eigenfunctions.
n=0;  % quantum number
psi = []; % initialize psi as an empty array.
   % (Without this, the first attempt to add a new row to psi would fail)
for i = 2:N_energies  % Note that we start at the second one! (why?)
   if (err(i)*err(i-1) < 0) % root bracketed
      n = n+1
      format long  % Display eigenvalues in maximum precision.
      eigenvalue = bisect(@(z) oneshot(z, v1, w, Npoints), E(i), E(i-1), ftol)
      format short
      EE(n) = eigenvalue;
      Eigenvalue_uncertainty = EE(n) * ftol % Report uncertainty
      DeltaE(n) = Eigenvalue_uncertainty; % store & return uncertainty
      % Now produce a version of psi(x), on an expanded domain, for output.
      [err_best, q, x] = oneshot(eigenvalue, v1, w*extend, Npoints*extend);
         % w*extend lets us get the tails of the wave function.
      % integral_psi_sq = sum( q(:,1).^2 ) * w/Npoints; % proper normalization
      % psi = [psi; q(:,1)'/integral_psi_sq];
      psi = [psi; q(:,1)'];
   endif
end

% And now for a complete plot. User can manually adjust axes: axis([a,b,c,d])
figure(2)
hold off
plot(x1, V1arr, 'k-');
title('Eigenvalues and Eigenfunctions', 'fontsize', 20)
xlabel('Position', 'fontsize', 18)
ylabel('Energy', 'fontsize', 18)
hold on
color = 'gcbm';
ncolors = numel(color);
nextcolor = 'r'; % red is used only for the ground state.
% size_psi = size(psi) % diagnostic output
for i = 1:n
   plot( [-w/2, w/2], [EE(i), EE(i)], [nextcolor, '-o'] )
   % Scale psi for plotting. Our soln may blow up in the evanescent region.
   evanescent = ( EE(i) < V1arr ); % Locate where soln is unstable.
   psiscale = (-E_lowest/n) / max( abs(psi(i,!evanescent)) );
   psi_scaled = psi(i,:) * psiscale + EE(i);
   plot( x, psi_scaled, nextcolor )
   nextcolor = color(1 + mod(i-1, ncolors));
end
legend('{V_1(x)}', '{E_0}', '{\psi_0 (x)}')
legend('boxon')
legend('location','northeast outside')
hold off