Merge pull request #165 from astro-group-bristol/fergus/visibility-cr…

…iteria

Fix: try to improve the visibility criteria

src/corona/samplers.jl

@@ -63,6 +63,14 @@ function _cart_to_spher_jacobian(θ, ϕ)
     ]
 end
 
+function _spher_to_cart_jacobian(θ, ϕ, r)
+    @SMatrix [
+        sin(θ)*cos(ϕ) r*cos(θ)*cos(ϕ) -r*sin(θ)*sin(ϕ)
+        sin(θ)*sin(ϕ) r*cos(θ)*sin(ϕ) r*sin(θ)*cos(ϕ)
+        cos(θ) -r*sin(θ) 0
+    ]
+end
+
 function _cart_local_direction(θ, ϕ)
     @SVector [sin(θ) * cos(ϕ), sin(θ) * sin(ϕ), cos(θ)]
 end

src/geometry/discs.jl

@@ -45,7 +45,7 @@ end
 cross_section(d::AbstractThickAccretionDisc, x4) =
     error("Not implemented for $(typeof(d)).")
 r_cross_section(d::AbstractThickAccretionDisc, r::Number) =
-    cross_section(d, SVector(0.0, r, π / 2, 0.0))
+    cross_section(d, SVector(0, r, π / 2, 0))
 
 struct EllipticalDisc{T} <: AbstractAccretionDisc{T}
     inner_radius::T

src/geometry/discs/thick-disc.jl

@@ -61,5 +61,27 @@ function distance_to_disc(d::AbstractThickAccretionDisc, x4; gtol)
     abs(z) - height - (gtol * x4[2])
 end
 
+function _cartesian_tangent_vector(ρ, d::AbstractThickAccretionDisc)
+    function _parameterization(r)
+        xz_parameterize(d, r)
+    end
+    ∇f = ForwardDiff.derivative(_parameterization, ρ)
+    v = SVector(∇f[1], 0, ∇f[2])
+    v ./ √(v[1]^2 + v[2]^2 + v[3]^2)
+end
+
+function _cartesian_surface_normal(ρ, d::AbstractThickAccretionDisc)
+    ∇f = _cartesian_tangent_vector(ρ, d)
+    # rotate 90 degrees in about ϕ̂
+    SVector(-∇f[3], ∇f[2], ∇f[1])
+end
+
+function _rotate_cartesian_about_z(n, ϕ)
+    SVector(n[1] * cos(ϕ), n[1] * sin(ϕ), n[3])
+end
+
+function _cartesian_surface_normal(ρ, ϕ, d::AbstractThickAccretionDisc)
+    _rotate_cartesian_about_z(_cartesian_surface_normal(ρ, d), ϕ)
+end
 
 export ThickDisc

src/plotting-recipes.jl

@@ -1,7 +1,14 @@
 using RecipesBase
 
-function _extract_path(sol, n_points; projection = :none, t_span = 100.0)
-    mid_i = max(1, length(sol.u) ÷ 2)
+function _extract_path(sol, n_points; projection = :none, t_span = 30.0)
+    status = Gradus.get_status_code(sol.prob.p)
+    mid_i =
+        if (status == StatusCodes.IntersectedWithGeometry) ||
+           (status == StatusCodes.WithinInnerBoundary)
+            lastindex(sol.t)
+        else
+            max(1, length(sol.u) ÷ 2)
+        end
 
     start_t = max(sol.t[mid_i] - t_span, sol.t[1])
     end_t = min(sol.t[mid_i] + t_span, sol.t[end])
@@ -31,13 +38,13 @@ end
     zlims --> _range
     aspect_ratio --> 1
 
-    itr = if !(typeof(sol) <: SciMLBase.EnsembleSolution)
-        (; u = (sol,))
+    itr = if typeof(sol) <: SciMLBase.EnsembleSolution
+        sol.u
     else
         sol
     end
 
-    for s in itr.u
+    for s in itr
         path = _extract_path(s, n_points, projection = :none, t_span = t_span)
         @series begin
             path

src/solution-processing.jl

@@ -114,6 +114,7 @@ function unpack_solution_full(
             ui = SVector{4,T}(us[i][1:4])
             vi = SVector{4,T}(us[i][5:8])
             ti = ts[i]
+            aux = unpack_auxiliary(us[i])
             GeodesicPoint(
                 get_status_code(sol.prob.p),
                 t_start,
@@ -122,6 +123,7 @@ function unpack_solution_full(
                 ui,
                 v_start,
                 vi,
+                aux,
             )
         end
     end

src/tracing/precision-solvers.jl

@@ -182,7 +182,6 @@ function impact_parameters_for_radius_obscured(
     x::AbstractVector{T},
     d::AbstractAccretionDisc,
     radius;
-    visible_tolerance = 1e-3,
     N = 500,
     β₀ = zero(T),
     α₀ = zero(T),
@@ -192,26 +191,43 @@ function impact_parameters_for_radius_obscured(
     β = zeros(T, N)
     impact_parameters_for_radius!(α, β, m, x, d, radius; β₀ = β₀, α₀ = α₀, kwargs...)
 
-    I = map(
-        i ->
-            !_trace_is_visible(
-                m,
-                x,
-                α[i],
-                β[i],
-                d,
-                radius;
-                visible_tolerance = visible_tolerance,
-                kwargs...,
-            ),
-        eachindex(α),
+    gps = tracegeodesics(
+        m,
+        x,
+        i -> map_impact_parameters(m, x, α[i], β[i]),
+        datumplane(d, radius),
+        2x[2];
+        ensemble = EnsembleEndpointThreads(),
+        save_on = false,
+        trajectories = length(α),
     )
-
+    n = _cartesian_surface_normal(radius, d)
+    I = [!_is_visible(m, d, gp, n) for gp in gps]
     α[I] .= NaN
     β[I] .= NaN
     α, β
 end
 
+"""
+If dot product between surface normal and the final velocity vector is positive, 
+the point is visible.
+"""
+function _is_visible(m::AbstractMetric, d, gp::AbstractGeodesicPoint, n::SVector{3})
+    gp_new =
+        tracegeodesics(m, gp.x_init, gp.v_init, d, gp.λ_max; save_on = false) |>
+        unpack_solution
+    # geodesics sufficiently close, test the angle
+    v = SVector(gp_new.v[2], gp_new.v[3], gp_new.v[4]) ./ gp_new.v[1]
+    v_geodesic = _spher_to_cart_jacobian(gp_new.x[3], gp_new.x[4], gp_new.x[2]) * v
+    n_rot = _rotate_cartesian_about_z(n, gp_new.x[4])
+
+    X1 = to_cartesian(gp.x)
+    X2 = to_cartesian(gp_new.x)
+    dist = sum(i -> i^2, X1 .- X2)
+
+    return _fast_dot(v_geodesic, n_rot) < 0 && isapprox(dist, 0, atol = 1e-10)
+end
+
 function jacobian_∂αβ_∂gr(
     m,
     x,

src/transfer-functions/cunningham-transfer-functions.jl

@@ -147,11 +147,12 @@ function _rear_workhorse_with_impact_parameters(
     d::AbstractAccretionDisc,
     rₑ;
     max_time = 2 * x[2],
-    redshift_pf = ConstPointFunctions.redshift(m, x),
     offset_max = 0.4rₑ + 10,
     zero_atol = 1e-7,
     β₀ = 0,
     α₀ = 0,
+    redshift_pf = ConstPointFunctions.redshift(m, x),
+    jacobian_disc = d,
     tracer_kwargs...,
 )
     function _workhorse(θ)
@@ -180,22 +181,22 @@ function _rear_workhorse_with_impact_parameters(
         J = jacobian_∂αβ_∂gr(
             m,
             x,
-            d,
+            jacobian_disc,
             α,
             β,
             max_time;
             redshift_pf = redshift_pf,
             tracer_kwargs...,
         )
-        (g, J, gp.x[1], α, β)
+        (g, J, gp, α, β)
     end
     (_workhorse, tracer_kwargs)
 end
 function _rear_workhorse(m::AbstractMetric, x, d::AbstractAccretionDisc, rₑ; kwargs...)
     workhorse, _ = _rear_workhorse_with_impact_parameters(m, x, d, rₑ; kwargs...)
     function _disc_workhorse(θ::T)::NTuple{3,T} where {T}
-        g, J, t, _, _ = workhorse(θ)
-        g, J, t
+        g, J, gp, _, _ = workhorse(θ)
+        g, J, gp.x[1]
     end
 end
 
@@ -205,7 +206,6 @@ function _rear_workhorse(
     d::AbstractThickAccretionDisc,
     rₑ;
     max_time = 2 * x[2],
-    visible_tolerance = 1e-3,
     kwargs...,
 )
     plane = datumplane(d, rₑ)
@@ -215,49 +215,17 @@ function _rear_workhorse(
         plane,
         rₑ;
         max_time = max_time,
+        jacobian_disc = d,
         kwargs...,
     )
+    n = _cartesian_surface_normal(rₑ, d)
     function _thick_workhorse(θ::T)::NTuple{4,T} where {T}
-        g, J, t, α, β = datum_workhorse(θ)
-        # check if the location is visible or not
-        is_visible = _trace_is_visible(
-            m,
-            x,
-            α,
-            β,
-            d,
-            rₑ;
-            visible_tolerance = visible_tolerance,
-            tracer_kwargs...,
-        )
-        (g, J, t, is_visible)
+        g, J, gp, _, _ = datum_workhorse(θ)
+        is_visible = _is_visible(m, d, gp, n)
+        (g, J, gp.x[1], is_visible)
     end
 end
 
-function _trace_is_visible(
-    m,
-    x,
-    α,
-    β,
-    d,
-    rₑ;
-    max_time = 2x[2],
-    visible_tolerance = 1e-3,
-    kwargs...,
-)
-    sol = tracegeodesics(
-        m,
-        x,
-        map_impact_parameters(m, x, α, β),
-        d,
-        max_time;
-        save_on = false,
-        kwargs...,
-    )
-    gp = unpack_solution(sol)
-    abs(rₑ - _equatorial_project(gp.x)) < visible_tolerance
-end
-
 function _cunningham_transfer_function!(
     data::_TransferDataAccumulator,
     workhorse,

test/transfer-functions/test-thick-disc.jl

@@ -8,4 +8,4 @@ d = ShakuraSunyaev(m)
 tf = cunningham_transfer_function(m, x, d, 3.0; β₀ = 1.0)
 
 total = sum(filter(!isnan, tf.f))
-@test total ≈ 13.419539069804333 atol = 1e-4
+@test total ≈ 12.148257594740231 atol = 1e-4