Use Veros as another Ocean component in NumericalEarth.jl

.github/workflows/ci.yml

@@ -58,6 +58,7 @@ jobs:
           - checkpointing
           - speedy_weather
           - downloading
+          - veros
     steps:
       - name: Show available storage before cleanup
         run: |

Project.toml

@@ -35,6 +35,8 @@ Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
 ZipFile = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea"
 
 [weakdeps]
+PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
+CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab"
 CopernicusMarine = "cd43e856-93a3-40c8-bc9e-6146cdce14fa"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9"
@@ -45,13 +47,15 @@ NumericalEarthCopernicusClimateDataStoreExt = "CopernicusClimateDataStore"
 NumericalEarthCopernicusMarineExt = "CopernicusMarine"
 NumericalEarthReactantExt = "Reactant"
 NumericalEarthSpeedyWeatherExt = ["SpeedyWeather", "XESMF"]
+NumericalEarthVerosExt = ["PythonCall", "CondaPkg"]
 
 [compat]
 Adapt = "4"
 CFTime = "0.1, 0.2"
 CUDA = "5.9.5"
 CUDA_Compiler_jll = "v0.3.1"
 ClimaSeaIce = "0.4.2"
+CondaPkg = "0.2.33"
 CopernicusClimateDataStore = "0.1"
 CopernicusMarine = "0.1.1"
 CubicSplines = "0.2"
@@ -68,6 +72,7 @@ NCDatasets = "0.12, 0.13, 0.14"
 Oceananigans = "0.104.2"
 OffsetArrays = "1.14"
 PrecompileTools = "1"
+PythonCall = "0.9.28"
 Reactant = "0.2.45"
 Scratch = "1"
 SeawaterPolynomials = "0.3.5"
@@ -87,4 +92,4 @@ MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "Reactant", "CopernicusClimateDataStore", "CopernicusMarine", "XESMF", "SpeedyWeather"]
+test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "Reactant", "CopernicusClimateDataStore", "CopernicusMarine", "XESMF", "SpeedyWeather", "PythonCall", "CondaPkg"]

docs/Project.toml

@@ -12,6 +12,8 @@ NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 NumericalEarth = "904d977b-046a-4731-8b86-9235c0d1ef02"
 Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40"
+PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
+CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab"
 SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9"
 XESMF = "2e0b0046-e7a1-486f-88de-807ee8ffabe5"
 
@@ -22,4 +24,3 @@ NumericalEarth = {path = ".."}
 Documenter = "1"
 DocumenterCitations = "1.3"
 Literate = "2.2"
-XESMF = "0.1.6"

docs/make.jl

@@ -33,6 +33,7 @@ examples = [
     Example("One-degree ocean--sea ice simulation", "one_degree_simulation", false),
     Example("Near-global ocean simulation", "near_global_ocean_simulation", false),
     Example("Global climate simulation", "global_climate_simulation", false),
+    Example("Veros ocean simulation", "veros_ocean_forced_simulation", true), # For the moment just to test, then put it to false
 ]
 
 # Developer examples from docs/src/developers/ directory
@@ -99,8 +100,9 @@ pages = [
 
 modules = Module[]
 NumericalEarthSpeedyWeatherExt = isdefined(Base, :get_extension) ? Base.get_extension(NumericalEarth, :NumericalEarthSpeedyWeatherExt) : NumericalEarth.NumericalEarthSpeedyWeatherExt
+NumericalEarthVerosExt = isdefined(Base, :get_extension) ? Base.get_extension(NumericalEarth, :NumericalEarthVerosExt) : NumericalEarth.NumericalEarthVerosExt
 
-for m in [NumericalEarth, NumericalEarthSpeedyWeatherExt]
+for m in [NumericalEarth, NumericalEarthSpeedyWeatherExt, NumericalEarthVerosExt]
     if !isnothing(m)
         push!(modules, m)
     end

examples/veros_ocean_forced_simulation.jl

@@ -0,0 +1,145 @@
+# # An Ocean Simulation at 4ᵒ Resolution Forced by JRA55 Reanalysis 
+#
+# This example showcases the use of NumericalEarth's PythonCall extension to run a
+# near-global ocean simulation at 4-degree resolution using the Veros ocean model.
+# The ocean is forced by the JRA55 reanalysis data
+#
+# For this example, we need Oceananigans, NumericalEarth, Dates, CUDA, and
+# CairoMakie to visualize the simulation.
+
+using NumericalEarth
+using PythonCall
+using Oceananigans, Oceananigans.Units
+using CairoMakie
+using Printf
+
+# We import the Veros 4 degree ocean simulation setup, which consists of a near-global ocean
+# with a uniform resolution of 4 degrees in both latitude and longitude and a latitude range spanning
+# from 80S to 80N. The setup is defined in the `veros.setups.global_4deg` module.
+
+# Before importing the setup, we need to ensure that the Veros module is installed and loaded
+# and that every output is removed to avoid conflicts.
+
+VerosModule = Base.get_extension(NumericalEarth, :NumericalEarthVerosExt)
+
+VerosModule.install_veros()
+VerosModule.remove_outputs(:global_4deg)
+
+# Actually loading and instantiating the Veros setup in the variable `ocean`.
+
+ocean = VerosModule.VerosOceanSimulation("global_4deg", :GlobalFourDegreeSetup)
+
+# The loaded Veros setup contains a `set_forcing` method which computes the fluxes as restoring from climatology.
+# We replace it with a custom function that only computes the TKE forcing (which depends on the wind stresses
+# that we set in NumericalEarth). This way our u, v, T, S forcings are not overwritten. 
+# The `set_forcing_tke_only` method defined below is modified from the `set_forcing` method defined in 
+# https://github.com/team-ocean/veros/blob/main/veros/setups/global_4deg/global_4deg.py
+
+pyexec("""
+def set_forcing_tke_only(state):
+    from veros.core.operators import numpy as npx, update, at
+    from veros import KernelOutput
+
+    vs = state.variables
+    settings = state.settings
+
+    if settings.enable_tke:
+        vs.forc_tke_surface = update(
+            vs.forc_tke_surface,
+            at[1:-1, 1:-1],
+            npx.sqrt(
+                (0.5 * (vs.surface_taux[1:-1, 1:-1] + vs.surface_taux[:-2, 1:-1]) / settings.rho_0) ** 2
+                + (0.5 * (vs.surface_tauy[1:-1, 1:-1] + vs.surface_tauy[1:-1, :-2]) / settings.rho_0) ** 2
+            ) ** 1.5,
+        )
+
+    return KernelOutput(
+        surface_taux=vs.surface_taux,
+        surface_tauy=vs.surface_tauy,
+        forc_tke_surface=vs.forc_tke_surface,
+        forc_temp_surface=vs.forc_temp_surface,
+        forc_salt_surface=vs.forc_salt_surface,
+    )
+
+ocean.set_forcing = set_forcing_tke_only
+""", Main, (ocean=ocean.setup,))
+
+# We force the 4-degree setup with a prescribed atmosphere based on the JRA-55 reanalysis data.
+# This includes 2-meter wind velocity, temperature, humidity, downwelling longwave and shortwave
+# radiation, as well as freshwater fluxes.
+
+atmos = JRA55PrescribedAtmosphere(; backend = JRA55NetCDFBackend(10))
+
+# The coupled ocean--atmosphere model.
+# We use the default radiation model and we do not couple an ice model for simplicity.
+
+radiation = Radiation()
+coupled_model = OceanSeaIceModel(ocean, nothing; atmosphere=atmos, radiation)
+simulation = Simulation(coupled_model; Δt = 1800, stop_time = 60days)
+
+# We set up a progress callback that will print the current time, iteration, and maximum velocities
+# every 10days. We also set up another callback that collects the surface prognostic variables
+# into arrays for later visualization.
+
+wall_time = Ref(time_ns())
+
+function progress(sim)
+    ocean = sim.model.ocean
+    umax = maximum(PyArray(ocean.setup.state.variables.u))
+    vmax = maximum(PyArray(ocean.setup.state.variables.v))
+    wmax = maximum(PyArray(ocean.setup.state.variables.w))
+
+    step_time = 1e-9 * (time_ns() - wall_time[])
+
+    msg1 = @sprintf("time: %s, iteration: %d, Δt: %s, ", prettytime(sim), iteration(sim), prettytime(sim.Δt))
+    msg5 = @sprintf("maximum(u): (%.2f, %.2f, %.2f) m/s, ", umax, vmax, wmax)
+    msg6 = @sprintf("wall time: %s \n", prettytime(step_time))
+
+    @info msg1 * msg5 * msg6
+
+    wall_time[] = time_ns()
+
+    return nothing
+end
+
+u = []
+v = []
+
+function save_variables(sim)
+    push!(u, deepcopy(sim.model.interfaces.exchanger.ocean.state.u))
+    push!(v, deepcopy(sim.model.interfaces.exchanger.ocean.state.v))
+end
+
+add_callback!(simulation, progress, TimeInterval(10days))
+add_callback!(simulation, save_variables, IterationInterval(10))
+
+# Let's run the simulation!
+
+run!(simulation)
+
+iter = Observable(1)
+ui = @lift(u[$iter])
+vi = @lift(v[$iter])
+Nt = length(u)
+
+fig = Figure(resolution = (1000, 700))
+ax1 = Axis(fig[1, 1]; title = "Surface zonal velocity (m/s)", xlabel = "", ylabel = "Latitude")
+ax2 = Axis(fig[2, 1]; title = "Surface meridional velocity (m/s)", xlabel = "", ylabel = "Latitude")
+
+grid = coupled_model.interfaces.exchanger.grid
+
+λ = λnodes(grid, Center())
+φ = φnodes(grid, Center())
+
+hm1 = heatmap!(ax1, λ, φ, ui, colormap = :bwr,     colorrange = (-0.2, 0.2))
+hm2 = heatmap!(ax2, λ, φ, vi, colormap = :bwr,     colorrange = (-0.2, 0.2))
+
+Colorbar(fig[1, 2], hm1)
+Colorbar(fig[2, 2], hm2)
+
+CairoMakie.record(fig, "veros_ocean_surface.mp4", 1:Nt, framerate = 8) do nn
+    iter[] = nn
+end
+nothing #hide
+
+# ![](veros_ocean_surface.mp4)

ext/NumericalEarthCopernicusMarineExt.jl

@@ -11,7 +11,6 @@ using NumericalEarth.DataWrangling.GLORYS: GLORYSMetadata, GLORYSMetadatum
 
 import NumericalEarth.DataWrangling: download_dataset
 
-
 # Download each date individually, instead of downloading the entire dataset at once.
 # This is useful for a possible extension of the temporal horizon of the dataset.
 function download_dataset(metadata::GLORYSMetadata; kwargs...)

ext/NumericalEarthVerosExt/NumericalEarthVerosExt.jl

@@ -0,0 +1,11 @@
+module NumericalEarthVerosExt
+
+using NumericalEarth
+using CondaPkg
+using PythonCall
+using Oceananigans
+
+include("veros_ocean_simulation.jl")
+include("veros_state_exchanger.jl")
+
+end # module NumericalEarthVerosExt