veros_routine and veros_kernel¶

veros_routine¶

veros.veros_routine(function=None, *, dist_safe=True, local_variables=())[source]¶

Note

This decorator should be applied to all functions that access the Veros state object (even when subclassing veros.VerosSetup).

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.

Parameters:

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).

Example

>>> 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_kernel¶

veros.veros_kernel(function=None, *, static_args=())[source]¶

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.

Parameters:: static_args (Tuple[str]) – Names of kernel arguments that should be static.

Example

>>> 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)

KernelOutput¶

veros.KernelOutput(**kwargs)¶