Open tomchor opened 10 months ago
This might be implementable as a Reduction
using cumsum
and GridMetricOperation
(basically the same as Integral
, but this time CumulativeIntegral
--- which I think is the right term rather than antiderivative or indefinite integral; this integral does have one definite limit).
But I think we looked into that, and it didn't work on the GPU? If so that would leave manually writing the kernels (as you've done here) as the only option.
Do you have an example of how to wrap a custom-written kernel in a Field
? (I'm assuming it can't be a KernelFunctionOperation
due to the non-locality of the calculation.)
You can wrap KernelFunctionOperation
in Field
.
You can wrap
KernelFunctionOperation
inField
.
I meant I think you can't wrap a non-local (due to the integration) custom-built kernel in a KernelFunctionOperation
.
Do you have examples on how I can wrap this integral in a Field?
The interface for implementing reductions is here:
https://github.com/CliMA/Oceananigans.jl/blob/main/src/Fields/field_reductions.jl
an example of computing a Reduction
with Field
is in the docstring.
Some reductions depend on the grid metric (like an integral, as opposed to, for example, a maximum). Examples of extending Reduction
for these cases are here:
https://github.com/CliMA/Oceananigans.jl/blob/main/src/AbstractOperations/metric_field_reductions.jl
On the CPU, cumsum!
is already sufficient:
help?> cumsum!
search: cumsum! cumsum
cumsum!(B, A; dims::Integer)
Cumulative sum of A along the dimension dims, storing the result in B. See also cumsum.
One difference between "cumulative reductions" versus non-cumulative ones is that dims
cannot be inferred from size(reduction.operand)
. Therefore a new struct is needed:
struct CumulativeSum
dims :: Int
end
(cs::CumulativeSum)(B, A) = cumsum!(B, C, cs.dims)
It makes sense to provide a convenience constructor:
Reduction(::typeof(cumsum!), operand; dims) = Reduction(CumulativeSum(dims), operand, dims=())
For cumulative integration, one can reuse CumulativeSum
. All that is needed is an appropriate transformation using GridMetricOperation
, same as Integral
.
For the GPU I don't think cumsum!
works (unless it does now...)
One needs to write custom kernels as in the OP for that case probably.
I think Oceananigans doesn't have an indifinite integral operator, only definite integrals. Ideally we'd do it in a way that can be plugged into KernelFunctionOperator, but I suspect that this isn't possible to be the fact that you need to unroll loops. For example, here's a (downwards) buoyancy integration used to calculate the hydrostatic pressure in Ocenanigans:
https://github.com/CliMA/Oceananigans.jl/blob/fa5e280115f619d01a460f012328bd7e6d253b38/src/Models/NonhydrostaticModels/update_hydrostatic_pressure.jl#L4-L18
The basic form (assuming a direct integration of
buoyancy
) is something like:Not sure how to best implement this here, but it would be useful. Maybe it's easier to just do these calculations on the CPU?