Plunging interpolation and point function changes (#26)

* modified GeodesicPoint to keep start and end-point

* tests passing with new geodesic point

* fix vertical orientation of first order methods

* isco interpolations

* dilaton fixes

* getgeodesicpoint changes

* circular orbit interpolations + plunging region in SAS metrics

* isco tests

Project.toml

@@ -16,6 +16,7 @@ OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
+Roots = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 ThreadsX = "ac1d9e8a-700a-412c-b207-f0111f4b6c0d"

src/AccretionFormulae/orbit-interpolations.jl

@@ -52,7 +52,7 @@ end
 
 function interpolate_plunging_velocities(
     m::AbstractMetricParams{T};
-    max_time = 10_000,
+    max_time = 50_000,
     contra_rotating = false,
     prograde = true,
     reltol = 1e-9,

src/AccretionFormulae/redshift.jl

@@ -228,8 +228,8 @@ function _redshift_guard(
     gp::FirstOrderGeodesicPoint{T},
     max_time,
 ) where {T}
-    RedshiftFunctions.redshift_function(m, gp.u, gp.p, gp.t)
+    RedshiftFunctions.redshift_function(m, gp.u2, gp.p, gp.t2)
 end
 function _redshift_guard(m::AbstractMetricParams{T}, gp, max_time) where {T}
-    RedshiftFunctions.redshift_function(m, gp.u, gp.v)
+    RedshiftFunctions.redshift_function(m, gp.u2, gp.v2)
 end

src/AccretionGeometry/circular-orbits.jl

@@ -6,11 +6,7 @@ import ..AccretionGeometry:
     metric_jacobian,
     inverse_metric_components
 
-function Ω(
-    m::AbstractAutoDiffStaticAxisSymmetricParams{T},
-    rθ::AbstractVector{T},
-    pos,
-) where {T}
+function Ω(m::AbstractAutoDiffStaticAxisSymmetricParams{T}, rθ, pos) where {T}
     jacs = metric_jacobian(m, rθ)
     ∂rg = jacs[:, 1]
 
@@ -22,18 +18,18 @@ function Ω(
     end
 end
 
-function __energy(g, Ωϕ) where {T}
-    -(g[1] + g[5] * Ωϕ) / √(-g[1] - 2g[5] * Ωϕ - g[4] * Ωϕ^2)
+function __energy(g, Ωϕ)
+    @inbounds -(g[1] + g[5] * Ωϕ) / √(-g[1] - 2g[5] * Ωϕ - g[4] * Ωϕ^2)
 end
 
-energy(m, r; kwargs...) =
+energy(m, r; contra_rotating = false) =
     let rθ = @SVector([r, π / 2])
         Ωϕ = Ω(m, rθ, contra_rotating)
-        __energy(metric_components(m), Ωϕ; kwargs...)
+        __energy(metric_components(m, rθ), Ωϕ)
     end
 
 function __angmom(g, Ωϕ, prograde) where {T}
-    res = (g[5] + g[4] * Ωϕ) / √(-g[1] - 2g[5] * Ωϕ - g[4] * Ωϕ^2)
+    @inbounds res = (g[5] + g[4] * Ωϕ) / √(-g[1] - 2g[5] * Ωϕ - g[4] * Ωϕ^2)
     if prograde
         res
     else
@@ -44,7 +40,7 @@ end
 angmom(m, r; contra_rotating = false, prograde = true) =
     let rθ = @SVector([r, π / 2])
         Ωϕ = Ω(m, rθ, contra_rotating)
-        __angmom(metric_components(m), Ωϕ, prograde)
+        __angmom(metric_components(m, rθ), Ωϕ, prograde)
     end
 
 function g_ginv_energy_angmom(
@@ -79,12 +75,7 @@ function __vt(ginv, E, L)
     -E * ginv[1] + L * ginv[5]
 end
 
-function vt(
-    m::AbstractAutoDiffStaticAxisSymmetricParams{T},
-    r;
-    contra_rotating = false,
-    prograde = true,
-) where {T}
+function vt(m::AbstractAutoDiffStaticAxisSymmetricParams{T}, r; kwargs...) where {T}
     _, ginv, E, L = g_ginv_energy_angmom(m, r; kwargs...)
     __vt(ginv, E, L)
 end

src/AccretionGeometry/geometry.jl

@@ -13,10 +13,10 @@ end
 """
     to_cartesian(gp::GradusBase.AbstractGeodesicPoint{T})
 
-Return the position of `gp` in Cartesian coordinates.
+Return the end position of `gp` in Cartesian coordinates.
 """
 function to_cartesian(gp::GradusBase.AbstractGeodesicPoint{T}) where {T}
-    @inbounds let r = gp.u[2], ϕ = gp.u[4]
+    @inbounds let r = gp.u2[2], ϕ = gp.u2[4]
         x = r * cos(ϕ)
         y = r * sin(ϕ)
         SVector{2,T}(x, y)

src/DiscProfiles/DiscProfiles.jl

@@ -6,7 +6,7 @@ import StaticArrays: @SVector, SVector
 
 
 import ..GradusBase:
-    AbstractMetricParams, GeodesicPoint, vector_to_local_sky, getendpoint, inner_radius
+    AbstractMetricParams, GeodesicPoint, vector_to_local_sky, getgeodesicpoint, inner_radius
 import ..GeodesicTracer: tracegeodesics
 import ..AccretionGeometry:
     AbstractAccretionGeometry,

src/DiscProfiles/disc-profiles.jl

@@ -35,9 +35,9 @@ end
 function VoronoiDiscProfile(
     m::AbstractMetricParams{T},
     d::AbstractAccretionDisc{T},
-    endpoints::AbstractVector{GeodesicPoint{T}};
+    endpoints::AbstractVector{GeodesicPoint{T,V}};
     padding = 1,
-) where {T}
+) where {T,V}
     dim = d.outer_radius + padding
     rect = VoronoiCells.Rectangle(
         GeometryBasics.Point2{T}(-dim, -dim),
@@ -64,7 +64,7 @@ function VoronoiDiscProfile(
     d::AbstractAccretionDisc{T},
     sols::AbstractArray{S},
 ) where {T,S<:SciMLBase.AbstractODESolution}
-    VoronoiDiscProfile(m, d, map(sol -> getendpoint(m, sol), sols))
+    VoronoiDiscProfile(m, d, map(sol -> getgeodesicpoint(m, sol), sols))
 end
 
 function VoronoiDiscProfile(

src/FirstOrderMethods/FirstOrderMethods.jl

@@ -12,10 +12,9 @@ import ..GradusBase:
     AbstractMetricParams,
     inner_radius,
     AbstractGeodesicPoint,
-    getendpoint,
-    geodesic_point_type,
     unpack_solution,
-    SciMLBase
+    SciMLBase,
+    getgeodesicpoint
 
 import ..GeodesicTracer:
     DiscreteCallback,

src/FirstOrderMethods/implementation.jl

@@ -1,26 +1,37 @@
 
-@with_kw struct FirstOrderGeodesicPoint{T,P} <: AbstractGeodesicPoint{T}
+@with_kw struct FirstOrderGeodesicPoint{T,V,P} <: AbstractGeodesicPoint{T}
     retcode::Symbol
-    t::T
-    u::AbstractVector{T}
-    v::AbstractVector{T}
+    "Start time"
+    t1::T
+    "End time"
+    t2::T
+    "Start position"
+    u1::V
+    "End position"
+    u2::V
+    "Start velocity"
+    v1::V
+    "End velocity"
+    v2::V
+    "First order extra parameter"
     p::P
 end
 
-function geodesic_point_type(m::AbstractFirstOrderMetricParams{T}) where {T}
-    p_type = typeof(make_parameters(T(0), T(0), 1, T))
-    FirstOrderGeodesicPoint{T,p_type}
-end
-
-function getendpoint(
+@inbounds function getgeodesicpoint(
     m::AbstractFirstOrderMetricParams{T},
     sol::SciMLBase.AbstractODESolution{T,N,S},
 ) where {T,N,S}
     us, ts, p = unpack_solution(sol)
-    u = us[end]
-    v = four_velocity(u, m, p)
-    t = ts[end]
-    FirstOrderGeodesicPoint(sol.retcode, t, u, convert_velocity_type(u, v), p)
+
+    u_start = SVector{4,T}(us[1][1:4])
+    v_start = SVector{4,T}(four_velocity(u_start, m, p))
+    t_start = ts[1]
+
+    u_end = SVector{4,T}(us[end][1:4])
+    v_end = SVector{4,T}(four_velocity(u_end, m, p))
+    t_end = ts[end]
+
+    FirstOrderGeodesicPoint(sol.retcode, t_start, t_end, u_start, u_end, v_start, v_end, p)
 end
 
 function metric_callback(

src/Gradus.jl

@@ -18,13 +18,14 @@ using StaticArrays
 using Parameters
 using DocStringExtensions
 
+import ForwardDiff
+import Roots
+
 # GradusBase
 export AbstractMetricParams,
     metric_params,
     metric,
-    getendpoint,
-    getstartpoint,
-    getpoint,
+    getgeodesicpoint,
     GeodesicPoint,
     vector_to_local_sky,
     AbstractMetricParams,

src/GradusBase/GradusBase.jl

@@ -3,26 +3,22 @@ module GradusBase
 import Parameters: @with_kw
 import SciMLBase
 import Base
+import StaticArrays: SVector
 
 include("metric-params.jl")
 include("physical-quantities.jl")
 include("geodesic-solutions.jl")
 include("geometry.jl")
 
-export AbstractMetricParams,
-    metric_params,
-    metric,
-    getendpoint,
-    getstartpoint,
-    getpoint,
-    GeodesicPoint,
-    vector_to_local_sky,
-    AbstractMetricParams,
-    geodesic_eq,
-    geodesic_eq!,
-    constrain,
-    onchart,
-    inner_radius,
-    metric_type
+export AbstractMetricParams, metric_params, metric, getgeodesicpoint
+GeodesicPoint,
+vector_to_local_sky,
+AbstractMetricParams,
+geodesic_eq,
+geodesic_eq!,
+constrain,
+onchart,
+inner_radius,
+metric_type
 
 end # module

src/GradusBase/geodesic-solutions.jl

@@ -1,18 +1,40 @@
 abstract type AbstractGeodesicPoint{T} end
 
-@with_kw struct GeodesicPoint{T} <: AbstractGeodesicPoint{T}
+@with_kw struct GeodesicPoint{T,V<:AbstractVector} <: AbstractGeodesicPoint{T}
     retcode::Symbol
-    t::T
-    u::AbstractVector{T}
-    v::AbstractVector{T}
+    "Start time"
+    t1::T
+    "End time"
+    t2::T
+    "Start position"
+    u1::V
+    "End position"
+    u2::V
+    "Start velocity"
+    v1::V
+    "End velocity"
+    v2::V
     # we don't store the problem parameters
     # and can create a specialistion for the carter method
     # then provide dispatched accessor methods
     # p::P
 end
 
-function geodesic_point_type(m::AbstractMetricParams{T}) where {T}
-    GeodesicPoint{T}
+@inbounds function getgeodesicpoint(
+    m::AbstractMetricParams{T},
+    sol::SciMLBase.AbstractODESolution{T,N,S},
+) where {T,N,S}
+    us, ts, _ = unpack_solution(sol)
+
+    u_start = SVector{4,T}(us[1][1:4])
+    v_start = SVector{4,T}(us[1][5:8])
+    t_start = ts[1]
+
+    u_end = SVector{4,T}(us[end][1:4])
+    v_end = SVector{4,T}(us[end][5:8])
+    t_end = ts[end]
+
+    GeodesicPoint(sol.retcode, t_start, t_end, u_start, u_end, v_start, v_end)
 end
 
 # TODO: GeodesicPath structure for the full geodesic path
@@ -28,38 +50,3 @@ end
 function unpack_solution(simsol::SciMLBase.AbstractEnsembleSolution{T,N,V}) where {T,N,V}
     map(unpack_solution, simsol)
 end
-
-function getpoint(
-    m::AbstractMetricParams{T},
-    sol::SciMLBase.AbstractODESolution{T,N,S},
-    i::Int,
-) where {T,N,S}
-    us, ts, _ = unpack_solution(sol)
-    if i == -1
-        u = us[end][1:4]
-        v = us[end][5:8]
-        t = ts[end]
-        GeodesicPoint(sol.retcode, t, u, v)
-    else
-        u = us[i][1:4]
-        v = us[i][5:8]
-        t = ts[i]
-        GeodesicPoint(sol.retcode, t, u, v)
-    end
-end
-
-function getpoint(
-    m::AbstractMetricParams{T},
-    simsol::Union{SciMLBase.AbstractEnsembleSolution{T,N,S},AbstractArray{P}},
-    i,
-) where {T,N,S,P<:SciMLBase.AbstractODESolution}
-    map(sol -> getpoint(m, sol, i), simsol)
-end
-
-function getendpoint(m::AbstractMetricParams{T}, sol_or_simsols) where {T}
-    getpoint(m, sol_or_simsols, -1)
-end
-
-function getstartpoint(m::AbstractMetricParams{T}, sol_or_simsols) where {T}
-    getpoint(m, sol_or_simsols, 1)
-end

src/Rendering/Rendering.jl

@@ -21,8 +21,8 @@ function rendergeodesics(
     m::AbstractMetricParams{T},
     init_pos,
     max_time;
-    pf = PointFunction((m, gp, mt) -> gp.t) ∘
-         FilterPointFunction((m, gp, max_time; kwargs...) -> gp.t < max_time, NaN),
+    pf = PointFunction((m, gp, mt) -> gp.t2) ∘
+         FilterPointFunction((m, gp, max_time; kwargs...) -> gp.t2 < max_time, NaN),
     image_width = 350,
     image_height = 250,
     fov_factor = 3.0,

src/Rendering/point-functions.jl

@@ -24,7 +24,7 @@ end
     kwargs...,
 ) where {M,T,G}
     ThreadsX.map(
-        sol -> pf.f(rc.m, getendpoint(m, sol), rc.max_time; kwargs...),
+        sol -> pf.f(rc.m, getgeodesicpoint(m, sol), rc.max_time; kwargs...),
         rc.geodesics,
     )
 end

src/Rendering/render.jl

@@ -13,7 +13,7 @@ end
 
 function apply_to_image!(m::AbstractMetricParams{T}, image, sols, pf, max_time) where {T}
     @inbounds @threads for i = 1:length(sols)
-        image[i] = pf(m, getendpoint(m, sols[i]), max_time)
+        image[i] = pf(m, getgeodesicpoint(m, sols[i]), max_time)
     end
 end
 
@@ -64,7 +64,7 @@ function __prerendergeodesics(
         kwargs...,
     )
 
-    point_cache = getendpoint(m, simsols)
+    point_cache = getgeodesicpoint(m, simsols)
 
     EndpointRenderCache(
         m,

src/const-point-functions.jl

@@ -3,12 +3,12 @@ import ..Rendering: PointFunction, FilterPointFunction
 import ..AccretionFormulae: _redshift_guard
 
 const filter_early_term =
-    FilterPointFunction((m, gp, max_time; kwargs...) -> gp.t < max_time, NaN)
+    FilterPointFunction((m, gp, max_time; kwargs...) -> gp.t2 < max_time, NaN)
 
 const filter_intersected =
     FilterPointFunction((m, gp, max_time; kwargs...) -> gp.retcode == :Intersected, NaN)
 
-const affine_time = PointFunction((m, gp, max_time; kwargs...) -> gp.t)
+const affine_time = PointFunction((m, gp, max_time; kwargs...) -> gp.t2)
 
 const shadow = affine_time ∘ filter_early_term