%N-body gravity solver. Determines fixed-timestep x,y,z trajectories and
%vx,vy,vz velocities of N masses, interacting via Newtonian gravitation
% LAST UPDATED by Andy French July 2020
function varargout = gravity( varargin )
if nargin==0
    %If no inputs, run demo simulation
    binary_stars_and_planet;
else
    %Get inputs
    dt = varargin{1}; tmax = varargin{2};
    x0 = varargin{3}; y0 = varargin{4}; z0 = varargin{5};
    vx0 = varargin{6}; vy0 = varargin{7}; vz0 = varargin{8};
    M = varargin{9};
    
    %Run gravity solver
    [t,x,y,z,vx,vy,vz] = gravity_solver( dt,tmax,x0,y0,z0,vx0,vy0,vz0,M );
    
    %Return outputs of gravity solver
    varargout{1} = t; varargout{2} = x; varargout{3} = y; varargout{4} = z;
    varargout{5} = vx; varargout{6} = vy; varargout{7} = vz;
end
    
%%

% x,y,z in AU, t in Yr, velocities in AU/Yr, M in solar masses.
% Gravity code using a Verlet (constant acceleration between time steps) method.
% Outputs are NxK matrices. N is the number of time steps, K is the number
% of masses. Input initial positions and velocities are 1xK vectors.
function [t,x,y,z,vx,vy,vz] = gravity_solver( dt,tmax,x0,y0,z0,vx0,vy0,vz0,M )

%Determine time t
K = length(M); t = ( 0 : dt : tmax ).'; N = length(t); t = repmat(t,[1,K] );

%Initialize outputs
x = repmat( x0, [N,1] ); y = repmat( y0, [N,1] ); z = repmat( z0, [N,1] );
vx = repmat( vx0, [N,1] ); vy = repmat( vy0, [N,1] ); vz = repmat( vz0, [N,1] );

%%

%Separate out zero masses and non-zero masses
zm = find(M==0); nzm = find(M~=0);

%Step through time array and determine positions and velocities from
%accelerations using the Verlet constant acceleration between time-steps
%numerical method
for n=1:N-1
    [ax,ay,az] = newton( x(n,:),y(n,:),z(n,:),M,zm,nzm );
    x(n+1,:) = x(n,:) + vx(n,:)*dt + 0.5*ax*(dt^2);
    y(n+1,:) = y(n,:) + vy(n,:)*dt + 0.5*ay*(dt^2);
    z(n+1,:) = z(n,:) + vz(n,:)*dt + 0.5*az*(dt^2);
    [aax,aay,aaz] = newton( x(n+1,:),y(n+1,:),z(n+1,:),M,zm,nzm );
    vx(n+1,:) = vx(n,:) + 0.5*dt*( ax + aax );
    vy(n+1,:) = vy(n,:) + 0.5*dt*( ay + aay );
    vz(n+1,:) = vz(n,:) + 0.5*dt*( az + aaz );  
end

%%

%Determine acceleration based upon Newton's universal gravitation
function [ax,ay,az] = newton( x,y,z,M,zm,nzm )
dmin = 0.1; %Set minimum separation /AU for Newtonian gravity to operate

%Initialize output accelerations / AU*Yr^-2
s = size(x); ax = zeros(s); ay = zeros(s); az = zeros(s);

%Determine number of non-zero masses (K1) and zero masses (K2)
K1 = length(nzm); K2 = length(zm);
if K1==0; return; end

%% Calculate accelerations felt by zero-masses %%
if K2>0
    
    %Determine a matrix of separations between the zero masses and the non-zero
    %masses
    dij = distance( [ x(nzm);y(nzm);z(nzm) ], [ x(zm);y(zm);z(zm) ] );
    
    %Make an associated mass matrix
    Mj = repmat( M(nzm).', [1,K2] );
    
    %Determine the x,y,z acceleration components using Newton's law
    %of Universal Gravitation.
    
    %x
    xi = repmat( x(zm), [K1,1] ); xj = repmat( x(nzm).', [1,K2] );
    axij = (-4*pi^2)*Mj.*( xi - xj ).*(dij.^-3);
    axij( dij < dmin ) = 0;
    ax(zm) = sum( axij, 1 );
    
    %y
    yi = repmat( y(zm), [K1,1] ); yj = repmat( y(nzm).', [1,K2] );
    ayij = (-4*pi^2)*Mj.*( yi - yj ).*(dij.^-3);
    ayij( dij < dmin ) = 0;
    ay(zm) = sum( ayij, 1 );
    
    %z
    zi = repmat( z(zm), [K1,1] ); zj = repmat( z(nzm).', [1,K2] );
    azij = (-4*pi^2)*Mj.*( zi - zj ).*(dij.^-3);
    azij( dij < dmin ) = 0;
    az(zm) = sum( azij, 1 );
end

%% Calculate accelerations felt by non-zero-masses %%

%Determine a matrix of separations between the zero masses and the non-zero
%masses
dij = distance( [ x(nzm);y(nzm);z(nzm) ], [ x(nzm);y(nzm);z(nzm) ] );

%Make an associated mass matrix
Mj = repmat( M(nzm).', [1,K1] );

%Determine the x,y,z acceleration components using Newton's law
%of Universal Gravitation.

%x
xi = repmat( x(nzm), [K1,1] ); xj = repmat( x(nzm).', [1,K1] );
axij = (-4*pi^2)*Mj.*( xi - xj ).*(dij.^-3);
axij( dij < dmin ) = 0;
ax(nzm) = sum( axij, 1 );

%y
yi = repmat( y(nzm), [K1,1] ); yj = repmat( y(nzm).', [1,K1] );
ayij = (-4*pi^2)*Mj.*( yi - yj ).*(dij.^-3);
ayij( dij < dmin ) = 0;
ay(nzm) = sum( ayij, 1 );

%z
zi = repmat( z(nzm), [K1,1] ); zj = repmat( z(nzm).', [1,K1] );
azij = (-4*pi^2)*Mj.*( zi - zj ).*(dij.^-3);
azij( dij < dmin ) = 0;
az(nzm) = sum( azij, 1 );

%%

%For a set of K1 column vectors of D dimensions (a), and K2 column vectors
%of D dimensions (b), compute a K1xK2 matrix of Euclidean distances between
%each vector. 
function d = distance( a,b )
K1 = size(a,2); K2 = size(b,2);
aa =sum( a.*a,1 ); aa = repmat( aa',[1,K2] );
bb =sum( b.*b,1 ); bb = repmat( bb,[K1,1] );
adotb = a'*b;
d = aa + bb - 2*adotb;
d = sqrt(d);

%This uses |a - b|^2 = |a|^2 + |b|^2 - 2*a.b where a and b are
%vectors. Inspired by code by Roland Bunschoten (University of Amsterdam)
%written in 1999.

%%

% Simulates the circular orbit of a planet around two massive stars.
% Ignore mass of the planet in the gravitational interaction between the
% stars.
function binary_stars_and_planet
a = 3; %Semi-major axis of mutual star orbit in AU
ap = a/3.4; %Planet (initial) circular orbit radius about star 1
theta0 = pi/4; %Initial angle from x axis (anticlockwise) of planet /radians
M1 = 4; M2 = 4; %Masses of stars in solar masses
k = 1.2;  %Initial vy velocity multiplier from circular

num_periods = 50; %Number of orbital periods
dt = 0.005; %Timestep in years
fsize = 18; %Fontsize
limit = 2*a; %Axes limits

%Determine period in years via Kepler III for binary stars, and hence time
%limit for simulation
T = sqrt( (1/(M1+M2))*a^3 ); tmax = num_periods*T;

%Set initial x,y,vx,vy,t parameters

%Stars 1 and 2
X1 = -M2*a/(M1+M2); Y1 = 0; VX1 = 0;  VY1 = k*2*pi*X1/T;
X2 = M1*a/(M1+M2); Y2 = 0; VX2 = 0;  VY2 = k*2*pi*X2/T;

%Planet
Xp = X1 + ap*cos(theta0); Yp = ap*sin(theta0); 
VXp = -2*pi*sqrt( M1/ap)*sin(theta0); 
VYp = 2*pi*sqrt( M1/ap)*cos(theta0) + VY1;

%Determine trajectories and velocities
M = [M1,M2,0]; x0 = [X1,X2,Xp]; y0 = [Y1,Y2,Yp]; z0 = [0,0,0];
vx0 = [VX1,VX2,VXp]; vy0 = [VY1,VY2,VYp]; vz0 = [0,0,0];
[t,x,y,z,vx,vy,vz] = gravity_solver( dt,tmax,x0,y0,z0,vx0,vy0,vz0,M );
X1 = x(:,1); X2 = x(:,2); x = x(:,3);
Y1 = y(:,1); Y2 = y(:,2); y = y(:,3);

%Plot trajectory
figure('color',[1 1 1],'name','star and planet','units','normalized',...
    'position',[0 0 1 1],'renderer','painters');
plot(x,y,'g'); hold on
plot(X1,Y1,'r'); plot(X2,Y2,'b');
s1 = plot( X1(1),Y1(1),'r*','markersize',30 );
s2 = plot( X2(1),Y2(1),'b*','markersize',30 );
p = plot(x(1),y(1),'ko','markersize',10,'markerfacecolor','k');
axis equal; axis manual; grid on;
xlim([-limit,limit]); ylim([-limit,limit]);
xlabel('x /AU','fontsize',fsize); ylabel('y /AU','fontsize',fsize);
title(['M1=',num2str(M1),', M2=',num2str(M2),...
    ' T=',num2str(T,3),' years, a=',num2str(a,3),'AU, k=',num2str(k),...
    ', a_p=',num2str(ap,3),'AU.'],'fontsize',fsize)
set(gca,'fontsize',fsize);

%Print a PNG file
print( gcf,'binary stars & planet.png','-dpng','-r300' );

%Animation
stop = 0; N = length(x); n=1;
while stop == 0
    try
        set( p, 'xdata', x(n), 'ydata', y(n) );
        set( s1, 'xdata', X1(n), 'ydata', Y1(n) );
        set( s2, 'xdata', X2(n), 'ydata', Y2(n) );
        drawnow; n = n+1; if n>N; n=1; end
    catch
        return
    end
end

%End of code