diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8b9337d4..1b841248 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -59,6 +59,7 @@ jobs: - speedy_weather - breeze - downloading + - veros steps: - name: Show available storage before cleanup run: | diff --git a/Project.toml b/Project.toml index 39179581..02048cce 100644 --- a/Project.toml +++ b/Project.toml @@ -36,6 +36,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" @@ -47,6 +49,7 @@ NumericalEarthCopernicusClimateDataStoreExt = "CopernicusClimateDataStore" NumericalEarthCopernicusMarineExt = "CopernicusMarine" NumericalEarthReactantExt = "Reactant" NumericalEarthSpeedyWeatherExt = ["SpeedyWeather", "XESMF"] +NumericalEarthVerosExt = ["PythonCall", "CondaPkg"] [compat] Adapt = "4" @@ -54,6 +57,7 @@ CFTime = "0.1, 0.2" CUDA = "5.9.5" CUDA_Compiler_jll = "v0.3.1" ClimaSeaIce = "0.4.2, 0.5" +CondaPkg = "0.2.33" CopernicusClimateDataStore = "0.1" CopernicusMarine = "0.1.1" CubicSplines = "0.2" @@ -70,6 +74,7 @@ NCDatasets = "0.12, 0.13, 0.14" Oceananigans = "0.104.2, 0.105" OffsetArrays = "1.14" PrecompileTools = "1" +PythonCall = "0.9.28" Reactant = "0.2.45" Scratch = "1" SeawaterPolynomials = "0.3.5" @@ -90,4 +95,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", "Breeze"] +test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "Reactant", "CopernicusClimateDataStore", "CopernicusMarine", "XESMF", "SpeedyWeather", "Breeze", "PythonCall", "CondaPkg"] diff --git a/docs/Project.toml b/docs/Project.toml index ee5c422e..d252de0e 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -13,6 +13,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" @@ -23,4 +25,3 @@ NumericalEarth = {path = ".."} Documenter = "1" DocumenterCitations = "1.3" Literate = "2.2" -XESMF = "0.1.6" diff --git a/docs/make.jl b/docs/make.jl index f1c35a95..7c614774 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -33,7 +33,8 @@ 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("Breeze over two oceans", "breeze_over_two_oceans", true), + Example("Veros ocean simulation", "veros_ocean_forced_simulation", true), # For the moment just to test, then put it to false + Example("Breeze over two oceans", "breeze_over_two_oceans", false), ] # Developer examples from docs/src/developers/ directory @@ -99,10 +100,11 @@ pages = [ ] modules = Module[] -NumericalEarthSpeedyWeatherExt = isdefined(Base, :get_extension) ? Base.get_extension(NumericalEarth, :NumericalEarthSpeedyWeatherExt) : nothing +NumericalEarthSpeedyWeatherExt = isdefined(Base, :get_extension) ? Base.get_extension(NumericalEarth, :NumericalEarthSpeedyWeatherExt) : NumericalEarth.NumericalEarthSpeedyWeatherExt +NumericalEarthVerosExt = isdefined(Base, :get_extension) ? Base.get_extension(NumericalEarth, :NumericalEarthVerosExt) : NumericalEarth.NumericalEarthVerosExt NumericalEarthBreezeExt = isdefined(Base, :get_extension) ? Base.get_extension(NumericalEarth, :NumericalEarthBreezeExt) : nothing -for m in [NumericalEarth, NumericalEarthSpeedyWeatherExt, NumericalEarthBreezeExt] +for m in [NumericalEarth, NumericalEarthSpeedyWeatherExt, NumericalEarthBreezeExt, NumericalEarthVerosExt] if !isnothing(m) push!(modules, m) end diff --git a/examples/veros_ocean_forced_simulation.jl b/examples/veros_ocean_forced_simulation.jl new file mode 100644 index 00000000..bb95573f --- /dev/null +++ b/examples/veros_ocean_forced_simulation.jl @@ -0,0 +1,139 @@ +# # 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) + +n = Observable(1) +un = @lift(u[$n]) +vn = @lift(v[$n]) +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") +hm1 = heatmap!(ax1, un, colormap = :bwr, colorrange = (-0.2, 0.2)) +hm2 = heatmap!(ax2, vn, 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 + n[] = nn +end +nothing #hide + +# ![](veros_ocean_surface.mp4) diff --git a/ext/NumericalEarthCopernicusMarineExt.jl b/ext/NumericalEarthCopernicusMarineExt.jl index a6de5d6e..896c3464 100644 --- a/ext/NumericalEarthCopernicusMarineExt.jl +++ b/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...) diff --git a/ext/NumericalEarthVerosExt/NumericalEarthVerosExt.jl b/ext/NumericalEarthVerosExt/NumericalEarthVerosExt.jl new file mode 100644 index 00000000..2a563e5b --- /dev/null +++ b/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 diff --git a/ext/NumericalEarthVerosExt/veros_ocean_simulation.jl b/ext/NumericalEarthVerosExt/veros_ocean_simulation.jl new file mode 100644 index 00000000..7e464c88 --- /dev/null +++ b/ext/NumericalEarthVerosExt/veros_ocean_simulation.jl @@ -0,0 +1,301 @@ +using CondaPkg + +using Oceananigans.Grids: topology +using OffsetArrays: OffsetArray + +import Oceananigans.Fields: set! +import Oceananigans.TimeSteppers: time_step!, initialize! + +import Oceananigans.Architectures: architecture + +using NumericalEarth.EarthSystemModels: EarthSystemModel +import NumericalEarth.EarthSystemModels: default_nan_checker +import NumericalEarth.EarthSystemModels: reference_density, + heat_capacity, + ocean_temperature, + ocean_salinity, + ocean_surface_salinity, + ocean_surface_velocities + +import NumericalEarth.Oceans: ocean_simulation + +import Base: eltype + +""" + install_veros() + +Install the Veros ocean model Marine CLI using CondaPkg. +Returns a NamedTuple containing package information if successful. +Also patches Veros's signal handling to work with PythonCall. +""" +function install_veros() + CondaPkg.add_pip("veros", version="@ https://github.com/team-ocean/veros/archive/refs/heads/main.zip") + cli = CondaPkg.which("veros") + + # Patch signal handling as early as possible + # This ensures it's patched before any Veros modules are used + patch_veros_signal_handling() + + @info "... the veros CLI has been installed at $(cli)." + return cli +end + +""" + patch_veros_signal_handling() + +Patch Veros's signal handling to work with PythonCall. +This prevents TypeError when Veros tries to set signal handlers from within Julia/PythonCall context. +The issue is that Veros tries to set signal handlers, but PythonCall's wrapped Python objects aren't recognized +as valid callable handlers by Python's signal.signal() function. We work around this by monkey-patching +signal.signal() to accept any handler and convert invalid ones to SIG_DFL. +""" +function patch_veros_signal_handling() + pyexec(""" + import signal + + # Monkey-patch signal.signal() to handle invalid handlers gracefully + # This is needed because PythonCall-wrapped objects aren't recognized + # as valid callable handlers by Python's signal.signal() + if not hasattr(signal, '_NumericalEarth_patched'): + _original_signal = signal.signal + + def _patched_signal(signum, handler): + # Check if handler is valid according to Python's rules + if handler in (signal.SIG_IGN, signal.SIG_DFL): + # Standard handlers are always valid + return _original_signal(signum, handler) + elif callable(handler): + # Try to use the original handler if it's callable + try: + return _original_signal(signum, handler) + except TypeError: + # If Python rejects it (e.g., PythonCall wrapper), use SIG_DFL + return _original_signal(signum, signal.SIG_DFL) + else: + # Invalid handler type - use SIG_DFL instead + return _original_signal(signum, signal.SIG_DFL) + + signal.signal = _patched_signal + signal._NumericalEarth_patched = True + + # Also patch Veros's signal wrapper to skip signal handling entirely + try: + import veros.signals + # Only patch if not already patched + if not hasattr(veros.signals, '_NumericalEarth_patched'): + def _patched_dnd_wrapper(): + # Skip signal handling entirely when running from PythonCall + # The signal handlers aren't needed when Veros is embedded in Julia + pass + + veros.signals.dnd_wrapper = _patched_dnd_wrapper + veros.signals._NumericalEarth_patched = True + except (ImportError, AttributeError): + # If veros.signals doesn't exist or can't be patched, that's okay + # This might happen if Veros version doesn't have signal handling + pass + """, Main) +end + +struct VerosOceanSimulation{S} + setup :: S +end + +default_nan_checker(model::EarthSystemModel{<:Any, <:Any, <:VerosOceanSimulation}) = nothing +initialize!(::NumericalEarthVerosExt.VerosOceanSimulation{Py}) = nothing + +function time_step!(ocean::VerosOceanSimulation, Δt) + # Align the timesteps + set!(ocean, "dt_tracer", Δt; path=:settings) + set!(ocean, "dt_mom", Δt; path=:settings) + + # Time-step + ocean.setup.step(ocean.setup.state) +end + +architecture(model::EarthSystemModel{<:Any, <:Any, <:VerosOceanSimulation}) = CPU() +eltype(model::EarthSystemModel{<:Any, <:Any, <:VerosOceanSimulation}) = Float64 + +reference_density(ocean::VerosOceanSimulation) = pyconvert(eltype(ocean), ocean.setup.state.settings.rho_0) +heat_capacity(ocean::VerosOceanSimulation) = convert(eltype(ocean), 3995) + +function ocean_surface_velocities(ocean::VerosOceanSimulation) + u = PyArray(ocean.setup.state.variables.u) + v = PyArray(ocean.setup.state.variables.v) + Nxu, Nyu, Nzu = size(u) + Nxv, Nyv, Nzv = size(v) + grid = surface_grid(ocean) + TX, TY, _ = topology(grid) + i_indices = TX == Periodic ? UnitRange(3, Nxu-2) : UnitRange(2, Nxu-2) + + u_view = view(u, i_indices, 3:Nyu-2, Nzu, 1) + v_view = view(v, 3:Nxv-2, 2:Nyv-2, Nzv, 1) + + return u_view, v_view +end + +function ocean_surface_salinity(ocean::VerosOceanSimulation) + S = ocean_salinity(ocean) + Nx, Ny, Nz = size(S) + return view(S, :, :, Nz) +end + +# Veros hardcodes 2 halos in the x and y direction, +# and each prognostic variable is 4 dimensional, where the first three dimensions +# are x, y, z, and the last index differentiate between variable, tendency at n and tendency at n-1 +function ocean_temperature(ocean::VerosOceanSimulation) + T = PyArray(ocean.setup.state.variables.temp) + Nx, Ny, Nz = size(T) + return view(T, 3:Nx-2, 3:Ny-2, 1:Nz, 1) +end + +# Veros hardcodes 2 halos in the x and y direction, +# and each prognostic variable is 4 dimensional, where the first three dimensions +# are x, y, z, and the last index differentiate between variable, tendency at n and tendency at n-1 +function ocean_salinity(ocean::VerosOceanSimulation) + S = PyArray(ocean.setup.state.variables.salt) + Nx, Ny, Nz = size(S) + return view(S, 3:Nx-2, 3:Ny-2, 1:Nz, 1) +end + +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::String, setup_name::Symbol) + # Patch signal handling BEFORE importing any Veros modules + # Veros may set up signal handlers during module import + patch_veros_signal_handling() + + setups = pyimport("veros.setups." * setup) + ocean = @eval $setups.$setup_name() + + # Patch again before setup() in case Veros imports more modules + patch_veros_signal_handling() + + # Remove any existing output files to avoid "file exists" errors from Veros diagnostics + remove_outputs(Symbol(setup)) + + # instantiate the setup + ocean.setup() + + return VerosOceanSimulation(ocean) +end + +# We assume that if we pass a python object, this is a veros simulation +function ocean_simulation(ocean::Py) + # Patch Veros's signal handling before initializing + patch_veros_signal_handling() + + # instantiate the setup + ocean.setup() + + return VerosOceanSimulation(ocean) +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 diff --git a/ext/NumericalEarthVerosExt/veros_state_exchanger.jl b/ext/NumericalEarthVerosExt/veros_state_exchanger.jl new file mode 100644 index 00000000..3c4485c9 --- /dev/null +++ b/ext/NumericalEarthVerosExt/veros_state_exchanger.jl @@ -0,0 +1,80 @@ +using NumericalEarth.Oceans + +import NumericalEarth.EarthSystemModels.InterfaceComputations: + net_fluxes, + sea_ice_ocean_interface, + atmosphere_ocean_interface, + initialize!, + ComponentExchanger, + exchange_grid + +import NumericalEarth.EarthSystemModels: + interpolate_state!, + update_net_fluxes! + +import NumericalEarth.Oceans: get_radiative_forcing + +function ComponentExchanger(ocean::VerosOceanSimulation, grid) + state = (; u = Field{Face, Center, Nothing}(grid), + v = Field{Center, Face, Nothing}(grid), + T = Field{Center, Center, Nothing}(grid), + S = Field{Center, Center, Nothing}(grid)) + + return ComponentExchanger(state, nothing) +end + +exchange_grid(atmosphere, ocean::VerosOceanSimulation, sea_ice) = surface_grid(ocean) + +@inline function net_fluxes(ocean::VerosOceanSimulation) + grid = surface_grid(ocean) + u = Field{Face, Center, Nothing}(grid) + v = Field{Center, Face, Nothing}(grid) + T = Field{Center, Center, Nothing}(grid) + S = Field{Center, Center, Nothing}(grid) + + return (; u, v, T, S) +end + +function interpolate_state!(exchanger, exchange_grid, ocean::VerosOceanSimulation, coupled_model) + u = exchanger.state.u + v = exchanger.state.v + T = exchanger.state.T + S = exchanger.state.S + + set!(u, ocean.setup.state.variables.u) + set!(v, ocean.setup.state.variables.v) + set!(T, ocean.setup.state.variables.temp) + set!(S, ocean.setup.state.variables.salt) + + return nothing +end + +initialize!(exchanger::ComponentExchanger, grid, ::VerosOceanSimulation) = nothing + +get_radiative_forcing(ocean::VerosOceanSimulation) = nothing + +function update_net_fluxes!(coupled_model, ocean::VerosOceanSimulation) + + # Update the flux containers + Oceans.update_net_ocean_fluxes!(coupled_model, ocean, coupled_model.interfaces.exchanger.grid) + net_ocean_fluxes = coupled_model.interfaces.net_fluxes.ocean + + # Pass the flux values to the python ocean + nx = pyconvert(Int, ocean.setup.state.settings.nx) + 4 + ny = pyconvert(Int, ocean.setup.state.settings.ny) + 4 + + ρₒ = pyconvert(eltype(ocean), ocean.setup.state.settings.rho_0) + taux = view(parent(net_ocean_fluxes.u), 1:nx, 1:ny, 1) .* ρₒ + tauy = view(parent(net_ocean_fluxes.v), 1:nx, 1:ny, 1) .* ρₒ + + set!(ocean, "surface_taux", taux; path=:variables) + set!(ocean, "surface_tauy", tauy; path=:variables) + + temp_flux = view(parent(net_ocean_fluxes.T), 1:nx, 1:ny, 1) + salt_flux = view(parent(net_ocean_fluxes.S), 1:nx, 1:ny, 1) + + set!(ocean, "forc_temp_surface", temp_flux; path=:variables) + set!(ocean, "forc_salt_surface", salt_flux; path=:variables) + + return nothing +end diff --git a/src/EarthSystemModels/InterfaceComputations/component_interfaces.jl b/src/EarthSystemModels/InterfaceComputations/component_interfaces.jl index f26b7a3f..35db6b72 100644 --- a/src/EarthSystemModels/InterfaceComputations/component_interfaces.jl +++ b/src/EarthSystemModels/InterfaceComputations/component_interfaces.jl @@ -294,7 +294,7 @@ Keyword Arguments - `gravitational_acceleration`: gravitational acceleration. Default: `default_gravitational_acceleration`. """ function ComponentInterfaces(atmosphere, ocean, sea_ice=nothing; - exchange_grid = exchange_grid(ocean), + exchange_grid = exchange_grid(atmosphere, ocean, sea_ice), radiation = Radiation(), freshwater_density = default_freshwater_density, atmosphere_ocean_fluxes = SimilarityTheoryFluxes(eltype(exchange_grid)), diff --git a/src/EarthSystemModels/components.jl b/src/EarthSystemModels/components.jl index 0c3cbd41..0ea202eb 100644 --- a/src/EarthSystemModels/components.jl +++ b/src/EarthSystemModels/components.jl @@ -12,6 +12,12 @@ const celsius_to_kelvin = 273.15 @inline convert_from_kelvin(::DegreesCelsius, T::FT) where FT = T - convert(FT, celsius_to_kelvin) @inline convert_from_kelvin(::DegreesKelvin, T) = T +##### +##### generic defaults +##### + +exchange_grid(atmosphere, ocean, sea_ice) = ocean.model.grid + ##### ##### Functions extended by sea-ice and ocean models ##### @@ -23,7 +29,6 @@ ocean_salinity(ocean) = ZeroField() ocean_surface_temperature(ocean) = ZeroField() ocean_surface_salinity(ocean) = ZeroField() ocean_surface_velocities(ocean) = ZeroField(), ZeroField() -exchange_grid(ocean) = ocean.model.grid temperature_units(ocean) = DegreesCelsius() ##### diff --git a/src/Oceans/assemble_net_ocean_fluxes.jl b/src/Oceans/assemble_net_ocean_fluxes.jl index 645bb516..83661133 100644 --- a/src/Oceans/assemble_net_ocean_fluxes.jl +++ b/src/Oceans/assemble_net_ocean_fluxes.jl @@ -45,15 +45,14 @@ function update_net_ocean_fluxes!(coupled_model, ocean_model, grid) freshwater_flux = atmosphere_fields.Mp.data ice_concentration = sea_ice_concentration(sea_ice) - ocean_salinity = EarthSystemModels.ocean_salinity(ocean_model) + ocean_surface_salinity = EarthSystemModels.ocean_surface_salinity(ocean_model) atmos_ocean_properties = coupled_model.interfaces.atmosphere_ocean_interface.properties ocean_properties = coupled_model.interfaces.ocean_properties - kernel_parameters = interface_kernel_parameters(grid) ocean_surface_temperature = coupled_model.interfaces.atmosphere_ocean_interface.temperature penetrating_radiation = get_radiative_forcing(ocean_model) - launch!(arch, grid, kernel_parameters, + launch!(arch, grid, :xy, _assemble_net_ocean_fluxes!, net_ocean_fluxes, penetrating_radiation, @@ -61,7 +60,7 @@ function update_net_ocean_fluxes!(coupled_model, ocean_model, grid) clock, atmos_ocean_fluxes, sea_ice_ocean_fluxes, - ocean_salinity, + ocean_surface_salinity, ocean_surface_temperature, ice_concentration, downwelling_radiation, @@ -78,7 +77,7 @@ end clock, atmos_ocean_fluxes, sea_ice_ocean_fluxes, - ocean_salinity, + ocean_surface_salinity, ocean_surface_temperature, sea_ice_concentration, downwelling_radiation, @@ -96,7 +95,7 @@ end @inbounds begin ℵᵢ = sea_ice_concentration[i, j, 1] - Sₒ = ocean_salinity[i, j, kᴺ] + Sₒ = ocean_surface_salinity[i, j, 1] Tₛ = ocean_surface_temperature[i, j, 1] Tₛ = convert_to_kelvin(ocean_properties.temperature_units, Tₛ) diff --git a/src/Oceans/ocean_simulation.jl b/src/Oceans/ocean_simulation.jl index 52c0986e..8ac44299 100644 --- a/src/Oceans/ocean_simulation.jl +++ b/src/Oceans/ocean_simulation.jl @@ -43,14 +43,16 @@ function estimate_maximum_Δt(grid) Δy = mean(yspacings(grid)) Δθ = rad2deg(mean([Δx, Δy])) / grid.radius - # The maximum Δt is roughly 30minutes / Δθ, giving: - # - 30 minutes for a 1 degree ocean + # The maximum Δt is roughly 1hours * Δθ, giving: + # - 60 minutes for a 1 degree ocean + # - 30 minutes for a 0.5 degree ocean # - 15 minutes for a 1/4 degree ocean # - 7.5 minutes for a 1/8 degree ocean # - 3.75 minutes for a 1/16 degree ocean # - 1.875 minutes for a 1/32 degree ocean - Δt = 30minutes / Δθ + # We set the maximum Δt to 1 hour + Δt = min(1hours, 1hours * Δθ) return all_reduce(min, Δt, arch) end diff --git a/src/Oceans/slab_ocean.jl b/src/Oceans/slab_ocean.jl index 31305e8e..2ccbff70 100644 --- a/src/Oceans/slab_ocean.jl +++ b/src/Oceans/slab_ocean.jl @@ -94,7 +94,7 @@ Base.eltype(::SlabOcean{FT}) where FT = FT reference_density(ocean::SlabOcean) = ocean.density heat_capacity(ocean::SlabOcean) = ocean.heat_capacity -exchange_grid(ocean::SlabOcean) = ocean.grid +exchange_grid(atmosphere, ocean::SlabOcean, sea_ice) = ocean.grid temperature_units(::SlabOcean) = DegreesKelvin() ocean_temperature(ocean::SlabOcean) = ocean.temperature ocean_salinity(ocean::SlabOcean{FT}) where FT = ConstantField(convert(FT, 35)) diff --git a/src/SeaIces/assemble_net_sea_ice_fluxes.jl b/src/SeaIces/assemble_net_sea_ice_fluxes.jl index b3ae6d19..4fce5ccf 100644 --- a/src/SeaIces/assemble_net_sea_ice_fluxes.jl +++ b/src/SeaIces/assemble_net_sea_ice_fluxes.jl @@ -31,12 +31,10 @@ function update_net_fluxes!(coupled_model, sea_ice::Simulation{<:SeaIceModel}) atmos_sea_ice_properties = coupled_model.interfaces.atmosphere_sea_ice_interface.properties sea_ice_properties = coupled_model.interfaces.sea_ice_properties - kernel_parameters = interface_kernel_parameters(grid) - sea_ice_surface_temperature = coupled_model.interfaces.atmosphere_sea_ice_interface.temperature ice_concentration = sea_ice_concentration(sea_ice) - launch!(arch, grid, kernel_parameters, + launch!(arch, grid, :xy, _assemble_net_sea_ice_fluxes!, top_fluxes, bottom_heat_flux, diff --git a/test/runtests.jl b/test/runtests.jl index 26d33019..27bef5dd 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -140,6 +140,10 @@ if test_group == :speedy_weather || test_group == :all include("test_speedy_coupling.jl") end +if test_group == :veros || test_group == :all + include("test_veros.jl") +end + if test_group == :breeze || test_group == :all include("test_breeze_coupling.jl") end diff --git a/test/runtests_setup.jl b/test/runtests_setup.jl index 29bb489d..056be354 100644 --- a/test/runtests_setup.jl +++ b/test/runtests_setup.jl @@ -104,6 +104,7 @@ function test_timestepping_with_dataset(arch, dataset, start_date, inpainting; end ocean = ocean_simulation(grid; tracers=fldnames, verbose=false) + set!(ocean.model, T=20, S=35) @test begin time_step!(ocean) @@ -213,7 +214,8 @@ function test_timestepping_with_dataset_restoring(arch, dataset, dates, inpainti restoring = DatasetRestoring(metadata, arch; inpainting, rate=1/1000) forcing = NamedTuple{tuple(fldnames[end])}(tuple(restoring)) ocean = ocean_simulation(grid; tracers=fldnames, forcing, verbose=false) - + set!(ocean.model, T=20, S=35) + @test begin time_step!(ocean) time_step!(ocean) @@ -250,6 +252,7 @@ function test_cycling_dataset_restoring(arch, dataset, dates, inpainting; times = native_times(forcing[1].field_time_series.backend.metadata) ocean = ocean_simulation(grid, tracers=fldnames, forcing=forcing) + set!(ocean.model, T=20, S=35) # start a bit after time_index time_index = 3 diff --git a/test/test_ecco2_daily.jl b/test/test_ecco2_daily.jl index 4c9262ac..4d41c54b 100644 --- a/test/test_ecco2_daily.jl +++ b/test/test_ecco2_daily.jl @@ -11,8 +11,8 @@ using Oceananigans.Units using CUDA: @allowscalar -# Inpaint only the first two cells inside the missing mask -inpainting = NearestNeighborInpainting(2) +# Inpaint only the first ten cells inside the missing mask +inpainting = NearestNeighborInpainting(10) dataset = ECCO2Daily() start_date = DateTime(1993, 1, 1) diff --git a/test/test_ecco2_monthly.jl b/test/test_ecco2_monthly.jl index 0440ab32..80df535f 100644 --- a/test/test_ecco2_monthly.jl +++ b/test/test_ecco2_monthly.jl @@ -12,8 +12,8 @@ using Oceananigans.Units using CUDA: @allowscalar -# Inpaint only the first two cells inside the missing mask -inpainting = NearestNeighborInpainting(2) +# Inpaint only the first ten cells inside the missing mask +inpainting = NearestNeighborInpainting(10) test_ecco_datasets = tuple((ds for ds in test_datasets if occursin(r"^ECCO2.*Monthly",string(typeof(ds)),))...) diff --git a/test/test_ecco4_en4.jl b/test/test_ecco4_en4.jl index 4b1e8dc7..61e2e220 100644 --- a/test/test_ecco4_en4.jl +++ b/test/test_ecco4_en4.jl @@ -14,8 +14,8 @@ using Oceananigans.Units using CUDA: @allowscalar -# Inpaint only the first two cells inside the missing mask -inpainting = NearestNeighborInpainting(2) +# Inpaint only the first ten cells inside the missing mask +inpainting = NearestNeighborInpainting(10) start_date = DateTime(1993, 1, 1) test_ecco_en4_datasets = tuple((ds for ds in test_datasets if occursin(r"^E.*4",string(typeof(ds)),))...) diff --git a/test/test_veros.jl b/test/test_veros.jl new file mode 100644 index 00000000..455f0207 --- /dev/null +++ b/test/test_veros.jl @@ -0,0 +1,66 @@ +include("runtests_setup.jl") + +using Test +using NumericalEarth +using Oceananigans +using PythonCall, CondaPkg + +@testset "Veros ocean model interface" begin + # Test that the Veros extension is available + VerosModule = Base.get_extension(NumericalEarth, :NumericalEarthVerosExt) + @test !isnothing(VerosModule) + + # Test Veros installation + @test begin + VerosModule.install_veros() + true + end + + # Test creating a VerosOceanSimulation + # Using a small setup for testing + @test begin + ocean = VerosModule.VerosOceanSimulation("global_4deg", :GlobalFourDegreeSetup) + @test ocean isa VerosModule.VerosOceanSimulation + @test !isnothing(ocean.setup) + true + end + + # Test basic interface functions + ocean = VerosModule.VerosOceanSimulation("global_4deg", :GlobalFourDegreeSetup) + + # Test interface + ρₒ = NumericalEarth.EarthSystemModels.reference_density(ocean) + cₚ = NumericalEarth.EarthSystemModels.heat_capacity(ocean) + @test ρₒ isa Real + @test cₚ isa Real + + T = NumericalEarth.EarthSystemModels.ocean_temperature(ocean) + S = NumericalEarth.EarthSystemModels.ocean_salinity(ocean) + + @test !isnothing(T) + @test !isnothing(S) + + @test length(size(T)) == 3 + @test length(size(S)) == 3 + + S_surf = NumericalEarth.EarthSystemModels.ocean_surface_salinity(ocean) + u_surf, v_surf = NumericalEarth.EarthSystemModels.ocean_surface_velocities(ocean) + + @test !isnothing(S_surf) + @test !isnothing(u_surf) + @test !isnothing(v_surf) + + # Test surface grid + grid = VerosModule.surface_grid(ocean) + @test grid isa Oceananigans.Grids.LatitudeLongitudeGrid + + # Test time stepping + @test begin + Δt = 60.0 + VerosModule.time_step!(ocean, Δt) + true + end + + # Clean up output files + VerosModule.remove_outputs(:global_4deg) +end