veros_routine and veros_kernel#
- veros.veros_routine(function=None, *, dist_safe=True, local_variables=())#
This decorator should be applied to all functions that access the Veros state object (even when subclassing
The first argument to the decorated function must be a VerosState instance.
Veros routines cannot return anything. All changes must be applied to the passed state object.
dist_safe (bool) – If set to False, all variables specified in local_variables are synced to the root process before execution and synced back after. This means that the routine will only be executed on rank 0. Has no effect in non-distributed contexts.
local_variables (Tuple[str]) – List of variable names to be synced if dist_safe=False. This must include all variables retrieved from the state object throughout the routine (inputs and outputs).
>>> from veros import VerosSetup, veros_routine >>> >>> class MyModel(VerosSetup): >>> @veros_routine >>> def set_topography(self, state): >>> vs = state.variables >>> settings = state.settings >>> vs.kbot = npx.random.randint(0, settings.nz, size=vs.kbot.shape)
- veros.veros_kernel(function=None, *, static_args=())#
Decorator that marks a function as a kernel that can be JIT compiled if supported by the backend.
Kernels cannot modify the Veros state object. Instead, all modifications have to be returned explicitly.
static_args (Tuple[str]) – Names of kernel arguments that should be static.
>>> from veros import veros_kernel, KernelOutput >>> >>> @veros_kernel >>> def double_psi(state): >>> vs = state.variables >>> vs.psi = 2 * vs.psi >>> return KernelOutput(psi=vs.psi)