diff --git a/Manifest.toml b/Manifest.toml index 214c0faf27..79ab4b1fea 100644 --- a/Manifest.toml +++ b/Manifest.toml @@ -318,9 +318,9 @@ version = "4.15.0" [[deps.LLVMExtra_jll]] deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "Pkg", "TOML"] -git-tree-sha1 = "771bfe376249626d3ca12bcd58ba243d3f961576" +git-tree-sha1 = "7718cf44439c676bc0ec66a87099f41015a522d6" uuid = "dad2f222-ce93-54a1-a47d-0025e8a3acab" -version = "0.0.16+0" +version = "0.0.16+2" [[deps.LazyArtifacts]] deps = ["Artifacts", "Pkg"] diff --git a/docs/src/appendix/library.md b/docs/src/appendix/library.md index ccfa0c1fbc..1685e2daba 100644 --- a/docs/src/appendix/library.md +++ b/docs/src/appendix/library.md @@ -37,7 +37,7 @@ Modules = [Oceananigans.BoundaryConditions] Private = false ``` -## BuoyancyModels +## Buoyancy models ```@autodocs Modules = [Oceananigans.BuoyancyModels] @@ -127,7 +127,6 @@ Private = false ```@autodocs Modules = [Oceananigans.Models.HydrostaticFreeSurfaceModels] Private = false -] ``` ### Shallow-water models @@ -149,7 +148,6 @@ Private = false ```@autodocs Modules = [Oceananigans.Operators] Private = false -] ``` ## Output readers @@ -157,7 +155,6 @@ Private = false ```@autodocs Modules = [Oceananigans.OutputReaders] Private = false -] ``` ## Output writers @@ -165,7 +162,6 @@ Private = false ```@autodocs Modules = [Oceananigans.OutputWriters] Private = false -] ``` ## Simulations diff --git a/examples/baroclinic_adjustment.jl b/examples/baroclinic_adjustment.jl index 64a2430c1b..a0c6c4309e 100644 --- a/examples/baroclinic_adjustment.jl +++ b/examples/baroclinic_adjustment.jl @@ -199,6 +199,8 @@ b_timeserieses = (east = FieldTimeSeries(slice_filenames.east, "b"), avg_b_timeseries = FieldTimeSeries(filename * "_zonal_average.jld2", "b") +times = avg_b_timeseries.times + nothing #hide # We build the coordinates. We rescale horizontal coordinates so that they correspond to kilometers. @@ -270,14 +272,15 @@ contour!(ax, y, z, avg_b; transformation = (:yz, zonal_slice_displacement * x[en Colorbar(fig[2, 2], sf, label = "m s⁻²", height = 200, tellheight=false) -# Finally, we add a figure title with the time of the snapshot and then record a movie. - -times = avg_b_timeseries.times - title = @lift "Buoyancy at t = " * string(round(times[$n] / day, digits=1)) * " days" fig[1, 1:2] = Label(fig, title; fontsize = 24, tellwidth = false, padding = (0, 0, -120, 0)) +current_figure() # hide +fig + +# Finally, we add a figure title with the time of the snapshot and then record a movie. + frames = 1:length(times) record(fig, filename * ".mp4", frames, framerate=8) do i diff --git a/examples/convecting_plankton.jl b/examples/convecting_plankton.jl index 13d95d32db..0c1f547596 100644 --- a/examples/convecting_plankton.jl +++ b/examples/convecting_plankton.jl @@ -37,7 +37,7 @@ # ```julia # using Pkg -# pkg"add Oceananigans, Measures, CairoMakie" +# pkg"add Oceananigans, CairoMakie" # ``` # ## The grid @@ -65,7 +65,7 @@ buoyancy_flux_bc = FluxBoundaryCondition(buoyancy_flux, parameters = buoyancy_fl # constant during the first phase of the simulation. We produce a plot of this time-dependent # buoyancy flux for the visually-oriented, -using CairoMakie, Measures +using CairoMakie set_theme!(Theme(fontsize = 24, linewidth=2)) times = range(0, 12hours, length=100) @@ -76,6 +76,9 @@ ax = Axis(fig[1, 1]; xlabel = "Time (hours)", ylabel = "Surface buoyancy flux (m flux_time_series = [buoyancy_flux(0, 0, t, buoyancy_flux_parameters) for t in times] lines!(ax, times ./ hour, flux_time_series) +current_figure() # hide +fig + # The buoyancy flux effectively shuts off after 6 hours of simulation time. # # !!! info "The flux convention in Oceananigans.jl" @@ -261,6 +264,9 @@ b_flux_point = @lift Point2(times[$n] / hour, buoyancy_flux_time_series[$n]) scatter!(ax_b, b_flux_point; marker = :circle, markersize = 16, color = :black) lines!(ax_avg_P, avg_Pₙ, zp) +current_figure() # hide +fig + # And, finally, we record a movie. frames = 1:length(times) diff --git a/examples/horizontal_convection.jl b/examples/horizontal_convection.jl index 2d00a680c5..9d6acf4130 100644 --- a/examples/horizontal_convection.jl +++ b/examples/horizontal_convection.jl @@ -13,7 +13,7 @@ # ```julia # using Pkg -# pkg"add Oceananigans, JLD2, CairoMakie" +# pkg"add Oceananigans, CairoMakie" # ``` # ## Horizontal convection @@ -163,7 +163,7 @@ run!(simulation) # ## Load saved output, process, visualize # -# We animate the results by opening the JLD2 file, extracting data for the iterations we ended +# We animate the results by loading the saved output, extracting data for the iterations we ended # up saving at, and ploting the saved fields. From the saved buoyancy field we compute the # buoyancy dissipation, ``\chi = \kappa |\boldsymbol{\nabla} b|^2``, and plot that also. # diff --git a/examples/internal_wave.jl b/examples/internal_wave.jl index dd70b4fae6..6fe40398e6 100644 --- a/examples/internal_wave.jl +++ b/examples/internal_wave.jl @@ -89,11 +89,11 @@ nothing # hide # observe wave propagation. ## Some Gaussian parameters -A = 1e-9 -δ = grid.Lx / 15 +gaussian_amplitude = 1e-9 +gaussian_width = grid.Lx / 15 ## A Gaussian envelope centered at ``(x, z) = (0, 0)``. -a(x, z) = A * exp( -( x^2 + z^2 ) / 2δ^2 ) +a(x, z) = gaussian_amplitude * exp( -( x^2 + z^2 ) / 2gaussian_width^2 ) nothing # hide # An inertia-gravity wave is a linear solution to the Boussinesq equations. diff --git a/examples/kelvin_helmholtz_instability.jl b/examples/kelvin_helmholtz_instability.jl index f067e04704..1063894da3 100644 --- a/examples/kelvin_helmholtz_instability.jl +++ b/examples/kelvin_helmholtz_instability.jl @@ -59,6 +59,9 @@ lines!(ax, [stratification(0, 0, z, 0, (Ri=Ri, h=h)) for z in zC], zC; linewidth ax = Axis(fig[1, 3], xlabel = "Ri(z)") lines!(ax, [Ri * sech(z / h)^2 / sech(z)^2 for z in zF], zF; linewidth = 3, color = :black) # Ri(z)= ∂_z B / (∂_z U)²; derivatives computed by hand +current_figure() # hide +fig + # In unstable flows it is often useful to determine the dominant spatial structure of the # instability and the growth rate at which the instability grows. # If the simulation idealizes a physical flow, this can be used to make diff --git a/examples/langmuir_turbulence.jl b/examples/langmuir_turbulence.jl index 48928d62ba..45c10afac1 100644 --- a/examples/langmuir_turbulence.jl +++ b/examples/langmuir_turbulence.jl @@ -355,6 +355,9 @@ ax_uxz = heatmap!(ax_uxz, xu, zu, uxzₙ; Colorbar(fig[3, 3], ax_uxz; label = "m s⁻¹") +current_figure() # hide +fig + # And, finally, we record a movie. frames = 1:length(times) diff --git a/examples/ocean_wind_mixing_and_convection.jl b/examples/ocean_wind_mixing_and_convection.jl index 2fbcfc6259..687bd449be 100644 --- a/examples/ocean_wind_mixing_and_convection.jl +++ b/examples/ocean_wind_mixing_and_convection.jl @@ -35,7 +35,10 @@ using Oceananigans.Units: minute, minutes, hour # maintains relatively constant vertical spacing in the mixed layer, which # is desirable from a numerical standpoint: +Nx = Ny = 32 # number of points in each of horizontal directions Nz = 24 # number of points in the vertical direction + +Lx = Ly = 64 # (m) domain horizontal extents Lz = 32 # (m) domain depth refinement = 1.2 # controls spacing near surface (higher means finer spaced) @@ -53,16 +56,16 @@ h(k) = (k - 1) / Nz ## Generating function z_faces(k) = Lz * (ζ₀(k) * Σ(k) - 1) -grid = RectilinearGrid(CPU(); - size = (32, 32, Nz), - x = (0, 64), - y = (0, 64), +grid = RectilinearGrid(size = (Nx, Nx, Nz), + x = (0, Lx), + y = (0, Ly), z = z_faces) # We plot vertical spacing versus depth to inspect the prescribed grid stretching: fig = Figure(resolution=(1200, 800)) ax = Axis(fig[1, 1], ylabel = "Depth (m)", xlabel = "Vertical spacing (m)") + lines!(ax, zspacings(grid, Center()), znodes(grid, Center())) scatter!(ax, zspacings(grid, Center()), znodes(grid, Center())) @@ -288,6 +291,9 @@ Colorbar(fig[3, 4], hm_νₑ; label = "m s⁻²") fig[1, 1:4] = Label(fig, title, fontsize=24, tellwidth=false) +current_figure() # hide +fig + # And now record a movie. frames = intro:length(times) diff --git a/examples/one_dimensional_diffusion.jl b/examples/one_dimensional_diffusion.jl index 294da46dc5..9d649682b2 100644 --- a/examples/one_dimensional_diffusion.jl +++ b/examples/one_dimensional_diffusion.jl @@ -88,8 +88,8 @@ current_figure() # hide # Next we set-up a `Simulation` that time-steps the model forward and manages output. ## Time-scale for diffusion across a grid cell -using Oceananigans.Grids: min_Δz -diffusion_time_scale = min_Δz(model.grid)^2 / model.closure.κ.T +min_Δz = minimum_zspacing(model.grid) +diffusion_time_scale = min_Δz^2 / model.closure.κ.T simulation = Simulation(model, Δt = 0.1 * diffusion_time_scale, stop_iteration = 1000) @@ -141,7 +141,9 @@ lines!(T, z) label = @lift "t = " * string(round(times[$n], digits=3)) Label(fig[1, 1], label, tellwidth=false) + current_figure() # hide +fig # Finally, we record a movie. diff --git a/examples/shallow_water_Bickley_jet.jl b/examples/shallow_water_Bickley_jet.jl index 84cc74654e..9a3bd28966 100644 --- a/examples/shallow_water_Bickley_jet.jl +++ b/examples/shallow_water_Bickley_jet.jl @@ -195,6 +195,9 @@ Colorbar(fig[2, 4], hm_ω′) title = @lift @sprintf("t = %.1f", times[$n]) fig[1, 1:4] = Label(fig, title, fontsize=24, tellwidth=false) +current_figure() # hide +fig + # Finally, we record a movie. frames = 1:length(times) diff --git a/examples/tilted_bottom_boundary_layer.jl b/examples/tilted_bottom_boundary_layer.jl index 487047c976..f9e6b2c985 100644 --- a/examples/tilted_bottom_boundary_layer.jl +++ b/examples/tilted_bottom_boundary_layer.jl @@ -142,9 +142,8 @@ model = NonhydrostaticModel(; grid, buoyancy, coriolis, closure, # conservatively, based on the smallest grid size of our domain and set-up a using Oceananigans.Units -using Oceananigans.Grids: min_Δz -simulation = Simulation(model, Δt = 0.5 * min_Δz(grid) / V∞, stop_time = 1days) +simulation = Simulation(model, Δt = 0.5 * minimum_zspacing(grid) / V∞, stop_time = 2days) # We use `TimeStepWizard` to adapt our time-step and print a progress message, @@ -227,6 +226,9 @@ times = collect(ds["time"]) title = @lift "t = " * string(prettytime(times[$n])) fig[1, :] = Label(fig, title, fontsize=20, tellwidth=false) +current_figure() # hide +fig + # Finally, we record a movie. frames = 1:length(times) diff --git a/examples/two_dimensional_turbulence.jl b/examples/two_dimensional_turbulence.jl index fa20199d87..4f1a0b5471 100644 --- a/examples/two_dimensional_turbulence.jl +++ b/examples/two_dimensional_turbulence.jl @@ -147,6 +147,9 @@ heatmap!(ax_s, xs, ys, s; colormap = :speed, colorrange = (0, 0.2)) title = @lift "t = " * string(round(times[$n], digits=2)) Label(fig[1, 1:2], title, fontsize=24, tellwidth=false) +current_figure() # hide +fig + # Finally, we record a movie. frames = 1:length(times) diff --git a/src/AbstractOperations/at.jl b/src/AbstractOperations/at.jl index 3dd3f73065..07537472b2 100644 --- a/src/AbstractOperations/at.jl +++ b/src/AbstractOperations/at.jl @@ -77,8 +77,11 @@ end interpolate_index(::Colon, ::Colon, args...) = Colon() interpolate_index(::Colon, b::UnitRange, args...) = b +instantiate(T::Type) = T() +instantiate(t) = t + function interpolate_index(a::UnitRange, ::Colon, loc, new_loc) - a = corrected_index(a, loc, new_loc) + a = corrected_index(a, instantiate(loc), instantiate(new_loc)) # Abstract operations that require an interpolation of a sliced fields are not supported! first(a) > last(a) && throw(ArgumentError("Cannot interpolate a sliced field from $loc to $(new_loc)!")) @@ -86,7 +89,7 @@ function interpolate_index(a::UnitRange, ::Colon, loc, new_loc) end function interpolate_index(a::UnitRange, b::UnitRange, loc, new_loc) - a = corrected_index(a, loc, new_loc) + a = corrected_index(a, instantiate(loc), instantiate(new_loc)) # Abstract operations that require an interpolation of a sliced fields are not supported! first(a) > last(a) && throw(ArgumentError("Cannot interpolate a sliced field from $loc to $(new_loc)!")) @@ -100,7 +103,7 @@ end # Windowed fields interpolated from `Center`s to `Face`s lose the first index. # Viceverse, windowed fields interpolated from `Face`s to `Center`s lose the last index -corrected_index(a, ::Type{Face}, ::Type{Face}) = UnitRange(first(a), last(a)) -corrected_index(a, ::Type{Center}, ::Type{Center}) = UnitRange(first(a), last(a)) -corrected_index(a, ::Type{Face}, ::Type{Center}) = UnitRange(first(a), last(a)-1) -corrected_index(a, ::Type{Center}, ::Type{Face}) = UnitRange(first(a)+1, last(a)) +corrected_index(a, ::Face, ::Face) = UnitRange(first(a), last(a) ) +corrected_index(a, ::Center, ::Center) = UnitRange(first(a), last(a) ) +corrected_index(a, ::Face, ::Center) = UnitRange(first(a), last(a)-1) +corrected_index(a, ::Center, ::Face) = UnitRange(first(a)+1, last(a) ) diff --git a/src/Advection/Advection.jl b/src/Advection/Advection.jl index 7f98c2cf63..40217081a0 100644 --- a/src/Advection/Advection.jl +++ b/src/Advection/Advection.jl @@ -64,5 +64,6 @@ include("topologically_conditional_interpolation.jl") include("momentum_advection_operators.jl") include("tracer_advection_operators.jl") include("positivity_preserving_tracer_advection_operators.jl") +include("cell_advection_timescale.jl") end # module diff --git a/src/Advection/cell_advection_timescale.jl b/src/Advection/cell_advection_timescale.jl new file mode 100644 index 0000000000..77374a6e35 --- /dev/null +++ b/src/Advection/cell_advection_timescale.jl @@ -0,0 +1,19 @@ +using Oceananigans.AbstractOperations: KernelFunctionOperation +using Oceananigans.Grids: topology, min_Δx, min_Δy, min_Δz + +function cell_advection_timescale(grid, velocities) + u, v, w = velocities + τ = KernelFunctionOperation{Center, Center, Center}(cell_advection_timescaleᶜᶜᶜ, grid, u, v, w) + return minimum(τ) +end + +@inline function cell_advection_timescaleᶜᶜᶜ(i, j, k, grid, u, v, w) + Δx = Δxᶠᶜᶜ(i, j, k, grid) + Δy = Δyᶜᶠᶜ(i, j, k, grid) + Δz = Δzᶜᶜᶠ(i, j, k, grid) + + return @inbounds min(Δx / abs(u[i, j, k]), + Δy / abs(v[i, j, k]), + Δz / abs(w[i, j, k])) +end + diff --git a/src/CubedSpheres/conformal_cubed_sphere_grid.jl b/src/CubedSpheres/conformal_cubed_sphere_grid.jl index ca0d14286c..a1e86d7e94 100644 --- a/src/CubedSpheres/conformal_cubed_sphere_grid.jl +++ b/src/CubedSpheres/conformal_cubed_sphere_grid.jl @@ -229,10 +229,11 @@ const OSSG = OrthogonalSphericalShellGrid ##### Grid utils ##### -Base.size(grid::ConformalCubedSphereGrid) = (size(grid.faces[1])..., length(grid.faces)) -Base.size(loc, grid::ConformalCubedSphereGrid) = size(loc, grid.faces[1]) -Base.size(grid::ConformalCubedSphereGrid, i) = size(grid)[i] -halo_size(ccsg::ConformalCubedSphereGrid) = halo_size(first(ccsg.faces)) # hack +Base.size(grid::ConformalCubedSphereGrid) = (size(grid.faces[1])..., length(grid.faces)) +Base.size(grid::ConformalCubedSphereGrid, loc::Tuple) = size(grid.faces[1], loc) +Base.size(grid::ConformalCubedSphereGrid, i::Int) = size(grid)[i] + +halo_size(ccsg::ConformalCubedSphereGrid) = halo_size(first(ccsg.faces)) # hack Base.eltype(grid::ConformalCubedSphereGrid{FT}) where FT = FT diff --git a/src/CubedSpheres/cubed_sphere_utils.jl b/src/CubedSpheres/cubed_sphere_utils.jl index 69b4637037..8466c6588f 100644 --- a/src/CubedSpheres/cubed_sphere_utils.jl +++ b/src/CubedSpheres/cubed_sphere_utils.jl @@ -11,95 +11,98 @@ using Oceananigans.Grids: ##### Viewing halos ##### +instantiate(T::Type) = T() +instantiate(t) = t + 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)) + include_corners ? view(f.data, left_halo_indices(instantiate(LX), instantiate(topology(f, 1)), f.grid.Nx, f.grid.Hx), :, :) : + view(f.data, left_halo_indices(instantiate(LX), instantiate(topology(f, 1)), f.grid.Nx, f.grid.Hx), + interior_indices(instantiate(LY), instantiate(topology(f, 2)), f.grid.Ny), + interior_indices(instantiate(LZ), instantiate(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)) + include_corners ? view(f.data, right_halo_indices(instantiate(LX), instantiate(topology(f, 1)), f.grid.Nx, f.grid.Hx), :, :) : + view(f.data, right_halo_indices(instantiate(LX), instantiate(topology(f, 1)), f.grid.Nx, f.grid.Hx), + interior_indices(instantiate(LY), instantiate(topology(f, 2)), f.grid.Ny), + interior_indices(instantiate(LZ), instantiate(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)) + include_corners ? view(f.data, :, left_halo_indices(instantiate(LY), instantiate(topology(f, 2)), f.grid.Ny, f.grid.Hy), :) : + view(f.data, interior_indices(instantiate(LX), instantiate(topology(f, 1)), f.grid.Nx), + left_halo_indices(instantiate(LY), instantiate(topology(f, 2)), f.grid.Ny, f.grid.Hy), + interior_indices(instantiate(LZ), instantiate(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)) + include_corners ? view(f.data, :, right_halo_indices(instantiate(LY), instantiate(topology(f, 2)), f.grid.Ny, f.grid.Hy), :) : + view(f.data, interior_indices(instantiate(LX), instantiate(topology(f, 1)), f.grid.Nx), + right_halo_indices(instantiate(LY), instantiate(topology(f, 2)), f.grid.Ny, f.grid.Hy), + interior_indices(instantiate(LZ), instantiate(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)) + include_corners ? view(f.data, :, :, left_halo_indices(instantiate(LZ), instantiate(topology(f, 3)), f.grid.Nz, f.grid.Hz)) : + view(f.data, interior_indices(instantiate(LX), instantiate(topology(f, 1)), f.grid.Nx), + interior_indices(instantiate(LY), instantiate(topology(f, 2)), f.grid.Ny), + left_halo_indices(instantiate(LZ), instantiate(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)) + include_corners ? view(f.data, :, :, right_halo_indices(instantiate(LZ), instantiate(topology(f, 3)), f.grid.Nz, f.grid.Hz)) : + view(f.data, interior_indices(instantiate(LX), instantiate(topology(f, 1)), f.grid.Nx), + interior_indices(instantiate(LY), instantiate(topology(f, 2)), f.grid.Ny), + right_halo_indices(instantiate(LZ), instantiate(topology(f, 3)), f.grid.Nz, f.grid.Hz)) underlying_west_halo(f, grid, location, topo=topology(grid, 1)) = - view(f.parent, underlying_left_halo_indices(location, topo, grid.Nx, grid.Hx), :, :) + view(f.parent, underlying_left_halo_indices(instantiate(location), instantiate(topo), grid.Nx, grid.Hx), :, :) underlying_east_halo(f, grid, location, topo=topology(grid, 1)) = - view(f.parent, underlying_right_halo_indices(location, topo, grid.Nx, grid.Hx), :, :) + view(f.parent, underlying_right_halo_indices(instantiate(location), instantiate(topo), grid.Nx, grid.Hx), :, :) underlying_south_halo(f, grid, location, topo=topology(grid, 2)) = - view(f.parent, :, underlying_left_halo_indices(location, topo, grid.Ny, grid.Hy), :) + view(f.parent, :, underlying_left_halo_indices(instantiate(location), instantiate(topo), grid.Ny, grid.Hy), :) underlying_north_halo(f, grid, location, topo=topology(grid, 2)) = - view(f.parent, :, underlying_right_halo_indices(location, topo, grid.Ny, grid.Hy), :) + view(f.parent, :, underlying_right_halo_indices(instantiate(location), instantiate(topo), grid.Ny, grid.Hy), :) underlying_bottom_halo(f, grid, location, topo=topology(grid, 3)) = - view(f.parent, :, :, underlying_left_halo_indices(location, topo, grid.Nz, grid.Hz)) + view(f.parent, :, :, underlying_left_halo_indices(instantiate(location), instantiate(topo), grid.Nz, grid.Hz)) underlying_top_halo(f, grid, location, topo=topology(grid, 3)) = - view(f.parent, :, :, underlying_right_halo_indices(location, topo, grid.Nz, grid.Hz)) + view(f.parent, :, :, underlying_right_halo_indices(instantiate(location), instantiate(topo), 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 +left_boundary_indices(::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 +right_boundary_indices(::Face, ::Bounded, N, H) = N-H:N+1 +right_boundary_indices(::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_left_boundary_indices(::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_right_boundary_indices(::Face, ::Bounded, N, H) = N+2:N+H+1 +underlying_right_boundary_indices(::Nothing, topo, N, H) = 1:0 # empty underlying_west_boundary(f, grid, location, topo=topology(grid, 1)) = - view(f.parent, underlying_left_boundary_indices(location, topo, grid.Nx, grid.Hx), :, :) + view(f.parent, underlying_left_boundary_indices(instantiate(location), instantiate(topo), grid.Nx, grid.Hx), :, :) underlying_east_boundary(f, grid, location, topo=topo = topology(grid, 1)) = - view(f.parent, underlying_right_boundary_indices(location, topo, grid.Nx, grid.Hx), :, :) + view(f.parent, underlying_right_boundary_indices(instantiate(location), instantiate(topo), grid.Nx, grid.Hx), :, :) underlying_south_boundary(f, grid, location, topo=topology(grid, 2)) = - view(f.parent, :, underlying_left_boundary_indices(location, topo, grid.Ny, grid.Hy), :) + view(f.parent, :, underlying_left_boundary_indices(instantiate(location), instantiate(topo), grid.Ny, grid.Hy), :) underlying_north_boundary(f, grid, location, topo=topology(grid, 2)) = - view(f.parent, :, underlying_right_boundary_indices(location, topo, grid.Ny, grid.Hy), :) + view(f.parent, :, underlying_right_boundary_indices(instantiate(location), instantiate(topo), grid.Ny, grid.Hy), :) underlying_bottom_boundary(f, grid, location, topo=topology(grid, 3)) = - view(f.parent, :, :, underlying_left_boundary_indices(location, topo, grid.Nz, grid.Hz)) + view(f.parent, :, :, underlying_left_boundary_indices(instantiate(location), instantiate(topo), grid.Nz, grid.Hz)) underlying_top_boundary(f, grid, location, topo=topology(grid, 3)) = - view(f.parent, :, :, underlying_right_boundary_indices(location, topo, grid.Nz, grid.Hz)) + view(f.parent, :, :, underlying_right_boundary_indices(instantiate(location), instantiate(topo), grid.Nz, grid.Hz)) ##### ##### Convenience functions diff --git a/src/Diagnostics/cfl.jl b/src/Diagnostics/cfl.jl index 0cac003061..045491ce66 100644 --- a/src/Diagnostics/cfl.jl +++ b/src/Diagnostics/cfl.jl @@ -1,4 +1,4 @@ -using Oceananigans.Utils: cell_advection_timescale +using Oceananigans.Advection: cell_advection_timescale using Oceananigans.TurbulenceClosures: cell_diffusion_timescale """ diff --git a/src/Distributed/distributed_fields.jl b/src/Distributed/distributed_fields.jl index cc29b5f955..1b365ac95b 100644 --- a/src/Distributed/distributed_fields.jl +++ b/src/Distributed/distributed_fields.jl @@ -12,6 +12,7 @@ function Field((LX, LY, LZ)::Tuple, grid::DistributedGrid, data, old_bcs, indice indices = validate_indices(indices, (LX, LY, LZ), grid) new_bcs = inject_halo_communication_boundary_conditions(old_bcs, arch.local_rank, arch.connectivity) buffers = FieldBoundaryBuffers(grid, data, new_bcs) + return Field{LX, LY, LZ}(grid, data, new_bcs, indices, op, status, buffers) end @@ -19,4 +20,4 @@ const DistributedField = Field{<:Any, <:Any, <:Any, <:Any, <:DistributedGri const DistributedFieldTuple = NamedTuple{S, <:NTuple{N, DistributedField}} where {S, N} # TODO: make sure the definition of architecture is consistent -architecture(f::DistributedField) = child_architecture(architecture(f.grid)) \ No newline at end of file +architecture(f::DistributedField) = child_architecture(architecture(f.grid)) diff --git a/src/Distributed/distributed_grids.jl b/src/Distributed/distributed_grids.jl index 3c72f99943..9480caa9ba 100644 --- a/src/Distributed/distributed_grids.jl +++ b/src/Distributed/distributed_grids.jl @@ -62,9 +62,9 @@ function RectilinearGrid(arch::DistributedArch, yl = partition(y, ny, Ry, rj) zl = partition(z, nz, Rz, rk) - Lx, xᶠᵃᵃ, xᶜᵃᵃ, Δxᶠᵃᵃ, Δxᶜᵃᵃ = generate_coordinate(FT, topology[1], nx, Hx, xl, child_architecture(arch)) - Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, topology[2], ny, Hy, yl, child_architecture(arch)) - Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, topology[3], nz, Hz, zl, child_architecture(arch)) + Lx, xᶠᵃᵃ, xᶜᵃᵃ, Δxᶠᵃᵃ, Δxᶜᵃᵃ = generate_coordinate(FT, topology[1](), nx, Hx, xl, child_architecture(arch)) + Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, topology[2](), ny, Hy, yl, child_architecture(arch)) + Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, topology[3](), nz, Hz, zl, child_architecture(arch)) architecture = DistributedArch(child_architecture(arch), topology = topology, @@ -122,18 +122,18 @@ function LatitudeLongitudeGrid(arch::DistributedArch, # Calculate all direction (which might be stretched) # A direction is regular if the domain passed is a Tuple{<:Real, <:Real}, # it is stretched if being passed is a function or vector (as for the VerticallyStretchedRectilinearGrid) - Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, TX, nλ, Hλ, λl, arch.child_architecture) - Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ, nz, Hz, zl, arch.child_architecture) + Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, TX(), nλ, Hλ, λl, arch.child_architecture) + Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ(), nz, Hz, zl, arch.child_architecture) # The Latitudinal direction is _special_ : # Preconmpute metrics assumes that `length(φᵃᶠᵃ) = length(φᵃᶜᵃ) + 1`, which is always the case in a # serial grid because `LatitudeLongitudeGrid` should be always `Bounded`, but it is not true for a # partitioned `DistributedGrid` with Ry > 1 (one rank will hold a `RightConnected` topology) # But we need an extra point to precompute the Y direction in case of only one halo so we disregard the topology # when constructing the metrics! - Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, Bounded, nφ, Hφ, φl, arch.child_architecture) + Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, Bounded(), nφ, Hφ, φl, arch.child_architecture) architecture = DistributedArch(child_architecture(arch); - topology = topology, + topology = topology, ranks = arch.ranks, communicator = arch.communicator, use_buffers = using_buffered_communication(arch)) @@ -185,9 +185,9 @@ function reconstruct_global_grid(grid::DistributedRectilinearGrid) FT = eltype(grid) - Lx, xᶠᵃᵃ, xᶜᵃᵃ, Δxᶠᵃᵃ, Δxᶜᵃᵃ = generate_coordinate(FT, TX, Nx, Hx, xG, child_arch) - Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, TY, Ny, Hy, yG, child_arch) - Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ, Nz, Hz, zG, child_arch) + Lx, xᶠᵃᵃ, xᶜᵃᵃ, Δxᶠᵃᵃ, Δxᶜᵃᵃ = generate_coordinate(FT, TX(), Nx, Hx, xG, child_arch) + Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, TY(), Ny, Hy, yG, child_arch) + Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ(), Nz, Hz, zG, child_arch) return RectilinearGrid{TX, TY, TZ}(child_arch, Nx, Ny, Nz, @@ -231,9 +231,9 @@ function reconstruct_global_grid(grid::DistributedLatitudeLongitudeGrid) # Calculate all direction (which might be stretched) # A direction is regular if the domain passed is a Tuple{<:Real, <:Real}, # it is stretched if being passed is a function or vector (as for the VerticallyStretchedRectilinearGrid) - Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, TX, Nλ, Hλ, λG, child_arch) - Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, TY, Nφ, Hφ, φG, child_arch) - Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ, Nz, Hz, zG, child_arch) + Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, TX(), Nλ, Hλ, λG, child_arch) + Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, TY(), Nφ, Hφ, φG, child_arch) + Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ(), Nz, Hz, zG, child_arch) precompute_metrics = metrics_precomputed(grid) diff --git a/src/Distributed/distributed_utils.jl b/src/Distributed/distributed_utils.jl index 7aed8d07e9..b91a951522 100644 --- a/src/Distributed/distributed_utils.jl +++ b/src/Distributed/distributed_utils.jl @@ -11,83 +11,86 @@ using Oceananigans.Grids: ##### 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)) + include_corners ? view(f.data, left_halo_indices(instantiate(LX), instantiate(topology(f, 1)), f.grid.Nx, f.grid.Hx), :, :) : + view(f.data, left_halo_indices(instantiate(LX), instantiate(topology(f, 1)), f.grid.Nx, f.grid.Hx), + interior_indices(instantiate(LY), instantiate(topology(f, 2)), f.grid.Ny), + interior_indices(instantiate(LZ), instantiate(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)) + include_corners ? view(f.data, right_halo_indices(instantiate(LX), instantiate(topology(f, 1)), f.grid.Nx, f.grid.Hx), :, :) : + view(f.data, right_halo_indices(instantiate(LX), instantiate(topology(f, 1)), f.grid.Nx, f.grid.Hx), + interior_indices(instantiate(LY), instantiate(topology(f, 2)), f.grid.Ny), + interior_indices(instantiate(LZ), instantiate(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)) + include_corners ? view(f.data, :, left_halo_indices(instantiate(LY), instantiate(topology(f, 2)), f.grid.Ny, f.grid.Hy), :) : + view(f.data, interior_indices(instantiate(LX), instantiate(topology(f, 1)), f.grid.Nx), + left_halo_indices(instantiate(LY), instantiate(topology(f, 2)), f.grid.Ny, f.grid.Hy), + interior_indices(instantiate(LZ), instantiate(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)) + include_corners ? view(f.data, :, right_halo_indices(instantiate(LY), instantiate(topology(f, 2)), f.grid.Ny, f.grid.Hy), :) : + view(f.data, interior_indices(instantiate(LX), instantiate(topology(f, 1)), f.grid.Nx), + right_halo_indices(instantiate(LY), instantiate(topology(f, 2)), f.grid.Ny, f.grid.Hy), + interior_indices(instantiate(LZ), instantiate(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)) + include_corners ? view(f.data, :, :, left_halo_indices(instantiate(LZ), instantiate(topology(f, 3)), f.grid.Nz, f.grid.Hz)) : + view(f.data, interior_indices(instantiate(LX), instantiate(topology(f, 1)), f.grid.Nx), + interior_indices(instantiate(LY), instantiate(topology(f, 2)), f.grid.Ny), + left_halo_indices(instantiate(LZ), instantiate(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)) + include_corners ? view(f.data, :, :, right_halo_indices(instantiate(LZ), instantiate(topology(f, 3)), f.grid.Nz, f.grid.Hz)) : + view(f.data, interior_indices(instantiate(LX), instantiate(topology(f, 1)), f.grid.Nx), + interior_indices(instantiate(LY), instantiate(topology(f, 2)), f.grid.Ny), + right_halo_indices(instantiate(LZ), instantiate(topology(f, 3)), f.grid.Nz, f.grid.Hz)) + +instantiate(T::Type) = T() +instantiate(t) = t underlying_west_halo(f, grid, location) = - view(f.parent, underlying_left_halo_indices(location, topology(grid, 1), grid.Nx, grid.Hx), :, :) + view(f.parent, underlying_left_halo_indices(instantiate(location), instantiate(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), :, :) + view(f.parent, underlying_right_halo_indices(instantiate(location), instantiate(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), :) + view(f.parent, :, underlying_left_halo_indices(instantiate(location), instantiate(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), :) + view(f.parent, :, underlying_right_halo_indices(instantiate(location), instantiate(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)) + view(f.parent, :, :, underlying_left_halo_indices(instantiate(location), instantiate(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)) + view(f.parent, :, :, underlying_right_halo_indices(instantiate(location), instantiate(topology(grid, 3)), grid.Nz, grid.Hz)) ##### ##### Viewing boundary grid points (used to fill other halos) ##### underlying_left_boundary_indices(loc, topo, N, H) = 1+H:2H -underlying_left_boundary_indices(::Type{Nothing}, topo, N, H) = 1:0 # empty +underlying_left_boundary_indices(::Nothing, topo, N, H) = 1:0 # empty underlying_right_boundary_indices(loc, topo, N, H) = N+1:N+H -underlying_right_boundary_indices(::Type{Nothing}, topo, N, H) = 1:0 # empty +underlying_right_boundary_indices(::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), :, :) + view(f.parent, underlying_left_boundary_indices(instantiate(location), instantiate(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), :, :) + view(f.parent, underlying_right_boundary_indices(instantiate(location), instantiate(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), :) + view(f.parent, :, underlying_left_boundary_indices(instantiate(location), instantiate(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), :) + view(f.parent, :, underlying_right_boundary_indices(instantiate(location), instantiate(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)) + view(f.parent, :, :, underlying_left_boundary_indices(instantiate(location), instantiate(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)) + view(f.parent, :, :, underlying_right_boundary_indices(instantiate(location), instantiate(topology(grid, 3)), grid.Nz, grid.Hz)) diff --git a/src/Distributed/halo_communication.jl b/src/Distributed/halo_communication.jl index 1c6aa531b1..1cc8d3a3dd 100644 --- a/src/Distributed/halo_communication.jl +++ b/src/Distributed/halo_communication.jl @@ -2,14 +2,14 @@ using KernelAbstractions: @kernel, @index, Event, MultiEvent using OffsetArrays: OffsetArray using CUDA: synchronize import Oceananigans.Utils: sync_device! -using Oceananigans.Fields: fill_west_and_east_send_buffers!, - fill_south_and_north_send_buffers!, +using Oceananigans.Fields: fill_west_and_east_send_buffers!, + fill_south_and_north_send_buffers!, fill_west_send_buffers!, fill_east_send_buffers!, fill_south_send_buffers!, fill_north_send_buffers!, - recv_from_buffers!, - reduced_dimensions, + recv_from_buffers!, + reduced_dimensions, instantiated_location import Oceananigans.Fields: tupled_fill_halo_regions! @@ -20,11 +20,11 @@ using Oceananigans.BoundaryConditions: permute_boundary_conditions, PBCT, DCBCT, DCBC -import Oceananigans.BoundaryConditions: +import Oceananigans.BoundaryConditions: fill_halo_regions!, fill_first, fill_halo_event!, fill_west_halo!, fill_east_halo!, fill_south_halo!, fill_north_halo!, fill_bottom_halo!, fill_top_halo!, - fill_west_and_east_halo!, + fill_west_and_east_halo!, fill_south_and_north_halo!, fill_bottom_and_top_halo! @@ -76,7 +76,7 @@ end ##### Filling halos for halo communication boundary conditions ##### -function tupled_fill_halo_regions!(full_fields, grid::DistributedGrid, args...; kwargs...) +function tupled_fill_halo_regions!(full_fields, grid::DistributedGrid, args...; kwargs...) for field in full_fields fill_halo_regions!(field, args...; kwargs...) end @@ -100,12 +100,12 @@ function fill_halo_regions!(c::OffsetArray, bcs, indices, loc, grid::Distributed arch = architecture(grid) child_arch = child_architecture(arch) halo_tuple = permute_boundary_conditions(bcs) - + for task = 1:3 barrier = device_event(child_arch) fill_halo_event!(task, halo_tuple, c, indices, loc, arch, barrier, grid, buffers, args...; kwargs...) end - + barrier = device_event(child_arch) fill_eventual_corners!(halo_tuple, c, indices, loc, arch, barrier, grid, buffers, args...; kwargs...) @@ -119,13 +119,13 @@ function fill_eventual_corners!(halo_tuple, c, indices, loc, arch, barrier, grid hbc_right = filter(bc -> bc isa DCBC, halo_tuple[3]) # 2D/3D Parallelization when `length(hbc_left) > 1 || length(hbc_right) > 1` - if length(hbc_left) > 1 + if length(hbc_left) > 1 idx = findfirst(bc -> bc isa DCBC, halo_tuple[2]) fill_halo_event!(idx, halo_tuple, c, indices, loc, arch, barrier, grid, buffers, args...; kwargs...) return nothing end - if length(hbc_right) > 1 + if length(hbc_right) > 1 idx = findfirst(bc -> bc isa DCBC, halo_tuple[3]) fill_halo_event!(idx, halo_tuple, c, indices, loc, arch, barrier, grid, buffers, args...; kwargs...) return nothing @@ -146,12 +146,12 @@ function fill_halo_event!(task, halo_tuple, c, indices, loc, arch::DistributedAr offset = fill_halo_offset(size, fill_halo!, indices) events_and_requests = fill_halo!(c, bc_left, bc_right, size, offset, loc, arch, barrier, grid, buffers, args...; kwargs...) - + if events_and_requests isa Event - wait(device(child_architecture(arch)), events_and_requests) + wait(device(child_architecture(arch)), events_and_requests) return nothing end - + MPI.Waitall(events_and_requests) buffer_side = mpi_communication_side(Val(fill_halo!)) @@ -208,7 +208,7 @@ for (side, opposite_side, dir) in zip([:west, :south, :bottom], [:east, :north, event = $fill_opposite_side_halo!(c, bc_opposite_side, size, offset, loc, arch, barrier, grid, buffers, args...; kwargs...) wait(device(child_arch), event) - + $fill_side_send_buffers!(c, buffers, grid) sync_device!(child_arch) @@ -293,4 +293,4 @@ for side in sides @inline $get_side_recv_buffer(c, grid, side_location, buffers, ::ViewsDistributedArch) = $underlying_side_halo(c, grid, side_location) @inline $get_side_recv_buffer(c, grid, side_location, buffers, arch) = buffers.$side.recv end -end \ No newline at end of file +end diff --git a/src/Fields/abstract_field.jl b/src/Fields/abstract_field.jl index 977f6ac00e..98db5efb95 100644 --- a/src/Fields/abstract_field.jl +++ b/src/Fields/abstract_field.jl @@ -53,7 +53,7 @@ Returns the size of an `AbstractField{LX, LY, LZ}` located at `LX, LY, LZ`. This is a 3-tuple of integers corresponding to the number of interior nodes of `f` along `x, y, z`. """ -Base.size(f::AbstractField) = size(location(f), f.grid) +Base.size(f::AbstractField) = size(f.grid, location(f)) Base.length(f::AbstractField) = prod(size(f)) Base.parent(f::AbstractField) = f @@ -63,7 +63,7 @@ Base.parent(f::AbstractField) = f Returns a 3-tuple that gives the "total" size of a field including both interior points and halo points. """ -total_size(f::AbstractField) = total_size(location(f), f.grid) +total_size(f::AbstractField) = total_size(f.grid, location(f)) interior(f::AbstractField) = f diff --git a/src/Fields/field.jl b/src/Fields/field.jl index 4ca8be5635..db2f00f4a3 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -35,7 +35,7 @@ end ##### function validate_field_data(loc, data, grid, indices) - Fx, Fy, Fz = total_size(loc, grid, indices) + Fx, Fy, Fz = total_size(grid, loc, indices) if size(data) != (Fx, Fy, Fz) LX, LY, LZ = loc @@ -84,9 +84,9 @@ end # Common outer constructor for all field flavors that performs input validation function Field(loc::Tuple, grid::AbstractGrid, data, bcs, indices, op=nothing, status=nothing) + @apply_regionally indices = validate_indices(indices, loc, grid) @apply_regionally validate_field_data(loc, data, grid, indices) @apply_regionally validate_boundary_conditions(loc, grid, bcs) - @apply_regionally indices = validate_indices(indices, loc, grid) buffers = FieldBoundaryBuffers(grid, data, bcs) LX, LY, LZ = loc return Field{LX, LY, LZ}(grid, data, bcs, indices, op, status, buffers) @@ -234,18 +234,20 @@ Return an `OffsetArray` of a `view` of `parent(data)` with `indices`. If `indices === (:, :, :)`, return an `OffsetArray` of `parent(data)`. """ -function offset_windowed_data(data, loc, grid, indices) +function offset_windowed_data(data, Loc, grid, indices) halo = halo_size(grid) - topo = topology(grid) + topo = map(instantiate, topology(grid)) + loc = map(instantiate, Loc) if indices isa typeof(default_indices(3)) windowed_parent = parent(data) else - parent_indices = parent_index_range.(indices, loc, topo, halo) + parent_indices = map(parent_index_range, indices, loc, topo, halo) windowed_parent = view(parent(data), parent_indices...) end sz = size(grid) + return offset_data(windowed_parent, loc, topo, sz, halo, indices) end @@ -345,9 +347,7 @@ function boundary_conditions(f::Field) end immersed_boundary_condition(f::Field) = f.boundary_conditions.immersed - data(field::Field) = field.data - indices(obj, i=default_indices(3)) = i indices(f::Field, i=default_indices(3)) = f.indices indices(a::SubArray, i=default_indices(ndims(a))) = a.indices @@ -357,15 +357,20 @@ indices(a::OffsetArray, i=default_indices(ndims(a))) = indices(parent(a), i) interior_view_indices(field_indices, interior_indices) = Colon() interior_view_indices(::Colon, interior_indices) = interior_indices +instantiate(T::Type) = T() +instantiate(t) = t + function interior(a::OffsetArray, - loc::Tuple, - topo::Tuple, - size::NTuple{N, Int}, - halo_size::NTuple{N, Int}, + Loc::Tuple, + Topo::Tuple, + sz::NTuple{N, Int}, + halo_sz::NTuple{N, Int}, ind::Tuple=default_indices(3)) where N - i_interior = interior_parent_indices.(loc, topo, size, halo_size) - i_view = interior_view_indices.(ind, i_interior) + loc = map(instantiate, Loc) + topo = map(instantiate, Topo) + i_interior = map(interior_parent_indices, loc, topo, sz, halo_sz) + i_view = map(interior_view_indices, ind, i_interior) return view(parent(a), i_view...) end @@ -399,11 +404,8 @@ Base.fill!(f::Field, val) = fill!(parent(f), val) Base.parent(f::Field) = parent(f.data) Adapt.adapt_structure(to, f::Field) = Adapt.adapt(to, f.data) -length_indices(N, i::Colon) = N -length_indices(N, i::UnitRange) = length(i) - -total_size(f::Field) = length_indices.(total_size(location(f), f.grid), f.indices) -Base.size(f::Field) = length_indices.( size(location(f), f.grid), f.indices) +total_size(f::Field) = total_size(f.grid, location(f), f.indices) +Base.size(f::Field) = size(f.grid, location(f), f.indices) ==(f::Field, a) = interior(f) == a ==(a, f::Field) = a == interior(f) @@ -723,3 +725,4 @@ function fill_halo_regions!(field::Field, args...; kwargs...) return nothing end + diff --git a/src/Fields/interpolate.jl b/src/Fields/interpolate.jl index e6386470a4..4c3f8f4904 100644 --- a/src/Fields/interpolate.jl +++ b/src/Fields/interpolate.jl @@ -58,13 +58,13 @@ end @inline fractional_y_index(y::FT, ::Nothing, grid) where FT = one(FT) @inline fractional_z_index(z::FT, ::Nothing, grid) where FT = one(FT) -@inline fractional_x_index(x::FT, ::Center, grid) where FT = fractional_index(length(Center, topology(grid)[1], grid.Nx), x, xnodes(grid, Center())) -@inline fractional_y_index(y::FT, ::Center, grid) where FT = fractional_index(length(Center, topology(grid)[2], grid.Ny), y, ynodes(grid, Center())) -@inline fractional_z_index(z::FT, ::Center, grid) where FT = fractional_index(length(Center, topology(grid)[3], grid.Nz), z, znodes(grid, Center())) +@inline fractional_x_index(x::FT, ::Center, grid) where FT = fractional_index(length(Center(), topology(grid, 1)(), size(grid, 1)), x, xnodes(grid, Center())) +@inline fractional_y_index(y::FT, ::Center, grid) where FT = fractional_index(length(Center(), topology(grid, 2)(), size(grid, 2)), y, ynodes(grid, Center())) +@inline fractional_z_index(z::FT, ::Center, grid) where FT = fractional_index(length(Center(), topology(grid, 3)(), size(grid, 3)), z, znodes(grid, Center())) -@inline fractional_x_index(x::FT, ::Face, grid) where FT = fractional_index(length(Face, topology(grid)[1], grid.Nx), x, xnodes(grid, Face())) - 1 -@inline fractional_y_index(y::FT, ::Face, grid) where FT = fractional_index(length(Face, topology(grid)[2], grid.Ny), y, ynodes(grid, Face())) - 1 -@inline fractional_z_index(z::FT, ::Face, grid) where FT = fractional_index(length(Face, topology(grid)[3], grid.Nz), z, znodes(grid, Face())) - 1 +@inline fractional_x_index(x::FT, ::Face, grid) where FT = fractional_index(length(Face(), topology(grid, 1)(), size(grid, 1)), x, xnodes(grid, Face())) - 1 +@inline fractional_y_index(y::FT, ::Face, grid) where FT = fractional_index(length(Face(), topology(grid, 2)(), size(grid, 2)), y, ynodes(grid, Face())) - 1 +@inline fractional_z_index(z::FT, ::Face, grid) where FT = fractional_index(length(Face(), topology(grid, 3)(), size(grid, 3)), z, znodes(grid, Face())) - 1 @inline fractional_x_index(x::FT, ::Face, grid::XRegRectilinearGrid) where FT = @inbounds FT((x - grid.xᶠᵃᵃ[1]) / grid.Δxᶠᵃᵃ) @inline fractional_x_index(x::FT, ::Center, grid::XRegRectilinearGrid) where FT = @inbounds FT((x - grid.xᶜᵃᵃ[1]) / grid.Δxᶜᵃᵃ) diff --git a/src/Fields/regridding_fields.jl b/src/Fields/regridding_fields.jl index 18af9b92c7..34a199d7db 100644 --- a/src/Fields/regridding_fields.jl +++ b/src/Fields/regridding_fields.jl @@ -88,7 +88,8 @@ end function regrid_in_y!(a, target_grid, source_grid, b) arch = architecture(a) source_y_faces = ynodes(source_grid, f) - event = launch!(arch, target_grid, :xz, _regrid_in_y!, a, b, target_grid, source_grid, source_y_faces) + Nx_source_faces = size(source_grid, (Face, Center, Center), 1) + event = launch!(arch, target_grid, :xz, _regrid_in_y!, a, b, target_grid, source_grid, source_y_faces, Nx_source_faces) wait(device(arch), event) return a end @@ -96,7 +97,8 @@ end function regrid_in_x!(a, target_grid, source_grid, b) arch = architecture(a) source_x_faces = xnodes(source_grid, f) - event = launch!(arch, target_grid, :yz, _regrid_in_x!, a, b, target_grid, source_grid, source_x_faces) + Ny_source_faces = size(source_grid, (Center, Face, Center), 2) + event = launch!(arch, target_grid, :yz, _regrid_in_x!, a, b, target_grid, source_grid, source_x_faces, Ny_source_faces) wait(device(arch), event) return a end @@ -125,7 +127,7 @@ function regrid!(a, target_grid, source_grid, b) end ##### -##### Regridding for single column grids +##### Regridding for all grids ##### @kernel function _regrid_in_z!(target_field, source_field, target_grid, source_grid, source_z_faces) @@ -181,7 +183,7 @@ end end end -@kernel function _regrid_in_y!(target_field, source_field, target_grid, source_grid, source_y_faces) +@kernel function _regrid_in_y!(target_field, source_field, target_grid, source_grid, source_y_faces, Nx_source_faces) i, k = @index(Global, NTuple) Nx_target, Ny_target, Nz_target = size(target_grid) @@ -189,7 +191,6 @@ end i_src = ifelse(Nx_target == Nx_source, i, 1) k_src = ifelse(Nz_target == Nz_source, k, 1) - Nx_source_faces = size((Face, Center, Center), source_grid, 1) i⁺_src = min(Nx_source_faces, i_src + 1) fo = ForwardOrdering() @@ -226,7 +227,7 @@ end if j₋_src > 1 j_left = j₋_src - 1 - x₁ = xnode(i_src, source_grid, f) + x₁ = xnode(i_src, source_grid, f) x₂ = xnode(i⁺_src, source_grid, f) Az_left = fractional_horizontal_area(source_grid, x₁, x₂, y₋, yj₋_src) @@ -237,7 +238,7 @@ end if j₊_src < source_grid.Ny+1 j_right = j₊_src - x₁ = xnode(i_src, source_grid, f) + x₁ = xnode(i_src, source_grid, f) x₂ = xnode(i⁺_src, source_grid, f) Az_right = fractional_horizontal_area(source_grid, x₁, x₂, yj₊_src, y₊) @@ -249,7 +250,7 @@ end end end -@kernel function _regrid_in_x!(target_field, source_field, target_grid, source_grid, source_x_faces) +@kernel function _regrid_in_x!(target_field, source_field, target_grid, source_grid, source_x_faces, Ny_source_faces) j, k = @index(Global, NTuple) Nx_target, Ny_target, Nz_target = size(target_grid) @@ -257,7 +258,6 @@ end j_src = ifelse(Ny_target == Ny_source, j, 1) k_src = ifelse(Nz_target == Nz_source, k, 1) - Ny_source_faces = size((Center, Face, Center), source_grid, 2) j⁺_src = min(Ny_source_faces, j_src + 1) fo = ForwardOrdering() @@ -267,7 +267,7 @@ end # Integrate source field from x₋ to x₊ x₋ = xnode(i, j, k, target_grid, f, c, c) - x₊ = xnode(i+1, j, k, target_grid, f, c, c) + x₊ = xnode(i+1, j, k, target_grid, f, c, c) # The first face on the source grid that appears inside the target cell i₋_src = searchsortedfirst(source_x_faces, x₋, 1, Nx_source+1, fo) @@ -302,7 +302,7 @@ end if i₋_src > 1 i_left = i₋_src - 1 - y₁ = ynode(j_src, source_grid, f) + y₁ = ynode(j_src, source_grid, f) y₂ = ynode(j⁺_src, source_grid, f) Az_left = fractional_horizontal_area(source_grid, x₋, xi₋_src, y₁, y₂) @@ -313,8 +313,8 @@ end if i₊_src < source_grid.Nx+1 i_right = i₊_src - y₁ = ynode(j_src, source_grid, f) - y₂ = ynode(j⁺_src, source_grid, f) + y₁ = ynode(j_src, source_grid, f) + y₂ = ynode(j⁺_src, source_grid, f) Az_right = fractional_horizontal_area(source_grid, xi₊_src, x₊, y₁, y₂) @inbounds target_field[i, j, k] += source_field[i_right, j_src, k_src] * Az_right diff --git a/src/Grids/Grids.jl b/src/Grids/Grids.jl index 7b2fc868ab..4feb604f26 100644 --- a/src/Grids/Grids.jl +++ b/src/Grids/Grids.jl @@ -10,7 +10,8 @@ export AbstractCurvilinearGrid, AbstractHorizontallyCurvilinearGrid export LatitudeLongitudeGrid, XRegLatLonGrid, YRegLatLonGrid, ZRegLatLonGrid export OrthogonalSphericalShellGrid, ConformalCubedSphereGrid export node, xnode, ynode, znode, xnodes, ynodes, znodes, nodes -export xspacings, yspacings, zspacings +export xspacings, yspacings, zspacings, xspacing, yspacing, zspacing +export minimum_xspacing, minimum_yspacing, minimum_zspacing export offset_data, new_data export on_architecture diff --git a/src/Grids/grid_generation.jl b/src/Grids/grid_generation.jl index 1828a33c34..f929449666 100644 --- a/src/Grids/grid_generation.jl +++ b/src/Grids/grid_generation.jl @@ -12,20 +12,21 @@ get_coord_face(coord::Nothing, i) = 1 get_coord_face(coord::Function, i) = coord(i) get_coord_face(coord::AbstractVector, i) = CUDA.@allowscalar coord[i] -lower_exterior_Δcoordᶠ(topology, Fi, Hcoord) = [Fi[end - Hcoord + i] - Fi[end - Hcoord + i - 1] for i = 1:Hcoord] -lower_exterior_Δcoordᶠ(::Type{<:BoundedTopology}, Fi, Hcoord) = [Fi[2] - Fi[1] for i = 1:Hcoord] +const AT = AbstractTopology +lower_exterior_Δcoordᶠ(::AT, Fi, Hcoord) = [Fi[end - Hcoord + i] - Fi[end - Hcoord + i - 1] for i = 1:Hcoord] +lower_exterior_Δcoordᶠ(::BoundedTopology, Fi, Hcoord) = [Fi[2] - Fi[1] for i = 1:Hcoord] -upper_exterior_Δcoordᶠ(topology, Fi, Hcoord) = [Fi[i + 1] - Fi[i] for i = 1:Hcoord] -upper_exterior_Δcoordᶠ(::Type{<:BoundedTopology}, Fi, Hcoord) = [Fi[end] - Fi[end - 1] for i = 1:Hcoord] +upper_exterior_Δcoordᶠ(::AT, Fi, Hcoord) = [Fi[i + 1] - Fi[i] for i = 1:Hcoord] +upper_exterior_Δcoordᶠ(::BoundedTopology, Fi, Hcoord) = [Fi[end] - Fi[end - 1] for i = 1:Hcoord] -upper_interior_F(topology, coord, Δ) = coord - Δ -upper_interior_F(::Type{<:BoundedTopology}, coord) = coord +upper_interior_F(::AT, coord, Δ) = coord - Δ +upper_interior_F(::BoundedTopology, coord) = coord -total_interior_length(topology, N) = N -total_interior_length(::Type{<:BoundedTopology}, N) = N + 1 +total_interior_length(::AT, N) = N +total_interior_length(::BoundedTopology, N) = N + 1 # generate a stretched coordinate passing the explicit coord faces as vector of functionL -function generate_coordinate(FT, topology, N, H, coord, arch) +function generate_coordinate(FT, topo::AT, N, H, coord, arch) # Ensure correct type for F and derived quantities interiorF = zeros(FT, N+1) @@ -37,8 +38,8 @@ function generate_coordinate(FT, topology, N, H, coord, arch) L = interiorF[N+1] - interiorF[1] # Build halo regions - Δᶠ₋ = lower_exterior_Δcoordᶠ(topology, interiorF, H) - Δᶠ₊ = reverse(upper_exterior_Δcoordᶠ(topology, interiorF, H)) + Δᶠ₋ = lower_exterior_Δcoordᶠ(topo, interiorF, H) + Δᶠ₊ = reverse(upper_exterior_Δcoordᶠ(topo, interiorF, H)) c¹, cᴺ⁺¹ = interiorF[1], interiorF[N+1] @@ -48,12 +49,12 @@ function generate_coordinate(FT, topology, N, H, coord, arch) F = vcat(F₋, interiorF, F₊) # Build cell centers, cell center spacings, and cell interface spacings - TC = total_length(Center, topology, N, H) + TC = total_length(Center(), topo, N, H) C = [ (F[i + 1] + F[i]) / 2 for i = 1:TC ] Δᶠ = [ C[i] - C[i - 1] for i = 2:TC ] # Trim face locations for periodic domains - TF = total_length(Face, topology, N, H) + TF = total_length(Face(), topo, N, H) F = F[1:TF] Δᶜ = [F[i + 1] - F[i] for i = 1:TF-1] @@ -77,7 +78,7 @@ function generate_coordinate(FT, topology, N, H, coord, arch) end # generate a regular coordinate passing the domain extent (2-tuple) and number of points -function generate_coordinate(FT, topology, N, H, coord::Tuple{<:Number, <:Number}, arch) +function generate_coordinate(FT, topo::AT, N, H, coord::Tuple{<:Number, <:Number}, arch) @assert length(coord) == 2 @@ -89,13 +90,13 @@ function generate_coordinate(FT, topology, N, H, coord::Tuple{<:Number, <:Number Δᶠ = Δᶜ = Δ = L / N F₋ = c₁ - H * Δ - F₊ = F₋ + total_extent(topology, H, Δ, L) + F₊ = F₋ + total_extent(topo, H, Δ, L) C₋ = F₋ + Δ / 2 C₊ = C₋ + L + Δ * (2H - 1) - TF = total_length(Face, topology, N, H) - TC = total_length(Center, topology, N, H) + TF = total_length(Face(), topo, N, H) + TC = total_length(Center(), topo, N, H) F = range(FT(F₋), FT(F₊), length = TF) C = range(FT(C₋), FT(C₊), length = TC) @@ -107,6 +108,6 @@ function generate_coordinate(FT, topology, N, H, coord::Tuple{<:Number, <:Number end # Flat domains -function generate_coordinate(FT, ::Type{Flat}, N, H, coord::Tuple{<:Number, <:Number}, arch) +function generate_coordinate(FT, ::Flat, N, H, coord::Tuple{<:Number, <:Number}, arch) return FT(1), range(1, 1, length=N), range(1, 1, length=N), FT(1), FT(1) end diff --git a/src/Grids/grid_utils.jl b/src/Grids/grid_utils.jl index a7a6b16250..3e27701cc9 100644 --- a/src/Grids/grid_utils.jl +++ b/src/Grids/grid_utils.jl @@ -4,21 +4,11 @@ using Base.Ryu: writeshortest using LinearAlgebra: dot, cross using OffsetArrays: IdOffsetRange -##### -##### Convenience functions -##### +# Define default indices +default_indices(n) = Tuple(Colon() for i=1:n) const BoundedTopology = Union{Bounded, LeftConnected} -Base.length(::Type{Face}, topo, N) = N -Base.length(::Type{Face}, ::Type{<:BoundedTopology}, N) = N+1 -Base.length(::Type{Center}, topo, N) = N -Base.length(::Type{Nothing}, topo, N) = 1 - -Base.length(::Type{Nothing}, ::Type{Flat}, N) = N -Base.length(::Type{Face}, ::Type{Flat}, N) = N -Base.length(::Type{Center}, ::Type{Flat}, N) = N - """ topology(grid) @@ -54,49 +44,84 @@ function Base.:(==)(grid1::AbstractGrid, grid2::AbstractGrid) CUDA.@allowscalar return x1 == x2 && y1 == y2 && z1 == z2 end +const AT = AbstractTopology + +Base.length(::Face, ::BoundedTopology, N) = N + 1 +Base.length(::Nothing, ::AT, N) = 1 +Base.length(::Face, ::AT, N) = N +Base.length(::Center, ::AT, N) = N +Base.length(::Nothing, ::Flat, N) = N +Base.length(::Face, ::Flat, N) = N +Base.length(::Center, ::Flat, N) = N + +# "Indices-aware" length +Base.length(loc, topo::AT, N, ::Colon) = length(loc, topo, N) +Base.length(loc, topo::AT, N, ind::UnitRange) = min(length(loc, topo, N), length(ind)) + """ - size(loc, grid) + size(grid) -Return the size of a `grid` at `loc`, not including halos. -This is a 3-tuple of integers corresponding to the number of interior nodes -along `x, y, z`. +Return a 3-tuple of the number of "center" cells on a grid in (x, y, z). +Center cells have the location (Center, Center, Center). """ -Base.size(loc, grid::AbstractGrid) = (length(loc[1], topology(grid, 1), grid.Nx), - length(loc[2], topology(grid, 2), grid.Ny), - length(loc[3], topology(grid, 3), grid.Nz)) +Base.size(grid::AbstractGrid) = (grid.Nx, grid.Ny, grid.Nz) -Base.size(grid::AbstractGrid) = size((Center, Center, Center), grid) -Base.size(grid::AbstractGrid, d) = size(grid)[d] -Base.size(loc, grid, d) = size(loc, grid)[d] +""" + halo_size(grid) -total_size(a) = size(a) # fallback +Return a 3-tuple with the number of halo cells on either side of the +domain in (x, y, z). +""" +halo_size(grid) = (grid.Hx, grid.Hy, grid.Hz) +halo_size(grid, d) = halo_size(grid)[d] + +Base.size(grid::AbstractGrid, d::Int) = size(grid)[d] + +Base.size(grid::AbstractGrid, loc::Tuple, indices=default_indices(length(loc))) = + size(loc, topology(grid), size(grid), indices) + +function Base.size(loc, topo, sz, indices=default_indices(length(loc))) + D = length(loc) + return Tuple(length(instantiate(loc[d]), instantiate(topo[d]), sz[d], indices[d]) for d = 1:D) +end + +Base.size(grid::AbstractGrid, loc::Tuple, d::Int) = size(grid, loc)[d] """ - total_size(loc, grid) + total_length(loc, topo, N, H=0, ind=Colon()) -Return the "total" size of a `grid` at `loc`. This is a 3-tuple of integers -corresponding to the number of grid points along `x, y, z`. +Return the total length of a field at `loc`ation along +one dimension of `topo`logy with `N` centered cells and +`H` halo cells. If `ind` is provided the total_length +is restricted by `length(ind)`. """ -total_size(loc, grid) = (total_length(loc[1], topology(grid, 1), grid.Nx, grid.Hx), - total_length(loc[2], topology(grid, 2), grid.Ny, grid.Hy), - total_length(loc[3], topology(grid, 3), grid.Nz, grid.Hz)) +total_length(::Face, ::AT, N, H=0) = N + 2H +total_length(::Center, ::AT, N, H=0) = N + 2H +total_length(::Face, ::BoundedTopology, N, H=0) = N + 1 + 2H +total_length(::Nothing, ::AT, N, H=0) = 1 +total_length(::Nothing, ::Flat, N, H=0) = N +total_length(::Face, ::Flat, N, H=0) = N +total_length(::Center, ::Flat, N, H=0) = N -function total_size(loc, grid, indices::Tuple) - sz = total_size(loc, grid) - return Tuple(ind isa Colon ? sz[i] : min(length(ind), sz[i]) for (i, ind) in enumerate(indices)) -end +# "Indices-aware" total length +total_length(loc, topo, N, H, ::Colon) = total_length(loc, topo, N, H) +total_length(loc, topo, N, H, ind::UnitRange) = min(total_length(loc, topo, N, H), length(ind)) -function Base.size(loc, grid::AbstractGrid, indices::Tuple) - sz = size(loc, grid) - return Tuple(ind isa Colon ? sz[i] : min(length(ind), sz[i]) for (i, ind) in enumerate(indices)) -end +total_size(a) = size(a) # fallback """ - halo_size(grid) + total_size(grid, loc) -Return a tuple with the size of the halo in each dimension. +Return the "total" size of a `grid` at `loc`. This is a 3-tuple of integers +corresponding to the number of grid points along `x, y, z`. """ -halo_size(grid) = (grid.Hx, grid.Hy, grid.Hz) +function total_size(loc, topo, sz, halo_sz, indices=default_indices(length(loc))) + D = length(loc) + return Tuple(total_length(instantiate(loc[d]), instantiate(topo[d]), sz[d], halo_sz[d], indices[d]) for d = 1:D) +end + +total_size(grid::AbstractGrid, loc, indices=default_indices(length(loc))) = + total_size(loc, topology(grid), size(grid), halo_size(grid), indices) """ total_extent(topology, H, Δ, L) @@ -105,30 +130,16 @@ Return the total extent, including halo regions, of constant-spaced `Periodic` and `Flat` dimensions with number of halo points `H`, constant grid spacing `Δ`, and interior extent `L`. """ -@inline total_extent(topology, H, Δ, L) = L + (2H - 1) * Δ -@inline total_extent(::Type{<:BoundedTopology}, H, Δ, L) = L + 2H * Δ - -""" - total_length(loc, topo, N, H=0) - -Return the total length of a field at `loc`ation along -one dimension of `topo`logy with `N` centered cells and -`H` halo cells. -""" -@inline total_length(loc, topo, N, H=0) = N + 2H -@inline total_length(::Type{Face}, ::Type{<:BoundedTopology}, N, H=0) = N + 1 + 2H -@inline total_length(::Type{Nothing}, topo, N, H=0) = 1 -@inline total_length(::Type{Nothing}, ::Type{Flat}, N, H=0) = N -@inline total_length(::Type{Face}, ::Type{Flat}, N, H=0) = N -@inline total_length(::Type{Center}, ::Type{Flat}, N, H=0) = N +@inline total_extent(topo, H, Δ, L) = L + (2H - 1) * Δ +@inline total_extent(::BoundedTopology, H, Δ, L) = L + 2H * Δ # Grid domains @inline domain(topo, N, ξ) = CUDA.@allowscalar ξ[1], ξ[N+1] -@inline domain(::Type{Flat}, N, ξ) = CUDA.@allowscalar ξ[1], ξ[1] +@inline domain(::Flat, N, ξ) = CUDA.@allowscalar ξ[1], ξ[1] -@inline x_domain(grid) = domain(topology(grid, 1), grid.Nx, grid.xᶠᵃᵃ) -@inline y_domain(grid) = domain(topology(grid, 2), grid.Ny, grid.yᵃᶠᵃ) -@inline z_domain(grid) = domain(topology(grid, 3), grid.Nz, grid.zᵃᵃᶠ) +@inline x_domain(grid) = domain(topology(grid, 1)(), grid.Nx, grid.xᶠᵃᵃ) +@inline y_domain(grid) = domain(topology(grid, 2)(), grid.Ny, grid.yᵃᶠᵃ) +@inline z_domain(grid) = domain(topology(grid, 3)(), grid.Nz, grid.zᵃᵃᶠ) regular_dimensions(grid) = () @@ -136,127 +147,127 @@ regular_dimensions(grid) = () ##### << Indexing >> ##### -@inline left_halo_indices(loc, topo, N, H) = 1-H:0 -@inline left_halo_indices(::Type{Nothing}, topo, N, H) = 1:0 # empty +@inline left_halo_indices(loc, ::AT, N, H) = 1-H:0 +@inline left_halo_indices(::Nothing, ::AT, N, H) = 1:0 # empty -@inline right_halo_indices(loc, topo, N, H) = N+1:N+H -@inline right_halo_indices(::Type{Face}, ::Type{<:BoundedTopology}, N, H) = N+2:N+1+H -@inline right_halo_indices(::Type{Nothing}, topo, N, H) = 1:0 # empty +@inline right_halo_indices(loc, ::AT, N, H) = N+1:N+H +@inline right_halo_indices(::Face, ::BoundedTopology, N, H) = N+2:N+1+H +@inline right_halo_indices(::Nothing, ::AT, N, H) = 1:0 # empty -@inline underlying_left_halo_indices(loc, topo, N, H) = 1:H -@inline underlying_left_halo_indices(::Type{Nothing}, topo, N, H) = 1:0 # empty +@inline underlying_left_halo_indices(loc, ::AT, N, H) = 1:H +@inline underlying_left_halo_indices(::Nothing, ::AT, N, H) = 1:0 # empty -@inline underlying_right_halo_indices(loc, topo, N, H) = N+1+H:N+2H -@inline underlying_right_halo_indices(::Type{Face}, ::Type{<:BoundedTopology}, N, H) = N+2+H:N+1+2H -@inline underlying_right_halo_indices(::Type{Nothing}, topo, N, H) = 1:0 # empty +@inline underlying_right_halo_indices(loc, ::AT, N, H) = N+1+H:N+2H +@inline underlying_right_halo_indices(::Face, ::BoundedTopology, N, H) = N+2+H:N+1+2H +@inline underlying_right_halo_indices(::Nothing, ::AT, N, H) = 1:0 # empty -@inline interior_indices(loc, topo, N) = 1:N -@inline interior_indices(::Type{Face}, ::Type{<:BoundedTopology}, N) = 1:N+1 -@inline interior_indices(::Type{Nothing}, topo, N) = 1:1 +@inline interior_indices(loc, ::AT, N) = 1:N +@inline interior_indices(::Face, ::BoundedTopology, N) = 1:N+1 +@inline interior_indices(::Nothing, ::AT, N) = 1:1 -@inline interior_indices(::Type{Nothing}, topo::Type{Flat}, N) = 1:N -@inline interior_indices(::Type{Face}, topo::Type{Flat}, N) = 1:N -@inline interior_indices(::Type{Center}, topo::Type{Flat}, N) = 1:N +@inline interior_indices(::Nothing, ::Flat, N) = 1:N +@inline interior_indices(::Face, ::Flat, N) = 1:N +@inline interior_indices(::Center, ::Flat, N) = 1:N -@inline interior_x_indices(loc, grid) = interior_indices(loc, topology(grid, 1), grid.Nx) -@inline interior_y_indices(loc, grid) = interior_indices(loc, topology(grid, 2), grid.Ny) -@inline interior_z_indices(loc, grid) = interior_indices(loc, topology(grid, 3), grid.Nz) +@inline interior_x_indices(grid, loc) = interior_indices(loc[1], topology(grid, 1)(), size(grid, 1)) +@inline interior_y_indices(grid, loc) = interior_indices(loc[2], topology(grid, 2)(), size(grid, 2)) +@inline interior_z_indices(grid, loc) = interior_indices(loc[3], topology(grid, 3)(), size(grid, 3)) -@inline interior_parent_offset(loc, topo, H) = H -@inline interior_parent_offset(::Type{Nothing}, topo, H) = 0 +@inline interior_parent_offset(loc, ::AT, H) = H +@inline interior_parent_offset(::Nothing, ::AT, H) = 0 -@inline interior_parent_indices(loc, topo, N, H) = 1+H:N+H -@inline interior_parent_indices(::Type{Face}, ::Type{<:BoundedTopology}, N, H) = 1+H:N+1+H -@inline interior_parent_indices(::Type{Nothing}, topo, N, H) = 1:1 +@inline interior_parent_indices(::Nothing, ::AT, N, H) = 1:1 +@inline interior_parent_indices(::Face, ::BoundedTopology, N, H) = 1+H:N+1+H +@inline interior_parent_indices(loc, ::AT, N, H) = 1+H:N+H -@inline interior_parent_indices(::Type{Nothing}, ::Type{Flat}, N, H) = 1:N -@inline interior_parent_indices(::Type{Face}, ::Type{Flat}, N, H) = 1:N -@inline interior_parent_indices(::Type{Center}, ::Type{Flat}, N, H) = 1:N +@inline interior_parent_indices(::Nothing, ::Flat, N, H) = 1:N +@inline interior_parent_indices(::Face, ::Flat, N, H) = 1:N +@inline interior_parent_indices(::Center, ::Flat, N, H) = 1:N # All indices including halos. -@inline all_indices(loc, topo, N, H) = 1-H:N+H -@inline all_indices(::Type{Face}, ::Type{<:BoundedTopology}, N, H) = 1-H:N+1+H -@inline all_indices(::Type{Nothing}, topo, N, H) = 1:1 +@inline all_indices(::Nothing, ::AT, N, H) = 1:1 +@inline all_indices(::Face, ::BoundedTopology, N, H) = 1-H:N+1+H +@inline all_indices(loc, ::AT, N, H) = 1-H:N+H -@inline all_indices(::Type{Nothing}, ::Type{Flat}, N, H) = 1:N -@inline all_indices(::Type{Face}, ::Type{Flat}, N, H) = 1:N -@inline all_indices(::Type{Center}, ::Type{Flat}, N, H) = 1:N +@inline all_indices(::Nothing, ::Flat, N, H) = 1:N +@inline all_indices(::Face, ::Flat, N, H) = 1:N +@inline all_indices(::Center, ::Flat, N, H) = 1:N -@inline all_x_indices(loc, grid) = all_indices(loc, topology(grid, 1), grid.Nx, grid.Hx) -@inline all_y_indices(loc, grid) = all_indices(loc, topology(grid, 2), grid.Ny, grid.Hy) -@inline all_z_indices(loc, grid) = all_indices(loc, topology(grid, 3), grid.Nz, grid.Hz) +@inline all_x_indices(grid, loc) = all_indices(loc[1](), topology(grid, 1)(), size(grid, 1), halo_size(grid, 1)) +@inline all_y_indices(grid, loc) = all_indices(loc[2](), topology(grid, 2)(), size(grid, 2), halo_size(grid, 2)) +@inline all_z_indices(grid, loc) = all_indices(loc[3](), topology(grid, 3)(), size(grid, 3), halo_size(grid, 3)) -@inline all_parent_indices(loc, topo, N, H) = 1:N+2H -@inline all_parent_indices(::Type{Face}, ::Type{<:BoundedTopology}, N, H) = 1:N+1+2H -@inline all_parent_indices(::Type{Nothing}, topo, N, H) = 1:1 +@inline all_parent_indices(loc, ::AT, N, H) = 1:N+2H +@inline all_parent_indices(::Face, ::BoundedTopology, N, H) = 1:N+1+2H +@inline all_parent_indices(::Nothing, ::AT, N, H) = 1:1 -@inline all_parent_indices(::Type{Nothing}, ::Type{Flat}, N, H) = 1:N -@inline all_parent_indices(::Type{Face}, ::Type{Flat}, N, H) = 1:N -@inline all_parent_indices(::Type{Center}, ::Type{Flat}, N, H) = 1:N +@inline all_parent_indices(::Nothing, ::Flat, N, H) = 1:N +@inline all_parent_indices(::Face, ::Flat, N, H) = 1:N +@inline all_parent_indices(::Center, ::Flat, N, H) = 1:N -@inline all_parent_x_indices(loc, grid) = all_parent_indices(loc, topology(grid, 1), grid.Nx, grid.Hx) -@inline all_parent_y_indices(loc, grid) = all_parent_indices(loc, topology(grid, 2), grid.Ny, grid.Hy) -@inline all_parent_z_indices(loc, grid) = all_parent_indices(loc, topology(grid, 3), grid.Nz, grid.Hz) +@inline all_parent_x_indices(grid, loc) = all_parent_indices(loc[1](), topology(grid, 1)(), size(grid, 1), halo_size(grid, 1)) +@inline all_parent_y_indices(grid, loc) = all_parent_indices(loc[2](), topology(grid, 2)(), size(grid, 2), halo_size(grid, 2)) +@inline all_parent_z_indices(grid, loc) = all_parent_indices(loc[3](), topology(grid, 3)(), size(grid, 3), halo_size(grid, 3)) parent_index_range(::Colon, loc, topo, halo) = Colon() parent_index_range(::Base.Slice{<:IdOffsetRange}, loc, topo, halo) = Colon() parent_index_range(index::UnitRange, loc, topo, halo) = index .+ interior_parent_offset(loc, topo, halo) -parent_index_range(index::UnitRange, ::Type{Nothing}, ::Type{Flat}, halo) = index -parent_index_range(index::UnitRange, ::Type{Nothing}, topo, halo) = 1:1 # or Colon() +parent_index_range(index::UnitRange, ::Nothing, ::Flat, halo) = index +parent_index_range(index::UnitRange, ::Nothing, ::AT, halo) = 1:1 # or Colon() index_range_offset(index::UnitRange, loc, topo, halo) = index[1] - interior_parent_offset(loc, topo, halo) index_range_offset(::Colon, loc, topo, halo) = - interior_parent_offset(loc, topo, halo) -@inline cpu_face_constructor_x(grid) = Array(xnodes(grid, Face(); with_halos=true)[1:grid.Nx+1]) -@inline cpu_face_constructor_y(grid) = Array(ynodes(grid, Face(); with_halos=true)[1:grid.Ny+1]) -@inline cpu_face_constructor_z(grid) = Array(znodes(grid, Face(); with_halos=true)[1:grid.Nz+1]) +@inline cpu_face_constructor_x(grid) = Array(xnodes(grid, Face(); with_halos=true)[1:size(grid, 1)+1]) +@inline cpu_face_constructor_y(grid) = Array(ynodes(grid, Face(); with_halos=true)[1:size(grid, 2)+1]) +@inline cpu_face_constructor_z(grid) = Array(znodes(grid, Face(); with_halos=true)[1:size(grid, 3)+1]) ##### ##### << Nodes >> ##### -@inline node(i, j, k, grid, LX, LY, LZ) = (xnode(i, j, k, grid, LX, LY, LZ), - ynode(i, j, k, grid, LX, LY, LZ), - znode(i, j, k, grid, LX, LY, LZ)) +@inline node(i, j, k, grid, ℓx, ℓy, ℓz) = (xnode(i, j, k, grid, ℓx, ℓy, ℓz), + ynode(i, j, k, grid, ℓx, ℓy, ℓz), + znode(i, j, k, grid, ℓx, ℓy, ℓz)) -@inline node(i, j, k, grid, LX::Nothing, LY, LZ) = (ynode(i, j, k, grid, LX, LY, LZ), znode(i, j, k, grid, LX, LY, LZ)) -@inline node(i, j, k, grid, LX, LY::Nothing, LZ) = (xnode(i, j, k, grid, LX, LY, LZ), znode(i, j, k, grid, LX, LY, LZ)) -@inline node(i, j, k, grid, LX, LY, LZ::Nothing) = (xnode(i, j, k, grid, LX, LY, LZ), ynode(i, j, k, grid, LX, LY, LZ)) +@inline node(i, j, k, grid, ℓx::Nothing, ℓy, ℓz) = (ynode(i, j, k, grid, ℓx, ℓy, ℓz), znode(i, j, k, grid, ℓx, ℓy, ℓz)) +@inline node(i, j, k, grid, ℓx, ℓy::Nothing, ℓz) = (xnode(i, j, k, grid, ℓx, ℓy, ℓz), znode(i, j, k, grid, ℓx, ℓy, ℓz)) +@inline node(i, j, k, grid, ℓx, ℓy, ℓz::Nothing) = (xnode(i, j, k, grid, ℓx, ℓy, ℓz), ynode(i, j, k, grid, ℓx, ℓy, ℓz)) -@inline node(i, j, k, grid, LX, LY::Nothing, LZ::Nothing) = tuple(xnode(i, j, k, grid, LX, LY, LZ)) -@inline node(i, j, k, grid, LX::Nothing, LY, LZ::Nothing) = tuple(ynode(i, j, k, grid, LX, LY, LZ)) -@inline node(i, j, k, grid, LX::Nothing, LY::Nothing, LZ) = tuple(znode(i, j, k, grid, LX, LY, LZ)) +@inline node(i, j, k, grid, ℓx, ℓy::Nothing, ℓz::Nothing) = tuple(xnode(i, j, k, grid, ℓx, ℓy, ℓz)) +@inline node(i, j, k, grid, ℓx::Nothing, ℓy, ℓz::Nothing) = tuple(ynode(i, j, k, grid, ℓx, ℓy, ℓz)) +@inline node(i, j, k, grid, ℓx::Nothing, ℓy::Nothing, ℓz) = tuple(znode(i, j, k, grid, ℓx, ℓy, ℓz)) xnodes(grid, ::Nothing; kwargs...) = 1:1 ynodes(grid, ::Nothing; kwargs...) = 1:1 znodes(grid, ::Nothing; kwargs...) = 1:1 """ - xnodes(grid, LX, LY, LZ, with_halos=false) + xnodes(grid, ℓx, ℓy, ℓz, with_halos=false) -Return the positions over the interior nodes on `grid` in the ``x``-direction for the location `LX`, -`LY`, `LZ`. For `Bounded` directions, `Face` nodes include the boundary points. +Return the positions over the interior nodes on `grid` in the ``x``-direction for the location `ℓx`, +`ℓy`, `ℓz`. For `Bounded` directions, `Face` nodes include the boundary points. See [`znodes`](@ref) for examples. """ -@inline xnodes(grid, LX, LY, LZ; kwargs...) = xnodes(grid, LX; kwargs...) +@inline xnodes(grid, ℓx, ℓy, ℓz; kwargs...) = xnodes(grid, ℓx; kwargs...) """ - ynodes(grid, LX, LY, LZ, with_halos=false) + ynodes(grid, ℓx, ℓy, ℓz, with_halos=false) -Return the positions over the interior nodes on `grid` in the ``y``-direction for the location `LX`, -`LY`, `LZ`. For `Bounded` directions, `Face` nodes include the boundary points. +Return the positions over the interior nodes on `grid` in the ``y``-direction for the location `ℓx`, +`ℓy`, `ℓz`. For `Bounded` directions, `Face` nodes include the boundary points. See [`znodes`](@ref) for examples. """ -@inline ynodes(grid, LX, LY, LZ; kwargs...) = ynodes(grid, LY; kwargs...) +@inline ynodes(grid, ℓx, ℓy, ℓz; kwargs...) = ynodes(grid, ℓy; kwargs...) """ - znodes(grid, LX, LY, LZ; with_halos=false) + znodes(grid, ℓx, ℓy, ℓz; with_halos=false) -Return the positions over the interior nodes on `grid` in the ``z``-direction for the location `LX`, -`LY`, `LZ`. For `Bounded` directions, `Face` nodes include the boundary points. +Return the positions over the interior nodes on `grid` in the ``z``-direction for the location `ℓx`, +`ℓy`, `ℓz`. For `Bounded` directions, `Face` nodes include the boundary points. ```jldoctest znodes julia> using Oceananigans @@ -280,14 +291,14 @@ julia> zC = znodes(horz_periodic_grid, Center(), Center(), Center(), with_halos= -1.1666666666666667:0.3333333333333333:0.16666666666666666 with indices 0:4 ``` """ -@inline znodes(grid, LX, LY, LZ; kwargs...) = znodes(grid, LZ; kwargs...) +@inline znodes(grid, ℓx, ℓy, ℓz; kwargs...) = znodes(grid, ℓz; kwargs...) """ - nodes(grid, (LX, LY, LZ); reshape=false, with_halos=false) - nodes(grid, LX, LY, LZ; reshape=false, with_halos=false) + nodes(grid, (ℓx, ℓy, ℓz); reshape=false, with_halos=false) + nodes(grid, ℓx, ℓy, ℓz; reshape=false, with_halos=false) Return a 3-tuple of views over the interior nodes -at the locations in `loc=(LX, LY, LZ)` in `x, y, z`. +at the locations in `loc=(ℓx, ℓy, ℓz)` in `x, y, z`. If `reshape=true`, the views are reshaped to 3D arrays with non-singleton dimensions 1, 2, 3 for `x, y, z`, respectively. @@ -296,14 +307,13 @@ or arrays. See [`xnodes`](@ref), [`ynodes`](@ref), and [`znodes`](@ref). """ -function nodes(grid::AbstractGrid, LX, LY, LZ; reshape=false, with_halos=false) - x = xnodes(grid, LX, LY, LZ; with_halos) - y = ynodes(grid, LX, LY, LZ; with_halos) - z = znodes(grid, LX, LY, LZ; with_halos) +function nodes(grid::AbstractGrid, ℓx, ℓy, ℓz; reshape=false, with_halos=false) + x = xnodes(grid, ℓx, ℓy, ℓz; with_halos) + y = ynodes(grid, ℓx, ℓy, ℓz; with_halos) + z = znodes(grid, ℓx, ℓy, ℓz; with_halos) if reshape N = (length(x), length(y), length(z)) - x = Base.reshape(x, N[1], 1, 1) y = Base.reshape(y, 1, N[2], 1) z = Base.reshape(z, 1, 1, N[3]) @@ -312,18 +322,23 @@ function nodes(grid::AbstractGrid, LX, LY, LZ; reshape=false, with_halos=false) return (x, y, z) end -nodes(grid::AbstractGrid, (LX, LY, LZ); reshape=false, with_halos=false) = nodes(grid, LX, LY, LZ; reshape, with_halos) +nodes(grid::AbstractGrid, (ℓx, ℓy, ℓz); reshape=false, with_halos=false) = nodes(grid, ℓx, ℓy, ℓz; reshape, with_halos) ##### ##### << Spacings >> ##### +# placeholders; see Oceananigans.Operators for x/y/zspacing definitions +function xspacing end +function yspacing end +function zspacing end + """ - xspacings(grid, LX, LY, LZ; with_halos=true) + xspacings(grid, ℓx, ℓy, ℓz; with_halos=true) -Return the spacings over the interior nodes on `grid` in the ``x``-direction for the location `LX`, -`LY`, `LZ`. For `Bounded` directions, `Face` nodes include the boundary points. +Return the spacings over the interior nodes on `grid` in the ``x``-direction for the location `ℓx`, +`ℓy`, `ℓz`. For `Bounded` directions, `Face` nodes include the boundary points. ```jldoctest xspacings julia> using Oceananigans @@ -349,14 +364,14 @@ julia> xspacings(grid, Center(), Face(), Center()) 714747.2110712599 ``` """ -@inline xspacings(grid, LX, LY, LZ; with_halos=true) = xspacings(grid, LX; with_halos) +@inline xspacings(grid, ℓx, ℓy, ℓz; with_halos=true) = xspacings(grid, ℓx; with_halos) """ - yspacings(grid, LX, LY, LZ; with_halos=true) + yspacings(grid, ℓx, ℓy, ℓz; with_halos=true) -Return the spacings over the interior nodes on `grid` in the ``y``-direction for the location `LX`, -`LY`, `LZ`. For `Bounded` directions, `Face` nodes include the boundary points. +Return the spacings over the interior nodes on `grid` in the ``y``-direction for the location `ℓx`, +`ℓy`, `ℓz`. For `Bounded` directions, `Face` nodes include the boundary points. ```jldoctest yspacings julia> using Oceananigans @@ -367,13 +382,13 @@ julia> yspacings(grid, Center(), Center(), Center()) 222389.85328911748 ``` """ -@inline yspacings(grid, LX, LY, LZ; with_halos=true) = yspacings(grid, LY; with_halos) +@inline yspacings(grid, ℓx, ℓy, ℓz; with_halos=true) = yspacings(grid, ℓy; with_halos) """ - zspacings(grid, LX, LY, LZ; with_halos=true) + zspacings(grid, ℓx, ℓy, ℓz; with_halos=true) -Return the spacings over the interior nodes on `grid` in the ``z``-direction for the location `LX`, -`LY`, `LZ`. For `Bounded` directions, `Face` nodes include the boundary points. +Return the spacings over the interior nodes on `grid` in the ``z``-direction for the location `ℓx`, +`ℓy`, `ℓz`. For `Bounded` directions, `Face` nodes include the boundary points. ```jldoctest zspacings julia> using Oceananigans @@ -384,8 +399,77 @@ julia> zspacings(grid, Center(), Center(), Center()) 10.0 ``` """ -@inline zspacings(grid, LX, LY, LZ; with_halos=true) = zspacings(grid, LZ; with_halos) +@inline zspacings(grid, ℓx, ℓy, ℓz; with_halos=true) = zspacings(grid, ℓz; with_halos) + +destantiate(::Face) = Face +destantiate(::Center) = Center + +function minimum_spacing(dir, grid, ℓx, ℓy, ℓz) + spacing = eval(Symbol(dir, :spacing)) + LX, LY, LZ = map(destantiate, (ℓx, ℓy, ℓz)) + Δ = KernelFunctionOperation{LX, LY, LZ}(spacing, grid, ℓx, ℓy, ℓz) + + return minimum(Δ) +end + +""" + minimum_xspacing(grid, ℓx, ℓy, ℓz) + minimum_xspacing(grid) = minimum_xspacing(grid, Center(), Center(), Center()) + +Return the minimum spacing for `grid` in ``x`` direction at location `ℓx, ℓy, ℓz`. + +Examples +======== +```jldoctest +julia> using Oceananigans + +julia> grid = RectilinearGrid(size=(2, 4, 8), extent=(1, 1, 1)); + +julia> minimum_xspacing(grid, Center(), Center(), Center()) +0.5 +``` +""" +minimum_xspacing(grid, ℓx, ℓy, ℓz) = minimum_spacing(:x, grid, ℓx, ℓy, ℓz) +minimum_xspacing(grid) = minimum_spacing(:x, grid, Center(), Center(), Center()) +""" + minimum_yspacing(grid, ℓx, ℓy, ℓz) + minimum_yspacing(grid) = minimum_yspacing(grid, Center(), Center(), Center()) + +Return the minimum spacing for `grid` in ``y`` direction at location `ℓx, ℓy, ℓz`. + +Examples +======== +```jldoctest +julia> using Oceananigans + +julia> grid = RectilinearGrid(size=(2, 4, 8), extent=(1, 1, 1)); + +julia> minimum_yspacing(grid, Center(), Center(), Center()) +0.25 +``` +""" +minimum_yspacing(grid, ℓx, ℓy, ℓz) = minimum_spacing(:y, grid, ℓx, ℓy, ℓz) +minimum_yspacing(grid) = minimum_spacing(:y, grid, Center(), Center(), Center()) +""" + minimum_zspacing(grid, ℓx, ℓy, ℓz) + minimum_zspacing(grid) = minimum_zspacing(grid, Center(), Center(), Center()) + +Return the minimum spacing for `grid` in ``z`` direction at location `ℓx, ℓy, ℓz`. + +Examples +======== +```jldoctest +julia> using Oceananigans + +julia> grid = RectilinearGrid(size=(2, 4, 8), extent=(1, 1, 1)); + +julia> minimum_zspacing(grid, Center(), Center(), Center()) +0.125 +``` +""" +minimum_zspacing(grid, ℓx, ℓy, ℓz) = minimum_spacing(:z, grid, ℓx, ℓy, ℓz) +minimum_zspacing(grid) = minimum_spacing(:z, grid, Center(), Center(), Center()) ##### ##### Convenience functions diff --git a/src/Grids/input_validation.jl b/src/Grids/input_validation.jl index 23ee67bc40..28a500bcaf 100644 --- a/src/Grids/input_validation.jl +++ b/src/Grids/input_validation.jl @@ -43,10 +43,10 @@ function validate_topology(topology) return topology end -function validate_size(TX, TY, TZ, size) - size = tupleit(size) - validate_tupled_argument(size, Integer, "size", topological_tuple_length(TX, TY, TZ)) - return inflate_tuple(TX, TY, TZ, size, default=1) +function validate_size(TX, TY, TZ, sz) + sz = tupleit(sz) + validate_tupled_argument(sz, Integer, "size", topological_tuple_length(TX, TY, TZ)) + return inflate_tuple(TX, TY, TZ, sz, default=1) end # Note that the default halo size is specified to be 1 in the following function. @@ -193,7 +193,7 @@ function validate_index(idx, loc, topo, N, H) end validate_index(::Colon, loc, topo, N, H) = Colon() -validate_index(idx::UnitRange, ::Type{Nothing}, topo, N, H) = UnitRange(1, 1) +validate_index(idx::UnitRange, ::Nothing, topo, N, H) = UnitRange(1, 1) function validate_index(idx::UnitRange, loc, topo, N, H) all_idx = all_indices(loc, topo, N, H) @@ -204,4 +204,7 @@ end validate_index(idx::Int, args...) = validate_index(UnitRange(idx, idx), args...) validate_indices(indices, loc, grid::AbstractGrid) = - validate_index.(indices, loc, topology(grid), size(loc, grid), halo_size(grid)) + validate_indices(indices, loc, topology(grid), size(grid, loc), halo_size(grid)) + +validate_indices(indices, loc, topo, sz, halo_sz) = + map(validate_index, indices, map(instantiate, loc), map(instantiate, topo), sz, halo_sz) diff --git a/src/Grids/latitude_longitude_grid.jl b/src/Grids/latitude_longitude_grid.jl index 13300e21cf..0a49a7e235 100644 --- a/src/Grids/latitude_longitude_grid.jl +++ b/src/Grids/latitude_longitude_grid.jl @@ -199,9 +199,9 @@ function LatitudeLongitudeGrid(architecture::AbstractArchitecture = CPU(), # it is stretched if being passed is a function or vector (as for the VerticallyStretchedRectilinearGrid) TX, TY, TZ = topology - Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, TX, Nλ, Hλ, longitude, architecture) - Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, TY, Nφ, Hφ, latitude, architecture) - Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ, Nz, Hz, z, architecture) + Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, TX(), Nλ, Hλ, longitude, architecture) + Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, TY(), Nφ, Hφ, latitude, architecture) + Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ(), Nz, Hz, z, architecture) preliminary_grid = LatitudeLongitudeGrid{TX, TY, TZ}(architecture, Nλ, Nφ, Nz, @@ -296,9 +296,9 @@ end function Base.show(io::IO, grid::LatitudeLongitudeGrid, withsummary=true) TX, TY, TZ = topology(grid) - λ₁, λ₂ = domain(topology(grid, 1), grid.Nx, grid.λᶠᵃᵃ) - φ₁, φ₂ = domain(topology(grid, 2), grid.Ny, grid.φᵃᶠᵃ) - z₁, z₂ = domain(topology(grid, 3), grid.Nz, grid.zᵃᵃᶠ) + λ₁, λ₂ = domain(TX(), size(grid, 1), grid.λᶠᵃᵃ) + φ₁, φ₂ = domain(TY(), size(grid, 2), grid.φᵃᶠᵃ) + z₁, z₂ = domain(TZ(), size(grid, 3), grid.zᵃᵃᶠ) x_summary = domain_summary(TX(), "λ", λ₁, λ₂) y_summary = domain_summary(TY(), "φ", φ₁, φ₂) @@ -546,7 +546,7 @@ function allocate_metrics(grid::LatitudeLongitudeGrid) arch = grid.architecture - if typeof(grid) <: XRegLatLonGrid + if grid isa XRegLatLonGrid offsets = grid.φᵃᶜᵃ.offsets[1] metric_size = length(grid.φᵃᶜᵃ) else @@ -560,7 +560,7 @@ function allocate_metrics(grid::LatitudeLongitudeGrid) @eval $metric = OffsetArray(arch_array($arch, $parentM), $offsets...) end - if typeof(grid) <: YRegLatLonGrid + if grid isa YRegLatLonGrid Δyᶠᶜ = FT(0.0) Δyᶜᶠ = FT(0.0) else @@ -584,18 +584,18 @@ return_metrics(::LatitudeLongitudeGrid) = (:λᶠᵃᵃ, :λᶜᵃᵃ, :φᵃᶠ ##### Grid nodes ##### -@inline xnodes(grid::LatLonGrid, LX::Face ; with_halos=false) = with_halos ? grid.λᶠᵃᵃ : view(grid.λᶠᵃᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx)) -@inline xnodes(grid::LatLonGrid, LX::Center; with_halos=false) = with_halos ? grid.λᶜᵃᵃ : view(grid.λᶜᵃᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx)) +@inline xnodes(grid::LatLonGrid, ℓx::Face ; with_halos=false) = with_halos ? grid.λᶠᵃᵃ : view(grid.λᶠᵃᵃ, interior_indices(ℓx, topology(grid, 1)(), size(grid, 1))) +@inline xnodes(grid::LatLonGrid, ℓx::Center; with_halos=false) = with_halos ? grid.λᶜᵃᵃ : view(grid.λᶜᵃᵃ, interior_indices(ℓx, topology(grid, 1)(), size(grid, 1))) -@inline ynodes(grid::LatLonGrid, LY::Face ; with_halos=false) = with_halos ? grid.φᵃᶠᵃ : view(grid.φᵃᶠᵃ, interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline ynodes(grid::LatLonGrid, LY::Center; with_halos=false) = with_halos ? grid.φᵃᶜᵃ : view(grid.φᵃᶜᵃ, interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) +@inline ynodes(grid::LatLonGrid, ℓy::Face ; with_halos=false) = with_halos ? grid.φᵃᶠᵃ : view(grid.φᵃᶠᵃ, interior_indices(ℓy, topology(grid, 2)(), size(grid, 2))) +@inline ynodes(grid::LatLonGrid, ℓy::Center; with_halos=false) = with_halos ? grid.φᵃᶜᵃ : view(grid.φᵃᶜᵃ, interior_indices(ℓy, topology(grid, 2)(), size(grid, 2))) -@inline znodes(grid::LatLonGrid, LZ::Face ; with_halos=false) = with_halos ? grid.zᵃᵃᶠ : view(grid.zᵃᵃᶠ, interior_indices(typeof(LZ), topology(grid, 3), grid.Nz)) -@inline znodes(grid::LatLonGrid, LZ::Center; with_halos=false) = with_halos ? grid.zᵃᵃᶜ : view(grid.zᵃᵃᶜ, interior_indices(typeof(LZ), topology(grid, 3), grid.Nz)) +@inline znodes(grid::LatLonGrid, ℓz::Face ; with_halos=false) = with_halos ? grid.zᵃᵃᶠ : view(grid.zᵃᵃᶠ, interior_indices(ℓz, topology(grid, 3)(), size(grid, 3))) +@inline znodes(grid::LatLonGrid, ℓz::Center; with_halos=false) = with_halos ? grid.zᵃᵃᶜ : view(grid.zᵃᵃᶜ, interior_indices(ℓz, topology(grid, 3)(), size(grid, 3))) -@inline xnodes(grid::LatLonGrid, LX, LY, LZ; with_halos=false) = xnodes(grid, LX; with_halos) -@inline ynodes(grid::LatLonGrid, LX, LY, LZ; with_halos=false) = ynodes(grid, LY; with_halos) -@inline znodes(grid::LatLonGrid, LX, LY, LZ; with_halos=false) = znodes(grid, LZ; with_halos) +@inline xnodes(grid::LatLonGrid, ℓx, ℓy, ℓz; with_halos=false) = xnodes(grid, ℓx; with_halos) +@inline ynodes(grid::LatLonGrid, ℓx, ℓy, ℓz; with_halos=false) = ynodes(grid, ℓy; with_halos) +@inline znodes(grid::LatLonGrid, ℓx, ℓy, ℓz; with_halos=false) = znodes(grid, ℓz; with_halos) @inline xnode(i, grid::LatLonGrid, ::Center) = @inbounds grid.λᶜᵃᵃ[i] @inline xnode(i, grid::LatLonGrid, ::Face) = @inbounds grid.λᶠᵃᵃ[i] @@ -606,49 +606,49 @@ return_metrics(::LatitudeLongitudeGrid) = (:λᶠᵃᵃ, :λᶜᵃᵃ, :φᵃᶠ @inline znode(k, grid::LatLonGrid, ::Center) = @inbounds grid.zᵃᵃᶜ[k] @inline znode(k, grid::LatLonGrid, ::Face) = @inbounds grid.zᵃᵃᶠ[k] -@inline xnode(i, j, k, grid::LatLonGrid, LX, LY, LZ) = xnode(i, grid, LX) -@inline ynode(i, j, k, grid::LatLonGrid, LX, LY, LZ) = ynode(j, grid, LY) -@inline znode(i, j, k, grid::LatLonGrid, LX, LY, LZ) = znode(k, grid, LZ) +@inline xnode(i, j, k, grid::LatLonGrid, ℓx, ℓy, ℓz) = xnode(i, grid, ℓx) +@inline ynode(i, j, k, grid::LatLonGrid, ℓx, ℓy, ℓz) = ynode(j, grid, ℓy) +@inline znode(i, j, k, grid::LatLonGrid, ℓx, ℓy, ℓz) = znode(k, grid, ℓz) ##### ##### Grid spacings in x, y, z (in meters) ##### -@inline xspacings(grid::LatLonGrid, LX::Center, LY::Center; with_halos=false) = with_halos ? grid.Δxᶜᶜᵃ : - view(grid.Δxᶜᶜᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline xspacings(grid::LatLonGrid, LX::Center, LY::Face; with_halos=false) = with_halos ? grid.Δxᶜᶠᵃ : - view(grid.Δxᶜᶠᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline xspacings(grid::LatLonGrid, LX::Face, LY::Center; with_halos=false) = with_halos ? grid.Δxᶠᶜᵃ : - view(grid.Δxᶠᶜᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline xspacings(grid::LatLonGrid, LX::Face, LY::Face; with_halos=false) = with_halos ? grid.Δxᶠᶠᵃ : - view(grid.Δxᶠᶠᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) - -@inline xspacings(grid::HRegLatLonGrid, LX::Center, LY::Center; with_halos=false) = with_halos ? grid.Δxᶜᶜᵃ : - view(grid.Δxᶜᶜᵃ, interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline xspacings(grid::HRegLatLonGrid, LX::Center, LY::Face; with_halos=false) = with_halos ? grid.Δxᶜᶠᵃ : - view(grid.Δxᶜᶠᵃ, interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline xspacings(grid::HRegLatLonGrid, LX::Face, LY::Center; with_halos=false) = with_halos ? grid.Δxᶠᶜᵃ : - view(grid.Δxᶠᶜᵃ, interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline xspacings(grid::HRegLatLonGrid, LX::Face, LY::Face; with_halos=false) = with_halos ? grid.Δxᶠᶠᵃ : - view(grid.Δxᶠᶠᵃ, interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) - -@inline yspacings(grid::YNonRegLatLonGrid, LX::Center, LY::Face; with_halos=false) = with_halos ? grid.Δyᶜᶠᵃ : - view(grid.Δyᶜᶠᵃ, interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline yspacings(grid::YNonRegLatLonGrid, LX::Face, LY::Center; with_halos=false) = with_halos ? grid.Δyᶠᶜᵃ : - view(grid.Δyᶠᶜᵃ, interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) - -@inline yspacings(grid::YRegLatLonGrid, LX, LY; with_halos=false) = yspacings(grid, LY; with_halos) -@inline yspacings(grid, LY::Center; kwargs...) = grid.Δyᶠᶜᵃ -@inline yspacings(grid, LY::Face; kwargs...) = grid.Δyᶜᶠᵃ - -@inline zspacings(grid::LatLonGrid, LZ::Center; with_halos=false) = with_halos ? grid.Δzᵃᵃᶜ : view(grid.Δzᵃᵃᶜ, interior_indices(typeof(LZ), topology(grid, 3), grid.Nz)) -@inline zspacings(grid::ZRegLatLonGrid, LZ::Center; with_halos=false) = grid.Δzᵃᵃᶜ -@inline zspacings(grid::LatLonGrid, LZ::Face; with_halos=false) = with_halos ? grid.Δzᵃᵃᶠ : view(grid.Δzᵃᵃᶠ, interior_indices(typeof(LZ), topology(grid, 3), grid.Nz)) -@inline zspacings(grid::ZRegLatLonGrid, LZ::Face; with_halos=false) = grid.Δzᵃᵃᶠ - -@inline xspacings(grid::LatLonGrid, LX, LY, LZ; kwargs...) = xspacings(grid, LX, LY; kwargs...) -@inline yspacings(grid::LatLonGrid, LX, LY, LZ; kwargs...) = yspacings(grid, LX, LY; kwargs...) -@inline zspacings(grid::LatLonGrid, LX, LY, LZ; kwargs...) = zspacings(grid, LZ; kwargs...) +@inline xspacings(grid::LatLonGrid, ℓx::Center, ℓy::Center; with_halos=false) = with_halos ? grid.Δxᶜᶜᵃ : + view(grid.Δxᶜᶜᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), size(grid, 2))) +@inline xspacings(grid::LatLonGrid, ℓx::Center, ℓy::Face; with_halos=false) = with_halos ? grid.Δxᶜᶠᵃ : + view(grid.Δxᶜᶠᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), size(grid, 2))) +@inline xspacings(grid::LatLonGrid, ℓx::Face, ℓy::Center; with_halos=false) = with_halos ? grid.Δxᶠᶜᵃ : + view(grid.Δxᶠᶜᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), size(grid, 2))) +@inline xspacings(grid::LatLonGrid, ℓx::Face, ℓy::Face; with_halos=false) = with_halos ? grid.Δxᶠᶠᵃ : + view(grid.Δxᶠᶠᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), size(grid, 2))) + +@inline xspacings(grid::HRegLatLonGrid, ℓx::Center, ℓy::Center; with_halos=false) = with_halos ? grid.Δxᶜᶜᵃ : + view(grid.Δxᶜᶜᵃ, interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline xspacings(grid::HRegLatLonGrid, ℓx::Center, ℓy::Face; with_halos=false) = with_halos ? grid.Δxᶜᶠᵃ : + view(grid.Δxᶜᶠᵃ, interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline xspacings(grid::HRegLatLonGrid, ℓx::Face, ℓy::Center; with_halos=false) = with_halos ? grid.Δxᶠᶜᵃ : + view(grid.Δxᶠᶜᵃ, interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline xspacings(grid::HRegLatLonGrid, ℓx::Face, ℓy::Face; with_halos=false) = with_halos ? grid.Δxᶠᶠᵃ : + view(grid.Δxᶠᶠᵃ, interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) + +@inline yspacings(grid::YNonRegLatLonGrid, ℓx::Center, ℓy::Face; with_halos=false) = with_halos ? grid.Δyᶜᶠᵃ : + view(grid.Δyᶜᶠᵃ, interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline yspacings(grid::YNonRegLatLonGrid, ℓx::Face, ℓy::Center; with_halos=false) = with_halos ? grid.Δyᶠᶜᵃ : + view(grid.Δyᶠᶜᵃ, interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) + +@inline yspacings(grid::YRegLatLonGrid, ℓx, ℓy; with_halos=false) = yspacings(grid, ℓy; with_halos) +@inline yspacings(grid, ℓy::Center; kwargs...) = grid.Δyᶠᶜᵃ +@inline yspacings(grid, ℓy::Face; kwargs...) = grid.Δyᶜᶠᵃ + +@inline zspacings(grid::LatLonGrid, ℓz::Center; with_halos=false) = with_halos ? grid.Δzᵃᵃᶜ : view(grid.Δzᵃᵃᶜ, interior_indices(ℓz, topology(grid, 3)(), size(grid, 3))) +@inline zspacings(grid::ZRegLatLonGrid, ℓz::Center; with_halos=false) = grid.Δzᵃᵃᶜ +@inline zspacings(grid::LatLonGrid, ℓz::Face; with_halos=false) = with_halos ? grid.Δzᵃᵃᶠ : view(grid.Δzᵃᵃᶠ, interior_indices(ℓz, topology(grid, 3)(), size(grid, 3))) +@inline zspacings(grid::ZRegLatLonGrid, ℓz::Face; with_halos=false) = grid.Δzᵃᵃᶠ + +@inline xspacings(grid::LatLonGrid, ℓx, ℓy, ℓz; kwargs...) = xspacings(grid, ℓx, ℓy; kwargs...) +@inline yspacings(grid::LatLonGrid, ℓx, ℓy, ℓz; kwargs...) = yspacings(grid, ℓx, ℓy; kwargs...) +@inline zspacings(grid::LatLonGrid, ℓx, ℓy, ℓz; kwargs...) = zspacings(grid, ℓz; kwargs...) min_Δx(grid::LatLonGrid) = topology(grid)[1] == Flat ? Inf : minimum(xspacings(grid, Center(), Center())) min_Δy(grid::LatLonGrid) = topology(grid)[2] == Flat ? Inf : minimum((yspacings(grid, Face(), Center()), yspacings(grid, Center(), Face()))) @@ -659,18 +659,18 @@ min_Δz(grid::LatLonGrid) = topology(grid)[3] == Flat ? Inf : minimum(zspacings( ##### Grid spacings in λ, φ (in degrees) ##### -@inline λspacings(grid::LatLonGrid, LX::Center; with_halos=false) = with_halos ? grid.Δλᶜᵃᵃ : view(grid.Δλᶜᵃᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx)) -@inline λspacings(grid::XRegLatLonGrid, LX::Center; with_halos=false) = grid.Δλᶜᵃᵃ -@inline λspacings(grid::LatLonGrid, LX::Face; with_halos=false) = with_halos ? grid.Δλᶠᵃᵃ : view(grid.Δλᶠᵃᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx)) -@inline λspacings(grid::XRegLatLonGrid, LX::Face; with_halos=false) = grid.Δλᶠᵃᵃ +@inline λspacings(grid::LatLonGrid, ℓx::Center; with_halos=false) = with_halos ? grid.Δλᶜᵃᵃ : view(grid.Δλᶜᵃᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx)) +@inline λspacings(grid::LatLonGrid, ℓx::Face; with_halos=false) = with_halos ? grid.Δλᶠᵃᵃ : view(grid.Δλᶠᵃᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx)) +@inline λspacings(grid::XRegLatLonGrid, ℓx::Center; with_halos=false) = grid.Δλᶜᵃᵃ +@inline λspacings(grid::XRegLatLonGrid, ℓx::Face; with_halos=false) = grid.Δλᶠᵃᵃ -@inline φspacings(grid::LatLonGrid, LY::Center; with_halos=false) = with_halos ? grid.Δφᵃᶜᵃ : view(grid.Δφᵃᶜᵃ, interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline φspacings(grid::YRegLatLonGrid, LY::Center; with_halos=false) = grid.Δφᵃᶜᵃ -@inline φspacings(grid::LatLonGrid, LY::Face; with_halos=false) = with_halos ? grid.Δφᵃᶠᵃ : view(grid.Δφᵃᶠᵃ, interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline φspacings(grid::YRegLatLonGrid, LY::Face; with_halos=false) = grid.Δφᵃᶠᵃ +@inline φspacings(grid::LatLonGrid, ℓy::Center; with_halos=false) = with_halos ? grid.Δφᵃᶜᵃ : view(grid.Δφᵃᶜᵃ, interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline φspacings(grid::LatLonGrid, ℓy::Face; with_halos=false) = with_halos ? grid.Δφᵃᶠᵃ : view(grid.Δφᵃᶠᵃ, interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline φspacings(grid::YRegLatLonGrid, ℓy::Center; with_halos=false) = grid.Δφᵃᶜᵃ +@inline φspacings(grid::YRegLatLonGrid, ℓy::Face; with_halos=false) = grid.Δφᵃᶠᵃ -@inline λspacings(grid::LatLonGrid, LX, LY, LZ; with_halos=false) = λspacings(grid, LX; with_halos) -@inline φspacings(grid::LatLonGrid, LX, LY, LZ; with_halos=false) = φspacings(grid, LY; with_halos) +@inline λspacings(grid::LatLonGrid, ℓx, ℓy, ℓz; with_halos=false) = λspacings(grid, ℓx; with_halos) +@inline φspacings(grid::LatLonGrid, ℓx, ℓy, ℓz; with_halos=false) = φspacings(grid, ℓy; with_halos) @inline λspacing(i, grid::LatLonGrid, ::Center) = @inbounds grid.Δλᶜᵃᵃ[i] @inline λspacing(i, grid::LatLonGrid, ::Face) = @inbounds grid.Δλᶠᵃᵃ[i] @@ -682,5 +682,5 @@ min_Δz(grid::LatLonGrid) = topology(grid)[3] == Flat ? Inf : minimum(zspacings( @inline φspacing(j, grid::YRegLatLonGrid, ::Center) = grid.Δφᵃᶜᵃ @inline φspacing(j, grid::YRegLatLonGrid, ::Face) = grid.Δφᵃᶠᵃ -@inline λspacing(i, j, k, grid::LatLonGrid, LX, LY, LZ) = λspacing(i, grid, LX) -@inline φspacing(i, j, k, grid::LatLonGrid, LX, LY, LZ) = φspacing(j, grid, LY) +@inline λspacing(i, j, k, grid::LatLonGrid, ℓx, ℓy, ℓz) = λspacing(i, grid, ℓx) +@inline φspacing(i, j, k, grid::LatLonGrid, ℓx, ℓy, ℓz) = φspacing(j, grid, ℓy) diff --git a/src/Grids/new_data.jl b/src/Grids/new_data.jl index 10af1fb530..d96f0f7365 100644 --- a/src/Grids/new_data.jl +++ b/src/Grids/new_data.jl @@ -2,7 +2,6 @@ using Oceananigans.Grids: total_length, topology using OffsetArrays: OffsetArray -default_indices(n) = Tuple(Colon() for i=1:n) ##### ##### Creating offset arrays for field data by dispatching on architecture. @@ -19,19 +18,24 @@ offset_indices(loc, topo, N, H=0) = 1 - H : N + H Return a range of indices for a field located at cell `Face`s along a grid dimension which is `Bounded` and has length `N` and with halo points `H`. """ -offset_indices(::Type{Face}, ::Type{<:BoundedTopology}, N, H=0) = 1 - H : N + H + 1 +offset_indices(::Face, ::BoundedTopology, N, H=0) = 1 - H : N + H + 1 """ Return a range of indices for a field along a 'reduced' dimension. """ -offset_indices(::Type{Nothing}, topo, N, H=0) = 1 : 1 +offset_indices(::Nothing, topo, N, H=0) = 1:1 -offset_indices(L, T, N, H, ::Colon) = offset_indices(L, T, N, H) -offset_indices(L, T, N, H, r::UnitRange) = r -offset_indices(::Type{Nothing}, T, N, H, ::UnitRange) = 1:1 +offset_indices(ℓ, topo, N, H, ::Colon) = offset_indices(ℓ, topo, N, H) +offset_indices(ℓ, topo, N, H, r::UnitRange) = r +offset_indices(::Nothing, topo, N, H, ::UnitRange) = 1:1 + +instantiate(T::Type) = T() +instantiate(t) = t function offset_data(underlying_data::AbstractArray, loc, topo, N, H, indices=default_indices(length(loc))) - ii = offset_indices.(loc, topo, N, H, indices) + loc = map(instantiate, loc) + topo = map(instantiate, topo) + ii = map(offset_indices, loc, topo, N, H, indices) # Add extra indices for arrays of higher dimension than loc, topo, etc. extra_ii = Tuple(axes(underlying_data, d) for d in length(ii)+1:ndims(underlying_data)) return OffsetArray(underlying_data, ii..., extra_ii...) @@ -48,18 +52,19 @@ offset_data(underlying_data::AbstractArray, grid::AbstractGrid, loc, indices=def offset_data(underlying_data, loc, topology(grid), size(grid), halo_size(grid), indices) """ - new_data(FT, grid, loc, indices) + new_data(FT, arch, loc, topo, sz, halo_sz, indices) Returns an `OffsetArray` of zeros of float type `FT` on `arch`itecture, with indices corresponding to a field on a `grid` of `size(grid)` and located at `loc`. """ -function new_data(FT::DataType, grid::AbstractGrid, loc, indices=default_indices(length(loc))) - arch = architecture(grid) - Tx, Ty, Tz = total_size(loc, grid, indices) +function new_data(FT::DataType, arch, loc, topo, sz, halo_sz, indices=default_indices(length(loc))) + Tx, Ty, Tz = total_size(loc, topo, sz, halo_sz, indices) underlying_data = zeros(FT, arch, Tx, Ty, Tz) - indices = validate_indices(indices, loc, grid) - return offset_data(underlying_data, grid, loc, indices) + indices = validate_indices(indices, loc, topo, sz, halo_sz) + return offset_data(underlying_data, loc, topo, sz, halo_sz, indices) end -new_data(grid, loc, indices=default_indices) = new_data(eltype(grid), grid, loc, indices) +new_data(FT::DataType, grid::AbstractGrid, loc, indices=default_indices(length(loc))) = + new_data(FT, architecture(grid), loc, topology(grid), size(grid), halo_size(grid), indices) +new_data(grid::AbstractGrid, loc, indices=default_indices) = new_data(eltype(grid), grid, loc, indices) diff --git a/src/Grids/orthogonal_spherical_shell_grid.jl b/src/Grids/orthogonal_spherical_shell_grid.jl index 38b5383acc..8eb30239f0 100644 --- a/src/Grids/orthogonal_spherical_shell_grid.jl +++ b/src/Grids/orthogonal_spherical_shell_grid.jl @@ -598,17 +598,17 @@ OrthogonalSphericalShellGrid(FT::DataType; kwargs...) = OrthogonalSphericalShell function load_and_offset_cubed_sphere_data(file, FT, arch, field_name, loc, topo, N, H) - ii = interior_indices(loc[1], topo[1], N[1]) - jj = interior_indices(loc[2], topo[2], N[2]) + ii = interior_indices(loc[1](), topo[1](), N[1]) + jj = interior_indices(loc[2](), topo[2](), N[2]) interior_data = arch_array(arch, file[field_name][ii, jj]) underlying_data = zeros(FT, arch, - total_length(loc[1], topo[1], N[1], H[1]), - total_length(loc[2], topo[2], N[2], H[2])) + total_length(loc[1](), topo[1](), N[1], H[1]), + total_length(loc[2](), topo[2](), N[2], H[2])) - ip = interior_parent_indices(loc[1], topo[1], N[1], H[1]) - jp = interior_parent_indices(loc[2], topo[2], N[2], H[2]) + ip = interior_parent_indices(loc[1](), topo[1](), N[1], H[1]) + jp = interior_parent_indices(loc[2](), topo[2](), N[2], H[2]) view(underlying_data, ip, jp) .= interior_data @@ -671,10 +671,10 @@ function OrthogonalSphericalShellGrid(filepath::AbstractString, architecture = C Azᶠᶠᵃ = load_and_offset_cubed_sphere_data(file, FT, architecture, "Azᶠᶠᵃ", loc_ff, topo_bbb, N, H) ## Maybe we won't need these? - Txᶠᶜ = total_length(loc_fc[1], topology[1], N[1], H[1]) - Txᶜᶠ = total_length(loc_cf[1], topology[1], N[1], H[1]) - Tyᶠᶜ = total_length(loc_fc[2], topology[2], N[2], H[2]) - Tyᶜᶠ = total_length(loc_cf[2], topology[2], N[2], H[2]) + Txᶠᶜ = total_length(loc_fc[1](), topology[1](), N[1], H[1]) + Txᶜᶠ = total_length(loc_cf[1](), topology[1](), N[1], H[1]) + Tyᶠᶜ = total_length(loc_fc[2](), topology[2](), N[2], H[2]) + Tyᶜᶠ = total_length(loc_cf[2](), topology[2](), N[2], H[2]) λᶠᶜᵃ = offset_data(zeros(FT, architecture, Txᶠᶜ, Tyᶠᶜ), loc_fc, topology[1:2], N[1:2], H[1:2]) λᶜᶠᵃ = offset_data(zeros(FT, architecture, Txᶜᶠ, Tyᶜᶠ), loc_cf, topology[1:2], N[1:2], H[1:2]) @@ -777,13 +777,69 @@ function Base.summary(grid::OrthogonalSphericalShellGrid) " and ", metric_computation) end +""" + get_center_and_extents_of_shell(grid::OSSG) + +Return the latitude-longitude coordinates of the center of the shell `(λ_center, φ_center)` +and also the longitudinal and latitudinal extend of the shell `(extend_λ, extend_φ)`. +""" +function get_center_and_extents_of_shell(grid::OSSG) + Nx, Ny, _ = size(grid) + + # find the indices that correspond to the center of the shell + # ÷ ensures that expressions below work for both odd and even + i_center = Nx÷2 + 1 + j_center = Ny÷2 + 1 + + if mod(Nx, 2) == 0 + ℓx = Face() + elseif mod(Nx, 2) == 1 + ℓx = Center() + end + + if mod(Ny, 2) == 0 + ℓy = Face() + elseif mod(Ny, 2) == 1 + ℓy = Center() + end + + # latitude and longitudes of the shell's center + λ_center = xnode(i_center, j_center, 1, grid, ℓx, ℓy, Center()) + φ_center = ynode(i_center, j_center, 1, grid, ℓx, ℓy, Center()) + + # the Δλ, Δφ are approximate if ξ, η are not symmetric about 0 + if mod(Ny, 2) == 0 + extend_λ = rad2deg(sum(grid.Δxᶜᶠᵃ[:, j_center])) / grid.radius + elseif mod(Ny, 2) == 1 + extend_λ = rad2deg(sum(grid.Δxᶜᶜᵃ[:, j_center])) / grid.radius + end + + if mod(Nx, 2) == 0 + extend_φ = rad2deg(sum(grid.Δyᶠᶜᵃ[i_center, :])) / grid.radius + elseif mod(Nx, 2) == 1 + extend_φ = rad2deg(sum(grid.Δyᶜᶜᵃ[i_center, :])) / grid.radius + end + + return (λ_center, φ_center), (extend_λ, extend_φ) +end + function Base.show(io::IO, grid::OrthogonalSphericalShellGrid, withsummary=true) TX, TY, TZ = topology(grid) Nx, Ny, Nz = size(grid) λ₁, λ₂ = minimum(grid.λᶠᶠᵃ[1:Nx+1, 1:Ny+1]), maximum(grid.λᶠᶠᵃ[1:Nx+1, 1:Ny+1]) φ₁, φ₂ = minimum(grid.φᶠᶠᵃ[1:Nx+1, 1:Ny+1]), maximum(grid.φᶠᶠᵃ[1:Nx+1, 1:Ny+1]) - z₁, z₂ = domain(topology(grid, 3), grid.Nz, grid.zᵃᵃᶠ) + z₁, z₂ = domain(topology(grid, 3)(), Nz, grid.zᵃᵃᶠ) + + (λc, φc), (Δλ, Δφ) = get_center_and_extents_of_shell(grid) + + λc = round(λc, digits=4) + φc = round(φc, digits=4) + + center_str = "centered at (λ, φ) = (" * prettysummary(λc) * ", " * prettysummary(φc) * ")" + + if abs(φc) ≈ 90; center_str = "centered at: North Pole, (λ, φ) = (" * prettysummary(λc) * ", " * prettysummary(φc) * ")"; end + if abs(φc) ≈ -90; center_str = "centered at: South Pole, (λ, φ) = (" * prettysummary(λc) * ", " * prettysummary(φc) * ")"; end x_summary = domain_summary(TX(), "λ", λ₁, λ₂) y_summary = domain_summary(TY(), "φ", φ₁, φ₂) @@ -808,30 +864,30 @@ end ##### Grid nodes ##### -@inline xnodes(grid::OSSG, LX::Face, LY::Face, ; with_halos=false) = with_halos ? grid.λᶠᶠᵃ : - view(grid.λᶠᶠᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline xnodes(grid::OSSG, LX::Face, LY::Center; with_halos=false) = with_halos ? grid.λᶠᶜᵃ : - view(grid.λᶠᶜᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline xnodes(grid::OSSG, LX::Center, LY::Face, ; with_halos=false) = with_halos ? grid.λᶜᶠᵃ : - view(grid.λᶜᶠᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline xnodes(grid::OSSG, LX::Center, LY::Center; with_halos=false) = with_halos ? grid.λᶜᶜᵃ : - view(grid.λᶜᶜᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) - -@inline ynodes(grid::OSSG, LX::Face, LY::Face, ; with_halos=false) = with_halos ? grid.φᶠᶠᵃ : - view(grid.φᶠᶠᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline ynodes(grid::OSSG, LX::Face, LY::Center; with_halos=false) = with_halos ? grid.φᶠᶜᵃ : - view(grid.φᶠᶜᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline ynodes(grid::OSSG, LX::Center, LY::Face, ; with_halos=false) = with_halos ? grid.φᶜᶠᵃ : - view(grid.φᶜᶠᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline ynodes(grid::OSSG, LX::Center, LY::Center; with_halos=false) = with_halos ? grid.φᶜᶜᵃ : - view(grid.φᶜᶜᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) - -@inline znodes(grid::OSSG, LZ::Face ; with_halos=false) = with_halos ? grid.zᵃᵃᶠ : view(grid.zᵃᵃᶠ, interior_indices(typeof(LZ), topology(grid, 3), grid.Nz)) -@inline znodes(grid::OSSG, LZ::Center; with_halos=false) = with_halos ? grid.zᵃᵃᶜ : view(grid.zᵃᵃᶜ, interior_indices(typeof(LZ), topology(grid, 3), grid.Nz)) - -@inline xnodes(grid::OSSG, LX, LY, LZ; with_halos=false) = xnodes(grid, LX, LY; with_halos) -@inline ynodes(grid::OSSG, LX, LY, LZ; with_halos=false) = ynodes(grid, LX, LY; with_halos) -@inline znodes(grid::OSSG, LX, LY, LZ; with_halos=false) = znodes(grid, LZ ; with_halos) +@inline xnodes(grid::OSSG, ℓx::Face, ℓy::Face, ; with_halos=false) = with_halos ? grid.λᶠᶠᵃ : + view(grid.λᶠᶠᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline xnodes(grid::OSSG, ℓx::Face, ℓy::Center; with_halos=false) = with_halos ? grid.λᶠᶜᵃ : + view(grid.λᶠᶜᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline xnodes(grid::OSSG, ℓx::Center, ℓy::Face, ; with_halos=false) = with_halos ? grid.λᶜᶠᵃ : + view(grid.λᶜᶠᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline xnodes(grid::OSSG, ℓx::Center, ℓy::Center; with_halos=false) = with_halos ? grid.λᶜᶜᵃ : + view(grid.λᶜᶜᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) + +@inline ynodes(grid::OSSG, ℓx::Face, ℓy::Face, ; with_halos=false) = with_halos ? grid.φᶠᶠᵃ : + view(grid.φᶠᶠᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline ynodes(grid::OSSG, ℓx::Face, ℓy::Center; with_halos=false) = with_halos ? grid.φᶠᶜᵃ : + view(grid.φᶠᶜᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline ynodes(grid::OSSG, ℓx::Center, ℓy::Face, ; with_halos=false) = with_halos ? grid.φᶜᶠᵃ : + view(grid.φᶜᶠᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline ynodes(grid::OSSG, ℓx::Center, ℓy::Center; with_halos=false) = with_halos ? grid.φᶜᶜᵃ : + view(grid.φᶜᶜᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) + +@inline znodes(grid::OSSG, ℓz::Face ; with_halos=false) = with_halos ? grid.zᵃᵃᶠ : view(grid.zᵃᵃᶠ, interior_indices(ℓz, topology(grid, 3)(), grid.Nz)) +@inline znodes(grid::OSSG, ℓz::Center; with_halos=false) = with_halos ? grid.zᵃᵃᶜ : view(grid.zᵃᵃᶜ, interior_indices(ℓz, topology(grid, 3)(), grid.Nz)) + +@inline xnodes(grid::OSSG, ℓx, ℓy, ℓz; with_halos=false) = xnodes(grid, ℓx, ℓy; with_halos) +@inline ynodes(grid::OSSG, ℓx, ℓy, ℓz; with_halos=false) = ynodes(grid, ℓx, ℓy; with_halos) +@inline znodes(grid::OSSG, ℓx, ℓy, ℓz; with_halos=false) = znodes(grid, ℓz ; with_halos) @inline xnode(i, j, grid::OSSG, ::Center, ::Center) = @inbounds grid.λᶜᶜᵃ[i, j] @inline xnode(i, j, grid::OSSG, ::Face , ::Center) = @inbounds grid.λᶠᶜᵃ[i, j] @@ -846,38 +902,37 @@ end @inline znode(k, grid::OSSG, ::Center) = @inbounds grid.zᵃᵃᶜ[k] @inline znode(k, grid::OSSG, ::Face ) = @inbounds grid.zᵃᵃᶠ[k] -@inline xnode(i, j, k, grid::OSSG, LX, LY, LZ) = xnode(i, j, grid, LX, LY) -@inline ynode(i, j, k, grid::OSSG, LX, LY, LZ) = ynode(i, j, grid, LX, LY) -@inline znode(i, j, k, grid::OSSG, LX, LY, LZ) = znode(k, grid, LZ) +@inline xnode(i, j, k, grid::OSSG, ℓx, ℓy, ℓz) = xnode(i, j, grid, ℓx, ℓy) +@inline ynode(i, j, k, grid::OSSG, ℓx, ℓy, ℓz) = ynode(i, j, grid, ℓx, ℓy) +@inline znode(i, j, k, grid::OSSG, ℓx, ℓy, ℓz) = znode(k, grid, ℓz) ##### ##### Grid spacings in x, y, z (in meters) ##### -@inline xspacings(grid::OSSG, LX::Center, LY::Center; with_halos=false) = - with_halos ? grid.Δxᶜᶜᵃ : view(grid.Δxᶜᶜᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline xspacings(grid::OSSG, LX::Face , LY::Center; with_halos=false) = - with_halos ? grid.Δxᶠᶜᵃ : view(grid.Δxᶠᶜᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline xspacings(grid::OSSG, LX::Center, LY::Face ; with_halos=false) = - with_halos ? grid.Δxᶜᶠᵃ : view(grid.Δxᶜᶠᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline xspacings(grid::OSSG, LX::Face , LY::Face ; with_halos=false) = - with_halos ? grid.Δxᶠᶠᵃ : view(grid.Δxᶠᶠᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) - -@inline yspacings(grid::OSSG, LX::Center, LY::Center; with_halos=false) = - with_halos ? grid.Δyᶜᶜᵃ : view(grid.Δyᶜᶜᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline yspacings(grid::OSSG, LX::Face , LY::Center; with_halos=false) = - with_halos ? grid.Δyᶠᶜᵃ : view(grid.Δyᶠᶜᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline yspacings(grid::OSSG, LX::Center, LY::Face ; with_halos=false) = - with_halos ? grid.Δyᶜᶠᵃ : view(grid.Δyᶜᶠᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline yspacings(grid::OSSG, LX::Face , LY::Face ; with_halos=false) = - with_halos ? grid.Δyᶠᶠᵃ : view(grid.Δyᶠᶠᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx), interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) - -@inline zspacings(grid::OSSG, LZ::Center; with_halos=false) = grid.Δz -@inline zspacings(grid::OSSG, LZ::Face; with_halos=false) = grid.Δz - -@inline xspacings(grid::OSSG, LX, LY, LZ; with_halos=false) = xspacings(grid, LX, LY; with_halos) -@inline yspacings(grid::OSSG, LX, LY, LZ; with_halos=false) = yspacings(grid, LX, LY; with_halos) -@inline zspacings(grid::OSSG, LX, LY, LZ; with_halos=false) = zspacings(grid, LZ; with_halos) +@inline xspacings(grid::OSSG, ℓx::Center, ℓy::Center; with_halos=false) = + with_halos ? grid.Δxᶜᶜᵃ : view(grid.Δxᶜᶜᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline xspacings(grid::OSSG, ℓx::Face , ℓy::Center; with_halos=false) = + with_halos ? grid.Δxᶠᶜᵃ : view(grid.Δxᶠᶜᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline xspacings(grid::OSSG, ℓx::Center, ℓy::Face ; with_halos=false) = + with_halos ? grid.Δxᶜᶠᵃ : view(grid.Δxᶜᶠᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline xspacings(grid::OSSG, ℓx::Face , ℓy::Face ; with_halos=false) = + with_halos ? grid.Δxᶠᶠᵃ : view(grid.Δxᶠᶠᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) + +@inline yspacings(grid::OSSG, ℓx::Center, ℓy::Center; with_halos=false) = + with_halos ? grid.Δyᶜᶜᵃ : view(grid.Δyᶜᶜᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline yspacings(grid::OSSG, ℓx::Face , ℓy::Center; with_halos=false) = + with_halos ? grid.Δyᶠᶜᵃ : view(grid.Δyᶠᶜᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline yspacings(grid::OSSG, ℓx::Center, ℓy::Face ; with_halos=false) = + with_halos ? grid.Δyᶜᶠᵃ : view(grid.Δyᶜᶠᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) +@inline yspacings(grid::OSSG, ℓx::Face , ℓy::Face ; with_halos=false) = + with_halos ? grid.Δyᶠᶠᵃ : view(grid.Δyᶠᶠᵃ, interior_indices(ℓx, topology(grid, 1)(), grid.Nx), interior_indices(ℓy, topology(grid, 2)(), grid.Ny)) + +@inline zspacings(grid::OSSG, ℓz; with_halos=false) = grid.Δz + +@inline xspacings(grid::OSSG, ℓx, ℓy, ℓz; with_halos=false) = xspacings(grid, ℓx, ℓy; with_halos) +@inline yspacings(grid::OSSG, ℓx, ℓy, ℓz; with_halos=false) = yspacings(grid, ℓx, ℓy; with_halos) +@inline zspacings(grid::OSSG, ℓx, ℓy, ℓz; with_halos=false) = zspacings(grid, ℓz; with_halos) min_Δx(grid::OSSG) = topology(grid)[1] == Flat ? Inf : minimum(xspacings(grid, Center(), Center())) min_Δy(grid::OSSG) = topology(grid)[2] == Flat ? Inf : minimum(yspacings(grid, Center(), Center())) diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl index 0ddb8d9dd5..024c77f7b5 100644 --- a/src/Grids/rectilinear_grid.jl +++ b/src/Grids/rectilinear_grid.jl @@ -265,9 +265,9 @@ function RectilinearGrid(architecture::AbstractArchitecture = CPU(), Nx, Ny, Nz = size Hx, Hy, Hz = halo - Lx, xᶠᵃᵃ, xᶜᵃᵃ, Δxᶠᵃᵃ, Δxᶜᵃᵃ = generate_coordinate(FT, topology[1], Nx, Hx, x, architecture) - Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, topology[2], Ny, Hy, y, architecture) - Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, topology[3], Nz, Hz, z, architecture) + Lx, xᶠᵃᵃ, xᶜᵃᵃ, Δxᶠᵃᵃ, Δxᶜᵃᵃ = generate_coordinate(FT, topology[1](), Nx, Hx, x, architecture) + Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, topology[2](), Ny, Hy, y, architecture) + Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, topology[3](), Nz, Hz, z, architecture) return RectilinearGrid{TX, TY, TZ}(architecture, Nx, Ny, Nz, @@ -294,9 +294,9 @@ end ##### Showing grids ##### -x_domain(grid::RectilinearGrid) = domain(topology(grid, 1), grid.Nx, grid.xᶠᵃᵃ) -y_domain(grid::RectilinearGrid) = domain(topology(grid, 2), grid.Ny, grid.yᵃᶠᵃ) -z_domain(grid::RectilinearGrid) = domain(topology(grid, 3), grid.Nz, grid.zᵃᵃᶠ) +x_domain(grid::RectilinearGrid) = domain(topology(grid, 1)(), grid.Nx, grid.xᶠᵃᵃ) +y_domain(grid::RectilinearGrid) = domain(topology(grid, 2)(), grid.Ny, grid.yᵃᶠᵃ) +z_domain(grid::RectilinearGrid) = domain(topology(grid, 3)(), grid.Nz, grid.zᵃᵃᶠ) # architecture = CPU() default, assuming that a DataType positional arg # is specifying the floating point type. @@ -314,9 +314,9 @@ end function Base.show(io::IO, grid::RectilinearGrid, withsummary=true) TX, TY, TZ = topology(grid) - x₁, x₂ = domain(topology(grid, 1), grid.Nx, grid.xᶠᵃᵃ) - y₁, y₂ = domain(topology(grid, 2), grid.Ny, grid.yᵃᶠᵃ) - z₁, z₂ = domain(topology(grid, 3), grid.Nz, grid.zᵃᵃᶠ) + x₁, x₂ = domain(TX(), grid.Nx, grid.xᶠᵃᵃ) + y₁, y₂ = domain(TY(), grid.Ny, grid.yᵃᶠᵃ) + z₁, z₂ = domain(TZ(), grid.Nz, grid.zᵃᵃᶠ) x_summary = domain_summary(TX(), "x", x₁, x₂) y_summary = domain_summary(TY(), "y", y₁, y₂) @@ -410,18 +410,18 @@ return_metrics(::RectilinearGrid) = (:xᶠᵃᵃ, :xᶜᵃᵃ, :yᵃᶠᵃ, :y ##### Grid nodes ##### -@inline xnodes(grid::RectilinearGrid, LX::Face ; with_halos=false) = with_halos ? grid.xᶠᵃᵃ : view(grid.xᶠᵃᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx)) -@inline xnodes(grid::RectilinearGrid, LX::Center; with_halos=false) = with_halos ? grid.xᶜᵃᵃ : view(grid.xᶜᵃᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx)) +@inline xnodes(grid::RectilinearGrid, ℓx::Face ; with_halos=false) = with_halos ? grid.xᶠᵃᵃ : view(grid.xᶠᵃᵃ, interior_indices(ℓx, topology(grid, 1)(), size(grid, 1))) +@inline xnodes(grid::RectilinearGrid, ℓx::Center; with_halos=false) = with_halos ? grid.xᶜᵃᵃ : view(grid.xᶜᵃᵃ, interior_indices(ℓx, topology(grid, 1)(), size(grid, 1))) -@inline ynodes(grid::RectilinearGrid, LY::Face ; with_halos=false) = with_halos ? grid.yᵃᶠᵃ : view(grid.yᵃᶠᵃ, interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline ynodes(grid::RectilinearGrid, LY::Center; with_halos=false) = with_halos ? grid.yᵃᶜᵃ : view(grid.yᵃᶜᵃ, interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) +@inline ynodes(grid::RectilinearGrid, ℓy::Face ; with_halos=false) = with_halos ? grid.yᵃᶠᵃ : view(grid.yᵃᶠᵃ, interior_indices(ℓy, topology(grid, 2)(), size(grid, 2))) +@inline ynodes(grid::RectilinearGrid, ℓy::Center; with_halos=false) = with_halos ? grid.yᵃᶜᵃ : view(grid.yᵃᶜᵃ, interior_indices(ℓy, topology(grid, 2)(), size(grid, 2))) -@inline znodes(grid::RectilinearGrid, LZ::Face ; with_halos=false) = with_halos ? grid.zᵃᵃᶠ : view(grid.zᵃᵃᶠ, interior_indices(typeof(LZ), topology(grid, 3), grid.Nz)) -@inline znodes(grid::RectilinearGrid, LZ::Center; with_halos=false) = with_halos ? grid.zᵃᵃᶜ : view(grid.zᵃᵃᶜ, interior_indices(typeof(LZ), topology(grid, 3), grid.Nz)) +@inline znodes(grid::RectilinearGrid, ℓz::Face ; with_halos=false) = with_halos ? grid.zᵃᵃᶠ : view(grid.zᵃᵃᶠ, interior_indices(ℓz, topology(grid, 3)(), size(grid, 3))) +@inline znodes(grid::RectilinearGrid, ℓz::Center; with_halos=false) = with_halos ? grid.zᵃᵃᶜ : view(grid.zᵃᵃᶜ, interior_indices(ℓz, topology(grid, 3)(), size(grid, 3))) -@inline xnodes(grid::RectilinearGrid, LX, LY, LZ; with_halos=false) = xnodes(grid, LX; with_halos) -@inline ynodes(grid::RectilinearGrid, LX, LY, LZ; with_halos=false) = ynodes(grid, LY; with_halos) -@inline znodes(grid::RectilinearGrid, LX, LY, LZ; with_halos=false) = znodes(grid, LZ; with_halos) +@inline xnodes(grid::RectilinearGrid, ℓx, ℓy, ℓz; with_halos=false) = xnodes(grid, ℓx; with_halos) +@inline ynodes(grid::RectilinearGrid, ℓx, ℓy, ℓz; with_halos=false) = ynodes(grid, ℓy; with_halos) +@inline znodes(grid::RectilinearGrid, ℓx, ℓy, ℓz; with_halos=false) = znodes(grid, ℓz; with_halos) @inline xnode(i, grid::RectilinearGrid, ::Center) = @inbounds grid.xᶜᵃᵃ[i] @inline xnode(i, grid::RectilinearGrid, ::Face) = @inbounds grid.xᶠᵃᵃ[i] @@ -432,33 +432,33 @@ return_metrics(::RectilinearGrid) = (:xᶠᵃᵃ, :xᶜᵃᵃ, :yᵃᶠᵃ, :y @inline znode(k, grid::RectilinearGrid, ::Center) = @inbounds grid.zᵃᵃᶜ[k] @inline znode(k, grid::RectilinearGrid, ::Face) = @inbounds grid.zᵃᵃᶠ[k] -@inline xnode(i, j, k, grid::RectilinearGrid, LX, LY, LZ) = xnode(i, grid, LX) -@inline ynode(i, j, k, grid::RectilinearGrid, LX, LY, LZ) = ynode(j, grid, LY) -@inline znode(i, j, k, grid::RectilinearGrid, LX, LY, LZ) = znode(k, grid, LZ) +@inline xnode(i, j, k, grid::RectilinearGrid, ℓx, ℓy, ℓz) = xnode(i, grid, ℓx) +@inline ynode(i, j, k, grid::RectilinearGrid, ℓx, ℓy, ℓz) = ynode(j, grid, ℓy) +@inline znode(i, j, k, grid::RectilinearGrid, ℓx, ℓy, ℓz) = znode(k, grid, ℓz) ##### ##### Grid spacings ##### -@inline xspacings(grid::RectilinearGrid, LX::Center; with_halos=false) = with_halos ? grid.Δxᶜᵃᵃ : view(grid.Δxᶜᵃᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx)) -@inline xspacings(grid::XRegRectilinearGrid, LX::Center; with_halos=false) = grid.Δxᶜᵃᵃ -@inline xspacings(grid::RectilinearGrid, LX::Face; with_halos=false) = with_halos ? grid.Δxᶠᵃᵃ : view(grid.Δxᶠᵃᵃ, interior_indices(typeof(LX), topology(grid, 1), grid.Nx)) -@inline xspacings(grid::XRegRectilinearGrid, LX::Face; with_halos=false) = grid.Δxᶠᵃᵃ +@inline xspacings(grid::RectilinearGrid, ℓx::Center; with_halos=false) = with_halos ? grid.Δxᶜᵃᵃ : view(grid.Δxᶜᵃᵃ, interior_indices(ℓx, topology(grid, 1)(), size(grid, 1))) +@inline xspacings(grid::XRegRectilinearGrid, ℓx::Center; with_halos=false) = grid.Δxᶜᵃᵃ +@inline xspacings(grid::RectilinearGrid, ℓx::Face; with_halos=false) = with_halos ? grid.Δxᶠᵃᵃ : view(grid.Δxᶠᵃᵃ, interior_indices(ℓx, topology(grid, 1)(), size(grid, 1))) +@inline xspacings(grid::XRegRectilinearGrid, ℓx::Face; with_halos=false) = grid.Δxᶠᵃᵃ -@inline yspacings(grid::RectilinearGrid, LY::Center; with_halos=false) = with_halos ? grid.Δyᵃᶜᵃ : view(grid.Δyᵃᶜᵃ, interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline yspacings(grid::YRegRectilinearGrid, LY::Center; with_halos=false) = grid.Δyᵃᶜᵃ -@inline yspacings(grid::RectilinearGrid, LY::Face; with_halos=false) = with_halos ? grid.Δyᵃᶠᵃ : view(grid.Δyᵃᶠᵃ, interior_indices(typeof(LY), topology(grid, 2), grid.Ny)) -@inline yspacings(grid::YRegRectilinearGrid, LY::Face; with_halos=false) = grid.Δyᵃᶠᵃ +@inline yspacings(grid::RectilinearGrid, ℓy::Center; with_halos=false) = with_halos ? grid.Δyᵃᶜᵃ : view(grid.Δyᵃᶜᵃ, interior_indices(ℓy, topology(grid, 2)(), size(grid, 2))) +@inline yspacings(grid::YRegRectilinearGrid, ℓy::Center; with_halos=false) = grid.Δyᵃᶜᵃ +@inline yspacings(grid::RectilinearGrid, ℓy::Face; with_halos=false) = with_halos ? grid.Δyᵃᶠᵃ : view(grid.Δyᵃᶠᵃ, interior_indices(ℓy, topology(grid, 2)(), size(grid, 2))) +@inline yspacings(grid::YRegRectilinearGrid, ℓy::Face; with_halos=false) = grid.Δyᵃᶠᵃ -@inline zspacings(grid::RectilinearGrid, LZ::Center; with_halos=false) = with_halos ? grid.Δzᵃᵃᶜ : view(grid.Δzᵃᵃᶜ, interior_indices(typeof(LZ), topology(grid, 3), grid.Nz)) -@inline zspacings(grid::ZRegRectilinearGrid, LZ::Center; with_halos=false) = grid.Δzᵃᵃᶜ -@inline zspacings(grid::RectilinearGrid, LZ::Face; with_halos=false) = with_halos ? grid.Δzᵃᵃᶠ : view(grid.Δzᵃᵃᶠ, interior_indices(typeof(LZ), topology(grid, 3), grid.Nz)) -@inline zspacings(grid::ZRegRectilinearGrid, LZ::Face; with_halos=false) = grid.Δzᵃᵃᶠ +@inline zspacings(grid::RectilinearGrid, ℓz::Center; with_halos=false) = with_halos ? grid.Δzᵃᵃᶜ : view(grid.Δzᵃᵃᶜ, interior_indices(ℓz, topology(grid, 3)(), size(grid, 3))) +@inline zspacings(grid::ZRegRectilinearGrid, ℓz::Center; with_halos=false) = grid.Δzᵃᵃᶜ +@inline zspacings(grid::RectilinearGrid, ℓz::Face; with_halos=false) = with_halos ? grid.Δzᵃᵃᶠ : view(grid.Δzᵃᵃᶠ, interior_indices(ℓz, topology(grid, 3)(), size(grid, 3))) +@inline zspacings(grid::ZRegRectilinearGrid, ℓz::Face; with_halos=false) = grid.Δzᵃᵃᶠ -@inline xspacings(grid::RectilinearGrid, LX, LY, LZ; kwargs...) = xspacings(grid, LX; kwargs...) -@inline yspacings(grid::RectilinearGrid, LX, LY, LZ; kwargs...) = yspacings(grid, LY; kwargs...) -@inline zspacings(grid::RectilinearGrid, LX, LY, LZ; kwargs...) = zspacings(grid, LZ; kwargs...) +@inline xspacings(grid::RectilinearGrid, ℓx, ℓy, ℓz; kwargs...) = xspacings(grid, ℓx; kwargs...) +@inline yspacings(grid::RectilinearGrid, ℓx, ℓy, ℓz; kwargs...) = yspacings(grid, ℓy; kwargs...) +@inline zspacings(grid::RectilinearGrid, ℓx, ℓy, ℓz; kwargs...) = zspacings(grid, ℓz; kwargs...) min_Δx(grid::RectilinearGrid) = topology(grid)[1] == Flat ? Inf : minimum(xspacings(grid, Center())) min_Δy(grid::RectilinearGrid) = topology(grid)[2] == Flat ? Inf : minimum(yspacings(grid, Center())) diff --git a/src/ImmersedBoundaries/ImmersedBoundaries.jl b/src/ImmersedBoundaries/ImmersedBoundaries.jl index ee23b55ed8..268a8ceeec 100644 --- a/src/ImmersedBoundaries/ImmersedBoundaries.jl +++ b/src/ImmersedBoundaries/ImmersedBoundaries.jl @@ -42,7 +42,7 @@ using Oceananigans.Advection: advective_tracer_flux_z import Base: show, summary -import Oceananigans.Utils: cell_advection_timescale +import Oceananigans.Advection: cell_advection_timescale import Oceananigans.Grids: cpu_face_constructor_x, @@ -233,7 +233,6 @@ const c = Center() const f = Face() @inline Base.zero(ibg::IBG) = zero(ibg.underlying_grid) -@inline cell_advection_timescale(u, v, w, ibg::IBG) = cell_advection_timescale(u, v, w, ibg.underlying_grid) @inline φᶠᶠᵃ(i, j, k, ibg::IBG) = φᶠᶠᵃ(i, j, k, ibg.underlying_grid) @inline xnode(i, ibg::IBG, LX; kwargs...) = xnode(i, ibg.underlying_grid, LX; kwargs...) diff --git a/src/ImmersedBoundaries/grid_fitted_immersed_boundaries.jl b/src/ImmersedBoundaries/grid_fitted_immersed_boundaries.jl index a91aea922c..46c0418b14 100644 --- a/src/ImmersedBoundaries/grid_fitted_immersed_boundaries.jl +++ b/src/ImmersedBoundaries/grid_fitted_immersed_boundaries.jl @@ -87,7 +87,7 @@ function ImmersedBoundaryGrid(grid, ib::AbstractGridFittedBottom{<:OffsetArray}) end function validate_ib_size(grid, ib) - bottom_height_size = total_size((Center, Center, Nothing), grid)[1:2] + bottom_height_size = total_size(grid, (Center, Center, Nothing))[1:2] size(ib.bottom_height) != bottom_height_size && throw(ArgumentError("The dimensions of the immersed boundary $(size(ib.bottom_height)) do not match the grid size $(bottom_height_size)")) diff --git a/src/ImmersedBoundaries/immersed_grid_metrics.jl b/src/ImmersedBoundaries/immersed_grid_metrics.jl index 305b7d9a3f..e7b983c79b 100644 --- a/src/ImmersedBoundaries/immersed_grid_metrics.jl +++ b/src/ImmersedBoundaries/immersed_grid_metrics.jl @@ -1,7 +1,6 @@ using Oceananigans.AbstractOperations: GridMetricOperation import Oceananigans.Grids: return_metrics, min_Δx, min_Δy, min_Δz -import Oceananigans.Operators: xspacing, yspacing, zspacing const c = Center() const f = Face() @@ -39,6 +38,3 @@ return_metrics(grid::IBG) = return_metrics(grid.underlying_grid) xspacings(X, grid::IBG) = xspacings(X, grid.underlying_grid) yspacings(Y, grid::IBG) = yspacings(Y, grid.underlying_grid) zspacings(Z, grid::IBG) = zspacings(Z, grid.underlying_grid) -min_Δx(grid::IBG) = min_Δx(grid.underlying_grid) -min_Δy(grid::IBG) = min_Δy(grid.underlying_grid) -min_Δz(grid::IBG) = min_Δz(grid.underlying_grid) diff --git a/src/ImmersedBoundaries/mask_immersed_field.jl b/src/ImmersedBoundaries/mask_immersed_field.jl index a169a59b07..64320c5901 100644 --- a/src/ImmersedBoundaries/mask_immersed_field.jl +++ b/src/ImmersedBoundaries/mask_immersed_field.jl @@ -4,7 +4,8 @@ using Statistics using Oceananigans.Architectures: architecture, device_event using Oceananigans.Fields: location, Field -instantiate(X) = X() +instantiate(T::Type) = T() +instantiate(t) = t ##### ##### Outer functions diff --git a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl index a5eb54d8dc..f9f83faa57 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl @@ -16,6 +16,7 @@ using DocStringExtensions import Oceananigans: fields, prognostic_fields import Oceananigans.Models: initialize_model! +import Oceananigans.Advection: cell_advection_timescale abstract type AbstractFreeSurface{E, G} end @@ -55,9 +56,11 @@ include("show_hydrostatic_free_surface_model.jl") include("set_hydrostatic_free_surface_model.jl") ##### -##### Time-stepping HydrostaticFreeSurfaceModels +##### AbstractModel interface ##### +cell_advection_timescale(model::HydrostaticFreeSurfaceModel) = cell_advection_timescale(model.grid, model.velocities) + """ fields(model::HydrostaticFreeSurfaceModel) diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl index 9cbfafabb3..9165be215a 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl @@ -13,6 +13,7 @@ using Oceananigans.Forcings: model_forcing using Oceananigans.Grids: halo_size, inflate_halo_size, with_halo, AbstractRectilinearGrid using Oceananigans.Grids: AbstractCurvilinearGrid, AbstractHorizontallyCurvilinearGrid, architecture using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid +using Oceananigans.Models: validate_model_halo using Oceananigans.Models.NonhydrostaticModels: extract_boundary_conditions using Oceananigans.TimeSteppers: Clock, TimeStepper, update_state! using Oceananigans.TurbulenceClosures: validate_closure, with_tracers, DiffusivityFields, add_closure_specific_boundary_conditions @@ -222,16 +223,5 @@ validate_momentum_advection(momentum_advection, grid) = momentum_advection validate_momentum_advection(momentum_advection, grid::AbstractHorizontallyCurvilinearGrid) = momentum_advection_squawk(momentum_advection, grid) validate_momentum_advection(momentum_advection::Union{VectorInvariant, Nothing}, grid::AbstractHorizontallyCurvilinearGrid) = momentum_advection -function validate_model_halo(grid, momentum_advection, tracer_advection, closure) - user_halo = halo_size(grid) - required_halo = inflate_halo_size(1, 1, 1, grid, - momentum_advection, - tracer_advection, - closure) - - any(user_halo .< required_halo) && - throw(ArgumentError("The grid halo $user_halo must be at least equal to $required_halo. Note that an ImmersedBoundaryGrid requires an extra halo point in all non-flat directions compared to a non-immersed boundary grid.")) -end - initialize_model!(model::HydrostaticFreeSurfaceModel) = initialize_free_surface!(model.free_surface, model.grid, model.velocities) initialize_free_surface!(free_surface, grid, velocities) = nothing diff --git a/src/Models/Models.jl b/src/Models/Models.jl index 804b7311fc..ce3bb4c967 100644 --- a/src/Models/Models.jl +++ b/src/Models/Models.jl @@ -8,6 +8,7 @@ export PrescribedVelocityFields, PressureField using Oceananigans: AbstractModel +using Oceananigans.Grids: halo_size, inflate_halo_size import Oceananigans.Architectures: device_event, architecture @@ -39,6 +40,17 @@ end abstract type AbstractNonhydrostaticModel{TS} <: AbstractModel{TS} end +function validate_model_halo(grid, momentum_advection, tracer_advection, closure) + user_halo = halo_size(grid) + required_halo = inflate_halo_size(1, 1, 1, grid, + momentum_advection, + tracer_advection, + closure) + + any(user_halo .< required_halo) && + throw(ArgumentError("The grid halo $user_halo must be at least equal to $required_halo. Note that an ImmersedBoundaryGrid requires an extra halo point in all non-flat directions compared to a non-immersed boundary grid.")) +end + include("NonhydrostaticModels/NonhydrostaticModels.jl") include("HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl") include("ShallowWaterModels/ShallowWaterModels.jl") diff --git a/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl b/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl index d5bf02808c..dd20c5d20d 100644 --- a/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl +++ b/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl @@ -14,6 +14,7 @@ using Oceananigans.Distributed: DistributedArch, DistributedFFTBasedPoissonSolve using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid import Oceananigans: fields, prognostic_fields +import Oceananigans.Advection: cell_advection_timescale function PressureSolver(arch::DistributedArch, local_grid::RegRectilinearGrid) global_grid = reconstruct_global_grid(local_grid) @@ -39,9 +40,11 @@ include("show_nonhydrostatic_model.jl") include("set_nonhydrostatic_model.jl") ##### -##### Time-stepping NonhydrostaticModels +##### AbstractModel interface ##### +cell_advection_timescale(model::NonhydrostaticModel) = cell_advection_timescale(model.grid, model.velocities) + """ fields(model::NonhydrostaticModel) diff --git a/src/Models/ShallowWaterModels/shallow_water_cell_advection_timescale.jl b/src/Models/ShallowWaterModels/shallow_water_cell_advection_timescale.jl index 9c9f1f0e20..7b2f58bb97 100644 --- a/src/Models/ShallowWaterModels/shallow_water_cell_advection_timescale.jl +++ b/src/Models/ShallowWaterModels/shallow_water_cell_advection_timescale.jl @@ -1,34 +1,16 @@ "Returns the time-scale for advection on a regular grid across a single grid cell for ShallowWaterModel." -import Oceananigans.Utils: cell_advection_timescale +import Oceananigans.Advection: cell_advection_timescale -function shallow_water_cell_advection_timescale(::ConservativeFormulation, uh, vh, h, grid::RectilinearGrid) - - Δxmin = minimum(grid.Δxᶜᵃᵃ) - Δymin = minimum(grid.Δyᵃᶜᵃ) - - uhmax = maximum(abs, uh) - vhmax = maximum(abs, vh) - hmin = minimum(abs, h) - - return min(Δxmin / uhmax, Δymin / vhmax) * hmin +function cell_advection_timescale(model::ShallowWaterModel) + u, v, _ = shallow_water_velocities(model) + τ = KernelFunctionOperation{Center, Center, Nothing}(shallow_water_cell_advection_timescaleᶜᶜᵃ, model.grid, u, v) + return minimum(τ) end -function shallow_water_cell_advection_timescale(::VectorInvariantFormulation, u, v, h, grid::RectilinearGrid) - - Δxmin = minimum(grid.Δxᶜᵃᵃ) - Δymin = minimum(grid.Δyᵃᶜᵃ) - - umax = maximum(abs, u) - vmax = maximum(abs, v) - - return min(Δxmin / umax, Δymin / vmax) +@inline function shallow_water_cell_advection_timescaleᶜᶜᵃ(i, j, k, grid, u, v) + Δx = Δxᶠᶜᶜ(i, j, k, grid) + Δy = Δyᶜᶠᶜ(i, j, k, grid) + return @inbounds min(Δx / abs(u[i, j, k]), Δy / abs(v[i, j, k])) end - -cell_advection_timescale(model::ShallowWaterModel) = shallow_water_cell_advection_timescale(model.formulation, - model.solution[1], - model.solution[2], - model.solution.h, - model.grid) - diff --git a/src/Models/ShallowWaterModels/shallow_water_model.jl b/src/Models/ShallowWaterModels/shallow_water_model.jl index b7aa1961a4..f0a5d62978 100644 --- a/src/Models/ShallowWaterModels/shallow_water_model.jl +++ b/src/Models/ShallowWaterModels/shallow_water_model.jl @@ -1,7 +1,7 @@ using Oceananigans: AbstractModel, AbstractOutputWriter, AbstractDiagnostic using Oceananigans.Architectures: AbstractArchitecture, CPU -using Oceananigans.AbstractOperations: @at +using Oceananigans.AbstractOperations: @at, KernelFunctionOperation using Oceananigans.Distributed using Oceananigans.Advection: CenteredSecondOrder, VectorInvariant using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions @@ -12,11 +12,13 @@ using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid using Oceananigans.TimeSteppers: Clock, TimeStepper, update_state! using Oceananigans.TurbulenceClosures: with_tracers, DiffusivityFields using Oceananigans.Utils: tupleit +using Oceananigans.Models: validate_model_halo using Oceananigans.Models.HydrostaticFreeSurfaceModels: validate_tracer_advection using Oceananigans.Models.NonhydrostaticModels: inflate_grid_halo_size + import Oceananigans.Architectures: architecture -const RectilinearGrids = Union{RectilinearGrid, ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:RectilinearGrid}} +const RectilinearGrids = Union{RectilinearGrid, ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:RectilinearGrid}} function ShallowWaterTendencyFields(grid, tracer_names, prognostic_names) u = XFaceField(grid) @@ -138,7 +140,8 @@ function ShallowWaterModel(; throw(ArgumentError("`ConservativeFormulation()` requires a rectilinear `grid`. \n" * "Use `VectorInvariantFormulation()` or change your grid to a rectilinear one.")) - grid = inflate_grid_halo_size(grid, momentum_advection, tracer_advection, mass_advection, closure) + # Check halos and throw an error if the grid's halo is too small + validate_model_halo(grid, momentum_advection, tracer_advection, closure) prognostic_field_names = formulation isa ConservativeFormulation ? (:uh, :vh, :h, tracers...) : (:u, :v, :h, tracers...) default_boundary_conditions = NamedTuple{prognostic_field_names}(Tuple(FieldBoundaryConditions() @@ -190,7 +193,7 @@ function ShallowWaterModel(; clock, eltype(grid)(gravitational_acceleration), advection, - shallow_water_velocities(solution, formulation), + shallow_water_velocities(formulation, solution), coriolis, forcing, closure, @@ -216,19 +219,16 @@ formulation(model::ShallowWaterModel) = model.formulation architecture(model::ShallowWaterModel) = model.architecture # The w velocity is needed to use generic TurbulenceClosures methods, therefore it is set to nothing -function shallow_water_velocities(solution, formulation) - if formulation isa VectorInvariantFormulation - return (u = solution.u, v = solution.v, w = nothing) - else - u = Field(@at (Face, Center, Center) solution.uh / solution.h) - v = Field(@at (Center, Face, Center) solution.vh / solution.h) - - compute!(u) - compute!(v) +shallow_water_velocities(::VectorInvariantFormulation, solution) = (u = solution.u, v = solution.v, w = nothing) - return (; u, v, w = nothing) - end +# TODO: convert u and v into binary operations +function shallow_water_velocities(::ConservativeFormulation, solution) + u = compute!(Field(solution.uh / solution.h)) + v = compute!(Field(solution.vh / solution.h)) + return (; u, v, w=nothing) end +shallow_water_velocities(model::ShallowWaterModel) = shallow_water_velocities(model.formulation, model.solution) + shallow_water_fields(velocities, solution, tracers, ::ConservativeFormulation) = merge(velocities, solution, tracers) -shallow_water_fields(velocities, solution, tracers, ::VectorInvariantFormulation) = merge(solution, (; w = velocities.w), tracers) \ No newline at end of file +shallow_water_fields(velocities, solution, tracers, ::VectorInvariantFormulation) = merge(solution, (; w = velocities.w), tracers) diff --git a/src/Models/ShallowWaterModels/update_shallow_water_state.jl b/src/Models/ShallowWaterModels/update_shallow_water_state.jl index b5fd8b43a4..6cf1dd7bb7 100644 --- a/src/Models/ShallowWaterModels/update_shallow_water_state.jl +++ b/src/Models/ShallowWaterModels/update_shallow_water_state.jl @@ -36,4 +36,5 @@ compute_velocities!(U, ::VectorInvariantFormulation) = nothing function compute_velocities!(U, ::ConservativeFormulation) compute!(U.u) compute!(U.v) + return nothing end diff --git a/src/MultiRegion/multi_region_models.jl b/src/MultiRegion/multi_region_models.jl index 959fb76b51..f56d6717b7 100644 --- a/src/MultiRegion/multi_region_models.jl +++ b/src/MultiRegion/multi_region_models.jl @@ -6,9 +6,7 @@ using Oceananigans.Models: PrescribedVelocityFields using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization using Oceananigans.Advection: AbstractAdvectionScheme -import Oceananigans.Simulations: new_time_step -import Oceananigans.Diagnostics: accurate_cell_advection_timescale -import Oceananigans.Advection: WENO +import Oceananigans.Advection: WENO, cell_advection_timescale import Oceananigans.Models.HydrostaticFreeSurfaceModels: build_implicit_step_solver, validate_tracer_advection import Oceananigans.TurbulenceClosures: implicit_diffusion_solver @@ -71,12 +69,8 @@ WENO(mrg::MultiRegionGrid, args...; kwargs...) = construct_regionally(WENO, mrg, t.divergence_stencil, getregion(t.vertical_scheme, r)) -function accurate_cell_advection_timescale(grid::MultiRegionGrid, velocities) - Δt = construct_regionally(accurate_cell_advection_timescale, grid, velocities) +function cell_advection_timescale(grid::MultiRegionGrid, velocities) + Δt = construct_regionally(cell_advection_timescale, grid, velocities) return minimum(Δt.regional_objects) end -function new_time_step(old_Δt, wizard, model::MultiRegionModel) - Δt = construct_regionally(new_time_step, old_Δt, wizard, model) - return minimum(Δt.regional_objects) -end diff --git a/src/Oceananigans.jl b/src/Oceananigans.jl index 531a34bea7..d85c933c77 100644 --- a/src/Oceananigans.jl +++ b/src/Oceananigans.jl @@ -24,6 +24,7 @@ export OrthogonalSphericalShellGrid, xnodes, ynodes, znodes, nodes, xspacings, yspacings, zspacings, + minimum_xspacing, minimum_yspacing, minimum_zspacing, # Immersed boundaries ImmersedBoundaryGrid, GridFittedBoundary, GridFittedBottom, ImmersedBoundaryCondition, diff --git a/src/Operators/Operators.jl b/src/Operators/Operators.jl index b877f2a1d5..b670ddc009 100644 --- a/src/Operators/Operators.jl +++ b/src/Operators/Operators.jl @@ -1,7 +1,6 @@ module Operators # Spacings -export xspacing, yspacing, zspacing export Δxᶠᶠᶠ, Δxᶠᶠᶜ, Δxᶠᶜᶠ, Δxᶠᶜᶜ, Δxᶜᶠᶠ, Δxᶜᶠᶜ, Δxᶜᶜᶠ, Δxᶜᶜᶜ export Δyᶠᶠᶠ, Δyᶠᶠᶜ, Δyᶠᶜᶠ, Δyᶠᶜᶜ, Δyᶜᶠᶠ, Δyᶜᶠᶜ, Δyᶜᶜᶠ, Δyᶜᶜᶜ export Δzᶠᶠᶠ, Δzᶠᶠᶜ, Δzᶠᶜᶠ, Δzᶠᶜᶜ, Δzᶜᶠᶠ, Δzᶜᶠᶜ, Δzᶜᶜᶠ, Δzᶜᶜᶜ @@ -60,6 +59,8 @@ export ℑxyzᶜᶜᶠ, ℑxyzᶜᶠᶜ, ℑxyzᶠᶜᶜ, ℑxyzᶜᶠᶠ, ℑxy using Oceananigans.Grids +import Oceananigans.Grids: xspacing, yspacing, zspacing + ##### ##### Convenient aliases ##### @@ -70,6 +71,10 @@ const RCG = RectilinearGrid const ACG = AbstractCurvilinearGrid const AHCG = AbstractHorizontallyCurvilinearGrid +const Δx = xspacing +const Δy = yspacing +const Δz = zspacing + include("difference_operators.jl") include("interpolation_operators.jl") include("interpolation_utils.jl") @@ -82,8 +87,4 @@ include("divergence_operators.jl") include("vorticity_operators.jl") include("laplacian_operators.jl") -const xspacing = Δx -const yspacing = Δy -const zspacing = Δz - end # module diff --git a/src/OutputReaders/field_time_series.jl b/src/OutputReaders/field_time_series.jl index 6c5b98ca9b..b86c135e31 100644 --- a/src/OutputReaders/field_time_series.jl +++ b/src/OutputReaders/field_time_series.jl @@ -49,8 +49,8 @@ struct UnspecifiedBoundaryConditions end iterations = nothing, times = nothing) -Returns a `FieldTimeSeries` for the field `name` describing a field's time history from a JLD2 file -located at `path`. +Return a `FieldTimeSeries` containing a time-series of the field `name` +load from JLD2 output located at `path`. Keyword arguments ================= @@ -69,6 +69,8 @@ Keyword arguments """ FieldTimeSeries(path, name; backend=InMemory(), kw...) = FieldTimeSeries(path, name, backend; kw...) +instantiate(T::Type) = T() + function FieldTimeSeries(path, name, backend; architecture = nothing, grid = nothing, @@ -82,7 +84,7 @@ function FieldTimeSeries(path, name, backend; # Defaults isnothing(iterations) && (iterations = parse.(Int, keys(file["timeseries/t"]))) isnothing(times) && (times = [file["timeseries/t/$i"] for i in iterations]) - isnothing(location) && (location = file["timeseries/$name/serialized/location"]) + isnothing(location) && (Location = file["timeseries/$name/serialized/location"]) if boundary_conditions isa UnspecifiedBoundaryConditions boundary_conditions = file["timeseries/$name/serialized/boundary_conditions"] @@ -104,13 +106,14 @@ function FieldTimeSeries(path, name, backend; # This should be removed in a month or two (4/5/2022). grid = on_architecture(architecture, grid) - LX, LY, LZ = location + LX, LY, LZ = Location + loc = map(instantiate, Location) if backend isa InMemory Nt = length(times) - space_size = total_size(location, grid, indices) + space_size = total_size(grid, loc, indices) underlying_data = zeros(eltype(grid), architecture, space_size..., Nt) - data = offset_data(underlying_data, grid, location, indices) + data = offset_data(underlying_data, grid, loc, indices) elseif backend isa OnDisk data = OnDiskData(path, name) else @@ -233,8 +236,8 @@ function set!(fts::FieldTimeSeries, fields_vector::AbstractVector{<:AbstractFiel end function interior(fts::FieldTimeSeries) - loc = location(fts) - topo = topology(fts.grid) + loc = instantiate.(location(fts)) + topo = instantiate.(topology(fts.grid)) sz = size(fts.grid) halo_sz = halo_size(fts.grid) @@ -275,7 +278,7 @@ end ##### # Include the time dimension. -@inline Base.size(fts::FieldTimeSeries) = (size(location(fts), fts.grid, fts.indices)..., length(fts.times)) +@inline Base.size(fts::FieldTimeSeries) = (size(fts.grid, location(fts), fts.indices)..., length(fts.times)) Base.setindex!(fts::FieldTimeSeries, val, inds...) = Base.setindex!(fts.data, val, inds...) diff --git a/src/OutputWriters/netcdf_output_writer.jl b/src/OutputWriters/netcdf_output_writer.jl index 9ba2f4b213..f4e25422dd 100644 --- a/src/OutputWriters/netcdf_output_writer.jl +++ b/src/OutputWriters/netcdf_output_writer.jl @@ -25,31 +25,33 @@ ext(::Type{NetCDFOutputWriter}) = ".nc" dictify(outputs) = outputs dictify(outputs::NamedTuple) = Dict(string(k) => dictify(v) for (k, v) in zip(keys(outputs), values(outputs))) -xdim(::Type{Face}) = ("xF",) -ydim(::Type{Face}) = ("yF",) -zdim(::Type{Face}) = ("zF",) +xdim(::Face) = ("xF",) +ydim(::Face) = ("yF",) +zdim(::Face) = ("zF",) -xdim(::Type{Center}) = ("xC",) -ydim(::Type{Center}) = ("yC",) -zdim(::Type{Center}) = ("zC",) +xdim(::Center) = ("xC",) +ydim(::Center) = ("yC",) +zdim(::Center) = ("zC",) -xdim(::Type{Nothing}) = () -ydim(::Type{Nothing}) = () -zdim(::Type{Nothing}) = () +xdim(::Nothing) = () +ydim(::Nothing) = () +zdim(::Nothing) = () netcdf_spatial_dimensions(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = - tuple(xdim(LX)..., ydim(LY)..., zdim(LZ)...) + tuple(xdim(instantiate(LX))..., ydim(instantiate(LY))..., zdim(instantiate(LZ))...) function default_dimensions(output, grid, indices, with_halos) Hx, Hy, Hz = halo_size(grid) TX, TY, TZ = topo = topology(grid) - locs = Dict("xC" => (Center, Center, Center), - "xF" => ( Face, Center, Center), - "yC" => (Center, Center, Center), - "yF" => (Center, Face, Center), - "zC" => (Center, Center, Center), - "zF" => (Center, Center, Face)) + locs = Dict("xC" => (Center(), Center(), Center()), + "xF" => (Face(), Center(), Center()), + "yC" => (Center(), Center(), Center()), + "yF" => (Center(), Face(), Center()), + "zC" => (Center(), Center(), Center()), + "zF" => (Center(), Center(), Face() )) + + topo = map(instantiate, topology(grid)) indices = Dict(name => validate_indices(indices, locs[name], grid) for name in keys(locs)) @@ -58,12 +60,12 @@ function default_dimensions(output, grid, indices, with_halos) for name in keys(locs)) end - dims = Dict("xC" => parent(xnodes(grid, Center(); with_halos=true))[parent_index_range(indices["xC"][1], Center, TX, Hx)], - "xF" => parent(xnodes(grid, Face(); with_halos=true))[parent_index_range(indices["xF"][1], Face, TX, Hx)], - "yC" => parent(ynodes(grid, Center(); with_halos=true))[parent_index_range(indices["yC"][2], Center, TY, Hy)], - "yF" => parent(ynodes(grid, Face(); with_halos=true))[parent_index_range(indices["yF"][2], Face, TY, Hy)], - "zC" => parent(znodes(grid, Center(); with_halos=true))[parent_index_range(indices["zC"][3], Center, TZ, Hz)], - "zF" => parent(znodes(grid, Face(); with_halos=true))[parent_index_range(indices["zF"][3], Face, TZ, Hz)]) + dims = Dict("xC" => parent(xnodes(grid, Center(); with_halos=true))[parent_index_range(indices["xC"][1], Center(), TX(), Hx)], + "xF" => parent(xnodes(grid, Face(); with_halos=true))[parent_index_range(indices["xF"][1], Face(), TX(), Hx)], + "yC" => parent(ynodes(grid, Center(); with_halos=true))[parent_index_range(indices["yC"][2], Center(), TY(), Hy)], + "yF" => parent(ynodes(grid, Face(); with_halos=true))[parent_index_range(indices["yF"][2], Face(), TY(), Hy)], + "zC" => parent(znodes(grid, Center(); with_halos=true))[parent_index_range(indices["zC"][3], Center(), TZ(), Hz)], + "zF" => parent(znodes(grid, Face(); with_halos=true))[parent_index_range(indices["zF"][3], Face(), TZ(), Hz)]) return dims end diff --git a/src/OutputWriters/output_construction.jl b/src/OutputWriters/output_construction.jl index df662dc605..542e393d68 100644 --- a/src/OutputWriters/output_construction.jl +++ b/src/OutputWriters/output_construction.jl @@ -3,8 +3,8 @@ using Oceananigans.AbstractOperations: AbstractOperation, ComputedField using Oceananigans.Grids: default_indices restrict_to_interior(::Colon, loc, topo, N) = interior_indices(loc, topo, N) -restrict_to_interior(::Colon, ::Type{Nothing}, topo, N) = UnitRange(1, 1) -restrict_to_interior(index::UnitRange, ::Type{Nothing}, topo, N) = UnitRange(1, 1) +restrict_to_interior(::Colon, ::Nothing, topo, N) = UnitRange(1, 1) +restrict_to_interior(index::UnitRange, ::Nothing, topo, N) = UnitRange(1, 1) function restrict_to_interior(index::UnitRange, loc, topo, N) from = max(first(index), 1) @@ -29,13 +29,15 @@ end ##### Support for Field, Reduction, and AbstractOperation outputs ##### +instantiate(T::Type) = T() + function output_indices(output::Union{AbstractField, Reduction}, grid, indices, with_halos) indices = validate_indices(indices, location(output), grid) if !with_halos # Maybe chop those indices - loc = location(output) - topo = topology(grid) - indices = restrict_to_interior.(indices, loc, topo, size(grid)) + loc = map(instantiate, location(output)) + topo = map(instantiate, topology(grid)) + indices = map(restrict_to_interior, indices, loc, topo, size(grid)) end return indices diff --git a/src/Simulations/Simulations.jl b/src/Simulations/Simulations.jl index d965514d92..8f9a21f6e7 100644 --- a/src/Simulations/Simulations.jl +++ b/src/Simulations/Simulations.jl @@ -8,17 +8,19 @@ export iteration export stopwatch export erroring_NaNChecker! -import Base: show - -using OrderedCollections: OrderedDict -using Oceananigans: AbstractDiagnostic, AbstractOutputWriter, fields - using Oceananigans.Models using Oceananigans.Diagnostics using Oceananigans.OutputWriters using Oceananigans.TimeSteppers using Oceananigans.Utils +using Oceananigans.Advection: cell_advection_timescale +using Oceananigans: AbstractDiagnostic, AbstractOutputWriter, fields + +using OrderedCollections: OrderedDict + +import Base: show + include("callback.jl") include("time_step_wizard.jl") include("nan_checker.jl") diff --git a/src/TurbulenceClosures/turbulence_closure_diagnostics.jl b/src/TurbulenceClosures/turbulence_closure_diagnostics.jl index e82623bed0..f2389bb129 100644 --- a/src/TurbulenceClosures/turbulence_closure_diagnostics.jl +++ b/src/TurbulenceClosures/turbulence_closure_diagnostics.jl @@ -5,19 +5,19 @@ using Oceananigans.Grids: topology, min_Δx, min_Δy, min_Δz function min_Δxyz(grid, ::ThreeDimensionalFormulation) - Δx = min_Δx(grid) - Δy = min_Δy(grid) - Δz = min_Δz(grid) + Δx = minimum_xspacing(grid, Center(), Center(), Center()) + Δy = minimum_yspacing(grid, Center(), Center(), Center()) + Δz = minimum_zspacing(grid, Center(), Center(), Center()) return min(Δx, Δy, Δz) end function min_Δxyz(grid, ::HorizontalFormulation) - Δx = min_Δx(grid) - Δy = min_Δy(grid) + Δx = minimum_xspacing(grid, Center(), Center(), Center()) + Δy = minimum_yspacing(grid, Center(), Center(), Center()) return min(Δx, Δy) end -min_Δxyz(grid, ::VerticalFormulation) = min_Δz(grid) +min_Δxyz(grid, ::VerticalFormulation) = minimum_zspacing(grid, Center(), Center(), Center()) cell_diffusion_timescale(model) = cell_diffusion_timescale(model.closure, model.diffusivity_fields, model.grid) diff --git a/src/Utils/Utils.jl b/src/Utils/Utils.jl index 8d99dc9970..a1a006a12a 100644 --- a/src/Utils/Utils.jl +++ b/src/Utils/Utils.jl @@ -1,8 +1,6 @@ module Utils export launch_config, work_layout, launch! -export cell_advection_timescale -export TimeStepWizard, update_Δt! export prettytime, pretty_filesize export tupleit, parenttuple, datatuple, datatuples export validate_intervals, time_to_run @@ -20,8 +18,8 @@ import CUDA # To avoid name conflicts ##### Misc. small utils ##### -instantiate(x) = x -instantiate(X::DataType) = X() +instantiate(T::Type) = T() +instantiate(t) = t getnamewrapper(type) = typeof(type).name.wrapper @@ -31,7 +29,6 @@ getnamewrapper(type) = typeof(type).name.wrapper include("prettysummary.jl") include("kernel_launching.jl") -include("cell_advection_timescale.jl") include("prettytime.jl") include("pretty_filesize.jl") include("tuple_utils.jl") diff --git a/src/Utils/cell_advection_timescale.jl b/src/Utils/cell_advection_timescale.jl deleted file mode 100644 index 79aa115e7d..0000000000 --- a/src/Utils/cell_advection_timescale.jl +++ /dev/null @@ -1,21 +0,0 @@ -using Oceananigans.Grids: topology, min_Δx, min_Δy, min_Δz - -"Returns the time-scale for advection on a regular grid across a single grid cell." -function cell_advection_timescale(u, v, w, grid) - - umax = maximum(abs, u) - vmax = maximum(abs, v) - wmax = maximum(abs, w) - - Δx = min_Δx(grid) - Δy = min_Δy(grid) - Δz = min_Δz(grid) - - return min(Δx/umax, Δy/vmax, Δz/wmax) -end - -cell_advection_timescale(model) = - cell_advection_timescale(model.velocities.u.data.parent, - model.velocities.v.data.parent, - model.velocities.w.data.parent, - model.grid) diff --git a/test/regression_tests/shallow_water_bickley_jet_regression.jl b/test/regression_tests/shallow_water_bickley_jet_regression.jl index 2103db9410..3a5461c96a 100644 --- a/test/regression_tests/shallow_water_bickley_jet_regression.jl +++ b/test/regression_tests/shallow_water_bickley_jet_regression.jl @@ -12,7 +12,8 @@ function run_shallow_water_regression(arch, formulation; regenerate_data = false grid = RectilinearGrid(arch, size = (Nx, Ny), x = (0, Lx), y = (-Ly/2, Ly/2), - topology = (Periodic, Bounded, Flat)) + topology = (Periodic, Bounded, Flat), + halo = (4, 4)) gravitational_acceleration = 1 coriolis = FPlane(f=1) diff --git a/test/test_grids.jl b/test/test_grids.jl index 488d5c1d53..f09f907b81 100644 --- a/test/test_grids.jl +++ b/test/test_grids.jl @@ -1,13 +1,12 @@ include("dependencies_for_runtests.jl") include("data_dependencies.jl") -using Oceananigans.Grids: total_extent, min_Δx, min_Δy, min_Δz, +using Oceananigans.Grids: total_extent, xspacings, yspacings, zspacings, xnode, ynode, znode, λspacings, φspacings, λspacing, φspacing -using Oceananigans.Operators: xspacing, yspacing, zspacing, - Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δxᶜᶜᵃ, Δyᶠᶜᵃ, Δyᶜᶠᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, Azᶜᶜᵃ +using Oceananigans.Operators: Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δxᶜᶜᵃ, Δyᶠᶜᵃ, Δyᶜᶠᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, Azᶜᶜᵃ ##### ##### Regular rectilinear grids @@ -187,7 +186,7 @@ function test_regular_rectilinear_xnode_ynode_znode_and_spacings(arch, FT) grids = [regular_spaced_grid, variably_spaced_grid] for (grid_type, grid) in zip(grids_types, grids) - @info " Testing on $grid_type grid...." + @info " Testing grid utils on $grid_type grid...." @test xnode(2, grid, Center()) ≈ FT(π/2) @test ynode(2, grid, Center()) ≈ FT(π/2) @@ -197,9 +196,9 @@ function test_regular_rectilinear_xnode_ynode_znode_and_spacings(arch, FT) @test ynode(2, grid, Face()) ≈ FT(π/3) @test znode(2, grid, Face()) ≈ FT(π/3) - @test min_Δx(grid) ≈ FT(π/3) - @test min_Δy(grid) ≈ FT(π/3) - @test min_Δz(grid) ≈ FT(π/3) + @test minimum_xspacing(grid) ≈ FT(π/3) + @test minimum_yspacing(grid) ≈ FT(π/3) + @test minimum_zspacing(grid) ≈ FT(π/3) @test all(xspacings(grid, Center()) .≈ FT(π/N)) @test all(yspacings(grid, Center()) .≈ FT(π/N)) @@ -413,7 +412,7 @@ function test_rectilinear_grid_correct_spacings(FT, N) @test all(isapprox.(zspacings(grid, Center(), with_halos=true), grid.Δzᵃᵃᶜ)) @test zspacing(1, 1, 2, grid, Center(), Center(), Face()) == grid.Δzᵃᵃᶠ[2] - @test min_Δz(grid) ≈ minimum(grid.Δzᵃᵃᶜ[1:grid.Nz]) + @test minimum_zspacing(grid, Center(), Center(), Center()) ≈ minimum(grid.Δzᵃᵃᶜ[1:grid.Nz]) # Note that Δzᵃᵃᶠ[1] involves a halo point, which is not directly determined by # the user-supplied zᵃᵃᶠ @@ -719,13 +718,12 @@ end ##### @testset "Grids" begin - @info "Testing grids..." + @info "Testing AbstractGrids..." @testset "Grid utils" begin @info " Testing grid utilities..." - - @test total_extent(Periodic, 1, 0.2, 1.0) == 1.2 - @test total_extent(Bounded, 1, 0.2, 1.0) == 1.4 + @test total_extent(Periodic(), 1, 0.2, 1.0) == 1.2 + @test total_extent(Bounded(), 1, 0.2, 1.0) == 1.4 end @testset "Regular rectilinear grid" begin @@ -753,7 +751,6 @@ end @testset "Grid dimensions" begin @info " Testing grid constructor errors..." - for FT in float_types test_regular_rectilinear_constructor_errors(FT) end @@ -761,7 +758,6 @@ end @testset "Grids with flat dimensions" begin @info " Testing construction of grids with Flat dimensions..." - for FT in float_types test_flat_size_regular_rectilinear_grid(FT) end @@ -878,7 +874,7 @@ end @test grid isa LatitudeLongitudeGrid end - + @testset "Conformal cubed sphere face grid" begin @info " Testing OrthogonalSphericalShellGrid grid..." diff --git a/test/test_internal_wave_dynamics.jl b/test/test_internal_wave_dynamics.jl index c4be67130b..9427ca2b5f 100644 --- a/test/test_internal_wave_dynamics.jl +++ b/test/test_internal_wave_dynamics.jl @@ -8,7 +8,7 @@ function internal_wave_solution(; L, background_stratification=false) k = 1 f = 0.2 ℕ = 1.0 - σ = sqrt( (ℕ^2*k^2 + f^2*m^2) / (k^2 + m^2) ) + σ = sqrt( (ℕ^2 * k^2 + f^2 * m^2) / (k^2 + m^2) ) # Numerical parameters Δt = 0.01 * 1/σ @@ -60,7 +60,7 @@ function internal_wave_dynamics_test(model, solution, Δt) v₀(x, y, z) = solution.v(x, y, z, 0) w₀(x, y, z) = solution.w(x, y, z, 0) b₀(x, y, z) = solution.b(x, y, z, 0) - + set!(model, u=u₀, v=v₀, w=w₀, b=b₀) simulation = Simulation(model, stop_iteration=10, Δt=Δt) @@ -71,4 +71,3 @@ function internal_wave_dynamics_test(model, solution, Δt) return nothing end - diff --git a/test/test_multi_region_advection_diffusion.jl b/test/test_multi_region_advection_diffusion.jl index eae98ca3f1..1639cb4fbc 100644 --- a/test/test_multi_region_advection_diffusion.jl +++ b/test/test_multi_region_advection_diffusion.jl @@ -1,6 +1,5 @@ using Oceananigans.MultiRegion using Oceananigans.MultiRegion: reconstruct_global_field -using Oceananigans.Grids: min_Δx, min_Δy using Oceananigans.Operators: hack_cosd # Tracer patch for visualization @@ -9,8 +8,8 @@ Gaussian(x, y, L) = exp(-(x^2 + y^2) / 2L^2) prescribed_velocities() = PrescribedVelocityFields(u=(λ, ϕ, z, t=0) -> 0.1 * hack_cosd(ϕ)) function Δ_min(grid) - Δx_min = min_Δx(grid) - Δy_min = min_Δy(grid) + Δx_min = minimum_xspacing(grid, Center(), Center(), Center()) + Δy_min = minimum_yspacing(grid, Center(), Center(), Center()) return min(Δx_min, Δy_min) end diff --git a/test/test_shallow_water_models.jl b/test/test_shallow_water_models.jl index 82e39f437e..ad1986786a 100644 --- a/test/test_shallow_water_models.jl +++ b/test/test_shallow_water_models.jl @@ -57,12 +57,14 @@ function test_shallow_water_diffusion_cosine(grid, formulation, fieldname, ξ) momentum_advection = nothing tracer_advection = nothing mass_advection = nothing + model = ShallowWaterModel(; grid, closure, gravitational_acceleration=1.0, momentum_advection, tracer_advection, mass_advection, formulation) field = model.velocities[fieldname] + interior(field) .= arch_array(architecture(grid), cos.(m * ξ)) update_state!(model) @@ -233,14 +235,19 @@ end @testset "ShallowWaterModels with ImmersedBoundaryGrid [$arch]" begin @info "Testing ShallowWaterModels with ImmersedBoundaryGrid [$arch]" - grid = RectilinearGrid(arch, size=(8, 8), x=(-10, 10), y=(0, 5), topology=(Periodic, Bounded, Flat)) - # Gaussian bump of width "1" bump(x, y, z) = y < exp(-x^2) - + + grid = RectilinearGrid(arch, size=(8, 8), x=(-10, 10), y=(0, 5), topology=(Periodic, Bounded, Flat)) grid_with_bump = ImmersedBoundaryGrid(grid, GridFittedBoundary(bump)) + + @test_throws ArgumentError model = ShallowWaterModel(grid=grid_with_bump, gravitational_acceleration=1) + + grid = RectilinearGrid(arch, size=(8, 8), x=(-10, 10), y=(0, 5), topology=(Periodic, Bounded, Flat), halo=(4, 4)) + grid_with_bump = ImmersedBoundaryGrid(grid, GridFittedBoundary(bump)) + model = ShallowWaterModel(grid=grid_with_bump, gravitational_acceleration=1) - + set!(model, h=1) simulation = Simulation(model, Δt=1.0, stop_iteration=1) run!(simulation) diff --git a/validation/advection/plot_rates_convergence_advection.jl b/validation/advection/plot_rates_convergence_advection.jl index 18c9a7f101..932ba53edf 100644 --- a/validation/advection/plot_rates_convergence_advection.jl +++ b/validation/advection/plot_rates_convergence_advection.jl @@ -6,7 +6,6 @@ using LinearAlgebra using OffsetArrays using Oceananigans -using Oceananigans.Grids: min_Δx using Oceananigans.Models: ShallowWaterModel rate_of_convergence(::UpwindBiased) = 5 @@ -19,7 +18,7 @@ rate_of_convergence(::WENO{6}) = 11 labels(::Centered) = "Center4ᵗʰ" labels(::UpwindBiased) = "Upwind5ᵗʰ" -labels(::WENO) = "WENOᵗʰ " +labels(::WENO) = "WENOᵗʰ " shapes(::Centered) = :diamond shapes(::UpwindBiased) = :square @@ -52,31 +51,32 @@ W = 0.1 Ns = 2 .^ (5:8) pnorm = 1 -c(x, y, z, t, U, W) = exp( - (x - U * t)^2 / W^2 ) - h(x, y, z) = 1 - uh(x, y, z) = U * h(x, y, z) + c(x, y, z, t, U, W) = exp( - (x - U * t)^2 / W^2 ) + h(x, y, z) = 1 +uh(x, y, z) = U * h(x, y, z) schemes = ( - CenteredFourthOrder(), - UpwindBiasedFifthOrder(), - WENO(order = 3), - WENO(order = 5), - WENO(order = 7), - WENO(order = 9), - WENO(order = 11) -); + CenteredFourthOrder(), + UpwindBiasedFifthOrder(), + WENO(order = 3), + WENO(order = 5), + WENO(order = 7), + WENO(order = 9), + WENO(order = 11) +)s error = Dict() ROC = Dict() for N in Ns, (adv, scheme) in enumerate(schemes) - grid = RectilinearGrid(Float64; size=N, - x=(-1, 1), - halo=(halos(scheme)), - topology=(Periodic, Flat, Flat)) + grid = RectilinearGrid(Float64; + size = N, + x = (-1, 1), + halo = (halos(scheme)), + topology = (Periodic, Flat, Flat)) - Δt = 0.1 * min_Δx(grid) + Δt = 0.1 * minimum_xspacing(grid, Center(), Center(), Center()) model = ShallowWaterModel(grid = grid, momentum_advection = scheme, diff --git a/validation/advection/validate_one_dimensional_advection.jl b/validation/advection/validate_one_dimensional_advection.jl index 8b03da502c..93360d5ffd 100644 --- a/validation/advection/validate_one_dimensional_advection.jl +++ b/validation/advection/validate_one_dimensional_advection.jl @@ -84,7 +84,7 @@ for (gr, grid) in enumerate([grid_str]) @info "testing grid number $gr" - Δt_max = 0.2 * min_Δx(grid) + Δt_max = 0.2 * minimum_xspacing(grid) end_time = 2.0 @show tot_iter = end_time ÷ Δt_max diff --git a/validation/stratified_couette_flow/stratified_couette_flow.jl b/validation/stratified_couette_flow/stratified_couette_flow.jl index 80c5eee894..ca8a7ca600 100644 --- a/validation/stratified_couette_flow/stratified_couette_flow.jl +++ b/validation/stratified_couette_flow/stratified_couette_flow.jl @@ -7,6 +7,8 @@ using Oceananigans.OutputWriters using Oceananigans.Diagnostics using Oceananigans.Utils +using Oceananigans.Advection: cell_advection_timescale + """ Friction velocity. See equation (16) of Vreugdenhil & Taylor (2018). """ function uτ(model, Uavg, U_wall, n) Nz, Hz, Δz = model.grid.Nz, model.grid.Hz, model.grid.Δzᵃᵃᶜ