%% 6.581 Problem set 2
%% Question 1. Molecular dynamics of a trapped ion
%% moldyn.m

%% Note: you need to fill in the section indicated below

%% Variable parameters:
dt   = 0.01;         % time step of integration
t    = 0.0;          % starting time of simulation
Tend = 2.0;         % length of simulation
atomtype = 1;        % 1 = F-, 2 = Cl-, 3 = Br-
xMol = [ -5.0 0.0 ]; % initial position of ion (X, Y)
vMol = [ 0.0 0.0 ] ; % initial velocity of ion (dX/dt, dY/dt)
marker = 'x';


%% Fixed system parameters:
b      = 10;      % size of box (-b to b in X and Y)
n      = 6;       % number of charges along one side of box
qPin   = -1.0;    % magnitude of pinned charges
sigPin = 2.73295; % VDW sigma for pinned charges
epsPin = 0.72000; % VDW epsilon for pinned charges

%% Constants:
eps = 2.397e-4;          % epsilon naught in AKMA units 
time_unit = 4.888821e-2; % time unit in ps

%% Atomic energy parameters:
%%    m is mass of molecule
%%    qMol is charge of molecule
%%    sigMol is sigma for molecule (VDW)
%%    epsMol is epsilon for molecule (VDW)
if (atomtype == 1)
  % F-
  mMol   = 18.998;
  qMol   = -1.0;
  sigMol = 2.73295;
  epsMol = 0.72000;
elseif (atomtype == 2)
  % Cl-
  mMol   = 35.453;   
  qMol   = -1.0;   
  sigMol = 4.41724;
  epsMol = 0.11779;
elseif (atomtype == 3)
  % Br-
  mMol   = 79.904;   
  qMol   = -1.0;  
  sigMol = 4.62376;
  epsMol = 0.09000;
end

%% Convert time
t = t / time_unit;
dt = dt / time_unit;
Tend = Tend / time_unit;

%% Set up pinned charges
x = linspace(-b,b,n);
[X,Y] = meshgrid(x);
x = [X(1,1:n-1) X(1:n-1,n)' X(n,n:-1:2) X(n:-1:2,1)'];
y = [Y(1,1:n-1) Y(1:n-1,n)' Y(n,n:-1:2) Y(n:-1:2,1)'];
chargeCoord = [x' y'];
nCharge = size(chargeCoord,1);  % nCharge = 4*(n-1)

%% Variable used for computing the force and energy
r      = zeros(length(chargeCoord), 2);
d      = zeros(length(chargeCoord), 1);
f      = zeros(1,2);
fElec  = zeros(1,2);
fVdw   = zeros(1,2);

%% Trajectory variables for plotting later
time       = zeros(floor(Tend/dt) + 1, 1);
xTraj      = zeros(length(time), 2);
vTraj      = zeros(length(time), 2);
e_potTraj  = zeros(length(time), 1);
e_kinTraj  = zeros(length(time), 1);
i          = 1;

%% Calculate initial energy
r(:,1) = xMol(1)-chargeCoord(:,1);
r(:,2) = xMol(2)-chargeCoord(:,2);
d(:)   = sqrt(r(:,1).^2 + r(:,2).^2);

e_Elec = qMol*qPin/(4*pi*eps).*sum(d.^-1);
e_Vdw  = 4*sqrt(epsMol*epsPin) * ...
  sum(((sqrt(sigMol*sigPin)./(d)).^12 - (sqrt(sigMol*sigPin)./(d)).^6));

xTraj(i,:)   = xMol;
vTraj(i,:)   = vMol;

e_potTraj(i) = e_Elec + e_Vdw;
e_kinTraj(i) = 0.5 * mMol * norm(vMol)^2;
e_pot_0      = e_potTraj(1);
e_kin_0      = e_kinTraj(1);


%% Visualization set up
subplot(2,2,1)
title('Trajectory')
hold on
subplot(2,2,3)
title('Energy vs Time')
hold on
subplot(2,2,2)
title('X-position vs Time')
hold on
subplot(2,2,4)
title('Y-position vs Time')
hold on

subplot(2,2,1)
for index = 1:nCharge
  plot([chargeCoord(index,1)-0.5 chargeCoord(index,1)+0.5], ...
      [chargeCoord(index,2) chargeCoord(index,2)],'r-','LineWidth',2,'MarkerSize',8);
  plot(chargeCoord(index,1),chargeCoord(index,2),'ro','LineWidth',2,'MarkerSize',12);
end
axis([-(b+1) (b+1) -(b+1) (b+1)])
axis square
plot(xMol(1),xMol(2),'r.','MarkerSize',24)

subplot(2,2,2)
plot([t*time_unit Tend*time_unit],[0 0],'k-');

subplot(2,2,3)
plot([t*time_unit Tend*time_unit],[0 0],'k-');
plot([t*time_unit Tend*time_unit],[(e_pot_0 + e_kin_0) (e_pot_0 + e_kin_0)],'m-');
plot([t*time_unit],[e_kin_0],'rx');
plot([t*time_unit],[e_pot_0],'gx');

subplot(2,2,4)
plot([t*time_unit Tend*time_unit],[0 0],'k-');
%% END Visualization set up

%%
%% INTEGRATION LOOP
%%

while t < Tend
  %% Force evaluation
  %% DO NOT CHANGE
  r(:,1) = xMol(1)-chargeCoord(:,1);
  r(:,2) = xMol(2)-chargeCoord(:,2);
  d(:)   = sqrt(r(:,1).^2 + r(:,2).^2);
  
  fElec(1) = (qMol*qPin/(4*pi*eps)) * sum(r(:,1)./(d(:).^3));
  fElec(2) = (qMol*qPin/(4*pi*eps)) * sum(r(:,2)./(d(:).^3));
  fVdw(1)  = 4*sqrt(epsMol*epsPin) *                                 ...
      sum((12 * ((sigMol*sigPin)^6)*(d.^-14)                         ...
          - 6 * ((sigMol*sigPin)^3)*(d.^ -8))   .*r(:,1));
  fVdw(2)  = 4*sqrt(epsMol*epsPin) *                                 ...
      sum((12 * ((sigMol*sigPin)^6)*(d(:).^-14)                      ...
          - 6 * ((sigMol*sigPin)^3)*(d(:).^ -8)).*r(:,2));
  f = fElec + fVdw; 
  %% END Energy and force evaluation
  %% The force on the particle at it's current position is stored
  %% in "f".
  
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  % INSERT AN APPROPRIATE INTEGRATION SCHEME HERE %
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  
  %%
  %% END OF INTEGRATION
  %% 
  
  %% Energy evaluation
  %% DO NOT CHANGE
  r(:,1) = xMol(1)-chargeCoord(:,1);
  r(:,2) = xMol(2)-chargeCoord(:,2);
  d(:)   = sqrt(r(:,1).^2 + r(:,2).^2);
  
  e_Elec = qMol*qPin/(4*pi*eps).*sum(d.^-1);
  e_Vdw  = 4*sqrt(epsMol*epsPin) * ...
      sum(((sqrt(sigMol*sigPin)./(d)).^12 - (sqrt(sigMol*sigPin)./(d)).^6));
      
  %% END Energy evaluation
  
  %% Store Quantites we want to plot at the end
  %% DO NOT CHANGE
  i            = i + 1;
  time(i)      = t * time_unit;
  xTraj(i,:)   = xMol;
  vTraj(i,:)   = vMol;
  e_potTraj(i) = e_Elec + e_Vdw;
  e_kinTraj(i) = 0.5 * mMol * norm(vMol)^2;  
end   
%% END INTEGRATION LOOP

%%
%% Trajectory Visualization
%%
% Plot X,Y vs time on model
subplot(2,2,1)
plot(xTraj(:, 1),xTraj(:, 2),'b.','MarkerSize',12)

% Plot X vs time
subplot(2,2,2)
plot(time,xTraj(:,1),'b.','MarkerSize',12)

% Plot energy vs time
subplot(2,2,3)
plot(time,e_potTraj,'r.','MarkerSize',12);
plot(time,e_kinTraj,'g.');
plot(time,e_potTraj+e_kinTraj,'k.');

% Plot Y vs time
subplot(2,2,4)
plot(time,xTraj(:, 2),'b.','MarkerSize',12)
%% END Trajectory Visualization