diff --git a/lib/GradusSpectralModels/src/GradusSpectralModels.jl b/lib/GradusSpectralModels/src/GradusSpectralModels.jl index 96f11489..28daf21b 100644 --- a/lib/GradusSpectralModels/src/GradusSpectralModels.jl +++ b/lib/GradusSpectralModels/src/GradusSpectralModels.jl @@ -28,7 +28,7 @@ function LineProfile( E₀ = FitParam(1.0), kwargs..., ) - setup = integration_setup(profile, table((get_value(θ), get_value(a))); kwargs...) + setup = integration_setup(profile, table((get_value(a), get_value(θ))); kwargs...) LineProfile((; setup = setup, table = table), K, a, θ, rin, rout, E₀) end diff --git a/src/corona/profiles/radial.jl b/src/corona/profiles/radial.jl index ecb635e0..9e903e1c 100644 --- a/src/corona/profiles/radial.jl +++ b/src/corona/profiles/radial.jl @@ -1,7 +1,7 @@ -struct RadialDiscProfile{T,I} <: AbstractDiscProfile - radii::Vector{T} - ε::Vector{T} - t::Vector{T} +struct RadialDiscProfile{V<:AbstractVector,I} <: AbstractDiscProfile + radii::V + ε::V + t::V interp_ε::I interp_t::I end diff --git a/src/interpolations.jl b/src/interpolations.jl index c4420339..179e2a2a 100644 --- a/src/interpolations.jl +++ b/src/interpolations.jl @@ -94,7 +94,7 @@ function _tuple_set(tuple::NTuple{N}, index, v)::NTuple{N} where {N} elseif index == N (tuple[1:end-1]..., v) else - (tuple[1:index-1], v, tuple[index+1:end]) + (tuple[1:index-1]..., v, tuple[index+1:end]...) end end diff --git a/src/reverberation.jl b/src/reverberation.jl index 8d90ee63..cd67d0ec 100644 --- a/src/reverberation.jl +++ b/src/reverberation.jl @@ -57,7 +57,7 @@ function lag_frequency( kwargs..., ) where {T} other_kwargs, em_setup = EmissivityProfileSetup(T, spectrum; kwargs...) - solver_kwargs, tf_setup = _TransferFunctionSetup(T; other_kwargs...) + solver_kwargs, tf_setup = _TransferFunctionSetup(m, d; other_kwargs...) prof = emissivity_profile(em_setup, m, d, model; solver_kwargs...) t0 = continuum_time(m, x, model; solver_kwargs...) diff --git a/src/tracing/precision-solvers.jl b/src/tracing/precision-solvers.jl index f9d8ea6a..5ddfda29 100644 --- a/src/tracing/precision-solvers.jl +++ b/src/tracing/precision-solvers.jl @@ -16,6 +16,7 @@ function _find_offset_for_measure( d::AbstractAccretionGeometry, θₒ; zero_atol = 1e-7, + root_solver = DEFAULT_ROOT_SOLVER(), offset_max = 20.0, initial_r = offset_max / 2, max_time = 2 * x[2], @@ -48,6 +49,7 @@ function _find_offset_for_measure( best = eltype(x)[0.0, 1.0] r0_candidate, resid = root_solve( + root_solver, _offset_objective, initial_r, (measure, _velfunc, _solve_geodesic, best); diff --git a/src/transfer-functions/cunningham-transfer-functions.jl b/src/transfer-functions/cunningham-transfer-functions.jl index 1264fb0a..e993d1b8 100644 --- a/src/transfer-functions/cunningham-transfer-functions.jl +++ b/src/transfer-functions/cunningham-transfer-functions.jl @@ -4,7 +4,7 @@ function _promote_disc_for_transfer_functions(::ThinDisc{T}) where {T} plane, plane end -struct _TransferFunctionSetup{T} +struct _TransferFunctionSetup{T,A} h::T θ_offset::T "Tolerance for root finding" @@ -15,10 +15,12 @@ struct _TransferFunctionSetup{T} β₀::T N::Int N_extrema::Int + root_solver::A end function _TransferFunctionSetup( - m::AbstractMetric{T}; + m::AbstractMetric{T}, + d::AbstractAccretionGeometry; θ_offset = T(0.6), zero_atol = T(1e-7), N = 80, @@ -26,9 +28,20 @@ function _TransferFunctionSetup( α₀ = 0, β₀ = 0, h = T(1e-6), + root_solver = nothing, kwargs..., ) where {T} - setup = _TransferFunctionSetup{T}( + # specialize the algorithm depending on whether we are calculating for a thin or thick disc + _alg = if isnothing(root_solver) + if d isa AbstractThickAccretionDisc + RootsAlg() + else + NonLinearAlg() + end + else + root_solver + end + setup = _TransferFunctionSetup{T,typeof(_alg)}( h, θ_offset, zero_atol, @@ -37,6 +50,7 @@ function _TransferFunctionSetup( convert(T, β₀), N, N_extrema, + _alg, ) kwargs, setup end @@ -196,6 +210,7 @@ function _setup_workhorse_jacobian_with_kwargs( zero_atol = setup.zero_atol, offset_max = offset_max, max_time = max_time, + root_solver = setup.root_solver, β₀ = setup.β₀, α₀ = setup.α₀, tracer_kwargs..., @@ -260,22 +275,31 @@ function _rear_workhorse( ) function _thick_workhorse(θ::T)::NTuple{4,T} where {T} g, gp, r = datum_workhorse(θ) - r₊, _ = _find_offset_for_radius( - m, - x, - d, - rₑ, - θ; - initial_r = r, - zero_atol = setup.zero_atol, - offset_max = offset_max, - max_time = max_time, - β₀ = setup.β₀, - α₀ = setup.α₀, - tracer_kwargs..., - # don't echo warnings - warn = false, - ) + r₊ = try + r_thick, _ = _find_offset_for_radius( + m, + x, + d, + rₑ, + θ; + initial_r = r, + zero_atol = setup.zero_atol, + root_solver = setup.root_solver, + offset_max = offset_max, + max_time = max_time, + β₀ = setup.β₀, + α₀ = setup.α₀, + tracer_kwargs..., + # don't echo warnings + warn = false, + ) + r_thick + catch + # if we fail, for whatever reason, to root solve on the thick discs, + # we don't care, we just need a NaN value and then set that point to + # "not visible" + NaN + end is_visible, J = if !isnan(r₊) && isapprox(r, r₊, atol = 1e-3) # trace jacobian on updated impact parameters α, β = _rθ_to_αβ(r₊, θ; α₀ = setup.α₀, β₀ = setup.β₀) @@ -290,11 +314,12 @@ end function _cunningham_transfer_function!( data::_TransferDataAccumulator, + setup::_TransferFunctionSetup, workhorse, θiterator, - θ_offset, rₑ, ) + θ_offset = setup.θ_offset for (i, θ) in enumerate(θiterator) θ_corrected = θ + 1e-4 insert_data!(data, i, θ_corrected, workhorse(θ)) @@ -328,7 +353,7 @@ function cunningham_transfer_function( rₑ::T; kwargs..., ) where {Q,T} - solver_kwargs, setup = _TransferFunctionSetup(m; kwargs...) + solver_kwargs, setup = _TransferFunctionSetup(m, d; kwargs...) cunningham_transfer_function(setup, m, x, d, rₑ; solver_kwargs...) end @@ -362,8 +387,7 @@ function cunningham_transfer_function( chart = chart, solver_kwargs..., ) - gmin, gmax = - _cunningham_transfer_function!(data, workhorse, θiterator, setup.θ_offset, rₑ) + gmin, gmax = _cunningham_transfer_function!(data, setup, workhorse, θiterator, rₑ) CunninghamTransferData( data.data[2, :], data.data[3, :], @@ -424,7 +448,7 @@ function interpolated_transfer_branches( radii; kwargs..., ) - solver_kwargs, setup = _TransferFunctionSetup(m; kwargs...) + solver_kwargs, setup = _TransferFunctionSetup(m, d; kwargs...) interpolated_transfer_branches(setup, m, x, d, radii; solver_kwargs...) end diff --git a/src/transfer-functions/integration.jl b/src/transfer-functions/integration.jl index 11f3402a..743f4505 100644 --- a/src/transfer-functions/integration.jl +++ b/src/transfer-functions/integration.jl @@ -312,6 +312,7 @@ function _integrate_transfer_problem!( transfer_function_radial_interpolation, r_limits, g_grid; + pure_radial = setup.pure_radial, g_scale = 1, ) where {T} g_grid_view = @views g_grid[1:end-1] @@ -327,7 +328,7 @@ function _integrate_transfer_problem!( Δrₑ = rₑ - r_prev # integration weight for this annulus - θ = Δrₑ * rₑ * setup.pure_radial(rₑ) * π / (branch.gmax - branch.gmin) + θ = Δrₑ * rₑ * pure_radial(rₑ) * π / (branch.gmax - branch.gmin) @inbounds for j in eachindex(g_grid_view) glo = g_grid[j] / g_scale @@ -355,6 +356,7 @@ function _integrate_transfer_problem!( r_limits, g_grid, t_grid; + pure_radial = setup.pure_radial, g_scale = 1, ) where {T} g_grid_view = @views g_grid[1:end-1] @@ -371,7 +373,7 @@ function _integrate_transfer_problem!( Δrₑ = rₑ - r_prev # integration weight for this annulus - θ = Δrₑ * rₑ * setup.pure_radial(rₑ) * π / (branch.gmax - branch.gmin) + θ = Δrₑ * rₑ * pure_radial(rₑ) * π / (branch.gmax - branch.gmin) # time delay for this annuli t_source_disc = setup.time(rₑ) diff --git a/src/utils.jl b/src/utils.jl index b1677827..a309c490 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -1,16 +1,49 @@ +abstract type AbstractRootAlgorithm end +struct RootsAlg <: AbstractRootAlgorithm end +Base.@kwdef struct NonLinearAlg{A} <: AbstractRootAlgorithm + alg::A = SimpleNonlinearSolve.SimpleBroyden() +end + +DEFAULT_ROOT_SOLVER() = NonLinearAlg() + """ root_solve(f_objective, initial_value, args) Wrapper to different root solving backends to make root solve fast and efficient """ function root_solve( + f_objective, + initial_value::T, + args; + kwargs..., +) where {T<:Union{<:Number,<:SVector{1}}} + root_solve(DEFAULT_ROOT_SOLVER(), f_objective, initial_value, args; kwargs...) +end +function root_solve( + ::RootsAlg, + f_objective, + initial_value::T, + args; + abstol = 1e-9, + kwargs..., +) where {T<:Union{<:Number,<:SVector{1}}} + x0 = Roots.find_zero( + r -> f_objective(r, args), + initial_value, + Roots.Order0(); + atol = abstol, + ) + resid = f_objective(x0, args) + x0, resid +end +function root_solve( + alg::NonLinearAlg, f_objective, initial_value::T, args; abstol = 1e-9, kwargs..., ) where {T<:Union{<:Number,<:SVector{1}}} - # Roots.find_zero(r -> f_objective(r, args), initial_value, Roots.Order0(); atol = abstol) x0, f = if T <: Number function _obj_wrapper(x::SVector, p) @inbounds SVector{1,eltype(x)}(f_objective(x[1], p)) @@ -20,14 +53,7 @@ function root_solve( initial_value, f_objective end prob = SimpleNonlinearSolve.NonlinearProblem{false}(f, x0, args) - sol = solve( - prob, - SimpleNonlinearSolve.SimpleBroyden(); - abstol = abstol, - reltol = abstol, - maxiters = 500, - kwargs..., - ) + sol = solve(prob, alg.alg, abstol = abstol, reltol = abstol, maxiters = 500, kwargs...) sol.u[1], sol.resid[1] end diff --git a/test/transfer-functions/test-thick-disc.jl b/test/transfer-functions/test-thick-disc.jl index 29f2ee93..b8637761 100644 --- a/test/transfer-functions/test-thick-disc.jl +++ b/test/transfer-functions/test-thick-disc.jl @@ -8,4 +8,17 @@ d = ShakuraSunyaev(m) tf = cunningham_transfer_function(m, x, d, 3.0; β₀ = 1.0) total = sum(filter(!isnan, tf.f)) -@test total ≈ 12.245276038643347 atol = 1e-4 +@test total ≈ 12.253422180875667 atol = 1e-4 + +m = KerrMetric(1.0, 0.2) +x = SVector(0.0, 10_000, deg2rad(20), 0.0) +d = ShakuraSunyaev(m; eddington_ratio = 0.2) + +tf = cunningham_transfer_function(m, x, d, 5.469668466100368; β₀ = 1.0) +total = sum(filter(!isnan, tf.f)) +@test total ≈ 20.83469 atol = 1e-4 + +# the transfer function here is pretty horrible as it's almost impossible to actually +# see; this is a test to make sure it doesn't error +# an offset to the isco of 4-e2 resolves this, but that's quite a lot +tf = cunningham_transfer_function(m, x, d, Gradus.isco(m) + 1e-2; β₀ = 1.0) diff --git a/test/unit/interpolations.jl b/test/unit/interpolations.jl index cdaeb36f..895a809d 100644 --- a/test/unit/interpolations.jl +++ b/test/unit/interpolations.jl @@ -32,13 +32,13 @@ ff(x, y) = 3x^2 + x * y - sin(y) ff(x) = SVector(ff(x[1], x[2])) X1 = collect(range(0, 1, 1000)) X2 = collect(range(0, 1, 1000)) -vals = reshape([ff(x, y) for y in X1, x in X2], (length(X1), length(X2))) +vals = reshape([ff(x, y) for x in X1, y in X2], (length(X1), length(X2))) cache = MultilinearInterpolator{2}(vals) # check that dual cache works too function _interpolate_wrapper(cache, X1, X2, vals) function _f(x) - x2, x1 = x + x1, x2 = x SVector(interpolate!(cache, (X1, X2), vals, (x1, x2))) end end @@ -89,3 +89,22 @@ intp = Gradus.interpolate!(cache, (X1, X2), vals, (0.0, 1.5)) intp = Gradus.interpolate!(cache, (X1, X2), vals, (0.0, 1.5)) @test intp.a == [1.5] @test intp.b == [1.5 0; 0 1.5] + +# try higher dimensional interpolation + +X1 = range(0.0, 1.0, 3) +X2 = range(1.0, 2.0, 4) +X3 = range(-1.0, 0.0, 2) + +f3(x, y, z) = 2x + 3y + 7z +vals = [f3(x, y, z) for x in X1, y in X2, z in X3] + +cache = MultilinearInterpolator{3}(vals) +intp = Gradus.interpolate!(cache, (X1, X2, X3), vals, (0.5, 1.5, -0.5)) +@test intp == f3(0.5, 1.5, -0.5) + +intp = Gradus.interpolate!(cache, (X1, X2, X3), vals, (1.0, 1.1, -0.1)) +@test intp == f3(1.0, 1.1, -0.1) + +intp = Gradus.interpolate!(cache, (X1, X2, X3), vals, (1.0, 1.9, -0.1)) +@test intp == f3(1.0, 1.9, -0.1)