Operations and averaging

Fields are more than mere vessels for data. They come equipped with a powerful infrastructure for assembling expression trees that represent arithmetic and discrete calculus operations. We also supply a system for computing reductions (averages, integrals, and cumulative integrals) of Fields. This infrastructure can be used to construct initial conditions, set up diagnostic calculations that are performed and saved to disk while a simulation runs, and also for post-processing.

We start by constructing a CenterField on a simple grid,

using Oceananigans

grid = RectilinearGrid(topology = (Periodic, Flat, Bounded),
                       size = (4, 4),
                       x = (0, 2π),
                       z = (-4, 0))

c = CenterField(grid)

periodic_but_decaying(x, z) = sin(x) * exp(z)
set!(c, periodic_but_decaying)

# output
4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU
├── grid: 4×1×4 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 3×0×3 halo
├── boundary conditions: FieldBoundaryConditions
│   └── west: Periodic, east: Periodic, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 10×1×10 OffsetArray(::Array{Float64, 3}, -2:7, 1:1, -2:7) with eltype Float64 with indices -2:7×1:1×-2:7
    └── max=0.428882, min=-0.428882, mean=1.04083e-17

An AbstractOperation (or operation for short) differs from a Field in that only represents a computation. Unlike Fields, AbstractOperations do not have any data, and are associated only with minimal memory allocation. AbstractOperations are generated by inflicting Fields with ordinary arithmetic expressions,

two_c = 2 * c

# output
BinaryOperation at (Center, Center, Center)
├── grid: 4×1×4 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 3×0×3 halo
└── tree:
    * at (Center, Center, Center)
    ├── 2
    └── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU

and even by chaining expressions together, which may themselves include AbstractOperations,

quadratic = c^2 + two_c + 1

# output
BinaryOperation at (Center, Center, Center)
├── grid: 4×1×4 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 3×0×3 halo
└── tree:
    + at (Center, Center, Center)
    ├── + at (Center, Center, Center)
    │   ├── ^ at (Center, Center, Center)
    │   │   ├── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU
    │   │   └── 2
    │   └── * at (Center, Center, Center)
        │   ├── 2
        │   └── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU
    └── 1

Like Fields, AbstractOperations have a location and a grid. In addition to BinaryOperations like the kind above, UnaryOperations and MultiaryOperations are also supported,

cos_c = cos(c)

# output
UnaryOperation at (Center, Center, Center)
├── grid: 4×1×4 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 3×0×3 halo
└── tree:
    cos at (Center, Center, Center) via identity
    └── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU

four_c = c + c + c + c

# output
MultiaryOperation at (Center, Center, Center)
├── grid: 4×1×4 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 3×0×3 halo
└── tree:
    + at (Center, Center, Center)
    ├── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU
    ├── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU
    ├── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU
    └── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU

UnaryOperation, BinaryOperation and MultiaryOperation all have both an "operator", and between 1 and many. Last, and definitely not least, the fourth flavor of AbstractOperation represents a derivative,

dx_c = ∂x(c)

# output
Derivative at (Face, Center, Center)
├── grid: 4×1×4 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 3×0×3 halo
└── tree:
    ∂xᶠᶜᶜ at (Face, Center, Center) via identity
    └── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU

Like Fields, AbstractOperations are evaluated by indexing,

@show c[1, 1, 1]
@show quadratic[1, 1, 1]
nothing

# output
c[1, 1, 1] = 0.02135277459201165
quadratic[1, 1, 1] = 1.0431614901668005

Also like Fields, AbstractOperations have a location. For example,

using Oceananigans: location

@show location(c)
@show location(dx_c)
nothing

# output
location(c) = (Center, Center, Center)
location(dx_c) = (Face, Center, Center)

Notice that the location of dx_c is shifted in x relative to c. Likewise, y-derivatives are shifted in y and z-derivatives are shifted in z.

Locations and interpolation

Reconstruction of Fields from one location to another is intrinsic to arithmetic on the staggered grid. Consider the magnitude of the gradient of c:

∇c² = ∂x(c)^2 + ∂z(c)^2

# output
BinaryOperation at (Face, Center, Center)
├── grid: 4×1×4 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 3×0×3 halo
└── tree:
    + at (Face, Center, Center)
    ├── ^ at (Face, Center, Center)
    │   ├── ∂xᶠᶜᶜ at (Face, Center, Center) via identity
    │   │   └── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU
    │   └── 2
    └── ^ at (Center, Center, Face)
        ├── ∂zᶜᶜᶠ at (Center, Center, Face) via identity
        │   └── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU
        └── 2

Because ∂x(c)^2 is located at (Face, Center, Center) and ∂z(c)^2 is located at (Center, Center, Face), a decision has to be made to compute ∇c². By default, AbstractOperations are reconstructed at the location of the first object in the expression. So

∇c²_ccf = ∂z(c)^2 + ∂x(c)^2

# output
BinaryOperation at (Center, Center, Face)
├── grid: 4×1×4 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 3×0×3 halo
└── tree:
    + at (Center, Center, Face)
    ├── ^ at (Center, Center, Face)
    │   ├── ∂zᶜᶜᶠ at (Center, Center, Face) via identity
    │   │   └── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU
    │   └── 2
    └── ^ at (Face, Center, Center)
        ├── ∂xᶠᶜᶜ at (Face, Center, Center) via identity
        │   └── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU
        └── 2

ends up at (Center, Center, Face). To control the location of an operation we use the macro @at,

∇c²_ccc = @at (Center, Center, Center) ∂x(c)^2 + ∂z(c)^2

# output
BinaryOperation at (Center, Center, Center)
├── grid: 4×1×4 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 3×0×3 halo
└── tree:
    + at (Center, Center, Center)
    ├── ^ at (Center, Center, Center)
    │   ├── ∂xᶠᶜᶜ at (Center, Center, Center) via ℑxᶜᵃᵃ
    │   │   └── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU
    │   └── 2
    └── ^ at (Center, Center, Center)
        ├── ∂zᶜᶜᶠ at (Center, Center, Center) via ℑzᵃᵃᶜ
        │   └── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU
        └── 2