astro-group-bristol / Gradus.jl

Extensible spacetime agnostic general relativistic ray-tracing (GRRT).
https://astro-group-bristol.github.io/Gradus.jl/dev/
GNU General Public License v3.0
18 stars 2 forks source link

Speed up transfer function integration #192

Closed fjebaker closed 3 months ago

fjebaker commented 3 months ago

The integration routines were written pretty lazily, and need to be reworked now that we want to fit to data. It's fast enough, but it could be much faster.

fjebaker commented 3 months ago
using Gradus, BenchmarkTools

d = ThinDisc(0.0, 400.0)
x = SVector(0.0, 1000.0, deg2rad(40), 0.0)
m = KerrMetric(1.0, 0.4)

# maximal integration radius
maxrₑ = 50.0

# emissivity function
ε(r) = r^(-3)

gs = range(0.0, 1.2, 200)
radii = Gradus.Grids._inverse_grid(Gradus.isco(m) + 1e-2, 100.0, 100)
tf_grid = @time Gradus.transfer_function_grid(m, x, d, radii; verbose = true)

@btime Gradus.integrate_lineprofile($ε, $tf_grid, $radii, $gs)

Currently 771.390 ms (1337754 allocations: 226.90 MiB)

fjebaker commented 3 months ago

Fixed a horrendous runtime dispatch with a Hail Mary use of @inline. Can now do the computation in about 170ms without allocating.

Profiling the integration call, it looks like most of the time is just spent on finding the index to interpolate with. Since we know the approximate access pattern into the transfer functions (monotonic in g), there is room for big optimization here:

Sat 22 Jun 17:03:48 CEST 2024

fjebaker commented 3 months ago

Proof of concept implementation achieves good performance 25.855 ms (0 allocations: 0 bytes)

Sat 22 Jun 17:41:31 CEST 2024