Add the possibility to use the Veros ocean model as the ocean component in ClimaOcean

Project.toml

@@ -1,8 +1,8 @@
 name = "ClimaOcean"
 uuid = "0376089a-ecfe-4b0e-a64f-9c555d74d754"
 license = "MIT"
-authors = ["Climate Modeling Alliance and contributors"]
 version = "0.8.7"
+authors = ["Climate Modeling Alliance and contributors"]
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -33,13 +33,16 @@ 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"
 XESMF = "2e0b0046-e7a1-486f-88de-807ee8ffabe5"
 
 [extensions]
 ClimaOceanCopernicusMarineExt = "CopernicusMarine"
+ClimaOceanPythonCallExt = ["PythonCall", "CondaPkg"]
 ClimaOceanReactantExt = "Reactant"
 ClimaOceanSpeedyWeatherExt = ["SpeedyWeather", "XESMF"]
 
@@ -48,6 +51,7 @@ Adapt = "4"
 CFTime = "0.1, 0.2"
 CUDA = "4, 5"
 ClimaSeaIce = "0.3.9"
+CondaPkg = "0.2.33"
 CopernicusMarine = "0.1.1"
 CubicSplines = "0.2"
 DataDeps = "0.7"
@@ -63,6 +67,7 @@ NCDatasets = "0.12, 0.13, 0.14"
 Oceananigans = "0.101"
 OffsetArrays = "1.14"
 PrecompileTools = "1"
+PythonCall = "0.9.28"
 Reactant = "0.2.45"
 Scratch = "1"
 SeawaterPolynomials = "0.3.5"
@@ -81,4 +86,4 @@ MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "Reactant", "CopernicusMarine", "XESMF", "SpeedyWeather"]
+test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "Reactant", "CopernicusMarine", "XESMF", "SpeedyWeather" "PythonCall", "CondaPkg"]

docs/Project.toml

@@ -10,6 +10,7 @@ Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40"
+PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9"
 XESMF = "2e0b0046-e7a1-486f-88de-807ee8ffabe5"
 

docs/make.jl

@@ -18,10 +18,11 @@ const EXAMPLES_DIR = joinpath(@__DIR__, "..", "examples")
 const OUTPUT_DIR   = joinpath(@__DIR__, "src/literated")
 
 to_be_literated = [
-    "single_column_os_papa_simulation.jl",
-    "one_degree_simulation.jl",
-    "near_global_ocean_simulation.jl",
-    "global_climate_simulation.jl",
+    # "single_column_os_papa_simulation.jl",
+    # "one_degree_simulation.jl",
+    # "near_global_ocean_simulation.jl",
+    "python_ocean_forced_simulation.jl",
+    # "global_climate_simulation.jl",
 ]
 
 for file in to_be_literated
@@ -45,10 +46,11 @@ pages = [
     "Home" => "index.md",
 
     "Examples" => [
-        "Single-column ocean simulation" => "literated/single_column_os_papa_simulation.md",
-        "One-degree ocean--sea ice simulation" => "literated/one_degree_simulation.md",
-        "Near-global ocean simulation" => "literated/near_global_ocean_simulation.md",
-        "Global climate simulation" => "literated/global_climate_simulation.md",
+        # "Single-column ocean simulation" => "literated/single_column_os_papa_simulation.md",
+        # "One-degree ocean--sea ice simulation" => "literated/one_degree_simulation.md",
+        # "Near-global ocean simulation" => "literated/near_global_ocean_simulation.md",
+        "Python ocean forced simulation" => "literated/python_ocean_forced_simulation.md",],
+        # "Global climate simulation" => "literated/global_climate_simulation.md",
         ],
 
     "Vertical grids" => "vertical_grids.md",

examples/python_ocean_forced_simulation.jl

@@ -0,0 +1,125 @@
+# # A Python Ocean Simulation at 4ᵒ Resolution Forced by JRA55 Reanalysis and Initialized from ECCO
+#
+# This example showcases the use of ClimaOcean'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 and initialized from the ECCO
+# state estimate.
+#
+# For this example, we need Oceananigans, ClimaOcean, Dates, CUDA, and
+# CairoMakie to visualize the simulation.
+
+using ClimaOcean
+using PythonCall
+using Oceananigans
+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(ClimaOcean, :ClimaOceanPythonCallExt)
+
+VerosModule.install_veros()
+VerosModule.remove_outputs(:global_4deg)
+
+# Actually loading and instantiating the Veros setup in the variable `ocean`.
+# This setup uses by default a different time-step for tracers and momentum, 
+# so we set it to the same value (1800 seconds) for both.
+
+ocean = VerosModule.VerosOceanSimulation("global_4deg", :GlobalFourDegreeSetup)
+
+set!(ocean, "dt_tracer", 1800.0; path=:settings)
+set!(ocean, "dt_mom",    1800.0; path=:settings)
+
+# 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
+# at every 5 iterations. It also collects the surface velocity fields and the net fluxes
+# into the arrays `s`, `tx`, and `ty` for later visualization.
+
+wall_time = Ref(time_ns())
+
+s  = []
+tx = []
+ty = []
+
+us = coupled_model.interfaces.exchanger.exchange_ocean_state.u
+vs = coupled_model.interfaces.exchanger.exchange_ocean_state.v
+
+stmp = Field(sqrt(us^2 + vs^2))
+
+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()
+
+    compute!(stmp)
+    push!(s,  deepcopy(interior(stmp, :, :, 1)))
+    push!(tx, deepcopy(interior(coupled_model.interfaces.net_fluxes.ocean_surface.u, :, :, 1) .* 1020))
+    push!(ty, deepcopy(interior(coupled_model.interfaces.net_fluxes.ocean_surface.v, :, :, 1) .* 1020))
+
+    return nothing
+end
+
+add_callback!(simulation, progress, IterationInterval(5))
+
+# Let's run the simulation!
+
+run!(simulation)
+
+# We can now visualize the surface speed and wind stress at the ocean surface
+# over the course of the simulation.
+
+iter = Observable(1)
+si   = @lift(s[$iter])
+txi  = @lift(tx[$iter])
+tyi  = @lift(ty[$iter])
+Nt   = length(s)
+
+fig = Figure(resolution = (1200, 300))
+ax1 = Axis(fig[1, 1]; title = "Surface speed (m/s)", xlabel = "Longitude", ylabel = "Latitude")
+ax2 = Axis(fig[1, 2]; title = "Zonal wind stress (N/m²)", xlabel = "Longitude")
+ax3 = Axis(fig[1, 3]; title = "Meridional wind stress (N/m²)", xlabel = "Longitude")
+
+grid = coupled_model.interfaces.exchanger.exchange_grid
+
+λ = λnodes(grid, Center())
+φ = φnodes(grid, Center())
+
+heatmap!(ax1, λ, φ, si,  colormap = :ice, colorrange = (0, 0.15))
+heatmap!(ax2, λ, φ, txi, colormap = :bwr, colorrange = (-0.2, 0.2))
+heatmap!(ax3, λ, φ, tyi, colormap = :bwr, colorrange = (-0.2, 0.2))
+
+CairoMakie.record(fig, "veros_ocean_surface.mp4", 1:Nt, framerate = 8) do nn
+    iter[] = nn
+end
+nothing #hide
+
+# ![](veros_ocean_surface.mp4)

ext/ClimaOceanPythonCallExt/ClimaOceanPythonCallExt.jl

@@ -0,0 +1,14 @@
+module ClimaOceanVerosExt
+
+using ClimaOcean
+using CondaPkg
+using PythonCall
+using Oceananigans
+using Oceananigans.DistributedComputations: @root
+
+using Dates: DateTime
+
+include("VerosOceanSimulations/veros_ocean_simulation.jl")
+include("veros_state_exchanger.jl")
+
+end # module ClimaOceanVerosExt

ext/ClimaOceanPythonCallExt/VerosOceanSimulations/VerosOceanSimulations.jl

@@ -0,0 +1,32 @@
+module VerosOceanSimulations
+
+using CondaPkg
+
+using Oceananigans.Grids: topology
+using ClimaOcean.OceanSeaIceModels: reference_density, heat_capacity, SeaIceSimulation
+
+import Oceananigans.Fields: set!
+import Oceananigans.TimeSteppers: time_step!, initialize!
+
+import ClimaOcean.OceanSeaIceModels: OceanSeaIceModel, default_nan_checker
+import Oceananigans.Architectures: architecture
+
+import Base: eltype
+
+"""
+    install_veros()
+
+Install the Veros ocean model Marine CLI using CondaPkg.
+Returns a NamedTuple containing package information if successful.
+"""
+function install_veros()
+    CondaPkg.add_pip("veros")
+    cli = CondaPkg.which("veros")
+    @info "... the veros CLI has been installed at $(cli)."
+    return cli
+end
+
+include("veros_ocean_simulation.jl")
+include("veros_state_exchanger.jl")
+
+end # module VerosOceanSimulations