buffer preallocation in hcubature

src/HCubature.jl

@@ -20,7 +20,7 @@ module HCubature
 using StaticArrays, LinearAlgebra
 import Combinatorics, DataStructures, QuadGK
 
-export hcubature, hquadrature, alloc_buf
+export hcubature, hquadrature, hcubature_buffer
 
 include("genz-malik.jl")
 include("gauss-kronrod.jl")
@@ -47,20 +47,36 @@ cubrule(::Val{0}, ::Type{T}) where {T} = Trivial()
 countevals(::Trivial) = 1
 
 """
-    alloc_buf(;dimension[, domain_type, range_type, error_type])
+    hcubature_buffer(;dimension[, domain_type, range_type, error_type])
 
 Allocate a buffer that can be used in calls to [`hcubature`](@ref).
 
 # Examples:
 
 ```julia
- buffer = alloc_buf(;dimension=2, range_type=ComplexF64, domain_type=Float64)
+ buffer = hcubature_buffer(;dimension=2, range_type=ComplexF64, domain_type=Float64)
  I, E = hcubature(x -> 2+im, (0,0), (2pi, pi); buffer))
  ```
-
 """
-function alloc_buf(;dimension, domain_type=Float64, range_type=Float64, error_type=real(range_type))
-    return DataStructures.BinaryMaxHeap{Box{dimension,domain_type,range_type, error_type}}()
+function hcubature_buffer(f,a,b;norm=norm)
+    hcubature_buffer_(f,a,b,norm)
+end
+
+function hcubature_buffer_(f,a::SVector{N,T},b::SVector{N,T},norm) where {N,T}
+    rule = cubrule(Val{N}(), T)
+    I, E, _ = rule(f, a, b, norm)
+    firstbox = Box(a, b, I, E, 0)
+    DataStructures.BinaryMaxHeap{typeof(firstbox)}()
+end
+
+function hcubature_buffer_(f, a::AbstractVector{T}, b::AbstractVector{S},norm) where {T<:Real, S<:Real}
+    length(a) == length(b) || throw(DimensionMismatch("endpoints $a and $b must have the same length"))
+    F = float(promote_type(T, S))
+    return hcubature_buffer_(f, SVector{length(a),F}(a), SVector{length(a),F}(b), norm)
+end
+
+function hcubature_buffer_(f, a::Tuple{Vararg{Real,n}}, b::Tuple{Vararg{Real,n}}, norm) where {n}
+    hcubature_buffer_(f, SVector{n}(float.(a)), SVector{n}(float.(b)), norm)
 end
 
 function hcubature_(f::F, a::SVector{n,T}, b::SVector{n,T}, norm, rtol_, atol, maxevals, initdiv, buf) where {F, n, T<:Real}
@@ -195,13 +211,14 @@ returns a vector of integrands with different scalings.)
 
 In normal usage, `hcubature(...)` will allocate a buffer for internal
 computations. You can instead pass a preallocated buffer allocated using
-`alloc_buf' as the `buffer` argument. This buffer can be used across
+[`hcubature_buffer'](@ref) as the `buffer` argument. This buffer can be used across
 multiple calls to avoid repeated allocation.
 """
 hcubature(f, a, b; norm=norm, rtol::Real=0, atol::Real=0,
                    maxevals::Integer=typemax(Int), initdiv::Integer=1, buffer=nothing) =
     hcubature_(f, a, b, norm, rtol, atol, maxevals, initdiv, buffer)
 
+
 """
     hquadrature(f, a, b; norm=norm, rtol=sqrt(eps), atol=0, maxevals=typemax(Int), initdiv=1)
 

test/runtests.jl

@@ -75,17 +75,28 @@ end
       @test hcubature(x -> x[2] < 0 ? Inf : x[1]*x[2], [-1, -1], [1, 1]) === (Inf, NaN)
 end
 
-@testset "alloc_buf" begin
-    buffer = alloc_buf(;dimension=1)
-    @test @inferred(hcubature(x -> cos(x[1]), (0,), (1,);buffer=buffer))[1] ≈ sin(1) ≈
-    @inferred(hquadrature(cos, 0, 1; buffer=buffer))[1]
-    buffer = alloc_buf(;dimension=2)
-    @test hcubature(x -> cos(x[1])*cos(x[2]), [0,0], [1,1]; buffer=buffer)[1] ≈ sin(1)^2 ≈
-    @inferred(hcubature(x -> cos(x[1])*cos(x[2]), (0,0), (1,1);buffer=buffer))[1]
-    buffer = alloc_buf(;dimension=1, range_type=Float32, domain_type=Float32)
-    @test @inferred(hcubature(x -> cos(x[1]), (0.0f0,), (1.0f0,);buffer=buffer))[1] ≈ sin(1.0f0)
-    buffer = alloc_buf(;dimension=2, range_type=Float64, domain_type=Float64)
-    @test @inferred(hcubature(x -> 2, (0,0), (2pi, pi);buffer=buffer)[1]) ≈ 4pi^2
-    buffer = alloc_buf(;dimension=2, range_type=ComplexF64, domain_type=Float64)
-    @test @inferred(hcubature(x -> 2+im, (0,0), (2pi, pi);buffer=buffer))[1] ≈ 4pi^2 + im*2*pi^2
+@testset "hcubature_buffer" begin
+    # 1d
+    f = x->cos(x[1])
+    a,b = (0,), (1,)
+    buffer = hcubature_buffer(f,a,b)
+    @test @inferred(hcubature(f,a,b;buffer=buffer))[1] ≈ sin(1) ≈
+    @inferred(hquadrature(f, 0, 1; buffer=buffer))[1]
+    # 2d
+    f = x -> cos(x[1])*cos(x[2])
+    a,b = [0,0], [1,1]
+    at,bt = (0,0), (1,1)
+    buffer = hcubature_buffer(f,a,b)
+    @test hcubature(f,a,b; buffer=buffer)[1] ≈ sin(1)^2 ≈
+    @inferred(hcubature(f, Tuple(a), Tuple(b);buffer=buffer))[1]
+    # 1d single precision
+    f = x -> cos(x[1])
+    a,b = (0.0f0,), (1.0f0,)
+    buffer = hcubature_buffer(f,a,b)
+    @test @inferred(hcubature(f,a,b;buffer=buffer))[1] ≈ sin(1.0f0)
+    # 2d complex entries
+    f = x -> (1+im)*cos(x[1])*cos(x[2])
+    a,b = (0,0), (1, 1)
+    buffer = hcubature_buffer(f,a,b)
+    @test @inferred(hcubature(f,a,b;buffer=buffer))[1] ≈ (1+im)*sin(1)^2
 end