Merge pull request #146 from astro-group-bristol/fergus/equatorial-pr…

…oject

Equatorial projection function

src/const-point-functions.jl

@@ -64,7 +64,7 @@ end
 
 Returns a [`PointFunction`](@ref).
 
-Calculate the analytic redshift at a given geodesic point, assuming equitorial, geometrically
+Calculate the analytic redshift at a given geodesic point, assuming equatorial, geometrically
 thin accretion disc. Implementation depends on the metric type. Currently implemented for
 
 - [`Gradus.KerrMetric`](@ref)

src/corona/emissivity.jl

@@ -25,10 +25,10 @@ function source_to_disc_emissivity(m::AbstractStaticAxisSymmetric, N, A, x, g; 
 end
 
 """
-    point_source_equitorial_disc_emissivity(θ, g, A, γ; Γ = 2)
+    point_source_equatorial_disc_emissivity(θ, g, A, γ; Γ = 2)
 
 Calculate the emissivity of a point illuminating source on the spin axis for an annulus of the
-equitorial accretion disc with (proper) area `A`. The precise formulation follows from Dauser et al. (2013),
+equatorial accretion disc with (proper) area `A`. The precise formulation follows from Dauser et al. (2013),
 with the emissivity calculated as
 ```math
 \\varepsilon = \\frac{\\sin \\theta}{A g^\\Gamma \\gamma}
@@ -43,7 +43,7 @@ in order to maintain point density. It may be regarded as the PDF that samples `
 
 Dauser et al. (2013)
 """
-point_source_equitorial_disc_emissivity(θ, g, A, γ; Γ = 2) = sin(θ) / (g^Γ * A * γ)
+point_source_equatorial_disc_emissivity(θ, g, A, γ; Γ = 2) = sin(θ) / (g^Γ * A * γ)
 
 """
     function emissivity_profile(
@@ -185,7 +185,7 @@ function _point_source_symmetric_emissivity_profile(
     I = [i.status == StatusCodes.IntersectedWithGeometry for i in gps]
     points = gps[I]
     δs = δs[I]
-    J = sortperm(points, by = i -> i.x[2])
+    J = sortperm(points, by = i -> _equatorial_project(i.x))
     points = points[J]
     δs = δs[J]
 
@@ -212,7 +212,7 @@ function _point_source_emissivity(
     Δr = diff(r)
     @. A = A * Δr
     r = r[1:end-1]
-    ε = point_source_equitorial_disc_emissivity.(@views(δs[1:end-1]), gs, A, γ)
+    ε = point_source_equatorial_disc_emissivity.(@views(δs[1:end-1]), gs, A, γ)
     r, ε
 end
 

src/corona/flux-calculations.jl

@@ -28,7 +28,12 @@ position `x` via
 ```
 """
 lorentz_factor(m::AbstractMetric, ::AbstractAccretionGeometry, x; kwargs...) =
-    lorentz_factor(m, x, CircularOrbits.fourvelocity(m, x[2]); kwargs...)
+    lorentz_factor(
+        m,
+        x,
+        CircularOrbits.fourvelocity(m, _equatorial_project(x[2]));
+        kwargs...,
+    )
 function lorentz_factor(m::AbstractMetric, x, v; component = 4)
     𝒱ϕ = local_velocity(m, x, v, component)
     inv(√(1 - 𝒱ϕ^2))
@@ -76,7 +81,7 @@ function flux_source_to_disc(
         @tullio E_s := -g_1[i, j] * gp.v_init[i] * v_source[j]
 
         # energy at disc
-        v_disc = disc_velocity(gp.x[2])
+        v_disc = disc_velocity(_equatorial_project(gp.x))
         @tullio E_d := -g_2[i, j] * gp.v[i] * v_disc[j]
 
         # relative redshift source to disc

src/corona/profiles/radial.jl

@@ -13,7 +13,7 @@ function RadialDiscProfile(
     intensities = nothing;
     kwargs...,
 )
-    radii = map(i -> i.x[2], gps)
+    radii = map(i -> _equatorial_project(i.x), gps)
     # ensure sorted: let the user sort so that everything is sure to be
     # in order
     if !issorted(radii)
@@ -33,7 +33,10 @@ function RadialDiscProfile(rs::AbstractArray, ts::AbstractArray, εs::AbstractAr
     t = DataInterpolations.LinearInterpolation(ts, rs)
     ε = DataInterpolations.LinearInterpolation(εs, rs)
     # wrap geodesic point wrappers
-    RadialDiscProfile(gp -> ε(gp.x[2]), gp -> t(gp.x[2]) + gp.x[1])
+    RadialDiscProfile(
+        gp -> ε(_equatorial_project(gp.x)),
+        gp -> t(_equatorial_project(gp.x)) + gp.x[1],
+    )
 end
 
 function RadialDiscProfile(rdp::RadialDiscProfile)
@@ -97,7 +100,7 @@ function RadialDiscProfile(
 end
 
 function RadialDiscProfile(ce::CoronaGeodesics; kwargs...)
-    J = sortperm(ce.geodesic_points; by = i -> i.x[2])
+    J = sortperm(ce.geodesic_points; by = i -> _equatorial_project(i.x))
     @views RadialDiscProfile(
         ce.metric,
         ce.model,
@@ -108,7 +111,7 @@ function RadialDiscProfile(ce::CoronaGeodesics; kwargs...)
 end
 
 function RadialDiscProfile(ce::CoronaGeodesics{<:TraceRadiativeTransfer}; kwargs...)
-    J = sortperm(ce.geodesic_points; by = i -> i.x[2])
+    J = sortperm(ce.geodesic_points; by = i -> _equatorial_project(i.x))
     @views RadialDiscProfile(
         ce.metric,
         ce.model,
@@ -120,18 +123,18 @@ function RadialDiscProfile(ce::CoronaGeodesics{<:TraceRadiativeTransfer}; kwargs
 end
 
 function RadialDiscProfile(ε, ce::CoronaGeodesics; kwargs...)
-    J = sortperm(ce.geodesic_points; by = i -> i.x[2])
-    radii = @views [i.x[2] for i in ce.geodesic_points[J]]
+    J = sortperm(ce.geodesic_points; by = i -> _equatorial_project(i.x))
+    radii = @views [_equatorial_project(i.x) for i in ce.geodesic_points[J]]
     times = @views [i.x[1] for i in ce.geodesic_points[J]]
 
     t = DataInterpolations.LinearInterpolation(times, radii)
 
     function _emissivity_wrapper(gp)
-        ε(gp.x[2])
+        ε(_equatorial_project(gp.x))
     end
 
     function _delay_wrapper(gp)
-        t(gp.x[2]) + gp.x[1]
+        t(_equatorial_project(gp.x)) + gp.x[1]
     end
 
     RadialDiscProfile(_emissivity_wrapper, _delay_wrapper; kwargs...)

src/corona/profiles/voronoi.jl

@@ -12,7 +12,7 @@ struct VoronoiDiscProfile{D,V,G} <: AbstractDiscProfile
     ) where {D<:AbstractAccretionDisc,V<:AbstractArray,G}
         if !isapprox(d.inclination, π / 2)
             return error(
-                "Currently only supported for discs in the equitorial plane (θ=π/2).",
+                "Currently only supported for discs in the equatorial plane (θ=π/2).",
             )
         end
         new{D,V,G}(d, polys, gen, gps)
@@ -119,7 +119,7 @@ getareas(vdp::VoronoiDiscProfile) = getarea.(vdp.polys)
 function getproperarea(poly::AbstractArray, m::AbstractMetric)
     A = getarea(poly)
     c = getbarycenter(poly)
-    # get value of metric at the radius of the polygon's barycenter, and in the equitorial plane
+    # get value of metric at the radius of the polygon's barycenter, and in the equatorial plane
     m_params = metric_components(m, (sqrt(c[1]^2 + c[2]^2), π / 2))
     # need radial and azimuthal components of the metric
     sqrt(m_params[2] * m_params[4]) * A

src/geometry/discs.jl

@@ -32,7 +32,7 @@ struct TopHatDisc{T} <: AbstractThickAccretionDisc{T}
 end
 
 function Gradus.cross_section(d::TopHatDisc, u)
-    # project u into equitorial plane
+    # project u into equatorial plane
     r = u[2] * sin(u[3])
     if (r < d.inner_r) || (r > d.outer_r)
         return -1.0

src/geometry/geometry.jl

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

src/line-profiles.jl

@@ -49,11 +49,12 @@ function lineprofile(
     numrₑ = 100,
     verbose = false,
     h = 2e-8,
+    Nr = 1000,
     kwargs...,
 )
     radii = Grids._inverse_grid(minrₑ, maxrₑ, numrₑ)
     itfs = interpolated_transfer_branches(m, u, d, radii; verbose = verbose, kwargs...)
-    flux = integrate_lineprofile(ε, itfs, radii, bins; h = h)
+    flux = integrate_lineprofile(ε, itfs, radii, bins; h = h, Nr = Nr)
     bins, flux
 end
 
@@ -69,7 +70,10 @@ function lineprofile(
     redshift_pf = ConstPointFunctions.redshift(m, u),
     verbose = false,
     minrₑ = isco(m),
-    maxrₑ = 50,
+    maxrₑ = T(50),
+    # this 5maxrₑ is a heuristic that is insulting
+    # todo: make it friendly
+    plane = PolarPlane(GeometricGrid(); Nr = 450, Nθ = 1300, r_max = 5maxrₑ),
     solver_args...,
 )
     progress_bar = init_progress_bar("Lineprofile: ", trajectory_count(plane), verbose)
@@ -87,13 +91,13 @@ function lineprofile(
         solver_args...,
     )
 
-    I = intersected_with_geometry(gps, x -> (minrₑ <= x[2] <= maxrₑ))
+    I = intersected_with_geometry(gps, x -> (minrₑ <= _equatorial_project(x) <= maxrₑ))
     points = @views gps[I]
     areas = unnormalized_areas(plane)[I]
 
     # calculate physical flux
     g = redshift_pf.(m, points, λ_max)
-    r = [p.x[2] for p in points]
+    r = map(p -> _equatorial_project(p.x), points)
 
     f = @. ε(r) * g^3 * areas
     # bin

src/orbits/orbit-solving.jl

@@ -1,5 +1,5 @@
 function trace_single_orbit(m, r, vϕ; max_time = 300.0, μ = 1.0, θ₀ = π / 2, tracer_args...)
-    # fixed in equitorial plane
+    # fixed in equatorial plane
     u = @SVector [0.0, r, θ₀, 0.0]
     v = @SVector [0.0, 0.0, 0.0, vϕ]
     tracegeodesics(m, u, v, (0.0, max_time); μ = μ, tracer_args...)
@@ -12,7 +12,7 @@ function measure_stability(m::AbstractMetric, r, vϕ; tracer_args...)
     sum(((rs .- r) ./ r) .^ 2) / length(rs)
 end
 
-function __solve_equitorial_circular_orbit(
+function __solve_equatorial_circular_orbit(
     m::AbstractMetric,
     r,
     optimizer,
@@ -29,15 +29,15 @@ function __solve_equitorial_circular_orbit(
     Optim.minimizer(res)
 end
 
-function solve_equitorial_circular_orbit(
+function solve_equatorial_circular_orbit(
     m::AbstractMetric,
     r::Number;
     lower_bound = 0.0,
     upper_bound = 1.0,
     optimizer = GoldenSection(),
     tracer_args...,
 )
-    __solve_equitorial_circular_orbit(
+    __solve_equatorial_circular_orbit(
         m,
         r,
         optimizer,
@@ -56,7 +56,7 @@ function sliding_window(func, N, lower_bound, upper_bound, lower_rate, upper_rat
     end
 end
 
-function solve_equitorial_circular_orbit(
+function solve_equatorial_circular_orbit(
     m::AbstractMetric,
     r_range::Union{<:AbstractRange,<:AbstractArray};
     lower_bound = 0.0,
@@ -74,7 +74,7 @@ function solve_equitorial_circular_orbit(
         upper_rate,
     ) do (i, lower_bound, upper_bound)
         r = r_range_reverse[i]
-        solve_equitorial_circular_orbit(
+        solve_equatorial_circular_orbit(
             m,
             r,
             ;
@@ -86,13 +86,13 @@ function solve_equitorial_circular_orbit(
     reverse!(candidate_vϕ)
 end
 
-function trace_equitorial_circular_orbit(m::AbstractMetric, rs; kwargs...)
-    map(zip(rs, solve_equitorial_circular_orbit(m, rs; kwargs...))) do (r, vϕ)
+function trace_equatorial_circular_orbit(m::AbstractMetric, rs; kwargs...)
+    map(zip(rs, solve_equatorial_circular_orbit(m, rs; kwargs...))) do (r, vϕ)
         trace_single_orbit(m, r, vϕ; kwargs...)
     end
 end
-function trace_equitorial_circular_orbit(m::AbstractMetric, r::Number; kwargs...)
-    vϕ = solve_equitorial_circular_orbit(m, r; kwargs...)
+function trace_equatorial_circular_orbit(m::AbstractMetric, r::Number; kwargs...)
+    vϕ = solve_equatorial_circular_orbit(m, r; kwargs...)
     trace_single_orbit(m, r, vϕ; kwargs...)
 end
 
@@ -166,5 +166,5 @@ function interpolate_plunging_velocities(
     PlungingInterpolation(m, sol)
 end
 
-export solve_equitorial_circular_orbit,
-    trace_equitorial_circular_orbit, PlungingInterpolation, interpolate_plunging_velocities
+export solve_equatorial_circular_orbit,
+    trace_equatorial_circular_orbit, PlungingInterpolation, interpolate_plunging_velocities

src/precompile.jl

@@ -250,7 +250,7 @@ Base.precompile(
     },
 )   # time: 0.007825662
 Base.precompile(
-    Tuple{typeof(point_source_equitorial_disc_emissivity),Float64,Any,Float64,Float64},
+    Tuple{typeof(point_source_equatorial_disc_emissivity),Float64,Any,Float64,Float64},
 )   # time: 0.006162999
 Base.precompile(
     Tuple{

src/redshift.jl

@@ -198,7 +198,7 @@ end
     # fixed stationary observer velocity
     v_obs = @SVector [1.0, 0.0, 0.0, 0.0]
 
-    r = gp.x[2] * sin(gp.x[3])
+    r = Gradus._equatorial_project(gp.x)
     v_disc = if r < isco
         # plunging region
         SVector(uᵗ(m.M, isco, r, m.a), -uʳ(m.M, isco, r), 0, uᶲ(m.M, isco, r, m.a))
@@ -247,7 +247,7 @@ function interpolate_redshift(plunging_interpolation, u)
     v_obs = @SVector [1.0, 0.0, 0.0, 0.0]
     closure =
         (m, gp, max_time) -> begin
-            let r = gp.x[2]
+            let r = _equatorial_project(gp.x)
                 v_disc = if r < isco
                     # plunging region
                     vtemp = plunging_interpolation(r)

src/tracing/radiative-transfer-problem.jl

@@ -11,7 +11,7 @@ function radiative_transfer(m::AbstractMetric, x, k, geometry, I, ν₀, r_isco,
 end
 
 function fluid_velocity(m::AbstractMetric, ::AbstractAccretionGeometry, x, r_isco, λ)
-    r = x[2] * sin(x[3])
+    r = _equatorial_project(x)
     if r > r_isco
         CircularOrbits.fourvelocity(m, r)
     else

src/utils.jl

@@ -1,4 +1,37 @@
+struct NaNLinearInterpolator{X,Y}
+    x::Vector{X}
+    y::Vector{Y}
+    default::Y
+end
+
+function _interpolate(interp::NaNLinearInterpolator{X}, x::X) where {X}
+    idx = clamp(searchsortedlast(interp.x, x), 1, length(interp.x) - 1)
+    x1, x2 = interp.x[idx], interp.x[idx+1]
+    y1, y2 = interp.y[idx], interp.y[idx+1]
+
+    w = (x - x1) / (x2 - x1)
+    y = (1 - w) * y1 + w * y2
+
+    if isnan(y)
+        if (w < 0.5)
+            isnan(y1) && return interp.default
+            return y1
+        else
+            isnan(y2) && return interp.default
+            return y2
+        end
+    end
+
+    y
+end
+
+function (interp::NaNLinearInterpolator)(x)
+    _interpolate(interp, x)
+end
+
 function _make_interpolation(x, y)
+    @assert issorted(x) "x must be sorted!"
+    # NaNLinearInterpolator(x, y, zero(eltype(y)))
     DataInterpolations.LinearInterpolation(y, x)
 end
 
@@ -80,4 +113,8 @@ end
 Lz(m::AbstractMetric, u, v) = Lz(metric(m, u), v)
 Lz(m::AbstractMetric, gp::AbstractGeodesicPoint) = Lz(m, gp.x, gp.v)
 
+_equatorial_project(x::SVector{4}) = x[2] * sin(abs(x[3]))
+
+_zero_if_nan(x::T) where {T} = isnan(x) ? zero(T) : x
+
 export cartesian_squared_distance, cartesian_distance, spherical_to_cartesian