the best I could get for now

examples/bickley_jet.jl

@@ -5,15 +5,19 @@ using OrthogonalSphericalShellGrids
 
 Nx = 360
 Ny = 180
-Nb = 30
+Nb = 40
 
 underlying_grid = TripolarGrid(size = (Nx, Ny, 1), halo = (5, 5, 5))
 
 bottom_height = zeros(Nx, Ny)
 
 bottom_height[1:Nb+1, end-Nb:end]                .= 1
 bottom_height[end-Nb:end, end-Nb:end]            .= 1
-bottom_height[(Nx-Nb)÷2:(Nx+Nb)÷2+1, end-Nb:end] .= 1
+bottom_height[(Nx-2Nb)÷2:(Nx+2Nb)÷2, end-Nb:end] .= 1
+
+@show size(underlying_grid)
+@show size(bottom_height)
+
 grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height))
 
 Ψ(y) = - tanh(y) * 10
@@ -69,7 +73,7 @@ set!(model, u=uᵢ, v=vᵢ, c=cᵢ)
 
 wizard = TimeStepWizard(cfl=0.3, max_change=1.1, max_Δt=1hour)
 
-simulation = Simulation(model, Δt=Δt, stop_time=15hours)
+simulation = Simulation(model, Δt=Δt, stop_time=500days)
 
 simulation.output_writers[:surface_tracer] = JLD2OutputWriter(model, merge(model.velocities, model.tracers, (; ζ)),
                                                               filename = "orca025_bickley.jld2", 
@@ -81,6 +85,6 @@ progress(sim) = @info @sprintf("%s with %s, velocity: %.2e %.2e", prettytime(tim
 simulation.callbacks[:progress] = Callback(progress, IterationInterval(10))
 simulation.callbacks[:wizard]   = Callback(wizard,   IterationInterval(10))
 
-# run!(simulation)
+run!(simulation)
 
 # Let's visualize the fields!

examples/rossby_haurwitz.jl

@@ -0,0 +1,101 @@
+using Oceananigans
+using Oceananigans.Units
+using Printf
+using OrthogonalSphericalShellGrids
+
+Nx = 360
+Ny = 180
+Nb = 30
+
+underlying_grid = TripolarGrid(size = (Nx, Ny, 1), halo = (5, 5, 5))
+
+bottom_height = zeros(Nx, Ny)
+
+bottom_height[1:Nb+1, end-Nb:end]                .= 1
+bottom_height[end-Nb:end, end-Nb:end]            .= 1
+bottom_height[(Nx-Nb)÷2:(Nx+Nb)÷2+1, end-Nb:end] .= 1
+
+@show size(underlying_grid)
+@show size(bottom_height)
+
+grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height))
+
+Ψ(y) = - tanh(y) * 10
+U(y) = sech(y)^2 
+
+# A sinusoidal tracer
+C(y, L) = sin(2π * y / L)
+
+# Slightly off-center vortical perturbations
+ψ̃(x, y, ℓ, k) = exp(-(y + ℓ/10)^2 / 2ℓ^2) * cos(k * x) * cos(k * y)
+
+# Vortical velocity fields (ũ, ṽ) = (-∂_y, +∂_x) ψ̃
+ũ(x, y, ℓ, k) = + ψ̃(x, y, ℓ, k) * (k * tan(k * y) + y / ℓ^2) 
+ṽ(x, y, ℓ, k) = - ψ̃(x, y, ℓ, k) * k * tan(k * x) 
+
+free_surface = SplitExplicitFreeSurface(grid; substeps = 30)
+
+@info "Building a model..."; start=time_ns()
+
+tracer_advection = Oceananigans.Advection.TracerAdvection(WENO(; order = 5), WENO(; order = 5), Centered())
+momentum_advection = WENOVectorInvariant(vorticity_order = 5)
+
+model = HydrostaticFreeSurfaceModel(; grid, free_surface,
+                                      momentum_advection,
+                                      tracer_advection,
+                                      coriolis = HydrostaticSphericalCoriolis(),
+                                      buoyancy = nothing,
+                                      tracers = :c)
+
+K = 7.848e-6
+ω = 0
+n = 4
+
+R = grid.radius
+g = model.free_surface.gravitational_acceleration
+Ω = model.coriolis.rotation_rate
+
+A(θ) = ω/2 * (2 * Ω + ω) * cos(θ)^2 + 1/4 * K^2 * cos(θ)^(2*n) * ((n+1) * cos(θ)^2 + (2 * n^2 - n - 2) - 2 * n^2 * sec(θ)^2)
+B(θ) = 2 * K * (Ω + ω) * ((n+1) * (n+2))^(-1) * cos(θ)^(n) * ( n^2 + 2*n + 2 - (n+1)^2 * cos(θ)^2) # Why not (n+1)^2 sin(θ)^2 + 1?
+C(θ)  = 1/4 * K^2 * cos(θ)^(2 * n) * ( (n+1) * cos(θ)^2 - (n+2))
+
+ψ_function(θ, ϕ) = -R^2 * ω * sin(θ)^2 + R^2 * K * cos(θ)^n * sin(θ) * cos(n*ϕ)
+
+u_function(θ, ϕ) =  R * ω * cos(θ) + R * K * cos(θ)^(n-1) * (n * sin(θ)^2 - cos(θ)^2) * cos(n*ϕ)
+v_function(θ, ϕ) = -n * K * R * cos(θ)^(n-1) * sin(θ) * sin(n*ϕ)
+
+h_function(θ, ϕ) = H + R^2/g * (A(θ) + B(θ) * cos(n * ϕ) + C(θ) * cos(2n * ϕ))
+
+rescale¹(λ) = (λ + 180) / 360 * 2π # λ to θ
+rescale²(ϕ) = ϕ / 180 * π # θ to ϕ
+
+# Arguments were u(θ, ϕ), λ |-> ϕ, θ |-> ϕ
+#=
+u₀(λ, ϕ, z) = u_function(rescale²(ϕ), rescale¹(λ))
+v₀(λ, ϕ, z) = v_function(rescale²(ϕ), rescale¹(λ))
+=#
+η₀(λ, ϕ, z) = h_function(rescale²(ϕ), rescale¹(λ))
+
+#=
+set!(model, u=u₀, v=v₀, η = η₀)
+=#
+
+ψ₀(λ, φ, z) = ψ_function(rescale²(φ), rescale¹(λ))
+
+ψ = Field{Face, Face, Center}(grid)
+
+# Note that set! fills only interior points; to compute u and v we need information in the halo regions.
+set!(ψ, ψ₀)
+
+fill_halo_regions!(ψ)
+
+u = Field(- ∂y(ψ))
+v = Field(∂x(ψ))
+
+compute!(u)
+compute!(v)
+
+fill_halo_regions!((u, v))
+
+set!(model, u = u, v = v, η = η₀)
+initialize!(model)

src/generate_tripolar_coordinates.jl

@@ -1,4 +1,26 @@
 
+@kernel function _compute_tripolar_coordinates!(λ, φ, Jeq, λ₀, Δλ, φ_grid, f_curve, xnum, ynum, jnum, Nλ, loc)
+    i, j = @index(Global, NTuple)
+
+    @inbounds begin
+        if j < Jeq + 1
+            h = (λ₀ - Δλ * i) * 2π / 360
+            x = - f_curve(φ_grid[j]) * cos(h)
+            y = - f_curve(φ_grid[j]) * sin(h)
+        elseif loc == Center() && j == size(φ, 2)
+            x = 0
+            y = linear_interpolate(j, jnum[i, :], ynum[i, :])
+        else
+            x = linear_interpolate(j, jnum[i, :], xnum[i, :])
+            y = linear_interpolate(j, jnum[i, :], ynum[i, :])
+        end
+
+        λ[i, j] = - 180 / π * ifelse(x == 0, ifelse(y == 0, 0, atan(Base.sign(y) * Inf)), atan(y / x))
+        φ[i, j] = 90 - 360 / π * atan(sqrt(y^2 + x^2)) 
+        λ[i, j] += ifelse(i ≤ Nλ÷2, -90, 90)
+    end
+end
+
 """
     compute_coords!(jnum, xnum, ynum, Δλᶠᵃᵃ, Jeq, f_interpolator, g_interpolator)
 
@@ -17,32 +39,32 @@ Compute the coordinates for an orthogonal spherical shell grid.
 This function computes the coordinates for an orthogonal spherical shell grid using the given parameters. 
 It uses a secant root finding method to find the value of `jnum` and an Adams-Bashforth-2 integrator to find the perpendicular to the circle.
 """
-@kernel function compute_tripolar_coords!(jnum, xnum, ynum, λᵢ, Δλ, Δj, Jeq, Nφ, a_interpolator, b_interpolator, c_interpolator)
+@kernel function compute_tripolar_coords!(jnum, xnum, ynum, λᵢ, Δλ, Δj, Jeq, Nφ, a_interpolator, b_interpolator)
     i = @index(Global, Linear)
     N = size(xnum, 2)
     @inbounds begin
         h = (λᵢ - Δλ * i) * 2π / 360
-        xnum[i, 1], ynum[i, 1] = cos(h), sin(h) # Starting always from a circumpherence at the equator
+        xnum[i, 1] = - a_interpolator(Jeq) * cos(h) 
+        ynum[i, 1] = - a_interpolator(Jeq) * sin(h)  # Starting always from a circumpherence at the equator
 
         Δx = xnum[i, 1] / N
 
         myfunc(x) = (xnum[i, 1] / a_interpolator(x)) ^2 + (ynum[i, 1] / b_interpolator(x))^2 - 1 
         jnum[i, 1] = bisection_root_find(myfunc, Jeq-1.0, Nφ+1-Δj, Δj)
 
-        xnum[i, 2] = xnum[i, 1] - Δx
-        ynum[i, 2] = ynum[i, 1] - Δx * tan(h)
-
-        for n in 3:N
+        for n in 2:N
             # Great circles
-            func(x) = (xnum[i, n-1] / a_interpolator(x)) ^2 + (ynum[i, n-1] / b_interpolator(x))^2 - 1 
-            jnum[i, n-1] = bisection_root_find(func, Jeq-1.0, Nφ+1-Δj, Δj)
-            xnum[i, n] = xnum[i, n-1] - Δx
 
-            # Euler forward integrator to find the perpendicular to the circle
-            G = ynum[i, n-1] * a_interpolator(jnum[i, n-1])^2 / b_interpolator(jnum[i, n-1])^2 / xnum[i, n-1]
+            func(j) = (xnum[i, n-1] / a_interpolator(j)) ^2 + (ynum[i, n-1] / b_interpolator(j))^2 - 1 
+            jnum[i, n-1] = bisection_root_find(func, Jeq-1.0, Nφ+1.0, Δj)
 
-            ynum[i, n] = ynum[i, n-1] - Δx * G
+            # Semi-Implicit integrator
+            G¹ = a_interpolator(jnum[i, n-1])^2 / b_interpolator(jnum[i, n-1])^2 / xnum[i, n-1]
+
+            ynum[i, n] = ynum[i, n-1] / (1 + G¹ * Δx) 
+            xnum[i, n] = xnum[i, n-1] - Δx
         end
+
         @show i
     end
 end
@@ -73,14 +95,21 @@ Generate tripolar metrics for a spherical shell grid.
 # Returns
 - `nothing`
 """
-function generate_tripolar_metrics!(λFF, φFF, λFC, φFC, λCF, φCF, λCC, φCC;
-                                    FT, size, halo, latitude, longitude,
-                                    Nproc, Nnum, a_curve, b_curve, c_curve,
+function generate_tripolar_metrics!(Nλ, Nφ, Hλ, Hφ;
+                                    FT, latitude, longitude,
+                                    Nproc, Nnum, a_curve, b_curve,
                                     first_pole_longitude)
+
+    λFF = zeros(Nλ, Nφ)
+    φFF = zeros(Nλ, Nφ)
+    λFC = zeros(Nλ, Nφ)
+    φFC = zeros(Nλ, Nφ)
+
+    λCF = zeros(Nλ, Nφ)
+    φCF = zeros(Nλ, Nφ)
+    λCC = zeros(Nλ, Nφ)
+    φCC = zeros(Nλ, Nφ)
 
-    Nλ, Nφ, Nz = size
-    Hλ, Hφ, Hz = halo
-
     Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, Periodic(), Nφ, Hφ, latitude,  :φ, CPU())
     Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, Periodic(), Nλ, Hλ, longitude, :λ, CPU())
 
@@ -114,58 +143,54 @@ function generate_tripolar_metrics!(λFF, φFF, λFC, φFC, λCF, φCF, λCC, φ
     for (j, φ) in enumerate(φproc)
         aᶠⱼ[j] = a_curve(φ)
         bᶠⱼ[j] = b_curve(φ) 
-        cᶠⱼ[j] = c_curve(φ) 
     end
 
     a_face_interp(j) = linear_interpolate(j, fx, aᶠⱼ)
     b_face_interp(j) = linear_interpolate(j, fx, bᶠⱼ)
-    c_face_interp(j) = linear_interpolate(j, fx, cᶠⱼ)
 
     φproc = range(φᵃᶜᵃ[1], 90, length = Nproc) 
 
     # calculate the eccentricities of the ellipse
     for (j, φ) in enumerate(φproc)
         aᶜⱼ[j] = a_curve(φ)
         bᶜⱼ[j] = b_curve(φ) 
-        cᶜⱼ[j] = c_curve(φ) 
     end
 
     a_center_interp(j) = linear_interpolate(j, fx, aᶜⱼ)
     b_center_interp(j) = linear_interpolate(j, fx, bᶜⱼ)
-    c_center_interp(j) = linear_interpolate(j, fx, cᶜⱼ)
 
     # X - Face coordinates
     λ₀ = 90 # ᵒ degrees  
 
     # Face - Face coordinates
     loop! = compute_tripolar_coords!(device(CPU()), min(256, Nλ+1), Nλ+1)
-    loop!(jnum, xnum, ynum, λ₀, Δλᶠᵃᵃ, 1/Nnum, Jeq, Nφ, a_face_interp, b_face_interp, c_face_interp) 
+    loop!(jnum, xnum, ynum, λ₀, Δλᶠᵃᵃ, 1/Nnum, Jeq, Nφ, a_face_interp, b_face_interp) 
 
-    loop! = _compute_tripolar_coordinates!(device(CPU()), (16, 16), (Nλ, Nφ+1))
+    loop! = _compute_tripolar_coordinates!(device(CPU()), (16, 16), (Nλ, Nφ))
     loop!(λFF, φFF, Jeq, λ₀, Δλᶠᵃᵃ, φᵃᶠᵃ, a_curve, xnum, ynum, jnum, Nλ, Face())
 
     # Face - Center coordinates
     loop! = compute_tripolar_coords!(device(CPU()), min(256, Nλ+1), Nλ+1)
-    loop!(jnum, xnum, ynum, λ₀, Δλᶠᵃᵃ, 1/Nnum, Jeq, Nφ, a_center_interp, b_center_interp, c_center_interp) 
+    loop!(jnum, xnum, ynum, λ₀, Δλᶠᵃᵃ, 1/Nnum, Jeq, Nφ, a_center_interp, b_center_interp) 
 
-    loop! = _compute_tripolar_coordinates!(device(CPU()), (16, 16), (Nλ, Nφ+1))
+    loop! = _compute_tripolar_coordinates!(device(CPU()), (16, 16), (Nλ, Nφ))
     loop!(λFC, φFC, Jeq, λ₀, Δλᶠᵃᵃ, φᵃᶜᵃ, a_curve, xnum, ynum, jnum, Nλ, Center())
 
     # X - Center coordinates
     λ₀ = 90 + Δλᶜᵃᵃ / 2 # ᵒ degrees  
 
     # Center - Face  
     loop! = compute_tripolar_coords!(device(CPU()), min(256, Nλ+1), Nλ+1)
-    loop!(jnum, xnum, ynum, λ₀, Δλᶜᵃᵃ, 1/Nnum, Jeq, Nφ, a_face_interp, b_face_interp, c_face_interp) 
+    loop!(jnum, xnum, ynum, λ₀, Δλᶜᵃᵃ, 1/Nnum, Jeq, Nφ, a_face_interp, b_face_interp) 
 
-    loop! = _compute_tripolar_coordinates!(device(CPU()), (16, 16), (Nλ, Nφ+1))
+    loop! = _compute_tripolar_coordinates!(device(CPU()), (16, 16), (Nλ, Nφ))
     loop!(λCF, φCF, Jeq, λ₀, Δλᶜᵃᵃ, φᵃᶠᵃ, a_curve, xnum, ynum, jnum, Nλ, Face())
 
     # Face - Center coordinates
     loop! = compute_tripolar_coords!(device(CPU()), min(256, Nλ+1), Nλ+1)
-    loop!(jnum, xnum, ynum, λ₀, Δλᶜᵃᵃ, 1/Nnum, Jeq, Nφ, a_center_interp, b_center_interp, c_center_interp) 
+    loop!(jnum, xnum, ynum, λ₀, Δλᶜᵃᵃ, 1/Nnum, Jeq, Nφ, a_center_interp, b_center_interp) 
 
-    loop! = _compute_tripolar_coordinates!(device(CPU()), (16, 16), (Nλ, Nφ+1))
+    loop! = _compute_tripolar_coordinates!(device(CPU()), (16, 16), (Nλ, Nφ))
     loop!(λCC, φCC, Jeq, λ₀, Δλᶜᵃᵃ, φᵃᶜᵃ, a_curve, xnum, ynum, jnum, Nλ, Center())
 
     λmFC = deepcopy(λFC)
@@ -181,7 +206,51 @@ function generate_tripolar_metrics!(λFF, φFF, λFC, φFC, λCF, φCF, λCC, φ
         λ .=  convert_to_0_360.(λ)
     end
 
-    return nothing
+    return λFF, φFF, λFC, φFC, λCF, φCF, λCC, φCC
+end
+
+function insert_midpoint(λ, φ, λ₀, Δλ, Jeq, φ_grid, a_curve, Nλ, ::Center)
+    λadd = zeros(Nλ)
+    φadd = zeros(Nλ)
+
+    for i in 1:Nλ
+        j = Jeq + 1
+        h = (λ₀ - Δλ * i) * 2π / 360
+        x = - a_curve(φ_grid[j]) * cos(h)
+        y = - a_curve(φ_grid[j]) * sin(h)
+
+        λadd[i] =  - 180 / π * atan(y / x)
+        φadd[i] =  90 - 360 / π * atan(sqrt(y^2 + x^2)) 
+
+        λadd[i] += ifelse(i ≤ Nλ÷2, -90, 90)
+    end
+
+    λ = hcat(λ[:, 1:Jeq], λadd, λ[:, Jeq+1:end])
+    φ = hcat(φ[:, 1:Jeq], φadd, φ[:, Jeq+1:end])
+
+    return λ, φ
+end
+
+function insert_midpoint(λ, φ, λ₀, Δλ, Jeq, φ_grid, a_curve, Nλ, ::Face)
+    λadd = zeros(Nλ)
+    φadd = zeros(Nλ)
+
+    for i in 1:Nλ
+        j = Jeq + 1
+        h = (λ₀ - Δλ * i) * 2π / 360
+        x = - a_curve(φ_grid[j]) * cos(h)
+        y = - a_curve(φ_grid[j]) * sin(h)
+
+        λadd[i] =  - 180 / π * atan(y / x)
+        φadd[i] =  90 - 360 / π * atan(sqrt(y^2 + x^2)) 
+
+        λadd[i] += ifelse(i ≤ Nλ÷2, -90, 90)
+    end
+
+    λ = hcat(λ[:, 1:Jeq], λadd, λ[:, Jeq+1:end])
+    φ = hcat(φ[:, 1:Jeq], φadd, φ[:, Jeq+1:end])
+
+    return λ, φ
 end
 
 import Oceananigans.Grids: lat_lon_to_cartesian, lat_lon_to_x, lat_lon_to_y, lat_lon_to_z