Improved and simplified BinaryOperation with "stubborn" location inference

[glwagner]

This PR changes the structure of BinaryOperations and the way that locations are inferred when constructing them.

Before this PR, BinaryOperations stored three interpolation operators: two for each argument, and a third interpolation function applied after the binary operation was computed. The motivation for this design was to ensure that products of fields were always computed at their shared location, while allowing abstract operations to be built at user-specified locations. In other words, users might want something like

using Oceananigans
grid = RegularRectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1))
model = IncompressibleModel(architecture=CPU(), grid=grid)
u, v, w = model.velocities

uu = @at (Center, Center, Center) u * u

# output
BinaryOperation at (Center, Center, Center)
├── grid: RegularRectilinearGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
└── tree: 
    * at (Center, Center, Center) via ℑxᶜᵃᵃ
    ├── Field located at (Face, Center, Center)
    └── Field located at (Face, Center, Center)

which would compute u * u at (Face, Center, Center), and subsequently interpolate to cell centers. The object uu would then be defined at cell centers.

The main issue of this design in terms of functionality is that it produces expression trees that interpolate "too eagerly". A common example is a turbulent kinetic energy computation:

U = AveragedField(u, dims=(1, 2))
V = AveragedField(v, dims=(1, 2))

tke = @at (Center, Center, Center) ((u - U)^2 + (v - V)^2 + w^2)

Inspection of this object reveals that while u - U does not interpolate, the squaring (u - U)^2 would be performed at cell centers:

julia> tke.args[1].b
2

julia> tke.args[1].a
BinaryOperation at (Center, Center, Center)
├── grid: RegularRectilinearGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
└── tree: 
    - at (Center, Center, Center) via identity
    ├── Field located at (Face, Center, Center)
    └── AveragedField over dims=(1, 2) located at (⋅, ⋅, Center) of Field located at (Face, Center, Center)

This isn't what we want (usually): instead, we want both u - U and subsequent squaring in (u - U)^2 to be performed at (Face, Center, Center).

This PR addresses this issue by getting rid of the third interpolation in BinaryOperation, and changing how locations for BinaryOperation are inferred. Now, when locations are specified via @at, they are taken as a "suggestion" that only acts if the two elements of the binary operation have different concrete locations such that a difference needs to be resolved. In other cases (such as a binary operation between fields at common locations, or a binary operation between a field and a number), the location of the members of the operation is preserved. In this way, BinaryOperations are "stubborn".

We thus have results like

julia> *((Center, Center, Center), u, u)
BinaryOperation at (Face, Center, Center)
├── grid: RegularRectilinearGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
└── tree: 
    * at (Face, Center, Center)
    ├── Field located at (Face, Center, Center)
    └── Field located at (Face, Center, Center)

and

julia> tke = @at (Center, Center, Center) ((u - U)^2 + (v - V)^2 + w^2);

julia> tke.arg.args[1]
BinaryOperation at (Face, Center, Center)
├── grid: RegularRectilinearGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
└── tree: 
    ^ at (Face, Center, Center)
    ├── - at (Face, Center, Center)
    │   ├── Field located at (Face, Center, Center)
    │   └── AveragedField over dims=(1, 2) located at (⋅, ⋅, Center) of Field located at (Face, Center, Center)
    └── 2

An added benefit is that the BinaryOperation object is simpler. This could help compilation as well...

Of course, users do want to be able to specify the location of BinaryOperations for output. For this we change how @at works: now we wrap the entire user-prescribed expression in an interpolation operator that interpolates the result to the user-specified location. If the location of the underlying expression is already at the user-specified location, this is just an identity. But when the underlying operation is a "stubborn" BinaryOperator, we interpolate:

julia> uu = @at (Center, Center, Center) u * u
UnaryOperation at (Center, Center, Center)
├── grid: RegularRectilinearGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
└── tree: 
    identity at (Center, Center, Center) via ℑxᶜᵃᵃ
    └── * at (Face, Center, Center)
        ├── Field located at (Face, Center, Center)
        └── Field located at (Face, Center, Center)

julia> uu.arg
BinaryOperation at (Face, Center, Center)
├── grid: RegularRectilinearGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
└── tree: 
    * at (Face, Center, Center)
    ├── Field located at (Face, Center, Center)
    └── Field located at (Face, Center, Center)

This also means that operations between ReducedField and Field also end up at the right place, and we have no need to throw an error if a BinaryOperation has a Nothing location (in fact, this might be a useful abstraction for binary operations between ReducedField).

[ali-ramadhan]

🤞 this helps with compilation! I guess this is technically a breaking change? Sounds like it might be fine for LESbrary.jl though.

[glwagner]

It is very much breaking! And for the better. I'll bump version. But yeah, since it mostly changes complex abstract operations that previously haven't compiled on the GPU, most code isn't affected.

[tomchor]

Sorry for not following, but I don't see what breaks here for the end user! Could someone give a quick example/explanation?

[glwagner]

Sorry for not following, but I don't see what breaks here for the end user! Could someone give a quick example/explanation?

Apologies that my explanation was not clear. It's a breaking change because the same user input, such as

tke = @at (Center, Center, Center) ((u - U)^2 + (v - V)^2 + w^2) / 2

produces a different object after this update:

• Before this PR, u - U would be interpolated to (Center, Center, Center), and then the binary operation ^(2, u - U) would be calculated at (Center, Center, Center).
• After this PR, both u - U and ^(2, u - U) = (u - U)^2 are calculated at (Face, Center, Center). Interpolation is then performed to (Center, Center, Center) after the exponentiation to form the three-part sum.

Let me know if that makes sense or if another example would be useful.

[tomchor]

@glwagner Thanks! I actually had gotten that, but I was under the impression that the previous syntax wasn't gonna work. (Which would possibly break many scripts I have...)

Thanks for the explanation though!

(Btw, this PR was super timely. I was planning on opening an issue this week about the order of operations in abstract operations following our discussion out it the past week 👍)

[glwagner]

@glwagner Thanks! I actually had gotten that, but I was under the impression that the previous syntax wasn't gonna work. (Which would possibly break many scripts I have...)

Thanks for the explanation though!

(Btw, this PR was super timely. I was planning on opening an issue this week about the order of operations in abstract operations following our discussion out it the past week 👍)

This was a hard problem to solve. We'll see if this algorithm is sufficient for all cases. Other solutions are welcome.

[glwagner]

under the impression that the previous syntax wasn't gonna work

I believe this qualifies as a "breaking change" because any code that relied on consistent output from an abstract operation could, in theory, break (for example, a regression test that passes only when output remains constant). It does not change the API however, just the results.

[glwagner]

After some tweaks it does indeed look like

tke_ccc  = @at (Center, Center, Center) ((u - U)^2  + (v - V)^2 + w^2) / 2

compiles on the GPU! What a journey.