% Find eigenfunctions (normal modes) of the particle in a box. The PDE is *not* provided
% by particle_box.m, because the dispersion relation of the Schrödinger equation has
% the unfortunate property that the highest frequencies have the highest phase speed---
% leading to an extra power of N in scaling the resource requirements.
%
% Instead, we use the scalar wave equation which has a constant phase speed.
% CCK 2015-Apr-22

% Setup
physics = @waveq_box  % system of equations to solve
N = 128                % resolution (number of points)
L = 1.0               % length of domain
c = 1.0               % wave speed
Ncrossings = 256      % number of wave crossing times to run forward in time
Safety = 16            % safety factor (as compared to Courant condition)
   % ?? I must be doing something wrong. Need at least 8 for a reasonable run at N=16!
dx = L/(N-1);         % uniform sampling pitch
x = 0:dx:L;           % x-coordinate, for plotting
t0 = 0;               % initial time
tcrossing = L/c;      % time required for a wave to cross the domain
t1 = tcrossing * Ncrossings;  % final time
dt = dx/c/Safety      % Courant condition with Safety factor thrown in
Nt = ceil((t1 - t0)/dt)     % Number of timesteps.
%epsilon = 0.5*dt       % for stability of rk2plus.
%solver  = @(ode, x0, t0, h) rk2plus(ode, x0, t0, h, epsilon)   % ODE integrator to use (see ode_fixed.m)
solver = @rk2

% Initial conditions. We use white noise for the initial wave function, so that presumably
% it will contain components of all eigenmodes.
z0 = normrnd(0,1,1,N); 
% Dirichlet BCs on the wave function:
z0(1) = 0;
z0(N) = 0;
v0 = zeros(1,N);
q0 = [z0, v0]; % Total initial state

% Integrate.
[q, t] = ode_fixed(solver, physics, q0, t0, t1, Nt-1);

% Fourier transforms wrt t of z and v
zpower = zeros(Nt,1); % storage for power spectrum
vpower = zpower;
omega = k_array(Nt, dt);
window = hanning(Nt);

% Get Fourier transform of the motion, and estimate power spectrum.
z = q(:,1:N);
v = q(:,N+1:2*N);
zf = zeros(Nt,N);
vf = zeros(Nt,N);
for n=1:N
   zf(:,n) = fft(z(:,n) .* window);
   vf(:,n) = fft(v(:,n) .* window);
   zpower += abs( zf(:,n) ).^2;
   vpower += abs( vf(:,n) ./ (omega') ).^2;
end

% Identify boundaries between modes as positive-going roots in the power spectrum slope,
% i.e. as local minima in the power spectrum.
Nfreq = Nt/2/Safety; % Number of "interesting" non-negative frequencies.
ps_slope = (zpower(2:Nfreq)-zpower(1:Nfreq-1))'./(omega(2:Nfreq)-omega(1:Nfreq-1));
   % Slope of power spectrum. 
mode_bounds = 2; % first mode begins at DC+1.
m=1; % use this index for counting subsequent mode moundaries.
for n=3:Nfreq-1
   if (ps_slope(n)*ps_slope(n-1) < 0) 
      if (ps_slope(n) < 0) % found a peak!
         mode_peaks(m) = n;
      elseif (ps_slope(n) > 0) % found a local minimum, a boundary between modes.
         m++;
         mode_bounds(m) = n;
      end
   end
end
Nmodes = m % Number of modes identified.
mode_bounds(m+1) = mode_peaks(m) + mode_bounds(m) - mode_bounds(m-1);
   % A reasonable guess of where to bound the highest mode.

% Plot power vs. temporal frequency, with an overlay of the mode peaks and bounds.
figure(1)
hold off;
semilogy(omega(1:Nfreq),zpower(1:Nfreq),'k','linewidth',2);
hold on;
title('Wave equation modes in 1D box, c=1, L=1')
xlabel('\omega (cycles per crossing time)');
ylabel('power (arbitrary units)');
% Label the individual modes...
y0 = max(zpower(1:Nfreq))*2;
for n=1:Nmodes-1
   x0 = omega(mode_peaks(n));
   x1 = omega(mode_peaks(n+1));
   y1 = zpower(mode_peaks(n));
   y2 = zpower(mode_peaks(n+1));
   plot ( [x0,x0,x1,x1], [y1,y0,y0,y2], 'b' );
end
% Now label the individual mode boundaries...
y0 = min(zpower(1:Nfreq))/2;
for n=1:Nmodes
   x0 = omega(mode_bounds(n));
   x1 = omega(mode_bounds(n+1));
   y1 = zpower(mode_bounds(n));
   y2 = zpower(mode_bounds(n+1));
   plot( [x0,x0,x1,x1], [y1,y0,y0,y2], 'g' );
end
print('particle_box_modes.pdf','-dpdf');

% Determine & plot the shape of each mode.
for n=1:Nmodes
   mode = real(zf(mode_peaks(n),:));
   mode /= max(abs(mode));
   figure(2)
   hold off
   plot(x,mode);
   print(['particle_box_modes/mode_', num2str(n), '.pdf'],'-dpdf');
end


