Potential energy `KernelFunctionOperation`

[jbisits]

Thanks for this great package! This is more a discussion/question than an issue.

With Oceananigans.jl v0.89 a seawater_density function was added.
In my own work, one way I used this was to calculate the gravitational potential energy as a KernelFunctionOperation during a simulation

using Oceananigans: Models.seawater_density

@inline function potential_energy(model)

    grid = model.grid
    ρ = seawater_density(model)
    Z = Oceananigans.Models.model_geopotential_height(model)
    g = tuple(model.buoyancy.model.gravitational_acceleration)

    return KernelFunctionOperation{Center, Center, Center}(Ep, grid, ρ, Z, g)
end
@inline Ep(i, j, k, grid, ρ, Z, g) = g[1] * ρ[i, j, k] * Z[i, j, k]

The volume integral of this can then be saved during a simulation

Eₚ = potential_energy(model)
∫Eₚ = Integral(Eₚ)

This is a bit rough around the edges but if it is something that you think would fit in to this package let me know and I can tidy it up and add some tests. No problem at all if you think this package is not the place for it though!

My goal with this was to calculate the LHS terms in the energy budget for a closed system using equations (21), (22), (23) from Winters et al. (1995), your KineticEnergy function and the above potential energy function. One thing I could not quite get to work was calculating a background state for the potential energy (equation (22)) which depends on a resorting and scalar indexing on GPU's. Being able to calculate the energy budget like in Winters et al. (1995) would be a neat functionality though!

[tomchor]

I think that's a great addition! If you submit a PR for it I'm happy to review it for you.

Your function is looks great, it just needs a couple more bells and whistles that you can see in the other functions (like a way to pass (and check) node location and a way to input a user-defined density, for example), but that's easy.

In line with #35 and #138, I'd create a file src/DensityEquationTerms.jl and include it there, but it could also go in https://github.com/tomchor/Oceanostics.jl/blob/main/src/FlowDiagnostics.jl for now.

[tomchor]

One thing I could not quite get to work was calculating a background state for the potential energy (equation (22)) which depends on a resorting and scalar indexing on GPU's. Being able to calculate the energy budget like in Winters et al. (1995) would be a neat functionality though!

That's a great goal, and one that I've tried tacking in the past but failed. I think we can't use KernelFunctionOperations for that since that calculation is global by definition, and KernelFunctionOperations are designed to take a local function. @glwagner and @simone-silvestri know way more about this than me and can probably clarify/correct me, but there's some relevant discussion in #150.

As a side note, in the distant past I've implemented the sorting for the BPE calculation with the snippet below

    function flattenedsort(A, dim_order::Union{Tuple, AbstractVector})
        return reshape(sort(Array(permutedims(A, dim_order)[:])), (grid.Nx, grid.Ny, grid.Nz))
    end

    function sort_b(model; average=false)
        b = model.tracers.b
        sorted_B = flattenedsort(interior(b), [3,2,1])
        if !average
            return sorted_B
        else
            return dropdims(mean(sorted_B, dims=(1,2)), dims=(1,2))
        end
    end
    mean_sort_b = (mod)->sort_b(mod; average=true)

But that's not very modular so I never implemented it in Oceanostics. (BTW I think these days it's not necessary to convert to Array anymore.)

Another side note: the resorting strategy only works for regular grids. It'd be good to implement Equation (11) of Winters et al. (1995) directly, which can be (I think trivially) extended to stretched grids, but that entails a volume integration for each individual point in the domain, which makes it an $O(n^6)$ operation for 3D simulations (oof), so it might be impossible to be made performant.

[simone-silvestri]

It probably would be difficult to do the resorting during the time-stepping.
These are the type of operations you want to leave for postprocessing.
If you need a script to calculate the resorted state, this is what I used

function calculate_z★_diagnostics(b::Field)

    vol = VolumeField(b.grid)
    z★  = similar(b)

    total_area = sum(AreaField(b.grid))

    calculate_z★!(z★, b, vol, total_area)
        
    return z★
end

function calculate_z★!(z★::Field, b::Field, vol, total_area)
    grid = b.grid
    arch = architecture(grid)

    b_arr = Array(interior(b))[:]
    v_arr = Array(interior(vol))[:]

    perm           = sortperm(b_arr)
    sorted_b_field = b_arr[perm]
    sorted_v_field = v_arr[perm]
    integrated_v   = cumsum(sorted_v_field)    

    sorted_b_field = arch_array(arch, sorted_b_field)
    integrated_v   = arch_array(arch, integrated_v)

    launch!(arch, grid, :xyz, _calculate_z★, z★, b, sorted_b_field, integrated_v)
    
    z★ ./= total_area

    return nothing
end

@kernel function _calculate_z★(z★, b, b_sorted, integrated_v)
    i, j, k = @index(Global, NTuple)
    bl  = b[i, j, k]
    i₁  = searchsortedfirst(b_sorted, bl)
    z★[i, j, k] = integrated_v[i₁] 
end

you can probably take this and come up with a GPU-friendly KernelFunctionOperation

[simone-silvestri]

where

function MetricField(loc, grid, metric)
    metric_operation = GridMetricOperation(loc, metric, grid)
    metric_field = Field(metric_operation)
    return compute!(metric_field)
end

VolumeField(grid, loc=(Center, Center, Center)) = MetricField(loc, grid, Oceananigans.AbstractOperations.volume)
  AreaField(grid, loc=(Center, Center, Nothing)) = MetricField(loc, grid, Oceananigans.AbstractOperations.Az)

[simone-silvestri]

sortperm should work on GPUs JuliaGPU/CUDA.jl#1217
so the steps to make the above script fully gpu ready is:

• a kernel to flatten 3D fields
• a permute kernel that takes perm and fills the sorted_b_field array
• a cumsum function for GPU
  everything should be GPU-ready then.

[glwagner]

It probably would be difficult to do the resorting during the time-stepping.

I disagree with @simone-silvestri here. I think you definitely can implement the resorting during time-stepping. Moreover, there are major advantages to this for very large simulations and I think this could be an important feature.

One thing I could not quite get to work was calculating a background state for the potential energy (equation (22)) which depends on a resorting and scalar indexing on GPU's.

You should be able to sort on the GPU. The key is that you will need to implement the background potential energy feature using more than one field. You will need to allocate an intermediate field to store the background state. Then you can use this background state to compute the potential energy.

I think the best way to do this is to design a custom operand that can be used with Field. Then you can define a special Field type:

const BackgroundPotentialEnergyField = Field{Center, Center, Center, <:BackgroundPotentialEnergyOperand}

Next you have to define a custom compute!

import Oceananigans.Fields: compute!

function compute!(bpe::BackgroundPotentialEnergyField, time=nothing)
    # the hard part
end

[glwagner]

This design leverages the fourth type parameter of Oceananigans Field:

https://github.com/CliMA/Oceananigans.jl/blob/427c92fac963ee0883fa2766df15ba6a48e62736/src/Fields/field.jl#L17-L23

which represents the operand. We only support AbstractOperations as operands internally in Oceananigans, but the intent of this design was precisely to permit users to build custom operands with custom compute! that would seamlessly integrate with the rest of Oceananigans AbstractOperations infrastructure. I think you should use it!

For reference, compute! is defined for Field with non-trivial operands here:

https://github.com/CliMA/Oceananigans.jl/blob/427c92fac963ee0883fa2766df15ba6a48e62736/src/AbstractOperations/computed_field.jl#L65-L75

[jbisits]

I think that's a great addition! If you submit a PR for it I'm happy to review it for you.

Your function is looks great, it just needs a couple more bells and whistles that you can see in the other functions (like a way to pass (and check) node location and a way to input a user-defined density, for example), but that's easy.

In line with #35 and #138, I'd create a file src/DensityEquationTerms.jl and include it there, but it could also go in https://github.com/tomchor/Oceanostics.jl/blob/main/src/FlowDiagnostics.jl for now.

Great, I will get to it in the coming days! Depending on how many changes there are it may be better to open another PR with the background potential energy state but I think with the advice above it looks achievable!

[glwagner]

Yes for sure, I would focus on one thing at a time. First task is to define a new field that computes the background state. I think it would be reasonable to do this in two PRs.