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

CondaPkg.toml

@@ -1,7 +1,8 @@
 
 [pip.deps]
-copernicusmarine = ">=2.0.0"
 xarray = ">=2024.7.0"
-numpy = ">=2.0.0"
 jax = ">=0.6"
+copernicusmarine = ">=2.0.0"
+numpy = "==2.3.1"
 tensorflow = ">=2.17"
+veros = ""

experiments/veros_forced_simulation.jl

@@ -0,0 +1,113 @@
+using ClimaOcean
+using PythonCall
+using Oceananigans
+using Printf
+
+#####
+##### A Prognostic Python Ocean (Veros) Simulation
+#####
+
+# 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 loaded
+# and that every output is removed to avoid conflicts.
+
+VerosModule = Base.get_extension(ClimaOcean, :ClimaOceanPythonCallExt)
+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)
+
+##### 
+##### A Prescribed Atmosphere (JRA55)
+##### 
+
+atmos = JRA55PrescribedAtmosphere(; backend = JRA55NetCDFBackend(10))
+
+#####
+##### An ice-free ocean forced by a prescribed atmosphere
+#####
+
+radiation = Radiation()
+coupled_model = OceanSeaIceModel(ocean, nothing; atmosphere=atmos, radiation)
+simulation = Simulation(coupled_model; Δt = 1800, stop_iteration = 100000)
+
+#####
+##### A simple progress callback
+#####
+
+# 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 go!!
+#####
+
+run!(simulation)
+
+#####
+##### Visualize
+#####
+
+iter = Observable(1)
+si  = @lift(s[$iter])
+txi = @lift(tx[$iter])
+tyi = @lift(ty[$iter])
+
+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))

ext/ClimaOceanPythonCallExt/ClimaOceanPythonCallExt.jl

@@ -0,0 +1,15 @@
+module ClimaOceanPythonCallExt
+
+using ClimaOcean
+using CondaPkg
+using PythonCall
+using Oceananigans
+using Oceananigans.DistributedComputations: @root
+
+using Dates: DateTime
+
+include("copernicus.jl")
+include("veros_ocean_simulation.jl")
+include("veros_state_exchanger.jl")
+
+end # module ClimaOceanPythonCallExt

ext/ClimaOceanPythonCallExt.jl → ext/ClimaOceanPythonCallExt/copernicus.jl

@@ -119,5 +119,3 @@ function depth_bounds_kw(z)
     maximum_depth = - z[1]
     return (; minimum_depth, maximum_depth)
 end
-
-end # module ClimaOceanPythonCallExt

ext/ClimaOceanPythonCallExt/veros_ocean_simulation.jl

@@ -0,0 +1,202 @@
+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("veros"; channel = "conda-forge")
+    cli = CondaPkg.which("veros")
+    @info "... the veros CLI has been installed at $(cli)."
+    return cli
+end
+
+struct VerosOceanSimulation{S}
+    setup :: S
+end
+
+default_nan_checker(model::OceanSeaIceModel{<:Any, <:Any, <:VerosOceanSimulation}) = nothing
+
+initialize!(::ClimaOceanPythonCallExt.VerosOceanSimulation{Py}) = nothing
+time_step!(ocean::VerosOceanSimulation, Δt) = ocean.setup.step(ocean.setup.state)
+architecture(model::OceanSeaIceModel{<:Any, <:Any, <:VerosOceanSimulation}) = CPU()
+eltype(model::OceanSeaIceModel{<:Any, <:Any, <:VerosOceanSimulation}) = Float64
+
+function remove_outputs(setup::Symbol)
+    rm("$(setup).averages.nc", force=true)
+    rm("$(setup).energy.nc", force=true)
+    rm("$(setup).overturning.nc", force=true)
+    rm("$(setup).snapshot.nc", force=true)
+    return nothing
+end
+
+const CCField2D = Field{<:Center, <:Center, <:Nothing}
+const FCField2D = Field{<:Face,   <:Center, <:Nothing}
+const CFField2D = Field{<:Center, <:Face,   <:Nothing}
+
+function set!(field::CCField2D, pyarray::Py, k=pyconvert(Int, pyarray.shape[2]))
+    array = PyArray(pyarray)
+    Nx, Ny, Nz = size(array)
+    set!(field, view(array, 3:Nx-2, 3:Ny-2, k, 1))
+    return field
+end
+
+function set!(field::FCField2D, pyarray::Py, k=pyconvert(Int, pyarray.shape[2]))
+    array = PyArray(pyarray)
+    Nx, Ny, Nz = size(array)
+    TX, TY, _  = topology(field.grid)
+    i_indices  = TX == Periodic ? UnitRange(3, Nx-2) : UnitRange(2, Nx-2)
+    set!(field, view(array, i_indices, 3:Ny-2, k, 1))
+    return field
+end
+
+function set!(field::CFField2D, pyarray::Py, k=pyconvert(Int, pyarray.shape[2]))
+    array = PyArray(pyarray)
+    Nx, Ny, Nz = size(array)
+    set!(field, view(array, 3:Nx-2, 2:Ny-2, k, 1))
+    return field
+end
+
+"""
+    VerosOceanSimulation(setup, setup_name::Symbol)
+
+Creates and initializes a preconfigured Veros ocean simulation using the 
+specified setup module and setup name.
+
+Arguments
+==========
+- `setup::AbstractString`: The name of the Veros setup module to import (e.g., `"global_4deg"`).
+- `setup_name::Symbol`: The name of the setup class or function within the module to instantiate (e.g., `:GlobalFourDegreeSetup`).
+"""
+function VerosOceanSimulation(setup, setup_name::Symbol)
+    setups = pyimport("veros.setups." * setup)
+    setup  = @eval $setups.$setup_name()
+
+    # instantiate the setup
+    setup.setup()
+
+    return VerosOceanSimulation(setup) 
+end
+
+"""
+    surface_grid(ocean::VerosOceanSimulation)
+
+Constructs a `LatitudeLongitudeGrid` representing the surface grid of the given `VerosOceanSimulation` object.
+Notes: Veros always uses a LatitudeLongitudeGrid with 2 halos in both the latitude and longitude directions.
+Both latitude and longitude can be either stretched or uniform, depending on the setup, and while the meridional
+direction (latitude) is always Bounded, the zonal direction (longitude) can be either Periodic or Bounded.
+
+Arguments
+==========
+- `ocean::VerosOceanSimulation`: The ocean simulation object containing the grid state variables.
+"""
+function surface_grid(ocean::VerosOceanSimulation)
+
+    xf = Array(PyArray(ocean.setup.state.variables.xu))
+    yf = Array(PyArray(ocean.setup.state.variables.yu))
+
+    xc = Array(PyArray(ocean.setup.state.variables.xt))
+    yc = Array(PyArray(ocean.setup.state.variables.yt))
+
+    xf = xf[2:end-2]
+    yf = yf[2:end-2]
+
+    xc = xc[3:end-2]
+    yc = yc[3:end-2]
+
+    xf[1] = xf[2] - 2xc[1]
+    yf[1] = sign(yf[2]) * (yf[2] - 2yc[1])
+
+    TX = if xf[1] == 0 && xf[end] == 360
+        Periodic
+    else
+        Bounded
+    end
+
+    Nx = length(xc) 
+    Ny = length(yc) 
+
+    return LatitudeLongitudeGrid(size=(Nx, Ny), longitude=xf, latitude=yf, topology=(TX, Bounded, Flat), halo=(2, 2))
+end
+
+"""
+    set!(ocean, v, x; path = :variable)
+
+Set the `v` variable in the `ocean` model to the value of `x`.
+the path corresponds to the path inside the class where to locate the 
+variable `v` to set. It can be either `:variables` or `:settings`.
+"""
+function set!(ocean::VerosOceanSimulation, v, x; path = :variables)
+    setup = ocean.setup
+    if path == :variables
+        pyexec("""
+        with setup.state.variables.unlock():
+            setup.state.variables.__setattr__(y, t)
+        """, Main, (y=v, t=x, setup=setup))
+    elseif path == :settings
+        pyexec("""
+        with setup.state.settings.unlock():
+            setup.state.settings.__setattr__(y, t)
+        """, Main, (y=v, t=x, setup=setup))
+    else
+        error("path must be either :variable or :settings.")
+    end
+end
+
+function OceanSeaIceModel(ocean::VerosOceanSimulation, sea_ice=nothing;
+                          atmosphere = nothing,
+                          radiation = Radiation(),
+                          clock = Clock(time=0),
+                          ocean_reference_density = 1020.0,
+                          ocean_heat_capacity = 3998.0, 
+                          sea_ice_reference_density = reference_density(sea_ice),
+                          sea_ice_heat_capacity = heat_capacity(sea_ice),
+                          interfaces = nothing)
+
+    if sea_ice isa SeaIceSimulation
+        if !isnothing(sea_ice.callbacks)
+            pop!(sea_ice.callbacks, :stop_time_exceeded, nothing)
+            pop!(sea_ice.callbacks, :stop_iteration_exceeded, nothing)
+            pop!(sea_ice.callbacks, :wall_time_limit_exceeded, nothing)
+            pop!(sea_ice.callbacks, :nan_checker, nothing)
+        end
+    end
+
+    # Contains information about flux contributions: bulk formula, prescribed fluxes, etc.
+    if isnothing(interfaces) && !(isnothing(atmosphere) && isnothing(sea_ice))
+        interfaces = ComponentInterfaces(atmosphere, ocean, sea_ice;
+                                         ocean_reference_density,
+                                         ocean_heat_capacity,
+                                         sea_ice_reference_density,
+                                         sea_ice_heat_capacity,
+                                         radiation)
+    end
+
+    arch = CPU()
+
+    ocean_sea_ice_model = OceanSeaIceModel(arch,
+                                           clock,
+                                           atmosphere,
+                                           sea_ice,
+                                           ocean,
+                                           interfaces)
+
+    # Make sure the initial temperature of the ocean
+    # is not below freezing and above melting near the surface
+    initialization_update_state!(ocean_sea_ice_model)
+
+    return ocean_sea_ice_model
+end