Add multi-segment eval

helgee · Mar 24, 2017 · 98e42c0 · 98e42c0
1 parent 4754ee2
commit 98e42c0
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 14 deletions.
diff --git a/src/spk.jl b/src/spk.jl
@@ -1,4 +1,5 @@
-import AstroDynBase: Ephemeris, position, velocity, state
+using AstroDynBase
+import AstroDynBase: position, velocity, state
 
 export SPK, position, velocity, state, segments, print_segments
 
@@ -203,32 +204,70 @@ function state(spk::SPK, seg::Segment, tdb::Float64, tdb2::Float64=0.0)
     [r;v]
 end
 
+function findsegment(segments, origin, target)
+    if origin in keys(segments) && target in keys(segments[origin])
+        factor = 1.0
+        return segments[origin][target], factor
+    elseif target in keys(segments) && origin in keys(segments[target])
+        factor = -1.0
+        return segments[target][origin], factor
+    else
+        error("No segment '$origin'->'$target' available.")
+    end
+end
+
+function findpath(origin, target)
+    if target == parent(origin) || parent(target) == origin
+        return [origin, target]
+    elseif parent(target) == parent(origin)
+        return [origin, parent(origin), target]
+    elseif parent(parent(target)) == parent(origin) ||
+        parent(target) == parent(parent(origin))
+        return [origin, parent(origin), parent(target), target]
+    elseif parent(parent(target)) == parent(parent(origin))
+        return [origin, parent(origin), parent(parent(origin)), parent(target), target]
+    end
+end
+
 for (f, n) in zip((:state, :velocity, :position), (6, 3, 3))
     @eval begin
-        function ($f)(spk::SPK, center::Int, target::Int, tdb::Float64, tdb2::Float64=0.0)
-            seg = spk.segments[center][target]
-            ($f)(spk, seg, tdb, tdb2)
+        function $f(spk::SPK, ep::TDBEpoch, from::Type{C1}, to::Type{C2}) where {C1<:CelestialBody, C2<:CelestialBody}
+            path = findpath(from, to)
+            jd1 = julian1(ep)
+            jd2 = julian2(ep)
+            length(path) == 2 && $f(spk, naif_id(from), naif_id(to), jd1, jd2)
+
+            res = zeros($n)
+            for (origin, target) in zip(path, path[2:end])
+                res += $f(spk, naif_id(origin), naif_id(target), jd1, jd2)
+            end
+            res
+        end
+
+        function $f(spk::SPK, center::Int, target::Int, tdb::Float64, tdb2::Float64=0.0)
+            seg, factor = findsegment(spk.segments, center, target)
+            $f(spk, seg, tdb, tdb2) * factor
         end
 
-        function ($f)(spk::SPK, target::Int, tdb::Float64, tdb2::Float64=0.0)
+        function $f(spk::SPK, target::Int, tdb::Float64, tdb2::Float64=0.0)
             seg = spk.segments[0][target]
-            ($f)(spk, seg, tdb, tdb2)
+            $f(spk, seg, tdb, tdb2)
         end
 
-        function ($f)(spk::SPK, target::AbstractString, tdb::Float64, tdb2::Float64=0.0)
-            ($f)(spk, naifid(target), tdb, tdb2)
+        function $f(spk::SPK, target::AbstractString, tdb::Float64, tdb2::Float64=0.0)
+            $f(spk, naifid(target), tdb, tdb2)
         end
 
-        function ($f)(spk::SPK, center::AbstractString, target::AbstractString, tdb::Float64, tdb2::Float64=0.0)
-            ($f)(spk, naifid(center), naifid(target), tdb, tdb2)
+        function $f(spk::SPK, center::AbstractString, target::AbstractString, tdb::Float64, tdb2::Float64=0.0)
+            $f(spk, naifid(center), naifid(target), tdb, tdb2)
         end
 
-        function ($f)(spk::SPK, center::Int, target::AbstractString, tdb::Float64, tdb2::Float64=0.0)
-            ($f)(spk, center, naifid(target), tdb, tdb2)
+        function $f(spk::SPK, center::Int, target::AbstractString, tdb::Float64, tdb2::Float64=0.0)
+            $f(spk, center, naifid(target), tdb, tdb2)
         end
 
-        function ($f)(spk::SPK, center::AbstractString, target::Int, tdb::Float64, tdb2::Float64=0.0)
-            ($f)(spk, naifid(center), target, tdb, tdb2)
+        function $f(spk::SPK, center::AbstractString, target::Int, tdb::Float64, tdb2::Float64=0.0)
+            $f(spk, naifid(center), target, tdb, tdb2)
         end
     end
 end
diff --git a/test/basic.jl b/test/basic.jl
@@ -1,3 +1,5 @@
+import AstroDynBase: TDBEpoch, MercuryBarycenter, SSB, Earth, Moon, EarthBarycenter, Mercury
+
 # Reference value from CSPICE
 rvm = [4.250906022073639e7,2.3501057648129586e7,8.158467467032234e6,-34.160008371029825,37.844059275357594,23.756128199757867]
 
@@ -18,9 +20,23 @@ de430segments = [
     "EARTH-MOON BARYCENTER (3) => EARTH (399)",
 ]
 jd = Dates.datetime2julian(DateTime(2016,1,1,0,0,0))
+ep = TDBEpoch(jd)
 spk = SPK("$path/de430.bsp")
 
 @testset "API" begin
+    @testset "Path" begin
+        @test JPLEphemeris.findpath(EarthBarycenter, SSB) == [EarthBarycenter, SSB]
+        @test JPLEphemeris.findpath(SSB, EarthBarycenter) == [SSB, EarthBarycenter]
+        @test JPLEphemeris.findpath(Earth, Mercury) == [Earth, EarthBarycenter, SSB, MercuryBarycenter, Mercury]
+        @test JPLEphemeris.findpath(Earth, Moon) == [Earth, EarthBarycenter, Moon]
+        @test JPLEphemeris.findpath(EarthBarycenter, MercuryBarycenter) == [EarthBarycenter, SSB, MercuryBarycenter]
+        @test JPLEphemeris.findpath(Earth, MercuryBarycenter) == [Earth, EarthBarycenter, SSB, MercuryBarycenter]
+        @test JPLEphemeris.findpath(EarthBarycenter, Mercury) == [EarthBarycenter, SSB, MercuryBarycenter, Mercury]
+    end
+    @testset "Segments" begin
+        @test JPLEphemeris.findsegment(spk.segments, 0, 3) == (spk.segments[0][3], 1.0)
+        @test JPLEphemeris.findsegment(spk.segments, 3, 0) == (spk.segments[0][3], -1.0)
+    end
     @testset for (a, b) in zip(JPLEphemeris.list_segments(spk), de430segments)
         @test a == b
     end
@@ -53,5 +69,17 @@ spk = SPK("$path/de430.bsp")
         @test res ≈ rvm[range]
         res = func.(spk, 0, 1, [jd; jd], [0.0, 0.0])
         @test all(map(x -> x ≈ rvm[range], res))
+
+        res = func(spk, ep, SSB, MercuryBarycenter)
+        @test res ≈ rvm[range]
+        res = func(spk, ep, MercuryBarycenter, SSB)
+        @test res ≈ -rvm[range]
+
+        res = func(spk, ep, Earth, Mercury)
+        exp = func(spk, ep, Earth, EarthBarycenter) +
+            func(spk, ep, EarthBarycenter, SSB) +
+            func(spk, ep, SSB, MercuryBarycenter) +
+            func(spk, ep, MercuryBarycenter, Mercury)
+        @test res ≈ exp
     end
 end