GitHub

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, the Veros Python package must be available in the active CondaPkg environment. In the documentation CI this is installed ahead of time in the workflow. For a fresh local environment, run

VerosModule = Base.get_extension(NumericalEarth, :NumericalEarthVerosExt)
VerosModule.install_veros()

once before executing this example.

VerosModule = Base.get_extension(NumericalEarth, :NumericalEarthVerosExt)

VerosModule.remove_outputs(:global_4deg)

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 0x7fa048b7f920>)

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()
radiation = JRA55PrescribedRadiation()
640×320×1×2920 PrescribedRadiation on Oceananigans.Grids.LatitudeLongitudeGrid:
├── times: 2920-element StepRangeLen{Float32, Float64, Float64, Int64}
├── stefan_boltzmann_constant: 5.67037e-8
└── surface_properties: (:ocean, :sea_ice)

The coupled ocean–atmosphere model. We do not couple an ice model for simplicity.

coupled_model = OceanSeaIceModel(ocean, nothing; atmosphere=atmos, radiation)
simulation = Simulation(coupled_model; Δt = 30minutes, 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 4 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)
│   └── nan_checker => Callback of NaNChecker for T_ocean on IterationInterval(100)
└── output_writers: OrderedDict with no entries

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)
[ Info: Initializing simulation...
[ Info: time: 0 seconds, iteration: 0, Δt: 30 minutes, maximum(u): (0.03, 0.03, 0.00) m/s, wall time: 21.043 seconds 
[ Info:     ... simulation initialization complete (991.793 ms)
[ Info: Executing initial time step...
[ Info:     ... initial time step complete (2.045 seconds).
[ Info: time: 10 days, iteration: 480, Δt: 30 minutes, maximum(u): (0.44, 0.41, 0.00) m/s, wall time: 3.541 minutes 
[ Info: time: 20 days, iteration: 960, Δt: 30 minutes, maximum(u): (0.32, 0.44, 0.00) m/s, wall time: 3.356 minutes 
[ Info: time: 30 days, iteration: 1440, Δt: 30 minutes, maximum(u): (0.33, 0.37, 0.00) m/s, wall time: 3.455 minutes 
[ Info: time: 40 days, iteration: 1920, Δt: 30 minutes, maximum(u): (0.47, 0.32, 0.00) m/s, wall time: 3.439 minutes 
[ Info: time: 50 days, iteration: 2400, Δt: 30 minutes, maximum(u): (0.50, 0.32, 0.00) m/s, wall time: 3.444 minutes 
[ Info: Simulation is stopping after running for 20.676 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.28, 0.39, 0.00) m/s, wall time: 3.434 minutes

After the simulation is done, we can visualize the surface zonal and meridional velocities as a function of latitude and time.

n  = Observable(1)
un = @lift(u[$n])
vn = @lift(v[$n])
Nt = length(u)

fig = Figure(size = (900, 630))
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

This page was generated using Literate.jl.