Merge pull request #144 from astro-group-bristol/fergus/thick-tf-inte…

…gration

Thick disc transfer function integration

src/corona/flux-calculations.jl

@@ -28,12 +28,7 @@ position `x` via
 ```
 """
 lorentz_factor(m::AbstractMetric, ::AbstractAccretionGeometry, x; kwargs...) =
-    lorentz_factor(
-        m,
-        x,
-        CircularOrbits.fourvelocity(m, _equatorial_project(x[2]));
-        kwargs...,
-    )
+    lorentz_factor(m, x, CircularOrbits.fourvelocity(m, _equatorial_project(x)); kwargs...)
 function lorentz_factor(m::AbstractMetric, x, v; component = 4)
     𝒱ϕ = local_velocity(m, x, v, component)
     inv(√(1 - 𝒱ϕ^2))

src/line-profiles.jl

@@ -61,11 +61,10 @@ end
 function lineprofile(
     bins,
     ε::Function,
-    m::AbstractMetric,
+    m::AbstractMetric{T},
     u,
     d::AbstractAccretionGeometry,
     ::BinnedLineProfile;
-    plane = PolarPlane(GeometricGrid(); Nr = 450, Nθ = 1300, r_max = 50.0),
     λ_max = 2 * u[2],
     redshift_pf = ConstPointFunctions.redshift(m, u),
     verbose = false,
@@ -75,7 +74,7 @@ function lineprofile(
     # todo: make it friendly
     plane = PolarPlane(GeometricGrid(); Nr = 450, Nθ = 1300, r_max = 5maxrₑ),
     solver_args...,
-)
+) where {T}
     progress_bar = init_progress_bar("Lineprofile: ", trajectory_count(plane), verbose)
 
     gps = tracegeodesics(
@@ -87,6 +86,7 @@ function lineprofile(
         save_on = false,
         verbose = verbose,
         progress_bar = progress_bar,
+        callback = domain_upper_hemisphere(),
         ensemble = EnsembleEndpointThreads(),
         solver_args...,
     )
@@ -101,8 +101,8 @@ function lineprofile(
 
     f = @. ε(r) * g^3 * areas
     # bin
-    F = bucket(Simple(), g, f, bins)
-    bins, F ./ sum(F)
+    flux = bucket(Simple(), g, f, bins)
+    bins, flux ./ sum(flux)
 end
 
 export AbstractLineProfileAlgorithm, BinnedLineProfile, CunninghamLineProfile, lineprofile

src/redshift.jl

@@ -191,12 +191,12 @@ end
     end
 end
 
-@inline function redshift_function(m::KerrMetric, gp)
+@inline function redshift_function(m::KerrMetric{T}, gp) where {T}
     isco = Gradus.isco(m)
     # metric matrix at observer
     g_obs = metric(m, gp.x_init)
     # fixed stationary observer velocity
-    v_obs = @SVector [1.0, 0.0, 0.0, 0.0]
+    v_obs = SVector{4,T}(1, 0, 0, 0)
 
     r = Gradus._equatorial_project(gp.x)
     v_disc = if r < isco
@@ -239,12 +239,12 @@ end
 For a full, annotated derivation of this method, see
 [the following blog post](https://fjebaker.github.io/blog/pages/2022-05-plunging-orbits/).
 """
-function interpolate_redshift(plunging_interpolation, u)
+function interpolate_redshift(plunging_interpolation, u::SVector{4,T}) where {T}
     isco = Gradus.isco(plunging_interpolation.m)
     # metric matrix at observer
     g_obs = metric(plunging_interpolation.m, u)
     # fixed stationary observer velocity
-    v_obs = @SVector [1.0, 0.0, 0.0, 0.0]
+    v_obs = SVector{4,T}(1, 0, 0, 0)
     closure =
         (m, gp, max_time) -> begin
             let r = _equatorial_project(gp.x)
@@ -253,7 +253,7 @@ function interpolate_redshift(plunging_interpolation, u)
                     vtemp = plunging_interpolation(r)
                     # we have to reverse radial velocity due to backwards tracing convention
                     # see https://github.com/astro-group-bristol/Gradus.jl/issues/3
-                    @SVector [vtemp[1], -vtemp[2], vtemp[3], vtemp[4]]
+                    SVector{4}(vtemp[1], -vtemp[2], vtemp[3], vtemp[4])
                 else
                     # regular circular orbit
                     CircularOrbits.fourvelocity(plunging_interpolation.m, r)
@@ -269,7 +269,7 @@ function interpolate_redshift(plunging_interpolation, u)
     PointFunction(closure)
 end
 
-interpolate_redshift(m::AbstractMetric, u; kwargs...) =
+interpolate_redshift(m::AbstractMetric, u::SVector{4}; kwargs...) =
     interpolate_redshift(interpolate_plunging_velocities(m; kwargs...), u)
 
 export RedshiftFunctions, interpolate_redshift

src/tracing/precision-solvers.jl

@@ -78,7 +78,7 @@ function find_offset_for_radius(
 )
     function _measure(gp::GeodesicPoint{T}) where {T}
         r = if gp.status == StatusCodes.IntersectedWithGeometry
-            gp.x[2] * sin(gp.x[3])
+            _equatorial_project(gp.x)
         else
             zero(T)
         end
@@ -174,7 +174,6 @@ function impact_parameters_for_radius(
     α = zeros(T, N)
     β = zeros(T, N)
     impact_parameters_for_radius!(α, β, m, u, d, radius; kwargs...)
-    # (filter(!isnan, α), filter(!isnan, β))
     α, β
 end
 
@@ -200,7 +199,7 @@ function jacobian_∂αβ_∂gr(
         gp = unpack_solution(sol)
         g = redshift_pf(m, gp, max_time)
         # return disc r and redshift g
-        SVector(gp.x[2], g)
+        SVector(_equatorial_project(gp.x), g)
     end
 
     j = ForwardDiff.jacobian(_jacobian_f, SVector(α, β))
@@ -275,13 +274,14 @@ end
 function optimize_for_target(
     target::SVector,
     m::AbstractMetric,
-    x0,
+    x0::SVector{4,T},
     args...;
     optimizer = NelderMead(),
+    p0 = zeros(T, 2),
     kwargs...,
-)
+) where {T}
     f, solver = _make_target_objective(target, m, x0, args...; kwargs...)
-    res = optimize(f, [0.0, 0.0], optimizer)
+    res = optimize(f, p0, optimizer)
     out = Optim.minimizer(res)
     # return α, β, accuracy
     α, β = out[1], out[2]

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

@@ -206,7 +206,7 @@ function _rear_workhorse(
             tracer_kwargs...,
         )
         gp = unpack_solution(sol)
-        is_visible = abs(rₑ - gp.x[2] * sin(gp.x[3])) < visible_tolerance
+        is_visible = abs(rₑ - _equatorial_project(gp.x)) < visible_tolerance
         (g, J, t, is_visible)
     end
 end

src/transfer-functions/integration.jl

@@ -57,22 +57,24 @@ function _upper_branch(c::_IntegrationClosures)
     function _upper_branch_integrand(g)
         g✶ = g_to_g✶(g, c.branch.gmin, c.branch.gmax)
         f = c.branch.upper_f(g✶)
-        integrand(c.setup, f, c.branch.rₑ, g, g✶)
+        integrand(c.setup, _zero_if_nan(f), c.branch.rₑ, g, g✶)
     end
 end
 
 function _lower_branch(c::_IntegrationClosures)
     function _lower_branch_integrand(g)
         g✶ = g_to_g✶(g, c.branch.gmin, c.branch.gmax)
         f = c.branch.lower_f(g✶)
-        integrand(c.setup, f, c.branch.rₑ, g, g✶)
+        integrand(c.setup, _zero_if_nan(f), c.branch.rₑ, g, g✶)
     end
 end
 
 function _both_branches(c::_IntegrationClosures)
     function _both_branches_integrand(g)
         g✶ = g_to_g✶(g, c.branch.gmin, c.branch.gmax)
-        f = c.branch.upper_f(g✶) + c.branch.lower_f(g✶)
+        fu = c.branch.upper_f(g✶)
+        fl = c.branch.lower_f(g✶)
+        f = _zero_if_nan(fu) + _zero_if_nan(fl)
         integrand(c.setup, f, c.branch.rₑ, g, g✶)
     end
 end

src/transfer-functions/utils.jl

@@ -43,8 +43,11 @@ function insert_data!(data::_TransferDataAccumulator, i, θ, vals::NTuple{3})
     data
 end
 function insert_data!(data::_TransferDataAccumulator, i, θ, vals::NTuple{4})
-    data.mask[i] = vals[4]
+    data.mask[i] = true
     data.data[:, i] .= (θ, vals[1:3]...)
+    if vals[4] == 0
+        data.Js[i] = NaN
+    end
     data
 end
 

test/smoke-tests/disc-profiles.jl

@@ -40,15 +40,15 @@ using StaticArrays
         # values computed under visual inspection
         # last computed 21/06/23: updated sampling method
         expected_areas = [
-            75.68214227687605,
-            114.7677610122124,
-            166.49995048864642,
-            263.6075724481634,
-            700.3226456903473,
-            739.4256113746189,
-            645.1547948341391,
-            478.14652288719276,
-            358.69318748667644,
+            75.67457406264847,
+            114.75628423611069,
+            166.48330049359774,
+            263.58121169091913,
+            700.3117449256706,
+            739.4180350022357,
+            645.1747233625393,
+            478.1614271222751,
+            358.7139818627764,
         ]
 
         @test areas1 ≈ expected_areas atol = 1e-3

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 ≈ 14.149627898685342 atol = 1e-4
+@test total ≈ 13.419539069804333 atol = 1e-4

test/unit/emissivity.jl

@@ -33,14 +33,14 @@ profile = emissivity_profile(
 r, em = get_emissivity(profile)
 
 @test em ≈ [
-    1.4536471470344825,
-    3.2159234810191997,
-    1.9624518676973357,
-    0.8654137966466704,
-    0.23973494196571282,
-    0.046232848740113686,
-    0.00789393186371828,
-    0.0012974087513837867,
-    0.0001872022563608107,
-    3.825853506250807e-5,
+    1.4531056469477812,
+    3.216028944655333,
+    1.9625919375522862,
+    0.8654885033740061,
+    0.2397572736938699,
+    0.04623730854934404,
+    0.007894706545002786,
+    0.0012975372018286873,
+    0.0001872208766790682,
+    3.826234997699e-5,
 ] atol = 1e-5