NORi: An ML-Augmented Ocean Boundary Layer Parameterization

NORi is a novel, physics-based machine learning parameterization that combines neural ordinary differential equations with a Richardson number-dependent closure to accurately and stably simulate ocean boundary layer turbulence and entrainment dynamics in climate models, outperforming traditional methods while requiring minimal training data and ensuring long-term numerical stability.

The Gist

The Big Problem: The Ocean is Too Small to See

Imagine you are trying to predict the weather for the entire planet. You have a giant map, but it's made of huge tiles, each 10 kilometers wide. On this map, you can see big storms and ocean currents. However, the ocean is full of tiny, chaotic swirls and mixing processes that happen in the top few hundred meters—like the foam on a wave or the way cold air chills the surface water. These tiny swirls are too small to fit on your giant map tiles.

In climate science, we call these tiny processes "subgrid-scale" processes. To make our big maps work, scientists have to guess what these tiny swirls are doing. They use "parameterizations"—which are basically simplified rulebooks or formulas that say, "When the wind blows this hard, the water mixes this much."

The Old Rulebooks vs. The New Hybrid

For decades, scientists have used physics-based rulebooks. Think of these like a manual written by a strict engineer. They are based on known laws (like how heat moves from hot to cold).

• The Good: They are fast and stable.
• The Bad: They miss some tricky physics. Specifically, they struggle to explain entrainment.

What is Entrainment? Imagine a pot of soup on the stove. If you cool the top, the cold soup sinks, but it doesn't just stop at the surface. It dives down like a plunger, dragging the warm soup from below up into the cold layer. This "plunging" action is non-local; it happens at the bottom of the mixed layer but is caused by what's happening at the top. Old rulebooks are like a recipe that only knows how to stir the pot gently; they don't know how to simulate that deep, plunging dive.

Enter NORi: The "Smart Assistant"

The authors created a new tool called NORi (Neural Ordinary differential equations Richardson number). Think of NORi not as a replacement for the old rulebook, but as a smart assistant attached to it.

1. The Base (The Engineer): NORi starts with a simple, physics-based formula (the "Base Closure"). This part handles the easy stuff: the gentle stirring caused by wind and local temperature differences. It's like the engine of a car—it does the heavy lifting.
2. The Neural Network (The AI Co-Pilot): The authors added a small, highly expressive Artificial Intelligence (AI) brain. This AI doesn't try to learn the whole ocean from scratch. Instead, it only learns the missing piece: the deep plunging (entrainment) that the engineer's formula misses.

The Analogy: Imagine you are driving a car (the ocean model). The engine (physics) gets you moving. But sometimes, you need to navigate a tricky, winding mountain road (entrainment). The AI is a co-pilot who only takes the wheel when the road gets twisty, guiding the car through the turns that the engine alone would miss.

How They Trained It: Learning from the "Future"

Usually, when you train an AI, you show it a snapshot and ask, "What is the answer right now?" (e.g., "Here is the wind speed; what is the mixing flux?"). The authors found this made the AI unstable. It was like teaching a student to pass a test by memorizing answers to single questions, but when they took the final exam (running the model for years), they failed because they didn't understand the flow of the story.

Instead, they used A Posteriori Training (learning from the outcome).

• The Method: They ran a super-detailed, high-resolution simulation (the "Ground Truth") that captured every tiny swirl. Then, they let their simple NORi model run alongside it.
• The Lesson: They didn't ask the AI to match the flux at one specific second. They asked, "After running for 2 days, did your temperature and salinity match the high-resolution simulation?"
• The Result: The AI learned to adjust its behavior over time to ensure the entire journey was correct, not just a single step. This is like teaching a student by saying, "Don't just get the right answer for question 1; make sure you can solve the whole story problem correctly."

Why It's a Game-Changer

The paper claims NORi solves three big problems at once:

1. Accuracy: In tests, NORi matched the high-resolution "ground truth" simulations much better than the old rulebooks, especially when the ocean was cooling and plunging (convection). It performed just as well as the most complex, expensive models (like the $k-\epsilon$ model) but was much simpler.
2. Stability: This is the biggest win. Many AI models crash or blow up when run for a long time (like a video game character glitching out after 10 hours). Because NORi was trained to keep the whole timeline stable, it ran for 100 years in a simulation without crashing, even though it was only trained on 2-day snapshots.
3. Speed: NORi is a "zero-equation" model, meaning it doesn't need to solve extra complex math equations like the heavy-duty models do. It can run with much larger time steps (up to 1 hour), making it much faster for global climate simulations.

The Real-World Test

The authors tested NORi against real-world data from Ocean Weather Station Papa in the Pacific Ocean. They ran the model for 120 days (fall to winter) using real weather data.

• The Result: NORi predicted the temperature and salinity of the ocean almost perfectly, matching the observations just as well as the state-of-the-art models.
• The Surprise: Even though NORi was trained on idealized, constant weather, it handled the messy, changing real-world weather perfectly. It knew when to "turn on" its AI brain (during strong cooling) and when to let the simple physics engine take over (during calm winds).

Summary

NORi is a new way to model the ocean's surface mixing. Instead of trying to build a giant, complex AI to replace physics, the authors built a simple physics engine and gave it a small, smart AI assistant to fix its blind spots. By training this assistant to care about the long-term journey rather than just the immediate moment, they created a model that is fast, stable for a century, and highly accurate. It's a "best of both worlds" approach that bridges the gap between simple physics and powerful machine learning.

Technical Summary

Technical Summary: NORi: An ML-Augmented Ocean Boundary Layer Parameterization

Problem Statement
Global ocean climate models typically operate at resolutions of 10 km or coarser, rendering them unable to explicitly resolve subgrid-scale turbulent mixing processes in the upper ocean boundary layer (BL), which occur on scales of meters to centimeters. These mixing processes, driven by wind, cooling, and evaporation, regulate the exchange of heat, carbon, and tracers between the atmosphere and the ocean interior. Traditional parameterizations face a trade-off: simple first-order diffusive closures (e.g., Pacanowski-Philander) are computationally efficient but fail to capture nonlocal entrainment at the base of the BL, while more complex two-equation models (e.g., $k$-$\epsilon$) or bulk schemes (e.g., KPP) offer better accuracy but often incur high computational costs or struggle with numerical stability over long integration times. Furthermore, purely data-driven machine learning (ML) approaches often suffer from numerical instability and physical inconsistency when deployed in long-term climate simulations, particularly when trained a priori on instantaneous fluxes rather than temporal trajectories.

Methodology
The authors introduce NORi (Neural Ordinary Differential equations Richardson number), a hybrid parameterization that augments a simple physics-based closure with neural networks trained via an a posteriori strategy.

1. Hybrid Architecture:

   • Base Closure ("Ri"): A first-order diffusive closure based on the local gradient Richardson number ($Ri$). This component captures local mixing driven by shear and convection, ensuring the model adheres to known physical scaling laws and maintains static stability ($N^2 \geq 0$).
   • Neural Network Component ("NO"): Neural Ordinary Differential Equations (NODEs) are employed to predict nonlocal entrainment fluxes of temperature and salinity at the base of the mixed layer. These networks are "convolved" vertically, sharing weights across grid points, and are active only within a specific "entrainment zone" (diagnosed via the $Ri$ criterion). Inputs include local tracer gradients, the Richardson number, and the nonlocal surface buoyancy flux. The networks predict residual fluxes that the base closure misses, specifically the deepening of the BL due to convective plumes.

2. Training Strategy (A Posteriori Calibration):

   • Unlike traditional a priori training which matches instantaneous subgrid fluxes, NORi is trained a posteriori. The loss function depends on the full time-integrated trajectories of coarse-grained variables (temperature, salinity, velocity) rather than instantaneous fluxes.
   • The training utilizes a 1D column model integrated forward in time using fully differentiable ODE solvers (implemented in Julia with Lux.jl and Enzyme.jl). Gradients are backpropagated through the ODE solver to update network weights.
   • Curriculum Learning: Training proceeds in stages with progressively longer integration windows (15, 23.3, and 43.3 hours) to ensure numerical stability and generalization.
   • Data Source: High-fidelity Large-Eddy Simulations (LES) generated with Oceananigans.jl serve as ground truth. The dataset spans diverse regimes including shear-driven, convection-driven, and mixed turbulence, with varying rotation, stratification, and thermodynamic conditions (using the TEOS-10 equation of state).

3. Physical Constraints:

   • The architecture enforces tracer conservation by setting surface and bottom fluxes of the neural network component to zero.
   • The base closure ensures the model respects the hydrostatic approximation and prevents spurious density inversions that often plague purely ML models in hydrostatic frameworks.

Key Results

• Accuracy vs. LES: In single-column configurations, NORi matches LES solutions with high fidelity, significantly outperforming the base closure alone in convective regimes where entrainment is dominant. It performs comparably to the state-of-the-art $k$-$\epsilon$ closure and the recently developed CATKE (Convective Adjustment TKE) parameterization.
• Generalization: NORi demonstrates strong generalization to unseen validation cases, including scenarios with nonlinear thermodynamics (Southern Ocean regimes) and time-varying surface fluxes, despite being trained on constant flux scenarios.
• Observational Validation: When tested against 70 years of data from Ocean Weather Station Papa (OWS Papa), NORi reproduces the seasonal evolution of the boundary layer with skill comparable to $k$-$\epsilon$ and CATKE.
• Numerical Stability: A critical finding is NORi's stability in long-term integrations. While trained on 2-day horizons, NORi remains numerically stable for at least 100 years in an idealized double-gyre simulation. In contrast, purely ML approaches trained a priori often suffer from finite-time blowup.
• Time Step Independence: NORi maintains accuracy with time steps up to 1 hour, whereas the $k$-$\epsilon$ model requires time steps below 30 minutes for convergence. This suggests NORi can offer computational advantages in large-scale models that rely on longer time steps.
• Computational Cost: In a naive implementation, NORi is approximately 15% slower per time step than $k$-$\epsilon$. However, given its ability to support time steps twice as long as $k$-$\epsilon$, the authors argue it offers a net computational advantage for large-scale simulations.

Significance and Claims
The paper claims that NORi represents a new paradigm for designing parameterizations for climate models by successfully bridging physics-based modeling and machine learning. The authors argue that:

1. Augmentation over Replacement: Starting with a simple, physically rigorous base closure and augmenting it with ML to capture missing physics (residuals) is more effective than training a neural network to learn the entire parameterization from scratch. This approach drastically reduces data requirements and model size while improving generalization.
2. Stability via A Posteriori Training: The use of a posteriori calibration, where the loss is defined by the evolution of state variables rather than instantaneous fluxes, is essential for achieving numerical stability in long-term climate integrations. This aligns the training distribution with the inference conditions and implicitly promotes stability.
3. Practicality: NORi is designed to be compatible with existing ocean models (demonstrated via integration into Oceananigans.jl) and supports realistic nonlinear thermodynamics (TEOS-10), making it a viable candidate for improving mixed-layer depth biases in global climate models without the prohibitive cost of two-equation closures.

The authors acknowledge limitations, such as the current fixed grid resolution (8 m) and the exclusion of deep convection physics and restratification processes, noting that future work will focus on variable grid support and joint calibration within global ocean simulations.