% Solves the Schrodinger equation 
%  y_t = i y_{xx} - i V(x) y
% by a 'split-step' spectral method. Each term, on its own, can be
% integrated exactly by using a Fourier transform in t and, for the first
% one, in x. Error comes in through the discretisation of x.
% Jonny Tsang (2016)

xmax = 32;
nx = 512;
dx = xmax / nx;
xs = linspace(-xmax/2, xmax/2 - dx, nx);

kmax = pi/(dx);
dk = pi/(xmax);
ks = [ 0:dk:(kmax/2-dk), (-kmax/2):dk:-dk ];

tmax = 20;
dt = 0.005;
ts = 0:dt:tmax;
ws = 0:(pi/(2*tmax)):(pi/(2*dt));

V = @(x) 7*sech(x).^2;
psi0 = @(x) 1./(0.1 + (x+7).^2) .* exp(4i*x);

%%
Vs = arrayfun(V, xs);
psis = zeros([length(ts), length(xs)]);
psis(1,:) = arrayfun(psi0, xs);
psis(1,:) = psis(1,:) / sqrt(sum(abs(psis(1,:)).^2));
for p = 1:length(ts)-1
    % First deal with the _{xx} term by FTing in x
    psi_intermediate = ifft ( fft( psis(p,:)) .* exp(-1i * ks.^2 * dt) );
    % Then time-integrate the V part exactly. 
    psis(p+1,:) = psi_intermediate .* exp(- 1i * dt * Vs);
end

%% 
subplot(2,1,1);
mesh(xs,ts,  log ( abs(psis).^2 ));
view(2); colormap jet; colorbar; caxis([-15,-3]);
title('log-probability distribution of particle position');
xlabel('x');
ylabel('t');
axis tight;

%%
subplot(2,1,2);
ks_rot = circshift(ks, nx/2, 2);
spectrum = abs( fft(psis,[],2) ).^2;
spectrum_rot = circshift(spectrum, nx/2, 2);
mesh(ks_rot, ts, log(spectrum_rot));
xlim([-5,5]);
view(2); colormap jet; colorbar; caxis([-5,4]);
title('log-probability distribution of wavenumber (particle momentum)');
xlabel('wavenumber');
ylabel('t');