Closed jbisits closed 3 months ago
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 https://github.com/tomchor/Oceanostics.jl/issues/35 and https://github.com/tomchor/Oceanostics.jl/issues/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.
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 KernelFunctionOperation
s for that since that calculation is global by definition, and KernelFunctionOperation
s 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 https://github.com/tomchor/Oceanostics.jl/issues/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.
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
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)
sortperm
should work on GPUs https://github.com/JuliaGPU/CUDA.jl/pull/1217
so the steps to make the above script fully gpu ready is:
perm
and fills the sorted_b_field
arraycumsum
function for GPU
everything should be GPU-ready then. 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
This design leverages the fourth type parameter of Oceananigans Field
:
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 operand
s here:
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!
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.
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 aKernelFunctionOperation
during a simulationThe volume integral of this can then be saved during a
simulation
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!