Merge pull request #217 from astro-group-bristol/fergus/adaptive

Brand new root-finder for solving emission radii

Author: fjebaker

lib/GradusSpectralModels/src/GradusSpectralModels.jl

@@ -8,13 +8,13 @@ struct LineProfile{D,T} <: AbstractTableModel{T,Additive}
     "Spin"
     a::T
     "Observer inclination (degrees off of the spin axis)."
-    θ::T
+    θ_obs::T
     "Inner radius of the accretion disc."
-    rin::T
+    inner_r::T
     "Outer radius of the accretion disc."
-    rout::T
+    outer_r::T
     "Central emission line energy (keV)."
-    E₀::T
+    lineE::T
 end
 
 function Base.copy(m::LineProfile)
@@ -39,37 +39,29 @@ function LineProfile(
     profile,
     table::Gradus.CunninghamTransferTable;
     K = FitParam(1.0),
-    a = FitParam(0.998),
-    θ = FitParam(45.0),
-    rin = FitParam(1.0),
-    rout = FitParam(100.0, upper_limit = 100.0),
-    E₀ = FitParam(1.0),
+    a = FitParam(0.998, upper_limit = 1.0),
+    θ_obs = FitParam(45.0, upper_limit = 90.0),
+    inner_r = FitParam(1.0),
+    outer_r = FitParam(100.0, upper_limit = 1000.0),
+    lineE = FitParam(6.4),
     kwargs...,
 )
-    setup = integration_setup(profile, table((get_value(a), get_value(θ))); kwargs...)
-    LineProfile((; setup = setup, table = table), K, a, θ, rin, rout, E₀)
+    setup = integration_setup(profile, table((get_value(a), get_value(θ_obs))); kwargs...)
+    LineProfile((; setup = setup, table = table), K, a, θ_obs, inner_r, outer_r, lineE)
 end
 
 function SpectralFitting.invoke!(output, domain, model::LineProfile)
-    grid = model.table.table((model.a, model.θ))
-    rmin = if model.rin < grid.r_grid[1]
-        grid.r_grid[1]
-    else
-        model.rin
-    end
-    rout = if model.rout < rmin
-        rmin
-    else
-        model.rout
-    end
+    grid = model.table.table((model.a, model.θ_obs))
+    rmin = (model.inner_r < grid.r_grid[1]) ? grid.r_grid[1] : model.inner_r
+    outer_r = (model.outer_r < rmin) ? rmin : model.outer_r
     Gradus.integrate_lineprofile!(
         output,
         model.table.setup,
         grid,
         domain;
         rmin = rmin,
-        rmax = rout,
-        g_scale = model.E₀,
+        rmax = outer_r,
+        g_scale = model.lineE,
     )
     output
 end

src/corona/adaptive-sample.jl

@@ -230,7 +230,8 @@ function bin_emissivity_grid(
     d::AbstractAccretionDisc,
     r_bins,
     ϕ_bins,
-    sky::AdaptiveSky{T,<:CoronaGridValues},
+    sky::AdaptiveSky{T,<:CoronaGridValues};
+    kwargs...,
 ) where {T}
     solid_angle, output = bin_emissivity_grid!(
         zeros(eltype(r_bins), (length(r_bins), length(ϕ_bins))),
@@ -239,7 +240,8 @@ function bin_emissivity_grid(
         d,
         r_bins,
         ϕ_bins,
-        sky,
+        sky;
+        kwargs...,
     )
 
     for i in eachindex(solid_angle)
@@ -273,7 +275,8 @@ function bin_emissivity_grid!(
     d::AbstractAccretionDisc,
     r_bins,
     ϕ_bins,
-    sky::AdaptiveSky{T,V},
+    sky::AdaptiveSky{T,V};
+    Γ = 2,
 ) where {T,V<:CoronaGridValues}
     @assert size(output) == size(solid_angle)
     @assert size(output) == (length(r_bins), length(ϕ_bins))
@@ -302,7 +305,7 @@ function bin_emissivity_grid!(
         ϕ_i = searchsortedlast(ϕ_bins, mod2pi(v.ϕ))
         if (r_i != 0) && (ϕ_i != 0)
             # TODO: allow generic spectrum
-            output[r_i, ϕ_i] += ΔΩ * v.J * (v.g^2 * area(v.r))
+            output[r_i, ϕ_i] += ΔΩ * v.J * (v.g^Γ * area(v.r))
             solid_angle[r_i, ϕ_i] += ΔΩ
         end
     end
@@ -392,6 +395,12 @@ function evaluate_refinement_metric(
     [_eval_metric(i, j) for i in axis_itt, j in axis_itt]
 end
 
+@enum StepBlockCodes begin
+    converged
+    not_converged
+    limit_exceeded
+end
+
 """
     function step_block!(
         m::AbstractMetric,
@@ -435,6 +444,7 @@ function step_block!(
     top = 3,
     split = 6,
     percentage = 10,
+    limit = 200_000,
     metric = empty_fraction,
     kwargs...,
 )
@@ -457,7 +467,7 @@ function step_block!(
     if check_threshold
         _counts = filter(i -> i.score < threshold, _counts)
         if length(_counts) == 0
-            return false
+            return converged
         end
     end
 
@@ -478,15 +488,24 @@ function step_block!(
         false
     end
 
-    trace_step!(sky; check_refine = refiner, kwargs...)
-    true
+    to_refine = find_need_refine(sky, refiner)
+    if length(to_refine) > limit
+        @warn (
+            "Convergence failed. Limit $limit exceeded: too many cells to refine (N = $(length(to_refine)))."
+        )
+        return limit_exceeded
+    end
+
+    refine_and_trace!(sky, to_refine; kwargs...)
+    not_converged
 end
 
 """
     refine_function(f)
 
 Used to define a new `check_refine` function for an [`AdaptiveSky`](@ref). This
-can be passed e.g. to [`refine_all`](@ref).
+can be passed e.g. to [`find_need_refine`](@ref) or
+[`refine_and_trace!`](@ref).
 
 The function `f` is given each cell's value and should return `true` if the
 cell should be refined, else `false`.
@@ -572,43 +591,58 @@ function adaptive_solve!(
     end
 
     # trace the initial sky
-    Gradus.trace_initial!(sky)
+    trace_initial!(sky)
     # this happens outside the loop as the first is so fast it tends to ruin
     # the output of the progress bar
-    Gradus.trace_step!(sky)
+    trace_step!(sky)
     i += 1
 
     for _ = 1:trace_calls
-        Gradus.trace_step!(sky; progress_bar = progress, showvalues = showvals)
+        trace_step!(sky; progress_bar = progress, showvalues = showvals)
         i += 1
     end
 
     # ensure the distant radii are well refined
-    Gradus.trace_step!(
+    trace_step!(
         sky;
         check_refine = fine_refine_function(v -> v.r > 800; percentage = 10),
         progress_bar = progress,
         showvalues = showvals,
     )
     i += 1
 
-    while step_block!(
-        m,
-        d,
+    r_isco = isco(m)
+
+    trace_step!(
         sky;
-        split = 5,
-        top = 4,
+        check_refine = fine_refine_function(
+            v -> (v.r > r_isco) && (v.r < 2.0);
+            percentage = 50,
+        ),
         progress_bar = progress,
         showvalues = showvals,
-        kwargs...,
     )
+    i += 1
+
+    status = not_converged
+    while status == not_converged
+        status = step_block!(
+            m,
+            d,
+            sky;
+            split = 5,
+            top = 4,
+            progress_bar = progress,
+            showvalues = showvals,
+            kwargs...,
+        )
         i += 1
         if i >= limit
             break
         end
     end
 
-    Gradus.ProgressMeter.finish!(progress)
+    ProgressMeter.finish!(progress)
 
     i
 end

src/corona/models/ring.jl

@@ -539,7 +539,9 @@ function _branch_emissivity(m::AbstractMetric, slice::PointSlice, I::Vector{Int}
         Δr = (abs(slice.r[i1] - slice.r[i2]))
         Δθ = ang_diff(slice.θ[i1], slice.θ[i2]) / 2
 
-        A = _proper_area(m, slice.r[i] + Δr / 2, π / 2) * (abs(slice.drdθ[i1]) + abs(slice.drdθ[i2])) / 2
+        A =
+            _proper_area(m, slice.r[i] + Δr / 2, π / 2) *
+            (abs(slice.drdθ[i1]) + abs(slice.drdθ[i2])) / 2
         point_source_equatorial_disc_emissivity(
             spectrum,
             slice.θ[i],

src/image-planes/adaptive-sky.jl

@@ -42,11 +42,11 @@ AdaptiveSky(V::Type, calc_v, check_refine; pole_offset = 1e-6) = AdaptiveSky(
 )
 
 """
-refine_all!(sky::AdaptiveSky, to_refine::Vector{Int})
+trace_all!(sky::AdaptiveSky, to_refine::Vector{Int})
 
 Refine all cells in `to_refine` using a multi-threaded loop.
 """
-function refine_all!(
+function trace_all!(
     sky::AdaptiveSky,
     to_refine::Vector{Int};
     verbose = false,
@@ -120,25 +120,16 @@ function trace_initial!(sky::AdaptiveSky; level = 3)
 end
 
 """
-    trace_step!(sky::AdaptiveSky; check_refine = sky.check_refine, verbose = false)
-
-Apply the refinement metric across each cell boundary and refine the cells
-where the metric is `true`.
-
-A different refinemenet metric from the default can be used by passing the
-`check_refine` kwarg, using the same function signature as documented in
-[`AdaptiveSky`](@ref).
+    find_need_refine(sky::AdaptiveSky, check_refine::Function)::Vector{Int}
 
-If `verbose` is true, a progress bar will be displayed during refinement.
+Check all cells with `check_refine`, and return those cell IDs that would need
+to refined.
 """
-function trace_step!(
-    sky::AdaptiveSky{T,V};
-    check_refine = sky.check_refine,
-    kwargs...,
-) where {T,V}
+function find_need_refine(sky::AdaptiveSky, check_refine::Function)::Vector{Int}
     N = lastindex(sky.grid.cells)
 
     to_refine = Set{Int}()
+
     # TODO: multi-thread using a divide and conquer approach
     for cell_index = 1:N
         # skip those we've already refined
@@ -155,32 +146,70 @@ function trace_step!(
         end
     end
 
+    sort!(collect(to_refine))
+end
+
+"""
+    function refine_and_trace!(
+        sky::AdaptiveSky,
+        cell_ids::Vector{Int};
+        kwargs...,
+    )
+
+Given a list of cell IDs, refine them using [`refine!`](@ref), and then trace
+the children to populate their values.
+
+All kwargs are forwarded to [`trace_all!`](@ref).
+"""
+function refine_and_trace!(
+    sky::AdaptiveSky{T,V},
+    cell_ids::Vector{Int};
+    kwargs...,
+) where {T,V}
+    N = lastindex(sky.grid.cells)
+
+    to_trace = Int[]
+    sizehint!(to_trace, length(cell_ids) * 8)
+
     # now that all have been reaped, apply refining
-    to_trace = map(collect(to_refine)) do index
-        if Grids.has_children(sky.grid, index)
-            Int[]
-        else
-            _children = Grids.refine!(sky.grid, index)
-            for (j, _) in enumerate(_children)
-                if j == 5
-                    push!(sky.values, sky.values[index])
-                else
-                    push!(sky.values, make_null(V))
-                end
+    for index in cell_ids
+        @assert !Grids.has_children(sky.grid, index)
+        _children = Grids.refine!(sky.grid, index)
+        for (j, child_id) in enumerate(_children)
+            if j == 5
+                push!(sky.values, sky.values[index])
+            else
+                push!(to_trace, child_id)
+                push!(sky.values, make_null(V))
             end
-            # return all but the middle for tracing
-            Int[_children[[1, 2, 3, 4, 6, 7, 8, 9]]...]
         end
     end
 
     if !isempty(to_trace)
-        all_trace = reduce(vcat, to_trace)
-        refine_all!(sky, all_trace; kwargs...)
+
+        trace_all!(sky, to_trace; kwargs...)
     end
 
     sky
 end
 
+"""
+    trace_step!(sky::AdaptiveSky; check_refine = sky.check_refine, verbose = false)
+
+Apply the refinement metric across each cell boundary and refine the cells
+where the metric is `true`.
+
+A different refinemenet metric from the default can be used by passing the
+`check_refine` kwarg, using the same function signature as documented in
+[`AdaptiveSky`](@ref).
+
+If `verbose` is true, a progress bar will be displayed during refinement.
+"""
+function trace_step!(sky::AdaptiveSky; check_refine = sky.check_refine, kwargs...)
+    to_refine = find_need_refine(sky, check_refine)
+    refine_and_trace!(sky, to_refine; kwargs...)
+end
+
 vector_average(distances, values) = sum(i -> i[1] * i[2], zip(distances, values))
 
 # TODO: use a buffer

src/line-profiles.jl

@@ -126,7 +126,7 @@ function lineprofile(
     d::AbstractAccretionGeometry,
     ::TransferFunctionMethod,
     ;
-    minrₑ = isco(m) + 1e-2, # delta to avoid numerical instabilities
+    minrₑ = isco(m),
     maxrₑ = 50,
     numrₑ = 100,
     verbose = false,