fix: enormously reduce allocations in composite disc tracing

astro-group-bristol · Jun 8, 2024 · 245cf7c · 245cf7c
1 parent 38c5e06
commit 245cf7c
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/src/tracing/radiative-transfer-problem.jl b/src/tracing/radiative-transfer-problem.jl
@@ -65,20 +65,21 @@ function _intensity_delta(
     m,
     x,
     k::AbstractArray{T},
-    geometry::CompositeGeometry,
+    geometry::CompositeGeometry{K,D},
     within,
     I,
     ν₀,
     r_isco,
     λ,
-) where {T}
-    total::T = zero(T)
-    for i in eachindex(geometry.geometry)
-        total += _intensity_delta(
+) where {T,K,D}
+    indexed_geom = (enumerate(geometry.geometry)...,)
+    sum(indexed_geom) do args
+        i, geom = args
+        _intensity_delta(
             m,
             x,
             k,
-            geometry.geometry[i],
+            geom,
             within,
             I,
             ν₀,
@@ -87,7 +88,6 @@ function _intensity_delta(
             index = i,
         )
     end
-    total
 end
 
 function _intensity_delta(