Add TaylorInterpolant JLD2 serialization methods (#36)

* Update CI.yml

* Move TaylorInterpolantNSerialization over from NEOs

* Add missing files

* Update tests

* Scattered fixes

* Uncomment tests

* Update tests

.github/workflows/CI.yml

@@ -4,8 +4,6 @@ on:
     paths-ignore:
       - 'LICENSE.md'
       - 'README.md'
-    tags: ['*']
-  pull_request:
 concurrency:
   # Skip intermediate builds: always.
   # Cancel intermediate builds: only if it is a pull request build.

Project.toml

@@ -1,7 +1,7 @@
 name = "PlanetaryEphemeris"
 uuid = "d83715d0-7e5f-11e9-1a59-4137b20d8363"
 authors = ["Jorge A. Pérez Hernández", "Luis Benet", "Luis Eduardo Ramírez Montoya"]
-version = "0.8.1"
+version = "0.8.2"
 
 [deps]
 ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"

src/PlanetaryEphemeris.jl

@@ -26,7 +26,9 @@ import JLD2: writeas
 include("constants.jl")
 include("jpl-de-430-431-earth-orientation-model.jl")
 include("initial_conditions.jl")
-include("interpolation.jl")
+include("interpolation/TaylorInterpolant.jl")
+include("interpolation/PlanetaryEphemerisSerialization.jl")
+include("interpolation/TaylorInterpolantNSerialization.jl")
 include("propagation.jl")
 include("osculating.jl")
 include("barycenter.jl")

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

src/interpolation/PlanetaryEphemerisSerialization.jl

@@ -0,0 +1,70 @@
+# Custom serialization
+
+@doc raw"""
+    PlanetaryEphemerisSerialization{T}
+
+Custom serialization struct to save a `TaylorInterpolant{T, T, 2}` to a `.jld2` file.
+
+# Fields
+- `order::Int`: order of Taylor polynomials.
+- `dims::Tuple{Int, Int}`: matrix dimensions.
+- `t0::T`: initial time.
+- `t::Vector{T}`: vector of times.
+- `x::Vector{T}`: vector of coefficients.
+"""
+struct PlanetaryEphemerisSerialization{T}
+    order::Int
+    dims::Tuple{Int, Int}
+    t0::T
+    t::Vector{T}
+    x::Vector{T}
+end
+
+# Tell JLD2 to save TaylorInterpolant{T, T, 2} as PlanetaryEphemerisSerialization{T}
+function writeas(::Type{<:TaylorInterpolant{T, T, 2, Vector{T}, Matrix{Taylor1{T}}}}) where {T<:Real}
+    return PlanetaryEphemerisSerialization{T}
+end
+
+# Convert method to write .jld2 files
+function convert(::Type{PlanetaryEphemerisSerialization{T}}, eph::TaylorInterpolant{T, T, 2, Vector{T}, Matrix{Taylor1{T}}}) where {T <: Real}
+    # Taylor polynomials order
+    order = eph.x[1, 1].order
+    # Number of coefficients in each polynomial
+    k = order + 1
+    # Matrix dimensions
+    dims = size(eph.x)
+    # Number of elements in matrix
+    N = dims[1] * dims[2]
+    # Vector of coefficients
+    x = Vector{T}(undef, k * N)
+    # Save coefficients
+    for i in 1:N
+        x[(i-1)*k+1 : i*k] = eph.x[i].coeffs
+    end
+
+    return PlanetaryEphemerisSerialization{T}(order, dims, eph.t0, eph.t, x)
+end
+
+# Convert method to read .jld2 files
+function convert(::Type{TaylorInterpolant{T, T, 2, Vector{T}, Matrix{Taylor1{T}}}}, eph::PlanetaryEphemerisSerialization{T}) where {T<:Real}
+    # Taylor polynomials order
+    order = eph.order
+    # Number of coefficients in each polynomial
+    k = order + 1
+    # Matrix dimensions
+    dims = eph.dims
+    # Number of elements in matrix
+    N = dims[1] * dims[2]
+    # Matrix of Taylor polynomials
+    x = Matrix{Taylor1{T}}(undef, dims[1], dims[2])
+    # Reconstruct Taylor polynomials
+    for i in 1:N
+        x[i] = Taylor1{T}(eph.x[(i-1)*k+1 : i*k], order)
+    end
+
+    return TaylorInterpolant{T, T, 2}(eph.t0, eph.t, x)
+end
+
+function convert(::Type{TaylorInterpolant{T, T, 2}}, eph::PlanetaryEphemerisSerialization{T}) where {T<:Real}
+    return convert(TaylorInterpolant{T, T, 2, Vector{T}, Matrix{Taylor1{T}}}, eph)
+end

src/interpolation.jl → src/interpolation/TaylorInterpolant.jl

@@ -206,7 +206,7 @@ function selecteph(eph::TaylorInterpolant, bodyind::Union{Int, AbstractVector{In
     end
     # Degrees of freedom
     N = numberofbodies(eph)
-    @assert all(bodyind .< N) "bodyind .< $N" 
+    @assert all(bodyind .< N) "bodyind .< $N"
     idxs = nbodyind(N, bodyind)
     if euler
         idxs = vcat(idxs, 6N+1:6N+12)
@@ -256,81 +256,3 @@ function auday2kmsec(pv)
     pv[4:6] /= daysec # (km, km/day) -> (km, km/sec)
     return pv
 end
-
-# Custom serialization
-
-@doc raw"""
-    PlanetaryEphemerisSerialization{T}
-
-Custom serialization struct to save a `TaylorInterpolant{T, T, 2}` to a `.jld2` file.
-
-# Fields
-- `order::Int`: order of Taylor polynomials.
-- `dims::Tuple{Int, Int}`: matrix dimensions.
-- `t0::T`: initial time.
-- `t::Vector{T}`: vector of times.
-- `x::Vector{T}`: vector of coefficients.
-"""
-struct PlanetaryEphemerisSerialization{T}
-    order::Int
-    dims::Tuple{Int, Int}
-    t0::T
-    t::Vector{T}
-    x::Vector{T}
-end
-
-# Tell JLD2 to save TaylorInterpolant{T, T, 2} as PlanetaryEphemerisSerialization{T}
-writeas(::Type{TaylorInterpolant{T, T, 2}}) where {T<:Real} = PlanetaryEphemerisSerialization{T}
-
-# Convert method to write .jld2 files
-function convert(::Type{PlanetaryEphemerisSerialization{T}}, eph::TaylorInterpolant{T, T, 2}) where {T <: Real}
-    # Taylor polynomials order
-    order = eph.x[1, 1].order
-    # Number of coefficients in each polynomial
-    k = order + 1
-    # Matrix dimensions
-    dims = size(eph.x)
-    # Number of elements in matrix
-    N = dims[1] * dims[2]
-    # Vector of coefficients
-    x = Vector{T}(undef, k * N)
-    # Save coefficients
-    for i in 1:N
-        x[(i-1)*k+1 : i*k] = eph.x[i].coeffs
-    end
-
-    return PlanetaryEphemerisSerialization{T}(order, dims, eph.t0, eph.t, x)
-end
-
-# Convert method to read .jld2 files
-function convert(::Type{TaylorInterpolant{T, T, 2}}, eph::PlanetaryEphemerisSerialization{T}) where {T<:Real}
-    # Taylor polynomials order
-    order = eph.order
-    # Number of coefficients in each polynomial
-    k = order + 1
-    # Matrix dimensions
-    dims = eph.dims
-    # Number of elements in matrix
-    N = dims[1] * dims[2]
-    # Matrix of Taylor polynomials
-    x = Matrix{Taylor1{T}}(undef, dims[1], dims[2])
-    # Reconstruct Taylor polynomials
-    for i in 1:N
-        x[i] = Taylor1{T}(eph.x[(i-1)*k+1 : i*k], order)
-    end
-
-    return TaylorInterpolant{T, T, 2}(eph.t0, eph.t, x)
-end
-
-@doc raw"""
-    Taylor1Serialization{T}
-
-Custom serialization struct used in previous versions (<= 0.4) of `PlanetaryEphemeris`. Currently, it is only used to
-read old `.jld2` files.
-"""
-struct Taylor1Serialization{T}
-    x::Vector{T}
-end
-
-# Convert method to read .jld2 files
-convert(::Type{Taylor1{T}}, a::Taylor1Serialization{T}) where {T} = Taylor1{T}(a.x, length(a.x) - 1)

src/interpolation/TaylorInterpolantNSerialization.jl

@@ -0,0 +1,128 @@
+@doc raw"""
+    TaylorInterpolantNSerialization{T}
+
+Custom serialization struct to save a `TaylorInterpolant{T, TaylorN{T}, 2}` to a `.jld2` file.
+
+# Fields
+- `vars::Vector{String}`: jet transport variables.
+- `order::Int`: order of Taylor polynomials w.r.t time.
+- `varorder::Int`: order of jet transport perturbations.
+- `dims::Tuple{Int, Int}`: matrix dimensions.
+- `t0::T`: initial time.
+- `t::Vector{T}`: vector of times.
+- `x::Vector{T}`: vector of coefficients.
+"""
+struct TaylorInterpolantNSerialization{T}
+    vars::Vector{String}
+    order::Int
+    varorder::Int
+    dims::Tuple{Int, Int}
+    t0::T
+    t::Vector{T}
+    x::Vector{T}
+end
+
+# Tell JLD2 to save <:TaylorInterpolant{T, TaylorN{T}, 2} as TaylorInterpolantNSerialization{T}
+function writeas(::Type{<:TaylorInterpolant{T, TaylorN{T}, 2, Vector{T}, Matrix{Taylor1{TaylorN{T}}}}}) where {T<:Real}
+    return TaylorInterpolantNSerialization{T}
+end
+
+# Convert method to write .jld2 files
+function convert(::Type{TaylorInterpolantNSerialization{T}}, eph::TaylorInterpolant{T, TaylorN{T}, 2, Vector{T}, Matrix{Taylor1{TaylorN{T}}}}) where {T}
+    # Variables
+    vars = TS.get_variable_names()
+    # Number of variables
+    n = length(vars)
+    # Matrix dimensions
+    dims = size(eph.x)
+    # Number of elements in matrix
+    N = length(eph.x)
+    # Taylor1 order
+    order = eph.x[1, 1].order
+    # Number of coefficients in each Taylor1
+    k = order + 1
+    # TaylorN order
+    varorder = eph.x[1, 1].coeffs[1].order
+    # Number of coefficients in each TaylorN
+    L = varorder + 1
+    # Number of coefficients in each HomogeneousPolynomial
+    M = binomial(n + varorder, varorder)
+
+    # Vector of coefficients
+    x = Vector{T}(undef, N * k * M)
+
+    # Save coefficients
+    i = 1
+    # Iterate over matrix elements
+    for i_1 in 1:N
+        # Iterate over Taylor1 coefficients
+        for i_2 in 1:k
+            # Iterate over TaylorN coefficients
+            for i_3 in 0:varorder
+                # Iterate over i_3 order HomogeneousPolynomial
+                for i_4 in 1:binomial(n + i_3 - 1, i_3)
+                    x[i] = eph.x[i_1].coeffs[i_2].coeffs[i_3+1].coeffs[i_4]
+                    i += 1
+                end
+            end
+        end
+    end
+
+    return TaylorInterpolantNSerialization{T}(vars, order, varorder, dims, eph.t0, eph.t, x)
+end
+
+# Convert method to read .jld2 files
+function convert(::Type{TaylorInterpolant{T, TaylorN{T}, 2, Vector{T}, Matrix{Taylor1{TaylorN{T}}}}}, eph::TaylorInterpolantNSerialization{T}) where {T}
+    # Variables
+    vars = eph.vars
+    # Number of variables
+    n = length(vars)
+    # Matrix dimensions
+    dims = eph.dims
+    # Number of elements in matrix
+    N = dims[1] * dims[2]
+    # Taylor1 order
+    order = eph.order
+    # Number of coefficients in each Taylor1
+    k = order + 1
+    # TaylorN order
+    varorder = eph.varorder
+    # Number of coefficients in each TaylorN
+    L = varorder + 1
+    # Number of coefficients in each HomogeneousPolynomial
+    M = binomial(n + varorder, varorder)
+    # M = sum(binomial(n + i_3 - 1, i_3) for i_3 in 0:varorder)
+
+    # Set variables
+    if TS.get_variable_names() != vars
+        TS.set_variables(T, vars, order = varorder)
+    end
+
+    # Matrix of Taylor polynomials
+    x = Matrix{Taylor1{TaylorN{T}}}(undef, dims[1], dims[2])
+
+    # Reconstruct Taylor polynomials
+    i = 1
+    # Iterate over matrix elements
+    for i_1 in 1:N
+        # Reconstruct Taylor1s
+        Taylor1_coeffs = Vector{TaylorN{T}}(undef, k)
+        for i_2 in 1:k
+            # Reconstruct TaylorNs
+            TaylorN_coeffs = Vector{HomogeneousPolynomial{T}}(undef, L)
+            # Reconstruct HomogeneousPolynomials
+            for i_3 in 0:varorder
+                TaylorN_coeffs[i_3 + 1] = HomogeneousPolynomial(eph.x[i : i + binomial(n + i_3 - 1, i_3)-1], i_3)
+                i += binomial(n + i_3 - 1, i_3)
+            end
+            Taylor1_coeffs[i_2] = TaylorN(TaylorN_coeffs, varorder)
+        end
+        x[i_1] = Taylor1{TaylorN{T}}(Taylor1_coeffs, order)
+    end
+
+    return TaylorInterpolant{T, TaylorN{T}, 2}(eph.t0, eph.t, x)
+end
+
+function convert(::Type{TaylorInterpolant{T, TaylorN{T}, 2}}, eph::PlanetaryEphemerisSerialization{T}) where {T<:Real}
+    return convert(TaylorInterpolant{T, TaylorN{T}, 2, Vector{T}, Matrix{Taylor1{TaylorN{T}}}}, eph)
+end

src/propagation.jl

@@ -195,7 +195,7 @@ for V_dense in (:(Val{true}), :(Val{false}))
                                     maxsteps, parse_eqs)
 
             if $V_dense == Val{true}
-                return TaylorInterpolant{T, T, 2, SubArray{T,1}, SubArray{Taylor1{T},2}}(jd0 - J2000, sol[1], sol[3])
+                return TaylorInterpolant{T, T, 2}(jd0 - J2000, sol[1], sol[3])
             else
                 return sol
             end

test/propagation.jl

@@ -68,19 +68,35 @@ 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
+        @test JLD2.writeas(typeof(sol)) == PlanetaryEphemeris.PlanetaryEphemerisSerialization{Float64}
+        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))
+        @test JLD2.writeas(typeof(sol1N)) == PlanetaryEphemeris.TaylorInterpolantNSerialization{Float64}
+        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