diff --git a/examples/bickley_jet.jl b/examples/bickley_jet.jl index 904aa84..b4c313d 100644 --- a/examples/bickley_jet.jl +++ b/examples/bickley_jet.jl @@ -5,7 +5,7 @@ using OrthogonalSphericalShellGrids Nx = 360 Ny = 180 -Nb = 30 +Nb = 40 underlying_grid = TripolarGrid(size = (Nx, Ny, 1), halo = (5, 5, 5)) @@ -13,7 +13,11 @@ 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! diff --git a/examples/rossby_haurwitz.jl b/examples/rossby_haurwitz.jl new file mode 100644 index 0000000..1275f07 --- /dev/null +++ b/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) diff --git a/examples/tracer_transport.jl b/examples/tracer_transport.jl deleted file mode 100644 index 9028edc..0000000 --- a/examples/tracer_transport.jl +++ /dev/null @@ -1,24 +0,0 @@ -using OrthogonalSphericalShellGrids - -using Oceananigans -using Oceananigans.Units -using Printf -using OrthogonalSphericalShellGrids - -underlying_grid = TripolarGrid(size = (360, 180, 1), halo = (5, 5, 5)) - -bottom_height = zeros(360, 180) - -bottom_height[1:10, end-10:end] .= 1 -bottom_height[end-10:end, end-10:end] .= 1 -bottom_height[170:190, end-10:end] .= 1 -grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height)) - -ψ(x, y, z) = - tanh((x .- 180) ./ 10) * 10 - -P = Field((Face, Face, Center), grid) -set!(P, ψ) - -uI = compute!(Field(- ∂y(P))) -vI = compute!(Field(∂x(P))) - diff --git a/src/generate_tripolar_coordinates.jl b/src/generate_tripolar_coordinates.jl index e85a07c..e89214d 100644 --- a/src/generate_tripolar_coordinates.jl +++ b/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,12 +143,10 @@ 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) @@ -127,28 +154,26 @@ 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_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 @@ -156,16 +181,16 @@ function generate_tripolar_metrics!(λFF, φFF, λFC, φFC, λCF, φCF, λCC, φ # 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 diff --git a/src/generate_tripolar_coordinates_new.jl b/src/generate_tripolar_coordinates_new.jl new file mode 100644 index 0000000..05854ee --- /dev/null +++ b/src/generate_tripolar_coordinates_new.jl @@ -0,0 +1,207 @@ + +""" + compute_coords!(jnum, xnum, ynum, Δλᶠᵃᵃ, Jeq, f_interpolator, g_interpolator) + +Compute the coordinates for an orthogonal spherical shell grid. + +# Arguments +- `jnum`: An array to store the computed values of `jnum`. +- `xnum`: An array to store the computed values of `xnum`. +- `ynum`: An array to store the computed values of `ynum`. +- `Δλᶠᵃᵃ`: The angular step size. +- `Jeq`: The value of j at the equator. +- `f_interpolator`: A function that interpolates the value of f. +- `g_interpolator`: A function that interpolates the value of g. + +# Details +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. +""" + +const γ¹ = 8 // 15 +const γ² = 5 // 12 +const γ³ = 3 // 4 + +const ζ² = -17 // 60 +const ζ³ = -5 // 12 + +@kernel function _compute_tripolar_coords!(λ, φ, λᵢ, φᵢ, Δλ, Δφ, Jeq, Nφ, Nλ, Nnum, a_curve, b_curve, a_integrator, b_integrator) + i = @index(Global, Linear) + h = (λᵢ - Δλ * i) * 2π / 360 + @inbounds begin + for j = 1:Jeq + φ[i, j] = φᵢ + Δφ * (j - 1) + x = - a_curve(φ[i, Jeq]) * cos(h) + y = - a_curve(φ[i, Jeq]) * sin(h) + λ[i, j] = - 180 / π * ifelse(x == 0, ifelse(y == 0, 0, - atan(Base.sign(y) * Inf)), atan(y / x)) + λ[i, j] += ifelse(i ≤ Nλ÷2, -90, 90) + end + + # Starting from Jeq + 1 + x = - a_curve(φ[i, Jeq]) * cos(h) + y = - a_curve(φ[i, Jeq]) * sin(h) + + Neq = Nφ - Jeq + 1 + + # Number of integration steps per step + Δx = x / (Neq * Nnum) + + for j in Jeq:Nφ + for n in 1:Nnum + # Great circles + func(j) = (x / a_integrator(j)) ^2 + (y / b_integrator(j))^2 - 1 + jn = bisection_root_find(func, Jeq-1.0, Nφ+1.0, 1/Nnum) + + # Semi-Implicit integrator + # φₙ = 90 - 360 / π * atan(sqrt(y^2 + x^2)) + G¹ = a_integrator(jn)^2 / b_integrator(jn)^2 + + y = y / (1 + G¹ / x * Δx) + x -= Δx + + # Runge Kutta integrator to find the perpendicular to the circle + + # # Fist substep! + # G² = y * a_curve(φₙ)^2 / b_curve(φₙ)^2 / x + # x -= Δx * γ¹ + # y -= Δx * γ¹ * G² + # φₙ = 90 - 360 / π * atan(sqrt(y^2 + x^2)) + + # # Second substep! + # G¹ = G² + # G² = y * a_curve(φₙ)^2 / b_curve(φₙ)^2/ x + # x -= Δx * (γ² + ζ²) + # y -= Δx * (γ² * G² + ζ² * G¹) + # φₙ = 90 - 360 / π * atan(sqrt(y^2 + x^2)) + + # # Third substep! + # G¹ = G² + # G² = y * a_curve(φₙ)^2 / b_curve(φₙ)^2 / x + # x -= Δx * (γ³ + ζ³) + # y -= Δx * (γ³ * G² + ζ³ * G¹) + + if i == Nλ + @show x, y, G¹ + end + 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) + if i == Nλ || i == Nλ-1 + @show λ[i, j], φ[i, j], i + end + end + + # @show i, x + end +end + +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φ) + + Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, Periodic(), Nφ, Hφ, latitude, :φ, CPU()) + Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, Periodic(), Nλ, Hλ, longitude, :λ, CPU()) + + # Identify equator line + J = Ref(0) + for j in 1:Nφ+1 + if φᵃᶠᵃ[j] < 0 + J[] = j + end + end + + Jeq = J[] + 1 + + aᶠⱼ = zeros(Nproc) + bᶠⱼ = zeros(Nproc) + cᶠⱼ = zeros(Nproc) + + aᶜⱼ = zeros(Nproc) + bᶜⱼ = zeros(Nproc) + cᶜⱼ = zeros(Nproc) + + fx = Float64.(collect(0:Nproc-1) ./ (Nproc-1) * Nφ .+ 1) + + φFproc = range(φᵃᶠᵃ[1], 90 - Δφᵃᶠᵃ / 2, length = Nproc) + Δφᵃᶠᵃ = (φFproc[2] - φFproc[1]) .* Nproc / Nφ + + # calculate the eccentricities of the ellipse + for (j, φ) in enumerate(φproc) + aᶠⱼ[j] = a_curve(φ) + bᶠⱼ[j] = b_curve(φ) + end + + a_face_interp(j) = linear_interpolate(j, fx, aᶠⱼ) + b_face_interp(j) = linear_interpolate(j, fx, bᶠⱼ) + + φCproc = range(φᵃᶜᵃ[1], 90, length = Nproc) + Δφᵃᶜᵃ = (φCproc[2] - φCproc[1]) .* Nproc / Nφ + + # calculate the eccentricities of the ellipse + for (j, φ) in enumerate(φproc) + aᶜⱼ[j] = a_curve(φ) + bᶜⱼ[j] = b_curve(φ) + end + + a_center_interp(j) = linear_interpolate(j, fx, aᶜⱼ) + b_center_interp(j) = linear_interpolate(j, fx, bᶜⱼ) + + integration_kernel! = _compute_tripolar_coords!(device(CPU()), min(256, Nλ+1), Nλ+1) + + # # X - Face coordinates + λ₀ = 90 # ᵒ degrees + + # # Face - Face coordinates + φ₀ = φFproc[1] + integration_kernel!(λFF, φFF, λ₀, φ₀, Δλᶠᵃᵃ, Δφᵃᶠᵃ, Jeq, Nφ, Nλ, 100, a_curve, b_curve, a_face_interp, b_face_interp) + + # # Face - Center coordinates + φ₀ = φCproc[1] + integration_kernel!(λFC, φFC, λ₀, φ₀, Δλᶠᵃᵃ, Δφᵃᶜᵃ, Jeq, Nφ, Nλ, 100, a_curve, b_curve, a_face_interp, b_face_interp) + + # X - Center coordinates + λ₀ = 90 + Δλᶜᵃᵃ / 2 # ᵒ degrees + + # Center - Face + φ₀ = φFproc[1] + integration_kernel!(λCF, φCF, λ₀, φ₀, Δλᶜᵃᵃ, Δφᵃᶠᵃ, Jeq, Nφ, Nλ, 100, a_curve, b_curve, a_center_interp, b_center_interp) + + # Face - Center coordinates + φ₀ = φCproc[1] + integration_kernel!(λCC, φCC, λ₀, φ₀, Δλᶜᵃᵃ, Δφᵃᶜᵃ, Jeq, Nφ, Nλ, 100, a_curve, b_curve, a_center_interp, b_center_interp) + + for λ in (λFF, λFC, λCF, λCC) + λ .+= first_pole_longitude + λ .= convert_to_0_360.(λ) + end + + return λFF, φFF, λFC, φFC, λCF, φCF, λCC, φCC +end + +# function calculate_face_coords!(λFF, φFF, λFC, φFC, λCF, φCF, λCC, φCC) +# x, y = stereographic_projection.(λCF, φCF) + + + +# import Oceananigans.Grids: lat_lon_to_cartesian, lat_lon_to_x, lat_lon_to_y, lat_lon_to_z + +# function stereographic_projection(λ, φ) + +# end + +function lat_lon_to_cartesian(lat, lon, radius) + return [lat_lon_to_x(lat, lon, radius), lat_lon_to_y(lat, lon, radius), lat_lon_to_z(lat, lon, radius)] +end \ No newline at end of file diff --git a/src/grid_utils.jl b/src/grid_utils.jl index 343b107..b01d975 100644 --- a/src/grid_utils.jl +++ b/src/grid_utils.jl @@ -143,26 +143,3 @@ end λ[i, j] += ifelse(i ≤ Nλ÷2, -90, 90) end end - -@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 j == size(λ, 2) - 1 && loc == Center() - 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 diff --git a/src/tripolar_grid.jl b/src/tripolar_grid.jl index 6a61e6e..8b949ef 100644 --- a/src/tripolar_grid.jl +++ b/src/tripolar_grid.jl @@ -12,10 +12,10 @@ TODO: put here information about the grid, i.e.: """ struct Tripolar end -@inline tripolar_stretching_function(φ; d = 0.4) = (φ / 90)^4 * d +@inline tripolar_stretching_function(φ; d = 0.4) = d / exp(-1) * exp( - 1 / ((pi / 4 - ((90 - φ) / 180) * pi/2)/pi*4)) -@inline cosine_a_curve(φ) = - equator_fcurve(φ) -@inline cosine_b_curve(φ; d = 0.4) = - equator_fcurve(φ) + ifelse(φ > 0, tripolar_stretching_function(φ; d), 0) +@inline tan_a_curve(φ) = - equator_fcurve(φ) +@inline exp_b_curve(φ; d = 0.4) = - equator_fcurve(φ) + ifelse(φ > 0, tripolar_stretching_function(φ; d), 0) @inline zero_c_curve(φ) = 0 @@ -67,43 +67,33 @@ function TripolarGrid(arch = CPU(), FT::DataType = Float64; first_pole_longitude = 75, # The second pole will be at `λ = first_pole_longitude + 180ᵒ` Nproc = 1000, Nnum = 1000, - a_curve = cosine_a_curve, - initial_b_curve = cosine_b_curve, + a_curve = tan_a_curve, + initial_b_curve = exp_b_curve, c_curve = zero_c_curve) + latitude = (southermost_latitude, 90) + longitude = (-180, 180) + # For now, only for domains Periodic in λ (from -180 to 180 degrees) and Bounded in φ. # φ has to reach the north pole.` # For all the rest we can easily use a `LatitudeLongitudeGrid` without warping final_b = sqrt((tan((90 - poles_latitude) / 360 * π))^2) b_curve(φ) = initial_b_curve(φ; d = final_b) - latitude = (southermost_latitude, 90) - longitude = (-180, 180) - Nλ, Nφ, Nz = size Hλ, Hφ, Hz = halo # the Z coordinate is the same as for the other grids Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, Bounded(), Nz, Hz, z, :z, CPU()) - λFF = zeros(Nλ, Nφ+1) - φFF = zeros(Nλ, Nφ+1) - λFC = zeros(Nλ, Nφ+1) - φFC = zeros(Nλ, Nφ+1) - - λCF = zeros(Nλ, Nφ+1) - φCF = zeros(Nλ, Nφ+1) - λCC = zeros(Nλ, Nφ+1) - φCC = zeros(Nλ, Nφ+1) - - generate_tripolar_metrics!(λFF, φFF, λFC, φFC, λCF, φCF, λCC, φCC; - FT, size, halo, latitude, longitude, - Nproc, Nnum, a_curve, b_curve, c_curve, - first_pole_longitude) + λFF, φFF, λFC, φFC, λCF, φCF, λCC, φCC = generate_tripolar_metrics!(Nλ, Nφ, Hλ, Hφ; + FT, latitude, longitude, + Nproc, Nnum, a_curve, b_curve, + first_pole_longitude) Nx = Nλ Ny = Nφ - + # Helper grid to fill halo grid = RectilinearGrid(; size = (Nx, Ny, 1), halo, topology = (Periodic, RightConnected, Bounded), z = (0, 1), x = (0, 1), y = (0, 1)) @@ -338,10 +328,10 @@ using Oceananigans.Fields: architecture, FieldBoundaryBuffers sign(LX, LY) = 1 -sign(::Type{Face}, ::Type{Face}) = -1 -sign(::Type{Face}, ::Type{Center}) = -1 -sign(::Type{Center}, ::Type{Face}) = -1 -sign(::Type{Center}, ::Type{Center}) = 1 +sign(::Type{Face}, ::Type{Face}) = 1 +sign(::Type{Face}, ::Type{Center}) = -1 # u-velocity type +sign(::Type{Center}, ::Type{Face}) = -1 # v-velocity type +sign(::Type{Center}, ::Type{Center}) = 1 function Field((LX, LY, LZ)::Tuple, grid::TRG, data, old_bcs, indices::Tuple, op, status) arch = architecture(grid) diff --git a/src/zipper_boundary_condition.jl b/src/zipper_boundary_condition.jl index 9e0d984..112d974 100644 --- a/src/zipper_boundary_condition.jl +++ b/src/zipper_boundary_condition.jl @@ -28,13 +28,14 @@ validate_boundary_condition_location(bc::Zipper, loc::Face, side) = @inline function fold_north_face_face!(i, k, grid, sign, c) Nx, Ny, _ = size(grid) - i′ = Nx - i + 2 + i′ = Nx - i + 2 # Remember! element Nx + 1 does not exist! + s = ifelse(i′ > Nx , abs(sign), sign) # for periodic elements we change the sign i′ = ifelse(i′ > Nx, i′ - Nx, i′) # Periodicity is hardcoded in the x-direction!! Hy = grid.Hy for j = 1 : Hy @inbounds begin - c[i, Ny + j, k] = sign * c[i′, Ny - j + 1, k] + c[i, Ny + j, k] = s * c[i′, Ny - j + 1, k] end end @@ -44,13 +45,14 @@ end @inline function fold_north_face_center!(i, k, grid, sign, c) Nx, Ny, _ = size(grid) - i′ = Nx - i + 2 + i′ = Nx - i + 2 # Remember! element Nx + 1 does not exist! + s = ifelse(i′ > Nx , abs(sign), sign) # for periodic elements we change the sign i′ = ifelse(i′ > Nx, i′ - Nx, i′) # Periodicity is hardcoded in the x-direction!! Hy = grid.Hy for j = 1 : Hy @inbounds begin - c[i, Ny + j, k] = sign * c[i′, Ny - j, k] # The Ny line is duplicated so we substitute starting Ny-1 + c[i, Ny + j, k] = s * c[i′, Ny - j, k] # The Ny line is duplicated so we substitute starting Ny-1 end end @@ -61,7 +63,6 @@ end Nx, Ny, _ = size(grid) i′ = Nx - i + 1 - i′ = ifelse(i′ > Nx, i′ - Nx, i′) # Periodicity is hardcoded in the x-direction!! Hy = grid.Hy for j = 1 : Hy @@ -77,7 +78,6 @@ end Nx, Ny, _ = size(grid) i′ = Nx - i + 1 - i′ = ifelse(i′ > Nx, i′ - Nx, i′) # Periodicity is hardcoded in the x-direction!! Hy = grid.Hy for j = 1 : Hy