Merge pull request #197 from astro-group-bristol/fergus/speedup-integ…

…ration Speedup transfer function integration
astro-group-bristol · Jun 23, 2024 · 1d420d5 · 1d420d5
2 parents 42a8dab + 0bf4b93
commit 1d420d5
Show file tree

Hide file tree

Showing 9 changed files with 352 additions and 196 deletions.
diff --git a/src/line-profiles.jl b/src/line-profiles.jl
@@ -69,7 +69,15 @@ function lineprofile(
     _warn_disc_integration_limits(d, minrₑ, maxrₑ)
     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, Nr = Nr)
+    flux = integrate_lineprofile(
+        ε,
+        itfs,
+        bins;
+        h = h,
+        n_radii = Nr,
+        rmin = minimum(radii),
+        rmax = maximum(radii),
+    )
     bins, flux
 end
 

diff --git a/src/tracing/precision-solvers.jl b/src/tracing/precision-solvers.jl
@@ -75,7 +75,7 @@ function _find_offset_for_radius(
     function _measure(gp::GeodesicPoint{T}) where {T}
         r = _equatorial_project(gp.x)
         if gp.status != StatusCodes.IntersectedWithGeometry
-            r = -r
+            r = -2r
         end
         rₑ - r
     end
@@ -270,10 +270,11 @@ function jacobian_∂αβ_∂gr(
 
     # these type hints are crucial for forward diff to be type stable
     function _jacobian_f(impact_params::SVector{2,T})::SVector{2,T} where {T}
+        v = map_impact_parameters(m, x, impact_params[1], impact_params[2])
         sol = tracegeodesics(
             m,
             x,
-            map_impact_parameters(m, x, impact_params[1], impact_params[2]),
+            v,
             d,
             max_time;
             save_on = false,

diff --git a/src/transfer-functions/cunningham-transfer-functions.jl b/src/transfer-functions/cunningham-transfer-functions.jl
@@ -150,7 +150,7 @@ function interpolate_branches(ctf::CunninghamTransferData{T}; h = 1e-6) where {T
     lower_time_branch = _make_interpolation(lower_g✶, lower_t)
     upper_branch = _make_interpolation(upper_g✶, upper_f)
     upper_time_branch = _make_interpolation(upper_g✶, upper_t)
-    TransferBranches(
+    TransferBranches{false}(
         upper_branch,
         lower_branch,
         upper_time_branch,
@@ -260,31 +260,22 @@ function _rear_workhorse(
         )
     function _thick_workhorse(θ::T)::NTuple{4,T} where {T}
         g, gp, r = datum_workhorse(θ)
-        r₊, _ = try
-            _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,
-            )
-        catch err
-            if err isa Roots.ConvergenceFailed
-                # return gp for type stability
-                NaN, gp
-            else
-                throw(err)
-            end
-        end
+        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,
+        )
         is_visible, J = if !isnan(r₊) && isapprox(r, r₊, atol = 1e-3)
             # trace jacobian on updated impact parameters
             α, β = _rθ_to_αβ(r₊, θ; α₀ = setup.α₀, β₀ = setup.β₀)
@@ -501,24 +492,6 @@ function transfer_function_grid(itfs::InterpolatingTransferBranches, Ng::Int)
     )
 end
 
-function (grid::CunninghamTransferGrid)(r)
-    idx = max(
-        1,
-        min(
-            DataInterpolations.searchsortedlastcorrelated(grid.r_grid, r, 0),
-            length(grid.r_grid) - 1,
-        ),
-    )
-    r1, r2 = grid.r_grid[idx], grid.r_grid[idx+1]
-    # interpolation weight
-    θ = (r - r1) / (r2 - r1)
-
-    gmin = _linear_interpolate(grid.g_min, idx, θ)
-    gmax = _linear_interpolate(grid.g_max, idx, θ)
-    upper_f, lower_f, upper_t, lower_t = _lazy_interpolate(grid, idx, θ)
-    TransferBranches(upper_f, lower_f, upper_t, lower_t, gmin, gmax, r)
-end
-
 export CunninghamTransferData,
     TransferBranches,
     InterpolatingTransferBranches,