Skip to content

Commit

Permalink
trying with the warped one
Browse files Browse the repository at this point in the history
  • Loading branch information
simone-silvestri committed Apr 21, 2024
1 parent 155afbb commit bfc512c
Show file tree
Hide file tree
Showing 2 changed files with 82 additions and 28 deletions.
32 changes: 31 additions & 1 deletion src/tripolar_grid.jl
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,7 @@ function TripolarGrid(arch = CPU(), FT::DataType = Float64;
grid = RectilinearGrid(; size = (Nx, Ny, 1), halo, topology = (Periodic, RightConnected, Bounded), z = (0, 1), x = (0, 1), y = (0, 1))

default_boundary_conditions = FieldBoundaryConditions(north = ZipperBoundaryCondition(),
south = FluxBoundaryCondition(nothing),
south = nothing, # The south should be `continued`
west = Oceananigans.PeriodicBoundaryCondition(),
east = Oceananigans.PeriodicBoundaryCondition(),
top = nothing,
Expand Down Expand Up @@ -227,6 +227,28 @@ function TripolarGrid(arch = CPU(), FT::DataType = Float64;

Hx, Hy, Hz = halo

latitude_longitude_grid = LatitudeLongitudeGrid(; size,
latitude,
longitude,
z,
halo,
radius)

continue_south!(Δxᶠᶠᵃ, latitude_longitude_grid.Δxᶠᶠᵃ)
continue_south!(Δxᶠᶜᵃ, latitude_longitude_grid.Δxᶠᶜᵃ)
continue_south!(Δxᶜᶠᵃ, latitude_longitude_grid.Δxᶜᶠᵃ)
continue_south!(Δxᶜᶜᵃ, latitude_longitude_grid.Δxᶜᶜᵃ)

continue_south!(Δyᶠᶠᵃ, latitude_longitude_grid.Δyᶠᶜᵃ)
continue_south!(Δyᶠᶜᵃ, latitude_longitude_grid.Δyᶠᶜᵃ)
continue_south!(Δyᶜᶠᵃ, latitude_longitude_grid.Δyᶜᶠᵃ)
continue_south!(Δyᶜᶜᵃ, latitude_longitude_grid.Δyᶜᶠᵃ)

continue_south!(Azᶠᶠᵃ, latitude_longitude_grid.Azᶠᶠᵃ)
continue_south!(Azᶠᶜᵃ, latitude_longitude_grid.Azᶠᶜᵃ)
continue_south!(Azᶜᶠᵃ, latitude_longitude_grid.Azᶜᶠᵃ)
continue_south!(Azᶜᶜᵃ, latitude_longitude_grid.Azᶜᶜᵃ)

# Final grid with correct metrics
grid = OrthogonalSphericalShellGrid{Periodic, RightConnected, Bounded}(arch,
Nx, Ny, Nz,
Expand All @@ -242,7 +264,15 @@ function TripolarGrid(arch = CPU(), FT::DataType = Float64;
return grid
end

function continue_south!(new_metric, lat_lon_metric)
Hx, Hy = new_metric.offsets
Nx, Ny = size(new_metric)
for i in -Hx:Nx-Hx, j in -Hy:1
new_metric[i, j] = lat_lon_metric[i, j]
end

return nothing
end

const TRG = Union{TripolarGrid, ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:TripolarGrid}}

Expand Down
78 changes: 51 additions & 27 deletions src/warped_latitude_longitude.jl
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,8 @@ function WarpedLatitudeLongitudeGrid(arch = CPU(), FT::DataType = Float64;
halo = (4, 4, 4),
radius = R_Earth,
z = (0, 1),
Nnum = 10000,
Nnum = 1000,
Nproc = 1000,
singularity_longitude = 230,
f_curve = quadratic_f_curve,
g_curve = quadratic_g_curve)
Expand Down Expand Up @@ -139,50 +140,73 @@ function WarpedLatitudeLongitudeGrid(arch = CPU(), FT::DataType = Float64;

Jeq = J[] + 1

fⱼ = zeros(1:+1)
gⱼ = zeros(1:+1)
fᶠⱼ = zeros(Nproc+1)
gᶠⱼ = zeros(Nproc+1)

fᶜⱼ = zeros(Nproc+1)
gᶜⱼ = zeros(Nproc+1)

xnum = zeros(Nλ+1, Nnum)
ynum = zeros(Nλ+1, Nnum)
jnum = zeros(Nλ+1, Nnum)

# Shif pole upwards
for j in 1:+1
fⱼ[j] = f_curve(φᵃᶠᵃ[j])
gⱼ[j] = g_curve(φᵃᶠᵃ[j])
φproc = range(φᵃᶠᵃ[1], φᵃᶠᵃ[Nφ+1], length = Nproc + 1)

# calculate the eccentricities of the ellipse
for (j, φ) in enumerate(φproc)
fᶠⱼ[j] = f_curve(φ)
gᶠⱼ[j] = g_curve(φ)
end

fy = fⱼ
gy = gⱼ
fx = Float64.(collect(1:+1))

f_interpolator(j) = linear_interpolate(j, fx, fy)
g_interpolator(j) = linear_interpolate(j, fx, gy)
fᶠx = Float64.(collect(0:Nproc) ./ Nproc *.+ 1)

f_face_interp(j) = linear_interpolate(j, fᶠx, fᶠⱼ)
g_face_interp(j) = linear_interpolate(j, fᶠx, gᶠⱼ)

φproc = range(φᵃᶜᵃ[1], φᵃᶜᵃ[Nφ+1], length = Nproc + 1)

# calculate the eccentricities of the ellipse
for (j, φ) in enumerate(φproc)
fᶜⱼ[j] = f_curve(φ)
gᶜⱼ[j] = g_curve(φ)
end

xnum = zeros(1:+1, Nnum)
ynum = zeros(1:+1, Nnum)
jnum = zeros(1:+1, Nnum)
f_center_interp(j) = linear_interpolate(j, fᶠx, fᶜⱼ)
g_center_interp(j) = linear_interpolate(j, fᶠx, gᶜⱼ)

xnum = zeros(Nλ+1, Nnum)
ynum = zeros(Nλ+1, Nnum)
jnum = zeros(Nλ+1, Nnum)

# X - Face coordinates
λ₀ = 90 # ᵒ degrees

loop! = compute_coords!(device(CPU()), min(256, Nλ+1), Nλ+1)
loop!(jnum, xnum, ynum, λ₀, Δλᶠᵃᵃ, Jeq, Nφ, f_interpolator, g_interpolator)
loop!(jnum, xnum, ynum, λ₀, Δλᶠᵃᵃ, Jeq, Nφ, f_face_interp, g_face_interp)

loop! = _compute_coordinates!(device(CPU()), (16, 16), (Nλ, Nφ+1))
loop!(λFF, φFF, Jeq, Δλᶠᵃᵃ, φᵃᶠᵃ, f_curve, xnum, ynum, jnum, Nλ)
loop!(λFF, φFF, Jeq, λ₀, Δλᶠᵃᵃ, φᵃᶠᵃ, f_curve, xnum, ynum, jnum, Nλ)

loop! = compute_coords!(device(CPU()), min(256, Nλ+1), Nλ+1)
loop!(jnum, xnum, ynum, λ₀, Δλᶠᵃᵃ, Jeq, Nφ, f_center_interp, g_center_interp)

loop! = _compute_coordinates!(device(CPU()), (16, 16), (Nλ, Nφ))
loop!(λFC, φFC, Jeq, Δλᶠᵃᵃ, φᵃᶠᵃ, f_curve, xnum, ynum, jnum, Nλ)
loop!(λFC, φFC, Jeq, λ₀, Δλᶠᵃᵃ, φᵃᶜᵃ, f_curve, xnum, ynum, jnum, Nλ)

# X - Face coordinates
λ₀ = 90 + Δλᶜᵃᵃ / 2 # ᵒ degrees

loop! = compute_coords!(device(CPU()), min(256, Nλ+1), Nλ+1)
loop!(jnum, xnum, ynum, λ₀, Δλᶜᵃᵃ, Jeq, Nφ, f_interpolator, g_interpolator)
loop!(jnum, xnum, ynum, λ₀, Δλᶜᵃᵃ, Jeq, Nφ, f_face_interp, g_face_interp)

loop! = _compute_coordinates!(device(CPU()), (16, 16), (Nλ, Nφ+1))
loop!(λCF, φCF, Jeq, Δλᶜᵃᵃ, φᵃᶜᵃ, f_curve, xnum, ynum, jnum, Nλ)
loop!(λCF, φCF, Jeq, λ₀, Δλᶜᵃᵃ, φᵃᶠᵃ, f_curve, xnum, ynum, jnum, Nλ)

loop! = compute_coords!(device(CPU()), min(256, Nλ+1), Nλ+1)
loop!(jnum, xnum, ynum, λ₀, Δλᶜᵃᵃ, Jeq, Nφ, f_center_interp, g_center_interp)

loop! = _compute_coordinates!(device(CPU()), (16, 16), (Nλ, Nφ))
loop!(λCC, φCC, Jeq, Δλᶜᵃᵃ, φᵃᶜᵃ, f_curve, xnum, ynum, jnum, Nλ)
loop!(λCC, φCC, Jeq, λ₀, Δλᶜᵃᵃ, φᵃᶜᵃ, f_curve, xnum, ynum, jnum, Nλ)

Nx =
Ny =
Expand Down Expand Up @@ -291,11 +315,11 @@ function WarpedLatitudeLongitudeGrid(arch = CPU(), FT::DataType = Float64;
Nx, Ny, Nz,
Hx, Hy, Hz,
convert(eltype(radius), Lz),
arch_array(arch, λᶜᶜᵃ), arch_array(arch, λᶠᶜᵃ), arch_array(arch, λᶜᶠᵃ), arch_array(arch, λᶠᶠᵃ),
arch_array(arch, φᶜᶜᵃ), arch_array(arch, φᶠᶜᵃ), arch_array(arch, φᶜᶠᵃ), arch_array(arch, φᶠᶠᵃ), arch_array(arch, zᵃᵃᶜ), arch_array(arch, zᵃᵃᶠ),
arch_array(arch, Δxᶜᶜᵃ), arch_array(arch, Δxᶠᶜᵃ), arch_array(arch, Δxᶜᶠᵃ), arch_array(arch, Δxᶠᶠᵃ),
arch_array(arch, Δyᶜᶜᵃ), arch_array(arch, Δyᶜᶠᵃ), arch_array(arch, Δyᶠᶜᵃ), arch_array(arch, Δyᶠᶠᵃ), arch_array(arch, Δzᵃᵃᶜ), arch_array(arch, Δzᵃᵃᶠ),
arch_array(arch, Azᶜᶜᵃ), arch_array(arch, Azᶠᶜᵃ), arch_array(arch, Azᶜᶠᵃ), arch_array(arch, Azᶠᶠᵃ),
on_architecture(arch, λᶜᶜᵃ), on_architecture(arch, λᶠᶜᵃ), on_architecture(arch, λᶜᶠᵃ), on_architecture(arch, λᶠᶠᵃ),
on_architecture(arch, φᶜᶜᵃ), on_architecture(arch, φᶠᶜᵃ), on_architecture(arch, φᶜᶠᵃ), on_architecture(arch, φᶠᶠᵃ), on_architecture(arch, zᵃᵃᶜ), on_architecture(arch, zᵃᵃᶠ),
on_architecture(arch, Δxᶜᶜᵃ), on_architecture(arch, Δxᶠᶜᵃ), on_architecture(arch, Δxᶜᶠᵃ), on_architecture(arch, Δxᶠᶠᵃ),
on_architecture(arch, Δyᶜᶜᵃ), on_architecture(arch, Δyᶜᶠᵃ), on_architecture(arch, Δyᶠᶜᵃ), on_architecture(arch, Δyᶠᶠᵃ), on_architecture(arch, Δzᵃᵃᶜ), on_architecture(arch, Δzᵃᵃᶠ),
on_architecture(arch, Azᶜᶜᵃ), on_architecture(arch, Azᶠᶜᵃ), on_architecture(arch, Azᶜᶠᵃ), on_architecture(arch, Azᶠᶠᵃ),
radius, WarpedLatitudeLongitude())

return grid
Expand Down

0 comments on commit bfc512c

Please sign in to comment.