% A test of adaptive rk23 (rk2.m, twostep.m, adaptive.m, adaptive_stepper.m)
% as compared to LSODE, doing the N-body problem with randomized initial
% conditions. Based on Nbody_test1.
%
% AS LSODE DOES NOT EXIST IN MATLAB, THIS EXAMPLE RUNS ONLY IN OCTAVE.
%
% OUTPUT FROM A TYPICAL RUN:
%    octave:4> Nbody_test3(3)
%    mass =
%
%       0.30435   0.78222   0.50446
%
%    K =  1.2263
%    U = -0.91073
%    E =  0.31554
%    Px =    2.7756e-17
%    Py =   -2.2204e-16
%    time_seconds_lsode =  2.0119
%    time_seconds_CCK =  18.088
%
% CCK 2009-Feb-26
%


function Nbody_test3(Nbodies)

% Setup for the particular ODE problem
ODE = @Nbody;
m = exp( randn(1,Nbodies) );  % random masses, log-normally distributed
Nbody(0,0,m);
t0 = 0; % initial time
tfinal = 50; % final time

% Initialize the state vector
q0 = randn(1,4*Nbodies);  % random initial state
[K, U, E, Px, Py] = NbodyInvariants(q0, m); % calculate momentum & energy

% Modify the state vector by transforming into the center-of-mass frame
u_cm = Px / sum(m); % center of mass velocity (x component)
v_cm = Py / sum(m); % center of mass velocity (y component)
q0( (1:Nbodies) + (2*Nbodies) ) = q0( (1:Nbodies) + (2*Nbodies) ) - u_cm;
q0( (1:Nbodies) + (3*Nbodies) ) = q0( (1:Nbodies) + (3*Nbodies) ) - v_cm;
[K, U, E, Px, Py] = NbodyInvariants(q0, m) % recalculate momentum & energy



% Setup for CCK's adaptive solver
stepper = @rk2;
order = 2;
% These next 3 parameters I keep changing to try and push to higher Nbodies...
Nsteps = 10000;
tolerance = [1e-3, 1e-3];
maximum_recursion_depth = 20;


% Setup and integrate with LSODE
time_begin_run = cputime;
tlsode = tfinal * (0:Nsteps)/Nsteps;
qlsode = lsode2(ODE, q0, tlsode);
time_seconds_lsode = cputime - time_begin_run

% Work out time series of system energies and momenta.
for n=1:Nsteps+1  % includes initial conditions + subsequent time steps.
   [Klsode(n), Ulsode(n), Elsode(n), Pxlsode(n), Pylsode(n)] = ...
            NbodyInvariants( qlsode(n,:), m );
end

% Plots
figure(1)
hold off
style = ['b-';'k-';'r-';'g-';'m-';'c-';'y-']; % Array of line styles
for body = 1:Nbodies
   plot(qlsode(1,body), qlsode(1,body+Nbodies), [style(body) '+'])
   hold on
   plot(qlsode(:,body), qlsode(:,body+Nbodies), style(body))
end
title('Gravitational N-body simulation','fontsize',18)
% legend('10 quality-controlled RK23 steps','Analytic soln')
xlabel('x','fontsize',18)
ylabel('y','fontsize',18)
axis("equal")  % isotropy is desirable.
hold off

figure(2)
plot(tlsode, Ulsode, 'b-', tlsode, Klsode, 'g-', tlsode, Elsode, 'r-')
title('Gravitational N-body simulation','fontsize',18)
xlabel('Time','fontsize',18)
ylabel('Energy','fontsize',18)
legend('Potential','Kinetic','Total')

figure(3)
plot(tlsode, Elsode, 'r-')
title('Gravitational N-body simulation','fontsize',18)
xlabel('Time','fontsize',18)
ylabel('Total Energy','fontsize',18)

figure(4)
plot(tlsode, Pxlsode, 'b-', tlsode, Pylsode, 'g-')
title('Gravitational N-body simulation','fontsize',18)
xlabel('Time','fontsize',18)
ylabel('Momentum','fontsize',18)
legend('P_x','P_y')




% Integrate with adaptive RK23.
time_begin_run = cputime;
[q,t, greatest_depth] = adaptive_stepper(ODE, q0, t0, tfinal, Nsteps,... 
   stepper, order, tolerance, maximum_recursion_depth);
time_seconds_CCK = cputime - time_begin_run

% Work out time series of system energies and momenta.
for n=1:Nsteps+1  % includes initial conditions + subsequent time steps.
   [K(n), U(n), E(n), Px(n), Py(n)] = NbodyInvariants( q(n,:), m );
end

% Plots
figure(5)
hold off
style = ['b-';'k-';'r-';'g-';'m-';'c-';'y-']; % Array of line styles
for body = 1:Nbodies
   plot(q(1,body), q(1,body+Nbodies), [style(body) '+'])
   hold on
   plot(q(:,body), q(:,body+Nbodies), style(body))
end
title('Gravitational N-body simulation','fontsize',18)
% legend('10 quality-controlled RK23 steps','Analytic soln')
xlabel('x','fontsize',18)
ylabel('y','fontsize',18)
axis("equal")  % isotropy is desirable.
hold off

figure(6)
plot(t, U, 'b-', t, K, 'g-', t, E, 'r-')
title('Gravitational N-body simulation','fontsize',18)
xlabel('Time','fontsize',18)
ylabel('Energy','fontsize',18)
legend('Potential','Kinetic','Total')

figure(7)
plot(t, E, 'r-')
title('Gravitational N-body simulation','fontsize',18)
xlabel('Time','fontsize',18)
ylabel('Total Energy','fontsize',18)

figure(8)
plot(t, Px, 'b-', t, Py, 'g-')
title('Gravitational N-body simulation','fontsize',18)
xlabel('Time','fontsize',18)
ylabel('Momentum','fontsize',18)
legend('P_x','P_y')
