% Momentum, kinetic and potential energy for Gravitational N-body problem.
% CCK 2009-Feb-23
function [K, U, E, Px, Py] = NbodyInvariants(q, m)
% q is the state, m is the N-element array of masses.
% The state vector is arranged as follows:
%     q = [ x1, x2,... xN,...
%           y1, y2,... yN,...
%           u1, u2,... uN,...
%           v1, v2,... vN ]
% The gravitational constant is taken to be G=1.
%
% The results are:
%  K = total kinetic energy of the system.
%  U = total potential energy.
%  E = K + U.

N = numel(m); % number of bodies

% velocity
u = q( (1:N)+(2*N) );
v = q( (1:N)+(3*N) );

% kinetic energy and momentum
K  = sum( m .* (u.^2 + v.^2) ) / 2;
Px = sum( m .* u );
Py = sum( m .* v );

% potential energy
U = 0; % initialize
for n = 1:N % Loop through all N bodies.
   % Find sum of forces on nth body:
   rx = q(n) - q(1:N);          % r vector x-components
   ry = q(n+N) - q( (1:N)+N );  % r vector y-components
   r = sqrt(rx.^2 + ry.^2);     % |r|
   r(n) = inf;                  % eliminate the self-energy term
   U = U - sum( m(n) * m ./ r )/2; % each term is counted twice, hence the /2.
end

E = U + K; % total energy

