team-ocean/veros

The versatile ocean simulator, in pure Python, powered by JAX.

What it solves

Veros is a high-performance ocean simulator designed to replace legacy code and complicated setup workflows common in ocean modeling. It allows researchers to run everything from simple idealized toy models to high-resolution, global ocean simulations without needing to struggle with ancient programming environments.

How it works

Written in pure Python, the simulator uses a primitive equation ocean model. It supports two backends: a NumPy backend for small-scale problems and a high-performance JAX backend for CPU and GPU acceleration. The system is fully parallelized via MPI, enabling distributed execution across multiple nodes and multi-GPU architectures.

Who it’s for

Oceanographers and climate scientists who need a flexible, accessible, and extensible ocean modeling tool that can be scaled from a laptop to a high-performance computing cluster.

Highlights

  • Flexible Geometry: Supports a fully staggered 3-D C-grid geometry in Cartesian or pseudo-spherical coordinates.
  • Advanced Parameterizations: Includes isoneutral mixing, eddy-kinetic energy, turbulent kinetic energy, and internal wave energy parameterizations.
  • Built-in Diagnostics: Pre-implemented diagnostics for energy fluxes, variable time averages, and vertical overturning stream functions.
  • Easy Setup: Provides pre-configured idealized and realistic setups that are ready to run via simple Python scripts.

Related

  • Project
    google/brax

    A fast, fully differentiable physics engine written in JAX for large-scale rigid body simulation and reinforcement learning training.

  • Project
    dynamiqs/dynamiqs

    A GPU-accelerated and differentiable Python library for high-performance simulation of quantum systems, designed for large-scale problems and gradient-based parameter estimation.

  • Project
    google-deepmind/mujoco_playground

    A GPU-accelerated suite of robot learning environments built with MuJoCo MJX, designed for efficient reinforcement learning and sim-to-real research.

  • Project
    CliMA/Oceananigans.jl

    Oceananigans.jl is a Julia package for fast, GPU‑accelerated finite‑volume simulations of the incompressible Boussinesq equations on Cartesian or spherical‑shell grids. It provides ready‑made non‑hydrostatic and hydrostatic models, many advection schemes, turbulence closures, and differentiable‑programming support, all with a simple, switch‑able CPU/GPU interface.

  • Project
    dfm/tinygp

    A lightweight Gaussian Process library built on JAX that supports GPU acceleration and automatic differentiation.