Add indefinite integral / antiderivative operator to Oceanostics

[tomchor]

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:

@kernel function antiderivative!(B, grid, buoyancy)
    i, j = @index(Global, NTuple)

    @inbounds B[i, j, grid.Nz] = buoyancy[i, j, grid.Nz+1] * Δzᶜᶜᶠ(i, j, grid.Nz+1, grid)

    @unroll for k in grid.Nz-1 : -1 : 1
        @inbounds B[i, j, k] = B[i, j, k+1] - buoyancy[i, j, k+1] * Δzᶜᶜᶠ(i, j, k+1, grid)
    end
end

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?

[glwagner]

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.

[tomchor]

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.)

[glwagner]

You can wrap KernelFunctionOperation in Field.

[tomchor]

You can wrap KernelFunctionOperation in Field.

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?

[glwagner]

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.