NORiOceanParameterization.jl

Neural ODEs ("NO") and Richardson Number ("Ri")-Based Ocean Parameterization

A ML-augmented ocean turbulence parameterization that combines physics-based Richardson number closures with neural networks to improve vertical mixing predictions in realistic ocean models. Designed for integration with Oceananigans.jl.

Trained entirely end-to-end (a posteriori) with fully-differentiable ODE solvers and Ensemble Kalman Inversion.

Overview

NORiOceanParameterization.jl provides hybrid ocean parameterizations that:

• Use physics-based Richardson number closures to model vertical diffusivity in different turbulent regimes (convective vs. shear-driven)
• Apply neural networks to learn entrainment processes during convection, acting as a residual flux
• Provide an end-to-end training framework using fully differentiable-(Neural ODEs) and ensemble-based (Ensemble Kalman Inversion) methods
• Provide calibrated closures trained on large-eddy simulations (LES) generated with Oceananigans
• Integrate seamlessly with Oceananigans.jl for use in ocean simulations on CPUs and GPUs

Repository Structure

The repository is organized into several directories:

Core Package

• src/: Core package source code for training and data processing
  • Closures/: Base closure implementations used during training
  • DataWrangling/: Tools for loading and processing LES data
  • ODEOperations/: ODE solvers and operations used during training
  • Operators/: Differential operators for the parameterizations
  • TrainingOperations/: Training utilities, neural differential equation and loss function definition
  • Plotting/: Visualization utilities
  • Utils/: Model I/O utilities (saving/loading NN weights and parameters)

Inference Package

• NORiImplementation/: Standalone package for using trained closures in Oceananigans
  • src/NORiBaseVerticalDiffusivity.jl: Richardson number-based vertical diffusivity closure
  • src/NORiNNFluxClosure.jl: Neural network closure for residual flux corrections
  • src/closure_builder.jl: Helper functions for combining closures correctly
  • Uses latest Oceananigans (v0.104+) and Lux (v1.x)

Pre-Trained Models

• calibrated_parameters/: Pre-trained parameters for base Richardson number closure and final neural network weights

Training and Validation

• training/: Complete training pipeline scripts
  • run_LES_boundary_layer.jl: Generate LES training data with Oceananigans.jl
  • train_NDE.jl: Neural Differential Equations training for NN closure
  • train_localbaseclosure_convection.jl: EKI calibration for convective regime of the base closure
  • train_localbaseclosure_shear.jl: EKI calibration for shear regime of the base closure
  • Pins old dependencies (Enzyme v0.11.20, Lux v0.5.31) for compatibility

Inference and Examples

• inference/: Example scripts for using NORi closures in Oceananigans models

  • column_model_nori_closures_example.jl: 1D column model example
  • doublegyre_nori_closures_example.jl: 3D double-gyre circulation example

• experiments/: Validation and benchmark experiments

  • Column model simulations comparing NORi against LES, CATKE, and k-epsilon closures
  • Scripts for timestep dependence studies and long integration tests

Visualization

• figure_scripts/: Scripts to generate all publication figures
  • Training results visualization (NDE losses, training metrics)
  • Validation plots comparing NORi against LES, CATKE, and k-epsilon closures
  • 3D LES field visualizations
  • Double-gyre circulation analysis plots
  • Data required to plot the figures are hosted on this Zenodo data companion
  • Data loading from Zenodo is handled automatically with DataDeps.jl

Key Features

Hybrid Physics-ML Approach

NORiBaseVerticalDiffusivity: A Richardson number-based vertical diffusivity closure with

• Separate parameterizations for convective (Ri < 0) and shear (Ri > 0) regimes
• Distinct Prandtl numbers for convective and shear regimes
• Pre-trained parameters loaded automatically from calibration

NORiNNFluxClosure: Neural network closure for residual flux corrections with

• Separate networks for temperature and salinity fluxes
• 21 input features: Richardson number, temperature/salinity/density gradients at 5 vertical levels, surface buoyancy flux
• Only active near base of the mixed layer (where entrainment mixing happens)
• Pre-trained Lux.jl models included

This hybrid approach combines interpretable physical closures with highly-expressive neural networks NORiBaseVerticalDiffusivity + NORiNNFluxClosure:

Designed for Realistic Ocean Equation of State (TEOS-10)

NORi is specifically trained on and designed for the nonlinear equation of state based on TEOS-10 (Roquet et al., 2015), as implemented in SeawaterPolynomials.jl. This becomes important in cold regions (e.g. Southern Ocean) where the seawater's equation of state is highly nonlinear.

Complete ML Training Pipeline

The repository includes a full training infrastructure:

• Neural Differential Equations (NDEs) for fully-differentiable neural network calibration over ODE solvers
• Ensemble Kalman Inversion (EKI) for base closure calibration
• Free and Open Source training dataset SOBLLES: Salty Ocean Boundary Layer Large-Eddy Simulations covering a wide range of realistic ocean scenarios
• Cutting-edge automatic differentiation via Enzyme.jl
• Highly efficient gradient-free methods using EnsembleKalmanProcesses.jl
• Multi-initial condition training with adaptive loss weighting

Important Caveats (for now!)

Package Compatibility

This repository is organized into separate environments to handle conflicting dependency versions:

• NORiImplementation/ — For running pre-trained closures in Oceananigans (uses latest Lux and Oceananigans v0.104+)
• training/ — For training new closures (pins Enzyme v0.11.20 and Lux v0.5.31)

Using other versions may lead to potential software issues, especially for the end-to-end training with Enzyme autodiff.

Using NORi in production

NORi can be readily used in the HydrostaticFreeSurfaceModel in Oceananigans. Details and examples can be found in the NORiImplementation/ and inference/ directories. However, it currently only supports a grid with uniform vertical spacing of Δz = 8 m. Support for nonuniform grids in the vertical will be coming soon!

Installation

For inference (using pre-trained closures)

From the repository root:

julia --project=NORiImplementation -e 'using Pkg; Pkg.develop(path="."); Pkg.instantiate()'

For training

The training scripts auto-bootstrap the environment on first run. From the repository root:

julia training/train_NDE.jl

Or set up manually:

julia --project=training -e 'using Pkg; Pkg.develop(path="."); Pkg.instantiate()'

Quick Start

Using Pre-Trained Closures

using NORiImplementation
using Oceananigans
using Oceananigans.Units

# Create closures with trained parameters (automatically loaded)
closures = NORiClosureWithNN()  # For CPU

# For GPU, use:
# using CUDA
# closures = NORiClosureWithNN(arch=GPU())

# Set up a simple ocean column
grid = RectilinearGrid(size = (32, 32, 32),
                       x = (0, 1),
                       y = (0, 1),
                       z = (-128, 0),
                       topology = (Periodic, Periodic, Bounded))

model = HydrostaticFreeSurfaceModel(grid = grid,
                                    closure = closures,
                                    tracers = (:T, :S),
                                    buoyancy = SeawaterBuoyancy(equation_of_state=TEOS10.TEOS10EquationOfState()))

# Run simulation
simulation = Simulation(model, Δt=5minutes, stop_time=1day)
run!(simulation)

Training Your Own Closures

Training scripts live in training/ and use a separate environment with pinned dependency versions (Enzyme v0.11.20, Lux v0.5.31). Each script auto-activates this environment and dev-installs the core package on first run.

# Run NDE training
julia training/train_NDE.jl

# Or run with explicit project activation
julia --project=training training/train_NDE.jl

License

This project is licensed under the MIT License.