KernelFunctionOperation produces incorrect result when given a NamedTuple with an AbstractOperation

[tomchor]

I noticed that when I pass a NamedTuple to a KernelFunctionOperation, it'll generally give me an incorrect result if one of the variables used in the calculation is in the tuple and it's an operation.

Here's a MWE where I'm creating a model with a constant u=2, setting up a simple KFO that returns u from velocities:

using Oceananigans
using Oceananigans.AbstractOperations: KernelFunctionOperation
using Oceananigans.Fields: @compute

grid = RectilinearGrid(size = (1, 1, 4), extent = (1,1,1));

model = NonhydrostaticModel(; grid,)
u₀ = 2;
set!(model, u=u₀)

@inline ufunc(i, j, k, grid, velocities) = velocities.u[i, j, k]

If I just pass model.velocities to it, it works fine and returns u. If I try to pass, for example, the perturbation u - Average(u), then for some reason the KFO returns u:

julia> velocities = (u = (model.velocities.u - Field(Average(model.velocities.u))),);

julia> @compute f1 = Field(KernelFunctionOperation{Face, Center, Center}(ufunc, model.grid, velocities));

julia> interior(f1) # This should be all zero
1×1×4 view(::Array{Float64, 3}, 4:4, 4:4, 4:7) with eltype Float64:
[:, :, 1] =
 2.0

[:, :, 2] =
 2.0

[:, :, 3] =
 2.0

[:, :, 4] =
 2.0

I expected the above to either work (i.e. produce the correct result) or throw an error, but returning u is unexpected. If, however, I define velocities = (u = Field(model.velocities.u - Field(Average(model.velocities.u))),) (i.e. wrap the operation in a Field) then the result is correct.

Curiously, if I use u₀ as the average in the snippet above (without wrapping the operation in Field), it appears to work:

julia> velocities = (u = (model.velocities.u - u₀),);

julia> @compute f2 = Field(KernelFunctionOperation{Face, Center, Center}(ufunc, model.grid, velocities));

julia> interior(f2)
1×1×4 view(::Array{Float64, 3}, 4:4, 4:4, 4:7) with eltype Float64:
[:, :, 1] =
 0.0

[:, :, 2] =
 0.0

[:, :, 3] =
 0.0

[:, :, 4] =
 0.0

Also if I bypass tuples altogether the result seems to be correct:

julia> @inline ufunc2(i, j, k, grid, u) = u[i, j, k];

julia> u′ = model.velocities.u - Field(Average(model.velocities.u));

julia> @compute f3 = Field(KernelFunctionOperation{Face, Center, Center}(ufunc2, model.grid, u′));

julia> interior(f3)
1×1×4 view(::Array{Float64, 3}, 4:4, 4:4, 4:7) with eltype Float64:
[:, :, 1] =
 0.0

[:, :, 2] =
 0.0

[:, :, 3] =
 0.0

[:, :, 4] =
 0.0

Are KFOs meant to be used only with Fields? Or is this a bug?

[glwagner]

I'm having trouble following your example. Can you come up with a simpler illustration of the issue and avoid constructions with many operations on a line?

[glwagner]

Does this illustrate the problem?

using Oceananigans
using Oceananigans.AbstractOperations: KernelFunctionOperation
using Oceananigans.Fields: @compute

grid = RectilinearGrid(size = (1, 1, 1), extent = (1, 1, 1))
model = NonhydrostaticModel(; grid)
set!(model, u=2)
get_u(i, j, k, grid, velocities) = velocities.u[i, j, k]

f_op = KernelFunctionOperation{Face, Center, Center}(get_u, model.grid, velocities)
# f_op[1, 1, 1] works?

u_avg = Field(Average(u))
u_prime = u - u_avg
perturbation_velocities = (; u = u_prime)
f_perturbation_op = KernelFunctionOperation{Face, Center, Center}(get_u, model.grid, perturbation_velocities)
# f_perturbation_op[1, 1, 1] doesn't work?

If so, check that u_avg is actually computed.

[tomchor]

Yes, that's the core of it. I modified your example a bit to make it run and I'm posting it here for completeness:

using Oceananigans

grid = RectilinearGrid(size = (1, 1, 1), extent = (1, 1, 1))
model = NonhydrostaticModel(; grid)
set!(model, u=2)
get_u(i, j, k, grid, velocities) = velocities.u[i, j, k]

f_op = KernelFunctionOperation{Face, Center, Center}(get_u, model.grid, model.velocities)

u_avg = Field(Average(model.velocities.u))
u_prime = model.velocities.u - u_avg
perturbation_velocities = (; u = u_prime)
f_perturbation_op = KernelFunctionOperation{Face, Center, Center}(get_u, model.grid, perturbation_velocities)

Then f_op gives the correct answer, but f_perturbation_op does not.

And yes, I agree the core of the issue is that u_avg isn't computed before. If I first compute u_avg, then f_perturbation_op also produces the correct result.

But is that the expected behavior? I'd expect either u_avg to be computed before computing f_perturbation_op or for a warning/error to be thrown.

[glwagner]

Hmm, well it looks like we do compute arguments:

Oceananigans.jl/src/AbstractOperations/kernel_function_operation.jl

Line 75 in abd4479

compute_at!(κ::KernelFunctionOperation, time) = Tuple(compute_at!(d, time) for d in κ.arguments)

But if the arguments are themselves wrapped inside a NamedTuple --- or any other object --- then they won't be computed. In other words compute!(perturbation_velocities) does not compute the elements of perturbation_velocities.

We could add a method compute!(tup::Tuple) = Tuple(compute!(t) for t in tup) and also for NamedTuple. But I'm not sure this is the best API. Maybe it's better to require that arguments that need to be computed should be included directly as arguments. Seems like its up for debate.

PS it does seem to test this correctly then we need to further evaluate

f_perturbation = Field(f_perturbation_op)
compute!(f_perturbation)

[glwagner]

So basically we could add two methods

compute!(tup::Tuple) = Tuple(compute!(t) for t in tup)
compute!(nt::NamedTuple) = NamedTuple(name => compute!(nt[name]) for name in keys(nt))

as a convenience to auto-compute Fields that are embedded in tuples or named tuples that are arguments to KernelFunctionOperation (which I think is the main application that I can see)

[glwagner]

This seems reasonable. I'm not sure if there are performance implications. @navidcy @simone-silvestri any thoughts?

[simone-silvestri]

I like the idea adding a method for tuples and named tuples. It seems to be the behavior we want when compute!ing on collection of fields.

[glwagner]

Perhaps a more concise and extendable implementation would be

compute!(collection::Union{Tuple, NamedTuple}) = map(compute!, collection)

inspired by Adapt

[glwagner]

test

grid = RectilinearGrid(size = (1, 1, 1), extent = (1, 1, 1))
c = CenterField(grid)
set!(c, 1)
at_ijk(i, j, k, grid, a) = a[i, j, k]

one_c = Field(1 * c)
two_c = tuple(Field(2 * c))
ten_c = (; ten_c = Field(10 * c))

one_c_op = KernelFunctionOperation{Center, Center, Center}(at_ijk, grid, one_c)
two_c_op = KernelFunctionOperation{Center, Center, Center}(at_ijk, grid, two_c)
ten_c_op = KernelFunctionOperation{Center, Center, Center}(at_ijk, grid, ten_c)

one_c_field = Field(one_c_op)
two_c_field = Field(two_c_op)
ten_c_field = Field(ten_c_op)

compute!(one_c_field)
compute!(two_c_field)
compute!(ten_c_field)

@test all(interior(one_c) .== 1)
@test all(interior(two_c) .== 2)
@test all(interior(ten_c) .== 10)