Add JLD2 extension (#354)

* Add JLD2 extension (move over from NEOs)

* Add missing files

* Bump patch version

* Update Project.toml (order fields alphabetically)

* Cleanup

* Remove leftover comment

Project.toml

@@ -1,7 +1,7 @@
 name = "TaylorSeries"
 uuid = "6aa5eb33-94cf-58f4-a9d0-e4b2c4fc25ea"
 repo = "https://github.com/JuliaDiff/TaylorSeries.jl.git"
-version = "0.17.3"
+version = "0.17.4"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -11,15 +11,18 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [weakdeps]
 IntervalArithmetic = "d1acc4aa-44c8-5952-acd4-ba5d80a2a253"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [extensions]
 TaylorSeriesIAExt = "IntervalArithmetic"
+TaylorSeriesJLD2Ext = "JLD2"
 TaylorSeriesSAExt = "StaticArrays"
 
 [compat]
 Aqua = "0.8"
 IntervalArithmetic = "0.15 - 0.20"
+JLD2 = "0.4"
 LinearAlgebra = "<0.0.1, 1"
 Markdown = "<0.0.1, 1"
 Requires = "0.5.2, 1"
@@ -31,10 +34,11 @@ julia = "1"
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 IntervalArithmetic = "d1acc4aa-44c8-5952-acd4-ba5d80a2a253"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["IntervalArithmetic", "LinearAlgebra", "SparseArrays", "StaticArrays", "Test", "Aqua"]
+test = ["IntervalArithmetic", "JLD2", "LinearAlgebra", "SparseArrays", "StaticArrays", "Test", "Aqua"]

ext/TaylorSeriesJLD2Ext.jl

@@ -0,0 +1,91 @@
+module TaylorSeriesJLD2Ext
+
+import Base: convert
+using TaylorSeries
+
+if isdefined(Base, :get_extension)
+    import JLD2: writeas
+else
+    import ..JLD2: writeas
+end
+
+@doc raw"""
+    TaylorNSerialization{T}
+
+Custom serialization struct to save a `TaylorN{T}` to a `.jld2` file.
+
+# Fields
+- `vars::Vector{String}`: jet transport variables.
+- `varorder::Int`: order of jet transport perturbations.
+- `x::Vector{T}`: vector of coefficients.
+"""
+struct TaylorNSerialization{T}
+    vars::Vector{String}
+    varorder::Int
+    x::Vector{T}
+end
+
+# Tell JLD2 to save TaylorN{T} as TaylorNSerialization{T}
+writeas(::Type{TaylorN{T}}) where {T} = TaylorNSerialization{T}
+
+# Convert method to write .jld2 files
+function convert(::Type{TaylorNSerialization{T}}, eph::TaylorN{T}) where {T}
+    # Variables
+    vars = TS.get_variable_names()
+    # Number of variables
+    n = length(vars)
+    # TaylorN order
+    varorder = eph.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, M)
+
+    # Save coefficients
+    i = 1
+    for i_1 in 0:varorder
+        # Iterate over i_1 order HomogeneousPolynomial
+        for i_2 in 1:binomial(n + i_1 - 1, i_1)
+            x[i] = eph.coeffs[i_1+1].coeffs[i_2]
+            i += 1
+        end
+    end
+
+    return TaylorNSerialization{T}(vars, varorder, x)
+end
+
+# Convert method to read .jld2 files
+function convert(::Type{TaylorN{T}}, eph::TaylorNSerialization{T}) where {T}
+    # Variables
+    vars = eph.vars
+    # Number of variables
+    n = length(vars)
+    # 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)
+
+    # Set variables
+    if TS.get_variable_names() != vars
+        TS.set_variables(T, vars, order = varorder)
+    end
+
+    # Reconstruct TaylorN
+    i = 1
+    TaylorN_coeffs = Vector{HomogeneousPolynomial{T}}(undef, L)
+    for i_1 in 0:varorder
+        # Reconstruct HomogeneousPolynomials
+        TaylorN_coeffs[i_1 + 1] = HomogeneousPolynomial(eph.x[i : i + binomial(n + i_1 - 1, i_1)-1], i_1)
+        i += binomial(n + i_1 - 1, i_1)
+    end
+    x = TaylorN{T}(TaylorN_coeffs, varorder)
+
+    return x
+end
+
+end

src/TaylorSeries.jl

@@ -89,6 +89,9 @@ function __init__()
         @require StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" begin
             include("../ext/TaylorSeriesSAExt.jl")
         end
+        @require JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" begin
+            include("../ext/TaylorSeriesJLD2Ext.jl")
+        end
     end
 end
 

test/jld2.jl

@@ -0,0 +1,15 @@
+# This file is part of TaylorSeries.jl, MIT licensed
+#
+
+using TaylorSeries, JLD2
+
+using Test
+
+@testset "Test TaylorSeries JLD2 extension" begin
+    dq = set_variables("q", order=4, numvars=6)
+    random_TaylorN = [cos(sum(dq .* rand(6))), sin(sum(dq .* rand(6))), tan(sum(dq .* rand(6)))]
+    jldsave("test.jld2"; random_TaylorN = random_TaylorN)
+    recovered_taylorN = JLD2.load("test.jld2", "random_TaylorN")
+    @test recovered_taylorN == random_TaylorN
+    rm("test.jld2")
+end

test/runtests.jl

@@ -11,7 +11,8 @@ testfiles = (
     "broadcasting.jl",
     "identities_Euler.jl",
     "fateman40.jl",
-    "staticarrays.jl"
+    "staticarrays.jl",
+    "jld2.jl"
     )
 
 for file in testfiles

test/staticarrays.jl

@@ -4,7 +4,6 @@
 using TaylorSeries, StaticArrays
 
 using Test
-# eeuler = Base.MathConstants.e
 
 @testset "Tests TaylorSeries operations over StaticArrays types" begin
     q = set_variables("q", order=2, numvars=2)