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[ Info: Oceananigans will use 8 threads
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) CondaPkg Found dependencies: /storage5/github-action-runners/actions-runner-3/_work/NumericalEarth.jl/NumericalEarth.jl/docs/CondaPkg.toml
CondaPkg Found dependencies: /storage5/github-action-runners/julia-depot/packages/XESMF/cFsfs/CondaPkg.toml
CondaPkg Found dependencies: /storage5/github-action-runners/julia-depot/packages/CondaPkg/8GjrP/CondaPkg.toml
CondaPkg Found dependencies: /storage5/github-action-runners/julia-depot/packages/PythonCall/83z4q/CondaPkg.toml
CondaPkg Found dependencies: /storage5/github-action-runners/julia-depot/packages/CopernicusClimateDataStore/gQCKi/CondaPkg.toml
CondaPkg Found dependencies: /storage5/github-action-runners/julia-depot/packages/CondaPkg/0UqYV/CondaPkg.toml
CondaPkg Resolving changes
+ veros (pip)
CondaPkg Initialising pixi
│ /storage5/github-action-runners/julia-depot/artifacts/cefba4912c2b400756d043a2563ef77a0088866b/bin/pixi
│ init
│ --format pixi
└ /storage5/github-action-runners/actions-runner-3/_work/NumericalEarth.jl/NumericalEarth.jl/docs/.CondaPkg
✔ Created /storage5/github-action-runners/actions-runner-3/_work/NumericalEarth.jl/NumericalEarth.jl/docs/.CondaPkg/pixi.toml
CondaPkg Wrote /storage5/github-action-runners/actions-runner-3/_work/NumericalEarth.jl/NumericalEarth.jl/docs/.CondaPkg/pixi.toml
│ [dependencies]
│ netcdf4 = "*"
│ libstdcxx = ">=3.4,<15.0"
│ openssl = ">=3, <3.6"
│ libstdcxx-ng = ">=3.4,<15.0"
│ uv = ">=0.4"
│ hdf5 = "*"
│
│ [dependencies.xesmf]
│ channel = "conda-forge"
│ version = "*"
│
│ [dependencies.esmpy]
│ channel = "conda-forge"
│ version = "*"
│
│ [dependencies.esmf]
│ channel = "conda-forge"
│ version = "*"
│
│ [dependencies.python]
│ channel = "conda-forge"
│ build = "*cp*"
│ version = ">=3.10,<3.14, >=3.10,!=3.14.0,!=3.14.1,<4"
│
│ [project]
│ name = ".CondaPkg"
│ platforms = ["linux-64"]
│ channels = ["conda-forge"]
│ channel-priority = "strict"
│ description = "automatically generated by CondaPkg.jl"
│
│ [pypi-dependencies]
│ era5cli = ">=2.0.1"
│
│ [pypi-dependencies.veros]
└ url = "https://github.com/team-ocean/veros/archive/refs/heads/main.zip"
CondaPkg Installing packages
│ /storage5/github-action-runners/julia-depot/artifacts/cefba4912c2b400756d043a2563ef77a0088866b/bin/pixi
│ install
└ --manifest-path /storage5/github-action-runners/actions-runner-3/_work/NumericalEarth.jl/NumericalEarth.jl/docs/.CondaPkg/pixi.toml
✔ The default environment has been installed.
[ Info: ... the veros CLI has been installed at /storage5/github-action-runners/actions-runner-3/_work/NumericalEarth.jl/NumericalEarth.jl/docs/.CondaPkg/.pixi/envs/default/bin/veros.
Actually loading and instantiating the Veros setup in the variable ocean.
ocean = VerosModule.VerosOceanSimulation("global_4deg", :GlobalFourDegreeSetup)NumericalEarthVerosExt.VerosOceanSimulation{PythonCall.Py}(<py veros.setups.global_4deg.global_4deg.GlobalFourDegreeSetup object at 0x7f1ceecfea50>)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/global4deg/global4deg.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))640×320×1×2920 PrescribedAtmosphere{Float32} on Oceananigans.Grids.LatitudeLongitudeGrid:
├── times: 2920-element StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}
├── surface_layer_height: 10.0
└── boundary_layer_height: 512.0The 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)Simulation of EarthSystemModel{CPU}(time = 0 seconds, iteration = 0)
├── Next time step: 30 minutes
├── run_wall_time: 0 seconds
├── run_wall_time / iteration: NaN days
├── stop_time: 60 days
├── stop_iteration: Inf
├── wall_time_limit: Inf
├── minimum_relative_step: 0.0
├── callbacks: OrderedDict with 3 entries:
│ ├── stop_time_exceeded => Callback of stop_time_exceeded on IterationInterval(1)
│ ├── stop_iteration_exceeded => Callback of stop_iteration_exceeded on IterationInterval(1)
│ └── wall_time_limit_exceeded => Callback of wall_time_limit_exceeded on IterationInterval(1)
└── output_writers: OrderedDict with no entriesWe 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[ Info: Initializing simulation...
[ Info: time: 0 seconds, iteration: 0, Δt: 30 minutes, maximum(u): (0.03, 0.03, 0.00) m/s, wall time: 10.143 seconds
[ Info: ... simulation initialization complete (762.718 ms)
[ Info: Executing initial time step...
[ Info: ... initial time step complete (1.287 seconds).
[ Info: time: 10 days, iteration: 480, Δt: 30 minutes, maximum(u): (0.44, 0.41, 0.00) m/s, wall time: 2.950 minutes
[ Info: time: 20 days, iteration: 960, Δt: 30 minutes, maximum(u): (0.32, 0.44, 0.00) m/s, wall time: 2.949 minutes
[ Info: time: 30 days, iteration: 1440, Δt: 30 minutes, maximum(u): (0.32, 0.37, 0.00) m/s, wall time: 2.919 minutes
[ Info: time: 40 days, iteration: 1920, Δt: 30 minutes, maximum(u): (0.47, 0.32, 0.00) m/s, wall time: 2.969 minutes
[ Info: time: 50 days, iteration: 2400, Δt: 30 minutes, maximum(u): (0.50, 0.34, 0.00) m/s, wall time: 2.936 minutes
[ Info: Simulation is stopping after running for 17.695 minutes.
[ Info: Simulation time 60 days equals or exceeds stop time 60 days.
[ Info: time: 60 days, iteration: 2880, Δt: 30 minutes, maximum(u): (0.27, 0.38, 0.00) m/s, wall time: 2.970 minutes
┌ Warning: Found `resolution` in the theme when creating a `Scene`. The `resolution` keyword for `Scene`s and `Figure`s has been deprecated. Use `Figure(; size = ...` or `Scene(; size = ...)` instead, which better reflects that this is a unitless size and not a pixel resolution. The key could also come from `set_theme!` calls or related theming functions.
└ @ Makie /storage5/github-action-runners/julia-depot/packages/Makie/Vn16E/src/scenes.jl:264
This page was generated using Literate.jl.