Making changes in Veros#
When contributing to Veros, please adhere to the following general guidelines:
Your first guide should be the surrounding Veros code. Look around, and be consistent with your modifications.
Unless you have a very good reason not to do so, please stick to the PEP8 style guide throughout your code. One exception we make in Veros is in regard to the maximum line length - since numerical operations can take up quite a lot of horizontal space, you may use longer lines if it increases readability.
In particular, please follow the PEP8 naming conventions, and use meaningful, telling names for your variables, functions, and classes. The variable name
stretching_factoris infinitely more meaningful than
k. This is especially important for settings and generic helper functions.
Document your functions using Google-style docstrings. This is especially important if you are implementing a user-facing API (such as a diagnostic, a setup, or tools that are meant to be called from setups).
flake8for linting and
blackfor code formatting. We automatically validate all changes to Veros through a pre-commit hook. To install it, run:
$ pip install pre-commit $ pre-commit install After this, black and flake8 will run automatically on every commit.
Distributed memory support#
By default, all core routines should support distributed execution via MPI.
In this case, every processor only operates on a chunk of the total data.
veros.variables.allocate(), you can make sure that allocated data always has the right shape.
Since none of the processes have access to the global data, you need to take special care during reductions (e.g.
sum) and accumulations (e.g.
cumsum) along horizontal dimensions.
Use functions from
veros.distributed.global_max()) where appropriate.
The dist_safe keyword#
If you are not comfortable writing code that is safe for distributed execution, you can use the
dist_safe keyword to
@veros_routine(dist_safe=False, local_variables=["temp"]) def my_function(state): # this function is now guaranteed to be executed on the main process vs = state.variables # since temp is declared as a local variable, we have access to all of the data vs.temp = update(vs.temp, at[2:-2, 2:-2], np.max(vs.temp)) # this would throw an error, since salt is not in local_variables # vs.salt = vs.salt * 0 # after execution, the updated contents of vs.temp are scattered to all processes, # and distributed execution continues
When encountering a
veros_routine that is marked as not safe for distributed execution (
dist_safe=False), Veros gathers all relevant data from the worker processes,
copies it to the main process, and executes the function there.
This ensures that you can write your code exactly as in the non-distributed case (but it comes with a performance penalty).
Running tests and benchmarks#
If you want to make sure that your changes did not break anything, you should run our test suite that compares the results of each subroutine to pyOM2. To do that, you will need to compile the Python interface of pyOM2 on your machine, and then point the testing suite to the library location, e.g. through:
$ pytest -v . --pyom2-lib /path/to/pyOM2/py_src/pyOM_code.so
from the main folder of the Veros repository.
If you deliberately introduced breaking changes, you can disable them during testing by prefixing them with:
from veros import runtime_settings if not runtime_settings.pyom_compatibility_mode: # your changes
Veros also provides automated benchmarks in a similar fashion. The benchmarks run some dummy problems with varying problem sizes and all available computational backends:
fortran (pyOM2), and
fortran-mpi (parallel pyOM2). For options and further information run:
$ python run_benchmarks.py --help
from the repository root.
If your changes to Veros turn out to have a negative effect on the runtime of the model, there several ways to investigate and solve performance problems:
Run your model with the
-v trace, and / or
--profile-modeoptions to get additional debugging output (such as timings for each time step, and a timing summary after the run has finished).
You should try and avoid explicit loops over arrays at all cost (but if you have to, you can use
veros.core.operators.for_loop(), which is reasonably efficient in JAX). You should always try to work on the whole array at once.
If you are still having trouble, don’t hesitate to ask for help (e.g. on GitHub).