From ff3fe6a8d8866734d5c3815933d492f8076ba25a Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Thu, 23 Jun 2022 15:39:59 -0400 Subject: [PATCH] format SciML Style --- .JuliaFormatter.toml | 1 + .github/workflows/FormatCheck.yml | 42 +++++++++++ examples/double_pendulum.jl | 115 ++++++++++++++++-------------- examples/pendulum.jl | 43 ++++++----- src/hamiltonian.jl | 36 ++++++---- src/plot.jl | 18 ++--- test/hamiltonian_test.jl | 30 ++++---- test/nbody_test.jl | 31 +++++--- 8 files changed, 192 insertions(+), 124 deletions(-) create mode 100644 .JuliaFormatter.toml create mode 100644 .github/workflows/FormatCheck.yml diff --git a/.JuliaFormatter.toml b/.JuliaFormatter.toml new file mode 100644 index 0000000..453925c --- /dev/null +++ b/.JuliaFormatter.toml @@ -0,0 +1 @@ +style = "sciml" \ No newline at end of file diff --git a/.github/workflows/FormatCheck.yml b/.github/workflows/FormatCheck.yml new file mode 100644 index 0000000..2a3517a --- /dev/null +++ b/.github/workflows/FormatCheck.yml @@ -0,0 +1,42 @@ +name: format-check + +on: + push: + branches: + - 'master' + - 'release-' + tags: '*' + pull_request: + +jobs: + build: + runs-on: ${{ matrix.os }} + strategy: + matrix: + julia-version: [1] + julia-arch: [x86] + os: [ubuntu-latest] + steps: + - uses: julia-actions/setup-julia@latest + with: + version: ${{ matrix.julia-version }} + + - uses: actions/checkout@v1 + - name: Install JuliaFormatter and format + # This will use the latest version by default but you can set the version like so: + # + # julia -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter", version="0.13.0"))' + run: | + julia -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter"))' + julia -e 'using JuliaFormatter; format(".", verbose=true)' + - name: Format check + run: | + julia -e ' + out = Cmd(`git diff --name-only`) |> read |> String + if out == "" + exit(0) + else + @error "Some files have not been formatted !!!" + write(stdout, out) + exit(1) + end' diff --git a/examples/double_pendulum.jl b/examples/double_pendulum.jl index 04e029e..1699ad5 100644 --- a/examples/double_pendulum.jl +++ b/examples/double_pendulum.jl @@ -2,7 +2,6 @@ using DifferentialEquations, Plots #Reference: https://diego.assencio.com/?index=e5ac36fcb129ce95a61f8e8ce0572dbf - #Solving the double pendulum with a traditional ODE method #==========================================================# function doublependulum(du, u, params, t) @@ -10,46 +9,47 @@ function doublependulum(du, u, params, t) l2 = params[2] m1 = params[3] m2 = params[4] - g = params[5] + g = params[5] p1 = u[1] p2 = u[2] θ1 = u[3] θ2 = u[4] - h1 = p1*p2*sin(θ1-θ2)/(l1*l2*(m1+m2*sin(θ1-θ2)^2)) - h2 = (m2*l2^2*p1^2 + (m1 + m2) * l1^2 * p2^2 - 2*m2*l1*l2*p1*p2*cos(θ1-θ2))/(2 * l1^2 * l2^2 *(m1 + m2*sin(θ1-θ2)^2)^2) - - d_p1 = -(m1 + m2)* g*l1*sin(θ1) - h1 + h2*sin(2*(θ1-θ2)) - d_θ1 = (l2*p1 - l1*p2*cos(θ1-θ2))/(l1^2 * l2*(m1 + m2*sin(θ1 - θ2)^2)) - d_p2 = -m2*g*l2*sin(θ2) + h1 - h2*sin(2*(θ1-θ2)) - d_θ2 = (-m2*l2*p1*cos(θ1 - θ2) + (m1 + m2)*l1*p2) / (m2*l1*l2^2*(m1 + m2*sin(θ1-θ2)^2)) + h1 = p1 * p2 * sin(θ1 - θ2) / (l1 * l2 * (m1 + m2 * sin(θ1 - θ2)^2)) + h2 = (m2 * l2^2 * p1^2 + (m1 + m2) * l1^2 * p2^2 - + 2 * m2 * l1 * l2 * p1 * p2 * cos(θ1 - θ2)) / + (2 * l1^2 * l2^2 * (m1 + m2 * sin(θ1 - θ2)^2)^2) + d_p1 = -(m1 + m2) * g * l1 * sin(θ1) - h1 + h2 * sin(2 * (θ1 - θ2)) + d_θ1 = (l2 * p1 - l1 * p2 * cos(θ1 - θ2)) / (l1^2 * l2 * (m1 + m2 * sin(θ1 - θ2)^2)) + d_p2 = -m2 * g * l2 * sin(θ2) + h1 - h2 * sin(2 * (θ1 - θ2)) + d_θ2 = (-m2 * l2 * p1 * cos(θ1 - θ2) + (m1 + m2) * l1 * p2) / + (m2 * l1 * l2^2 * (m1 + m2 * sin(θ1 - θ2)^2)) du .= [d_p1, d_p2, d_θ1, d_θ2] return nothing end -l1 = 1. #length of pendulum1 -l2 = 2. #length of pendulum2 -m1 = 1. #mass of pendulum1 -m2 = 1. #mass of pendulum2 +l1 = 1.0 #length of pendulum1 +l2 = 2.0 #length of pendulum2 +m1 = 1.0 #mass of pendulum1 +m2 = 1.0 #mass of pendulum2 g = 9.81 #gravity params = [l1, l2, m1, m2, g] -times = (0., 25.) -u0 = [1.,1.,1.,1.] +times = (0.0, 25.0) +u0 = [1.0, 1.0, 1.0, 1.0] prob = ODEProblem(doublependulum, u0, times, params) -sol1 = solve(prob, AutoVern7(Rodas5()), dt = .005) +sol1 = solve(prob, AutoVern7(Rodas5()), dt = 0.005) #plot solution -plot(sol1, vars=1, xlim=(0,20), label="Momentum1") -plot!(sol1, vars=2, xlim=(0,20), label="Momentum2") -plot!(sol1, vars=3, xlim=(0,20), label="theta1") -plot!(sol1, vars=4, xlim=(0,20), label="theta2") +plot(sol1, vars = 1, xlim = (0, 20), label = "Momentum1") +plot!(sol1, vars = 2, xlim = (0, 20), label = "Momentum2") +plot!(sol1, vars = 3, xlim = (0, 20), label = "theta1") +plot!(sol1, vars = 4, xlim = (0, 20), label = "theta2") #==========================================================# - #Now with HamiltonianProblem() #==========================================================# function H(p, θ, params) @@ -57,58 +57,67 @@ function H(p, θ, params) l2 = params[2] m1 = params[3] m2 = params[4] - g = params[5] + g = params[5] - return (m2*l2^2*p[1]^2 + (m1+m2)*l1^2*p[2]^2 - 2*m2*l1*l2*p[1]*p[2]*cos(θ[1]-θ[2]) ) / - (2*m2*l1^2*l2^2*(m1+m2*sin(θ[1]-θ[2])^2)) - - (m1+m2)*g*l1*cos(θ[1]) - m2*g*l2*cos(θ[2]) + return (m2 * l2^2 * p[1]^2 + (m1 + m2) * l1^2 * p[2]^2 - + 2 * m2 * l1 * l2 * p[1] * p[2] * cos(θ[1] - θ[2])) / + (2 * m2 * l1^2 * l2^2 * (m1 + m2 * sin(θ[1] - θ[2])^2)) - + (m1 + m2) * g * l1 * cos(θ[1]) - m2 * g * l2 * cos(θ[2]) end -l1 = 1. #length of pendulum1 -l2 = 2. #length of pendulum2 -m1 = 1. #mass of pendulum1 -m2 = 1. #mass of pendulum2 +l1 = 1.0 #length of pendulum1 +l2 = 2.0 #length of pendulum2 +m1 = 1.0 #mass of pendulum1 +m2 = 1.0 #mass of pendulum2 g = 9.81 #gravity params = [l1, l2, m1, m2, g] -q0 = [1.0,1.0] -p0 = [1.0,1.0] -times = (0.,25.) +q0 = [1.0, 1.0] +p0 = [1.0, 1.0] +times = (0.0, 25.0) prob = HamiltonianProblem(H, q0, p0, times, params) -sol2 = solve(prob, SofSpa10(), dt = .05) +sol2 = solve(prob, SofSpa10(), dt = 0.05) -plot(sol2, vars=1, xlim=(0,20), label="Momentum1") -plot!(sol2, vars=2, xlim=(0,20), label="Momentum2") -plot!(sol2, vars=3, xlim=(0,20), label="theta1") -plot!(sol2, vars=4, xlim=(0,20), label="theta2") +plot(sol2, vars = 1, xlim = (0, 20), label = "Momentum1") +plot!(sol2, vars = 2, xlim = (0, 20), label = "Momentum2") +plot!(sol2, vars = 3, xlim = (0, 20), label = "theta1") +plot!(sol2, vars = 4, xlim = (0, 20), label = "theta2") #==========================================================# - #Animation #==========================================================# function make_pretty_gif(sol) timepoints = sol.t - x1 = l1*sin.(sol[3,:]) - y1 = -l1*cos.(sol[3,:]) - x2 = x1 + l2*sin.(sol[4,:]) - y2 = y1 - l2*cos.(sol[4,:]) + x1 = l1 * sin.(sol[3, :]) + y1 = -l1 * cos.(sol[3, :]) + x2 = x1 + l2 * sin.(sol[4, :]) + y2 = y1 - l2 * cos.(sol[4, :]) - axis_lim = (l1+l2)*1.2 + axis_lim = (l1 + l2) * 1.2 anim = Animation() - for i =1:length(timepoints) - str = string("Time = ", round(timepoints[i],digits=1), " sec") - plot([0,x1[i]], [0,y1[i]], size=(400,300), xlim=(-axis_lim,axis_lim), ylim=(-axis_lim,1), markersize = 10, markershape = :circle,label ="",axis = []) - plot!([x1[i],x2[i]], [y1[i],y2[i]], markersize = 10, markershape = :circle,label ="",title = str, title_location = :left) + for i in 1:length(timepoints) + str = string("Time = ", round(timepoints[i], digits = 1), " sec") + plot([0, x1[i]], [0, y1[i]], size = (400, 300), xlim = (-axis_lim, axis_lim), + ylim = (-axis_lim, 1), markersize = 10, markershape = :circle, label = "", + axis = []) + plot!([x1[i], x2[i]], [y1[i], y2[i]], markersize = 10, markershape = :circle, + label = "", title = str, title_location = :left) if i > 8 #rainbow trail - plot!([x2[i-2:i]], [y2[i-2:i]], alpha = 0.15, linewidth = 2, color = :red, label=nothing) - plot!([x2[i-3:i-2]], [y2[i-3:i-2]],alpha = 0.15, linewidth = 2, color = :orange, label=nothing) - plot!([x2[i-4:i-3]], [y2[i-4:i-3]],alpha = 0.15, linewidth = 2, color = :yellow, label=nothing) - plot!([x2[i-6:i-4]], [y2[i-6:i-4]],alpha = 0.15, linewidth = 2, color = :green, label=nothing) - plot!([x2[i-7:i-6]], [y2[i-7:i-6]],alpha = 0.15, linewidth = 2, color = :blue, label=nothing) - plot!([x2[i-8:i-7]], [y2[i-8:i-7]],alpha = 0.15, linewidth = 2, color = :purple, label=nothing) + plot!([x2[(i - 2):i]], [y2[(i - 2):i]], alpha = 0.15, linewidth = 2, + color = :red, label = nothing) + plot!([x2[(i - 3):(i - 2)]], [y2[(i - 3):(i - 2)]], alpha = 0.15, linewidth = 2, + color = :orange, label = nothing) + plot!([x2[(i - 4):(i - 3)]], [y2[(i - 4):(i - 3)]], alpha = 0.15, linewidth = 2, + color = :yellow, label = nothing) + plot!([x2[(i - 6):(i - 4)]], [y2[(i - 6):(i - 4)]], alpha = 0.15, linewidth = 2, + color = :green, label = nothing) + plot!([x2[(i - 7):(i - 6)]], [y2[(i - 7):(i - 6)]], alpha = 0.15, linewidth = 2, + color = :blue, label = nothing) + plot!([x2[(i - 8):(i - 7)]], [y2[(i - 8):(i - 7)]], alpha = 0.15, linewidth = 2, + color = :purple, label = nothing) end frame(anim) end diff --git a/examples/pendulum.jl b/examples/pendulum.jl index febef5e..30ff590 100644 --- a/examples/pendulum.jl +++ b/examples/pendulum.jl @@ -3,7 +3,7 @@ using DifferentialEquations, Plots #Solving the simple pendulum with a traditional ODE method #==========================================================# function pendulum(du, u, params, t) -#reference: http://www.pgccphy.net/ref/advmech.pdf page 6 + #reference: http://www.pgccphy.net/ref/advmech.pdf page 6 g = params[1] #gravitational acceleration m = params[2] #mass @@ -12,26 +12,24 @@ function pendulum(du, u, params, t) θ = u[1] ℒ = u[2] - d_θ = ℒ/(m*l^2) - d_ℒ = -m*g*l*sin(θ) + d_θ = ℒ / (m * l^2) + d_ℒ = -m * g * l * sin(θ) du .= [d_θ, d_ℒ] return nothing end g = 9.81 -m = 2. -l = 1. -params = [g,m,l] -u0 = [1.0,1.0] +m = 2.0 +l = 1.0 +params = [g, m, l] +u0 = [1.0, 1.0] -prob1 = ODEProblem(pendulum, u0, (0., 100.), params) -sol1 = solve(prob1, AutoVern7(Rodas5()), dt = .05) +prob1 = ODEProblem(pendulum, u0, (0.0, 100.0), params) +sol1 = solve(prob1, AutoVern7(Rodas5()), dt = 0.05) #==========================================================# - - #Solving the simple pendulum with the DiffEqPhysics.jl HamiltonianProblem() #==========================================================# function H(ℒ, θ, params, t) @@ -39,23 +37,22 @@ function H(ℒ, θ, params, t) m = params[2] #mass l = params[3] #length - return ℒ^2/(2*m*l^2) + m*g*l*(1-cos(θ)) + return ℒ^2 / (2 * m * l^2) + m * g * l * (1 - cos(θ)) end g = 9.81 -m = 2. -l = 1. -params = [g,m,l] -θ₀ = 1. -ℒ₀ = 1. - -prob2 = HamiltonianProblem(H, ℒ₀, θ₀, (0., 100.), params) -sol2 = solve(prob2, SofSpa10(), dt = .05); +m = 2.0 +l = 1.0 +params = [g, m, l] +θ₀ = 1.0 +ℒ₀ = 1.0 + +prob2 = HamiltonianProblem(H, ℒ₀, θ₀, (0.0, 100.0), params) +sol2 = solve(prob2, SofSpa10(), dt = 0.05); #==========================================================# - #Plotting #==========================================================# -plot(sol1, vars=1, tspan=(0,20), label="ODE Method") -plot!(sol2.t, sol2[2,:], xlim=(0,20), label="Hamiltonian Method") +plot(sol1, vars = 1, tspan = (0, 20), label = "ODE Method") +plot!(sol2.t, sol2[2, :], xlim = (0, 20), label = "Hamiltonian Method") #They produce the same solution! diff --git a/src/hamiltonian.jl b/src/hamiltonian.jl index 5e0a082..5b070de 100644 --- a/src/hamiltonian.jl +++ b/src/hamiltonian.jl @@ -25,9 +25,9 @@ AD (`ForwardDiff`). `H` may be defined with or without time as fourth argument. If both methods are defined, that with 4 arguments is used. """ -function HamiltonianProblem(H, p0::S, q0::T, tspan, param=nothing; kwargs...) where {S,T} - iip = T <: AbstractArray && !(T <: SArray) && S <: AbstractArray && !(S <: SArray) - HamiltonianProblem{iip}(H, p0, q0, tspan, param; kwargs...) +function HamiltonianProblem(H, p0::S, q0::T, tspan, param = nothing; kwargs...) where {S, T} + iip = T <: AbstractArray && !(T <: SArray) && S <: AbstractArray && !(S <: SArray) + HamiltonianProblem{iip}(H, p0, q0, tspan, param; kwargs...) end struct PhysicsTag end @@ -40,14 +40,18 @@ function generic_derivative(q0::Number, hami, x) ForwardDiff.derivative(hami, x) end -function HamiltonianProblem{false}((dp, dq)::Tuple{Any,Any}, p0, q0, tspan, param=nothing; kwargs...) - return ODEProblem(DynamicalODEFunction{false}(dp, dq), ArrayPartition(p0, q0), tspan, param; kwargs...) +function HamiltonianProblem{false}((dp, dq)::Tuple{Any, Any}, p0, q0, tspan, + param = nothing; kwargs...) + return ODEProblem(DynamicalODEFunction{false}(dp, dq), ArrayPartition(p0, q0), tspan, + param; kwargs...) end -function HamiltonianProblem{true}((dp, dq)::Tuple{Any,Any}, p0, q0, tspan, param=nothing; kwargs...) - return ODEProblem(DynamicalODEFunction{true}(dp, dq), ArrayPartition(p0, q0), tspan, param; kwargs...) +function HamiltonianProblem{true}((dp, dq)::Tuple{Any, Any}, p0, q0, tspan, param = nothing; + kwargs...) + return ODEProblem(DynamicalODEFunction{true}(dp, dq), ArrayPartition(p0, q0), tspan, + param; kwargs...) end -function HamiltonianProblem{false}(H, p0, q0, tspan, param=nothing; kwargs...) +function HamiltonianProblem{false}(H, p0, q0, tspan, param = nothing; kwargs...) if DiffEqBase.numargs(H) == 4 dp = (p, q, param, t) -> generic_derivative(q0, q -> -H(p, q, param, t), q) dq = (p, q, param, t) -> generic_derivative(q0, p -> H(p, q, param, t), p) @@ -59,18 +63,22 @@ function HamiltonianProblem{false}(H, p0, q0, tspan, param=nothing; kwargs...) return HamiltonianProblem{false}((dp, dq), p0, q0, tspan, param; kwargs...) end -function HamiltonianProblem{true}(H, p0, q0, tspan, param=nothing; kwargs...) +function HamiltonianProblem{true}(H, p0, q0, tspan, param = nothing; kwargs...) let cp = ForwardDiff.GradientConfig(PhysicsTag(), p0), cq = ForwardDiff.GradientConfig(PhysicsTag(), q0), vfalse = Val(false) - + if DiffEqBase.numargs(H) == 4 - dp = (Δp, p, q, param, t) -> ForwardDiff.gradient!(Δp, q->-H(p, q, param, t), q, cq, vfalse) - dq = (Δq, p, q, param, t) -> ForwardDiff.gradient!(Δq, p-> H(p, q, param, t), p, cp, vfalse) + dp = (Δp, p, q, param, t) -> ForwardDiff.gradient!(Δp, q -> -H(p, q, param, t), + q, cq, vfalse) + dq = (Δq, p, q, param, t) -> ForwardDiff.gradient!(Δq, p -> H(p, q, param, t), + p, cp, vfalse) else issue_depwarn() - dp = (Δp, p, q, param, t) -> ForwardDiff.gradient!(Δp, q->-H(p, q, param), q, cq, vfalse) - dq = (Δq, p, q, param, t) -> ForwardDiff.gradient!(Δq, p-> H(p, q, param), p, cp, vfalse) + dp = (Δp, p, q, param, t) -> ForwardDiff.gradient!(Δp, q -> -H(p, q, param), q, + cq, vfalse) + dq = (Δq, p, q, param, t) -> ForwardDiff.gradient!(Δq, p -> H(p, q, param), p, + cp, vfalse) end return HamiltonianProblem{true}((dp, dq), p0, q0, tspan, param; kwargs...) end diff --git a/src/plot.jl b/src/plot.jl index a0e5dce..ac350b4 100644 --- a/src/plot.jl +++ b/src/plot.jl @@ -1,7 +1,7 @@ struct OrbitPlot - sol - body_names - dim + sol::Any + body_names::Any + dim::Any end @recipe function f(p::OrbitPlot) sol = p.sol @@ -13,7 +13,7 @@ end ind = i -> i:N:length(sol.u[1].x[1]) for i in 1:N @series begin - vars=(ind(i)...,) + vars = (ind(i)...,) label --> body_names[i] vars --> vars sol @@ -21,20 +21,20 @@ end end end -function orbitplot(sol::DESolution;body_names=nothing,dim=3,kwargs...) - RecipesBase.plot(OrbitPlot(sol,body_names,dim);kwargs...) +function orbitplot(sol::DESolution; body_names = nothing, dim = 3, kwargs...) + RecipesBase.plot(OrbitPlot(sol, body_names, dim); kwargs...) end export orbitplot -function plot_orbits(sol;body_names=nothing,dim=3,kwargs...) +function plot_orbits(sol; body_names = nothing, dim = 3, kwargs...) @assert dim ∈ (2, 3) N = length(sol.u[1].x[1].x[1]) body_names = body_names === nothing ? ["orbit $i" for i in 1:N] : body_names ind = i -> i:N:length(sol.u[1].x[1]) - p = plot(sol, vars=(ind(1)...,), lab=body_names[1], kwargs...) + p = plot(sol, vars = (ind(1)...,), lab = body_names[1], kwargs...) for i in 2:N - plot!(p, sol, vars=(ind(i)...,), lab=body_names[i], kwargs...) + plot!(p, sol, vars = (ind(i)...,), lab = body_names[i], kwargs...) end p end diff --git a/test/hamiltonian_test.jl b/test/hamiltonian_test.jl index 8f21edb..5546d61 100644 --- a/test/hamiltonian_test.jl +++ b/test/hamiltonian_test.jl @@ -2,18 +2,18 @@ using Test using DiffEqPhysics, ForwardDiff, OrdinaryDiffEq using StaticArrays, LinearAlgebra, Random -test_solve(prob...) = mapreduce(p->solve(p, Tsit5(), dt=1//2).u, ==, prob) +test_solve(prob...) = mapreduce(p -> solve(p, Tsit5(), dt = 1 // 2).u, ==, prob) -p0, q0 = rand(2) +p0, q0 = rand(2) H(dθ, θ, p) = dθ / 2 - 9.8 * cos(θ) -dp(dθ, θ, p, t) = - 9.8 * sin(θ) +dp(dθ, θ, p, t) = -9.8 * sin(θ) dq(dθ, θ, p, t) = 0.5 -acc = (v, x, p, t) -> ForwardDiff.derivative(x -> -H(v[1], x, p), x[1]) -vel = (v, x, p, t) -> ForwardDiff.derivative(v -> H(v, x[1], p), v[1]) -prob_1 = DynamicalODEProblem(acc, vel, p0, q0, (0., 10.)) +acc = (v, x, p, t) -> ForwardDiff.derivative(x -> -H(v[1], x, p), x[1]) +vel = (v, x, p, t) -> ForwardDiff.derivative(v -> H(v, x[1], p), v[1]) +prob_1 = DynamicalODEProblem(acc, vel, p0, q0, (0.0, 10.0)) @testset "prob1($h)" for h in (H, (dp, dq)) - prob1 = HamiltonianProblem(h, p0, q0, (0., 10.)) + prob1 = HamiltonianProblem(h, p0, q0, (0.0, 10.0)) @test test_solve(prob1, prob_1) end @@ -22,12 +22,12 @@ q0 = @SVector rand(2) H4(dθ, θ, p, t) = dθ[1] / 2 + dθ[2] / 2 - 9.8 * cos(θ[1]) - 9.8 * cos(θ[2]) + t dq(dθ, θ, p, t) = @SVector [0.5, 0.5] dp(dθ, θ, p, t) = @SVector [-9.8 * sin(θ[1]), -9.8 * sin(θ[2])] -acc = (v, x, p, t) -> ForwardDiff.gradient(x -> -H4(v, x, p, t), x) -vel = (v, x, p, t) -> ForwardDiff.gradient(v -> H4(v, x, p, t), v) -prob_2 = DynamicalODEProblem(acc, vel, p0, q0, (0., 10.)) +acc = (v, x, p, t) -> ForwardDiff.gradient(x -> -H4(v, x, p, t), x) +vel = (v, x, p, t) -> ForwardDiff.gradient(v -> H4(v, x, p, t), v) +prob_2 = DynamicalODEProblem(acc, vel, p0, q0, (0.0, 10.0)) @testset "prob2($h)" for h in (H4, (dp, dq)) - prob2 = HamiltonianProblem(h, p0, q0, (0., 10.)) + prob2 = HamiltonianProblem(h, p0, q0, (0.0, 10.0)) @test test_solve(prob2, prob_2) end @@ -37,10 +37,10 @@ dq(dx, dΘ, θ, p, t) = (dx .= [0.5, 0, 0, 0, 0]) dp(dv, dθ, θ, p, t) = (dv .= [(-9.8 * sin.(θ))[1], 0, 0, 0, 0]) acc = (dv, v, x, p, t) -> ForwardDiff.gradient!(dv, x -> -H(v, x, p), x) vel = (dx, v, x, p, t) -> ForwardDiff.gradient!(dx, v -> H(v, x, p), v) -prob_3 = DynamicalODEProblem(acc, vel, p0, q0, (0., 10.)) +prob_3 = DynamicalODEProblem(acc, vel, p0, q0, (0.0, 10.0)) @testset "prob3($h)" for h in (H, (dp, dq)) - prob3 = HamiltonianProblem(h, p0, q0, (0., 10.)) + prob3 = HamiltonianProblem(h, p0, q0, (0.0, 10.0)) @test test_solve(prob3, prob_3) end @@ -50,9 +50,9 @@ dq(dx, dΘ, θ, p, t) = (dx .= [0.5, 0, 0, 0, 0]) dp(dv, dθ, θ, p, t) = (dv .= [(-9.8 * sin.(θ))[1], 0, 0, 0, 0]) acc = (dv, v, x, p, t) -> ForwardDiff.gradient!(dv, x -> -H4(v, x, p, t), x) vel = (dx, v, x, p, t) -> ForwardDiff.gradient!(dx, v -> H4(v, x, p, t), v) -prob_4 = DynamicalODEProblem(acc, vel, p0, q0, (0., 10.)) +prob_4 = DynamicalODEProblem(acc, vel, p0, q0, (0.0, 10.0)) @testset "prob4($h)" for h in (H4, (dp, dq)) - prob4 = HamiltonianProblem(h, p0, q0, (0., 10.)) + prob4 = HamiltonianProblem(h, p0, q0, (0.0, 10.0)) @test test_solve(prob4, prob_4) end diff --git a/test/nbody_test.jl b/test/nbody_test.jl index b5cd843..7140375 100644 --- a/test/nbody_test.jl +++ b/test/nbody_test.jl @@ -2,27 +2,38 @@ using RecursiveArrayTools @testset "NBodyProblem Test" begin G = 2.95912208286e-4 - M = [1.00000597682, 0.000954786104043, 0.000285583733151, 0.0000437273164546, 0.0000517759138449, 1 / 1.3e8] + M = [ + 1.00000597682, + 0.000954786104043, + 0.000285583733151, + 0.0000437273164546, + 0.0000517759138449, + 1 / 1.3e8, + ] invM = inv.(M) planets = ["Sun", "Jupiter", "Saturn", "Uranus", "Neptune", "Pluto"] - pos_x = [0.0,-3.5023653,9.0755314,8.3101420,11.4707666,-15.5387357] - pos_y = [0.0,-3.8169847,-3.0458353,-16.2901086,-25.7294829,-25.2225594] - pos_z = [0.0,-1.5507963,-1.6483708,-7.2521278,-10.8169456,-3.1902382] + pos_x = [0.0, -3.5023653, 9.0755314, 8.3101420, 11.4707666, -15.5387357] + pos_y = [0.0, -3.8169847, -3.0458353, -16.2901086, -25.7294829, -25.2225594] + pos_z = [0.0, -1.5507963, -1.6483708, -7.2521278, -10.8169456, -3.1902382] pos = ArrayPartition(pos_x, pos_y, pos_z) - vel_x = [0.0,0.00565429,0.00168318,0.00354178,0.00288930,0.00276725] - vel_y = [0.0,-0.00412490,0.00483525,0.00137102,0.00114527,-0.00170702] - vel_z = [0.0,-0.00190589,0.00192462,0.00055029,0.00039677,-0.00136504] + vel_x = [0.0, 0.00565429, 0.00168318, 0.00354178, 0.00288930, 0.00276725] + vel_y = [0.0, -0.00412490, 0.00483525, 0.00137102, 0.00114527, -0.00170702] + vel_z = [0.0, -0.00190589, 0.00192462, 0.00055029, 0.00039677, -0.00136504] vel = ArrayPartition(vel_x, vel_y, vel_z) - tspan = (0., 200_000) + tspan = (0.0, 200_000) ∑ = sum N = 6 - potential(p, t, x, y, z, M) = -G * ∑(i -> ∑(j -> (M[i] * M[j]) / sqrt((x[i] - x[j])^2 + (y[i] - y[j])^2 + (z[i] - z[j])^2), 1:i - 1), 2:N) + function potential(p, t, x, y, z, M) + -G * ∑(i -> ∑(j -> (M[i] * M[j]) / + sqrt((x[i] - x[j])^2 + (y[i] - y[j])^2 + (z[i] - z[j])^2), 1:(i - 1)), + 2:N) + end nprob = NBodyProblem(potential, M, vel, pos, tspan) - sol = solve(nprob, Yoshida6(), dt=100) + sol = solve(nprob, Yoshida6(), dt = 100) # Make sure the distances from the sun stay small enough f = (x, y, z, j) -> sqrt((x[1] - x[j])^2 + (y[1] - y[j])^2 + (z[1] - z[j])^2)