Advection of tracer with prescribed or constant velocities over a `ConformalCubedSphereGrid`

[navidcy]

Add a validation script and/or tests.

[navidcy]

After lot of fiddling, the script below seems to work OK on #main @ b99b467, but only after we comment out

https://github.com/CliMA/Oceananigans.jl/blob/b99b467e163efc64170c2a68398aead6cd5c4478/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl#L54

using Oceananigans

using Oceananigans.Grids: φnode, λnode, halo_size
using Oceananigans.MultiRegion: getregion
using Oceananigans.BoundaryConditions: fill_halo_regions!
using Oceananigans.Fields: replace_horizontal_velocity_halos!

using GLMakie
GLMakie.activate!()

Nx, Ny, Nz = 180, 180, 1

radius = 1
grid = ConformalCubedSphereGrid(; panel_size=(Nx, Ny, Nz),
                                  z = (-1, 0),
                                  radius, 
                                  horizontal_direction_halo = 4,
                                  partition = CubedSpherePartition(; R = 1))

u = XFaceField(grid)
v = YFaceField(grid)

R = getregion(grid, 1).radius  # radius of the sphere (m)

α  = 0
u₀ = 1

ψ(λ, φ, z) = - R * u₀ * (sind(φ) * cosd(α) - cosd(λ) * cosd(φ) * sind(α))

Hx, Hy, Hz = halo_size(grid)

Ψ = Field{Center, Center, Center}(grid)

for region in 1:6
    region_grid = getregion(grid, region)
    for k in -Hz+1:region_grid.Nz+Hz, j in -Hy+1:grid.Ny+Hy, i in -Hx+1:grid.Nx+Hx
        λ = λnode(i, j, k, getregion(grid, region), Face(), Face(), Center())
        φ = φnode(i, j, k, getregion(grid, region), Face(), Face(), Center())
        getregion(Ψ, region)[i, j, k] = ψ(λ, φ, 0)
    end
end

u = Field{Face,   Center, Center}(grid)
v = Field{Center, Face,   Center}(grid)

for region in 1:6
    region_grid = getregion(grid, region)
    for k in -Hz+2:region_grid.Nz+Hz-1, j in -Hy+2:region_grid.Ny+Hy-1, i in -Hx+2:region_grid.Nx+Hx-1
        getregion(u, region)[i, j, k] = - (getregion(Ψ, region)[i, j+1, k] - getregion(Ψ, region)[i, j, k]) / getregion(grid, region).Δyᶜᶠᵃ[i, j]
        getregion(v, region)[i, j, k] =   (getregion(Ψ, region)[i+1, j, k] - getregion(Ψ, region)[i, j, k]) / getregion(grid, region).Δxᶠᶜᵃ[i, j]
    end
end

for _ in 1:2
    fill_halo_regions!(u)
    fill_halo_regions!(v)
    @apply_regionally replace_horizontal_velocity_halos!((; u = u, v = v, w = nothing), grid)
end

velocities = PrescribedVelocityFields(; u , v)

model = HydrostaticFreeSurfaceModel(; grid,
                                      velocities,
                                      momentum_advection = nothing,
                                      tracer_advection = WENO(),
                                      tracers = :θ,
                                      buoyancy = nothing)

# initial condition for tracer

panel = 1
λ₁ = λnode(3Nx÷4+1, Ny÷4+1, getregion(grid, panel), Center(), Center())
φ₁ = φnode(3Nx÷4+1, Ny÷4+1, getregion(grid, panel), Center(), Center())

panel = 4
λ₂ = λnode(Nx÷4+1, 3Ny÷4+1, getregion(grid, panel), Center(), Center())
φ₂ = φnode(Nx÷4+1, 3Ny÷4+1, getregion(grid, panel), Center(), Center())

panel = 3
λ₃ = λnode(3Nx÷4+1, 3Ny÷4+1, getregion(grid, panel), Center(), Center())
φ₃ = φnode(3Nx÷4+1, 3Ny÷4+1, getregion(grid, panel), Center(), Center())

panel = 6
λ₄ = λnode(3Nx÷4+1, 3Ny÷4+1, getregion(grid, panel), Center(), Center())
φ₄ = φnode(3Nx÷4+1, 3Ny÷4+1, getregion(grid, panel), Center(), Center())

δR = 2
θ₀ = 1

θᵢ(λ, φ, z) =  θ₀ * exp(-((λ - λ₁)^2 + (φ - φ₁)^2) / 2δR^2) +
               θ₀ * exp(-((λ - λ₂)^2 + (φ - φ₂)^2) / 2δR^2) +
               θ₀ * exp(-((λ - λ₃)^2 + (φ - φ₃)^2) / 2δR^2) + 
               θ₀ * exp(-((λ - λ₄)^2 + (φ - φ₄)^2) / 2δR^2)

set!(model, θ = θᵢ)

θ = model.tracers.θ
fill_halo_regions!(θ)

Δt = 0.0015
stop_iteration = 8000

simulation = Simulation(model; Δt, stop_iteration)

# Print a progress message
using Printf

progress_message(sim) = @printf("Iteration: %04d, time: %s, Δt: %s, wall time: %s\n",
                                iteration(sim), prettytime(sim), prettytime(sim.Δt),
                                prettytime(sim.run_wall_time))

simulation.callbacks[:progress] = Callback(progress_message, IterationInterval(100))

tracer_fields = Field[]

function save_tracer(sim)
    push!(tracer_fields, deepcopy(sim.model.tracers.θ))
end

simulation.callbacks[:save_tracer] = Callback(save_tracer, IterationInterval(20))

run!(simulation)

@info "Making an animation from the saved data..."

n = Observable(1)

Θₙ = []
for region in 1:6
    push!(Θₙ, @lift parent(getregion(tracer_fields[$n], region)[:, :, grid.Nz]))
end

function where_to_plot(region)
    region == 1 && return (3, 1)
    region == 2 && return (3, 2)
    region == 3 && return (2, 2)
    region == 4 && return (2, 3)
    region == 5 && return (1, 3)
    region == 6 && return (1, 4)
end

function heatlatlon!(ax::Axis, field, k=1; kwargs...)

    LX, LY, LZ = location(field)

    grid = field.grid
    _, (λvertices, φvertices) = get_lat_lon_nodes_and_vertices(grid, LX(), LY(), LZ())

    quad_points = vcat([Point2.(λvertices[:, i, j], φvertices[:, i, j]) 
                        for i in axes(λvertices, 2), j in axes(λvertices, 3)]...)
    quad_faces = vcat([begin; j = (i-1) * 4 + 1; [j j+1  j+2; j+2 j+3 j]; end for i in 1:length(quad_points)÷4]...)

    colors_per_point = vcat(fill.(vec(interior(field, :, :, k)), 4)...)

    mesh!(ax, quad_points, quad_faces; color = colors_per_point, shading = false, kwargs...)

    xlims!(ax, (-180, 180))
    ylims!(ax, (-90, 90))

    return ax
end

heatlatlon!(ax::Axis, field::CubedSphereField, k=1; kwargs...) = apply_regionally!(heatlatlon!, ax, field, k; kwargs...)
heatlatlon!(ax::Axis, field::Observable{<:CubedSphereField}, k=1; kwargs...) = apply_regionally!(heatlatlon!, ax, field.val, k; kwargs...)

using GeoMakie
using Oceananigans.Utils: Iterate, get_lat_lon_nodes_and_vertices, get_cartesian_nodes_and_vertices, apply_regionally!

n = Observable(1)

Θₙ = @lift tracer_fields[$n]

fig = Figure(resolution = (1600, 1200), fontsize=30)
ax = GeoAxis(fig[1, 1], coastlines = true, lonlims = automatic)
heatlatlon!(ax, Θₙ, colorrange=(0, 0.5θ₀))

fig

frames = 1:length(tracer_fields)

GLMakie.record(fig, "multi_region_tracer_advection_latlon.mp4", frames, framerate = 48) do i
    @info string("Plotting frame ", i, " of ", frames[end])
    Θₙ[] = tracer_fields[i]
    heatlatlon!(ax, Θₙ, colorrange=(0, 0.5θ₀))
end

https://github.com/CliMA/Oceananigans.jl/assets/7112768/d83c1652-400a-4d9b-925a-95923ffbd46d

(PR is coming!)

[glwagner]

The PR should fix prescribed velocity fields on the cubed sphere without commenting out things that are not for the cubed sphere?

[navidcy]

The PR should fix prescribed velocity fields on the cubed sphere without commenting out things that are not for the cubed sphere?

Of course!

i was just mostly excited with the video and wanted to post about the progress…

[glwagner]

It's very exciting! I was just remarking about "we have to comment X out". Extending a method that doesn't work on the cubed sphere is practically the same amount of work (in this case), and illustrates the proper workflow when developing a new feature (often methods will have to be extended, since they may make assumptions that are no longer valid). So I wanted to point that out in case people run into this in the future and use it as a template for their own development. Also if you are working with others, being clear about the development that's needed can help coordinate efforts.