function [r,v,a,istop] = newtonRK_JMS( force, mass, r0, v0, ...
                                   dt, tmax, stopCheck )

% Solves Newton's equation, using fourth-order Runge-Kutta method,
% for np particles in nd dimensions. J.M.Soler, Jan.2011 & May.2019
% 
% Input:
%   force     : function to calculate the force as f=force(x,v,t)
%   mass      : particle mass(es)
%               If size(mass)=[1,1] => all masses are equal
%               If size(mass)=[1,np] => different particle masses
%               If size(mass)=[nd,np] => different coordinate masses
%               If size(mass)=[nd*np,nd*np] => mass tensor of generalized
%               coordinates => sum_j(mass(i,j)*a(j))=force(i)
%   r0(nd,np) : intial positions of np particles in nd dimensions
%   v0(nd,np) : initial velocities
%   dt        : integration time interval
%   tmax      : final integration time
%   stopCheck : stopping function called as istop=stopCheck(x,v,t)
%               where istop~=0 will stop the trajectory (optional)
% Output:
%   r(nd,np,nt) : position vectors at each time (0:dt:tmax)
%   v(nd,np,nt) : velocities at each time
%   a(nd,np,nt) : accelerations at each time
%                 If nd=1, the output sizes of x,v,a will be (np,nt)
%                 If np=1, the output sizes will be (nd,nt)
%   istop       : stop code returned by stopCheck, or zero, if not present
% Units:
%   Units are arbitrary, but units of input variables must be consistent
%   (e.g. all in MKS system). Generalized coordinates (e.g. lengths and 
%   angles) may have different units but they must be consistent with
%   corresponding masses (e.g. masses and moments of inertia).
%   Output variables will then be consistent with units of r0 and dt.
%
% Algorithms:
% First we rewrite the second-order Newton equation as a first order
% differential equation, by doubling the number of variables: y=[x,v].
% Then we call function rungeKutta to solve dy/dt=[v,force/m]

% Check array sizes: nd=num_spacial_dimensions, np=num_particles
[nd,np] = size(r0);
if (size(v0)~=size(r0))
  error('newtonRK: ERROR: size(v0)~=size(r0)')
end
if numel(mass)==1             % assume all particles have the same mass
  m = mass*eye(nd*np);        % mass tensor
elseif numel(mass)==np        % different mass for each particle
  m = repmat(mass(:)',nd,1);  % same mass for each coordinate
  m = diag(m(:));
elseif numel(mass)==nd        % different mass for each space coordinate
  m = repmat(mass(:),1,np);   % same mass for each particle
  m = diag(m(:));
elseif numel(mass)==nd*np     % different mass for each generalized coord
  m = diag(mass(:));
elseif numel(mass)==(nd*np)^2 % input mass tensor
  m = reshape(mass,nd*np,nd*np);
else
  error('newtonRK ERROR: wrong size(mass)')
end
mi = inv(m);                  % inverse of mass tensor

% Combine positions and velocities into a single column vector
nx = nd*np;                                  % total number of coordinates
y0 = [ reshape(r0,nx,1); reshape(v0,nx,1) ]; % y combines r and v
y2r = @(y)reshape(y(1:nx),nd,np);            % extracts r from y
y2v = @(y)reshape(y(nx+1:end),nd,np);        % extracts v from y
myderiv = @(t,y)[ y(nx+1:end); mi*reshape( force(y2r(y),y2v(y),t), nx,1) ];

% Integrate trajectory as dy/dt=myderiv(t,y)
if nargin>6   % stopCheck present in input
  mystop = @(t,y)stopCheck(y2r(y),y2v(y),t);
  [y,dydt,istop] = rungeKutta_JMS( myderiv, y0, 0, dt, tmax, mystop );
else          % there is no stop condition
  [y,dydt,istop] = rungeKutta_JMS( myderiv, y0, 0, dt, tmax );
end

% Extract positions, velocities, and accelerations
nt = size(y,2);    % number of times
r = reshape(y(1:nx,:),nd,np,nt);
v = reshape(y(nx+1:end,:),nd,np,nt);
a = reshape(dydt(nx+1:end,:),nd,np,nt);

% Remove dimensions of size 1
r = squeeze(r);
v = squeeze(v);
a = squeeze(a);

end % function newtonRK
