Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[RFC] Better integral docs #609

Merged
merged 3 commits into from
Feb 10, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 6 additions & 5 deletions src/problems/basic_problems.jl
Original file line number Diff line number Diff line change
Expand Up @@ -448,7 +448,7 @@
p::P
kwargs::K
@add_kwonly function IntegralProblem{iip}(f::AbstractIntegralFunction{iip}, domain,
p = NullParameters();
p = NullParameters(); nout = nothing, batch = nothing,
kwargs...) where {iip}
warn_paramtype(p)
new{iip, typeof(p), typeof(f), typeof(domain), typeof(kwargs)}(f,
Expand All @@ -465,17 +465,18 @@
IntegralProblem{isinplace(f)}(f, domain, p; kwargs...)
end

@deprecate IntegralProblem(f::AbstractIntegralFunction,
@deprecate IntegralProblem{iip}(f::AbstractIntegralFunction,
lb::Union{Number,AbstractVector{<:Number}},
ub::Union{Number,AbstractVector{<:Number}},
p = NullParameters(); kwargs...) IntegralProblem(f, (lb, ub), p; kwargs...)
p = NullParameters(); kwargs...) where {iip} IntegralProblem{iip}(f, (lb, ub), p; kwargs...)

function IntegralProblem(f, args...; nout = nothing, batch = nothing, kwargs...)
IntegralProblem(f, args...; kwargs...) = IntegralProblem{isinplace(f, 3)}(f, args...; kwargs...)
function IntegralProblem{iip}(f, args...; nout = nothing, batch = nothing, kwargs...) where {iip}

Check warning on line 474 in src/problems/basic_problems.jl

View check run for this annotation

Codecov / codecov/patch

src/problems/basic_problems.jl#L473-L474

Added lines #L473 - L474 were not covered by tests
if nout !== nothing || batch !== nothing
@warn "`nout` and `batch` keywords are deprecated in favor of inplace `IntegralFunction`s or `BatchIntegralFunction`s. See the updated Integrals.jl documentation for details."
end

g = if isinplace(f, 3)
g = if iip

Check warning on line 479 in src/problems/basic_problems.jl

View check run for this annotation

Codecov / codecov/patch

src/problems/basic_problems.jl#L479

Added line #L479 was not covered by tests
if batch === nothing
output_prototype = nout === nothing ? Array{Float64, 0}(undef) : Vector{Float64}(undef, nout)
IntegralFunction(f, output_prototype)
Expand Down
84 changes: 50 additions & 34 deletions src/scimlfunctions.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2088,12 +2088,12 @@ present, `f` is interpreted as in-place, and otherwise `f` is assumed to be out-
## iip: In-Place vs Out-Of-Place

Out-of-place functions must be of the form ``y = f(u, p)`` and in-place functions of the form
``f(y, u, p)``. Since `f` is allowed to return any type (e.g. real or complex numbers or
``f(y, u, p)``, where `y` is a number or array containing the output. Since `f` is allowed to return any type (e.g. real or complex numbers or
arrays), in-place functions must provide a container `integrand_prototype` that is of the
right type for the variable ``y``, and the result is written to this container in-place.
right type and size for the variable ``y``, and the result is written to this container in-place.
When in-place forms are used, in-place array operations, i.e. broadcasting, may be used by
algorithms to reduce allocations. If `integrand_prototype` is not provided, `f` is assumed
to be out-of-place and quadrature is performed assuming immutable return types.
to be out-of-place.

## specialize

Expand All @@ -2112,59 +2112,75 @@ end
TruncatedStacktraces.@truncate_stacktrace IntegralFunction 1 2

@doc doc"""
BatchIntegralFunction{iip,specialize,F,T} <: AbstractIntegralFunction{iip}
BatchIntegralFunction{iip,specialize,F,T} <: AbstractIntegralFunction{iip}

A representation of an integrand `f` that can be evaluated at multiple points simultaneously
using threads, the gpu, or distributed memory defined by:
A batched representation of an (non-batched) integrand `f(u, p)` that can be
evaluated at multiple points simultaneously using threads, the gpu, or
distributed memory defined by:

```math
y = f(u, p)
by = bf(bu, p)
```

``u`` is a vector whose elements correspond to distinct evaluation points to `f`, whose
output must be returned as an array whose last "batching" dimension corresponds to integrand
evaluations at the different points in ``u``. In general, the integration algorithm is
allowed to vary the number of evaluation points between subsequent calls to `f`.
Here we prefix variables with `b` to indicate they are batched variables, which
implies that they are arrays whose **last** dimension is reserved for batching
different evaluation points, or function values, and may be of a variable
length. ``bu`` is an array whose elements correspond to distinct evaluation
points to `f`, and `bf` is a function to evaluate `f` 'point-wise' so that
`f(bu[..., i], p) == bf(bu, p)[..., i]`. For example, a simple batching implementation
of a scalar, univariate function is via broadcasting: `bf(bu, p) = f.(bu, Ref(p))`,
although this interface exists in order to allow user parallelization.
In general, the integration algorithm is allowed to vary the number of
evaluation points between subsequent calls to `bf`.

For an in-place form of `f` see the `iip` section below for details on in-place or
out-of-place handling.
For an in-place form of `bf` see the `iip` section below for details on in-place
or out-of-place handling.

```julia
BatchIntegralFunction{iip,specialize}(f, [integrand_prototype];
BatchIntegralFunction{iip,specialize}(bf, [integrand_prototype];
max_batch=typemax(Int))
```
Note that only `f` is required, and in the case of inplace integrands a mutable container
`integrand_prototype` to store a batch of integrand evaluations, with a last "batching"
dimension.
Note that only `bf` is required, and in the case of inplace integrands a mutable
container `integrand_prototype` to store a batch of integrand evaluations, with
a last "batching" dimension.

The keyword `max_batch` is used to set a soft limit on the number of points to batch at the
same time so that memory usage is controlled.
The keyword `max_batch` is used to set a soft limit on the number of points to
batch at the same time so that memory usage is controlled.

If `integrand_prototype` is present, `f` is interpreted as in-place, and otherwise `f` is
assumed to be out-of-place.
If `integrand_prototype` is present, `bf` is interpreted as in-place, and
otherwise `bf` is assumed to be out-of-place.

## iip: In-Place vs Out-Of-Place

Out-of-place functions must be of the form ``y = f(u,p)`` and in-place functions of the form
``f(y, u, p)``. Since `f` is allowed to return any type (e.g. real or complex numbers or
arrays), in-place functions must provide a container `integrand_prototype` of the right type
for ``y``. The only assumption that is enforced is that the last axes of `the `y`` and ``u``
arrays are the same length and correspond to distinct batched points. The algorithm will
then allocate arrays `similar` to ``y`` to pass to the integrand. Since the algorithm may
vary the number of points to batch, the length of the batching dimension of ``y`` may vary
between subsequent calls to `f`. To reduce allocations, views of ``y`` may also be passed to
the integrand. In the out-of-place case, the algorithm may infer the type
of ``y`` by passing `f` an empty array of input points. When in-place forms are used,
in-place array operations may be used by algorithms to reduce allocations. If
`integrand_prototype` is not provided, `f` is assumed to be out-of-place.
Out-of-place functions must be of the form `by = bf(bu, p)` and in-place
functions of the form `bf(by, bu, p)` where `by` is a batch array containing the
output. Since the algorithm may vary the number of points to batch, the batching
dimension can be of any length, including zero, and since `bf` is allowed to
return arrays of any type (e.g. real or complex) or size, in-place functions
must provide a container `integrand_prototype` of the desired type and size for
`by`. If `integrand_prototype` is not provided, `bf` is assumed to be
out-of-place.

In the out-of-place case, we require `f(bu[..., i], p) == bf(bu, p)[..., i]`,
and certain algorithms, such as those implemented in C, may infer the type or
shape of `by` by calling `bf` with an empty array of input points, i.e. `bu`
with `size(bu)[end] == 0`. Then it is expected for the resulting `by` to have
the same type and `size(by)[begin:end-1]` for all subsequent calls.

When the in-place form is used, we require `f(by[..., i], bu[..., i], p) ==
bf(by, bu, p)[..., i]` and `size(by)[begin:end-1] ==
size(integrand_prototype)[begin:end-1]`. The algorithm should always pass the
integrand `by` arrays that are `similar` to `integrand_prototype`, and may use
views and in-place array operations to reduce allocations.

## specialize

This field is currently unused

## Fields

The fields of the BatchIntegralFunction type directly match the names of the inputs.
The fields of the BatchIntegralFunction type are `f`, corresponding to `bf`
above, and `integrand_prototype`.
"""
struct BatchIntegralFunction{iip, specialize, F, T} <:
AbstractIntegralFunction{iip}
Expand Down
26 changes: 19 additions & 7 deletions test/function_building_error_messages.jl
Original file line number Diff line number Diff line change
Expand Up @@ -461,15 +461,27 @@ NonlinearFunction(nfiip, vjp = nvjp)
NonlinearFunction(nfoop, vjp = nvjp)

# Integrals
intf(u) = 1.0
@test_throws SciMLBase.TooFewArgumentsError IntegralProblem(intf, (0.0, 1.0))
intfew(u) = 1.0
@test_throws SciMLBase.TooFewArgumentsError IntegralProblem(intfew, (0.0, 1.0))
@test_throws SciMLBase.TooFewArgumentsError IntegralFunction(intfew)
@test_throws SciMLBase.TooFewArgumentsError IntegralFunction(intfew, zeros(3))
@test_throws SciMLBase.TooFewArgumentsError BatchIntegralFunction(intfew)
@test_throws SciMLBase.TooFewArgumentsError BatchIntegralFunction(intfew, zeros(3))
intf(u, p) = 1.0
p = 2.0

IntegralProblem(intf, (0.0, 1.0))
IntegralProblem(intf, (0.0, 1.0), p)
IntegralProblem(intf, ([0.0], [1.0]))
IntegralProblem(intf, ([0.0], [1.0]), p)
intfiip(y, u, p) = y .= 1.0

for (f, kws, iip) in (
(intf, (;), false),
(IntegralFunction(intf), (;), false),
(intfiip, (; nout=3), true),
(IntegralFunction(intfiip, zeros(3)), (;), true),
), domain in (((0.0, 1.0),), (([0.0], [1.0]),), (0.0, 1.0), ([0.0], [1.0],))
IntegralProblem(f, domain...; kws...)
IntegralProblem(f, domain..., p; kws...)
IntegralProblem{iip}(f, domain...; kws...)
IntegralProblem{iip}(f, domain..., p; kws...)
end

x = [1.0, 2.0]
y = rand(2, 2)
Expand Down
Loading