formatting and clean

astro-group-bristol · Sep 30, 2022 · 49035c6 · 49035c6
1 parent 76e993a
commit 49035c6
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 37 deletions.
diff --git a/src/metrics/dilaton-axion-ad.jl b/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̂(δ, a, Δ̂, W, θ) = δ^2 - a^2 * Δ̂ * W^2 * sin(θ)^2
 @fastmath δ(r, b, a) = r^2 - 2b * r + a^2

diff --git a/src/orbits/circular-orbits.jl b/src/orbits/circular-orbits.jl
@@ -6,12 +6,16 @@ import ..Gradus:
     metric_jacobian,
     inverse_metric_components
 
-function Ω(m::AbstractAutoDiffStaticAxisSymmetricParams{T}, rθ, pos) where {T}
+function Ω(
+    m::AbstractAutoDiffStaticAxisSymmetricParams,
+    rθ;
+    contra_rotating = false,
+)
     _, jacs = metric_jacobian(m, rθ)
     ∂rg = jacs[:, 1]
 
     Δ = √(∂rg[5]^2 - ∂rg[1] * ∂rg[4])
-    if pos
+    if contra_rotating
         -(∂rg[5] + Δ) / ∂rg[4]
     else
         -(∂rg[5] - Δ) / ∂rg[4]
@@ -22,11 +26,12 @@ function __energy(g, Ωϕ)
     @inbounds -(g[1] + g[5] * Ωϕ) / √(-g[1] - 2g[5] * Ωϕ - g[4] * Ωϕ^2)
 end
 
-energy(m, r; contra_rotating = false) =
-    let rθ = @SVector([r, π / 2])
-        Ωϕ = Ω(m, rθ, contra_rotating)
-        __energy(metric_components(m, rθ), Ωϕ)
-    end
+function energy(m, rθ; contra_rotating = false)
+    Ωϕ = Ω(m, rθ; contra_rotating = contra_rotating)
+    __energy(metric_components(m, rθ), Ωϕ)
+end
+energy(m::AbstractAutoDiffStaticAxisSymmetricParams, r::Number; kwargs...) =
+    energy(m, @SVector([r, π / 2]); kwargs...)
 
 function __angmom(g, Ωϕ, prograde)
     @inbounds res = (g[5] + g[4] * Ωϕ) / √(-g[1] - 2g[5] * Ωϕ - g[4] * Ωϕ^2)
@@ -37,70 +42,74 @@ function __angmom(g, Ωϕ, prograde)
     end
 end
 
-angmom(m, r; contra_rotating = false, prograde = true) =
-    let rθ = @SVector([r, π / 2])
-        Ωϕ = Ω(m, rθ, contra_rotating)
-        __angmom(metric_components(m, rθ), Ωϕ, prograde)
-    end
+function angmom(m, rθ; contra_rotating = false, prograde = true)
+    Ωϕ = Ω(m, rθ; contra_rotating = contra_rotating)
+    __angmom(metric_components(m, rθ), Ωϕ, prograde)
+end
+angmom(m::AbstractAutoDiffStaticAxisSymmetricParams, r::Number; kwargs...) =
+    angmom(m, @SVector([r, π / 2]); kwargs...)
 
 function g_ginv_energy_angmom(
-    m::AbstractAutoDiffStaticAxisSymmetricParams{T},
-    r;
+    m::AbstractAutoDiffStaticAxisSymmetricParams,
+    rθ;
     contra_rotating = false,
     prograde = true,
-) where {T}
-    let rθ = @SVector([r, π / 2])
-
-        g = metric_components(m, rθ)
-        ginv = inverse_metric_components(g)
+)
+    g = metric_components(m, rθ)
+    ginv = inverse_metric_components(g)
 
-        Ωϕ = Ω(m, rθ, contra_rotating)
-        E = __energy(g, Ωϕ)
-        L = __angmom(g, Ωϕ, prograde)
+    Ωϕ = Ω(m, rθ; contra_rotating = contra_rotating)
+    E = __energy(g, Ωϕ)
+    L = __angmom(g, Ωϕ, prograde)
 
-        (g, ginv, E, L)
-    end
+    (g, ginv, E, L)
 end
 
 function __vϕ(ginv, E, L)
     -E * ginv[5] + L * ginv[4]
 end
 
-function vϕ(m::AbstractAutoDiffStaticAxisSymmetricParams{T}, r; kwargs...) where {T}
-    _, ginv, E, L = g_ginv_energy_angmom(m, r; kwargs...)
+function vϕ(m::AbstractAutoDiffStaticAxisSymmetricParams, rθ; kwargs...)
+    _, ginv, E, L = g_ginv_energy_angmom(m, rθ; kwargs...)
     __vϕ(ginv, E, L)
 end
+vϕ(m::AbstractAutoDiffStaticAxisSymmetricParams, r::Number; kwargs...) =
+    vϕ(m, @SVector([r, π / 2]); kwargs...)
 
 # this component doesn't actually seem to correctly constrain the geodesic
 # to being light- / null-like, or timelike. maybe revist? else call constrain before returning
 function __vt(ginv, E, L)
     -E * ginv[1] + L * ginv[5]
 end
 
-function vt(m::AbstractAutoDiffStaticAxisSymmetricParams{T}, r; kwargs...) where {T}
-    _, ginv, E, L = g_ginv_energy_angmom(m, r; kwargs...)
+function vt(m::AbstractAutoDiffStaticAxisSymmetricParams, rθ; kwargs...)
+    _, ginv, E, L = g_ginv_energy_angmom(m, rθ; kwargs...)
     __vt(ginv, E, L)
 end
+vt(m::AbstractAutoDiffStaticAxisSymmetricParams, r::Number; kwargs...) =
+    vt(m, @SVector([r, π / 2]); kwargs...)
 
 function fourvelocity(
-    m::AbstractAutoDiffStaticAxisSymmetricParams{T},
-    r;
+    m::AbstractAutoDiffStaticAxisSymmetricParams,
+    rθ;
     kwargs...,
-) where {T}
-    _, ginv, E, L = g_ginv_energy_angmom(m, r; kwargs...)
+)
+    _, ginv, E, L = g_ginv_energy_angmom(m, rθ; kwargs...)
 
     vt = __vt(ginv, E, L)
     vϕ = __vϕ(ginv, E, L)
 
     @SVector [vt, 0.0, 0.0, vϕ]
 end
+fourvelocity(m::AbstractAutoDiffStaticAxisSymmetricParams, r::Number; kwargs...) =
+    fourvelocity(m, @SVector([r, π / 2]); kwargs...)
 
 function plunging_fourvelocity(
-    m::AbstractAutoDiffStaticAxisSymmetricParams{T},
-    r;
+    m::AbstractAutoDiffStaticAxisSymmetricParams,
+    rθ;
     kwargs...,
-) where {T}
-    g, ginv, E, L = g_ginv_energy_angmom(m, r; kwargs...)
+)
+    g, ginv, E, L = g_ginv_energy_angmom(m, rθ; kwargs...)
     vt = __vt(ginv, E, L)
     vϕ = __vϕ(ginv, E, L)
 
@@ -109,6 +118,8 @@ function plunging_fourvelocity(
 
     @SVector[vt, -sqrt(abs(nom / denom)), 0.0, vϕ]
 end
+plunging_fourvelocity(m::AbstractAutoDiffStaticAxisSymmetricParams, r::Number; kwargs...) =
+    plunging_fourvelocity(m, @SVector([r, π / 2]); kwargs...)
 
 end # module