Mercurial > lbo > hg > lattice2d
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