Skip to content

Commit

Permalink
format and include
Browse files Browse the repository at this point in the history
  • Loading branch information
lxvm committed Sep 19, 2023
1 parent 21b7895 commit 6a2038a
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 13 deletions.
1 change: 1 addition & 0 deletions src/SciMLBase.jl
Original file line number Diff line number Diff line change
Expand Up @@ -714,6 +714,7 @@ include("ensemble/ensemble_analysis.jl")

include("solve.jl")
include("interpolation.jl")
include("integrand_interface.jl")
include("integrator_interface.jl")
include("tabletraits.jl")
include("remake.jl")
Expand Down
33 changes: 20 additions & 13 deletions src/integrand_interface.jl
Original file line number Diff line number Diff line change
Expand Up @@ -8,13 +8,12 @@ written to `result`, so use the IntegralSolution immediately after the calculati
the result, and don't expect it to persist if the same integrand is used for another
calculation.
"""
struct InplaceIntegrand{F,T<:AbstractArray}
struct InplaceIntegrand{F, T <: AbstractArray}
# in-place function f!(y, x, p) that takes one x value and outputs an array of results in-place
f!::F
I::T
end


"""
BatchIntegrand(f!, y::AbstractArray, x::AbstractVector, max_batch=typemax(Int))
Expand All @@ -24,35 +23,39 @@ the GPU, or distributed-memory. The `max_batch` keyword is a soft limit on the n
nodes passed to the integrand. The buffers `y,x` must both be `resize!`-able since the
number of evaluation points may vary between calls to `f!`.
"""
struct BatchIntegrand{F,Y,X}
struct BatchIntegrand{F, Y, X}
# in-place function f!(y, x, p) that takes an array of x values and outputs an array of results in-place
f!::F
y::Y
x::X
max_batch::Int # maximum number of x to supply in parallel
function BatchIntegrand(f!, y::AbstractVector, x::AbstractVector, max_batch::Integer=typemax(Int))
function BatchIntegrand(f!,
y::AbstractVector,
x::AbstractVector,
max_batch::Integer = typemax(Int))
max_batch > 0 || throw(ArgumentError("maximum batch size must be positive"))
return new{typeof(f!),typeof(y),typeof(x)}(f!, y, x, max_batch)
return new{typeof(f!), typeof(y), typeof(x)}(f!, y, x, max_batch)
end
end


"""
BatchIntegrand(f!, y, x; max_batch=typemax(Int))
Constructor for a `BatchIntegrand` with pre-allocated buffers.
"""
BatchIntegrand(f!, y, x; max_batch::Integer=typemax(Int)) =
function BatchIntegrand(f!, y, x; max_batch::Integer = typemax(Int))
BatchIntegrand(f!, y, x, max_batch)
end

"""
BatchIntegrand(f!, y::Type, x::Type=Nothing; max_batch=typemax(Int))
Constructor for a `BatchIntegrand` whose range type is known. The domain type is optional.
Array buffers for those types are allocated internally.
"""
BatchIntegrand(f!, Y::Type, X::Type=Nothing; max_batch::Integer=typemax(Int)) =
function BatchIntegrand(f!, Y::Type, X::Type = Nothing; max_batch::Integer = typemax(Int))
BatchIntegrand(f!, Y[], X[], max_batch)
end

"""
InplaceBatchIntegrand(f!, result::AbstractArray, y::AbstractArray, x::AbstractVector, max_batch=typemax(Int))
Expand All @@ -65,24 +68,28 @@ buffers `y,x` must both be `resize!`-able since the number of evaluation points
between calls to `f!`. In particular, for a resizeable `y` buffer see ElasticArrays.jl . The
solution is written inplace to `result`.
"""
struct InplaceBatchIntegrand{F,T,Y,X}
struct InplaceBatchIntegrand{F, T, Y, X}
# in-place function f!(y, x, p) that takes an array of x values and outputs an array of results in-place
f!::F
I::T
y::Y
x::X
max_batch::Int # maximum number of x to supply in parallel
function InplaceBatchIntegrand(f!, I::AbstractArray, y::AbstractArray, x::AbstractVector, max_batch::Integer=typemax(Int))
function InplaceBatchIntegrand(f!,
I::AbstractArray,
y::AbstractArray,
x::AbstractVector,
max_batch::Integer = typemax(Int))
max_batch > 0 || throw(ArgumentError("maximum batch size must be positive"))
return new{typeof(f!),typeof(I),typeof(y),typeof(x)}(f!, I, y, x, max_batch)
return new{typeof(f!), typeof(I), typeof(y), typeof(x)}(f!, I, y, x, max_batch)
end
end


"""
InplaceBatchIntegrand(f!, result, y, x; max_batch=typemax(Int))
Constructor for a `InplaceBatchIntegrand` with pre-allocated buffers.
"""
InplaceBatchIntegrand(f!, result, y, x; max_batch::Integer=typemax(Int)) =
function InplaceBatchIntegrand(f!, result, y, x; max_batch::Integer = typemax(Int))
InplaceBatchIntegrand(f!, result, y, x, max_batch)
end

0 comments on commit 6a2038a

Please sign in to comment.