clear all
% Fourier analysis of an N level system

N = 10;
% parameters of the Hamiltonian

En = ones(1,N); % On-site energy in position
b = -6% hopping between the levels

bc = 1;

%discretization parameters
dt = pi/abs(4*b); % step it must be < pi/abs(2*b)
Nt = 40; % number of steps
t = [0:Nt-1]*dt; % time vector



% The Hamiltonian
% The Hamiltonian
H = diag(En) +...
    diag(b*ones(1,N-1),1) + diag(b*ones(1,N-1),-1)

if bc == 1
    H = H + diag(b*ones(1,1),N-1) + diag(b*ones(1,1),-N+1);
end

[U,Delta] = eig(H);

%initial condition (vector -> pure state)
l_1 = 0;
l_2 = 1;

% psi_Pos = exp(i*2*pi/N*l_1*[1:N])+exp(i*2*pi/N*l_2*[1:N]);
% rho_Pos = psi_Pos'*psi_Pos;

% initial condition (energy statistical mixture)
p = rand(1,N);
% I'm  taking advantage that U is real since it diagonalizes an Hermitian
% matrix -> ' = transpose
% rho_Pos = zeros(N);
% for n = 1:N
% rho_Pos = rho_Pos + p(n)*(U(:,n)*U(:,n)');
% end

%initial condition (density in position space)
 rho_Pos = diag([1,zeros(1,N-1)]); %|1><1| initialization
% rho_Pos = rand(N) + i*rand(N); % random initialization



% this procedure takes care that the 
% input for the initial condition 
% is Hermitian and normalized
rho_Pos = (rho_Pos + rho_Pos')/(2*trace(rho_Pos));

%important variables for the analytical solution
rho_En0 = U'*rho_Pos*U; %initial condition 
Bohr = repmat(diag(Delta),1,N) - repmat(diag(Delta)',N,1); % Bohr frequencies 


% generation of the sphere coordinates necessary for the movie 
[X,Y,Z] = sphere(100);

% coordinates of the position of the atoms
R = [cos([1:N]'*2*pi/N),sin([1:N]'*2*pi/N),zeros(N,1)];

% time evolution both in the position and energy basis
for n = 1:Nt 
    rhoEn = rho_En0.*exp(-i*Bohr*t(n));
    rhoPos = U*rhoEn*U';
    % vector of probabilities
    P(n,:) = real(diag(rhoPos));
end

% vector of frequencies that can be detected 
w = linspace(-pi/dt,pi/dt,Nt);

% calculation of the Fourier transformation
for m = 1:length(w)   
   P_w(m,:) = abs(exp(-i*w(m)*t)*P).^2/(Nt)^2;
   m
end

%%
figure(1)
subplot(1,2,1)
plot(t/(2*pi)*abs(b),P(:,1))
ylim([0 1])
set(gca,'Fontsize',18)
xlabel('t [h/|b|]')

subplot(1,2,2)
plot(w/abs(b),P_w(:,1))
hold on
plot(reshape(Bohr,1,N^2)/abs(b),repmat([0;max(max(P_w))],1,N^2),'ko')
xlim([-5 5])
set(gca,'Fontsize',18)
xlabel('\omega [2\pi |b|/h]')
hold off







