-
Notifications
You must be signed in to change notification settings - Fork 186
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
Abstraction for "equation" for performance, code clarity, memory footprint reduction, and powerful user interface #259
Comments
A simpler (but less general) solution than what is proposed above is simply to introduce a new struct: struct BoussinesqEquations{A, B, C, D, F, P} <: AbstractEquation
advection :: A
buoyancy :: B
coriolis :: C
diffusivity :: D
forcing :: F
pressure :: P
end This will reduce the number of fields in Less general than a tuple of tendency kernels, but probably sufficient for anything we want to do in the forseeable future. We can then have a function eg Could also have an independent struct for tracer equations (or tuple of structs for each tracer). By the way, the more I think about it, the more I think that 'closure' should actually be called 'diffusivity', for clarity... |
I like this idea, it would definitely work towards having an equations abstraction. Although I might be wary of doing this while we only support one equation set. Might make sense to consider an |
Yesterday, @ali-ramadhan and I discussed a new idea for improving the abstraction of equations. The idea is to abstract a Each This would greatly simplify the time-stepping routines, which currently involve long function signatures. There'd be no need for 'unpacking', because each term in the equation would perform unpacking upon instantiation. It would also probably be easier for users to extend / add terms to an equation. As an example, we can consider a simple implementation struct AdvectiveFlux{S, D, C, G}
u :: D
v :: D
w :: D
c :: C
grid :: G
scheme :: S
end
getindex(adv::AdvectiveFlux{Centered}, i, j, k) = # centered advective flux calculation
struct IsotropicDiffusiveFlux{N, D, G}
ν :: N
ψ :: D
grid :: G
end
getindex(diff::IsotropicDiffusiveFlux, i, j, k) = # calculates flux due to isotropic diffusion by ν
struct RightHandSide{F, V}
fluxes :: F
volume_terms :: V
end
advection = AdvectiveFlux(velocities..., tracers.c)
diffusion = IsotropicDiffusiveFlux(ν, tracers.c)
tracer_rhs = RightHandSide((advection, diffusion), nothing) We'd have functions that look something like function x_flux_divergence(i, j, k, grid, fluxes...)
incoming_flux = add_fluxes(i, j, k, grid, fluxes...)
outgoing_flux = add_fluxes(i+1, j, k, grid, fluxes...)
return (incoming_flux - outgoing_flux) * grid.Ax / grid.V
end ... for example. Obviously questions of performance are paramount, though in the case that everything is inlined I think there is hope. A downside of this approach is that we can't use shared memory stencils on the GPU. Shared memory stencils on the GPU require functions for all terms that avoid carrying around internal references to data (since we need to be able to pass them references to shared memory blocks, rather than references to global memory. This is a fair rewrite of the code internals. For example, each term in the would get it's own type. It's worth brainstorming ways to implement such abstraction incrementally so we might avoid possibly time-consuming total-code-demolishment. This is mostly food for thought at this point. I don't think we should take action without substantial consideration. |
I don't think we want this anymore: we currently use |
I'm thinking about the future and hoping to start a discussion about the future of our equation abstraction system.
The problem
We need an abstraction for the concept of an equation, so that we can make the model as performant and lightweight as possible for a given use case. For example, I think we should require that
The memory footprint of our model is no larger than it needs to be for a given problem (no 'extra' allocation of memory for unused tracers, for hydrostatic pressure when running in non hydrostatic mode, etc).
We do not perform unnecessary floating-point computations or indexing into arrays (the latter is especially important in GPU code) for unused tracers or hydrostatic pressure fields.
Equations are constructed / specified clearly and concisely (both in source code and user scripts).
Users can specify arbitrary types of forcing, including numbers, arrays, or functions (solving Allow forcing to be an array. #110).
We can support arbitrary tracers with various features, such as sinking/rising velocities, or reaction systems for biological/chemical tracer systems.
When I have talked to various people about this, there was a concern that this system would be 'inelegant' or 'complex'. However I believe an equation abstraction system provides the opposite: with an abstraction system, equations are 'written down' in some logical place (like a file
equations.jl
in thesrc
directory where they can be easily read and modified, rather than buried inside a time-stepping loop. Correspondly, our time-stepping code becomes shorter and more concise. Using multiple dispatch correctly, we avoid theinfinite if-statement
problem. This abstraction may also make the code more modular such that we move closer to supporting multiple time-steppers.Below I provide one example of an implementation that would solve some of the problems I listed. However, this is not the only solution, and I think we should expend some intellectual effort and have a discussion about what the best solution might be, so that we design something that is nice, easy to extend, performant, and powerful.
A list of kernel equations in a named tuple
The simplest solution for this abstraction is probably just to add new fields to
Model
(model.equations.velocities
andmodel.equations.tracers
) that are named tuples of kernel equations. An example of how this might work while demonstrating hierarchical multiple dispatch is:Equations as new types
Even more flexibility / clarity might be provided by a callable
Equation
type implementation with appropriate type parameters (Linear
,Hydrostatic
, etc) which would avoid the need forif
statements as on the last line of the code above. I'm not sure exactly what that would look like. The design problem for a newEquation
type does seem a bit more difficult and would require care --- thinking about how we would implement a nonlinear biological model might provide a nice use case to help use make our ideas concrete.Hopefully this provides food for thought and also motivates us to think carefully about our design decisions.
The text was updated successfully, but these errors were encountered: