% [u,x,t,kk,amp,M,H] = implicit_schroedinger(t_0,t_f,dt,NPrint,N)
%
% this solves the periodic Schroedinger equation 
%
%         i u_t + u_xx + |u|^2 u = 0
%
% with periodic boundary conditions using spectral derivatives in
% space and implicit time-stepping.  t_0 is the initial time, t_f is the
% final time, dt is the time-step size, N is the number of mesh-points, 
% and M is the number of printouts.  It also returns the power spectrum of the
% solution, where kk is the collection of wave numbers and amp = log(f_k^2)
% is the magnitude of the kth Fourier mode.
%
% Also, it returns M which is the mass \int |u|^2
% and H which is the hamiltonian \int |u_x|^2 - 1/2 |u|^4
%
% N SHOULD BE A POWER OF 2!  
% (It will work no matter what, but will go much faster if N is a power
% of 2.)

function [u,x,t,kk,amp,M,H] = implicit_schroedinger(t_0,t_f,dt,NPrint,N)

% define the mesh in space
dx = 2*pi/N;
x = 0:dx:2*pi-dx;
x = x';
% the wave-numbers...
kk = (-N/2+1):1:(N/2-1);

% total number of steps to be taken:
Ntot = (t_f-t_0)/dt;
% test to see if Ntot is an integer, to see if your choice of dt
% is compatible with your choice of t_0 and t_f
if abs(Ntot - round(Ntot)) > 100000*eps
  disp(' ')
  disp('dt is not consistent with your choice of t_0 and t_f.  Choose another dt!')
  disp(' ')
  break
end
Ntot = round(Ntot);
% number of steps to be taken before a (screen) printout
NRun = Ntot/NPrint;
% test to see if NRun is an integer to see if your choice of dt is
% consistent with NPrint
if abs(NRun - round(NRun)) > 1000*eps
  disp(' ')
  disp('NPrint is not consistent with dt.  Choose another NPrint!')
  disp(' ')
  break
end
NRun = round(NRun);
% define the vector t that is to be printed out...
t = t_0:NRun*dt:t_f;

% choose initial data
%u_now = cos(x) + 3*i*sin(2*x) - 8*exp(i*3*x);
u_now = exp(i*3*x);

% We're computing solutions of
%
%    u_t = i u_xx + i |u|^2 u
%
% With implicit time-stepping, this would be
%
%  u_new - u_now = dt*i*u_new_xx + i*dt*|u_now|^2 u_now
%
%  u_new - i*dt*u_new_xx = u_now + i*dt*|u_now|^2 u_now
%
% That is, we want to solve
%
%      L u_new = RHS
%
% This is especially easily done spectrally since
%
%  (1 + i dt k^2) u_new_hat(k) = RHS_hat(k)
%
% Hence
%
%  u_new_hat(k) = RHS_hat(k)/(1 + i dt k^2) 

% save first printout:
u(:,1) = u_now;
[M(1),H(1)] = conserved(u_now,dx,N);
v = fft(u_now); 
amp(N/2:N-1,1) = log10( abs(v(1:N/2)) + 1.e-16);
amp(1:N/2-1,1) = log10( abs(v(N/2+2:N)) + 1.e-16);

for jj=1:NPrint

  % take NRun time-steps
  for j=1:NRun

    % compute solution with one step of size dt
    u_coarse = tstep(u_now,dt,N);
    % compute solution with two steps of size dt/2
    u_fine = tstep(u_now,dt/2,N);
    u_fine = tstep(u_fine,dt/2,N);
    % use Richardson extrapolation to combine u_coarse and
    % u_fine to get a solution with local truncation
    % error of O(dt^3)...
    u_now = 2*u_fine - u_coarse;

  end

  % save printout
  u(:,jj+1) = u_now;
  [M(jj+1),H(jj+1)] = conserved(u_now,dx,N);
  v = fft(u_now); 
  amp(N/2:N-1,jj+1) = log10( abs(v(1:N/2)) + 1.e-16);
  amp(1:N/2-1,jj+1) = log10( abs(v(N/2+2:N)) + 1.e-16);
end
x(N+1) = 2*pi;
u(N+1,:) = u(1,:);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function u_now = tstep(u_now,dt,N);

    % start by computing power spectrum of u_now...
    v = fft(u_now);
    % compute i dt |u_now|^2 u_now
    wrk = i*dt*abs(u_now).^2.*u_now;
    % take its FFT and form the right-hand-side:
    RHS = fft(wrk) + v;

    % we now time-step the PDE at the spectral level
    % by solving L v = RHS...
    ii = 1:N/2;
    k = (ii-1)';
    v(ii) = RHS(ii)./(1 + i*dt*k.^2);
    ii = N/2+2:N;
    k = (ii-N-1)';
    v(ii) = RHS(ii)./(1 + i*dt*k.^2);

    % take the inverse fourier transform.  
    u_now = ifft(v);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function [M,H] = conserved(u,dx,N);
%
% use the Trapezoidal rule to compute the mass M and hamiltonian H
%             M = \int |u|^2
%             H = \int |u_x|^2 - 1/2 |u|^4
%

arg = abs(u(1:N)).^2 ;
M = dx*sum(arg);

v = fft(u);
ii = 1:N/2;
k = (ii-1)';
v(ii) = (i*k).*v(ii);
ii = N/2+2:N;
k = (ii-N-1)';
v(ii) = (i*k).*v(ii);
% take the inverse fourier transform.  This gives u_x.
v = ifft(v);

arg = abs(v(1:N)).^2 - 1/2*abs(u(1:N)).^4;
H = dx*sum(arg);