Generic source-to-disc flux (#53)

* julia formatting

* bugfixes in tetrad + new basis function

* formatting and clean

* LNRF test

Also removed Dilaton-Axion from tests (see issue #52).

* lnrf tests and correct domain

* use correct analytic expressions

* divide flux by time * energy bin size

* allow type to be inferred

* switch default solver + inference

* revert to Tsit5 since other default too noisy

* slightly altered api

* update exports

* fix tetrad tests

* voronoi disc profile stores geodesic points as well

* julia formatting

* source-to-disc flux model

* formatting

* bump version

* typo fixes, includes, and exports

Project.toml

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

src/Gradus.jl

@@ -46,7 +46,14 @@ import .GradusBase:
     metric_type,
     metric_components,
     inverse_metric_components,
-    unpack_solution
+    unpack_solution,
+    dotproduct,
+    propernorm,
+    tetradframe,
+    lnrbasis,
+    lnrframe,
+    lowerindices,
+    raiseindices
 
 export AbstractMetricParams,
     getgeodesicpoint,
@@ -57,7 +64,14 @@ export AbstractMetricParams,
     constrain,
     inner_radius,
     metric_components,
-    inverse_metric_components
+    inverse_metric_components,
+    dotproduct,
+    propernorm,
+    tetradframe,
+    lnrbasis,
+    lnrframe,
+    lowerindices,
+    raiseindices
 
 """
     abstract type AbstractPointFunction
@@ -137,6 +151,7 @@ include("corona-to-disc/sky-geometry.jl")
 include("corona-to-disc/corona-models.jl")
 include("corona-to-disc/disc-profiles.jl")
 include("corona-to-disc/transfer-functions.jl")
+include("corona-to-disc/flux-calculations.jl")
 
 include("metrics/boyer-lindquist-ad.jl")
 include("metrics/boyer-lindquist-fo.jl")

src/GradusBase/GradusBase.jl

@@ -1,10 +1,12 @@
 module GradusBase
 
-import Parameters: @with_kw
-import SciMLBase
 import Base
-import StaticArrays: SVector, MMatrix, SMatrix, @SVector
-import Tullio: @tullio
+import SciMLBase
+
+using Parameters: @with_kw
+using StaticArrays: SVector, MMatrix, SMatrix, @SVector
+using Tullio: @tullio
+using LinearAlgebra: norm, inv
 
 # for doc bindings
 import ..Gradus
@@ -27,6 +29,13 @@ constrain,
 inner_radius,
 metric_type,
 metric_components,
-inverse_metric_components
+inverse_metric_components,
+dotproduct,
+propernorm,
+tetradframe,
+lnrbasis,
+lnrframe,
+lowerindices,
+raiseindices
 
 end # module

src/GradusBase/geometry.jl

@@ -1,17 +1,17 @@
-function vector_to_local_sky(m::AbstractMetricParams{T}, u, θ, ϕ) where {T}
+function vector_to_local_sky(m::AbstractMetricParams, u, θ, ϕ)
     error("Not implemented for $(typeof(m))")
 end
 
-mdot(g, v1, v2) = @tullio r := g[i, j] * v1[i] * v2[j]
-mnorm(g, v) = mdot(g, v, v)
+dotproduct(g::AbstractMatrix, v1, v2) = @tullio r := g[i, j] * v1[i] * v2[j]
+propernorm(g::AbstractMatrix, v) = dotproduct(g, v, v)
 
 """
     mproject(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(g, v, u) = mdot(g, v, u) / mnorm(g, u)
+mproject(g, v, u) = dotproduct(g, v, u) / propernorm(g, u)
 
 function projectbasis(g, basis, v)
     s = zero(SVector{4,Float64})
@@ -21,7 +21,7 @@ function projectbasis(g, basis, v)
     s
 end
 
-function grammschmidt(v, basis, g; tol = 4eps(Float64))
+function gramschmidt(v, basis, g; tol = 4eps(Float64))
     p = projectbasis(g, basis, v)
 
     while sum(p) > tol
@@ -30,19 +30,45 @@ function grammschmidt(v, basis, g; tol = 4eps(Float64))
     end
 
     v = v .- p
-    vnorm = √mnorm(g, v)
+    vnorm = √abs(propernorm(g, v))
     v / vnorm
 end
 
-function tetradframe(m::AbstractMetricParams{T}, u, v) where {T}
-    g = metric(m, u)
-
+# TODO: this presupposes static and axis symmetric
+@inline function tetradframe(g::AbstractMatrix, v)
+    vt = v ./ √abs(propernorm(g, v))
     # start procedure with ϕ, which has zero for r and θ
-    vϕ = grammschmidt(@SVector[1.0, 0.0, 0.0, 1.0], (v,), g)
+    vϕ = gramschmidt(@SVector[1.0, 0.0, 0.0, 1.0], (vt,), g)
     # then do r, which has zero for θ
-    vr = grammschmidt(@SVector[1.0, 1.0, 0.0, 1.0], (v, vϕ), g)
+    vr = gramschmidt(@SVector[0.0, 1.0, 0.0, 0.0], (vt, vϕ), g)
     # then finally θ
-    vθ = grammschmidt(@SVector[1.0, 1.0, 1.0, 1.0], (v, vϕ, vr), g)
+    vθ = gramschmidt(@SVector[0.0, 0.0, 1.0, 0.0], (vt, vϕ, vr), g)
+    (vt, vr, vθ, vϕ)
+end
 
-    (v, vr, vθ, vϕ)
+tetradframe(m::AbstractMetricParams, u, v) = tetradframe(metric(m, u), v)
+
+# TODO: this presupposes static and axis symmetric
+# tetrad with indices down: frame
+function lnrframe(g::AbstractMatrix)
+    ω = -g[1, 4] / g[4, 4]
+    v = @SVector [1.0, 0.0, 0.0, ω]
+    tetradframe(g, v)
 end
+lnrframe(m::AbstractMetricParams, u) = lnrframe(metric(m, u))
+
+# tetrad with indices up: basis
+function lnrbasis(g::AbstractMatrix)
+    ω = -g[1, 4] / g[4, 4]
+    v = @SVector [-ω, 0.0, 0.0, 1.0]
+    (vϕ, vr, vθ, vt) = tetradframe(inv(g), v)
+    # rearrange
+    (vt, vr, vθ, vϕ)
+end
+lnrbasis(m::AbstractMetricParams, u) = lnrbasis(metric(m, u))
+
+lowerindices(g::AbstractMatrix, v) = g * v
+lowerindices(m::AbstractMetricParams, u, v) = lowerindices(metric(m, u), v)
+
+raiseindices(ginv::AbstractMatrix, v) = ginv * v
+raiseindices(m::AbstractMetricParams, u, v) = raiseindices(inv(metric(m, u)), v)

src/corona-to-disc/corona-models.jl

@@ -29,4 +29,15 @@ function sample_velocity(
     end
 end
 
+function source_velocity(m::AbstractMetricParams, model::AbstractCoronaModel)
+    error("Not implemented for $(typeof(model)).")
+end
+
+function source_velocity(m::AbstractMetricParams, model::LampPostModel)
+    # stationary source
+    rθ = @SVector[model.h, model.θ]
+    gcomp = metric_components(m, rθ)
+    inv(√(-gcomp[1])) * @SVector[1.0, 0.0, 0.0, 0.0]
+end
+
 export LampPostModel

src/corona-to-disc/disc-profiles.jl

@@ -41,22 +41,24 @@ end
 # evaluate(aap::AbstractAccretionProfile, p) =
 #     error("Not implemented for `$(typeof(aap))` yet.")
 
-struct VoronoiDiscProfile{D,V} <: AbstractDiscProfile
+struct VoronoiDiscProfile{D,V,G} <: AbstractDiscProfile
     disc::D
     polys::Vector{Vector{V}}
     generators::Vector{V}
+    geodesic_points::Vector{G}
 
     function VoronoiDiscProfile(
         d::D,
         polys::Vector{Vector{V}},
         gen::Vector{V},
-    ) where {V<:AbstractArray{T}} where {D,T}
+        gps::Vector{G},
+    ) where {D<:AbstractAccretionDisc,V<:AbstractArray,G}
         if !isapprox(d.inclination, π / 2)
             return error(
                 "Currently only supported for discs in the equitorial plane (θ=π/2).",
             )
         end
-        new{D,V}(d, polys, gen)
+        new{D,V,G}(d, polys, gen, gps)
     end
 end
 
@@ -67,9 +69,9 @@ end
 function VoronoiDiscProfile(
     m::AbstractMetricParams{T},
     d::AbstractAccretionDisc{T},
-    endpoints::AbstractVector{GeodesicPoint{T,V}};
+    endpoints::AbstractVector{<:GeodesicPoint{T}};
     padding = 1,
-) where {T,V}
+) where {T}
     dim = d.outer_radius + padding
     rect = VoronoiCells.Rectangle(
         GeometryBasics.Point2{T}(-dim, -dim),
@@ -88,7 +90,7 @@ function VoronoiDiscProfile(
 
     polys_vecs = unpack_polys(polys)
 
-    VoronoiDiscProfile(d, polys_vecs, generators)
+    VoronoiDiscProfile(d, polys_vecs, generators, endpoints)
 end
 
 function VoronoiDiscProfile(

src/corona-to-disc/flux-calculations.jl

@@ -0,0 +1,79 @@
+"""
+    lorentz_factor(g::AbstractMatrix, isco, u, v) 
+
+Calculate Lorentz factor in LNRF of `u`.
+"""
+function lorentz_factor(g::AbstractMatrix, isco_r, u, v)
+    frame = Gradus.GradusBase.lnrbasis(g)
+    B = reduce(hcat, frame)
+    denom = B[:, 1] ⋅ v 
+
+    𝒱ϕ = (B[:, 4] ⋅ v) / denom
+
+    if u[2] < isco_r
+        𝒱r = (B[:, 2] ⋅ v) / denom
+        inv(√(1 - 𝒱r^2 - 𝒱ϕ^2))
+    else
+        inv(√(1 - 𝒱ϕ^2))
+    end
+end
+
+function flux_source_to_disc(
+    m::AbstractMetricParams,
+    model::AbstractCoronaModel,
+    vdp::AbstractDiscProfile,
+)
+    error(
+        "Not implemented for metric $(typeof(m)) with model $(typeof(model)) and disc profile $(typeof(vdp)).",
+    )
+end
+
+function flux_source_to_disc(
+    m::AbstractMetricParams,
+    model::LampPostModel,
+    vdp::VoronoiDiscProfile;
+    α = 1.0,
+)
+    v_source = source_velocity(m, model)
+
+    intensity = inv.(getareas(vdp))
+    total_intensity = sum(intensity)
+
+    isco_r = isco(m)
+    intp = interpolate_plunging_velocities(m)
+
+    disc_velocity(r) =
+        if r < isco_r
+            vtemp = intp(r)
+            @SVector [vtemp[1], -vtemp[2], vtemp[3], vtemp[4]]
+        else
+            CircularOrbits.fourvelocity(m, r)
+        end
+
+    flux = args -> begin
+        (i, gp) = args
+        g_1 = metric(m, gp.u1)
+        g_2 = metric(m, gp.u2)
+
+        # energy at source
+        @tullio E_s := -g_1[i, j] * gp.v1[i] * v_source[j]
+
+        # energy at disc
+        v_disc = disc_velocity(gp.u2[2])
+        @tullio E_d := -g_2[i, j] * gp.v2[i] * v_disc[j]
+
+        # relative redshift source to disc
+        g_sd = E_d / E_s
+
+        # area element
+        dA = inv(√(g_2[2, 2] * g_2[4, 4]))
+
+        γ = lorentz_factor(g_2, isco_r, gp.u2, v_disc)
+        f_sd = intensity[i] / total_intensity
+        # total reflected flux 
+        g_sd^(1 + α) * E_d^(-α) * dA * f_sd / γ
+    end
+    map(flux, enumerate(vdp.geodesic_points))
+end
+
+export flux_source_to_disc

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

@@ -18,7 +18,7 @@ function bin_transfer_function(
         e_mask = @. (e ≤ energy) & (energy < (e + de))
         sub_flux = @views(flux[t_mask.&e_mask])
         if length(sub_flux) > 0
-            sum(sub_flux)
+            sum(sub_flux) / (dt * de)
         else
             NaN
         end

src/metrics/dilaton-axion-ad.jl

@@ -10,7 +10,8 @@ using ..StaticArrays
 @fastmath A(r, a, Δ, θ) = (r^2 + a^2)^2 - a^2 * Δ * sin(θ)^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, 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