Merge pull request #1522 from CliMA/ali/full-cubed-sphere

.buildkite/pipeline.yml

@@ -58,7 +58,7 @@ steps:
 ##### Unit tests
 #####
 
-  - label: "🐀 gpu unit tests"
+  - label: "🐿️ gpu unit tests"
     env:
       JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER"
       TEST_GROUP: "unit"
@@ -166,7 +166,7 @@ steps:
 ##### HydrostaticFreeSurfaceModel
 #####
 
-  - label: "🐡 gpu hydrostatic free surface model tests"
+  - label: "🐙 gpu hydrostatic free surface model tests"
     env:
       JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER"
       TEST_GROUP: "hydrostatic_free_surface"
@@ -177,7 +177,7 @@ steps:
       architecture: GPU
     depends_on: "init_gpu"
 
-  - label: "🐠 cpu hydrostatic free surface model tests"
+  - label: "🦑 cpu hydrostatic free surface model tests"
     env:
       JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER"
       TEST_GROUP: "hydrostatic_free_surface"
@@ -193,7 +193,7 @@ steps:
 ##### ShallowWaterModel
 #####
 
-  - label: "🦑 gpu shallow water model tests"
+  - label: "🦢 gpu shallow water model tests"
     env:
       JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER"
       TEST_GROUP: "shallow_water"
@@ -204,7 +204,7 @@ steps:
       architecture: GPU
     depends_on: "init_gpu"
 
-  - label: "🦐 cpu shallow water model tests"
+  - label: "🦆 cpu shallow water model tests"
     env:
       JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER"
       TEST_GROUP: "shallow_water"
@@ -243,6 +243,33 @@ steps:
       architecture: CPU
     depends_on: "init_cpu"
 
+#####
+##### Cubed sphere
+#####
+
+  - label: "🐡 gpu cubed sphere tests"
+    env:
+      JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER"
+      TEST_GROUP: "cubed_sphere"
+    commands:
+      - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'"
+    agents:
+      queue: Oceananigans
+      architecture: GPU
+    depends_on: "init_gpu"
+
+  - label: "🦔 cpu cubed sphere tests"
+    env:
+      JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER"
+      TEST_GROUP: "cubed_sphere"
+      CUDA_VISIBLE_DEVICES: "-1"
+    commands:
+      - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'"
+    agents:
+      queue: Oceananigans
+      architecture: CPU
+    depends_on: "init_cpu"
+
 #####
 ##### Distributed/MPI
 #####

src/BoundaryConditions/fill_halo_regions.jl

@@ -35,7 +35,7 @@ function fill_halo_regions!(c::AbstractArray, fieldbcs, arch, grid, args...; kwa
 
     # Wait at the end
     events = [west_event, east_event, south_event, north_event, bottom_event, top_event]
-    events = filter(e -> typeof(e) <: Event, events)
+    events = filter(e -> e isa Event, events)
     wait(device(arch), MultiEvent(Tuple(events)))
 
     return nothing
@@ -47,7 +47,7 @@ end
 
   fill_west_halo!(c, ::Nothing, args...; kwargs...) = nothing
   fill_east_halo!(c, ::Nothing, args...; kwargs...) = nothing
- fill_south_halo!(c, ::Nothing, args...; kwargs...) = nothing 
+ fill_south_halo!(c, ::Nothing, args...; kwargs...) = nothing
  fill_north_halo!(c, ::Nothing, args...; kwargs...) = nothing
    fill_top_halo!(c, ::Nothing, args...; kwargs...) = nothing
 fill_bottom_halo!(c, ::Nothing, args...; kwargs...) = nothing

src/BoundaryConditions/fill_halo_regions_flux.jl

@@ -64,5 +64,3 @@ end
  fill_north_halo!(c, bc::FBC, arch, dep, grid, args...; kwargs...) = launch!(arch, grid, :xz, _fill_north_halo!,  c, bc, grid.Hy, grid.Ny; dependencies=dep, kwargs...)
 fill_bottom_halo!(c, bc::FBC, arch, dep, grid, args...; kwargs...) = launch!(arch, grid, :xy, _fill_bottom_halo!, c, bc, grid.Hz, grid.Nz; dependencies=dep, kwargs...)
    fill_top_halo!(c, bc::FBC, arch, dep, grid, args...; kwargs...) = launch!(arch, grid, :xy, _fill_top_halo!,    c, bc, grid.Hz, grid.Nz; dependencies=dep, kwargs...)
-
-

src/Coriolis/hydrostatic_spherical_coriolis.jl

@@ -30,7 +30,7 @@ HydrostaticSphericalCoriolis(FT::DataType=Float64; rotation_rate=Ω_Earth, schem
 @inline φᶠᶠᵃ(i, j, k, grid::RegularLatitudeLongitudeGrid) = @inbounds grid.φᵃᶠᵃ[j]
 @inline φᶠᶠᵃ(i, j, k, grid::ConformalCubedSphereFaceGrid) = @inbounds grid.φᶠᶠᵃ[i, j]
 
-@inline fᶠᶠᵃ(i, j, k, grid::RegularLatitudeLongitudeGrid, coriolis::HydrostaticSphericalCoriolis) =
+@inline fᶠᶠᵃ(i, j, k, grid, coriolis::HydrostaticSphericalCoriolis) =
     2 * coriolis.rotation_rate * hack_sind(φᶠᶠᵃ(i, j, k, grid))
 
 @inline z_f_cross_U(i, j, k, grid::AbstractGrid{FT}, coriolis::HydrostaticSphericalCoriolis, U) where FT = zero(FT)

src/CubedSpheres/CubedSpheres.jl

@@ -0,0 +1,85 @@
+module CubedSpheres
+
+export
+    ConformalCubedSphereGrid,
+    ConformalCubedSphereField,
+    λnodes, φnodes
+
+include("cubed_sphere_utils.jl")
+include("conformal_cubed_sphere_grid.jl")
+include("cubed_sphere_exchange_bcs.jl")
+include("cubed_sphere_fields.jl")
+include("cubed_sphere_set!.jl")
+include("cubed_sphere_halo_filling.jl")
+include("cubed_sphere_kernel_launching.jl")
+
+#####
+##### Proper launch! when `ExplicitFreeSurface` is an argument
+#####
+
+using Oceananigans.Models.HydrostaticFreeSurfaceModels: ExplicitFreeSurface, PrescribedVelocityFields
+
+maybe_replace_with_face(free_surface::ExplicitFreeSurface, cubed_sphere_grid, face_number) =
+    ExplicitFreeSurface(free_surface.η.faces[face_number], free_surface.gravitational_acceleration)
+
+maybe_replace_with_face(velocities::PrescribedVelocityFields, cubed_sphere_grid, face_number) =
+    PrescribedVelocityFields(velocities.u.faces[face_number], velocities.v.faces[face_number], velocities.w.faces[face_number], velocities.parameters)
+
+#####
+##### NaN checker for cubed sphere fields
+#####
+
+import Oceananigans.Diagnostics: error_if_nan_in_field
+
+function error_if_nan_in_field(field::AbstractCubedSphereField, name, clock)
+    for (face_number, field_face) in enumerate(field.faces)
+        error_if_nan_in_field(field_face, string(name) * " (face $face_number)", clock)
+    end
+end
+
+#####
+##### CFL for cubed sphere fields
+#####
+
+import Oceananigans.Diagnostics: accurate_cell_advection_timescale
+
+function accurate_cell_advection_timescale(grid::ConformalCubedSphereGrid, velocities)
+
+    min_timescale_on_faces = []
+
+    for (face_number, grid_face) in enumerate(grid.faces)
+        velocities_face = maybe_replace_with_face(velocities, grid, face_number)
+        min_timescale_on_face = accurate_cell_advection_timescale(grid_face, velocities_face)
+        push!(min_timescale_on_faces, min_timescale_on_face)
+    end
+
+    return minimum(min_timescale_on_faces)
+end
+
+#####
+##### Output writing for cubed sphere fields
+#####
+
+import Oceananigans.OutputWriters: fetch_output
+
+fetch_output(field::AbstractCubedSphereField, model, field_slicer) =
+    Tuple(fetch_output(field_face, model, field_slicer) for field_face in field.faces)
+
+#####
+##### StateChecker for each face is useful for debugging
+#####
+
+import Oceananigans.Diagnostics: state_check
+
+function state_check(field::AbstractCubedSphereField, name, pad)
+    Nf = length(field.faces)
+    for (face_number, field_face) in enumerate(field.faces)
+        face_str = " face $face_number"
+        state_check(field_face, string(name) * face_str, pad + length(face_str))
+
+        # Leave empty line between fields for easier visual inspection.
+        face_number == Nf && @info ""
+    end
+end
+
+end # module