diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 6ecaa5a8e4..f40fdbf814 100644 --- a/.buildkite/pipeline.yml +++ b/.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 ##### diff --git a/sandbox/plot_cubed_sphere_checkerboards.jl b/sandbox/plot_cubed_sphere_checkerboards.jl deleted file mode 100644 index 9541d0eb63..0000000000 --- a/sandbox/plot_cubed_sphere_checkerboards.jl +++ /dev/null @@ -1,50 +0,0 @@ -using PyCall - -plt = pyimport("matplotlib.pyplot") -ccrs = pyimport("cartopy.crs") - -include("conformal_cubed_sphere_grid.jl") - -Nx_face = Ny_face = 21 -grid = ConformalCubedSphereGrid(face_size=(Nx_face, Ny_face, 1), z=(-1, 0)) - -## Plot staggered grid points on a checkerboard for one face - -transform = ccrs.PlateCarree() - -fig = plt.figure(figsize=(16, 9)) - -cmaps = ("coolwarm", "viridis", "cividis", "plasma", "PiYG", "bone") - -ax = fig.add_subplot(1, 2, 1, projection=ccrs.NearsidePerspective(central_longitude=270, central_latitude=45)) - -for (n, face) in enumerate(grid.faces) - lons = face.λᶜᶜᶜ[1:Nx_face, 1:Ny_face] - lats = face.ϕᶜᶜᶜ[1:Nx_face, 1:Ny_face] - checkerboard = [(i+j)%2 for i in 1:Nx_face, j in 1:Ny_face] - - ax.pcolormesh(lons, lats, checkerboard, transform=transform, cmap=cmaps[n], alpha=0.4) -end - -ax.gridlines(xlocs=range(-180, 180, step=20), ylocs=range(-80, 80, step=10)) -ax.coastlines(resolution="50m") -ax.set_global() - -ax = fig.add_subplot(1, 2, 2, projection=ccrs.NearsidePerspective(central_longitude=90, central_latitude=-45)) - -for (n, face) in enumerate(grid.faces) - lons = face.λᶜᶜᶜ[1:Nx_face, 1:Ny_face] - lats = face.ϕᶜᶜᶜ[1:Nx_face, 1:Ny_face] - checkerboard = [(i+j)%2 for i in 1:Nx_face, j in 1:Ny_face] - - ax.pcolormesh(lons, lats, checkerboard, transform=transform, cmap=cmaps[n], alpha=0.4) -end - -ax.gridlines(xlocs=range(-180, 180, step=20), ylocs=range(-80, 80, step=10)) -ax.coastlines(resolution="50m") -ax.set_global() - -plt.show() -# plt.savefig("cubed_sphere_checkerboards.png", dpi=200) - -plt.close("all") diff --git a/sandbox/plot_cubed_sphere_faces.jl b/sandbox/plot_cubed_sphere_faces.jl deleted file mode 100644 index 2a73f210de..0000000000 --- a/sandbox/plot_cubed_sphere_faces.jl +++ /dev/null @@ -1,32 +0,0 @@ -using PyCall - -plt = pyimport("matplotlib.pyplot") -ccrs = pyimport("cartopy.crs") - -include("conformal_cubed_sphere_grid.jl") - -Nx_face = Ny_face = 20 -grid = ConformalCubedSphereGrid(face_size=(Nx_face, Ny_face, 1), z=(-1, 0)) - -## Plot points of each face - -projection = ccrs.Robinson() -transform = ccrs.PlateCarree() - -fig = plt.figure(figsize=(16, 9)) -ax = fig.add_subplot(1, 1, 1, projection=projection) - -for (n, face) in enumerate(grid.faces) - lons = face.λᶜᶜᶜ[1:Nx_face, 1:Ny_face] - lats = face.ϕᶜᶜᶜ[1:Nx_face, 1:Ny_face] - ax.scatter(lons, lats, s=4, label="face $n", transform=transform) -end - -ax.legend(loc="lower right") -ax.coastlines(resolution="50m") -ax.set_global() - -# plt.show() -plt.savefig("cubed_sphere_points.png", dpi=200) - -plt.close("all") diff --git a/sandbox/plot_cubed_sphere_staggered_grid.jl b/sandbox/plot_cubed_sphere_staggered_grid.jl deleted file mode 100644 index 3e693303ff..0000000000 --- a/sandbox/plot_cubed_sphere_staggered_grid.jl +++ /dev/null @@ -1,37 +0,0 @@ -using PyCall - -plt = pyimport("matplotlib.pyplot") -ccrs = pyimport("cartopy.crs") - -include("conformal_cubed_sphere_grid.jl") - -Nx_face = Ny_face = 10 -grid = ConformalCubedSphereGrid(face_size=(Nx_face, Ny_face, 1), z=(-1, 0)) - -## Plot staggered grid points on a checkerboard for one face - -# projection = ccrs.Orthographic(central_longitude=270) -projection = ccrs.Robinson() -transform = ccrs.PlateCarree() - -fig = plt.figure(figsize=(16, 9)) -ax = fig.add_subplot(1, 1, 1, projection=projection) - -face = grid.faces[2] -lons = face.λᶜᶜᶜ[1:Nx_face, 1:Ny_face] -lats = face.ϕᶜᶜᶜ[1:Nx_face, 1:Ny_face] -checkerboard = [(i+j)%2 for i in 1:Nx_face, j in 1:Ny_face] - -ax.pcolormesh(lons, lats, checkerboard, transform=transform, cmap="coolwarm", alpha=0.2) - -ax.scatter(face.λᶜᶜᶜ[1:Nx_face, 1:Ny_face ], face.ϕᶜᶜᶜ[1:Nx_face, 1:Ny_face ], s=4, label="ccc (tracer)", transform=transform) -ax.scatter(face.λᶠᶜᶜ[1:Nx_face+1, 1:Ny_face ], face.ϕᶠᶜᶜ[1:Nx_face+1, 1:Ny_face ], s=4, label="fcc (u)", transform=transform) -ax.scatter(face.λᶜᶠᶜ[1:Nx_face, 1:Ny_face+1], face.ϕᶜᶠᶜ[1:Nx_face, 1:Ny_face+1], s=4, label="cfc (v)", transform=transform) -ax.scatter(face.λᶠᶠᶜ[1:Nx_face+1, 1:Ny_face+1], face.ϕᶠᶠᶜ[1:Nx_face+1, 1:Ny_face+1], s=4, label="ffc (vorticity)", transform=transform) - -ax.legend(loc="lower right") -ax.coastlines(resolution="50m") -ax.set_global() - -# plt.show() -plt.savefig("cubed_sphere_staggered_grid.png", dpi=200) diff --git a/src/BoundaryConditions/fill_halo_regions.jl b/src/BoundaryConditions/fill_halo_regions.jl index b583bcacdd..fd0c6ede29 100644 --- a/src/BoundaryConditions/fill_halo_regions.jl +++ b/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 diff --git a/src/BoundaryConditions/fill_halo_regions_flux.jl b/src/BoundaryConditions/fill_halo_regions_flux.jl index 96ac01fe41..1e5802ebcd 100644 --- a/src/BoundaryConditions/fill_halo_regions_flux.jl +++ b/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...) - - diff --git a/src/Coriolis/hydrostatic_spherical_coriolis.jl b/src/Coriolis/hydrostatic_spherical_coriolis.jl index efa59e178a..9be00470e4 100644 --- a/src/Coriolis/hydrostatic_spherical_coriolis.jl +++ b/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) diff --git a/src/CubedSpheres/CubedSpheres.jl b/src/CubedSpheres/CubedSpheres.jl new file mode 100644 index 0000000000..cebc9dc6d5 --- /dev/null +++ b/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 diff --git a/src/CubedSpheres/conformal_cubed_sphere_grid.jl b/src/CubedSpheres/conformal_cubed_sphere_grid.jl new file mode 100644 index 0000000000..ab9dc16141 --- /dev/null +++ b/src/CubedSpheres/conformal_cubed_sphere_grid.jl @@ -0,0 +1,296 @@ +using Rotations +using Oceananigans.Grids +using Oceananigans.Grids: R_Earth, interior_indices + +import Base: show, size, eltype +import Oceananigans.Grids: topology, domain_string + +struct CubedSphereFaceConnectivityDetails{F, S} + face :: F + side :: S +end + +short_string(deets::CubedSphereFaceConnectivityDetails) = "face $(deets.face) $(deets.side) side" + +Base.show(io::IO, deets::CubedSphereFaceConnectivityDetails) = + print(io, "CubedSphereFaceConnectivityDetails: $(short_string(deets))") + +struct CubedSphereFaceConnectivity{W, E, S, N} + west :: W + east :: E + south :: S + north :: N +end + +CubedSphereFaceConnectivity(; west, east, south, north) = + CubedSphereFaceConnectivity(west, east, south, north) + +function Base.show(io::IO, connectivity::CubedSphereFaceConnectivity) + print(io, "CubedSphereFaceConnectivity:\n", + "├── west: $(short_string(connectivity.west))\n", + "├── east: $(short_string(connectivity.east))\n", + "├── south: $(short_string(connectivity.south))\n", + "└── north: $(short_string(connectivity.north))") +end + +function default_face_connectivity() + # See figure 8.4 of https://mitgcm.readthedocs.io/en/latest/phys_pkgs/exch2.html?highlight=cube%20sphere#fig-6tile + # + # face F5 face F6 + # +----------+----------+ + # | ↑↑ | ↑↑ | + # | 1W | 1S | + # |←3N F5 6W→|←5E F6 2S→| + # | 4N | 4E | + # face F3 | ↓↓ | ↓↓ | + # +----------+----------+----------+ + # | ↑↑ | ↑↑ | + # | 5W | 5S | + # |←1N F3 4W→|←3E F4 6S→| + # | 2N | 2E | + # | ↓↓ | ↓↓ | + # +----------+----------+----------+ + # | ↑↑ | ↑↑ | face F4 + # | 3W | 3S | + # |←5N F1 2W→|←1E F2 4S→| + # | 6N | 6E | + # | ↓↓ | ↓↓ | + # +----------+----------+ + # face F1 face F2 + + face1_connectivity = CubedSphereFaceConnectivity( + west = CubedSphereFaceConnectivityDetails(5, :north), + east = CubedSphereFaceConnectivityDetails(2, :west), + south = CubedSphereFaceConnectivityDetails(6, :north), + north = CubedSphereFaceConnectivityDetails(3, :west), + ) + + face2_connectivity = CubedSphereFaceConnectivity( + west = CubedSphereFaceConnectivityDetails(1, :east), + east = CubedSphereFaceConnectivityDetails(4, :south), + south = CubedSphereFaceConnectivityDetails(6, :east), + north = CubedSphereFaceConnectivityDetails(3, :south), + ) + + face3_connectivity = CubedSphereFaceConnectivity( + west = CubedSphereFaceConnectivityDetails(1, :north), + east = CubedSphereFaceConnectivityDetails(4, :west), + south = CubedSphereFaceConnectivityDetails(2, :north), + north = CubedSphereFaceConnectivityDetails(5, :west), + ) + + face4_connectivity = CubedSphereFaceConnectivity( + west = CubedSphereFaceConnectivityDetails(3, :east), + east = CubedSphereFaceConnectivityDetails(6, :south), + south = CubedSphereFaceConnectivityDetails(2, :east), + north = CubedSphereFaceConnectivityDetails(5, :south), + ) + + face5_connectivity = CubedSphereFaceConnectivity( + west = CubedSphereFaceConnectivityDetails(3, :north), + east = CubedSphereFaceConnectivityDetails(6, :west), + south = CubedSphereFaceConnectivityDetails(4, :north), + north = CubedSphereFaceConnectivityDetails(1, :west), + ) + + + face6_connectivity = CubedSphereFaceConnectivity( + west = CubedSphereFaceConnectivityDetails(5, :east), + east = CubedSphereFaceConnectivityDetails(2, :south), + south = CubedSphereFaceConnectivityDetails(4, :east), + north = CubedSphereFaceConnectivityDetails(1, :south), + ) + + face_connectivity = ( + face1_connectivity, + face2_connectivity, + face3_connectivity, + face4_connectivity, + face5_connectivity, + face6_connectivity + ) + + return face_connectivity +end + +# Note: I think we want to keep faces and face_connectivity tuples +# so it's easy to support an arbitrary number of faces. + +struct ConformalCubedSphereGrid{FT, F, C} + faces :: F + face_connectivity :: C +end + +function ConformalCubedSphereGrid(FT=Float64; face_size, z, radius=R_Earth) + @warn "ConformalCubedSphereGrid is experimental: use with caution!" + + # +z face (face 1) + z⁺_face_grid = ConformalCubedSphereFaceGrid(FT, size=face_size, z=z, radius=radius, rotation=nothing) + + # +x face (face 2) + x⁺_face_grid = ConformalCubedSphereFaceGrid(FT, size=face_size, z=z, radius=radius, rotation=RotX(π/2)) + + # +y face (face 3) + y⁺_face_grid = ConformalCubedSphereFaceGrid(FT, size=face_size, z=z, radius=radius, rotation=RotY(π/2)) + + # -x face (face 4) + x⁻_face_grid = ConformalCubedSphereFaceGrid(FT, size=face_size, z=z, radius=radius, rotation=RotX(-π/2)) + + # -y face (face 5) + y⁻_face_grid = ConformalCubedSphereFaceGrid(FT, size=face_size, z=z, radius=radius, rotation=RotY(-π/2)) + + # -z face (face 6) + z⁻_face_grid = ConformalCubedSphereFaceGrid(FT, size=face_size, z=z, radius=radius, rotation=RotX(π)) + + faces = ( + z⁺_face_grid, + x⁺_face_grid, + y⁺_face_grid, + x⁻_face_grid, + y⁻_face_grid, + z⁻_face_grid + ) + + face_connectivity = default_face_connectivity() + + return ConformalCubedSphereGrid{FT, typeof(faces), typeof(face_connectivity)}(faces, face_connectivity) +end + +function ConformalCubedSphereGrid(filepath::AbstractString, FT=Float64; Nz, z, radius = R_Earth, halo = (1, 1, 1)) + @warn "ConformalCubedSphereGrid is experimental: use with caution!" + + face_topo = (Connected, Connected, Bounded) + face_kwargs = (Nz=Nz, z=z, topology=face_topo, radius=radius, halo=halo) + + faces = Tuple(ConformalCubedSphereFaceGrid(filepath, FT; face=n, face_kwargs...) for n in 1:6) + + face_connectivity = default_face_connectivity() + + grid = ConformalCubedSphereGrid{FT, typeof(faces), typeof(face_connectivity)}(faces, face_connectivity) + + fill_grid_metric_halos!(grid) + + return grid +end + +function Base.show(io::IO, grid::ConformalCubedSphereGrid{FT}) where FT + Nx, Ny, Nz, Nf = size(grid) + print(io, "ConformalCubedSphereGrid{$FT}: $Nf faces with size = ($Nx, $Ny, $Nz)") +end + +##### +##### Nodes for ConformalCubedSphereFaceGrid +##### + +λnode(LX::Face, LY::Face, LZ, i, j, k, grid::ConformalCubedSphereFaceGrid) = grid.λᶠᶠᵃ[i, j] +λnode(LX::Center, LY::Center, LZ, i, j, k, grid::ConformalCubedSphereFaceGrid) = grid.λᶜᶜᵃ[i, j] +φnode(LX::Face, LY::Face, LZ, i, j, k, grid::ConformalCubedSphereFaceGrid) = grid.φᶠᶠᵃ[i, j] +φnode(LX::Center, LY::Center, LZ, i, j, k, grid::ConformalCubedSphereFaceGrid) = grid.φᶜᶜᵃ[i, j] + +znode(LX, LY, LZ::Face, i, j, k, grid::ConformalCubedSphereFaceGrid) = grid.zᵃᵃᶠ[k] +znode(LX, LY, LZ::Center, i, j, k, grid::ConformalCubedSphereFaceGrid) = grid.zᵃᵃᶜ[k] + +# FIXME! +λnode(LX::Face, LY::Center, LZ, i, j, k, grid::ConformalCubedSphereFaceGrid) = grid.λᶠᶠᵃ[i, j] +λnode(LX::Center, LY::Face, LZ, i, j, k, grid::ConformalCubedSphereFaceGrid) = grid.λᶠᶠᵃ[i, j] +φnode(LX::Face, LY::Center, LZ, i, j, k, grid::ConformalCubedSphereFaceGrid) = grid.φᶠᶠᵃ[i, j] +φnode(LX::Center, LY::Face, LZ, i, j, k, grid::ConformalCubedSphereFaceGrid) = grid.φᶠᶠᵃ[i, j] + +λnodes(LX::Face, LY::Face, LZ, grid::ConformalCubedSphereFaceGrid{TX, TY}) where {TX, TY} = + view(grid.λᶠᶠᵃ, interior_indices(LX, TX, grid.Nx), interior_indices(LY, TY, grid.Ny)) + +λnodes(LX::Center, LY::Center, LZ, grid::ConformalCubedSphereFaceGrid{TX, TY}) where {TX, TY} = + view(grid.λᶜᶜᵃ, interior_indices(LX, TX, grid.Nx), interior_indices(LY, TY, grid.Ny)) + +φnodes(LX::Face, LY::Face, LZ, grid::ConformalCubedSphereFaceGrid{TX, TY}) where {TX, TY} = + view(grid.φᶠᶠᵃ, interior_indices(LX, TX, grid.Nx), interior_indices(LY, TY, grid.Ny)) + +φnodes(LX::Center, LY::Center, LZ, grid::ConformalCubedSphereFaceGrid{TX, TY}) where {TX, TY} = + view(grid.φᶜᶜᵃ, interior_indices(LX, TX, grid.Nx), interior_indices(LY, TY, grid.Ny)) + +# Nodes for ::ConformalCubedSphereGrid +# Not sure how to best represent these so will concatenate along dim 3 for now. + +λnodes(LX, LY, LZ, grid::ConformalCubedSphereGrid) = cat(Tuple(λnodes(LX, LY, LZ, grid_face) for grid_face in grid.faces)..., dims=3) +φnodes(LX, LY, LZ, grid::ConformalCubedSphereGrid) = cat(Tuple(φnodes(LX, LY, LZ, grid_face) for grid_face in grid.faces)..., dims=3) + +##### +##### Grid utils +##### + +Base.size(grid::ConformalCubedSphereGrid) = (size(grid.faces[1])..., length(grid.faces)) + +Base.eltype(grid::ConformalCubedSphereGrid{FT}) where FT = FT + +topology(::ConformalCubedSphereGrid) = (Bounded, Bounded, Bounded) + +# Not sure what to put. Gonna leave it blank so that Base.show(io::IO, operation::AbstractOperation) doesn't error. +domain_string(grid::ConformalCubedSphereFaceGrid) = "" +domain_string(grid::ConformalCubedSphereGrid) = "" + +##### +##### filling grid halos +##### + +function grid_metric_halo(grid_metric, grid, location, side) + LX, LY = location + side == :west && return underlying_west_halo(grid_metric, grid, LX) + side == :east && return underlying_east_halo(grid_metric, grid, LX) + side == :south && return underlying_south_halo(grid_metric, grid, LY) + side == :north && return underlying_north_halo(grid_metric, grid, LY) +end + +function grid_metric_boundary(grid_metric, grid, location, side) + LX, LY = location + side == :west && return underlying_west_boundary(grid_metric, grid, LX) + side == :east && return underlying_east_boundary(grid_metric, grid, LX) + side == :south && return underlying_south_boundary(grid_metric, grid, LY) + side == :north && return underlying_north_boundary(grid_metric, grid, LY) +end + +function fill_grid_metric_halos!(grid) + + loc_cc = (Center, Center) + loc_cf = (Center, Face ) + loc_fc = (Face, Center) + loc_ff = (Face, Face ) + + for face_number in 1:6, side in (:west, :east, :south, :north) + + connectivity_info = getproperty(grid.face_connectivity[face_number], side) + src_face_number = connectivity_info.face + src_side = connectivity_info.side + + grid_face = grid.faces[face_number] + src_grid_face = grid.faces[src_face_number] + + if sides_in_the_same_dimension(side, src_side) + grid_metric_halo(grid_face.Δxᶜᶜᵃ, grid_face, loc_cc, side) .= grid_metric_boundary(grid_face.Δxᶜᶜᵃ, src_grid_face, loc_cc, src_side) + grid_metric_halo(grid_face.Δyᶜᶜᵃ, grid_face, loc_cc, side) .= grid_metric_boundary(grid_face.Δyᶜᶜᵃ, src_grid_face, loc_cc, src_side) + + grid_metric_halo(grid_face.Δxᶜᶠᵃ, grid_face, loc_cf, side) .= grid_metric_boundary(grid_face.Δxᶜᶠᵃ, src_grid_face, loc_cf, src_side) + grid_metric_halo(grid_face.Δyᶜᶠᵃ, grid_face, loc_cf, side) .= grid_metric_boundary(grid_face.Δyᶜᶠᵃ, src_grid_face, loc_cf, src_side) + + grid_metric_halo(grid_face.Δxᶠᶜᵃ, grid_face, loc_fc, side) .= grid_metric_boundary(grid_face.Δxᶠᶜᵃ, src_grid_face, loc_fc, src_side) + grid_metric_halo(grid_face.Δyᶠᶜᵃ, grid_face, loc_fc, side) .= grid_metric_boundary(grid_face.Δyᶠᶜᵃ, src_grid_face, loc_fc, src_side) + + grid_metric_halo(grid_face.Δxᶠᶠᵃ, grid_face, loc_ff, side) .= grid_metric_boundary(grid_face.Δxᶠᶠᵃ, src_grid_face, loc_ff, src_side) + grid_metric_halo(grid_face.Δyᶠᶠᵃ, grid_face, loc_ff, side) .= grid_metric_boundary(grid_face.Δyᶠᶠᵃ, src_grid_face, loc_ff, src_side) + else + reverse_dim = src_side in (:west, :east) ? 1 : 2 + grid_metric_halo(grid_face.Δxᶜᶜᵃ, grid_face, loc_cc, side) .= reverse(permutedims(grid_metric_boundary(grid_face.Δyᶜᶜᵃ, src_grid_face, loc_cc, src_side), (2, 1, 3)), dims=reverse_dim) + grid_metric_halo(grid_face.Δyᶜᶜᵃ, grid_face, loc_cc, side) .= reverse(permutedims(grid_metric_boundary(grid_face.Δxᶜᶜᵃ, src_grid_face, loc_cc, src_side), (2, 1, 3)), dims=reverse_dim) + + grid_metric_halo(grid_face.Δxᶜᶠᵃ, grid_face, loc_cf, side) .= reverse(permutedims(grid_metric_boundary(grid_face.Δyᶠᶜᵃ, src_grid_face, loc_fc, src_side), (2, 1, 3)), dims=reverse_dim) + grid_metric_halo(grid_face.Δyᶜᶠᵃ, grid_face, loc_cf, side) .= reverse(permutedims(grid_metric_boundary(grid_face.Δxᶠᶜᵃ, src_grid_face, loc_fc, src_side), (2, 1, 3)), dims=reverse_dim) + + grid_metric_halo(grid_face.Δxᶠᶜᵃ, grid_face, loc_fc, side) .= reverse(permutedims(grid_metric_boundary(grid_face.Δyᶜᶠᵃ, src_grid_face, loc_cf, src_side), (2, 1, 3)), dims=reverse_dim) + grid_metric_halo(grid_face.Δyᶠᶜᵃ, grid_face, loc_fc, side) .= reverse(permutedims(grid_metric_boundary(grid_face.Δxᶜᶠᵃ, src_grid_face, loc_cf, src_side), (2, 1, 3)), dims=reverse_dim) + + grid_metric_halo(grid_face.Δxᶠᶠᵃ, grid_face, loc_ff, side) .= reverse(permutedims(grid_metric_boundary(grid_face.Δyᶠᶠᵃ, src_grid_face, loc_ff, src_side), (2, 1, 3)), dims=reverse_dim) + grid_metric_halo(grid_face.Δyᶠᶠᵃ, grid_face, loc_ff, side) .= reverse(permutedims(grid_metric_boundary(grid_face.Δxᶠᶠᵃ, src_grid_face, loc_ff, src_side), (2, 1, 3)), dims=reverse_dim) + end + end + + return nothing +end diff --git a/src/CubedSpheres/cubed_sphere_exchange_bcs.jl b/src/CubedSpheres/cubed_sphere_exchange_bcs.jl new file mode 100644 index 0000000000..84c807faf3 --- /dev/null +++ b/src/CubedSpheres/cubed_sphere_exchange_bcs.jl @@ -0,0 +1,70 @@ +using Oceananigans.BoundaryConditions +using Oceananigans.BoundaryConditions: BCType + +import Base: show +import Oceananigans.BoundaryConditions: bctype_str, print_condition + +struct CubedSphereExchange <: BCType end + +const CubedSphereExchangeBC = BoundaryCondition{<:CubedSphereExchange} + +bctype_str(::CubedSphereExchangeBC) ="CubedSphereExchange" + +CubedSphereExchangeBoundaryCondition(val; kwargs...) = BoundaryCondition(CubedSphereExchange, val; kwargs...) + +struct CubedSphereExchangeInformation{F, S} + from_face :: F + to_face :: F + from_side :: S + to_side :: S +end + +CubedSphereExchangeInformation(; from_face, to_face, from_side, to_side) = + CubedSphereExchangeInformation(from_face, to_face, from_side, to_side) + +Base.show(io::IO, ex::CubedSphereExchangeInformation) = + print(io, "CubedSphereExchangeInformation: (from: face $(ex.from_face) $(ex.from_side) side, to: face $(ex.to_face) $(ex.to_side) side)") + +print_condition(info::CubedSphereExchangeInformation) = + "(from: face $(info.from_face) $(info.from_side) side, to: face $(info.to_face) $(info.to_side) side)" + +function inject_cubed_sphere_exchange_boundary_conditions(field_bcs, face_number, face_connectivity) + + west_exchange_info = CubedSphereExchangeInformation( + from_face = face_number, + from_side = :west, + to_face = face_connectivity[face_number].west.face, + to_side = face_connectivity[face_number].west.side + ) + + east_exchange_info = CubedSphereExchangeInformation( + from_face = face_number, + from_side = :east, + to_face = face_connectivity[face_number].east.face, + to_side = face_connectivity[face_number].east.side + ) + + south_exchange_info = CubedSphereExchangeInformation( + from_face = face_number, + from_side = :south, + to_face = face_connectivity[face_number].south.face, + to_side = face_connectivity[face_number].south.side + ) + + north_exchange_info = CubedSphereExchangeInformation( + from_face = face_number, + from_side = :north, + to_face = face_connectivity[face_number].north.face, + to_side = face_connectivity[face_number].north.side + ) + + west_exchange_bc = CubedSphereExchangeBoundaryCondition(west_exchange_info) + east_exchange_bc = CubedSphereExchangeBoundaryCondition(east_exchange_info) + south_exchange_bc = CubedSphereExchangeBoundaryCondition(south_exchange_info) + north_exchange_bc = CubedSphereExchangeBoundaryCondition(north_exchange_info) + + x_bcs = CoordinateBoundaryConditions(west_exchange_bc, east_exchange_bc) + y_bcs = CoordinateBoundaryConditions(south_exchange_bc, north_exchange_bc) + + return FieldBoundaryConditions(x_bcs, y_bcs, field_bcs.z) +end diff --git a/src/CubedSpheres/cubed_sphere_fields.jl b/src/CubedSpheres/cubed_sphere_fields.jl new file mode 100644 index 0000000000..14eda5adbf --- /dev/null +++ b/src/CubedSpheres/cubed_sphere_fields.jl @@ -0,0 +1,131 @@ +using Statistics + +using Oceananigans.Fields: AbstractField, call_func + +import Base: getindex, size, show, minimum, maximum +import Statistics: mean +import Oceananigans.Fields: Field, CenterField, XFaceField, YFaceField, ZFaceField, FunctionField, ReducedField, interior +import Oceananigans.BoundaryConditions: fill_halo_regions! + +abstract type AbstractCubedSphereField{X, Y, Z, A, G} <: AbstractField{X, Y, Z, A, G} end + +struct ConformalCubedSphereField{X, Y, Z, A, G, F, B} <: AbstractCubedSphereField{X, Y, Z, A, G} + grid :: G + faces :: F + boundary_conditions :: B +end + +##### +##### Regular fields +##### + +Field(X, Y, Z, arch, grid::ConformalCubedSphereGrid, ::Nothing, data) = + Field(X, Y, Z, arch, grid, FieldBoundaryConditions(grid, (X, Y, Z)), data) + +function Field(X, Y, Z, arch, grid::ConformalCubedSphereGrid, bcs, data) + + faces = Tuple( + Field(X, Y, Z, arch, face_grid, inject_cubed_sphere_exchange_boundary_conditions(bcs, face_number, grid.face_connectivity)) + for (face_number, face_grid) in enumerate(grid.faces) + ) + + # This field needs BCs otherwise errors happen so I'll assume all faces have + # the same boundary conditions. A very bad assumption... + cubed_sphere_bcs = faces[1].boundary_conditions + + return ConformalCubedSphereField{X, Y, Z, typeof(arch), typeof(grid), typeof(faces), typeof(cubed_sphere_bcs)}(grid, faces, cubed_sphere_bcs) +end + +##### +##### Convinient field constructors +##### + +# Gotta short-circuit for `::ConformalCubedSphereGrid` because these +# Field constructors allocate `data` in the function signature. + +CenterField(FT::DataType, arch, grid::ConformalCubedSphereGrid, bcs=nothing, data=nothing) = Field(Center, Center, Center, arch, grid, bcs, data) + XFaceField(FT::DataType, arch, grid::ConformalCubedSphereGrid, bcs=nothing, data=nothing) = Field(Face, Center, Center, arch, grid, bcs, data) + YFaceField(FT::DataType, arch, grid::ConformalCubedSphereGrid, bcs=nothing, data=nothing) = Field(Center, Face, Center, arch, grid, bcs, data) + ZFaceField(FT::DataType, arch, grid::ConformalCubedSphereGrid, bcs=nothing, data=nothing) = Field(Center, Center, Face, arch, grid, bcs, data) + +##### +##### Utils +##### + +Base.size(field::AbstractCubedSphereField) = (size(field.faces[1])..., length(field.faces)) + +Base.minimum(field::AbstractCubedSphereField; dims=:) = minimum(minimum(field_face; dims) for field_face in field.faces) +Base.maximum(field::AbstractCubedSphereField; dims=:) = maximum(maximum(field_face; dims) for field_face in field.faces) +Statistics.mean(field::AbstractCubedSphereField; dims=:) = mean(mean(field_face; dims) for field_face in field.faces) + +Base.minimum(f, field::AbstractCubedSphereField; dims=:) = minimum(minimum(f, field_face; dims) for field_face in field.faces) +Base.maximum(f, field::AbstractCubedSphereField; dims=:) = maximum(maximum(f, field_face; dims) for field_face in field.faces) +Statistics.mean(f, field::AbstractCubedSphereField; dims=:) = mean(mean(f, field_face; dims) for field_face in field.faces) + +interior(field::AbstractCubedSphereField) = cat(Tuple(interior(field_face) for field_face in field.faces)..., dims=4) + +const ConformalCubedSphereFaceField{LX, LY, LZ, A} = AbstractField{LX, LY, LZ, A, <:ConformalCubedSphereFaceGrid} + +λnodes(field::ConformalCubedSphereFaceField{LX, LY, LZ}) where {LX, LY, LZ} = λnodes(LX(), LY(), LZ(), field.grid) +φnodes(field::ConformalCubedSphereFaceField{LX, LY, LZ}) where {LX, LY, LZ} = φnodes(LX(), LY(), LZ(), field.grid) + +λnodes(field::ConformalCubedSphereField{LX, LY, LZ}) where {LX, LY, LZ} = λnodes(LX(), LY(), LZ(), field.grid) +φnodes(field::ConformalCubedSphereField{LX, LY, LZ}) where {LX, LY, LZ} = φnodes(LX(), LY(), LZ(), field.grid) + +##### +##### Function fields +##### + +struct ConformalCubedSphereFunctionField{X, Y, Z, F, G} <: AbstractCubedSphereField{X, Y, Z, F, G} + faces :: F +end + +function FunctionField{X, Y, Z}(func, grid::ConformalCubedSphereGrid; clock=nothing, parameters=nothing) where {X, Y, Z} + + faces = Tuple(FunctionField{X, Y, Z}(func, face_grid; clock, parameters) for face_grid in grid.faces) + + return ConformalCubedSphereFunctionField{X, Y, Z, typeof(faces), typeof(grid)}(faces) +end + +const ConformalCubedSphereFaceFunctionField = FunctionField{X, Y, Z, C, P, F, <:ConformalCubedSphereFaceGrid} where {X, Y, Z, C, P, F} + +@inline Base.getindex(f::ConformalCubedSphereFaceFunctionField{X, Y, Z}, i, j, k) where {X, Y, Z} = + call_func(f.clock, f.parameters, f.func, + λnode(X(), Y(), Z(), i, j, k, f.grid), + φnode(X(), Y(), Z(), i, j, k, f.grid), + znode(Z(), Y(), Z(), i, j, k, f.grid)) + +fill_halo_regions!(::ConformalCubedSphereFunctionField, args...) = nothing + +##### +##### Reduced fields +##### + +struct ConformalCubedSphereReducedField{X, Y, Z, A, G, F, B} <: AbstractCubedSphereField{X, Y, Z, A, G} + grid :: G + faces :: F + boundary_conditions :: B +end + +function ReducedField(X, Y, Z, arch, grid::ConformalCubedSphereGrid; dims, data=nothing, boundary_conditions=nothing) + + faces = Tuple( + ReducedField(X, Y, Z, arch, grid_face, dims=dims, data=data, boundary_conditions=inject_cubed_sphere_exchange_boundary_conditions(FieldBoundaryConditions(grid_face, (X, Y, Z)), face_number, grid.face_connectivity)) + for (face_number, grid_face) in enumerate(grid.faces) + ) + + # This field needs BCs otherwise errors happen so I'll assume all faces have + # the same boundary conditions. A very bad assumption... + cubed_sphere_bcs = faces[1].boundary_conditions + + return ConformalCubedSphereReducedField{X, Y, Z, typeof(arch), typeof(grid), typeof(faces), typeof(cubed_sphere_bcs)}(grid, faces, cubed_sphere_bcs) +end + +##### +##### Pretty printing +##### + +function Base.show(io::IO, field::AbstractCubedSphereField{X, Y, Z}) where {X, Y, Z} + Nx, Ny, Nz, Nf = size(field) + print(io, "$(Base.typename(typeof(field))){$X, $Y, $Z}: $Nf faces with size = ($Nx, $Ny, $Nz)") +end diff --git a/src/CubedSpheres/cubed_sphere_halo_filling.jl b/src/CubedSpheres/cubed_sphere_halo_filling.jl new file mode 100644 index 0000000000..9917b12256 --- /dev/null +++ b/src/CubedSpheres/cubed_sphere_halo_filling.jl @@ -0,0 +1,135 @@ +import Oceananigans.BoundaryConditions: + fill_halo_regions!, fill_west_halo!, fill_east_halo!, fill_south_halo!, fill_north_halo! + +import Oceananigans.Models.HydrostaticFreeSurfaceModels: fill_horizontal_velocity_halos! + +# These filling functions won't work so let's not use them. + + fill_west_halo!(c, bc::CubedSphereExchangeBC, args...; kwargs...) = nothing + fill_east_halo!(c, bc::CubedSphereExchangeBC, args...; kwargs...) = nothing +fill_south_halo!(c, bc::CubedSphereExchangeBC, args...; kwargs...) = nothing +fill_north_halo!(c, bc::CubedSphereExchangeBC, args...; kwargs...) = nothing + +function fill_halo_regions!(field::AbstractCubedSphereField{LX, LY, LZ}, arch, args...) where {LX, LY, LZ} + + for field_face in field.faces + # Fill the top and bottom halos the usual way. + fill_halo_regions!(field_face, arch, args...) + + # Deal with halo exchanges. + fill_west_halo!(field_face, field) + fill_east_halo!(field_face, field) + fill_south_halo!(field_face, field) + fill_north_halo!(field_face, field) + end + + return nothing +end + +function fill_west_halo!(field::ConformalCubedSphereFaceField{LX, LY, LZ}, cubed_sphere_field) where {LX, LY, LZ} + location = (LX, LY, LZ) + dest_halo = underlying_west_halo(field.data, field.grid, LX) + + exchange_info = field.boundary_conditions.west.condition + src_face_number = exchange_info.to_face + src_side = exchange_info.to_side + src_boundary = cubed_sphere_boundary(cubed_sphere_field, location, src_face_number, src_side) + + if sides_in_the_same_dimension(:west, src_side) + dest_halo .= src_boundary + else + dest_halo .= reverse(permutedims(src_boundary, (2, 1, 3)), dims=2) + end + + return nothing +end + +function fill_east_halo!(field::ConformalCubedSphereFaceField{LX, LY, LZ}, cubed_sphere_field) where {LX, LY, LZ} + location = (LX, LY, LZ) + dest_halo = underlying_east_halo(field.data, field.grid, LX) + + exchange_info = field.boundary_conditions.east.condition + src_face_number = exchange_info.to_face + src_side = exchange_info.to_side + src_boundary = cubed_sphere_boundary(cubed_sphere_field, location, src_face_number, src_side) + + if sides_in_the_same_dimension(:east, src_side) + dest_halo .= src_boundary + else + dest_halo .= reverse(permutedims(src_boundary, (2, 1, 3)), dims=2) + end + + return nothing +end + +function fill_south_halo!(field::ConformalCubedSphereFaceField{LX, LY, LZ}, cubed_sphere_field) where {LX, LY, LZ} + location = (LX, LY, LZ) + dest_halo = underlying_south_halo(field.data, field.grid, LY) + + exchange_info = field.boundary_conditions.south.condition + src_face_number = exchange_info.to_face + src_side = exchange_info.to_side + src_boundary = cubed_sphere_boundary(cubed_sphere_field, location, src_face_number, src_side) + + if sides_in_the_same_dimension(:south, src_side) + dest_halo .= src_boundary + else + dest_halo .= reverse(permutedims(src_boundary, (2, 1, 3)), dims=1) + end + + return nothing +end + +function fill_north_halo!(field::ConformalCubedSphereFaceField{LX, LY, LZ}, cubed_sphere_field) where {LX, LY, LZ} + location = (LX, LY, LZ) + dest_halo = underlying_north_halo(field.data, field.grid, LY) + + exchange_info = field.boundary_conditions.north.condition + src_face_number = exchange_info.to_face + src_side = exchange_info.to_side + src_boundary = cubed_sphere_boundary(cubed_sphere_field, location, src_face_number, src_side) + + if sides_in_the_same_dimension(:north, src_side) + dest_halo .= src_boundary + else + dest_halo .= reverse(permutedims(src_boundary, (2, 1, 3)), dims=1) + end + + return nothing +end + +# Don't worry about this when not on a cubed sphere. +fill_horizontal_velocity_halos!(u, v, arch) = nothing + +function fill_horizontal_velocity_halos!(u::ConformalCubedSphereField, v::ConformalCubedSphereField, arch) + + ## Fill them like they're tracers to get the top and bottom filled. + fill_halo_regions!(u, arch) + fill_halo_regions!(v, arch) + + ## Now fill the horizontal halos. + + u_loc = (Face, Center, Center) + v_loc = (Center, Face, Center) + + for face_number in 1:6, side in (:west, :east, :south, :north) + exchange_info = getproperty(u.faces[face_number].boundary_conditions, side).condition + src_face_number = exchange_info.to_face + src_side = exchange_info.to_side + + if sides_in_the_same_dimension(side, src_side) + cubed_sphere_halo(u, u_loc, face_number, side) .= cubed_sphere_boundary(u, u_loc, src_face_number, src_side) + cubed_sphere_halo(v, v_loc, face_number, side) .= cubed_sphere_boundary(v, v_loc, src_face_number, src_side) + else + u_sign = (isodd(face_number) && side == :west ) || (iseven(face_number) && side == :east ) ? +1 : -1 + v_sign = (isodd(face_number) && side == :north) || (iseven(face_number) && side == :south) ? +1 : -1 + + reverse_dim = src_side in (:west, :east) ? 1 : 2 + + cubed_sphere_halo(u, u_loc, face_number, side) .= u_sign * reverse(permutedims(cubed_sphere_boundary(v, v_loc, src_face_number, src_side), (2, 1, 3)), dims=reverse_dim) + cubed_sphere_halo(v, v_loc, face_number, side) .= v_sign * reverse(permutedims(cubed_sphere_boundary(u, u_loc, src_face_number, src_side), (2, 1, 3)), dims=reverse_dim) + end + end + + return nothing +end diff --git a/src/CubedSpheres/cubed_sphere_kernel_launching.jl b/src/CubedSpheres/cubed_sphere_kernel_launching.jl new file mode 100644 index 0000000000..95c1a1b1cc --- /dev/null +++ b/src/CubedSpheres/cubed_sphere_kernel_launching.jl @@ -0,0 +1,31 @@ +using KernelAbstractions: Event, MultiEvent +using Oceananigans.Architectures: device + +using Oceananigans.Utils: launch! + +import Oceananigans.Utils: launch! + +maybe_replace_with_face(elem, cubed_sphere_grid, face_number) = elem + +maybe_replace_with_face(grid::ConformalCubedSphereGrid, cubed_sphere_grid, face_number) = grid.faces[face_number] +maybe_replace_with_face(field::AbstractCubedSphereField, cubed_sphere_grid, face_number) = field.faces[face_number] + +maybe_replace_with_face(t::Tuple, cubed_sphere_grid, face_number) = Tuple(maybe_replace_with_face(t_elem, cubed_sphere_grid, face_number) for t_elem in t) +maybe_replace_with_face(nt::NamedTuple, cubed_sphere_grid, face_number) = NamedTuple{keys(nt)}(maybe_replace_with_face(nt_elem, cubed_sphere_grid, face_number) for nt_elem in nt) + +function launch!(arch, grid::ConformalCubedSphereGrid, args...; kwargs...) + + events = [] + for (face_number, grid_face) in enumerate(grid.faces) + new_args = Tuple(maybe_replace_with_face(elem, grid, face_number) for elem in args) + event = launch!(arch, grid_face, new_args...; kwargs...) + push!(events, event) + end + + # We should return the events but let's just wait here because errors. + events = filter(e -> e isa Event, events) + wait(device(arch), MultiEvent(Tuple(events))) + + # TODO: Other function expect an `event` to be returned. + return events[1] +end diff --git a/src/CubedSpheres/cubed_sphere_set!.jl b/src/CubedSpheres/cubed_sphere_set!.jl new file mode 100644 index 0000000000..c1c9d7e00e --- /dev/null +++ b/src/CubedSpheres/cubed_sphere_set!.jl @@ -0,0 +1,59 @@ +using Oceananigans.Fields: AbstractField + +import Oceananigans.Fields: set! + +function set!(u::ConformalCubedSphereField, v::ConformalCubedSphereField) + for (u_face, v_face) in zip(u.faces, v.faces) + @. u_face.data.parent = v_face.data.parent + end + return nothing +end + +const ConformalCubedSphereFaceFieldᶜᶜᶜ = AbstractField{Center, Center, Center, A, <:ConformalCubedSphereFaceGrid} where A +const ConformalCubedSphereFaceFieldᶠᶜᶜ = AbstractField{Face, Center, Center, A, <:ConformalCubedSphereFaceGrid} where A +const ConformalCubedSphereFaceFieldᶜᶠᶜ = AbstractField{Center, Face, Center, A, <:ConformalCubedSphereFaceGrid} where A +const ConformalCubedSphereFaceFieldᶠᶠᶜ = AbstractField{Face, Face, Center, A, <:ConformalCubedSphereFaceGrid} where A + +const ConformalCubedSphereFaceFieldᶜᶜⁿ = AbstractField{Center, Center, Nothing, A, <:ConformalCubedSphereFaceGrid} where A + +function set!(field::ConformalCubedSphereFaceFieldᶜᶜᶜ, f::Function) + grid = field.grid + for i in 1:grid.Nx, j in 1:grid.Ny, k in 1:grid.Nz + field[i, j, k] = f(grid.λᶜᶜᵃ[i, j], grid.φᶜᶜᵃ[i, j], grid.zᵃᵃᶜ[k]) + end + return nothing +end + +function set!(field::ConformalCubedSphereFaceFieldᶠᶜᶜ, f::Function) + grid = field.grid + for i in 1:grid.Nx, j in 1:grid.Ny, k in 1:grid.Nz + field[i, j, k] = f(grid.λᶠᶜᵃ[i, j], grid.φᶠᶜᵃ[i, j], grid.zᵃᵃᶜ[k]) + end + return nothing +end + +function set!(field::ConformalCubedSphereFaceFieldᶜᶠᶜ, f::Function) + grid = field.grid + for i in 1:grid.Nx, j in 1:grid.Ny, k in 1:grid.Nz + field[i, j, k] = f(grid.λᶜᶠᵃ[i, j], grid.φᶜᶠᵃ[i, j], grid.zᵃᵃᶜ[k]) + end + return nothing +end + +function set!(field::ConformalCubedSphereFaceFieldᶠᶠᶜ, f::Function) + grid = field.grid + for i in 1:grid.Nx, j in 1:grid.Ny, k in 1:grid.Nz + field[i, j, k] = f(grid.λᶠᶠᵃ[i, j], grid.φᶠᶠᵃ[i, j], grid.zᵃᵃᶜ[k]) + end + return nothing +end + +function set!(field::ConformalCubedSphereFaceFieldᶜᶜⁿ, f::Function) + grid = field.grid + for i in 1:grid.Nx, j in 1:grid.Ny + field[i, j, 1] = f(grid.λᶜᶜᵃ[i, j], grid.φᶜᶜᵃ[i, j]) + end + return nothing +end + +set!(field::AbstractCubedSphereField, f::Function) = [set!(field_face, f) for field_face in field.faces] diff --git a/src/CubedSpheres/cubed_sphere_utils.jl b/src/CubedSpheres/cubed_sphere_utils.jl new file mode 100644 index 0000000000..7c51bf229c --- /dev/null +++ b/src/CubedSpheres/cubed_sphere_utils.jl @@ -0,0 +1,135 @@ +using Oceananigans.Fields: AbstractField +using Oceananigans.Grids: + Face, Bounded, + interior_indices, + left_halo_indices, right_halo_indices, + underlying_left_halo_indices, underlying_right_halo_indices + +# TODO: Move to Grids/grid_utils.jl + +##### +##### Viewing halos +##### + +west_halo(f::AbstractField{LX, LY, LZ}; include_corners=true) where {LX, LY, LZ} = + include_corners ? view(f.data, left_halo_indices(LX, topology(f, 1), f.grid.Nx, f.grid.Hx), :, :) : + view(f.data, left_halo_indices(LX, topology(f, 1), f.grid.Nx, f.grid.Hx), + interior_indices(LY, topology(f, 2), f.grid.Ny), + interior_indices(LZ, topology(f, 3), f.grid.Nz)) + +east_halo(f::AbstractField{LX, LY, LZ}; include_corners=true) where {LX, LY, LZ} = + include_corners ? view(f.data, right_halo_indices(LX, topology(f, 1), f.grid.Nx, f.grid.Hx), :, :) : + view(f.data, right_halo_indices(LX, topology(f, 1), f.grid.Nx, f.grid.Hx), + interior_indices(LY, topology(f, 2), f.grid.Ny), + interior_indices(LZ, topology(f, 3), f.grid.Nz)) + +south_halo(f::AbstractField{LX, LY, LZ}; include_corners=true) where {LX, LY, LZ} = + include_corners ? view(f.data, :, left_halo_indices(LY, topology(f, 2), f.grid.Ny, f.grid.Hy), :) : + view(f.data, interior_indices(LX, topology(f, 1), f.grid.Nx), + left_halo_indices(LY, topology(f, 2), f.grid.Ny, f.grid.Hy), + interior_indices(LZ, topology(f, 3), f.grid.Nz)) + +north_halo(f::AbstractField{LX, LY, LZ}; include_corners=true) where {LX, LY, LZ} = + include_corners ? view(f.data, :, right_halo_indices(LY, topology(f, 2), f.grid.Ny, f.grid.Hy), :) : + view(f.data, interior_indices(LX, topology(f, 1), f.grid.Nx), + right_halo_indices(LY, topology(f, 2), f.grid.Ny, f.grid.Hy), + interior_indices(LZ, topology(f, 3), f.grid.Nz)) + +bottom_halo(f::AbstractField{LX, LY, LZ}; include_corners=true) where {LX, LY, LZ} = + include_corners ? view(f.data, :, :, left_halo_indices(LZ, topology(f, 3), f.grid.Nz, f.grid.Hz)) : + view(f.data, interior_indices(LX, topology(f, 1), f.grid.Nx), + interior_indices(LY, topology(f, 2), f.grid.Ny), + left_halo_indices(LZ, topology(f, 3), f.grid.Nz, f.grid.Hz)) + +top_halo(f::AbstractField{LX, LY, LZ}; include_corners=true) where {LX, LY, LZ} = + include_corners ? view(f.data, :, :, right_halo_indices(LZ, topology(f, 3), f.grid.Nz, f.grid.Hz)) : + view(f.data, interior_indices(LX, topology(f, 1), f.grid.Nx), + interior_indices(LY, topology(f, 2), f.grid.Ny), + right_halo_indices(LZ, topology(f, 3), f.grid.Nz, f.grid.Hz)) + +underlying_west_halo(f, grid, location) = + view(f.parent, underlying_left_halo_indices(location, topology(grid, 1), grid.Nx, grid.Hx), :, :) + +underlying_east_halo(f, grid, location) = + view(f.parent, underlying_right_halo_indices(location, topology(grid, 1), grid.Nx, grid.Hx), :, :) + +underlying_south_halo(f, grid, location) = + view(f.parent, :, underlying_left_halo_indices(location, topology(grid, 2), grid.Ny, grid.Hy), :) + +underlying_north_halo(f, grid, location) = + view(f.parent, :, underlying_right_halo_indices(location, topology(grid, 2), grid.Ny, grid.Hy), :) + +underlying_bottom_halo(f, grid, location) = + view(f.parent, :, :, underlying_left_halo_indices(location, topology(grid, 3), grid.Nz, grid.Hz)) + +underlying_top_halo(f, grid, location) = + view(f.parent, :, :, underlying_right_halo_indices(location, topology(grid, 3), grid.Nz, grid.Hz)) + +##### +##### Viewing boundary grid points (used to fill other halos) +##### + +left_boundary_indices(loc, topo, N, H) = 1:H +left_boundary_indices(::Type{Nothing}, topo, N, H) = 1:0 # empty + +right_boundary_indices(loc, topo, N, H) = N-H+1:N +right_boundary_indices(::Type{Face}, ::Type{Bounded}, N, H) = N-H:N+1 +right_boundary_indices(::Type{Nothing}, topo, N, H) = 1:0 # empty + +underlying_left_boundary_indices(loc, topo, N, H) = 1+H:2H +underlying_left_boundary_indices(::Type{Nothing}, topo, N, H) = 1:0 # empty + +underlying_right_boundary_indices(loc, topo, N, H) = N+1:N+H +underlying_right_boundary_indices(::Type{Face}, ::Type{Bounded}, N, H) = N+2:N+H+1 +underlying_right_boundary_indices(::Type{Nothing}, topo, N, H) = 1:0 # empty + +underlying_west_boundary(f, grid, location) = + view(f.parent, underlying_left_boundary_indices(location, topology(grid, 1), grid.Nx, grid.Hx), :, :) + +underlying_east_boundary(f, grid, location) = + view(f.parent, underlying_right_boundary_indices(location, topology(grid, 1), grid.Nx, grid.Hx), :, :) + +underlying_south_boundary(f, grid, location) = + view(f.parent, :, underlying_left_boundary_indices(location, topology(grid, 2), grid.Ny, grid.Hy), :) + +underlying_north_boundary(f, grid, location) = + view(f.parent, :, underlying_right_boundary_indices(location, topology(grid, 2), grid.Ny, grid.Hy), :) + +underlying_bottom_boundary(f, grid, location) = + view(f.parent, :, :, underlying_left_boundary_indices(location, topology(grid, 3), grid.Nz, grid.Hz)) + +underlying_top_boundary(f, grid, location) = + view(f.parent, :, :, underlying_right_boundary_indices(location, topology(grid, 3), grid.Nz, grid.Hz)) + + +##### +##### Convinience functions +##### + +function sides_in_the_same_dimension(side1, side2) + x_sides = (:west, :east) + y_sides = (:south, :north) + z_sides = (:bottom, :top) + side1 in x_sides && side2 in x_sides && return true + side1 in y_sides && side2 in y_sides && return true + side1 in z_sides && side2 in z_sides && return true + return false +end + +function cubed_sphere_halo(cubed_sphere_field, location, face_number, side) + LX, LY, LZ = location + src_field = cubed_sphere_field.faces[face_number] + side == :west && return underlying_west_halo(src_field.data, src_field.grid, LX) + side == :east && return underlying_east_halo(src_field.data, src_field.grid, LX) + side == :south && return underlying_south_halo(src_field.data, src_field.grid, LY) + side == :north && return underlying_north_halo(src_field.data, src_field.grid, LY) +end + +function cubed_sphere_boundary(cubed_sphere_field, location, face_number, side) + LX, LY, LZ = location + src_field = cubed_sphere_field.faces[face_number] + side == :west && return underlying_west_boundary(src_field.data, src_field.grid, LX) + side == :east && return underlying_east_boundary(src_field.data, src_field.grid, LX) + side == :south && return underlying_south_boundary(src_field.data, src_field.grid, LY) + side == :north && return underlying_north_boundary(src_field.data, src_field.grid, LY) +end diff --git a/src/Diagnostics/Diagnostics.jl b/src/Diagnostics/Diagnostics.jl index 27b9b88a5e..1430c91a84 100644 --- a/src/Diagnostics/Diagnostics.jl +++ b/src/Diagnostics/Diagnostics.jl @@ -1,6 +1,6 @@ module Diagnostics -export NaNChecker, CFL, AdvectiveCFL, DiffusiveCFL, WindowedSpatialAverage +export NaNChecker, StateChecker, CFL, AdvectiveCFL, DiffusiveCFL, WindowedSpatialAverage using CUDA using Oceananigans @@ -9,9 +9,11 @@ using Oceananigans.Operators using Oceananigans: AbstractDiagnostic using Oceananigans.Utils: TimeInterval, IterationInterval, WallTimeInterval +import Base: show import Oceananigans: run_diagnostic! include("nan_checker.jl") +include("state_checker.jl") include("cfl.jl") include("windowed_spatial_average.jl") diff --git a/src/Diagnostics/cfl.jl b/src/Diagnostics/cfl.jl index 1df2b55489..8de03a770a 100644 --- a/src/Diagnostics/cfl.jl +++ b/src/Diagnostics/cfl.jl @@ -78,8 +78,9 @@ using Oceananigans.Models using Oceananigans.Grids: halo_size using Oceananigans.Operators: Δxᶠᶜᵃ, Δyᶜᶠᵃ, Δzᵃᵃᶠ -function accurate_cell_advection_timescale(model::Union{IncompressibleModel, HydrostaticFreeSurfaceModel}) - grid = model.grid +accurate_cell_advection_timescale(model) = accurate_cell_advection_timescale(model.grid, model.velocities) + +function accurate_cell_advection_timescale(grid, velocities) Nx, Ny, Nz = size(grid) Hx, Hy, Hz = halo_size(grid) @@ -87,9 +88,9 @@ function accurate_cell_advection_timescale(model::Union{IncompressibleModel, Hyd js = 1+Hy:Ny+Hy ks = 1+Hz:Nz+Hz - u = view(model.velocities.u.data.parent, is, js, ks) - v = view(model.velocities.v.data.parent, is, js, ks) - w = view(model.velocities.w.data.parent, is, js, ks) + u = view(velocities.u.data.parent, is, js, ks) + v = view(velocities.v.data.parent, is, js, ks) + w = view(velocities.w.data.parent, is, js, ks) min_timescale = minimum( @tullio (min) timescale[k] := 1 / ( abs(u[i, j, k]) / Δxᶠᶜᵃ(i, j, k, grid) diff --git a/src/Diagnostics/nan_checker.jl b/src/Diagnostics/nan_checker.jl index 6e6d2169f0..01b6d7fb59 100644 --- a/src/Diagnostics/nan_checker.jl +++ b/src/Diagnostics/nan_checker.jl @@ -13,9 +13,12 @@ NaNChecker(model=nothing; schedule, fields) = NaNChecker(schedule, fields) function run_diagnostic!(nc::NaNChecker, model) for (name, field) in pairs(nc.fields) - if any(isnan, field.data.parent) - t, i = model.clock.time, model.clock.iteration - error("time = $t, iteration = $i: NaN found in $name. Aborting simulation.") - end + error_if_nan_in_field(field, name, model.clock) + end +end + +function error_if_nan_in_field(field, name, clock) + if any(isnan, field.data.parent) + error("time = $(clock.time), iteration = $(clock.iteration): NaN found in field $name. Aborting simulation.") end end diff --git a/src/Diagnostics/state_checker.jl b/src/Diagnostics/state_checker.jl new file mode 100644 index 0000000000..ad03f7367d --- /dev/null +++ b/src/Diagnostics/state_checker.jl @@ -0,0 +1,45 @@ +using Printf +using Statistics + +using Oceananigans: short_show + +struct StateChecker{T, F} <: AbstractDiagnostic + schedule :: T + fields :: F +end + +""" + StateChecker(; schedule, fields) + +Returns a `StateChecker` that logs field information (minimum, maximum, mean) +for each field in a named tuple of `fields` when `schedule` actuates. +""" +StateChecker(model; schedule, fields=fields(model)) = StateChecker(schedule, fields) + +function run_diagnostic!(sc::StateChecker, model) + pad = keys(sc.fields) .|> string .|> length |> maximum + + @info "State check @ $(short_show(model.clock))" + + for (name, field) in pairs(sc.fields) + state_check(field, name, pad) + end + + return nothing +end + +function state_check(field, name, pad) + min_val = minimum(field) + max_val = maximum(field) + mean_val = mean(field) + + name = lpad(name, pad) + + @info @sprintf("%s | min = %+.15e | max = %+.15e | mean = %+.15e", name, min_val, max_val, mean_val) + return nothing +end + +(sc::StateChecker)(model) = run_diagnostic!(sc, model) + +Base.show(io::IO, sc::StateChecker) = + print(io, "StateChecker checking $(length(sc.fields)) fields: $(keys(sc.fields))") diff --git a/src/Fields/field.jl b/src/Fields/field.jl index 9637b6b5e3..17bc9a5a97 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -128,9 +128,9 @@ function ZFaceField(FT::DataType, arch, grid, end CenterField(arch::AbstractArchitecture, grid, args...) = CenterField(eltype(grid), arch, grid, args...) - XFaceField(arch::AbstractArchitecture, grid, args...) = XFaceField(eltype(grid), arch, grid, args...) - YFaceField(arch::AbstractArchitecture, grid, args...) = YFaceField(eltype(grid), arch, grid, args...) - ZFaceField(arch::AbstractArchitecture, grid, args...) = ZFaceField(eltype(grid), arch, grid, args...) + XFaceField(arch::AbstractArchitecture, grid, args...) = XFaceField(eltype(grid), arch, grid, args...) + YFaceField(arch::AbstractArchitecture, grid, args...) = YFaceField(eltype(grid), arch, grid, args...) + ZFaceField(arch::AbstractArchitecture, grid, args...) = ZFaceField(eltype(grid), arch, grid, args...) @propagate_inbounds Base.setindex!(f::Field, v, inds...) = @inbounds setindex!(f.data, v, inds...) diff --git a/src/Fields/set!.jl b/src/Fields/set!.jl index d15e5693ac..0271e323ca 100644 --- a/src/Fields/set!.jl +++ b/src/Fields/set!.jl @@ -63,22 +63,22 @@ set!(u::AbstractCPUField, f::Function) = interior(u) .= f.(nodes(u; reshape=true """ Set the GPU field `u` to the array `v`. """ function set!(u::AbstractGPUField, v::Array) v_field = similar_cpu_field(u) - + set!(v_field, v) set!(u, v_field) - + return nothing end - + """ Set the GPU field `u` to the CuArray `v`. """ function set!(u::AbstractGPUField, v::CuArray) - + launch!(GPU(), u.grid, :xyz, _set_gpu!, u.data, v, u.grid, include_right_boundaries=true, location=location(u)) - + return nothing end - + @kernel function _set_gpu!(u, v, grid) i, j, k = @index(Global, NTuple) @inbounds u[i, j, k] = v[i, j, k] @@ -86,18 +86,18 @@ set!(u::AbstractCPUField, f::Function) = interior(u) .= f.(nodes(u; reshape=true """ Set the CPU field `u` data to the GPU field data of `v`. """ set!(u::AbstractCPUField, v::AbstractGPUField) = u.data.parent .= Array(v.data.parent) - + """ Set the GPU field `u` data to the CPU field data of `v`. """ set!(u::AbstractGPUField, v::AbstractCPUField) = copyto!(u.data.parent, v.data.parent) - + """ Set the GPU field `u` data to the function `f(x, y, z)`. """ function set!(u::AbstractGPUField, f::Function) # Create a temporary field with bcs = nothing. v_field = similar_cpu_field(u) - + set!(v_field, f) set!(u, v_field) - + return nothing end end diff --git a/src/Forcings/model_forcing.jl b/src/Forcings/model_forcing.jl index 077b4b114c..1e2d5c4b22 100644 --- a/src/Forcings/model_forcing.jl +++ b/src/Forcings/model_forcing.jl @@ -23,7 +23,7 @@ function regularize_forcing(forcing::Function, field, field_name, model_field_na return ContinuousForcing{LX, LY, LZ}(forcing) end -regularize_forcing(::Nothing, field, field_name, model_field_names) = zeroforcing +regularize_forcing(::Nothing, field::AbstractField, field_name, model_field_names) = zeroforcing """ model_forcing(model_fields; forcings...) diff --git a/src/Grids/Grids.jl b/src/Grids/Grids.jl index d8a3e72824..f02711540d 100644 --- a/src/Grids/Grids.jl +++ b/src/Grids/Grids.jl @@ -2,7 +2,7 @@ module Grids export Center, Face, - AbstractTopology, Periodic, Bounded, Flat, topology, + AbstractTopology, Periodic, Bounded, Flat, Connected, topology, AbstractGrid, halo_size, AbstractRectilinearGrid, RegularRectilinearGrid, VerticallyStretchedRectilinearGrid, AbstractCurvilinearGrid, AbstractHorizontallyCurvilinearGrid, @@ -67,6 +67,13 @@ is uniform and does not vary. """ struct Flat <: AbstractTopology end +""" + Connected + +Grid topology for dimensions that are connected to other models or domains on both sides. +""" +const Connected = Periodic # Right now we just need them to behave like Periodic dimensions except we change the boundary conditions. + """ AbstractGrid{FT, TX, TY, TZ} diff --git a/src/Grids/conformal_cubed_sphere_face_grid.jl b/src/Grids/conformal_cubed_sphere_face_grid.jl index b792751891..653642c6c5 100644 --- a/src/Grids/conformal_cubed_sphere_face_grid.jl +++ b/src/Grids/conformal_cubed_sphere_face_grid.jl @@ -1,7 +1,6 @@ using CubedSphere using JLD2 using OffsetArrays -using Rotations using Oceananigans using Oceananigans.Grids @@ -150,16 +149,28 @@ function ConformalCubedSphereFaceGrid(FT = Float64; size, z, return ConformalCubedSphereFaceGrid{FT, TX, TY, TZ, typeof(λᶜᶜᵃ), typeof(zᵃᵃᶠ)}( Nξ, Nη, Nz, Hx, Hy, Hz, - λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ, φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ, zᵃᵃᶠ, zᵃᵃᶜ, - Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δyᶜᶜᵃ, Δyᶠᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶠᵃ, Δz, + λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ, φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ, zᵃᵃᶜ, zᵃᵃᶠ, + Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ, Δz, Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, radius) end -function offset_data(data, Hx, Hy) - Nx, Ny = size(data) .- 1 # Just count cell centers - offset_data = zeros(Nx + 1 + 2Hx, Ny + 1 + 2Hy) - offset_data[1+Hx:Nx+1+Hx, 1+Hy:Ny+1+Hy] .= data - return OffsetArray(offset_data, -Hx, -Hy) +function offset_data_2d(data, loc, topo, N, H) + LX, LY = loc + TX, TY = topo + Nx, Ny = N + Hx, Hy = H + + Tx = total_length(LX, TX, Nx, Hx) + Ty = total_length(LY, TY, Ny, Hy) + + new_data = zeros(Tx, Ty) + + ii = interior_parent_indices(LX, TX, Nx, Hx) + jj = interior_parent_indices(LY, TY, Ny, Hy) + + new_data[ii, jj] .= data + + return OffsetArray(new_data, -Hx, -Hy) end function ConformalCubedSphereFaceGrid(filepath::AbstractString, FT = Float64; face, Nz, z, @@ -167,7 +178,8 @@ function ConformalCubedSphereFaceGrid(filepath::AbstractString, FT = Float64; fa radius = R_Earth, halo = (1, 1, 1), rotation = nothing) - TX, TY, TZ = topology + + TX, TY, TZ = topo = topology Hx, Hy, Hz = halo ## Use a regular rectilinear grid for the vertical grid @@ -186,26 +198,39 @@ function ConformalCubedSphereFaceGrid(filepath::AbstractString, FT = Float64; fa Nξ, Nη = size(cubed_sphere_data["λᶠᶠᵃ"]) .- 1 - λᶜᶜᵃ = offset_data(cubed_sphere_data["λᶜᶜᵃ"], Hx, Hy) - λᶠᶠᵃ = offset_data(cubed_sphere_data["λᶠᶠᵃ"], Hx, Hy) + N = (Nξ, Nη, Nz) + H = halo + + loc_cc = (Center, Center) + loc_cf = (Center, Face) + loc_fc = (Face, Center) + loc_ff = (Face, Face) + + c_inds = 1:Nξ + f_inds = 1:Nξ+1 + + topo_2d = (Bounded, Bounded) + + λᶜᶜᵃ = offset_data_2d(cubed_sphere_data["λᶜᶜᵃ"][c_inds, c_inds], loc_cc, topo_2d, N, H) + λᶠᶠᵃ = offset_data_2d(cubed_sphere_data["λᶠᶠᵃ"][f_inds, f_inds], loc_ff, topo_2d, N, H) - φᶜᶜᵃ = offset_data(cubed_sphere_data["φᶜᶜᵃ"], Hx, Hy) - φᶠᶠᵃ = offset_data(cubed_sphere_data["φᶠᶠᵃ"], Hx, Hy) + φᶜᶜᵃ = offset_data_2d(cubed_sphere_data["φᶜᶜᵃ"][c_inds, c_inds], loc_cc, topo_2d, N, H) + φᶠᶠᵃ = offset_data_2d(cubed_sphere_data["φᶠᶠᵃ"][f_inds, f_inds], loc_ff, topo_2d, N, H) - Δxᶜᶜᵃ = offset_data(cubed_sphere_data["Δxᶜᶜᵃ"], Hx, Hy) - Δxᶠᶜᵃ = offset_data(cubed_sphere_data["Δxᶠᶜᵃ"], Hx, Hy) - Δxᶜᶠᵃ = offset_data(cubed_sphere_data["Δxᶜᶠᵃ"], Hx, Hy) - Δxᶠᶠᵃ = offset_data(cubed_sphere_data["Δxᶠᶠᵃ"], Hx, Hy) + Δxᶜᶜᵃ = offset_data_2d(cubed_sphere_data["Δxᶜᶜᵃ"][c_inds, c_inds], loc_cc, topo_2d, N, H) + Δxᶠᶜᵃ = offset_data_2d(cubed_sphere_data["Δxᶠᶜᵃ"][f_inds, c_inds], loc_fc, topo_2d, N, H) + Δxᶜᶠᵃ = offset_data_2d(cubed_sphere_data["Δxᶜᶠᵃ"][c_inds, f_inds], loc_cf, topo_2d, N, H) + Δxᶠᶠᵃ = offset_data_2d(cubed_sphere_data["Δxᶠᶠᵃ"][f_inds, f_inds], loc_ff, topo_2d, N, H) - Δyᶜᶜᵃ = offset_data(cubed_sphere_data["Δyᶜᶜᵃ"], Hx, Hy) - Δyᶠᶜᵃ = offset_data(cubed_sphere_data["Δyᶠᶜᵃ"], Hx, Hy) - Δyᶜᶠᵃ = offset_data(cubed_sphere_data["Δyᶜᶠᵃ"], Hx, Hy) - Δyᶠᶠᵃ = offset_data(cubed_sphere_data["Δyᶠᶠᵃ"], Hx, Hy) + Δyᶜᶜᵃ = offset_data_2d(cubed_sphere_data["Δyᶜᶜᵃ"][c_inds, c_inds], loc_cc, topo_2d, N, H) + Δyᶠᶜᵃ = offset_data_2d(cubed_sphere_data["Δyᶠᶜᵃ"][f_inds, c_inds], loc_fc, topo_2d, N, H) + Δyᶜᶠᵃ = offset_data_2d(cubed_sphere_data["Δyᶜᶠᵃ"][c_inds, f_inds], loc_cf, topo_2d, N, H) + Δyᶠᶠᵃ = offset_data_2d(cubed_sphere_data["Δyᶠᶠᵃ"][f_inds, f_inds], loc_ff, topo_2d, N, H) - Azᶜᶜᵃ = offset_data(cubed_sphere_data["Azᶜᶜᵃ"], Hx, Hy) - Azᶠᶜᵃ = offset_data(cubed_sphere_data["Azᶠᶜᵃ"], Hx, Hy) - Azᶜᶠᵃ = offset_data(cubed_sphere_data["Azᶜᶠᵃ"], Hx, Hy) - Azᶠᶠᵃ = offset_data(cubed_sphere_data["Azᶠᶠᵃ"], Hx, Hy) + Azᶜᶜᵃ = offset_data_2d(cubed_sphere_data["Azᶜᶜᵃ"][c_inds, c_inds], loc_cc, topo_2d, N, H) + Azᶠᶜᵃ = offset_data_2d(cubed_sphere_data["Azᶠᶜᵃ"][f_inds, c_inds], loc_fc, topo_2d, N, H) + Azᶜᶠᵃ = offset_data_2d(cubed_sphere_data["Azᶜᶠᵃ"][c_inds, f_inds], loc_cf, topo_2d, N, H) + Azᶠᶠᵃ = offset_data_2d(cubed_sphere_data["Azᶠᶠᵃ"][f_inds, f_inds], loc_ff, topo_2d, N, H) ## Maybe we won't need these? @@ -217,8 +242,8 @@ function ConformalCubedSphereFaceGrid(filepath::AbstractString, FT = Float64; fa return ConformalCubedSphereFaceGrid{FT, TX, TY, TZ, typeof(λᶜᶜᵃ), typeof(zᵃᵃᶠ)}( Nξ, Nη, Nz, Hx, Hy, Hz, - λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ, φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ, zᵃᵃᶠ, zᵃᵃᶜ, - Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δyᶜᶜᵃ, Δyᶠᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶠᵃ, Δz, + λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ, φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ, zᵃᵃᶜ, zᵃᵃᶠ, + Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ, Δz, Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, radius) end diff --git a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl index ccfa9b12cc..38dc0c347b 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl @@ -1,6 +1,9 @@ module HydrostaticFreeSurfaceModels -export HydrostaticFreeSurfaceModel, VectorInvariant, ExplicitFreeSurface, ImplicitFreeSurface +export + HydrostaticFreeSurfaceModel, VectorInvariant, + ExplicitFreeSurface, ImplicitFreeSurface, + PrescribedVelocityFields using KernelAbstractions: @index, @kernel, Event, MultiEvent using KernelAbstractions.Extras.LoopInfo: @unroll @@ -9,6 +12,9 @@ using Oceananigans.Utils: launch! import Oceananigans: fields +# This is only used by the cubed sphere for now. +fill_horizontal_velocity_halos!(args...) = nothing + ##### ##### HydrostaticFreeSurfaceModel definition ##### diff --git a/src/Models/HydrostaticFreeSurfaceModels/calculate_hydrostatic_free_surface_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/calculate_hydrostatic_free_surface_tendencies.jl index c826856fdd..a233415998 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/calculate_hydrostatic_free_surface_tendencies.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/calculate_hydrostatic_free_surface_tendencies.jl @@ -45,20 +45,19 @@ function calculate_hydrostatic_momentum_tendencies!(tendencies, velocities, arch free_surface, tracers, diffusivities, hydrostatic_pressure_anomaly, forcings, clock, barrier) - calculate_Gu_kernel! = calculate_hydrostatic_free_surface_Gu!(device(arch), work_layout(grid, :xyz)...) - calculate_Gv_kernel! = calculate_hydrostatic_free_surface_Gv!(device(arch), work_layout(grid, :xyz)...) - calculate_Gη_kernel! = calculate_hydrostatic_free_surface_Gη!(device(arch), work_layout(grid, :xy)...) + Gu_event = launch!(arch, grid, :xyz, calculate_hydrostatic_free_surface_Gu!, + tendencies.u, grid, advection.momentum, coriolis, closure, + velocities, free_surface, tracers, diffusivities, hydrostatic_pressure_anomaly, + forcings, clock; dependencies = barrier) - Gu_event = calculate_Gu_kernel!(tendencies.u, grid, advection.momentum, coriolis, closure, - velocities, free_surface, tracers, diffusivities, hydrostatic_pressure_anomaly, - forcings, clock; dependencies = barrier) + Gv_event = launch!(arch, grid, :xyz, calculate_hydrostatic_free_surface_Gv!, + tendencies.v, grid, advection.momentum, coriolis, closure, + velocities, free_surface, tracers, diffusivities, hydrostatic_pressure_anomaly, + forcings, clock; dependencies = barrier) - Gv_event = calculate_Gv_kernel!(tendencies.v, grid, advection.momentum, coriolis, closure, - velocities, free_surface, tracers, diffusivities, hydrostatic_pressure_anomaly, - forcings, clock; dependencies = barrier) - - Gη_event = calculate_Gη_kernel!(tendencies.η, grid, velocities, free_surface, tracers, - forcings, clock; dependencies = barrier) + Gη_event = launch!(arch, grid, :xy, calculate_hydrostatic_free_surface_Gη!, + tendencies.η, grid, velocities, free_surface, tracers, + forcings, clock; dependencies = barrier) events = [Gu_event, Gv_event, Gη_event] @@ -81,23 +80,23 @@ function calculate_hydrostatic_free_surface_interior_tendency_contributions!(ten forcings, clock) - calculate_Gc_kernel! = calculate_hydrostatic_free_surface_Gc!(device(arch), work_layout(grid, :xyz)...) - barrier = Event(device(arch)) events = calculate_hydrostatic_momentum_tendencies!(tendencies, velocities, arch, grid, advection, coriolis, closure, free_surface, tracers, diffusivities, hydrostatic_pressure_anomaly, forcings, clock, barrier) - + for (tracer_index, tracer_name) in enumerate(propertynames(tracers)) @inbounds c_tendency = tendencies[tracer_name] @inbounds c_advection = advection[tracer_name] @inbounds forcing = forcings[tracer_name] - Gc_event = calculate_Gc_kernel!(c_tendency, grid, Val(tracer_index), - c_advection, closure, buoyancy, - velocities, free_surface, tracers, diffusivities, - forcing, clock; dependencies = barrier) + Gc_event = launch!(arch, grid, :xyz, calculate_hydrostatic_free_surface_Gc!, + c_tendency, grid, Val(tracer_index), + c_advection, closure, buoyancy, + velocities, free_surface, tracers, diffusivities, + forcing, clock; + dependencies=barrier) push!(events, Gc_event) end diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl index 93aa718bb5..cc2c3352d2 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl @@ -7,23 +7,22 @@ combine_events(free_surface_event, tracer_events) = MultiEvent(tuple(free_surfac function ab2_step!(model::HydrostaticFreeSurfaceModel, Δt, χ) - workgroup, worksize = work_layout(model.grid, :xyz) - barrier = Event(device(model.architecture)) - step_field_kernel! = ab2_step_field!(device(model.architecture), workgroup, worksize) - # Launch velocity update kernels velocities_events = [] for name in (:u, :v) + model.velocities isa PrescribedVelocityFields && break + Gⁿ = model.timestepper.Gⁿ[name] G⁻ = model.timestepper.G⁻[name] velocity_field = model.velocities[name] - event = step_field_kernel!(velocity_field, Δt, χ, Gⁿ, G⁻, - dependencies=Event(device(model.architecture))) + event = launch!(model.architecture, model.grid, :xyz, ab2_step_field!, + velocity_field, Δt, χ, Gⁿ, G⁻, + dependencies=Event(device(model.architecture))) push!(velocities_events, event) end @@ -37,8 +36,9 @@ function ab2_step!(model::HydrostaticFreeSurfaceModel, Δt, χ) G⁻ = model.timestepper.G⁻[name] tracer_field = model.tracers[name] - event = step_field_kernel!(tracer_field, Δt, χ, Gⁿ, G⁻, - dependencies=Event(device(model.architecture))) + event = launch!(model.architecture, model.grid, :xyz, ab2_step_field!, + tracer_field, Δt, χ, Gⁿ, G⁻, + dependencies=Event(device(model.architecture))) push!(tracer_events, event) end @@ -88,10 +88,11 @@ end function ab2_step_free_surface!(free_surface::ImplicitFreeSurface, velocities_update, model, χ, Δt) ##### Implicit solver for η - + ## Need to wait for u* and v* to finish wait(device(model.architecture), velocities_update) - fill_halo_regions!(model.velocities, model.architecture, model.clock, fields(model) ) + fill_halo_regions!(model.velocities, model.architecture, model.clock, fields(model)) + fill_horizontal_velocity_halos!(model.velocities.u, model.velocities.v, model.architecture) ## Leaving this here for now. There may be some scenarios where stepping forward η and then using ## the stepped forward value as a guess is helpful. @@ -104,13 +105,14 @@ function ab2_step_free_surface!(free_surface::ImplicitFreeSurface, velocities_up wait(device(model.architecture), event) u=free_surface.barotropic_volume_flux.u v=free_surface.barotropic_volume_flux.v - fill_halo_regions!(u.data ,u.boundary_conditions, model.architecture, model.grid, model.clock, fields(model) ) - fill_halo_regions!(v.data ,v.boundary_conditions, model.architecture, model.grid, model.clock, fields(model) ) + fill_halo_regions!(u.data, u.boundary_conditions, model.architecture, model.grid, model.clock, fields(model)) + fill_halo_regions!(v.data, v.boundary_conditions, model.architecture, model.grid, model.clock, fields(model)) + fill_horizontal_velocity_halos!(model.velocities.u, model.velocities.v, model.architecture) ## Compute volume scaled divergence of the barotropic transport and put into solver RHS event = compute_volume_scaled_divergence!(free_surface, model) wait(device(model.architecture), event) - + ## Include surface pressure term into RHS RHS = free_surface.implicit_step_solver.solver.settings.RHS RHS .= RHS/(model.free_surface.gravitational_acceleration*Δt) diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_advection.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_advection.jl index def39a69db..afbc80bb76 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_advection.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_advection.jl @@ -38,4 +38,3 @@ using Oceananigans.Advection: div_Uu, div_Uv @inline U_dot_∇u(i, j, k, grid::AbstractGrid{FT}, advection::Nothing, U) where FT = zero(FT) @inline U_dot_∇v(i, j, k, grid::AbstractGrid{FT}, advection::Nothing, U) where FT = zero(FT) - diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_fields.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_fields.jl index 62010a8d22..a9af85c4bd 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_fields.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_fields.jl @@ -13,5 +13,3 @@ function HydrostaticFreeSurfaceTendencyFields(velocities, free_surface, arch, gr tracers = TracerFields(tracer_names, arch, grid) return merge((u=u, v=v, η=η), tracers) end - - diff --git a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl index 5d2a843bd6..4d0a883264 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl @@ -14,6 +14,7 @@ function update_state!(model::HydrostaticFreeSurfaceModel) # Fill halos for velocities and tracers fill_halo_regions!(fields(model), model.architecture, model.clock, fields(model)) + fill_horizontal_velocity_halos!(model.velocities.u, model.velocities.v, model.architecture) compute_w_from_continuity!(model) diff --git a/src/Models/Models.jl b/src/Models/Models.jl index 312123fa1a..e09ddcb64c 100644 --- a/src/Models/Models.jl +++ b/src/Models/Models.jl @@ -1,6 +1,8 @@ module Models -export IncompressibleModel, NonDimensionalIncompressibleModel, HydrostaticFreeSurfaceModel, ShallowWaterModel +export + IncompressibleModel, NonDimensionalIncompressibleModel, HydrostaticFreeSurfaceModel, ShallowWaterModel, + ExplicitFreeSurface, VectorInvariant, HydrostaticSphericalCoriolis, VectorInvariantEnstrophyConserving using Oceananigans: AbstractModel @@ -11,7 +13,10 @@ include("HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl") include("ShallowWaterModels/ShallowWaterModels.jl") using .IncompressibleModels: IncompressibleModel, NonDimensionalIncompressibleModel -using .HydrostaticFreeSurfaceModels: HydrostaticFreeSurfaceModel using .ShallowWaterModels: ShallowWaterModel +using .HydrostaticFreeSurfaceModels: + HydrostaticFreeSurfaceModel, ExplicitFreeSurface, VectorInvariant, + HydrostaticSphericalCoriolis, VectorInvariantEnstrophyConserving + end diff --git a/src/Oceananigans.jl b/src/Oceananigans.jl index 4830605d0c..91e60ff012 100644 --- a/src/Oceananigans.jl +++ b/src/Oceananigans.jl @@ -50,6 +50,7 @@ export IsotropicDiffusivity, AnisotropicDiffusivity, AnisotropicBiharmonicDiffusivity, ConstantSmagorinsky, AnisotropicMinimumDissipation, + HorizontallyCurvilinearAnisotropicDiffusivity, # Lagrangian particle tracking LagrangianParticles, @@ -57,6 +58,9 @@ export # Models IncompressibleModel, NonDimensionalIncompressibleModel, HydrostaticFreeSurfaceModel, fields, + # Hydrostatic free surface model + ExplicitFreeSurface, VectorInvariant, HydrostaticSphericalCoriolis, VectorInvariantEnstrophyConserving, + # Time stepping Clock, TimeStepWizard, time_step!, @@ -65,7 +69,7 @@ export iteration_limit_exceeded, stop_time_exceeded, wall_time_limit_exceeded, # Diagnostics - NaNChecker, FieldMaximum, + NaNChecker, StateChecker, CFL, AdvectiveCFL, DiffusiveCFL, # Output writers @@ -75,6 +79,9 @@ export # Abstract operations ∂x, ∂y, ∂z, @at, + # Cubed sphere + ConformalCubedSphereGrid, + # Utils prettytime @@ -166,6 +173,7 @@ include("Diagnostics/Diagnostics.jl") include("OutputWriters/OutputWriters.jl") include("Simulations/Simulations.jl") include("AbstractOperations/AbstractOperations.jl") +include("CubedSpheres/CubedSpheres.jl") include("Distributed/Distributed.jl") ##### @@ -192,6 +200,7 @@ using .Diagnostics using .OutputWriters using .Simulations using .AbstractOperations +using .CubedSpheres using .Distributed function __init__() diff --git a/src/TimeSteppers/clock.jl b/src/TimeSteppers/clock.jl index ed5525a4c5..15c520b92e 100644 --- a/src/TimeSteppers/clock.jl +++ b/src/TimeSteppers/clock.jl @@ -34,13 +34,10 @@ Returns a `Clock` initialized to the zeroth iteration and first time step stage. """ Clock(; time, iteration=0, stage=1) = Clock{typeof(time)}(time, iteration, stage) -short_show(clock::Clock) = string("Clock(time=", prettytime(clock.time), - ", iteration=", clock.iteration, ")") +short_show(clock::Clock) = string("Clock(time=$(prettytime(clock.time)), iteration=$(clock.iteration))") Base.show(io::IO, c::Clock{T}) where T = - println(io, "Clock{$T}: time = ", prettytime(c.time), - ", iteration = ", c.iteration, - ", stage = ", c.stage) + println(io, "Clock{$T}: time = $(prettytime(c.time)), iteration = $(c.iteration), stage = $(c.stage)") next_time(clock, Δt) = clock.time + Δt next_time(clock::Clock{<:AbstractTime}, Δt) = clock.time + Nanosecond(round(Int, 1e9 * Δt)) diff --git a/src/TimeSteppers/quasi_adams_bashforth_2.jl b/src/TimeSteppers/quasi_adams_bashforth_2.jl index 6966368da2..06ad165ade 100644 --- a/src/TimeSteppers/quasi_adams_bashforth_2.jl +++ b/src/TimeSteppers/quasi_adams_bashforth_2.jl @@ -108,4 +108,7 @@ Time step via end end -@kernel ab2_step_field!(ϕ::FunctionField, args...) = nothing +@kernel ab2_step_field!(::FunctionField, args...) = nothing + +# Needed for use with `HydrostaticFreeSurfaceModel` + `PrescribedVelocities` + prescribed `Field` (not `FunctionField`). +# @kernel ab2_step_field!(::Nothing, args...) = nothing diff --git a/src/TimeSteppers/store_tendencies.jl b/src/TimeSteppers/store_tendencies.jl index baea12ac64..fb334d0a8a 100644 --- a/src/TimeSteppers/store_tendencies.jl +++ b/src/TimeSteppers/store_tendencies.jl @@ -1,5 +1,7 @@ using Oceananigans.Grids: AbstractGrid +using Oceananigans.Utils: launch! + """ Store source terms for `u`, `v`, and `w`. """ @kernel function store_field_tendencies!(G⁻, grid::AbstractGrid{FT}, G⁰) where FT i, j, k = @index(Global, NTuple) @@ -11,20 +13,17 @@ function store_tendencies!(model) barrier = Event(device(model.architecture)) - workgroup, worksize = work_layout(model.grid, :xyz) - model_fields = fields(model) - store_field_tendencies_kernel! = store_field_tendencies!(device(model.architecture), workgroup, worksize) - events = [] for field_name in keys(model_fields) - field_event = store_field_tendencies_kernel!(model.timestepper.G⁻[field_name], - model.grid, - model.timestepper.Gⁿ[field_name], - dependencies=barrier) + field_event = launch!(model.architecture, model.grid, :xyz, store_field_tendencies!, + model.timestepper.G⁻[field_name], + model.grid, + model.timestepper.Gⁿ[field_name], + dependencies = barrier) push!(events, field_event) end diff --git a/src/Utils/kernel_launching.jl b/src/Utils/kernel_launching.jl index be62270772..0a2609c9a8 100644 --- a/src/Utils/kernel_launching.jl +++ b/src/Utils/kernel_launching.jl @@ -69,7 +69,7 @@ Returns an `event` token associated with the `kernel!` launch. The keyword argument `dependencies` is an `Event` or `MultiEvent` specifying prior kernels that must complete before `kernel!` is launched. """ -function launch!(arch, grid, dims, kernel!, args...; +function launch!(arch, grid, dims, kernel!, args...; dependencies = nothing, include_right_boundaries = false, reduced_dimensions = (), diff --git a/sandbox/cubed_sphere_grid_from_file.jl b/test/data_dependencies.jl similarity index 53% rename from sandbox/cubed_sphere_grid_from_file.jl rename to test/data_dependencies.jl index cf4be7a675..66244331f5 100644 --- a/sandbox/cubed_sphere_grid_from_file.jl +++ b/test/data_dependencies.jl @@ -1,16 +1,11 @@ using DataDeps -using JLD2 -using Oceananigans ENV["DATADEPS_ALWAYS_ACCEPT"] = "true" dd = DataDep("cubed_sphere_32_grid", "Conformal cubed sphere grid with 32×32 grid points on each face", "https://github.com/CliMA/OceananigansArtifacts.jl/raw/main/cubed_sphere_grids/cubed_sphere_32_grid.jld2", - "3cc5d86290c3af028cddfa47e61e095ee470fe6f8d779c845de09da2f1abeb15" # sha256sum + "b1dafe4f9142c59a2166458a2def743cd45b20a4ed3a1ae84ad3a530e1eff538" # sha256sum ) DataDeps.register(dd) - -cs32_filepath = datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2" -grid = ConformalCubedSphereFaceGrid(cs32_filepath, face=1, Nz=1, z=(-1, 0)) diff --git a/test/runtests.jl b/test/runtests.jl index e7907bfc7d..1da8f065a0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -72,6 +72,7 @@ closures = ( CUDA.allowscalar(true) +include("data_dependencies.jl") include("utils_for_runtests.jl") group = get(ENV, "TEST_GROUP", :all) |> Symbol @@ -138,6 +139,11 @@ group = get(ENV, "TEST_GROUP", :all) |> Symbol end end + if group == :cubed_sphere || group == :all + include("test_cubed_spheres.jl") + include("test_cubed_sphere_halo_exchange.jl") + end + if group == :distributed || group == :all MPI.Initialized() || MPI.Init() include("test_distributed_models.jl") diff --git a/test/test_cubed_sphere_halo_exchange.jl b/test/test_cubed_sphere_halo_exchange.jl new file mode 100644 index 0000000000..1e212254fc --- /dev/null +++ b/test/test_cubed_sphere_halo_exchange.jl @@ -0,0 +1,616 @@ +using Logging +using Printf +using Test +using DataDeps + +using Oceananigans +using Oceananigans.CubedSpheres + +using Oceananigans.BoundaryConditions: fill_halo_regions! +using Oceananigans.CubedSpheres: west_halo, east_halo, south_halo, north_halo, fill_horizontal_velocity_halos! + +# Opposite of the `Base.digits` function +# Source: https://stackoverflow.com/a/55529778 +function undigits(d; base=10) + (s, b) = promote(zero(eltype(d)), base) + mult = one(s) + for val in d + s += val * mult + mult *= b + end + return s +end + +cs32_filepath = datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2" + +@testset "Cubed sphere halo exchange" begin + arch = CPU() + grid = ConformalCubedSphereGrid(cs32_filepath, Nz=1, z=(-1, 0)) + field = CenterField(Float64, arch, grid) + + ## We will fill each grid point with a 5-digit integer "fiijj" where + ## the f digit is the face number, the ii digits are the i index, and + ## the jj digits are the j index. We then check that the halo exchange + ## happened correctly. + + face_digit(n) = digits(abs(Int(n)))[5] + i_digits(n) = digits(abs(Int(n)))[3:4] |> undigits + j_digits(n) = digits(abs(Int(n)))[1:2] |> undigits + + for (face_number, field_face) in enumerate(field.faces) + for i in 1:field_face.grid.Nx, j in 1:field_face.grid.Ny + field_face[i, j, 1] = parse(Int, @sprintf("%d%02d%02d", face_number, i, j)) + end + end + + fill_halo_regions!(field, arch) + + @testset "Source and destination faces are correct" begin + for (face_number, field_face) in enumerate(field.faces) + west_halo_vals = west_halo(field_face, include_corners=false) + east_halo_vals = east_halo(field_face, include_corners=false) + south_halo_vals = south_halo(field_face, include_corners=false) + north_halo_vals = north_halo(field_face, include_corners=false) + + @test all(face_digit.(west_halo_vals) .== grid.face_connectivity[face_number].west.face) + @test all(face_digit.(east_halo_vals) .== grid.face_connectivity[face_number].east.face) + @test all(face_digit.(south_halo_vals) .== grid.face_connectivity[face_number].south.face) + @test all(face_digit.(north_halo_vals) .== grid.face_connectivity[face_number].north.face) + end + end + + @testset "1W halo <- 5N boundary halo exchange" begin + # Grid point [i, j] = [0, 1] in 1W halo should be from [i, j] = [32, 32] in 5N boundary. + west_halo_south_value = field.faces[1][0, 1, 1] + @test face_digit(west_halo_south_value) == 5 + @test i_digits(west_halo_south_value) == 32 + @test j_digits(west_halo_south_value) == 32 + + # Grid point [i, j] = [0, 32] in 1W halo should be from [i, j] = [1, 32] in 5N boundary. + west_halo_north_value = field.faces[1][0, 32, 1] + @test face_digit(west_halo_north_value) == 5 + @test i_digits(west_halo_north_value) == 1 + @test j_digits(west_halo_north_value) == 32 + + west_halo_values = west_halo(field.faces[1], include_corners=false)[:] + @test all(face_digit.(west_halo_values) .== 5) + @test all(i_digits.(west_halo_values) .== reverse(1:32)) + @test all(j_digits.(west_halo_values) .== 32) + end + + @testset "1E halo <- 2W boundary halo exchange" begin + # Grid point [i, j] = [33, 1] in 1E halo should be from [i, j] = [1, 1] in 2W boundary. + east_halo_south_value = field.faces[1][33, 1, 1] + @test face_digit(east_halo_south_value) == 2 + @test i_digits(east_halo_south_value) == 1 + @test j_digits(east_halo_south_value) == 1 + + # Grid point [i, j] = [33, 32] in 1E halo should be from [i, j] = [1, 32] in 2W boundary. + east_halo_north_value = field.faces[1][33, 32, 1] + @test face_digit(east_halo_north_value) == 2 + @test i_digits(east_halo_north_value) == 1 + @test j_digits(east_halo_north_value) == 32 + + east_halo_values = east_halo(field.faces[1], include_corners=false)[:] + @test all(face_digit.(east_halo_values) .== 2) + @test all(i_digits.(east_halo_values) .== 1) + @test all(j_digits.(east_halo_values) .== 1:32) + end + + @testset "1S halo <- 6N boundary halo exchange" begin + # Grid point [i, j] = [1, 0] in 1S halo should be from [i, j] = [1, 32] in 6N boundary. + south_halo_west_value = field.faces[1][1, 0, 1] + @test face_digit(south_halo_west_value) == 6 + @test i_digits(south_halo_west_value) == 1 + @test j_digits(south_halo_west_value) == 32 + + # Grid point [i, j] = [32, 0] in 1S halo should be from [i, j] = [32, 32] in 6N boundary. + south_halo_east_value = field.faces[1][32, 0, 1] + @test face_digit(south_halo_east_value) == 6 + @test i_digits(south_halo_east_value) == 32 + @test j_digits(south_halo_east_value) == 32 + + south_halo_values = south_halo(field.faces[1], include_corners=false)[:] + @test all(face_digit.(south_halo_values) .== 6) + @test all(i_digits.(south_halo_values) .== 1:32) + @test all(j_digits.(south_halo_values) .== 32) + end + + @testset "1N halo <- 3W boundary halo exchange" begin + # Grid point [i, j] = [1, 33] in 1N halo should be from [i, j] = [1, 32] in 3W boundary. + north_halo_west_value = field.faces[1][1, 33, 1] + @test face_digit(north_halo_west_value) == 3 + @test i_digits(north_halo_west_value) == 1 + @test j_digits(north_halo_west_value) == 32 + + # Grid point [i, j] = [32, 33] in 1N halo should be from [i, j] = [1, 1] in 3W boundary. + north_halo_east_value = field.faces[1][32, 33, 1] + @test face_digit(north_halo_east_value) == 3 + @test i_digits(north_halo_east_value) == 1 + @test j_digits(north_halo_east_value) == 1 + + north_halo_values = north_halo(field.faces[1], include_corners=false)[:] + @test all(face_digit.(north_halo_values) .== 3) + @test all(i_digits.(north_halo_values) .== 1) + @test all(j_digits.(north_halo_values) .== reverse(1:32)) + end + + @testset "2W halo <- 1E boundary halo exchange" begin + # Grid point [i, j] = [0, 1] in 2W halo should be from [i, j] = [32, 1] in 1E boundary. + west_halo_south_value = field.faces[2][0, 1, 1] + @test face_digit(west_halo_south_value) == 1 + @test i_digits(west_halo_south_value) == 32 + @test j_digits(west_halo_south_value) == 1 + + # Grid point [i, j] = [0, 32] in 2W halo should be from [i, j] = [32, 32] in 1E boundary. + west_halo_north_value = field.faces[2][0, 32, 1] + @test face_digit(west_halo_north_value) == 1 + @test i_digits(west_halo_north_value) == 32 + @test j_digits(west_halo_north_value) == 32 + + west_halo_values = west_halo(field.faces[2], include_corners=false)[:] + @test all(face_digit.(west_halo_values) .== 1) + @test all(i_digits.(west_halo_values) .== 32) + @test all(j_digits.(west_halo_values) .== 1:32) + end + + @testset "2E halo <- 4S boundary halo exchange" begin + # Grid point [i, j] = [33, 1] in 2E halo should be from [i, j] = [32, 1] in 4S boundary. + east_halo_south_value = field.faces[2][33, 1, 1] + @test face_digit(east_halo_south_value) == 4 + @test i_digits(east_halo_south_value) == 32 + @test j_digits(east_halo_south_value) == 1 + + # Grid point [i, j] = [33, 32] in 2E halo should be from [i, j] = [1, 1] in 4S boundary. + east_halo_north_value = field.faces[2][33, 32, 1] + @test face_digit(east_halo_north_value) == 4 + @test i_digits(east_halo_north_value) == 1 + @test j_digits(east_halo_north_value) == 1 + + east_halo_values = east_halo(field.faces[2], include_corners=false)[:] + @test all(face_digit.(east_halo_values) .== 4) + @test all(i_digits.(east_halo_values) .== reverse(1:32)) + @test all(j_digits.(east_halo_values) .== 1) + end + + @testset "2S halo <- 6E boundary halo exchange" begin + # Grid point [i, j] = [1, 0] in 2S halo should be from [i, j] = [32, 32] in 6E boundary. + south_halo_west_value = field.faces[2][1, 0, 1] + @test face_digit(south_halo_west_value) == 6 + @test i_digits(south_halo_west_value) == 32 + @test j_digits(south_halo_west_value) == 32 + + # Grid point [i, j] = [32, 0] in 2S halo should be from [i, j] = [32, 1] in 6E boundary. + south_halo_east_value = field.faces[2][32, 0, 1] + @test face_digit(south_halo_east_value) == 6 + @test i_digits(south_halo_east_value) == 32 + @test j_digits(south_halo_east_value) == 1 + + south_halo_values = south_halo(field.faces[2], include_corners=false)[:] + @test all(face_digit.(south_halo_values) .== 6) + @test all(i_digits.(south_halo_values) .== 32) + @test all(j_digits.(south_halo_values) .== reverse(1:32)) + end + + @testset "2N halo <- 3S boundary halo exchange" begin + # Grid point [i, j] = [1, 33] in 2N halo should be from [i, j] = [1, 1] in 3S boundary. + north_halo_west_value = field.faces[2][1, 33, 1] + @test face_digit(north_halo_west_value) == 3 + @test i_digits(north_halo_west_value) == 1 + @test j_digits(north_halo_west_value) == 1 + + # Grid point [i, j] = [32, 33] in 2N halo should be from [i, j] = [32, 1] in 3S boundary. + north_halo_east_value = field.faces[2][32, 33, 1] + @test face_digit(north_halo_east_value) == 3 + @test i_digits(north_halo_east_value) == 32 + @test j_digits(north_halo_east_value) == 1 + + north_halo_values = north_halo(field.faces[2], include_corners=false)[:] + @test all(face_digit.(north_halo_values) .== 3) + @test all(i_digits.(north_halo_values) .== 1:32) + @test all(j_digits.(north_halo_values) .== 1) + end + + @testset "3W halo <- 1N boundary halo exchange" begin + # Grid point [i, j] = [0, 1] in 3W halo should be from [i, j] = [32, 32] in 1N boundary. + west_halo_south_value = field.faces[3][0, 1, 1] + @test face_digit(west_halo_south_value) == 1 + @test i_digits(west_halo_south_value) == 32 + @test j_digits(west_halo_south_value) == 32 + + # Grid point [i, j] = [0, 32] in 3W halo should be from [i, j] = [1, 32] in 1N boundary. + west_halo_north_value = field.faces[3][0, 32, 1] + @test face_digit(west_halo_north_value) == 1 + @test i_digits(west_halo_north_value) == 1 + @test j_digits(west_halo_north_value) == 32 + + west_halo_values = west_halo(field.faces[3], include_corners=false)[:] + @test all(face_digit.(west_halo_values) .== 1) + @test all(i_digits.(west_halo_values) .== reverse(1:32)) + @test all(j_digits.(west_halo_values) .== 32) + end + + @testset "3E halo <- 4W boundary halo exchange" begin + # Grid point [i, j] = [33, 1] in 3E halo should be from [i, j] = [1, 1] in 4W boundary. + east_halo_south_value = field.faces[3][33, 1, 1] + @test face_digit(east_halo_south_value) == 4 + @test i_digits(east_halo_south_value) == 1 + @test j_digits(east_halo_south_value) == 1 + + # Grid point [i, j] = [33, 32] in 3E halo should be from [i, j] = [1, 32] in 4W boundary. + east_halo_north_value = field.faces[3][33, 32, 1] + @test face_digit(east_halo_north_value) == 4 + @test i_digits(east_halo_north_value) == 1 + @test j_digits(east_halo_north_value) == 32 + + east_halo_values = east_halo(field.faces[3], include_corners=false)[:] + @test all(face_digit.(east_halo_values) .== 4) + @test all(i_digits.(east_halo_values) .== 1) + @test all(j_digits.(east_halo_values) .== 1:32) + end + + @testset "3S halo <- 2N boundary halo exchange" begin + # Grid point [i, j] = [1, 0] in 3S halo should be from [i, j] = [1, 32] in 2N boundary. + south_halo_west_value = field.faces[3][1, 0, 1] + @test face_digit(south_halo_west_value) == 2 + @test i_digits(south_halo_west_value) == 1 + @test j_digits(south_halo_west_value) == 32 + + # Grid point [i, j] = [32, 0] in 3S halo should be from [i, j] = [32, 32] in 2N boundary. + south_halo_east_value = field.faces[3][32, 0, 1] + @test face_digit(south_halo_east_value) == 2 + @test i_digits(south_halo_east_value) == 32 + @test j_digits(south_halo_east_value) == 32 + + south_halo_values = south_halo(field.faces[3], include_corners=false)[:] + @test all(face_digit.(south_halo_values) .== 2) + @test all(i_digits.(south_halo_values) .== 1:32) + @test all(j_digits.(south_halo_values) .== 32) + end + + @testset "3N halo <- 5W boundary halo exchange" begin + # Grid point [i, j] = [1, 33] in 3N halo should be from [i, j] = [1, 32] in 5W boundary. + north_halo_west_value = field.faces[3][1, 33, 1] + @test face_digit(north_halo_west_value) == 5 + @test i_digits(north_halo_west_value) == 1 + @test j_digits(north_halo_west_value) == 32 + + # Grid point [i, j] = [32, 33] in 3N halo should be from [i, j] = [1, 1] in 5W boundary. + north_halo_east_value = field.faces[3][32, 33, 1] + @test face_digit(north_halo_east_value) == 5 + @test i_digits(north_halo_east_value) == 1 + @test j_digits(north_halo_east_value) == 1 + + north_halo_values = north_halo(field.faces[3], include_corners=false)[:] + @test all(face_digit.(north_halo_values) .== 5) + @test all(i_digits.(north_halo_values) .== 1) + @test all(j_digits.(north_halo_values) .== reverse(1:32)) + end + + @testset "4W halo <- 3E boundary halo exchange" begin + # Grid point [i, j] = [0, 1] in 4W halo should be from [i, j] = [32, 1] in 3E boundary. + west_halo_south_value = field.faces[4][0, 1, 1] + @test face_digit(west_halo_south_value) == 3 + @test i_digits(west_halo_south_value) == 32 + @test j_digits(west_halo_south_value) == 1 + + # Grid point [i, j] = [0, 32] in 4W halo should be from [i, j] = [32, 32] in 1N boundary. + west_halo_north_value = field.faces[4][0, 32, 1] + @test face_digit(west_halo_north_value) == 3 + @test i_digits(west_halo_north_value) == 32 + @test j_digits(west_halo_north_value) == 32 + + west_halo_values = west_halo(field.faces[4], include_corners=false)[:] + @test all(face_digit.(west_halo_values) .== 3) + @test all(i_digits.(west_halo_values) .== 32) + @test all(j_digits.(west_halo_values) .== 1:32) + end + + @testset "4E halo <- 6S boundary halo exchange" begin + # Grid point [i, j] = [33, 1] in 4E halo should be from [i, j] = [32, 1] in 6S boundary. + east_halo_south_value = field.faces[4][33, 1, 1] + @test face_digit(east_halo_south_value) == 6 + @test i_digits(east_halo_south_value) == 32 + @test j_digits(east_halo_south_value) == 1 + + # Grid point [i, j] = [33, 32] in 4E halo should be from [i, j] = [1, 1] in 6S boundary. + east_halo_north_value = field.faces[4][33, 32, 1] + @test face_digit(east_halo_north_value) == 6 + @test i_digits(east_halo_north_value) == 1 + @test j_digits(east_halo_north_value) == 1 + + east_halo_values = east_halo(field.faces[4], include_corners=false)[:] + @test all(face_digit.(east_halo_values) .== 6) + @test all(i_digits.(east_halo_values) .== reverse(1:32)) + @test all(j_digits.(east_halo_values) .== 1) + end + + @testset "4S halo <- 2E boundary halo exchange" begin + # Grid point [i, j] = [1, 0] in 4S halo should be from [i, j] = [32, 32] in 2E boundary. + south_halo_west_value = field.faces[4][1, 0, 1] + @test face_digit(south_halo_west_value) == 2 + @test i_digits(south_halo_west_value) == 32 + @test j_digits(south_halo_west_value) == 32 + + # Grid point [i, j] = [32, 0] in 4S halo should be from [i, j] = [32, 1] in 2E boundary. + south_halo_east_value = field.faces[4][32, 0, 1] + @test face_digit(south_halo_east_value) == 2 + @test i_digits(south_halo_east_value) == 32 + @test j_digits(south_halo_east_value) == 1 + + south_halo_values = south_halo(field.faces[4], include_corners=false)[:] + @test all(face_digit.(south_halo_values) .== 2) + @test all(i_digits.(south_halo_values) .== 32) + @test all(j_digits.(south_halo_values) .== reverse(1:32)) + end + + @testset "4N halo <- 5S boundary halo exchange" begin + # Grid point [i, j] = [1, 33] in 4N halo should be from [i, j] = [1, 1] in 5S boundary. + north_halo_west_value = field.faces[4][1, 33, 1] + @test face_digit(north_halo_west_value) == 5 + @test i_digits(north_halo_west_value) == 1 + @test j_digits(north_halo_west_value) == 1 + + # Grid point [i, j] = [32, 33] in 4N halo should be from [i, j] = [32, 1] in 5S boundary. + north_halo_east_value = field.faces[4][32, 33, 1] + @test face_digit(north_halo_east_value) == 5 + @test i_digits(north_halo_east_value) == 32 + @test j_digits(north_halo_east_value) == 1 + + north_halo_values = north_halo(field.faces[4], include_corners=false)[:] + @test all(face_digit.(north_halo_values) .== 5) + @test all(i_digits.(north_halo_values) .== 1:32) + @test all(j_digits.(north_halo_values) .== 1) + end + + ## TODO: Test faces 5 and 6 (or generalize the test to all faces)! + +end + +@testset "Cubed sphere velocity halo exchange" begin + arch = CPU() + grid = ConformalCubedSphereGrid(cs32_filepath, Nz=1, z=(-1, 0)) + + u_field = XFaceField(Float64, arch, grid) + v_field = YFaceField(Float64, arch, grid) + + ## We will fill each grid point with a 6-digit integer "ufiijj" where + ## the u digit is 1 for u and 2 for v, the f digit is the face number, + ## the ii digits are the i index, and the jj digits are the j index. + ## We then check that the halo exchange happened correctly (including + ## any sign changes). + + U_DIGIT = 1 + V_DIGIT = 2 + + uv_digit(n) = digits(abs(Int(n)))[6] + face_digit(n) = digits(abs(Int(n)))[5] + i_digits(n) = digits(abs(Int(n)))[3:4] |> undigits + j_digits(n) = digits(abs(Int(n)))[1:2] |> undigits + + for (face_number, u_field_face) in enumerate(u_field.faces) + for i in 1:u_field_face.grid.Nx+1, j in 1:u_field_face.grid.Ny + u_field_face[i, j, 1] = parse(Int, @sprintf("%d%d%02d%02d", U_DIGIT, face_number, i, j)) + end + end + + for (face_number, v_field_face) in enumerate(v_field.faces) + for i in 1:v_field_face.grid.Nx, j in 1:v_field_face.grid.Ny+1 + v_field_face[i, j, 1] = parse(Int, @sprintf("%d%d%02d%02d", V_DIGIT, face_number, i, j)) + end + end + + fill_horizontal_velocity_halos!(u_field, v_field, arch) + + # Oceananigans.CubedSpheres.state_check(u_field, "u", 2) + # Oceananigans.CubedSpheres.state_check(v_field, "v", 2) + + @testset "Source and destination faces are correct" begin + for field in (u_field, v_field), (face_number, field_face) in enumerate(field.faces) + west_halo_vals = west_halo(field_face, include_corners=false) + east_halo_vals = east_halo(field_face, include_corners=false) + south_halo_vals = south_halo(field_face, include_corners=false) + north_halo_vals = north_halo(field_face, include_corners=false) + + @test all(face_digit.(west_halo_vals) .== grid.face_connectivity[face_number].west.face) + @test all(face_digit.(east_halo_vals) .== grid.face_connectivity[face_number].east.face) + @test all(face_digit.(south_halo_vals) .== grid.face_connectivity[face_number].south.face) + @test all(face_digit.(north_halo_vals) .== grid.face_connectivity[face_number].north.face) + end + end + + @testset "(1W halo <- 5N boundary) horizontal velocity halo exchange" begin + # Grid point u[i, j] = u[0, 1] in 1W halo should be from +v[i, j] = +v[32, 32] in 5N boundary. + u_west_halo_south_value = u_field.faces[1][0, 1, 1] + @test uv_digit(u_west_halo_south_value) == V_DIGIT + @test sign(u_west_halo_south_value) == +1 + @test face_digit(u_west_halo_south_value) == 5 + @test i_digits(u_west_halo_south_value) == 32 + @test j_digits(u_west_halo_south_value) == 32 + + # Grid point u[i, j] = u[0, 32] in 1W halo should be from +v[i, j] = +v[1, 32] in 5N boundary. + u_west_halo_north_value = u_field.faces[1][0, 32, 1] + @test uv_digit(u_west_halo_north_value) == V_DIGIT + @test sign(u_west_halo_north_value) == +1 + @test face_digit(u_west_halo_north_value) == 5 + @test i_digits(u_west_halo_north_value) == 1 + @test j_digits(u_west_halo_north_value) == 32 + + u_west_halo_values = west_halo(u_field.faces[1], include_corners=false)[:] + @test all(uv_digit.(u_west_halo_values) .== V_DIGIT) + @test all(sign.(u_west_halo_values) .== +1) + @test all(face_digit.(u_west_halo_values) .== 5) + @test all(i_digits.(u_west_halo_values) .== reverse(1:32)) + @test all(j_digits.(u_west_halo_values) .== 32) + + # Grid point v[i, j] = v[0, 1] in 1W halo should be from -u[i, j] = -u[32, 32] in 5N boundary. + v_west_halo_south_value = v_field.faces[1][0, 1, 1] + @test uv_digit(v_west_halo_south_value) == U_DIGIT + @test sign(v_west_halo_south_value) == -1 + @test face_digit(v_west_halo_south_value) == 5 + @test i_digits(v_west_halo_south_value) == 32 + @test j_digits(v_west_halo_south_value) == 32 + + # Grid point v[i, j] = u[0, 33] in 1W halo should be from -u[i, j] = -u[1, 32] in 5N boundary. + v_west_halo_north_value = v_field.faces[1][0, 32, 1] + @test uv_digit(v_west_halo_north_value) == U_DIGIT + @test sign(v_west_halo_north_value) == -1 + @test face_digit(v_west_halo_north_value) == 5 + @test i_digits(v_west_halo_north_value) == 1 + @test j_digits(v_west_halo_north_value) == 32 + + v_west_halo_values = west_halo(v_field.faces[1], include_corners=false)[:] + @test all(uv_digit.(v_west_halo_values) .== U_DIGIT) + @test all(sign.(v_west_halo_values) .== -1) + @test all(face_digit.(v_west_halo_values) .== 5) + @test all(i_digits.(v_west_halo_values) .== reverse(1:32)) + @test all(j_digits.(v_west_halo_values) .== 32) + end + + @testset "(1E halo <- 2W boundary) horizontal velocity halo exchange" begin + # Grid point u[i, j] = u[33, 1] in 1E halo should be from +u[i, j] = +u[1, 1] in 2W boundary. + u_east_halo_south_value = u_field.faces[1][33, 1, 1] + @test uv_digit(u_east_halo_south_value) == U_DIGIT + @test sign(u_east_halo_south_value) == +1 + @test face_digit(u_east_halo_south_value) == 2 + @test i_digits(u_east_halo_south_value) == 1 + @test j_digits(u_east_halo_south_value) == 1 + + # Grid point u[i, j] = u[33, 32] in 1E halo should be from +u[i, j] = +u[1, 32] in 2W boundary. + u_east_halo_north_value = u_field.faces[1][33, 32, 1] + @test uv_digit(u_east_halo_north_value) == U_DIGIT + @test sign(u_east_halo_north_value) == +1 + @test face_digit(u_east_halo_north_value) == 2 + @test i_digits(u_east_halo_north_value) == 1 + @test j_digits(u_east_halo_north_value) == 32 + + u_east_halo_values = east_halo(u_field.faces[1], include_corners=false)[:] + @test all(uv_digit.(u_east_halo_values) .== U_DIGIT) + @test all(sign.(u_east_halo_values) .== +1) + @test all(face_digit.(u_east_halo_values) .== 2) + @test all(i_digits.(u_east_halo_values) .== 1) + @test all(j_digits.(u_east_halo_values) .== 1:32) + + # Grid point v[i, j] = v[33, 1] in 1E halo should be from +v[i, j] = +v[1, 1] in 2W boundary. + v_east_halo_south_value = v_field.faces[1][33, 1, 1] + @test uv_digit(v_east_halo_south_value) == V_DIGIT + @test sign(v_east_halo_south_value) == +1 + @test face_digit(v_east_halo_south_value) == 2 + @test i_digits(v_east_halo_south_value) == 1 + @test j_digits(v_east_halo_south_value) == 1 + + # # Grid point v[i, j] = v[33, 32] in 1E halo should be from +v[i, j] = +v[1, 32] in 2W boundary. + v_east_halo_north_value = v_field.faces[1][33, 32, 1] + @test uv_digit(v_east_halo_north_value) == V_DIGIT + @test sign(v_east_halo_north_value) == +1 + @test face_digit(v_east_halo_north_value) == 2 + @test i_digits(v_east_halo_north_value) == 1 + @test j_digits(v_east_halo_north_value) == 32 + + v_east_halo_values = east_halo(v_field.faces[1], include_corners=false)[:] + @test all(uv_digit.(v_east_halo_values) .== V_DIGIT) + @test all(sign.(v_east_halo_values) .== +1) + @test all(face_digit.(v_east_halo_values) .== 2) + @test all(i_digits.(v_east_halo_values) .== 1) + @test all(j_digits.(v_east_halo_values) .== 1:32) + end + + @testset "(1S halo <- 6N boundary) horizontal velocity halo exchange" begin + # Grid point u[i, j] = u[1, 0] in 1S halo should be from +u[i, j] = +u[1, 32] in 6N boundary. + u_south_halo_west_value = u_field.faces[1][1, 0, 1] + @test uv_digit(u_south_halo_west_value) == U_DIGIT + @test sign(u_south_halo_west_value) == +1 + @test face_digit(u_south_halo_west_value) == 6 + @test i_digits(u_south_halo_west_value) == 1 + @test j_digits(u_south_halo_west_value) == 32 + + # Grid point u[i, j] = u[32, 0] in 1S halo should be from +u[i, j] = +u[32, 32] in 6N boundary. + u_south_halo_east_value = u_field.faces[1][32, 0, 1] + @test uv_digit(u_south_halo_east_value) == U_DIGIT + @test sign(u_south_halo_east_value) == +1 + @test face_digit(u_south_halo_east_value) == 6 + @test i_digits(u_south_halo_east_value) == 32 + @test j_digits(u_south_halo_east_value) == 32 + + u_south_halo_values = south_halo(u_field.faces[1], include_corners=false)[:] + @test all(uv_digit.(u_south_halo_values) .== U_DIGIT) + @test all(sign.(u_south_halo_values) .== +1) + @test all(face_digit.(u_south_halo_values) .== 6) + @test all(i_digits.(u_south_halo_values) .== 1:32) + @test all(j_digits.(u_south_halo_values) .== 32) + + # Grid point v[i, j] = v[1, 0] in 1S halo should be from +v[i, j] = +v[1, 32] in 6N boundary. + v_south_halo_west_value = v_field.faces[1][1, 0, 1] + @test uv_digit(v_south_halo_west_value) == V_DIGIT + @test sign(v_south_halo_west_value) == +1 + @test face_digit(v_south_halo_west_value) == 6 + @test i_digits(v_south_halo_west_value) == 1 + @test j_digits(v_south_halo_west_value) == 32 + + # # Grid point v[i, j] = v[32, 0] in 1S halo should be from +v[i, j] = +v[32, 32] in 6N boundary. + v_south_halo_east_value = v_field.faces[1][32, 0, 1] + @test uv_digit(v_south_halo_east_value) == V_DIGIT + @test sign(v_south_halo_east_value) == +1 + @test face_digit(v_south_halo_east_value) == 6 + @test i_digits(v_south_halo_east_value) == 32 + @test j_digits(v_south_halo_east_value) == 32 + + v_south_halo_values = south_halo(v_field.faces[1], include_corners=false)[:] + @test all(uv_digit.(v_south_halo_values) .== V_DIGIT) + @test all(sign.(v_south_halo_values) .== +1) + @test all(face_digit.(v_south_halo_values) .== 6) + @test all(i_digits.(v_south_halo_values) .== 1:32) + @test all(j_digits.(v_south_halo_values) .== 32) + end + + @testset "(1N halo <- 3W boundary) horizontal velocity halo exchange" begin + # Grid point u[i, j] = u[1, 33] in 1N halo should be from -v[i, j] = -v[1, 32] in 3W boundary. + u_nouth_halo_west_value = u_field.faces[1][1, 33, 1] + @test uv_digit(u_nouth_halo_west_value) == V_DIGIT + @test sign(u_nouth_halo_west_value) == -1 + @test face_digit(u_nouth_halo_west_value) == 3 + @test i_digits(u_nouth_halo_west_value) == 1 + @test j_digits(u_nouth_halo_west_value) == 32 + + # Grid point u[i, j] = u[32, 33] in 1N halo should be from -v[i, j] = -v[1, 1] in 3W boundary. + u_north_halo_east_value = u_field.faces[1][32, 33, 1] + @test uv_digit(u_north_halo_east_value) == V_DIGIT + @test sign(u_north_halo_east_value) == -1 + @test face_digit(u_north_halo_east_value) == 3 + @test i_digits(u_north_halo_east_value) == 1 + @test j_digits(u_north_halo_east_value) == 1 + + u_north_halo_values = north_halo(u_field.faces[1], include_corners=false)[:] + @test all(uv_digit.(u_north_halo_values) .== V_DIGIT) + @test all(sign.(u_north_halo_values) .== -1) + @test all(face_digit.(u_north_halo_values) .== 3) + @test all(i_digits.(u_north_halo_values) .== 1) + @test all(j_digits.(u_north_halo_values) .== reverse(1:32)) + + # Grid point v[i, j] = v[1, 33] in 1N halo should be from +u[i, j] = +u[1, 32] in 3W boundary. + v_north_halo_west_value = v_field.faces[1][1, 33, 1] + @test uv_digit(v_north_halo_west_value) == U_DIGIT + @test sign(v_north_halo_west_value) == +1 + @test face_digit(v_north_halo_west_value) == 3 + @test i_digits(v_north_halo_west_value) == 1 + @test j_digits(v_north_halo_west_value) == 32 + + # # Grid point v[i, j] = v[32, 33] in 1N halo should be from +u[i, j] = +u[1, 1] in 3W boundary. + v_north_halo_east_value = v_field.faces[1][32, 33, 1] + @test uv_digit(v_north_halo_east_value) == U_DIGIT + @test sign(v_north_halo_east_value) == +1 + @test face_digit(v_north_halo_east_value) == 3 + @test i_digits(v_north_halo_east_value) == 1 + @test j_digits(v_north_halo_east_value) == 1 + + v_north_halo_values = north_halo(v_field.faces[1], include_corners=false)[:] + @test all(uv_digit.(v_north_halo_values) .== U_DIGIT) + @test all(sign.(v_north_halo_values) .== +1) + @test all(face_digit.(v_north_halo_values) .== 3) + @test all(i_digits.(v_north_halo_values) .== 1) + @test all(j_digits.(v_north_halo_values) .== reverse(1:32)) + end + + ## TODO: Test the other faces, especially face 2. I hope we can generalize this... + +end diff --git a/test/test_cubed_spheres.jl b/test/test_cubed_spheres.jl new file mode 100644 index 0000000000..0ac57d7a58 --- /dev/null +++ b/test/test_cubed_spheres.jl @@ -0,0 +1,58 @@ +using Oceananigans.CubedSpheres +using Oceananigans.Models.HydrostaticFreeSurfaceModels + +@testset "Cubed spheres" begin + @testset "Conformal cubed sphere grid" begin + @info " Testing conformal cubed sphere grid..." + + # Test show function + grid = ConformalCubedSphereGrid(face_size=(10, 10, 1), z=(-1, 0)) + show(grid); println(); + @test grid isa ConformalCubedSphereGrid + end + + @testset "Constructing a grid from file" begin + cs32_filepath = datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2" + grid = ConformalCubedSphereGrid(cs32_filepath, Nz=1, z=(-1, 0)) + + @test grid isa ConformalCubedSphereGrid + end + + @testset "Constructing a HydrostaticFreeSurfaceModel" begin + cs32_filepath = datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2" + grid = ConformalCubedSphereGrid(cs32_filepath, Nz=1, z=(-1, 0)) + + model = HydrostaticFreeSurfaceModel( + architecture = CPU(), + grid = grid, + momentum_advection = VectorInvariant(), + free_surface = ExplicitFreeSurface(gravitational_acceleration=0.1), + coriolis = nothing, + closure = nothing, + tracers = nothing, + buoyancy = nothing + ) + + @test model isa HydrostaticFreeSurfaceModel + end + + @testset "Time stepping a HydrostaticFreeSurfaceModel" begin + cs32_filepath = datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2" + grid = ConformalCubedSphereGrid(cs32_filepath, Nz=1, z=(-1, 0)) + + model = HydrostaticFreeSurfaceModel( + architecture = CPU(), + grid = grid, + momentum_advection = VectorInvariant(), + free_surface = ExplicitFreeSurface(gravitational_acceleration=0.1), + coriolis = nothing, + closure = nothing, + tracers = nothing, + buoyancy = nothing + ) + + time_step!(model, 1) + + @test model isa HydrostaticFreeSurfaceModel + end +end diff --git a/test/test_grids.jl b/test/test_grids.jl index b64f2c9bd7..58c56de7f1 100644 --- a/test/test_grids.jl +++ b/test/test_grids.jl @@ -629,14 +629,6 @@ end @testset "Conformal cubed sphere face grid from file" begin @info " Testing conformal cubed sphere face grid construction from file..." - dd = DataDep("cubed_sphere_32_grid", - "Conformal cubed sphere grid with 32×32 grid points on each face", - "https://github.com/CliMA/OceananigansArtifacts.jl/raw/main/cubed_sphere_grids/cubed_sphere_32_grid.jld2", - "b1dafe4f9142c59a2166458a2def743cd45b20a4ed3a1ae84ad3a530e1eff538" # sha256sum - ) - - DataDeps.register(dd) - cs32_filepath = datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2" for face in 1:6 diff --git a/validation/cubed_sphere_rossby_haurwitz/animate_on_map_projection.jl b/validation/cubed_sphere_rossby_haurwitz/animate_on_map_projection.jl new file mode 100644 index 0000000000..81c532d119 --- /dev/null +++ b/validation/cubed_sphere_rossby_haurwitz/animate_on_map_projection.jl @@ -0,0 +1,109 @@ +using Printf +using Glob +using JLD2 +using PyCall + +using Oceananigans.Utils: prettytime + +np = pyimport("numpy") +ma = pyimport("numpy.ma") +plt = pyimport("matplotlib.pyplot") +mticker = pyimport("matplotlib.ticker") +ccrs = pyimport("cartopy.crs") +cmocean = pyimport("cmocean") + +function plot_cubed_sphere_tracer_field!(fig, ax, var, grid; add_colorbar, transform, cmap, vmin, vmax) + + Nf = grid["faces"] |> keys |> length + Nx = grid["faces/1/Nx"] + Ny = grid["faces/1/Ny"] + Hx = grid["faces/1/Hx"] + Hy = grid["faces/1/Hy"] + + for face in 1:Nf + λᶠᶠᵃ = grid["faces/$face/λᶠᶠᵃ"][1+Hx:Nx+2Hx, 1+Hy:Ny+2Hy] + φᶠᶠᵃ = grid["faces/$face/φᶠᶠᵃ"][1+Hx:Nx+2Hx, 1+Hy:Ny+2Hy] + + var_face = var[face][:, :, 1] + + # Remove very specific problematic grid cells near λ = 180° on face 6 that mess up the plot. + # May be related to https://github.com/SciTools/cartopy/issues/1151 + # Not needed for Orthographic or NearsidePerspective projection so let's use those. + # if face == 6 + # for i in 1:Nx+1, j in 1:Ny+1 + # if isapprox(λᶠᶠᵃ[i, j], -180, atol=15) + # var_face[min(i, Nx), min(j, Ny)] = NaN + # end + # end + # end + + var_face_masked = ma.masked_where(np.isnan(var_face), var_face) + + im = ax.pcolormesh(λᶠᶠᵃ, φᶠᶠᵃ, var_face_masked; transform, cmap, vmin, vmax) + + # Add colorbar below all the subplots. + if add_colorbar && face == Nf + ax_cbar = fig.add_axes([0.25, 0.1, 0.5, 0.02]) + fig.colorbar(im, cax=ax_cbar, orientation="horizontal") + end + + ax.set_global() + end + + return ax +end + +function animate_rossby_haurwitz(; projections=[ccrs.Robinson()]) + + ## Extract data + + file = jldopen("cubed_sphere_rossby_haurwitz.jld2") + + iterations = parse.(Int, keys(file["timeseries/t"])) + + ## Makie movie of tracer field h + + for (n, i) in enumerate(iterations) + @info "Plotting iteration $i/$(iterations[end]) (frame $n/$(length(iterations)))..." + + u = file["timeseries/u/$i"] + v = file["timeseries/v/$i"] + η = file["timeseries/η/$i"] + + t = prettytime(file["timeseries/t/$i"]) + plot_title = "Rossby-Haurwitz wave (mode 4): η(λ, φ) at t = $t" + + fig = plt.figure(figsize=(16, 9)) + n_subplots = length(projections) + subplot_kwargs = (transform=ccrs.PlateCarree(), cmap=cmocean.cm.balance, vmin=8000-80, vmax=8000+80) + + for (n, projection) in enumerate(projections) + ax = fig.add_subplot(1, n_subplots, n, projection=projection) + plot_cubed_sphere_tracer_field!(fig, ax, η, file["grid"]; add_colorbar = (n == n_subplots), subplot_kwargs...) + n_subplots == 1 && ax.set_title(plot_title) + + gl = ax.gridlines(color="gray", alpha=0.5, linestyle="--") + gl.xlocator = mticker.FixedLocator(-180:30:180) + gl.ylocator = mticker.FixedLocator(-80:20:80) + end + + n_subplots > 1 && fig.suptitle(plot_title, y=0.85) + + filename = @sprintf("cubed_sphere_rossby_haurwitz_%04d.png", n) + plt.savefig(filename, dpi=200, bbox_inches="tight") + plt.close(fig) + end + + close(file) + + filename_pattern = "cubed_sphere_rossby_haurwitz_%04d.png" + output_filename = "cubed_sphere_rossby_haurwitz.mp4" + + # Need extra crop video filter in case we end up with odd number of pixels in width or height. + # See: https://stackoverflow.com/a/29582287 + run(`ffmpeg -y -i $filename_pattern -c:v libx264 -vf "fps=10, crop=trunc(iw/2)*2:trunc(ih/2)*2" -pix_fmt yuv420p $output_filename`) + + [rm(f) for f in glob("cubed_sphere_rossby_haurwitz_*.png")] + + return nothing +end diff --git a/validation/cubed_sphere_rossby_haurwitz/animate_three_spheres.jl b/validation/cubed_sphere_rossby_haurwitz/animate_three_spheres.jl new file mode 100644 index 0000000000..ad68d5f936 --- /dev/null +++ b/validation/cubed_sphere_rossby_haurwitz/animate_three_spheres.jl @@ -0,0 +1,104 @@ +using Printf +using Glob +using JLD2 +using PyCall + +using Oceananigans.Utils: prettytime + +np = pyimport("numpy") +ma = pyimport("numpy.ma") +plt = pyimport("matplotlib.pyplot") +mticker = pyimport("matplotlib.ticker") +ccrs = pyimport("cartopy.crs") +cmocean = pyimport("cmocean") + +function plot_cubed_sphere_tracer_field!(fig, ax, var, grid; add_colorbar, transform, cmap, vmin, vmax) + + Nf = grid["faces"] |> keys |> length + Nx = grid["faces/1/Nx"] + Ny = grid["faces/1/Ny"] + Hx = grid["faces/1/Hx"] + Hy = grid["faces/1/Hy"] + + for face in 1:Nf + λᶠᶠᵃ = grid["faces/$face/λᶠᶠᵃ"][1+Hx:Nx+2Hx, 1+Hy:Ny+2Hy] + φᶠᶠᵃ = grid["faces/$face/φᶠᶠᵃ"][1+Hx:Nx+2Hx, 1+Hy:Ny+2Hy] + + var_face = var[face][:, :, 1] + + var_face_masked = ma.masked_where(np.isnan(var_face), var_face) + + im = ax.pcolormesh(λᶠᶠᵃ, φᶠᶠᵃ, var_face_masked; transform, cmap, vmin, vmax) + + # Add colorbar below all the subplots. + if add_colorbar && face == Nf + ax_cbar = fig.add_axes([0.25, 0.1, 0.5, 0.02]) + fig.colorbar(im, cax=ax_cbar, orientation="horizontal") + end + + ax.set_global() + end + + return ax +end + +function animate_rossby_haurwitz_three_spheres(; projection=ccrs.NearsidePerspective(central_longitude=0, central_latitude=0)) + + ## Extract data + + file = jldopen("cubed_sphere_rossby_haurwitz.jld2") + + iterations = parse.(Int, keys(file["timeseries/t"])) + + ## Makie movie of u, v, η + + for (n, i) in enumerate(iterations) + @info "Plotting iteration $i/$(iterations[end]) (frame $n/$(length(iterations)))..." + + u = file["timeseries/u/$i"] + v = file["timeseries/v/$i"] + η = file["timeseries/η/$i"] + + t = prettytime(file["timeseries/t/$i"]) + plot_title = "Rossby-Haurwitz wave (mode 4) at t = $t" + + fig = plt.figure(figsize=(16, 9)) + + ax = fig.add_subplot(1, 3, 1, projection=projection) + plot_cubed_sphere_tracer_field!(fig, ax, u, file["grid"], transform=ccrs.PlateCarree(), cmap=cmocean.cm.balance, vmin=-80, vmax=80, add_colorbar=false) + gl = ax.gridlines(color="gray", alpha=0.5, linestyle="--") + gl.xlocator = mticker.FixedLocator(-180:30:180) + gl.ylocator = mticker.FixedLocator(-80:20:80) + + ax = fig.add_subplot(1, 3, 2, projection=projection) + plot_cubed_sphere_tracer_field!(fig, ax, v, file["grid"], transform=ccrs.PlateCarree(), cmap=cmocean.cm.balance, vmin=-80, vmax=80, add_colorbar=false) + gl = ax.gridlines(color="gray", alpha=0.5, linestyle="--") + gl.xlocator = mticker.FixedLocator(-180:30:180) + gl.ylocator = mticker.FixedLocator(-80:20:80) + + ax = fig.add_subplot(1, 3, 3, projection=projection) + plot_cubed_sphere_tracer_field!(fig, ax, η, file["grid"], transform=ccrs.PlateCarree(), cmap=cmocean.cm.balance, vmin=8000, vmax=8250, add_colorbar=false) + gl = ax.gridlines(color="gray", alpha=0.5, linestyle="--") + gl.xlocator = mticker.FixedLocator(-180:30:180) + gl.ylocator = mticker.FixedLocator(-80:20:80) + + fig.suptitle(plot_title, y=0.75) + + filename = @sprintf("cubed_sphere_rossby_haurwitz_%04d.png", n) + plt.savefig(filename, dpi=200, bbox_inches="tight") + plt.close(fig) + end + + close(file) + + filename_pattern = "cubed_sphere_rossby_haurwitz_%04d.png" + output_filename = "cubed_sphere_rossby_haurwitz_three_spheres.mp4" + + # Need extra crop video filter in case we end up with odd number of pixels in width or height. + # See: https://stackoverflow.com/a/29582287 + run(`ffmpeg -y -i $filename_pattern -c:v libx264 -vf "fps=10, crop=trunc(iw/2)*2:trunc(ih/2)*2" -pix_fmt yuv420p $output_filename`) + + [rm(f) for f in glob("cubed_sphere_rossby_haurwitz_*.png")] + + return nothing +end diff --git a/validation/cubed_sphere_rossby_haurwitz/cubed_sphere_rossby_haurwitz.jl b/validation/cubed_sphere_rossby_haurwitz/cubed_sphere_rossby_haurwitz.jl new file mode 100644 index 0000000000..898a96aa41 --- /dev/null +++ b/validation/cubed_sphere_rossby_haurwitz/cubed_sphere_rossby_haurwitz.jl @@ -0,0 +1,216 @@ +using Statistics +using Logging +using Printf +using DataDeps +using JLD2 + +using Oceananigans +using Oceananigans.Units + +using Oceananigans.Diagnostics: accurate_cell_advection_timescale + +Logging.global_logger(OceananigansLogger()) + +##### +##### Progress monitor +##### + +mutable struct Progress + interval_start_time :: Float64 +end + +function (p::Progress)(sim) + wall_time = (time_ns() - p.interval_start_time) * 1e-9 + progress = 100 * (sim.model.clock.time / sim.stop_time) + + @info @sprintf("[%06.2f%%] Time: %s, iteration: %d, max(|u⃗|): (%.2e, %.2e) m/s, extrema(η): (min=%.2e, max=%.2e), CFL: %.2e, Δ(wall time): %s", + progress, + prettytime(sim.model.clock.time), + sim.model.clock.iteration, + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + minimum(sim.model.free_surface.η), + maximum(sim.model.free_surface.η), + sim.parameters.cfl(sim.model), + prettytime(wall_time)) + + p.interval_start_time = time_ns() + + return nothing +end + +##### +##### Script starts here +##### + +ENV["DATADEPS_ALWAYS_ACCEPT"] = "true" + +Logging.global_logger(OceananigansLogger()) + +dd32 = DataDep("cubed_sphere_32_grid", + "Conformal cubed sphere grid with 32×32 grid points on each face", + "https://github.com/CliMA/OceananigansArtifacts.jl/raw/main/cubed_sphere_grids/cubed_sphere_32_grid.jld2", + "b1dafe4f9142c59a2166458a2def743cd45b20a4ed3a1ae84ad3a530e1eff538" # sha256sum +) + +dd96 = DataDep("cubed_sphere_96_grid", + "Conformal cubed sphere grid with 96×96 grid points on each face", + "https://github.com/CliMA/OceananigansArtifacts.jl/raw/main/cubed_sphere_grids/cs96/cubed_sphere_96_grid.jld2" +) + +DataDeps.register(dd32) +DataDeps.register(dd96) + +function diagnose_velocities_from_streamfunction(ψ, grid) + ψᶠᶠᶜ = Field(Face, Face, Center, CPU(), grid, nothing, nothing) + uᶠᶜᶜ = Field(Face, Center, Center, CPU(), grid, nothing, nothing) + vᶜᶠᶜ = Field(Center, Face, Center, CPU(), grid, nothing, nothing) + + for (f, grid_face) in enumerate(grid.faces) + Nx, Ny, Nz = size(grid_face) + for i in 1:Nx+1, j in 1:Ny+1 + ψᶠᶠᶜ.faces[f][i, j, 1] = ψ(grid_face.λᶠᶠᵃ[i, j], grid_face.φᶠᶠᵃ[i, j]) + end + + for i in 1:Nx+1, j in 1:Ny + uᶠᶜᶜ.faces[f][i, j, 1] = (ψᶠᶠᶜ.faces[f][i, j, 1] - ψᶠᶠᶜ.faces[f][i, j+1, 1]) / grid.faces[f].Δyᶠᶜᵃ[i, j] + end + + for i in 1:Nx, j in 1:Ny+1 + vᶜᶠᶜ.faces[f][i, j, 1] = (ψᶠᶠᶜ.faces[f][i+1, j, 1] - ψᶠᶠᶜ.faces[f][i, j, 1]) / grid.faces[f].Δxᶜᶠᵃ[i, j] + end + end + + return uᶠᶜᶜ, vᶜᶠᶜ, ψᶠᶠᶜ +end + +function cubed_sphere_rossby_haurwitz(grid_filepath) + + ## Grid setup + + H = 8kilometers + grid = ConformalCubedSphereGrid(grid_filepath, Nz=1, z=(-H, 0)) + + ## Model setup + + model = HydrostaticFreeSurfaceModel( + architecture = CPU(), + grid = grid, + momentum_advection = nothing, + free_surface = ExplicitFreeSurface(gravitational_acceleration=100), + coriolis = HydrostaticSphericalCoriolis(scheme = VectorInvariantEnstrophyConserving()), + closure = nothing, + tracers = nothing, + buoyancy = nothing + ) + + ## Rossby-Haurwitz initial condition from Williamson et al. (§3.6, 1992) + ## # here: θ ∈ [-π/2, π/2] is latitude, ϕ ∈ [0, 2π) is longitude + + R = grid.faces[1].radius + K = 7.848e-6 + ω = 0 + n = 4 + + 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)) + + ψ(θ, ϕ) = -R^2 * ω * sin(θ)^2 + R^2 * K * cos(θ)^n * sin(θ) * cos(n*ϕ) + + u(θ, ϕ) = R * ω * cos(θ) + R * K * cos(θ)^(n-1) * (n * sin(θ)^2 - cos(θ)^2) * cos(n*ϕ) + v(θ, ϕ) = -n * K * R * cos(θ)^(n-1) * sin(θ) * sin(n*ϕ) + + h(θ, ϕ) = H + R^2/g * (A(θ) + B(θ) * cos(n * ϕ) + C(θ) * cos(2n * ϕ)) + + # Total initial conditions + # Previously: θ ∈ [-π/2, π/2] is latitude, ϕ ∈ [0, 2π) is longitude + # Oceananigans: ϕ ∈ [-90, 90], λ ∈ [-180, 180], + + rescale¹(λ) = (λ + 180)/ 360 * 2π # λ to θ + rescale²(ϕ) = ϕ / 180 * π # θ to ϕ + + # arguments were u(θ, ϕ), λ |-> ϕ, θ |-> ϕ + uᵢ(λ, ϕ, z) = u(rescale²(ϕ), rescale¹(λ)) + vᵢ(λ, ϕ, z) = v(rescale²(ϕ), rescale¹(λ)) + ηᵢ(λ, ϕ) = h(rescale²(ϕ), rescale¹(λ)) + + # set!(model, u=uᵢ, v=vᵢ, η = ηᵢ) + + ψ₀(λ, φ) = ψ(rescale²(φ), rescale¹(λ)) + + u₀, v₀, _ = diagnose_velocities_from_streamfunction(ψ₀, grid) + + Oceananigans.set!(model, u=u₀, v=v₀, η=ηᵢ) + + ## Simulation setup + + # Compute amount of time needed for a 45° rotation. + angular_velocity = (n * (3+n) * ω - 2Ω) / ((1+n) * (2+n)) + stop_time = deg2rad(45) / abs(angular_velocity) + @info "Stop time = $(prettytime(stop_time))" + + Δt = 20seconds + + gravity_wave_speed = sqrt(g * H) + min_spacing = filter(!iszero, grid.faces[1].Δyᶠᶠᵃ) |> minimum + wave_propagation_time_scale = min_spacing / gravity_wave_speed + gravity_wave_cfl = Δt / wave_propagation_time_scale + @info "Gravity wave CFL = $gravity_wave_cfl" + + if !isnothing(model.closure) + ν = model.closure.νh + diffusive_cfl = ν * Δt / min_spacing^2 + @info "Diffusive CFL = $diffusive_cfl" + end + + cfl = CFL(Δt, accurate_cell_advection_timescale) + + simulation = Simulation(model, + Δt = Δt, + stop_time = stop_time, + iteration_interval = 20, + progress = Progress(time_ns()), + parameters = (; cfl) + ) + + fields_to_check = ( + u = model.velocities.u, + v = model.velocities.v, + η = model.free_surface.η + ) + + simulation.diagnostics[:state_checker] = + StateChecker(model, fields=fields_to_check, schedule=IterationInterval(20)) + + output_fields = merge(model.velocities, (η=model.free_surface.η,)) + + simulation.output_writers[:fields] = + JLD2OutputWriter(model, output_fields, + schedule = TimeInterval(5minutes), + prefix = "cubed_sphere_rossby_haurwitz", + force = true) + + run!(simulation) + + return simulation +end + + +include("animate_on_map_projection.jl") + +function run_cubed_sphere_rossby_haurwitz_validation() + + cubed_sphere_rossby_haurwitz(datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2") + # cubed_sphere_rossby_haurwitz(datadep"cubed_sphere_96_grid/cubed_sphere_96_grid.jld2") + + projections = [ + ccrs.NearsidePerspective(central_longitude=0, central_latitude=30), + ccrs.NearsidePerspective(central_longitude=180, central_latitude=-30) + ] + + animate_rossby_haurwitz(projections=projections) +end diff --git a/validation/surface_gravity_waves_on_sphere/animate.jl b/validation/cubed_sphere_surface_gravity_waves/animate.jl similarity index 100% rename from validation/surface_gravity_waves_on_sphere/animate.jl rename to validation/cubed_sphere_surface_gravity_waves/animate.jl diff --git a/validation/cubed_sphere_surface_gravity_waves/animate_face.jl b/validation/cubed_sphere_surface_gravity_waves/animate_face.jl new file mode 100644 index 0000000000..0ee94d8ce5 --- /dev/null +++ b/validation/cubed_sphere_surface_gravity_waves/animate_face.jl @@ -0,0 +1,53 @@ +using Statistics +using JLD2 +using Printf +using GLMakie + +# using Oceananigans.Grids +# using Oceananigans.Utils: prettytime, hours, day, days, years + +function geographic2cartesian(λ, φ, r=1) + λ_azimuthal = λ .+ 180 # Convert to λ ∈ [0°, 360°] + φ_azimuthal = 90 .- φ # Convert to φ ∈ [0°, 180°] (0° at north pole) + + x = @. r * cosd(λ_azimuthal) * sind(φ_azimuthal) + y = @. r * sind(λ_azimuthal) * sind(φ_azimuthal) + z = @. r * cosd(φ_azimuthal) + + return x, y, z +end + +ds_cs = jldopen("cubed_sphere_face_waves.jld2") + +λᶜᶜᵃ = ds_cs["grid/λᶜᶜᵃ"][2:33, 2:33] +φᶜᶜᵃ = ds_cs["grid/φᶜᶜᵃ"][2:33, 2:33] + +xc, yc, zc = geographic2cartesian(λᶜᶜᵃ, φᶜᶜᵃ) + +iterations = parse.(Int, keys(ds_cs["timeseries/t"])) + +iter = Node(0) + +plot_title = @lift @sprintf("η′ on a cubed sphere face: time = %s", prettytime(ds_cs["timeseries/t/" * string($iter)])) + +η = @lift ds_cs["timeseries/η/" * string($iter)][:, :, 1] + +fig = Figure(resolution = (1920, 1080)) + +ax = fig[1, 1] = LScene(fig) # make plot area wider +wireframe!(ax, Sphere(Point3f0(0), 0.99f0), show_axis=false) +sf = surface!(ax, xc, yc, zc, color=η, colormap=:balance, colorrange=(-0.01, 0.01)) +rotate_cam!(ax.scene, (3π/4, π/6, 0)) +zoom!(ax.scene, (0, 0, 0), 5, false) +# fig[2, 2 + 3*(n-1)] = Label(fig, statenames[n], textsize = 50) # put names in center + +cb1 = fig[1, 2] = Colorbar(fig, sf, label="η′", width=30) + +supertitle = fig[0, :] = Label(fig, plot_title, textsize=50) + +record(fig, "cubed_sphere_waves.mp4", iterations, framerate=60) do i + @info "Animating iteration $i/$(iterations[end])..." + iter[] = i +end + +close(ds_cs) diff --git a/validation/cubed_sphere_surface_gravity_waves/animate_on_map_projection.jl b/validation/cubed_sphere_surface_gravity_waves/animate_on_map_projection.jl new file mode 100644 index 0000000000..2530839bbd --- /dev/null +++ b/validation/cubed_sphere_surface_gravity_waves/animate_on_map_projection.jl @@ -0,0 +1,107 @@ +using Printf +using Glob +using JLD2 +using PyCall + +using Oceananigans.Utils: prettytime + +np = pyimport("numpy") +ma = pyimport("numpy.ma") +plt = pyimport("matplotlib.pyplot") +mticker = pyimport("matplotlib.ticker") +ccrs = pyimport("cartopy.crs") +cmocean = pyimport("cmocean") + +function plot_cubed_sphere_tracer_field!(fig, ax, var, grid; add_colorbar, transform, cmap, vmin, vmax) + + Nf = grid["faces"] |> keys |> length + Nx = grid["faces/1/Nx"] + Ny = grid["faces/1/Ny"] + Hx = grid["faces/1/Hx"] + Hy = grid["faces/1/Hy"] + + for face in 1:Nf + λᶠᶠᵃ = grid["faces/$face/λᶠᶠᵃ"][1+Hx:Nx+2Hx, 1+Hy:Ny+2Hy] + φᶠᶠᵃ = grid["faces/$face/φᶠᶠᵃ"][1+Hx:Nx+2Hx, 1+Hy:Ny+2Hy] + + var_face = var[face][:, :, 1] + + # Remove very specific problematic grid cells near λ = 180° on face 6 that mess up the plot. + # May be related to https://github.com/SciTools/cartopy/issues/1151 + # Not needed for Orthographic or NearsidePerspective projection so let's use those. + # if face == 6 + # for i in 1:Nx+1, j in 1:Ny+1 + # if isapprox(λᶠᶠᵃ[i, j], -180, atol=15) + # var_face[min(i, Nx), min(j, Ny)] = NaN + # end + # end + # end + + var_face_masked = ma.masked_where(np.isnan(var_face), var_face) + + im = ax.pcolormesh(λᶠᶠᵃ, φᶠᶠᵃ, var_face_masked; transform, cmap, vmin, vmax) + + # Add colorbar below all the subplots. + if add_colorbar && face == Nf + ax_cbar = fig.add_axes([0.25, 0.1, 0.5, 0.02]) + fig.colorbar(im, cax=ax_cbar, orientation="horizontal") + end + + ax.set_global() + end + + return ax +end + +function animate_surface_gravity_waves(; face_number, projections=[ccrs.Robinson()]) + + ## Extract data + + file = jldopen("cubed_sphere_surface_gravity_waves_face$face_number.jld2") + + iterations = parse.(Int, keys(file["timeseries/t"])) + + ## Makie movie of tracer field h + + for (n, i) in enumerate(iterations) + @info "Plotting face $face_number iteration $i/$(iterations[end]) (frame $n/$(length(iterations)))..." + + η = file["timeseries/η/$i"] + + t = prettytime(file["timeseries/t/$i"]) + plot_title = "Surface gravity waves from face $face_number: η(λ, φ) at t = $t" + + fig = plt.figure(figsize=(16, 9)) + n_subplots = length(projections) + subplot_kwargs = (transform=ccrs.PlateCarree(), cmap=cmocean.cm.balance, vmin=-0.01, vmax=0.01) + + for (n, projection) in enumerate(projections) + ax = fig.add_subplot(1, n_subplots, n, projection=projection) + plot_cubed_sphere_tracer_field!(fig, ax, η, file["grid"]; add_colorbar = (n == n_subplots), subplot_kwargs...) + n_subplots == 1 && ax.set_title(plot_title) + + gl = ax.gridlines(color="gray", alpha=0.5, linestyle="--") + gl.xlocator = mticker.FixedLocator(-180:30:180) + gl.ylocator = mticker.FixedLocator(-80:20:80) + end + + n_subplots > 1 && fig.suptitle(plot_title, y=0.85) + + filename = @sprintf("cubed_sphere_surface_gravity_waves_face%d_%04d.png", face_number, n) + plt.savefig(filename, dpi=200, bbox_inches="tight") + plt.close(fig) + end + + close(file) + + filename_pattern = "cubed_sphere_surface_gravity_waves_face$(face_number)_%04d.png" + output_filename = "cubed_sphere_surface_gravity_waves_face$(face_number).mp4" + + # Need extra crop video filter in case we end up with odd number of pixels in width or height. + # See: https://stackoverflow.com/a/29582287 + run(`ffmpeg -y -i $filename_pattern -c:v libx264 -vf "fps=15, crop=trunc(iw/2)*2:trunc(ih/2)*2" -pix_fmt yuv420p $output_filename`) + + [rm(f) for f in glob("cubed_sphere_surface_gravity_waves_face$(face_number)_*.png")] + + return nothing +end diff --git a/validation/cubed_sphere_surface_gravity_waves/animate_sphere.jl b/validation/cubed_sphere_surface_gravity_waves/animate_sphere.jl new file mode 100644 index 0000000000..c0eabd173f --- /dev/null +++ b/validation/cubed_sphere_surface_gravity_waves/animate_sphere.jl @@ -0,0 +1,134 @@ +using Statistics +using JLD2 +using Printf +using GLMakie + +using Oceananigans.Utils: prettytime + +file = jldopen("full_cubed_sphere_gravity_waves.jld2") + +iterations = parse.(Int, keys(file["timeseries/t"])) + +# Azimuthal spherical coordinates: +# Convert to λ ∈ [0°, 360°] +# Convert to φ ∈ [0°, 180°] (0° at north pole) +geographic2x(λ, φ, r=1) = r * cosd(λ + 180) * sind(90 - φ) +geographic2y(λ, φ, r=1) = r * sind(λ + 180) * sind(90 - φ) +geographic2z(λ, φ, r=1) = r * cosd(90 - φ) + +Nf = file["grid/faces"] |> keys |> length +Nx = file["grid/faces/1/Nx"] +Ny = file["grid/faces/1/Ny"] +Nz = file["grid/faces/1/Nz"] + +##### +##### Plot η +##### + +size_2d = (Nx, Ny * Nf) + +flatten_cubed_sphere(field, size) = reshape(cat(field..., dims=4), size) + +λ = flatten_cubed_sphere((file["grid/faces/$n/λᶜᶜᵃ"][2:33, 2:33] for n in 1:6), size_2d) +φ = flatten_cubed_sphere((file["grid/faces/$n/φᶜᶜᵃ"][2:33, 2:33] for n in 1:6), size_2d) + +x = geographic2x.(λ, φ) +y = geographic2y.(λ, φ) +z = geographic2z.(λ, φ) + +iter = Node(0) + +plot_title = @lift @sprintf("Surface gravity waves on a cubed sphere: time = %s", prettytime(file["timeseries/t/" * string($iter)])) + +η = @lift flatten_cubed_sphere(file["timeseries/η/" * string($iter)], size_2d) + +fig = Figure(resolution = (1920, 1080)) +ax = fig[1, 1] = LScene(fig) # make plot area wider +wireframe!(ax, Sphere(Point3f0(0), 0.99f0), show_axis=false) +sf = surface!(ax, x, y, z, color=η, colormap=:balance, colorrange=(-0.01, 0.01)) + +rotate_cam!(ax.scene, (3π/4, 0, 0)) +zoom!(ax.scene, (0, 0, 0), 8, false) + +cb1 = fig[1, 2] = Colorbar(fig, sf, label="sea surface height η′ (m)", width=30) + +supertitle = fig[0, :] = Label(fig, plot_title, textsize=50) + +record(fig, "surface_gravity_waves_on_a_cubed_sphere.mp4", iterations, framerate=15) do i + @info "Animating iteration $i/$(iterations[end])..." + iter[] = i +end + +##### +##### Plot u +##### + +# size_2d = (Nx+1, Ny * Nf) + +# λ = flatten_cubed_sphere((file["grid/faces/$n/λᶠᶠᵃ"][2:34, 2:33] for n in 1:6), size_2d) +# φ = flatten_cubed_sphere((file["grid/faces/$n/φᶠᶠᵃ"][2:34, 2:33] for n in 1:6), size_2d) + +# x = geographic2x.(λ, φ) +# y = geographic2y.(λ, φ) +# z = geographic2z.(λ, φ) + +# iter = Node(0) + +# plot_title = @lift @sprintf("Surface gravity waves on a cubed sphere: time = %s", prettytime(file["timeseries/t/" * string($iter)])) + +# u = @lift flatten_cubed_sphere(file["timeseries/u/" * string($iter)]) + +# fig = Figure(resolution = (1920, 1080)) +# ax = fig[1, 1] = LScene(fig) # make plot area wider +# wireframe!(ax, Sphere(Point3f0(0), 0.99f0), show_axis=false) +# sf = surface!(ax, x, y, z, color=u, colormap=:balance, colorrange=(-1e-4, 1e-4)) + +# rotate_cam!(ax.scene, (3π/4, 0, 0)) +# zoom!(ax.scene, (0, 0, 0), 8, false) + +# cb1 = fig[1, 2] = Colorbar(fig, sf, label="u-velocity (m/s)", width=30) + +# supertitle = fig[0, :] = Label(fig, plot_title, textsize=50) + +# record(fig, "surface_gravity_waves_on_a_cubed_sphere_u.mp4", iterations[1:20], framerate=15) do i +# @info "Animating iteration $i/$(iterations[end])..." +# iter[] = i +# end + +##### +##### Plot v +##### + +# size_2d = (Nx, (Ny+1) * Nf) + +# λ = flatten_cubed_sphere((file["grid/faces/$n/λᶠᶠᵃ"][2:33, 2:34] for n in 1:6), size_2d) +# φ = flatten_cubed_sphere((file["grid/faces/$n/φᶠᶠᵃ"][2:33, 2:34] for n in 1:6), size_2d) + +# x = geographic2x.(λ, φ) +# y = geographic2y.(λ, φ) +# z = geographic2z.(λ, φ) + +# iter = Node(0) + +# plot_title = @lift @sprintf("Surface gravity waves on a cubed sphere: time = %s", prettytime(file["timeseries/t/" * string($iter)])) + +# v = @lift flatten_cubed_sphere(file["timeseries/v/" * string($iter)]) + +# fig = Figure(resolution = (1920, 1080)) +# ax = fig[1, 1] = LScene(fig) # make plot area wider +# wireframe!(ax, Sphere(Point3f0(0), 0.99f0), show_axis=false) +# sf = surface!(ax, x, y, z, color=v, colormap=:balance, colorrange=(-1e-4, 1e-4)) + +# rotate_cam!(ax.scene, (3π/4, 0, 0)) +# zoom!(ax.scene, (0, 0, 0), 8, false) + +# cb1 = fig[1, 2] = Colorbar(fig, sf, label="v-velocity (m/s)", width=30) + +# supertitle = fig[0, :] = Label(fig, plot_title, textsize=50) + +# record(fig, "surface_gravity_waves_on_a_cubed_sphere_v.mp4", iterations[1:20], framerate=15) do i +# @info "Animating iteration $i/$(iterations[end])..." +# iter[] = i +# end + +close(file) diff --git a/validation/cubed_sphere_surface_gravity_waves/cubed_sphere_surface_gravity_waves.jl b/validation/cubed_sphere_surface_gravity_waves/cubed_sphere_surface_gravity_waves.jl new file mode 100644 index 0000000000..0ba37a1759 --- /dev/null +++ b/validation/cubed_sphere_surface_gravity_waves/cubed_sphere_surface_gravity_waves.jl @@ -0,0 +1,147 @@ +using Statistics +using Logging +using Printf +using DataDeps +using JLD2 + +using Oceananigans +using Oceananigans.Units + +using Oceananigans.Diagnostics: accurate_cell_advection_timescale + +Logging.global_logger(OceananigansLogger()) + +##### +##### Progress monitor +##### + +mutable struct Progress + interval_start_time :: Float64 +end + +function (p::Progress)(sim) + wall_time = (time_ns() - p.interval_start_time) * 1e-9 + + @info @sprintf("Time: %s, iteration: %d, max(|u⃗|): (%.2e, %.2e) m/s, extrema(η): (min=%.2e, max=%.2e), CFL: %.2e, wall time: %s", + prettytime(sim.model.clock.time), + sim.model.clock.iteration, + maximum(abs, sim.model.velocities.u), + maximum(abs, sim.model.velocities.v), + minimum(sim.model.free_surface.η), + maximum(sim.model.free_surface.η), + sim.parameters.cfl(sim.model), + prettytime(wall_time)) + + p.interval_start_time = time_ns() + + return nothing +end + +##### +##### Script starts here +##### + +ENV["DATADEPS_ALWAYS_ACCEPT"] = "true" + +dd = DataDep("cubed_sphere_32_grid", + "Conformal cubed sphere grid with 32×32 grid points on each face", + "https://github.com/CliMA/OceananigansArtifacts.jl/raw/main/cubed_sphere_grids/cubed_sphere_32_grid.jld2", + "b1dafe4f9142c59a2166458a2def743cd45b20a4ed3a1ae84ad3a530e1eff538" # sha256sum +) + +DataDeps.register(dd) + +cs32_filepath = datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2" + +# Central (λ, φ) for each face of the cubed sphere +central_longitude = (0, 90, 0, 180, -90, 0) +central_latitude = (0, 0, 90, 0, 0, -90) + +function cubed_sphere_surface_gravity_waves(; face_number) + + H = 4kilometers + grid = ConformalCubedSphereGrid(cs32_filepath, Nz=1, z=(-H, 0)) + + ## Model setup + + model = HydrostaticFreeSurfaceModel( + architecture = CPU(), + grid = grid, + momentum_advection = nothing, + free_surface = ExplicitFreeSurface(gravitational_acceleration=0.1), + coriolis = nothing, + closure = nothing, + tracers = nothing, + buoyancy = nothing + ) + + ## Initial condition: + ## Very small sea surface height perturbation so the resulting dynamics are well-described + ## by a linear free surface. + + A = 1e-5 * H # Amplitude of the perturbation + λ₀ = central_longitude[face_number] + φ₀ = central_latitude[face_number] + Δλ = 15 # Longitudinal width + Δφ = 15 # Latitudinal width + + η′(λ, φ) = A * exp(- (λ - λ₀)^2 / Δλ^2) * exp(- (φ - φ₀)^2 / Δφ^2) + + Oceananigans.set!(model, η=η′) + + ## Simulation setup + + Δt = 10minutes + + cfl = CFL(Δt, accurate_cell_advection_timescale) + + simulation = Simulation(model, + Δt = Δt, + stop_time = 25days, + iteration_interval = 1, + progress = Progress(time_ns()), + parameters = (; cfl) + ) + + fields_to_check = ( + u = model.velocities.u, + v = model.velocities.v, + η = model.free_surface.η, + Gu = model.timestepper.Gⁿ.u, + Gv = model.timestepper.Gⁿ.v, + Gη = model.timestepper.Gⁿ.η + ) + + simulation.diagnostics[:state_checker] = + StateChecker(model, fields=fields_to_check, schedule=IterationInterval(1)) + + output_fields = merge(model.velocities, (η=model.free_surface.η,)) + + simulation.output_writers[:fields] = + JLD2OutputWriter(model, output_fields, + schedule = TimeInterval(1hour), + prefix = "cubed_sphere_surface_gravity_waves_face$face_number", + force = true) + + run!(simulation) + + return simulation +end + +include("animate_on_map_projection.jl") + +function run_cubed_sphere_surface_gravity_waves_validation() + + for f in 1:6 + cubed_sphere_surface_gravity_waves(face_number=f) + end + + projections = [ + ccrs.NearsidePerspective(central_longitude=0, central_latitude=30), + ccrs.NearsidePerspective(central_longitude=180, central_latitude=-30) + ] + + for f in 1:6 + animate_surface_gravity_waves(face_number=f, projections=projections) + end +end diff --git a/validation/surface_gravity_waves_on_sphere/surface_gravity_waves_on_sphere.jl b/validation/cubed_sphere_surface_gravity_waves/surface_gravity_waves_on_face.jl similarity index 93% rename from validation/surface_gravity_waves_on_sphere/surface_gravity_waves_on_sphere.jl rename to validation/cubed_sphere_surface_gravity_waves/surface_gravity_waves_on_face.jl index ae04c73312..7d5ad24165 100644 --- a/validation/surface_gravity_waves_on_sphere/surface_gravity_waves_on_sphere.jl +++ b/validation/cubed_sphere_surface_gravity_waves/surface_gravity_waves_on_face.jl @@ -17,7 +17,7 @@ Logging.global_logger(OceananigansLogger()) dd = DataDep("cubed_sphere_32_grid", "Conformal cubed sphere grid with 32×32 grid points on each face", "https://github.com/CliMA/OceananigansArtifacts.jl/raw/main/cubed_sphere_grids/cubed_sphere_32_grid.jld2", - "3cc5d86290c3af028cddfa47e61e095ee470fe6f8d779c845de09da2f1abeb15" # sha256sum + "b1dafe4f9142c59a2166458a2def743cd45b20a4ed3a1ae84ad3a530e1eff538" # sha256sum ) DataDeps.register(dd) @@ -68,13 +68,7 @@ A = 1e-5 * H # Amplitude of the perturbation η′(λ, φ, z) = A * exp(- (λ - λ₀)^2 / Δλ^2) * exp(- (φ - φ₀)^2 / Δφ^2) -## set! doesn't work for ConformalCubedSphereFaceGrid yet! - -# set!(model, η=η′) - -for i in 1:grid.Nx, j in 1:grid.Ny - model.free_surface.η[i, j, 1] = η′(grid.λᶜᶜᵃ[i, j], grid.φᶜᶜᵃ[i, j], 0) -end +set!(model, η=η′) # g = model.free_surface.gravitational_acceleration # gravity_wave_speed = sqrt(g * H) # hydrostatic (shallow water) gravity wave speed diff --git a/validation/cubed_sphere_tracer_advection/animate_faces.jl b/validation/cubed_sphere_tracer_advection/animate_faces.jl new file mode 100644 index 0000000000..98cb37e773 --- /dev/null +++ b/validation/cubed_sphere_tracer_advection/animate_faces.jl @@ -0,0 +1,76 @@ +using Printf +using Glob +using JLD2 +using PyCall + +plt = pyimport("matplotlib.pyplot") +gridspec = pyimport("matplotlib.gridspec") +ccrs = pyimport("cartopy.crs") +cmocean = pyimport("cmocean") + +## Extract data + +file = jldopen("tracer_advection_over_the_poles.jld2") + +iterations = parse.(Int, keys(file["timeseries/t"])) + +Nf = file["grid/faces"] |> keys |> length +Nx = file["grid/faces/1/Nx"] +Ny = file["grid/faces/1/Ny"] +Nz = file["grid/faces/1/Nz"] + +## Plot! + +for (n, i) in enumerate(iterations) + @info "Plotting iteration $i/$(iterations[end]) (frame $n/$(length(iterations)))..." + + # fig = plt.figure(constrained_layout=true) + fig = plt.figure() + + ax1 = plt.subplot2grid((3, 4), (2, 0)) + ax2 = plt.subplot2grid((3, 4), (2, 1)) + ax3 = plt.subplot2grid((3, 4), (1, 1)) + ax4 = plt.subplot2grid((3, 4), (1, 2)) + ax5 = plt.subplot2grid((3, 4), (0, 2)) + ax6 = plt.subplot2grid((3, 4), (0, 3)) + + ax1.set_title("face 1") + ax1.pcolormesh(file["timeseries/h/$i"][1][:, :, 1]', cmap=cmocean.cm.thermal, vmin=0, vmax=1000) + ax1.get_xaxis().set_visible(false) + ax1.get_yaxis().set_visible(false) + + ax2.set_title("face 2") + ax2.pcolormesh(file["timeseries/h/$i"][2][:, :, 1]', cmap=cmocean.cm.thermal, vmin=0, vmax=1000) + ax2.get_xaxis().set_visible(false) + ax2.get_yaxis().set_visible(false) + + ax3.set_title("face 3") + ax3.pcolormesh(file["timeseries/h/$i"][3][:, :, 1]', cmap=cmocean.cm.thermal, vmin=0, vmax=1000) + ax3.get_xaxis().set_visible(false) + ax3.get_yaxis().set_visible(false) + + ax4.set_title("face 4") + ax4.pcolormesh(file["timeseries/h/$i"][4][:, :, 1]', cmap=cmocean.cm.thermal, vmin=0, vmax=1000) + ax4.get_xaxis().set_visible(false) + ax4.get_yaxis().set_visible(false) + + ax5.set_title("face 5") + ax5.pcolormesh(file["timeseries/h/$i"][5][:, :, 1]', cmap=cmocean.cm.thermal, vmin=0, vmax=1000) + ax5.get_xaxis().set_visible(false) + ax5.get_yaxis().set_visible(false) + + ax6.set_title("face 6") + ax6.pcolormesh(file["timeseries/h/$i"][6][:, :, 1]', cmap=cmocean.cm.thermal, vmin=0, vmax=1000) + ax6.get_xaxis().set_visible(false) + ax6.get_yaxis().set_visible(false) + + filename = @sprintf("tracer_avection_over_the_poles_%04d.png", n) + plt.savefig(filename, dpi=200) + plt.close(fig) +end + +close(file) + +run(`ffmpeg -y -i tracer_avection_over_the_poles_%04d.png -c:v libx264 -vf fps=10 -pix_fmt yuv420p out.mp4`) + +# [rm(f) for f in glob("*.png")]; diff --git a/validation/cubed_sphere_tracer_advection/animate_on_map_projection.jl b/validation/cubed_sphere_tracer_advection/animate_on_map_projection.jl new file mode 100644 index 0000000000..02fd56ae8a --- /dev/null +++ b/validation/cubed_sphere_tracer_advection/animate_on_map_projection.jl @@ -0,0 +1,107 @@ +using Printf +using Glob +using JLD2 +using PyCall + +using Oceananigans.Utils: prettytime + +np = pyimport("numpy") +ma = pyimport("numpy.ma") +plt = pyimport("matplotlib.pyplot") +mticker = pyimport("matplotlib.ticker") +ccrs = pyimport("cartopy.crs") +cmocean = pyimport("cmocean") + +function plot_cubed_sphere_tracer_field!(fig, ax, var, grid; add_colorbar, transform, cmap, vmin, vmax) + + Nf = grid["faces"] |> keys |> length + Nx = grid["faces/1/Nx"] + Ny = grid["faces/1/Ny"] + Hx = grid["faces/1/Hx"] + Hy = grid["faces/1/Hy"] + + for face in 1:Nf + λᶠᶠᵃ = grid["faces/$face/λᶠᶠᵃ"][1+Hx:Nx+2Hx, 1+Hy:Ny+2Hy] + φᶠᶠᵃ = grid["faces/$face/φᶠᶠᵃ"][1+Hx:Nx+2Hx, 1+Hy:Ny+2Hy] + + var_face = var[face][:, :, 1] + + # Remove very specific problematic grid cells near λ = 180° on face 6 that mess up the plot. + # May be related to https://github.com/SciTools/cartopy/issues/1151 + # Not needed for Orthographic or NearsidePerspective projection so let's use those. + # if face == 6 + # for i in 1:Nx+1, j in 1:Ny+1 + # if isapprox(λᶠᶠᵃ[i, j], -180, atol=15) + # var_face[min(i, Nx), min(j, Ny)] = NaN + # end + # end + # end + + var_face_masked = ma.masked_where(np.isnan(var_face), var_face) + + im = ax.pcolormesh(λᶠᶠᵃ, φᶠᶠᵃ, var_face_masked; transform, cmap, vmin, vmax) + + # Add colorbar below all the subplots. + if add_colorbar && face == Nf + ax_cbar = fig.add_axes([0.25, 0.1, 0.5, 0.02]) + fig.colorbar(im, cax=ax_cbar, orientation="horizontal") + end + + ax.set_global() + end + + return ax +end + +function animate_tracer_advection(; face_number, α, projections=[ccrs.Robinson()]) + + ## Extract data + + file = jldopen("tracer_advection_over_the_poles_face$(face_number)_alpha$α.jld2") + + iterations = parse.(Int, keys(file["timeseries/t"])) + + ## Makie movie of tracer field h + + for (n, i) in enumerate(iterations) + @info "Plotting face $face_number iteration $i/$(iterations[end]) (frame $n/$(length(iterations)))..." + + h = file["timeseries/h/$i"] + + t = prettytime(file["timeseries/t/$i"]) + plot_title = "Tracer advection from face $face_number at α = $(α)°: tracer at t = $t" + + fig = plt.figure(figsize=(16, 9)) + n_subplots = length(projections) + subplot_kwargs = (transform=ccrs.PlateCarree(), cmap=cmocean.cm.matter, vmin=0, vmax=1000) + + for (n, projection) in enumerate(projections) + ax = fig.add_subplot(1, n_subplots, n, projection=projection) + plot_cubed_sphere_tracer_field!(fig, ax, h, file["grid"]; add_colorbar = (n == n_subplots), subplot_kwargs...) + n_subplots == 1 && ax.set_title(plot_title) + + gl = ax.gridlines(color="gray", alpha=0.5, linestyle="--") + gl.xlocator = mticker.FixedLocator(-180:30:180) + gl.ylocator = mticker.FixedLocator(-80:20:80) + end + + n_subplots > 1 && fig.suptitle(plot_title, y=0.85) + + filename = @sprintf("cubed_sphere_tracer_advection_face%d_alpha%d_%04d.png", face_number, α, n) + plt.savefig(filename, dpi=200, bbox_inches="tight") + plt.close(fig) + end + + close(file) + + filename_pattern = "cubed_sphere_tracer_advection_face$(face_number)_alpha$(α)_%04d.png" + output_filename = "cubed_sphere_tracer_advection_face$(face_number)_alpha$(α).mp4" + + # Need extra crop video filter in case we end up with odd number of pixels in width or height. + # See: https://stackoverflow.com/a/29582287 + run(`ffmpeg -y -i $filename_pattern -c:v libx264 -vf "fps=15, crop=trunc(iw/2)*2:trunc(ih/2)*2" -pix_fmt yuv420p $output_filename`) + + [rm(f) for f in glob("cubed_sphere_tracer_advection_face$(face_number)_alpha$(α)_*.png")] + + return nothing +end diff --git a/validation/cubed_sphere_tracer_advection/cubed_sphere_tracer_advection.jl b/validation/cubed_sphere_tracer_advection/cubed_sphere_tracer_advection.jl new file mode 100644 index 0000000000..2f3c773690 --- /dev/null +++ b/validation/cubed_sphere_tracer_advection/cubed_sphere_tracer_advection.jl @@ -0,0 +1,176 @@ +using Statistics +using Logging +using Printf +using DataDeps +using JLD2 + +using Oceananigans +using Oceananigans.Units + +using Oceananigans.Diagnostics: accurate_cell_advection_timescale + +Logging.global_logger(OceananigansLogger()) + +##### +##### Progress monitor +##### + +mutable struct Progress + interval_start_time :: Float64 +end + +function (p::Progress)(sim) + wall_time = (time_ns() - p.interval_start_time) * 1e-9 + + @info @sprintf("Time: %s, iteration: %d, extrema(h): (min=%.2e, max=%.2e), CFL: %.2e, wall time: %s", + prettytime(sim.model.clock.time), + sim.model.clock.iteration, + minimum(sim.model.tracers.h), + maximum(sim.model.tracers.h), + sim.parameters.cfl(sim.model), + prettytime(wall_time)) + + p.interval_start_time = time_ns() + + return nothing +end + +##### +##### Script starts here +##### + +ENV["DATADEPS_ALWAYS_ACCEPT"] = "true" + +dd = DataDep("cubed_sphere_32_grid", + "Conformal cubed sphere grid with 32×32 grid points on each face", + "https://github.com/CliMA/OceananigansArtifacts.jl/raw/main/cubed_sphere_grids/cubed_sphere_32_grid.jld2", + "b1dafe4f9142c59a2166458a2def743cd45b20a4ed3a1ae84ad3a530e1eff538" # sha256sum +) + +DataDeps.register(dd) + +cs32_filepath = datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2" + +# Central (λ, φ) for each face of the cubed sphere +central_longitude = (0, 90, 0, 180, -90, 0) +central_latitude = (0, 0, 90, 0, 0, -90) + +""" + tracer_advection_over_the_poles(; face_number, α) + +Run a tracer advection experiment that initializes a cosine bell on face `face_number` +and advects it around the sphere over 12 days. `α` is the angle between the axis of +solid body rotation and the polar axis (degrees). +""" +function tracer_advection_over_the_poles(; face_number, α) + + grid = ConformalCubedSphereGrid(cs32_filepath, Nz=1, z=(-1, 0)) + + ## Prescribed velocities and initial condition according to Williamson et al. (1992) §3.1 + + period = 12days # Time to make a full rotation (s) + R = grid.faces[1].radius # radius of the sphere (m) + u₀ = 2π*R / period # advecting velocity (m/s) + + # U(λ, φ, z) = u₀ * (cosd(φ) * cosd(α) + sind(φ) * cosd(λ) * sind(α)) + # V(λ, φ, z) = - u₀ * sind(λ) * sind(α) + + Ψ(λ, φ, z) = - R * u₀ * (sind(φ) * cosd(α) - cosd(λ) * cosd(φ) * sind(α)) + + Ψᶠᶠᶜ = Field(Face, Face, Center, CPU(), grid, nothing, nothing) + Uᶠᶜᶜ = Field(Face, Center, Center, CPU(), grid, nothing, nothing) + Vᶜᶠᶜ = Field(Center, Face, Center, CPU(), grid, nothing, nothing) + Wᶜᶜᶠ = Field(Center, Center, Face, CPU(), grid, nothing, nothing) # So we can use CFL + + for (f, grid_face) in enumerate(grid.faces) + for i in 1:grid_face.Nx+1, j in 1:grid_face.Ny+1 + Ψᶠᶠᶜ.faces[f][i, j, 1] = Ψ(grid_face.λᶠᶠᵃ[i, j], grid_face.φᶠᶠᵃ[i, j], 0) + end + end + + for (f, grid_face) in enumerate(grid.faces) + for i in 1:grid_face.Nx+1, j in 1:grid_face.Ny + Uᶠᶜᶜ.faces[f][i, j, 1] = (Ψᶠᶠᶜ.faces[f][i, j, 1] - Ψᶠᶠᶜ.faces[f][i, j+1, 1]) / grid.faces[f].Δyᶠᶜᵃ[i, j] + end + for i in 1:grid_face.Nx, j in 1:grid_face.Ny+1 + Vᶜᶠᶜ.faces[f][i, j, 1] = (Ψᶠᶠᶜ.faces[f][i+1, j, 1] - Ψᶠᶠᶜ.faces[f][i, j, 1]) / grid.faces[f].Δxᶜᶠᵃ[i, j] + end + end + + ## Model setup + + model = HydrostaticFreeSurfaceModel( + architecture = CPU(), + grid = grid, + tracers = :h, + velocities = PrescribedVelocityFields(u=Uᶠᶜᶜ, v=Vᶜᶠᶜ, w=Wᶜᶜᶠ), + free_surface = ExplicitFreeSurface(gravitational_acceleration=0.1), + coriolis = nothing, + closure = nothing, + buoyancy = nothing + ) + + ## Cosine bell initial condition according to Williamson et al. (1992) §3.1 + + h₀ = 1000 + λ₀ = central_longitude[face_number] + φ₀ = central_longitude[face_number] + + # Great circle distance between (λ, φ) and the center of the cosine bell (λ₀, φ₀) + # using the haversine formula + r(λ, φ) = 2R * asin(√(sind((φ - φ₀) / 2)^2 + cosd(φ) * cosd(φ₀) * sind((λ - λ₀) / 2)^2)) + + cosine_bell(λ, φ, z) = r(λ, φ) < R ? h₀/2 * (1 + cos(π * r(λ, φ) / R)) : 0 + + Oceananigans.set!(model, h=cosine_bell) + + ## Simulation setup + + Δt = 10minutes + + cfl = CFL(Δt, accurate_cell_advection_timescale) + + simulation = Simulation(model, + Δt = Δt, + stop_time = 1period, + iteration_interval = 1, + progress = Progress(time_ns()), + parameters = (; cfl) + ) + + outputs = (u=Uᶠᶜᶜ, v=Vᶜᶠᶜ, h=model.tracers.h) + + simulation.output_writers[:fields] = + JLD2OutputWriter(model, outputs, + schedule = TimeInterval(1hour), + prefix = "tracer_advection_over_the_poles_face$(face_number)_alpha$α", + force = true) + + run!(simulation) + + return simulation +end + +##### +##### Run all the experiments! +##### + +include("animate_on_map_projection.jl") + +function run_cubed_sphere_tracer_advection_validation() + + αs = (0, 45, 90) + + for f in 1:6, α in αs + tracer_advection_over_the_poles(face_number=f, α=α) + end + + projections = [ + ccrs.NearsidePerspective(central_longitude=0, central_latitude=30), + ccrs.NearsidePerspective(central_longitude=180, central_latitude=-30) + ] + + for f in 1:6, α in αs + animate_tracer_advection(face_number=f, α=α, projections=projections) + end +end