changeset 9:4dcde0992c2b

Add a few more DE solvers
author Lewin Bormann <lbo@spheniscida.de>
date Sun, 09 Jan 2022 11:27:04 +0100
parents a9a74f7cee66
children 9356247f3999
files .hgignore de/cho.jl de/manybody.jl
diffstat 3 files changed, 170 insertions(+), 0 deletions(-) [+]
line wrap: on
line diff
--- a/.hgignore	Sun Jan 09 11:11:05 2022 +0100
+++ b/.hgignore	Sun Jan 09 11:27:04 2022 +0100
@@ -1,1 +1,2 @@
 .gif
+.svg
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/de/cho.jl	Sun Jan 09 11:27:04 2022 +0100
@@ -0,0 +1,33 @@
+using DifferentialEquations
+using Plots
+theme(:ggplot2)
+
+function cho(du, u, p, t)
+    du[1] = u[3]
+    du[2] = u[4]
+    du[3] = -p[:k]/p[:m] * (2*u[1] - u[2])
+    du[4] = -p[:k]/p[:m] * (-u[1] + 2*u[2])
+end
+
+function plot_chain_animation(n::Int, solution)
+    xs = 2 .* collect(1:n)
+    ys = zeros(n)
+
+    anim = @animate for (i,u) in enumerate(solution.u)
+        pos = u[1:n]
+        plot(vcat([0], xs .+ pos, [6]), ones(n+2), xlim=(0, 2*(n+1)), ylim=(.9, 1.1), size=(800, 200), title="t = $(solution.t[i])", legend=false, markersize=8, markershape=:circle)
+    end
+    gif(anim, "anim_$(n).gif", fps=12)
+end
+
+p = Dict(:k => 1., :m => 1.)
+u0 = [1/sqrt(2), 1/sqrt(2), 0, 0]
+
+problem = ODEProblem(cho, u0, (0, 3 * 2pi/sqrt(1)), p)
+@time solution1 = solve(problem)
+
+plot(
+    plot(solution1, vars=[(0, 1), (0, 2)], xlabel="t", ylabel="x", label=["\$1: x_1(t)\$" "\$1: x_2(t)\$"]),
+    plot(solution2, vars=[(0, 1), (0, 2)], label=["\$2: x_1(t)\$" "\$2: x_2(t)\$"]),
+    layout=(2,1))
+#plot(solution2, vars=[(0, 1), (0, 2)], label=["\$2: x_1(t)\$" "\$2: x_2(t)\$"])
\ No newline at end of file
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/de/manybody.jl	Sun Jan 09 11:27:04 2022 +0100
@@ -0,0 +1,136 @@
+using DifferentialEquations
+using Plots
+
+
+# Many-body, 2D.
+
+const F = Float64
+const D = 2
+const G = 6.6743e-11  # m³/kg s²
+
+const M_S = 2e30
+const M_E = 6e24
+const M_M = 7.35e22
+
+const R_E = 150e9
+const R_M = 400e6
+
+mutable struct MBSystem
+    state::Matrix{F}
+    masses::Vector{F}
+    G::F
+end
+
+# TODO: Write a recipe for this.
+function plot_mbs(mbs::MBSystem)
+    scatter(mbs.state[1,:]', mbs.state[2,:]', markersize=3 .+ 10 * log.(mbs.masses'), legend=:false)
+    for i in 1:length(mbs.masses)
+        p0 = mbs.state[1:D, i]
+        p1 = p0 + mbs.state[D+1:end, i]
+
+        plot!([p0[1], p1[1]], [p0[2], p1[2]], width=2)
+    end
+    current()
+end
+
+function plot_solution(mbs::MBSystem, solution)
+    p0 = mbs.state[1:D, :]
+    vars = [(2D*(i-1) + 1, 2D*(i-1) + 2) for i in 1:length(mbs.masses)]
+    plot(solution, vars=vars, labels=[string(i) for j in [1], i in (1:length(mbs.masses))], size=(1000,1000))
+    current()
+end
+
+function animate_solution(mbs::MBSystem, solution; filename="orbit.gif")
+    p0 = mbs.state[1:D, :]
+    vars = [(2D*(i-1) + 1, 2D*(i-1) + 2) for i in 1:length(mbs.masses)]
+    anim = @animate for t in solution.t
+        plot(solution, vars=vars, tspan=(0, t), labels=[string(i) for j in [1], i in (1:length(mbs.masses))], size=(1000,1000), width=5)
+    end
+    gif(anim, filename)
+end
+
+function new_mb(n::Int; masses=nothing, G=G)::MBSystem
+    if !isnothing(masses) && length(masses) != n
+        error("Length of masses must equal n")
+    end
+    MBSystem(zeros(2D, n), isnothing(masses) ? zeros(n) : masses, G)
+end
+
+function mb_set(mbs::MBSystem, i::Int, pos::Vector, spd::Vector=[0,0])
+    mbs.state[1:D, i] .= pos;
+    mbs.state[D+1:end, i] .= spd;
+    nothing
+end
+
+function radial_equilibrium_speed(mbs::MBSystem, i::Int, r::F)
+    # v²/r = G M / r² → v = √(G M / r)
+    return sqrt(mbs.G * mbs.masses[i] / r) 
+end
+
+# Return ̂r / r^2 for bodies i, j where r = r_i - r_j
+function radial_term(m::Matrix{F}, i::Int, j::Int)::Vector
+    r = m[1:D, j] - m[1:D, i]
+    l = sqrt(sum(r.^2))
+    r ./ l^3
+end
+
+# Calculate vector of acceleration according to Newton's gravity law.
+# Considers all bodies j ≠ i.
+# m a = G (m M / r^2) * ̂r
+function acceleration(mbs::MBSystem, i::Int; state=nothing)::Vector{F}
+    n = length(mbs.masses)
+    a = zeros(D)
+    st = isnothing(state) ? mbs.state : state
+    for j = 1:n
+        if i == j
+            continue
+        end
+        rt = radial_term(st, i, j)
+        a .+= rt * mbs.G * mbs.masses[j]
+    end
+    a
+end
+
+function DE(du, u, p, t)
+    # Link velocities. d{x,y}/dt = v{x,y}
+    du[1:D, :] .= u[D+1:end, :]
+    # Calculate force application
+    for i in 1:length(p.masses)
+        du[D+1:end, i] .= acceleration(p, i; state=u)
+    end
+end
+
+function solve_sme()
+    MB_SME = new_mb(3, masses=[M_S, M_E, M_M], G=G)
+    mb_set(MB_SME, 1, [0., 0], [0., 0.])
+    mb_set(MB_SME, 2, [0., -R_E], [1.1*radial_equilibrium_speed(MB_SME, 1, R_E), 0])
+    mb_set(MBSME, 3, [0., -(R_E+R_M)], [1.1*radial_equilibrium_speed(MB_SME, 1, R_E) + radial_equilibrium_speed(MB0, 2, R_M), 0])
+
+    problem = ODEProblem(DE, MB_SME.state, (0, 365*24*3600), MB_SME)
+    solution = solve(problem)
+
+    plot_solution(MB_SME, solution)
+end
+
+function solve_me()
+    MB_ME = new_mb(2, masses=[M_E, M_M], G=G)
+    mb_set(MB_ME, 1, [0., 0], [0., 0])
+    mb_set(MB_ME, 2, [0., -R_M], [1.1*radial_equilibrium_speed(MB_ME, 1, R_M), 0])
+
+    problem = ODEProblem(DE, MB_ME.state, (0, 365*24*3600), MB_ME)
+    solution = solve(problem)
+
+    plot_solution(MB_ME, solution)
+end
+
+function solve_3body()
+    MB0 = new_mb(3, masses=[40., 2, 1], G=0.002)
+    mb_set(MB0, 1, [0., 0.], [0., 0.])
+    mb_set(MB0, 2, [0., 5.], [-0.15, 0.0])
+    mb_set(MB0, 3, [0., -5.], [0.1, 0.0])
+
+    problem = ODEProblem(DE, MB0.state, (0, 1000), MB0)
+    solution = solve(problem)
+
+    animate_solution(MB0, solution)
+end