Merge pull request #30 from dpad/dev/STM-convenience-functions

Add minor sensitivity functions for even simpler STM extraction.

.github/workflows/ci.yml

@@ -22,7 +22,7 @@ jobs:
       fail-fast: false
       matrix:
         julia_version:
-          - '1.5.3'
+          - 1.6
         os:
           - ubuntu-latest
           # - windows-latest
@@ -59,7 +59,7 @@ jobs:
           path: |
             ~/.julia/packages
             ~/.julia/compiled
-          key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('Project.toml', 'test/Project.toml') }}
+          key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('Project.toml') }}
           restore-keys: |
             ${{ runner.os }}-test-${{ env.cache-name }}-
 
@@ -78,14 +78,12 @@ jobs:
           file: lcov.info
 
   docs:
-    needs: test
-
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
         julia_version:
-          - '1.5.3'
+          - 1.6
         os:
           - ubuntu-latest
           # - windows-latest
@@ -122,7 +120,7 @@ jobs:
           path: |
             ~/.julia/packages
             ~/.julia/compiled
-          key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('Project.toml', 'test/Project.toml') }}
+          key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('test/Project.toml') }}
           restore-keys: |
             ${{ runner.os }}-test-${{ env.cache-name }}-
 
@@ -134,4 +132,5 @@ jobs:
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # For authentication with GitHub Actions token
           GITHUB_REPOSITORY: ${{ env.GITHUB_REPOSITORY }}
-        run: julia --project=docs/ docs/make.jl
+        run: |
+          xvfb-run julia --color=yes --project=docs/ -e 'include("docs/make.jl")' 

Project.toml

@@ -68,4 +68,4 @@ StatsPlots = "0.14.17"
 SymbolicUtils = "0.8.4, 0.9, 0.10, 0.11"
 Symbolics = "0.1.2"
 Unitful = "1.5.0"
-julia = "1.5.3"
+julia = "1.6"

docs/Project.toml

@@ -7,3 +7,6 @@ LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
 OrbitalTrajectories = "2b613a20-8d2a-5290-b19f-e06f4bcc2e7d"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
+
+[compat]
+Documenter = "0.27.0"

docs/make.jl

@@ -1,6 +1,9 @@
 using Documenter
 using OrbitalTrajectories
 
+get!(ENV, "JULIA_DEBUG", Documenter)  # Print @debug notices from Documenter
+get!(ENV, "GKSwstype", "100")  # Avoid Plots GR issues (see Documenter.jl#1583)
+
 makedocs(
     sitename = "OrbitalTrajectories",
     format = Documenter.HTML(),

docs/src/tutorials/QSOFamilies.md

@@ -49,7 +49,7 @@ of many similar orbits along the axis.
     orbits like DROs/QSOs, perturbing only along the initial ``x`` position.
 
 ```@example 1
-trajectories = continuation_simple(state; x_perturbation=3e-4, dc_tolerance=1e-4)
+trajectories = continuation_simple(state; x_perturbation=3e-4, dc_tolerance=1e-4, verbose=false)
 length(trajectories)
 ```
 

src/dynamics/correctors/axisymmetric.jl

@@ -51,7 +51,7 @@ end
 
 crossed(sol::Trajectory) = sol.retcode == :NCrossings
 
-function corrector_callback(::Abstract_AxisymmetricCorrector, system::EphemerisNBP)
+function corrector_callback(::Abstract_AxisymmetricCorrector, system::EphemerisNBP; interp_points=10)
     # TODO: Play with interp_points, interp_points=0 halves runtime/memory, but might cause issues due to oscillation around y-axis
     # TODO: idxs=[2] makes things very slow, even though it seems it should speed it up.
     return VectorContinuousCallback(
@@ -62,7 +62,7 @@ function corrector_callback(::Abstract_AxisymmetricCorrector, system::EphemerisN
                 terminate!(integrator, :Crashed)
             end
         end, 3;
-        interp_points=10, save_positions=(false, false)) do du, u, t, integrator
+        interp_points, save_positions=(false, false)) do du, u, t, integrator
             diam1 = maximum(bodvrd(String(primary_body(system)), "RADII"))  # km
             diam2 = maximum(bodvrd(String(secondary_body(system)), "RADII"))  # km
             du[1] = check_distance(u, t, system, primary_body(system), diam1)
@@ -79,14 +79,14 @@ function corrector_callback(::Abstract_AxisymmetricCorrector, system::EphemerisN
         end
 end
 
-function corrector_callback(::Abstract_AxisymmetricCorrector, ::Abstract_DynamicalModel)
+function corrector_callback(::Abstract_AxisymmetricCorrector, ::Abstract_DynamicalModel; interp_points=10)
     # TODO: Play with interp_points, interp_points=0 halves runtime/memory, but might cause issues due to oscillation around y-axis
     # TODO: idxs=[2] makes things very slow, even though it seems it should speed it up.
-    return ContinuousCallback(orbit_xcrossing, terminate_after_N_crossings!; interp_points=10, save_positions=(false, false))
+    return ContinuousCallback(orbit_xcrossing, terminate_after_N_crossings!; interp_points, save_positions=(false, false))
 end
 
-function corrector_solve(corrector::Abstract_AxisymmetricCorrector, state::State; strict=false, verbose=false, crossings=1, kwargs...)
-    callback = strict ? corrector_callback(corrector, state.model) : nothing
+function corrector_solve(corrector::Abstract_AxisymmetricCorrector, state::State; interp_points=2, strict=false, verbose=false, crossings=1, kwargs...)
+    callback = strict ? corrector_callback(corrector, state.model; interp_points) : nothing
     userdata = Dict{Symbol,Any}(:crossings => crossings)  # XXX: Dict type specified for easy merging
 
     sol = sensitivity_trace(AD, state, corrector.frame, corrector.alg; trace_time=true, abstol=1e-12, reltol=1e-12, callback, userdata, kwargs...)

src/dynamics/correctors/continuation.jl

@@ -1,6 +1,6 @@
 export continuation_simple
 
-function continuation_simple(state::State; x_perturbation=1e-2, ops=(:zero, :+, :-), dc_tolerance=1e-3, kwargs...)
+function continuation_simple(state::State; x_perturbation=1e-2, ops=(:zero, :+, :-), dc_tolerance=1e-3, verbose=true, kwargs...)
     # NOTE: Only designed for DiffCorrectAxisymmetric!
     corrector = DiffCorrectAxisymmetric()
 
@@ -19,11 +19,10 @@ function continuation_simple(state::State; x_perturbation=1e-2, ops=(:zero, :+,
     new_sols::Array{Trajectory} = []
 
     # Threads.@threads for id in ops
-    p = ProgressThresh(0., "Continuing orbits:")
     try
         for id in ops
             op = getfield(Base, id)
-            p = ProgressThresh(id == :+ ? max_x : min_x, "Continuing orbits along $(op):")
+            p = ProgressThresh(id == :+ ? max_x : min_x; desc="Continuing orbits along $(op):", enabled=verbose)
             last_state = orig_state
             while true
                 # Build a new solution by perturbing the last solution

src/dynamics/models/cr3bp.jl

@@ -31,18 +31,25 @@ end
 
 # TODO: Move hand-coded functions into :analytical of CR3BP
 CR3BP(args...; kwargs...) = CR3BP{Val{false}}(R3BPSystemProperties(args...; kwargs...))
-function CR3BP{nothandcoded}(props::R3BPSystemProperties) where {nothandcoded<:Val{false}} 
-    ode = CR3BP_ODEFunctions
+function CR3BP{nothandcoded}(props::R3BPSystemProperties; kwargs...) where {nothandcoded<:Val{false}} 
+    if length(kwargs) == 0
+        ode = CR3BP_ODEFunctions
+    else
+        ode = _CR3BP_ODEFunctions(; kwargs...)
+    end
     CR3BP{false, typeof(ode), typeof(props)}(ode, props)
 end
 
+# XXX: The parameters overload below gives us a performance boost.
+ModelingToolkit.parameters(model::CR3BP{false}) = SVector(model.props.μ)
+
 # Hand-coded CR3BP
 HandcodedCR3BP(args...; kwargs...) = CR3BP{Val{true}}(R3BPSystemProperties(args...; kwargs...))
 function CR3BP{handcoded}(props::R3BPSystemProperties) where {handcoded<:Val{true}} 
     ode = _CR3BP_ODEFunctions(nothing, handcodedCR3BP, handcodedCR3BP_withSTM)
     CR3BP{true, typeof(ode), typeof(props)}(ode, props)
 end
-ModelingToolkit.parameters(model::CR3BP{true}) = [model.props.μ]
+ModelingToolkit.parameters(model::CR3BP{true}) = SVector(model.props.μ)
 
 #---------#
 # METHODS #

src/dynamics/models/er3bp.jl

@@ -54,4 +54,7 @@ struct ER3BP{O<:_ER3BP_ODEFunctions,P<:R3BPSystemProperties} <: Abstract_R3BPMod
     ode   :: O
     props :: P
 end
-ER3BP(args...; kwargs...) = ER3BP(ER3BP_ODEFunctions, R3BPSystemProperties(args...; kwargs...))
+ER3BP(args...; kwargs...) = ER3BP(ER3BP_ODEFunctions, R3BPSystemProperties(args...; kwargs...))
+
+# XXX: The parameters overload below gives us a performance boost.
+ModelingToolkit.parameters(model::ER3BP) = SVector(model.props.μ, model.props.e)

src/dynamics/models/nbp.jl

@@ -89,7 +89,7 @@ function EphemerisNBP(bodies::Vararg{Symbol}; center=nothing, kwargs...)
 end
 
 Base.show(io::IO, x::EphemerisNBP) = print(io, "$(nameof(typeof(x)))$(x.props)")
-ModelingToolkit.parameters(model::EphemerisNBP) = []
+ModelingToolkit.parameters(model::EphemerisNBP) = SVector{0,Float64}()
 
 # HELPERS
 

src/dynamics/sensitivities.jl

@@ -64,6 +64,21 @@ end
     return convert_to_frame(traj_vareqns_unconverted, desired_frame)
 end
 
+@doc """ Sensitivity of the state at time t with respect to initial state. """
+function sensitivity(sol::Trajectory, t)
+    extract_STMs(sol, t)
+end
+
+@doc """ Sensitivity of the state at time t2 with respect to the state at time t1 <= t2. """
+function sensitivity(sol::Trajectory, t1, t2)
+    @assert t1 <= t2  "Expected t1 <= t2"
+    if t2 == t1
+        return Matrix(1.0 * I, 6, 6)  # TODO: Make this type and size-generic
+    else
+        return sensitivity(sol, t2) / sensitivity(sol, t1)
+    end
+end
+
 # TODO: Better functions for extracting STMs and stability index from Dual numbers
 function extract_STMs(sol::Trajectory, t)
     extract_STMs(sol(t))