Make integration gpu compatible

[edejong-caltech]

Purpose

Touches integrate_simpon_evenfast

To-do

Content

---

• I have read and checked the items on the review checklist.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 98.21%. Comparing base (29b59bb) to head (d816e5c).
Report is 11 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #162      +/-   ##
==========================================
+ Coverage   98.17%   98.21%   +0.04%     
==========================================
  Files           8        9       +1     
  Lines         493      505      +12     
==========================================
+ Hits          484      496      +12     
  Misses          9        9              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[sajjadazimi]

I remember we changed this from n_bins to n_bins_per_log_unit because in some of the examples where the range of scales is larger the results were not reasonable/satisfactory (because of inaccurate numerical integrations). Would it be possible to keep n_bins_per_log_unit?

[edejong-caltech]

Unfortunately not, because to create an SVector the size must use something in type space and cannot use something that is computed within the body

[charleskawczynski]

How much is @fastmath helping here? There are some issues with @fastmath that I'd rather avoid. e.g., JuliaLang/julia#36214, JuliaLang/julia#36246.

I think there's other ways that we can improve performance without @fastmath

[charleskawczynski]

Is 5:(K - 4) intended?

julia> collect(5:(4 - 4))
Int64[]

[charleskawczynski]

Rather than making large SVectors in moment_source_helper, doing things lazily may be faster (especially on the gpu with SVectors/Tuples of size 100), can we compare the performance against a lazy implementation? E.g., something like:

function integrate_SimpsonEvenFast(::Val{n_bins}, x_min::FT, dx::FT, y::F) where {FT <: Real, F, n_bins}
    n_bins ≥ 3 || error("n_bins must be at least 3")
    # dx = SVector{K - 1}(x[2:end] - x[1:(end - 1)])
    # need to figure out how to verify this if still applicable
    # is_uniform(dx) || error("x must be evenly spaced")
    e = n_bins + 1
    @inbounds begin
        retval =
            sum(j-> y(j), 5:(n_bins-3)) +
            (17 * (y(1) + y(e)) + 59 * (y(2) + y(e - 1)) + 43 * (y(3) + y(e - 2)) + 49 * (y(4) + y(e - 3))) / 48
        return (logx(x_min, 2, dx) - logx(x_min, 1, dx)) * retval
    end
end

logx(x_min::FT, j::Int, dx::FT) where {FT} = x_min + (j - 1) * dx
function moment_source_helper(
    dist::GammaPrimitiveParticleDistribution{FT},
    p1::FT,
    p2::FT,
    x_threshold::FT,
    ::Val{n_bins} = Val(100),
) where {n_bins, FT <: Real}
    (; n, θ, k) = dist

    f(x) = x^(p1 + k - 1) * exp(-x / θ) * gamma_inc(p2 + k, (x_threshold - x) / θ)[1] * gamma(p2 + k)

    x_lowerbound = FT(min(FT(1e-5), FT(1e-5) * x_threshold))
    x_min = log(x_lowerbound)
    dx = (log(x_threshold) - log(x_lowerbound)) / n_bins
    y_func(j) = j <= n_bins ? exp(logx(x_min, j, dx)) * f(exp(logx(x_min, j, dx))) : zero(typeof(dx))
    return n^2 * θ^(p2 - k) / gamma(k)^2 * integrate_SimpsonEvenFast(Val(n_bins), x_min, dx, y_func)
end

[juliasloan25]

Looks good to me!

[juliasloan25]

I don't think you need the outer FT if you're wrapping the inner expressions with it. i.e. this could be:
x_lowerbound = min(FT(1e-5), FT(1e-5) * x_threshold)

[juliasloan25]

same note as above, you can remove the outer FT now (it might be cleaner)

[juliasloan25]

Why did you add this? Just curious what it helps with :)

[edejong-caltech]

It's to avoid counting the allocations that to go the return value -- Charlie/Dennis' suggestion :)