diff --git a/CITATION.cff b/CITATION.cff new file mode 100644 index 00000000..e8112e31 --- /dev/null +++ b/CITATION.cff @@ -0,0 +1,12 @@ +cff-version: 1.2.0 +message: "If you use this software, please cite it as below." +authors: + - family-names: Baker + given-names: Fergus + - family-names: Young + given-names: Andrew +title: "Gradus.jl" +version: 0.1.0 +doi: +date-released: +url: "https://github.com/astro-group-bristol/Gradus.jl" \ No newline at end of file diff --git a/Project.toml b/Project.toml new file mode 100644 index 00000000..55a91aa0 --- /dev/null +++ b/Project.toml @@ -0,0 +1,19 @@ +name = "Gradus" +uuid = "c5b7b928-ea15-451c-ad6f-a331a0f3e4b7" +authors = ["fjebaker "] +version = "0.1.0" + +[deps] +Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" +DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def" +DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +GeometricalPredicates = "fd0ad045-b25c-564e-8f9c-8ef5c5f21267" +GeometryBasics = "5c1252a2-5f33-56bf-86c9-59e7332b4326" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" +Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a" +RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" +SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc" diff --git a/README.md b/README.md index 0798a6cd..6f498fd7 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,37 @@ # Gradus.jl -Spacetime generic relativistic ray tracing toolkit, leveraging AD and CAS, capable of calculating reverberations lags around accreting black holes, and other observational signatures. + +A spacetime generic, relativistic ray tracing toolkit, leveraging AD and CAS, capable of calculating reverberation lags around accreting black holes and observational signatures. + +

This package is in development. Please see the documentation.

+ +## About + +A pure Julia geodesic integration system build on [DifferentialEquations.jl](https://github.com/SciML/DifferentialEquations.jl) using automatic differentiation (AD) and computer algebra systems (CAS) to efficiently compute the geodesic equation. This package requires only a specification of the non-zero metric components in order to solve the 2nd order geodesic system. Alternatively, an implementation of the four velocity components may be specified to integrate a regular 1st order system. + +The motivation behind this package began with an interest in studying reverberation lags around accreting black holes, however the scope has since expanded to facilitate the exploration of generic metrics through time-like, space-like, and null geodesics. + +Our aim is to make testing modified Kerr metrics and alternative gravity theories fast. + +Gradus.jl allows for drastically different relativistic simulations to be computed with a composable and reusable API, permitting an end user to simply and expressively calculate physical formulae, create observational signatures, and interface with other popular astrophysics tools. Gradus.jl implements a number of high level abstractions, on the path towards a fully parallelized, high performance numerical relativity ecosystem, scalable from personal computers to super computers. + +## Usage + +We assume you already have Julia >1.6 installed, in which case you can just add the package from the GitHub URL: +```julia +import Pkg; +Pkg.add("https://github.com/astro-group-bristol/Gradus.jl") + +using Gradus +``` + +## Credits + +This work would not be possible without the incredible work of: + +- [The Julia programming language](https://github.com/JuliaLang/Julia) +- [DifferentialEquations.jl](https://github.com/SciML/DifferentialEquations.jl) +- [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) + +
+ +

Astrophysics Group Bristol

\ No newline at end of file diff --git a/src/AccretionFormulae/AccretionFormulae.jl b/src/AccretionFormulae/AccretionFormulae.jl new file mode 100644 index 00000000..1a3cf128 --- /dev/null +++ b/src/AccretionFormulae/AccretionFormulae.jl @@ -0,0 +1,15 @@ +module AccretionFormulae + +import ..Gradus +import ..GradusBase: AbstractMetricParams, metric + +using ..FirstOrderMethods: FirstOrderGeodesicPoint +using ..Rendering: PointFunction + +include("redshift.jl") + +const redshift = PointFunction(_redshift_guard) + +export redshift + +end # module diff --git a/src/AccretionFormulae/redshift.jl b/src/AccretionFormulae/redshift.jl new file mode 100644 index 00000000..1b3fdb83 --- /dev/null +++ b/src/AccretionFormulae/redshift.jl @@ -0,0 +1,235 @@ +module RedshiftFunctions +import ..AccretionFormulae.Gradus +import ..AccretionFormulae.Gradus: __BoyerLindquistFO +import ..AccretionFormulae: AbstractMetricParams, metric +using StaticArrays +using Tullio: @tullio + +""" + eⱽ(M, r, a, θ) +Modified from Cunningham et al. (1975) eq. (A2a): +```math +e^\\nu = \\sqrt{\\frac{\\Delta \\Sigma}{A}}. +``` +""" +eⱽ(M, r, a, θ) = √( + __BoyerLindquistFO.Σ(r, a, θ) * __BoyerLindquistFO.Δ(M, r, a) / + __BoyerLindquistFO.A(M, r, a, θ), +) + + +""" + eᶲ(M, r, a, θ) +Modified from Cunningham et al. (1975) eq. (A2b): +```math +e^\\Phi = \\sin \\theta \\sqrt{\\frac{A}{\\Sigma}}. +``` +""" +eᶲ(M, r, a, θ) = + sin(θ) * √(__BoyerLindquistFO.A(M, r, a, θ) / __BoyerLindquistFO.Σ(r, a, θ)) + + +""" + ω(M, r, a, θ) +From Cunningham et al. (1975) eq. (A2c): +```math +\\omega = \\frac{2 a M r}{A}. +``` +""" +ω(M, r, a, θ) = 2 * a * M * r / __BoyerLindquistFO.A(M, r, a, θ) + + +""" + Ωₑ(M, r, a) +Coordinate angular velocity of an accreting gas. +Taken from Cunningham et al. (1975) eq. (A7b): +```math +\\Omega_e = \\frac{\\sqrt{M}}{a \\sqrt{M} + r_e^{3/2}}. +``` +# Notes +Fanton et al. (1997) use +```math +\\Omega_e = \\frac{\\sqrt{M}}{a \\sqrt{M} \\pm r_e^{3/2}}, +``` +where the sign is dependent on co- or contra-rotation. This function may be extended in the future to support this definition. +""" +Ωₑ(M, r, a) = √M / (r^1.5 + a * √M) + + +""" + Vₑ(M, r, a, θ) +Velocity of an accreting gas in a locally non-rotating reference frame (see Bardeen et al. 1973). +Taken from Cunningham et al. (1975) eq. (A7b): +```math +V_e = (\\Omega_e - \\omega) e^{\\Phi - \\nu}. +``` +""" +Vₑ(M, r, a, θ) = (Ωₑ(M, r, a) - ω(M, r, a, θ)) * eᶲ(M, r, a, θ) / eⱽ(M, r, a, θ) + + +""" + Lₑ(M, rms, a) +Angular momentum of an accreting gas within ``r_ms``. +Taken from Cunningham et al. (1975) eq. (A11b): +```math +L_e = \\sqrt{M} \\frac{ + r_{\\text{ms}}^2 - 2 a \\sqrt{M r_{\\text{ms}}} + a^2 + }{ + r_{\\text{ms}}^{3/2} - 2 M \\sqrt{r_{\\text{ms}}} + a \\sqrt{M} + }. +``` +""" +Lₑ(M, rms, a) = √M * (rms^2 - 2 * a * √(M * rms) + a^2) / (rms^1.5 - 2 * M * √rms + a * √M) + + +""" + H(M, rms, r, a) +Taken from Cunningham et al. (1975) eq. (A12e): +```math +H = \\frac{2 M r_e - a \\lambda_e}{\\Delta}, +``` +where we distinguing ``r_e`` as the position of the accreting gas. +""" +H(M, rms, r, a) = (2 * M * r - a * Lₑ(M, rms, a)) / __BoyerLindquistFO.Δ(M, r, a) + + +""" + γₑ(M, rms) +Taken from Cunningham et al. (1975) eq. (A11c): +```math +\\gamma_e = \\sqrt{1 - \\frac{ + 2M + }{ + 3 r_{\\text{ms}} + }}. +``` +""" +γₑ(M, rms) = √(1 - (2 * M) / (3 * rms)) + + +""" + uʳ(M, rms, r) +Taken from Cunningham et al. (1975) eq. (A12b): +```math +u^r = - \\sqrt{\\frac{ + 2M + }{ + 3 r_{\\text{ms}} + }} \\left( + \\frac{ r_{\\text{ms}} }{r_e} - 1 + \\right)^{3/2}. +``` +""" +uʳ(M, rms, r) = -√((2 * M) / (3 * rms)) * (rms / r - 1)^1.5 + + +""" + uᶲ(M, rms, r, a) +Taken from Cunningham et al. (1975) eq. (A12c): +```math +u^\\phi = \\frac{\\gamma_e}{r_e^2} \\left( + L_e + aH + \\right). +``` +""" +uᶲ(M, rms, r, a) = γₑ(M, rms) / r^2 * (Lₑ(M, rms, a) + a * H(M, rms, r, a)) + + +""" + uᵗ(M, rms, r, a) +Taken from Cunningham et al. (1975) eq. (A12b): +```math +u^t = \\gamma_e \\left( + 1 + \\frac{2 M (1 + H)}{r_e} + \\right). +``` +""" +uᵗ(M, rms, r, a) = γₑ(M, rms) * (1 + 2 * M * (1 + H(M, rms, r, a)) / r) + + +""" +Experimental API: +""" +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} + 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] + 1 / g +end + +function plunging_p_dot_u(E, a, M, L, Q, rms, r, sign_r) + inv( + uᵗ(M, rms, r, a) - uᶲ(M, rms, r, a) * L - + sign_r * uʳ(M, rms, r) * __BoyerLindquistFO.Σδr_δλ(E, L, M, Q, r, a) / + __BoyerLindquistFO.Δ(M, r, a), + ) +end + +function plunging_p_dot_u(m::AbstractMetricParams{T}, u, v, rms) where {T} + metric_matrix = metric(m, u) + + # reverse signs of the velocity vector + # since we're integrating backwards + p = @inbounds @SVector [-v[1], v[2], 0, -v[4]] + + u_disc = @inbounds @SVector [ + uᵗ(m.M, rms, u[2], m.a), + uʳ(m.M, rms, u[2]), + 0, + uᶲ(m.M, rms, u[2], m.a), + ] + + @tullio g := metric_matrix[i, j] * u_disc[i] * p[j] + 1 / g +end + +@inline function redshift_function(m::Gradus.BoyerLindquistFO{T}, u, p, λ) where {T} + isco = Gradus.isco(m) + if u[2] > isco + @inbounds regular_pdotu_inv(p.L, m.M, u[2], m.a, u[3]) + else + # change sign if we're after the sign flip + # TODO: this isn't robust to multiple sign changes + # TODO: i feel like there should be a better way than checking this with two conditions + # used to have λ > p.changes[1] to make sure we're ahead of the time flip (but when wouldn't we be???) + # now p.changes[1] > 0.0 to make sure there was a time flip at all + sign_r = (p.changes[1] > 0.0 ? 1 : -1) * p.r + @inbounds plunging_p_dot_u(m.E, m.a, m.M, p.L, p.Q, isco, u[2], sign_r) + end +end + +@inline function redshift_function(m::AbstractMetricParams{T}, u, v) where {T} + isco = Gradus.isco(m) + if u[2] > isco + regular_pdotu_inv(m, u, v) + else + plunging_p_dot_u(m, u, v, isco) + end +end + +end # module + +# point functions exports +function _redshift_guard( + m::Gradus.BoyerLindquistFO{T}, + gp::FirstOrderGeodesicPoint{T}, + max_time, +) where {T} + RedshiftFunctions.redshift_function(m, gp.u, gp.p, gp.t) +end +function _redshift_guard(m::AbstractMetricParams{T}, gp, max_time) where {T} + RedshiftFunctions.redshift_function(m, gp.u, gp.v) +end diff --git a/src/AccretionGeometry/AccretionGeometry.jl b/src/AccretionGeometry/AccretionGeometry.jl new file mode 100644 index 00000000..cdaca2e7 --- /dev/null +++ b/src/AccretionGeometry/AccretionGeometry.jl @@ -0,0 +1,80 @@ +module AccretionGeometry + +import ..GeodesicTracer: tracegeodesics, DiscreteCallback, terminate! +import ..Rendering: rendergeodesics, prerendergeodesics +import Base: in +import RecursiveArrayTools: ArrayPartition +import GeometryBasics +import LinearAlgebra: ×, ⋅ + +using ..GradusBase +using StaticArrays + +include("geometry.jl") +include("meshes.jl") +include("intersections.jl") +include("discs.jl") + +function tracegeodesics( + m::AbstractMetricParams{T}, + init_positions, + init_velocities, + accretion_geometry, + time_domain::Tuple{T,T}; + callback = nothing, + kwargs..., +) where {T} + cbs = add_collision_callback(callback, accretion_geometry) + tracegeodesics( + m, + init_positions, + init_velocities, + time_domain; + callback = cbs, + kwargs..., + ) +end + +function rendergeodesics( + m::AbstractMetricParams{T}, + init_pos, + accretion_geometry, + max_time::T; + callback = nothing, + kwargs..., +) where {T} + cbs = add_collision_callback(callback, accretion_geometry) + rendergeodesics(m, init_pos, max_time; callback = cbs, kwargs...) +end + +function prerendergeodesics( + m::AbstractMetricParams{T}, + init_pos, + accretion_geometry, + max_time::T; + callback = nothing, + kwargs..., +) where {T} + cbs = add_collision_callback(callback, accretion_geometry) + prerendergeodesics(m, init_pos, max_time; callback = cbs, kwargs...) +end + +add_collision_callback(::Nothing, accretion_geometry) = + build_collision_callback(accretion_geometry) +add_collision_callback(callback::Base.AbstractVecOrTuple, accretion_geometry) = + (callback..., build_collision_callback(accretion_geometry)) +add_collision_callback(callback, accretion_geometry) = + (callback, build_collision_callback(accretion_geometry)) + +function build_collision_callback(geometry::AbstractAccretionGeometry{T}) where {T} + DiscreteCallback(collision_callback(geometry), i -> terminate!(i, :Intersected)) +end + +collision_callback(m::AbstractAccretionGeometry{T}) where {T} = + (u, λ, integrator) -> intersects_geometry(m, line_element(u, integrator)) + + +export AbstractAccretionGeometry, + AbstractAccretionDisc, MeshAccretionGeometry, GeometricThinDisc + +end # module diff --git a/src/AccretionGeometry/discs.jl b/src/AccretionGeometry/discs.jl new file mode 100644 index 00000000..f6930aed --- /dev/null +++ b/src/AccretionGeometry/discs.jl @@ -0,0 +1,22 @@ + +struct GeometricThinDisc{T} <: AbstractAccretionDisc{T} + inner_radius::T + outer_radius::T + inclination::T +end + +function in_nearby_region(d::GeometricThinDisc{T}, line_element) where {T} + p = line_element[2] + d.inner_radius < p[1] < d.outer_radius +end + +function has_intersect(d::GeometricThinDisc{T}, line_element) where {T} + u1, u2 = line_element + sinα = sin(d.inclination) + cosα = cos(d.inclination) + + s1 = u1[1] * (cosα * sin(u1[2]) * cos(u1[3]) + sinα * cos(u1[2])) + s2 = u2[1] * (cosα * sin(u2[2]) * cos(u2[3]) + sinα * cos(u2[2])) + + s1 * s2 ≤ 0 +end diff --git a/src/AccretionGeometry/geometry.jl b/src/AccretionGeometry/geometry.jl new file mode 100644 index 00000000..f3a54c5e --- /dev/null +++ b/src/AccretionGeometry/geometry.jl @@ -0,0 +1,27 @@ +abstract type AbstractAccretionGeometry{T} end +abstract type AbstractAccretionDisc{T} <: AbstractAccretionGeometry{T} end + +function to_cartesian(vec::AbstractVector{T}) where {T} + SVector{3,T}(to_cartesian(vec[2], vec[3], vec[4])...) +end + +function to_cartesian(r, ϕ, θ) + sinϕ = sin(ϕ) + (r * sinϕ * cos(θ), r * sinϕ * sin(θ), r * cos(ϕ)) +end + +function cartesian_line_element(u::ArrayPartition{F,T}, integrator) where {F,T} + (to_cartesian(integrator.uprev.x[2]), to_cartesian(u.x[2])) +end + +function cartesian_line_element(u, integrator) + (to_cartesian(integrator.uprev), to_cartesian(u)) +end + +function line_element(u::ArrayPartition{F,T}, integrator) where {F,T} + @inbounds (@view(integrator.uprev.x[2][2:4]), @view(u.x[2][2:4])) +end + +function line_element(u, integrator) + @inbounds (@view(integrator.uprev[2:4]), @view(u[2:4])) +end diff --git a/src/AccretionGeometry/intersections.jl b/src/AccretionGeometry/intersections.jl new file mode 100644 index 00000000..e70ecb31 --- /dev/null +++ b/src/AccretionGeometry/intersections.jl @@ -0,0 +1,50 @@ + +function intersects_geometry(m::AbstractAccretionGeometry{T}, line_element) where {T} + if in_nearby_region(m, line_element) + return has_intersect(m, line_element) + end + false +end + +in_nearby_region(d::AbstractAccretionGeometry{T}, line_element) where {T} = + error("Not implemented for $(typeof(d))") +has_intersect(d::AbstractAccretionGeometry{T}, line_element) where {T} = + error("Not implemented for $(typeof(d))") + +# Jiménez, Segura, Feito. Computation Geometry 43 (2010) 474-492 +function jsr_algorithm(V₁::T, V₂::T, V₃::T, Q₁::V, Q₂::V; ϵ = 1e-8) where {T,V} + A = Q₁ .- V₃ + B = V₁ .- V₃ + C = V₂ .- V₃ + W₁ = B × C + w = A ⋅ W₁ + if w > ϵ + D = Q₂ .- V₃ + s = D ⋅ W₁ + s > ϵ && return false + W₂ = A × D + t = W₂ ⋅ C + t < -ϵ && return false + u = -W₂ ⋅ B + u < -ϵ && return false + w < s + t + u && return false + elseif w < -ϵ + return false + else # w == 0 + D = Q₂ .- V₃ + s = D ⋅ W₁ + if s > ϵ + return false + elseif s < -ϵ + W₂ = D × A + t = W₂ ⋅ C + t > ϵ && return false + u = -W₂ ⋅ B + u > ϵ && return false + -s > t + u && return false + else + return false + end + end + true +end diff --git a/src/AccretionGeometry/meshes.jl b/src/AccretionGeometry/meshes.jl new file mode 100644 index 00000000..ba1b197e --- /dev/null +++ b/src/AccretionGeometry/meshes.jl @@ -0,0 +1,68 @@ + +struct MeshAccretionGeometry{T} <: AbstractAccretionGeometry{T} + mesh::Vector{Tuple{SVector{3,T},SVector{3,T},SVector{3,T}}} + x_extent::Tuple{T,T} + y_extent::Tuple{T,T} + z_extent::Tuple{T,T} +end + +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)...) +end + +# naive implementation +function bounding_box(mesh::Union{GeometryBasics.Mesh{3,T}}) where {T} + bounding_box(T, mesh) +end + +function bounding_box(mesh::Vector{Tuple{SVector{3,T},SVector{3,T},SVector{3,T}}}) where {T} + bounding_box(T, mesh) +end + +function bounding_box(T, mesh) + xmin = typemax(T) + xmax = -typemax(T) + ymin = typemax(T) + ymax = -typemax(T) + zmin = typemax(T) + zmax = -typemax(T) + for t in mesh + for p in t + let x = p[1], y = p[2], z = p[3] + (x > xmax) && (xmax = x) + (x < xmin) && (xmin = x) + (y > ymax) && (ymax = y) + (y < ymin) && (ymin = y) + (z > zmax) && (zmax = z) + (z < zmin) && (zmin = z) + end + end + end + (xmin, xmax), (ymin, ymax), (zmin, zmax) +end + +function in_nearby_region(m::MeshAccretionGeometry{T}, line_element) where {T} + p = line_element[2] + @inbounds m.x_extent[1] < p[1] < m.x_extent[2] && + m.y_extent[1] < p[2] < m.y_extent[2] && + m.z_extent[1] < p[3] < m.z_extent[2] +end + +function has_intersect(m::MeshAccretionGeometry{T}, line_element) where {T} + for triangle in m.mesh + dist_sq = sum(@.((triangle[1] - line_element[2])^2)) + # assume line element and mesh triangle are small; check if we're within a + # certain distance before running jsr + if dist_sq < 3.0 && jsr_algorithm(triangle..., line_element...) + return true + end + end + false +end + +function collision_callback(m::MeshAccretionGeometry{T}) where {T} + (u, λ, integrator) -> intersects_geometry(m, cartesian_line_element(u, integrator)) +end diff --git a/src/DiscProfiles/DiscProfiles.jl b/src/DiscProfiles/DiscProfiles.jl new file mode 100644 index 00000000..c2a20277 --- /dev/null +++ b/src/DiscProfiles/DiscProfiles.jl @@ -0,0 +1,71 @@ +module DiscProfiles + +using Parameters + +import StaticArrays: @SVector + + +import ..GradusBase: AbstractMetricParams, vector_to_local_sky +import ..GeodesicTracer: tracegeodesics +import ..AccretionGeometry: AbstractAccretionGeometry + +import GeometricalPredicates: Point2D, getx, gety, geta, getb, getc + +include("sky-geometry.jl") +include("corona-models.jl") + +function tracegeodesics( + m::AbstractMetricParams{T}, + model::AbstractCoronaModel{T}, + time_domain::Tuple{T,T}; + n_samples = 1024, + sampler = WeierstrassSampler(res = 100.0), + kwargs..., +) where {T} + us = sample_position(m, model, n_samples) + vs = sample_velocity(m, model, sampler, us, n_samples) + tracegeodesics(m, us, vs, time_domain; kwargs...) +end + + +function tracegeodesics( + m::AbstractMetricParams{T}, + model::AbstractCoronaModel{T}, + time_domain::Tuple{T,T}, + d::AbstractAccretionGeometry{T}, + ; + n_samples = 1024, + sampler = WeierstrassSampler(res = 100.0), + kwargs..., +) where {T} + us = sample_position(m, model, n_samples) + vs = sample_velocity(m, model, sampler, us, n_samples) + tracegeodesics(m, us, vs, d, time_domain; kwargs...) +end + + +function renderprofile( + m::AbstractMetricParams{T}, + model::AbstractCoronaModel{T}, + max_time::T, + d::AbstractAccretionGeometry{T}; + n_samples = 2048, + sampler = WeierstrassSampler(res = 100.0), + kwargs..., +) where {T} + __renderprofile(m, model, d, n_samples, (0.0, max_time); kwargs...) +end + +export AbstractCoronaModel, + LampPostModel, + tracegeodesics, + renderprofile, + LowerHemisphere, + BothHemispheres, + EvenSampler, + WeierstrassSampler, + RandomGenerator, + GoldenSpiralGenerator + + +end # module diff --git a/src/DiscProfiles/corona-models.jl b/src/DiscProfiles/corona-models.jl new file mode 100644 index 00000000..efde305d --- /dev/null +++ b/src/DiscProfiles/corona-models.jl @@ -0,0 +1,32 @@ +abstract type AbstractCoronaModel{T} end + +sample_position(m::AbstractMetricParams{T}, model::AbstractCoronaModel{T}, N) where {T} = + error("Not implemented for $(typeof(model)).") + +@with_kw struct LampPostModel{T} <: AbstractCoronaModel{T} + @deftype T + h = 5.0 + θ = 0.01 + ϕ = 0.0 +end + +function sample_position(::AbstractMetricParams{T}, model::LampPostModel{T}, N) where {T} + u = @SVector [T(0.0), model.h, model.θ, model.ϕ] + fill(u, N) +end + +function sample_velocity( + m::AbstractMetricParams{T}, + ::AbstractCoronaModel{T}, + sampler::AbstractDirectionSampler, + us, + N, +) where {T} + map(1:N) do index + # todo: sampler should have proper iterator interface + i = geti(sampler, index, N) + θ, ϕ = sample_angles(sampler, i, N) + r, t, p = vector_to_local_sky(m, us[index], θ, ϕ) + @SVector [T(0.0), r, t, p] + end +end diff --git a/src/DiscProfiles/disc-profiles.jl b/src/DiscProfiles/disc-profiles.jl new file mode 100644 index 00000000..1e79dbc8 --- /dev/null +++ b/src/DiscProfiles/disc-profiles.jl @@ -0,0 +1,53 @@ +# work in progress code +# not actually used yet + +abstract type AbstractAccretionProfile end +abstract type AbstractInterpolatedProfile <: AbstractAccretionProfile end +abstract type AbstractDelaunayProfile <: AbstractAccretionProfile end + +evaluate(aap::AbstractAccretionProfile, p) = + error("Not implemented for `$(typeof(aap))` yet.") + +struct WeightedPoint{T} <: Point2D + _x::Float64 + _y::Float64 + _weight::T +end +GeometricalPredicates.getx(p::WeightedPoint{T}) where {T} = p._x +GeometricalPredicates.gety(p::WeightedPoint{T}) where {T} = p._y +getw(p::WeightedPoint{T}) where {T} = p._weight + +struct DelaunayDiscProfile{T} <: AbstractDelaunayProfile + tesselation::DelaunayTessellation2D{WeightedPoint{T}} + radius::Float64 +end +xfm_forward(ddp::DelaunayDiscProfile{T}, p) where {T} = + ((p .+ ddp.radius) ./ (2 .* ddp.radius)) .* 0.99 .+ 1.005 +xfm_backward(ddp::DelaunayDiscProfile{T}, p) where {T} = + (((points .- 1.005) ./ 0.99) .* 2 .* ddp.radius) .- ddp.radius + + +function __renderprofile( + m::AbstractMetricParams{T}, + model::AbstractCoronaModel{T}, + d::AbstractAccretionGeometry{T}, + time_domain; + sampler, + kwargs..., +) where {T} + # TODO + error("Not implemented for $(typeof(d)).") +end + + +function __renderprofile( + m::AbstractMetricParams{T}, + model::AbstractCoronaModel{T}, + d::GeometricThinDisc{T}, + time_domain; + sampler, + kwargs..., +) where {T} + # TODO + error("Not implemented for $(typeof(d)).") +end diff --git a/src/DiscProfiles/sky-geometry.jl b/src/DiscProfiles/sky-geometry.jl new file mode 100644 index 00000000..894561b0 --- /dev/null +++ b/src/DiscProfiles/sky-geometry.jl @@ -0,0 +1,60 @@ +abstract type AbstractSkyDomain end + +struct LowerHemisphere <: AbstractSkyDomain end +struct BothHemispheres <: AbstractSkyDomain end + +abstract type AbstractGenerator end + +struct RandomGenerator <: AbstractGenerator end +struct GoldenSpiralGenerator <: AbstractGenerator end + +abstract type AbstractDirectionSampler{AbstractSkyDomain,AbstractGenerator} end + +struct EvenSampler{D,G} <: AbstractDirectionSampler{D,G} + EvenSampler(domain::AbstractSkyDomain, generator::AbstractGenerator) = + new{typeof(domain),typeof(generator)}() + EvenSampler(; + domain::AbstractSkyDomain = LowerHemisphere(), + generator::AbstractGenerator = GoldenSpiralGenerator(), + ) = EvenSampler(domain, generator) +end +struct WeierstrassSampler{D,G} <: AbstractDirectionSampler{D,G} + resolution::Float64 + WeierstrassSampler(res, domain::AbstractSkyDomain, generator::AbstractGenerator) = + new{typeof(domain),typeof(generator)}(res) + WeierstrassSampler(; + res = 100.0, + domain::AbstractSkyDomain = LowerHemisphere(), + generator::AbstractGenerator = GoldenSpiralGenerator(), + ) = WeierstrassSampler(res, domain, generator) +end + +@inline geti(::AbstractDirectionSampler{D,RandomGenerator}, i, N) where {D} = + rand(Float64) * N +@inline geti(::AbstractDirectionSampler{D,G}, i, N) where {D,G} = i + +@inline sample_azimuth(::AbstractDirectionSampler{D,GoldenSpiralGenerator}, i) where {D} = + π * (1 + √5) * i +@inline sample_azimuth(::AbstractDirectionSampler{D,G}, i) where {D,G} = 2π * i + + +sample_angles(sm::AbstractDirectionSampler{D,G}, i, N) where {D,G} = + (sample_elevation(sm, i / N), sample_azimuth(sm, i)) +@inline sample_angles(sm::WeierstrassSampler{D}, i, N) where {D} = + (sample_elevation(sm, i), sample_azimuth(sm, i)) + + +sample_elevation(sm::AbstractDirectionSampler{D,G}, i) where {D,G} = + error("Not implemented for $(typeof(sm)).") +@inline sample_elevation(::EvenSampler{LowerHemisphere}, i) = acos(1 - i) +@inline sample_elevation(::EvenSampler{BothHemispheres}, i) = acos(1 - 2i) +@inline sample_elevation(sm::WeierstrassSampler{LowerHemisphere}, i) = + π - 2atan(√(sm.resolution / i)) +@inline function sample_elevation(sm::WeierstrassSampler{BothHemispheres}, i) + ϕ = 2atan(√(sm.resolution / i)) + if iseven(i) + ϕ + else + π - ϕ + end +end diff --git a/src/FirstOrderMethods/FirstOrderMethods.jl b/src/FirstOrderMethods/FirstOrderMethods.jl new file mode 100644 index 00000000..2b904727 --- /dev/null +++ b/src/FirstOrderMethods/FirstOrderMethods.jl @@ -0,0 +1,36 @@ +module FirstOrderMethods + +using Accessors +using Parameters +using DocStringExtensions +using StaticArrays + +using SciMLBase +using DiffEqCallbacks + +import ..GradusBase: + AbstractMetricParams, + inner_radius, + AbstractGeodesicPoint, + get_endpoint, + geodesic_point_type, + unpack_solution, + SciMLBase + +import ..GeodesicTracer: + DiscreteCallback, + terminate!, + integrator_problem, + metric_callback, + create_callback_set, + constrain, + alpha_beta_to_vel + +abstract type AbstractFirstOrderMetricParams{T} <: AbstractMetricParams{T} end + +include("implementation.jl") +include("callbacks.jl") + +export AbstractFirstOrderMetricParams, FirstOrderGeodesicPoint, BoyerLindquistFO + +end # module diff --git a/src/FirstOrderMethods/callbacks.jl b/src/FirstOrderMethods/callbacks.jl new file mode 100644 index 00000000..b03c1099 --- /dev/null +++ b/src/FirstOrderMethods/callbacks.jl @@ -0,0 +1,24 @@ +function ensure_domain(m::AbstractFirstOrderMetricParams{T}) where {T} + min_r = inner_radius(m) + (u, λ, integrator) -> u[2] ≤ min_r * 1.1 || u[2] > 1200.0 +end + +function flip_radial_sign!(integrator) + p = integrator.p + integrator.p = @set p.r = -p.r + integrator.sol.prob.p.changes[1] = integrator.t[end] +end + +function flip_angular_sign!(integrator) + p = integrator.p + integrator.p = @set p.θ = -p.θ + integrator.sol.prob.p.changes[2] = integrator.t[end] +end + +function radial_negative_check(m::AbstractFirstOrderMetricParams{T}) where {T} + (u, λ, integrator) -> Vr(m, u, integrator.p) < 0 +end + +function angular_negative_check(m::AbstractFirstOrderMetricParams{T}) where {T} + (u, λ, integrator) -> Vθ(m, u, integrator.p) < 0 +end diff --git a/src/FirstOrderMethods/implementation.jl b/src/FirstOrderMethods/implementation.jl new file mode 100644 index 00000000..3355cf37 --- /dev/null +++ b/src/FirstOrderMethods/implementation.jl @@ -0,0 +1,73 @@ + +@with_kw struct FirstOrderGeodesicPoint{T,P} <: AbstractGeodesicPoint{T} + retcode::Symbol + t::T + u::AbstractVector{T} + v::AbstractVector{T} + 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 get_endpoint( + 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) +end + +function metric_callback(m::AbstractFirstOrderMetricParams{T}) where {T} + ( + DiscreteCallback(ensure_domain(m), terminate!), + DiscreteCallback(radial_negative_check(m), flip_radial_sign!), + DiscreteCallback(angular_negative_check(m), flip_angular_sign!), + ) +end + +four_velocity(u, m::AbstractFirstOrderMetricParams{T}, p) where {T} = + error("Not implmented for $(typeof(m)).") + +make_parameters(L, Q, sign_θ, T) = + (L = L, Q = Q, r = -1, θ = convert(Int, sign_θ), changes = T[0.0, 0.0]) + +function integrator_problem( + m::AbstractFirstOrderMetricParams{T}, + pos::StaticVector{S,T}, + vel::StaticVector{S,T}, + time_domain, +) where {S,T} + L, Q = calc_lq(m, pos, vel) + ODEProblem{false}(pos, time_domain, make_parameters(L, Q, vel[2], T)) do u, p, λ + SVector(four_velocity(u, m, p)...) + end +end + +function integrator_problem( + m::AbstractFirstOrderMetricParams{T}, + pos::AbstractVector{T}, + vel::AbstractVector{T}, + time_domain, +) where {T} + L, Q = calc_lq(m, pos, vel) + ODEProblem{true}(pos, time_domain, make_parameters(L, Q, vel[2], T)) do du, u, p, λ + du .= four_velocity(u, m, p) + end +end + +convert_velocity_type(u::StaticVector{S,T}, v) where {S,T} = convert(SVector{S,T}, v) +convert_velocity_type(u::AbstractVector{T}, v) where {T} = convert(typeof(u), collect(v)) + +Vr(m::AbstractFirstOrderMetricParams{T}, u, p) where {T} = + error("Not implmented for $(typeof(m)).") +Vθ(m::AbstractFirstOrderMetricParams{T}, u, p) where {T} = + error("Not implmented for $(typeof(m)).") + +calc_lq(m::AbstractFirstOrderMetricParams{T}, pos, vel) where {T} = + error("Not implmented for $(typeof(m)).") diff --git a/src/GeodesicTracer/GeodesicTracer.jl b/src/GeodesicTracer/GeodesicTracer.jl new file mode 100644 index 00000000..4b6bfd52 --- /dev/null +++ b/src/GeodesicTracer/GeodesicTracer.jl @@ -0,0 +1,76 @@ +module GeodesicTracer + +using SciMLBase +using OrdinaryDiffEq +using DiffEqCallbacks + +using StaticArrays +using DocStringExtensions +using Parameters + +import ..GradusBase: + AbstractMetricParams, geodesic_eq, constrain, on_chart, inner_radius, metric + +import ForwardDiff + +include("callbacks.jl") +include("problem.jl") +include("tracer.jl") +include("constraints.jl") +include("utility.jl") +include("auto-diff.jl") + +""" + tracegeodesics( + m::AbstractMetricParams{T}, + init_positions, init_velocities, + time_domain::Tuple{T,T} + ; + μ = 0.0f0, + callbacks=Nothing, + solver=Tsit5(), + solver_opts... + ) + +Integrate a geodesic for metric parameterised by `m`, for some initial positions and velocities. +The positions and velocities may be + + - a single position and velocity in the form of a vector of numbers + - a collection of positions and velocities, as either a vector of vectors, or as a matrix + +The matrix specification reads each corresponding column as the initial position and velocity. When a collection of +positions and velocities is supplied, this method dispatched `EnsembleProblem`, offering `ensemble` as a `solver_opts`, +specifying the ensemble method to use. + +`solver_opts` are the common solver options in DifferentialEquations. +""" +function tracegeodesics( + m::AbstractMetricParams{T}, + init_positions, + init_velocities, + time_domain::Tuple{T,T}; + solver = Tsit5(), + μ = 0.0, + solver_opts..., +) where {T} + __tracegeodesics( + m, + init_positions, + # ensure everything already normalised + constrain_all(m, init_positions, init_velocities, T(μ)), + time_domain, + solver; + callback = nothing, + abstol = 1e-8, + reltol = 1e-8, + solver_opts..., + ) +end + +export tracegeodesics, + map_impact_parameters, + AbstractAutoDiffMetricParams, + AbstractAutoDiffStaticAxisSymmetricParams, + metric_components + +end # module diff --git a/src/GeodesicTracer/auto-diff.jl b/src/GeodesicTracer/auto-diff.jl new file mode 100644 index 00000000..a979ed71 --- /dev/null +++ b/src/GeodesicTracer/auto-diff.jl @@ -0,0 +1,168 @@ +# for use with static, axis-symmetric metrics +# create new abstract type for easy re-definition +abstract type AbstractAutoDiffMetricParams{T} <: AbstractMetricParams{T} end +abstract type AbstractAutoDiffStaticAxisSymmetricParams{T} <: + AbstractAutoDiffMetricParams{T} end + +# interface +metric_components(m::AbstractAutoDiffStaticAxisSymmetricParams{T}, rθ) where {T} = + error("Not implemented for $(typeof(m)).") + +""" + inverse_metric_components(g_comp) + +Calculates ``g^{tt}``, ``g^{rr}``, ``g^{\\theta\\theta}``, ``g^{\\phi\\phi}``, ``g^{t\\phi}`` of a static, +axis-symmetric metric from ``g_{tt}``, ``g_{rr}``, ``g_{\\theta\\theta}``, ``g_{\\phi\\phi}``, ``g_{t\\phi}`` +using a symbolically computed inverse matrix method. + +## Developer notes + +To recreate: + +```julia +using Symbolics +@variables g[1:5] # non zero metric components +metric = [ + g[1] 0 0 g[5] + 0 g[2] 0 0 + 0 0 g[3] 0 + g[5] 0 0 g[4] +] +inv(metric) +``` +""" +@fastmath function inverse_metric_components(g_comp) + @inbounds let g1 = g_comp[1], + g2 = g_comp[2], + g3 = g_comp[3], + g4 = g_comp[4], + g5 = g_comp[5] + + ( + (g2 * g3 * g4) / (g1 * g2 * g3 * g4 - (g5^2) * g2 * g3), + (g1 * g3 * g4 - (g5^2) * g3) / (g1 * g2 * g3 * g4 - (g5^2) * g2 * g3), + (g1 * g2 * g4 - (g5^2) * g2) / (g1 * g2 * g3 * g4 - (g5^2) * g2 * g3), + (g1 * g2 * g3) / (g1 * g2 * g3 * g4 - (g5^2) * g2 * g3), + (-g2 * g3 * g5) / (g1 * g2 * g3 * g4 - (g5^2) * g2 * g3), + ) + end +end + + +""" + calc_symb_geodesic_equation(g1inv, j1, j2, v) + +Symbolically pre-computed geodesic equations for a static, axis-symmetric metric. + +## Developer notes + +To recreate: + +```julia +using Symbolics, Tullio +@variables ginv[1:5], j1[1:5], j2[1:5], v[1:4] # non zero metric components +inverse_metric = [ + ginv[1] 0 0 ginv[5] + 0 ginv[2] 0 0 + 0 0 ginv[3] 0 + ginv[5] 0 0 ginv[4] +] +j1_mat = [ + j1[1] 0 0 j1[5] + 0 j1[2] 0 0 + 0 0 j1[3] 0 + j1[5] 0 0 j1[4] +] +j2_mat = [ + j2[1] 0 0 j2[5] + 0 j2[2] 0 0 + 0 0 j2[3] 0 + j2[5] 0 0 j2[4] +] +j0 = zeros(Float64, (4, 4)) +jacobian = (j0, j1_mat, j2_mat, j0) +# christoffel symbols +@tullio Γ[i, k, l] := + 1 / 2 * + inverse_metric[i, m] * + (jacobian[l][m, k] + jacobian[k][m, l] - jacobian[m][k, l]) +# compute geodesic equation +@tullio δxδλ[i] := -v[j] * Γ[i, j, k] * v[k] +``` +""" +@fastmath function compute_geodesic_equation(ginv, j1, j2, v::AbstractArray{T}) where {T} + @inbounds @SVector [ + -2(0.5ginv[5] * j1[4] + 0.5ginv[1] * j1[5]) * v[2] * v[4] - + 2(0.5ginv[1] * j1[1] + 0.5ginv[5] * j1[5]) * v[1] * v[2] - + 2(0.5ginv[1] * j2[1] + 0.5ginv[5] * j2[5]) * v[1] * v[3] - + 2(0.5ginv[5] * j2[4] + 0.5ginv[1] * j2[5]) * v[3] * v[4], + 0.5(v[1]^2) * ginv[2] * j1[1] + + 0.5(v[3]^2) * ginv[2] * j1[3] + + 0.5(v[4]^2) * ginv[2] * j1[4] + + ginv[2] * j1[5] * v[1] * v[4] - 0.5(v[2]^2) * ginv[2] * j1[2] - + ginv[2] * j2[2] * v[2] * v[3], + 0.5(v[1]^2) * ginv[3] * j2[1] + + 0.5(v[2]^2) * ginv[3] * j2[2] + + 0.5(v[4]^2) * ginv[3] * j2[4] + + ginv[3] * j2[5] * v[1] * v[4] - 0.5(v[3]^2) * ginv[3] * j2[3] - + ginv[3] * j1[3] * v[2] * v[3], + -2(0.5ginv[4] * j1[4] + 0.5ginv[5] * j1[5]) * v[2] * v[4] - + 2(0.5ginv[4] * j2[4] + 0.5ginv[5] * j2[5]) * v[3] * v[4] - + 2(0.5ginv[5] * j1[1] + 0.5ginv[4] * j1[5]) * v[1] * v[2] - + 2(0.5ginv[4] * j2[5] + 0.5ginv[5] * j2[1]) * v[1] * v[3], + ] +end + +@fastmath function constrain_time(g_comp, v, μ = 0.0, positive::Bool = true) + @inbounds begin + discriminant = ( + -g_comp[1] * g_comp[2] * v[2]^2 - g_comp[1] * g_comp[3] * v[3]^2 - + g_comp[1] * μ^2 - (g_comp[1] * g_comp[4] - g_comp[5]^2) * v[4]^2 + ) + if positive + -(g_comp[5] * v[4] + √discriminant) / g_comp[1] + else + -(g_comp[5] * v[4] - √discriminant) / g_comp[1] + end + end +end + +function constrain( + m::AbstractAutoDiffStaticAxisSymmetricParams{T}, + u, + v; + μ::T = 0.0, +) where {T} + rθ = @SVector [u[2], u[3]] + g_comps = metric_components(m, rθ) + constrain_time(g_comps, v, μ) +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 = ForwardDiff.vector_mode_jacobian(x -> metric_components(m, x), rθ) + # calculate all non-zero inverse matric components + ginv_comps = inverse_metric_components(g_comps) + # calculate acceleration + compute_geodesic_equation(ginv_comps, jacs[:, 1], jacs[:, 2], v) +end + + +function metric(m::AbstractAutoDiffStaticAxisSymmetricParams{T}, u) where {T} + rθ = @SVector [u[2], u[3]] + comps = metric_components(m, rθ) + @SMatrix [ + comps[1] 0 0 comps[5] + 0 comps[2] 0 0 + 0 0 comps[3] 0 + comps[5] 0 0 comps[4] + ] +end diff --git a/src/GeodesicTracer/callbacks.jl b/src/GeodesicTracer/callbacks.jl new file mode 100644 index 00000000..25bce4c7 --- /dev/null +++ b/src/GeodesicTracer/callbacks.jl @@ -0,0 +1,22 @@ +function metric_callback(m::AbstractMetricParams{T}) where {T} + min_r = inner_radius(m) + DiscreteCallback((u, λ, integrator) -> u[6] ≤ min_r * 1.1 || u[6] > 1200.0, terminate!) +end + +function create_callback_set(m::AbstractMetricParams{T}, cb::Nothing) where {T} + mcb = metric_callback(m) + mcb isa Tuple ? CallbackSet(mcb...) : mcb +end + +function create_callback_set( + m::AbstractMetricParams{T}, + cb::NTuple{N,SciMLBase.DECallback}, +) where {T,N} + mcb = metric_callback(m) + mcb isa Tuple ? CallbackSet(mcb..., cb...) : CallbackSet(mcb, cb...) +end + +function create_callback_set(m::AbstractMetricParams{T}, cb::SciMLBase.DECallback) where {T} + mcb = metric_callback(m) + mcb isa Tuple ? CallbackSet(mcb..., cb) : CallbackSet(mcb, cb) +end diff --git a/src/GeodesicTracer/constraints.jl b/src/GeodesicTracer/constraints.jl new file mode 100644 index 00000000..2a03bc56 --- /dev/null +++ b/src/GeodesicTracer/constraints.jl @@ -0,0 +1,51 @@ +# inplace +function constrain_all( + m::AbstractMetricParams{T}, + us::AbstractArray{T,2}, + vs::AbstractArray{T,2}, + μ, +) where {T<:Number} + us_cols = eachcol(us) + vs_cols = eachcol(vs) + curried(u, v) = constrain(m, u, v, μ = μ) + @. vs[1, :] = curried(us_cols, vs_cols) + vs +end + +# SMatrix -- do we actually want to support this? +function constrain_all( + m::AbstractMetricParams{T}, + us::SMatrix{4,M,T,L}, + vs::SMatrix{4,M,T,L}, + μ, +) where {M,T<:Number,L} + alloc_vs = constrain_all(m, us, similar(vs), μ) + SMatrix{4,M}(alloc_vs) +end + +# vectors of vectors +function constrain_all(m::AbstractMetricParams{T}, us, vs, μ) where {T<:Number} + curried(u, v) = constrain_all(m, u, v, μ) + curried.(us, vs) +end + +function constrain_all( + m::AbstractMetricParams{T}, + u::AbstractVector{T}, + v::AbstractVector{T}, + μ, +) where {T<:Number} + v[1] = constrain(m, u, v, μ = μ) + v +end + +function constrain_all( + m::AbstractMetricParams{T}, + u::StaticVector{S,T}, + v::StaticVector{S,T}, + μ, +) where {S,T<:Number} + mut = MVector{S,T}(v) + mut[1] = constrain(m, u, v, μ = μ) + SVector{S,T}(mut) +end diff --git a/src/GeodesicTracer/problem.jl b/src/GeodesicTracer/problem.jl new file mode 100644 index 00000000..debc6d06 --- /dev/null +++ b/src/GeodesicTracer/problem.jl @@ -0,0 +1,21 @@ +function integrator_problem( + m::AbstractMetricParams{T}, + pos::StaticVector{S,T}, + vel::StaticVector{S,T}, + time_domain, +) where {S,T} + SecondOrderODEProblem{false}(vel, pos, time_domain, m) do v, u, p, λ + SVector(geodesic_eq(p, u, v)...) + end +end + +function integrator_problem( + m::AbstractMetricParams{T}, + pos::AbstractVector{T}, + vel::AbstractVector{T}, + time_domain, +) where {T} + SecondOrderODEProblem{true}(vel, pos, time_domain, m) do dv, v, u, p, λ + dv .= geodesic_eq(p, u, v) + end +end diff --git a/src/GeodesicTracer/tracer.jl b/src/GeodesicTracer/tracer.jl new file mode 100644 index 00000000..ff919ed2 --- /dev/null +++ b/src/GeodesicTracer/tracer.jl @@ -0,0 +1,91 @@ +# single position and single velocity +function __tracegeodesics( + m::AbstractMetricParams{T}, + init_pos::AbstractVector{T}, + init_vel::AbstractVector{T}, + time_domain::Tuple{T,T}, + solver; + callback, + kwargs..., +) where {T} + prob = integrator_problem(m, init_pos, init_vel, time_domain) + cbs = create_callback_set(m, callback) + solve_geodesic(solver, prob; callback = cbs, kwargs...) +end + +# columnar of positions and velocity +function __tracegeodesics( + m::AbstractMetricParams{T}, + init_positions::AbstractArray{T,2}, + init_velocities::AbstractArray{T,2}, + time_domain::Tuple{T,T}, + solver; + ensemble = EnsembleThreads(), + callback, + kwargs..., +) where {T} + prob = integrator_problem( + m, + @view(init_positions[:, 1]), + @view(init_velocities[:, 1]), + time_domain, + ) + ens_prob = EnsembleProblem( + prob, + prob_func = (prob, i, repeat) -> integrator_problem( + m, + @view(init_positions[:, i]), + @view(init_velocities[:, i]), + time_domain, + ), + safetycopy = false, + ) + + cbs = create_callback_set(m, callback) + solve_geodesic( + solver, + ens_prob, + ensemble; + trajectories = size(init_velocities, 2), + callback = cbs, + kwargs..., + ) +end + +# indexables +function __tracegeodesics( + m::AbstractMetricParams{T}, + init_positions, + init_velocities, + time_domain::Tuple{T,T}, + solver; + ensemble = EnsembleThreads(), + callback, + kwargs..., +) where {T} + prob = integrator_problem(m, init_positions[1], init_velocities[1], time_domain) + ens_prob = EnsembleProblem( + prob, + prob_func = (prob, i, repeat) -> + integrator_problem(m, init_positions[i], init_velocities[i], time_domain), + safetycopy = false, + ) + + cbs = create_callback_set(m, callback) + solve_geodesic( + solver, + ens_prob, + ensemble; + trajectories = length(init_velocities), + callback = cbs, + kwargs..., + ) +end + +function solve_geodesic(solver, prob, ensemble; solver_opts...) + solve(prob, solver, ensemble; solver_opts...) +end + +function solve_geodesic(solver, prob; solver_opts...) + solve(prob, solver, ; solver_opts...) +end diff --git a/src/GeodesicTracer/utility.jl b/src/GeodesicTracer/utility.jl new file mode 100644 index 00000000..0f95b745 --- /dev/null +++ b/src/GeodesicTracer/utility.jl @@ -0,0 +1,37 @@ +""" + map_impact_parameters(m::AbstractMetricParams{T}, u, α, β) + +Map impact parameters `α` and `β` to a velocity vector at some position `u` in the given metric `m`. +""" +function map_impact_parameters( + m::AbstractMetricParams{T}, + u::AbstractVector{T}, + α::N, + β::N, +) where {T,N<:Number} + [0.0, alpha_beta_to_vel(m, u, α, β)...] +end + +function map_impact_parameters( + m::AbstractMetricParams{T}, + u::SVector{S,T}, + α::N, + β::N, +) where {S,T,N<:Number} + SVector{S,T}(0.0, alpha_beta_to_vel(m, u, α, β)...) +end + +function map_impact_parameters( + m::AbstractMetricParams{T}, + u, + α::AbstractVector{P}, + β::AbstractVector{P}, +) where {T,P} + curried(_α, _β) = map_impact_parameters(m, u, _α, _β) + curried.(α, β) +end + +function alpha_beta_to_vel(m::AbstractMetricParams{T}, u, α, β) where {T} + reg = u[2]^2 + -1.0, -β / reg, -α / reg +end diff --git a/src/Gradus.jl b/src/Gradus.jl new file mode 100644 index 00000000..85ee0b2a --- /dev/null +++ b/src/Gradus.jl @@ -0,0 +1,79 @@ +module Gradus + +include("GradusBase/GradusBase.jl") +include("GeodesicTracer/GeodesicTracer.jl") +include("FirstOrderMethods/FirstOrderMethods.jl") +include("Rendering/Rendering.jl") +include("AccretionGeometry/AccretionGeometry.jl") +include("DiscProfiles/DiscProfiles.jl") + +using .GradusBase +using .GeodesicTracer +using .FirstOrderMethods +using .Rendering +using .AccretionGeometry +using .DiscProfiles + +using StaticArrays +using Parameters +using DocStringExtensions + +# GradusBase +export AbstractMetricParams, + metric_params, + metric, + get_endpoint, + GeodesicPoint, + vector_to_local_sky, + AbstractMetricParams, + geodesic_eq, + geodesic_eq!, + constrain, + on_chart, + inner_radius, + metric_type + +# GeodesicTracer +export tracegeodesics, + map_impact_parameters, + AbstractAutoDiffMetricParams, + AbstractAutoDiffStaticAxisSymmetricParams, + metric_components + +# FirstOrderMethods +export AbstractFirstOrderMetricParams, FirstOrderGeodesicPoint, BoyerLindquistFO + +# Rendering +export rendergeodesics, + prerendergeodesics, PointFunction, FilterPointFunction, ConstPointFunctions + +# AccretionGeometry +export AbstractAccretionGeometry, + AbstractAccretionDisc, MeshAccretionGeometry, GeometricThinDisc + +# DiscProfiles +export AbstractCoronaModel, + LampPostModel, + tracegeodesics, + renderprofile, + LowerHemisphere, + BothHemispheres, + EvenSampler, + WeierstrassSampler, + RandomGenerator, + GoldenSpiralGenerator + +# pre-defined metrics +include("metrics/metrics.jl") + +export BoyerLindquistAD, BoyerLindquistFO, JohannsenAD, MorrisThorneAD + +# downstream modules and work +include("special-radii.jl") + +include("AccretionFormulae/AccretionFormulae.jl") + +using .AccretionFormulae +export redshift + +end # module diff --git a/src/GradusBase/GradusBase.jl b/src/GradusBase/GradusBase.jl new file mode 100644 index 00000000..5c441e53 --- /dev/null +++ b/src/GradusBase/GradusBase.jl @@ -0,0 +1,25 @@ +module GradusBase + +import Parameters: @with_kw +import SciMLBase + +include("metric-params.jl") +include("physical-quantities.jl") +include("geodesic-solutions.jl") +include("geometry.jl") + +export AbstractMetricParams, + metric_params, + metric, + get_endpoint, + GeodesicPoint, + vector_to_local_sky, + AbstractMetricParams, + geodesic_eq, + geodesic_eq!, + constrain, + on_chart, + inner_radius, + metric_type + +end # module diff --git a/src/GradusBase/geodesic-solutions.jl b/src/GradusBase/geodesic-solutions.jl new file mode 100644 index 00000000..3813f890 --- /dev/null +++ b/src/GradusBase/geodesic-solutions.jl @@ -0,0 +1,48 @@ +abstract type AbstractGeodesicPoint{T} end + +@with_kw struct GeodesicPoint{T} <: AbstractGeodesicPoint{T} + retcode::Symbol + t::T + u::AbstractVector{T} + v::AbstractVector{T} + # 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} +end + +# TODO: GeodesicPath structure for the full geodesic path +# do we want to support this? + +function unpack_solution(sol::SciMLBase.AbstractODESolution{T,N,S}) where {T,N,S} + u = sol.u + p = sol.prob.p + t = sol.t + (u, t, p) +end + +function unpack_solution(simsol::SciMLBase.AbstractEnsembleSolution{T,N,V}) where {T,N,V} + map(unpack_solution, simsol) +end + +function get_endpoint( + m::AbstractMetricParams{T}, + sol::SciMLBase.AbstractODESolution{T,N,S}, +) where {T,N,S} + us, ts, _ = unpack_solution(sol) + u = us[end].x[2] + v = us[end].x[1] + t = ts[end] + GeodesicPoint(sol.retcode, t, u, v) +end + +function get_endpoint( + m::AbstractMetricParams{T}, + simsol::SciMLBase.AbstractEnsembleSolution{T,N,S}, +) where {T,N,S} + map(sol -> get_endpoint(m, sol), simsol) +end diff --git a/src/GradusBase/geometry.jl b/src/GradusBase/geometry.jl new file mode 100644 index 00000000..d3b00289 --- /dev/null +++ b/src/GradusBase/geometry.jl @@ -0,0 +1,3 @@ +function vector_to_local_sky(m::AbstractMetricParams{T}, u, θ, ϕ) where {T} + error("Not implemented for $(typeof(m))") +end diff --git a/src/GradusBase/metric-params.jl b/src/GradusBase/metric-params.jl new file mode 100644 index 00000000..19c28e8f --- /dev/null +++ b/src/GradusBase/metric-params.jl @@ -0,0 +1,80 @@ +""" + abstract type AbstractMetricParams{T} end + +Abstract type used to dispatch different geodesic problems. +""" +abstract type AbstractMetricParams{T} end + + +# contains the full metric components (this type needed for DiffGeoSymbolics) +abstract type AbstractMetric{T} <: AbstractMatrix{T} end + +metric_params(m::AbstractMetric{T}) where {T} = + error("Not implemented for metric $(typeof(m))") + +""" + geodesic_eq(m::AbstractMetricParams{T}, u, v) + geodesic_eq!(m::AbstractMetricParams{T}, u, v) + +Calculate the acceleration components of the geodesic equation given a position `u`, a velocity `v`, and a metric `m`. +""" +geodesic_eq(m::AbstractMetricParams{T}, u, v) where {T} = + error("Not implemented for metric parameters $(typeof(m))") +geodesic_eq!(m::AbstractMetricParams{T}, u, v) where {T} = + error("Not implemented for metric parameters $(typeof(m))") + +""" + constrain(m::AbstractMetricParams{T}, u, v; μ::T=0.0f0) + +Give time component which would constrain a velocity vector `v` at position `x` to be a +geodesic with mass `μ`. +""" +constrain(m::AbstractMetricParams{T}, u, v; μ::T = 0.0) where {T} = + error("Not implemented for metric parameters $(typeof(m))") + +""" + on_chart(m::AbstractMetricParams{T}, u) + +Check if point `u` is a valid point for the metric described by `m`. + +Returns false is `u` is a singularity. +""" +on_chart(m::AbstractMetricParams{T}, u) where {T} = !(sum(u) ≈ 0) + + +""" + inner_radius(m::AbstractMetricParams{T}) + +Return the innermost valid coordinate relative to the origin, for use in geodesic tracing. + +This usually represents some property of the metric, e.g. event horizon radius in Kerr/Schwarzschild metrics, or +throat diameter in worm hole metrics. +""" +inner_radius(m::AbstractMetricParams{T}) where {T} = convert(T, 0.0) + +""" + metric_type(m::AbstractMetricParams{T}) + +Return the [`AbstractMetric`](@ref) type associated with the metric parameters `m`. +""" +metric_type(m::AbstractMetricParams{T}) where {T} = + error("Not implemented for metric parameters $(typeof(m))") + + +""" + metric(m::AbstractMetricParams{T}, u) + +Numerically evaluate the corresponding metric for [`AbstractMetricParams`](@ref), given parameter values `m` +and some point `u`. +""" +metric(m::AbstractMetricParams{T}, u) where {T} = + error("Not implemented for metric $(typeof(m))") + +# do we actually want to support this? +# since if it's a symbolic matrix, you can subs other ways better? +#""" +# metric(m::AbstractMetric{T}, u) +# +#Evaluate the metric at a point `u`. +#""" +#metric(m::AbstractMetric{T}, u) where {T} = error("Not implemented for metric $(typeof(m))") diff --git a/src/GradusBase/physical-quantities.jl b/src/GradusBase/physical-quantities.jl new file mode 100644 index 00000000..11d95601 --- /dev/null +++ b/src/GradusBase/physical-quantities.jl @@ -0,0 +1,37 @@ + +# both energy and angular momentum +# assume time only coupled to radial coordinate +# need to think of a nice way to keep this efficient +# whilst allowing metrics with other couplings + +""" + E(m::AbstractMatrix{T}, v) + E(m::AbstractMetricParams{T}, u, v) + +Compute the energy for a numerically evaluated metric, and some velocity four vector `v`, +```math +E = - p_t = - g_{t\\nu} p^\\nu. +``` + +For null geodesics, the velocity is the momentum ``v^\\nu = p^\\nu``. For massive geodesics, +the mass ``\\mu`` needs to be known to compute ``\\mu v^\\nu = p^\\nu``. +""" +function E(metric::AbstractMatrix{T}, v) where {T} + T(@inbounds -(metric[1, 1] * v[1] + metric[1, 4] * v[4])) +end +E(m::AbstractMetricParams{T}, u, v) where {T} = E(metric(m, u), v) + + +""" + Lz(m::AbstractMatrix{T}, v) + Lz(m::AbstractMetricParams{T}, u, v) + +Compute the angular momentum for a numerically evaluated metric, and some velocity four vector `v`. +```math +L_z = p_\\phi = - g_{\\phi\\nu} p^\\nu. +``` +""" +function Lz(metric::AbstractMatrix{T}, v) where {T} + T(@inbounds metric[4, 4] * v[4] + metric[1, 4] * v[1]) +end +Lz(m::AbstractMetricParams{T}, u, v) where {T} = Lz(metric(m, u), v) diff --git a/src/Rendering/Rendering.jl b/src/Rendering/Rendering.jl new file mode 100644 index 00000000..8a765e03 --- /dev/null +++ b/src/Rendering/Rendering.jl @@ -0,0 +1,61 @@ +module Rendering + +import Base.Threads: @threads + +using Parameters + +import SciMLBase + +using ..GradusBase +using ..GeodesicTracer + +include("utility.jl") +include("render-cache.jl") +include("point-functions.jl") +include("render.jl") + + +function rendergeodesics( + m::AbstractMetricParams{T}, + init_pos, + max_time; + pf = ConstPointFunctions.shadow, + kwargs..., +) where {T} + __rendergeodesics( + m, + init_pos; + image_width = 350, + image_height = 250, + fov_factor = 3.0, + max_time = convert(T, max_time), + pf = pf, + kwargs..., + ) +end + +function prerendergeodesics( + m::AbstractMetricParams{T}, + init_pos, + max_time; + cache = DefaultCache(), + kwargs..., +) where {T} + __prerendergeodesics( + m, + init_pos, + cache; + image_width = 350, + image_height = 250, + fov_factor = 3.0, + max_time = convert(T, max_time), + kwargs..., + ) +end + + +export rendergeodesics, + prerendergeodesics, PointFunction, FilterPointFunction, ConstPointFunctions + + +end # module diff --git a/src/Rendering/point-functions.jl b/src/Rendering/point-functions.jl new file mode 100644 index 00000000..1be21546 --- /dev/null +++ b/src/Rendering/point-functions.jl @@ -0,0 +1,66 @@ +abstract type AbstractPointFunction end + +struct PointFunction{F} <: AbstractPointFunction + f::F +end + +struct FilterPointFunction{F,T} <: AbstractPointFunction + f::F + default::T +end + +@inline function (pf::AbstractPointFunction)( + m::AbstractMetricParams{T}, + gp::GradusBase.AbstractGeodesicPoint{T}, + max_time; + kwargs..., +)::T where {T} + convert(T, pf.f(m, gp, max_time; kwargs...)) +end + +@inline function apply( + pf::AbstractPointFunction, + rc::SolutionRenderCache{M,T,G}; + kwargs..., +) where {M,T,G} + map(sol -> pf.f(rc.m, get_endpoint(m, sol), rc.max_time; kwargs...), rc.geodesics) +end + +@inline function apply( + pf::AbstractPointFunction, + rc::EndpointRenderCache{M,T,P}; + kwargs..., +) where {M,T,P} + map(gp -> pf.f(rc.m, gp, rc.max_time; kwargs...), rc.points) +end + +@inline function Base.:∘(pf1::AbstractPointFunction, pf2::AbstractPointFunction) + PointFunction((m, gp, max_time; kwargs...) -> pf1.f(pf2.f(m, gp, max_time; kwargs...))) +end + +@inline function Base.:∘(pf1::AbstractPointFunction, pf2::FilterPointFunction{F}) where {F} + PointFunction( + (m, gp, max_time; kwargs...) -> begin + pass_on = pf2.f(m, gp, max_time; kwargs...) + if pass_on + pf1.f(m, gp, max_time; kwargs...) + else + pf2.default + end + end, + ) +end + +module ConstPointFunctions +import ..Rendering: PointFunction, FilterPointFunction + +const filter_early_term = + FilterPointFunction((m, gp, max_time; kwargs...) -> gp.t < 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 shadow = affine_time ∘ filter_early_term +end # module diff --git a/src/Rendering/render-cache.jl b/src/Rendering/render-cache.jl new file mode 100644 index 00000000..32ef96e8 --- /dev/null +++ b/src/Rendering/render-cache.jl @@ -0,0 +1,62 @@ +abstract type AbstractCacheStrategy end + +struct DefaultCache <: AbstractCacheStrategy end +struct SolutionCache <: AbstractCacheStrategy end +struct EndpointCache <: AbstractCacheStrategy end + +abstract type AbstractRenderCache{M,T} end + +struct SolutionRenderCache{M,T,G} <: AbstractRenderCache{M,T} + # metric + m::M + + # max time + max_time::T + + # size information + height::Int + width::Int + + # geodesics themselves in 2d array + geodesics::AbstractMatrix{G} + + function SolutionRenderCache( + m::AbstractMetricParams{T}, + max_time::T, + height, + width, + cache::AbstractVector{SciMLBase.EnsembleSolution{T,N,Vector{O}}}, + ) where {T,N,O} + + geodesics = Matrix{O}(undef, (height, width)) + + # populate store + for (col, simsol) in enumerate(cache) + for (row, sol) in enumerate(simsol) + geodesics[col, row] = sol + end + end + + # return instance + new{typeof(m),T,O}(m, max_time, height, width, geodesics) + end +end + +struct EndpointRenderCache{M,T,P} <: AbstractRenderCache{M,T} + # metric + m::M + + # max time + max_time::T + + # size information + height::Int + width::Int + + points::Matrix{P} +end + +function Base.show(io::IO, cache::AbstractRenderCache{M}) where {M} + repr = "$(Base.typename(typeof(cache)).name){$M} (dimensions $(cache.height)x$(cache.width))" + write(io, repr) +end diff --git a/src/Rendering/render.jl b/src/Rendering/render.jl new file mode 100644 index 00000000..44d99088 --- /dev/null +++ b/src/Rendering/render.jl @@ -0,0 +1,190 @@ +function __rendergeodesics( + m::AbstractMetricParams{T}, + init_pos; + image_width, + image_height, + kwargs..., +) where {T} + image = zeros(T, (image_height, image_width)) + render_into_image!( + m, + init_pos, + image; + image_width = image_width, + image_height = image_height, + kwargs..., + ) +end + +function __prerendergeodesics( + m::AbstractMetricParams{T}, + init_pos, + cache::AbstractRenderCache; + kwargs..., +) where {T} + error("Not implemented for render cache strategy '$typeof(cache)'.") +end + +function __prerendergeodesics( + m::AbstractMetricParams{T}, + init_pos, + cache::DefaultCache; + kwargs..., +) where {T} + # default is endpoints for now + __prerendergeodesics(m, init_pos, EndpointCache(); kwargs...) +end + +function __prerendergeodesics( + m::AbstractMetricParams{T}, + init_pos, + cache::SolutionCache; + kwargs..., +) where {T} + prerender_solutions(m, init_pos; kwargs...) +end + +function __prerendergeodesics( + m::AbstractMetricParams{T}, + init_pos, + cache::EndpointCache; + kwargs..., +) where {T} + prerender_endpoints(m, init_pos; kwargs...) +end + +#get_parameter_type(m, u, v, max_time) = +# typeof(GeodesicTracer.integrator_problem(m, u, v, (0.0, max_time)).p) + +function prerender_endpoints( + m::AbstractMetricParams{T}, + init_pos::V; + image_width, + image_height, + fov_factor, + max_time, + solver_opts..., +) where {T,V} + y_mid = image_height ÷ 2 + x_mid = image_width ÷ 2 + + #u_cache = Matrix{V}(undef, (image_height, image_width)) + #v_cache = Matrix{V}(undef, (image_height, image_width)) + #t_cache = Matrix{T}(undef, (image_height, image_width)) + #p_cache = Matrix{get_parameter_type(m, init_pos, init_pos, max_time)}(undef, (image_height, image_width)) + #retcode_cache = Matrix{Symbol}(undef, (image_height, image_width)) + + point_cache = + Matrix{GradusBase.geodesic_point_type(m)}(undef, (image_height, image_width)) + + vs = fill(init_pos, image_width) + us = fill(init_pos, image_width) + + for Y = 1:image_height + + β = T(y_to_β(Y, y_mid, fov_factor)) + α_generator_row = (T(x_to_α(X, x_mid, fov_factor)) for X = 1:image_width) + + calculate_velocities!(vs, m, init_pos, α_generator_row, β) + simsols = + tracegeodesics(m, us, vs, (T(0.0), max_time); save_on = false, solver_opts...) + println("+ $Y / $image_height ...") + + point_cache[Y, :] = get_endpoint(m, simsols) + end + + EndpointRenderCache(m, max_time, image_height, image_width, point_cache) +end + +function prerender_solutions( + m::AbstractMetricParams{T}, + init_pos; + image_width, + image_height, + fov_factor, + max_time, + solver_opts..., +) where {T} + y_mid = image_height ÷ 2 + x_mid = image_width ÷ 2 + + + # this hits the garbage collector like a truck + # need a better pre-cache for the velocity and position vectors + # maybe using a row by row renderer as before? + # or can we somehow repeat us without copying, as with FillArrays.jl? + # -- tried that, doesn't make enough of a difference + + vs = fill(init_pos, image_width) + us = fill(init_pos, image_width) + + simsol_array = map(1:image_height) do Y + + β = T(y_to_β(Y, y_mid, fov_factor)) + α_generator_row = (T(x_to_α(X, x_mid, fov_factor)) for X = 1:image_width) + + calculate_velocities!(vs, m, init_pos, α_generator_row, β) + simsols = tracegeodesics( + m, + us, + vs, + (T(0.0), max_time); + save_on = false, + solver_opts..., + ) + println("+ $Y / $image_height ...") + simsols + end + + SolutionRenderCache(m, max_time, image_height, image_width, simsol_array) +end + +function render_into_image!( + m::AbstractMetricParams{T}, + init_pos, + image; + image_width, + image_height, + fov_factor, + max_time, + pf, + verbose = true, + solver_opts..., +) where {T} + y_mid = image_height ÷ 2 + x_mid = image_width ÷ 2 + + vs = fill(init_pos, image_width) + us = fill(init_pos, image_width) + for Y = 1:image_height + + β = T(y_to_β(Y, y_mid, fov_factor)) + α_generator_row = (T(x_to_α(X, x_mid, fov_factor)) for X = 1:image_width) + + calculate_velocities!(vs, m, init_pos, α_generator_row, β) + + # do the render + simsols = tracegeodesics( + m, + us, + vs, + (T(0.0), max_time); + save_on = false, + verbose = verbose, + solver_opts..., + ) + apply_to_location!(m, @view(image[Y, :]), simsols, pf, max_time) + + if verbose + println("+ $Y / $image_height ...") + end + end + + image +end + +function apply_to_location!(m::AbstractMetricParams{T}, loc, sols, pf, max_time) where {T} + @threads for i = 1:length(sols) + loc[i] = pf(m, get_endpoint(m, sols[i]), max_time) + end +end diff --git a/src/Rendering/utility.jl b/src/Rendering/utility.jl new file mode 100644 index 00000000..eb5194a6 --- /dev/null +++ b/src/Rendering/utility.jl @@ -0,0 +1,36 @@ +function calculate_velocities( + m::AbstractMetricParams{T}, + init_pos, + α_generator, + β::T, +) where {T} + [map_impact_parameters(m, init_pos, α, β) for α in α_generator] +end + +function calculate_velocities( + m::AbstractMetricParams{T}, + init_pos, + α_genetator, + β_generator, +) where {T} + [map_impact_parameters(m, init_pos, α, β) for α in α_genetator, β in β_generator] +end + +""" +In-place specialisation +""" +function calculate_velocities!( + vs, + m::AbstractMetricParams{T}, + init_pos, + α_generator, + β::T, +) where {T} + for (i, α) in enumerate(α_generator) + vs[i] = map_impact_parameters(m, init_pos, α, β) + end +end + +# have to use a slight 0.001 offset to avoid integrating α=0.0 geodesics +x_to_α(X, x_mid, fov_factor) = (X + 0.001 - x_mid) / fov_factor +y_to_β(Y, y_mid, fov_factor) = (Y - y_mid) / fov_factor diff --git a/src/metrics/boyer-lindquist-ad.jl b/src/metrics/boyer-lindquist-ad.jl new file mode 100644 index 00000000..993d32f6 --- /dev/null +++ b/src/metrics/boyer-lindquist-ad.jl @@ -0,0 +1,55 @@ +module __BoyerLindquistAD +using ..StaticArrays + +@inline Σ(r, a, θ) = r^2 + a^2 * cos(θ)^2 +@inline Δ(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 +@fastmath function metric_components(M, a, rθ) + (r, θ) = rθ + R = 2M + Σ₀ = Σ(r, a, θ) + sinθ2 = sin(θ)^2 + + tt = -(1 - (R * r) / Σ₀) + rr = Σ₀ / Δ(r, R, a) + θθ = Σ₀ + ϕϕ = sinθ2 * (r^2 + a^2 + (sinθ2 * R * r * a^2) / Σ₀) + + tϕ = (-R * r * a * sinθ2) / Σ₀ + @SVector [tt, rr, θθ, ϕϕ, tϕ] +end + +end # module + +# new structure for our spacetime +@with_kw struct BoyerLindquistAD{T} <: AbstractAutoDiffStaticAxisSymmetricParams{T} + @deftype T + M = 1.0 + a = 0.0 +end + +# implementation +GeodesicTracer.metric_components(m::BoyerLindquistAD{T}, rθ) where {T} = + __BoyerLindquistAD.metric_components(m.M, m.a, rθ) +GradusBase.inner_radius(m::BoyerLindquistAD{T}) where {T} = m.M + √(m.M^2 - m.a^2) + +# additional utilities +function convert_angles(a, r, θ, ϕ, θ_obs, ϕ_obs) + ϕ̃ = ϕ - ϕ_obs + R = sqrt(r^2 + a^2) + o1 = r * R * sin(θ) * sin(θ_obs) * cos(ϕ̃) + R^2 * cos(θ) * cos(θ_obs) + o2 = R * cos(θ) * sin(θ_obs) * cos(ϕ̃) - r * sin(θ) * cos(θ_obs) + o3 = sin(θ_obs) * sin(ϕ̃) / sin(θ) + sigma = r^2 + a^2 * cos(θ)^2 + -o1 / sigma, -o2 / sigma, o3 / R +end + +# for disc profile models +function GradusBase.vector_to_local_sky(m::BoyerLindquistAD{T}, u, θ, ϕ) where {T} + convert_angles(m.a, u[2], u[3], u[4], θ, ϕ) +end + +# special radii +isco(m::BoyerLindquistAD{T}) where {T} = __BoyerLindquistFO.isco(m.M, m.a) diff --git a/src/metrics/boyer-lindquist-fo.jl b/src/metrics/boyer-lindquist-fo.jl new file mode 100644 index 00000000..22f16ab7 --- /dev/null +++ b/src/metrics/boyer-lindquist-fo.jl @@ -0,0 +1,369 @@ +module __BoyerLindquistFO +using ..FirstOrderMethods +using ..DocStringExtensions + +""" + $(TYPEDSIGNATURES) + +From Bardeen et al. (1972) eq. (2.10): + +```math +T = E \\left( r^2 + a^2 \\right) - L * a. +``` +""" +@inline T(E, L, r, a) = E * (r^2 + a^2) - L * a + + +""" + $(TYPEDSIGNATURES) + +From Bardeen et al. (1972) eq. (2.10), for the special case of a null-geodesic ``\\mu = 0``: + +```math +V_r = T^2 - \\Delta \\left[ (L - a E)^2 + Q \\right] +``` +""" +@inline Vr(E, L, M, Q, r, a) = T(E, L, r, a)^2 - Δ(M, r, a) * ((L - a * E)^2 + Q) + + +""" + $(TYPEDSIGNATURES) + +From Bardeen et al. (1972) eq. (2.10), for the special case of a null-geodesic ``\\mu = 0``: + +```math +V_\\theta = + Q + \\cos^2 (\\theta) \\left[ a^2 E^2 - \\frac{L^2}{\\sin^2 (\\theta) } \\right]. +``` +""" +@inline Vθ(E, L, Q, a, θ) = Q + cos(θ)^2 * ((a * E)^2 - (L / sin(θ))^2) + + +""" + $(TYPEDSIGNATURES) + +The ``t`` compontent of the equation of motion for a photon around a black hole, multiplied +by ``\\Sigma``. + +From Bardeen et al. (1972) eq. (2.9d): + +```math +\\Sigma \\frac{\\text{d}t}{\\text{d}\\lambda} = +- a \\left( a E \\sin^2 \\theta - L \\right) ++ \\frac{\\left( r^2 + a^2 \\right) T}{\\Delta}. +``` +""" +@inline function Σδt_δλ(E, L, M, r, a, θ) + -a * (a * E * sin(θ)^2 - L) + (r^2 + a^2) * T(E, L, r, a) / Δ(M, r, a) +end + + +""" + $(TYPEDSIGNATURES) + +The ``r`` compontent of the equation of motion for a photon around a black hole, multiplied +by ``\\Sigma``. + +Modified from Bardeen et al. (1972) eq. (2.9a): + +```math +\\Sigma \\frac{\\text{d}r}{\\text{d}\\lambda} = +\\pm \\sqrt{\\lvert V_r \\rvert}, +``` + +where, for implementation reason, the sign is always positive. Instead, the sign is applied +in [`δ`](@ref). +""" +@inline function Σδr_δλ(E, L, M, Q, r, a) + V = Vr(E, L, M, Q, r, a) + √abs(V) +end + + +""" + $(TYPEDSIGNATURES) + +The ``\\theta`` compontent of the equation of motion for a photon around a black hole, +multiplied by ``\\Sigma``. + +Modified from Bardeen et al. (1972) eq. (2.9b): + +```math +\\Sigma \\frac{\\text{d}\\theta}{\\text{d}\\lambda} = +\\pm \\sqrt{\\lvert V_\\theta \\rvert}, +``` + +where, for implementation reason, the sign is always positive. Instead, the sign is applied +in [`δ`](@ref). +""" +@inline function Σδθ_δλ(E, L, Q, a, θ) + V = Vθ(E, L, Q, a, θ) + √abs(V) +end + + +""" + $(TYPEDSIGNATURES) + +The ``\\phi`` compontent of the equation of motion for a photon around a black hole, +multiplied by ``\\Sigma``. + +From Bardeen et al. (1972) eq. (2.9c): + +```math +\\Sigma \\frac{\\text{d}\\phi}{\\text{d}\\lambda} = +- \\frac{L}{\\sin^2 \\theta} - aE ++ \\frac{aT}{\\Delta}. +``` +""" +@inline function Σδϕ_δλ(E, L, M, r, a, θ) + (L / sin(θ)^2) - (a * E) + a * T(E, L, r, a) / Δ(M, r, a) +end + + +""" +$(TYPEDSIGNATURES) + +From Fanton et al. (1997), eq. (76): + +```math +S = 1 + \\alpha \\omega \\sin \\theta. +``` +""" +S(θ, α, ω) = 1 + α * ω * sin(θ) + + +""" +$(TYPEDSIGNATURES) + +Calculates and returns the observer's angles ``\\sin \\Theta`` and ``\\sin \\Phi``, where +the parameters in the function signature correspond to the Bardeen et al. (1972) paper. + +From Fanton et al. (1997), eq. (74) and eq. (75): + +```math +\\begin{align*} +\\sin \\Theta &= +\\frac{\\alpha \\Sigma}{\\sqrt{A}} +\\left\\{ + \\beta^2 + \\left( \\alpha + a \\sin \\theta \\right)^2 + + \\frac{ + A S^2 - \\left( r^2 + a^2 + a \\alpha \\sin \\theta \\right) + }{\\Delta} +\\right\\}, \\\\ +\\sin \\Phi &= +-\\frac{\\alpha \\Sigma \\sqrt{\\Delta}}{A S \\sin\\Theta}. +\\end{align*} +``` +""" +function sinΦsinΨ(Σ₀, sinθ, A₀, Δ₀, S₀, r, a, α, β) + # calc 1 + sinΦ = + ((α * Σ₀) / √A₀) / + sqrt(β^2 + (α + a * sinθ)^2 + (A₀ * S₀^2 - (r^2 + a^2 + a * α * sinθ)^2) / Δ₀) + + # calc 2 + sinΨ = -(α * Σ₀ * √Δ₀) / (S₀ * A₀ * sinΦ) + + # return + (sinΦ, sinΨ) +end + +function sinΦsinΨ(M, r, a, θ, α, β) + # value cache + Σ₀ = Σ(r, a, θ) + sinθ = sin(θ) + A₀ = A(M, r, a, θ) + Δ₀ = Δ(M, r, a) + ω = 2.0 * a * r / A₀ + S₀ = S(θ, α, ω) + + sinΦsinΨ(Σ₀, sinθ, A₀, Δ₀, S₀, r, a, α, β) +end + + +""" +$(TYPEDSIGNATURES) + +Calculates conserved quantities +- angular momentum ``L`` +- Carter parameter ``Q`` + +for a photon described with position described by `x` in a Kerr spacetime given by +`p`. + +From Fanton et al. (1997), eq. (69): + +```math +L = \\frac{\\Upsilon_1}{\\Upsilon_2}, +``` + +where + +```math +\\begin{align*} +\\Upsilon_1 &= \\sin \\theta \\sin \\Phi \\sin \\Theta, \\\\ +\\Upsilon_2 &= \\frac{\\Sigma \\sqrt{\\Delta}}{A} + \\omega \\Upsilon_1, +\\omega = \\frac{2 a r}{A}, +\\end{align*} +``` + +taken from eq. (72) and eq. (73). + +From Fanton et al. (1997), eq. (70): + +```math +Q = +\\frac{P^2}{\\Delta} +- \\left( \\lambda - a \\right)^2 +- \\frac{\\Sigma^2}{A} \\left( \\frac{\\cos \\Phi}{\\Upsilon_2} \\right)^2, +``` + +and + +```math +P = \\left( r^2 + a^2 - a L \\right), +``` + +taken from eq. (71). +""" +function LQ(M, r, a, θ, sinΦ, sinΨ) + sinθ = sin(θ) + Σ₀ = Σ(r, a, θ) + A₀ = A(M, r, a, θ) + Δ₀ = Δ(M, r, a) + ω = 2.0 * a * r / A₀ + + Υ₁ = sinθ * sinΦ * sinΨ + Υ₂ = (Σ₀ * √Δ₀ / A₀) + ω * Υ₁ + + Υ₁ = sinθ * sinΦ * sinΨ + Υ₂ = (Σ₀ * √Δ₀ / A₀) + ω * Υ₁ + + L = Υ₁ / Υ₂ + P = (r^2 + a^2 - (a * L)) + Q = (P^2 / Δ₀) - (L - a)^2 - ((Σ₀^2 / A₀) * (cos(asin(sinΨ)) / Υ₂)^2) + L, Q +end + + +""" +$(TYPEDSIGNATURES) + +From Bardeen et al. (1972) eq. (2.3): + +```math +\\Sigma = r^2 + a^2 \\cos^2( \\theta ). +``` +""" +Σ(r, a, θ) = r^2 + (a * cos(θ))^2 + + +""" +$(TYPEDSIGNATURES) + +From Bardeen et al. (1972) eq. (2.3): + +```math +\\Delta = r^2 - 2 M r + a^2. +``` +""" +Δ(M, r, a) = r^2 - 2 * M * r + a^2 + + +""" +$(TYPEDSIGNATURES) + +From Bardeen et al. (1972) eq. (2.3): + +```math +A = (r^2 + a^2)^2 - a^2 \\Delta \\sin^2 ( \\theta ). +``` +""" +A(M, r, a, θ) = (r^2 + a^2)^2 - a^2 * Δ(M, r, a) * sin(θ)^2 + + +""" +$(TYPEDSIGNATURES) + +From Bardeen et al. (1972) eq. (2.21): + +```math +Z_1 = +1 + \\sqrt[3]{1 - \\frac{a^2}{M^2}} +\\left[ + \\sqrt[3]{1 + \\frac{a}{M}} + \\sqrt[3]{1 - \\frac{a}{M}} +\\right]. +``` +""" +Z₁(M, a) = 1 + ∛(1 - (a / M)^2) * (∛(1 + (a / M)) + ∛(1 - (a / M))) + + +""" +$(TYPEDSIGNATURES) + +From Bardeen et al. (1972) eq. (2.21): + +```math +Z_2 = \\sqrt{\\frac{3a^2}{M^2} + Z_1^2}. +``` +""" +Z₂(M, a) = √(3(a / M)^2 + Z₁(M, a)^2) + + +function four_velocity(u, E, M, a, p) + let r = u[2], θ = u[3], L = p.L, Q = p.Q + Σ₀ = Σ(r, a, θ) + ( + Σδt_δλ(E, L, M, r, a, θ) / Σ₀, + p.r * Σδr_δλ(E, L, M, Q, r, a) / Σ₀, + p.θ * Σδθ_δλ(E, L, Q, a, θ) / Σ₀, + Σδϕ_δλ(E, L, M, r, a, θ) / Σ₀, + ) + end +end + +""" +From Bardeen et al. (1972) eq. (2.21): + +```math +r_\\text{ms} = M \\left\\{ 3 + Z_2 \\pm \\sqrt{(3 - Z_1)(3 + Z_1 + 2 Z_2)} \\right\\}. +``` + +The choice of ``\\pm`` is chosen by the sign of ``a``. +""" +isco(M, a, ±) = M * (3 + Z₂(M, a) ± √((3 - Z₁(M, a)) * (3 + Z₁(M, a) + 2 * Z₂(M, a)))) +isco(M, a) = a > 0.0 ? isco(M, a, -) : isco(M, a, +) +end # module + +@with_kw struct BoyerLindquistFO{T} <: AbstractFirstOrderMetricParams{T} + @deftype T + M = 1.0 + a = 0.0 + E = 1.0 +end + +GradusBase.inner_radius(m::BoyerLindquistFO{T}) where {T} = m.M + √(m.M^2 - m.a^2) +GradusBase.constrain(::BoyerLindquistFO{T}, u, v; μ = 0.0) where {T} = v[1] + +FirstOrderMethods.four_velocity(u, m::BoyerLindquistFO{T}, p) where {T} = + __BoyerLindquistFO.four_velocity(u, m.E, m.M, m.a, p) +FirstOrderMethods.calc_lq(m::BoyerLindquistFO{T}, pos, vel) where {T} = + __BoyerLindquistFO.LQ(m.M, pos[2], m.a, pos[3], vel[3], vel[4]) + +function FirstOrderMethods.Vr(m::BoyerLindquistFO{T}, u, p) where {T} + L, Q, _, _ = p + __BoyerLindquistFO.Vr(m.E, L, m.M, Q, u[2], m.a) +end +function FirstOrderMethods.Vθ(m::BoyerLindquistFO{T}, u, p) where {T} + L, Q, _, _ = p + __BoyerLindquistFO.Vθ(m.E, L, Q, m.a, u[3]) +end + +function GeodesicTracer.alpha_beta_to_vel(m::BoyerLindquistFO{T}, u, α, β) where {T} + sinΦ, sinΨ = __BoyerLindquistFO.sinΦsinΨ(m.M, u[2], m.a, u[3], α, β) + (β < 0.0 ? -1.0 : 1.0, sinΦ, sinΨ) +end + +# special radii +isco(m::BoyerLindquistFO{T}) where {T} = __BoyerLindquistFO.isco(m.M, m.a) diff --git a/src/metrics/johannsen-ad.jl b/src/metrics/johannsen-ad.jl new file mode 100644 index 00000000..f5a91303 --- /dev/null +++ b/src/metrics/johannsen-ad.jl @@ -0,0 +1,50 @@ +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 + +@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 function metric_components(m, rθ) + (r, θ) = rθ + + A1 = A₁(m.M, r, m.α13) + A2 = A₂(m.M, r, m.α22) + A5 = A₅(m.M, r, m.α52) + Σ₀ = Σ(r, m.a, θ, m.M, m.ϵ3) + Δ₀ = Δ(r, m.M, m.a) + + # cache common components + r2a2 = r^2 + m.a^2 + sin_theta2 = sin(θ)^2 + + denom = ((r2a2) * A1 - m.a^2 * A2 * sin_theta2)^2 + + tt = -Σ₀ * (Δ₀ - m.a^2 * A2^2 * sin_theta2) + rr = Σ₀ / (Δ₀ * A5) + θθ = Σ₀ + ϕϕ = Σ₀ * sin_theta2 * ((r2a2)^2 * A1^2 - m.a^2 * Δ₀ * sin_theta2) + + tϕ = -m.a * Σ₀ * sin_theta2 * ((r2a2) * A1 * A2 - Δ₀) + @SVector [tt / denom, rr, θθ, ϕϕ / denom, tϕ / denom] +end + +end # module + +@with_kw struct JohannsenAD{T} <: AbstractAutoDiffStaticAxisSymmetricParams{T} + @deftype T + M = 1.0 + a = 0.0 + α13 = 0.0 + α22 = 0.0 + α52 = 0.0 + ϵ3 = 0.0 +end + +GeodesicTracer.metric_components(m::JohannsenAD{T}, rθ) where {T} = + __JohannsenAD.metric_components(m, rθ) +GradusBase.inner_radius(m::JohannsenAD{T}) where {T} = m.M + √(m.M^2 - m.a^2) diff --git a/src/metrics/metrics.jl b/src/metrics/metrics.jl new file mode 100644 index 00000000..5db91625 --- /dev/null +++ b/src/metrics/metrics.jl @@ -0,0 +1,4 @@ +include("boyer-lindquist-ad.jl") +include("boyer-lindquist-fo.jl") +include("johannsen-ad.jl") +include("morris-thorne-ad.jl") diff --git a/src/metrics/morris-thorne-ad.jl b/src/metrics/morris-thorne-ad.jl new file mode 100644 index 00000000..fa92391e --- /dev/null +++ b/src/metrics/morris-thorne-ad.jl @@ -0,0 +1,26 @@ +module __MorrisThorneAD +using ..StaticArrays +@fastmath function metric_components(b, lθ) + (l, θ) = lθ + + tt = -1 + rr = 1 + θθ = b^2 + l^2 + ϕϕ = (b^2 + l^2) * sin(θ) + + # no time-azimuth couple + @SVector [tt, rr, θθ, ϕϕ, 0.0] +end + +end # module + +# new structure for our spacetime +@with_kw struct MorrisThorneAD{T} <: AbstractAutoDiffStaticAxisSymmetricParams{T} + @deftype T + b = 1.0 +end + +# implementation +GeodesicTracer.metric_components(m::MorrisThorneAD{T}, rθ) where {T} = + __MorrisThorneAD.metric_components(m.b, rθ) +GradusBase.inner_radius(m::MorrisThorneAD{T}) where {T} = 0.0 diff --git a/src/special-radii.jl b/src/special-radii.jl new file mode 100644 index 00000000..c3e8eb07 --- /dev/null +++ b/src/special-radii.jl @@ -0,0 +1,6 @@ +""" +$(TYPEDSIGNATURES) + +Innermost stable circular orbit. +""" +isco(m::AbstractMetricParams{T}) where {T} = error("Not implemented for $(typeof(m)).") diff --git a/test/Project.toml b/test/Project.toml new file mode 100644 index 00000000..bd76ff16 --- /dev/null +++ b/test/Project.toml @@ -0,0 +1,3 @@ +[deps] +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl new file mode 100644 index 00000000..7a2a78de --- /dev/null +++ b/test/runtests.jl @@ -0,0 +1,3 @@ +include("smoke-tests/rendergeodesics.jl") +include("smoke-tests/tracegeodesics.jl") +include("smoke-tests/pointfunctions.jl") \ No newline at end of file diff --git a/test/smoke-tests/pointfunctions.jl b/test/smoke-tests/pointfunctions.jl new file mode 100644 index 00000000..fe4faaee --- /dev/null +++ b/test/smoke-tests/pointfunctions.jl @@ -0,0 +1,30 @@ +# Tests to make sure the basic pointfunctions work +using Test, Gradus, StaticArrays + +@testset "pointfunctions" begin + + function run_pointfunction(m, pf) + u = @SVector [0.0, 100.0, deg2rad(85), 0.0] + d = GeometricThinDisc(10.0, 40.0, deg2rad(90.0)) + img = rendergeodesics( + m, + u, + d, + 200.0, + fov_factor = 1.0, + image_width = 100, + image_height = 50, + pf = pf, + verbose = false, + ) + end + + @testset "redshift" begin + # only implemented for the BoyerLindquist metrics at the moment + for m in [BoyerLindquistAD(), BoyerLindquistFO()] + run_pointfunction(m, redshift) + end + # smoke test passed + @test true + end +end diff --git a/test/smoke-tests/rendergeodesics.jl b/test/smoke-tests/rendergeodesics.jl new file mode 100644 index 00000000..d23562ac --- /dev/null +++ b/test/smoke-tests/rendergeodesics.jl @@ -0,0 +1,54 @@ +# Tests to make sure the basic `rendergeodesics` function works for (ideally) all metrics. +using Test, Gradus, StaticArrays + +@testset "rendergeodesics" begin + + @testset "plain" begin + u = @SVector [0.0, 100.0, deg2rad(85), 0.0] + for (m, expectation) in zip( + [BoyerLindquistAD(), JohannsenAD(), BoyerLindquistFO(), MorrisThorneAD()], + # expectation values for the sum of the image + # last computed 20/04/2022: v0.1.0 + [8949.345629557964, 8949.36445058427, 8961.703066256105, 306.88361044534764], + ) + img = rendergeodesics( + m, + u, + 200.0, + fov_factor = 1.0, + image_width = 100, + image_height = 50, + verbose = false, + ) + image_fingerprint = sum(filter(!isnan, img)) + # have to be really coarse cus the first order method is so variable??? + # the rest are very resolute + @test isapprox(expectation, image_fingerprint; atol = 2.0, rtol = 0.0) + end + end + + @testset "thin-disc" begin + u = @SVector [0.0, 100.0, deg2rad(85), 0.0] + d = GeometricThinDisc(10.0, 40.0, deg2rad(90.0)) + for (m, expectation) in zip( + [BoyerLindquistAD(), JohannsenAD(), BoyerLindquistFO(), MorrisThorneAD()], + # expectation values for the sum of the image + # last computed 20/04/2022: v0.1.0 + [86795.10968236077, 86796.73210488849, 81959.05505753662, 36803.48713530963], + ) + img = rendergeodesics( + m, + u, + d, + 200.0, + fov_factor = 1.0, + image_width = 100, + image_height = 50, + verbose = false, + ) + image_fingerprint = sum(filter(!isnan, img)) + @test isapprox(expectation, image_fingerprint; atol = 2.0, rtol = 0.0) + end + end + +end diff --git a/test/smoke-tests/tracegeodesics.jl b/test/smoke-tests/tracegeodesics.jl new file mode 100644 index 00000000..56a46a7d --- /dev/null +++ b/test/smoke-tests/tracegeodesics.jl @@ -0,0 +1,65 @@ +# Tests to make sure the basic `tracegeodesics` function works for (ideally) all metrics. +using Test, Gradus, StaticArrays + +@testset "tracegeodesics" begin + # tests if a single geodesic can be integrated + function test_single(m, u, v) + sol = tracegeodesics(m, u, v, (0.0, 200.0)) + @test sol.retcode == :Terminated + sol + end + + # tests if a multiple geodesics can be integrated + function test_many(m, u, v) + us = [u, u, u] + vs = [v, v, v] + simsols = tracegeodesics(m, us, vs, (0.0, 200.0)) + @test all(i -> simsols[i].retcode == :Terminated, 1:3) + simsols + end + + @testset "static-arrays" begin + u = @SVector [0.0, 100.0, deg2rad(85), 0.0] + # arbitrary single velocity vector + v = @SVector [0.0, -1.0, -0.0, -3.5e-6] + for m in [BoyerLindquistAD(), JohannsenAD(), BoyerLindquistFO(), MorrisThorneAD()] + test_single(m, u, v) + test_many(m, u, v) + end + end + + @testset "regular-arrays" begin + u = Float64[0.0, 100.0, deg2rad(85), 0.0] + # arbitrary single velocity vector + v = Float64[0.0, -1.0, -0.0, -3.5e-6] + for m in [BoyerLindquistAD(), JohannsenAD(), BoyerLindquistFO(), MorrisThorneAD()] + test_single(m, u, v) + test_many(m, u, v) + end + end + + @testset "corona-models" begin + # only implemented for BoyerLindquistAD at the moment + # because of the vector to local sky methods + m = BoyerLindquistAD(M = 1.0, a = 0.0) + d = GeometricThinDisc(Gradus.isco(m), 50.0, deg2rad(90.0)) + model = LampPostModel(h = 10.0, θ = deg2rad(0.001)) + + for Sampler in [EvenSampler, WeierstrassSampler], + Generator in [GoldenSpiralGenerator, RandomGenerator], + Hemisphere in [LowerHemisphere, BothHemispheres] + + simsols = tracegeodesics( + m, + model, + (0.0, 200.0), + n_samples = 32, + d; + save_on = false, + sampler = Sampler(generator = Generator(), domain = Hemisphere()), + ) + # smoke test passed + @test true + end + end +end