Merge pull request #2 from simone-silvestri/ss/distributed_grid

Distributed Tripolar grid

Project.toml

@@ -8,10 +8,10 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterTools = "35a29f4d-8980-5a13-9543-d66fff28ecb8"
-FFMPEG = "c87230d0-a227-11e9-1b43-d7ebe4e7570a"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"

examples/convert_to_latlong_frame.jl

@@ -0,0 +1,84 @@
+using OrthogonalSphericalShellGrids
+using OrthogonalSphericalShellGrids.Oceananigans
+using OrthogonalSphericalShellGrids: Face, Center
+using OrthogonalSphericalShellGrids: get_cartesian_nodes_and_vertices
+using Oceananigans.Fields: interpolate!
+using Oceananigans.Grids: φnode
+using Oceananigans.Operators
+
+using GLMakie
+
+# Here we assume that the tripolar grid is locally orthogonal
+@inline function convert_to_latlong_frame(i, j, grid, uₒ, vₒ)
+
+    φᶜᶠᵃ₊ = φnode(i, j+1, 1, grid, Center(), Face(), Center())
+    φᶜᶠᵃ₋ = φnode(i,   j, 1, grid, Center(), Face(), Center())
+    Δyᶜᶜᵃ = Δyᶜᶜᶜ(i,   j, 1, grid)
+
+    ũ = deg2rad(φᶜᶠᵃ₊ - φᶜᶠᵃ₋) / Δyᶜᶜᵃ
+
+    φᶠᶜᵃ₊ = φnode(i+1, j, 1, grid, Face(), Center(), Center())
+    φᶠᶜᵃ₋ = φnode(i,   j, 1, grid, Face(), Center(), Center())
+    Δxᶜᶜᵃ = Δxᶜᶜᶜ(i,   j, 1, grid)
+
+    ṽ = - deg2rad(φᶠᶜᵃ₊ - φᶠᶜᵃ₋) / Δxᶜᶜᵃ
+
+    𝒰 = sqrt(ũ^2 + ṽ^2)
+
+    d₁ = ũ / 𝒰
+    d₂ = ṽ / 𝒰
+
+    return uₒ * d₁ - vₒ * d₂, uₒ * d₂ + vₒ * d₁
+end
+
+# Here we assume that the tripolar grid is locally orthogonal
+@inline function convert_to_native_frame(i, j, grid, uₒ, vₒ)
+
+    φᶜᶠᵃ₊ = φnode(i, j+1, 1, grid, Center(), Face(), Center())
+    φᶜᶠᵃ₋ = φnode(i,   j, 1, grid, Center(), Face(), Center())
+    Δyᶜᶜᵃ = Δyᶜᶜᶜ(i,   j, 1, grid)
+
+    ũ = deg2rad(φᶜᶠᵃ₊ - φᶜᶠᵃ₋) / Δyᶜᶜᵃ
+
+    φᶠᶜᵃ₊ = φnode(i+1, j, 1, grid, Face(), Center(), Center())
+    φᶠᶜᵃ₋ = φnode(i,   j, 1, grid, Face(), Center(), Center())
+    Δxᶜᶜᵃ = Δxᶜᶜᶜ(i,   j, 1, grid)
+
+    ṽ = - deg2rad(φᶠᶜᵃ₊ - φᶠᶜᵃ₋) / Δxᶜᶜᵃ
+
+    𝒰 = sqrt(ũ^2 + ṽ^2)
+
+    d₁ = ũ / 𝒰
+    d₂ = ṽ / 𝒰
+
+    return uₒ * d₁ + vₒ * d₂, uₒ * d₂ - vₒ * d₁
+end
+
+# Generate a Tripolar grid with a 2 degree resolution and ``north'' singularities at 20 degrees latitude
+grid = TripolarGrid(size = (180, 90, 1), north_poles_latitude = 35)
+
+uLL = CenterField(grid)
+vLL = CenterField(grid)
+
+# We want a purely zonal velocity 
+fill!(uLL, 1)
+
+# Correct interpolation of uLL and vLL
+@inline convert_x(i, j, k, grid, uₒ, vₒ) = convert_to_latlong_frame(i, j, grid, uₒ[i, j, k], vₒ[i, j, k])[1]
+@inline convert_y(i, j, k, grid, uₒ, vₒ) = convert_to_latlong_frame(i, j, grid, uₒ[i, j, k], vₒ[i, j, k])[2]
+
+uC = KernelFunctionOperation{Center, Center, Center}(convert_x, grid, uLL, vLL)
+vC = KernelFunctionOperation{Center, Center, Center}(convert_y, grid, uLL, vLL)
+
+uTR = compute!(Field(uC))
+vTR = compute!(Field(vC))
+
+# Correct interpolation of uLL and vLL
+@inline convert_x_back(i, j, k, grid, uₒ, vₒ) = convert_to_native_frame(i, j, grid, uₒ[i, j, k], vₒ[i, j, k])[1]
+@inline convert_y_back(i, j, k, grid, uₒ, vₒ) = convert_to_native_frame(i, j, grid, uₒ[i, j, k], vₒ[i, j, k])[2]
+
+uB = KernelFunctionOperation{Center, Center, Center}(convert_x_back, grid, uTR, vTR)
+vB = KernelFunctionOperation{Center, Center, Center}(convert_y_back, grid, uTR, vTR)
+
+uB = compute!(Field(uB))
+vB = compute!(Field(vB))

examples/distributed_bickley_jet.jl

@@ -0,0 +1,94 @@
+using Oceananigans
+using Oceananigans.Units
+using Printf
+using OrthogonalSphericalShellGrids
+using Oceananigans.DistributedComputations: Equal
+using Oceananigans.Utils: get_cartesian_nodes_and_vertices
+
+Nx = 320
+Ny = 240
+Nb = 20
+
+first_pole_longitude = λ¹ₚ = 45
+north_poles_latitude = φₚ  = 35
+
+λ²ₚ = λ¹ₚ + 180
+
+# Divide the y-direction in 4 CPU ranks
+arch = Distributed(CPU(), partition = Partition(y = Equal()))
+
+# Build a tripolar grid with singularities at
+# (0, -90), (45, 25), (225, 25)
+underlying_grid = TripolarGrid(arch; size = (Nx, Ny, 1), 
+                               halo = (5, 5, 5), 
+                               first_pole_longitude,
+                               north_poles_latitude)
+
+# We need a bottom height field that ``masks'' the singularities
+bottom_height(λ, φ) = ((abs(λ - λ¹ₚ) < 5) & (abs(φₚ - φ) < 5)) |
+                      ((abs(λ - λ²ₚ) < 5) & (abs(φₚ - φ) < 5)) | (φ < -78) ? 1 : 0
+
+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 = WENO()
+momentum_advection = WENOVectorInvariant(vorticity_order = 5)
+
+model = HydrostaticFreeSurfaceModel(; grid, free_surface,
+                                      momentum_advection,
+                                      tracer_advection,
+                                      buoyancy = nothing,
+                                      tracers = :c)
+
+ζ = Oceananigans.Models.HydrostaticFreeSurfaceModels.VerticalVorticityField(model)    
+
+# Parameters
+ϵ = 0.1 # perturbation magnitude
+ℓ = 0.5 # Gaussian width
+k = 2.5 # Sinusoidal wavenumber
+
+dr(x) = deg2rad(x)
+
+# Total initial conditions
+uᵢ(x, y, z) = U(dr(y)*8) + ϵ * ũ(dr(x)*2, dr(y)*8, ℓ, k)
+vᵢ(x, y, z) = ϵ * ṽ(dr(x)*2, dr(y)*4, ℓ, k)
+cᵢ(x, y, z) = C(dr(y)*8, 167.0)
+
+set!(model, u=uᵢ, v=vᵢ, c=cᵢ)
+
+Δt = 1minutes
+
+wizard = TimeStepWizard(cfl=0.3, max_change=1.1, max_Δt=3hours)
+
+simulation = Simulation(model, Δt=Δt, stop_time=1500days)
+
+rank = arch.local_rank
+
+simulation.output_writers[:surface_tracer] = JLD2OutputWriter(model, merge(model.velocities, model.tracers, (; ζ)),
+                                                              filename = "tripolar_bickley_$(rank).jld2", 
+                                                              schedule = TimeInterval(0.5day),
+                                                              with_halos = true,
+                                                              overwrite_existing = true)
+
+progress(sim) = @info @sprintf("rank %d, %s with %s, velocity: %.2e %.2e", rank, prettytime(time(sim)), prettytime(sim.Δt), maximum(sim.model.velocities.u), maximum(sim.model.velocities.v)) 
+
+simulation.callbacks[:progress] = Callback(progress, IterationInterval(10))
+simulation.callbacks[:wizard]   = Callback(wizard,   IterationInterval(10))
+
+run!(simulation)

src/OrthogonalSphericalShellGrids.jl

@@ -39,7 +39,8 @@ include("zipper_boundary_condition.jl")
 include("generate_tripolar_coordinates.jl")
 include("tripolar_grid.jl")
 include("grid_extensions.jl")
-include("split_explicit_free_surface.jl")
 include("with_halo.jl")
+include("distributed_tripolar_grid.jl")
+include("split_explicit_free_surface.jl")
 
 end