From 39aba75940e4595c29ba5f22c1b3b5831075a6e3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jorge=20A=2E=20P=C3=A9rez-Hern=C3=A1ndez?=
 <perez.hz@gmail.com>
Date: Mon, 1 Apr 2024 15:38:13 +0200
Subject: [PATCH] Update tests

---
 src/dynamics/trivial.jl | 18 ++++++++++++++++++
 test/propagation.jl     | 22 ++++++++++++++++++----
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/src/dynamics/trivial.jl b/src/dynamics/trivial.jl
index 4a62abb..35912a4 100644
--- a/src/dynamics/trivial.jl
+++ b/src/dynamics/trivial.jl
@@ -6,3 +6,21 @@
     end
     nothing
 end
+
+@taylorize function freeparticle!(dq, q, params, t)
+    # N: number of bodies
+    # jd0: initial Julian date
+    local N, jd0 = params
+    Threads.@threads for j in 1:N
+        # Fill first 3N elements of dq with velocities
+        dq[3j-2] = q[3(N+j)-2]
+        dq[3j-1] = q[3(N+j)-1]
+        dq[3j  ] = q[3(N+j)  ]
+    end
+    Threads.@threads for i in 1:N
+        dq[3(N+i)-2] = zero(q[3(N+i)-2])
+        dq[3(N+i)-1] = zero(q[3(N+i)-1])
+        dq[3(N+i)  ] = zero(q[3(N+i)  ])
+    end
+    nothing
+end
diff --git a/test/propagation.jl b/test/propagation.jl
index 313ed23..a964612 100644
--- a/test/propagation.jl
+++ b/test/propagation.jl
@@ -68,19 +68,33 @@ using LinearAlgebra: norm
         @test iszero(zero(TaylorInterpolant{T, T, 2, SubArray{T, 1}, SubArray{Taylor1{T}, 2}}))
         @test iszero(zero(TaylorInterpolant{T, U, 2, SubArray{T, 1}, SubArray{Taylor1{U}, 2}}))
         # Test integration
-        sol = propagate(5, jd0, nyears, dense; dynamics=PlanetaryEphemeris.trivialdynamics!, order, abstol)
-        @test sol isa TaylorInterpolant{Float64,Float64,2}
+        sol = propagate(5, jd0, nyears, dense; dynamics=PlanetaryEphemeris.freeparticle!, order, abstol)
+        @test sol isa TaylorInterpolant{Float64, Float64, 2}
         q0 = initialcond(N, jd0)
         @test sol(sol.t0) == q0
         @test sol.t0 == 0.0
         @test length(sol.t) == size(sol.x, 1) + 1
         @test length(q0) == size(sol.x, 2)
         dq = TaylorSeries.set_variables("dq", order=2, numvars=2)
-        tmid = sol.t0 + sol.t[1]/2
+        tmid = sol.t0 + sol.t[2]/2
         @test sol(tmid) isa Vector{Float64}
         @test sol(tmid + Taylor1(order)) isa Vector{Taylor1{Float64}}
         @test sol(tmid + dq[1] + dq[1]*dq[2]) isa Vector{TaylorN{Float64}}
-        @test sol(tmid + Taylor1([dq[1],dq[1]*dq[2]],order)) isa Vector{Taylor1{TaylorN{Float64}}}
+        @test sol(tmid + Taylor1([dq[1],dq[1]*dq[2]], order)) isa Vector{Taylor1{TaylorN{Float64}}}
+        sol1N = TaylorInterpolant(sol.t0, sol.t, sol.x .+ Taylor1(dq[1], 25))
+        @test sol1N(sol.t0)() == sol(sol.t0)
+        @test sol1N(tmid)() == sol(tmid)
+        # Test PlanetaryEphemerisSerialization
+        jldsave("test.jld2"; sol)
+        sol_file = JLD2.load("test.jld2", "sol")
+        rm("test.jld2")
+        @test sol_file == sol
+        # Test TaylorInterpolantNSerialization
+        sol1N = TaylorInterpolant(sol.t0, sol.t, sol.x .* Taylor1(one(dq[1]), 25))
+        jldsave("test.jld2"; sol1N)
+        sol1N_file = JLD2.load("test.jld2", "sol1N")
+        @test sol1N_file == sol1N
+        rm("test.jld2")
     end
 
     @testset "Propagation: DE430 dynamical model" begin