Tetrad basis and various fixes (#51)

* evaluate jacobian and value together

* cache sin eval

* calculate bounding box on converted mesh

* improve bin transfer function api

* fix gramm schmidt

* fast math on metric components

* allow different redshift functions to be used in cunningham

* julia formatting

* cleanup redshift function

* bugfix

* unit test for tetrad frame

* bump version

Project.toml

@@ -1,7 +1,7 @@
 name = "Gradus"
 uuid = "c5b7b928-ea15-451c-ad6f-a331a0f3e4b7"
 authors = ["fjebaker <fergusbkr@gmail.com>"]
-version = "0.1.10"
+version = "0.1.11"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"

src/GradusBase/GradusBase.jl

@@ -3,7 +3,7 @@ module GradusBase
 import Parameters: @with_kw
 import SciMLBase
 import Base
-import StaticArrays: SVector, MMatrix
+import StaticArrays: SVector, MMatrix, SMatrix, @SVector
 import Tullio: @tullio
 
 # for doc bindings

src/GradusBase/geometry.jl

@@ -5,54 +5,44 @@ end
 mdot(g, v1, v2) = @tullio r := g[i, j] * v1[i] * v2[j]
 mnorm(g, v) = mdot(g, v, v)
 
-# fallback methods
-mdot(::AbstractMetricParams{T}, g, v1, v2) where {T} = mdot(g, v1, v2)
-mnorm(m::AbstractMetricParams{T}, g, v) where {T} = mdot(m, g, v, v)
-
 """
     mproject(g, v, u)
-    mproject(m::AbstractMetricParams{T}, g, v, u)
 
 Project vector `v` onto `u` with metric `g`. Optional first argument may be
 [`AbstractMetricParams`](@ref) for more optimized methods, which fallback to an einsum.
 """
-mproject(m::AbstractMetricParams{T}, g, v, u) where {T} = mdot(m, g, v, u) / mnorm(m, g, u)
 mproject(g, v, u) = mdot(g, v, u) / mnorm(g, u)
 
-function projectbasis(m::AbstractMetricParams{T}, g, basis, v::AbstractArray{T}) where {T}
+function projectbasis(g, basis, v)
     s = zero(SVector{4,Float64})
     for e in basis
-        s += mproject(m, g, v, e) .* e
+        s += mproject(g, v, e) .* e
     end
     s
 end
 
-function grammschmidt(m::AbstractMetricParams{T}, g, basis; tol = eps(T)) where {T}
-    v = ones(SVector{4,Float64})
-    p = projectbasis(m, g, basis, v)
+function grammschmidt(v, basis, g; tol = 4eps(Float64))
+    p = projectbasis(g, basis, v)
 
-    while sum(p) > 4 * tol
+    while sum(p) > tol
         v = v .- p
-        p = projectbasis(m, g, basis, v)
+        p = projectbasis(g, basis, v)
     end
 
     v = v .- p
-    v ./ √mnorm(m, g, v)
+    vnorm = √mnorm(g, v)
+    v / vnorm
 end
 
 function tetradframe(m::AbstractMetricParams{T}, u, v) where {T}
     g = metric(m, u)
-    M = zero(MMatrix{4,4,T})
-    known_frame_vectors = minimalframe(m, g, u, v)
-    for (i, e) in enumerate(known_frame_vectors)
-        M[:, i] .= e
-    end
-    for i = length(known_frame_vectors)+1:4
-        M[:, i] .= grammschmidt(m, g, eachcol(M[:, 1:i-1]))
-    end
-    M
-end
 
-function minimalframe(m::AbstractMetricParams{T}, g, u, v) where {T}
-    (v,)
+    # start procedure with ϕ, which has zero for r and θ
+    vϕ = grammschmidt(@SVector[1.0, 0.0, 0.0, 1.0], (v,), g)
+    # then do r, which has zero for θ
+    vr = grammschmidt(@SVector[1.0, 1.0, 0.0, 1.0], (v, vϕ), g)
+    # then finally θ
+    vθ = grammschmidt(@SVector[1.0, 1.0, 1.0, 1.0], (v, vϕ, vr), g)
+
+    (v, vr, vθ, vϕ)
 end

src/accretion-geometry/discs.jl

@@ -44,7 +44,8 @@ end
         if r < m.inner_radius || r > m.outer_radius
             return 1.0
         end
-        @SVector [r * sin(θ) * cos(ϕ), r * sin(θ) * sin(ϕ), r * cos(θ)]
+        sinθ = sin(θ)
+        @SVector [r * sinθ * cos(ϕ), r * sinθ * sin(ϕ), r * cos(θ)]
     end
     n = @SVector [T(0.0), cos(m.inclination), sin(m.inclination)]
     # project u into normal vector n

src/accretion-geometry/meshes.jl

@@ -14,7 +14,7 @@ function MeshAccretionGeometry(mesh)
     static_mesh = map(mesh) do triangle
         Tuple(SVector(p[1], p[2], p[3]) for p in triangle)
     end
-    MeshAccretionGeometry(static_mesh, bounding_box(mesh)...)
+    MeshAccretionGeometry(static_mesh, bounding_box(static_mesh)...)
 end
 
 # naive implementation

src/corona-to-disc/transfer-functions.jl

@@ -1,6 +1,14 @@
-function bin_transfer_function(time_delays, energy, flux; N = 300)
-    energy_bins = range(extrema(energy)..., N)
-    time_bins = range(extrema(time_delays)..., N)
+function bin_transfer_function(
+    time_delays,
+    energy,
+    flux;
+    N_E = 300,
+    N_t = 300,
+    energy_lims = extrema(energy),
+    time_lims = extrema(time_delays),
+)
+    energy_bins = range(energy_lims..., N_E)
+    time_bins = range(time_lims..., N_t)
 
     de = step(energy_bins)
     dt = step(time_bins)

src/metrics/boyer-lindquist-ad.jl

@@ -1,8 +1,8 @@
 module __BoyerLindquistAD
 using ..StaticArrays
 
-@inline Σ(r, a, θ) = r^2 + a^2 * cos(θ)^2
-@inline Δ(r, R, a) = r^2 - R * r + a^2
+@fastmath Σ(r, a, θ) = r^2 + a^2 * cos(θ)^2
+@fastmath Δ(r, R, a) = r^2 - R * r + a^2
 
 # the way this function must be defined is a little complex
 # but helps with type-stability

src/metrics/dilaton-axion-ad.jl

@@ -5,15 +5,15 @@ García et al. (1995).
 module __DilatonAxionAD
 using ..StaticArrays
 
-Σ(r, a, θ) = r^2 + a^2 * cos(θ)^2
-Δ(r, a, rg) = r^2 - 2rg * r + a^2
-A(r, a, Δ, θ) = (r^2 + a^2)^2 - a^2 * Δ * sin(θ)^2
+@fastmath Σ(r, a, θ) = r^2 + a^2 * cos(θ)^2
+@fastmath Δ(r, a, rg) = r^2 - 2rg * r + a^2
+@fastmath A(r, a, Δ, θ) = (r^2 + a^2)^2 - a^2 * Δ * sin(θ)^2
 
-W(βab, θ, βa) = 1 + (βab * (2 * cos(θ) - βab) + βa^2) * csc(θ)^2
-Σ̂(Σ, β, b, r, βb, a, θ, rg) = Σ - (β^2 + 2b * r) + rg^2 * βb * (βb - 2 * (a / rg) * cos(θ))
-Δ̂(Δ, β, b, r, βb, rg) = Δ - (β^2 + 2b * r) - rg * (rg + 2b) * βb^2
-Â(δ, a, Δ̂, W, θ) = δ^2 - a^2 * Δ̂ * W^2 * sin(θ)^2
-δ(r, b, a) = r^2 - 2b * r + a^2
+@fastmath W(βab, θ, βa) = 1 + (βab * (2 * cos(θ) - βab) + βa^2) * csc(θ)^2
+@fastmath Σ̂(Σ, β, b, r, βb, a, θ, rg) = Σ - (β^2 + 2b * r) + rg^2 * βb * (βb - 2 * (a / rg) * cos(θ))
+@fastmath Δ̂(Δ, β, b, r, βb, rg) = Δ - (β^2 + 2b * r) - rg * (rg + 2b) * βb^2
+@fastmath Â(δ, a, Δ̂, W, θ) = δ^2 - a^2 * Δ̂ * W^2 * sin(θ)^2
+@fastmath δ(r, b, a) = r^2 - 2b * r + a^2
 
 @fastmath function metric_components(M, a, β, b, rθ)
     (r, θ) = rθ

src/metrics/johannsen-ad.jl

@@ -1,13 +1,13 @@
 module __JohannsenAD
 using ..StaticArrays
 
-@inline f(M, r, ϵ3) = ϵ3 * M^3 / r
-@inline Σ(r, a, θ, M, ϵ3) = r^2 + a^2 * cos(θ)^2 + f(M, r, ϵ3)
-@inline Δ(r, M, a) = r^2 - 2M * r + a^2
+@fastmath f(M, r, ϵ3) = ϵ3 * M^3 / r
+@fastmath Σ(r, a, θ, M, ϵ3) = r^2 + a^2 * cos(θ)^2 + f(M, r, ϵ3)
+@fastmath Δ(r, M, a) = r^2 - 2M * r + a^2
 
-@inline A₁(M, r, α13) = 1 + α13 * (M / r)^3
-@inline A₂(M, r, α22) = 1 + α22 * (M / r)^2
-@inline A₅(M, r, α52) = 1 + α52 * (M / r)^2
+@fastmath A₁(M, r, α13) = 1 + α13 * (M / r)^3
+@fastmath A₂(M, r, α22) = 1 + α22 * (M / r)^2
+@fastmath A₅(M, r, α52) = 1 + α52 * (M / r)^2
 
 @fastmath function metric_components(m, rθ)
     (r, θ) = rθ

src/metrics/johannsen-psaltis-ad.jl

@@ -1,9 +1,9 @@
 module __JohannsenPsaltisAD
 using ..StaticArrays
 
-@inline h(r, M, ϵ3, Σ) = ϵ3 * M^3 * r / Σ^2
-@inline Σ(r, a, θ) = r^2 + a^2 * cos(θ)^2
-@inline Δ(r, M, a) = r^2 - 2M * r + a^2
+@fastmath h(r, M, ϵ3, Σ) = ϵ3 * M^3 * r / Σ^2
+@fastmath Σ(r, a, θ) = r^2 + a^2 * cos(θ)^2
+@fastmath Δ(r, M, a) = r^2 - 2M * r + a^2
 
 @fastmath function metric_components(M, a, ϵ3, rθ)
     (r, θ) = rθ

src/metrics/kerr-refractive-ad.jl

@@ -1,8 +1,8 @@
 module __KerrRefractiveAD
 using ..StaticArrays
 
-@inline Σ(r, a, θ) = r^2 + a^2 * cos(θ)^2
-@inline Δ(r, R, a) = r^2 - R * r + a^2
+@fastmath Σ(r, a, θ) = r^2 + a^2 * cos(θ)^2
+@fastmath Δ(r, R, a) = r^2 - R * r + a^2
 
 @fastmath function metric_components(M, a, n, corona_radius, rθ)
     (r, θ) = rθ

src/metrics/morris-thorne-ad.jl

@@ -1,5 +1,6 @@
 module __MorrisThorneAD
 using ..StaticArrays
+
 @fastmath function metric_components(b, lθ)
     (l, θ) = lθ
 

src/orbits/circular-orbits.jl

@@ -7,7 +7,7 @@ import ..Gradus:
     inverse_metric_components
 
 function Ω(m::AbstractAutoDiffStaticAxisSymmetricParams{T}, rθ, pos) where {T}
-    jacs = metric_jacobian(m, rθ)
+    _, jacs = metric_jacobian(m, rθ)
     ∂rg = jacs[:, 1]
 
     Δ = √(∂rg[5]^2 - ∂rg[1] * ∂rg[4])

src/orbits/emission-radii.jl

@@ -78,6 +78,7 @@ function redshift_ratio(
     max_time,
     αs,
     βs;
+    redshift_pf = Gradus.ConstPointFunctions.redshift,
     kwargs...,
 )
     velfunc(i) = map_impact_parameters(m, u, αs[i], βs[i])
@@ -93,19 +94,30 @@ function redshift_ratio(
     )
     map(simsols) do sol
         gp = getgeodesicpoint(m, sol)
-        Gradus.ConstPointFunctions.redshift(m, gp, max_time)
+        redshift_pf(m, gp, max_time)
     end
 end
 
-function jacobian_∂αβ_∂gr(m, u, d, max_time, gs, αs, βs; order = 3, kwargs...)
+function jacobian_∂αβ_∂gr(
+    m,
+    u,
+    d,
+    max_time,
+    gs,
+    αs,
+    βs;
+    order = 3,
+    redshift_pf = Gradus.ConstPointFunctions.redshift,
+    kwargs...,
+)
     gmin, gmax = extrema(gs)
     gstar(g) = (g - gmin) / (gmax - gmin)
 
     f((α, β)) = begin
         v = map_impact_parameters(m, u, α, β)
         sol = tracegeodesics(m, u, v, d, (0.0, max_time); save_on = false, kwargs...)
         gp = getgeodesicpoint(m, sol)
-        g = Gradus.ConstPointFunctions.redshift(m, gp, 2000.0)
+        g = redshift_pf(m, gp, 2000.0)
         @SVector [gp.u2[2], gstar(g)]
     end
 
@@ -141,12 +153,22 @@ function cunningham_transfer_function(
     max_time;
     num_points = 1000,
     finite_diff_order = 3,
+    redshift_pf::PointFunction = Gradus.ConstPointFunctions.redshift,
     tracer_kwargs...,
 )
     αs, βs =
         Gradus.impact_parameters_for_radius(m, u, d, rₑ; N = num_points, tracer_kwargs...)
 
-    gs = redshift_ratio(m, u, d, max_time, αs, βs; tracer_kwargs...)
+    gs = redshift_ratio(
+        m,
+        u,
+        d,
+        max_time,
+        αs,
+        βs;
+        redshift_pf = redshift_pf,
+        tracer_kwargs...,
+    )
     gstars = gstar(gs)
 
     js = jacobian_∂αβ_∂gr(
@@ -158,6 +180,7 @@ function cunningham_transfer_function(
         αs,
         βs;
         order = finite_diff_order,
+        redshift_pf = redshift_pf,
         tracer_kwargs...,
     )
 

src/redshift.jl

@@ -1,6 +1,6 @@
 module RedshiftFunctions
 import ..Gradus
-import ..Gradus: __BoyerLindquistFO, AbstractMetricParams, metric
+import ..Gradus: __BoyerLindquistFO, AbstractMetricParams, BoyerLindquistAD, metric
 using StaticArrays
 using Tullio: @tullio
 
@@ -153,19 +153,16 @@ function regular_pdotu_inv(L, M, r, a, θ)
     (eⱽ(M, r, a, θ) * √(1 - Vₑ(M, r, a, θ)^2)) / (1 - L * Ωₑ(M, r, a))
 end
 
-@inline function regular_pdotu_inv(m::AbstractMetricParams{T}, u, v) where {T}
+@inline function regular_pdotu_inv(m::BoyerLindquistAD{T}, u, v) where {T}
     metric_matrix = metric(m, u)
-    # reverse signs of the velocity vector
-    # since we're integrating backwards
-    p = @inbounds @SVector [-v[1], 0, 0, -v[4]]
 
     # TODO: this only works for Kerr
     disc_norm = (eⱽ(m.M, u[2], m.a, u[3]) * √(1 - Vₑ(m.M, u[2], m.a, u[3])^2))
 
     u_disc = @SVector [1 / disc_norm, 0, 0, Ωₑ(m.M, u[2], m.a) / disc_norm]
 
     # use Tullio to do the einsum
-    @tullio g := metric_matrix[i, j] * (u_disc[i]) * p[j]
+    @tullio g := metric_matrix[i, j] * u_disc[i] * (-v[j])
     1 / g
 end
 

src/tracing/method-implementations/auto-diff.jl

@@ -243,25 +243,37 @@ end
 """
     metric_jacobian(m::AbstractAutoDiffStaticAxisSymmetricParams{T}, rθ)
 
-Calculate the Jacobian elements of the metric with respect to ``r`` and ``\\theta``.
+Calculate the value and Jacobian elements of the metric with respect to ``r`` and ``\\theta``.
 
 Limitations:
 - currenly pre-supposes static, axis-symmetric metric.
+    
+## Notes
+
+Function body is equivalent to
+```julia
+f = x -> metric_components(m, x)
+J = ForwardDiff.vector_mode_jacobian(f, rθ)
+f(rθ), J
+```
+but non-allocating.
 """
-@inline metric_jacobian(m::AbstractAutoDiffStaticAxisSymmetricParams{T}, rθ) where {T} =
-    ForwardDiff.vector_mode_jacobian(x -> metric_components(m, x), rθ)
+function metric_jacobian(m::AbstractAutoDiffStaticAxisSymmetricParams, rθ)
+    f = x -> metric_components(m, x)
+    T = typeof(ForwardDiff.Tag(f, eltype(rθ)))
+    ydual = ForwardDiff.static_dual_eval(T, f, rθ)
+    ForwardDiff.value.(T, ydual), ForwardDiff.extract_jacobian(T, ydual, rθ)
+end
 
 @inbounds function geodesic_eq(
     m::AbstractAutoDiffStaticAxisSymmetricParams{T},
     u,
     v,
 ) where {T}
     # get the only position components we need for this metric type
-    rθ = @SVector [u[2], u[3]]
-    # calculate all non-zero components
-    g_comps = metric_components(m, rθ)
-    # use AD to get derivatives
-    jacs = metric_jacobian(m, rθ)
+    rθ = SVector{2,Float64}(u[2], u[3])
+    # calculate all non-zero components, and use AD to get derivatives
+    g_comps, jacs = metric_jacobian(m, rθ)
     # calculate all non-zero inverse matric components
     ginv_comps = inverse_metric_components(g_comps)
     # calculate acceleration

src/tracing/tracing.jl

@@ -69,9 +69,10 @@ function integrator_problem(
 ) where {S,T}
     u_init = vcat(pos, vel)
     ODEProblem{false}(u_init, time_domain) do u, p, λ
-        @inbounds let x = @view(u[1:4]), v = @view(u[5:8])
+        @inbounds let x = SVector{4}(@view(u[1:4])), v = SVector{4}(@view(u[5:8]))
             dv = SVector{4}(geodesic_eq(m, x, v))
-            SVector{8}(v[1], v[2], v[3], v[4], dv[1], dv[2], dv[3], dv[4])
+            # SVector{8}(v[1], v[2], v[3], v[4], dv[1], dv[2], dv[3], dv[4])
+            vcat(v, dv)
         end
     end
 end

test/Project.toml

@@ -1,4 +1,5 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"