Skip to content

Commit

Permalink
Add TaylorInterpolant JLD2 serialization methods (#36)
Browse files Browse the repository at this point in the history
* Update CI.yml

* Move TaylorInterpolantNSerialization over from NEOs

* Add missing files

* Update tests

* Scattered fixes

* Uncomment tests

* Update tests
  • Loading branch information
PerezHz authored Apr 2, 2024
1 parent ee973de commit ed4165e
Show file tree
Hide file tree
Showing 9 changed files with 242 additions and 88 deletions.
2 changes: 0 additions & 2 deletions .github/workflows/CI.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
2 changes: 1 addition & 1 deletion Project.toml
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
4 changes: 3 additions & 1 deletion src/PlanetaryEphemeris.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
18 changes: 18 additions & 0 deletions src/dynamics/trivial.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
70 changes: 70 additions & 0 deletions src/interpolation/PlanetaryEphemerisSerialization.jl
Original file line number Diff line number Diff line change
@@ -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
80 changes: 1 addition & 79 deletions src/interpolation.jl → src/interpolation/TaylorInterpolant.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
128 changes: 128 additions & 0 deletions src/interpolation/TaylorInterpolantNSerialization.jl
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion src/propagation.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
24 changes: 20 additions & 4 deletions test/propagation.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

2 comments on commit ed4165e

@PerezHz
Copy link
Owner Author

@PerezHz PerezHz commented on ed4165e Apr 2, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/104068

Tip: Release Notes

Did you know you can add release notes too? Just add markdown formatted text underneath the comment after the text
"Release notes:" and it will be added to the registry PR, and if TagBot is installed it will also be added to the
release that TagBot creates. i.e.

@JuliaRegistrator register

Release notes:

## Breaking changes

- blah

To add them here just re-invoke and the PR will be updated.

Tagging

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v0.8.2 -m "<description of version>" ed4165e2e3246ccfbec68268ec6fd609ee6da598
git push origin v0.8.2

Please sign in to comment.