orbits.jl

# model two bodies orbiting each other with a bunch of near zero mass particles around them
# the first time you run this, Julia will ask you if you want to install them (say yes) and
# then spend a fair bit of time compiling them. That won't happen on subsequent loads.

"""
To load and run this code, start Julia and then "include" this file:
```
julia> include("orbits.jl")
┌ Info: start
│   size(u0) = (408,)
│   minimum(θ) = 0.0
└   maximum(θ) = 6.283185307179586
┌ Info: progress
└   t = 10.036284179201633
┌ Info: progress
└   t = 20.04042710154572
┌ Info: progress
└   t = 30.043443609743203
┌ Info: progress
└   t = 40.046944314371345
...
```
"""

using SciMLBase, LinearAlgebra, OrdinaryDiffEq, DiffEqCallbacks, ADTypes, OrdinaryDiffEq, Plots
"""
Calculate the acceleration on object i due to object j
"""
function F(u, i, j)
    if i == j
        return [0, 0]
    else
        dr = u[:, j] .- u[:, i]
        ur = dr / norm(dr)
        # force divided by mass of object i
        return (G * m[j] / dot(dr, dr) ) * ur
    end
end

max_t = 0

"""
gravity updates the left-hand side of the differential system governing
a bunch of gravitationally interacting masses.

du is the derivative of the current state (updated)
u  is the current state
p is unused, but required by the diff eq framework
t is time

The current state is made up of 2n values representing positions and 2n
values representing velocities in a single vector. These can be reshaped
for presentation.
"""
function gravity!(du, u, p, t)
    global max_t
    if t ≥ max_t + 10
        max_t = t
        @info "progress" t
    end
    du = reshape(du, 2, :)
    u = reshape(u, 2, :)
    n = size(u, 2) ÷ 2

    du[:, 1:n] = u[:, n+1:2n]
    for i in 1:n
        du[:, n+i] = sum(j -> F(u, i, j), 1:2)
    end
end

G = 1

# we have two big masses and lots of little ones 
m = [10, 100, 1e-6 * ones(15)...]

# the little ones are in a spiral that is in an orbit that will take them near the big ones
n = 100
θ = LinRange(0,2π,n)
asteroids = 6 * [cos.(θ) sin.(θ)]' .* LinRange(0.9, 1.1, n)'
asteroid_v = 0.7 * [0 -1; 1 0] * asteroids
u0 = [2.0;0.0; -0.2;0.0; vec(asteroids); 0.0;6.5; 0.0;-0.65; vec(asteroid_v)]
@info "start" size(u0) minimum(θ) maximum(θ)
tspan = (0.0, 400.0)
prob = ODEProblem(gravity!, u0, tspan)

# compute the orbits of everything. The state at any time
# `t` can be interpolated using `sol(t)`
sol = solve(prob, TsitPap8(); reltol=5e-6)

# now animate it. The heavy masses are red, the rest are green
col = [:red, :red, [:lightgreen for i in 1:100]...]

# this runs quite a while and produces an animation too big to upload to bsky
# adjust the frame interval or ending time to get different speeds and lengths
@gif for t in 1:0.025:200
    let z = reshape(sol(t), 2, :)
        scatter(z[1,1:102], z[2,1:102], xlim=(-20,20), ylim=(-20,20), size=(600,600), color=col, legend=false)
    end
end