Skip to content

Convert ErtelPotentialVorticity and its thermal-wind counterpart into a Diagnostic #194

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 5 commits into
base: main
Choose a base branch
from

Conversation

tomchor
Copy link
Owner

@tomchor tomchor commented Feb 27, 2025

This PR serves two purposes.

  1. First, I was thinking that it might be good to condense the different versions of the same quantity (that differ only by an assumption for example) into one to save on code. So this PR starts by doing that with ErtelPotentialVorticity and ThermalWindPotentialVorticity.
  2. I'm also using this draft PR to explore the idea of creating an AbstractDiagnostic type, and again I'm trying this out for PV first.

This is what it's looking like for PV:

struct ErtelPotentialVorticity <: AbstractDiagnostic
operation
thermal_wind
tracer
end

The idea is still a bit abstract in my head, so I'd appreciate some feedback. But basically AbstractDiagnostic would allow other types of operation (some types that don't fit the scope of KernelFunctionOperation) to become the same type, so that users can interact with them in the same way. For example, the show method for ErtelPV when thermal_wind=true would be something more friendly:

julia> ErtelPotentialVorticity(model, thermal_wind=true)
ErtelPotentialVorticity in thermal wind balance, calculated with
KernelFunctionOperation at (Face, Face, Face)
├── grid: 3×3×6 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
├── kernel_function: ertel_potential_vorticity_in_thermal_wind_fff (generic function with 1 method)
└── arguments: ("3×3×6 Field{Face, Center, Center} on RectilinearGrid on CPU", "3×3×6 Field{Center, Face, Center} on RectilinearGrid on CPU", "3×3×6 Field{Center, Center, Center} on RectilinearGrid on CPU", "0", "0", "0")

and we'd also have access to the tracer used in the PV definition.

Although at the moment there isn't much practical difference, I can envision this being more important as we move towards reorganizing Oceanostics by budget equations (see #138 and #35), where it would make it easy to keep track of assumptions that get be inherited from budget equations to each of their terms (again, in the future).

@glwagner
Copy link
Collaborator

glwagner commented Feb 27, 2025

Another approach to this might implement the kernel functions as callable objects, eg

struct ErtelPVKernelFunction{thermal_wind, C}
    tracer :: C
end
const ErtelPV = KernelFunctionOperation{<:Any, <:Any, <:Any, <:Any, <:Any, <:ErtelPVKernelFunction}

(pv::ErtelPVKernelFunction{false}(i, j, k, grid, args...) = # whatever needs to be
(pv::ErtelPVKernelFunction{true}(i, j, k, grid, args...) = # whatever needs to be

Then you can provide constructors like

function ErtelPV(model, thermal_wind=true)
    kernel_function = ErtelPVKernelFunction{thermal_wind, typeof(tracer)}(tracer)
    return KernelFunctionOperation{Face, Face, Face}(grid, kernel_function, args...)
end

This is a strategy of extension rather than building a wrapper, which may be simpler and also easier for users to understand.

https://github.com/CliMA/Oceananigans.jl/blob/1034140b080e106c02e895296d7ee7556235f994/src/AbstractOperations/kernel_function_operation.jl#L3

@tomchor
Copy link
Owner Author

tomchor commented Mar 5, 2025

That's an interesting approach, but it has the downside that users don't easily know whether the PV that's being calculated is in thermal or not, right? They'd need to investigate the function that's used in the kernel.

@glwagner
Copy link
Collaborator

glwagner commented Mar 6, 2025

That's an interesting approach, but it has the downside that users don't easily know whether the PV that's being calculated is in thermal or not, right? They'd need to investigate the function that's used in the kernel.

Huh. Types are easier, not harder. Functions are a bit more annoying. Not only for the function itself but also the operation. For example:

const ThermalWindErtelPVKernelFunction = ErtelPVKernelFunction{true}
const ErtelPVOperation = KernelFunctionOperation{<:Any, <:Any, <:Any, <:Any, <:Any, ErtelPVKernelFunction}
const ThermalWindErtelPVOperation = KernelFunctionOperation{<:Any, <:Any, <:Any, <:Any, <:Any, ThermalWindErtelPVKernelFunction}

That said I don't know what tools you are providing the user to understand what operations are being used. But also I don't know if this matters; at least at the REPL we are using prettysummary (which fallback to summary), and this can be extended for ErtelPVKernelFunction:

https://github.com/CliMA/Oceananigans.jl/blob/7afb05786bb4ae834fb32bc213d1ad95bf7d7130/src/AbstractOperations/kernel_function_operation.jl#L92

In other words there is no downside in terms of functionality or user interface. To a first approximation they are the same, in the details; explicit types are easier to work with than function-types. There is a bit more flexibility with callable objects because embedded data/parameters do not need to be arguments.

There are downsides or we would always use callback objects instead of functions. The downside is the more complex implementation, requiring both a struct and a function. But I would also argue that the functionality being implemented is pretty complex already.

@tomchor
Copy link
Owner Author

tomchor commented Mar 28, 2025

Sorry, I'm not sure I fully understand the pros and cons here. In both cases you have to introduce new structs, no? I personally think that introducing a struct directly for ErtelPotentialVorticity is easier to understand, also because I think new AbstractDiagnostic type makes it clearer what it is and can be shared by all other diagnostics implemented here (if we decide to go this route). So I'm confused about what the downside is for this approach. Do you mind clarifying?

@glwagner
Copy link
Collaborator

glwagner commented Mar 28, 2025

I can't remember all the reasoning on the thread but one way to think about it is that types should be used to represent concepts and functions represent actions or verbs. Here it seems we are really dealing with types, ie you want an object that embodies ErtelPotentialVorticity. You also want to compute it, but that happens under the hood really, and the functional interface for that is implemented in Oceananigans.

We do get an object here in KernelFunctionOperation. But that is not specific. It'd be nice to have a specific type for the ErtelPotentialVorticityOperation. The way to do that given that KFO is on the "outside" is to use a type alias, which is not merely a definition; it also shows up in the REPL and when users request type info.

Technically you can achieve type alias using the fact that functions are types too but I think its cleaner and more explicit to use a struct.

@tomchor
Copy link
Owner Author

tomchor commented Mar 29, 2025

I can't remember all the reasoning on the thread but one way to think about it is that types should be used to represent concepts and functions represent actions or verbs. Here it seems we are really dealing with types, ie you want an object that embodies ErtelPotentialVorticity. You also want to compute it, but that happens under the hood really, and the functional interface for that is implemented in Oceananigans.

Agreed.

We do get an object here in KernelFunctionOperation. But that is not specific. It'd be nice to have a specific type for the ErtelPotentialVorticityOperation. The way to do that given that KFO is on the "outside" is to use a type alias, which is not merely a definition; it also shows up in the REPL and when users request type info.

Technically you can achieve type alias using the fact that functions are types too but I think its cleaner and more explicit to use a struct.

Sorry, that's where I get a bit confused. The way I understand it, you're proposing this construct:

struct ErtelPVKernelFunction{thermal_wind, C}
    tracer :: C
end
const ErtelPV = KernelFunctionOperation{<:Any, <:Any, <:Any, <:Any, <:Any, <:ErtelPVKernelFunction}

right? So that when the user calls ErtelPV, they're still given a KernelFunctionOperation, correct? (Albeit with a specific new EPV type for the kernel function.)

In my view it's better to define a struct for the operation itself that either wraps around or uses KernelFunctionOp under the hood. Something like

struct ErtelPotentialVorticity <: AbstractDiagnostic
    operation # In this case, this would be a KernelFunctionOperation
    thermal_wind :: Bool
    tracer
end

where in this case AbstractDiagnostic would be a new Oceanostics type.

Not sure what's best, but my idea is to do it in a way that all diagnostics provided by Oceanostics, whether they use KernelFunctionOperations or not, would use the same type and be implemented in a similar way, providing a unified interface. Does that make sense?

It may be that your suggestion is still a better way to do things, I'm not sure. I just wanna make sure we're on the same page before moving on.

@tomchor
Copy link
Owner Author

tomchor commented Mar 29, 2025

As another example, under what I'm proposing, KineticEnergyDissipationRate would become:

struct KineticEnergyDissipationRate <: AbstractDiagnostic
    operation
    isotropic_viscosity :: Bool
end

which would condense both IsotropicKineticEnergyDissipationRate and KineticEnergyDissipationRate, and sdo on for other diagnostics provided here.

@glwagner
Copy link
Collaborator

In my view it's better to define a struct for the operation itself that either wraps around or uses KernelFunctionOp under the hood. Something like

But then you need to reimplement the entire interface for KernelFunctionOperation, including its AbstractArray interface. And even if you do that for this case specifically, this approach does not scale. You might want 10 such types. This is a lot of re-implementation.

@glwagner
Copy link
Collaborator

So no, I disagree. I think conceptually, your idea represents a new "kernel function". And the best way to implement this in Julia is indeed by implementing a new kernel function, and then using the type parameter of KernelFunctionOperation to control dispatch where needed.

This is the idea:

struct OuterType{T}
    inner :: T
end

struct FunInnerType{V}
    value :: V
end

const FunOuterType = OuterType{<:FunInnerType}
FunOuterType(val) = OuterType(FunInnerType(val))

function Base.show(io::IO, fot::FunOuterType)
    V = typeof(fot.inner.value)
    print(io, string("FunOuterType{$V}(", fot.inner.value, ")"))
end

fun_one = FunOuterType(1)
julia> include("alias.jl")
FunOuterType{Int64}(1)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants