From 251cc952ae1f2c09639eba7f2a4011117e277a90 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 10 Dec 2021 16:39:16 -0700 Subject: [PATCH 001/140] One Field to rule them all --- src/AbstractOperations/AbstractOperations.jl | 3 +- src/BuoyancyModels/buoyancy_field.jl | 19 +- src/CubedSpheres/cubed_sphere_faces.jl | 28 +- src/CubedSpheres/cubed_sphere_set!.jl | 20 +- src/Distributed/halo_communication.jl | 9 +- src/Fields/Fields.jl | 6 +- src/Fields/abstract_field.jl | 135 +----- src/Fields/averaged_field.jl | 2 +- src/Fields/broadcasting_abstract_fields.jl | 12 +- src/Fields/field.jl | 427 +++++++++++++++--- src/Fields/kernel_computed_field.jl | 2 +- src/Fields/reduced_field.jl | 84 ---- src/Fields/set!.jl | 17 +- src/Grids/new_data.jl | 17 +- src/ImmersedBoundaries/mask_immersed_field.jl | 10 +- .../pcg_implicit_free_surface_solver.jl | 10 +- .../prescribed_hydrostatic_velocity_fields.jl | 2 +- src/OutputReaders/field_time_series.jl | 2 +- test/test_field.jl | 93 ++-- 19 files changed, 502 insertions(+), 396 deletions(-) diff --git a/src/AbstractOperations/AbstractOperations.jl b/src/AbstractOperations/AbstractOperations.jl index d4526a94fc..0caab7b7b5 100644 --- a/src/AbstractOperations/AbstractOperations.jl +++ b/src/AbstractOperations/AbstractOperations.jl @@ -27,7 +27,8 @@ import Oceananigans.Fields: data, compute_at! ##### Basic functionality ##### -abstract type AbstractOperation{X, Y, Z, A, G, T} <: AbstractField{X, Y, Z, A, G, T, 3} end +using Oceananigans.Fields: AbstractOperation +# abstract type AbstractOperation{X, Y, Z, A, G, T} <: AbstractField{X, Y, Z, A, G, T, 3} end const AF = AbstractField diff --git a/src/BuoyancyModels/buoyancy_field.jl b/src/BuoyancyModels/buoyancy_field.jl index 261bca4c0a..12f5c43b4e 100644 --- a/src/BuoyancyModels/buoyancy_field.jl +++ b/src/BuoyancyModels/buoyancy_field.jl @@ -1,7 +1,22 @@ using Adapt using KernelAbstractions -using Oceananigans.Fields: AbstractDataField, FieldStatus, validate_field_data, conditional_compute! +# TODO: define buoyancy_operation(model) and be done with it. + +#= +function buoyancy_operation(model) + buoyancy = model.buoyancy + tracers = model.tracers + return KernelFunctionOperation{Center, Center, Center}(...) +end + +function BuoyancyField(model) + op = buoyancy_operation(model) + return Field(op) +end +=# + +using Oceananigans.Fields: AbstractField, FieldStatus, validate_field_data, conditional_compute! using Oceananigans.Fields: architecture, tracernames using Oceananigans.Architectures: device using Oceananigans.Utils: work_layout @@ -11,7 +26,7 @@ import Oceananigans.Fields: compute!, compute_at! import Oceananigans: short_show -struct BuoyancyField{B, S, A, D, G, T, C} <: AbstractDataField{Center, Center, Center, A, G, T, 3} +struct BuoyancyField{B, S, A, D, G, T, C} <: AbstractField{Center, Center, Center, A, G, T, 3} data :: D architecture :: A grid :: G diff --git a/src/CubedSpheres/cubed_sphere_faces.jl b/src/CubedSpheres/cubed_sphere_faces.jl index 7aa708d1c8..eb643f31d6 100644 --- a/src/CubedSpheres/cubed_sphere_faces.jl +++ b/src/CubedSpheres/cubed_sphere_faces.jl @@ -7,7 +7,7 @@ using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid import Base: getindex, size, show, minimum, maximum import Statistics: mean -import Oceananigans.Fields: AbstractField, AbstractDataField, AbstractReducedField, Field, ReducedField, minimum, maximum, mean, location, short_show, KernelComputedField +import Oceananigans.Fields: AbstractField, Field, minimum, maximum, mean, location, short_show, KernelComputedField import Oceananigans.Grids: new_data import Oceananigans.BoundaryConditions: FieldBoundaryConditions @@ -40,24 +40,14 @@ const ImmersedCubedSphereFaceField = AbstractField{X, Y, Z, A, <:ImmersedConf const CubedSphereFaceField = Union{NonImmersedCubedSphereFaceField{X, Y, Z, A}, ImmersedCubedSphereFaceField{X, Y, Z, A}} where {X, Y, Z, A} -const NonImmersedCubedSphereAbstractReducedFaceField = AbstractReducedField{X, Y, Z, A, D, <:ConformalCubedSphereFaceGrid} where {X, Y, Z, A, D} -const ImmersedCubedSphereAbstractReducedFaceField = AbstractReducedField{X, Y, Z, A, D, <:ImmersedConformalCubedSphereFaceGrid} where {X, Y, Z, A, D} - -const CubedSphereAbstractReducedFaceField = Union{NonImmersedCubedSphereAbstractReducedFaceField{X, Y, Z, A}, - ImmersedCubedSphereAbstractReducedFaceField{X, Y, Z, A}} where {X, Y, Z, A} - # CubedSphereField # Flavors of CubedSphereField const CubedSphereField = Field{X, Y, Z, A, <:CubedSphereData} where {X, Y, Z, A} -const CubedSphereReducedField = ReducedField{X, Y, Z, A, <:CubedSphereData} where {X, Y, Z, A} const CubedSphereAbstractField = AbstractField{X, Y, Z, A, <:ConformalCubedSphereGrid} where {X, Y, Z, A} -const CubedSphereAbstractDataField = AbstractDataField{X, Y, Z, A, <:ConformalCubedSphereGrid} where {X, Y, Z, A} -const AbstractCubedSphereField{X, Y, Z, A} = Union{ CubedSphereAbstractField{X, Y, Z, A}, - CubedSphereAbstractDataField{X, Y, Z, A}, - CubedSphereReducedField{X, Y, Z, A}, - CubedSphereField{X, Y, Z, A}} where {X, Y, Z, A} +const AbstractCubedSphereField{X, Y, Z, A} = Union{CubedSphereAbstractField{X, Y, Z, A}, + CubedSphereField{X, Y, Z, A}} where {X, Y, Z, A} ##### ##### new data @@ -123,17 +113,7 @@ Base.size(data::CubedSphereData) = (size(data.faces[1])..., length(data.faces)) get_face(field.boundary_conditions, face_index)) end -@inline function get_face(reduced_field::CubedSphereReducedField, face_index) - X, Y, Z = location(reduced_field) - - return ReducedField{X, Y, Z}(get_face(reduced_field.data, face_index), - reduced_field.architecture, - get_face(reduced_field.grid, face_index), - reduced_field.dims, - get_face(reduced_field.boundary_conditions, face_index)) -end - -@inline function get_face(kernel_field::CubedSphereAbstractDataField, face_index) +@inline function get_face(kernel_field::CubedSphereAbstractField, face_index) X, Y, Z = location(kernel_field) return Field{X, Y, Z}(get_face(kernel_field.data, face_index), diff --git a/src/CubedSpheres/cubed_sphere_set!.jl b/src/CubedSpheres/cubed_sphere_set!.jl index 332ef4893b..c1135dde3b 100644 --- a/src/CubedSpheres/cubed_sphere_set!.jl +++ b/src/CubedSpheres/cubed_sphere_set!.jl @@ -1,37 +1,31 @@ using Oceananigans.Architectures: CPU -using Oceananigans.Fields: AbstractField, ReducedField +using Oceananigans.Fields: AbstractField import Oceananigans.Fields: set! const CubedSphereCPUField = CubedSphereField{X, Y, Z, <:CPU} where {X, Y, Z} const CubedSphereGPUField = CubedSphereField{X, Y, Z, <:GPU} where {X, Y, Z} -const CubedSphereCPUReducedField = CubedSphereReducedField{X, Y, Z, <:CPU} where {X, Y, Z} -const CubedSphereGPUReducedField = CubedSphereReducedField{X, Y, Z, <:GPU} where {X, Y, Z} - -const CubedSphereCPUFields = Union{CubedSphereCPUField, CubedSphereCPUReducedField} -const CubedSphereGPUFields = Union{CubedSphereGPUField, CubedSphereGPUReducedField} - # We need to define the function once for CPU fields then again for GPU fields to avoid the method # ambiguity with Fields.set!. -function set!(u::CubedSphereCPUFields , v::CubedSphereCPUFields) +function set!(u::CubedSphereCPUField, v::CubedSphereCPUField) for (u_face, v_face) in zip(faces(u), faces(v)) @. u_face.data.parent = v_face.data.parent end return nothing end -function set!(u::CubedSphereGPUField , v::CubedSphereGPUField) +function set!(u::CubedSphereGPUField, v::CubedSphereGPUField) for (u_face, v_face) in zip(faces(u), faces(v)) @. u_face.data.parent = v_face.data.parent end return nothing end -set!(field::CubedSphereCPUFields, f::Function) = [set_face_field!(field_face, f) for field_face in faces(field)] -set!(field::CubedSphereCPUFields, f::Number) = [set_face_field!(field_face, f) for field_face in faces(field)] -set!(field::CubedSphereGPUFields, f::Function) = [set_face_field!(field_face, f) for field_face in faces(field)] -set!(field::CubedSphereGPUFields, f::Number) = [set_face_field!(field_face, f) for field_face in faces(field)] +set!(field::CubedSphereCPUField, f::Function) = [set_face_field!(field_face, f) for field_face in faces(field)] +set!(field::CubedSphereCPUField, f::Number) = [set_face_field!(field_face, f) for field_face in faces(field)] +set!(field::CubedSphereGPUField, f::Function) = [set_face_field!(field_face, f) for field_face in faces(field)] +set!(field::CubedSphereGPUField, f::Number) = [set_face_field!(field_face, f) for field_face in faces(field)] set_face_field!(field, a) = set!(field, a) diff --git a/src/Distributed/halo_communication.jl b/src/Distributed/halo_communication.jl index f8480bb079..25e3e54d39 100644 --- a/src/Distributed/halo_communication.jl +++ b/src/Distributed/halo_communication.jl @@ -1,4 +1,4 @@ -using Oceananigans.Fields: AbstractField, AbstractReducedField +using Oceananigans.Fields: AbstractField using KernelAbstractions: @kernel, @index, Event, MultiEvent using OffsetArrays: OffsetArray @@ -53,13 +53,6 @@ end ##### Filling halos for halo communication boundary conditions ##### -fill_halo_regions!(field::AbstractField{LX, LY, LZ}, arch::AbstractMultiArchitecture, args...) where {LX, LY, LZ} = - fill_halo_regions!(field.data, field.boundary_conditions, arch, field.grid, (LX, LY, LZ), args...) - - -fill_halo_regions!(field::AbstractReducedField, arch::AbstractMultiArchitecture, args...) = - fill_halo_regions!(field.data, field.boundary_conditions, arch, field.grid, args...; reduced_dimensions=field.dims) - function fill_halo_regions!(c::OffsetArray, bcs, arch::AbstractMultiArchitecture, grid, c_location, args...) barrier = Event(device(child_architecture(arch))) diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl index 912f215eb6..3ff61c567d 100644 --- a/src/Fields/Fields.jl +++ b/src/Fields/Fields.jl @@ -16,12 +16,12 @@ using Oceananigans.Grids using Oceananigans.BoundaryConditions include("abstract_field.jl") -include("reduced_getindex_setindex.jl") +# include("reduced_getindex_setindex.jl") include("field.jl") include("zero_field.jl") -include("reduced_field.jl") +# include("reduced_field.jl") include("averaged_field.jl") -include("computed_field.jl") +# include("computed_field.jl") include("kernel_computed_field.jl") include("pressure_field.jl") include("function_field.jl") diff --git a/src/Fields/abstract_field.jl b/src/Fields/abstract_field.jl index eef0b9f95e..9844bef44f 100644 --- a/src/Fields/abstract_field.jl +++ b/src/Fields/abstract_field.jl @@ -27,23 +27,16 @@ Abstract supertype for fields located at `(X, Y, Z)` on architecture `A` and defined on a grid `G` with eltype `T` and `N` dimensions. """ abstract type AbstractField{X, Y, Z, A <: ArchOrNothing, G <: GridOrNothing, T, N} <: AbstractArray{T, N} end - -""" - AbstractDataField{X, Y, Z, A, G, T, N} - -Abstract supertype for fields with concrete data in settable underlying arrays, -located at `(X, Y, Z)` on architecture `A` and defined on a grid `G` with eltype `T` -and `N` dimensions. -""" -abstract type AbstractDataField{X, Y, Z, A, G, T, N} <: AbstractField{X, Y, Z, A, G, T, N} end +abstract type AbstractOperation{X, Y, Z, A, G, T} <: AbstractField{X, Y, Z, A, G, T, 3} end Base.IndexStyle(::AbstractField) = IndexCartesian() -function validate_field_data(X, Y, Z, data, grid) - Tx, Ty, Tz = total_size((X, Y, Z), grid) +function validate_field_data(loc, data, grid) + Tx, Ty, Tz = total_size(loc, grid) if size(data) != (Tx, Ty, Tz) - e = "Cannot construct field at ($X, $Y, $Z) with size(data)=$(size(data)). " * + LX, LY, LZ = loc + e = "Cannot construct field at ($LX, $LY, $LZ) with size(data)=$(size(data)). " * "`data` must have size ($Tx, $Ty, $Tz)." throw(ArgumentError(e)) end @@ -54,63 +47,6 @@ end # Endpoint for recursive `datatuple` function: @inline datatuple(obj::AbstractField) = data(obj) -##### -##### Computing AbstractField -##### - -# Note: overload compute! for custom fields to produce non-default behavior - -""" - compute!(field) - -Computes `field.data`. -""" -compute!(field) = nothing - -""" - compute_at!(field, time) - -Computes `field.data` at `time`. Falls back to compute!(field). -""" -compute_at!(field, time) = compute!(field) - -mutable struct FieldStatus{T} - time :: T -end - -Adapt.adapt_structure(to, status::FieldStatus) = (time = status.time,) - -""" - conditional_compute!(field, time) - -Computes `field.data` if `time != field.status.time`. -""" -function conditional_compute!(field, time) - - if time == zero(time) || time != field.status.time - compute!(field, time) - field.status.time = time - end - - return nothing -end - -# This edge case occurs if `fetch_output` is called with `model::Nothing`. -# We do the safe thing here and always compute. -conditional_compute!(field, ::Nothing) = compute!(field, nothing) - -""" - @compute(exprs...) - -Call compute! on fields after defining them. -""" -macro compute(def) - expr = Expr(:block) - field = def.args[1] - push!(expr.args, :($(esc(def)))) - push!(expr.args, :(compute!($(esc(field))))) - return expr -end ##### ##### AbstractField functionality @@ -118,14 +54,14 @@ end @inline location(a) = (Nothing, Nothing, Nothing) -"Returns the location `(X, Y, Z)` of an `AbstractField{X, Y, Z}`." -@inline location(::AbstractField{X, Y, Z}) where {X, Y, Z} = (X, Y, Z) # note no instantiation +"Returns the location `(LX, LY, LZ)` of an `AbstractField{LX, LY, LZ}`." +@inline location(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = (LX, LY, LZ) # note no instantiation @inline location(f, i) = location(f)[i] @inline instantiated_location(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = (LX(), LY(), LZ()) "Returns the architecture of on which `f` is defined." -architecture(f::AbstractField) = f.architecture +architecture(f::AbstractField) = architecture(f.grid) "Returns the topology of a fields' `grid`." @inline topology(f, args...) = topology(f.grid, args...) @@ -150,30 +86,20 @@ both interior points and halo points. """ total_size(f::AbstractField) = total_size(location(f), f.grid) -Base.fill!(f::AbstractDataField, val) = fill!(parent(f), val) - ##### ##### Accessing wrapped arrays ##### "Returns `f.data` for `f::Field` or `f` for `f::AbstractArray." -@inline data(a) = nothing # fallback -@inline data(f::AbstractDataField) = f.data - -@inline cpudata(a) = data(a) +data(a) = nothing # fallback +cpudata(a) = data(a) -@inline cpudata(f::AbstractField{X, Y, Z, <:GPU}) where {X, Y, Z} = +cpudata(f::AbstractField{X, Y, Z, <:GPU}) where {X, Y, Z} = offset_data(Array(parent(f)), f.grid, location(f)) "Returns `f.data.parent` for `f::Field`." @inline Base.parent(f::AbstractField) = parent(data(f)) -"Returns a view of `f` that excludes halo points." -@inline interior(f::AbstractDataField{X, Y, Z}) where {X, Y, Z} = - view(parent(f), interior_parent_indices(X, topology(f, 1), f.grid.Nx, f.grid.Hx), - interior_parent_indices(Y, topology(f, 2), f.grid.Ny, f.grid.Hy), - interior_parent_indices(Z, topology(f, 3), f.grid.Nz, f.grid.Hz)) - @inline interior(f::AbstractField) = f @inline interior_copy(f::AbstractField{X, Y, Z}) where {X, Y, Z} = @@ -181,33 +107,6 @@ Base.fill!(f::AbstractDataField, val) = fill!(parent(f), val) interior_parent_indices(Y, topology(f, 2), f.grid.Ny, f.grid.Hy), interior_parent_indices(Z, topology(f, 3), f.grid.Nz, f.grid.Hz)] -##### -##### getindex -##### - -# Don't use axes(f) to checkbounds; use axes(f.data) -Base.checkbounds(f::AbstractDataField, I...) = Base.checkbounds(f.data, I...) - -@propagate_inbounds Base.getindex(f::AbstractDataField, inds...) = getindex(f.data, inds...) - -# Linear indexing -@propagate_inbounds Base.getindex(f::AbstractDataField, i::Int) = parent(f)[i] - -##### -##### setindex -##### - -@propagate_inbounds function Base.setindex!(f::AbstractDataField, val, i, j, k) - f.data[i, j, k] = val - return f -end - -# Linear indexing -@propagate_inbounds function Base.setindex!(f::AbstractDataField, val, i::Int) - parent(f)[i] = val - return f -end - ##### ##### Coordinates of fields ##### @@ -216,9 +115,6 @@ end @propagate_inbounds ynode(j, ψ::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = ynode(LY(), j, ψ.grid) @propagate_inbounds znode(k, ψ::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = znode(LZ(), k, ψ.grid) -@propagate_inbounds Base.lastindex(f::AbstractDataField) = lastindex(f.data) -@propagate_inbounds Base.lastindex(f::AbstractDataField, dim) = lastindex(f.data, dim) - xnodes(ψ::AbstractField) = xnodes(location(ψ, 1), ψ.grid) ynodes(ψ::AbstractField) = ynodes(location(ψ, 2), ψ.grid) znodes(ψ::AbstractField) = znodes(location(ψ, 3), ψ.grid) @@ -241,4 +137,11 @@ for f in (:+, :-) @eval Base.$f(ϕ::AbstractField, ψ::AbstractArray) = $f(interior(ϕ), ψ) end -Base.isapprox(ϕ::AbstractDataField, ψ::AbstractDataField; kw...) = isapprox(interior(ϕ), interior(ψ); kw...) +function Statistics.norm(a::AbstractField) + arch = architecture(a) + grid = a.grid + r = zeros(arch, grid, 1) + Base.mapreducedim!(x -> x * x, +, r, a) + return CUDA.@allowscalar sqrt(r[1]) +end + diff --git a/src/Fields/averaged_field.jl b/src/Fields/averaged_field.jl index 9339eebbd9..5ee0b8f6cf 100644 --- a/src/Fields/averaged_field.jl +++ b/src/Fields/averaged_field.jl @@ -3,7 +3,7 @@ using Statistics using Oceananigans.Grids using Oceananigans.Grids: interior_parent_indices -struct AveragedField{X, Y, Z, S, A, D, G, T, N, O} <: AbstractReducedField{X, Y, Z, A, G, T, N} +struct AveragedField{X, Y, Z, S, A, D, G, T, N, O} <: AbstractField{X, Y, Z, A, G, T, N} data :: D architecture :: A grid :: G diff --git a/src/Fields/broadcasting_abstract_fields.jl b/src/Fields/broadcasting_abstract_fields.jl index 1ef3f0a30f..7146db5525 100644 --- a/src/Fields/broadcasting_abstract_fields.jl +++ b/src/Fields/broadcasting_abstract_fields.jl @@ -59,13 +59,13 @@ broadcast_kernel(::AbstractField) = broadcast_xyz! const Loc = Union{Center, Face} -broadcast_kernel(::AbstractReducedField{Nothing, <:Loc, <:Loc}) = broadcast_yz! -broadcast_kernel(::AbstractReducedField{<:Loc, Nothing, <:Loc}) = broadcast_xz! -broadcast_kernel(::AbstractReducedField{<:Loc, <:Loc, Nothing}) = broadcast_xy! +broadcast_kernel(::Field{Nothing, <:Loc, <:Loc}) = broadcast_yz! +broadcast_kernel(::Field{<:Loc, Nothing, <:Loc}) = broadcast_xz! +broadcast_kernel(::Field{<:Loc, <:Loc, Nothing}) = broadcast_xy! -launch_configuration(::AbstractReducedField{Nothing, <:Loc, <:Loc}) = :yz -launch_configuration(::AbstractReducedField{<:Loc, Nothing, <:Loc}) = :xz -launch_configuration(::AbstractReducedField{<:Loc, <:Loc, Nothing}) = :xy +launch_configuration(::Field{Nothing, <:Loc, <:Loc}) = :yz +launch_configuration(::Field{<:Loc, Nothing, <:Loc}) = :xz +launch_configuration(::Field{<:Loc, <:Loc, Nothing}) = :xy broadcasted_to_abstract_operation(loc, grid, a) = a diff --git a/src/Fields/field.jl b/src/Fields/field.jl index 354a7ec89a..c6302ed65c 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -1,113 +1,426 @@ +using Oceananigans.Architectures: device_event + using Adapt +using KernelAbstractions: @kernel, @index + +using Base: @propagate_inbounds -struct Field{LX, LY, LZ, A, D, G, T, B} <: AbstractDataField{LX, LY, LZ, A, G, T, 3} - data :: D - architecture :: A - grid :: G +struct Field{LX, LY, LZ, O, A, G, T, D, B, S} <: AbstractField{LX, LY, LZ, A, G, T, 3} + grid :: G + data :: D boundary_conditions :: B + operand :: O + status :: S - function Field{LX, LY, LZ}(data::D, arch::A, grid::G, bcs::B) where {LX, LY, LZ, D, A, G, B} + # Inner constructor that does not validate _anything_! + function Field{LX, LY, LZ}(grid::G, data::D, bcs::B, op::O, status::S) where {LX, LY, LZ, G, D, B, O, S} T = eltype(grid) - return new{LX, LY, LZ, A, D, G, T, B}(data, arch, grid, bcs) + A = typeof(architecture(grid)) + return new{LX, LY, LZ, O, A, G, T, D, B, S}(grid, data, bcs, op, status) end end +# Common outer constructor for all field flavors that validates data and boundary conditions +function Field(loc::Tuple, grid::AbstractGrid, data, bcs, op, status) + validate_field_data(loc, data, grid) + # validate_boundary_conditions(loc, grid, bcs) + arch = architecture(grid) + LX, LY, LZ = loc + return Field{LX, LY, LZ}(grid, data, bcs, op, status) +end + +##### +##### Vanilla, non-computed Field +##### + """ - Field(LX, LY, LZ, [arch = CPU()], grid, - [ bcs = FieldBoundaryConditions(grid, (LX, LY, LZ)), - data = new_data(eltype(grid), arch, grid, (LX, LY, LZ))]) + Field{LX, LY, LZ}(grid; kwargs...) -Construct a `Field` on `grid` with `data` on architecture `arch` with -boundary conditions `bcs`. Each of `(LX, LY, LZ)` is either `Center` or `Face` and determines +Construct a `Field` on `grid` at the location `(LX, LY, LZ)`. +Each of `(LX, LY, LZ)` is either `Center` or `Face` and determines the field's location in `(x, y, z)`. +Keyword arguments +================= + +- data: + Example ======= ```jldoctest julia> using Oceananigans -julia> ω = Field(Face, Face, Center, CPU(), RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1))) +julia> ω = Field{Face, Face, Center}(RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1))) Field located at (Face, Face, Center) ├── data: OffsetArrays.OffsetArray{Float64, 3, Array{Float64, 3}}, size: (1, 1, 1) ├── grid: RectilinearGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1) └── boundary conditions: west=Periodic, east=Periodic, south=Periodic, north=Periodic, bottom=ZeroFlux, top=ZeroFlux, immersed=ZeroFlux ``` """ -function Field(LX, LY, LZ, - arch::AbstractArchitecture, +Field{LX, LY, LZ}(grid::AbstractGrid, T; kwargs...) where {LX, LY, LZ} = Field((LX, LY, LZ), grid, T; kwargs...) + +function Field(loc::Tuple, grid::AbstractGrid, - bcs = FieldBoundaryConditions(grid, (LX, LY, LZ)), - data = new_data(eltype(grid), arch, grid, (LX, LY, LZ))) + T = eltype(grid); + data = new_data(T, grid, loc), + boundary_conditions = FieldBoundaryConditions(grid, loc)) - # Make sure architecture is the same as the grid architecture - # if hasproperty(grid, :architecture) - # arch = grid.architecture - # end - - validate_field_data(LX, LY, LZ, data, grid) + return Field(loc, grid, data, boundary_conditions, nothing, nothing) +end - return Field{LX, LY, LZ}(data, arch, grid, bcs) +##### +##### Field utils +##### + +# Canonical `similar` for Field (doesn't transfer boundary conditions) +function Base.similar(f::Field, grid=f.grid) + loc = location(f) + return Field(loc, + grid, + new_data(eltype(parent(f)), grid, loc), + FieldBoundaryConditions(grid, loc), + f.operand, + deepcopy(f.status)) end -# Default CPU architecture -Field(LX, LY, LZ, grid::AbstractGrid, args...) = Field(LX, LY, LZ, CPU(), grid, args...) +data(f::Field) = f.data -# Canonical `similar` for AbstractField (doesn't transfer boundary conditions) -Base.similar(f::AbstractField{LX, LY, LZ, Arch}) where {LX, LY, LZ, Arch} = Field(LX, LY, LZ, Arch(), f.grid) +"Returns a view of `f` that excludes halo points." +function interior(f::Field) + LX, LY, LZ = location(f) + TX, TY, TZ = topology(f.grid) + ii = interior_parent_indices(LX, TX, f.grid.Nx, f.grid.Hx) + jj = interior_parent_indices(LY, TY, f.grid.Ny, f.grid.Hy) + kk = interior_parent_indices(LZ, TZ, f.grid.Nz, f.grid.Hz) + return view(parent(f), ii, jj, kk) +end -# Type "destantiation": convert Face() to Face and Center() to Center if needed. -destantiate(X) = typeof(X) -destantiate(X::DataType) = X +# Don't use axes(f) to checkbounds; use axes(f.data) +Base.checkbounds(f::Field, I...) = Base.checkbounds(data(f), I...) +@propagate_inbounds Base.getindex(f::Field, inds...) = getindex(data(f), inds...) +@propagate_inbounds Base.getindex(f::Field, i::Int) = parent(f)[i] +@propagate_inbounds Base.setindex!(f::Field, val, i, j, k) = setindex!(data(f), val, i, j, k) +@propagate_inbounds Base.lastindex(f::Field) = lastindex(data(f)) +@propagate_inbounds Base.lastindex(f::Field, dim) = lastindex(data(f), dim) +Base.fill!(f::Field, val) = fill!(parent(f), val) -""" - Field(L::Tuple, arch, grid, data, bcs) +Base.isapprox(ϕ::Field, ψ::Field; kw...) = isapprox(interior(ϕ), interior(ψ); kw...) -Construct a `Field` at the location defined by the 3-tuple `L`, -whose elements are `Center` or `Face`. -""" -Field(L::Tuple, args...) = Field(destantiate.(L)..., args...) +Adapt.adapt_structure(to, field::Field) = Adapt.adapt(to, field.data) ##### ##### Special constructors for tracers and velocity fields ##### """ - CenterField([arch=CPU()], grid, args...) + CenterField(grid; kwargs...) Returns `Field{Center, Center, Center}` on `arch`itecture and `grid`. -Additional arguments are passed to the `Field` constructor. +Additional keyword arguments are passed to the `Field` constructor. """ -CenterField(arch::AbstractArchitecture, grid, args...) = Field(Center, Center, Center, arch, grid, args...) +CenterField(grid::AbstractGrid, FT=eltype(grid); kw...) = Field{Center, Center, Center}(grid, FT; kw...) """ - XFaceField([arch=CPU()], grid, args...) + XFaceField(grid; kw...) -Returns `Field{Face, Center, Center}` on `arch`itecture and `grid`. -Additional arguments are passed to the `Field` constructor. +Returns `Field{Face, Center, Center}` on `grid`. +Additional keyword arguments are passed to the `Field` constructor. """ -XFaceField(arch::AbstractArchitecture, args...) = Field(Face, Center, Center, arch, args...) +XFaceField(grid::AbstractGrid, FT=eltype(grid); kw...) = Field{Face, Center, Center}(grid, FT; kw...) """ - YFaceField([arch=CPU()], grid, args...) + YFaceField(grid; kw...) -Returns `Field{Center, Face, Center}` on `arch`itecture and `grid`. -Additional arguments are passed to the `Field` constructor. +Returns `Field{Center, Face, Center}` on `grid`. +Additional keyword arguments are passed to the `Field` constructor. """ -YFaceField(arch::AbstractArchitecture, args...) = Field(Center, Face, Center, arch, args...) +YFaceField(grid::AbstractGrid, FT=eltype(grid); kw...) = Field{Center, Face, Center}(grid, FT; kw...) """ - ZFaceField([arch=CPU()], grid, args...) + ZFaceField(grid; kw...) -Returns `Field{Center, Center, Face}` on `arch`itecture and `grid`. -Additional arguments are passed to the `Field` constructor. +Returns `Field{Center, Center, Face}` on `grid`. +Additional keyword arguments are passed to the `Field` constructor. """ -ZFaceField(arch::AbstractArchitecture, args...) = Field(Center, Center, Face, arch, args...) +ZFaceField(grid::AbstractGrid, FT=eltype(grid); kw...) = Field{Center, Center, Face}(grid, FT; kw...) -# Default CPU architectures... -CenterField(grid::AbstractGrid, args...) = CenterField(CPU(), grid, args...) - XFaceField(grid::AbstractGrid, args...) = XFaceField(CPU(), grid, args...) - YFaceField(grid::AbstractGrid, args...) = YFaceField(CPU(), grid, args...) - ZFaceField(grid::AbstractGrid, args...) = ZFaceField(CPU(), grid, args...) +##### +##### Fields computed from AbstractOperation and associated utilities +##### + +const ComputedField = Field{<:Any, <:Any, <:Any, <:AbstractOperation} + +""" + compute!(field) + +Computes `field.data` from `field.operand`. +""" +compute!(field) = nothing # fallback + +""" + @compute(exprs...) + +Call compute! on fields after defining them. +""" +macro compute(def) + expr = Expr(:block) + field = def.args[1] + push!(expr.args, :($(esc(def)))) + push!(expr.args, :(compute!($(esc(field))))) + return expr +end + +""" + compute_at!(field, time) + +Computes `field.data` at `time`. Falls back to compute!(field). +""" +compute_at!(field, time) = compute!(field) + +""" + conditional_compute!(field, time) + +Computes `field.data` if `time != field.status.time`. +""" +function conditional_compute!(field, time) + if time == zero(time) || time != field.status.time + compute!(field, time) + field.status.time = time + end + return nothing +end + +# This edge case occurs if `fetch_output` is called with `model::Nothing`. +# We do the safe thing here and always compute. +conditional_compute!(field, ::Nothing) = compute!(field, nothing) + +mutable struct FieldStatus{T} + time :: T +end + +Adapt.adapt_structure(to, status::FieldStatus) = (; time = status.time) + +""" + Field(operand::AbstractOperation; kwargs...) + +Return `f::Field` where `f.data` is computed from `f.operand` by +calling compute!(f). + +Keyword arguments +================= + +data (AbstractArray): An offset Array or CuArray for storing the result of a computation. + Must have `total_size(location(operand), grid)`. + +boundary_conditions (FieldBoundaryConditions): Boundary conditions for `f`. + +recompute_safely (Bool): whether or not to _always_ "recompute" `f` if `f` is + nested within another computation via an `AbstractOperation`. + If `data` is not provided then `recompute_safely=false` and + recomputation is _avoided_. If `data` is provided, then + `recompute_safely=true` by default. +""" +function Field(operand::AbstractOperation; + data = nothing, + boundary_conditions = FieldBoundaryConditions(op.grid, location(op)), + recompute_safely = true) + + if isnothing(data) + data = new_data(op.grid, location(op)) + recompute_safely = false + end + + status = recompute_safely ? nothing : FieldStatus(0.0) + + return Field(location(op), op.grid, data, boundary_conditions, operand, status) +end + +""" + compute!(comp::ComputedField) + +Compute `comp.operand` and store the result in `comp.data`. +""" +function compute!(comp::ComputedField, time=nothing) + # First compute `dependencies`: + compute_at!(comp.operand, time) + + workgroup, worksize = + work_layout(comp.grid, :xyz, include_right_boundaries = true, location = location(comp)) + + arch = architecture(comp) + compute_kernel! = _compute!(device(arch), workgroup, worksize) + event = compute_kernel!(comp.data, comp.operand; dependencies = device_event(arch)) + wait(device(arch), event) + + fill_halo_regions!(comp, arch) + + return comp +end + +"""Compute an `operand` and store in `data`.""" +@kernel function _compute!(data, operand) + i, j, k = @index(Global, NTuple) + @inbounds data[i, j, k] = operand[i, j, k] +end + +function compute_at!(field::ComputedField, time) + if isnothing(field.status) + return compute!(field, time) + else + return conditional_compute!(field, time) + end +end + +##### +##### Fields that are reduced along one or more dimensions +##### + +const XReducedField = Field{Nothing} +const YReducedField = Field{Any, Nothing} +const ZReducedField = Field{Any, Any, Nothing} + +const YZReducedField = Field{Any, Nothing, Nothing} +const XZReducedField = Field{Nothing, Any, Nothing} +const XYReducedField = Field{Nothing, Nothing, Any} + +const XYZReducedField = Field{Nothing, Nothing, Nothing} + +const ReducedField = Union{XReducedField, + YReducedField, + ZReducedField, + YZReducedField, + XZReducedField, + XYReducedField, + XYZReducedField} + +reduced_dimensions(field::Field) = () + +reduced_dimensions(field::XReducedField) = tuple(1) +reduced_dimensions(field::YReducedField) = tuple(2) +reduced_dimensions(field::ZReducedField) = tuple(3) + +reduced_dimensions(field::YZReducedField) = (2, 3) +reduced_dimensions(field::XZReducedField) = (1, 3) +reduced_dimensions(field::XYReducedField) = (1, 2) + +reduced_dimensions(field::XYZReducedField) = (1, 2, 3) + +function fill_halo_regions!(field::Field, arch, args...; kwargs...) + reduced_dims = reduced_dimensions(field) + if reduced_dimensions === () # the field is not reduced! + return fill_halo_regions!(field.data, + field.boundary_conditions, + architecture(field), + field.grid, + args...; + kwargs...) + else + return fill_halo_regions!(field.data, + field.boundary_conditions, + architecture(field), + field.grid, + args...; + reduced_dimensions, + kwargs...) + end +end + +@propagate_inbounds Base.getindex(r::XReducedField, i, j, k) = getindex(r.data, 1, j, k) +@propagate_inbounds Base.getindex(r::YReducedField, i, j, k) = getindex(r.data, i, 1, k) +@propagate_inbounds Base.getindex(r::ZReducedField, i, j, k) = getindex(r.data, i, j, 1) + +@propagate_inbounds Base.setindex!(r::XReducedField, v, i, j, k) = setindex!(r.data, v, 1, j, k) +@propagate_inbounds Base.setindex!(r::YReducedField, v, i, j, k) = setindex!(r.data, v, i, 1, k) +@propagate_inbounds Base.setindex!(r::ZReducedField, v, i, j, k) = setindex!(r.data, v, i, j, 1) + +@propagate_inbounds Base.getindex(r::YZReducedField, i, j, k) = getindex(r.data, i, 1, 1) +@propagate_inbounds Base.getindex(r::XZReducedField, i, j, k) = getindex(r.data, 1, j, 1) +@propagate_inbounds Base.getindex(r::XYReducedField, i, j, k) = getindex(r.data, 1, 1, k) + +@propagate_inbounds Base.setindex!(r::YZReducedField, v, i, j, k) = setindex!(r.data, v, i, 1, 1) +@propagate_inbounds Base.setindex!(r::XZReducedField, v, i, j, k) = setindex!(r.data, v, 1, j, 1) +@propagate_inbounds Base.setindex!(r::XYReducedField, v, i, j, k) = setindex!(r.data, v, 1, 1, k) + +@propagate_inbounds Base.getindex(r::XYZReducedField, i, j, k) = getindex(r.data, 1, 1, 1) +@propagate_inbounds Base.setindex!(r::XYZReducedField, v, i, j, k) = setindex!(r.data, v, 1, 1, 1) + +# Preserve location when adapting fields reduced on one or more dimensions +function Adapt.adapt_structure(to, reduced_field::ReducedField) + LX, LY, LZ = location(reduced_field) + return Field{LX, LY, LZ}(nothing, + adapt(to, reduced_field.data), + nothing, nothing, nothing) +end + +##### +##### Field reductions +##### + +# Risky to use these without tests. Docs would also be nice. +Statistics.dot(a::Field, b::Field) = mapreduce((x, y) -> x * y, +, interior(a), interior(b)) + +# TODO: In-place allocations with function mappings need to be fixed in Julia Base... +const SumReduction = typeof(Base.sum!) +const ProdReduction = typeof(Base.prod!) +const MaximumReduction = typeof(Base.maximum!) +const MinimumReduction = typeof(Base.minimum!) +const AllReduction = typeof(Base.all!) +const AnyReduction = typeof(Base.any!) + +initialize_reduced_field!(::SumReduction, f, r::ReducedField, c) = Base.initarray!(interior(r), Base.add_sum, true, interior(c)) +initialize_reduced_field!(::ProdReduction, f, r::ReducedField, c) = Base.initarray!(interior(r), Base.mul_prod, true, interior(c)) +initialize_reduced_field!(::AllReduction, f, r::ReducedField, c) = Base.initarray!(interior(r), &, true, interior(c)) +initialize_reduced_field!(::AnyReduction, f, r::ReducedField, c) = Base.initarray!(interior(r), |, true, interior(c)) + +initialize_reduced_field!(::MaximumReduction, f, r::ReducedField, c) = Base.mapfirst!(f, interior(r), c) +initialize_reduced_field!(::MinimumReduction, f, r::ReducedField, c) = Base.mapfirst!(f, interior(r), c) + +filltype(f, grid) = eltype(grid) +filltype(::Union{AllReduction, AnyReduction}, grid) = Bool + +reduced_location(loc; dims) = Tuple(i ∈ dims ? Nothing : loc[i] for i in 1:3) + +# Allocating and in-place reductions +for reduction in (:sum, :maximum, :minimum, :all, :any) + + reduction! = Symbol(reduction, '!') + + @eval begin + + # In-place + Base.$(reduction!)(f::Function, r::ReducedField, a::AbstractArray; kwargs...) = + Base.$(reduction!)(f, interior(r), a; kwargs...) + + Base.$(reduction!)(r::ReducedField, a::AbstractArray; kwargs...) = + Base.$(reduction!)(identity, interior(r), a; kwargs...) + + # Allocating + function Base.$(reduction)(f::Function, c::AbstractField; dims=:) + if dims === (:) + r = zeros(architecture(c), c.grid, 1, 1, 1) + Base.$(reduction!)(f, r, c) + return CUDA.@allowscalar r[1, 1, 1] + else + T = filltype(Base.$(reduction!), c.grid) + loc = reduced_location(location(c); dims) + r = Field(loc, c.grid, T) + initialize_reduced_field!(Base.$(reduction!), f, r, c) + Base.$(reduction!)(f, r, c, init=false) + return r + end + end + + Base.$(reduction)(c::AbstractField; dims=:) = Base.$(reduction)(identity, c; dims) + end +end + +Statistics._mean(f, c::AbstractField, ::Colon) = sum(f, c) / length(c) + +function Statistics._mean(f, c::AbstractField, dims) + r = sum(f, c; dims) + n = mapreduce(i -> size(c, i), *, unique(dims); init=1) + parent(r) ./= n + return r +end -Adapt.adapt_structure(to, field::Field) = Adapt.adapt(to, field.data) diff --git a/src/Fields/kernel_computed_field.jl b/src/Fields/kernel_computed_field.jl index c4de51bd67..96b55f43e0 100644 --- a/src/Fields/kernel_computed_field.jl +++ b/src/Fields/kernel_computed_field.jl @@ -2,7 +2,7 @@ using Oceananigans: AbstractModel using Oceananigans.Grids using Oceananigans.Utils: tupleit -struct KernelComputedField{X, Y, Z, A, S, D, G, T, K, B, F, P} <: AbstractDataField{X, Y, Z, A, G, T, 3} +struct KernelComputedField{X, Y, Z, A, S, D, G, T, K, B, F, P} <: AbstractField{X, Y, Z, A, G, T, 3} data :: D architecture :: A grid :: G diff --git a/src/Fields/reduced_field.jl b/src/Fields/reduced_field.jl index 36bcdeda7c..3a3e80b396 100644 --- a/src/Fields/reduced_field.jl +++ b/src/Fields/reduced_field.jl @@ -118,87 +118,3 @@ Base.similar(r::AbstractReducedField{X, Y, Z, Arch}) where {X, Y, Z, Arch} = reduced_location(loc; dims) = Tuple(i ∈ dims ? Nothing : loc[i] for i in 1:3) -Adapt.adapt_structure(to, reduced_field::ReducedField{X, Y, Z}) where {X, Y, Z} = - ReducedField{X, Y, Z}(adapt(to, reduced_field.data), nothing, adapt(to, reduced_field.grid), reduced_field.dims, nothing) - -##### -##### Field reductions -##### - -# Risky to use these without tests. Docs would also be nice. -Statistics.norm(a::AbstractDataField) = sqrt(mapreduce(x -> x * x, +, interior(a))) -Statistics.dot(a::AbstractDataField, b::AbstractDataField) = mapreduce((x, y) -> x * y, +, interior(a), interior(b)) - -# The more general case, for AbstractOperations -function Statistics.norm(a::AbstractField) - arch = architecture(a) - grid = a.grid - - r = zeros(arch, grid, 1) - - Base.mapreducedim!(x -> x * x, +, r, a) - - return CUDA.@allowscalar sqrt(r[1]) -end - -# TODO: In-place allocations with function mappings need to be fixed in Julia Base... -const SumReduction = typeof(Base.sum!) -const ProdReduction = typeof(Base.prod!) -const MaximumReduction = typeof(Base.maximum!) -const MinimumReduction = typeof(Base.minimum!) -const AllReduction = typeof(Base.all!) -const AnyReduction = typeof(Base.any!) - -initialize_reduced_field!(::SumReduction, f, r::AbstractReducedField, c) = Base.initarray!(interior(r), Base.add_sum, true, interior(c)) -initialize_reduced_field!(::ProdReduction, f, r::AbstractReducedField, c) = Base.initarray!(interior(r), Base.mul_prod, true, interior(c)) -initialize_reduced_field!(::AllReduction, f, r::AbstractReducedField, c) = Base.initarray!(interior(r), &, true, interior(c)) -initialize_reduced_field!(::AnyReduction, f, r::AbstractReducedField, c) = Base.initarray!(interior(r), |, true, interior(c)) - -initialize_reduced_field!(::MaximumReduction, f, r::AbstractReducedField, c) = Base.mapfirst!(f, interior(r), c) -initialize_reduced_field!(::MinimumReduction, f, r::AbstractReducedField, c) = Base.mapfirst!(f, interior(r), c) - -filltype(f, grid) = eltype(grid) -filltype(::Union{AllReduction, AnyReduction}, grid) = Bool - -# Allocating and in-place reductions -for reduction in (:sum, :maximum, :minimum, :all, :any) - - reduction! = Symbol(reduction, '!') - - @eval begin - - # In-place - Base.$(reduction!)(f::Function, r::AbstractReducedField, a::AbstractArray; kwargs...) = - Base.$(reduction!)(f, interior(r), a; kwargs...) - - Base.$(reduction!)(r::AbstractReducedField, a::AbstractArray; kwargs...) = - Base.$(reduction!)(identity, interior(r), a; kwargs...) - - # Allocating - function Base.$(reduction)(f::Function, c::AbstractField; dims=:) - if dims === (:) - r = zeros(architecture(c), c.grid, 1, 1, 1) - Base.$(reduction!)(f, r, c) - return CUDA.@allowscalar r[1, 1, 1] - else - FT = filltype(Base.$(reduction!), c.grid) - r = ReducedField(FT, location(c), architecture(c), c.grid; dims) - initialize_reduced_field!(Base.$(reduction!), f, r, c) - Base.$(reduction!)(f, r, c, init=false) - return r - end - end - - Base.$(reduction)(c::AbstractField; dims=:) = Base.$(reduction)(identity, c; dims) - end -end - -Statistics._mean(f, c::AbstractField, ::Colon) = sum(f, c) / length(c) - -function Statistics._mean(f, c::AbstractField, dims) - r = sum(f, c; dims) - n = mapreduce(i -> size(c, i), *, unique(dims); init=1) - parent(r) ./= n - return r -end - diff --git a/src/Fields/set!.jl b/src/Fields/set!.jl index 4d5a3bde5a..ad22c8122c 100644 --- a/src/Fields/set!.jl +++ b/src/Fields/set!.jl @@ -18,13 +18,6 @@ set!(u::AbstractField, v) = u .= v # fallback # Niceties const AbstractCPUField = AbstractField{X, Y, Z, <:CPU} where {X, Y, Z} -const AbstractReducedCPUField = AbstractReducedField{X, Y, Z, <:CPU} where {X, Y, Z} - -""" Returns an AbstractReducedField on the CPU. """ -function similar_cpu_field(u::AbstractReducedField) - FieldType = typeof(u).name.wrapper - return FieldType(location(u), CPU(), u.grid; dims=u.dims) -end """ Set the CPU field `u` data to the function `f(x, y, z)`. """ function set!(u::AbstractCPUField, f::Function) @@ -38,17 +31,11 @@ end ##### const AbstractGPUField = AbstractField{X, Y, Z, <:GPU} where {X, Y, Z} -const AbstractReducedGPUField = AbstractReducedField{X, Y, Z, <:GPU} where {X, Y, Z} - -""" Returns a field on the CPU with `nothing` boundary conditions. """ -function similar_cpu_field(u) - FieldType = typeof(u).name.wrapper - return FieldType(location(u), CPU(), adapt_structure(CPU(), u.grid), nothing) -end """ Set the GPU field `u` to the array or function `v`. """ function set!(u::AbstractGPUField, v::Union{Array, Function}) - v_field = similar_cpu_field(u) + cpu_grid = adapt_structure(CPU(), u.grid) + v_field = similar(u, cpu_grid) set!(v_field, v) set!(u, v_field) return nothing diff --git a/src/Grids/new_data.jl b/src/Grids/new_data.jl index e88b4f181d..e4366a59f1 100644 --- a/src/Grids/new_data.jl +++ b/src/Grids/new_data.jl @@ -50,19 +50,18 @@ offset_data(underlying_data, grid::AbstractGrid, loc) = offset_data(underlying_data, loc, topology(grid), size(grid), halo_size(grid)) """ - new_data([FT=Float64], arch, grid, loc) + new_data(FT, grid, loc) 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, arch, grid, loc) - underlying_data = zeros(FT, arch, - 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)) - +function new_data(FT, grid, loc) + arch = architecture(grid) + Tx = total_length(loc[1], topology(grid, 1), grid.Nx, grid.Hx) + Ty = total_length(loc[2], topology(grid, 2), grid.Ny, grid.Hy) + Tz = total_length(loc[3], topology(grid, 3), grid.Nz, grid.Hz) + underlying_data = zeros(FT, arch, Tx, Ty, Tz) return offset_data(underlying_data, grid, loc) end -# Default to type of Grid -new_data(arch, grid, loc) = new_data(eltype(grid), arch, grid, loc) +new_data(grid, loc) = new_data(eltype(grid), grid, loc) diff --git a/src/ImmersedBoundaries/mask_immersed_field.jl b/src/ImmersedBoundaries/mask_immersed_field.jl index 68ef0e33e7..03b7f40b5a 100644 --- a/src/ImmersedBoundaries/mask_immersed_field.jl +++ b/src/ImmersedBoundaries/mask_immersed_field.jl @@ -1,16 +1,16 @@ using KernelAbstractions using Statistics using Oceananigans.Architectures: architecture, device_event -using Oceananigans.Fields: location, AbstractReducedField +using Oceananigans.Fields: location, ZReducedField, Field instantiate(X) = X() -mask_immersed_field!(field::AbstractDataField, value=zero(eltype(field.grid))) = +mask_immersed_field!(field::Field, value=zero(eltype(field.grid))) = mask_immersed_field!(field, field.grid, location(field), value) mask_immersed_field!(field, grid, loc, value) = NoneEvent() -function mask_immersed_field!(field::AbstractDataField, grid::ImmersedBoundaryGrid, loc, value) +function mask_immersed_field!(field::Field, grid::ImmersedBoundaryGrid, loc, value) arch = architecture(field) loc = instantiate.(loc) return launch!(arch, grid, :xyz, _mask_immersed_field!, field, loc, grid, value; dependencies = device_event(arch)) @@ -21,14 +21,14 @@ end @inbounds field[i, j, k] = scalar_mask(i, j, k, grid, grid.immersed_boundary, loc..., value, field) end -mask_immersed_reduced_field_xy!(field::AbstractReducedField, value=zero(eltype(field.grid)); k) = +mask_immersed_reduced_field_xy!(field::ZReducedField, value=zero(eltype(field.grid)); k) = mask_immersed_reduced_field_xy!(field, field.grid, location(field), value; k) mask_immersed_reduced_field_xy!(::Nothing, args...; kwargs...) = NoneEvent() mask_immersed_reduced_field_xy!(field, grid, loc, value; k) = NoneEvent() -function mask_immersed_reduced_field_xy!(field::AbstractReducedField, grid::ImmersedBoundaryGrid, loc, value; k) +function mask_immersed_reduced_field_xy!(field::ZReducedField, grid::ImmersedBoundaryGrid, loc, value; k) arch = architecture(field) loc = instantiate.(loc) return launch!(arch, grid, :xy, _mask_immersed_reduced_field_xy!, field, loc, grid, value, k; dependencies = device_event(arch)) diff --git a/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl b/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl index a45c778a5f..d8d685477c 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl @@ -1,7 +1,7 @@ using Oceananigans.Solvers using Oceananigans.Operators using Oceananigans.Architectures -using Oceananigans.Fields: ReducedField +using Oceananigans.Fields: Field, ZReducedField import Oceananigans.Solvers: solve! @@ -26,14 +26,14 @@ step `Δt`, gravitational acceleration `g`, and free surface at time-step `n` ` """ function PCGImplicitFreeSurfaceSolver(arch::AbstractArchitecture, grid, settings) # Initialize vertically integrated lateral face areas - ∫ᶻ_Axᶠᶜᶜ = ReducedField(Face, Center, Nothing, arch, grid; dims=3) - ∫ᶻ_Ayᶜᶠᶜ = ReducedField(Center, Face, Nothing, arch, grid; dims=3) + ∫ᶻ_Axᶠᶜᶜ = Field{Face, Center, Nothing}(grid) + ∫ᶻ_Ayᶜᶠᶜ = Field{Center, Face, Nothing}(grid) vertically_integrated_lateral_areas = (xᶠᶜᶜ = ∫ᶻ_Axᶠᶜᶜ, yᶜᶠᶜ = ∫ᶻ_Ayᶜᶠᶜ) compute_vertically_integrated_lateral_areas!(vertically_integrated_lateral_areas, grid, arch) - right_hand_side = ReducedField(Center, Center, Nothing, arch, grid; dims=3) + right_hand_side = Field{Center, Center, Nothing}(grid) # Set maximum iterations to Nx * Ny if not set settings = Dict{Symbol, Any}(settings) @@ -131,7 +131,7 @@ end Compute the horizontal divergence of vertically-uniform quantity using vertically-integrated face areas `∫ᶻ_Axᶠᶜᶜ` and `∫ᶻ_Ayᶜᶠᶜ`. """ -@inline Az_∇h²ᶜᶜᵃ(i, j, k, grid, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, η::ReducedField{X, Y, Nothing}) where {X, Y} = +@inline Az_∇h²ᶜᶜᵃ(i, j, k, grid, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, η::ZReducedField) = (δxᶜᵃᵃ(i, j, k, grid, ∫ᶻ_Ax_∂x_ηᶠᶜᶜ, ∫ᶻ_Axᶠᶜᶜ, η) + δyᵃᶜᵃ(i, j, k, grid, ∫ᶻ_Ay_∂y_ηᶜᶠᶜ, ∫ᶻ_Ayᶜᶠᶜ, η)) diff --git a/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl b/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl index ccf3d89274..1fcfada069 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl @@ -47,7 +47,7 @@ PrescribedVelocityFields(; u=zerofunc, v=zerofunc, w=zerofunc, parameters=nothin PrescribedField(X, Y, Z, f::Function, grid; kwargs...) = FunctionField{X, Y, Z}(f, grid; kwargs...) PrescribedField(X, Y, Z, f::AbstractField, grid; kwargs...) = f -function PrescribedField(X, Y, Z, f::AbstractDataField, grid; kwargs...) +function PrescribedField(X, Y, Z, f::Field, grid; kwargs...) fill_halo_regions!(f, architecture(f)) return f end diff --git a/src/OutputReaders/field_time_series.jl b/src/OutputReaders/field_time_series.jl index 23cc6c90b9..620943edf4 100644 --- a/src/OutputReaders/field_time_series.jl +++ b/src/OutputReaders/field_time_series.jl @@ -13,7 +13,7 @@ using Oceananigans.Fields: show_location import Oceananigans: short_show import Oceananigans.Fields: Field, set!, interior -struct FieldTimeSeries{X, Y, Z, K, A, T, D, G, B, χ} <: AbstractDataField{X, Y, Z, A, G, T, 4} +struct FieldTimeSeries{X, Y, Z, K, A, T, D, G, B, χ} <: AbstractField{X, Y, Z, A, G, T, 4} data :: D architecture :: A grid :: G diff --git a/test/test_field.jl b/test/test_field.jl index ccceb2bc85..69da8d6318 100644 --- a/test/test_field.jl +++ b/test/test_field.jl @@ -1,24 +1,21 @@ -using Statistics -using Oceananigans -using Oceananigans.Architectures: array_type, arch_array -using Oceananigans.Fields: cpudata, FieldSlicer, interior_copy, regrid!, ReducedField, has_velocities, VelocityFields, TracerFields, interpolate - -include("utils_for_runtests.jl") +include("dependencies_for_runtests.jl") -archs = test_architectures() +using Oceananigans.Fields: cpudata, FieldSlicer, interior_copy, regrid!, ReducedField, has_velocities, VelocityFields, TracerFields, interpolate """ - correct_field_size(arch, grid, FieldType, Tx, Ty, Tz) + correct_field_size(grid, FieldType, Tx, Ty, Tz) -Test that the field initialized by the FieldType constructor on `arch` and `grid` +Test that the field initialized by the FieldType constructor on `grid` has size `(Tx, Ty, Tz)`. """ -correct_field_size(a, g, FieldType, Tx, Ty, Tz) = size(parent(FieldType(a, g))) == (Tx, Ty, Tz) +correct_field_size(g, FieldType, Tx, Ty, Tz) = size(parent(FieldType(g))) == (Tx, Ty, Tz) function run_similar_field_tests(f) g = similar(f) @test typeof(f) == typeof(g) @test f.grid == g.grid + @test location(f) === location(g) + @test !(f.data === g.data) return nothing end @@ -29,8 +26,9 @@ Test that the field initialized by the field type function `ftf` on the grid g can be correctly filled with the value `val` using the `set!(f::AbstractField, v)` function. """ -function correct_field_value_was_set(arch, grid, FieldType, val::Number) - f = FieldType(arch, grid) +function correct_field_value_was_set(grid, FieldType, val::Number) + arch = architecture(grid) + f = FieldType(grid) set!(f, val) return all(interior(f) .≈ val * arch_array(arch, ones(size(f)))) end @@ -40,10 +38,10 @@ function run_field_reduction_tests(FT, arch) topo = (Bounded, Bounded, Bounded) grid = RectilinearGrid(arch, FT, topology=topo, size=(N, N, N), x=(-1, 1), y=(0, 2π), z=(-1, 1)) - u = XFaceField(arch, grid) - v = YFaceField(arch, grid) - w = ZFaceField(arch, grid) - c = CenterField(arch, grid) + u = XFaceField(grid) + v = YFaceField(grid) + w = ZFaceField(grid) + c = CenterField(grid) f(x, y, z) = 1 + exp(x) * sin(y) * tanh(z) @@ -107,12 +105,12 @@ function run_field_reduction_tests(FT, arch) return nothing end -function run_field_interpolation_tests(arch, FT) +function run_field_interpolation_tests(FT, arch) grid = RectilinearGrid(arch, size=(4, 5, 7), x=(0, 1), y=(-π, π), z=(-5.3, 2.7)) - velocities = VelocityFields(arch, grid) - tracers = TracerFields((:c,), arch, grid) + velocities = VelocityFields(grid) + tracers = TracerFields((:c,), grid) (u, v, w), c = velocities, tracers.c @@ -163,6 +161,10 @@ function run_field_interpolation_tests(arch, FT) return nothing end +##### +##### +##### + @testset "Fields" begin @info "Testing Fields..." @@ -175,28 +177,28 @@ end for arch in archs, FT in float_types grid = RectilinearGrid(arch , FT, size=N, extent=L, halo=H, topology=(Periodic, Periodic, Periodic)) - @test correct_field_size(arch, grid, CenterField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(arch, grid, XFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(arch, grid, YFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(arch, grid, ZFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, CenterField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, XFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, YFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, ZFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) grid = RectilinearGrid(arch, FT, size=N, extent=L, halo=H, topology=(Periodic, Periodic, Bounded)) - @test correct_field_size(arch, grid, CenterField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(arch, grid, XFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(arch, grid, YFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(arch, grid, ZFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3] + 1) + @test correct_field_size(grid, CenterField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, XFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, YFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, ZFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3] + 1) grid = RectilinearGrid(arch, FT, size=N, extent=L, halo=H, topology=(Periodic, Bounded, Bounded)) - @test correct_field_size(arch, grid, CenterField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(arch, grid, XFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(arch, grid, YFaceField, N[1] + 2 * H[1], N[2] + 1 + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(arch, grid, ZFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 1 + 2 * H[3]) + @test correct_field_size(grid, CenterField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, XFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, YFaceField, N[1] + 2 * H[1], N[2] + 1 + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, ZFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 1 + 2 * H[3]) grid = RectilinearGrid(arch, FT, size=N, extent=L, halo=H, topology=(Bounded, Bounded, Bounded)) - @test correct_field_size(arch, grid, CenterField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(arch, grid, XFaceField, N[1] + 1 + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(arch, grid, YFaceField, N[1] + 2 * H[1], N[2] + 1 + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(arch, grid, ZFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 1 + 2 * H[3]) + @test correct_field_size(grid, CenterField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, XFaceField, N[1] + 1 + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, YFaceField, N[1] + 2 * H[1], N[2] + 1 + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, ZFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 1 + 2 * H[3]) end end @@ -223,11 +225,11 @@ end grid = RectilinearGrid(arch, FT, size=N, extent=L, topology=(Periodic, Periodic, Bounded)) for FieldType in FieldTypes, val in vals - @test correct_field_value_was_set(arch, grid, FieldType, val) + @test correct_field_value_was_set(grid, FieldType, val) end for FieldType in FieldTypes - field = FieldType(arch, grid) + field = FieldType( grid) sz = size(field) A = rand(FT, sz...) set!(field, A) @@ -238,10 +240,10 @@ end topo = (Bounded, Bounded, Bounded) grid = RectilinearGrid(arch, FT, topology=topo, size=(Nx, Nx, Nx), x=(-1, 1), y=(0, 2π), z=(-1, 1)) - u = XFaceField(arch, grid) - v = YFaceField(arch, grid) - w = ZFaceField(arch, grid) - c = CenterField(arch, grid) + u = XFaceField(grid) + v = YFaceField(grid) + w = ZFaceField(grid) + c = CenterField(grid) f(x, y, z) = exp(x) * sin(y) * tanh(z) @@ -267,7 +269,7 @@ end @info " Testing field interpolation..." for arch in archs, FT in float_types - run_field_interpolation_tests(arch, FT) + run_field_interpolation_tests(FT, arch) end end @@ -296,17 +298,19 @@ end for X in (Center, Face), Y in (Center, Face), Z in (Center, Face) for arch in archs - f = Field(X, Y, Z, arch, grid) + f = Field{X, Y, Z}(grid) run_similar_field_tests(f) for dims in (3, (1, 2), (1, 2, 3)) - f = ReducedField(X, Y, Z, arch, grid, dims=dims) + loc = reduced_location(loc; dims) + f = Field(loc, grid) run_similar_field_tests(f) end end end end + #= @testset "Regridding" begin @info " Testing field regridding..." @@ -381,4 +385,5 @@ end end end end + =# end From 4ffaeb211441b7e26756493004eec903fdaf70ad Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 15 Dec 2021 10:27:43 -0700 Subject: [PATCH 002/140] test_field success --- src/Fields/field.jl | 17 +++---- src/Fields/field_tuples.jl | 96 +++++++++++++++++++------------------- test/test_field.jl | 41 +++++++++------- 3 files changed, 81 insertions(+), 73 deletions(-) diff --git a/src/Fields/field.jl b/src/Fields/field.jl index c6302ed65c..d227a92a88 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -34,7 +34,7 @@ end ##### """ - Field{LX, LY, LZ}(grid; kwargs...) + Field{LX, LY, LZ}(grid; kw...) Construct a `Field` on `grid` at the location `(LX, LY, LZ)`. Each of `(LX, LY, LZ)` is either `Center` or `Face` and determines @@ -58,11 +58,12 @@ Field located at (Face, Face, Center) └── boundary conditions: west=Periodic, east=Periodic, south=Periodic, north=Periodic, bottom=ZeroFlux, top=ZeroFlux, immersed=ZeroFlux ``` """ -Field{LX, LY, LZ}(grid::AbstractGrid, T; kwargs...) where {LX, LY, LZ} = Field((LX, LY, LZ), grid, T; kwargs...) +Field{LX, LY, LZ}(grid::AbstractGrid, T::DataType=eltype(grid); kw...) where {LX, LY, LZ} = + Field((LX, LY, LZ), grid, T; kw...) function Field(loc::Tuple, grid::AbstractGrid, - T = eltype(grid); + T::DataType = eltype(grid); data = new_data(T, grid, loc), boundary_conditions = FieldBoundaryConditions(grid, loc)) @@ -277,12 +278,12 @@ end ##### const XReducedField = Field{Nothing} -const YReducedField = Field{Any, Nothing} -const ZReducedField = Field{Any, Any, Nothing} +const YReducedField = Field{<:Any, Nothing} +const ZReducedField = Field{<:Any, <:Any, Nothing} -const YZReducedField = Field{Any, Nothing, Nothing} -const XZReducedField = Field{Nothing, Any, Nothing} -const XYReducedField = Field{Nothing, Nothing, Any} +const YZReducedField = Field{<:Any, Nothing, Nothing} +const XZReducedField = Field{Nothing, <:Any, Nothing} +const XYReducedField = Field{Nothing, Nothing, <:Any} const XYZReducedField = Field{Nothing, Nothing, Nothing} diff --git a/src/Fields/field_tuples.jl b/src/Fields/field_tuples.jl index c0209fef36..19712c40e1 100644 --- a/src/Fields/field_tuples.jl +++ b/src/Fields/field_tuples.jl @@ -5,13 +5,13 @@ using Oceananigans.BoundaryConditions: FieldBoundaryConditions, regularize_field ##### """ - VelocityFields(arch, grid, user_bcs = NamedTuple()) + VelocityFields(grid, user_bcs = NamedTuple()) -Return a `NamedTuple` with fields `u`, `v`, `w` initialized on the architecture `arch` -and `grid`. Boundary conditions `bcs` may be specified via a named tuple of +Return a `NamedTuple` with fields `u`, `v`, `w` initialized on `grid`. +Boundary conditions `bcs` may be specified via a named tuple of `FieldBoundaryCondition`s. """ -function VelocityFields(arch, grid, user_bcs = NamedTuple()) +function VelocityFields(grid, user_bcs = NamedTuple()) template = FieldBoundaryConditions() @@ -23,9 +23,9 @@ function VelocityFields(arch, grid, user_bcs = NamedTuple()) bcs = merge(default_bcs, user_bcs) - u = XFaceField(arch, grid, bcs.u) - v = YFaceField(arch, grid, bcs.v) - w = ZFaceField(arch, grid, bcs.w) + u = XFaceField(grid, boundary_conditions=bcs.u) + v = YFaceField(grid, boundary_conditions=bcs.v) + w = ZFaceField(grid, boundary_conditions=bcs.w) return (u=u, v=v, w=w) end @@ -35,49 +35,49 @@ end ##### """ - TracerFields(tracer_names, arch, grid, user_bcs) + TracerFields(tracer_names, grid, user_bcs) Return a `NamedTuple` with tracer fields specified by `tracer_names` initialized as -`CenterField`s on the architecture `arch` and `grid`. Boundary conditions `user_bcs` +`CenterField`s on `grid`. Boundary conditions `user_bcs` may be specified via a named tuple of `FieldBoundaryCondition`s. """ -function TracerFields(tracer_names, arch, grid, user_bcs) +function TracerFields(tracer_names, grid, user_bcs) default_bcs = NamedTuple(name => FieldBoundaryConditions(grid, (Center, Center, Center)) for name in tracer_names) bcs = merge(default_bcs, user_bcs) # provided bcs overwrite defaults - return NamedTuple(c => CenterField(arch, grid, bcs[c]) for c in tracer_names) + return NamedTuple(c => CenterField(grid, boundary_conditions=bcs[c]) for c in tracer_names) end """ - TracerFields(tracer_names, arch, grid; kwargs...) + TracerFields(tracer_names, grid; kwargs...) Return a `NamedTuple` with tracer fields specified by `tracer_names` initialized as -`CenterField`s on the architecture `arch` and `grid`. Fields may be passed via optional +`CenterField`s on `grid`. Fields may be passed via optional keyword arguments `kwargs` for each field. This function is used by `OutputWriters.Checkpointer` and `TendencyFields`. ``` """ -TracerFields(tracer_names, arch, grid; kwargs...) = - NamedTuple(c => c ∈ keys(kwargs) ? kwargs[c] : CenterField(arch, grid) for c in tracer_names) +TracerFields(tracer_names, grid; kwargs...) = + NamedTuple(c => c ∈ keys(kwargs) ? kwargs[c] : CenterField(grid) for c in tracer_names) # 'Nothing', or empty tracer fields -TracerFields(::Union{Tuple{}, Nothing}, arch, grid, bcs) = NamedTuple() +TracerFields(::Union{Tuple{}, Nothing}, grid, bcs) = NamedTuple() "Shortcut constructor for empty tracer fields." -TracerFields(::NamedTuple{(), Tuple{}}, arch, grid, bcs) = NamedTuple() +TracerFields(::NamedTuple{(), Tuple{}}, grid, bcs) = NamedTuple() ##### ##### Pressure fields tuples ##### """ - PressureFields(arch, grid, bcs::NamedTuple) + PressureFields(grid, bcs::NamedTuple) Return a `NamedTuple` with pressure fields `pHY′` and `pNHS` initialized as -`CenterField`s on the architecture `arch` and `grid`. Boundary conditions `bcs` +`CenterField`s on `grid`. Boundary conditions `bcs` may be specified via a named tuple of `FieldBoundaryCondition`s. """ -function PressureFields(arch, grid, bcs=NamedTuple()) +function PressureFields(grid, bcs=NamedTuple()) default_pressure_boundary_conditions = (pHY′ = FieldBoundaryConditions(grid, (Center, Center, Center)), @@ -85,43 +85,43 @@ function PressureFields(arch, grid, bcs=NamedTuple()) bcs = merge(default_pressure_boundary_conditions, bcs) - pHY′ = CenterField(arch, grid, bcs.pHY′) - pNHS = CenterField(arch, grid, bcs.pNHS) + pHY′ = CenterField(grid, boundary_conditions=bcs.pHY′) + pNHS = CenterField(grid, boundary_conditions=bcs.pNHS) return (pHY′=pHY′, pNHS=pNHS) end -function PressureFields(arch, grid::AbstractGrid{<:Any, <:Any, <:Any, <:Flat}, bcs=NamedTuple()) +function PressureFields(grid::AbstractGrid{<:Any, <:Any, <:Any, <:Flat}, bcs=NamedTuple()) default_pressure_boundary_conditions = (pHY′ = FieldBoundaryConditions(grid, (Center, Center, Center)), pNHS = FieldBoundaryConditions(grid, (Center, Center, Center))) bcs = merge(default_pressure_boundary_conditions, bcs) - pNHS = CenterField(arch, grid, bcs.pNHS) + pNHS = CenterField(grid, boundary_conditions=bcs.pNHS) return (; pHY′=nothing, pNHS=pNHS) end """ - TendencyFields(arch, grid, tracer_names; - u = XFaceField(arch, grid), - v = YFaceField(arch, grid), - w = ZFaceField(arch, grid), + TendencyFields(grid, tracer_names; + u = XFaceField(grid), + v = YFaceField(grid), + w = ZFaceField(grid), kwargs...) Return a `NamedTuple` with tendencies for all solution fields (velocity fields and -tracer fields), initialized on the architecture `arch` and `grid`. Optional `kwargs` +tracer fields), initialized on `grid`. Optional `kwargs` can be specified to assign data arrays to each tendency field. """ -function TendencyFields(arch, grid, tracer_names; - u = XFaceField(arch, grid), - v = YFaceField(arch, grid), - w = ZFaceField(arch, grid), +function TendencyFields(grid, tracer_names; + u = XFaceField(grid), + v = YFaceField(grid), + w = ZFaceField(grid), kwargs...) velocities = (u=u, v=v, w=w) - tracers = TracerFields(tracer_names, arch, grid; kwargs...) + tracers = TracerFields(tracer_names, grid; kwargs...) return merge(velocities, tracers) end @@ -130,53 +130,53 @@ end ##### Helper functions for NonhydrostaticModel constructor ##### -VelocityFields(::Nothing, arch, grid, bcs) = VelocityFields(arch, grid, bcs) -PressureFields(::Nothing, arch, grid, bcs) = PressureFields(arch, grid, bcs) +VelocityFields(::Nothing, grid, bcs) = VelocityFields(grid, bcs) +PressureFields(::Nothing, grid, bcs) = PressureFields(grid, bcs) """ - VelocityFields(proposed_velocities::NamedTuple{(:u, :v, :w)}, arch, grid, bcs) + VelocityFields(proposed_velocities::NamedTuple{(:u, :v, :w)}, grid, bcs) Return a `NamedTuple` of velocity fields, overwriting boundary conditions in `proposed_velocities` with corresponding fields in the `NamedTuple` `bcs`. """ -function VelocityFields(proposed_velocities::NamedTuple{(:u, :v, :w)}, arch, grid, bcs) +function VelocityFields(proposed_velocities::NamedTuple{(:u, :v, :w)}, grid, bcs) validate_field_tuple_grid("velocities", proposed_velocities, grid) - u = XFaceField(arch, grid, bcs.u, proposed_velocities.u.data) - v = YFaceField(arch, grid, bcs.v, proposed_velocities.v.data) - w = ZFaceField(arch, grid, bcs.w, proposed_velocities.w.data) + u = XFaceField(grid, boundary_conditions=bcs.u, data=proposed_velocities.u.data) + v = YFaceField(grid, boundary_conditions=bcs.v, data=proposed_velocities.v.data) + w = ZFaceField(grid, boundary_conditions=bcs.w, data=proposed_velocities.w.data) return (u=u, v=v, w=w) end """ - TracerFields(proposed_tracers::NamedTuple, arch, grid, bcs) + TracerFields(proposed_tracers::NamedTuple, grid, bcs) Return a `NamedTuple` of tracers, overwriting boundary conditions in `proposed_tracers` with corresponding fields in the `NamedTuple` `bcs`. """ -function TracerFields(proposed_tracers::NamedTuple, arch, grid, bcs) +function TracerFields(proposed_tracers::NamedTuple, grid, bcs) validate_field_tuple_grid("tracers", proposed_tracers, grid) tracer_names = propertynames(proposed_tracers) - tracer_fields = Tuple(CenterField(arch, grid, bcs[c], proposed_tracers[c].data) for c in tracer_names) + tracer_fields = Tuple(CenterField(grid, boundary_conditions=bcs[c], data=proposed_tracers[c].data) for c in tracer_names) return NamedTuple{tracer_names}(tracer_fields) end """ - PressureFields(proposed_pressures::NamedTuple{(:pHY′, :pNHS)}, arch, grid, bcs) + PressureFields(proposed_pressures::NamedTuple{(:pHY′, :pNHS)}, grid, bcs) Return a `NamedTuple` of pressure fields with, overwriting boundary conditions in `proposed_tracer_fields` with corresponding fields in the `NamedTuple` `bcs`. """ -function PressureFields(proposed_pressures::NamedTuple{(:pHY′, :pNHS)}, arch, grid, bcs) +function PressureFields(proposed_pressures::NamedTuple{(:pHY′, :pNHS)}, grid, bcs) validate_field_tuple_grid("pressures", proposed_pressures, grid) - pHY′ = CenterField(arch, grid, bcs.pHY′, proposed_pressures.pHY′.data) - pNHS = CenterField(arch, grid, bcs.pNHS, proposed_pressures.pNHS.data) + pHY′ = CenterField(grid, boundary_conditions=bcs.pHY′, data=proposed_pressures.pHY′.data) + pNHS = CenterField(grid, boundary_conditions=bcs.pNHS, data=proposed_pressures.pNHS.data) return (pHY′=pHY′, pNHS=pNHS) end diff --git a/test/test_field.jl b/test/test_field.jl index 69da8d6318..24d11b2aed 100644 --- a/test/test_field.jl +++ b/test/test_field.jl @@ -1,6 +1,9 @@ include("dependencies_for_runtests.jl") -using Oceananigans.Fields: cpudata, FieldSlicer, interior_copy, regrid!, ReducedField, has_velocities, VelocityFields, TracerFields, interpolate +using Oceananigans.Fields: cpudata, FieldSlicer, interior_copy +using Oceananigans.Fields: regrid!, ReducedField, has_velocities +using Oceananigans.Fields: VelocityFields, TracerFields, interpolate +using Oceananigans.Fields: reduced_location """ correct_field_size(grid, FieldType, Tx, Ty, Tz) @@ -229,7 +232,7 @@ end end for FieldType in FieldTypes - field = FieldType( grid) + field = FieldType(grid) sz = size(field) A = rand(FT, sz...) set!(field, A) @@ -250,10 +253,15 @@ end ϕs = (u, v, w, c) [set!(ϕ, f) for ϕ in ϕs] - @test u[1, 2, 3] == f(grid.xᶠᵃᵃ[1], grid.yᵃᶜᵃ[2], grid.zᵃᵃᶜ[3]) - @test v[1, 2, 3] == f(grid.xᶜᵃᵃ[1], grid.yᵃᶠᵃ[2], grid.zᵃᵃᶜ[3]) - @test w[1, 2, 3] == f(grid.xᶜᵃᵃ[1], grid.yᵃᶜᵃ[2], grid.zᵃᵃᶠ[3]) - @test c[1, 2, 3] == f(grid.xᶜᵃᵃ[1], grid.yᵃᶜᵃ[2], grid.zᵃᵃᶜ[3]) + xu, yu, zu = nodes(u) + xv, yv, zv = nodes(v) + xw, yw, zw = nodes(w) + xc, yc, zc = nodes(c) + + @test u[1, 2, 3] == f(xu[1], yu[2], zu[3]) + @test v[1, 2, 3] == f(xv[1], yv[2], zv[3]) + @test w[1, 2, 3] == f(xw[1], yw[2], zw[3]) + @test c[1, 2, 3] == f(xc[1], yc[2], zc[3]) end end @@ -282,11 +290,12 @@ end @test has_velocities((:u, :v, :w)) == true grid = RectilinearGrid(CPU(), size=(4, 6, 8), extent=(1, 1, 1)) - ϕ = CenterField(CPU(), grid) + ϕ = CenterField(grid) @test cpudata(ϕ).parent isa Array if CUDA.has_cuda() - ϕ = CenterField(GPU(), grid) + grid = RectilinearGrid(GPU(), size=(4, 6, 8), extent=(1, 1, 1)) + ϕ = CenterField(grid) @test cpudata(ϕ).parent isa Array end @@ -302,7 +311,7 @@ end run_similar_field_tests(f) for dims in (3, (1, 2), (1, 2, 3)) - loc = reduced_location(loc; dims) + loc = reduced_location((X, Y, Z); dims) f = Field(loc, grid) run_similar_field_tests(f) end @@ -310,7 +319,6 @@ end end end - #= @testset "Regridding" begin @info " Testing field regridding..." @@ -326,12 +334,12 @@ end super_fine_column_stretched_grid = RectilinearGrid(arch, size=4, z = [0, 0.1, 0.3, 0.65, Lz], topology=topology) super_fine_column_regular_grid = RectilinearGrid(arch, size=5, z=(0, Lz), topology=topology) - coarse_column_regular_c = CenterField(arch, coarse_column_regular_grid) - fine_column_regular_c = CenterField(arch, fine_column_regular_grid) - fine_column_stretched_c = CenterField(arch, fine_column_stretched_grid) - very_fine_column_stretched_c = CenterField(arch, very_fine_column_stretched_grid) - super_fine_column_stretched_c = CenterField(arch, super_fine_column_stretched_grid) - super_fine_column_regular_c = CenterField(arch, super_fine_column_regular_grid) + coarse_column_regular_c = CenterField(coarse_column_regular_grid) + fine_column_regular_c = CenterField(fine_column_regular_grid) + fine_column_stretched_c = CenterField(fine_column_stretched_grid) + very_fine_column_stretched_c = CenterField(very_fine_column_stretched_grid) + super_fine_column_stretched_c = CenterField(super_fine_column_stretched_grid) + super_fine_column_regular_c = CenterField(super_fine_column_regular_grid) # we initialize an array on the `fine_column_stretched_grid`, regrid it to the rest # grids, and check whether we get the anticipated results @@ -385,5 +393,4 @@ end end end end - =# end From a28c71f5e1dc1572944933acff8b18aa878b0e27 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 15 Dec 2021 10:55:40 -0700 Subject: [PATCH 003/140] Fix up nonhydrostatic model tests --- src/Fields/field.jl | 10 ++---- .../hydrostatic_free_surface_field_tuples.jl | 18 +++++----- .../hydrostatic_free_surface_model.jl | 14 ++++---- .../prescribed_hydrostatic_velocity_fields.jl | 10 +++--- .../nonhydrostatic_model.jl | 8 ++--- .../ShallowWaterModels/shallow_water_model.jl | 34 +++++++++---------- src/TimeSteppers/quasi_adams_bashforth_2.jl | 8 ++--- src/TimeSteppers/runge_kutta_3.jl | 8 ++--- test/test_nonhydrostatic_models.jl | 16 +++++---- 9 files changed, 62 insertions(+), 64 deletions(-) diff --git a/src/Fields/field.jl b/src/Fields/field.jl index d227a92a88..3dd9bb9379 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -287,12 +287,8 @@ const XYReducedField = Field{Nothing, Nothing, <:Any} const XYZReducedField = Field{Nothing, Nothing, Nothing} -const ReducedField = Union{XReducedField, - YReducedField, - ZReducedField, - YZReducedField, - XZReducedField, - XYReducedField, +const ReducedField = Union{XReducedField, YReducedField, ZReducedField, + YZReducedField, XZReducedField, XYReducedField, XYZReducedField} reduced_dimensions(field::Field) = () @@ -322,7 +318,7 @@ function fill_halo_regions!(field::Field, arch, args...; kwargs...) architecture(field), field.grid, args...; - reduced_dimensions, + reduced_dimensions=reduced_dims, kwargs...) end end diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_field_tuples.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_field_tuples.jl index cf51f6c225..cf542ac636 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_field_tuples.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_field_tuples.jl @@ -1,19 +1,19 @@ using Oceananigans.Grids: Center, Face using Oceananigans.Fields: XFaceField, YFaceField, ZFaceField, TracerFields -function HydrostaticFreeSurfaceVelocityFields(::Nothing, arch, grid, clock, bcs=NamedTuple()) - u = XFaceField(arch, grid, bcs.u) - v = YFaceField(arch, grid, bcs.v) - w = ZFaceField(arch, grid) +function HydrostaticFreeSurfaceVelocityFields(::Nothing, grid, clock, bcs=NamedTuple()) + u = XFaceField(grid, boundary_conditions=bcs.u) + v = YFaceField(grid, boundary_conditions=bcs.v) + w = ZFaceField(grid) return (u=u, v=v, w=w) end -function HydrostaticFreeSurfaceTendencyFields(velocities, free_surface, arch, grid, tracer_names) - u = XFaceField(arch, grid) - v = YFaceField(arch, grid) - η = FreeSurfaceDisplacementField(velocities, free_surface, arch, grid) - tracers = TracerFields(tracer_names, arch, grid) +function HydrostaticFreeSurfaceTendencyFields(velocities, free_surface, grid, tracer_names) + u = XFaceField(grid) + v = YFaceField(grid) + η = FreeSurfaceDisplacementField(velocities, free_surface, grid) + tracers = TracerFields(tracer_names, grid) return merge((u=u, v=v, η=η), tracers) end diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl index 93154cbbd7..eaca695286 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl @@ -145,10 +145,10 @@ function HydrostaticFreeSurfaceModel(; grid, closure = with_tracers(tracernames(tracers), closure) # Either check grid-correctness, or construct tuples of fields - velocities = HydrostaticFreeSurfaceVelocityFields(velocities, arch, grid, clock, boundary_conditions) - tracers = TracerFields(tracers, arch, grid, boundary_conditions) - pressure = PressureField(arch, grid) - diffusivity_fields = DiffusivityFields(diffusivity_fields, arch, grid, tracernames(tracers), boundary_conditions, closure) + velocities = HydrostaticFreeSurfaceVelocityFields(velocities, grid, clock, boundary_conditions) + tracers = TracerFields(tracers, grid, boundary_conditions) + pressure = PressureField(grid) + diffusivity_fields = DiffusivityFields(diffusivity_fields, grid, tracernames(tracers), boundary_conditions, closure) validate_velocity_boundary_conditions(velocities) @@ -156,10 +156,10 @@ function HydrostaticFreeSurfaceModel(; grid, # Instantiate timestepper if not already instantiated implicit_solver = implicit_diffusion_solver(time_discretization(closure), arch, grid) - timestepper = TimeStepper(:QuasiAdamsBashforth2, arch, grid, tracernames(tracers); + timestepper = TimeStepper(:QuasiAdamsBashforth2, grid, tracernames(tracers); implicit_solver = implicit_solver, - Gⁿ = HydrostaticFreeSurfaceTendencyFields(velocities, free_surface, arch, grid, tracernames(tracers)), - G⁻ = HydrostaticFreeSurfaceTendencyFields(velocities, free_surface, arch, grid, tracernames(tracers))) + Gⁿ = HydrostaticFreeSurfaceTendencyFields(velocities, free_surface, grid, tracernames(tracers)), + G⁻ = HydrostaticFreeSurfaceTendencyFields(velocities, free_surface, grid, tracernames(tracers))) # Regularize forcing for model tracer and velocity fields. model_fields = hydrostatic_prognostic_fields(velocities, free_surface, tracers) diff --git a/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl b/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl index 1fcfada069..57912acc28 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl @@ -52,7 +52,7 @@ function PrescribedField(X, Y, Z, f::Field, grid; kwargs...) return f end -function HydrostaticFreeSurfaceVelocityFields(velocities::PrescribedVelocityFields, arch, grid, clock, bcs) +function HydrostaticFreeSurfaceVelocityFields(velocities::PrescribedVelocityFields, grid, clock, bcs) u = PrescribedField(Face, Center, Center, velocities.u, grid; clock=clock, parameters=velocities.parameters) v = PrescribedField(Center, Face, Center, velocities.v, grid; clock=clock, parameters=velocities.parameters) @@ -61,7 +61,7 @@ function HydrostaticFreeSurfaceVelocityFields(velocities::PrescribedVelocityFiel return PrescribedVelocityFields(u, v, w, velocities.parameters) end -function HydrostaticFreeSurfaceTendencyFields(::PrescribedVelocityFields, free_surface, arch, grid, tracer_names) +function HydrostaticFreeSurfaceTendencyFields(::PrescribedVelocityFields, free_surface, grid, tracer_names) tracers = TracerFields(tracer_names, arch, grid) return merge((u = nothing, v = nothing, η = nothing), tracers) end @@ -76,9 +76,9 @@ compute_w_from_continuity!(::PrescribedVelocityFields, args...) = nothing validate_velocity_boundary_conditions(::PrescribedVelocityFields) = nothing extract_boundary_conditions(::PrescribedVelocityFields) = NamedTuple() -FreeSurfaceDisplacementField(::PrescribedVelocityFields, ::Nothing, arch, grid) = nothing -HorizontalVelocityFields(::PrescribedVelocityFields, arch, grid) = nothing, nothing -FreeSurface(free_surface::ExplicitFreeSurface{Nothing}, ::PrescribedVelocityFields, arch, grid) = nothing +FreeSurfaceDisplacementField(::PrescribedVelocityFields, ::Nothing, grid) = nothing +HorizontalVelocityFields(::PrescribedVelocityFields, grid) = nothing, nothing +FreeSurface(free_surface::ExplicitFreeSurface{Nothing}, ::PrescribedVelocityFields, grid) = nothing hydrostatic_prognostic_fields(::PrescribedVelocityFields, ::Nothing, tracers) = tracers calculate_hydrostatic_momentum_tendencies!(model, ::PrescribedVelocityFields; kwargs...) = [] diff --git a/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl b/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl index 0b1e160486..4fec96acca 100644 --- a/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl +++ b/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl @@ -151,10 +151,10 @@ function NonhydrostaticModel(; grid, closure = with_tracers(tracernames(tracers), closure) # Either check grid-correctness, or construct tuples of fields - velocities = VelocityFields(velocities, arch, grid, boundary_conditions) - tracers = TracerFields(tracers, arch, grid, boundary_conditions) - pressures = PressureFields(pressures, arch, grid, boundary_conditions) - diffusivity_fields = DiffusivityFields(diffusivity_fields, arch, grid, tracernames(tracers), boundary_conditions, closure) + velocities = VelocityFields(velocities, grid, boundary_conditions) + tracers = TracerFields(tracers, grid, boundary_conditions) + pressures = PressureFields(pressures, grid, boundary_conditions) + diffusivity_fields = DiffusivityFields(diffusivity_fields, grid, tracernames(tracers), boundary_conditions, closure) if isnothing(pressure_solver) pressure_solver = PressureSolver(arch, grid) diff --git a/src/Models/ShallowWaterModels/shallow_water_model.jl b/src/Models/ShallowWaterModels/shallow_water_model.jl index d6ee298961..0e8d843475 100644 --- a/src/Models/ShallowWaterModels/shallow_water_model.jl +++ b/src/Models/ShallowWaterModels/shallow_water_model.jl @@ -13,21 +13,22 @@ using Oceananigans.Utils: tupleit function ShallowWaterTendencyFields(arch, grid, tracer_names) - uh = XFaceField(arch, grid) - vh = YFaceField(arch, grid) - h = CenterField(arch, grid) + uh = XFaceField(grid) + vh = YFaceField(grid) + h = CenterField(grid) - tracers = TracerFields(tracer_names, arch, grid) + tracers = TracerFields(tracer_names, grid) + solution = (; uh, vh, h) - return merge((uh=uh, vh=vh, h=h), tracers) + return merge(solution, tracers) end -function ShallowWaterSolutionFields(arch, grid, bcs) - uh = XFaceField(arch, grid, bcs.uh) - vh = YFaceField(arch, grid, bcs.vh) - h = CenterField(arch, grid, bcs.h) +function ShallowWaterSolutionFields(grid, bcs) + uh = XFaceField(grid, bcs.uh) + vh = YFaceField(grid, bcs.vh) + h = CenterField(grid, bcs.h) - return (uh=uh, vh=vh, h=h) + return (; uh, vh, h) end mutable struct ShallowWaterModel{G, A<:AbstractArchitecture, T, V, R, F, E, B, Q, C, K, TS} <: AbstractModel{TS} @@ -114,15 +115,14 @@ function ShallowWaterModel(; boundary_conditions = regularize_field_boundary_conditions(boundary_conditions, grid, prognostic_field_names) - solution = ShallowWaterSolutionFields(arch, grid, boundary_conditions) - tracers = TracerFields(tracers, arch, grid, boundary_conditions) - diffusivity_fields = DiffusivityFields(diffusivity_fields, arch, grid, - tracernames(tracers), boundary_conditions, closure) + solution = ShallowWaterSolutionFields(grid, boundary_conditions) + tracers = TracerFields(tracers, grid, boundary_conditions) + diffusivity_fields = DiffusivityFields(diffusivity_fields, grid, tracernames(tracers), boundary_conditions, closure) # Instantiate timestepper if not already instantiated - timestepper = TimeStepper(timestepper, arch, grid, tracernames(tracers); - Gⁿ = ShallowWaterTendencyFields(arch, grid, tracernames(tracers)), - G⁻ = ShallowWaterTendencyFields(arch, grid, tracernames(tracers))) + timestepper = TimeStepper(timestepper, grid, tracernames(tracers); + Gⁿ = ShallowWaterTendencyFields(grid, tracernames(tracers)), + G⁻ = ShallowWaterTendencyFields(grid, tracernames(tracers))) # Regularize forcing and closure for model tracer and velocity fields. model_fields = merge(solution, tracers) diff --git a/src/TimeSteppers/quasi_adams_bashforth_2.jl b/src/TimeSteppers/quasi_adams_bashforth_2.jl index 835e9ff797..d39154a0d1 100644 --- a/src/TimeSteppers/quasi_adams_bashforth_2.jl +++ b/src/TimeSteppers/quasi_adams_bashforth_2.jl @@ -13,8 +13,8 @@ end """ QuasiAdamsBashforth2TimeStepper(arch, grid, tracers, χ=0.1; implicit_solver = nothing, - Gⁿ = TendencyFields(arch, grid, tracers), - G⁻ = TendencyFields(arch, grid, tracers)) + Gⁿ = TendencyFields(grid, tracers), + G⁻ = TendencyFields(grid, tracers)) Return an QuasiAdamsBashforth2TimeStepper object with tendency fields on `arch` and `grid` with AB2 parameter `χ`. The tendency fields can be specified via optional @@ -23,8 +23,8 @@ kwargs. function QuasiAdamsBashforth2TimeStepper(arch, grid, tracers, χ = 0.1; implicit_solver::IT = nothing, - Gⁿ = TendencyFields(arch, grid, tracers), - G⁻ = TendencyFields(arch, grid, tracers)) where IT + Gⁿ = TendencyFields(grid, tracers), + G⁻ = TendencyFields(grid, tracers)) where IT FT = eltype(grid) GT = typeof(Gⁿ) diff --git a/src/TimeSteppers/runge_kutta_3.jl b/src/TimeSteppers/runge_kutta_3.jl index 46500b0d28..4b48663aa0 100644 --- a/src/TimeSteppers/runge_kutta_3.jl +++ b/src/TimeSteppers/runge_kutta_3.jl @@ -19,16 +19,16 @@ end """ RungeKutta3TimeStepper(arch, grid, tracers, - Gⁿ = TendencyFields(arch, grid, tracers), - G⁻ = TendencyFields(arch, grid, tracers)) + Gⁿ = TendencyFields(grid, tracers), + G⁻ = TendencyFields(grid, tracers)) Return an `RungeKutta3TimeStepper` object with tendency fields on `arch` and `grid`. The tendency fields can be specified via optional kwargs. """ function RungeKutta3TimeStepper(arch, grid, tracers; implicit_solver::TI = nothing, - Gⁿ::TG = TendencyFields(arch, grid, tracers), - G⁻ = TendencyFields(arch, grid, tracers)) where {TI, TG} + Gⁿ::TG = TendencyFields(grid, tracers), + G⁻ = TendencyFields(grid, tracers)) where {TI, TG} !isnothing(implicit_solver) && @warn("Implicit-explicit time-stepping with RungeKutta3TimeStepper is not tested. " * diff --git a/test/test_nonhydrostatic_models.jl b/test/test_nonhydrostatic_models.jl index 1bbd037a79..f3eaf50915 100644 --- a/test/test_nonhydrostatic_models.jl +++ b/test/test_nonhydrostatic_models.jl @@ -1,3 +1,5 @@ +include("dependencies_for_runtests.jl") + @testset "Models" begin @info "Testing models..." @@ -121,11 +123,11 @@ Nx, Ny, Nz = size(model.grid) - u_cpu = XFaceField(CPU(), grid) - v_cpu = YFaceField(CPU(), grid) - w_cpu = ZFaceField(CPU(), grid) - T_cpu = CenterField(CPU(), grid) - S_cpu = CenterField(CPU(), grid) + u_cpu = XFaceField(grid) + v_cpu = YFaceField(grid) + w_cpu = ZFaceField(grid) + T_cpu = CenterField(grid) + S_cpu = CenterField(grid) set!(u_cpu, u) set!(v_cpu, v) @@ -156,12 +158,12 @@ @test all(abs.(interior(w_cpu)) .< ϵ) # Test setting the background_fields to a Field - U_field = XFaceField(arch, grid) + U_field = XFaceField(grid) U_field .= 1 model = NonhydrostaticModel(grid = grid, background_fields = (u=U_field,)) @test model.background_fields.velocities.u isa Field - U_field = CenterField(arch, grid) + U_field = CenterField(grid) @test_throws ArgumentError NonhydrostaticModel(grid=grid, background_fields = (u=U_field,)) end end From 39a8f65f3f214bd2a9e4e1a87165a3313dc073fa Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 15 Dec 2021 11:06:42 -0700 Subject: [PATCH 004/140] Moves reduced field tests into test_field --- test/runtests.jl | 1 - test/test_field.jl | 61 ++++++++---- test/test_kernel_computed_field.jl | 66 ------------- test/test_reduced_field.jl | 85 ----------------- test/test_reduced_fields.jl | 148 +++++++---------------------- 5 files changed, 76 insertions(+), 285 deletions(-) delete mode 100644 test/test_kernel_computed_field.jl delete mode 100644 test/test_reduced_field.jl diff --git a/test/runtests.jl b/test/runtests.jl index 3a144683b6..42e5be952c 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -8,7 +8,6 @@ include("dependencies_for_runtests.jl") include("test_boundary_conditions.jl") include("test_field.jl") include("test_reduced_fields.jl") - include("test_kernel_computed_field.jl") include("test_halo_regions.jl") include("test_coriolis.jl") include("test_buoyancy.jl") diff --git a/test/test_field.jl b/test/test_field.jl index 24d11b2aed..71f64cf87c 100644 --- a/test/test_field.jl +++ b/test/test_field.jl @@ -11,7 +11,7 @@ using Oceananigans.Fields: reduced_location Test that the field initialized by the FieldType constructor on `grid` has size `(Tx, Ty, Tz)`. """ -correct_field_size(g, FieldType, Tx, Ty, Tz) = size(parent(FieldType(g))) == (Tx, Ty, Tz) +correct_field_size(grid, loc, Tx, Ty, Tz) = size(parent(Field(loc, grid))) == (Tx, Ty, Tz) function run_similar_field_tests(f) g = similar(f) @@ -180,28 +180,40 @@ end for arch in archs, FT in float_types grid = RectilinearGrid(arch , FT, size=N, extent=L, halo=H, topology=(Periodic, Periodic, Periodic)) - @test correct_field_size(grid, CenterField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(grid, XFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(grid, YFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(grid, ZFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, (Center, Center, Center), N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, (Face, Center, Center), N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, (Center, Face, Center), N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, (Center, Center, Face), N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) grid = RectilinearGrid(arch, FT, size=N, extent=L, halo=H, topology=(Periodic, Periodic, Bounded)) - @test correct_field_size(grid, CenterField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(grid, XFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(grid, YFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(grid, ZFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3] + 1) + @test correct_field_size(grid, (Center, Center, Center), N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, (Face, Center, Center), N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, (Center, Face, Center), N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, (Center, Center, Face), N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3] + 1) grid = RectilinearGrid(arch, FT, size=N, extent=L, halo=H, topology=(Periodic, Bounded, Bounded)) - @test correct_field_size(grid, CenterField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(grid, XFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(grid, YFaceField, N[1] + 2 * H[1], N[2] + 1 + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(grid, ZFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 1 + 2 * H[3]) + @test correct_field_size(grid, (Center, Center, Center), N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, (Face, Center, Center), N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, (Center, Face, Center), N[1] + 2 * H[1], N[2] + 1 + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, (Center, Center, Face), N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 1 + 2 * H[3]) grid = RectilinearGrid(arch, FT, size=N, extent=L, halo=H, topology=(Bounded, Bounded, Bounded)) - @test correct_field_size(grid, CenterField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(grid, XFaceField, N[1] + 1 + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(grid, YFaceField, N[1] + 2 * H[1], N[2] + 1 + 2 * H[2], N[3] + 2 * H[3]) - @test correct_field_size(grid, ZFaceField, N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 1 + 2 * H[3]) + @test correct_field_size(grid, (Center, Center, Center), N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, (Face, Center, Center), N[1] + 1 + 2 * H[1], N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, (Center, Face, Center), N[1] + 2 * H[1], N[2] + 1 + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, (Center, Center, Face), N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 1 + 2 * H[3]) + + @test correct_reduced_field_size((Nothing, Center, Center), grid, 1, 1, N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_reduced_field_size((Nothing, Center, Center), grid, 1, 1, N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_reduced_field_size((Nothing, Face, Center), grid, 1, 1, N[2] + 2 * H[2] + 1, N[3] + 2 * H[3]) + @test correct_reduced_field_size((Nothing, Face, Face), grid, 1, 1, N[2] + 2 * H[2] + 1, N[3] + 2 * H[3] + 1) + @test correct_reduced_field_size((Center, Nothing, Center), grid, 2, N[1] + 2 * H[1], 1, N[3] + 2 * H[3]) + @test correct_reduced_field_size((Center, Nothing, Center), grid, 2, N[1] + 2 * H[1], 1, N[3] + 2 * H[3]) + @test correct_reduced_field_size((Center, Center, Nothing), grid, 3, N[1] + 2 * H[1], N[2] + 2 * H[2], 1) + @test correct_reduced_field_size((Nothing, Nothing, Center), grid, (1, 2), 1, 1, N[3] + 2 * H[3]) + @test correct_reduced_field_size((Center, Nothing, Nothing), grid, (2, 3), N[1] + 2 * H[1], 1, 1) + @test correct_reduced_field_size((Nothing, Nothing, Nothing), grid, (1, 2, 3), 1, 1, 1) + end end @@ -231,12 +243,21 @@ end @test correct_field_value_was_set(grid, FieldType, val) end - for FieldType in FieldTypes - field = FieldType(grid) + for loc in ((Center, Center, Center), + (Face, Center, Center), + (Center, Face, Center), + (Center, Center, Face), + (Nothing, Center, Center), + (Center, Nothing, Center), + (Center, Center, Nothing), + (Nothing, Nothing, Center), + (Nothing, Nothing, Nothing)) + + field = Field(loc, grid) sz = size(field) A = rand(FT, sz...) set!(field, A) - @test field.data[2, 4, 6] == A[2, 4, 6] + @test field.data[1, 1, 1] == A[1, 1, 1] end Nx = 8 diff --git a/test/test_kernel_computed_field.jl b/test/test_kernel_computed_field.jl deleted file mode 100644 index eee30db3fc..0000000000 --- a/test/test_kernel_computed_field.jl +++ /dev/null @@ -1,66 +0,0 @@ -using Oceananigans.Fields: KernelComputedField -using KernelAbstractions: @kernel, @index - -grid = RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1), - topology=(Periodic, Periodic, Bounded)) - -@kernel function double!(doubled_field, grid, single_field) - i, j, k = @index(Global, NTuple) - @inbounds doubled_field[i, j, k] = 2 * single_field[i, j, k] -end - -@kernel function multiply!(multiplied_field, grid, field, multiple) - i, j, k = @index(Global, NTuple) - @inbounds multiplied_field[i, j, k] = multiple * field[i, j, k] -end - -for arch in archs - @testset "KernelComputedField [$(typeof(arch))]" begin - @info " Testing KernelComputedField..." - - single_field = Field(Center, Center, Center, arch, grid) - - doubled_field = KernelComputedField(Center, Center, Center, - double!, arch, grid, - computed_dependencies = single_field) - - # Test that the constructor worked - @test doubled_field isa KernelComputedField - - multiple = 3 - multiplied_field = KernelComputedField(Center, Center, Center, - multiply!, arch, grid, - computed_dependencies = doubled_field, - parameters = multiple) - - # Test that the constructor worked - @test multiplied_field isa KernelComputedField - - set!(single_field, π) - @test single_field[1, 1, 1] == convert(eltype(single_field), π) - - compute!(doubled_field) - @test doubled_field[1, 1, 1] == 2π - - # Test boundary conditions / halo filling - @test doubled_field[0, 1, 1] == doubled_field[2, 1, 1] # periodic - @test doubled_field[1, 0, 1] == doubled_field[1, 2, 1] # periodic - @test doubled_field[1, 1, 0] == doubled_field[1, 1, 1] # no flux - @test doubled_field[1, 1, 1] == doubled_field[1, 1, 2] # no flux - - set!(doubled_field, 0) - compute!(multiplied_field) - - @test doubled_field[1, 1, 1] == 2π - @test multiplied_field[1, 1, 1] == multiple * 2 * π - - doubled_face_field = KernelComputedField(Center, Center, Face, - double!, arch, grid, - computed_dependencies = single_field) - - # Test that nothing happens for fields on faces in bounded directions - compute!(doubled_face_field) - @test doubled_face_field[1, 1, 0] == 0 - @test doubled_face_field[1, 1, 3] == 0 - end -end diff --git a/test/test_reduced_field.jl b/test/test_reduced_field.jl deleted file mode 100644 index 46210f35b0..0000000000 --- a/test/test_reduced_field.jl +++ /dev/null @@ -1,85 +0,0 @@ -""" - correct_reduced_field_size(loc, arch, grid, dims, Tx, Ty, Tz) - -Test that the ReducedField at `loc`ation on `arch`itecture and `grid` -and reduced along `dims` has size `(Tx, Ty, Tz)`. -""" -correct_reduced_field_size(loc, arch, grid, dims, Tx, Ty, Tz) = - size(parent(ReducedField(loc, arch, grid; dims=dims))) == (Tx, Ty, Tz) - -function correct_reduced_field_value_was_set(arch, grid, loc, dims, val::Number) - f = ReducedField(loc, arch, grid; dims=dims) - set!(f, val) - return all(interior(f) .≈ val * arch_array(arch, ones(size(f)))) -end - -@testset "ReducedFields" begin - @info "Testing ReducedFields..." - - N = (4, 6, 8) - L = (2π, 3π, 5π) - H = (1, 1, 1) - - @testset "ReducedField initialization" begin - @info " Testing ReducedField initialization..." - for arch in archs, FT in float_types - - grid = RectilinearGrid(arch, FT, size=N, extent=L, halo=H, topology=(Bounded, Bounded, Bounded)) - - @test correct_reduced_field_size((Center, Center, Center), arch, grid, 1, 1, N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_reduced_field_size((Face, Center, Center), arch, grid, 1, 1, N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_reduced_field_size((Center, Face, Center), arch, grid, 1, 1, N[2] + 2 * H[2] + 1, N[3] + 2 * H[3]) - @test correct_reduced_field_size((Center, Face, Face), arch, grid, 1, 1, N[2] + 2 * H[2] + 1, N[3] + 2 * H[3] + 1) - @test correct_reduced_field_size((Center, Center, Center), arch, grid, 2, N[1] + 2 * H[1], 1, N[3] + 2 * H[3]) - @test correct_reduced_field_size((Center, Center, Center), arch, grid, 2, N[1] + 2 * H[1], 1, N[3] + 2 * H[3]) - @test correct_reduced_field_size((Center, Center, Center), arch, grid, 3, N[1] + 2 * H[1], N[2] + 2 * H[2], 1) - @test correct_reduced_field_size((Center, Center, Center), arch, grid, (1, 2), 1, 1, N[3] + 2 * H[3]) - @test correct_reduced_field_size((Center, Center, Center), arch, grid, (2, 3), N[1] + 2 * H[1], 1, 1) - @test correct_reduced_field_size((Center, Center, Center), arch, grid, (1, 2, 3), 1, 1, 1) - end - end - - reduced_dims = (1, 2, 3, (1, 2), (2, 3), (1, 3), (1, 2, 3)) - - int_vals = Any[0, Int8(-1), Int16(2), Int32(-3), Int64(4)] - uint_vals = Any[6, UInt8(7), UInt16(8), UInt32(9), UInt64(10)] - float_vals = Any[0.0, -0.0, 6e-34, 1.0f10] - rational_vals = Any[1//11, -23//7] - other_vals = Any[π] - vals = vcat(int_vals, uint_vals, float_vals, rational_vals, other_vals) - - @testset "Setting ReducedFields" begin - @info " Testing ReducedField setting..." - for arch in archs, FT in float_types - - grid = RectilinearGrid(arch, FT, size=N, extent=L, halo=H, topology=(Periodic, Periodic, Bounded)) - - for dims in reduced_dims, val in vals - @test correct_reduced_field_value_was_set(arch, grid, (Center, Center, Center), dims, val) - end - - for dims in reduced_dims - reduced_field = ReducedField((Center, Center, Center), arch, grid, dims=dims) - sz = size(reduced_field) - A = rand(FT, sz...) - set!(reduced_field, A) - - @test reduced_field[1, 1, 1] == A[1, 1, 1] - - fill_halo_regions!(reduced_field, arch) - - # No-flux boundary conditions at top and bottom - @test reduced_field[1, 1, 0] == A[1, 1, 1] - @test reduced_field[1, 1, grid.Nz+1] == A[1, 1, end] - - # Periodic boundary conditions in the horizontal directions - @test reduced_field[1, 0, 1] == A[1, end, 1] - @test reduced_field[1, grid.Ny+1, 1] == A[1, 1, 1] - - @test reduced_field[0, 1, 1] == A[end, 1, 1] - @test reduced_field[grid.Nx+1, 1, 1] == A[1, 1, 1] - end - end - end -end - diff --git a/test/test_reduced_fields.jl b/test/test_reduced_fields.jl index b5defb04e0..6f4837200d 100644 --- a/test/test_reduced_fields.jl +++ b/test/test_reduced_fields.jl @@ -8,125 +8,46 @@ include("utils_for_runtests.jl") archs = test_architectures() -""" - correct_reduced_field_size(loc, arch, grid, dims, Tx, Ty, Tz) - -Test that the ReducedField at `loc`ation on `arch`itecture and `grid` -and reduced along `dims` has size `(Tx, Ty, Tz)`. -""" -correct_reduced_field_size(loc, arch, grid, dims, Tx, Ty, Tz) = - size(parent(ReducedField(loc, arch, grid; dims=dims))) == (Tx, Ty, Tz) - -function correct_reduced_field_value_was_set(arch, grid, loc, dims, val::Number) - f = ReducedField(loc, arch, grid; dims=dims) - set!(f, val) - return all(interior(f) .≈ val * arch_array(arch, ones(size(f)))) -end - -@testset "AbstractReducedFields" begin - @info "Testing AbstractReducedFields..." - - N = (4, 6, 8) - L = (2π, 3π, 5π) - H = (1, 1, 1) - - @testset "ReducedField initialization" begin - @info " Testing ReducedField initialization..." - for arch in archs, FT in float_types - - grid = RectilinearGrid(arch, FT, size=N, extent=L, halo=H, topology=(Bounded, Bounded, Bounded)) - - @test correct_reduced_field_size((Center, Center, Center), arch, grid, 1, 1, N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_reduced_field_size((Face, Center, Center), arch, grid, 1, 1, N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_reduced_field_size((Center, Face, Center), arch, grid, 1, 1, N[2] + 2 * H[2] + 1, N[3] + 2 * H[3]) - @test correct_reduced_field_size((Center, Face, Face), arch, grid, 1, 1, N[2] + 2 * H[2] + 1, N[3] + 2 * H[3] + 1) - @test correct_reduced_field_size((Center, Center, Center), arch, grid, 2, N[1] + 2 * H[1], 1, N[3] + 2 * H[3]) - @test correct_reduced_field_size((Center, Center, Center), arch, grid, 2, N[1] + 2 * H[1], 1, N[3] + 2 * H[3]) - @test correct_reduced_field_size((Center, Center, Center), arch, grid, 3, N[1] + 2 * H[1], N[2] + 2 * H[2], 1) - @test correct_reduced_field_size((Center, Center, Center), arch, grid, (1, 2), 1, 1, N[3] + 2 * H[3]) - @test correct_reduced_field_size((Center, Center, Center), arch, grid, (2, 3), N[1] + 2 * H[1], 1, 1) - @test correct_reduced_field_size((Center, Center, Center), arch, grid, (1, 2, 3), 1, 1, 1) - end - end - - reduced_dims = (1, 2, 3, (1, 2), (2, 3), (1, 3), (1, 2, 3)) - - int_vals = Any[0, Int8(-1), Int16(2), Int32(-3), Int64(4)] - uint_vals = Any[6, UInt8(7), UInt16(8), UInt32(9), UInt64(10)] - float_vals = Any[0.0, -0.0, 6e-34, 1.0f10] - rational_vals = Any[1//11, -23//7] - other_vals = Any[π] - vals = vcat(int_vals, uint_vals, float_vals, rational_vals, other_vals) - - @testset "Setting ReducedFields" begin - @info " Testing ReducedField setting..." - for arch in archs, FT in float_types - - grid = RectilinearGrid(arch, FT, size=N, extent=L, halo=H, topology=(Periodic, Periodic, Bounded)) - - for dims in reduced_dims, val in vals - @test correct_reduced_field_value_was_set(arch, grid, (Center, Center, Center), dims, val) - end - - for dims in reduced_dims - reduced_field = ReducedField((Center, Center, Center), arch, grid, dims=dims) - sz = size(reduced_field) - A = rand(FT, sz...) - set!(reduced_field, A) - - @test reduced_field[1, 1, 1] == A[1, 1, 1] - - fill_halo_regions!(reduced_field, arch) - - # No-flux boundary conditions at top and bottom - @test reduced_field[1, 1, 0] == A[1, 1, 1] - @test reduced_field[1, 1, grid.Nz+1] == A[1, 1, end] - - # Periodic boundary conditions in the horizontal directions - @test reduced_field[1, 0, 1] == A[1, end, 1] - @test reduced_field[1, grid.Ny+1, 1] == A[1, 1, 1] - - @test reduced_field[0, 1, 1] == A[end, 1, 1] - @test reduced_field[grid.Nx+1, 1, 1] == A[1, 1, 1] - end - end - end +@testset "AveragedField" begin + @info "Testing AveragedField..." for arch in archs + arch_str = string(typeof(arch)) - grid = RectilinearGrid(arch, topology = (Periodic, Periodic, Bounded), - size = (2, 2, 2), - x = (0, 2), - y = (0, 2), - z = (0, 2)) + @testset "Averaged fields [$arch_str]" begin + @info " Testing AveragedFields [$arch_str]" - w = ZFaceField(arch, grid) - T = CenterField(arch, grid) + grid = RectilinearGrid(arch, topology = (Periodic, Periodic, Bounded), + size = (2, 2, 2), + x = (0, 2), + y = (0, 2), + z = (0, 2)) - trilinear(x, y, z) = x + y + z + w = ZFaceField(grid) + T = CenterField(grid) - set!(T, trilinear) - set!(w, trilinear) + trilinear(x, y, z) = x + y + z - @compute Txyz = AveragedField(T, dims=(1, 2, 3)) + set!(T, trilinear) + set!(w, trilinear) - # Note: halo regions must be *filled* prior to computing an average - # if the average within halo regions is to be correct. - fill_halo_regions!(T, arch) - @compute Txy = AveragedField(T, dims=(1, 2)) + @compute Txyz = AveragedField(T, dims=(1, 2, 3)) - fill_halo_regions!(T, arch) - @compute Tx = AveragedField(T, dims=1) + # Note: halo regions must be *filled* prior to computing an average + # if the average within halo regions is to be correct. + fill_halo_regions!(T, arch) + @compute Txy = AveragedField(T, dims=(1, 2)) - @compute wxyz = AveragedField(w, dims=(1, 2, 3)) - @compute wxy = AveragedField(w, dims=(1, 2)) - @compute wx = AveragedField(w, dims=1) + fill_halo_regions!(T, arch) + @compute Tx = AveragedField(T, dims=1) - Nx, Ny, Nz = grid.Nx, grid.Ny, grid.Nz + @compute wxyz = AveragedField(w, dims=(1, 2, 3)) + @compute wxy = AveragedField(w, dims=(1, 2)) + @compute wx = AveragedField(w, dims=1) + + Nx, Ny, Nz = grid.Nx, grid.Ny, grid.Nz - @testset "Averaged fields [$arch_str]" begin - @info " Testing AveragedFields [$arch_str]" for FT in float_types @test Txyz[1, 1, 1] ≈ 3 @@ -140,13 +61,14 @@ end @test Array(interior(wx))[1, :, :] ≈ [[1.5, 2.5] [2.5, 3.5] [3.5, 4.5]] # Test whether a race condition gets hit for averages over large fields - big_grid = RectilinearGrid(arch, topology = (Periodic, Periodic, Bounded), - size = (256, 256, 128), - x = (0, 2), - y = (0, 2), - z = (0, 2)) - - c = CenterField(arch, big_grid) + big_grid = RectilinearGrid(arch, + topology = (Periodic, Periodic, Bounded), + size = (256, 256, 128), + x = (0, 2), + y = (0, 2), + z = (0, 2)) + + c = CenterField(big_grid) c .= 1 C = AveragedField(c, dims=(1, 2)) From 0e18182170cca99096d2f876685aaecfcaf79ef0 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 15 Dec 2021 11:09:09 -0700 Subject: [PATCH 005/140] Bugfix in test_field --- test/test_field.jl | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/test/test_field.jl b/test/test_field.jl index 71f64cf87c..82fe271c7a 100644 --- a/test/test_field.jl +++ b/test/test_field.jl @@ -203,16 +203,17 @@ end @test correct_field_size(grid, (Center, Face, Center), N[1] + 2 * H[1], N[2] + 1 + 2 * H[2], N[3] + 2 * H[3]) @test correct_field_size(grid, (Center, Center, Face), N[1] + 2 * H[1], N[2] + 2 * H[2], N[3] + 1 + 2 * H[3]) - @test correct_reduced_field_size((Nothing, Center, Center), grid, 1, 1, N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_reduced_field_size((Nothing, Center, Center), grid, 1, 1, N[2] + 2 * H[2], N[3] + 2 * H[3]) - @test correct_reduced_field_size((Nothing, Face, Center), grid, 1, 1, N[2] + 2 * H[2] + 1, N[3] + 2 * H[3]) - @test correct_reduced_field_size((Nothing, Face, Face), grid, 1, 1, N[2] + 2 * H[2] + 1, N[3] + 2 * H[3] + 1) - @test correct_reduced_field_size((Center, Nothing, Center), grid, 2, N[1] + 2 * H[1], 1, N[3] + 2 * H[3]) - @test correct_reduced_field_size((Center, Nothing, Center), grid, 2, N[1] + 2 * H[1], 1, N[3] + 2 * H[3]) - @test correct_reduced_field_size((Center, Center, Nothing), grid, 3, N[1] + 2 * H[1], N[2] + 2 * H[2], 1) - @test correct_reduced_field_size((Nothing, Nothing, Center), grid, (1, 2), 1, 1, N[3] + 2 * H[3]) - @test correct_reduced_field_size((Center, Nothing, Nothing), grid, (2, 3), N[1] + 2 * H[1], 1, 1) - @test correct_reduced_field_size((Nothing, Nothing, Nothing), grid, (1, 2, 3), 1, 1, 1) + # Reduced fields + @test correct_field_size(grid, (Nothing, Center, Center), 1, N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, (Nothing, Center, Center), 1, N[2] + 2 * H[2], N[3] + 2 * H[3]) + @test correct_field_size(grid, (Nothing, Face, Center), 1, N[2] + 2 * H[2] + 1, N[3] + 2 * H[3]) + @test correct_field_size(grid, (Nothing, Face, Face), 1, N[2] + 2 * H[2] + 1, N[3] + 2 * H[3] + 1) + @test correct_field_size(grid, (Center, Nothing, Center), N[1] + 2 * H[1], 1, N[3] + 2 * H[3]) + @test correct_field_size(grid, (Center, Nothing, Center), N[1] + 2 * H[1], 1, N[3] + 2 * H[3]) + @test correct_field_size(grid, (Center, Center, Nothing), N[1] + 2 * H[1], N[2] + 2 * H[2], 1) + @test correct_field_size(grid, (Nothing, Nothing, Center), 1, 1, N[3] + 2 * H[3]) + @test correct_field_size(grid, (Center, Nothing, Nothing), N[1] + 2 * H[1], 1, 1) + @test correct_field_size(grid, (Nothing, Nothing, Nothing), 1, 1, 1) end end From e844db410b0a76defd48728444c3c3552a785ccd Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 15 Dec 2021 11:22:43 -0700 Subject: [PATCH 006/140] Changing the Field API one constructor at a time --- test/runtests.jl | 5 +- test/test_batched_tridiagonal_solver.jl | 3 +- test/test_halo_regions.jl | 4 +- ...static_free_surface_immersed_boundaries.jl | 170 +++++++++++++++--- ...rface_immersed_boundaries_apply_surf_bc.jl | 71 -------- ..._immersed_boundaries_vertical_integrals.jl | 61 ------- ...reconditioned_conjugate_gradient_solver.jl | 21 ++- 7 files changed, 167 insertions(+), 168 deletions(-) delete mode 100644 test/test_hydrostatic_free_surface_immersed_boundaries_apply_surf_bc.jl delete mode 100644 test/test_hydrostatic_free_surface_immersed_boundaries_vertical_integrals.jl diff --git a/test/runtests.jl b/test/runtests.jl index 42e5be952c..1cfcf56eaf 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -7,7 +7,7 @@ include("dependencies_for_runtests.jl") include("test_operators.jl") include("test_boundary_conditions.jl") include("test_field.jl") - include("test_reduced_fields.jl") + # include("test_reduced_fields.jl") include("test_halo_regions.jl") include("test_coriolis.jl") include("test_buoyancy.jl") @@ -50,10 +50,9 @@ include("dependencies_for_runtests.jl") include("test_hydrostatic_free_surface_models.jl") include("test_ensemble_hydrostatic_free_surface_models.jl") include("test_hydrostatic_free_surface_immersed_boundaries.jl") + include("test_hydrostatic_free_surface_immersed_boundaries_vertical_integrals.jl") include("test_vertical_vorticity_field.jl") include("test_implicit_free_surface_solver.jl") - include("test_hydrostatic_free_surface_immersed_boundaries_apply_surf_bc.jl") - include("test_hydrostatic_free_surface_immersed_boundaries_vertical_integrals.jl") end end diff --git a/test/test_batched_tridiagonal_solver.jl b/test/test_batched_tridiagonal_solver.jl index a62ebfcccf..e823f850e5 100644 --- a/test/test_batched_tridiagonal_solver.jl +++ b/test/test_batched_tridiagonal_solver.jl @@ -1,5 +1,6 @@ -using LinearAlgebra +include("dependencies_for_runtests.jl") +using LinearAlgebra using Oceananigans.Architectures: array_type function can_solve_single_tridiagonal_system(arch, N) diff --git a/test/test_halo_regions.jl b/test/test_halo_regions.jl index e5e086ef38..a0e581b0f8 100644 --- a/test/test_halo_regions.jl +++ b/test/test_halo_regions.jl @@ -3,7 +3,7 @@ function halo_regions_initalized_correctly(arch, FT, Nx, Ny, Nz) Lx, Ly, Lz = 10, 20, 30 grid = RectilinearGrid(arch, FT, size=(Nx, Ny, Nz), extent=(Lx, Ly, Lz)) - field = CenterField(arch, grid) + field = CenterField(grid) # Fill the interior with random numbers. set!(field, rand(FT, Nx, Ny, Nz)) @@ -25,7 +25,7 @@ function halo_regions_correctly_filled(arch, FT, Nx, Ny, Nz) grid = RectilinearGrid(arch, FT, size=(Nx, Ny, Nz), extent=(Lx, Ly, Lz), topology=(Periodic, Periodic, Bounded)) - field = CenterField(arch, grid) + field = CenterField(grid) set!(field, rand(FT, Nx, Ny, Nz)) fill_halo_regions!(field, arch) diff --git a/test/test_hydrostatic_free_surface_immersed_boundaries.jl b/test/test_hydrostatic_free_surface_immersed_boundaries.jl index 561c0a7e5d..f90748ed2b 100644 --- a/test/test_hydrostatic_free_surface_immersed_boundaries.jl +++ b/test/test_hydrostatic_free_surface_immersed_boundaries.jl @@ -1,41 +1,169 @@ using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBoundary using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization +@inline surface_wind_stress(λ, φ, t, p) = p.τ₀ * cos(2π * (φ - p.φ₀) / p.Lφ) +@inline u_bottom_drag(i, j, grid, clock, fields, μ) = @inbounds - μ * fields.u[i, j, 1] +@inline v_bottom_drag(i, j, grid, clock, fields, μ) = @inbounds - μ * fields.v[i, j, 1] + @testset "Immersed boundaries with hydrostatic free surface models" begin @info "Testing immersed boundaries with hydrostatic free surface models..." for arch in archs - underlying_grid = RectilinearGrid(arch, size=(8, 8, 8), x = (-5, 5), y = (-5, 5), z = (0, 2)) - bump(x, y, z) = z < exp(-x^2 - y^2) + arch_str = string(typeof(arch)) - grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBoundary(bump)) - - for closure in (IsotropicDiffusivity(ν=1, κ=0.5), - IsotropicDiffusivity(ν=1, κ=0.5, time_discretization=VerticallyImplicitTimeDiscretization())) + @testset "GridFittedBoundary [$arch_str]" begin + @info "Testing GridFittedBoundary with HydrostaticFreeSurfaceModel [$arch_str]..." + + underlying_grid = RectilinearGrid(arch, size=(8, 8, 8), x = (-5, 5), y = (-5, 5), z = (0, 2)) + + bump(x, y, z) = z < exp(-x^2 - y^2) + + grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBoundary(bump)) + + for closure in (IsotropicDiffusivity(ν=1, κ=0.5), + IsotropicDiffusivity(ν=1, κ=0.5, time_discretization=VerticallyImplicitTimeDiscretization())) + + model = HydrostaticFreeSurfaceModel(grid = grid, + tracers = :b, + buoyancy = BuoyancyTracer(), + closure = closure) + + u = model.velocities.u + b = model.tracers.b + + # Linear stratification + set!(model, u = 1, b = (x, y, z) -> 4 * z) + + # Inside the bump + @test b[4, 4, 2] == 0 + @test u[4, 4, 2] == 0 + + simulation = Simulation(model, Δt = 1e-3, stop_iteration=2) - model = HydrostaticFreeSurfaceModel(grid = grid, - tracers = :b, - buoyancy = BuoyancyTracer(), - closure = closure) + run!(simulation) - u = model.velocities.u - b = model.tracers.b + # Inside the bump + @test b[4, 4, 2] == 0 + @test u[4, 4, 2] == 0 + end - # Linear stratification - set!(model, u = 1, b = (x, y, z) -> 4 * z) + @testset "Surface boundary conditions with immersed boundaries [$arch_str]" begin + @info " Testing surface boundary conditions with ImmersedBoundaries in HydrostaticFreeSurfaceModel [$arch_str]..." - # Inside the bump - @test b[4, 4, 2] == 0 - @test u[4, 4, 2] == 0 + Nx = 60 + Ny = 60 + + # A spherical domain + underlying_grid = LatitudeLongitudeGrid(arch, size = (Nx, Ny, 1), + longitude = (-30, 30), + latitude = (15, 75), + z = (-4000, 0)) + + bathymetry = zeros(Nx, Ny) .- 4000 + view(bathymetry, 31:34, 43:47) .= 0 + bathymetry = arch_array(arch, bathymetry) + + grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bathymetry)) + + free_surface = ImplicitFreeSurface(gravitational_acceleration=0.1) + coriolis = HydrostaticSphericalCoriolis(scheme = VectorInvariantEnstrophyConserving()) + + surface_wind_stress_parameters = (τ₀ = 1e-4, + Lφ = grid.Ly, + φ₀ = 15) + + surface_wind_stress_bc = FluxBoundaryCondition(surface_wind_stress, + parameters = surface_wind_stress_parameters) + + μ = 1 / 60days + + u_bottom_drag_bc = FluxBoundaryCondition(u_bottom_drag, + discrete_form = true, + parameters = μ) + + v_bottom_drag_bc = FluxBoundaryCondition(v_bottom_drag, + discrete_form = true, + parameters = μ) - simulation = Simulation(model, Δt = 1e-3, stop_iteration=2) + u_bcs = FieldBoundaryConditions(top = surface_wind_stress_bc, bottom = u_bottom_drag_bc) + v_bcs = FieldBoundaryConditions(bottom = v_bottom_drag_bc) + + νh₀ = 5e3 * (60 / grid.Nx)^2 + constant_horizontal_diffusivity = HorizontallyCurvilinearAnisotropicDiffusivity(νh=νh₀) + + model = HydrostaticFreeSurfaceModel(grid = grid, + momentum_advection = VectorInvariant(), + free_surface = free_surface, + coriolis = coriolis, + boundary_conditions = (u=u_bcs, v=v_bcs), + closure = constant_horizontal_diffusivity, + tracers = nothing, + buoyancy = nothing) + + + simulation = Simulation(model, Δt=3600, stop_iteration=1) run!(simulation) - # Inside the bump - @test b[4, 4, 2] == 0 - @test u[4, 4, 2] == 0 + # If reached here it didn't error, so pass for now! + @test true + end + + @testset "Correct vertically-integrated lateral face areas with immersed boundaries [$arch_str]" begin + @info " Testing correct vertically-integrated lateral face areas with immersed boundaries [$arch_str]..." + + Nx = 5 + Ny = 5 + + underlying_grid = RectilinearGrid(arch, + size = (Nx, Ny, 3), + extent = (Nx, Ny, 3), + topology = (Periodic, Periodic, Bounded)) + + # B for bathymetry + B = [-3. for i=1:Nx, j=1:Ny ] + B[2:Nx-1,2:Ny-1] .= [-2. for i=2:Nx-1, j=2:Ny-1 ] + B[3:Nx-2,3:Ny-2] .= [-1. for i=3:Nx-2, j=3:Ny-2 ] + + grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(B)) + + model = HydrostaticFreeSurfaceModel(grid = grid, + free_surface = ImplicitFreeSurface(), + tracer_advection = WENO5(), + buoyancy = nothing, + tracers = nothing, + closure = nothing) + + x_ref = [0.0 0.0 0.0 0.0 0.0 0.0 0.0 + 0.0 3.0 3.0 3.0 3.0 3.0 0.0 + 0.0 3.0 2.0 2.0 2.0 2.0 0.0 + 0.0 3.0 2.0 1.0 1.0 2.0 0.0 + 0.0 3.0 2.0 2.0 2.0 2.0 0.0 + 0.0 3.0 3.0 3.0 3.0 3.0 0.0 + 0.0 0.0 0.0 0.0 0.0 0.0 0.0]' + + y_ref = [0.0 0.0 0.0 0.0 0.0 0.0 0.0 + 0.0 3.0 3.0 3.0 3.0 3.0 0.0 + 0.0 3.0 2.0 2.0 2.0 3.0 0.0 + 0.0 3.0 2.0 1.0 2.0 3.0 0.0 + 0.0 3.0 2.0 1.0 2.0 3.0 0.0 + 0.0 3.0 2.0 2.0 2.0 3.0 0.0 + 0.0 0.0 0.0 0.0 0.0 0.0 0.0]' + + fs = model.free_surface + vertically_integrated_lateral_areas = fs.implicit_step_solver.vertically_integrated_lateral_areas + + ∫Axᶠᶜᶜ = vertically_integrated_lateral_areas.xᶠᶜᶜ + ∫Ayᶜᶠᶜ = vertically_integrated_lateral_areas.yᶜᶠᶜ + + ∫Axᶠᶜᶜ = Array(parent(∫Axᶠᶜᶜ)) + ∫Ayᶜᶠᶜ = Array(parent(∫Ayᶜᶠᶜ)) + + Ax_ok = ∫Axᶠᶜᶜ[:, :, 1] ≈ x_ref + Ay_ok = ∫Ayᶜᶠᶜ[:, :, 1] ≈ y_ref + + @test (Ax_ok & Ay_ok) end end end diff --git a/test/test_hydrostatic_free_surface_immersed_boundaries_apply_surf_bc.jl b/test/test_hydrostatic_free_surface_immersed_boundaries_apply_surf_bc.jl deleted file mode 100644 index cc25fd4025..0000000000 --- a/test/test_hydrostatic_free_surface_immersed_boundaries_apply_surf_bc.jl +++ /dev/null @@ -1,71 +0,0 @@ -using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBottom -using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization - -@inline surface_wind_stress(λ, φ, t, p) = p.τ₀ * cos(2π * (φ - p.φ₀) / p.Lφ) -@inline u_bottom_drag(i, j, grid, clock, fields, μ) = @inbounds - μ * fields.u[i, j, 1] -@inline v_bottom_drag(i, j, grid, clock, fields, μ) = @inbounds - μ * fields.v[i, j, 1] - -@testset "Immersed boundaries with hydrostatic free surface models" begin - @info "Testing immersed boundaries with hydrostatic free surface models..." - - for arch in archs - Nx = 60 - Ny = 60 - - # A spherical domain - underlying_grid = LatitudeLongitudeGrid(arch, size = (Nx, Ny, 1), - longitude = (-30, 30), - latitude = (15, 75), - z = (-4000, 0)) - - bathymetry = zeros(Nx, Ny) .- 4000 - view(bathymetry, 31:34, 43:47) .= 0 - bathymetry = arch_array(arch, bathymetry) - - grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bathymetry)) - - free_surface = ImplicitFreeSurface(gravitational_acceleration=0.1) - coriolis = HydrostaticSphericalCoriolis(scheme = VectorInvariantEnstrophyConserving()) - - surface_wind_stress_parameters = (τ₀ = 1e-4, - Lφ = grid.Ly, - φ₀ = 15) - - surface_wind_stress_bc = FluxBoundaryCondition(surface_wind_stress, - parameters = surface_wind_stress_parameters) - - μ = 1 / 60days - - u_bottom_drag_bc = FluxBoundaryCondition(u_bottom_drag, - discrete_form = true, - parameters = μ) - - v_bottom_drag_bc = FluxBoundaryCondition(v_bottom_drag, - discrete_form = true, - parameters = μ) - - u_bcs = FieldBoundaryConditions(top = surface_wind_stress_bc, bottom = u_bottom_drag_bc) - v_bcs = FieldBoundaryConditions(bottom = v_bottom_drag_bc) - - νh₀ = 5e3 * (60 / grid.Nx)^2 - constant_horizontal_diffusivity = HorizontallyCurvilinearAnisotropicDiffusivity(νh=νh₀) - - model = HydrostaticFreeSurfaceModel(grid = grid, - momentum_advection = VectorInvariant(), - free_surface = free_surface, - coriolis = coriolis, - boundary_conditions = (u=u_bcs, v=v_bcs), - closure = constant_horizontal_diffusivity, - tracers = nothing, - buoyancy = nothing) - - - simulation = Simulation(model, Δt = 3600, stop_time = 3600) - - run!(simulation) - - # If reached here it didn't error, so pass for now! - @test true - end -end - diff --git a/test/test_hydrostatic_free_surface_immersed_boundaries_vertical_integrals.jl b/test/test_hydrostatic_free_surface_immersed_boundaries_vertical_integrals.jl deleted file mode 100644 index c6649fa4da..0000000000 --- a/test/test_hydrostatic_free_surface_immersed_boundaries_vertical_integrals.jl +++ /dev/null @@ -1,61 +0,0 @@ -using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBottom -using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization - -@testset "Immersed boundaries with hydrostatic free surface models" begin - @info "Testing immersed boundaries vertical integrals" - - for arch in archs - Nx = 5 - Ny = 5 - - underlying_grid = RectilinearGrid(arch, - size = (Nx, Ny, 3), - extent = (Nx, Ny, 3), - topology = (Periodic, Periodic, Bounded)) - - # B for bathymetry - B = [-3. for i=1:Nx, j=1:Ny ] - B[2:Nx-1,2:Ny-1] .= [-2. for i=2:Nx-1, j=2:Ny-1 ] - B[3:Nx-2,3:Ny-2] .= [-1. for i=3:Nx-2, j=3:Ny-2 ] - - grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(B)) - - model = HydrostaticFreeSurfaceModel(grid = grid, - free_surface = ImplicitFreeSurface(), - tracer_advection = WENO5(), - buoyancy = nothing, - tracers = nothing, - closure = nothing) - - x_ref = [0.0 0.0 0.0 0.0 0.0 0.0 0.0 - 0.0 3.0 3.0 3.0 3.0 3.0 0.0 - 0.0 3.0 2.0 2.0 2.0 2.0 0.0 - 0.0 3.0 2.0 1.0 1.0 2.0 0.0 - 0.0 3.0 2.0 2.0 2.0 2.0 0.0 - 0.0 3.0 3.0 3.0 3.0 3.0 0.0 - 0.0 0.0 0.0 0.0 0.0 0.0 0.0]' - - y_ref = [0.0 0.0 0.0 0.0 0.0 0.0 0.0 - 0.0 3.0 3.0 3.0 3.0 3.0 0.0 - 0.0 3.0 2.0 2.0 2.0 3.0 0.0 - 0.0 3.0 2.0 1.0 2.0 3.0 0.0 - 0.0 3.0 2.0 1.0 2.0 3.0 0.0 - 0.0 3.0 2.0 2.0 2.0 3.0 0.0 - 0.0 0.0 0.0 0.0 0.0 0.0 0.0]' - - fs = model.free_surface - vertically_integrated_lateral_areas = fs.implicit_step_solver.vertically_integrated_lateral_areas - - ∫Axᶠᶜᶜ = vertically_integrated_lateral_areas.xᶠᶜᶜ - ∫Ayᶜᶠᶜ = vertically_integrated_lateral_areas.yᶜᶠᶜ - - ∫Axᶠᶜᶜ = Array(parent(∫Axᶠᶜᶜ)) - ∫Ayᶜᶠᶜ = Array(parent(∫Ayᶜᶠᶜ)) - - Ax_ok = ∫Axᶠᶜᶜ[:, :, 1] ≈ x_ref - Ay_ok = ∫Ayᶜᶠᶜ[:, :, 1] ≈ y_ref - - @test (Ax_ok & Ay_ok) - end -end - diff --git a/test/test_preconditioned_conjugate_gradient_solver.jl b/test/test_preconditioned_conjugate_gradient_solver.jl index d46f215ffe..313688165a 100644 --- a/test/test_preconditioned_conjugate_gradient_solver.jl +++ b/test/test_preconditioned_conjugate_gradient_solver.jl @@ -1,3 +1,5 @@ +include("dependencies_for_runtests.jl") + using Oceananigans.Solvers: solve! using Statistics @@ -6,8 +8,8 @@ function identity_operator!(b, x) return nothing end -function run_identity_operator_test(arch, grid) - b = Field(Center, Center, Center, arch, grid) +function run_identity_operator_test(grid) + b = CenterField(grid) solver = PreconditionedConjugateGradientSolver(identity_operator!, template_field = b) @@ -19,9 +21,10 @@ function run_identity_operator_test(arch, grid) @test norm(solution) .< solver.tolerance end -function run_poisson_equation_test(arch, grid) +function run_poisson_equation_test(grid) + arch = architecture(grid) # Solve ∇²ϕ = r - ϕ_truth = Field(Center, Center, Center, arch, grid) + ϕ_truth = CenterField(grid) # Initialize zero-mean "truth" solution with random numbers set!(ϕ_truth, (x, y, z) -> rand()) @@ -29,17 +32,17 @@ function run_poisson_equation_test(arch, grid) fill_halo_regions!(ϕ_truth, arch) # Calculate Laplacian of "truth" - ∇²ϕ = r = Field(Center, Center, Center, arch, grid) + ∇²ϕ = r = CenterField(grid) compute_∇²!(∇²ϕ, ϕ_truth, arch, grid) solver = PreconditionedConjugateGradientSolver(compute_∇²!, template_field=ϕ_truth) # Solve Poisson equation - ϕ_solution = Field(Center, Center, Center, arch, grid) + ϕ_solution = CenterField(grid) solve!(ϕ_solution, solver, r, arch, grid) # Diagnose Laplacian of solution - ∇²ϕ_solution = Field(Center, Center, Center, arch, grid) + ∇²ϕ_solution = CenterField(grid) compute_∇²!(∇²ϕ_solution, ϕ_solution, arch, grid) # Test @@ -63,7 +66,7 @@ end for arch in archs @info "Testing PreconditionedConjugateGradientSolver [$(typeof(arch))]..." grid = RectilinearGrid(arch, size=(4, 8, 4), extent=(1, 3, 1)) - run_identity_operator_test(arch, grid) - run_poisson_equation_test(arch, grid) + run_identity_operator_test(grid) + run_poisson_equation_test(grid) end end From d9df9df3b3dc1f3dc2730ad160458246b244a246 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 17 Dec 2021 15:55:09 -0700 Subject: [PATCH 007/140] Fixes for test_poisson_solvers --- src/Fields/field.jl | 16 ++++-- .../solve_for_pressure.jl | 6 +- src/Solvers/batched_tridiagonal_solver.jl | 20 ++++--- src/Solvers/discrete_transforms.jl | 17 ++++-- src/Solvers/fft_based_poisson_solver.jl | 21 ++++--- .../fourier_tridiagonal_poisson_solver.jl | 19 ++++--- src/Solvers/plan_transforms.jl | 48 ++++++++-------- test/test_poisson_solvers.jl | 56 ++++++++++--------- 8 files changed, 117 insertions(+), 86 deletions(-) diff --git a/src/Fields/field.jl b/src/Fields/field.jl index 3dd9bb9379..1fdbdcf1a4 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -58,8 +58,12 @@ Field located at (Face, Face, Center) └── boundary conditions: west=Periodic, east=Periodic, south=Periodic, north=Periodic, bottom=ZeroFlux, top=ZeroFlux, immersed=ZeroFlux ``` """ -Field{LX, LY, LZ}(grid::AbstractGrid, T::DataType=eltype(grid); kw...) where {LX, LY, LZ} = - Field((LX, LY, LZ), grid, T; kw...) +function Field{LX, LY, LZ}(grid::AbstractGrid, + T::DataType=eltype(grid); + kw...) where {LX, LY, LZ} + + return Field((LX, LY, LZ), grid, T; kw...) +end function Field(loc::Tuple, grid::AbstractGrid, @@ -120,7 +124,7 @@ Adapt.adapt_structure(to, field::Field) = Adapt.adapt(to, field.data) Returns `Field{Center, Center, Center}` on `arch`itecture and `grid`. Additional keyword arguments are passed to the `Field` constructor. """ -CenterField(grid::AbstractGrid, FT=eltype(grid); kw...) = Field{Center, Center, Center}(grid, FT; kw...) +CenterField(grid::AbstractGrid, T::DataType=eltype(grid); kw...) = Field{Center, Center, Center}(grid, T; kw...) """ XFaceField(grid; kw...) @@ -128,7 +132,7 @@ CenterField(grid::AbstractGrid, FT=eltype(grid); kw...) = Field{Center, Center, Returns `Field{Face, Center, Center}` on `grid`. Additional keyword arguments are passed to the `Field` constructor. """ -XFaceField(grid::AbstractGrid, FT=eltype(grid); kw...) = Field{Face, Center, Center}(grid, FT; kw...) +XFaceField(grid::AbstractGrid, T::DataType=eltype(grid); kw...) = Field{Face, Center, Center}(grid, T; kw...) """ YFaceField(grid; kw...) @@ -136,7 +140,7 @@ XFaceField(grid::AbstractGrid, FT=eltype(grid); kw...) = Field{Face, Center, Cen Returns `Field{Center, Face, Center}` on `grid`. Additional keyword arguments are passed to the `Field` constructor. """ -YFaceField(grid::AbstractGrid, FT=eltype(grid); kw...) = Field{Center, Face, Center}(grid, FT; kw...) +YFaceField(grid::AbstractGrid, T::DataType=eltype(grid); kw...) = Field{Center, Face, Center}(grid, T; kw...) """ ZFaceField(grid; kw...) @@ -144,7 +148,7 @@ YFaceField(grid::AbstractGrid, FT=eltype(grid); kw...) = Field{Center, Face, Cen Returns `Field{Center, Center, Face}` on `grid`. Additional keyword arguments are passed to the `Field` constructor. """ -ZFaceField(grid::AbstractGrid, FT=eltype(grid); kw...) = Field{Center, Center, Face}(grid, FT; kw...) +ZFaceField(grid::AbstractGrid, T::DataType=eltype(grid); kw...) = Field{Center, Center, Face}(grid, T; kw...) ##### ##### Fields computed from AbstractOperation and associated utilities diff --git a/src/Models/NonhydrostaticModels/solve_for_pressure.jl b/src/Models/NonhydrostaticModels/solve_for_pressure.jl index da80c05832..4d0c321400 100644 --- a/src/Models/NonhydrostaticModels/solve_for_pressure.jl +++ b/src/Models/NonhydrostaticModels/solve_for_pressure.jl @@ -24,7 +24,7 @@ end function solve_for_pressure!(pressure, solver::DistributedFFTBasedPoissonSolver, Δt, U★) rhs = first(solver.storage) - arch = solver.architecture + arch = architecture(solver) grid = solver.my_grid rhs_event = launch!(arch, grid, :xyz, calculate_pressure_source_term_fft_based_solver!, @@ -43,7 +43,7 @@ function solve_for_pressure!(pressure, solver::FFTBasedPoissonSolver, Δt, U★) # Calculate right hand side: rhs = solver.storage - arch = solver.architecture + arch = architecture(solver) grid = solver.grid rhs_event = launch!(arch, grid, :xyz, calculate_pressure_source_term_fft_based_solver!, @@ -61,7 +61,7 @@ function solve_for_pressure!(pressure, solver::FourierTridiagonalPoissonSolver, # Calculate right hand side: rhs = solver.source_term - arch = solver.architecture + arch = architecture(solver) grid = solver.grid rhs_event = launch!(arch, grid, :xyz, calculate_pressure_source_term_fourier_tridiagonal_solver!, diff --git a/src/Solvers/batched_tridiagonal_solver.jl b/src/Solvers/batched_tridiagonal_solver.jl index f51700ea81..2cad8e6407 100644 --- a/src/Solvers/batched_tridiagonal_solver.jl +++ b/src/Solvers/batched_tridiagonal_solver.jl @@ -1,20 +1,24 @@ using Oceananigans.Architectures: device_event, arch_array +import Oceananigans.Architectures: architecture + """ BatchedTridiagonalSolver A batched solver for large numbers of triadiagonal systems. """ -struct BatchedTridiagonalSolver{A, B, C, T, G, R, P} +struct BatchedTridiagonalSolver{A, B, C, T, G, P} a :: A b :: B c :: C t :: T grid :: G - architecture :: R parameters :: P end +architecture(solver::BatchedTridiagonalSolver) = architecture(solver.grid) + + """ BatchedTridiagonalSolver(grid; lower_diagonal, diagonal, upper_diagonal, parameters=nothing) @@ -38,15 +42,15 @@ where `a` is the `lower_diagonal`, `b` is the `diagonal`, and `c` is the `upper_ * If `!isnothing(parameters)` then `aⁱʲᵏ = a(i, j, k, grid, parameters, args...)`. where `args...` are `Varargs` passed to `solve_batched_tridiagonal_system!(ϕ, solver, args...)`. """ -function BatchedTridiagonalSolver(arch, grid; +function BatchedTridiagonalSolver(grid; lower_diagonal, diagonal, upper_diagonal, - scratch = arch_array(arch, zeros(eltype(grid), grid.Nx, grid.Ny, grid.Nz)), + scratch = arch_array(architecture(grid), zeros(eltype(grid), grid.Nx, grid.Ny, grid.Nz)), parameters = nothing) return BatchedTridiagonalSolver(lower_diagonal, diagonal, upper_diagonal, - scratch, grid, arch, parameters) + scratch, grid, parameters) end @inline get_coefficient(a::AbstractArray{T, 1}, i, j, k, grid, p, args...) where {T} = @inbounds a[k] @@ -67,16 +71,16 @@ The result is stored in `ϕ` which must have size `(grid.Nx, grid.Ny, grid.Nz)`. Reference implementation per Numerical Recipes, Press et. al 1992 (§ 2.4). """ -function solve!(ϕ, solver::BatchedTridiagonalSolver, rhs, args...; dependencies = device_event(solver.architecture)) +function solve!(ϕ, solver::BatchedTridiagonalSolver, rhs, args...; dependencies = device_event(architecture(solver))) a, b, c, t, parameters = solver.a, solver.b, solver.c, solver.t, solver.parameters grid = solver.grid - event = launch!(solver.architecture, grid, :xy, + event = launch!(architecture(solver), grid, :xy, solve_batched_tridiagonal_system_kernel!, ϕ, a, b, c, rhs, t, grid, parameters, args..., dependencies = dependencies) - wait(device(solver.architecture), event) + wait(device(architecture(solver)), event) return nothing end diff --git a/src/Solvers/discrete_transforms.jl b/src/Solvers/discrete_transforms.jl index bf6f43215a..e6e06085b8 100644 --- a/src/Solvers/discrete_transforms.jl +++ b/src/Solvers/discrete_transforms.jl @@ -1,11 +1,12 @@ +import Oceananigans.Architectures: architecture + abstract type AbstractTransformDirection end struct Forward <: AbstractTransformDirection end struct Backward <: AbstractTransformDirection end -struct DiscreteTransform{P, D, A, G, Δ, Ω, N, T, Σ} +struct DiscreteTransform{P, D, G, Δ, Ω, N, T, Σ} plan :: P - architecture :: A grid :: G direction :: D dims :: Δ @@ -15,6 +16,8 @@ struct DiscreteTransform{P, D, A, G, Δ, Ω, N, T, Σ} transpose_dims :: Σ end +architecture(transform::DiscreteTransform) = architecture(transform.grid) + ##### ##### Normalization factors ##### @@ -76,7 +79,9 @@ end NoTransform() = DiscreteTransform([nothing for _ in fieldnames(DiscreteTransform)]...) -function DiscreteTransform(plan, direction, arch, grid, dims) +function DiscreteTransform(plan, direction, grid, dims) + arch = architecture(grid) + isnothing(plan) && return NoTransform() N = size(grid) @@ -90,7 +95,7 @@ function DiscreteTransform(plan, direction, arch, grid, dims) dims = length(dims) == 1 ? dims[1] : dims - return DiscreteTransform(plan, arch, grid, direction, dims, topo, normalization, twiddle, transpose) + return DiscreteTransform(plan, grid, direction, dims, topo, normalization, twiddle, transpose) end ##### @@ -100,7 +105,7 @@ end (transform::DiscreteTransform{<:Nothing})(A, buffer) = nothing function (transform::DiscreteTransform{P, <:Forward})(A, buffer) where P - maybe_permute_indices!(A, buffer, transform.architecture, transform.grid, transform.dims, transform.topology) + maybe_permute_indices!(A, buffer, architecture(transform), transform.grid, transform.dims, transform.topology) apply_transform!(A, buffer, transform.plan, transform.transpose_dims) maybe_twiddle_forward!(A, transform.twiddle_factors) maybe_normalize!(A, transform.normalization) @@ -110,7 +115,7 @@ end function (transform::DiscreteTransform{P, <:Backward})(A, buffer) where P maybe_twiddle_backward!(A, transform.twiddle_factors) apply_transform!(A, buffer, transform.plan, transform.transpose_dims) - maybe_unpermute_indices!(A, buffer, transform.architecture, transform.grid, transform.dims, transform.topology) + maybe_unpermute_indices!(A, buffer, architecture(transform), transform.grid, transform.dims, transform.topology) maybe_normalize!(A, transform.normalization) return nothing end diff --git a/src/Solvers/fft_based_poisson_solver.jl b/src/Solvers/fft_based_poisson_solver.jl index 878884d1f5..b7550f7e3c 100644 --- a/src/Solvers/fft_based_poisson_solver.jl +++ b/src/Solvers/fft_based_poisson_solver.jl @@ -1,8 +1,9 @@ using Oceananigans.Architectures: device_event using Oceananigans: short_show -struct FFTBasedPoissonSolver{A, G, Λ, S, B, T} - architecture :: A +import Oceananigans.Architectures: architecture + +struct FFTBasedPoissonSolver{G, Λ, S, B, T} grid :: G eigenvalues :: Λ storage :: S @@ -10,6 +11,8 @@ struct FFTBasedPoissonSolver{A, G, Λ, S, B, T} transforms :: T end +architecture(solver::FFTBasedPoissonSolver) = architecture(solver.grid) + transform_str(transform) = string(typeof(transform).name.wrapper, ", ") function transform_list_str(transform_list) @@ -19,8 +22,8 @@ function transform_list_str(transform_list) return list end -Base.show(io::IO, solver::FFTBasedPoissonSolver{A, G}) where {A, G}= - print(io, "FFTBasedPoissonSolver{$A}: \n", +Base.show(io::IO, solver::FFTBasedPoissonSolver) = +print(io, "FFTBasedPoissonSolver on ", string(typeof(architecture(solver))), ": \n", "├── grid: $(short_show(solver.grid))\n", "├── storage: $(typeof(solver.storage))\n", "├── buffer: $(typeof(solver.buffer))\n", @@ -28,13 +31,15 @@ Base.show(io::IO, solver::FFTBasedPoissonSolver{A, G}) where {A, G}= " ├── forward: ", transform_list_str(solver.transforms.forward), "\n", " └── backward: ", transform_list_str(solver.transforms.backward)) -function FFTBasedPoissonSolver(arch, grid, planner_flag=FFTW.PATIENT) +function FFTBasedPoissonSolver(grid, planner_flag=FFTW.PATIENT) topo = (TX, TY, TZ) = topology(grid) λx = poisson_eigenvalues(grid.Nx, grid.Lx, 1, TX()) λy = poisson_eigenvalues(grid.Ny, grid.Ly, 2, TY()) λz = poisson_eigenvalues(grid.Nz, grid.Lz, 3, TZ()) + arch = architecture(grid) + eigenvalues = ( λx = arch_array(arch, λx), λy = arch_array(arch, λy), @@ -43,13 +48,13 @@ function FFTBasedPoissonSolver(arch, grid, planner_flag=FFTW.PATIENT) storage = arch_array(arch, zeros(complex(eltype(grid)), size(grid)...)) - transforms = plan_transforms(arch, grid, storage, planner_flag) + transforms = plan_transforms(grid, storage, planner_flag) # Need buffer for index permutations and transposes. buffer_needed = arch isa GPU && Bounded in topo buffer = buffer_needed ? similar(storage) : nothing - return FFTBasedPoissonSolver(arch, grid, eigenvalues, storage, buffer, transforms) + return FFTBasedPoissonSolver(grid, eigenvalues, storage, buffer, transforms) end """ @@ -71,7 +76,7 @@ Note: ``(∇² + m) ϕ = b`` is sometimes called the "screened Poisson" equation when ``m < 0``, or the Helmholtz equation when ``m > 0``. """ function solve!(ϕ, solver::FFTBasedPoissonSolver, b, m=0) - arch = solver.architecture + arch = architecture(solver) topo = TX, TY, TZ = topology(solver.grid) Nx, Ny, Nz = size(solver.grid) λx, λy, λz = solver.eigenvalues diff --git a/src/Solvers/fourier_tridiagonal_poisson_solver.jl b/src/Solvers/fourier_tridiagonal_poisson_solver.jl index 644307e377..b5dd757f99 100644 --- a/src/Solvers/fourier_tridiagonal_poisson_solver.jl +++ b/src/Solvers/fourier_tridiagonal_poisson_solver.jl @@ -1,8 +1,8 @@ using Oceananigans.Operators: Δzᵃᵃᶜ, Δzᵃᵃᶠ using Oceananigans.Architectures: device_event +import Oceananigans.Architectures: architecture -struct FourierTridiagonalPoissonSolver{A, G, B, R, S, β, T} - architecture :: A +struct FourierTridiagonalPoissonSolver{G, B, R, S, β, T} grid :: G batched_tridiagonal_solver :: B source_term :: R @@ -11,6 +11,8 @@ struct FourierTridiagonalPoissonSolver{A, G, B, R, S, β, T} transforms :: T end +architecture(solver::FourierTridiagonalPoissonSolver) = architecture(solver.grid) + @kernel function compute_main_diagonals!(D, grid, λx, λy) i, j = @index(Global, NTuple) Nz = grid.Nz @@ -25,7 +27,7 @@ end D[i, j, Nz] = -1 / Δzᵃᵃᶠ(i, j, Nz, grid) - Δzᵃᵃᶜ(i, j, Nz, grid) * (λx[i] + λy[j]) end -function FourierTridiagonalPoissonSolver(arch, grid, planner_flag=FFTW.PATIENT) +function FourierTridiagonalPoissonSolver(grid, planner_flag=FFTW.PATIENT) TX, TY, TZ = topology(grid) TZ != Bounded && error("FourierTridiagonalPoissonSolver can only be used with a Bounded z topology.") @@ -35,12 +37,13 @@ function FourierTridiagonalPoissonSolver(arch, grid, planner_flag=FFTW.PATIENT) λx = poisson_eigenvalues(grid.Nx, grid.Lx, 1, TX()) λy = poisson_eigenvalues(grid.Ny, grid.Ly, 2, TY()) + arch = architecture(grid) λx = arch_array(arch, λx) λy = arch_array(arch, λy) # Plan required transforms for x and y sol_storage = arch_array(arch, zeros(complex(eltype(grid)), size(grid)...)) - transforms = plan_transforms(arch, grid, sol_storage, planner_flag) + transforms = plan_transforms(grid, sol_storage, planner_flag) # Lower and upper diagonals are the same lower_diagonal = CUDA.@allowscalar [1 / Δzᵃᵃᶠ(1, 1, k, grid) for k in 2:Nz] @@ -53,7 +56,7 @@ function FourierTridiagonalPoissonSolver(arch, grid, planner_flag=FFTW.PATIENT) wait(device(arch), event) # Set up batched tridiagonal solver - btsolver = BatchedTridiagonalSolver(arch, grid; + btsolver = BatchedTridiagonalSolver(grid; lower_diagonal = lower_diagonal, diagonal = diagonal, upper_diagonal = upper_diagonal) @@ -65,13 +68,13 @@ function FourierTridiagonalPoissonSolver(arch, grid, planner_flag=FFTW.PATIENT) # Storage space for right hand side of Poisson equation rhs = arch_array(arch, zeros(complex(eltype(grid)), size(grid)...)) - return FourierTridiagonalPoissonSolver(arch, grid, btsolver, rhs, sol_storage, buffer, transforms) + return FourierTridiagonalPoissonSolver(grid, btsolver, rhs, sol_storage, buffer, transforms) end function solve!(x, solver::FourierTridiagonalPoissonSolver, b=nothing) !isnothing(b) && set_source_term!(solver, b) # otherwise, assume source term is set correctly - arch = solver.architecture + arch = architecture(solver) ϕ = solver.storage # Apply forward transforms in order @@ -105,7 +108,7 @@ to `source_term` by multiplying it by the vertical grid spacing at z cell center """ function set_source_term!(solver::FourierTridiagonalPoissonSolver, source_term) grid = solver.grid - arch = solver.architecture + arch = architecture(solver) solver.source_term .= source_term event = launch!(arch, grid, :xyz, multiply_by_Δzᵃᵃᶜ!, solver.source_term, grid, dependencies=Event(device(arch))) diff --git a/src/Solvers/plan_transforms.jl b/src/Solvers/plan_transforms.jl index b2e7bb0548..85dfe56ee8 100644 --- a/src/Solvers/plan_transforms.jl +++ b/src/Solvers/plan_transforms.jl @@ -80,7 +80,7 @@ backward_orders(::Type{Bounded}, ::Type{Bounded}, ::Type{Periodic}) = (3, 1, 2 backward_orders(::Type{Bounded}, ::Type{Bounded}, ::Type{Bounded}) = (1, 2, 3) " Used by FFTBasedPoissonSolver " -function plan_transforms(arch, grid::Regular, storage, planner_flag) +function plan_transforms(grid::Regular, storage, planner_flag) Nx, Ny, Nz = size(grid) topo = topology(grid) periodic_dims = findall(t -> t == Periodic, topo) @@ -90,6 +90,8 @@ function plan_transforms(arch, grid::Regular, storage, planner_flag) # Note that transforms are omitted in Flat directions. unflattened_topo = Tuple(T() isa Flat ? Bounded : T for T in topo) + arch = architecture(grid) + if arch isa GPU && !(unflattened_topo in batchable_GPU_topologies) rs_storage = reshape(storage, (Ny, Nx, Nz)) @@ -108,15 +110,15 @@ function plan_transforms(arch, grid::Regular, storage, planner_flag) b_order = backward_orders(unflattened_topo...) forward_transforms = ( - DiscreteTransform(forward_plans[f_order[1]], Forward(), arch, grid, [f_order[1]]), - DiscreteTransform(forward_plans[f_order[2]], Forward(), arch, grid, [f_order[2]]), - DiscreteTransform(forward_plans[f_order[3]], Forward(), arch, grid, [f_order[3]]) + DiscreteTransform(forward_plans[f_order[1]], Forward(), grid, [f_order[1]]), + DiscreteTransform(forward_plans[f_order[2]], Forward(), grid, [f_order[2]]), + DiscreteTransform(forward_plans[f_order[3]], Forward(), grid, [f_order[3]]) ) backward_transforms = ( - DiscreteTransform(backward_plans[b_order[1]], Backward(), arch, grid, [b_order[1]]), - DiscreteTransform(backward_plans[b_order[2]], Backward(), arch, grid, [b_order[2]]), - DiscreteTransform(backward_plans[b_order[3]], Backward(), arch, grid, [b_order[3]]) + DiscreteTransform(backward_plans[b_order[1]], Backward(), grid, [b_order[1]]), + DiscreteTransform(backward_plans[b_order[2]], Backward(), grid, [b_order[2]]), + DiscreteTransform(backward_plans[b_order[3]], Backward(), grid, [b_order[3]]) ) else @@ -130,16 +132,16 @@ function plan_transforms(arch, grid::Regular, storage, planner_flag) forward_bounded_plan = plan_forward_transform(storage, Bounded(), bounded_dims, planner_flag) forward_transforms = ( - DiscreteTransform(forward_bounded_plan, Forward(), arch, grid, bounded_dims), - DiscreteTransform(forward_periodic_plan, Forward(), arch, grid, periodic_dims) + DiscreteTransform(forward_bounded_plan, Forward(), grid, bounded_dims), + DiscreteTransform(forward_periodic_plan, Forward(), grid, periodic_dims) ) backward_periodic_plan = plan_backward_transform(storage, Periodic(), periodic_dims, planner_flag) backward_bounded_plan = plan_backward_transform(storage, Bounded(), bounded_dims, planner_flag) backward_transforms = ( - DiscreteTransform(backward_periodic_plan, Backward(), arch, grid, periodic_dims), - DiscreteTransform(backward_bounded_plan, Backward(), arch, grid, bounded_dims) + DiscreteTransform(backward_periodic_plan, Backward(), grid, periodic_dims), + DiscreteTransform(backward_bounded_plan, Backward(), grid, bounded_dims) ) end @@ -150,7 +152,7 @@ end """ Used by FourierTridiagonalPoissonSolver. """ -function plan_transforms(arch, grid::VerticallyStretched, storage, planner_flag) +function plan_transforms(grid::VerticallyStretched, storage, planner_flag) Nx, Ny, Nz = size(grid) TX, TY, TZ = topo = topology(grid) @@ -162,6 +164,8 @@ function plan_transforms(arch, grid::VerticallyStretched, storage, planner_flag) !(topo[3] === Bounded) && error("Cannot plan transforms on z-periodic RectilinearGrids.") + arch = architecture(grid) + if arch isa CPU # This is the case where batching transforms is possible. It's always possible on the CPU # since FFTW is awesome so it includes all topologies on the CPU. @@ -172,14 +176,14 @@ function plan_transforms(arch, grid::VerticallyStretched, storage, planner_flag) forward_periodic_plan = plan_forward_transform(storage, Periodic(), periodic_dims, planner_flag) forward_bounded_plan = plan_forward_transform(storage, Bounded(), bounded_dims, planner_flag) - forward_transforms = (DiscreteTransform(forward_bounded_plan, Forward(), arch, grid, bounded_dims), - DiscreteTransform(forward_periodic_plan, Forward(), arch, grid, periodic_dims)) + forward_transforms = (DiscreteTransform(forward_bounded_plan, Forward(), grid, bounded_dims), + DiscreteTransform(forward_periodic_plan, Forward(), grid, periodic_dims)) backward_periodic_plan = plan_backward_transform(storage, Periodic(), periodic_dims, planner_flag) backward_bounded_plan = plan_backward_transform(storage, Bounded(), bounded_dims, planner_flag) - backward_transforms = (DiscreteTransform(backward_periodic_plan, Backward(), arch, grid, periodic_dims), - DiscreteTransform(backward_bounded_plan, Backward(), arch, grid, bounded_dims)) + backward_transforms = (DiscreteTransform(backward_periodic_plan, Backward(), grid, periodic_dims), + DiscreteTransform(backward_bounded_plan, Backward(), grid, bounded_dims)) elseif !(Bounded in (TX, TY)) # We're on the GPU and either (Periodic, Periodic), (Flat, Periodic), or (Periodic, Flat) in xy. @@ -187,8 +191,8 @@ function plan_transforms(arch, grid::VerticallyStretched, storage, planner_flag) forward_periodic_plan = plan_forward_transform(storage, Periodic(), [1, 2], planner_flag) backward_periodic_plan = plan_backward_transform(storage, Periodic(), [1, 2], planner_flag) - forward_transforms = tuple(DiscreteTransform(forward_periodic_plan, Forward(), arch, grid, [1, 2])) - backward_transforms = tuple(DiscreteTransform(backward_periodic_plan, Backward(), arch, grid, [1, 2])) + forward_transforms = tuple(DiscreteTransform(forward_periodic_plan, Forward(), grid, [1, 2])) + backward_transforms = tuple(DiscreteTransform(backward_periodic_plan, Backward(), grid, [1, 2])) else # we are on the GPU and we cannot / should not batch! rs_storage = reshape(storage, (Ny, Nx, Nz)) @@ -210,11 +214,11 @@ function plan_transforms(arch, grid::VerticallyStretched, storage, planner_flag) f_order = Tuple(f_order[i] for i in findall(d -> d != 3, f_order)) b_order = Tuple(b_order[i] for i in findall(d -> d != 3, b_order)) - forward_transforms = (DiscreteTransform(forward_plans[f_order[1]], Forward(), arch, grid, [f_order[1]]), - DiscreteTransform(forward_plans[f_order[2]], Forward(), arch, grid, [f_order[2]])) + forward_transforms = (DiscreteTransform(forward_plans[f_order[1]], Forward(), grid, [f_order[1]]), + DiscreteTransform(forward_plans[f_order[2]], Forward(), grid, [f_order[2]])) - backward_transforms = (DiscreteTransform(backward_plans[b_order[1]], Backward(), arch, grid, [b_order[1]]), - DiscreteTransform(backward_plans[b_order[2]], Backward(), arch, grid, [b_order[2]])) + backward_transforms = (DiscreteTransform(backward_plans[b_order[1]], Backward(), grid, [b_order[1]]), + DiscreteTransform(backward_plans[b_order[2]], Backward(), grid, [b_order[2]])) end transforms = (forward=forward_transforms, backward=backward_transforms) diff --git a/test/test_poisson_solvers.jl b/test/test_poisson_solvers.jl index 05ceea485f..d80e59abe2 100644 --- a/test/test_poisson_solvers.jl +++ b/test/test_poisson_solvers.jl @@ -1,23 +1,27 @@ +include("dependencies_for_runtests.jl") + +using CUDA using Oceananigans.Solvers: solve!, set_source_term! using Oceananigans.Solvers: poisson_eigenvalues using Oceananigans.Models.NonhydrostaticModels: solve_for_pressure! using Oceananigans.Models.HydrostaticFreeSurfaceModels: _compute_w_from_continuity! using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions -function poisson_solver_instantiates(arch, grid, planner_flag) - solver = FFTBasedPoissonSolver(arch, grid, planner_flag) +function poisson_solver_instantiates(grid, planner_flag) + solver = FFTBasedPoissonSolver(grid, planner_flag) return true # Just making sure the FFTBasedPoissonSolver does not error/crash. end -function random_divergent_source_term(arch, grid) +function random_divergent_source_term(grid) + arch = architecture(grid) default_bcs = FieldBoundaryConditions() u_bcs = regularize_field_boundary_conditions(default_bcs, grid, :u) v_bcs = regularize_field_boundary_conditions(default_bcs, grid, :v) w_bcs = regularize_field_boundary_conditions(default_bcs, grid, :w) - Ru = CenterField(arch, grid, u_bcs) - Rv = CenterField(arch, grid, v_bcs) - Rw = CenterField(arch, grid, w_bcs) + Ru = CenterField(grid, boundary_conditions=u_bcs) + Rv = CenterField(grid, boundary_conditions=v_bcs) + Rw = CenterField(grid, boundary_conditions=w_bcs) U = (u=Ru, v=Rv, w=Rw) Nx, Ny, Nz = size(grid) @@ -39,16 +43,16 @@ function random_divergent_source_term(arch, grid) return R, U end -function random_divergence_free_source_term(arch, grid) +function random_divergence_free_source_term(grid) default_bcs = FieldBoundaryConditions() u_bcs = regularize_field_boundary_conditions(default_bcs, grid, :u) v_bcs = regularize_field_boundary_conditions(default_bcs, grid, :v) w_bcs = regularize_field_boundary_conditions(default_bcs, grid, :w) # Random right hand side - Ru = CenterField(arch, grid, u_bcs) - Rv = CenterField(arch, grid, v_bcs) - Rw = CenterField(arch, grid, w_bcs) + Ru = CenterField(grid, boundary_conditions=u_bcs) + Rv = CenterField(grid, boundary_conditions=v_bcs) + Rw = CenterField(grid, boundary_conditions=w_bcs) U = (u=Ru, v=Rv, w=Rw) Nx, Ny, Nz = size(grid) @@ -56,6 +60,7 @@ function random_divergence_free_source_term(arch, grid) set!(Rv, rand(Nx, Ny, Nz)) set!(Rw, zeros(Nx, Ny, Nz)) + arch = architecture(grid) fill_halo_regions!(Ru, arch, nothing, nothing) fill_halo_regions!(Rv, arch, nothing, nothing) fill_halo_regions!(Rw, arch, nothing, nothing) @@ -80,16 +85,17 @@ end ##### Regular rectilinear grid Poisson solver ##### -function divergence_free_poisson_solution(arch, grid, planner_flag=FFTW.MEASURE) +function divergence_free_poisson_solution(grid, planner_flag=FFTW.MEASURE) + arch = architecture(grid) ArrayType = array_type(arch) FT = eltype(grid) - solver = FFTBasedPoissonSolver(arch, grid, planner_flag) - R, U = random_divergent_source_term(arch, grid) + solver = FFTBasedPoissonSolver(grid, planner_flag) + R, U = random_divergent_source_term(grid) p_bcs = FieldBoundaryConditions(grid, (Center, Center, Center)) - ϕ = CenterField(arch, grid, p_bcs) # "kinematic pressure" - ∇²ϕ = CenterField(arch, grid, p_bcs) + ϕ = CenterField(grid, boundary_conditions=p_bcs) # "kinematic pressure" + ∇²ϕ = CenterField(grid, boundary_conditions=p_bcs) # Using Δt = 1 but it doesn't matter since velocities = 0. solve_for_pressure!(ϕ.data, solver, 1, U) @@ -111,7 +117,7 @@ k²(::Type{Periodic}, n) = n^2 function analytical_poisson_solver_test(arch, N, topo; FT=Float64, mode=1) grid = RectilinearGrid(arch, FT, topology=topo, size=(N, N, N), x=(0, 2π), y=(0, 2π), z=(0, 2π)) - solver = FFTBasedPoissonSolver(arch, grid) + solver = FFTBasedPoissonSolver(grid) xC, yC, zC = nodes((Center, Center, Center), grid, reshape=true) @@ -160,13 +166,13 @@ function vertically_stretched_poisson_solver_correct_answer(FT, arch, topo, Nx, sz = get_grid_size(topo..., Nx, Ny, Nz) xy_intervals = get_xy_interval_kwargs(topo...) vs_grid = RectilinearGrid(arch, FT; topology=topo, size=sz, z=zF, xy_intervals...) - solver = FourierTridiagonalPoissonSolver(arch, vs_grid) + solver = FourierTridiagonalPoissonSolver(vs_grid) p_bcs = FieldBoundaryConditions(vs_grid, (Center, Center, Center)) - ϕ = CenterField(arch, vs_grid, p_bcs) # "kinematic pressure" - ∇²ϕ = CenterField(arch, vs_grid, p_bcs) + ϕ = CenterField(vs_grid, boundary_conditions=p_bcs) # "kinematic pressure" + ∇²ϕ = CenterField(vs_grid, boundary_conditions=p_bcs) - R = random_divergence_free_source_term(arch, vs_grid) + R = random_divergence_free_source_term(vs_grid) set_source_term!(solver, R) ϕc = solver.storage @@ -214,8 +220,8 @@ two_dimensional_topologies = [(Flat, Bounded, Bounded), push!(grids, grids_3d..., grids_2d...) for grid in grids - @test poisson_solver_instantiates(arch, grid, FFTW.ESTIMATE) - @test poisson_solver_instantiates(arch, grid, FFTW.MEASURE) + @test poisson_solver_instantiates(grid, FFTW.ESTIMATE) + @test poisson_solver_instantiates(grid, FFTW.MEASURE) end end end @@ -239,7 +245,7 @@ two_dimensional_topologies = [(Flat, Bounded, Bounded), for grid in grids N == 7 && @info " Testing $(topology(grid)) topology on square grids [$(typeof(arch))]..." - @test divergence_free_poisson_solution(arch, grid) + @test divergence_free_poisson_solution(grid) end end end @@ -249,7 +255,7 @@ two_dimensional_topologies = [(Flat, Bounded, Bounded), @info " Testing $topo topology on rectangular grids with even and prime sizes [$(typeof(arch))]..." for Nx in Ns, Ny in Ns, Nz in Ns grid = RectilinearGrid(arch, topology=topo, size=(Nx, Ny, Nz), extent=(1, 1, 1)) - @test divergence_free_poisson_solution(arch, grid) + @test divergence_free_poisson_solution(grid) end end @@ -258,7 +264,7 @@ two_dimensional_topologies = [(Flat, Bounded, Bounded), RectilinearGrid(arch, Float32, topology=(Bounded, Bounded, Periodic), size=(7, 11, 13), extent=(1, 1, 1))] for grid in Float32_grids - @test divergence_free_poisson_solution(arch, grid) + @test divergence_free_poisson_solution(grid) end end From 8ea5331f58e9cddfa2a8f65a53baed0284b6d0c6 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 17 Dec 2021 17:40:48 -0700 Subject: [PATCH 008/140] Fix time-stepping tests and DiffusivityFields --- .../NonhydrostaticModels/NonhydrostaticModels.jl | 7 ++++--- src/TurbulenceClosures/diffusivity_fields.jl | 12 ++++++------ .../anisotropic_minimum_dissipation.jl | 6 +++--- .../convective_adjustment_vertical_diffusivity.jl | 6 +++--- .../isopycnal_skew_symmetric_diffusivity.jl | 2 -- .../leith_enstrophy_diffusivity.jl | 4 ++-- .../smagorinsky_lilly.jl | 4 ++-- test/test_time_stepping.jl | 9 ++++++--- 8 files changed, 26 insertions(+), 24 deletions(-) diff --git a/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl b/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl index 6b94f5756d..5dee55b824 100644 --- a/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl +++ b/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl @@ -18,14 +18,15 @@ import Oceananigans: fields, prognostic_fields function PressureSolver(arch::MultiArch, grid::RegRectilinearGrid) global_grid = reconstruct_global_grid(grid) if arch.ranks[1] == 1 # we would have to allow different settings - return DistributedFFTBasedPoissonSolver(arch, global_grid, grid) + return DistributedFFTBasedPoissonSolver(global_grid, grid) else @warn "A Distributed NonhydrostaticModel is allowed only when the x-direction is not parallelized" return nothing end end -PressureSolver(arch, grid::RegRectilinearGrid) = FFTBasedPoissonSolver(arch, grid) -PressureSolver(arch, grid::HRegRectilinearGrid) = FourierTridiagonalPoissonSolver(arch, grid) + +PressureSolver(arch, grid::RegRectilinearGrid) = FFTBasedPoissonSolver(grid) +PressureSolver(arch, grid::HRegRectilinearGrid) = FourierTridiagonalPoissonSolver(grid) # *Evil grin* PressureSolver(arch, ibg::ImmersedBoundaryGrid) = PressureSolver(arch, ibg.grid) diff --git a/src/TurbulenceClosures/diffusivity_fields.jl b/src/TurbulenceClosures/diffusivity_fields.jl index 72f6747874..915ade0615 100644 --- a/src/TurbulenceClosures/diffusivity_fields.jl +++ b/src/TurbulenceClosures/diffusivity_fields.jl @@ -4,22 +4,22 @@ using Oceananigans.Fields: validate_field_tuple_grid ##### Forms for NonhydrostaticModel constructor ##### -DiffusivityFields(diffusivities::NamedTuple, arch, grid, tracer_names, bcs, closure) = +DiffusivityFields(diffusivities::NamedTuple, grid, tracer_names, bcs, closure) = validate_field_tuple_grid("diffusivities", diffusivities, grid) -DiffusivityFields(::Nothing, arch, grid, tracer_names, bcs, closure) = - DiffusivityFields(arch, grid, tracer_names, bcs, closure) +DiffusivityFields(::Nothing, grid, tracer_names, bcs, closure) = + DiffusivityFields(grid, tracer_names, bcs, closure) ##### ##### Closures without precomputed diffusivities ##### -DiffusivityFields(arch, grid, tracer_names, bcs, closure) = nothing +DiffusivityFields(grid, tracer_names, bcs, closure) = nothing ##### ##### Closure tuples ##### -DiffusivityFields(arch, grid, tracer_names, bcs, closure_tuple::Tuple) = - Tuple(DiffusivityFields(arch, grid, tracer_names, bcs, closure) for closure in closure_tuple) +DiffusivityFields(grid, tracer_names, bcs, closure_tuple::Tuple) = + Tuple(DiffusivityFields(grid, tracer_names, bcs, closure) for closure in closure_tuple) diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/anisotropic_minimum_dissipation.jl b/src/TurbulenceClosures/turbulence_closure_implementations/anisotropic_minimum_dissipation.jl index b04783e247..1f5d6df82c 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/anisotropic_minimum_dissipation.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/anisotropic_minimum_dissipation.jl @@ -371,7 +371,7 @@ end ##### DiffusivityFields ##### -function DiffusivityFields(arch, grid, tracer_names, user_bcs, ::AMD) +function DiffusivityFields(grid, tracer_names, user_bcs, ::AMD) default_diffusivity_bcs = FieldBoundaryConditions(grid, (Center, Center, Center)) default_κₑ_bcs = NamedTuple(c => default_diffusivity_bcs for c in tracer_names) @@ -379,8 +379,8 @@ function DiffusivityFields(arch, grid, tracer_names, user_bcs, ::AMD) bcs = merge((; νₑ = default_diffusivity_bcs, κₑ = κₑ_bcs), user_bcs) - νₑ = CenterField(arch, grid, bcs.νₑ) - κₑ = NamedTuple(c => CenterField(arch, grid, bcs.κₑ[c]) for c in tracer_names) + νₑ = CenterField(grid, boundary_conditions=bcs.νₑ) + κₑ = NamedTuple(c => CenterField(grid, boundary_conditions=bcs.κₑ[c]) for c in tracer_names) return (; νₑ, κₑ) end diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/convective_adjustment_vertical_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/convective_adjustment_vertical_diffusivity.jl index a56ae07c6d..f418c5ce8b 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/convective_adjustment_vertical_diffusivity.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/convective_adjustment_vertical_diffusivity.jl @@ -77,11 +77,11 @@ function with_tracers(tracers, closure_array::CAVDArray) end # Note: computing diffusivities at cell centers for now. -function DiffusivityFields(arch, grid, tracer_names, bcs, closure::Union{CAVD, CAVDArray}) +function DiffusivityFields(grid, tracer_names, bcs, closure::Union{CAVD, CAVDArray}) ## If we can get away with only precomputing the "stability" of a cell: # data = new_data(Bool, arch, grid, (Center, Center, Center)) - κ = Field(Center, Center, Center, arch, grid) - ν = Field(Center, Center, Center, arch, grid) + κ = CenterField(grid) + ν = CenterField(grid) return (; κ, ν) end diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/isopycnal_skew_symmetric_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/isopycnal_skew_symmetric_diffusivity.jl index 936dfe0a75..06d92ec9d8 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/isopycnal_skew_symmetric_diffusivity.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/isopycnal_skew_symmetric_diffusivity.jl @@ -181,8 +181,6 @@ end calculate_diffusivities!(diffusivity_fields, closure::Union{ISSD, ISSDVector}, model) = nothing -DiffusivityFields(arch, grid, tracer_names, bcs, ::Union{ISSD, ISSDVector}) = nothing - ##### ##### Show ##### diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/leith_enstrophy_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/leith_enstrophy_diffusivity.jl index ed69b98a27..c766ce64b4 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/leith_enstrophy_diffusivity.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/leith_enstrophy_diffusivity.jl @@ -188,9 +188,9 @@ end "Return the filter width for a Leith Diffusivity on a regular rectilinear grid." @inline Δᶠ(i, j, k, grid::RectilinearGrid, ::TwoDimensionalLeith) = sqrt(Δxᶜᶜᵃ(i, j, k, grid) * Δyᶜᶜᵃ(i, j, k, grid)) -function DiffusivityFields(arch, grid, tracer_names, bcs, ::L2D) +function DiffusivityFields(grid, tracer_names, bcs, ::L2D) default_eddy_viscosity_bcs = (; νₑ = FieldBoundaryConditions(grid, (Center, Center, Center))) bcs = merge(default_eddy_viscosity_bcs, bcs) - νₑ = CenterField(arch, grid, bcs.νₑ) + νₑ = CenterField(grid, boundary_conditions=bcs.νₑ) return (; νₑ) end diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/smagorinsky_lilly.jl b/src/TurbulenceClosures/turbulence_closure_implementations/smagorinsky_lilly.jl index 768868846f..e2c5102897 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/smagorinsky_lilly.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/smagorinsky_lilly.jl @@ -202,11 +202,11 @@ Base.show(io::IO, closure::SmagorinskyLilly) = ##### For closures that only require an eddy viscosity νₑ field. ##### -function DiffusivityFields(arch, grid, tracer_names, bcs, closure::SmagorinskyLilly) +function DiffusivityFields(grid, tracer_names, bcs, closure::SmagorinskyLilly) default_eddy_viscosity_bcs = (; νₑ = FieldBoundaryConditions(grid, (Center, Center, Center))) bcs = merge(default_eddy_viscosity_bcs, bcs) - νₑ = CenterField(arch, grid, bcs.νₑ) + νₑ = CenterField(grid, boundary_conditions=bcs.νₑ) # Use AbstractOperations to write eddy diffusivities in terms of # eddy viscosity diff --git a/test/test_time_stepping.jl b/test/test_time_stepping.jl index ccacd254fa..b6c9eb2d1c 100644 --- a/test/test_time_stepping.jl +++ b/test/test_time_stepping.jl @@ -1,3 +1,5 @@ +include("dependencies_for_runtests.jl") + using Oceananigans.Grids: topological_tuple_length, total_size using Oceananigans.Fields: BackgroundField using Oceananigans.TimeSteppers: Clock @@ -98,13 +100,13 @@ end stepped. It just initializes a cube shaped hot bubble perturbation in the center of the 3D domain to induce a velocity field. """ -function incompressible_in_time(arch, grid, Nt, timestepper) +function incompressible_in_time(grid, Nt, timestepper) model = NonhydrostaticModel(grid=grid, timestepper=timestepper, buoyancy=SeawaterBuoyancy(), tracers=(:T, :S)) grid = model.grid u, v, w = model.velocities - div_U = CenterField(arch, grid) + div_U = CenterField(grid) # Just add a temperature perturbation so we get some velocity field. CUDA.@allowscalar interior(model.tracers.T)[8:24, 8:24, 8:24] .+= 0.01 @@ -113,6 +115,7 @@ function incompressible_in_time(arch, grid, Nt, timestepper) ab2_or_rk3_time_step!(model, 0.05, n) end + arch = architecture(grid) event = launch!(arch, grid, :xyz, divergence!, grid, u.data, v.data, w.data, div_U.data, dependencies=Event(device(arch))) wait(device(arch), event) @@ -343,7 +346,7 @@ timesteppers = (:QuasiAdamsBashforth2, :RungeKutta3) @info " Testing incompressibility [$FT, $(typeof(grid).name.wrapper)]..." for Nt in [1, 10, 100], timestepper in timesteppers - @test incompressible_in_time(arch, grid, Nt, timestepper) + @test incompressible_in_time(grid, Nt, timestepper) end end end From 79cdafb686dcb70c3f739a8ee89fa1bbf0935129 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 17 Dec 2021 21:47:23 -0700 Subject: [PATCH 009/140] Fixing up hydrostatic model tests --- src/Distributed/distributed_fields.jl | 4 ++-- .../HydrostaticFreeSurfaceModels.jl | 4 ++-- .../explicit_free_surface.jl | 4 ++-- .../hydrostatic_free_surface_model.jl | 4 ++-- .../implicit_free_surface.jl | 4 ++-- .../prescribed_hydrostatic_velocity_fields.jl | 2 +- .../single_column_model_mode.jl | 2 +- .../NonhydrostaticModels/nonhydrostatic_model.jl | 2 +- src/TimeSteppers/quasi_adams_bashforth_2.jl | 2 +- src/TimeSteppers/runge_kutta_3.jl | 14 +++++++------- .../CATKEVerticalDiffusivities.jl | 8 ++++---- test/test_hydrostatic_free_surface_models.jl | 4 ++-- 12 files changed, 27 insertions(+), 27 deletions(-) diff --git a/src/Distributed/distributed_fields.jl b/src/Distributed/distributed_fields.jl index ef20de9f27..bb093fae81 100644 --- a/src/Distributed/distributed_fields.jl +++ b/src/Distributed/distributed_fields.jl @@ -7,7 +7,7 @@ function Field(X, Y, Z, arch::AbstractMultiArchitecture, grid::AbstractGrid, bcs = FieldBoundaryConditions(grid, (X, Y, Z)), data = new_data(eltype(grid), arch, grid, (X, Y, Z))) - communicative_bcs = inject_halo_communication_boundary_conditions(bcs, arch.local_rank, arch.connectivity) + boudnary_conditions = inject_halo_communication_boundary_conditions(bcs, arch.local_rank, arch.connectivity) - return Field(X, Y, Z, child_architecture(arch), grid, communicative_bcs, data) + return Field{X, Y, Z}(grid; boundary_conditions, data) end diff --git a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl index 2dc1fef31b..d64905b0ba 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl @@ -19,8 +19,8 @@ fill_horizontal_velocity_halos!(args...) = nothing ##### HydrostaticFreeSurfaceModel definition ##### -FreeSurfaceDisplacementField(velocities, free_surface, arch, grid) = ReducedField(Center, Center, Nothing, arch, grid; dims=3) -FreeSurfaceDisplacementField(velocities, ::Nothing, arch, grid) = nothing +FreeSurfaceDisplacementField(velocities, free_surface, grid) = Field{Center, Center, Nothing}(grid) +FreeSurfaceDisplacementField(velocities, ::Nothing, grid) = nothing include("compute_w_from_continuity.jl") diff --git a/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl index c4c448ceae..be36d7115d 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl @@ -20,8 +20,8 @@ Adapt.adapt_structure(to, free_surface::ExplicitFreeSurface) = ##### Interface to HydrostaticFreeSurfaceModel ##### -function FreeSurface(free_surface::ExplicitFreeSurface{Nothing}, velocities, arch, grid) - η = FreeSurfaceDisplacementField(velocities, free_surface, arch, grid) +function FreeSurface(free_surface::ExplicitFreeSurface{Nothing}, velocities, grid) + η = FreeSurfaceDisplacementField(velocities, free_surface, grid) g = convert(eltype(grid), free_surface.gravitational_acceleration) return ExplicitFreeSurface(η, g) end diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl index eaca695286..e601047084 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl @@ -24,7 +24,7 @@ validate_tracer_advection(invalid_tracer_advection, grid) = error("$invalid_trac validate_tracer_advection(tracer_advection_tuple::NamedTuple, grid) = CenteredSecondOrder(), tracer_advection_tuple validate_tracer_advection(tracer_advection::AbstractAdvectionScheme, grid) = tracer_advection, NamedTuple() -PressureField(arch, grid) = (; pHY′ = CenterField(arch, grid)) +PressureField(grid) = (; pHY′ = CenterField(grid)) mutable struct HydrostaticFreeSurfaceModel{TS, E, A<:AbstractArchitecture, S, G, T, V, B, R, F, P, U, C, Φ, K, AF} <: AbstractModel{TS} @@ -152,7 +152,7 @@ function HydrostaticFreeSurfaceModel(; grid, validate_velocity_boundary_conditions(velocities) - free_surface = FreeSurface(free_surface, velocities, arch, grid) + free_surface = FreeSurface(free_surface, velocities, grid) # Instantiate timestepper if not already instantiated implicit_solver = implicit_diffusion_solver(time_discretization(closure), arch, grid) diff --git a/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl index 1f8e3ecb29..7d5f69acbe 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl @@ -50,8 +50,8 @@ Adapt.adapt_structure(to, free_surface::ImplicitFreeSurface) = nothing, nothing, nothing, nothing) # Internal function for HydrostaticFreeSurfaceModel -function FreeSurface(free_surface::ImplicitFreeSurface{Nothing}, velocities, arch, grid) - η = FreeSurfaceDisplacementField(velocities, free_surface, arch, grid) +function FreeSurface(free_surface::ImplicitFreeSurface{Nothing}, velocities, grid) + η = FreeSurfaceDisplacementField(velocities, free_surface, grid) g = convert(eltype(grid), free_surface.gravitational_acceleration) # Initialize barotropic volume fluxes diff --git a/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl b/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl index 57912acc28..7c8d5f3358 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl @@ -62,7 +62,7 @@ function HydrostaticFreeSurfaceVelocityFields(velocities::PrescribedVelocityFiel end function HydrostaticFreeSurfaceTendencyFields(::PrescribedVelocityFields, free_surface, grid, tracer_names) - tracers = TracerFields(tracer_names, arch, grid) + tracers = TracerFields(tracer_names, grid) return merge((u = nothing, v = nothing, η = nothing), tracers) end diff --git a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl index f88df214ba..71d063155c 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl @@ -28,7 +28,7 @@ const SingleColumnGrid = AbstractGrid{<:AbstractFloat, <:Flat, <:Flat, <:Bounded ##### Model constructor utils ##### -PressureField(arch, ::SingleColumnGrid) = (pHY′ = nothing,) +PressureField(::SingleColumnGrid) = (; pHY′ = nothing) FreeSurface(free_surface::ExplicitFreeSurface{Nothing}, velocities, arch, ::SingleColumnGrid) = nothing FreeSurface(free_surface::ImplicitFreeSurface{Nothing}, velocities, arch, ::SingleColumnGrid) = nothing diff --git a/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl b/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl index 4fec96acca..b211fa414a 100644 --- a/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl +++ b/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl @@ -165,7 +165,7 @@ function NonhydrostaticModel(; grid, # Instantiate timestepper if not already instantiated implicit_solver = implicit_diffusion_solver(time_discretization(closure), arch, grid) - timestepper = TimeStepper(timestepper, arch, grid, tracernames(tracers), implicit_solver=implicit_solver) + timestepper = TimeStepper(timestepper, grid, tracernames(tracers), implicit_solver=implicit_solver) # Regularize forcing for model tracer and velocity fields. model_fields = merge(velocities, tracers) diff --git a/src/TimeSteppers/quasi_adams_bashforth_2.jl b/src/TimeSteppers/quasi_adams_bashforth_2.jl index d39154a0d1..92ff5d7aa1 100644 --- a/src/TimeSteppers/quasi_adams_bashforth_2.jl +++ b/src/TimeSteppers/quasi_adams_bashforth_2.jl @@ -20,7 +20,7 @@ Return an QuasiAdamsBashforth2TimeStepper object with tendency fields on `arch` `grid` with AB2 parameter `χ`. The tendency fields can be specified via optional kwargs. """ -function QuasiAdamsBashforth2TimeStepper(arch, grid, tracers, +function QuasiAdamsBashforth2TimeStepper(grid, tracers, χ = 0.1; implicit_solver::IT = nothing, Gⁿ = TendencyFields(grid, tracers), diff --git a/src/TimeSteppers/runge_kutta_3.jl b/src/TimeSteppers/runge_kutta_3.jl index 4b48663aa0..8f3b8eb072 100644 --- a/src/TimeSteppers/runge_kutta_3.jl +++ b/src/TimeSteppers/runge_kutta_3.jl @@ -18,14 +18,14 @@ struct RungeKutta3TimeStepper{FT, TG, TI} <: AbstractTimeStepper end """ - RungeKutta3TimeStepper(arch, grid, tracers, + RungeKutta3TimeStepper(grid, tracers, Gⁿ = TendencyFields(grid, tracers), G⁻ = TendencyFields(grid, tracers)) -Return an `RungeKutta3TimeStepper` object with tendency fields on `arch` and -`grid`. The tendency fields can be specified via optional kwargs. +Return an `RungeKutta3TimeStepper` object with tendency fields on `grid`. +The tendency fields can be specified via optional kwargs. """ -function RungeKutta3TimeStepper(arch, grid, tracers; +function RungeKutta3TimeStepper(grid, tracers; implicit_solver::TI = nothing, Gⁿ::TG = TendencyFields(grid, tracers), G⁻ = TendencyFields(grid, tracers)) where {TI, TG} @@ -142,9 +142,9 @@ function rk3_substep!(model, Δt, γⁿ, ζⁿ) workgroup, worksize = work_layout(model.grid, :xyz) - barrier = Event(device(model.architecture)) + barrier = Event(device(architecture(model))) - substep_field_kernel! = rk3_substep_field!(device(model.architecture), workgroup, worksize) + substep_field_kernel! = rk3_substep_field!(device(architecture(model)), workgroup, worksize) model_fields = prognostic_fields(model) @@ -173,7 +173,7 @@ function rk3_substep!(model, Δt, γⁿ, ζⁿ) push!(events, field_event) end - wait(device(model.architecture), MultiEvent(Tuple(events))) + wait(device(architecture(model)), MultiEvent(Tuple(events))) return nothing end diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/CATKEVerticalDiffusivities/CATKEVerticalDiffusivities.jl b/src/TurbulenceClosures/turbulence_closure_implementations/CATKEVerticalDiffusivities/CATKEVerticalDiffusivities.jl index ed9d74e349..c69f58a85c 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/CATKEVerticalDiffusivities/CATKEVerticalDiffusivities.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/CATKEVerticalDiffusivities/CATKEVerticalDiffusivities.jl @@ -110,7 +110,7 @@ end ##### Diffusivities and diffusivity fields utilities ##### -function DiffusivityFields(arch, grid, tracer_names, bcs, closure::Union{CATKEVD, CATKEVDArray}) +function DiffusivityFields(grid, tracer_names, bcs, closure::Union{CATKEVD, CATKEVDArray}) default_diffusivity_bcs = (Kᵘ = FieldBoundaryConditions(grid, (Center, Center, Center)), Kᶜ = FieldBoundaryConditions(grid, (Center, Center, Center)), @@ -118,9 +118,9 @@ function DiffusivityFields(arch, grid, tracer_names, bcs, closure::Union{CATKEVD bcs = merge(default_diffusivity_bcs, bcs) - Kᵘ = CenterField(arch, grid, bcs.Kᵘ) - Kᶜ = CenterField(arch, grid, bcs.Kᶜ) - Kᵉ = CenterField(arch, grid, bcs.Kᵉ) + Kᵘ = CenterField(grid, boundary_conditions=bcs.Kᵘ) + Kᶜ = CenterField(grid, boundary_conditions=bcs.Kᶜ) + Kᵉ = CenterField(grid, boundary_conditions=bcs.Kᵉ) return (; Kᵘ, Kᶜ, Kᵉ) end diff --git a/test/test_hydrostatic_free_surface_models.jl b/test/test_hydrostatic_free_surface_models.jl index ef589dca9e..9310a25abe 100644 --- a/test/test_hydrostatic_free_surface_models.jl +++ b/test/test_hydrostatic_free_surface_models.jl @@ -1,9 +1,9 @@ -using Oceananigans: CPU, GPU +include("dependencies_for_runtests.jl") + using Oceananigans.Models.HydrostaticFreeSurfaceModels: VectorInvariant, PrescribedVelocityFields, PrescribedField, ExplicitFreeSurface using Oceananigans.Models.HydrostaticFreeSurfaceModels: ExplicitFreeSurface, ImplicitFreeSurface using Oceananigans.Coriolis: VectorInvariantEnergyConserving, VectorInvariantEnstrophyConserving using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization, ExplicitTimeDiscretization, CATKEVerticalDiffusivity -using Oceananigans.Grids: Periodic, Bounded function time_step_hydrostatic_model_works(grid; coriolis = nothing, From 13be16ac3ce39352a9e2aa875a7019b46bd85f66 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 29 Dec 2021 07:05:18 -0700 Subject: [PATCH 010/140] Fixes tests and implmements field Reductions --- src/AbstractOperations/AbstractOperations.jl | 1 + src/AbstractOperations/computed_field.jl | 83 ++++ .../kernel_function_operation.jl | 26 +- src/BuoyancyModels/buoyancy_field.jl | 147 +------ src/Distributed/distributed_fields.jl | 14 +- src/Distributed/distributed_grids.jl | 241 ++++------- src/Fields/Fields.jl | 6 +- src/Fields/computed_field.jl | 2 +- src/Fields/field.jl | 111 ++--- src/Fields/field_reductions.jl | 58 +++ src/Fields/metric_reductions.jl | 36 ++ src/Fields/pressure_field.jl | 5 - src/Fields/show_fields.jl | 16 - src/Grids/latitude_longitude_grid.jl | 401 ++++++++++-------- src/Grids/rectilinear_grid.jl | 197 +++++---- ...ute_vertically_integrated_lateral_areas.jl | 4 +- .../fft_based_implicit_free_surface_solver.jl | 21 +- .../implicit_free_surface.jl | 10 +- .../pcg_implicit_free_surface_solver.jl | 16 +- .../vertical_vorticity_field.jl | 19 +- .../NonhydrostaticModels.jl | 1 + .../nonhydrostatic_model.jl | 4 + .../NonhydrostaticModels/pressure_field.jl | 3 + src/Oceananigans.jl | 3 +- ...reconditioned_conjugate_gradient_solver.jl | 4 + src/TimeSteppers/runge_kutta_3.jl | 3 +- .../vertically_implicit_diffusion_solver.jl | 4 +- test/dependencies_for_runtests.jl | 2 + test/test_dynamics.jl | 19 +- test/test_hydrostatic_free_surface_models.jl | 5 +- test/test_time_stepping.jl | 6 +- test/utils_for_runtests.jl | 20 - 32 files changed, 727 insertions(+), 761 deletions(-) create mode 100644 src/AbstractOperations/computed_field.jl create mode 100644 src/Fields/field_reductions.jl create mode 100644 src/Fields/metric_reductions.jl delete mode 100644 src/Fields/pressure_field.jl create mode 100644 src/Models/NonhydrostaticModels/pressure_field.jl diff --git a/src/AbstractOperations/AbstractOperations.jl b/src/AbstractOperations/AbstractOperations.jl index 0caab7b7b5..a8eb449607 100644 --- a/src/AbstractOperations/AbstractOperations.jl +++ b/src/AbstractOperations/AbstractOperations.jl @@ -55,6 +55,7 @@ include("binary_operations.jl") include("multiary_operations.jl") include("derivatives.jl") include("kernel_function_operation.jl") +include("computed_field.jl") include("at.jl") include("broadcasting_abstract_operations.jl") include("show_abstract_operations.jl") diff --git a/src/AbstractOperations/computed_field.jl b/src/AbstractOperations/computed_field.jl new file mode 100644 index 0000000000..179fe1ed40 --- /dev/null +++ b/src/AbstractOperations/computed_field.jl @@ -0,0 +1,83 @@ +##### +##### Fields computed from abstract operations +##### + +using KernelAbstractions: @kernel, @index + +import Oceananigans: short_show +import Oceananigans.Fields: Field, compute! + +const ComputedField = Field{<:Any, <:Any, <:Any, <:AbstractOperation} + +""" + Field(operand::AbstractOperation; kwargs...) + +Return `f::Field` where `f.data` is computed from `f.operand` by +calling compute!(f). + +Keyword arguments +================= + +data (AbstractArray): An offset Array or CuArray for storing the result of a computation. + Must have `total_size(location(operand), grid)`. + +boundary_conditions (FieldBoundaryConditions): Boundary conditions for `f`. + +recompute_safely (Bool): whether or not to _always_ "recompute" `f` if `f` is + nested within another computation via an `AbstractOperation`. + If `data` is not provided then `recompute_safely=false` and + recomputation is _avoided_. If `data` is provided, then + `recompute_safely=true` by default. +""" +function Field(operand::AbstractOperation; + data = nothing, + boundary_conditions = FieldBoundaryConditions(op.grid, location(op)), + recompute_safely = true) + + if isnothing(data) + data = new_data(op.grid, location(op)) + recompute_safely = false + end + + status = recompute_safely ? nothing : FieldStatus(0.0) + + return Field(location(op), op.grid, data, boundary_conditions, operand, status) +end + +""" + compute!(comp::ComputedField) + +Compute `comp.operand` and store the result in `comp.data`. +""" +function compute!(comp::ComputedField, time=nothing) + # First compute `dependencies`: + compute_at!(comp.operand, time) + + workgroup, worksize = + work_layout(comp.grid, :xyz, include_right_boundaries = true, location = location(comp)) + + arch = architecture(comp) + compute_kernel! = _compute!(device(arch), workgroup, worksize) + event = compute_kernel!(comp.data, comp.operand; dependencies = device_event(arch)) + wait(device(arch), event) + + fill_halo_regions!(comp, arch) + + return comp +end + +"""Compute an `operand` and store in `data`.""" +@kernel function _compute!(data, operand) + i, j, k = @index(Global, NTuple) + @inbounds data[i, j, k] = operand[i, j, k] +end + +short_show(field::ComputedField) = string("ComputedField located at ", show_location(field), " of ", short_show(field.operand)) + +Base.show(io::IO, field::ComputedField) = + print(io, "$(short_show(field))\n", + "├── data: $(typeof(field.data)), size: $(size(field))\n", + "├── grid: $(short_show(field.grid))\n", + "├── operand: $(short_show(field.operand))\n", + "└── status: $(show_status(field.status))") + diff --git a/src/AbstractOperations/kernel_function_operation.jl b/src/AbstractOperations/kernel_function_operation.jl index a844c8a65c..fbd3f89465 100644 --- a/src/AbstractOperations/kernel_function_operation.jl +++ b/src/AbstractOperations/kernel_function_operation.jl @@ -6,16 +6,18 @@ struct KernelFunctionOperation{LX, LY, LZ, P, A, G, T, K, D} <: AbstractOperatio grid :: G function KernelFunctionOperation{LX, LY, LZ}(kernel_function::K, computed_dependencies::D, - parameters::P, architecture::A, grid::G) where {LX, LY, LZ, K, G, A, D, P} + parameters::P, grid::G) where {LX, LY, LZ, K, G, D, P} + arch = architecture(grid) T = eltype(grid) - return new{LX, LY, LZ, P, A, G, T, K, D}(kernel_function, computed_dependencies, parameters, architecture, grid) + A = typeof(arch) + return new{LX, LY, LZ, P, A, G, T, K, D}(kernel_function, computed_dependencies, parameters, grid) end end """ - KernelFunctionOperation{LX, LY, LZ}(kernel_function, grid; architecture=nothing, + KernelFunctionOperation{LX, LY, LZ}(kernel_function, grid; computed_dependencies=(), parameters=nothing) Constructs a `KernelFunctionOperation` at location `(LX, LY, LZ)` on `grid` an with @@ -42,7 +44,7 @@ Construct a kernel function operation that returns random numbers: ```julia random_kernel_function(i, j, k, grid) = rand() # use CUDA.rand on the GPU -kernel_op = KernelFunctionOperation{Center, Center, Center}(random_kernel_function, grid; architecture=CPU()) +kernel_op = KernelFunctionOperation{Center, Center, Center}(random_kernel_function, grid) ``` Construct a kernel function operation using the vertical vorticity operator @@ -57,19 +59,8 @@ u, v, w = model.velocities ζ_op = KernelFunctionOperation{Face, Face, Center}(ζ₃ᶠᶠᵃ, grid, computed_dependencies=(u, v)) ``` """ -function KernelFunctionOperation{LX, LY, LZ}(kernel_function, - grid; - architecture = nothing, - computed_dependencies = (), - parameters = nothing) where {LX, LY, LZ} - - arch = isnothing(architecture) ? - Oceananigans.Architectures.architecture(computed_dependencies...) : - architecture - - return KernelFunctionOperation{LX, LY, LZ}(kernel_function, computed_dependencies, parameters, arch, grid) -end - +KernelFunctionOperation{LX, LY, LZ}(kernel_function, grid; computed_dependencies = (), parameters = nothing) where {LX, LY, LZ} = + KernelFunctionOperation{LX, LY, LZ}(kernel_function, computed_dependencies, parameters, grid) @inline Base.getindex(κ::KernelFunctionOperation, i, j, k) = κ.kernel_function(i, j, k, κ.grid, κ.computed_dependencies..., κ.parameters) @inline Base.getindex(κ::KernelFunctionOperation{LX, LY, LZ, <:Nothing}, i, j, k) where {LX, LY, LZ} = κ.kernel_function(i, j, k, κ.grid, κ.computed_dependencies...) @@ -84,3 +75,4 @@ Adapt.adapt_structure(to, κ::KernelFunctionOperation{LX, LY, LZ}) where {LX, LY Adapt.adapt(to, κ.parameters), nothing, Adapt.adapt(to, κ.grid)) + diff --git a/src/BuoyancyModels/buoyancy_field.jl b/src/BuoyancyModels/buoyancy_field.jl index 12f5c43b4e..b640d60e62 100644 --- a/src/BuoyancyModels/buoyancy_field.jl +++ b/src/BuoyancyModels/buoyancy_field.jl @@ -1,152 +1,19 @@ using Adapt using KernelAbstractions -# TODO: define buoyancy_operation(model) and be done with it. - -#= function buoyancy_operation(model) - buoyancy = model.buoyancy + buoyancy_model = model.buoyancy.model tracers = model.tracers - return KernelFunctionOperation{Center, Center, Center}(...) + return buoyancy_operation(buoyancy_model, model.grid, tracers) end +buoyancy_operation(buoyancy_model, grid, tracers) = + KernelFunctionOperation{Center, Center, Center}(buoyancy_perturbation, model.grid, computed=dependencies=(buoyancy_model, tracers)) + +buoyancy_operation(::Nothing, grid, tracers) = nothing + function BuoyancyField(model) op = buoyancy_operation(model) return Field(op) end -=# - -using Oceananigans.Fields: AbstractField, FieldStatus, validate_field_data, conditional_compute! -using Oceananigans.Fields: architecture, tracernames -using Oceananigans.Architectures: device -using Oceananigans.Utils: work_layout -using Oceananigans.Grids: new_data - -import Oceananigans.Fields: compute!, compute_at! - -import Oceananigans: short_show - -struct BuoyancyField{B, S, A, D, G, T, C} <: AbstractField{Center, Center, Center, A, G, T, 3} - data :: D - architecture :: A - grid :: G - buoyancy :: B - tracers :: C - status :: S - - @doc """ - BuoyancyField(data, grid, buoyancy, tracers) - - Returns a `BuoyancyField` with `data` on `grid` corresponding to - `buoyancy` computed from `tracers`. - """ - function BuoyancyField(data::D, arch::A, grid::G, buoyancy::B, tracers::C, - recompute_safely::Bool) where {D, A, G, B, C} - - validate_field_data(Center, Center, Center, data, grid) - - status = recompute_safely ? nothing : FieldStatus(zero(eltype(grid))) - - S = typeof(status) - T = eltype(grid) - - return new{B, S, A, D, G, T, C}(data, arch, grid, buoyancy, tracers, status) - end - - function BuoyancyField(data::D, arch::A, grid::G, buoyancy::B, tracers::C, status::S) where {D, A, G, B, C, S} - validate_field_data(Center, Center, Center, data, grid) - T = eltype(grid) - return new{B, S, A, D, G, T, C}(data, grid, buoyancy, tracers, status) - end -end - -""" - BuoyancyField(model; data=nothing, recompute_safely=true) - -Returns a `BuoyancyField` corresponding to `model.buoyancy`. -Calling `compute!(b::BuoyancyField)` computes the current buoyancy field -associated with `model` and stores the result in `b.data`. -""" -BuoyancyField(model; data=nothing, recompute_safely=true) = - _buoyancy_field(model.buoyancy, model.tracers, model.architecture, model.grid, data, recompute_safely) - -# Convenience for buoyancy=nothing -_buoyancy_field(::Nothing, args...; kwargs...) = nothing - -##### -##### BuoyancyTracer -##### - -_buoyancy_field(buoyancy::BuoyancyTracerModel, tracers, arch, grid, args...) = - BuoyancyField(tracers.b.data, arch, grid, buoyancy, tracers, true) - -compute!(::BuoyancyField{<:BuoyancyTracerModel}, time=nothing) = nothing - -##### -##### Other buoyancy types -##### - -function _buoyancy_field(buoyancy::Buoyancy, tracers, arch, grid, - data, recompute_safely) - - if isnothing(data) - data = new_data(arch, grid, (Center, Center, Center)) - recompute_safely = false - end - - return BuoyancyField(data, arch, grid, buoyancy, tracers, recompute_safely) -end - -""" - compute!(buoyancy_field::BuoyancyField, time=nothing) - -Compute the current `buoyancy_field` associated with `buoyancy_field.tracers` and store -the result in `buoyancy_field.data`. -""" -function compute!(buoyancy_field::BuoyancyField, time=nothing) - - data = buoyancy_field.data - grid = buoyancy_field.grid - tracers = buoyancy_field.tracers - buoyancy = buoyancy_field.buoyancy - arch = architecture(buoyancy_field) - - workgroup, worksize = work_layout(grid, :xyz) - - compute_kernel! = compute_buoyancy!(device(arch), workgroup, worksize) - - event = compute_kernel!(data, grid, buoyancy, tracers; dependencies=Event(device(arch))) - - wait(device(arch), event) - - return nothing -end - -compute_at!(b::BuoyancyField{B, <:FieldStatus}, time) where B = - conditional_compute!(b, time) - -"""Compute an `operation` and store in `data`.""" -@kernel function compute_buoyancy!(data, grid, buoyancy, tracers) - i, j, k = @index(Global, NTuple) - @inbounds data[i, j, k] = buoyancy_perturbation(i, j, k, grid, buoyancy.model, tracers) -end - -##### -##### Adapt -##### - -Adapt.adapt_structure(to, buoyancy_field::BuoyancyField) = Adapt.adapt(to, buoyancy_field.data) - -##### -##### Show -##### - -short_show(field::BuoyancyField) = string("BuoyancyField for ", typeof(field.buoyancy)) -show(io::IO, field::BuoyancyField) = - print(io, "$(short_show(field))\n", - "├── data: $(typeof(field.data)), size: $(size(field.data))\n", - "├── grid: $(short_show(field.grid))", '\n', - "├── buoyancy: $(typeof(field.buoyancy))", '\n', - "├── tracers: $(tracernames(field.tracers))", '\n', - "└── status: ", show_status(field.status), '\n') diff --git a/src/Distributed/distributed_fields.jl b/src/Distributed/distributed_fields.jl index bb093fae81..1ffd532e1a 100644 --- a/src/Distributed/distributed_fields.jl +++ b/src/Distributed/distributed_fields.jl @@ -1,13 +1,7 @@ -using MPI -using Oceananigans.Fields: location import Oceananigans.Fields: Field -import Oceananigans.Grids: AbstractGrid, size -function Field(X, Y, Z, arch::AbstractMultiArchitecture, grid::AbstractGrid, - bcs = FieldBoundaryConditions(grid, (X, Y, Z)), - data = new_data(eltype(grid), arch, grid, (X, Y, Z))) - - boudnary_conditions = inject_halo_communication_boundary_conditions(bcs, arch.local_rank, arch.connectivity) - - return Field{X, Y, Z}(grid; boundary_conditions, data) +function Field((LX, LY, LZ)::Tuple, grid::DistributedGrid, data, bcs, op, status) + arch = architecture(grid) + boundary_conditions = inject_halo_communication_boundary_conditions(bcs, arch.local_rank, arch.connectivity) + return Field{LX, LY, LZ}(grid, data, bcs, op, status) end diff --git a/src/Distributed/distributed_grids.jl b/src/Distributed/distributed_grids.jl index 6f49b09bf4..97947baa28 100644 --- a/src/Distributed/distributed_grids.jl +++ b/src/Distributed/distributed_grids.jl @@ -1,27 +1,41 @@ using MPI using Oceananigans.Grids: topology, size, halo_size, architecture, pop_flat_elements -using Oceananigans.Grids: validate_halo, validate_rectilinear_domain, validate_size, validate_topology -using Oceananigans.Grids: generate_coordinate, cpu_face_constructor_x, cpu_face_constructor_y, cpu_face_constructor_z +using Oceananigans.Grids: validate_rectilinear_domain, validate_topology +using Oceananigans.Grids: validate_rectilinear_grid_args, validate_lat_lon_grid_args +using Oceananigans.Grids: generate_coordinate, with_precomputed_metrics +using Oceananigans.Grids: cpu_face_constructor_x, cpu_face_constructor_y, cpu_face_constructor_z import Oceananigans.Grids: RectilinearGrid, LatitudeLongitudeGrid, with_halo +const DistributedGrid{FT, TX, TY, TZ} = AbstractGrid{FT, TX, TY, TZ, <:MultiArch} + @inline get_local_coords(c::Tuple , nc, R, index) = (c[1] + (index-1) * (c[2] - c[1]) / R, c[1] + index * (c[2] - c[1]) / R) @inline get_local_coords(c::AbstractVector, nc, R, index) = c[1 + (index-1) * nc : 1 + nc * index] @inline get_global_coords(c::Tuple , nc, R, index, arch) = (c[2] - index * (c[2] - c[1]), c[2] - (index - R) * (c[2] - c[1])) -function get_global_coords(c::AbstractVector, nc, R, index, arch) - cG = zeros(eltype(c), nc*R+1) - cG[1 + (index-1) * nc : nc * index] .= c[1:end-1] - - if index == R - cG[end] = c[end] - end - MPI.Allreduce!(cG, +, arch.communicator) +""" + get_global_coords(c_local::AbstractVector, nc, R, index, arch) + +Build a linear global coordinate vector given a local coordinate vector `c_local` +a local number of elements `nc`, number of ranks `R`, rank `index`, +and `arch`itecture. +""" +function get_global_coords(c_local::AbstractVector, nc, R, index, arch) + c_global = zeros(eltype(c_local), nc*R+1) + c_global[1 + (index-1) * nc : nc * index] .= c[1:end-1] + index == R && (c_global[end] = c[end]) + + MPI.Allreduce!(c_global, +, arch.communicator) - return cG + return c_global end +""" + RectilinearGrid(arch::MultiArch, FT=Float64; kw...) + +Return the rank-local portion of `RectilinearGrid` on `arch`itecture. +""" function RectilinearGrid(arch::MultiArch, FT = Float64; size, x = nothing, @@ -31,17 +45,13 @@ function RectilinearGrid(arch::MultiArch, FT = Float64; extent = nothing, topology = (Periodic, Periodic, Bounded)) - TX, TY, TZ = validate_topology(topology) - size = validate_size(TX, TY, TZ, size) - halo = validate_halo(TX, TY, TZ, halo) - - # Validate the rectilinear domain - x, y, z = validate_rectilinear_domain(TX, TY, TZ, FT, extent, x, y, z) + TX, TY, TZ, size, halo, x, y, z = + validate_rectilinear_grid_args(topology, size, halo, FT, extent, x, y, z) Nx, Ny, Nz = size - hx, hy, hz = halo + Hx, Hy, Hz = halo - i, j, k = arch.local_index + ri, rj, rk = arch.local_index Rx, Ry, Rz = arch.ranks # Make sure we can put an integer number of grid points in each rank. @@ -50,64 +60,40 @@ function RectilinearGrid(arch::MultiArch, FT = Float64; @assert isinteger(Ny / Ry) @assert isinteger(Nz / Rz) + # Local sizes are denoted with lowercase `n` nx, ny, nz = local_size = Nx÷Rx, Ny÷Ry, Nz÷Rz - xl = get_local_coords(x, nx, Rx, i) - yl = get_local_coords(y, ny, Ry, j) - zl = get_local_coords(z, nz, Rz, k) + xl = get_local_coords(x, nx, Rx, ri) + yl = get_local_coords(y, ny, Ry, rj) + zl = get_local_coords(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)) - - FX = typeof(Δxᶠᵃᵃ) - FY = typeof(Δyᵃᶠᵃ) - FZ = typeof(Δzᵃᵃᶠ) - VX = typeof(xᶠᵃᵃ) - VY = typeof(yᵃᶠᵃ) - VZ = typeof(zᵃᵃᶠ) + 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 = MultiArch(child_architecture(arch), topology = topology, ranks = arch.ranks, communicator = arch.communicator) - Arch = typeof(arch) - - return RectilinearGrid{FT, TX, TY, TZ, FX, FY, FZ, VX, VY, VZ, Arch}(architecture, - nx, ny, nz, hx, hy, hz, Lx, Ly, Lz, Δxᶠᵃᵃ, Δxᶜᵃᵃ, xᶠᵃᵃ, xᶜᵃᵃ, Δyᵃᶜᵃ, Δyᵃᶠᵃ, yᵃᶠᵃ, yᵃᶜᵃ, Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ) + return RectilinearGrid{TX, TY, TZ}(architecture, + nx, ny, nz, + Hx, Hy, Hz, + Lx, Ly, Lz, + Δxᶠᵃᵃ, Δxᶜᵃᵃ, xᶠᵃᵃ, xᶜᵃᵃ, + Δyᵃᶜᵃ, Δyᵃᶠᵃ, yᵃᶠᵃ, yᵃᶜᵃ, + Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ) end function LatitudeLongitudeGrid(arch::MultiArch, - FT=Float64; - precompute_metrics=false, - size, - latitude, - longitude, - z, - radius=R_Earth, - halo=(1, 1, 1)) - - λ₁, λ₂ = get_domain_extent(longitude, size[1]) - @assert λ₁ < λ₂ && λ₂ - λ₁ ≤ 360 - - φ₁, φ₂ = get_domain_extent(latitude, size[2]) - @assert -90 <= φ₁ < φ₂ <= 90 - - (φ₁ == -90 || φ₂ == 90) && - @warn "Are you sure you want to use a latitude-longitude grid with a grid point at the pole?" - - Lλ = λ₂ - λ₁ - Lφ = φ₂ - φ₁ - - TX = Lλ == 360 ? Periodic : Bounded - TY = Bounded - TZ = Bounded - topo = (TX, TY, TZ) - - Nλ, Nφ, Nz = N = validate_size(TX, TY, TZ, size) - hλ, hφ, hz = H = validate_halo(TX, TY, TZ, halo) - - # 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) + FT = Float64; + precompute_metrics = false, + size, + latitude, + longitude, + z, + radius = R_Earth, + halo = (1, 1, 1)) + + Nλ, Nφ, Nz, Hλ, Hφ, Hz, latitude, longitude, topo = + validate_lat_lon_grid_args(latitude, longitude, size, halo) i, j, k = arch.local_index Rx, Ry, Rz = arch.ranks @@ -124,52 +110,28 @@ function LatitudeLongitudeGrid(arch::MultiArch, φl = get_local_coords(latitude , ny, Ry, j) zl = get_local_coords(z, nz, Rz, k) - Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, topo[1], nλ, hλ, λl, arch.child_architecture) - Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, topo[2], nφ, hφ, φl, arch.child_architecture) - Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, topo[3], nz, hz, zl, arch.child_architecture) - - FX = typeof(Δλᶠᵃᵃ) - FY = typeof(Δφᵃᶠᵃ) - FZ = typeof(Δzᵃᵃᶠ) - VX = typeof(λᶠᵃᵃ) - VY = typeof(φᵃᶠᵃ) - VZ = typeof(zᵃᵃᶠ) + # 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, topo[1], nλ, Hλ, λl, arch.child_architecture) + Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, topo[2], nφ, Hφ, φl, arch.child_architecture) + Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, topo[3], nz, Hz, zl, arch.child_architecture) architecture = MultiArch(child_architecture(arch), grid = grid, ranks = arch.ranks, communicator = arch.communicator) - Arch = typeof(architecture) + preliminary_grid = LatitudeLongitudeGrid(architecture, + Nλ, Nφ, Nz, + Hλ, Hφ, Hz, + Lλ, Lφ, Lz, + Δλᶠᵃᵃ, Δλᶜᵃᵃ, λᶠᵃᵃ, λᶜᵃᵃ, + Δφᵃᶠᵃ, Δφᵃᶜᵃ, φᵃᶠᵃ, φᵃᶜᵃ, + Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ, + (nothing for i=1:10)..., radius) - - if precompute_metrics == true - grid = LatitudeLongitudeGrid{FT, TX, TY, TZ, Nothing, Nothing, FX, FY, FZ, VX, VY, VZ, Arch}(architecture, - nλ, nφ, nz, hλ, hφ, hz, Lλ, Lφ, Lz, Δλᶠᵃᵃ, Δλᶜᵃᵃ, λᶠᵃᵃ, λᶜᵃᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ, φᵃᶠᵃ, φᵃᶜᵃ, Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ, - nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, radius) - - Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Δyᶠᶜ, Δyᶜᶠ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ = allocate_metrics(FT, grid) - wait(device_event(architecture)) - - precompute_curvilinear_metrics!(grid, Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ ) - wait(device_event(architecture)) - - Δyᶠᶜ, Δyᶜᶠ = precompute_Δy_metrics(grid, Δyᶠᶜ, Δyᶜᶠ) - - M = typeof(Δxᶠᶜ) - MY = typeof(Δyᶠᶜ) - else - metrics = (:Δxᶠᶜ, :Δxᶜᶠ, :Δxᶠᶠ, :Δxᶜᶜ, :Δyᶠᶜ, :Δyᶜᶠ, :Azᶠᶜ, :Azᶜᶠ, :Azᶠᶠ, :Azᶜᶜ) - for metric in metrics - @eval $metric = nothing - end - M = Nothing - MY = Nothing - end - - return LatitudeLongitudeGrid{FT, TX, TY, TZ, M, MY, FX, FY, FZ, VX, VY, VZ, Arch}(architecture, - nλ, nφ, nz, hλ, hφ, hz, Lλ, Lφ, Lz, Δλᶠᵃᵃ, Δλᶜᵃᵃ, λᶠᵃᵃ, λᶜᵃᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ, φᵃᶠᵃ, φᵃᶜᵃ, Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ, - Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Δyᶠᶜ, Δyᶜᶠ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ, radius) + return !precompute_metrics ? preliminary_grid : with_precomputed_metrics(preliminary_grid) end -function reconstruct_global_grid(grid) +function reconstruct_global_grid(grid::RectilinearGrid) arch = grid.architecture i, j, k = arch.local_index @@ -198,55 +160,40 @@ function reconstruct_global_grid(grid) Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, TY, Ny, Hy, yG, architecture) Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ, Nz, Hz, zG, architecture) - FX = typeof(Δxᶠᵃᵃ) - FY = typeof(Δyᵃᶠᵃ) - FZ = typeof(Δzᵃᵃᶠ) - VX = typeof(xᶠᵃᵃ) - VY = typeof(yᵃᶠᵃ) - VZ = typeof(zᵃᵃᶠ) - Arch = typeof(architecture) - - return RectilinearGrid{FT, TX, TY, TZ, FX, FY, FZ, VX, VY, VZ, Arch}(architecture, - Nx, Ny, Nz, Hx, Hy, Hz, Lx, Ly, Lz, Δxᶠᵃᵃ, Δxᶜᵃᵃ, xᶠᵃᵃ, xᶜᵃᵃ, Δyᵃᶜᵃ, Δyᵃᶠᵃ, yᵃᶠᵃ, yᵃᶜᵃ, Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ) - + return RectilinearGrid{TX, TY, TZ}(architecture, + Nx, Ny, Nz, + Hx, Hy, Hz, + Lx, Ly, Lz, + Δxᶠᵃᵃ, Δxᶜᵃᵃ, xᶠᵃᵃ, xᶜᵃᵃ, + Δyᵃᶜᵃ, Δyᵃᶠᵃ, yᵃᶠᵃ, yᵃᶜᵃ, + Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ) + end -const LocalGrid = Union{RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:MultiArch}, - LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:MultiArch}} - -function with_halo(new_halo, grid::LocalGrid) - new_grid = with_halo(new_halo, reconstruct_global_grid(grid)) +function with_halo(new_halo, grid::DistributedGrid) + new_grid = with_halo(new_halo, reconstruct_global_grid(grid)) return scatter_local_grids(architecture(grid), new_grid) end -function scatter_local_grids(arch::MultiArch, grid::RectilinearGrid) - +function scatter_grid_properties(global_grid) # Pull out face grid constructors - x = cpu_face_constructor_x(grid) - y = cpu_face_constructor_y(grid) - z = cpu_face_constructor_z(grid) + x = cpu_face_constructor_x(global_grid) + y = cpu_face_constructor_y(global_grid) + z = cpu_face_constructor_z(global_grid) topo = topology(grid) - N = pop_flat_elements(size(grid), topo) - halo = pop_flat_elements(halo_size(grid), topo) - - local_grid = RectilinearGrid(arch, eltype(grid); size = N, x = x, y = y, z = z, halo = halo, topology = topo) + sz = pop_flat_elements(size(global_grid), topo) + halo = pop_flat_elements(halo_size(global_grid), topo) - return local_grid + return x, y, z, topo, sz, halo end -function scatter_local_grids(arch::MultiArch, grid::LatitudeLongitudeGrid) - - # Pull out face grid constructors - x = cpu_face_constructor_x(grid) - y = cpu_face_constructor_y(grid) - z = cpu_face_constructor_z(grid) - - topo = topology(grid) - N = pop_flat_element(size(grid), topo) - halo = pop_flat_elements(halo_size(grid), topo) +function scatter_local_grids(arch::MultiArch, global_grid::RectilinearGrid) + x, y, z, topo, sz, halo = scatter_grid_properties(global_grid) + return RectilinearGrid(arch, eltype(grid); size=sz, x=x, y=y, z=z, halo=halo, topology=topo) +end - local_grid = LatitudeLongitudeGrid(arch, eltype(grid); size = N, longitude = x, latitude = y, z = z, halo = halo) - - return local_grid +function scatter_local_grids(arch::MultiArch, global_grid::LatitudeLongitudeGrid) + x, y, z, topo, sz, halo = scatter_grid_properties(global_grid) + return LatitudeLongitudeGrid(arch, eltype(global_grid); size=sz, longitude=x, latitude=y, z=z, halo=halo) end diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl index 3ff61c567d..b4333dca9b 100644 --- a/src/Fields/Fields.jl +++ b/src/Fields/Fields.jl @@ -18,12 +18,12 @@ using Oceananigans.BoundaryConditions include("abstract_field.jl") # include("reduced_getindex_setindex.jl") include("field.jl") +include("field_reductions.jl") include("zero_field.jl") # include("reduced_field.jl") include("averaged_field.jl") # include("computed_field.jl") include("kernel_computed_field.jl") -include("pressure_field.jl") include("function_field.jl") include("regridding_fields.jl") include("set!.jl") @@ -36,8 +36,4 @@ include("field_slicer.jl") include("show_fields.jl") include("broadcasting_abstract_fields.jl") -# Fallback: cannot infer boundary conditions. -boundary_conditions(field) = nothing -boundary_conditions(f::Union{Field, ReducedField, ComputedField, KernelComputedField}) = f.boundary_conditions - end diff --git a/src/Fields/computed_field.jl b/src/Fields/computed_field.jl index 1235e1b458..0ac16a21a0 100644 --- a/src/Fields/computed_field.jl +++ b/src/Fields/computed_field.jl @@ -22,7 +22,7 @@ struct ComputedField{X, Y, Z, S, O, A, D, G, T, C} <: AbstractDataField{X, Y, Z, validate_field_data(X, Y, Z, data, grid) # Use FieldStatus if we want to avoid always recomputing - status = recompute_safely ? nothing : FieldStatus(0.0) + status = recompute_safely ? nothing : FieldStatus() S = typeof(status) T = eltype(grid) diff --git a/src/Fields/field.jl b/src/Fields/field.jl index 1fdbdcf1a4..443afefa19 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -91,6 +91,10 @@ end data(f::Field) = f.data +# Fallback: cannot infer boundary conditions. +boundary_conditions(field) = nothing +boundary_conditions(f::Field) = f.boundary_conditions + "Returns a view of `f` that excludes halo points." function interior(f::Field) LX, LY, LZ = location(f) @@ -151,11 +155,9 @@ Additional keyword arguments are passed to the `Field` constructor. ZFaceField(grid::AbstractGrid, T::DataType=eltype(grid); kw...) = Field{Center, Center, Face}(grid, T; kw...) ##### -##### Fields computed from AbstractOperation and associated utilities +##### Interface for field computations ##### -const ComputedField = Field{<:Any, <:Any, <:Any, <:AbstractOperation} - """ compute!(field) @@ -176,106 +178,45 @@ macro compute(def) return expr end -""" - compute_at!(field, time) - -Computes `field.data` at `time`. Falls back to compute!(field). -""" -compute_at!(field, time) = compute!(field) - -""" - conditional_compute!(field, time) - -Computes `field.data` if `time != field.status.time`. -""" -function conditional_compute!(field, time) - if time == zero(time) || time != field.status.time - compute!(field, time) - field.status.time = time - end - return nothing -end - -# This edge case occurs if `fetch_output` is called with `model::Nothing`. -# We do the safe thing here and always compute. -conditional_compute!(field, ::Nothing) = compute!(field, nothing) +##### +##### Conditional computation +##### mutable struct FieldStatus{T} time :: T end +FieldStatus() = FieldStatus(0.0) Adapt.adapt_structure(to, status::FieldStatus) = (; time = status.time) """ - Field(operand::AbstractOperation; kwargs...) - -Return `f::Field` where `f.data` is computed from `f.operand` by -calling compute!(f). - -Keyword arguments -================= - -data (AbstractArray): An offset Array or CuArray for storing the result of a computation. - Must have `total_size(location(operand), grid)`. - -boundary_conditions (FieldBoundaryConditions): Boundary conditions for `f`. + compute_at!(field, time) -recompute_safely (Bool): whether or not to _always_ "recompute" `f` if `f` is - nested within another computation via an `AbstractOperation`. - If `data` is not provided then `recompute_safely=false` and - recomputation is _avoided_. If `data` is provided, then - `recompute_safely=true` by default. +Computes `field.data` at `time`. Falls back to compute!(field). """ -function Field(operand::AbstractOperation; - data = nothing, - boundary_conditions = FieldBoundaryConditions(op.grid, location(op)), - recompute_safely = true) - - if isnothing(data) - data = new_data(op.grid, location(op)) - recompute_safely = false - end - - status = recompute_safely ? nothing : FieldStatus(0.0) - - return Field(location(op), op.grid, data, boundary_conditions, operand, status) -end +compute_at!(field, time) = compute!(field) """ - compute!(comp::ComputedField) + compute_at!(field, time) -Compute `comp.operand` and store the result in `comp.data`. +Computes `field.data` if `time != field.status.time`. """ -function compute!(comp::ComputedField, time=nothing) - # First compute `dependencies`: - compute_at!(comp.operand, time) - - workgroup, worksize = - work_layout(comp.grid, :xyz, include_right_boundaries = true, location = location(comp)) - - arch = architecture(comp) - compute_kernel! = _compute!(device(arch), workgroup, worksize) - event = compute_kernel!(comp.data, comp.operand; dependencies = device_event(arch)) - wait(device(arch), event) - - fill_halo_regions!(comp, arch) +function compute_at!(field::Field, time) + if isnothing(field.status) # then always compute: + compute!(field, time) - return comp -end + # Otherwise, compute only on initialization or if field.status.time is not current, + elseif time == zero(time) || time != field.status.time + compute!(field, time) + field.status.time = time + end -"""Compute an `operand` and store in `data`.""" -@kernel function _compute!(data, operand) - i, j, k = @index(Global, NTuple) - @inbounds data[i, j, k] = operand[i, j, k] + return field end -function compute_at!(field::ComputedField, time) - if isnothing(field.status) - return compute!(field, time) - else - return conditional_compute!(field, time) - end -end +# This edge case occurs if `fetch_output` is called with `model::Nothing`. +# We do the safe thing here and always compute. +compute_at!(field::Field, ::Nothing) = compute!(field, nothing) ##### ##### Fields that are reduced along one or more dimensions diff --git a/src/Fields/field_reductions.jl b/src/Fields/field_reductions.jl new file mode 100644 index 0000000000..2f80b93004 --- /dev/null +++ b/src/Fields/field_reductions.jl @@ -0,0 +1,58 @@ +##### +##### Reductions of AbstractField +##### + +struct Reduction{R, O, D} + reduce! :: R + operand :: O + dims :: D +end + +""" + Reduction(reduce!, operand; dims) + +Return a `Reduction` of `operand` with `reduce!`, along `dims`. + +Example +======= + +julia> grid = RectilinearGrid(size=(3, 3, 3), x=(0, 1), y=(0, 1), z=(0, 1), + topology=(Periodic, Periodic, Periodic)) + +julia> c = CenterField(grid) + +julia> set!(c, (x, y, z) -> x + y + z) + +julia> max_c² = Field(Reduction(maximum!, c^2, dims=3)) + +julia> compute!(max_c²) +""" +Reduction(reduce!, operand; dims) = Reduction(reduce!, operand, dims) + +function Field(reduction::Reduction; + data = nothing, + recompute_safely = false) + + operand = reduction.operand + grid = operand.grid + LX, LY, LZ = loc = reduced_location(location(operand); dims) + + if isnothing(data) + data = new_data(grid, loc) + recompute_safely = false + end + + boundary_conditions = FieldBoundaryConditions(grid, loc) + status = recompute_safely ? nothing : FieldStatus() + return Field(loc, grid, data, boundary_conditions, reduction, status) +end + +const ReducedComputedField = Field{<:Any, <:Any, <:Any, <:Reduction} + +function compute!(field::ReducedComputedField, time=nothing) + reduction = field.operand + compute_at!(reduction.operand, time) + reduction.reduce!(field, reduction.operand) + return nothing +end + diff --git a/src/Fields/metric_reductions.jl b/src/Fields/metric_reductions.jl new file mode 100644 index 0000000000..ac8a6e92dc --- /dev/null +++ b/src/Fields/metric_reductions.jl @@ -0,0 +1,36 @@ +import Oceananigans.Fields: Reduction + +##### +##### Metric inference +##### + +reduction_grid_metric(dims::Number) reduction_grid_metric(tuple(dims)) +reduction_grid_metric(dims) = dims === tuple(1) ? Δx : + dims === tuple(2) ? Δy : + dims === tuple(3) ? Δz : + dims === (1, 2) ? Az : + dims === (1, 3) ? Ay : + dims === (2, 3) ? Ax : + dims === (1, 2, 3) ? volume + +##### +##### Metric reductions +##### + +abstract type AbstractMetricReduction end +struct Average <: AbstractMetricReduction end +struct Integral <: AbstractMetricReduction end + +reduction(::Average) = mean! +reduction(::Integral) = sum! + +function Reduction(r::AbstractMetricReductions, operand; dims) + dx = reduction_grid_metric(dims) + field_dx = field * dx + return Reduction(reduction(r), field_dx; dims) +end + +# Convenience +Average(field::AbstractField; dims) = Reduction(Average(), field; dims) +Integral(field::AbstractField; dims) = Reduction(Integral(), field; dims) + diff --git a/src/Fields/pressure_field.jl b/src/Fields/pressure_field.jl deleted file mode 100644 index 494f26086c..0000000000 --- a/src/Fields/pressure_field.jl +++ /dev/null @@ -1,5 +0,0 @@ -using Oceananigans.Fields: ComputedField - -PressureField(model, data=model.pressures.pHY′.data, recompute_safely=false) = - ComputedField(model.pressures.pHY′ + model.pressures.pNHS, data=data, - recompute_safely=recompute_safely) diff --git a/src/Fields/show_fields.jl b/src/Fields/show_fields.jl index 2e54027d5e..9963e953b6 100644 --- a/src/Fields/show_fields.jl +++ b/src/Fields/show_fields.jl @@ -12,7 +12,6 @@ short_show(m::Missing) = "$m" short_show(field::AbstractField) = string(typeof(field).name.wrapper, " located at ", show_location(field)) short_show(field::AveragedField) = string("AveragedField over dims=$(field.dims) located at ", show_location(field), " of ", short_show(field.operand)) -short_show(field::ComputedField) = string("ComputedField located at ", show_location(field), " of ", short_show(field.operand)) Base.show(io::IO, field::AbstractField{X, Y, Z, A}) where {X, Y, Z, A} = print(io, "$(short_show(field))\n", @@ -36,21 +35,6 @@ Base.show(io::IO, field::AveragedField) = "├── operand: $(short_show(field.operand))\n", "└── status: ", show_status(field.status)) -Base.show(io::IO, field::ComputedField) = - print(io, "$(short_show(field))\n", - "├── data: $(typeof(field.data)), size: $(size(field))\n", - "├── grid: $(short_show(field.grid))\n", - "├── operand: $(short_show(field.operand))\n", - "└── status: $(show_status(field.status))") - -Base.show(io::IO, field::KernelComputedField) = - print(io, "$(short_show(field))\n", - "├── data: $(typeof(field.data)), size: $(size(field))\n", - "├── grid: $(short_show(field.grid))\n", - "├── computed_dependencies: $(Tuple(short_show(d) for d in field.computed_dependencies))\n", - "├── kernel: $(short_show(field.kernel))\n", - "└── status: $(show_status(field.status))") - Base.show(io::IO, field::ZeroField) = print(io, "ZeroField") short_show(array::OffsetArray{T, D, A}) where {T, D, A} = string("OffsetArray{$T, $D, $A}") diff --git a/src/Grids/latitude_longitude_grid.jl b/src/Grids/latitude_longitude_grid.jl index 5bb7cc2ac6..41191be2ba 100644 --- a/src/Grids/latitude_longitude_grid.jl +++ b/src/Grids/latitude_longitude_grid.jl @@ -3,42 +3,72 @@ using KernelAbstractions: @kernel, @index const R_Earth = 6371.0e3 # Mean radius of the Earth [m] https://en.wikipedia.org/wiki/Earth struct LatitudeLongitudeGrid{FT, TX, TY, TZ, M, MY, FX, FY, FZ, VX, VY, VZ, Arch} <: AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ, Arch} - architecture::Arch - Nx :: Int - Ny :: Int - Nz :: Int - Hx :: Int - Hy :: Int - Hz :: Int - Lx :: FT - Ly :: FT - Lz :: FT - # All directions can be either regular (FX, FY, FZ) <: Number or stretched (FX, FY, FZ)<: AbstractVector - Δλᶠᵃᵃ :: FX - Δλᶜᵃᵃ :: FX - λᶠᵃᵃ :: VX - λᶜᵃᵃ :: VX - Δφᵃᶠᵃ :: FY - Δφᵃᶜᵃ :: FY - φᵃᶠᵃ :: VY - φᵃᶜᵃ :: VY - Δzᵃᵃᶠ :: FZ - Δzᵃᵃᶜ :: FZ - zᵃᵃᶠ :: VZ - zᵃᵃᶜ :: VZ - # Precomputed metrics M <: Nothing means the metrics will be computed on the fly - Δxᶠᶜᵃ :: M - Δxᶜᶠᵃ :: M - Δxᶠᶠᵃ :: M - Δxᶜᶜᵃ :: M - Δyᶠᶜᵃ :: MY - Δyᶜᶠᵃ :: MY - Azᶠᶜᵃ :: M - Azᶜᶠᵃ :: M - Azᶠᶠᵃ :: M - Azᶜᶜᵃ :: M - # Radius of the spherical - radius :: FT + architecture :: Arch + Nx :: Int + Ny :: Int + Nz :: Int + Hx :: Int + Hy :: Int + Hz :: Int + Lx :: FT + Ly :: FT + Lz :: FT + # All directions can be either regular (FX, FY, FZ) <: Number + # or stretched (FX, FY, FZ) <: AbstractVector + Δλᶠᵃᵃ :: FX + Δλᶜᵃᵃ :: FX + λᶠᵃᵃ :: VX + λᶜᵃᵃ :: VX + Δφᵃᶠᵃ :: FY + Δφᵃᶜᵃ :: FY + φᵃᶠᵃ :: VY + φᵃᶜᵃ :: VY + Δzᵃᵃᶠ :: FZ + Δzᵃᵃᶜ :: FZ + zᵃᵃᶠ :: VZ + zᵃᵃᶜ :: VZ + # Precomputed metrics M <: Nothing means metrics will be computed on the fly + Δxᶠᶜᵃ :: M + Δxᶜᶠᵃ :: M + Δxᶠᶠᵃ :: M + Δxᶜᶜᵃ :: M + Δyᶠᶜᵃ :: MY + Δyᶜᶠᵃ :: MY + Azᶠᶜᵃ :: M + Azᶜᶠᵃ :: M + Azᶠᶠᵃ :: M + Azᶜᶜᵃ :: M + # Spherical radius + radius :: FT + + function LatitudeLongitudeGrid{TX, TY, TZ}(architecture::Arch, + Nλ, Nφ, Nz, + Hλ, Hφ, Hz, + Lλ::FT, Lφ::FT, Lz::FT, + Δλᶠᵃᵃ :: FX, Δλᶜᵃᵃ :: FX, + λᶠᵃᵃ :: VX, λᶜᵃᵃ :: VX, + Δφᵃᶠᵃ :: FY, Δφᵃᶜᵃ :: FY, + φᵃᶠᵃ :: VY, φᵃᶜᵃ :: VY, + Δzᵃᵃᶠ :: FZ, Δzᵃᵃᶜ :: FZ, + zᵃᵃᶠ :: VZ, zᵃᵃᶜ :: VZ, + Δxᶠᶜ::M, Δxᶜᶠ::M, + Δxᶠᶠ::M, Δxᶜᶜ::M, + Δyᶠᶜ::MY, Δyᶜᶠ::MY, + Azᶠᶜ::M, Azᶜᶠ::M, Azᶠᶠ::M, Azᶜᶜ::M, + radius::FT) where {Arch <: AbstractArchitecture, + FT, TX, TY, TZ, FX, FY, FZ, VX, VY, VZ, M, MY} + + return new{FT, TX, TY, TZ, M, MY, FX, FY, FZ, VX, VY, VZ, Arch}(architecture, + Nλ, Nφ, Nz, + Hλ, Hφ, Hz, + Lλ, Lφ, Lz, + Δλᶠᵃᵃ, Δλᶜᵃᵃ, λᶠᵃᵃ, λᶜᵃᵃ, + Δφᵃᶠᵃ, Δφᵃᶜᵃ, φᵃᶠᵃ, φᵃᶜᵃ, + Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ, + Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, + Δyᶠᶜ, Δyᶜᶠ, + Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ, radius) + end end const XRegLatLonGrid = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number} @@ -62,7 +92,61 @@ function LatitudeLongitudeGrid(architecture=CPU(), z, radius=R_Earth, halo=(1, 1, 1)) - + + Nλ, Nφ, Nz, Hλ, Hφ, Hz, latitude, longitude, topo = + validate_lat_lon_grid_args(latitude, longitude, size, halo) + + # 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) + + TX, TY, TZ = topo + + 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, + Hλ, Hφ, Hz, + Lλ, Lφ, Lz, + Δλᶠᵃᵃ, Δλᶜᵃᵃ, λᶠᵃᵃ, λᶜᵃᵃ, + Δφᵃᶠᵃ, Δφᵃᶜᵃ, φᵃᶠᵃ, φᵃᶜᵃ, + Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ, + (nothing for i=1:10)..., radius) + + return !precompute_metrics ? preliminary_grid : with_precomputed_metrics(preliminary_grid) +end + +""" Return a reproduction of `grid` with precomputed metric terms. """ +function with_precomputed_metrics(grid) + arch = architecture(grid) + Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Δyᶠᶜ, Δyᶜᶠ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ = allocate_metrics(grid) + wait(device_event(arch)) + + precompute_curvilinear_metrics!(grid, Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ ) + wait(device_event(arch)) + + Δyᶠᶜ, Δyᶜᶠ = precompute_Δy_metrics(grid, Δyᶠᶜ, Δyᶜᶠ) + + Nλ, Nφ, Nz = size(grid) + Hλ, Hφ, Hz = halo_size(grid) + TX, TY, TZ = topology(grid) + + return LatitudeLongitudeGrid{TX, TY, TZ}(architecture(grid), + Nλ, Nφ, Nz, + Hλ, Hφ, Hz, + grid.Lx, grid.Ly, grid.Lz, + grid.Δλᶠᵃᵃ, grid.Δλᶜᵃᵃ, grid.λᶠᵃᵃ, grid.λᶜᵃᵃ, + grid.Δφᵃᶠᵃ, grid.Δφᵃᶜᵃ, grid.φᵃᶠᵃ, grid.φᵃᶜᵃ, + grid.Δzᵃᵃᶠ, grid.Δzᵃᵃᶜ, grid.zᵃᵃᶠ, grid.zᵃᵃᶜ, + Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, + Δyᶠᶜ, Δyᶜᶠ, + Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ, grid.radius) +end + +function validate_lat_lon_grid_args(latitude, longitude, size, halo) + λ₁, λ₂ = get_domain_extent(longitude, size[1]) @assert λ₁ < λ₂ && λ₂ - λ₁ ≤ 360 @@ -82,58 +166,8 @@ function LatitudeLongitudeGrid(architecture=CPU(), Nλ, Nφ, Nz = N = validate_size(TX, TY, TZ, size) Hλ, Hφ, Hz = H = validate_halo(TX, TY, TZ, halo) - - # 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, topo[1], Nλ, Hλ, longitude, architecture) - Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, topo[2], Nφ, Hφ, latitude, architecture) - Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, topo[3], Nz, Hz, z, architecture) - - FX = typeof(Δλᶠᵃᵃ) - FY = typeof(Δφᵃᶠᵃ) - FZ = typeof(Δzᵃᵃᶠ) - VX = typeof(λᶠᵃᵃ) - VY = typeof(φᵃᶠᵃ) - VZ = typeof(zᵃᵃᶠ) - Arch = typeof(architecture) - - - if precompute_metrics == true - grid = LatitudeLongitudeGrid{FT, TX, TY, TZ, Nothing, Nothing, FX, FY, FZ, VX, VY, VZ, Arch}(architecture, - Nλ, Nφ, Nz, Hλ, Hφ, Hz, Lλ, Lφ, Lz, Δλᶠᵃᵃ, Δλᶜᵃᵃ, λᶠᵃᵃ, λᶜᵃᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ, φᵃᶠᵃ, φᵃᶜᵃ, Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ, - nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, radius) - - Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Δyᶠᶜ, Δyᶜᶠ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ = allocate_metrics(FT, grid) - wait(device_event(architecture)) - - precompute_curvilinear_metrics!(grid, Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ ) - wait(device_event(architecture)) - - Δyᶠᶜ, Δyᶜᶠ = precompute_Δy_metrics(grid, Δyᶠᶜ, Δyᶜᶠ) - - M = typeof(Δxᶠᶜ) - MY = typeof(Δyᶠᶜ) - else - Δxᶠᶜ = nothing - Δxᶜᶠ = nothing - Δxᶠᶠ = nothing - Δxᶜᶜ = nothing - Δyᶠᶜ = nothing - Δyᶜᶠ = nothing - Azᶠᶜ = nothing - Azᶜᶠ = nothing - Azᶠᶠ = nothing - Azᶜᶜ = nothing - - M = Nothing - MY = Nothing - end - return LatitudeLongitudeGrid{FT, TX, TY, TZ, M, MY, FX, FY, FZ, VX, VY, VZ, Arch}(architecture, - Nλ, Nφ, Nz, Hλ, Hφ, Hz, Lλ, Lφ, Lz, Δλᶠᵃᵃ, Δλᶜᵃᵃ, λᶠᵃᵃ, λᶜᵃᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ, φᵃᶠᵃ, φᵃᶜᵃ, Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ, - Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Δyᶠᶜ, Δyᶜᶠ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ, radius) + return Nλ, Nφ, Nz, Hλ, Hφ, Hz, latitude, longitude, topo end function domain_string(grid::LatitudeLongitudeGrid) @@ -206,57 +240,47 @@ function with_arch(new_arch, old_grid::LatitudeLongitudeGrid) size = pop_flat_elements(size, topo) halo = pop_flat_elements(halo_size(old_grid), topo) - new_grid = LatitudeLongitudeGrid(new_arch, eltype(old_grid); - size = size, - longitude = x, + return LatitudeLongitudeGrid(new_arch, eltype(old_grid); + size = size, + longitude = x, latitude = y, - z = z, - halo = halo) - return new_grid + z = z, + halo = halo) end - -Adapt.adapt_structure(to, grid::LatitudeLongitudeGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = - LatitudeLongitudeGrid{FT, TX, TY, TZ, - typeof(Adapt.adapt(to, grid.Δxᶠᶜᵃ)), - typeof(Adapt.adapt(to, grid.Δyᶠᶜᵃ)), - typeof(Adapt.adapt(to, grid.Δλᶠᵃᵃ)), - typeof(Adapt.adapt(to, grid.Δφᵃᶠᵃ)), - typeof(Adapt.adapt(to, grid.Δzᵃᵃᶠ)), - typeof(Adapt.adapt(to, grid.λᶠᵃᵃ)), - typeof(Adapt.adapt(to, grid.φᵃᶠᵃ)), - typeof(Adapt.adapt(to, grid.zᵃᵃᶠ)), - Nothing}( - nothing, - grid.Nx, grid.Ny, grid.Nz, - grid.Hx, grid.Hy, grid.Hz, - grid.Lx, grid.Ly, grid.Lz, - Adapt.adapt(to, grid.Δλᶠᵃᵃ), - Adapt.adapt(to, grid.Δλᶜᵃᵃ), - Adapt.adapt(to, grid.λᶠᵃᵃ), - Adapt.adapt(to, grid.λᶜᵃᵃ), - Adapt.adapt(to, grid.Δφᵃᶠᵃ), - Adapt.adapt(to, grid.Δφᵃᶜᵃ), - Adapt.adapt(to, grid.φᵃᶠᵃ), - Adapt.adapt(to, grid.φᵃᶜᵃ), - Adapt.adapt(to, grid.Δzᵃᵃᶠ), - Adapt.adapt(to, grid.Δzᵃᵃᶜ), - Adapt.adapt(to, grid.zᵃᵃᶠ), - Adapt.adapt(to, grid.zᵃᵃᶜ), - Adapt.adapt(to, grid.Δxᶠᶜᵃ), - Adapt.adapt(to, grid.Δxᶜᶠᵃ), - Adapt.adapt(to, grid.Δxᶠᶠᵃ), - Adapt.adapt(to, grid.Δxᶜᶜᵃ), - Adapt.adapt(to, grid.Δyᶠᶜᵃ), - Adapt.adapt(to, grid.Δyᶜᶠᵃ), - Adapt.adapt(to, grid.Azᶠᶜᵃ), - Adapt.adapt(to, grid.Azᶜᶠᵃ), - Adapt.adapt(to, grid.Azᶠᶠᵃ), - Adapt.adapt(to, grid.Azᶜᶜᵃ), - grid.radius) +function Adapt.adapt_structure(to, grid::LatitudeLongitudeGrid) + TX, TY, TZ = topology(grid) + return LatitudeLongitudeGrid{TX, TY, TZ}(nothing, + grid.Nx, grid.Ny, grid.Nz, + grid.Hx, grid.Hy, grid.Hz, + grid.Lx, grid.Ly, grid.Lz, + Adapt.adapt(to, grid.Δλᶠᵃᵃ), + Adapt.adapt(to, grid.Δλᶜᵃᵃ), + Adapt.adapt(to, grid.λᶠᵃᵃ), + Adapt.adapt(to, grid.λᶜᵃᵃ), + Adapt.adapt(to, grid.Δφᵃᶠᵃ), + Adapt.adapt(to, grid.Δφᵃᶜᵃ), + Adapt.adapt(to, grid.φᵃᶠᵃ), + Adapt.adapt(to, grid.φᵃᶜᵃ), + Adapt.adapt(to, grid.Δzᵃᵃᶠ), + Adapt.adapt(to, grid.Δzᵃᵃᶜ), + Adapt.adapt(to, grid.zᵃᵃᶠ), + Adapt.adapt(to, grid.zᵃᵃᶜ), + Adapt.adapt(to, grid.Δxᶠᶜᵃ), + Adapt.adapt(to, grid.Δxᶜᶠᵃ), + Adapt.adapt(to, grid.Δxᶠᶠᵃ), + Adapt.adapt(to, grid.Δxᶜᶜᵃ), + Adapt.adapt(to, grid.Δyᶠᶜᵃ), + Adapt.adapt(to, grid.Δyᶜᶠᵃ), + Adapt.adapt(to, grid.Azᶠᶜᵃ), + Adapt.adapt(to, grid.Azᶜᶠᵃ), + Adapt.adapt(to, grid.Azᶠᶠᵃ), + Adapt.adapt(to, grid.Azᶜᶜᵃ), + grid.radius) +end ##### -##### Pre compute LatitudeLongitudeGrid metrics +##### On-the-fly computation of LatitudeLongitudeGrid metrics ##### @inline hack_cosd(φ) = cos(π * φ / 180) @@ -284,68 +308,76 @@ Adapt.adapt_structure(to, grid::LatitudeLongitudeGrid{FT, TX, TY, TZ}) where {FT @inline Azᶠᶠᵃ(i, j, k, grid::XRegLatLonGrid) = @inbounds grid.radius^2 * deg2rad(grid.Δλᶠᵃᵃ) * (hack_sind(grid.φᵃᶜᵃ[j]) - hack_sind(grid.φᵃᶜᵃ[j-1])) @inline Azᶜᶜᵃ(i, j, k, grid::XRegLatLonGrid) = @inbounds grid.radius^2 * deg2rad(grid.Δλᶜᵃᵃ) * (hack_sind(grid.φᵃᶠᵃ[j+1]) - hack_sind(grid.φᵃᶠᵃ[j])) -####### -####### Utilities to precompute Metrics -####### +##### +##### Utilities to precompute metrics +##### @inline metrics_precomputed(::LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, Nothing}) = false @inline metrics_precomputed(::LatitudeLongitudeGrid) = true -####### Kernels that precompute the z- and x-metric +##### +##### Kernels that precompute the z- and x-metric +##### -@inline metric_worksize(grid::LatitudeLongitudeGrid) = (length(grid.Δλᶜᵃᵃ), length(grid.φᵃᶜᵃ) - 1) -@inline metric_workgroup(grid::LatitudeLongitudeGrid) = (16, 16) +@inline metric_worksize(grid::LatitudeLongitudeGrid) = (length(grid.Δλᶜᵃᵃ), length(grid.φᵃᶜᵃ) - 1) +@inline metric_workgroup(grid::LatitudeLongitudeGrid) = (16, 16) @inline metric_worksize(grid::XRegLatLonGrid) = length(grid.φᵃᶜᵃ) - 1 @inline metric_workgroup(grid::XRegLatLonGrid) = 16 - -function precompute_curvilinear_metrics!(grid, Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ) +function precompute_curvilinear_metrics!(grid, Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ) arch = grid.architecture + workgroup, worksize = metric_workgroup(grid), metric_worksize(grid) curvilinear_metrics! = precompute_metrics_kernel!(Architectures.device(arch), workgroup, worksize) - event = curvilinear_metrics!(grid, Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ; dependencies=device_event(arch)) - + + event = curvilinear_metrics!(grid, Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ; dependencies=device_event(arch)) wait(event) + return nothing end @kernel function precompute_metrics_kernel!(grid::LatitudeLongitudeGrid, Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ) i, j = @index(Global, NTuple) - i += grid.Δλᶜᵃᵃ.offsets[1] - j += grid.φᵃᶜᵃ.offsets[1] + 1 + + # Manually offset x- and y-index + i′ = i + grid.Δλᶜᵃᵃ.offsets[1] + j′ = j + grid.φᵃᶜᵃ.offsets[1] + 1 + @inbounds begin - Δxᶠᶜ[i, j] = Δxᶠᶜᵃ(i, j, 1, grid) - Δxᶜᶠ[i, j] = Δxᶜᶠᵃ(i, j, 1, grid) - Δxᶠᶠ[i, j] = Δxᶠᶠᵃ(i, j, 1, grid) - Δxᶜᶜ[i, j] = Δxᶜᶜᵃ(i, j, 1, grid) - Azᶠᶜ[i, j] = Azᶠᶜᵃ(i, j, 1, grid) - Azᶜᶠ[i, j] = Azᶜᶠᵃ(i, j, 1, grid) - Azᶠᶠ[i, j] = Azᶠᶠᵃ(i, j, 1, grid) - Azᶜᶜ[i, j] = Azᶜᶜᵃ(i, j, 1, grid) + Δxᶠᶜ[i′, j′] = Δxᶠᶜᵃ(i′, j′, 1, grid) + Δxᶜᶠ[i′, j′] = Δxᶜᶠᵃ(i′, j′, 1, grid) + Δxᶠᶠ[i′, j′] = Δxᶠᶠᵃ(i′, j′, 1, grid) + Δxᶜᶜ[i′, j′] = Δxᶜᶜᵃ(i′, j′, 1, grid) + Azᶠᶜ[i′, j′] = Azᶠᶜᵃ(i′, j′, 1, grid) + Azᶜᶠ[i′, j′] = Azᶜᶠᵃ(i′, j′, 1, grid) + Azᶠᶠ[i′, j′] = Azᶠᶠᵃ(i′, j′, 1, grid) + Azᶜᶜ[i′, j′] = Azᶜᶜᵃ(i′, j′, 1, grid) end end @kernel function precompute_metrics_kernel!(grid::XRegLatLonGrid, Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ) j = @index(Global, Linear) - j += grid.φᵃᶜᵃ.offsets[1] + 1 + + # Manually offset y-index + j′ = j + grid.φᵃᶜᵃ.offsets[1] + 1 + @inbounds begin - Δxᶠᶜ[j] = Δxᶠᶜᵃ(1, j, 1, grid) - Δxᶜᶠ[j] = Δxᶜᶠᵃ(1, j, 1, grid) - Δxᶠᶠ[j] = Δxᶠᶠᵃ(1, j, 1, grid) - Δxᶜᶜ[j] = Δxᶜᶜᵃ(1, j, 1, grid) - Azᶠᶜ[j] = Azᶠᶜᵃ(1, j, 1, grid) - Azᶜᶠ[j] = Azᶜᶠᵃ(1, j, 1, grid) - Azᶠᶠ[j] = Azᶠᶠᵃ(1, j, 1, grid) - Azᶜᶜ[j] = Azᶜᶜᵃ(1, j, 1, grid) + Δxᶠᶜ[j′] = Δxᶠᶜᵃ(1, j′, 1, grid) + Δxᶜᶠ[j′] = Δxᶜᶠᵃ(1, j′, 1, grid) + Δxᶠᶠ[j′] = Δxᶠᶠᵃ(1, j′, 1, grid) + Δxᶜᶜ[j′] = Δxᶜᶜᵃ(1, j′, 1, grid) + Azᶠᶜ[j′] = Azᶠᶜᵃ(1, j′, 1, grid) + Azᶜᶠ[j′] = Azᶜᶠᵃ(1, j′, 1, grid) + Azᶠᶠ[j′] = Azᶠᶠᵃ(1, j′, 1, grid) + Azᶜᶜ[j′] = Azᶜᶜᵃ(1, j′, 1, grid) end end ####### Kernels that precompute the y-metric -function precompute_Δy_metrics(grid::LatitudeLongitudeGrid, Δyᶠᶜ, Δyᶜᶠ) - +function precompute_Δy_metrics(grid::LatitudeLongitudeGrid, Δyᶠᶜ, Δyᶜᶠ) arch = grid.architecture precompute_Δy! = precompute_Δy_kernel!(Architectures.device(arch), 16, length(grid.Δφᵃᶜᵃ) - 1) event = precompute_Δy!(grid, Δyᶠᶜ, Δyᶜᶠ; dependencies=device_event(arch)) @@ -362,21 +394,24 @@ end @kernel function precompute_Δy_kernel!(grid, Δyᶠᶜ, Δyᶜᶠ) j = @index(Global, Linear) - j += grid.Δφᵃᶜᵃ.offsets[1] + 1 + + # Manually offset y-index + j′ = j + grid.Δφᵃᶜᵃ.offsets[1] + 1 + @inbounds begin - Δyᶜᶠ[j] = Δyᶜᶠᵃ(1, j, 1, grid) - Δyᶠᶜ[j] = Δyᶜᶠᵃ(1, j, 1, grid) + Δyᶜᶠ[j′] = Δyᶜᶠᵃ(1, j′, 1, grid) + Δyᶠᶜ[j′] = Δyᶜᶠᵃ(1, j′, 1, grid) end end -####### -####### Preallocation kernel for metrics -####### +##### +##### Metric memory allocation +##### -function allocate_metrics(FT, grid::LatitudeLongitudeGrid) +function allocate_metrics(grid::LatitudeLongitudeGrid) + FT = eltype(grid) # preallocate quantities to ensure correct type and size - grid_metrics = (:Δxᶠᶜ, :Δxᶜᶠ, :Δxᶠᶠ, @@ -389,11 +424,11 @@ function allocate_metrics(FT, grid::LatitudeLongitudeGrid) arch = grid.architecture if typeof(grid) <: XRegLatLonGrid - offsets = grid.φᵃᶜᵃ.offsets[1] - metric_size = length(grid.φᵃᶜᵃ) + offsets = grid.φᵃᶜᵃ.offsets[1] + metric_size = length(grid.φᵃᶜᵃ) else - offsets = (grid.Δλᶜᵃᵃ.offsets[1], grid.φᵃᶜᵃ.offsets[1]) - metric_size = (length(grid.Δλᶜᵃᵃ) , length(grid.φᵃᶜᵃ)) + offsets = (grid.Δλᶜᵃᵃ.offsets[1], grid.φᵃᶜᵃ.offsets[1]) + metric_size = (length(grid.Δλᶜᵃᵃ) , length(grid.φᵃᶜᵃ)) end for metric in grid_metrics @@ -403,13 +438,13 @@ function allocate_metrics(FT, grid::LatitudeLongitudeGrid) end if typeof(grid) <: YRegLatLonGrid - Δyᶠᶜ = FT(0.0) - Δyᶜᶠ = FT(0.0) + Δyᶠᶜ = FT(0.0) + Δyᶜᶠ = FT(0.0) else - parentC = zeros(FT, length(grid.Δφᵃᶜᵃ)) - parentF = zeros(FT, length(grid.Δφᵃᶜᵃ)) - Δyᶠᶜ = OffsetArray(arch_array(arch, parentC), grid.Δφᵃᶜᵃ.offsets[1]) - Δyᶜᶠ = OffsetArray(arch_array(arch, parentF), grid.Δφᵃᶜᵃ.offsets[1]) + parentC = zeros(FT, length(grid.Δφᵃᶜᵃ)) + parentF = zeros(FT, length(grid.Δφᵃᶜᵃ)) + Δyᶠᶜ = OffsetArray(arch_array(arch, parentC), grid.Δφᵃᶜᵃ.offsets[1]) + Δyᶜᶠ = OffsetArray(arch_array(arch, parentF), grid.Δφᵃᶜᵃ.offsets[1]) end return Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Δyᶠᶜ, Δyᶜᶠ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl index 2c03e5a8bd..3fa537e79c 100644 --- a/src/Grids/rectilinear_grid.jl +++ b/src/Grids/rectilinear_grid.jl @@ -5,31 +5,52 @@ A rectilinear grid with with either constant or varying grid spacings between ce in all directions. Grid elements of type `FT`, topology `{TX, TY, TZ}`, grid spacings of type `{FX, FY, FZ}` and coordinates in each direction of type `{VX, VY, VZ}`. """ - struct RectilinearGrid{FT, TX, TY, TZ, FX, FY, FZ, VX, VY, VZ, Arch} <: AbstractRectilinearGrid{FT, TX, TY, TZ, Arch} - architecture::Arch - Nx :: Int - Ny :: Int - Nz :: Int - Hx :: Int - Hy :: Int - Hz :: Int - Lx :: FT - Ly :: FT - Lz :: FT - # All directions can be either regular (FX, FY, FZ) <: Number or stretched (FX, FY, FZ) <: AbstractVector - Δxᶠᵃᵃ :: FX - Δxᶜᵃᵃ :: FX - xᶠᵃᵃ :: VX - xᶜᵃᵃ :: VX - Δyᵃᶠᵃ :: FY - Δyᵃᶜᵃ :: FY - yᵃᶠᵃ :: VY - yᵃᶜᵃ :: VY - Δzᵃᵃᶠ :: FZ - Δzᵃᵃᶜ :: FZ - zᵃᵃᶠ :: VZ - zᵃᵃᶜ :: VZ + architecture :: Arch + Nx :: Int + Ny :: Int + Nz :: Int + Hx :: Int + Hy :: Int + Hz :: Int + Lx :: FT + Ly :: FT + Lz :: FT + # All directions can be either regular (FX, FY, FZ) <: Number + # or stretched (FX, FY, FZ) <: AbstractVector + Δxᶠᵃᵃ :: FX + Δxᶜᵃᵃ :: FX + xᶠᵃᵃ :: VX + xᶜᵃᵃ :: VX + Δyᵃᶠᵃ :: FY + Δyᵃᶜᵃ :: FY + yᵃᶠᵃ :: VY + yᵃᶜᵃ :: VY + Δzᵃᵃᶠ :: FZ + Δzᵃᵃᶜ :: FZ + zᵃᵃᶠ :: VZ + zᵃᵃᶜ :: VZ + + function RectilinearGrid{TX, TY, TZ}(arch::Arch, + Nx, Ny, Nz, + Hx, Hy, Hz, + Lx::FT, Ly::FT, Lz::FT, + Δxᶠᵃᵃ :: FX, Δxᶜᵃᵃ :: FX, + xᶠᵃᵃ :: VX, xᶜᵃᵃ :: VX, + Δyᵃᶠᵃ :: FY, Δyᵃᶜᵃ :: FY, + yᵃᶠᵃ :: VY, yᵃᶜᵃ :: VY, + Δzᵃᵃᶠ :: FZ, Δzᵃᵃᶜ :: FZ, + zᵃᵃᶠ :: VZ, zᵃᵃᶜ :: VZ) where {Arch <: AbstractArchitecture, + TX, TY, TZ, FT, + FX, VX, FY, VY, + FZ, VZ} + + return new{FT, TX, TY, TZ, FX, FY, FZ, VX, VY, VZ, Arch}(arch, Nx, Ny, Nz, + Hx, Hy, Hz, Lx, Ly, Lz, + Δxᶠᵃᵃ, Δxᶜᵃᵃ, xᶠᵃᵃ, xᶜᵃᵃ, + Δyᵃᶠᵃ, Δyᵃᶜᵃ, yᵃᶠᵃ, yᵃᶜᵃ, + Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ) + end end const XRegRectilinearGrid = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number} @@ -53,10 +74,10 @@ Creates a `RectilinearGrid` with `size = (Nx, Ny, Nz)` grid points. Positional arguments ================= -- `architecture`: Specifies whether the array of coordinates, interfaces, and spacings are stored +- `architecture`: Specifies whether arrays of coordinates and spacings are stored on the CPU or GPU. Default: `architecture = CPU()`. -- `FT` : Floating point precision, might be Float32 or Float64 +- `FT` : Floating point data type. Default: `FT = Float64`. Keyword arguments ================= @@ -146,7 +167,7 @@ RectilinearGrid{Float64, Periodic, Periodic, Bounded} ```jldoctest julia> using Oceananigans -julia> grid = RectilinearGrid(CPU(), Float32; size=(32, 32, 16), x=(0, 8), y=(-10, 10), z=(-π, π)) +julia> grid = RectilinearGrid(Float32; size=(32, 32, 16), x=(0, 8), y=(-10, 10), z=(-π, π)) RectilinearGrid{Float32, Periodic, Periodic, Bounded} architecture: CPU() domain: x ∈ [0.0, 8.0], y ∈ [-10.0, 10.0], z ∈ [-3.1415927, 3.1415927] @@ -205,10 +226,8 @@ julia> Lz = 32; # depth (m) julia> hyperbolically_spaced_faces(k) = - Lz * (1 - tanh(σ * (k - 1) / Nz) / tanh(σ)); -julia> grid = RectilinearGrid(size = (32, 32, Nz), - x = (0, 64), - y = (0, 64), - z = hyperbolically_spaced_faces) +julia> grid = RectilinearGrid(size = (32, 32, Nz), x = (0, 64), + y = (0, 64), z = hyperbolically_spaced_faces) RectilinearGrid{Float64, Periodic, Periodic, Bounded} architecture: CPU() domain: x ∈ [0.0, 64.0], y ∈ [0.0, 64.0], z ∈ [-32.0, -0.0] @@ -263,12 +282,7 @@ function RectilinearGrid(architecture = CPU(), extent = nothing, topology = (Periodic, Periodic, Bounded)) - TX, TY, TZ = validate_topology(topology) - size = validate_size(TX, TY, TZ, size) - halo = validate_halo(TX, TY, TZ, halo) - - # Validate the rectilinear domain - x, y, z = validate_rectilinear_domain(TX, TY, TZ, FT, extent, x, y, z) + TX, TY, TZ, size, halo, x, y, z = validate_rectilinear_grid_args(topology, size, halo, FT, extent, x, y, z) Nx, Ny, Nz = size Hx, Hy, Hz = halo @@ -277,21 +291,34 @@ function RectilinearGrid(architecture = CPU(), 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) - FX = typeof(Δxᶠᵃᵃ) - FY = typeof(Δyᵃᶠᵃ) - FZ = typeof(Δzᵃᵃᶠ) - VX = typeof(xᶠᵃᵃ) - VY = typeof(yᵃᶠᵃ) - VZ = typeof(zᵃᵃᶠ) - Arch = typeof(architecture) - - return RectilinearGrid{FT, TX, TY, TZ, FX, FY, FZ, VX, VY, VZ, Arch}(architecture, - Nx, Ny, Nz, Hx, Hy, Hz, Lx, Ly, Lz, Δxᶠᵃᵃ, Δxᶜᵃᵃ, xᶠᵃᵃ, xᶜᵃᵃ, Δyᵃᶜᵃ, Δyᵃᶠᵃ, yᵃᶠᵃ, yᵃᶜᵃ, Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ) + return RectilinearGrid{TX, TY, TZ}(architecture, + Nx, Ny, Nz, + Hx, Hy, Hz, + FT(Lx), FT(Ly), FT(Lz), + Δxᶠᵃᵃ, Δxᶜᵃᵃ, xᶠᵃᵃ, xᶜᵃᵃ, + Δyᵃᶜᵃ, Δyᵃᶠᵃ, yᵃᶠᵃ, yᵃᶜᵃ, + Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ) end -@inline x_domain(grid::RectilinearGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = domain(TX, grid.Nx, grid.xᶠᵃᵃ) -@inline y_domain(grid::RectilinearGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = domain(TY, grid.Ny, grid.yᵃᶠᵃ) -@inline z_domain(grid::RectilinearGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = domain(TZ, grid.Nz, grid.zᵃᵃᶠ) +""" Validate user input arguments to the `RectilinearGrid` constructor. """ +function validate_rectilinear_grid_args(topology, size, halo, FT, extent, x, y, z) + TX, TY, TZ = validate_topology(topology) + size = validate_size(TX, TY, TZ, size) + halo = validate_halo(TX, TY, TZ, halo) + + # Validate the rectilinear domain + x, y, z = validate_rectilinear_domain(TX, TY, TZ, FT, extent, x, y, z) + + return TX, TY, TZ, size, halo, x, y, z +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ᵃᵃᶠ) short_show(grid::RectilinearGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = "RectilinearGrid{$FT, $TX, $TY, $TZ}(Nx=$(grid.Nx), Ny=$(grid.Ny), Nz=$(grid.Nz))" @@ -315,32 +342,30 @@ function show(io::IO, g::RectilinearGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} " grid in z: ", show_coordinate(g.Δzᵃᵃᶜ, TZ)) end +##### +##### Utilities +##### + +function Adapt.adapt_structure(to, grid::RectilinearGrid) + TX, TY, TZ = topology(grid) + return RectilinearGrid{TX, TY, TZ}(nothing, + grid.Nx, grid.Ny, grid.Nz, + grid.Hx, grid.Hy, grid.Hz, + grid.Lx, grid.Ly, grid.Lz, + Adapt.adapt(to, grid.Δxᶠᵃᵃ), + Adapt.adapt(to, grid.Δxᶜᵃᵃ), + Adapt.adapt(to, grid.xᶠᵃᵃ), + Adapt.adapt(to, grid.xᶜᵃᵃ), + Adapt.adapt(to, grid.Δyᵃᶠᵃ), + Adapt.adapt(to, grid.Δyᵃᶜᵃ), + Adapt.adapt(to, grid.yᵃᶠᵃ), + Adapt.adapt(to, grid.yᵃᶜᵃ), + Adapt.adapt(to, grid.Δzᵃᵃᶠ), + Adapt.adapt(to, grid.Δzᵃᵃᶜ), + Adapt.adapt(to, grid.zᵃᵃᶠ), + Adapt.adapt(to, grid.zᵃᵃᶜ)) +end -Adapt.adapt_structure(to, grid::RectilinearGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = - RectilinearGrid{FT, TX, TY, TZ, - typeof(Adapt.adapt(to, grid.Δxᶜᵃᵃ)), - typeof(Adapt.adapt(to, grid.Δyᵃᶠᵃ)), - typeof(Adapt.adapt(to, grid.Δzᵃᵃᶠ)), - typeof(Adapt.adapt(to, grid.xᶠᵃᵃ)), - typeof(Adapt.adapt(to, grid.yᵃᶠᵃ)), - typeof(Adapt.adapt(to, grid.zᵃᵃᶠ)), - Nothing}( - nothing, - grid.Nx, grid.Ny, grid.Nz, - grid.Hx, grid.Hy, grid.Hz, - grid.Lx, grid.Ly, grid.Lz, - Adapt.adapt(to, grid.Δxᶠᵃᵃ), - Adapt.adapt(to, grid.Δxᶜᵃᵃ), - Adapt.adapt(to, grid.xᶠᵃᵃ), - Adapt.adapt(to, grid.xᶜᵃᵃ), - Adapt.adapt(to, grid.Δyᵃᶠᵃ), - Adapt.adapt(to, grid.Δyᵃᶜᵃ), - Adapt.adapt(to, grid.yᵃᶠᵃ), - Adapt.adapt(to, grid.yᵃᶜᵃ), - Adapt.adapt(to, grid.Δzᵃᵃᶠ), - Adapt.adapt(to, grid.Δzᵃᵃᶜ), - Adapt.adapt(to, grid.zᵃᵃᶠ), - Adapt.adapt(to, grid.zᵃᵃᶜ)) @inline xnode(::Center, i, grid::RectilinearGrid) = @inbounds grid.xᶜᵃᵃ[i] @inline xnode(::Face , i, grid::RectilinearGrid) = @inbounds grid.xᶠᵃᵃ[i] @@ -377,10 +402,10 @@ function with_halo(new_halo, old_grid::RectilinearGrid) new_halo = pop_flat_elements(new_halo, topo) new_grid = RectilinearGrid(architecture(old_grid), eltype(old_grid); - size = size, - x = x, y = y, z = z, - topology = topo, - halo = new_halo) + size = size, + x = x, y = y, z = z, + topology = topo, + halo = new_halo) return new_grid end @@ -400,17 +425,17 @@ function with_arch(new_arch, old_grid::RectilinearGrid) halo = pop_flat_elements(halo_size(old_grid), topo) new_grid = RectilinearGrid(new_arch, eltype(old_grid); - size = size, - x = x, y = y, z = z, - topology = topo, - halo = halo) + size = size, + x = x, y = y, z = z, + topology = topo, + halo = halo) return new_grid end -# -# Get minima of grid -# +##### +##### Get minima of grid +##### function min_Δx(grid::RectilinearGrid) topo = topology(grid) diff --git a/src/Models/HydrostaticFreeSurfaceModels/compute_vertically_integrated_lateral_areas.jl b/src/Models/HydrostaticFreeSurfaceModels/compute_vertically_integrated_lateral_areas.jl index 85be87b5aa..a7c4c9c999 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/compute_vertically_integrated_lateral_areas.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/compute_vertically_integrated_lateral_areas.jl @@ -12,7 +12,9 @@ end end -function compute_vertically_integrated_lateral_areas!(∫ᶻ_A, grid, arch) +function compute_vertically_integrated_lateral_areas!(∫ᶻ_A, grid) + + arch = architecture(grid) event = launch!(arch, grid, diff --git a/src/Models/HydrostaticFreeSurfaceModels/fft_based_implicit_free_surface_solver.jl b/src/Models/HydrostaticFreeSurfaceModels/fft_based_implicit_free_surface_solver.jl index c8e0c531ce..0c6dff4ad5 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/fft_based_implicit_free_surface_solver.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/fft_based_implicit_free_surface_solver.jl @@ -16,7 +16,7 @@ struct FFTImplicitFreeSurfaceSolver{S, G3, G2, R} end """ - function FFTImplicitFreeSurfaceSolver(arch, grid, settings) + function FFTImplicitFreeSurfaceSolver(grid, settings) Return a solver based on the fast Fourier transform for the elliptic equation @@ -28,7 +28,7 @@ representing an implicit time discretization of the linear free surface evolutio for a fluid with constant depth `H`, horizontal areas `Az`, barotropic volume flux `Q★`, time step `Δt`, gravitational acceleration `g`, and free surface at time-step `n`, `ηⁿ`. """ -function FFTImplicitFreeSurfaceSolver(arch, grid, settings) +function FFTImplicitFreeSurfaceSolver(grid, settings) grid isa RegRectilinearGrid || grid isa HRegRectilinearGrid || throw(ArgumentError("FFTImplicitFreeSurfaceSolver requires horizontally-regular rectilinear grids.")) @@ -53,19 +53,20 @@ function FFTImplicitFreeSurfaceSolver(arch, grid, settings) # Even if the three dimensional grid is vertically stretched, we can only use # FFTImplicitFreeSurfaceSolver with grids that are regularly spaced in the # horizontal direction. - horizontal_grid = RectilinearGrid(arch; topology = (TX, TY, Flat), - size = sz, - halo = halo, - domain...) + horizontal_grid = RectilinearGrid(architecture(grid); + topology = (TX, TY, Flat), + size = sz, + halo = halo, + domain...) - solver = FFTBasedPoissonSolver(arch, horizontal_grid) + solver = FFTBasedPoissonSolver(horizontal_grid) right_hand_side = solver.storage return FFTImplicitFreeSurfaceSolver(solver, grid, horizontal_grid, right_hand_side) end -build_implicit_step_solver(::Val{:FastFourierTransform}, arch, grid, settings) = - FFTImplicitFreeSurfaceSolver(arch, grid, settings) +build_implicit_step_solver(::Val{:FastFourierTransform}, grid, settings) = + FFTImplicitFreeSurfaceSolver(grid, settings) ##### ##### Solve... @@ -89,7 +90,7 @@ function compute_implicit_free_surface_right_hand_side!(rhs, g, Δt, ∫ᶻQ, η) solver = implicit_solver.fft_based_poisson_solver - arch = solver.architecture + arch = architecture(solver) grid = implicit_solver.three_dimensional_grid H = grid.Lz diff --git a/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl index 7d5f69acbe..928bbf2570 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl @@ -55,12 +55,12 @@ function FreeSurface(free_surface::ImplicitFreeSurface{Nothing}, velocities, gri g = convert(eltype(grid), free_surface.gravitational_acceleration) # Initialize barotropic volume fluxes - barotropic_x_volume_flux = ReducedField(Face, Center, Nothing, arch, grid; dims=3) - barotropic_y_volume_flux = ReducedField(Center, Face, Nothing, arch, grid; dims=3) + barotropic_x_volume_flux = Field{Face, Center, Nothing}(grid) + barotropic_y_volume_flux = Field{Center, Face, Nothing}(grid) barotropic_volume_flux = (u=barotropic_x_volume_flux, v=barotropic_y_volume_flux) solver_method = free_surface.solver_method - solver = build_implicit_step_solver(Val(solver_method), arch, grid, free_surface.solver_settings) + solver = build_implicit_step_solver(Val(solver_method), grid, free_surface.solver_settings) return ImplicitFreeSurface(η, g, barotropic_volume_flux, @@ -72,9 +72,9 @@ end is_horizontally_regular(grid) = false is_horizontally_regular(::RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number, <:Number}) = true -function build_implicit_step_solver(::Val{:Default}, arch, grid, settings) +function build_implicit_step_solver(::Val{:Default}, grid, settings) default_method = is_horizontally_regular(grid) ? :FastFourierTransform : :PreconditionedConjugateGradient - return build_implicit_step_solver(Val(default_method), arch, grid, settings) + return build_implicit_step_solver(Val(default_method), grid, settings) end @inline explicit_barotropic_pressure_x_gradient(i, j, k, grid, ::ImplicitFreeSurface) = 0 diff --git a/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl b/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl index d8d685477c..0c0c91395a 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl @@ -4,6 +4,7 @@ using Oceananigans.Architectures using Oceananigans.Fields: Field, ZReducedField import Oceananigans.Solvers: solve! +import Oceananigans.Architectures: architecture struct PCGImplicitFreeSurfaceSolver{V, S, R} vertically_integrated_lateral_areas :: V @@ -11,8 +12,11 @@ struct PCGImplicitFreeSurfaceSolver{V, S, R} right_hand_side :: R end +architecture(solver::PCGImplicitFreeSurfaceSolver) = + architecture(solver.preconditioned_conjugate_gradient_solver) + """ - PCGImplicitFreeSurfaceSolver(arch::AbstractArchitecture, grid, settings) + PCGImplicitFreeSurfaceSolver(grid, settings) Return a solver based on a preconditioned conjugate gradient method for the elliptic equation @@ -24,14 +28,14 @@ representing an implicit time discretization of the linear free surface evolutio for a fluid with variable depth `H`, horizontal areas `Az`, barotropic volume flux `Q★`, time step `Δt`, gravitational acceleration `g`, and free surface at time-step `n` `ηⁿ`. """ -function PCGImplicitFreeSurfaceSolver(arch::AbstractArchitecture, grid, settings) +function PCGImplicitFreeSurfaceSolver(grid, settings) # Initialize vertically integrated lateral face areas ∫ᶻ_Axᶠᶜᶜ = Field{Face, Center, Nothing}(grid) ∫ᶻ_Ayᶜᶠᶜ = Field{Center, Face, Nothing}(grid) vertically_integrated_lateral_areas = (xᶠᶜᶜ = ∫ᶻ_Axᶠᶜᶜ, yᶜᶠᶜ = ∫ᶻ_Ayᶜᶠᶜ) - compute_vertically_integrated_lateral_areas!(vertically_integrated_lateral_areas, grid, arch) + compute_vertically_integrated_lateral_areas!(vertically_integrated_lateral_areas, grid) right_hand_side = Field{Center, Center, Nothing}(grid) @@ -47,8 +51,8 @@ function PCGImplicitFreeSurfaceSolver(arch::AbstractArchitecture, grid, settings return PCGImplicitFreeSurfaceSolver(vertically_integrated_lateral_areas, solver, right_hand_side) end -build_implicit_step_solver(::Val{:PreconditionedConjugateGradient}, arch, grid, settings) = - PCGImplicitFreeSurfaceSolver(arch, grid, settings) +build_implicit_step_solver(::Val{:PreconditionedConjugateGradient}, grid, settings) = + PCGImplicitFreeSurfaceSolver(grid, settings) ##### ##### Solve... @@ -76,7 +80,7 @@ function compute_implicit_free_surface_right_hand_side!(rhs, g, Δt, ∫ᶻQ, η) solver = implicit_solver.preconditioned_conjugate_gradient_solver - arch = solver.architecture + arch = architecture(solver) grid = solver.grid event = launch!(arch, grid, :xy, diff --git a/src/Models/HydrostaticFreeSurfaceModels/vertical_vorticity_field.jl b/src/Models/HydrostaticFreeSurfaceModels/vertical_vorticity_field.jl index 808fe1dc3c..dd3fdb3a07 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/vertical_vorticity_field.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/vertical_vorticity_field.jl @@ -1,17 +1,11 @@ using Oceananigans.Grids: Face, Face, Center using Oceananigans.Operators: ζ₃ᶠᶠᵃ -using Oceananigans.Fields: KernelComputedField -using KernelAbstractions: @kernel - -@kernel function compute_vertical_vorticity!(ζ, grid, u, v) - i, j, k, = @index(Global, NTuple) - @inbounds ζ[i, j, k] = ζ₃ᶠᶠᵃ(i, j, k, grid, u, v) -end +using Oceananigans.AbstractOperations: KernelFunctionOperation """ VerticalVorticityField(model) -Returns a KernelComputedField that `compute!`s vertical vorticity in a +Returns a Field that `compute!`s vertical vorticity in a manner consistent with the `VectorInvariant` momentum advection scheme for curvilinear grids. @@ -20,5 +14,10 @@ the vertical vorticity by first integrating the velocity field around the border of the vorticity cell to find the vertical circulation, and then dividing by the area of the vorticity cell to compute vertical vorticity. """ -VerticalVorticityField(model) = KernelComputedField(Face, Face, Center, compute_vertical_vorticity!, model, - computed_dependencies = (model.velocities.u, model.velocities.v)) +function VerticalVorticityField(model; kw...) + grid = model.grid + u, v, w = model.velocities + vorticity_operation = KernelFunctionOperation{Face, Face, Center}(ζ₃ᶠᶠᵃ, grid, computed_dependencies=(u, v)) + return Field(vorticity_operation; kw...) +end + diff --git a/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl b/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl index 5dee55b824..f5d7cdfe21 100644 --- a/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl +++ b/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl @@ -36,6 +36,7 @@ PressureSolver(arch, ibg::ImmersedBoundaryGrid) = PressureSolver(arch, ibg.grid) ##### include("nonhydrostatic_model.jl") +include("pressure_field.jl") include("show_nonhydrostatic_model.jl") include("set_nonhydrostatic_model.jl") diff --git a/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl b/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl index b211fa414a..f36475d4ee 100644 --- a/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl +++ b/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl @@ -18,6 +18,8 @@ using Oceananigans.LagrangianParticleTracking: LagrangianParticles using Oceananigans.Utils: tupleit using Oceananigans.Grids: topology +import Oceananigans.Architectures: architecture + const ParticlesOrNothing = Union{Nothing, LagrangianParticles} mutable struct NonhydrostaticModel{TS, E, A<:AbstractArchitecture, G, T, B, R, SD, U, C, Φ, F, @@ -181,6 +183,8 @@ function NonhydrostaticModel(; grid, return model end +architecture(model::NonhydrostaticModel) = model.architecture + ##### ##### Recursive util for building NamedTuples of boundary conditions from NamedTuples of fields ##### diff --git a/src/Models/NonhydrostaticModels/pressure_field.jl b/src/Models/NonhydrostaticModels/pressure_field.jl new file mode 100644 index 0000000000..804a84c9a2 --- /dev/null +++ b/src/Models/NonhydrostaticModels/pressure_field.jl @@ -0,0 +1,3 @@ +PressureField(model, data=model.pressures.pHY′.data; kw...) = + Field(model.pressures.pHY′ + model.pressures.pNHS; data, kw...) + diff --git a/src/Oceananigans.jl b/src/Oceananigans.jl index d209752550..361381119e 100644 --- a/src/Oceananigans.jl +++ b/src/Oceananigans.jl @@ -178,8 +178,8 @@ include("Logger.jl") include("Operators/Operators.jl") include("BoundaryConditions/BoundaryConditions.jl") include("Fields/Fields.jl") -include("Advection/Advection.jl") include("AbstractOperations/AbstractOperations.jl") +include("Advection/Advection.jl") include("Solvers/Solvers.jl") include("Distributed/Distributed.jl") @@ -203,6 +203,7 @@ include("Simulations/Simulations.jl") # Abstractions for distributed and multi-region models include("CubedSpheres/CubedSpheres.jl") + ##### ##### Needed so we can export names from sub-modules at the top-level ##### diff --git a/src/Solvers/preconditioned_conjugate_gradient_solver.jl b/src/Solvers/preconditioned_conjugate_gradient_solver.jl index 9a2b62cfa7..862f62c5be 100644 --- a/src/Solvers/preconditioned_conjugate_gradient_solver.jl +++ b/src/Solvers/preconditioned_conjugate_gradient_solver.jl @@ -3,6 +3,8 @@ using Oceananigans.Grids: interior_parent_indices using Statistics: norm, dot using LinearAlgebra +import Oceananigans.Architectures: architecture + mutable struct PreconditionedConjugateGradientSolver{A, G, L, T, F, M, P} architecture :: A grid :: G @@ -18,6 +20,8 @@ mutable struct PreconditionedConjugateGradientSolver{A, G, L, T, F, M, P} preconditioner_product :: P end +architecture(solver::PreconditionedConjugateGradientSolver) = solver.architecture + no_precondition!(args...) = nothing initialize_precondition_product(precondition, template_field) = similar(template_field) diff --git a/src/TimeSteppers/runge_kutta_3.jl b/src/TimeSteppers/runge_kutta_3.jl index 8f3b8eb072..68d47a88bc 100644 --- a/src/TimeSteppers/runge_kutta_3.jl +++ b/src/TimeSteppers/runge_kutta_3.jl @@ -1,3 +1,4 @@ +using Oceananigans.Architectures: architecture using Oceananigans: fields """ @@ -51,7 +52,7 @@ end ##### """ - time_step!(model::AbstractModel{<:RungeKutta3TimeStepper}, Δt; euler=false) + time_step!(model::AbstractModel{<:RungeKutta3TimeStepper}, Δt) Step forward `model` one time step `Δt` with a 3rd-order Runge-Kutta method. The 3rd-order Runge-Kutta method takes three intermediate substep stages to diff --git a/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl b/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl index 271185b1fb..2279b2f68e 100644 --- a/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl +++ b/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl @@ -110,12 +110,12 @@ function implicit_diffusion_solver(::VerticallyImplicitTimeDiscretization, arch, topo[3] == Periodic && error("VerticallyImplicitTimeDiscretization can only be specified on " * "grids that are Bounded in the z-direction.") - z_center_solver = BatchedTridiagonalSolver(arch, grid; + z_center_solver = BatchedTridiagonalSolver(grid; lower_diagonal = ivd_lower_diagonalᵃᵃᶜ, diagonal = ivd_diagonalᵃᵃᶜ, upper_diagonal = ivd_upper_diagonalᵃᵃᶜ) - z_face_solver = BatchedTridiagonalSolver(arch, grid; + z_face_solver = BatchedTridiagonalSolver(grid; lower_diagonal = ivd_lower_diagonalᵃᵃᶠ, diagonal = ivd_diagonalᵃᵃᶠ, upper_diagonal = ivd_upper_diagonalᵃᵃᶠ, diff --git a/test/dependencies_for_runtests.jl b/test/dependencies_for_runtests.jl index c094d8a97e..4bec3a30ec 100644 --- a/test/dependencies_for_runtests.jl +++ b/test/dependencies_for_runtests.jl @@ -71,6 +71,8 @@ closures = ( CUDA.allowscalar(true) +float_types = (Float32, Float64) + include("data_dependencies.jl") include("utils_for_runtests.jl") diff --git a/test/test_dynamics.jl b/test/test_dynamics.jl index c006f8a39d..f6494a8fdc 100644 --- a/test/test_dynamics.jl +++ b/test/test_dynamics.jl @@ -1,3 +1,5 @@ +include("dependencies_for_runtests.jl") + using Oceananigans.TurbulenceClosures: ExplicitTimeDiscretization, VerticallyImplicitTimeDiscretization, z_viscosity using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBoundary @@ -21,8 +23,10 @@ function test_diffusion_simple(fieldname, timestepper, time_discretization) value = π interior(field) .= value + update_state!(model) + for n in 1:10 - ab2_or_rk3_time_step!(model, 1, n) + time_step!(model, 1) end field_data = interior(field) @@ -53,9 +57,11 @@ end function test_diffusion_budget(fieldname, field, model, κ, Δ, order=2) init_mean = mean(interior(field)) + update_state!(model) + for n in 1:10 # Very small time-steps required to bring error under machine precision - ab2_or_rk3_time_step!(model, 1e-4 * Δ^order / κ, n) + time_step!(model, 1e-4 * Δ^order / κ) end final_mean = mean(interior(field)) @@ -84,8 +90,11 @@ function test_diffusion_cosine(fieldname, timestepper, grid, time_discretization # Step forward with small time-step relative to diff. time-scale Δt = 1e-6 * grid.Lz^2 / κ + + update_state!(model) + for n in 1:10 - ab2_or_rk3_time_step!(model, Δt, n) + time_step!(model, Δt) end numerical = interior(field) @@ -124,7 +133,7 @@ function passive_tracer_advection_test(timestepper; N=128, κ=1e-12, Nt=100, bac set!(model, u=u₀, v=v₀, T=T₀) for n in 1:Nt - ab2_or_rk3_time_step!(model, Δt, n) + time_step!(model, Δt) end # Error tolerance is a bit arbitrary @@ -162,7 +171,7 @@ function taylor_green_vortex_test(arch, timestepper, time_discretization; FT=Flo set!(model, u=u₀, v=v₀) for n in 1:Nt - ab2_or_rk3_time_step!(model, Δt, n) + time_step!(model, Δt) end xF, yC, zC = nodes(model.velocities.u, reshape=true) diff --git a/test/test_hydrostatic_free_surface_models.jl b/test/test_hydrostatic_free_surface_models.jl index 9310a25abe..817b0f1a1a 100644 --- a/test/test_hydrostatic_free_surface_models.jl +++ b/test/test_hydrostatic_free_surface_models.jl @@ -55,7 +55,6 @@ function hydrostatic_free_surface_model_tracers_and_forcings_work(arch) return nothing end - topo_1d = (Flat, Flat, Bounded) topos_2d = ((Periodic, Flat, Bounded), @@ -224,8 +223,8 @@ topos_3d = ((Periodic, Periodic, Bounded), u₀, v₀ = 0.1, 0.2 - U = Field(Face, Center, Center, arch, grid) - V = Field(Center, Face, Center, arch, grid) + U = Field{Face, Center, Center}(grid) + V = Field{Center, Face, Center}(grid) CUDA.@allowscalar begin parent(U)[2, 1, 1] = u₀ diff --git a/test/test_time_stepping.jl b/test/test_time_stepping.jl index b6c9eb2d1c..2978281fe0 100644 --- a/test/test_time_stepping.jl +++ b/test/test_time_stepping.jl @@ -111,8 +111,9 @@ function incompressible_in_time(grid, Nt, timestepper) # Just add a temperature perturbation so we get some velocity field. CUDA.@allowscalar interior(model.tracers.T)[8:24, 8:24, 8:24] .+= 0.01 + update_state!(model) for n in 1:Nt - ab2_or_rk3_time_step!(model, 0.05, n) + time_step!(model, 0.05) end arch = architecture(grid) @@ -164,8 +165,9 @@ function tracer_conserved_in_channel(arch, FT, Nt) Tavg0 = CUDA.@allowscalar mean(interior(model.tracers.T)) + update_state!(model) for n in 1:Nt - ab2_or_rk3_time_step!(model, 600, n) + time_step!(model, 600) end Tavg = CUDA.@allowscalar mean(interior(model.tracers.T)) diff --git a/test/utils_for_runtests.jl b/test/utils_for_runtests.jl index 3e5af8bdcc..0353e49fc0 100644 --- a/test/utils_for_runtests.jl +++ b/test/utils_for_runtests.jl @@ -1,19 +1,15 @@ using Oceananigans - using Statistics using KernelAbstractions: @kernel, @index, Event using CUDA using Test using Printf - using Test - using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper, RungeKutta3TimeStepper, update_state! import Oceananigans.Fields: interior test_architectures() = CUDA.has_cuda() ? tuple(GPU()) : tuple(CPU()) -float_types = (Float32, Float64) function summarize_regression_test(fields, correct_fields) for (field_name, φ, φ_c) in zip(keys(fields), fields, correct_fields) @@ -60,22 +56,6 @@ end ##### Useful utilities ##### -const AB2Model = NonhydrostaticModel{<:QuasiAdamsBashforth2TimeStepper} -const RK3Model = NonhydrostaticModel{<:RungeKutta3TimeStepper} - -# For time-stepping without a Simulation -function ab2_or_rk3_time_step!(model::AB2Model, Δt, n) - n == 1 && update_state!(model) - time_step!(model, Δt, euler=n==1) - return nothing -end - -function ab2_or_rk3_time_step!(model::RK3Model, Δt, n) - n == 1 && update_state!(model) - time_step!(model, Δt) - return nothing -end - interior(a, grid) = view(a, grid.Hx+1:grid.Nx+grid.Hx, grid.Hy+1:grid.Ny+grid.Hy, grid.Hz+1:grid.Nz+grid.Hz) From 983d2688e9655645b786b5af1e64b98de2aa6e16 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 30 Dec 2021 06:43:50 -0700 Subject: [PATCH 011/140] Widespread bugfixes --- src/AbstractOperations/AbstractOperations.jl | 2 + src/AbstractOperations/computed_field.jl | 12 +- .../metric_field_reductions.jl} | 14 +- .../distributed_fft_based_poisson_solver.jl | 21 ++- src/Fields/Fields.jl | 13 +- src/Fields/computed_field.jl | 112 ------------ src/Fields/field.jl | 3 +- src/Fields/reduced_field.jl | 120 ------------ src/Fields/reduced_getindex_setindex.jl | 39 ---- src/Grids/latitude_longitude_grid.jl | 2 +- .../grid_fitted_immersed_boundaries.jl | 4 +- src/Oceananigans.jl | 2 +- src/OutputWriters/netcdf_output_writer.jl | 12 +- test/runtests.jl | 44 +++-- test/test_buoyancy.jl | 14 +- test/test_cubed_sphere_circulation.jl | 10 +- test/test_distributed_poisson_solvers.jl | 10 +- test/test_dynamics.jl | 173 +++++++----------- ...ced_fields.jl => test_field_reductions.jl} | 48 ++--- ...static_free_surface_immersed_boundaries.jl | 6 +- test/test_turbulence_closures.jl | 22 +-- 21 files changed, 189 insertions(+), 494 deletions(-) rename src/{Fields/metric_reductions.jl => AbstractOperations/metric_field_reductions.jl} (57%) delete mode 100644 src/Fields/computed_field.jl delete mode 100644 src/Fields/reduced_field.jl delete mode 100644 src/Fields/reduced_getindex_setindex.jl rename test/{test_reduced_fields.jl => test_field_reductions.jl} (76%) diff --git a/src/AbstractOperations/AbstractOperations.jl b/src/AbstractOperations/AbstractOperations.jl index a8eb449607..833cdb7edb 100644 --- a/src/AbstractOperations/AbstractOperations.jl +++ b/src/AbstractOperations/AbstractOperations.jl @@ -1,6 +1,7 @@ module AbstractOperations export ∂x, ∂y, ∂z, @at, @unary, @binary, @multiary +export AveragedField, IntegratedField export KernelFunctionOperation using Base: @propagate_inbounds @@ -50,6 +51,7 @@ at(loc, f) = f # fallback include("grid_validation.jl") include("grid_metrics.jl") +include("metric_field_reductions.jl") include("unary_operations.jl") include("binary_operations.jl") include("multiary_operations.jl") diff --git a/src/AbstractOperations/computed_field.jl b/src/AbstractOperations/computed_field.jl index 179fe1ed40..2833332c1b 100644 --- a/src/AbstractOperations/computed_field.jl +++ b/src/AbstractOperations/computed_field.jl @@ -3,6 +3,8 @@ ##### using KernelAbstractions: @kernel, @index +using Oceananigans.Fields: FieldStatus +using Oceananigans.Utils: work_layout import Oceananigans: short_show import Oceananigans.Fields: Field, compute! @@ -31,17 +33,19 @@ recompute_safely (Bool): whether or not to _always_ "recompute" `f` if `f` is """ function Field(operand::AbstractOperation; data = nothing, - boundary_conditions = FieldBoundaryConditions(op.grid, location(op)), + boundary_conditions = FieldBoundaryConditions(operand.grid, location(operand)), recompute_safely = true) + grid = operand.grid + if isnothing(data) - data = new_data(op.grid, location(op)) + data = new_data(grid, location(operand)) recompute_safely = false end - status = recompute_safely ? nothing : FieldStatus(0.0) + status = recompute_safely ? nothing : FieldStatus() - return Field(location(op), op.grid, data, boundary_conditions, operand, status) + return Field(location(operand), grid, data, boundary_conditions, operand, status) end """ diff --git a/src/Fields/metric_reductions.jl b/src/AbstractOperations/metric_field_reductions.jl similarity index 57% rename from src/Fields/metric_reductions.jl rename to src/AbstractOperations/metric_field_reductions.jl index ac8a6e92dc..90acd4c0ff 100644 --- a/src/Fields/metric_reductions.jl +++ b/src/AbstractOperations/metric_field_reductions.jl @@ -4,14 +4,16 @@ import Oceananigans.Fields: Reduction ##### Metric inference ##### -reduction_grid_metric(dims::Number) reduction_grid_metric(tuple(dims)) +reduction_grid_metric(dims::Number) = reduction_grid_metric(tuple(dims)) + reduction_grid_metric(dims) = dims === tuple(1) ? Δx : dims === tuple(2) ? Δy : dims === tuple(3) ? Δz : dims === (1, 2) ? Az : dims === (1, 3) ? Ay : dims === (2, 3) ? Ax : - dims === (1, 2, 3) ? volume + dims === (1, 2, 3) ? volume : + throw(ArgumentError("Cannot determine grid metric for reducing over dims = $dims")) ##### ##### Metric reductions @@ -24,13 +26,15 @@ struct Integral <: AbstractMetricReduction end reduction(::Average) = mean! reduction(::Integral) = sum! -function Reduction(r::AbstractMetricReductions, operand; dims) +function Reduction(r::AbstractMetricReduction, operand; dims) dx = reduction_grid_metric(dims) field_dx = field * dx return Reduction(reduction(r), field_dx; dims) end # Convenience -Average(field::AbstractField; dims) = Reduction(Average(), field; dims) -Integral(field::AbstractField; dims) = Reduction(Integral(), field; dims) +Average(field::AbstractField; dims=:) = Reduction(Average(), field; dims) +Integral(field::AbstractField; dims=:) = Reduction(Integral(), field; dims) +const AveragedField = Field{<:Any, <:Any, <:Any, <:Reduction{<:Average}} +const IntegratedField = Field{<:Any, <:Any, <:Any, <:Reduction{<:Integral}} diff --git a/src/Distributed/distributed_fft_based_poisson_solver.jl b/src/Distributed/distributed_fft_based_poisson_solver.jl index 5770db79ed..900dcc4b66 100644 --- a/src/Distributed/distributed_fft_based_poisson_solver.jl +++ b/src/Distributed/distributed_fft_based_poisson_solver.jl @@ -1,20 +1,25 @@ import PencilFFTs -import Oceananigans.Solvers: poisson_eigenvalues, solve! using Oceananigans.Solvers: copy_real_component! using Oceananigans.Distributed: rank2index -struct DistributedFFTBasedPoissonSolver{A, P, F, L, λ, S} - architecture :: A +import Oceananigans.Solvers: poisson_eigenvalues, solve! +import Oceananigans.Architectures: architecture + +struct DistributedFFTBasedPoissonSolver{P, F, L, λ, S} plan :: P global_grid :: F - my_grid :: L + local_grid :: L eigenvalues :: λ storage :: S end -function DistributedFFTBasedPoissonSolver(arch, global_grid, local_grid) +architecture(solver::DistributedFFTBasedPoissonSolver) = + architecture(solver.global_grid) + +function DistributedFFTBasedPoissonSolver(global_grid, local_grid) + arch = architecture(global_grid) topo = (TX, TY, TZ) = topology(global_grid) λx = poisson_eigenvalues(global_grid.Nx, global_grid.Lx, 1, TX()) @@ -41,11 +46,11 @@ function DistributedFFTBasedPoissonSolver(arch, global_grid, local_grid) plan = PencilFFTs.PencilFFTPlan(size(global_grid), transform, proc_dims, MPI.COMM_WORLD) storage = PencilFFTs.allocate_input(plan) - return DistributedFFTBasedPoissonSolver(arch, plan, global_grid, local_grid, eigenvalues, storage) + return DistributedFFTBasedPoissonSolver(plan, global_grid, local_grid, eigenvalues, storage) end function solve!(x, solver::DistributedFFTBasedPoissonSolver) - arch = solver.architecture + arch = architecture(solver) λx, λy, λz = solver.eigenvalues # Apply forward transforms. @@ -68,7 +73,7 @@ function solve!(x, solver::DistributedFFTBasedPoissonSolver) solver.plan \ solver.storage xc_transposed = first(solver.storage) - copy_event = launch!(arch, solver.my_grid, :xyz, copy_real_component!, x, xc_transposed, dependencies=device_event(arch)) + copy_event = launch!(arch, solver.local_grid, :xyz, copy_real_component!, x, xc_transposed, dependencies=device_event(arch)) wait(device(arch), copy_event) return x diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl index b4333dca9b..9bd990d1e4 100644 --- a/src/Fields/Fields.jl +++ b/src/Fields/Fields.jl @@ -1,13 +1,11 @@ module Fields export Face, Center -export AbstractField, AbstractDataField, Field +export AbstractField, Field, Average, Integral, Reduction export CenterField, XFaceField, YFaceField, ZFaceField -export ReducedField, AveragedField, ComputedField, KernelComputedField, BackgroundField -export interior, data -export xnode, ynode, znode, location -export regrid! -export set!, compute!, @compute +export BackgroundField +export interior, data, xnode, ynode, znode, location +export set!, compute!, @compute, regrid! export VelocityFields, TracerFields, tracernames, PressureFields, TendencyFields export interpolate, FieldSlicer @@ -16,13 +14,10 @@ using Oceananigans.Grids using Oceananigans.BoundaryConditions include("abstract_field.jl") -# include("reduced_getindex_setindex.jl") include("field.jl") include("field_reductions.jl") include("zero_field.jl") -# include("reduced_field.jl") include("averaged_field.jl") -# include("computed_field.jl") include("kernel_computed_field.jl") include("function_field.jl") include("regridding_fields.jl") diff --git a/src/Fields/computed_field.jl b/src/Fields/computed_field.jl deleted file mode 100644 index 0ac16a21a0..0000000000 --- a/src/Fields/computed_field.jl +++ /dev/null @@ -1,112 +0,0 @@ -using Adapt -using Statistics -using KernelAbstractions: @kernel, @index, Event -using Oceananigans.Grids -using Oceananigans.BoundaryConditions: fill_halo_regions! - -struct ComputedField{X, Y, Z, S, O, A, D, G, T, C} <: AbstractDataField{X, Y, Z, A, G, T, 3} - data :: D - architecture :: A - grid :: G - operand :: O - boundary_conditions :: C - status :: S - - function ComputedField{X, Y, Z}(data::D, - arch::A, - grid::G, - operand::O, - boundary_conditions::C; - recompute_safely=true) where {X, Y, Z, D, A, G, O, C} - - validate_field_data(X, Y, Z, data, grid) - - # Use FieldStatus if we want to avoid always recomputing - status = recompute_safely ? nothing : FieldStatus() - - S = typeof(status) - T = eltype(grid) - - return new{X, Y, Z, S, O, A, D, G, T, C}(data, arch, grid, operand, boundary_conditions, status) - end -end - -""" - ComputedField(operand [, arch=nothing]; data = nothing, recompute_safely = true, - boundary_conditions = ComputedFieldBoundaryConditions(operand.grid, location(operand)) - -Return a field whose data is `computed` from `operand`. If `arch`itecture is not supplied it -is inferred from `operand`. - -If the keyword argument `data` is not provided, memory is allocated to store -the result. The `arch`itecture of `data` is inferred from `operand`. - -If `data` is provided and `recompute_safely=false`, then "recomputation" of the `ComputedField` -is avoided if possible. -""" -function ComputedField(operand, arch=nothing; kwargs...) - - loc = location(operand) - grid = operand.grid - - return ComputedField(loc..., operand, arch, grid; kwargs...) -end - -function ComputedField(LX, LY, LZ, operand, arch, grid; - data = nothing, - recompute_safely = true, - boundary_conditions = FieldBoundaryConditions(grid, (LX, LY, LZ))) - - # Architecturanigans - operand_arch = architecture(operand) - arch = isnothing(operand_arch) ? arch : operand_arch - isnothing(arch) && error("The architecture must be provided, or inferrable from `operand`!") - - if isnothing(data) - data = new_data(arch, grid, (LX, LY, LZ)) - recompute_safely = false - end - - return ComputedField{LX, LY, LZ}(data, arch, grid, operand, boundary_conditions; recompute_safely=recompute_safely) -end - -""" - compute!(comp::ComputedField) - -Compute `comp.operand` and store the result in `comp.data`. -""" -function compute!(comp::ComputedField{LX, LY, LZ}, time=nothing) where {LX, LY, LZ} - compute_at!(comp.operand, time) # ensures any 'dependencies' of the computation are computed first - - arch = architecture(comp) - - workgroup, worksize = work_layout(comp.grid, - :xyz, - include_right_boundaries=true, - location=(LX, LY, LZ)) - - compute_kernel! = _compute!(device(arch), workgroup, worksize) - - event = compute_kernel!(comp.data, comp.operand; dependencies=Event(device(arch))) - - wait(device(arch), event) - - fill_halo_regions!(comp, arch) - - return nothing -end - -compute_at!(field::ComputedField{LX, LY, LZ, <:FieldStatus}, time) where {LX, LY, LZ} = - conditional_compute!(field, time) - -"""Compute an `operand` and store in `data`.""" -@kernel function _compute!(data, operand) - i, j, k = @index(Global, NTuple) - @inbounds data[i, j, k] = operand[i, j, k] -end - -##### -##### Adapt -##### - -Adapt.adapt_structure(to, computed_field::ComputedField) = Adapt.adapt(to, computed_field.data) diff --git a/src/Fields/field.jl b/src/Fields/field.jl index 443afefa19..95110fa89b 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -24,7 +24,6 @@ end function Field(loc::Tuple, grid::AbstractGrid, data, bcs, op, status) validate_field_data(loc, data, grid) # validate_boundary_conditions(loc, grid, bcs) - arch = architecture(grid) LX, LY, LZ = loc return Field{LX, LY, LZ}(grid, data, bcs, op, status) end @@ -163,7 +162,7 @@ ZFaceField(grid::AbstractGrid, T::DataType=eltype(grid); kw...) = Field{Center, Computes `field.data` from `field.operand`. """ -compute!(field) = nothing # fallback +compute!(field, time=nothing) = nothing # fallback """ @compute(exprs...) diff --git a/src/Fields/reduced_field.jl b/src/Fields/reduced_field.jl deleted file mode 100644 index 3a3e80b396..0000000000 --- a/src/Fields/reduced_field.jl +++ /dev/null @@ -1,120 +0,0 @@ -using Adapt -using Statistics - -import Oceananigans.BoundaryConditions: fill_halo_regions! - -##### -##### AbstractReducedField stuff -##### - -abstract type AbstractReducedField{X, Y, Z, A, G, T, N} <: AbstractDataField{X, Y, Z, A, G, T, 3} end - -const ARF = AbstractReducedField - -fill_halo_regions!(field::AbstractReducedField, arch, args...) = - fill_halo_regions!(field.data, field.boundary_conditions, arch, field.grid, args...; reduced_dimensions=field.dims) - -const DimsType = NTuple{N, Int} where N - -function validate_reduced_dims(dims) - dims = Tuple(dims) - - # Check type - dims isa DimsType || error("Reduced dims must be an integer or tuple of integers.") - - # Check length - length(dims) > 3 && error("Models are 3-dimensional. Cannot reduce over 4+ dimensions.") - - # Check values - all(1 <= d <= 3 for d in dims) || error("Dimensions must be one of 1, 2, 3.") - - return dims -end - -function validate_reduced_locations(X, Y, Z, dims) - loc = (X, Y, Z) - - # Check reduced locations - for i in dims - isnothing(loc[i]) || ArgumentError("The location of reduced dimensions must be Nothing") - end - - return nothing -end - -##### -##### Concrete ReducedField -##### - -struct ReducedField{X, Y, Z, A, D, G, T, N, B} <: AbstractReducedField{X, Y, Z, A, G, T, N} - data :: D - architecture :: A - grid :: G - dims :: NTuple{N, Int} - boundary_conditions :: B - - @doc """ - ReducedField{X, Y, Z}(data, grid, dims) - - Returns a `ReducedField` at location `(X, Y, Z)` with `data` on `grid` - that is reduced over the dimensions in `dims`. - """ - function ReducedField{X, Y, Z}(data::D, arch::A, grid::G, dims, bcs::B) where {X, Y, Z, A, D, G, B} - - dims = validate_reduced_dims(dims) - validate_reduced_locations(X, Y, Z, dims) - validate_field_data(X, Y, Z, data, grid) - - N = length(dims) - T = eltype(grid) - - return new{X, Y, Z, A, D, G, T, N, B}(data, arch, grid, dims, bcs) - end -end - -""" - ReducedField(X, Y, Z, arch, grid; dims, data=nothing, boundary_conditions=nothing) - -Returns a `ReducedField` reduced over `dims` on `grid` and `arch`itecture with `boundary_conditions`. - -The location `(X, Y, Z)` may be the parent, three-dimension location or the reduced location. - -If `data` is specified, it should be an `OffsetArray` with singleton reduced dimensions; -otherwise `data` is allocated. - -If `boundary_conditions` are not provided, default boundary conditions are constructed -using the reduced location. -""" -function ReducedField(FT::DataType, Xr, Yr, Zr, arch, grid::AbstractGrid; dims, data=nothing, - boundary_conditions=nothing) - - dims = validate_reduced_dims(dims) - - # Reduce non-reduced dimensions - X, Y, Z = reduced_location((Xr, Yr, Zr); dims=dims) - - if isnothing(data) - data = new_data(FT, arch, grid, (X, Y, Z)) - end - - if isnothing(boundary_conditions) - boundary_conditions = FieldBoundaryConditions(grid, (X, Y, Z)) - end - - return ReducedField{X, Y, Z}(data, arch, grid, dims, boundary_conditions) -end - -ReducedField(Xr, Yr, Zr, arch, grid::AbstractGrid; dims, kw...) = ReducedField(eltype(grid), Xr, Yr, Zr, arch, grid; dims, kw...) -ReducedField(Lr::Tuple, arch, grid::AbstractGrid; dims, kw...) = ReducedField(Lr..., arch, grid; dims, kw...) -ReducedField(FT::DataType, Lr::Tuple, arch, grid::AbstractGrid; dims, kw...) = ReducedField(FT, Lr..., arch, grid; dims, kw...) - -# Canonical `similar` for AbstractReducedField -Base.similar(r::AbstractReducedField{X, Y, Z, Arch}) where {X, Y, Z, Arch} = - ReducedField(X, Y, Z, Arch(), r.grid; dims=r.dims, boundary_conditions=r.boundary_conditions) - -##### -##### ReducedField utils -##### - -reduced_location(loc; dims) = Tuple(i ∈ dims ? Nothing : loc[i] for i in 1:3) - diff --git a/src/Fields/reduced_getindex_setindex.jl b/src/Fields/reduced_getindex_setindex.jl deleted file mode 100644 index b9f591eed8..0000000000 --- a/src/Fields/reduced_getindex_setindex.jl +++ /dev/null @@ -1,39 +0,0 @@ -const ADF = AbstractDataField - -# Two-dimensional fields -@inline Base.getindex( r::ADF{Nothing, Y, Z}, i, j, k) where {Y, Z} = @inbounds r.data[1, j, k] -@inline Base.setindex!(r::ADF{Nothing, Y, Z}, d, i, j, k) where {Y, Z} = @inbounds r.data[1, j, k] = d -@inline Base.getindex( r::ADF{X, Nothing, Z}, i, j, k) where {X, Z} = @inbounds r.data[i, 1, k] -@inline Base.setindex!(r::ADF{X, Nothing, Z}, d, i, j, k) where {X, Z} = @inbounds r.data[i, 1, k] = d -@inline Base.getindex( r::ADF{X, Y, Nothing}, i, j, k) where {X, Y} = @inbounds r.data[i, j, 1] -@inline Base.setindex!(r::ADF{X, Y, Nothing}, d, i, j, k) where {X, Y} = @inbounds r.data[i, j, 1] = d - -@inline Base.getindex( r::ADF{Nothing, Y, Z}, j, k) where {Y, Z} = @inbounds r.data[1, j, k] -@inline Base.setindex!(r::ADF{Nothing, Y, Z}, d, j, k) where {Y, Z} = @inbounds r.data[1, j, k] = d -@inline Base.getindex( r::ADF{X, Nothing, Z}, i, k) where {X, Z} = @inbounds r.data[i, 1, k] -@inline Base.setindex!(r::ADF{X, Nothing, Z}, d, i, k) where {X, Z} = @inbounds r.data[i, 1, k] = d -@inline Base.getindex( r::ADF{X, Y, Nothing}, i, j) where {X, Y} = @inbounds r.data[i, j, 1] -@inline Base.setindex!(r::ADF{X, Y, Nothing}, d, i, j) where {X, Y} = @inbounds r.data[i, j, 1] = d - -# One-dimensional fields -@inline Base.getindex( r::ADF{X, Nothing, Nothing}, i, j, k) where X = @inbounds r.data[i, 1, 1] -@inline Base.setindex!(r::ADF{X, Nothing, Nothing}, d, i, j, k) where X = @inbounds r.data[i, 1, 1] = d -@inline Base.getindex( r::ADF{Nothing, Y, Nothing}, i, j, k) where Y = @inbounds r.data[1, j, 1] -@inline Base.setindex!(r::ADF{Nothing, Y, Nothing}, d, i, j, k) where Y = @inbounds r.data[1, j, 1] = d -@inline Base.getindex( r::ADF{Nothing, Nothing, Z}, i, j, k) where Z = @inbounds r.data[1, 1, k] -@inline Base.setindex!(r::ADF{Nothing, Nothing, Z}, d, i, j, k) where Z = @inbounds r.data[1, 1, k] = d - -@inline Base.getindex( r::ADF{X, Nothing, Nothing}, i) where X = @inbounds r.data[i, 1, 1] -@inline Base.setindex!(r::ADF{X, Nothing, Nothing}, d, i) where X = @inbounds r.data[i, 1, 1] = d -@inline Base.getindex( r::ADF{Nothing, Y, Nothing}, j) where Y = @inbounds r.data[1, j, 1] -@inline Base.setindex!(r::ADF{Nothing, Y, Nothing}, d, j) where Y = @inbounds r.data[1, j, 1] = d -@inline Base.getindex( r::ADF{Nothing, Nothing, Z}, k) where Z = @inbounds r.data[1, 1, k] -@inline Base.setindex!(r::ADF{Nothing, Nothing, Z}, d, k) where Z = @inbounds r.data[1, 1, k] = d - -# Zero-dimensional fields -@inline Base.getindex( r::ADF{Nothing, Nothing, Nothing}, i, j, k) = @inbounds r.data[1, 1, 1] -@inline Base.setindex!(r::ADF{Nothing, Nothing, Nothing}, d, i, j, k) = @inbounds r.data[1, 1, 1] = d - -@inline Base.getindex( r::ADF{Nothing, Nothing, Nothing} ) = @inbounds r.data[1, 1, 1] -@inline Base.setindex!(r::ADF{Nothing, Nothing, Nothing}, d) = @inbounds r.data[1, 1, 1] = d - diff --git a/src/Grids/latitude_longitude_grid.jl b/src/Grids/latitude_longitude_grid.jl index 41191be2ba..903d5f7f0c 100644 --- a/src/Grids/latitude_longitude_grid.jl +++ b/src/Grids/latitude_longitude_grid.jl @@ -113,7 +113,7 @@ function LatitudeLongitudeGrid(architecture=CPU(), Δλᶠᵃᵃ, Δλᶜᵃᵃ, λᶠᵃᵃ, λᶜᵃᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ, φᵃᶠᵃ, φᵃᶜᵃ, Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ, - (nothing for i=1:10)..., radius) + (nothing for i=1:10)..., FT(radius)) return !precompute_metrics ? preliminary_grid : with_precomputed_metrics(preliminary_grid) end diff --git a/src/ImmersedBoundaries/grid_fitted_immersed_boundaries.jl b/src/ImmersedBoundaries/grid_fitted_immersed_boundaries.jl index 79697c4f94..94e918c339 100644 --- a/src/ImmersedBoundaries/grid_fitted_immersed_boundaries.jl +++ b/src/ImmersedBoundaries/grid_fitted_immersed_boundaries.jl @@ -1,6 +1,6 @@ using Adapt using CUDA: CuArray -using Oceananigans.Fields: ReducedField, fill_halo_regions! +using Oceananigans.Fields: fill_halo_regions! using Oceananigans.Architectures: arch_array abstract type AbstractGridFittedBoundary <: AbstractImmersedBoundary end @@ -44,7 +44,7 @@ const CuArrayGridFittedBottom = GridFittedBottom{<:CuArray} function ImmersedBoundaryGrid(grid, ib::Union{ArrayGridFittedBottom, CuArrayGridFittedBottom}) # Wrap bathymetry in an OffsetArray with halos arch = grid.architecture - bottom_field = ReducedField(Center, Center, Nothing, arch, grid; dims=3) + bottom_field = Field{Center, Center, Nothing}(grid) bottom_data = arch_array(arch, ib.bottom) bottom_field .= bottom_data fill_halo_regions!(bottom_field, arch) diff --git a/src/Oceananigans.jl b/src/Oceananigans.jl index 361381119e..b12a8f8a84 100644 --- a/src/Oceananigans.jl +++ b/src/Oceananigans.jl @@ -33,7 +33,7 @@ export # Fields and field manipulation Field, CenterField, XFaceField, YFaceField, ZFaceField, - AveragedField, ComputedField, KernelComputedField, BackgroundField, + Average, Integral, Reduction, BackgroundField, interior, set!, compute!, regrid!, # Forcing functions diff --git a/src/OutputWriters/netcdf_output_writer.jl b/src/OutputWriters/netcdf_output_writer.jl index 8ef1deb3c0..20eea7dfd3 100644 --- a/src/OutputWriters/netcdf_output_writer.jl +++ b/src/OutputWriters/netcdf_output_writer.jl @@ -10,7 +10,7 @@ using Oceananigans.Grids: topology, halo_size, all_x_nodes, all_y_nodes, all_z_n using Oceananigans.Utils: versioninfo_with_gpu, oceananigans_versioninfo using Oceananigans.TimeSteppers: float_or_date_time using Oceananigans.Diagnostics: WindowedSpatialAverage -using Oceananigans.Fields: reduced_location, location, FieldSlicer, parent_slice_indices +using Oceananigans.Fields: reduced_dimensions, reduced_location, location, FieldSlicer, parent_slice_indices dictify(outputs) = outputs dictify(outputs::NamedTuple) = Dict(string(k) => dictify(v) for (k, v) in zip(keys(outputs), values(outputs))) @@ -141,7 +141,7 @@ end Construct a `NetCDFOutputWriter` that writes `(label, output)` pairs in `outputs` (which should be a `Dict`) to a NetCDF file, where `label` is a string that labels the output and `output` is -either a `Field` (e.g. `model.velocities.u` or an `AveragedField`) or a function `f(model)` that +either a `Field` (e.g. `model.velocities.u`) or a function `f(model)` that returns something to be written to disk. Custom output requires the spatial `dimensions` (a `Dict`) to be manually specified (see examples). @@ -415,7 +415,7 @@ Base.close(nc::NetCDFOutputWriter) = close(nc.dataset) function save_output!(ds, output, model, ow, time_index, name) data = fetch_and_convert_output(output, model, ow) - data = drop_averaged_dims(output, data) + data = drop_output_dims(output, data) colons = Tuple(Colon() for _ in 1:ndims(data)) ds[name][colons..., time_index] = data @@ -478,9 +478,9 @@ function write_output!(ow::NetCDFOutputWriter, model) return nothing end -drop_averaged_dims(output, data) = data # fallback -drop_averaged_dims(output::AveragedField, data) = dropdims(data, dims=output.dims) -drop_averaged_dims(output::WindowedTimeAverage{<:AveragedField}, data) = dropdims(data, dims=output.operand.dims) +drop_output_dims(output, data) = data # fallback +drop_output_dims(output::Field, data) = dropdims(data, dims=reduced_dimensions(output)) +drop_output_dims(output::WindowedTimeAverage{<:Field}, data) = dropdims(data, dims=reduced_dimensions(output.operand)) ##### ##### Show diff --git a/test/runtests.jl b/test/runtests.jl index 1cfcf56eaf..fc75cc11d8 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,13 +1,14 @@ include("dependencies_for_runtests.jl") @testset "Oceananigans" begin + # Core Oceananigans if group == :unit || group == :all @testset "Unit tests" begin include("test_grids.jl") include("test_operators.jl") include("test_boundary_conditions.jl") include("test_field.jl") - # include("test_reduced_fields.jl") + include("test_field_reductions.jl") include("test_halo_regions.jl") include("test_coriolis.jl") include("test_buoyancy.jl") @@ -17,6 +18,14 @@ include("dependencies_for_runtests.jl") end end + if group == :abstract_operations || group == :all + @testset "AbstractOperations and broadcasting tests" begin + include("test_abstract_operations.jl") + include("test_computed_field.jl") + include("test_broadcasting.jl") + end + end + if group == :solvers || group == :all @testset "Solvers" begin include("test_batched_tridiagonal_solver.jl") @@ -25,6 +34,18 @@ include("dependencies_for_runtests.jl") end end + # Simulations + if group == :simulation || group == :all + @testset "Simulation tests" begin + include("test_simulations.jl") + include("test_diagnostics.jl") + include("test_output_writers.jl") + include("test_output_readers.jl") + include("test_lagrangian_particle_tracking.jl") + end + end + + # Models if group == :time_stepping_1 || group == :all @testset "Model and time stepping tests (part 1)" begin include("test_nonhydrostatic_models.jl") @@ -55,25 +76,8 @@ include("dependencies_for_runtests.jl") include("test_implicit_free_surface_solver.jl") end end - - if group == :abstract_operations || group == :all - @testset "AbstractOperations and broadcasting tests" begin - include("test_abstract_operations.jl") - include("test_computed_field.jl") - include("test_broadcasting.jl") - end - end - - if group == :simulation || group == :all - @testset "Simulation tests" begin - include("test_simulations.jl") - include("test_diagnostics.jl") - include("test_output_writers.jl") - include("test_output_readers.jl") - include("test_lagrangian_particle_tracking.jl") - end - end - + + # Model enhancements: cubed sphere, distributed, etc if group == :cubed_sphere || group == :all @testset "Cubed sphere tests" begin include("test_cubed_spheres.jl") diff --git a/test/test_buoyancy.jl b/test/test_buoyancy.jl index 2acdc5b927..97b7c3e65c 100644 --- a/test/test_buoyancy.jl +++ b/test/test_buoyancy.jl @@ -1,3 +1,5 @@ +include("dependencies_for_runtests.jl") + using Oceananigans.Fields: TracerFields using Oceananigans.BuoyancyModels: @@ -17,35 +19,35 @@ end function density_perturbation_works(arch, FT, eos) grid = RectilinearGrid(arch, FT, size=(3, 3, 3), extent=(1, 1, 1)) - C = datatuple(TracerFields((:T, :S), arch, grid)) + C = datatuple(TracerFields((:T, :S), grid)) density_anomaly = ρ′(2, 2, 2, grid, eos, C.T, C.S) return true end function ∂x_b_works(arch, FT, buoyancy) grid = RectilinearGrid(arch, FT, size=(3, 3, 3), extent=(1, 1, 1)) - C = datatuple(TracerFields(required_tracers(buoyancy), arch, grid)) + C = datatuple(TracerFields(required_tracers(buoyancy), grid)) dbdx = ∂x_b(2, 2, 2, grid, buoyancy, C) return true end function ∂y_b_works(arch, FT, buoyancy) grid = RectilinearGrid(arch, FT, size=(3, 3, 3), extent=(1, 1, 1)) - C = datatuple(TracerFields(required_tracers(buoyancy), arch, grid)) + C = datatuple(TracerFields(required_tracers(buoyancy), grid)) dbdy = ∂y_b(2, 2, 2, grid, buoyancy, C) return true end function ∂z_b_works(arch, FT, buoyancy) grid = RectilinearGrid(arch, FT, size=(3, 3, 3), extent=(1, 1, 1)) - C = datatuple(TracerFields(required_tracers(buoyancy), arch, grid)) + C = datatuple(TracerFields(required_tracers(buoyancy), grid)) dbdz = ∂z_b(2, 2, 2, grid, buoyancy, C) return true end function thermal_expansion_works(arch, FT, eos) grid = RectilinearGrid(arch, FT, size=(3, 3, 3), extent=(1, 1, 1)) - C = datatuple(TracerFields((:T, :S), arch, grid)) + C = datatuple(TracerFields((:T, :S), grid)) α = thermal_expansionᶜᶜᶜ(2, 2, 2, grid, eos, C.T, C.S) α = thermal_expansionᶠᶜᶜ(2, 2, 2, grid, eos, C.T, C.S) α = thermal_expansionᶜᶠᶜ(2, 2, 2, grid, eos, C.T, C.S) @@ -55,7 +57,7 @@ end function haline_contraction_works(arch, FT, eos) grid = RectilinearGrid(arch, FT, size=(3, 3, 3), extent=(1, 1, 1)) - C = datatuple(TracerFields((:T, :S), arch, grid)) + C = datatuple(TracerFields((:T, :S), grid)) β = haline_contractionᶜᶜᶜ(2, 2, 2, grid, eos, C.T, C.S) β = haline_contractionᶠᶜᶜ(2, 2, 2, grid, eos, C.T, C.S) β = haline_contractionᶜᶠᶜ(2, 2, 2, grid, eos, C.T, C.S) diff --git a/test/test_cubed_sphere_circulation.jl b/test/test_cubed_sphere_circulation.jl index 23a5208307..578562aa48 100644 --- a/test/test_cubed_sphere_circulation.jl +++ b/test/test_cubed_sphere_circulation.jl @@ -1,9 +1,9 @@ using Oceananigans.Operators: Γᶠᶠᵃ, ζ₃ᶠᶠᵃ function diagnose_velocities_from_streamfunction(ψ, arch, grid) - ψᶠᶠᶜ = Field(Face, Face, Center, arch, grid) - uᶠᶜᶜ = Field(Face, Center, Center, arch, grid) - vᶜᶠᶜ = Field(Center, Face, Center, arch, grid) + ψᶠᶠᶜ = Field{Face, Face, Center}(grid) + uᶠᶜᶜ = Field{Face, Center, Center}(grid) + vᶜᶠᶜ = Field{Center, Face, Center}(grid) for (f, grid_face) in enumerate(grid.faces) Nx, Ny, Nz = size(grid_face) @@ -45,8 +45,8 @@ for arch in archs Nx, Ny, Nz, Nf = size(grid) R = grid.faces[1].radius - u_field = XFaceField(arch, grid) - v_field = YFaceField(arch, grid) + u_field = XFaceField(grid) + v_field = YFaceField(grid) ψ(λ, φ) = R * sind(φ) CUDA.@allowscalar set_velocities_from_streamfunction!(u_field, v_field, ψ, arch, grid) diff --git a/test/test_distributed_poisson_solvers.jl b/test/test_distributed_poisson_solvers.jl index 45ba5cfff3..400f12f407 100644 --- a/test/test_distributed_poisson_solvers.jl +++ b/test/test_distributed_poisson_solvers.jl @@ -1,11 +1,12 @@ +include("dependencies_for_runtests.jl") using Oceananigans.Distributed: reconstruct_global_grid -function random_divergent_source_term(arch, grid) +function random_divergent_source_term(grid) # Generate right hand side from a random (divergent) velocity field. - Ru = XFaceField(arch, grid) - Rv = YFaceField(arch, grid) - Rw = ZFaceField(arch, grid) + Ru = XFaceField(grid) + Rv = YFaceField(grid) + Rw = ZFaceField(grid) U = (u=Ru, v=Rv, w=Rw) Nx, Ny, Nz = size(grid) @@ -13,6 +14,7 @@ function random_divergent_source_term(arch, grid) set!(Rv, (x, y, z) -> rand()) set!(Rw, (x, y, z) -> rand()) + arch = architecture(grid) fill_halo_regions!(Ru, arch) fill_halo_regions!(Rv, arch) fill_halo_regions!(Rw, arch) diff --git a/test/test_dynamics.jl b/test/test_dynamics.jl index f6494a8fdc..4191fc5d97 100644 --- a/test/test_dynamics.jl +++ b/test/test_dynamics.jl @@ -4,65 +4,52 @@ using Oceananigans.TurbulenceClosures: ExplicitTimeDiscretization, VerticallyImp using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBoundary function relative_error(u_num, u, time) - u_ans = Field(location(u_num), architecture(u_num), u_num.grid, nothing) + u_ans = Field(location(u_num), u_num.grid) set!(u_ans, (x, y, z) -> u(x, y, z, time)) return mean((interior(u_num) .- interior(u_ans)).^2 ) / mean(interior(u_ans).^2) end function test_diffusion_simple(fieldname, timestepper, time_discretization) - model = NonhydrostaticModel(timestepper = timestepper, - grid = RectilinearGrid(CPU(), size=(1, 1, 16), extent=(1, 1, 1)), - closure = IsotropicDiffusivity(ν=1, κ=1, time_discretization=time_discretization), - coriolis = nothing, - tracers = :c, - buoyancy = nothing) - - field = get_model_field(fieldname, model) + model = NonhydrostaticModel(; timestepper, + grid = RectilinearGrid(CPU(), size=(1, 1, 16), extent=(1, 1, 1)), + closure = IsotropicDiffusivity(ν=1, κ=1, time_discretization=time_discretization), + coriolis = nothing, + tracers = :c, + buoyancy = nothing) value = π + field = get_model_field(fieldname, model) interior(field) .= value - update_state!(model) - for n in 1:10 - time_step!(model, 1) - end + [time_step!(model, 1) for n = 1:10] field_data = interior(field) - return !any(@. !isapprox(value, field_data)) end function test_isotropic_diffusion_budget(fieldname, model) set!(model; u=0, v=0, w=0, c=0) set!(model; Dict(fieldname => (x, y, z) -> rand())...) - field = get_model_field(fieldname, model) - ν = z_viscosity(model.closure, nothing) # for generalizing to isotropic AnisotropicDiffusivity - return test_diffusion_budget(fieldname, field, model, ν, model.grid.Δzᵃᵃᶜ) end function test_biharmonic_diffusion_budget(fieldname, model) set!(model; u=0, v=0, w=0, c=0) set!(model; Dict(fieldname => (x, y, z) -> rand())...) - field = get_model_field(fieldname, model) - return test_diffusion_budget(fieldname, field, model, model.closure.νz, model.grid.Δzᵃᵃᶜ, 4) end function test_diffusion_budget(fieldname, field, model, κ, Δ, order=2) init_mean = mean(interior(field)) - update_state!(model) - for n in 1:10 - # Very small time-steps required to bring error under machine precision - time_step!(model, 1e-4 * Δ^order / κ) - end + # Very small time-steps required to bring error under machine precision + [time_step!(model, 1e-4 * Δ^order / κ) for n = 1:10] final_mean = mean(interior(field)) @@ -75,30 +62,25 @@ end function test_diffusion_cosine(fieldname, timestepper, grid, time_discretization) κ, m = 1, 2 # diffusivity and cosine wavenumber - model = NonhydrostaticModel(timestepper = timestepper, - grid = grid, + model = NonhydrostaticModel(; timestepper, grid, closure = IsotropicDiffusivity(ν=κ, κ=κ, time_discretization=time_discretization), tracers = (:T, :S), buoyancy = nothing) field = get_model_field(fieldname, model) - zC = znodes(Center, grid, reshape=true) - interior(field) .= cos.(m * zC) - - diffusing_cosine(κ, m, z, t) = exp(-κ * m^2 * t) * cos(m * z) + z = znodes(Center, grid, reshape=true) + interior(field) .= cos.(m * z) + update_state!(model) # Step forward with small time-step relative to diff. time-scale Δt = 1e-6 * grid.Lz^2 / κ + [time_step!(model, Δt) for n = 1:10] - update_state!(model) - - for n in 1:10 - time_step!(model, Δt) - end + diffusing_cosine(κ, m, z, t) = exp(-κ * m^2 * t) * cos(m * z) numerical = interior(field) - analytical = diffusing_cosine.(κ, m, zC, model.clock.time) + analytical = diffusing_cosine.(κ, m, z, model.clock.time) return !any(@. !isapprox(numerical, analytical, atol=1e-6, rtol=1e-6)) end @@ -126,15 +108,12 @@ function passive_tracer_advection_test(timestepper; N=128, κ=1e-12, Nt=100, bac grid = RectilinearGrid(size=(N, N, 2), extent=(L, L, L)) closure = IsotropicDiffusivity(ν=κ, κ=κ) - model = NonhydrostaticModel(timestepper=timestepper, grid=grid, closure=closure, + model = NonhydrostaticModel(; grid, closure, timestepper, buoyancy=SeawaterBuoyancy(), tracers=(:T, :S), background_fields=background_fields) set!(model, u=u₀, v=v₀, T=T₀) - - for n in 1:Nt - time_step!(model, Δt) - end + [time_step!(model, Δt) for n = 1:Nt] # Error tolerance is a bit arbitrary return relative_error(model.tracers.T, T, model.clock.time) < 1e-4 @@ -207,13 +186,10 @@ function stratified_fluid_remains_at_rest_with_tilted_gravity_buoyancy_tracer(ar z_bc = GradientBoundaryCondition(N² * g̃[3]) b_bcs = FieldBoundaryConditions(bottom=z_bc, top=z_bc, south=y_bc, north=y_bc) - model = NonhydrostaticModel( - grid = grid, - buoyancy = buoyancy, - tracers = :b, - closure = nothing, - boundary_conditions = (b=b_bcs,) - ) + model = NonhydrostaticModel(; grid, buoyancy, + tracers = :b, + closure = nothing, + boundary_conditions = (; b=b_bcs)) b₀(x, y, z) = N² * (x*g̃[1] + y*g̃[2] + z*g̃[3]) set!(model, b=b₀) @@ -221,11 +197,8 @@ function stratified_fluid_remains_at_rest_with_tilted_gravity_buoyancy_tracer(ar simulation = Simulation(model, Δt=10minutes, stop_time=1hour) run!(simulation) - ∂y_b = ComputedField(∂y(model.tracers.b)) - ∂z_b = ComputedField(∂z(model.tracers.b)) - - compute!(∂y_b) - compute!(∂z_b) + @compute ∂y_b = Field(∂y(model.tracers.b)) + @compute ∂z_b = Field(∂z(model.tracers.b)) mean_∂y_b = mean(∂y_b) mean_∂z_b = mean(∂z_b) @@ -262,13 +235,10 @@ function stratified_fluid_remains_at_rest_with_tilted_gravity_temperature_tracer z_bc = GradientBoundaryCondition(∂T∂z * g̃[3]) T_bcs = FieldBoundaryConditions(bottom=z_bc, top=z_bc, south=y_bc, north=y_bc) - model = NonhydrostaticModel( - grid = grid, - buoyancy = buoyancy, - tracers = (:T, :S), - closure = nothing, - boundary_conditions = (T=T_bcs,) - ) + model = NonhydrostaticModel(; grid, buoyancy, + tracers = (:T, :S), + closure = nothing, + boundary_conditions = (; T=T_bcs)) T₀(x, y, z) = ∂T∂z * (x*g̃[1] + y*g̃[2] + z*g̃[3]) set!(model, T=T₀) @@ -276,11 +246,8 @@ function stratified_fluid_remains_at_rest_with_tilted_gravity_temperature_tracer simulation = Simulation(model, Δt=10minute, stop_time=1hour) run!(simulation) - ∂y_T = ComputedField(∂y(model.tracers.T)) - ∂z_T = ComputedField(∂z(model.tracers.T)) - - compute!(∂y_T) - compute!(∂z_T) + @compute ∂y_T = Field(∂y(model.tracers.T)) + @compute ∂z_T = Field(∂z(model.tracers.T)) mean_∂y_T = mean(∂y_T) mean_∂z_T = mean(∂z_T) @@ -302,43 +269,35 @@ function stratified_fluid_remains_at_rest_with_tilted_gravity_temperature_tracer return nothing end - - function inertial_oscillations_work_with_rotation_in_different_axis(arch, FT) grid = RectilinearGrid(arch, FT, size=(), topology=(Flat, Flat, Flat)) - f₀ = 1 ū = 1 Δt = 1e-3 T_inertial = 2π/f₀ stop_time = T_inertial / 2 - zcoriolis = FPlane(f=f₀) xcoriolis = ConstantCartesianCoriolis(f=f₀, rotation_axis=(1,0,0)) - model_x = NonhydrostaticModel(grid=grid, buoyancy=nothing, tracers=nothing, closure=nothing, - timestepper = :RungeKutta3, - coriolis=xcoriolis, - ) + model_x = NonhydrostaticModel(; grid, buoyancy=nothing, tracers=nothing, closure=nothing, + timestepper = :RungeKutta3, coriolis = xcoriolis) set!(model_x, v=ū) simulation_x = Simulation(model_x, Δt=Δt, stop_time=stop_time) run!(simulation_x) - model_z = NonhydrostaticModel(grid=grid, buoyancy=nothing, tracers=nothing, closure=nothing, - timestepper = :RungeKutta3, - coriolis=zcoriolis, - ) + model_z = NonhydrostaticModel(; grid, buoyancy=nothing, tracers=nothing, closure=nothing, + timestepper = :RungeKutta3, coriolis = zcoriolis) set!(model_z, u=ū) simulation_z = Simulation(model_z, Δt=Δt, stop_time=stop_time) run!(simulation_z) - u_x = model_x.velocities.u.data[1,1,1] - v_x = model_x.velocities.v.data[1,1,1] - w_x = model_x.velocities.w.data[1,1,1] + u_x = model_x.velocities.u[1, 1, 1] + v_x = model_x.velocities.v[1, 1, 1] + w_x = model_x.velocities.w[1, 1, 1] - u_z = model_z.velocities.u.data[1,1,1] - v_z = model_z.velocities.v.data[1,1,1] - w_z = model_z.velocities.w.data[1,1,1] + u_z = model_z.velocities.u[1, 1, 1] + v_z = model_z.velocities.v[1, 1, 1] + w_z = model_z.velocities.w[1, 1, 1] @test w_z == 0 @test u_x == 0 @@ -374,31 +333,27 @@ timesteppers = (:QuasiAdamsBashforth2, :RungeKutta3) (Periodic, Bounded, Bounded), (Bounded, Bounded, Bounded)) - if topology !== (Periodic, Periodic, Periodic) # can't use implicit time-stepping in vertically-periodic domains right now - time_discretizations = (ExplicitTimeDiscretization(), VerticallyImplicitTimeDiscretization()) - else - time_discretizations = (ExplicitTimeDiscretization(),) - end + time_discretizations = Any[ExplicitTimeDiscretization()] - for time_discretization in time_discretizations + topology !== (Periodic, Periodic, Periodic) && # can't use implicit time-stepping in vertically-periodic domains right now + push!(time_discretizations, VerticallyImplicitTimeDiscretization()) + for time_discretization in time_discretizations for closure in (IsotropicDiffusivity(ν=1, κ=1, time_discretization=time_discretization), AnisotropicDiffusivity(νh=1, νz=1, κh=1, κz=1, time_discretization=time_discretization)) fieldnames = [:c] - topology[1] === Periodic && push!(fieldnames, :u) topology[2] === Periodic && push!(fieldnames, :v) topology[3] === Periodic && push!(fieldnames, :w) grid = RectilinearGrid(size=(4, 4, 4), extent=(1, 1, 1), topology=topology) - model = NonhydrostaticModel(timestepper = timestepper, - grid = grid, - closure = closure, - tracers = :c, - coriolis = nothing, - buoyancy = nothing) + model = NonhydrostaticModel(; grid, closure, + timestepper = timestepper, + tracers = :c, + coriolis = nothing, + buoyancy = nothing) td = typeof(time_discretization).name.wrapper closurename = typeof(closure).name.wrapper @@ -422,19 +377,18 @@ timesteppers = (:QuasiAdamsBashforth2, :RungeKutta3) (Bounded, Bounded, Bounded)) fieldnames = [:c] - topology[1] === Periodic && push!(fieldnames, :u) topology[2] === Periodic && push!(fieldnames, :v) topology[3] === Periodic && push!(fieldnames, :w) grid = RectilinearGrid(size=(2, 2, 2), extent=(1, 1, 1), topology=topology) - model = NonhydrostaticModel(timestepper = timestepper, - grid = grid, - closure = AnisotropicBiharmonicDiffusivity(νh=1, νz=1, κh=1, κz=1), - coriolis = nothing, - tracers = :c, - buoyancy = nothing) + model = NonhydrostaticModel(; grid, + timestepper = timestepper, + closure = AnisotropicBiharmonicDiffusivity(νh=1, νz=1, κh=1, κz=1), + coriolis = nothing, + tracers = :c, + buoyancy = nothing) for fieldname in fieldnames @info " [$timestepper] Testing $fieldname budget in a $topology domain with biharmonic diffusion..." @@ -478,20 +432,20 @@ timesteppers = (:QuasiAdamsBashforth2, :RungeKutta3) # Regular grid with no flat dimension y_periodic_regular_grid = RectilinearGrid(topology=(Periodic, Periodic, Bounded), - size=(Nx, 1, Nz), x=(0, Lx), y=(0, Lx), z=(-Lz, 0)) + size=(Nx, 1, Nz), x=(0, Lx), y=(0, Lx), z=(-Lz, 0)) # Regular grid with a flat y-dimension y_flat_regular_grid = RectilinearGrid(topology=(Periodic, Flat, Bounded), - size=(Nx, Nz), x=(0, Lx), z=(-Lz, 0)) + size=(Nx, Nz), x=(0, Lx), z=(-Lz, 0)) # Vertically stretched grid with regular spacing and no flat dimension z_faces = collect(znodes(Face, y_periodic_regular_grid)) y_periodic_regularly_spaced_vertically_stretched_grid = RectilinearGrid(topology=(Periodic, Periodic, Bounded), - size=(Nx, 1, Nz), x=(0, Lx), y=(0, Lx), z=z_faces) + size=(Nx, 1, Nz), x=(0, Lx), y=(0, Lx), z=z_faces) # Vertically stretched grid with regular spacing and no flat dimension y_flat_regularly_spaced_vertically_stretched_grid = RectilinearGrid(topology=(Periodic, Flat, Bounded), - size=(Nx, Nz), x=(0, Lx), z=z_faces) + size=(Nx, Nz), x=(0, Lx), z=z_faces) solution, kwargs, background_fields, Δt, σ = internal_wave_solution(L=Lx, background_stratification=false) @@ -548,10 +502,9 @@ timesteppers = (:QuasiAdamsBashforth2, :RungeKutta3) # Regular grid with no flat dimension y_periodic_regular_grid = RectilinearGrid(topology=(Periodic, Periodic, Bounded), - size=(Nx, 1, Nz), x=(0, Lx), y=(0, Lx), z=(-Lz, 0)) + size=(Nx, 1, Nz), x=(0, Lx), y=(0, Lx), z=(-Lz, 0)) solution, kwargs, background_fields, Δt, σ = internal_wave_solution(L=Lx, background_stratification=true) - model = NonhydrostaticModel(; grid=y_periodic_regular_grid, background_fields=background_fields, kwargs...) internal_wave_dynamics_test(model, solution, Δt) end @@ -573,5 +526,5 @@ timesteppers = (:QuasiAdamsBashforth2, :RungeKutta3) inertial_oscillations_work_with_rotation_in_different_axis(arch, Float64) end end - end + diff --git a/test/test_reduced_fields.jl b/test/test_field_reductions.jl similarity index 76% rename from test/test_reduced_fields.jl rename to test/test_field_reductions.jl index 6f4837200d..92b0d8c413 100644 --- a/test/test_reduced_fields.jl +++ b/test/test_field_reductions.jl @@ -1,28 +1,28 @@ +include("dependencies_for_runtests.jl") + using Statistics using Oceananigans.Architectures: arch_array using Oceananigans.Fields: ReducedField, CenterField, ZFaceField, compute_at!, @compute using Oceananigans.BoundaryConditions: fill_halo_regions! using Oceananigans.Grids: halo_size -include("utils_for_runtests.jl") - -archs = test_architectures() - -@testset "AveragedField" begin - @info "Testing AveragedField..." +@testset "Fields computed by Reduction" begin + @info "Testing Fields computed by reductions..." for arch in archs - arch_str = string(typeof(arch)) @testset "Averaged fields [$arch_str]" begin - @info " Testing AveragedFields [$arch_str]" + @info " Testing averaged Fields [$arch_str]" + + grid = RectilinearGrid(arch, + topology = (Periodic, Periodic, Bounded), + size = (2, 2, 2), + x = (0, 2), + y = (0, 2), + z = (0, 2)) - grid = RectilinearGrid(arch, topology = (Periodic, Periodic, Bounded), - size = (2, 2, 2), - x = (0, 2), - y = (0, 2), - z = (0, 2)) + Nx, Ny, Nz = size(grid) w = ZFaceField(grid) T = CenterField(grid) @@ -32,21 +32,21 @@ archs = test_architectures() set!(T, trilinear) set!(w, trilinear) - @compute Txyz = AveragedField(T, dims=(1, 2, 3)) + #@compute Txyz = AveragedField(T, dims=(1, 2, 3)) + @compute Txyz = Field(Average(T, dims=(1, 2, 3))) # Note: halo regions must be *filled* prior to computing an average # if the average within halo regions is to be correct. fill_halo_regions!(T, arch) - @compute Txy = AveragedField(T, dims=(1, 2)) + #@compute Txy = AveragedField(T, dims=(1, 2)) + @compute Txy = Field(Average(T, dims=(1, 2))) fill_halo_regions!(T, arch) @compute Tx = AveragedField(T, dims=1) - @compute wxyz = AveragedField(w, dims=(1, 2, 3)) - @compute wxy = AveragedField(w, dims=(1, 2)) - @compute wx = AveragedField(w, dims=1) - - Nx, Ny, Nz = grid.Nx, grid.Ny, grid.Nz + @compute wxyz = Field(Average(w, dims=(1, 2, 3))) + @compute wxy = Field(Average(w, dims=(1, 2))) + @compute wx = Field(Average(w, dims=1)) for FT in float_types @@ -71,7 +71,7 @@ archs = test_architectures() c = CenterField(big_grid) c .= 1 - C = AveragedField(c, dims=(1, 2)) + C = Field(Average(c, dims=(1, 2))) # Test that the mean consistently returns 1 at every z for many evaluations results = [all(interior(mean!(C, C.operand)) .== 1) for i = 1:10] # warm up... @@ -102,11 +102,11 @@ archs = test_architectures() end end - @testset "Conditional computation of AveragedFields [$(typeof(arch))]" begin - @info " Testing conditional computation of AveragedFields [$(typeof(arch))]" + @testset "Conditional computation of averaged Fields [$(typeof(arch))]" begin + @info " Testing conditional computation of averaged Fields [$(typeof(arch))]" for FT in float_types grid = RectilinearGrid(arch, FT, size=(2, 2, 2), extent=(1, 1, 1)) - c = CenterField(arch, grid) + c = CenterField(grid) for dims in (1, 2, 3, (1, 2), (2, 3), (1, 3), (1, 2, 3)) C = AveragedField(c, dims=dims) diff --git a/test/test_hydrostatic_free_surface_immersed_boundaries.jl b/test/test_hydrostatic_free_surface_immersed_boundaries.jl index f90748ed2b..98fd258bde 100644 --- a/test/test_hydrostatic_free_surface_immersed_boundaries.jl +++ b/test/test_hydrostatic_free_surface_immersed_boundaries.jl @@ -1,4 +1,6 @@ -using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBoundary +include("dependencies_for_runtests.jl") + +using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBoundary, GridFittedBottom using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization @inline surface_wind_stress(λ, φ, t, p) = p.τ₀ * cos(2π * (φ - p.φ₀) / p.Lφ) @@ -18,7 +20,6 @@ using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization underlying_grid = RectilinearGrid(arch, size=(8, 8, 8), x = (-5, 5), y = (-5, 5), z = (0, 2)) bump(x, y, z) = z < exp(-x^2 - y^2) - grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBoundary(bump)) for closure in (IsotropicDiffusivity(ν=1, κ=0.5), @@ -47,6 +48,7 @@ using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization @test b[4, 4, 2] == 0 @test u[4, 4, 2] == 0 end + end @testset "Surface boundary conditions with immersed boundaries [$arch_str]" begin @info " Testing surface boundary conditions with ImmersedBoundaries in HydrostaticFreeSurfaceModel [$arch_str]..." diff --git a/test/test_turbulence_closures.jl b/test/test_turbulence_closures.jl index 08a47128f2..068bfa7301 100644 --- a/test/test_turbulence_closures.jl +++ b/test/test_turbulence_closures.jl @@ -26,8 +26,8 @@ function run_constant_isotropic_diffusivity_fluxdiv_tests(FT=Float64; ν=FT(0.3) arch = CPU() closure = IsotropicDiffusivity(FT, κ=(T=κ, S=κ), ν=ν) grid = RectilinearGrid(FT, size=(3, 1, 4), extent=(3, 1, 4)) - velocities = VelocityFields(arch, grid) - tracers = TracerFields((:T, :S), arch, grid) + velocities = VelocityFields(grid) + tracers = TracerFields((:T, :S), grid) clock = Clock(time=0.0) u, v, w = velocities @@ -59,8 +59,8 @@ function anisotropic_diffusivity_fluxdiv(FT=Float64; νh=FT(0.3), κh=FT(0.7), grid = RectilinearGrid(arch, FT, size=(3, 1, 4), extent=(3, 1, 4)) eos = LinearEquationOfState(FT) buoyancy = SeawaterBuoyancy(FT, gravitational_acceleration=1, equation_of_state=eos) - velocities = VelocityFields(arch, grid) - tracers = TracerFields((:T, :S), arch, grid) + velocities = VelocityFields(grid) + tracers = TracerFields((:T, :S), grid) clock = Clock(time=0.0) u, v, w, T, S = merge(velocities, tracers) @@ -93,15 +93,12 @@ function anisotropic_diffusivity_fluxdiv(FT=Float64; νh=FT(0.3), κh=FT(0.7), end function time_step_with_variable_isotropic_diffusivity(arch) - + grid = RectilinearGrid(arch, size=(1, 1, 1), extent=(1, 2, 3)) closure = IsotropicDiffusivity(ν = (x, y, z, t) -> exp(z) * cos(x) * cos(y) * cos(t), κ = (x, y, z, t) -> exp(z) * cos(x) * cos(y) * cos(t)) - model = NonhydrostaticModel(closure=closure, - grid=RectilinearGrid(arch, size=(1, 1, 1), extent=(1, 2, 3))) - + model = NonhydrostaticModel(; grid, closure) time_step!(model, 1, euler=true) - return true end @@ -136,15 +133,12 @@ function compute_closure_specific_diffusive_cfl(closurename) grid = RectilinearGrid(CPU(), size=(1, 1, 1), extent=(1, 2, 3)) closure = getproperty(TurbulenceClosures, closurename)() - model = NonhydrostaticModel(grid=grid, closure=closure) + model = NonhydrostaticModel(; grid, closure) dcfl = DiffusiveCFL(0.1) @test dcfl(model) isa Number - tracerless_model = NonhydrostaticModel(grid=grid, closure=closure, - buoyancy=nothing, tracers=nothing) - + tracerless_model = NonhydrostaticModel(; grid, closure, buoyancy=nothing, tracers=nothing) dcfl = DiffusiveCFL(0.2) - @test dcfl(tracerless_model) isa Number return nothing From 057f573f4139fe9eb3efac24528de3e02ee8c24c Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 30 Dec 2021 17:09:40 -0700 Subject: [PATCH 012/140] Fixes bug with Average field reduction --- src/AbstractOperations/AbstractOperations.jl | 3 +- .../metric_field_reductions.jl | 33 ++++-- src/Fields/field_reductions.jl | 3 +- src/Fields/zero_field.jl | 1 + test/test_field_reductions.jl | 112 +++++++++--------- 5 files changed, 82 insertions(+), 70 deletions(-) diff --git a/src/AbstractOperations/AbstractOperations.jl b/src/AbstractOperations/AbstractOperations.jl index 833cdb7edb..cfaea3f9f3 100644 --- a/src/AbstractOperations/AbstractOperations.jl +++ b/src/AbstractOperations/AbstractOperations.jl @@ -1,8 +1,7 @@ module AbstractOperations export ∂x, ∂y, ∂z, @at, @unary, @binary, @multiary -export AveragedField, IntegratedField -export KernelFunctionOperation +export Average, Integral, KernelFunctionOperation using Base: @propagate_inbounds diff --git a/src/AbstractOperations/metric_field_reductions.jl b/src/AbstractOperations/metric_field_reductions.jl index 90acd4c0ff..0005aaa055 100644 --- a/src/AbstractOperations/metric_field_reductions.jl +++ b/src/AbstractOperations/metric_field_reductions.jl @@ -1,3 +1,5 @@ +using Statistics: mean!, sum! + import Oceananigans.Fields: Reduction ##### @@ -19,22 +21,33 @@ reduction_grid_metric(dims) = dims === tuple(1) ? Δx : ##### Metric reductions ##### -abstract type AbstractMetricReduction end -struct Average <: AbstractMetricReduction end -struct Integral <: AbstractMetricReduction end +"""docstring...""" +struct Average end + +function Reduction(avg::Average, field::AbstractField; dims) + dx = reduction_grid_metric(dims) + + # Compute "size" (length, area, or volume) of averaging region + metric = GridMetricOperation(location(field), dx, field.grid) + L = sum(metric; dims) + L⁻¹_field_dx = field * dx / L + + return Reduction(sum!, L⁻¹_field_dx; dims) +end + +Average(field::AbstractField; dims=:) = Reduction(Average(), field; dims) + +const AveragedField = Field{<:Any, <:Any, <:Any, <:Reduction{<:Average}} -reduction(::Average) = mean! -reduction(::Integral) = sum! +"""docstring...""" +struct Integral end -function Reduction(r::AbstractMetricReduction, operand; dims) +function Reduction(int::Integral, field::AbstractField; dims) dx = reduction_grid_metric(dims) - field_dx = field * dx - return Reduction(reduction(r), field_dx; dims) + return Reduction(sum!, field * dx; dims) end # Convenience -Average(field::AbstractField; dims=:) = Reduction(Average(), field; dims) Integral(field::AbstractField; dims=:) = Reduction(Integral(), field; dims) -const AveragedField = Field{<:Any, <:Any, <:Any, <:Reduction{<:Average}} const IntegratedField = Field{<:Any, <:Any, <:Any, <:Reduction{<:Integral}} diff --git a/src/Fields/field_reductions.jl b/src/Fields/field_reductions.jl index 2f80b93004..c7025c4f7e 100644 --- a/src/Fields/field_reductions.jl +++ b/src/Fields/field_reductions.jl @@ -35,7 +35,7 @@ function Field(reduction::Reduction; operand = reduction.operand grid = operand.grid - LX, LY, LZ = loc = reduced_location(location(operand); dims) + LX, LY, LZ = loc = reduced_location(location(operand); dims=reduction.dims) if isnothing(data) data = new_data(grid, loc) @@ -44,6 +44,7 @@ function Field(reduction::Reduction; boundary_conditions = FieldBoundaryConditions(grid, loc) status = recompute_safely ? nothing : FieldStatus() + return Field(loc, grid, data, boundary_conditions, reduction, status) end diff --git a/src/Fields/zero_field.jl b/src/Fields/zero_field.jl index cc6561e149..7b903c9925 100644 --- a/src/Fields/zero_field.jl +++ b/src/Fields/zero_field.jl @@ -1,3 +1,4 @@ struct ZeroField <: AbstractField{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, 3} end @inline Base.getindex(::ZeroField, i, j, k) = 0 + diff --git a/test/test_field_reductions.jl b/test/test_field_reductions.jl index 92b0d8c413..c09ececc4d 100644 --- a/test/test_field_reductions.jl +++ b/test/test_field_reductions.jl @@ -12,77 +12,75 @@ using Oceananigans.Grids: halo_size for arch in archs arch_str = string(typeof(arch)) - @testset "Averaged fields [$arch_str]" begin - @info " Testing averaged Fields [$arch_str]" + grid = RectilinearGrid(arch, size = (2, 2, 2), + x = (0, 2), y = (0, 2), z = (0, 2), + topology = (Periodic, Periodic, Bounded)) - grid = RectilinearGrid(arch, - topology = (Periodic, Periodic, Bounded), - size = (2, 2, 2), - x = (0, 2), - y = (0, 2), - z = (0, 2)) + Nx, Ny, Nz = size(grid) - Nx, Ny, Nz = size(grid) + w = ZFaceField(grid) + T = CenterField(grid) - w = ZFaceField(grid) - T = CenterField(grid) + trilinear(x, y, z) = x + y + z - trilinear(x, y, z) = x + y + z + set!(T, trilinear) + set!(w, trilinear) - set!(T, trilinear) - set!(w, trilinear) + @compute Txyz = Field(Average(T, dims=(1, 2, 3))) - #@compute Txyz = AveragedField(T, dims=(1, 2, 3)) - @compute Txyz = Field(Average(T, dims=(1, 2, 3))) + # Note: halo regions must be *filled* prior to computing an average + # if the average within halo regions is to be correct. + fill_halo_regions!(T, arch) + @compute Txy = Field(Average(T, dims=(1, 2))) - # Note: halo regions must be *filled* prior to computing an average - # if the average within halo regions is to be correct. - fill_halo_regions!(T, arch) - #@compute Txy = AveragedField(T, dims=(1, 2)) - @compute Txy = Field(Average(T, dims=(1, 2))) + fill_halo_regions!(T, arch) + @compute Tx = Field(Average(T, dims=1)) - fill_halo_regions!(T, arch) - @compute Tx = AveragedField(T, dims=1) + @compute wxyz = Field(Average(w, dims=(1, 2, 3))) + @compute wxy = Field(Average(w, dims=(1, 2))) + @compute wx = Field(Average(w, dims=1)) - @compute wxyz = Field(Average(w, dims=(1, 2, 3))) - @compute wxy = Field(Average(w, dims=(1, 2))) - @compute wx = Field(Average(w, dims=1)) + @testset "Averaged fields [$arch_str]" begin + @info " Testing averaged Fields [$arch_str]" - for FT in float_types - - @test Txyz[1, 1, 1] ≈ 3 - - @test Array(interior(Txy))[1, 1, :] ≈ [2.5, 3.5] - @test Array(interior(Tx))[1, :, :] ≈ [[2, 3] [3, 4]] - - @test wxyz[1, 1, 1] ≈ 3 - - @test Array(interior(wxy))[1, 1, :] ≈ [2, 3, 4] - @test Array(interior(wx))[1, :, :] ≈ [[1.5, 2.5] [2.5, 3.5] [3.5, 4.5]] - - # Test whether a race condition gets hit for averages over large fields - big_grid = RectilinearGrid(arch, - topology = (Periodic, Periodic, Bounded), - size = (256, 256, 128), - x = (0, 2), - y = (0, 2), - z = (0, 2)) - - c = CenterField(big_grid) - c .= 1 - - C = Field(Average(c, dims=(1, 2))) - - # Test that the mean consistently returns 1 at every z for many evaluations - results = [all(interior(mean!(C, C.operand)) .== 1) for i = 1:10] # warm up... - results = [all(interior(mean!(C, C.operand)) .== 1) for i = 1:10] # the real deal - @test mean(results) == 1.0 + @test Txyz[1, 1, 1] ≈ 3 + @test Array(interior(Txy))[1, 1, :] ≈ [2.5, 3.5] + @test Array(interior(Tx))[1, :, :] ≈ [[2, 3] [3, 4]] + @test wxyz[1, 1, 1] ≈ 3 + @test Array(interior(wxy))[1, 1, :] ≈ [2, 3, 4] + @test Array(interior(wx))[1, :, :] ≈ [[1.5, 2.5] [2.5, 3.5] [3.5, 4.5]] + + # Test whether a race condition gets hit for averages over large fields + big_grid = RectilinearGrid(arch, + topology = (Periodic, Periodic, Bounded), + size = (256, 256, 128), + x = (0, 2), + y = (0, 2), + z = (0, 2)) + + c = CenterField(big_grid) + c .= 1 + + C = Field(Average(c, dims=(1, 2))) + + # Test that the mean consistently returns 1 at every z for many evaluations + # Warm up + for i = 1:10 + sum!(C, C.operand.operand) + end + + results = [] + for i = 1:10 + sum!(C, C.operand.operand) + push!(results, all(interior(C) .== 1)) end + + @test mean(results) == 1.0 end @testset "Allocating reductions [$arch_str]" begin @info " Testing allocating reductions" - + # Mean @test Txyz[1, 1, 1] == mean(T) @test interior(Txy) == interior(mean(T, dims=(1, 2))) @@ -109,7 +107,7 @@ using Oceananigans.Grids: halo_size c = CenterField(grid) for dims in (1, 2, 3, (1, 2), (2, 3), (1, 3), (1, 2, 3)) - C = AveragedField(c, dims=dims) + C = Field(Average(c, dims=dims)) @test !isnothing(C.status) From 86eff52ba1bfff5442fc0a9ac3ce48cf3056eddb Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 30 Dec 2021 18:18:44 -0700 Subject: [PATCH 013/140] Updates computed Field and BuoyancyField tests --- .../kernel_function_operation.jl | 16 +- src/BuoyancyModels/buoyancy_field.jl | 4 +- test/test_computed_field.jl | 156 +++++++++--------- 3 files changed, 90 insertions(+), 86 deletions(-) diff --git a/src/AbstractOperations/kernel_function_operation.jl b/src/AbstractOperations/kernel_function_operation.jl index fbd3f89465..0327b64813 100644 --- a/src/AbstractOperations/kernel_function_operation.jl +++ b/src/AbstractOperations/kernel_function_operation.jl @@ -2,14 +2,12 @@ struct KernelFunctionOperation{LX, LY, LZ, P, A, G, T, K, D} <: AbstractOperatio kernel_function :: K computed_dependencies :: D parameters :: P - architecture :: A grid :: G function KernelFunctionOperation{LX, LY, LZ}(kernel_function::K, computed_dependencies::D, parameters::P, grid::G) where {LX, LY, LZ, K, G, D, P} - arch = architecture(grid) T = eltype(grid) - A = typeof(arch) + A = typeof(architecture(grid)) return new{LX, LY, LZ, P, A, G, T, K, D}(kernel_function, computed_dependencies, parameters, grid) end @@ -23,8 +21,7 @@ end Constructs a `KernelFunctionOperation` at location `(LX, LY, LZ)` on `grid` an with an optional iterable of `computed_dependencies` and arbitrary `parameters`. -With `isnothing(parameters)` (the default), `kernel_function` is called -with +With `isnothing(parameters)` (the default), `kernel_function` is called with ```julia kernel_function(i, j, k, grid, computed_dependencies...) @@ -59,8 +56,12 @@ u, v, w = model.velocities ζ_op = KernelFunctionOperation{Face, Face, Center}(ζ₃ᶠᶠᵃ, grid, computed_dependencies=(u, v)) ``` """ -KernelFunctionOperation{LX, LY, LZ}(kernel_function, grid; computed_dependencies = (), parameters = nothing) where {LX, LY, LZ} = - KernelFunctionOperation{LX, LY, LZ}(kernel_function, computed_dependencies, parameters, grid) +function KernelFunctionOperation{LX, LY, LZ}(kernel_function, grid; + computed_dependencies = (), + parameters = nothing) where {LX, LY, LZ} + + return KernelFunctionOperation{LX, LY, LZ}(kernel_function, computed_dependencies, parameters, grid) +end @inline Base.getindex(κ::KernelFunctionOperation, i, j, k) = κ.kernel_function(i, j, k, κ.grid, κ.computed_dependencies..., κ.parameters) @inline Base.getindex(κ::KernelFunctionOperation{LX, LY, LZ, <:Nothing}, i, j, k) where {LX, LY, LZ} = κ.kernel_function(i, j, k, κ.grid, κ.computed_dependencies...) @@ -73,6 +74,5 @@ Adapt.adapt_structure(to, κ::KernelFunctionOperation{LX, LY, LZ}) where {LX, LY KernelFunctionOperation{LX, LY, LZ}(Adapt.adapt(to, κ.kernel_function), Adapt.adapt(to, κ.computed_dependencies), Adapt.adapt(to, κ.parameters), - nothing, Adapt.adapt(to, κ.grid)) diff --git a/src/BuoyancyModels/buoyancy_field.jl b/src/BuoyancyModels/buoyancy_field.jl index b640d60e62..736e2b68cc 100644 --- a/src/BuoyancyModels/buoyancy_field.jl +++ b/src/BuoyancyModels/buoyancy_field.jl @@ -1,3 +1,5 @@ +using Oceananigans.AbstractOperations: KernelFunctionOperation +using Oceananigans.Fields: Field using Adapt using KernelAbstractions @@ -8,7 +10,7 @@ function buoyancy_operation(model) end buoyancy_operation(buoyancy_model, grid, tracers) = - KernelFunctionOperation{Center, Center, Center}(buoyancy_perturbation, model.grid, computed=dependencies=(buoyancy_model, tracers)) + KernelFunctionOperation{Center, Center, Center}(buoyancy_perturbation, grid, computed_dependencies=(buoyancy_model, tracers)) buoyancy_operation(::Nothing, grid, tracers) = nothing diff --git a/test/test_computed_field.jl b/test/test_computed_field.jl index 83e7854fd9..4ff1a62f99 100644 --- a/test/test_computed_field.jl +++ b/test/test_computed_field.jl @@ -1,13 +1,15 @@ +include("dependencies_for_runtests.jl") + using Oceananigans.AbstractOperations: UnaryOperation, Derivative, BinaryOperation, MultiaryOperation using Oceananigans.AbstractOperations: KernelFunctionOperation using Oceananigans.Operators: ℑxyᶜᶠᵃ, ℑxyᶠᶜᵃ -using Oceananigans.Fields: PressureField, compute_at! +using Oceananigans.Fields: compute_at! using Oceananigans.BuoyancyModels: BuoyancyField function compute_derivative(model, ∂) T, S = model.tracers S.data.parent .= π - @compute ∂S = ComputedField(∂(S)) + @compute ∂S = Field(∂(S)) result = Array(interior(∂S)) return all(result .≈ zero(eltype(model.grid))) end @@ -15,7 +17,7 @@ end function compute_unary(unary, model) set!(model; S=π) T, S = model.tracers - @compute uS = ComputedField(unary(S), data=model.pressures.pHY′.data) + @compute uS = Field(unary(S), data=model.pressures.pHY′.data) result = Array(interior(uS)) return all(result .≈ unary(eltype(model.grid)(π))) end @@ -23,7 +25,7 @@ end function compute_plus(model) set!(model; S=π, T=42) T, S = model.tracers - @compute ST = ComputedField(S + T, data=model.pressures.pHY′.data) + @compute ST = Field(S + T, data=model.pressures.pHY′.data) result = Array(interior(ST)) return all(result .≈ eltype(model.grid)(π+42)) end @@ -32,7 +34,7 @@ function compute_many_plus(model) set!(model; u=2, S=π, T=42) T, S = model.tracers u, v, w = model.velocities - @compute uTS = ComputedField(@at((Center, Center, Center), u + T + S)) + @compute uTS = Field(@at((Center, Center, Center), u + T + S)) result = Array(interior(uTS)) return all(result .≈ eltype(model.grid)(2+π+42)) end @@ -40,7 +42,7 @@ end function compute_minus(model) set!(model; S=π, T=42) T, S = model.tracers - @compute ST = ComputedField(S - T, data=model.pressures.pHY′.data) + @compute ST = Field(S - T, data=model.pressures.pHY′.data) result = Array(interior(ST)) return all(result .≈ eltype(model.grid)(π-42)) end @@ -48,7 +50,7 @@ end function compute_times(model) set!(model; S=π, T=42) T, S = model.tracers - @compute ST = ComputedField(S * T, data=model.pressures.pHY′.data) + @compute ST = Field(S * T, data=model.pressures.pHY′.data) result = Array(interior(ST)) return all(result .≈ eltype(model.grid)(π*42)) end @@ -60,7 +62,7 @@ function compute_kinetic_energy(model) set!(w, 3) kinetic_energy_operation = @at (Center, Center, Center) (u^2 + v^2 + w^2) / 2 - @compute kinetic_energy = ComputedField(kinetic_energy_operation, data=model.pressures.pHY′.data) + @compute kinetic_energy = Field(kinetic_energy_operation, data=model.pressures.pHY′.data) result = Array(interior(kinetic_energy))[2:3, 2:3, 2:3] return all(result .≈ 7) @@ -74,7 +76,7 @@ function horizontal_average_of_plus(model) set!(model; S=S₀, T=T₀) T, S = model.tracers - @compute ST = AveragedField(S + T, dims=(1, 2)) + @compute ST = Field(Average(S + T, dims=(1, 2))) zC = znodes(Center, model.grid) correct_profile = @. sin(π * zC) + 42 * zC @@ -92,7 +94,7 @@ function zonal_average_of_plus(model) set!(model; S=S₀, T=T₀) T, S = model.tracers - @compute ST = AveragedField(S + T, dims=1) + @compute ST = Field(Average(S + T, dims=1)) yC = ynodes(Center, model.grid, reshape=true) zC = znodes(Center, model.grid, reshape=true) @@ -111,7 +113,7 @@ function volume_average_of_times(model) set!(model; S=S₀, T=T₀) T, S = model.tracers - @compute ST = AveragedField(S * T, dims=(1, 2, 3)) + @compute ST = Field(Average(S * T, dims=(1, 2, 3))) result = Array(interior(ST)) @@ -126,7 +128,7 @@ function horizontal_average_of_minus(model) set!(model; S=S₀, T=T₀) T, S = model.tracers - @compute ST = AveragedField(S - T, dims=(1, 2)) + @compute ST = Field(Average(S - T, dims=(1, 2))) zC = znodes(Center, model.grid) correct_profile = @. sin(π * zC) - 42 * zC @@ -144,7 +146,7 @@ function horizontal_average_of_times(model) set!(model; S=S₀, T=T₀) T, S = model.tracers - @compute ST = AveragedField(S * T, dims=(1, 2)) + @compute ST = Field(Average(S * T, dims=(1, 2))) zC = znodes(Center, model.grid) correct_profile = @. sin(π * zC) * 42 * zC @@ -164,7 +166,7 @@ function multiplication_and_derivative_ccf(model) w = model.velocities.w T = model.tracers.T - @compute wT = AveragedField(w * ∂z(T), dims=(1, 2)) + @compute wT = Field(Average(w * ∂z(T), dims=(1, 2))) zF = znodes(Face, model.grid) correct_profile = @. 42 * sin(π * zF) @@ -189,7 +191,7 @@ function multiplication_and_derivative_ccc(model) T = model.tracers.T wT_ccc = @at (C, C, C) w * ∂z(T) - @compute wT_ccc_avg = AveragedField(wT_ccc, dims=(1, 2)) + @compute wT_ccc_avg = Field(Average(wT_ccc, dims=(1, 2))) zF = znodes(Face, model.grid) sinusoid = sin.(π * zF) @@ -213,32 +215,34 @@ function computation_including_boundaries(arch) @. w.data = 3 + rand() op = @at (Center, Face, Face) u * v * w - @compute uvw = ComputedField(op) + @compute uvw = Field(op) return all(interior(uvw) .!= 0) end function operations_with_computed_field(model) u, v, w = model.velocities - uv = ComputedField(u * v) - @compute uvw = ComputedField(uv * w) + uv = Field(u * v) + @compute uvw = Field(uv * w) return true end function operations_with_averaged_field(model) u, v, w = model.velocities - UV = AveragedField(u * v, dims=(1, 2)) - wUV = ComputedField(w * UV) + UV = Field(Average(u * v, dims=(1, 2))) + wUV = Field(w * UV) compute!(wUV) return true end +#= function pressure_field(model) p = PressureField(model) u, v, w = model.velocities - @compute up = ComputedField(u * p) + @compute up = Field(u * p) return true end +=# function computations_with_buoyancy_field(arch, buoyancy) grid = RectilinearGrid(arch, size=(1, 1, 1), extent=(1, 1, 1)) @@ -251,9 +255,9 @@ function computations_with_buoyancy_field(arch, buoyancy) compute!(b) - ub = ComputedField(b * u) - vb = ComputedField(b * v) - wb = ComputedField(b * w) + ub = Field(b * u) + vb = Field(b * v) + wb = Field(b * w) compute!(ub) compute!(vb) @@ -268,11 +272,11 @@ function computations_with_averaged_fields(model) set!(model, enforce_incompressibility = false, u = (x, y, z) -> z, v = 2, w = 3) # Two ways to compute turbulent kinetic energy - U = AveragedField(u, dims=(1, 2)) - V = AveragedField(v, dims=(1, 2)) + U = Field(Average(u, dims=(1, 2))) + V = Field(Average(v, dims=(1, 2))) tke_op = @at (Center, Center, Center) ((u - U)^2 + (v - V)^2 + w^2) / 2 - tke = ComputedField(tke_op) + tke = Field(tke_op) compute!(tke) return all(interior(tke)[2:3, 2:3, 2:3] .== 9/2) @@ -285,12 +289,12 @@ function computations_with_averaged_field_derivative(model) u, v, w, T, S = fields(model) # Two ways to compute turbulent kinetic energy - U = AveragedField(u, dims=(1, 2)) - V = AveragedField(v, dims=(1, 2)) + U = Field(Average(u, dims=(1, 2))) + V = Field(Average(v, dims=(1, 2))) # This tests a vertical derivative of an AveragedField shear_production_op = @at (Center, Center, Center) u * w * ∂z(U) - shear = ComputedField(shear_production_op) + shear = Field(shear_production_op) compute!(shear) set!(model, T = (x, y, z) -> 3 * z) @@ -304,32 +308,30 @@ function computations_with_computed_fields(model) set!(model, enforce_incompressibility = false, u = (x, y, z) -> z, v = 2, w = 3) # Two ways to compute turbulent kinetic energy - U = AveragedField(u, dims=(1, 2)) - V = AveragedField(v, dims=(1, 2)) + U = Field(Average(u, dims=(1, 2))) + V = Field(Average(v, dims=(1, 2))) - u′ = ComputedField(u - U) - v′ = ComputedField(v - V) + u′ = Field(u - U) + v′ = Field(v - V) tke_op = @at (Center, Center, Center) (u′^2 + v′^2 + w^2) / 2 - tke = ComputedField(tke_op) + tke = Field(tke_op) compute!(tke) return all(interior(tke)[2:3, 2:3, 2:3] .== 9/2) end for arch in archs - @testset "ComputedFields [$(typeof(arch))]" begin - @info " Testing ComputedFields [$(typeof(arch))]..." + @testset "Computed Fields [$(typeof(arch))]" begin + @info " Testing computed Fields [$(typeof(arch))]..." grid = RectilinearGrid(arch, size=(4, 4, 4), extent=(1, 1, 1), - topology=(Periodic, Periodic, Bounded)) + topology=(Periodic, Periodic, Bounded)) buoyancy = SeawaterBuoyancy(gravitational_acceleration = 1, - equation_of_state = LinearEquationOfState(α=1, β=1)) + equation_of_state = LinearEquationOfState(α=1, β=1)) - model = NonhydrostaticModel(grid = grid, buoyancy = buoyancy, - tracers = (:T, :S) - ) + model = NonhydrostaticModel(; grid, buoyancy, tracers = (:T, :S)) @testset "Derivative computations [$(typeof(arch))]" begin @info " Testing compute! derivatives..." @@ -353,7 +355,7 @@ for arch in archs # Basic compilation test for nested BinaryOperations... u, v, w = model.velocities - @test try compute!(ComputedField(u + v - w)); true; catch; false; end + @test try compute!(Field(u + v - w)); true; catch; false; end end @testset "Multiary computations [$(typeof(arch))]" begin @@ -364,38 +366,38 @@ for arch in archs @test compute_kinetic_energy(model) end - @testset "Computations with KernelComputedField [$(typeof(arch))]" begin + @testset "Computations with KernelFunctionOperation [$(typeof(arch))]" begin @test begin @inline trivial_kernel_function(i, j, k, grid) = 1 op = KernelFunctionOperation{Center, Center, Center}(trivial_kernel_function, grid) - f = ComputedField(op, arch) + f = Field(op) compute!(f) - f isa ComputedField && f.operand === op + f isa Field && f.operand === op end @test begin @inline trivial_parameterized_kernel_function(i, j, k, grid, μ) = μ op = KernelFunctionOperation{Center, Center, Center}(trivial_parameterized_kernel_function, grid, parameters=0.1) - f = ComputedField(op, arch) + f = Field(op) compute!(f) - f isa ComputedField && f.operand === op + f isa Field && f.operand === op end @test begin u, v, w = model.velocities ζ_op = KernelFunctionOperation{Face, Face, Center}(ζ₃ᶠᶠᵃ, grid, computed_dependencies=(u, v)) - ζ = ComputedField(ζ_op) # identical to `VerticalVorticityField` + ζ = Field(ζ_op) # identical to `VerticalVorticityField` compute!(ζ) - ζ isa ComputedField && ζ.operand.kernel_function === ζ₃ᶠᶠᵃ + ζ isa Field && ζ.operand.kernel_function === ζ₃ᶠᶠᵃ end end - @testset "Operations with ComputedField and PressureField [$(typeof(arch))]" begin - @info " Testing operations with ComputedField..." + @testset "Operations with Field and PressureField [$(typeof(arch))]" begin + @info " Testing operations with Field..." @test operations_with_computed_field(model) - @info " Testing PressureField..." - @test pressure_field(model) + # @info " Testing PressureField..." + # @test pressure_field(model) end @testset "Horizontal averages of operations [$(typeof(arch))]" begin @@ -418,13 +420,13 @@ for arch in archs @test volume_average_of_times(model) end - @testset "ComputedField boundary conditions [$(typeof(arch))]" begin - @info " Testing boundary conditions for ComputedField..." + @testset "Field boundary conditions [$(typeof(arch))]" begin + @info " Testing boundary conditions for Field..." set!(model; S=π, T=42) T, S = model.tracers - @compute ST = ComputedField(S + T, data=model.pressures.pHY′.data) + @compute ST = Field(S + T, data=model.pressures.pHY′.data) Nx, Ny, Nz = size(model.grid) @@ -433,7 +435,7 @@ for arch in archs @test all(ST.data[1:Nx, 1:Ny, 0] .== ST.data[1:Nx, 1:Ny, 1]) @test all(ST.data[1:Nx, 1:Ny, Nz] .== ST.data[1:Nx, 1:Ny, Nz+1]) - @compute ST_face = ComputedField(@at (Center, Center, Face) S * T) + @compute ST_face = Field(@at (Center, Center, Face) S * T) @test all(ST_face.data[1:Nx, 1:Ny, 0] .== 0) @test all(ST_face.data[1:Nx, 1:Ny, Nz+2] .== 0) @@ -443,7 +445,7 @@ for arch in archs @info " Testing operations with AveragedField..." T, S = model.tracers - TS = AveragedField(T * S, dims=(1, 2)) + TS = Field(Average(T * S, dims=(1, 2))) @test operations_with_averaged_field(model) end @@ -477,8 +479,8 @@ for arch in archs set!(model, enforce_incompressibility = false, u = (x, y, z) -> z, v = 2, w = 3) # Two ways to compute turbulent kinetic energy - U = AveragedField(u, dims=(1, 2)) - V = AveragedField(v, dims=(1, 2)) + U = Field(Average(u, dims=(1, 2))) + V = Field(Average(v, dims=(1, 2))) # Build up compilation tests incrementally... u_prime = u - U @@ -492,35 +494,35 @@ for arch in archs tke = ((u - U)^2 + (v - V)^2 + w^2) / 2 tke_ccc = @at (Center, Center, Center) ((u - U)^2 + (v - V)^2 + w^2) / 2 - @test try compute!(ComputedField(u_prime )); true; catch; false; end - @test try compute!(ComputedField(u_prime_ccc )); true; catch; false; end - @test try compute!(ComputedField(u_prime_squared )); true; catch; false; end - @test try compute!(ComputedField(u_prime_squared_ccc )); true; catch; false; end - @test try compute!(ComputedField(horizontal_twice_tke)); true; catch; false; end - @test try compute!(ComputedField(horizontal_tke )); true; catch; false; end - @test try compute!(ComputedField(twice_tke )); true; catch; false; end + @test try compute!(Field(u_prime )); true; catch; false; end + @test try compute!(Field(u_prime_ccc )); true; catch; false; end + @test try compute!(Field(u_prime_squared )); true; catch; false; end + @test try compute!(Field(u_prime_squared_ccc )); true; catch; false; end + @test try compute!(Field(horizontal_twice_tke)); true; catch; false; end + @test try compute!(Field(horizontal_tke )); true; catch; false; end + @test try compute!(Field(twice_tke )); true; catch; false; end - @test try compute!(ComputedField(horizontal_tke_ccc )); true; catch; false; end - @test try compute!(ComputedField(tke )); true; catch; false; end - @test try compute!(ComputedField(tke_ccc )); true; catch; false; end + @test try compute!(Field(horizontal_tke_ccc )); true; catch; false; end + @test try compute!(Field(tke )); true; catch; false; end + @test try compute!(Field(tke_ccc )); true; catch; false; end - computed_tke = ComputedField(tke_ccc) + computed_tke = Field(tke_ccc) @test (compute!(computed_tke); all(interior(computed_tke)[2:3, 2:3, 2:3] .== 9/2)) end - @testset "Computations with ComputedFields [$(typeof(arch))]" begin - @info " Testing computations with ComputedField [$(typeof(arch))]..." + @testset "Computations with Fields [$(typeof(arch))]" begin + @info " Testing computations with Field [$(typeof(arch))]..." @test computations_with_computed_fields(model) end - @testset "Conditional computation of ComputedField and BuoyancyField [$(typeof(arch))]" begin - @info " Testing conditional computation of ComputedField and BuoyancyField " * + @testset "Conditional computation of Field and BuoyancyField [$(typeof(arch))]" begin + @info " Testing conditional computation of Field and BuoyancyField " * "[$(typeof(arch))]..." set!(model, u=2, v=0, w=0, T=3, S=0) u, v, w, T, S = fields(model) - uT = ComputedField(u * T) + uT = Field(u * T) α = model.buoyancy.model.equation_of_state.α g = model.buoyancy.model.gravitational_acceleration From e0c9761141320f312ee06e9dfbe8e383f53edc34 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 30 Dec 2021 18:25:04 -0700 Subject: [PATCH 014/140] Fix broadcasting tests --- test/test_broadcasting.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/test/test_broadcasting.jl b/test/test_broadcasting.jl index 7707bc58f7..de3dc22bab 100644 --- a/test/test_broadcasting.jl +++ b/test/test_broadcasting.jl @@ -8,7 +8,7 @@ ##### grid = RectilinearGrid(arch, size=(1, 1, 1), extent=(1, 1, 1)) - a, b, c = [CenterField(arch, grid) for i = 1:3] + a, b, c = [CenterField(grid) for i = 1:3] Nx, Ny, Nz = size(a) @@ -34,10 +34,10 @@ three_point_grid = RectilinearGrid(arch, size=(1, 1, 3), extent=(1, 1, 1)) - a2 = CenterField(arch, three_point_grid) + a2 = CenterField(three_point_grid) b2_bcs = FieldBoundaryConditions(grid, (Center, Center, Face), top=OpenBoundaryCondition(0), bottom=OpenBoundaryCondition(0)) - b2 = ZFaceField(arch, three_point_grid, b2_bcs) + b2 = ZFaceField(three_point_grid, boundary_conditions=b2_bcs) b2 .= 1 fill_halo_regions!(b2, arch) # sets b2[1, 1, 1] = b[1, 1, 4] = 0 @@ -61,7 +61,7 @@ ##### Broadcasting with ReducedField ##### - r, p, q = [ReducedField(Center, Center, Nothing, arch, grid, dims=3) for i = 1:3] + r, p, q = [Field{Center, Center, Nothing}(grid) for i = 1:3] r .= 2 @test all(r .== 2) @@ -80,7 +80,7 @@ two_two_two_grid = RectilinearGrid(arch, size=(2, 2, 2), extent=(1, 1, 1)) - c = CenterField(arch, two_two_two_grid) + c = CenterField(two_two_two_grid) random_column = arch_array(arch, reshape(rand(2), 1, 1, 2)) c .= random_column # broadcast to every horizontal column in c From 82595b7694aec7272cf9d352f1a992ae58ef7bc8 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 31 Dec 2021 06:00:48 -0700 Subject: [PATCH 015/140] Updates simulation tests --- test/test_simulations.jl | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/test/test_simulations.jl b/test/test_simulations.jl index 2411576469..842b31cb9e 100644 --- a/test/test_simulations.jl +++ b/test/test_simulations.jl @@ -1,17 +1,11 @@ -using Test +include("dependencies_for_runtests.jl") -using Oceananigans -using Oceananigans.Units using Oceananigans.Simulations: stop_iteration_exceeded, stop_time_exceeded, wall_time_limit_exceeded, TimeStepWizard, new_time_step, reset! using Dates: DateTime -include("utils_for_runtests.jl") - -archs = test_architectures() - function wall_time_step_wizard_tests(arch) grid = RectilinearGrid(arch, size=(1, 1, 1), extent=(1, 1, 1)) model = NonhydrostaticModel(grid=grid) From 3bcbcffd1593c2e84ec02abfd3b054a013ea9662 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 31 Dec 2021 07:38:08 -0700 Subject: [PATCH 016/140] Bugfix in show mmethod for lat-lon grid --- src/Grids/grid_utils.jl | 2 -- src/Grids/latitude_longitude_grid.jl | 9 ++++++--- test/test_grids.jl | 5 +++-- 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/src/Grids/grid_utils.jl b/src/Grids/grid_utils.jl index 95a55146a0..a3e49b9f8e 100644 --- a/src/Grids/grid_utils.jl +++ b/src/Grids/grid_utils.jl @@ -66,8 +66,6 @@ Base.size(loc, grid, d) = size(loc, grid)[d] total_size(a) = size(a) # fallback -halo_size(grid) = (grid.Hx, grid.Hy, grid.Hz) - """ total_size(loc, grid) diff --git a/src/Grids/latitude_longitude_grid.jl b/src/Grids/latitude_longitude_grid.jl index ace4955a41..913346313f 100644 --- a/src/Grids/latitude_longitude_grid.jl +++ b/src/Grids/latitude_longitude_grid.jl @@ -181,9 +181,12 @@ function domain_string(grid::LatitudeLongitudeGrid) return "longitude λ ∈ [$λ₁, $λ₂], latitude ∈ [$φ₁, $φ₂], z ∈ [$z₁, $z₂]" end -function show(io::IO, g::LatitudeLongitudeGrid{FT, TX, TY, TZ, M}) - show_metrics = M === Nothing ? "metrics are computed on the fly" : - "metrics are pre-computed" +function show(io::IO, g::LatitudeLongitudeGrid) + FT = eltype(g) + TX, TY, TZ = topology(g) + + show_metrics = isnothing(g.Δxᶠᶜᵃ) ? "metrics are computed on the fly" : + "metrics are pre-computed" return print(io, "LatitudeLongitudeGrid{$FT, $TX, $TY, $TZ} \n", " architecture: $(g.architecture)\n", diff --git a/test/test_grids.jl b/test/test_grids.jl index 21c5a5a0ac..009366ff8b 100644 --- a/test/test_grids.jl +++ b/test/test_grids.jl @@ -1,5 +1,6 @@ -using Oceananigans.Architectures -using Oceananigans.Grids: total_extent, halo_size + +include("dependencies_for_runtests.jl") + using Oceananigans.Operators: Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δxᶜᶜᵃ, Δyᶠᶜᵃ, Δyᶜᶠᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, Azᶜᶜᵃ ##### From e7aaf093053047279e8bdb3d2509256276181a87 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 31 Dec 2021 09:15:06 -0700 Subject: [PATCH 017/140] Import total_extent into test_grids --- src/Grids/grid_utils.jl | 1 - test/test_grids.jl | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/src/Grids/grid_utils.jl b/src/Grids/grid_utils.jl index a3e49b9f8e..013a8ec30f 100644 --- a/src/Grids/grid_utils.jl +++ b/src/Grids/grid_utils.jl @@ -30,7 +30,6 @@ Return the architecture (CPU or GPU) that the `grid` lives on. """ @inline architecture(grid::AbstractGrid) = grid.architecture - """ Constant Grid Definitions """ diff --git a/test/test_grids.jl b/test/test_grids.jl index 009366ff8b..40ea894a2b 100644 --- a/test/test_grids.jl +++ b/test/test_grids.jl @@ -1,6 +1,6 @@ - include("dependencies_for_runtests.jl") +using Oceananigans.Grids: total_extent using Oceananigans.Operators: Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δxᶜᶜᵃ, Δyᶠᶜᵃ, Δyᶜᶠᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, Azᶜᶜᵃ ##### From 30952c4036ca422600e91247dd3e842e72b482bb Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 31 Dec 2021 10:04:17 -0700 Subject: [PATCH 018/140] Fixs abstract operatons tests --- src/AbstractOperations/AbstractOperations.jl | 2 ++ .../field_boundary_conditions.jl | 3 ++- src/Fields/field_tuples.jl | 2 +- src/Models/Models.jl | 4 +-- src/Oceananigans.jl | 1 + test/test_abstract_operations.jl | 26 +++++++++---------- 6 files changed, 20 insertions(+), 18 deletions(-) diff --git a/src/AbstractOperations/AbstractOperations.jl b/src/AbstractOperations/AbstractOperations.jl index cfaea3f9f3..19c70961d8 100644 --- a/src/AbstractOperations/AbstractOperations.jl +++ b/src/AbstractOperations/AbstractOperations.jl @@ -1,7 +1,9 @@ module AbstractOperations export ∂x, ∂y, ∂z, @at, @unary, @binary, @multiary +export Δx, Δy, Δz, Ax, Ay, Az, volume export Average, Integral, KernelFunctionOperation +export UnaryOperation, Derivative, BinaryOperation, MultiaryOperation using Base: @propagate_inbounds diff --git a/src/BoundaryConditions/field_boundary_conditions.jl b/src/BoundaryConditions/field_boundary_conditions.jl index aedb75dc85..c1c37b2ffc 100644 --- a/src/BoundaryConditions/field_boundary_conditions.jl +++ b/src/BoundaryConditions/field_boundary_conditions.jl @@ -126,7 +126,8 @@ boundary conditions for prognostic model field boundary conditions. Currently, there is no support `ContinuousBoundaryFunction` for immersed boundary conditions. """ -function regularize_field_boundary_conditions(bcs::FieldBoundaryConditions, grid, field_name, prognostic_field_names=nothing) +function regularize_field_boundary_conditions(bcs::FieldBoundaryConditions, grid::AbstractGrid, field_name, + prognostic_field_names=nothing) topo = topology(grid) loc = assumed_field_location(field_name) diff --git a/src/Fields/field_tuples.jl b/src/Fields/field_tuples.jl index 19712c40e1..5c4016009c 100644 --- a/src/Fields/field_tuples.jl +++ b/src/Fields/field_tuples.jl @@ -11,7 +11,7 @@ Return a `NamedTuple` with fields `u`, `v`, `w` initialized on `grid`. Boundary conditions `bcs` may be specified via a named tuple of `FieldBoundaryCondition`s. """ -function VelocityFields(grid, user_bcs = NamedTuple()) +function VelocityFields(grid::AbstractGrid, user_bcs = NamedTuple()) template = FieldBoundaryConditions() diff --git a/src/Models/Models.jl b/src/Models/Models.jl index 1c459d4051..2a53e598e6 100644 --- a/src/Models/Models.jl +++ b/src/Models/Models.jl @@ -5,7 +5,7 @@ export HydrostaticFreeSurfaceModel, VectorInvariant, ExplicitFreeSurface, ImplicitFreeSurface, HydrostaticSphericalCoriolis, VectorInvariantEnstrophyConserving, - PrescribedVelocityFields + PrescribedVelocityFields, PressureField using Oceananigans: AbstractModel @@ -19,7 +19,7 @@ include("NonhydrostaticModels/NonhydrostaticModels.jl") include("HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl") include("ShallowWaterModels/ShallowWaterModels.jl") -using .NonhydrostaticModels: NonhydrostaticModel +using .NonhydrostaticModels: NonhydrostaticModel, PressureField using .HydrostaticFreeSurfaceModels: HydrostaticFreeSurfaceModel, VectorInvariant, diff --git a/src/Oceananigans.jl b/src/Oceananigans.jl index 36e5c01d40..9a2df20e5b 100644 --- a/src/Oceananigans.jl +++ b/src/Oceananigans.jl @@ -67,6 +67,7 @@ export NonhydrostaticModel, HydrostaticFreeSurfaceModel, ShallowWaterModel, + PressureField, fields, # Hydrostatic free surface model stuff diff --git a/test/test_abstract_operations.jl b/test/test_abstract_operations.jl index 9e2c6aad0e..876d78474f 100644 --- a/test/test_abstract_operations.jl +++ b/test/test_abstract_operations.jl @@ -1,11 +1,9 @@ -using Oceananigans.AbstractOperations: UnaryOperation, Derivative, BinaryOperation, MultiaryOperation -using Oceananigans.AbstractOperations: KernelFunctionOperation +include("dependencies_for_runtests.jl") + using Oceananigans.Operators: ℑxyᶜᶠᵃ, ℑxyᶠᶜᵃ -using Oceananigans.Fields: PressureField, compute_at! +using Oceananigans.Fields: compute_at! using Oceananigans.BuoyancyModels: BuoyancyField -using Oceananigans.AbstractOperations: Δx, Δy, Δz, Ax, Ay, Az, volume - function simple_binary_operation(op, a, b, num1, num2) a_b = op(a, b) interior(a) .= num1 @@ -68,7 +66,7 @@ end function x_derivative_cell(arch) grid = RectilinearGrid(arch, size=(3, 3, 3), extent=(3, 3, 3)) - a = Field(Center, Center, Center, arch, grid, nothing) + a = Field{Center, Center, Center}(grid) dx_a = ∂x(a) one_four_four = arch_array(arch, [1, 4, 4]) @@ -92,8 +90,8 @@ for arch in archs @info "Testing abstract operations [$(typeof(arch))]..." grid = RectilinearGrid(arch, size=(3, 3, 3), extent=(3, 3, 3)) - u, v, w = VelocityFields(arch, grid) - c = Field(Center, Center, Center, arch, grid, nothing) + u, v, w = VelocityFields(grid) + c = Field{Center, Center, Center}(grid) @testset "Unary operations and derivatives [$(typeof(arch))]" begin for ψ in (u, v, w, c) @@ -149,8 +147,8 @@ for arch in archs num2 = Float64(42) grid = RectilinearGrid(arch, size=(3, 3, 3), extent=(3, 3, 3)) - u, v, w = VelocityFields(arch, grid) - T, S = TracerFields((:T, :S), arch, grid) + u, v, w = VelocityFields(grid) + T, S = TracerFields((:T, :S), grid) for op in (+, *, -, /) @test simple_binary_operation(op, u, v, num1, num2) @@ -166,8 +164,8 @@ for arch in archs grid = RectilinearGrid(arch, size=(3, 3, 3), extent=(3, 3, 3), topology=(Periodic, Periodic, Periodic)) - u, v, w = VelocityFields(arch, grid) - T, S = TracerFields((:T, :S), arch, grid) + u, v, w = VelocityFields(grid) + T, S = TracerFields((:T, :S), grid) for a in (u, v, w, T) @test x_derivative(a) @test y_derivative(a) @@ -182,7 +180,7 @@ for arch in archs arch = CPU() Nx = 3 # Δx=1, xC = 0.5, 1.5, 2.5 grid = RectilinearGrid(arch, size=(Nx, Nx, Nx), extent=(Nx, Nx, Nx)) - a, b = (Field(Center, Center, Center, arch, grid, nothing) for i in 1:2) + a, b = (Field{Center, Center, Center}(grid) for i in 1:2) set!(b, 2) set!(a, (x, y, z) -> x < 2 ? 3x : 6) @@ -268,7 +266,7 @@ for arch in archs @testset "BinaryOperations with GridMetricOperation [$(typeof(arch))]" begin grid = RectilinearGrid(arch, size=(1, 1, 1), extent=(2, 3, 4)) - c = CenterField(arch, grid) + c = CenterField(grid) c .= 1 # Δx, Δy, Δz = 2, 3, 4 From 8d9672f2a7252aad305d046dd24c11878c580ba0 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Sun, 2 Jan 2022 18:56:59 -0700 Subject: [PATCH 019/140] Multifarous updates --- src/Distributed/multi_architectures.jl | 3 +- src/Grids/Grids.jl | 2 +- src/Grids/grid_utils.jl | 2 + src/Grids/latitude_longitude_grid.jl | 2 +- src/Grids/rectilinear_grid.jl | 4 +- src/Grids/zeros.jl | 3 +- .../pcg_implicit_free_surface_solver.jl | 2 +- src/OutputReaders/field_time_series.jl | 91 ++++++++++--------- src/OutputWriters/jld2_output_writer.jl | 2 +- test/dependencies_for_runtests.jl | 3 +- test/test_jld2_output_writer.jl | 17 +--- test/test_netcdf_output_writer.jl | 13 +-- test/test_output_readers.jl | 39 +++----- test/test_output_writers.jl | 12 +-- 14 files changed, 83 insertions(+), 112 deletions(-) diff --git a/src/Distributed/multi_architectures.jl b/src/Distributed/multi_architectures.jl index 12e2871193..824d38501f 100644 --- a/src/Distributed/multi_architectures.jl +++ b/src/Distributed/multi_architectures.jl @@ -1,5 +1,5 @@ using Oceananigans.Architectures -using Oceananigans.Grids: topology, validate_tupled_argument, with_arch +using Oceananigans.Grids: topology, validate_tupled_argument import Oceananigans.Architectures: device, device_event, arch_array import Oceananigans.Grids: zeros @@ -13,7 +13,6 @@ struct MultiArch{A, R, I, ρ, C, γ} <: AbstractMultiArchitecture communicator :: γ end - ##### ##### Constructors ##### diff --git a/src/Grids/Grids.jl b/src/Grids/Grids.jl index cd126e4cb8..3d3d52c502 100644 --- a/src/Grids/Grids.jl +++ b/src/Grids/Grids.jl @@ -11,6 +11,7 @@ export LatitudeLongitudeGrid, XRegLatLonGrid, YRegLatLonGrid, ZRegLatLonGrid export ConformalCubedSphereFaceGrid, ConformalCubedSphereGrid export node, xnode, ynode, znode, xnodes, ynodes, znodes, nodes export offset_data, new_data +export on_architecture using CUDA using Adapt @@ -113,7 +114,6 @@ Abstract supertype for horizontally-curvilinear grids with elements of type `FT` """ abstract type AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ, Arch} <: AbstractCurvilinearGrid{FT, TX, TY, TZ, Arch} end - include("grid_utils.jl") include("zeros.jl") include("new_data.jl") diff --git a/src/Grids/grid_utils.jl b/src/Grids/grid_utils.jl index 013a8ec30f..f2cc458712 100644 --- a/src/Grids/grid_utils.jl +++ b/src/Grids/grid_utils.jl @@ -4,6 +4,8 @@ using CUDA ##### Convenience functions ##### +on_architecture(::Nothing, grid) = grid + Base.length(::Type{Face}, topo, N) = N Base.length(::Type{Face}, ::Type{Bounded}, N) = N+1 Base.length(::Type{Center}, topo, N) = N diff --git a/src/Grids/latitude_longitude_grid.jl b/src/Grids/latitude_longitude_grid.jl index 913346313f..b39cd6baba 100644 --- a/src/Grids/latitude_longitude_grid.jl +++ b/src/Grids/latitude_longitude_grid.jl @@ -223,7 +223,7 @@ all_z_nodes(::Type{Face}, grid::LatitudeLongitudeGrid) = grid.zᵃᵃᶠ @inline cpu_face_constructor_y(grid::YRegLatLonGrid) = y_domain(grid) @inline cpu_face_constructor_z(grid::ZRegLatLonGrid) = z_domain(grid) -function with_arch(new_arch, old_grid::LatitudeLongitudeGrid) +function on_architecture(new_arch, old_grid::LatitudeLongitudeGrid) size = (old_grid.Nx, old_grid.Ny, old_grid.Nz) topo = topology(old_grid) diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl index 6cf71114a4..b3b0f633f6 100644 --- a/src/Grids/rectilinear_grid.jl +++ b/src/Grids/rectilinear_grid.jl @@ -370,8 +370,6 @@ function Adapt.adapt_structure(to, grid::RectilinearGrid) Adapt.adapt(to, grid.zᵃᵃᶜ)) end - - @inline xnode(::Center, i, grid::RectilinearGrid) = @inbounds grid.xᶜᵃᵃ[i] @inline xnode(::Face , i, grid::RectilinearGrid) = @inbounds grid.xᶠᵃᵃ[i] @@ -415,7 +413,7 @@ function with_halo(new_halo, old_grid::RectilinearGrid) return new_grid end -function with_arch(new_arch, old_grid::RectilinearGrid) +function on_architecture(new_arch, old_grid::RectilinearGrid) size = (old_grid.Nx, old_grid.Ny, old_grid.Nz) topo = topology(old_grid) diff --git a/src/Grids/zeros.jl b/src/Grids/zeros.jl index 2fef74ca1c..b5399a20a0 100644 --- a/src/Grids/zeros.jl +++ b/src/Grids/zeros.jl @@ -3,9 +3,8 @@ using Oceananigans.Architectures: CPU, GPU, AbstractMultiArchitecture import Base: zeros - - zeros(FT, ::CPU, N...) = zeros(FT, N...) zeros(FT, ::GPU, N...) = CUDA.zeros(FT, N...) zeros(arch::AbstractArchitecture, grid, N...) = zeros(eltype(grid), arch, N...) +zeros(grid::AbstractGrid, N...) = zeros(eltype(grid), architecture(grid), N...) diff --git a/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl b/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl index 60c6efa9b4..428ccb8df9 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl @@ -28,7 +28,7 @@ representing an implicit time discretization of the linear free surface evolutio for a fluid with variable depth `H`, horizontal areas `Az`, barotropic volume flux `Q★`, time step `Δt`, gravitational acceleration `g`, and free surface at time-step `n` `ηⁿ`. """ -function PCGImplicitFreeSurfaceSolver(grid, gravitational_acceleration, settings) +function PCGImplicitFreeSurfaceSolver(grid::AbstractGrid, gravitational_acceleration::Number, settings) # Initialize vertically integrated lateral face areas ∫ᶻ_Axᶠᶜᶜ = Field{Face, Center, Nothing}(grid) ∫ᶻ_Ayᶜᶠᶜ = Field{Center, Face, Nothing}(grid) diff --git a/src/OutputReaders/field_time_series.jl b/src/OutputReaders/field_time_series.jl index 620943edf4..fe02105216 100644 --- a/src/OutputReaders/field_time_series.jl +++ b/src/OutputReaders/field_time_series.jl @@ -12,46 +12,44 @@ using Oceananigans.Fields: show_location import Oceananigans: short_show import Oceananigans.Fields: Field, set!, interior +import Oceananigans.Architectures: architecture struct FieldTimeSeries{X, Y, Z, K, A, T, D, G, B, χ} <: AbstractField{X, Y, Z, A, G, T, 4} data :: D - architecture :: A grid :: G boundary_conditions :: B times :: χ - function FieldTimeSeries{X, Y, Z, K}(data::D, arch::A, grid::G, bcs::B, times::χ) where {X, Y, Z, K, D, A, G, B, χ} + function FieldTimeSeries{X, Y, Z, K}(data::D, grid::G, bcs::B, times::χ) where {X, Y, Z, K, D, G, B, χ} T = eltype(grid) - return new{X, Y, Z, K, A, T, D, G, B, χ}(data, arch, grid, bcs, times) + A = typeof(architecture(grid)) + return new{X, Y, Z, K, A, T, D, G, B, χ}(data, grid, bcs, times) end end +architecture(fts::FieldTimeSeries) = architecture(fts.grid) + ##### ##### Constructors ##### """ - FieldTimeSeries{LX, LY, LZ}([architecture = CPU()], grid, times, boundary_conditions=nothing) + FieldTimeSeries{LX, LY, LZ}(grid, times, boundary_conditions=nothing) Return `FieldTimeSeries` at location `(LX, LY, LZ)`, on `grid`, at `times`, with `boundary_conditions`, and initialized with zeros of `eltype(grid)`. """ -function FieldTimeSeries{LX, LY, LZ}(architecture, grid, times, boundary_conditions=nothing) where {LX, LY, LZ} +function FieldTimeSeries{LX, LY, LZ}(grid, times; boundary_conditions=nothing) where {LX, LY, LZ} location = (LX, LY, LZ) Nt = length(times) data_size = total_size(location, grid) - raw_data = zeros(architecture, grid, data_size..., Nt) + raw_data = zeros(grid, data_size..., Nt) data = offset_data(raw_data, grid, location) - return FieldTimeSeries{LX, LY, LZ, InMemory}(data, architecture, grid, boundary_conditions, times) + return FieldTimeSeries{LX, LY, LZ, InMemory}(data, grid, boundary_conditions, times) end -# CPU() default -FieldTimeSeries{LX, LY, LZ}(grid::AbstractGrid, times, bcs=nothing) where {LX, LY, LZ} = - FieldTimeSeries{LX, LY, LZ}(CPU(), grid, times, bcs) - """ FieldTimeSeries(path, name; - architecture = CPU(), backend = InMemory(), grid = nothing, iterations = nothing, @@ -63,8 +61,6 @@ located at `path`. Keyword arguments ================= -- `architecture`: The architecture on which to store time series data. `CPU()` by default. - - `backend`: `InMemory()` to load data into a 4D array or `OnDisk()` to lazily load data from disk when indexing into `FieldTimeSeries`. @@ -77,8 +73,8 @@ Keyword arguments comparison to recorded save times. Defaults to times associated with `iterations`. Takes precedence over `iterations` if `times` is specified. """ -FieldTimeSeries(path, name; architecture=CPU(), backend=InMemory(), kwargs...) = - FieldTimeSeries(path, name, architecture, backend; kwargs...) +FieldTimeSeries(path, name; backend=InMemory(), kwargs...) = + FieldTimeSeries(path, name, backend; kwargs...) ##### ##### InMemory time serieses @@ -88,7 +84,8 @@ const InMemoryFieldTimeSeries{X, Y, Z} = FieldTimeSeries{X, Y, Z, InMemory} struct UnspecifiedBoundaryConditions end -function FieldTimeSeries(path, name, architecture, backend::InMemory; +function FieldTimeSeries(path, name, backend::InMemory; + architecture = nothing, grid = nothing, location = nothing, boundary_conditions = UnspecifiedBoundaryConditions(), @@ -97,55 +94,63 @@ function FieldTimeSeries(path, name, architecture, backend::InMemory; file = jldopen(path) - # Non-defaults - isnothing(grid) && (grid = file["serialized/grid"]) - 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"]) + # 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"]) + + if isnothing(grid) + grid = on_architecture(architecture, file["serialized/grid"]) + end if boundary_conditions isa UnspecifiedBoundaryConditions boundary_conditions = file["timeseries/$name/serialized/boundary_conditions"] end LX, LY, LZ = location - - time_series = FieldTimeSeries{LX, LY, LZ}(architecture, grid, times, boundary_conditions) - + time_series = FieldTimeSeries{LX, LY, LZ}(grid, times; boundary_conditions) set!(time_series, path, name) - + return time_series end Base.getindex(fts::InMemoryFieldTimeSeries{LX, LY, LZ}, n::Int) where {LX, LY, LZ} = - Field(LX, LY, LZ, fts.architecture, fts.grid, fts.boundary_conditions, view(fts.data, :, :, :, n)) + Field{LX, LY, LZ}(fts.grid, boundary_conditions=fts.boundary_conditions, data=view(fts.data, :, :, :, n)) ##### ##### set! ##### """ - Field(path::String, name::String, iter; architecture=GPU(), grid=nothing) + Field(path::String, name::String, iter; grid=nothing) Load a Field saved in JLD2 file at `path`, with `name` and at `iter`ation. -`architecture = CPU()` by default, and `grid` is loaded from `path` if not specified. +`grid` is loaded from `path` if not specified. """ function Field(location, path::String, name::String, iter; - architecture = CPU(), grid = nothing, + architecture = nothing, boundary_conditions = nothing) file = jldopen(path) + # Default to CPU if neither architecture nor grid is specified + architecture = isnothing(architecture) ? + (isnothing(grid) ? CPU() : Architectures.architecture(grid)) : + architecture + + grid = isnothing(grid) ? + on_architecture(architecture, file["serialized/grid"]) : grid + raw_data = arch_array(architecture, file["timeseries/$name/$iter"]) - isnothing(grid) && (grid = file["serialized/grid"]) close(file) try data = offset_data(raw_data, grid, location) - return Field(location, architecture, grid, boundary_conditions, data) + return Field(location, grid; boundary_conditions, data) catch - field = Field(location, architecture, grid, boundary_conditions) + field = Field(location, grid; boundary_conditions) interior(field) .= raw_data return field end @@ -174,7 +179,6 @@ end function set!(time_series::FieldTimeSeries, fields_vector::AbstractVector{<:AbstractField}) raw_data = parent(time_series.data) - ArrayType = array_type(time_series.architecture) file = jldopen(path) @@ -191,7 +195,7 @@ end # is there a better way? # FieldTimeSeries[i] returns ViewField -const ViewField = Field{<:Any, <:Any, <:Any, <:Any, <:SubArray} +const ViewField = Field{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:SubArray} using OffsetArrays: IdOffsetRange @@ -218,11 +222,11 @@ struct OnDiskData name :: String end -function FieldTimeSeries(path, name, architecture, backend::OnDisk; grid=nothing) +function FieldTimeSeries(path, name, backend::OnDisk; architecture=nothing, grid=nothing) file = jldopen(path) if isnothing(grid) - grid = file["serialized/grid"] + grid = on_architecture(architecture, file["serialized/grid"]) end iterations = parse.(Int, keys(file["timeseries/t"])) @@ -234,20 +238,20 @@ function FieldTimeSeries(path, name, architecture, backend::OnDisk; grid=nothing close(file) - return FieldTimeSeries{LX, LY, LZ, OnDisk}(data, architecture, grid, bcs, times) + return FieldTimeSeries{LX, LY, LZ, OnDisk}(data, grid, bcs, times) end # For creating an empty `FieldTimeSeries`. -function FieldTimeSeries(grid, location, times; architecture=CPU(), ArrayType=array_type(architecture), name="", filepath="", bcs=nothing) +function FieldTimeSeries(grid, location, times; name="", filepath="", bcs=nothing) LX, LY, LZ = location Nt = length(times) data_size = total_size(location, grid) - raw_data = zeros(data_size..., Nt) |> ArrayType + raw_data = zeros(grid, data_size..., Nt) data = offset_data(raw_data, grid, location) - return FieldTimeSeries{LX, LY, LZ}(InMemory(), data, architecture, grid, bcs, times, name, filepath, 4) + return FieldTimeSeries{LX, LY, LZ}(InMemory(), data, grid, bcs, times, name, filepath, 4) end ##### @@ -261,16 +265,17 @@ end function Base.getindex(fts::FieldTimeSeries{X, Y, Z, OnDisk}, n::Int) where {X, Y, Z} # Load data + arch = architecture(fts) file = jldopen(fts.data.path) iter = keys(file["timeseries/t"])[n] - raw_data = file["timeseries/$(fts.data.name)/$iter"] |> array_type(fts.architecture) + raw_data = arch_array(architecture(fts), file["timeseries/$(fts.data.name)/$iter"]) close(file) # Wrap Field loc = (X, Y, Z) field_data = offset_data(raw_data, fts.grid, loc) - return Field(loc..., fts.architecture, fts.grid, fts.boundary_conditions, field_data) + return Field(loc, fts.grid; boundary_conditions=fts.boundary_conditions, data=field_data) end Base.setindex!(fts::FieldTimeSeries, val, inds...) = Base.setindex!(fts.data, val, inds...) diff --git a/src/OutputWriters/jld2_output_writer.jl b/src/OutputWriters/jld2_output_writer.jl index f8ccd10c21..22299aba4a 100644 --- a/src/OutputWriters/jld2_output_writer.jl +++ b/src/OutputWriters/jld2_output_writer.jl @@ -123,7 +123,7 @@ function init_save_some_metadata!(file, model) return nothing end -c_avg = AveragedField(model.tracers.c, dims=(1, 2)) +c_avg = Field(Average(model.tracers.c, dims=(1, 2))) # Note that model.velocities is NamedTuple simulation.output_writers[:velocities] = JLD2OutputWriter(model, model.velocities, diff --git a/test/dependencies_for_runtests.jl b/test/dependencies_for_runtests.jl index bb5c16251d..08d58620a6 100644 --- a/test/dependencies_for_runtests.jl +++ b/test/dependencies_for_runtests.jl @@ -33,8 +33,9 @@ using Oceananigans.Distributed using Oceananigans.Logger using Oceananigans.Units using Oceananigans.Utils -using Oceananigans.Architectures: device # to resolve conflict with CUDA.device +using Oceananigans.Architectures: device, array_type # to resolve conflict with CUDA.device +using Oceananigans: Clock using Dates: DateTime, Nanosecond using TimesDates: TimeDate using Statistics: mean diff --git a/test/test_jld2_output_writer.jl b/test/test_jld2_output_writer.jl index b3c9c9c72b..b58e2c9682 100644 --- a/test/test_jld2_output_writer.jl +++ b/test/test_jld2_output_writer.jl @@ -1,13 +1,4 @@ -using Test -using CUDA -using JLD2 -using Oceananigans -using Oceananigans.Units -using Oceananigans: Clock - -include("utils_for_runtests.jl") - -archs = test_architectures() +include("dependencies_for_runtests.jl") ##### ##### JLD2OutputWriter tests @@ -105,9 +96,9 @@ function test_jld2_time_averaging_of_horizontal_averages(model) Δt = 0.1 simulation = Simulation(model, Δt=Δt, stop_iteration=5) - average_fluxes = (wu = AveragedField(w * u, dims=(1, 2)), - uv = AveragedField(u * v, dims=(1, 2)), - wT = AveragedField(w * T, dims=(1, 2))) + average_fluxes = (wu = Field(Average(w * u, dims=(1, 2))), + uv = Field(Average(u * v, dims=(1, 2))), + wT = Field(Average(w * T, dims=(1, 2)))) simulation.output_writers[:fluxes] = JLD2OutputWriter(model, average_fluxes, schedule = AveragedTimeInterval(4Δt, window=2Δt), diff --git a/test/test_netcdf_output_writer.jl b/test/test_netcdf_output_writer.jl index 8ac820b368..c1e72c0af6 100644 --- a/test/test_netcdf_output_writer.jl +++ b/test/test_netcdf_output_writer.jl @@ -1,16 +1,11 @@ -using Test +include("dependencies_for_runtests.jl") + using Dates: DateTime, Nanosecond, Millisecond using TimesDates: TimeDate using CUDA using NCDatasets -using Oceananigans -using Oceananigans.Units using Oceananigans: Clock -include("utils_for_runtests.jl") - -archs = test_architectures() - ##### ##### NetCDFOutputWriter tests ##### @@ -522,8 +517,8 @@ function test_netcdf_time_averaging(arch) Δt = 1/64 # Nice floating-point number simulation = Simulation(model, Δt=Δt, stop_time=50Δt) - ∫c1_dxdy = AveragedField(model.tracers.c1, dims=(1, 2)) - ∫c2_dxdy = AveragedField(model.tracers.c2, dims=(1, 2)) + ∫c1_dxdy = Field(Average(model.tracers.c1, dims=(1, 2))) + ∫c2_dxdy = Field(Average(model.tracers.c2, dims=(1, 2))) nc_outputs = Dict("c1" => ∫c1_dxdy, "c2" => ∫c2_dxdy) nc_dimensions = Dict("c1" => ("zC",), "c2" => ("zC",)) diff --git a/test/test_output_readers.jl b/test/test_output_readers.jl index e04c2cec58..bd1f7dd277 100644 --- a/test/test_output_readers.jl +++ b/test/test_output_readers.jl @@ -1,11 +1,4 @@ -using Test -using Statistics -using JLD2 - -using Oceananigans -using Oceananigans.Units -using Oceananigans.Architectures: array_type -using Oceananigans.Fields: location +include("dependencies_for_runtests.jl") function generate_some_interesting_simulation_data(Nx, Ny, Nz; architecture=CPU()) grid = RectilinearGrid(architecture, size=(Nx, Ny, Nz), extent=(64, 64, 32)) @@ -17,12 +10,8 @@ function generate_some_interesting_simulation_data(Nx, Ny, Nz; architecture=CPU( evaporation_bc = FluxBoundaryCondition(Qˢ, field_dependencies=:S, parameters=3e-7) S_bcs = FieldBoundaryConditions(top=evaporation_bc) - model = NonhydrostaticModel( - grid = grid, - tracers = (:T, :S), - buoyancy = SeawaterBuoyancy(), - boundary_conditions = (u=u_bcs, T=T_bcs, S=S_bcs) - ) + model = NonhydrostaticModel(; grid, tracers = (:T, :S), buoyancy = SeawaterBuoyancy(), + boundary_conditions = (u=u_bcs, T=T_bcs, S=S_bcs)) dTdz = 0.01 Tᵢ(x, y, z) = 20 + dTdz * z + 1e-6 * randn() @@ -37,27 +26,27 @@ function generate_some_interesting_simulation_data(Nx, Ny, Nz; architecture=CPU( computed_fields = ( b = BuoyancyField(model), - ζ = ComputedField(∂x(v) - ∂y(u)), - ke = ComputedField(√(u^2 + v^2)) + ζ = Field(∂x(v) - ∂y(u)), + ke = Field(√(u^2 + v^2)) ) fields_to_output = merge(model.velocities, model.tracers, computed_fields) simulation.output_writers[:jld2_3d_with_halos] = JLD2OutputWriter(model, fields_to_output, - prefix = "test_3d_output_with_halos", - field_slicer = FieldSlicer(with_halos=true), - schedule = TimeInterval(30seconds), - force = true) + prefix = "test_3d_output_with_halos", + field_slicer = FieldSlicer(with_halos=true), + schedule = TimeInterval(30seconds), + force = true) - profiles = NamedTuple{keys(fields_to_output)}(AveragedField(f, dims=(1, 2)) for f in fields_to_output) + profiles = NamedTuple{keys(fields_to_output)}(Field(Average(f, dims=(1, 2))) for f in fields_to_output) simulation.output_writers[:jld2_1d_with_halos] = JLD2OutputWriter(model, profiles, - prefix = "test_1d_output_with_halos", - field_slicer = FieldSlicer(with_halos=true), - schedule = TimeInterval(30seconds), - force = true) + prefix = "test_1d_output_with_halos", + field_slicer = FieldSlicer(with_halos=true), + schedule = TimeInterval(30seconds), + force = true) run!(simulation) diff --git a/test/test_output_writers.jl b/test/test_output_writers.jl index 2596e2fd0f..76c54f6a46 100644 --- a/test/test_output_writers.jl +++ b/test/test_output_writers.jl @@ -1,21 +1,13 @@ +include("dependencies_for_runtests.jl") + using Statistics using NCDatasets -using Test - -using Oceananigans -using Oceananigans.Diagnostics -using Oceananigans.Fields -using Oceananigans.OutputWriters using Dates: Millisecond using Oceananigans: write_output! using Oceananigans.BoundaryConditions: PBC, FBC, ZFBC, ContinuousBoundaryFunction using Oceananigans.TimeSteppers: update_state! -include("utils_for_runtests.jl") - -archs = test_architectures() - ##### ##### WindowedTimeAverage tests ##### From f6fbf49fc1834b828e73241e4063912fa96d1c53 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Sun, 2 Jan 2022 20:07:05 -0700 Subject: [PATCH 020/140] Fixes some tests --- test/dependencies_for_runtests.jl | 2 +- test/test_abstract_operations.jl | 23 ++++++--- test/test_lagrangian_particle_tracking.jl | 61 ++++++++++++----------- test/test_output_readers.jl | 8 +++ 4 files changed, 56 insertions(+), 38 deletions(-) diff --git a/test/dependencies_for_runtests.jl b/test/dependencies_for_runtests.jl index 08d58620a6..9180d747c8 100644 --- a/test/dependencies_for_runtests.jl +++ b/test/dependencies_for_runtests.jl @@ -19,6 +19,7 @@ using Oceananigans.Operators using Oceananigans.Advection using Oceananigans.BoundaryConditions using Oceananigans.Fields +using Oceananigans.AbstractOperations using Oceananigans.Coriolis using Oceananigans.BuoyancyModels using Oceananigans.Forcings @@ -28,7 +29,6 @@ using Oceananigans.Simulations using Oceananigans.Diagnostics using Oceananigans.OutputWriters using Oceananigans.TurbulenceClosures -using Oceananigans.AbstractOperations using Oceananigans.Distributed using Oceananigans.Logger using Oceananigans.Units diff --git a/test/test_abstract_operations.jl b/test/test_abstract_operations.jl index 876d78474f..bfe47e589d 100644 --- a/test/test_abstract_operations.jl +++ b/test/test_abstract_operations.jl @@ -251,7 +251,14 @@ for arch in archs end end - for metric in (Δx, Δy, Δz, Ax, Ay, Az, volume) + for metric in (AbstractOperations.Δx, + AbstractOperations.Δy, + AbstractOperations.Δz, + AbstractOperations.Ax, + AbstractOperations.Ay, + AbstractOperations.Az, + AbstractOperations.volume) + @test location(metric * ϕ) == location(ϕ) end end @@ -272,13 +279,13 @@ for arch in archs # Δx, Δy, Δz = 2, 3, 4 # Ax, Ay, Az = 12, 8, 6 # volume = 24 - op = c * Δx; @test op[1, 1, 1] == 2 - op = c * Δy; @test op[1, 1, 1] == 3 - op = c * Δz; @test op[1, 1, 1] == 4 - op = c * Ax; @test op[1, 1, 1] == 12 - op = c * Ay; @test op[1, 1, 1] == 8 - op = c * Az; @test op[1, 1, 1] == 6 - op = c * volume; @test op[1, 1, 1] == 24 + op = c * AbstractOperations.Δx; @test op[1, 1, 1] == 2 + op = c * AbstractOperations.Δy; @test op[1, 1, 1] == 3 + op = c * AbstractOperations.Δz; @test op[1, 1, 1] == 4 + op = c * AbstractOperations.Ax; @test op[1, 1, 1] == 12 + op = c * AbstractOperations.Ay; @test op[1, 1, 1] == 8 + op = c * AbstractOperations.Az; @test op[1, 1, 1] == 6 + op = c * AbstractOperations.volume; @test op[1, 1, 1] == 24 end end end diff --git a/test/test_lagrangian_particle_tracking.jl b/test/test_lagrangian_particle_tracking.jl index 45ae81ca1e..196c9750a1 100644 --- a/test/test_lagrangian_particle_tracking.jl +++ b/test/test_lagrangian_particle_tracking.jl @@ -1,3 +1,5 @@ +include("dependencies_for_runtests.jl") + using NCDatasets using StructArrays @@ -17,9 +19,9 @@ function run_simple_particle_tracking_tests(arch, timestepper) grid = RectilinearGrid(arch, topology=topo, size=(5, 5, 5); domain...) P = 10 - xs = convert(array_type(arch), zeros(P)) - ys = convert(array_type(arch), zeros(P)) - zs = convert(array_type(arch), 0.5*ones(P)) + xs = arch_array(arch, zeros(P)) + ys = arch_array(arch, zeros(P)) + zs = arch_array(arch, 0.5*ones(P)) # Test first constructor lagrangian_particles = LagrangianParticles(x=xs, y=ys, z=zs) @@ -32,11 +34,11 @@ function run_simple_particle_tracking_tests(arch, timestepper) particles = StructArray{TestParticle}((xs, ys, zs, us, vs, ws, ss)) - velocities = VelocityFields(arch, grid) + velocities = VelocityFields(grid) u, v, w = velocities - speed = ComputedField(√(u*u + v*v + w*w)) + speed = Field(√(u*u + v*v + w*w)) - tracked_fields = merge(velocities, (s=speed,)) + tracked_fields = merge(velocities, (; s=speed)) # Test second constructor lagrangian_particles = LagrangianParticles(particles; tracked_fields) @@ -51,7 +53,8 @@ function run_simple_particle_tracking_tests(arch, timestepper) jld2_filepath = "test_particles.jld2" sim.output_writers[:particles_jld2] = - JLD2OutputWriter(model, (particles=model.particles,), prefix="test_particles", schedule=IterationInterval(1)) + JLD2OutputWriter(model, (; particles=model.particles), + prefix="test_particles", schedule=IterationInterval(1)) nc_filepath = "test_particles.nc" sim.output_writers[:particles_nc] = @@ -63,7 +66,7 @@ function run_simple_particle_tracking_tests(arch, timestepper) run!(sim) @test length(model.particles) == P - @test size(model.particles) == (P,) + @test size(model.particles) == tuple(P) @test propertynames(model.particles.properties) == (:x, :y, :z, :u, :v, :w, :s) x = convert(array_type(arch), model.particles.properties.x) @@ -74,13 +77,13 @@ function run_simple_particle_tracking_tests(arch, timestepper) w = convert(array_type(arch), model.particles.properties.w) s = convert(array_type(arch), model.particles.properties.s) - @test size(x) == (P,) - @test size(y) == (P,) - @test size(z) == (P,) - @test size(u) == (P,) - @test size(v) == (P,) - @test size(w) == (P,) - @test size(s) == (P,) + @test size(x) == tuple(P) + @test size(y) == tuple(P) + @test size(z) == tuple(P) + @test size(u) == tuple(P) + @test size(v) == tuple(P) + @test size(w) == tuple(P) + @test size(s) == tuple(P) @test all(x .≈ 0.01) @test all(y .≈ 0.01) @@ -120,13 +123,13 @@ function run_simple_particle_tracking_tests(arch, timestepper) @test haskey(file["timeseries/particles"], "0") @test haskey(file["timeseries/particles"], "0") - @test size(file["timeseries/particles/1"].x) == (P,) - @test size(file["timeseries/particles/1"].y) == (P,) - @test size(file["timeseries/particles/1"].z) == (P,) - @test size(file["timeseries/particles/1"].u) == (P,) - @test size(file["timeseries/particles/1"].v) == (P,) - @test size(file["timeseries/particles/1"].w) == (P,) - @test size(file["timeseries/particles/1"].s) == (P,) + @test size(file["timeseries/particles/1"].x) == tuple(P) + @test size(file["timeseries/particles/1"].y) == tuple(P) + @test size(file["timeseries/particles/1"].z) == tuple(P) + @test size(file["timeseries/particles/1"].u) == tuple(P) + @test size(file["timeseries/particles/1"].v) == tuple(P) + @test size(file["timeseries/particles/1"].w) == tuple(P) + @test size(file["timeseries/particles/1"].s) == tuple(P) @test all(file["timeseries/particles/1"].x .≈ 0.01) @test all(file["timeseries/particles/1"].y .≈ 0.01) @@ -160,13 +163,13 @@ function run_simple_particle_tracking_tests(arch, timestepper) @test model.particles.properties isa StructArray - @test size(x) == (P,) - @test size(y) == (P,) - @test size(z) == (P,) - @test size(u) == (P,) - @test size(v) == (P,) - @test size(w) == (P,) - @test size(s) == (P,) + @test size(x) == tuple(P) + @test size(y) == tuple(P) + @test size(z) == tuple(P) + @test size(u) == tuple(P) + @test size(v) == tuple(P) + @test size(w) == tuple(P) + @test size(s) == tuple(P) @test all(x .≈ 0.01) @test all(y .≈ 0.01) diff --git a/test/test_output_readers.jl b/test/test_output_readers.jl index bd1f7dd277..2303cf6fe5 100644 --- a/test/test_output_readers.jl +++ b/test/test_output_readers.jl @@ -76,6 +76,14 @@ end b3 = FieldTimeSeries(filepath3d, "b", architecture=arch) ζ3 = FieldTimeSeries(filepath3d, "ζ", architecture=arch) + # This behavior ensures that set! works + # but perhaps should be changed in the future + @test parent(u3) isa SubArray + @test parent(v3) isa SubArray + @test parent(w3) isa SubArray + @test parent(T3) isa SubArray + @test parent(b3) isa SubArray + @test location(u3) == (Face, Center, Center) @test location(v3) == (Center, Face, Center) @test location(w3) == (Center, Center, Face) From ceec2dc9514ea8cf871bf020db6b173bc104bbb2 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Sun, 2 Jan 2022 20:34:32 -0700 Subject: [PATCH 021/140] Fix shallow water tests --- src/Models/ShallowWaterModels/shallow_water_model.jl | 12 ++++++++---- test/test_shallow_water_models.jl | 5 ++--- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/src/Models/ShallowWaterModels/shallow_water_model.jl b/src/Models/ShallowWaterModels/shallow_water_model.jl index 0e8d843475..3996a33105 100644 --- a/src/Models/ShallowWaterModels/shallow_water_model.jl +++ b/src/Models/ShallowWaterModels/shallow_water_model.jl @@ -11,7 +11,9 @@ using Oceananigans.TimeSteppers: Clock, TimeStepper, update_state! using Oceananigans.TurbulenceClosures: with_tracers, DiffusivityFields using Oceananigans.Utils: tupleit -function ShallowWaterTendencyFields(arch, grid, tracer_names) +import Oceananigans.Architectures: architecture + +function ShallowWaterTendencyFields(grid, tracer_names) uh = XFaceField(grid) vh = YFaceField(grid) @@ -24,9 +26,9 @@ function ShallowWaterTendencyFields(arch, grid, tracer_names) end function ShallowWaterSolutionFields(grid, bcs) - uh = XFaceField(grid, bcs.uh) - vh = YFaceField(grid, bcs.vh) - h = CenterField(grid, bcs.h) + uh = XFaceField(grid, boundary_conditions=bcs.uh) + vh = YFaceField(grid, boundary_conditions=bcs.vh) + h = CenterField(grid, boundary_conditions=bcs.h) return (; uh, vh, h) end @@ -147,3 +149,5 @@ function ShallowWaterModel(; return model end + +architecture(model::ShallowWaterModel) = model.architecture diff --git a/test/test_shallow_water_models.jl b/test/test_shallow_water_models.jl index cb31c8cdd1..556888b360 100644 --- a/test/test_shallow_water_models.jl +++ b/test/test_shallow_water_models.jl @@ -1,6 +1,5 @@ -using Oceananigans -using Oceananigans.Models -using Oceananigans.Grids +include("dependencies_for_runtests.jl") + using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBoundary function time_stepping_shallow_water_model_works(arch, topo, coriolis, advection; timestepper=:RungeKutta3) From 09d3aa64ef0c4e41c4d3102b33836cd7970f9798 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Sun, 2 Jan 2022 20:53:17 -0700 Subject: [PATCH 022/140] ComputedField -> Field in validation --- validation/advection/validate_weno_scheme.jl | 12 ++++++------ validation/bickley_jet/bickley_jet.jl | 2 +- .../immersed_boundaries/immersed_bickley_jet.jl | 2 +- .../shallow_water_flow_past_cape.jl | 2 +- .../shallow_water_flow_past_cylinder.jl | 2 +- .../geostrophic_adjustment_test.jl | 2 +- validation/lid_driven_cavity/lid_driven_cavity.jl | 2 +- .../mesoscale_turbulence/abernathey_channel.jl | 12 ++++++------ validation/mesoscale_turbulence/eddying_channel.jl | 2 +- .../zonally_averaged_abernathey_channel.jl | 2 +- .../zonally_averaged_baroclinic_adjustment.jl | 2 +- .../mesoscale_turbulence/zonally_averaged_channel.jl | 6 +++--- .../mpi_shallow_water_turbulence.jl | 2 +- .../many_tke_based_free_convection.jl | 2 +- 14 files changed, 26 insertions(+), 26 deletions(-) diff --git a/validation/advection/validate_weno_scheme.jl b/validation/advection/validate_weno_scheme.jl index 2ae427d2fc..3b06d04cf9 100644 --- a/validation/advection/validate_weno_scheme.jl +++ b/validation/advection/validate_weno_scheme.jl @@ -46,9 +46,9 @@ coord = Dict() time = Dict() # 1D grid constructions -grid_reg = RectilinearGrid(size = (N,), x = Freg, halo = (3,), topology = (Periodic, Flat, Flat), architecture = arch) -grid_str = RectilinearGrid(size = (N,), x = Fsaw, halo = (3,), topology = (Periodic, Flat, Flat), architecture = arch) -grid_str2 = RectilinearGrid(size = (N,), x = Fstr2, halo = (3,), topology = (Periodic, Flat, Flat), architecture = arch) +grid_reg = RectilinearGrid(arch, size = N, x = Freg, halo = 3, topology = (Periodic, Flat, Flat)) +grid_str = RectilinearGrid(arch, size = N, x = Fsaw, halo = 3, topology = (Periodic, Flat, Flat)) +grid_str2 = RectilinearGrid(arch, size = N, x = Fstr2, halo = 3, topology = (Periodic, Flat, Flat)) # placeholder for the four different advection schemes # (1) WENO5(), @@ -67,7 +67,7 @@ c₀_2D(x, y, z) = mask(x) * mask(y) # Checking the accuracy of different schemes with different settings for (gr, grid) in enumerate([grid_reg, grid_str, grid_str2]) - U = Field(Face, Center, Center, arch, grid) + U = Field{Face, Center, Center}(grid) parent(U) .= 1 Δt_max = 0.2 * min_Δx(grid) @@ -139,8 +139,8 @@ grid_str2 = RectilinearGrid(size = (N, N), x = Fstr2, y = Fstr2, halo = (3, 3), for (gr, grid) in enumerate([grid_reg, grid_str, grid_str2]) - U = Field(Face, Center, Center, arch, grid) - V = Field(Center, Face, Center, arch, grid) + U = Field{Face, Center, Center}(grid) + V = Field{Center, Face, Center}(grid) parent(U) .= 1 parent(V) .= 0.3 diff --git a/validation/bickley_jet/bickley_jet.jl b/validation/bickley_jet/bickley_jet.jl index 0204a655a9..d46b2cc424 100644 --- a/validation/bickley_jet/bickley_jet.jl +++ b/validation/bickley_jet/bickley_jet.jl @@ -93,7 +93,7 @@ function run_bickley_jet(; output_time_interval = 2, stop_time = 200, arch = CPU # Output: primitive fields + computations u, v, w, c = merge(model.velocities, model.tracers) - ζ = ComputedField(∂x(v) - ∂y(u)) + ζ = Field(∂x(v) - ∂y(u)) outputs = merge(model.velocities, model.tracers, (ζ=ζ, η=model.free_surface.η)) diff --git a/validation/immersed_boundaries/immersed_bickley_jet.jl b/validation/immersed_boundaries/immersed_bickley_jet.jl index 6b49875ef1..032074ca2b 100644 --- a/validation/immersed_boundaries/immersed_bickley_jet.jl +++ b/validation/immersed_boundaries/immersed_bickley_jet.jl @@ -110,7 +110,7 @@ function run_bickley_jet(; output_time_interval = 2, stop_time = 200, arch = CPU # Output: primitive fields + computations u, v, w, c = merge(m.velocities, m.tracers) - ζ = ComputedField(∂x(v) - ∂y(u)) + ζ = Field(∂x(v) - ∂y(u)) outputs = merge(m.velocities, m.tracers, (ζ=ζ,)) output_name = m.grid isa ImmersedBoundaryGrid ? diff --git a/validation/immersed_boundaries/shallow_water_flow_past_cape.jl b/validation/immersed_boundaries/shallow_water_flow_past_cape.jl index cc1a27e143..5c7356e2be 100644 --- a/validation/immersed_boundaries/shallow_water_flow_past_cape.jl +++ b/validation/immersed_boundaries/shallow_water_flow_past_cape.jl @@ -49,7 +49,7 @@ simulation = Simulation(model, Δt=wizard, stop_time=1, progress=progress, itera uh, vh, h = model.solution -ζ = ComputedField(∂x(vh / h) - ∂y(uh / h)) +ζ = Field(∂x(vh / h) - ∂y(uh / h)) outputs = merge(model.solution, (ζ=ζ,)) diff --git a/validation/immersed_boundaries/shallow_water_flow_past_cylinder.jl b/validation/immersed_boundaries/shallow_water_flow_past_cylinder.jl index 9882fa333a..1584be39b4 100644 --- a/validation/immersed_boundaries/shallow_water_flow_past_cylinder.jl +++ b/validation/immersed_boundaries/shallow_water_flow_past_cylinder.jl @@ -52,7 +52,7 @@ simulation = Simulation(model, Δt=wizard, stop_time=1, progress=progress, itera uh, vh, h = model.solution -ζ = ComputedField(∂x(vh / h) - ∂y(uh / h)) +ζ = Field(∂x(vh / h) - ∂y(uh / h)) outputs = merge(model.solution, (ζ=ζ,)) diff --git a/validation/implicit_free_surface/geostrophic_adjustment_test.jl b/validation/implicit_free_surface/geostrophic_adjustment_test.jl index 759b9a5c2e..54c0d7b4f1 100644 --- a/validation/implicit_free_surface/geostrophic_adjustment_test.jl +++ b/validation/implicit_free_surface/geostrophic_adjustment_test.jl @@ -50,7 +50,7 @@ function run_and_analyze(simulation) u, v, w = simulation.model.velocities Δt = simulation.Δt - ηx = ComputedField(∂x(η)) + ηx = Field(∂x(η)) compute!(ηx) u₀ = interior(u)[:, 1, 1] diff --git a/validation/lid_driven_cavity/lid_driven_cavity.jl b/validation/lid_driven_cavity/lid_driven_cavity.jl index bcebd5ace8..12b6c1dab7 100644 --- a/validation/lid_driven_cavity/lid_driven_cavity.jl +++ b/validation/lid_driven_cavity/lid_driven_cavity.jl @@ -25,7 +25,7 @@ function simulate_lid_driven_cavity(; Re, N, end_time) ) u, v, w = model.velocities - ζ = ComputedField(∂y(w) - ∂z(v)) + ζ = Field(∂y(w) - ∂z(v)) fields = (; v, w, ζ) global_attributes = Dict("Re" => Re) diff --git a/validation/mesoscale_turbulence/abernathey_channel.jl b/validation/mesoscale_turbulence/abernathey_channel.jl index d2075dbe12..3bea32186e 100644 --- a/validation/mesoscale_turbulence/abernathey_channel.jl +++ b/validation/mesoscale_turbulence/abernathey_channel.jl @@ -203,18 +203,18 @@ simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterv u, v, w = model.velocities b = model.tracers.b -ζ = ComputedField(∂x(v) - ∂y(u)) +ζ = Field(∂x(v) - ∂y(u)) -B = AveragedField(b, dims=1) -V = AveragedField(v, dims=1) -W = AveragedField(w, dims=1) +B = Field(Average(b, dims=1)) +V = Field(Average(v, dims=1)) +W = Field(Average(w, dims=1)) b′ = b - B v′ = v - V w′ = w - W -v′b′ = AveragedField(v′ * b′, dims=1) -w′b′ = AveragedField(w′ * b′, dims=1) +v′b′ = Field(Average(v′ * b′, dims=1)) +w′b′ = Field(Average(w′ * b′, dims=1)) outputs = (; b, ζ, w) diff --git a/validation/mesoscale_turbulence/eddying_channel.jl b/validation/mesoscale_turbulence/eddying_channel.jl index 68c7b99959..7971a6e4fb 100644 --- a/validation/mesoscale_turbulence/eddying_channel.jl +++ b/validation/mesoscale_turbulence/eddying_channel.jl @@ -239,7 +239,7 @@ simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterv u, v, w = model.velocities b, c = model.tracers.b, model.tracers.c -ζ = ComputedField(∂x(v) - ∂y(u)) +ζ = Field(∂x(v) - ∂y(u)) B = AveragedField(b, dims=1) U = AveragedField(u, dims=1) diff --git a/validation/mesoscale_turbulence/zonally_averaged_abernathey_channel.jl b/validation/mesoscale_turbulence/zonally_averaged_abernathey_channel.jl index c7780d9b56..7062cb98f5 100644 --- a/validation/mesoscale_turbulence/zonally_averaged_abernathey_channel.jl +++ b/validation/mesoscale_turbulence/zonally_averaged_abernathey_channel.jl @@ -131,7 +131,7 @@ K = FunctionField{Center, Center, Center}(K_func, grid) f² = FunctionField{Center, Center, Center}(f²_func, grid) ν_op = @at (Center, Center, Center) K * f² / ∂z(b) -ν = ComputedField(ν_op) +ν = Field(ν_op) closure = AnisotropicDiffusivity(νh = 100, νz = 10, κh = 10, κz = 10, time_discretization = VerticallyImplicitTimeDiscretization()) diff --git a/validation/mesoscale_turbulence/zonally_averaged_baroclinic_adjustment.jl b/validation/mesoscale_turbulence/zonally_averaged_baroclinic_adjustment.jl index 0a6ab37f8a..1050b52054 100644 --- a/validation/mesoscale_turbulence/zonally_averaged_baroclinic_adjustment.jl +++ b/validation/mesoscale_turbulence/zonally_averaged_baroclinic_adjustment.jl @@ -167,7 +167,7 @@ using Oceananigans.TurbulenceClosures: ∇_dot_qᶜ computed_dependencies = dependencies) # R(b) eg the Redi operator applied to buoyancy -Rb = ComputedField(∇_q_op) +Rb = Field(∇_q_op) outputs = merge(fields(model), (; Rb)) diff --git a/validation/mesoscale_turbulence/zonally_averaged_channel.jl b/validation/mesoscale_turbulence/zonally_averaged_channel.jl index afb4ec8990..dc92829825 100644 --- a/validation/mesoscale_turbulence/zonally_averaged_channel.jl +++ b/validation/mesoscale_turbulence/zonally_averaged_channel.jl @@ -252,9 +252,9 @@ vb_op = KernelFunctionOperation{Center, Face, Center}(diffusive_flux_y, grid, a wb_op = KernelFunctionOperation{Center, Center, Face}(diffusive_flux_z, grid, architecture=architecture, computed_dependencies=dependencies) ∇_q_op = KernelFunctionOperation{Center, Center, Center}(∇_dot_qᶜ, grid, architecture=architecture, computed_dependencies=dependencies) -vb = ComputedField(vb_op) -wb = ComputedField(wb_op) -∇_q = ComputedField(∇_q_op) +vb = Field(vb_op) +wb = Field(wb_op) +∇_q = Field(∇_q_op) outputs = merge(fields(model), (; vb, wb, ∇_q)) diff --git a/validation/mpi_shallow_water/mpi_shallow_water_turbulence.jl b/validation/mpi_shallow_water/mpi_shallow_water_turbulence.jl index 344d0c4a83..00a0c449ec 100644 --- a/validation/mpi_shallow_water/mpi_shallow_water_turbulence.jl +++ b/validation/mpi_shallow_water/mpi_shallow_water_turbulence.jl @@ -34,7 +34,7 @@ progress(sim) = @info "Iteration: $(sim.model.clock.iteration), time: $(sim.mode simulation = Simulation(model, Δt=0.001, stop_time=100.0, iteration_interval=1, progress=progress) uh, vh, h = model.solution -outputs = (ζ=ComputedField(∂x(vh/h) - ∂y(uh/h)),) +outputs = (ζ=Field(∂x(vh/h) - ∂y(uh/h)),) filepath = "mpi_shallow_water_turbulence_rank$(local_rank).nc" simulation.output_writers[:fields] = NetCDFOutputWriter(model, outputs, filepath=filepath, schedule=TimeInterval(1.0), mode="c") diff --git a/validation/vertical_mixing_closures/many_tke_based_free_convection.jl b/validation/vertical_mixing_closures/many_tke_based_free_convection.jl index 54a0b25a0f..e7824143a3 100644 --- a/validation/vertical_mixing_closures/many_tke_based_free_convection.jl +++ b/validation/vertical_mixing_closures/many_tke_based_free_convection.jl @@ -43,7 +43,7 @@ z = znodes(model.tracers.b) simulation = Simulation(model, Δt = 1minute/2, stop_time = 0.0) b = model.tracers.b -bz = ComputedField(∂z(b)) +bz = Field(∂z(b)) function column_bz(j) compute!(bz) From e5543f67b658db8280fe239dd69ecad317859adb Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Sun, 2 Jan 2022 21:00:26 -0700 Subject: [PATCH 023/140] Touch up FieldTimeSeries --- src/OutputReaders/field_time_series.jl | 4 ++++ test/test_output_readers.jl | 11 ++++++----- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/src/OutputReaders/field_time_series.jl b/src/OutputReaders/field_time_series.jl index fe02105216..eab17ad435 100644 --- a/src/OutputReaders/field_time_series.jl +++ b/src/OutputReaders/field_time_series.jl @@ -114,6 +114,8 @@ function FieldTimeSeries(path, name, backend::InMemory; return time_series end +Base.parent(fts::FieldTimeSeries) = parent(fts.data) + Base.getindex(fts::InMemoryFieldTimeSeries{LX, LY, LZ}, n::Int) where {LX, LY, LZ} = Field{LX, LY, LZ}(fts.grid, boundary_conditions=fts.boundary_conditions, data=view(fts.data, :, :, :, n)) @@ -280,6 +282,8 @@ end Base.setindex!(fts::FieldTimeSeries, val, inds...) = Base.setindex!(fts.data, val, inds...) +Base.parent(fts::FieldTimeSeries{X, Y, Z, OnDisk}) = nothing + ##### ##### Show methods ##### diff --git a/test/test_output_readers.jl b/test/test_output_readers.jl index 2303cf6fe5..29e6321bc3 100644 --- a/test/test_output_readers.jl +++ b/test/test_output_readers.jl @@ -78,11 +78,12 @@ end # This behavior ensures that set! works # but perhaps should be changed in the future - @test parent(u3) isa SubArray - @test parent(v3) isa SubArray - @test parent(w3) isa SubArray - @test parent(T3) isa SubArray - @test parent(b3) isa SubArray + @test parent(u3[1]) isa SubArray + @test parent(v3[1]) isa SubArray + @test parent(w3[1]) isa SubArray + @test parent(T3[1]) isa SubArray + @test parent(b3[1]) isa SubArray + @test parent(ζ3[1]) isa SubArray @test location(u3) == (Face, Center, Center) @test location(v3) == (Center, Face, Center) From b1bd858485cb334643c774e8e70341033fcbd37f Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Mon, 3 Jan 2022 10:57:19 -0700 Subject: [PATCH 024/140] Fixes bug in FieldTimeSeries --- src/OutputReaders/field_time_series.jl | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/src/OutputReaders/field_time_series.jl b/src/OutputReaders/field_time_series.jl index eab17ad435..a8736e171d 100644 --- a/src/OutputReaders/field_time_series.jl +++ b/src/OutputReaders/field_time_series.jl @@ -14,16 +14,16 @@ import Oceananigans: short_show import Oceananigans.Fields: Field, set!, interior import Oceananigans.Architectures: architecture -struct FieldTimeSeries{X, Y, Z, K, A, T, D, G, B, χ} <: AbstractField{X, Y, Z, A, G, T, 4} +struct FieldTimeSeries{LX, LY, LZ, K, A, T, D, G, B, χ} <: AbstractField{LX, LY, LZ, A, G, T, 4} data :: D grid :: G boundary_conditions :: B times :: χ - function FieldTimeSeries{X, Y, Z, K}(data::D, grid::G, bcs::B, times::χ) where {X, Y, Z, K, D, G, B, χ} + function FieldTimeSeries{LX, LY, LZ, K}(data::D, grid::G, bcs::B, times::χ) where {LX, LY, LZ, K, D, G, B, χ} T = eltype(grid) A = typeof(architecture(grid)) - return new{X, Y, Z, K, A, T, D, G, B, χ}(data, grid, bcs, times) + return new{LX, LY, LZ, K, A, T, D, G, B, χ}(data, grid, bcs, times) end end @@ -80,7 +80,7 @@ FieldTimeSeries(path, name; backend=InMemory(), kwargs...) = ##### InMemory time serieses ##### -const InMemoryFieldTimeSeries{X, Y, Z} = FieldTimeSeries{X, Y, Z, InMemory} +const InMemoryFieldTimeSeries{LX, LY, LZ} = FieldTimeSeries{LX, LY, LZ, InMemory} struct UnspecifiedBoundaryConditions end @@ -208,7 +208,7 @@ parent_indices(idx::Base.Slice{<:IdOffsetRange}) = Colon() Base.parent(vf::ViewField) = view(parent(parent(vf.data)), parent_indices.(vf.data.indices)...) "Returns a view of `f` that excludes halo points." -@inline interior(f::FieldTimeSeries{X, Y, Z}) where {X, Y, Z} = +@inline interior(f::FieldTimeSeries{LX, LY, LZ}) where {LX, LY, LZ} = view(parent(f.data), interior_parent_indices(X, topology(f, 1), f.grid.Nx, f.grid.Hx), interior_parent_indices(Y, topology(f, 2), f.grid.Ny, f.grid.Hy), @@ -263,9 +263,9 @@ end # Include the time dimension. @inline Base.size(fts::FieldTimeSeries) = (size(location(fts), fts.grid)..., length(fts.times)) -@propagate_inbounds Base.getindex(f::FieldTimeSeries{X, Y, Z, InMemory}, i, j, k, n) where {X, Y, Z} = f.data[i, j, k, n] +@propagate_inbounds Base.getindex(f::FieldTimeSeries{LX, LY, LZ, InMemory}, i, j, k, n) where {LX, LY, LZ} = f.data[i, j, k, n] -function Base.getindex(fts::FieldTimeSeries{X, Y, Z, OnDisk}, n::Int) where {X, Y, Z} +function Base.getindex(fts::FieldTimeSeries{LX, LY, LZ, OnDisk}, n::Int) where {LX, LY, LZ} # Load data arch = architecture(fts) file = jldopen(fts.data.path) @@ -274,7 +274,7 @@ function Base.getindex(fts::FieldTimeSeries{X, Y, Z, OnDisk}, n::Int) where {X, close(file) # Wrap Field - loc = (X, Y, Z) + loc = (LX, LY, LZ) field_data = offset_data(raw_data, fts.grid, loc) return Field(loc, fts.grid; boundary_conditions=fts.boundary_conditions, data=field_data) @@ -282,7 +282,7 @@ end Base.setindex!(fts::FieldTimeSeries, val, inds...) = Base.setindex!(fts.data, val, inds...) -Base.parent(fts::FieldTimeSeries{X, Y, Z, OnDisk}) = nothing +Base.parent(fts::FieldTimeSeries{LX, LY, LZ, OnDisk}) where {LX, LY, LZ} = nothing ##### ##### Show methods @@ -295,10 +295,10 @@ backend_str(::OnDisk) = "OnDisk" ##### show ##### -short_show(fts::FieldTimeSeries{X, Y, Z, K}) where {X, Y, Z, K} = +short_show(fts::FieldTimeSeries{LX, LY, LZ, K}) where {LX, LY, LZ, K} = string("$(join(size(fts), "×")) FieldTimeSeries{$(backend_str(K()))} located at $(show_location(fts))") -Base.show(io::IO, fts::FieldTimeSeries{X, Y, Z, K, A}) where {X, Y, Z, K, A} = +Base.show(io::IO, fts::FieldTimeSeries{LX, LY, LZ, K, A}) where {LX, LY, LZ, K, A} = print(io, "$(short_show(fts))\n", "├── architecture: $A\n", "└── grid: $(short_show(fts.grid))") From 94fceabd3de0b03da42358151cedee342fcbf1e0 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Mon, 3 Jan 2022 16:27:35 -0700 Subject: [PATCH 025/140] Fixes tridiagonal solver test --- test/test_batched_tridiagonal_solver.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test/test_batched_tridiagonal_solver.jl b/test/test_batched_tridiagonal_solver.jl index e823f850e5..f78b4a9dad 100644 --- a/test/test_batched_tridiagonal_solver.jl +++ b/test/test_batched_tridiagonal_solver.jl @@ -22,7 +22,7 @@ function can_solve_single_tridiagonal_system(arch, N) grid = RectilinearGrid(arch, size=(1, 1, N), extent=(1, 1, 1)) - btsolver = BatchedTridiagonalSolver(arch, grid; + btsolver = BatchedTridiagonalSolver(grid; lower_diagonal = a, diagonal = b, upper_diagonal = c) @@ -55,7 +55,7 @@ function can_solve_single_tridiagonal_system_with_functions(arch, N) ϕ = reshape(zeros(N), (1, 1, N)) |> ArrayType - btsolver = BatchedTridiagonalSolver(arch, grid; + btsolver = BatchedTridiagonalSolver(grid; lower_diagonal = a, diagonal = b, upper_diagonal = c) @@ -85,7 +85,7 @@ function can_solve_batched_tridiagonal_system_with_3D_RHS(arch, Nx, Ny, Nz) a, b, c, f = ArrayType.([a, b, c, f]) grid = RectilinearGrid(arch, size=(Nx, Ny, Nz), extent=(1, 1, 1)) - btsolver = BatchedTridiagonalSolver(arch, grid; + btsolver = BatchedTridiagonalSolver(grid; lower_diagonal = a, diagonal = b, upper_diagonal = c) @@ -122,7 +122,7 @@ function can_solve_batched_tridiagonal_system_with_3D_functions(arch, Nx, Ny, Nz # Convert to CuArray if needed. a, c = ArrayType.([a, c]) - btsolver = BatchedTridiagonalSolver(arch, grid; + btsolver = BatchedTridiagonalSolver(grid; lower_diagonal = a, diagonal = b, upper_diagonal = c) From 236dc46cbfa8015af5bf844771c92dcad4ad9b6b Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Mon, 3 Jan 2022 16:29:09 -0700 Subject: [PATCH 026/140] No longer need both arch and grid --- .../HydrostaticFreeSurfaceModels/implicit_free_surface.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl index 195ebb1773..b9d8ba777b 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl @@ -51,7 +51,7 @@ Adapt.adapt_structure(to, free_surface::ImplicitFreeSurface) = # Internal function for HydrostaticFreeSurfaceModel function FreeSurface(free_surface::ImplicitFreeSurface{Nothing}, velocities, grid) - η = FreeSurfaceDisplacementField(velocities, free_surface, arch, grid) + η = FreeSurfaceDisplacementField(velocities, free_surface, grid) gravitational_acceleration = convert(eltype(grid), free_surface.gravitational_acceleration) # Initialize barotropic volume fluxes @@ -74,7 +74,7 @@ is_horizontally_regular(::RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number, function build_implicit_step_solver(::Val{:Default}, grid, gravitational_acceleration, settings) default_method = is_horizontally_regular(grid) ? :FastFourierTransform : :PreconditionedConjugateGradient - return build_implicit_step_solver(Val(default_method), arch, grid, gravitational_acceleration, settings) + return build_implicit_step_solver(Val(default_method), grid, gravitational_acceleration, settings) end @inline explicit_barotropic_pressure_x_gradient(i, j, k, grid, ::ImplicitFreeSurface) = 0 From 6f86969274ccabefb3ad8054cfae49e160ada3bc Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Mon, 3 Jan 2022 16:31:15 -0700 Subject: [PATCH 027/140] Updates stratified couette flow script --- .../stratified_couette_flow.jl | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/validation/stratified_couette_flow/stratified_couette_flow.jl b/validation/stratified_couette_flow/stratified_couette_flow.jl index a5ea943869..9ee97ad810 100644 --- a/validation/stratified_couette_flow/stratified_couette_flow.jl +++ b/validation/stratified_couette_flow/stratified_couette_flow.jl @@ -192,12 +192,12 @@ function simulate_stratified_couette_flow(; Nxy, Nz, arch=GPU(), h=1, U_wall=1, ##### Set up profile output writer ##### - Uavg = AveragedField(model.velocities.u, dims=(1, 2)) - Vavg = AveragedField(model.velocities.v, dims=(1, 2)) - Wavg = AveragedField(model.velocities.w, dims=(1, 2)) - Tavg = AveragedField(model.tracers.T, dims=(1, 2)) - νavg = AveragedField(model.diffusivity_fields.νₑ, dims=(1, 2)) - κavg = AveragedField(model.diffusivity_fields.κₑ.T, dims=(1, 2)) + Uavg = Field(Average(model.velocities.u, dims=(1, 2))) + Vavg = Field(Average(model.velocities.v, dims=(1, 2))) + Wavg = Field(Average(model.velocities.w, dims=(1, 2))) + Tavg = Field(Average(model.tracers.T, dims=(1, 2))) + νavg = Field(Average(model.diffusivity_fields.νₑ, dims=(1, 2))) + κavg = Field(Average(model.diffusivity_fields.κₑ.T, dims=(1, 2))) profiles = Dict( :u => Uavg, From 5174310bb4901238a7e61802992c8ae05d3dedf4 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Mon, 10 Jan 2022 10:15:05 -0700 Subject: [PATCH 028/140] Bugfix in building PCG implicit free surface solver --- .../pcg_implicit_free_surface_solver.jl | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl b/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl index 428ccb8df9..b1300c26ab 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/pcg_implicit_free_surface_solver.jl @@ -52,7 +52,7 @@ function PCGImplicitFreeSurfaceSolver(grid::AbstractGrid, gravitational_accelera end build_implicit_step_solver(::Val{:PreconditionedConjugateGradient}, grid, gravitational_acceleration, settings) = - PCGImplicitFreeSurfaceSolver(arch, grid, gravitational_acceleration, settings) + PCGImplicitFreeSurfaceSolver(grid, gravitational_acceleration, settings) ##### ##### Solve... @@ -102,8 +102,6 @@ end end """ - implicit_free_surface_linear_operation!(result, x, arch, grid, bcs; args...) - Returns `L(ηⁿ)`, where `ηⁿ` is the free surface displacement at time step `n` and `L` is the linear operator that arises in an implicit time step for the free surface displacement `η`. From f9e7f8cda5d387ab2bd0a1b518cf8ed66287e89a Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Mon, 10 Jan 2022 10:40:13 -0700 Subject: [PATCH 029/140] Fixes bug in DistributedField constructor --- src/Distributed/distributed_fields.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/Distributed/distributed_fields.jl b/src/Distributed/distributed_fields.jl index 1ffd532e1a..d2ff6502e9 100644 --- a/src/Distributed/distributed_fields.jl +++ b/src/Distributed/distributed_fields.jl @@ -1,7 +1,7 @@ import Oceananigans.Fields: Field -function Field((LX, LY, LZ)::Tuple, grid::DistributedGrid, data, bcs, op, status) +function Field((LX, LY, LZ)::Tuple, grid::DistributedGrid, data, old_bcs, op, status) arch = architecture(grid) - boundary_conditions = inject_halo_communication_boundary_conditions(bcs, arch.local_rank, arch.connectivity) - return Field{LX, LY, LZ}(grid, data, bcs, op, status) + new_bcs = inject_halo_communication_boundary_conditions(old_bcs, arch.local_rank, arch.connectivity) + return Field{LX, LY, LZ}(grid, data, new_bcs, op, status) end From d6338b4dcec04dd602cb0c86fd3f123cd82fc055 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Mon, 10 Jan 2022 17:29:36 -0700 Subject: [PATCH 030/140] Bugfixes in Matrix implicit solver --- .../matrix_implicit_free_surface_solver.jl | 29 +++++++++---------- 1 file changed, 14 insertions(+), 15 deletions(-) diff --git a/src/Models/HydrostaticFreeSurfaceModels/matrix_implicit_free_surface_solver.jl b/src/Models/HydrostaticFreeSurfaceModels/matrix_implicit_free_surface_solver.jl index c2472a7c88..0d4aee6f3b 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/matrix_implicit_free_surface_solver.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/matrix_implicit_free_surface_solver.jl @@ -1,7 +1,7 @@ using Oceananigans.Solvers using Oceananigans.Operators using Oceananigans.Architectures -using Oceananigans.Architectures: architecture +using Oceananigans.Grids: AbstractGrid using Oceananigans.Fields: ReducedField using Oceananigans.Solvers: MatrixIterativeSolver using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, solid_cell @@ -14,7 +14,7 @@ struct MatrixImplicitFreeSurfaceSolver{V, S, R} end """ - PCGImplicitFreeSurfaceSolver(arch::AbstractArchitecture, grid, settings) + MatrixImplicitFreeSurfaceSolver(grid, gravitational_acceleration, settings) Return a the framework for solving the elliptic equation with one of the iterative solvers of IterativeSolvers.jl with a sparse matrix formulation. @@ -28,17 +28,18 @@ for a fluid with variable depth `H`, horizontal areas `Az`, barotropic volume fl step `Δt`, gravitational acceleration `g`, and free surface at time-step `n` `ηⁿ`. """ -function MatrixImplicitFreeSurfaceSolver(arch::AbstractArchitecture, grid, gravitational_acceleration, settings) +function MatrixImplicitFreeSurfaceSolver(grid::AbstractGrid, gravitational_acceleration, settings) # Initialize vertically integrated lateral face areas - ∫ᶻ_Axᶠᶜᶜ = ReducedField(Face, Center, Nothing, arch, grid; dims=3) - ∫ᶻ_Ayᶜᶠᶜ = ReducedField(Center, Face, Nothing, arch, grid; dims=3) + ∫ᶻ_Axᶠᶜᶜ = Field{Face, Center, Nothing}(grid) + ∫ᶻ_Ayᶜᶠᶜ = Field{Center, Face, Nothing}(grid) vertically_integrated_lateral_areas = (xᶠᶜᶜ = ∫ᶻ_Axᶠᶜᶜ, yᶜᶠᶜ = ∫ᶻ_Ayᶜᶠᶜ) - compute_vertically_integrated_lateral_areas!(vertically_integrated_lateral_areas, grid, arch) + compute_vertically_integrated_lateral_areas!(vertically_integrated_lateral_areas, grid) - right_hand_side = arch_array(arch, zeros(eltype(grid), grid.Nx * grid.Ny)) # linearized RHS for matrix operations + arch = architecture(grid) + right_hand_side = zeros(grid, grid.Nx * grid.Ny) # linearized RHS for matrix operations # Set maximum iterations to Nx * Ny if not set settings = Dict{Symbol, Any}(settings) @@ -47,16 +48,14 @@ function MatrixImplicitFreeSurfaceSolver(arch::AbstractArchitecture, grid, gravi coeffs = compute_matrix_coefficients(vertically_integrated_lateral_areas, grid, gravitational_acceleration) - solver = MatrixIterativeSolver(coeffs; - reduced_dim = (false, false, true), - grid = grid, - settings...) + solver = MatrixIterativeSolver(coeffs; reduced_dim = (false, false, true), + grid = grid, settings...) return MatrixImplicitFreeSurfaceSolver(vertically_integrated_lateral_areas, solver, right_hand_side) end -build_implicit_step_solver(::Val{:MatrixIterativeSolver}, arch, grid, gravitational_acceleration, settings) = - MatrixImplicitFreeSurfaceSolver(arch, grid, gravitational_acceleration, settings) +build_implicit_step_solver(::Val{:MatrixIterativeSolver}, grid, gravitational_acceleration, settings) = + MatrixImplicitFreeSurfaceSolver(grid, gravitational_acceleration, settings) ##### ##### Solve... @@ -76,8 +75,8 @@ function compute_implicit_free_surface_right_hand_side!(rhs, g, Δt, ∫ᶻQ, η) solver = implicit_solver.matrix_iterative_solver - arch = architecture(solver.matrix) grid = solver.grid + arch = architecture(grid) event = launch!(arch, grid, :xy, implicit_linearized_free_surface_right_hand_side!, @@ -127,4 +126,4 @@ end Ax[i, j, 1] = ∫Ax[i, j, 1] / Δxᶠᶜᵃ(i, j, 1, grid) diag[i, j, 1] = - Azᶜᶜᵃ(i, j, 1, grid) / g end -end \ No newline at end of file +end From 2eb44ae08fdb96784d80a3d8be3c4fa3b059cf88 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Mon, 10 Jan 2022 21:25:05 -0700 Subject: [PATCH 031/140] Need halo=(3, 3, 3) with WENO advection --- ...static_free_surface_immersed_boundaries.jl | 28 ++++++++++--------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/test/test_hydrostatic_free_surface_immersed_boundaries.jl b/test/test_hydrostatic_free_surface_immersed_boundaries.jl index 98fd258bde..fc371ce6b4 100644 --- a/test/test_hydrostatic_free_surface_immersed_boundaries.jl +++ b/test/test_hydrostatic_free_surface_immersed_boundaries.jl @@ -57,10 +57,11 @@ using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization Ny = 60 # A spherical domain - underlying_grid = LatitudeLongitudeGrid(arch, size = (Nx, Ny, 1), - longitude = (-30, 30), - latitude = (15, 75), - z = (-4000, 0)) + underlying_grid = LatitudeLongitudeGrid(arch, + size = (Nx, Ny, 1), + longitude = (-30, 30), + latitude = (15, 75), + z = (-4000, 0)) bathymetry = zeros(Nx, Ny) .- 4000 view(bathymetry, 31:34, 43:47) .= 0 @@ -94,14 +95,14 @@ using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization νh₀ = 5e3 * (60 / grid.Nx)^2 constant_horizontal_diffusivity = HorizontallyCurvilinearAnisotropicDiffusivity(νh=νh₀) - model = HydrostaticFreeSurfaceModel(grid = grid, - momentum_advection = VectorInvariant(), - free_surface = free_surface, - coriolis = coriolis, - boundary_conditions = (u=u_bcs, v=v_bcs), - closure = constant_horizontal_diffusivity, - tracers = nothing, - buoyancy = nothing) + model = HydrostaticFreeSurfaceModel(; grid, + momentum_advection = VectorInvariant(), + free_surface = free_surface, + coriolis = coriolis, + boundary_conditions = (u=u_bcs, v=v_bcs), + closure = constant_horizontal_diffusivity, + tracers = nothing, + buoyancy = nothing) simulation = Simulation(model, Δt=3600, stop_iteration=1) @@ -120,6 +121,7 @@ using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization underlying_grid = RectilinearGrid(arch, size = (Nx, Ny, 3), + halo = (3, 3, 3), extent = (Nx, Ny, 3), topology = (Periodic, Periodic, Bounded)) @@ -130,7 +132,7 @@ using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(B)) - model = HydrostaticFreeSurfaceModel(grid = grid, + model = HydrostaticFreeSurfaceModel(; grid, free_surface = ImplicitFreeSurface(), tracer_advection = WENO5(), buoyancy = nothing, From e63ad80092023948a6628d712f99497fd8f3637a Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Mon, 10 Jan 2022 21:27:48 -0700 Subject: [PATCH 032/140] Fix matrix poisson solver tests --- test/test_matrix_poisson_solver.jl | 39 ++++++++++++++++-------------- 1 file changed, 21 insertions(+), 18 deletions(-) diff --git a/test/test_matrix_poisson_solver.jl b/test/test_matrix_poisson_solver.jl index 930e4db475..1623d9dbae 100644 --- a/test/test_matrix_poisson_solver.jl +++ b/test/test_matrix_poisson_solver.jl @@ -5,7 +5,8 @@ using Oceananigans.Architectures: arch_array using KernelAbstractions: @kernel, @index using Statistics, LinearAlgebra, SparseArrays -function calc_∇²!(∇²ϕ, ϕ, arch, grid) +function calc_∇²!(∇²ϕ, ϕ, grid) + arch = architecture(grid) fill_halo_regions!(ϕ, arch) event = launch!(arch, grid, :xyz, ∇²!, ∇²ϕ, grid, ϕ) wait(event) @@ -19,7 +20,8 @@ function identity_operator!(b, x) return nothing end -function run_identity_operator_test(arch, grid) +function run_identity_operator_test(grid) + arch = architecture(grid) N = size(grid) M = prod(N) @@ -34,7 +36,7 @@ function run_identity_operator_test(arch, grid) fill!(b, rand()) - initial_guess = solution = CenterField(arch, grid) + initial_guess = solution = CenterField(grid) set!(initial_guess, (x, y, z) -> rand()) solve!(initial_guess, solver, b, 1.0) @@ -54,7 +56,7 @@ function compute_poisson_weights(grid) Ay = zeros(N...) Az = zeros(N...) C = zeros(N...) - D = arch_array(grid.architecture, zeros(N...)) + D = arch_array(architecture(grid), zeros(N...)) for i =1:grid.Nx, j = 1:grid.Ny, k = 1:grid.Nz Ax[i, j, k] = Δzᵃᵃᶜ(i, j, k, grid) * Δyᶠᶜᵃ(i, j, k, grid) / Δxᶠᶜᵃ(i, j, k, grid) Ay[i, j, k] = Δzᵃᵃᶜ(i, j, k, grid) * Δxᶜᶠᵃ(i, j, k, grid) / Δyᶜᶠᵃ(i, j, k, grid) @@ -65,14 +67,15 @@ end function poisson_rhs!(r, grid) - event = launch!(grid.architecture, grid, :xyz, _multiply_by_volume!, r, grid) + event = launch!(architecture(grid), grid, :xyz, _multiply_by_volume!, r, grid) wait(event) return end -function run_poisson_equation_test(arch, grid) +function run_poisson_equation_test(grid) # Solve ∇²ϕ = r - ϕ_truth = Field(Center, Center, Center, arch, grid) + ϕ_truth = Field{Center, Center, Center}(grid) + arch = architecture(grid) # Initialize zero-mean "truth" solution with random numbers set!(ϕ_truth, (x, y, z) -> rand()) @@ -80,8 +83,8 @@ function run_poisson_equation_test(arch, grid) fill_halo_regions!(ϕ_truth, arch) # Calculate Laplacian of "truth" - ∇²ϕ = Field(Center, Center, Center, arch, grid) - calc_∇²!(∇²ϕ, ϕ_truth, arch, grid) + ∇²ϕ = Field{Center, Center, Center}(grid) + calc_∇²!(∇²ϕ, ϕ_truth, grid) rhs = deepcopy(∇²ϕ) poisson_rhs!(rhs, grid) @@ -90,13 +93,13 @@ function run_poisson_equation_test(arch, grid) solver = MatrixIterativeSolver(weights, grid = grid) # Solve Poisson equation - ϕ_solution = CenterField(arch, grid) + ϕ_solution = CenterField(grid) solve!(ϕ_solution, solver, rhs, 1.0) # Diagnose Laplacian of solution - ∇²ϕ_solution = CenterField(arch, grid) - calc_∇²!(∇²ϕ_solution, ϕ_solution, arch, grid) + ∇²ϕ_solution = CenterField(grid) + calc_∇²!(∇²ϕ_solution, ϕ_solution, grid) parent(ϕ_solution) .-= mean(ϕ_solution) @@ -114,16 +117,16 @@ end @info "Testing 2D MatrixIterativeSolver [$(typeof(arch)) $topo]..." grid = RectilinearGrid(arch, size=(4, 8), extent=(1, 3), topology = topo) - run_identity_operator_test(arch, grid) - run_poisson_equation_test(arch, grid) + run_identity_operator_test(grid) + run_poisson_equation_test(grid) end topologies = [(Periodic, Periodic, Periodic), (Bounded, Bounded, Periodic), (Periodic, Bounded, Periodic), (Bounded, Periodic, Bounded)] for arch in archs, topo in topologies @info "Testing 3D MatrixIterativeSolver [$(typeof(arch)) $topo]..." grid = RectilinearGrid(arch, size=(4, 8, 6), extent=(1, 3, 4), topology = topo) - run_identity_operator_test(arch, grid) - run_poisson_equation_test(arch, grid) + run_identity_operator_test(grid) + run_poisson_equation_test(grid) end stretch_coord = [0, 1.5, 3, 7, 8.5, 10] for arch in archs @@ -133,7 +136,7 @@ end for grid in grids @info "Testing stretched grid MatrixIterativeSolver [$(typeof(arch))]..." - run_poisson_equation_test(arch, grid) + run_poisson_equation_test(grid) end end @@ -146,4 +149,4 @@ end @test all(M .≈ A⁻¹) -end \ No newline at end of file +end From 52226ab16e0d5ae4917d991bf2197c1b192318e3 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Mon, 10 Jan 2022 22:18:57 -0700 Subject: [PATCH 033/140] Different fix for immersed boundary test issue --- test/test_hydrostatic_free_surface_immersed_boundaries.jl | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/test/test_hydrostatic_free_surface_immersed_boundaries.jl b/test/test_hydrostatic_free_surface_immersed_boundaries.jl index fc371ce6b4..3ed530b4a2 100644 --- a/test/test_hydrostatic_free_surface_immersed_boundaries.jl +++ b/test/test_hydrostatic_free_surface_immersed_boundaries.jl @@ -121,20 +121,18 @@ using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization underlying_grid = RectilinearGrid(arch, size = (Nx, Ny, 3), - halo = (3, 3, 3), extent = (Nx, Ny, 3), topology = (Periodic, Periodic, Bounded)) # B for bathymetry B = [-3. for i=1:Nx, j=1:Ny ] - B[2:Nx-1,2:Ny-1] .= [-2. for i=2:Nx-1, j=2:Ny-1 ] - B[3:Nx-2,3:Ny-2] .= [-1. for i=3:Nx-2, j=3:Ny-2 ] + B[2:Nx-1, 2:Ny-1] .= [-2 for i=2:Nx-1, j=2:Ny-1 ] + B[3:Nx-2, 3:Ny-2] .= [-1 for i=3:Nx-2, j=3:Ny-2 ] grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(B)) model = HydrostaticFreeSurfaceModel(; grid, free_surface = ImplicitFreeSurface(), - tracer_advection = WENO5(), buoyancy = nothing, tracers = nothing, closure = nothing) From f48c6fc8f13be0eeece420076638fd7252abe99e Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Mon, 10 Jan 2022 22:37:35 -0700 Subject: [PATCH 034/140] Updates to matrix poisson solver tests --- test/test_matrix_poisson_solver.jl | 43 +++++++++++++++++------------- 1 file changed, 25 insertions(+), 18 deletions(-) diff --git a/test/test_matrix_poisson_solver.jl b/test/test_matrix_poisson_solver.jl index 1623d9dbae..4e635ee4ad 100644 --- a/test/test_matrix_poisson_solver.jl +++ b/test/test_matrix_poisson_solver.jl @@ -1,3 +1,5 @@ +include("dependencies_for_runtests.jl") + using Oceananigans.Solvers: solve!, MatrixIterativeSolver, spai_preconditioner using Oceananigans.Fields: interior_copy using Oceananigans.Operators: volume, Δyᶠᶜᵃ, Δyᶜᶠᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δyᵃᶜᵃ, Δxᶜᵃᵃ, Δzᵃᵃᶠ, Δzᵃᵃᶜ, ∇²ᶜᶜᶜ @@ -14,7 +16,6 @@ function calc_∇²!(∇²ϕ, ϕ, grid) return nothing end - function identity_operator!(b, x) parent(b) .= parent(x) return nothing @@ -25,12 +26,12 @@ function run_identity_operator_test(grid) N = size(grid) M = prod(N) + b = zeros(arch, M) + A = zeros(arch, N...) + D = zeros(grid, N...) - b = arch_array(arch, zeros(M)) - - A = zeros(N...) - C = ones(N...) - D = arch_array(arch, zeros(N...)) + C = zeros(arch, N...) + fill!(C, 1) solver = MatrixIterativeSolver((A, A, A, C, D), grid = grid) @@ -42,6 +43,7 @@ function run_identity_operator_test(grid) solve!(initial_guess, solver, b, 1.0) b = reshape(b, size(grid)...) + @test norm(interior(solution) .- b) .< solver.tolerance end @@ -56,8 +58,8 @@ function compute_poisson_weights(grid) Ay = zeros(N...) Az = zeros(N...) C = zeros(N...) - D = arch_array(architecture(grid), zeros(N...)) - for i =1:grid.Nx, j = 1:grid.Ny, k = 1:grid.Nz + D = zeros(grid, N...) + for i = 1:grid.Nx, j = 1:grid.Ny, k = 1:grid.Nz Ax[i, j, k] = Δzᵃᵃᶜ(i, j, k, grid) * Δyᶠᶜᵃ(i, j, k, grid) / Δxᶠᶜᵃ(i, j, k, grid) Ay[i, j, k] = Δzᵃᵃᶜ(i, j, k, grid) * Δxᶜᶠᵃ(i, j, k, grid) / Δyᶜᶠᵃ(i, j, k, grid) Az[i, j, k] = Δxᶜᵃᵃ(i, j, k, grid) * Δyᵃᶜᵃ(i, j, k, grid) / Δzᵃᵃᶠ(i, j, k, grid) @@ -65,7 +67,6 @@ function compute_poisson_weights(grid) return (Ax, Ay, Az, C, D) end - function poisson_rhs!(r, grid) event = launch!(architecture(grid), grid, :xyz, _multiply_by_volume!, r, grid) wait(event) @@ -73,17 +74,18 @@ function poisson_rhs!(r, grid) end function run_poisson_equation_test(grid) - # Solve ∇²ϕ = r - ϕ_truth = Field{Center, Center, Center}(grid) arch = architecture(grid) + # Solve ∇²ϕ = r + ϕ_truth = CenterField(grid) + # Initialize zero-mean "truth" solution with random numbers set!(ϕ_truth, (x, y, z) -> rand()) parent(ϕ_truth) .-= mean(ϕ_truth) fill_halo_regions!(ϕ_truth, arch) # Calculate Laplacian of "truth" - ∇²ϕ = Field{Center, Center, Center}(grid) + ∇²ϕ = CenterField(grid) calc_∇²!(∇²ϕ, ϕ_truth, grid) rhs = deepcopy(∇²ϕ) @@ -112,35 +114,41 @@ function run_poisson_equation_test(grid) end @testset "MatrixIterativeSolver" begin + topologies = [(Periodic, Periodic, Flat), (Bounded, Bounded, Flat), (Periodic, Bounded, Flat), (Bounded, Periodic, Flat)] + for arch in archs, topo in topologies - @info "Testing 2D MatrixIterativeSolver [$(typeof(arch)) $topo]..." + @info " Testing 2D MatrixIterativeSolver [$(typeof(arch)) $topo]..." grid = RectilinearGrid(arch, size=(4, 8), extent=(1, 3), topology = topo) run_identity_operator_test(grid) run_poisson_equation_test(grid) end + topologies = [(Periodic, Periodic, Periodic), (Bounded, Bounded, Periodic), (Periodic, Bounded, Periodic), (Bounded, Periodic, Bounded)] + for arch in archs, topo in topologies - @info "Testing 3D MatrixIterativeSolver [$(typeof(arch)) $topo]..." + @info " Testing 3D MatrixIterativeSolver [$(typeof(arch)) $topo]..." - grid = RectilinearGrid(arch, size=(4, 8, 6), extent=(1, 3, 4), topology = topo) + grid = RectilinearGrid(arch, size=(4, 8, 6), extent=(1, 3, 4), topology=topo) run_identity_operator_test(grid) run_poisson_equation_test(grid) end + stretch_coord = [0, 1.5, 3, 7, 8.5, 10] + for arch in archs grids = [RectilinearGrid(arch, size=(5, 5, 5), x = stretch_coord, y = (0, 10), z = (0, 10), topology = (Periodic, Periodic, Periodic)), RectilinearGrid(arch, size=(5, 5, 5), x = (0, 10), y = stretch_coord, z = (0, 10), topology = (Periodic, Periodic, Periodic)), RectilinearGrid(arch, size=(5, 5, 5), x = (0, 10), y = (0, 10), z = stretch_coord, topology = (Periodic, Periodic, Periodic))] for grid in grids - @info "Testing stretched grid MatrixIterativeSolver [$(typeof(arch))]..." + @info " Testing stretched grid MatrixIterativeSolver [$(typeof(arch))]..." run_poisson_equation_test(grid) end end - @info "Testing Sparse Approximate Inverse Preconditioner" + @info " Testing Sparse Approximate Inverse Preconditioner" A = sprand(100, 100, 0.1) A = A + A' + 1I @@ -148,5 +156,4 @@ end M = spai_preconditioner(A, ε = 0.0, nzrel = size(A, 1)) @test all(M .≈ A⁻¹) - end From 5de6818ee31ea5917b1ded2617338ed8b526f72c Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Mon, 10 Jan 2022 22:55:16 -0700 Subject: [PATCH 035/140] Random fix to matrix poisson solver test --- test/test_matrix_poisson_solver.jl | 36 ++++++++++++++++-------------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/test/test_matrix_poisson_solver.jl b/test/test_matrix_poisson_solver.jl index 4e635ee4ad..f4fd78c215 100644 --- a/test/test_matrix_poisson_solver.jl +++ b/test/test_matrix_poisson_solver.jl @@ -26,11 +26,11 @@ function run_identity_operator_test(grid) N = size(grid) M = prod(N) - b = zeros(arch, M) - A = zeros(arch, N...) - D = zeros(grid, N...) - C = zeros(arch, N...) + b = zeros(grid, M) + A = zeros(grid, N...) + D = zeros(grid, N...) + C = zeros(grid, N...) fill!(C, 1) solver = MatrixIterativeSolver((A, A, A, C, D), grid = grid) @@ -57,12 +57,12 @@ function compute_poisson_weights(grid) Ax = zeros(N...) Ay = zeros(N...) Az = zeros(N...) - C = zeros(N...) + C = zeros(grid, N...) D = zeros(grid, N...) for i = 1:grid.Nx, j = 1:grid.Ny, k = 1:grid.Nz Ax[i, j, k] = Δzᵃᵃᶜ(i, j, k, grid) * Δyᶠᶜᵃ(i, j, k, grid) / Δxᶠᶜᵃ(i, j, k, grid) Ay[i, j, k] = Δzᵃᵃᶜ(i, j, k, grid) * Δxᶜᶠᵃ(i, j, k, grid) / Δyᶜᶠᵃ(i, j, k, grid) - Az[i, j, k] = Δxᶜᵃᵃ(i, j, k, grid) * Δyᵃᶜᵃ(i, j, k, grid) / Δzᵃᵃᶠ(i, j, k, grid) + Az[i, j, k] = Δxᶜᶜᵃ(i, j, k, grid) * Δyᶜᶜᵃ(i, j, k, grid) / Δzᵃᵃᶠ(i, j, k, grid) end return (Ax, Ay, Az, C, D) end @@ -118,9 +118,9 @@ end topologies = [(Periodic, Periodic, Flat), (Bounded, Bounded, Flat), (Periodic, Bounded, Flat), (Bounded, Periodic, Flat)] for arch in archs, topo in topologies - @info " Testing 2D MatrixIterativeSolver [$(typeof(arch)) $topo]..." - - grid = RectilinearGrid(arch, size=(4, 8), extent=(1, 3), topology = topo) + @info " Testing 2D MatrixIterativeSolver [$(typeof(arch)), $topo]..." + + grid = RectilinearGrid(arch, size=(4, 8), extent=(1, 3), topology=topo) run_identity_operator_test(grid) run_poisson_equation_test(grid) end @@ -128,27 +128,29 @@ end topologies = [(Periodic, Periodic, Periodic), (Bounded, Bounded, Periodic), (Periodic, Bounded, Periodic), (Bounded, Periodic, Bounded)] for arch in archs, topo in topologies - @info " Testing 3D MatrixIterativeSolver [$(typeof(arch)) $topo]..." + @info " Testing 3D MatrixIterativeSolver [$(typeof(arch)), $topo]..." grid = RectilinearGrid(arch, size=(4, 8, 6), extent=(1, 3, 4), topology=topo) run_identity_operator_test(grid) run_poisson_equation_test(grid) end - stretch_coord = [0, 1.5, 3, 7, 8.5, 10] + stretched_faces = [0, 1.5, 3, 7, 8.5, 10] + topo = (Periodic, Periodic, Periodic) + sz = (5, 5, 5) for arch in archs - grids = [RectilinearGrid(arch, size=(5, 5, 5), x = stretch_coord, y = (0, 10), z = (0, 10), topology = (Periodic, Periodic, Periodic)), - RectilinearGrid(arch, size=(5, 5, 5), x = (0, 10), y = stretch_coord, z = (0, 10), topology = (Periodic, Periodic, Periodic)), - RectilinearGrid(arch, size=(5, 5, 5), x = (0, 10), y = (0, 10), z = stretch_coord, topology = (Periodic, Periodic, Periodic))] + grids = [RectilinearGrid(arch, size = sz, x = stretched_faces, y = (0, 10), z = (0, 10), topology = topo), + RectilinearGrid(arch, size = sz, x = (0, 10), y = stretched_faces, z = (0, 10), topology = topo), + RectilinearGrid(arch, size = sz, x = (0, 10), y = (0, 10), z = stretched_faces, topology = topo)] - for grid in grids - @info " Testing stretched grid MatrixIterativeSolver [$(typeof(arch))]..." + for (grid, stretched_direction) in zip(grids, [:x, :y, :z]) + @info " Testing MatrixIterativeSolver [stretched in $stretched_direction, $(typeof(arch))]..." run_poisson_equation_test(grid) end end - @info " Testing Sparse Approximate Inverse Preconditioner" + @info " Testing Sparse Approximate Inverse Preconditioner..." A = sprand(100, 100, 0.1) A = A + A' + 1I From de58b4f77dce94144bc658c33781ff2b5198db0b Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Mon, 10 Jan 2022 23:02:05 -0700 Subject: [PATCH 036/140] Update new_data for ConformalCubedSphereGrid --- src/CubedSpheres/cubed_sphere_faces.jl | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/src/CubedSpheres/cubed_sphere_faces.jl b/src/CubedSpheres/cubed_sphere_faces.jl index eb643f31d6..83e0a92074 100644 --- a/src/CubedSpheres/cubed_sphere_faces.jl +++ b/src/CubedSpheres/cubed_sphere_faces.jl @@ -53,13 +53,8 @@ const AbstractCubedSphereField{X, Y, Z, A} = Union{CubedSphereAbstractField{X, Y ##### new data ##### -function new_data(FT, arch::CPU, grid::ConformalCubedSphereGrid, (X, Y, Z)) - faces = Tuple(new_data(FT, arch, face_grid, (X, Y, Z)) for face_grid in grid.faces) - return CubedSphereFaces{typeof(faces[1]), typeof(faces)}(faces) -end - -function new_data(FT, arch::GPU, grid::ConformalCubedSphereGrid, (X, Y, Z)) - faces = Tuple(new_data(FT, arch, face_grid, (X, Y, Z)) for face_grid in grid.faces) +function new_data(FT, grid::ConformalCubedSphereGrid, (X, Y, Z)) + faces = Tuple(new_data(FT, face_grid, (X, Y, Z)) for face_grid in grid.faces) return CubedSphereFaces{typeof(faces[1]), typeof(faces)}(faces) end From 34c798a3ba3ee0a48b9851247c320b4620681c7c Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 06:02:47 -0700 Subject: [PATCH 037/140] Update validate_field_data for cubed sphere fields --- src/CubedSpheres/CubedSpheres.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/CubedSpheres/CubedSpheres.jl b/src/CubedSpheres/CubedSpheres.jl index c17f4d30f7..122cf4dee7 100644 --- a/src/CubedSpheres/CubedSpheres.jl +++ b/src/CubedSpheres/CubedSpheres.jl @@ -18,10 +18,10 @@ include("immersed_conformal_cubed_sphere_grid.jl") import Oceananigans.Fields: validate_field_data import Oceananigans.Models.HydrostaticFreeSurfaceModels: validate_vertical_velocity_boundary_conditions -function validate_field_data(X, Y, Z, data, grid::ConformalCubedSphereGrid) +function validate_field_data(loc, data, grid::ConformalCubedSphereGrid) for (face_data, face_grid) in zip(data.faces, grid.faces) - validate_field_data(X, Y, Z, face_data, face_grid) + validate_field_data(loc, face_data, face_grid) end return nothing From c8799746769eac19c1eb20949cdbd3a7c8a02f71 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 06:03:06 -0700 Subject: [PATCH 038/140] Dont use WENO5 advection for conjugate gradient solver test --- ...rface_immersed_boundaries_congrad_solve.jl | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/test/test_hydrostatic_free_surface_immersed_boundaries_congrad_solve.jl b/test/test_hydrostatic_free_surface_immersed_boundaries_congrad_solve.jl index 96ddfb5a8c..961f821227 100644 --- a/test/test_hydrostatic_free_surface_immersed_boundaries_congrad_solve.jl +++ b/test/test_hydrostatic_free_surface_immersed_boundaries_congrad_solve.jl @@ -40,7 +40,6 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels: pressure_correct_velocit model = HydrostaticFreeSurfaceModel(grid = grid, free_surface = free_surface, - tracer_advection = WENO5(), buoyancy = nothing, tracers = nothing, closure = nothing) @@ -48,10 +47,10 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels: pressure_correct_velocit # Now create a divergent flow field and solve for # pressure correction u, v, w = model.velocities - u[imm1,jmm1,1:Nz ] .= 1. - u[imp1,jmm1,1:Nz ] .= -1. - v[imm1,jmm1,1:Nz ] .= 1. - v[imm1,jmp1,1:Nz ] .= -1. + u[imm1, jmm1, 1:Nz ] .= 1 + u[imp1, jmm1, 1:Nz ] .= -1 + v[imm1, jmm1, 1:Nz ] .= 1 + v[imm1, jmp1, 1:Nz ] .= -1 η = model.free_surface.η g = model.free_surface.gravitational_acceleration @@ -67,23 +66,23 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels: pressure_correct_velocit fill_halo_regions!(η, arch) - # println(typeof(model)) - + #= println("model.free_surface.gravitational_acceleration = ",model.free_surface.gravitational_acceleration) println("∫ᶻQ.u") - show(stdout,"text/plain",∫ᶻQ.u.data[1:Nx,1:Ny,1]) + show(stdout,"text/plain", ∫ᶻQ.u.data[1:Nx, 1:Ny, 1]) println("") println("η") - show(stdout,"text/plain",η.data[1:Nx,1:Ny,1]) + show(stdout,"text/plain", η.data[1:Nx, 1:Ny, 1]) println("") pressure_correct_velocities!(model, Δt) fill_halo_regions!(u, arch) println("u") - show(stdout,"text/plain",u.data[1:Nx,1:Ny,1]) + show(stdout,"text/plain", u.data[1:Nx, 1:Ny, 1]) println("") + =# fs = model.free_surface vertically_integrated_lateral_areas = fs.implicit_step_solver.vertically_integrated_lateral_areas From dbac16d3497b6ec28cdeac7f510ba55a38ce11ad Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 06:11:43 -0700 Subject: [PATCH 039/140] Import fill_halo_regions into field.jl plus cosmetics --- src/BoundaryConditions/fill_halo_regions.jl | 32 +++++++++------------ src/Distributed/halo_communication.jl | 5 ++-- src/Fields/field.jl | 3 +- 3 files changed, 19 insertions(+), 21 deletions(-) diff --git a/src/BoundaryConditions/fill_halo_regions.jl b/src/BoundaryConditions/fill_halo_regions.jl index e48ac6eb2c..2c1c47c0e7 100644 --- a/src/BoundaryConditions/fill_halo_regions.jl +++ b/src/BoundaryConditions/fill_halo_regions.jl @@ -26,7 +26,7 @@ end fill_halo_regions!(c::OffsetArray, ::Nothing, args...; kwargs...) = nothing "Fill halo regions in x, y, and z for a given field's data." -function fill_halo_regions!(c::OffsetArray, field_bcs, arch, grid, args...; kwargs...) +function fill_halo_regions!(c::OffsetArray, boundary_conditions, arch, grid, args...; kwargs...) fill_halos! = [ fill_west_and_east_halo!, @@ -34,39 +34,35 @@ function fill_halo_regions!(c::OffsetArray, field_bcs, arch, grid, args...; kwar fill_bottom_and_top_halo!, ] - field_bcs_array_left = [ - field_bcs.west, - field_bcs.south, - field_bcs.bottom, + boundary_conditions_array_left = [ + boundary_conditions.west, + boundary_conditions.south, + boundary_conditions.bottom, ] - field_bcs_array_right = [ - field_bcs.east, - field_bcs.north, - field_bcs.top, + boundary_conditions_array_right = [ + boundary_conditions.east, + boundary_conditions.north, + boundary_conditions.top, ] - perm = sortperm(field_bcs_array_left, lt=fill_first) + perm = sortperm(boundary_conditions_array_left, lt=fill_first) fill_halos! = fill_halos![perm] - field_bcs_array_left = field_bcs_array_left[perm] - field_bcs_array_right = field_bcs_array_right[perm] + boundary_conditions_array_left = boundary_conditions_array_left[perm] + boundary_conditions_array_right = boundary_conditions_array_right[perm] for task = 1:3 - barrier = device_event(arch) fill_halo! = fill_halos![task] - bc_left = field_bcs_array_left[task] - bc_right = field_bcs_array_right[task] + bc_left = boundary_conditions_array_left[task] + bc_right = boundary_conditions_array_right[task] events = fill_halo!(c, bc_left, bc_right, arch, barrier, grid, args...; kwargs...) wait(device(arch), events) - end - - return nothing end diff --git a/src/Distributed/halo_communication.jl b/src/Distributed/halo_communication.jl index 25e3e54d39..934a8d46ce 100644 --- a/src/Distributed/halo_communication.jl +++ b/src/Distributed/halo_communication.jl @@ -1,4 +1,3 @@ -using Oceananigans.Fields: AbstractField using KernelAbstractions: @kernel, @index, Event, MultiEvent using OffsetArrays: OffsetArray @@ -107,7 +106,9 @@ for (side, opposite_side) in zip([:west, :south, :bottom], [:east, :north, :top] recv_and_fill_opposite_side_halo! = Symbol("recv_and_fill_$(opposite_side)_halo!") @eval begin - function $fill_both_halos!(c, bc_side::HaloCommunicationBC, bc_opposite_side::HaloCommunicationBC, arch, barrier, grid, c_location, args...) + function $fill_both_halos!(c, bc_side::HaloCommunicationBC, bc_opposite_side::HaloCommunicationBC, + arch, barrier, grid, c_location, args...) + @assert bc_side.condition.from == bc_opposite_side.condition.from # Extra protection in case of bugs local_rank = bc_side.condition.from diff --git a/src/Fields/field.jl b/src/Fields/field.jl index 95110fa89b..07cdb949c8 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -2,9 +2,10 @@ using Oceananigans.Architectures: device_event using Adapt using KernelAbstractions: @kernel, @index - using Base: @propagate_inbounds +import Oceananigans.BoundaryConditions: fill_halo_regions! + struct Field{LX, LY, LZ, O, A, G, T, D, B, S} <: AbstractField{LX, LY, LZ, A, G, T, 3} grid :: G data :: D From e3f0c44501e7ac3b3958a40cf7e52b718498cbf7 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 06:32:02 -0700 Subject: [PATCH 040/140] Change ComputedField to Field in docs --- docs/src/model_setup/output_writers.md | 2 +- docs/src/simulation_tips.md | 93 +++++++++++--------------- 2 files changed, 39 insertions(+), 56 deletions(-) diff --git a/docs/src/model_setup/output_writers.md b/docs/src/model_setup/output_writers.md index 4303a581ef..6c2e90639a 100644 --- a/docs/src/model_setup/output_writers.md +++ b/docs/src/model_setup/output_writers.md @@ -180,7 +180,7 @@ function init_save_some_metadata!(file, model) return nothing end -c_avg = AveragedField(model.tracers.c, dims=(1, 2)) +c_avg = Field(Average(model.tracers.c, dims=(1, 2))) # Note that model.velocities is NamedTuple simulation.output_writers[:velocities] = JLD2OutputWriter(model, model.velocities, diff --git a/docs/src/simulation_tips.md b/docs/src/simulation_tips.md index fcaec07e59..83b5a165fe 100644 --- a/docs/src/simulation_tips.md +++ b/docs/src/simulation_tips.md @@ -1,19 +1,15 @@ - # Simulation tips -In Oceananigans we try to do most of the optimizing behind the scenes, that way the average user -doesn't have to worry about details when setting up a simulation. However, there's just so much -optimization that can be done in the source code. Because of Oceananigans script-based interface, -the user has to be aware of some things when writing the simulation script in order to take full -advantage of Julia's speed. Furthermore, in case of more complex GPU runs, some details could +Oceananigans attemps to optimize as much of a computation as possible "behind the scenes". +Yet Oceananigans' flexibility places some responsibility on users to ensure high performance simulations, +especially for complex setups with user-defined forcing functions, boundary condition functions, and diagnostics. +Furthermore, in case of more complex GPU runs, some details could sometimes prevent your simulation from running altogether. While Julia knowledge is obviously desirable here, a user that is unfamiliar with Julia can get away with efficient simulations by learning a few rules of thumb. It is nonetheless recommended that users go through Julia's [performance tips](https://docs.julialang.org/en/v1/manual/performance-tips/), which contains more in-depth explanations of some of the aspects discussed here. - - ## General (CPU/GPU) simulation tips ### Avoid global variables whenever possible @@ -37,7 +33,6 @@ GPU kernels (such as functions defining boundary conditions and forcings). Other compiler can fail with obscure errors. This is explained in more detail in the GPU simulation tips section below. - ### Consider inlining small functions Inlining is when the compiler [replaces a function call with the body of the function that is being @@ -58,8 +53,6 @@ certainty_, since Julia and KernelAbstractions.jl (needed for GPU runs) already functions automatically. However, it is generally a good idea to at least investigate this aspect in your code as the benefits can potentially be significant. - - ## GPU simulation tips Running on GPUs can be very different from running on CPUs. Oceananigans makes most of the necessary @@ -69,7 +62,6 @@ for more complex simulations some care needs to be taken on the part of the user GPU computing (and Julia) is again desirable, an inexperienced user can also achieve high efficiency in GPU simulations by following a few simple principles. - ### Global variables that need to be used in GPU computations need to be defined as constants or passed as parameters Any global variable that needs to be accessed by the GPU needs to be a constant or the simulation @@ -102,9 +94,9 @@ surface_temperature(x, y, t, p) = p.T₀ * sin(2π / 86400 * t) T_bcs = FieldBoundaryConditions(bottom = GradientBoundaryCondition(surface_temperature, parameters=(T₀=T₀,))) ``` -### Complex diagnostics using `ComputedField`s may not work on GPUs +### Complex diagnostics using computed `Field`s may not work on GPUs -`ComputedField`s are the most convenient way to calculate diagnostics for your simulation. They will +`Field`s are the most convenient way to calculate diagnostics for your simulation. They will always work on CPUs, but when their complexity is high (in terms of number of abstract operations) the compiler can't translate them into GPU code and they fail for GPU runs. (This limitation is summarized in [this Github issue](https://github.com/CliMA/Oceananigans.jl/issues/1886) and contributions are welcome.) @@ -117,40 +109,33 @@ grid = RectilinearGrid(size=(4, 4, 4), extent=(1, 1, 1)) model = NonhydrostaticModel(grid=grid, closure=IsotropicDiffusivity(ν=1e-6)) u, v, w = model.velocities ν = model.closure.ν -u² = ComputedField(u^2) -ε = ComputedField(ν*(∂x(u)^2 + ∂x(v)^2 + ∂x(w)^2 + ∂y(u)^2 + ∂y(v)^2 + ∂y(w)^2 + ∂z(u)^2 + ∂z(v)^2 + ∂z(w)^2)) +u² = Field(u^2) +ε = Field(ν*(∂x(u)^2 + ∂x(v)^2 + ∂x(w)^2 + ∂y(u)^2 + ∂y(v)^2 + ∂y(w)^2 + ∂z(u)^2 + ∂z(v)^2 + ∂z(w)^2)) compute!(u²) compute!(ε) ``` -There are two approaches to -bypass this issue. The first is to nest `ComputedField`s. For example, -we can make `ε` be successfully computed on GPUs by defining it as +There are a few ways to work around this issue. +One is to compute `ε` in steps by nesting computed `Field`s, ```julia -ddx² = ComputedField(∂x(u)^2 + ∂x(v)^2 + ∂x(w)^2) -ddy² = ComputedField(∂y(u)^2 + ∂y(v)^2 + ∂y(w)^2) -ddz² = ComputedField(∂z(u)^2 + ∂z(v)^2 + ∂z(w)^2) -ε = ComputedField(ν*(ddx² + ddy² + ddz²)) +ddx² = Field(∂x(u)^2 + ∂x(v)^2 + ∂x(w)^2) +ddy² = Field(∂y(u)^2 + ∂y(v)^2 + ∂y(w)^2) +ddz² = Field(∂z(u)^2 + ∂z(v)^2 + ∂z(w)^2) +ε = Field(ν * (ddx² + ddy² + ddz²)) compute!(ε) ``` -This is a simple workaround that is especially suited for the development stage of a simulation. -However, when running this, the code will iterate over the whole domain 4 times to calculate `ε` -(one for each computed field defined), which is not very efficient and may slow down your simulation -if this diagnostic is being calculated very often. - -A different way to calculate `ε` is by using `KernelFunctionOperations`s, where the -user manually specifies the computing kernel function to the compiler. The advantage of this method is that -it's more efficient (the code will only iterate once over the domain in order to calculate `ε`), -but the disadvantage is that this method requires some knowledge of Oceananigans operations -and how they should be performed on a C-grid. For example calculating `ε` with this approach would -look like this: +This method is expensive because it requires computing and storing 3 intermediate terms. +`ε` may also be calculated via `KernelFunctionOperations`s, which +requires explicitly building a "kernel function" from low-level Oceananigans +operators. ```julia using Oceananigans.Operators using Oceananigans.AbstractOperations: KernelFunctionOperation @inline fψ_plus_gφ²(i, j, k, grid, f, ψ, g, φ) = @inbounds (f(i, j, k, grid, ψ) + g(i, j, k, grid, φ))^2 + function isotropic_viscous_dissipation_rate_ccc(i, j, k, grid, u, v, w, ν) Σˣˣ² = ∂xᶜᵃᵃ(i, j, k, grid, u)^2 Σʸʸ² = ∂yᵃᶜᵃ(i, j, k, grid, v)^2 @@ -162,45 +147,46 @@ function isotropic_viscous_dissipation_rate_ccc(i, j, k, grid, u, v, w, ν) return ν * 2 * (Σˣˣ² + Σʸʸ² + Σᶻᶻ² + 2 * (Σˣʸ² + Σˣᶻ² + Σʸᶻ²)) end -ε = ComputedField(KernelFunctionOperation{Center, Center, Center}(isotropic_viscous_dissipation_rate_ccc, grid; - computed_dependencies=(u, v, w, ν))) + +ε_op = KernelFunctionOperation{Center, Center, Center}(isotropic_viscous_dissipation_rate_ccc, + grid; + computed_dependencies=(u, v, w, ν)) + +ε = Field(ε_op) + compute!(ε) ``` +Writing kernel functions like `isotropic_viscous_dissipation_rate_ccc` +requires understanding the C-grid, but incurs only one iteration over the domain. -It may be useful to know that there are some kernels already defined for commonly-used diagnostics -in packages that are companions to Oceananigans. For example -[Oceanostics.jl](https://github.com/tomchor/Oceanostics.jl/blob/3b8f67338656557877ef8ef5ebe3af9e7b2974e2/src/TurbulentKineticEnergyTerms.jl#L35-L57) -and -[LESbrary.jl](https://github.com/CliMA/LESbrary.jl/blob/master/src/TurbulenceStatistics/shear_production.jl). -Users should first look there before writing any kernel by hand and are always welcome to [start an -issue on Github](https://github.com/CliMA/Oceananigans.jl/issues/new) if they need help to write a -different kernel. As an illustration, the calculation of `ε` using Oceanostics.jl (after installing the package) -which works on both CPUs and GPUs is simply +`KernelFunctionOperation`s for some diagnostics common to large eddy simulation are defined in +[Oceanostics.jl](https://github.com/tomchor/Oceanostics.jl/blob/3b8f67338656557877ef8ef5ebe3af9e7b2974e2/src/TurbulentKineticEnergyTerms.jl#L35-L57), ```julia using Oceanostics: IsotropicPseudoViscousDissipationRate ε = IsotropicViscousDissipationRate(model, u, v, w, ν) compute!(ε) ``` +[Start an issue on Github](https://github.com/CliMA/Oceananigans.jl/issues/new) if more help is needed. ### Try to decrease the memory-use of your runs -GPU runs are generally memory-limited. As an example, a state-of-the-art Tesla V100 GPU has 32GB of -memory, which is enough to fit, on average, a simulation with about 100 million points --- a bit -smaller than a 512-cubed simulation. (The precise number depends on many other things, such as the -number of tracers simulated, as well as the diagnostics that are calculated.) This means that it is -especially important to be mindful of the size of your runs when running Oceananigans on GPUs and it -is generally good practice to decrease the memory required for your runs. Below are some useful tips -to achieve this +GPU runs are sometimes memory-limited. A state-of-the-art Tesla V100 GPU has 32GB of +memory --- enough memory for simulations with about 100 million points, or grids a bit smaller +than 512 × 512 × 512. (The maximum grid size depends on some user-specified factors, +like the number of passive tracers or computed diagnostics.) +For large simulations on the GPU, careful management of memory allocation may be required: - Use the [`nvidia-smi`](https://developer.nvidia.com/nvidia-system-management-interface) command line utility to monitor the memory usage of the GPU. It should tell you how much memory there is on your GPU and how much of it you're using and you can run it from Julia with the command ``run(`nvidia-smi`)``. + - Try to use higher-order advection schemes. In general when you use a higher-order scheme you need fewer grid points to achieve the same accuracy that you would with a lower-order one. Oceananigans provides two high-order advection schemes: 5th-order WENO method (WENO5) and 3rd-order upwind. + - Manually define scratch space to be reused in diagnostics. By default, every time a user-defined diagnostic is calculated the compiler reserves a new chunk of memory for that calculation, usually called scratch space. In general, the more diagnostics, the more scratch space needed and the bigger @@ -212,8 +198,6 @@ to achieve this and then being used in calculations [here](https://github.com/CliMA/LESbrary.jl/blob/cf31b0ec20219d5ad698af334811d448c27213b0/src/TurbulenceStatistics/first_through_third_order.jl#L109-L112). - - ### Arrays in GPUs are usually different from arrays in CPUs Oceananigans.jl uses [`CUDA.CuArray`](https://cuda.juliagpu.org/stable/usage/array/) to store @@ -227,7 +211,6 @@ which is very slow and can result in huge slowdowns. For this reason, Oceananiga scalar indexing by default. See the [scalar indexing](https://juliagpu.github.io/CUDA.jl/dev/usage/workflow/#UsageWorkflowScalar) section of the CUDA.jl documentation for more information on scalar indexing. - For example if can be difficult to just view a `CuArray` since Julia needs to access its elements to do that. Consider the example below: From 3b28967027367a3d33580868fe6527b23df544b1 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 06:33:35 -0700 Subject: [PATCH 041/140] Use architecture(field) --- src/CubedSpheres/cubed_sphere_faces.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/CubedSpheres/cubed_sphere_faces.jl b/src/CubedSpheres/cubed_sphere_faces.jl index 83e0a92074..91a0c35f35 100644 --- a/src/CubedSpheres/cubed_sphere_faces.jl +++ b/src/CubedSpheres/cubed_sphere_faces.jl @@ -103,7 +103,7 @@ Base.size(data::CubedSphereData) = (size(data.faces[1])..., length(data.faces)) # Should we define a new lower-level constructor for Field that doesn't call validate_field_data? return Field{X, Y, Z}(get_face(field.data, face_index), - field.architecture, + architecture(field) get_face(field.grid, face_index), get_face(field.boundary_conditions, face_index)) end @@ -112,9 +112,9 @@ end X, Y, Z = location(kernel_field) return Field{X, Y, Z}(get_face(kernel_field.data, face_index), - kernel_field.architecture, - get_face(kernel_field.grid, face_index), - get_face(kernel_field.boundary_conditions, face_index)) + architecture(kernel_field) + get_face(kernel_field.grid, face_index), + get_face(kernel_field.boundary_conditions, face_index)) end faces(field::AbstractCubedSphereField) = Tuple(get_face(field, face_index) for face_index in 1:length(field.data.faces)) From f6f20b3a19ffc1fe58e80ef942330d61a5a548cb Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 06:42:01 -0700 Subject: [PATCH 042/140] Missing comma --- src/CubedSpheres/cubed_sphere_faces.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/CubedSpheres/cubed_sphere_faces.jl b/src/CubedSpheres/cubed_sphere_faces.jl index 91a0c35f35..69b8830180 100644 --- a/src/CubedSpheres/cubed_sphere_faces.jl +++ b/src/CubedSpheres/cubed_sphere_faces.jl @@ -103,7 +103,7 @@ Base.size(data::CubedSphereData) = (size(data.faces[1])..., length(data.faces)) # Should we define a new lower-level constructor for Field that doesn't call validate_field_data? return Field{X, Y, Z}(get_face(field.data, face_index), - architecture(field) + architecture(field), get_face(field.grid, face_index), get_face(field.boundary_conditions, face_index)) end From 6e161b7828aec9db95aad09b850d8ffe1a50a368 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 06:42:11 -0700 Subject: [PATCH 043/140] Figuring out fill_halo_regions for distributed fields --- src/BoundaryConditions/fill_halo_regions.jl | 2 -- src/Distributed/distributed_fields.jl | 12 +++++++++++- src/Distributed/halo_communication.jl | 2 +- 3 files changed, 12 insertions(+), 4 deletions(-) diff --git a/src/BoundaryConditions/fill_halo_regions.jl b/src/BoundaryConditions/fill_halo_regions.jl index 2c1c47c0e7..25a6354cf6 100644 --- a/src/BoundaryConditions/fill_halo_regions.jl +++ b/src/BoundaryConditions/fill_halo_regions.jl @@ -14,11 +14,9 @@ Fill halo regions for each field in the tuple `fields` according to their bounda conditions, possibly recursing into `fields` if it is a nested tuple-of-tuples. """ function fill_halo_regions!(fields::Union{Tuple, NamedTuple}, arch, args...) - for field in fields fill_halo_regions!(field, arch, args...) end - return nothing end diff --git a/src/Distributed/distributed_fields.jl b/src/Distributed/distributed_fields.jl index d2ff6502e9..5d54923f78 100644 --- a/src/Distributed/distributed_fields.jl +++ b/src/Distributed/distributed_fields.jl @@ -1,7 +1,17 @@ -import Oceananigans.Fields: Field +import Oceananigans.Fields: Field, fill_halo_regions! function Field((LX, LY, LZ)::Tuple, grid::DistributedGrid, data, old_bcs, op, status) arch = architecture(grid) new_bcs = inject_halo_communication_boundary_conditions(old_bcs, arch.local_rank, arch.connectivity) return Field{LX, LY, LZ}(grid, data, new_bcs, op, status) end + +const DistributedField = AbstractField{<:Any, <:Any, <:Any, <:MultiArch} + +fill_halo_regions!(field::DistributedField, arch, args...; kwargs...) = + fill_halo_regions!(field.data, + field.boundary_conditions, + arch, + field.grid, + location(field), + args...; kwargs...) diff --git a/src/Distributed/halo_communication.jl b/src/Distributed/halo_communication.jl index 934a8d46ce..6194fd656a 100644 --- a/src/Distributed/halo_communication.jl +++ b/src/Distributed/halo_communication.jl @@ -52,7 +52,7 @@ end ##### Filling halos for halo communication boundary conditions ##### -function fill_halo_regions!(c::OffsetArray, bcs, arch::AbstractMultiArchitecture, grid, c_location, args...) +function fill_halo_regions!(c::OffsetArray, bcs, arch::MultiArch, grid, c_location, args...) barrier = Event(device(child_architecture(arch))) From 733d228cd38eda63b5e2bae0a443355ae7a736ac Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 06:45:22 -0700 Subject: [PATCH 044/140] Gets rid of ReducedField --- test/test_implicit_free_surface_solver.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/test_implicit_free_surface_solver.jl b/test/test_implicit_free_surface_solver.jl index f08d21b784..e8247a6eb5 100644 --- a/test/test_implicit_free_surface_solver.jl +++ b/test/test_implicit_free_surface_solver.jl @@ -58,7 +58,7 @@ function run_pcg_implicit_free_surface_solver_tests(arch, grid) ∫ᶻ_Axᶠᶜᶜ = model.free_surface.implicit_step_solver.vertically_integrated_lateral_areas.xᶠᶜᶜ ∫ᶻ_Ayᶜᶠᶜ = model.free_surface.implicit_step_solver.vertically_integrated_lateral_areas.yᶜᶠᶜ - left_hand_side = ReducedField(Center, Center, Nothing, arch, grid; dims=3) + left_hand_side = Field{Center, Center, Nothing}(grid) implicit_free_surface_linear_operation!(left_hand_side, η, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, g, Δt) # Compare @@ -101,7 +101,7 @@ function run_matrix_implicit_free_surface_solver_tests(arch, grid) ∫ᶻ_Axᶠᶜᶜ = model.free_surface.implicit_step_solver.vertically_integrated_lateral_areas.xᶠᶜᶜ ∫ᶻ_Ayᶜᶠᶜ = model.free_surface.implicit_step_solver.vertically_integrated_lateral_areas.yᶜᶠᶜ - left_hand_side = ReducedField(Center, Center, Nothing, arch, grid; dims=3) + left_hand_side = Field{Center, Center, Nothing}(grid) implicit_free_surface_linear_operation!(left_hand_side, η, ∫ᶻ_Axᶠᶜᶜ, ∫ᶻ_Ayᶜᶠᶜ, g, Δt) # Compare From 0956f8494681a6ff5764d1f43cf9a2b1d2b494a8 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 07:50:49 -0700 Subject: [PATCH 045/140] Missing comma --- src/CubedSpheres/cubed_sphere_faces.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/CubedSpheres/cubed_sphere_faces.jl b/src/CubedSpheres/cubed_sphere_faces.jl index 69b8830180..e31fcfb35a 100644 --- a/src/CubedSpheres/cubed_sphere_faces.jl +++ b/src/CubedSpheres/cubed_sphere_faces.jl @@ -112,7 +112,7 @@ end X, Y, Z = location(kernel_field) return Field{X, Y, Z}(get_face(kernel_field.data, face_index), - architecture(kernel_field) + architecture(kernel_field), get_face(kernel_field.grid, face_index), get_face(kernel_field.boundary_conditions, face_index)) end From 9ac2d1be2654fbb0679b7cb41b873488747694f9 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 08:08:56 -0700 Subject: [PATCH 046/140] Dispatch shenanigans --- src/CubedSpheres/cubed_sphere_faces.jl | 27 +++++++++----------------- src/Distributed/halo_communication.jl | 8 ++++---- src/Fields/abstract_field.jl | 8 -------- 3 files changed, 13 insertions(+), 30 deletions(-) diff --git a/src/CubedSpheres/cubed_sphere_faces.jl b/src/CubedSpheres/cubed_sphere_faces.jl index e31fcfb35a..03e0ee0d8f 100644 --- a/src/CubedSpheres/cubed_sphere_faces.jl +++ b/src/CubedSpheres/cubed_sphere_faces.jl @@ -98,24 +98,15 @@ end Base.size(data::CubedSphereData) = (size(data.faces[1])..., length(data.faces)) -@inline function get_face(field::CubedSphereField, face_index) - X, Y, Z = location(field) - - # Should we define a new lower-level constructor for Field that doesn't call validate_field_data? - return Field{X, Y, Z}(get_face(field.data, face_index), - architecture(field), - get_face(field.grid, face_index), - get_face(field.boundary_conditions, face_index)) -end - -@inline function get_face(kernel_field::CubedSphereAbstractField, face_index) - X, Y, Z = location(kernel_field) - - return Field{X, Y, Z}(get_face(kernel_field.data, face_index), - architecture(kernel_field), - get_face(kernel_field.grid, face_index), - get_face(kernel_field.boundary_conditions, face_index)) -end +@inline get_face(field::CubedSphereField, face_index) = + Field(location(field), get_face(field.grid, face_index); + data = get_face(field.data, face_index), + boundary_conditions = get_face(field.boundary_conditions, face_index)) + +@inline get_face(field::CubedSphereAbstractField, face_index) = + Field(location(field), get_face(field.grid, face_index); + data = get_face(field.data, face_index), + boundary_conditions = get_face(field.boundary_conditions, face_index)) faces(field::AbstractCubedSphereField) = Tuple(get_face(field, face_index) for face_index in 1:length(field.data.faces)) diff --git a/src/Distributed/halo_communication.jl b/src/Distributed/halo_communication.jl index 6194fd656a..73c915b710 100644 --- a/src/Distributed/halo_communication.jl +++ b/src/Distributed/halo_communication.jl @@ -52,13 +52,13 @@ end ##### Filling halos for halo communication boundary conditions ##### -function fill_halo_regions!(c::OffsetArray, bcs, arch::MultiArch, grid, c_location, args...) +function fill_halo_regions!(c::OffsetArray, bcs, arch::MultiArch, grid, c_location, args...; kwargs...) barrier = Event(device(child_architecture(arch))) - x_events_requests = fill_west_and_east_halos!(c, bcs.west, bcs.east, arch, barrier, grid, c_location, args...) - y_events_requests = fill_south_and_north_halos!(c, bcs.south, bcs.north, arch, barrier, grid, c_location, args...) - z_events_requests = fill_bottom_and_top_halos!(c, bcs.bottom, bcs.top, arch, barrier, grid, c_location, args...) + x_events_requests = fill_west_and_east_halos!(c, bcs.west, bcs.east, arch, barrier, grid, c_location, args...; kwargs...) + y_events_requests = fill_south_and_north_halos!(c, bcs.south, bcs.north, arch, barrier, grid, c_location, args...; kwargs...) + z_events_requests = fill_bottom_and_top_halos!(c, bcs.bottom, bcs.top, arch, barrier, grid, c_location, args...; kwargs...) events_and_requests = [x_events_requests..., y_events_requests..., z_events_requests...] diff --git a/src/Fields/abstract_field.jl b/src/Fields/abstract_field.jl index 9844bef44f..b80cdc735a 100644 --- a/src/Fields/abstract_field.jl +++ b/src/Fields/abstract_field.jl @@ -12,7 +12,6 @@ import Base: minimum, maximum, extrema import Statistics: mean import Oceananigans: location, instantiated_location import Oceananigans.Architectures: architecture -import Oceananigans.BoundaryConditions: fill_halo_regions! import Oceananigans.Grids: interior_x_indices, interior_y_indices, interior_z_indices import Oceananigans.Grids: total_size, topology, nodes, xnodes, ynodes, znodes, xnode, ynode, znode import Oceananigans.Utils: datatuple @@ -121,13 +120,6 @@ znodes(ψ::AbstractField) = znodes(location(ψ, 3), ψ.grid) nodes(ψ::AbstractField; kwargs...) = nodes(location(ψ), ψ.grid; kwargs...) -##### -##### fill_halo_regions! -##### - -fill_halo_regions!(field::AbstractField, arch, args...) = - fill_halo_regions!(field.data, field.boundary_conditions, arch, field.grid, args...) - ##### ##### Some conveniences ##### From df643d11ca998a73e8fa1e7b522db9aa20c6365c Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 10:12:44 -0700 Subject: [PATCH 047/140] CubedSphere updates --- src/CubedSpheres/cubed_sphere_faces.jl | 6 +- test/test_averaged_field.jl | 98 ------------------------- test/test_cubed_sphere_halo_exchange.jl | 2 +- test/test_cubed_spheres.jl | 2 +- 4 files changed, 4 insertions(+), 104 deletions(-) delete mode 100644 test/test_averaged_field.jl diff --git a/src/CubedSpheres/cubed_sphere_faces.jl b/src/CubedSpheres/cubed_sphere_faces.jl index 03e0ee0d8f..d0b73b20cc 100644 --- a/src/CubedSpheres/cubed_sphere_faces.jl +++ b/src/CubedSpheres/cubed_sphere_faces.jl @@ -43,8 +43,8 @@ const CubedSphereFaceField = Union{NonImmersedCubedSphereFaceField{X, Y, Z, A}, # CubedSphereField # Flavors of CubedSphereField -const CubedSphereField = Field{X, Y, Z, A, <:CubedSphereData} where {X, Y, Z, A} -const CubedSphereAbstractField = AbstractField{X, Y, Z, A, <:ConformalCubedSphereGrid} where {X, Y, Z, A} +const CubedSphereField = Field{X, Y, Z, O, A, G, T, <:CubedSphereData} where {X, Y, Z, O, A, G, T} +const CubedSphereAbstractField = AbstractField{X, Y, Z, A, <:ConformalCubedSphereGrid} where {X, Y, Z, A} const AbstractCubedSphereField{X, Y, Z, A} = Union{CubedSphereAbstractField{X, Y, Z, A}, CubedSphereField{X, Y, Z, A}} where {X, Y, Z, A} @@ -89,8 +89,6 @@ function Base.show(io::IO, field::AbstractCubedSphereField) "└── grid: $(short_show(field.grid))") end - - @inline function interior(field::AbstractCubedSphereField) faces = Tuple(interior(face_field_face) for face_field in faces(field)) return CubedSphereFaces(faces) diff --git a/test/test_averaged_field.jl b/test/test_averaged_field.jl deleted file mode 100644 index 02765b325a..0000000000 --- a/test/test_averaged_field.jl +++ /dev/null @@ -1,98 +0,0 @@ -using Statistics - -using Oceananigans.Fields: CenterField, ZFaceField, compute_at! -using Oceananigans.Grids: halo_size - -@testset "Averaged fields" begin - @info "Testing averaged fields..." - - for arch in archs - @testset "Averaged fields [$(typeof(arch))]" begin - @info " Testing AveragedFields [$(typeof(arch))]" - for FT in float_types - - grid = RectilinearGrid(topology = (Periodic, Periodic, Bounded), - size = (2, 2, 2), - x = (0, 2), y = (0, 2), z = (0, 2)) - - w = ZFaceField(arch, grid) - T = CenterField(arch, grid) - - trilinear(x, y, z) = x + y + z - - set!(T, trilinear) - set!(w, trilinear) - - @compute T̃ = AveragedField(T, dims=(1, 2, 3)) - - # Note: halo regions must be *filled* prior to computing an average - # if the average within halo regions is to be correct. - fill_halo_regions!(T, arch) - @compute T̅ = AveragedField(T, dims=(1, 2)) - - fill_halo_regions!(T, arch) - @compute T̂ = AveragedField(T, dims=1) - - @compute w̃ = AveragedField(w, dims=(1, 2, 3)) - @compute w̅ = AveragedField(w, dims=(1, 2)) - @compute ŵ = AveragedField(w, dims=1) - - Nx, Ny, Nz = grid.Nx, grid.Ny, grid.Nz - - @test T̃[1, 1, 1] ≈ 3 - - @test Array(interior(T̅))[1, 1, :] ≈ [2.5, 3.5] - @test Array(interior(T̂))[1, :, :] ≈ [[2, 3] [3, 4]] - - @test w̃[1, 1, 1] ≈ 3 - - @test Array(interior(w̅))[1, 1, :] ≈ [2, 3, 4] - @test Array(interior(ŵ))[1, :, :] ≈ [[1.5, 2.5] [2.5, 3.5] [3.5, 4.5]] - - # Test whether a race condition gets hit for averages over large fields - big_grid = RectilinearGrid(topology = (Periodic, Periodic, Bounded), - size = (256, 256, 128), - x = (0, 2), y = (0, 2), z = (0, 2)) - - c = CenterField(arch, big_grid) - c .= 1 - - C = AveragedField(c, dims=(1, 2)) - - # Test that the mean consistently returns 1 at every z for many evaluations - results = [all(interior(mean!(C, C.operand)) .== 1) for i = 1:10] # warm up... - results = [all(interior(mean!(C, C.operand)) .== 1) for i = 1:10] # the real deal - @test mean(results) == 1.0 - end - end - - @testset "Conditional computation of AveragedFields [$(typeof(arch))]" begin - @info " Testing conditional computation of AveragedFields [$(typeof(arch))]" - for FT in float_types - grid = RectilinearGrid(FT, size=(2, 2, 2), extent=(1, 1, 1)) - c = CenterField(arch, grid) - - for dims in (1, 2, 3, (1, 2), (2, 3), (1, 3), (1, 2, 3)) - C = AveragedField(c, dims=dims) - - @test !isnothing(C.status) - - # Test conditional computation - set!(c, 1) - compute_at!(C, FT(1)) # will compute - @test all(interior(C) .== 1) - @test C.status.time == FT(1) - - set!(c, 2) - compute_at!(C, FT(1)) # will not compute because status == 1 - @test C.status.time == FT(1) - @test all(interior(C) .== 1) - - compute_at!(C, FT(2)) - @test C.status.time == FT(2) - @test all(interior(C) .== 2) - end - end - end - end -end diff --git a/test/test_cubed_sphere_halo_exchange.jl b/test/test_cubed_sphere_halo_exchange.jl index 29a056e0a7..810c31a58e 100644 --- a/test/test_cubed_sphere_halo_exchange.jl +++ b/test/test_cubed_sphere_halo_exchange.jl @@ -21,7 +21,7 @@ for arch in archs @info " Testing cubed sphere tracer halo exchange [$(typeof(arch))]..." grid = ConformalCubedSphereGrid(cs32_filepath, arch, Nz=1, z=(-1, 0)) - field = CenterField(arch, grid) + field = CenterField(grid) ## We will fill each grid point with a 5-digit integer "fiijj" where ## the f digit is the face number, the ii digits are the i index, and diff --git a/test/test_cubed_spheres.jl b/test/test_cubed_spheres.jl index 96ec5774e0..50613213ef 100644 --- a/test/test_cubed_spheres.jl +++ b/test/test_cubed_spheres.jl @@ -79,7 +79,7 @@ include("data_dependencies.jl") @info "Testing KernelComputedField on a ConformalCubedSphereGrid [$(typeof(arch))]..." ζ = VerticalVorticityField(model) - @test ζ isa KernelComputedField + @test ζ isa Field set!(model, u = (x, y, z) -> rand()) From 45b28d691dd6ee9a6368ada828f2923d219a9cff Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 10:14:24 -0700 Subject: [PATCH 048/140] Accept kwargs in distributed halo filling functions --- src/Distributed/halo_communication.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/Distributed/halo_communication.jl b/src/Distributed/halo_communication.jl index 73c915b710..fefc4d0b10 100644 --- a/src/Distributed/halo_communication.jl +++ b/src/Distributed/halo_communication.jl @@ -84,8 +84,8 @@ for (side, opposite_side) in zip([:west, :south, :bottom], [:east, :north, :top] fill_opposite_side_halo! = Symbol("fill_$(opposite_side)_halo!") @eval begin - function $fill_both_halos!(c, bc_side, bc_opposite_side, arch, barrier, grid, args...) - event_side = $fill_side_halo!(c, bc_side, child_architecture(arch), barrier, grid, args...) + function $fill_both_halos!(c, bc_side, bc_opposite_side, arch, barrier, grid, args...; kwargs...) + event_side = $fill_side_halo!(c, bc_side, child_architecture(arch), barrier, grid, args...; kwargs...) event_opposite_side = $fill_opposite_side_halo!(c, bc_opposite_side, child_architecture(arch), barrier, grid, args...) return event_side, event_opposite_side end @@ -107,7 +107,7 @@ for (side, opposite_side) in zip([:west, :south, :bottom], [:east, :north, :top] @eval begin function $fill_both_halos!(c, bc_side::HaloCommunicationBC, bc_opposite_side::HaloCommunicationBC, - arch, barrier, grid, c_location, args...) + arch, barrier, grid, c_location, args...; kwargs...) @assert bc_side.condition.from == bc_opposite_side.condition.from # Extra protection in case of bugs local_rank = bc_side.condition.from From ee685734fe8e27e5eedea631383d7fbbde91ddf4 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 13:16:50 -0700 Subject: [PATCH 049/140] Typing issues plus better show for computed Fields --- src/AbstractOperations/computed_field.jl | 4 ++-- src/CubedSpheres/cubed_sphere_faces.jl | 23 ++++++++++++----------- src/CubedSpheres/cubed_sphere_set!.jl | 4 ++-- test/test_cubed_spheres.jl | 23 ++++++++--------------- 4 files changed, 24 insertions(+), 30 deletions(-) diff --git a/src/AbstractOperations/computed_field.jl b/src/AbstractOperations/computed_field.jl index 2833332c1b..15bb68fb5d 100644 --- a/src/AbstractOperations/computed_field.jl +++ b/src/AbstractOperations/computed_field.jl @@ -3,7 +3,7 @@ ##### using KernelAbstractions: @kernel, @index -using Oceananigans.Fields: FieldStatus +using Oceananigans.Fields: FieldStatus, show_status using Oceananigans.Utils: work_layout import Oceananigans: short_show @@ -76,7 +76,7 @@ end @inbounds data[i, j, k] = operand[i, j, k] end -short_show(field::ComputedField) = string("ComputedField located at ", show_location(field), " of ", short_show(field.operand)) +short_show(field::ComputedField) = string("Field located at ", show_location(field), " computed from ", short_show(field.operand)) Base.show(io::IO, field::ComputedField) = print(io, "$(short_show(field))\n", diff --git a/src/CubedSpheres/cubed_sphere_faces.jl b/src/CubedSpheres/cubed_sphere_faces.jl index d0b73b20cc..1d134389ec 100644 --- a/src/CubedSpheres/cubed_sphere_faces.jl +++ b/src/CubedSpheres/cubed_sphere_faces.jl @@ -1,5 +1,6 @@ using Statistics using Oceananigans.Architectures +using Oceananigans.AbstractOperations: AbstractOperation using OffsetArrays: OffsetArray using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid @@ -43,11 +44,15 @@ const CubedSphereFaceField = Union{NonImmersedCubedSphereFaceField{X, Y, Z, A}, # CubedSphereField # Flavors of CubedSphereField -const CubedSphereField = Field{X, Y, Z, O, A, G, T, <:CubedSphereData} where {X, Y, Z, O, A, G, T} -const CubedSphereAbstractField = AbstractField{X, Y, Z, A, <:ConformalCubedSphereGrid} where {X, Y, Z, A} +const CubedSphereField{LX, LY, LZ, A} = + Union{Field{LX, LY, LZ, <:Nothing, A, <:ConformalCubedSphereGrid}, + Field{LX, LY, LZ, <:AbstractOperation, A, <:ConformalCubedSphereGrid}} -const AbstractCubedSphereField{X, Y, Z, A} = Union{CubedSphereAbstractField{X, Y, Z, A}, - CubedSphereField{X, Y, Z, A}} where {X, Y, Z, A} +const CubedSphereAbstractField{LX, LY, LZ, A} = AbstractField{LX, LY, LZ, A, <:ConformalCubedSphereGrid} + +const AbstractCubedSphereField{LX, LY, LZ, A} = + Union{CubedSphereAbstractField{LX, LY, LZ, A}, + CubedSphereField{LX, LY, LZ, A}} ##### ##### new data @@ -80,7 +85,7 @@ end ##### Utils ##### -function Base.show(io::IO, field::AbstractCubedSphereField) +function Base.show(io::IO, field::Union{CubedSphereField, AbstractCubedSphereField}) LX, LY, LZ = location(field) arch = architecture(field) A = typeof(arch) @@ -94,6 +99,7 @@ end return CubedSphereFaces(faces) end +Base.size(field::CubedSphereField) = size(field.data) Base.size(data::CubedSphereData) = (size(data.faces[1])..., length(data.faces)) @inline get_face(field::CubedSphereField, face_index) = @@ -101,18 +107,13 @@ Base.size(data::CubedSphereData) = (size(data.faces[1])..., length(data.faces)) data = get_face(field.data, face_index), boundary_conditions = get_face(field.boundary_conditions, face_index)) -@inline get_face(field::CubedSphereAbstractField, face_index) = - Field(location(field), get_face(field.grid, face_index); - data = get_face(field.data, face_index), - boundary_conditions = get_face(field.boundary_conditions, face_index)) - faces(field::AbstractCubedSphereField) = Tuple(get_face(field, face_index) for face_index in 1:length(field.data.faces)) minimum(field::AbstractCubedSphereField; dims=:) = minimum(minimum(face_field; dims) for face_field in faces(field)) maximum(field::AbstractCubedSphereField; dims=:) = maximum(maximum(face_field; dims) for face_field in faces(field)) mean(field::AbstractCubedSphereField; dims=:) = mean(mean(face_field; dims) for face_field in faces(field)) -minimum(f::Function, field::AbstractCubedSphereField; dims=:) = minimum(minimum(f, face_field; dims) for face_field in faces(field)) +inimum(f::Function, field::AbstractCubedSphereField; dims=:) = minimum(minimum(f, face_field; dims) for face_field in faces(field)) maximum(f::Function, field::AbstractCubedSphereField; dims=:) = maximum(maximum(f, face_field; dims) for face_field in faces(field)) mean(f::Function, field::AbstractCubedSphereField; dims=:) = mean(mean(f, face_field; dims) for face_field in faces(field)) diff --git a/src/CubedSpheres/cubed_sphere_set!.jl b/src/CubedSpheres/cubed_sphere_set!.jl index c1135dde3b..8dd94f7206 100644 --- a/src/CubedSpheres/cubed_sphere_set!.jl +++ b/src/CubedSpheres/cubed_sphere_set!.jl @@ -3,8 +3,8 @@ using Oceananigans.Fields: AbstractField import Oceananigans.Fields: set! -const CubedSphereCPUField = CubedSphereField{X, Y, Z, <:CPU} where {X, Y, Z} -const CubedSphereGPUField = CubedSphereField{X, Y, Z, <:GPU} where {X, Y, Z} +const CubedSphereCPUField = CubedSphereField{LX, LY, LZ, <:CPU} where {LX, LY, LZ} +const CubedSphereGPUField = CubedSphereField{LX, LY, LZ, <:GPU} where {LX, LY, LZ} # We need to define the function once for CPU fields then again for GPU fields to avoid the method # ambiguity with Fields.set!. diff --git a/test/test_cubed_spheres.jl b/test/test_cubed_spheres.jl index 50613213ef..d1de5641e8 100644 --- a/test/test_cubed_spheres.jl +++ b/test/test_cubed_spheres.jl @@ -1,15 +1,10 @@ -using Test +include("dependencies_for_runtests.jl") using Statistics: mean -using CUDA - -using Oceananigans using Oceananigans.CubedSpheres using Oceananigans.Models.HydrostaticFreeSurfaceModels using Oceananigans.Models.HydrostaticFreeSurfaceModels: VerticalVorticityField -include("data_dependencies.jl") - @testset "Cubed spheres" begin @testset "Conformal cubed sphere grid" begin @@ -29,15 +24,13 @@ include("data_dependencies.jl") @info "Constructing a HydrostaticFreeSurfaceModel on a ConformalCubedSphereGrid [$(typeof(arch))]..." - model = HydrostaticFreeSurfaceModel( - grid = grid, - momentum_advection = VectorInvariant(), - free_surface = ExplicitFreeSurface(gravitational_acceleration=0.1), - coriolis = nothing, - closure = nothing, - tracers = :c, - buoyancy = nothing - ) + free_surface = ExplicitFreeSurface(gravitational_acceleration=0.1) + model = HydrostaticFreeSurfaceModel(; grid, free_surface, + momentum_advection = VectorInvariant(), + coriolis = nothing, + closure = nothing, + tracers = :c, + buoyancy = nothing) @testset "Constructing a grid from file [$(typeof(arch))]" begin @test grid isa ConformalCubedSphereGrid From 62ea28b9dfa7fd2e41b64f52bd9670909029aeab Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 13:48:53 -0700 Subject: [PATCH 050/140] Adds get_face for KernelFunctionOperation --- src/AbstractOperations/computed_field.jl | 15 ++++++++------- src/CubedSpheres/cubed_sphere_faces.jl | 2 +- src/CubedSpheres/cubed_sphere_kernel_launching.jl | 12 ++++++++++++ 3 files changed, 21 insertions(+), 8 deletions(-) diff --git a/src/AbstractOperations/computed_field.jl b/src/AbstractOperations/computed_field.jl index 15bb68fb5d..a63ec61a69 100644 --- a/src/AbstractOperations/computed_field.jl +++ b/src/AbstractOperations/computed_field.jl @@ -3,8 +3,8 @@ ##### using KernelAbstractions: @kernel, @index -using Oceananigans.Fields: FieldStatus, show_status -using Oceananigans.Utils: work_layout +using Oceananigans.Fields: FieldStatus, show_status, reduced_dimensions +using Oceananigans.Utils: launch! import Oceananigans: short_show import Oceananigans.Fields: Field, compute! @@ -57,12 +57,13 @@ function compute!(comp::ComputedField, time=nothing) # First compute `dependencies`: compute_at!(comp.operand, time) - workgroup, worksize = - work_layout(comp.grid, :xyz, include_right_boundaries = true, location = location(comp)) - arch = architecture(comp) - compute_kernel! = _compute!(device(arch), workgroup, worksize) - event = compute_kernel!(comp.data, comp.operand; dependencies = device_event(arch)) + + event = launch!(arch, comp.grid, :xyz, _compute!, comp.data, comp.operand; + include_right_boundaries = true, + location = location(comp), + reduced_dimensions = reduced_dimensions(comp)) + wait(device(arch), event) fill_halo_regions!(comp, arch) diff --git a/src/CubedSpheres/cubed_sphere_faces.jl b/src/CubedSpheres/cubed_sphere_faces.jl index 1d134389ec..1938861b1a 100644 --- a/src/CubedSpheres/cubed_sphere_faces.jl +++ b/src/CubedSpheres/cubed_sphere_faces.jl @@ -113,7 +113,7 @@ minimum(field::AbstractCubedSphereField; dims=:) = minimum(minimum(face_field; d maximum(field::AbstractCubedSphereField; dims=:) = maximum(maximum(face_field; dims) for face_field in faces(field)) mean(field::AbstractCubedSphereField; dims=:) = mean(mean(face_field; dims) for face_field in faces(field)) -inimum(f::Function, field::AbstractCubedSphereField; dims=:) = minimum(minimum(f, face_field; dims) for face_field in faces(field)) +minimum(f::Function, field::AbstractCubedSphereField; dims=:) = minimum(minimum(f, face_field; dims) for face_field in faces(field)) maximum(f::Function, field::AbstractCubedSphereField; dims=:) = maximum(maximum(f, face_field; dims) for face_field in faces(field)) mean(f::Function, field::AbstractCubedSphereField; dims=:) = mean(mean(f, face_field; dims) for face_field in faces(field)) diff --git a/src/CubedSpheres/cubed_sphere_kernel_launching.jl b/src/CubedSpheres/cubed_sphere_kernel_launching.jl index 1fd81a70f5..eac45d0e26 100644 --- a/src/CubedSpheres/cubed_sphere_kernel_launching.jl +++ b/src/CubedSpheres/cubed_sphere_kernel_launching.jl @@ -1,5 +1,6 @@ using KernelAbstractions: Event, MultiEvent +using Oceananigans.AbstractOperations: KernelFunctionOperation using Oceananigans.Architectures: device using Oceananigans.Models.HydrostaticFreeSurfaceModels: ExplicitFreeSurface, PrescribedVelocityFields @@ -21,6 +22,17 @@ get_face(velocities::PrescribedVelocityFields, face_index) = get_face(velocities.w, face_index), velocities.parameters) +function get_face(op::KernelFunctionOperation, face_index) + LX, LY, LZ = location(op) + computed_dependencies = get_face(op.computed_dependencies, face_index) + parameters = get_face(op.parameters, face_index) + face_grid = get_face(op.grid, face_index) + return KernelFunctionOperation{LX, LY, LZ}(op.kernel_function, + computed_dependencies, + parameters, + face_grid) +end + function launch!(arch, grid::ConformalCubedSphereGrid, dims, kernel!, args...; kwargs...) events = [] From e229e58dc87099a6ac6a441001a86ce4a7680664 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 13:59:05 -0700 Subject: [PATCH 051/140] Bugfix in halo exchange test --- test/test_cubed_sphere_halo_exchange.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/test_cubed_sphere_halo_exchange.jl b/test/test_cubed_sphere_halo_exchange.jl index 810c31a58e..e4eea8ec0b 100644 --- a/test/test_cubed_sphere_halo_exchange.jl +++ b/test/test_cubed_sphere_halo_exchange.jl @@ -373,8 +373,8 @@ for arch in archs grid = ConformalCubedSphereGrid(cs32_filepath, arch, Nz=1, z=(-1, 0)) - u_field = XFaceField(arch, grid) - v_field = YFaceField(arch, grid) + u_field = XFaceField(grid) + v_field = YFaceField(grid) ## We will fill each grid point with a 6-digit integer "ufiijj" where ## the u digit is 1 for u and 2 for v, the f digit is the face number, From ff7a55157dc8dadec494bc60cae3c1e21c5d82b7 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 14:56:23 -0700 Subject: [PATCH 052/140] Fiddling with dispatch for distributed halo filling --- src/Distributed/distributed_fields.jl | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/Distributed/distributed_fields.jl b/src/Distributed/distributed_fields.jl index 5d54923f78..17fda0d5f5 100644 --- a/src/Distributed/distributed_fields.jl +++ b/src/Distributed/distributed_fields.jl @@ -1,4 +1,6 @@ -import Oceananigans.Fields: Field, fill_halo_regions! +import Oceananigans.Fields: Field +import Oceananigans.BoundaryConditions: fill_halo_regions! + function Field((LX, LY, LZ)::Tuple, grid::DistributedGrid, data, old_bcs, op, status) arch = architecture(grid) @@ -6,7 +8,7 @@ function Field((LX, LY, LZ)::Tuple, grid::DistributedGrid, data, old_bcs, op, st return Field{LX, LY, LZ}(grid, data, new_bcs, op, status) end -const DistributedField = AbstractField{<:Any, <:Any, <:Any, <:MultiArch} +const DistributedField = Field{<:Any, <:Any, <:Any, <:Any, <:MultiArch} fill_halo_regions!(field::DistributedField, arch, args...; kwargs...) = fill_halo_regions!(field.data, From 65925f9dcf83cbc5f34e71294cabc339ada4134e Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 14:57:59 -0700 Subject: [PATCH 053/140] loosen type restriction on RectilinearGrid architecture --- src/Grids/rectilinear_grid.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl index b3b0f633f6..6fc1a401aa 100644 --- a/src/Grids/rectilinear_grid.jl +++ b/src/Grids/rectilinear_grid.jl @@ -40,7 +40,7 @@ struct RectilinearGrid{FT, TX, TY, TZ, FX, FY, FZ, VX, VY, VZ, Arch} <: Abstract Δyᵃᶠᵃ :: FY, Δyᵃᶜᵃ :: FY, yᵃᶠᵃ :: VY, yᵃᶜᵃ :: VY, Δzᵃᵃᶠ :: FZ, Δzᵃᵃᶜ :: FZ, - zᵃᵃᶠ :: VZ, zᵃᵃᶜ :: VZ) where {Arch <: AbstractArchitecture, + zᵃᵃᶠ :: VZ, zᵃᵃᶜ :: VZ) where {Arch <: Union{Nothing, AbstractArchitecture}, TX, TY, TZ, FT, FX, VX, FY, VY, FZ, VZ} From e72fc8a0aad1b255d23062f26827f190a1c9de32 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 15:14:07 -0700 Subject: [PATCH 054/140] Ok completely release type restriction in grid constructors --- src/Grids/latitude_longitude_grid.jl | 7 +++++-- src/Grids/rectilinear_grid.jl | 8 ++++---- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/src/Grids/latitude_longitude_grid.jl b/src/Grids/latitude_longitude_grid.jl index b39cd6baba..7c56847808 100644 --- a/src/Grids/latitude_longitude_grid.jl +++ b/src/Grids/latitude_longitude_grid.jl @@ -55,8 +55,11 @@ struct LatitudeLongitudeGrid{FT, TX, TY, TZ, M, MY, FX, FY, FZ, VX, VY, VZ, Arch Δxᶠᶠ::M, Δxᶜᶜ::M, Δyᶠᶜ::MY, Δyᶜᶠ::MY, Azᶠᶜ::M, Azᶜᶠ::M, Azᶠᶠ::M, Azᶜᶜ::M, - radius::FT) where {Arch <: AbstractArchitecture, - FT, TX, TY, TZ, FX, FY, FZ, VX, VY, VZ, M, MY} + radius::FT) where {Arch, FT, + TX, TY, TZ, + FX, FY, FZ, + VX, VY, VZ, + M, MY} return new{FT, TX, TY, TZ, M, MY, FX, FY, FZ, VX, VY, VZ, Arch}(architecture, Nλ, Nφ, Nz, diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl index 6fc1a401aa..7983daeb8d 100644 --- a/src/Grids/rectilinear_grid.jl +++ b/src/Grids/rectilinear_grid.jl @@ -40,10 +40,10 @@ struct RectilinearGrid{FT, TX, TY, TZ, FX, FY, FZ, VX, VY, VZ, Arch} <: Abstract Δyᵃᶠᵃ :: FY, Δyᵃᶜᵃ :: FY, yᵃᶠᵃ :: VY, yᵃᶜᵃ :: VY, Δzᵃᵃᶠ :: FZ, Δzᵃᵃᶜ :: FZ, - zᵃᵃᶠ :: VZ, zᵃᵃᶜ :: VZ) where {Arch <: Union{Nothing, AbstractArchitecture}, - TX, TY, TZ, FT, - FX, VX, FY, VY, - FZ, VZ} + zᵃᵃᶠ :: VZ, zᵃᵃᶜ :: VZ) where {Arch, FT, + TX, TY, TZ, + FX, VX, FY, + VY, FZ, VZ} return new{FT, TX, TY, TZ, FX, FY, FZ, VX, VY, VZ, Arch}(arch, Nx, Ny, Nz, Hx, Hy, Hz, Lx, Ly, Lz, From 365b87e79251219b7d978f9f8afc61c55000e00a Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 15:25:13 -0700 Subject: [PATCH 055/140] Bugfix in distributed fill_halo_regions! --- src/Distributed/distributed_fields.jl | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/Distributed/distributed_fields.jl b/src/Distributed/distributed_fields.jl index 17fda0d5f5..4b732ede6c 100644 --- a/src/Distributed/distributed_fields.jl +++ b/src/Distributed/distributed_fields.jl @@ -1,7 +1,6 @@ -import Oceananigans.Fields: Field +import Oceananigans.Fields: Field, location import Oceananigans.BoundaryConditions: fill_halo_regions! - function Field((LX, LY, LZ)::Tuple, grid::DistributedGrid, data, old_bcs, op, status) arch = architecture(grid) new_bcs = inject_halo_communication_boundary_conditions(old_bcs, arch.local_rank, arch.connectivity) @@ -13,7 +12,7 @@ const DistributedField = Field{<:Any, <:Any, <:Any, <:Any, <:MultiArch} fill_halo_regions!(field::DistributedField, arch, args...; kwargs...) = fill_halo_regions!(field.data, field.boundary_conditions, - arch, + architecture(field), field.grid, location(field), args...; kwargs...) From 50f7a69e66698332e30d83faaabd25b5e62a7e7a Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 15:43:45 -0700 Subject: [PATCH 056/140] Fixes bug in reconstruct_global_grid --- src/Distributed/distributed_grids.jl | 13 ++++++------- .../NonhydrostaticModels/NonhydrostaticModels.jl | 6 +++--- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/src/Distributed/distributed_grids.jl b/src/Distributed/distributed_grids.jl index 97947baa28..379657f340 100644 --- a/src/Distributed/distributed_grids.jl +++ b/src/Distributed/distributed_grids.jl @@ -133,7 +133,7 @@ end function reconstruct_global_grid(grid::RectilinearGrid) - arch = grid.architecture + arch = grid.architecture i, j, k = arch.local_index Rx, Ry, Rz = R = arch.ranks @@ -152,22 +152,21 @@ function reconstruct_global_grid(grid::RectilinearGrid) yG = get_global_coords(y, ny, Ry, j, arch) zG = get_global_coords(z, nz, Rz, k, arch) - architecture = child_architecture(arch) + child_arch = child_architecture(arch) FT = eltype(grid) - Lx, xᶠᵃᵃ, xᶜᵃᵃ, Δxᶠᵃᵃ, Δxᶜᵃᵃ = generate_coordinate(FT, TX, Nx, Hx, xG, architecture) - Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, TY, Ny, Hy, yG, architecture) - Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ, Nz, Hz, zG, architecture) + 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}(architecture, + return RectilinearGrid{TX, TY, TZ}(arch, Nx, Ny, Nz, Hx, Hy, Hz, Lx, Ly, Lz, Δxᶠᵃᵃ, Δxᶜᵃᵃ, xᶠᵃᵃ, xᶜᵃᵃ, Δyᵃᶜᵃ, Δyᵃᶠᵃ, yᵃᶠᵃ, yᵃᶜᵃ, Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ) - end function with_halo(new_halo, grid::DistributedGrid) diff --git a/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl b/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl index f5d7cdfe21..30332273da 100644 --- a/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl +++ b/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl @@ -15,10 +15,10 @@ using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid import Oceananigans: fields, prognostic_fields -function PressureSolver(arch::MultiArch, grid::RegRectilinearGrid) - global_grid = reconstruct_global_grid(grid) +function PressureSolver(arch::MultiArch, local_grid::RegRectilinearGrid) + global_grid = reconstruct_global_grid(local_grid) if arch.ranks[1] == 1 # we would have to allow different settings - return DistributedFFTBasedPoissonSolver(global_grid, grid) + return DistributedFFTBasedPoissonSolver(global_grid, local_grid) else @warn "A Distributed NonhydrostaticModel is allowed only when the x-direction is not parallelized" return nothing From 30a79d0d2ff7f4299c92c21331c3d76725d70944 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 15:45:19 -0700 Subject: [PATCH 057/140] Fix doctest for grid_metrics --- src/AbstractOperations/grid_metrics.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/AbstractOperations/grid_metrics.jl b/src/AbstractOperations/grid_metrics.jl index 46642da906..589031f1e3 100644 --- a/src/AbstractOperations/grid_metrics.jl +++ b/src/AbstractOperations/grid_metrics.jl @@ -66,7 +66,7 @@ julia> using Oceananigans julia> using Oceananigans.AbstractOperations: Δz -julia> grid = RectilinearGrid(size=(1, 1, 1), extent=(1, 2, 3)); c = CenterField(CPU(), grid); +julia> grid = RectilinearGrid(size=(1, 1, 1), extent=(1, 2, 3)); c = CenterField(grid); julia> c_dz = c * Δz # returns BinaryOperation between Field and GridMetricOperation BinaryOperation at (Center, Center, Center) @@ -105,7 +105,7 @@ julia> using Oceananigans julia> using Oceananigans.AbstractOperations: volume -julia> grid = RectilinearGrid(size=(2, 2, 2), extent=(1, 2, 3)); c = CenterField(CPU(), grid); +julia> grid = RectilinearGrid(size=(2, 2, 2), extent=(1, 2, 3)); c = CenterField(grid); julia> c .= 1; From 781e1f25b9d42cd92959b6c0c9e2f9d87a10fa6d Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 15:46:53 -0700 Subject: [PATCH 058/140] Fix doctest for binary operation --- src/AbstractOperations/binary_operations.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/AbstractOperations/binary_operations.jl b/src/AbstractOperations/binary_operations.jl index 7eb4af7d9b..d96071ef0d 100644 --- a/src/AbstractOperations/binary_operations.jl +++ b/src/AbstractOperations/binary_operations.jl @@ -157,7 +157,7 @@ Set{Any} with 6 elements: :* :plus_or_times -julia> c, d = (Field(Center, Center, Center, CPU(), RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1))) for i = 1:2); +julia> c, d = (CenterField(RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1))) for i = 1:2); julia> plus_or_times(c, d) BinaryOperation at (Center, Center, Center) From 79d7dd7d75ff794bbfa1987dde3d8887e878f884 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 15:47:32 -0700 Subject: [PATCH 059/140] Fix doctest in UnaryOperation --- src/AbstractOperations/unary_operations.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/AbstractOperations/unary_operations.jl b/src/AbstractOperations/unary_operations.jl index 1b8329fbf5..3360f99a28 100644 --- a/src/AbstractOperations/unary_operations.jl +++ b/src/AbstractOperations/unary_operations.jl @@ -66,7 +66,7 @@ Set{Any} with 8 elements: :tanh :sin -julia> c = Field(Center, Center, Center, CPU(), RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1))); +julia> c = CenterField(RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1))); julia> square_it(c) UnaryOperation at (Center, Center, Center) From 3ab3e550b71497ba6ac780605762778f4e2eccd3 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 15:48:11 -0700 Subject: [PATCH 060/140] Fix doctest in multiary operation --- src/AbstractOperations/multiary_operations.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/AbstractOperations/multiary_operations.jl b/src/AbstractOperations/multiary_operations.jl index 503db5b8a5..d07995d87e 100644 --- a/src/AbstractOperations/multiary_operations.jl +++ b/src/AbstractOperations/multiary_operations.jl @@ -76,7 +76,7 @@ julia> using Oceananigans, Oceananigans.AbstractOperations julia> harmonic_plus(a, b, c) = 1/3 * (1/a + 1/b + 1/c) harmonic_plus (generic function with 1 method) -julia> c, d, e = Tuple(Field(Center, Center, Center, CPU(), RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1))) for i = 1:3); +julia> c, d, e = Tuple(CenterField(RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1))) for i = 1:3); julia> harmonic_plus(c, d, e) # before magic @multiary transformation BinaryOperation at (Center, Center, Center) From 99314c7e52eaee9912f3c3c4ef5e602a8e3c2d3a Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 16:02:42 -0700 Subject: [PATCH 061/140] my_grid -> local_grid --- src/Models/NonhydrostaticModels/solve_for_pressure.jl | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/Models/NonhydrostaticModels/solve_for_pressure.jl b/src/Models/NonhydrostaticModels/solve_for_pressure.jl index 4d0c321400..622fa9a72d 100644 --- a/src/Models/NonhydrostaticModels/solve_for_pressure.jl +++ b/src/Models/NonhydrostaticModels/solve_for_pressure.jl @@ -21,11 +21,10 @@ end ##### Solve for pressure ##### - function solve_for_pressure!(pressure, solver::DistributedFFTBasedPoissonSolver, Δt, U★) rhs = first(solver.storage) arch = architecture(solver) - grid = solver.my_grid + grid = solver.local_grid rhs_event = launch!(arch, grid, :xyz, calculate_pressure_source_term_fft_based_solver!, rhs, grid, Δt, U★, dependencies = device_event(arch)) From 4946540cd84c2f8939f0a0b681c6d7c544009327 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 16:15:33 -0700 Subject: [PATCH 062/140] Only replace shallow water model grid if necessary --- src/Models/ShallowWaterModels/shallow_water_model.jl | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/Models/ShallowWaterModels/shallow_water_model.jl b/src/Models/ShallowWaterModels/shallow_water_model.jl index 3996a33105..e11d9e0b8e 100644 --- a/src/Models/ShallowWaterModels/shallow_water_model.jl +++ b/src/Models/ShallowWaterModels/shallow_water_model.jl @@ -106,15 +106,18 @@ function ShallowWaterModel(; tracers = tupleit(tracers) # supports tracers=:c keyword argument (for example) - @assert topology(grid, 3) === Flat "ShallowWaterModel requires `topology(grid, 3) === Flat`. Use `topology = ($(topology(grid, 1)), $(topology(grid, 2)), Flat)` when constructing `grid`." + @assert topology(grid, 3) === Flat "ShallowWaterModel requires `topology(grid, 3) === Flat`. " * + "Use `topology = ($(topology(grid, 1)), $(topology(grid, 2)), Flat)` " * + "when constructing `grid`." Hx, Hy, Hz = inflate_halo_size(grid.Hx, grid.Hy, 0, topology(grid), advection, closure) - grid = with_halo((Hx, Hy, 0), grid) + grid.Hx, grid.Hy, grid.Hz != (Hx, Hy, 0) && (grid = with_halo((Hx, Hy, 0), grid)) prognostic_field_names = (:uh, :vh, :h, tracers...) - default_boundary_conditions = NamedTuple{prognostic_field_names}(Tuple(FieldBoundaryConditions() for name in prognostic_field_names)) - boundary_conditions = merge(default_boundary_conditions, boundary_conditions) + default_boundary_conditions = NamedTuple{prognostic_field_names}(Tuple(FieldBoundaryConditions() + for name in prognostic_field_names)) + boundary_conditions = merge(default_boundary_conditions, boundary_conditions) boundary_conditions = regularize_field_boundary_conditions(boundary_conditions, grid, prognostic_field_names) solution = ShallowWaterSolutionFields(grid, boundary_conditions) From cce7dde8c93dad0e69edd042765cdff2f592dfb1 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 17:27:55 -0700 Subject: [PATCH 063/140] Update docstring for rectilinear grid --- src/Grids/rectilinear_grid.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl index 7983daeb8d..0dc9c6f4c3 100644 --- a/src/Grids/rectilinear_grid.jl +++ b/src/Grids/rectilinear_grid.jl @@ -268,7 +268,7 @@ RectilinearGrid{Float64, Periodic, Bounded, Bounded} size (Nx, Ny, Nz): (32, 30, 24) halo (Hx, Hy, Hz): (1, 1, 1) spacing in x: Regular, with spacing 6.25 - spacing in y: Stretched, with spacing min=0.2739052315863333, max=5.226423163382681 + spacing in y: Stretched, with spacing min=0.2739052315863262, max=5.22642316338267 spacing in z: Stretched, with spacing min=0.6826950100338962, max=1.8309085743885056 ``` """ From b300ebbf0b65bac2ed5a87562f847d05a35d63f6 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 23:44:18 -0500 Subject: [PATCH 064/140] Use on_architecture rather than adapt_structure in similar --- src/Fields/field.jl | 10 ++++++---- src/Fields/set!.jl | 2 +- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/Fields/field.jl b/src/Fields/field.jl index 07cdb949c8..e7b6cbaa49 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -14,9 +14,10 @@ struct Field{LX, LY, LZ, O, A, G, T, D, B, S} <: AbstractField{LX, LY, LZ, A, G, status :: S # Inner constructor that does not validate _anything_! - function Field{LX, LY, LZ}(grid::G, data::D, bcs::B, op::O, status::S) where {LX, LY, LZ, G, D, B, O, S} - T = eltype(grid) - A = typeof(architecture(grid)) + function Field{LX, LY, LZ}(grid::G, data::D, bcs::B, op::O, status::S, + arch::A=architecture(grid)) where {LX, LY, LZ, + G, D, B, O, S, A} + T = eltype(data) return new{LX, LY, LZ, O, A, G, T, D, B, S}(grid, data, bcs, op, status) end end @@ -292,7 +293,8 @@ function Adapt.adapt_structure(to, reduced_field::ReducedField) LX, LY, LZ = location(reduced_field) return Field{LX, LY, LZ}(nothing, adapt(to, reduced_field.data), - nothing, nothing, nothing) + nothing, nothing, nothing, + adapt(to, architecture(reduced_field))) end ##### diff --git a/src/Fields/set!.jl b/src/Fields/set!.jl index ad22c8122c..70434489ed 100644 --- a/src/Fields/set!.jl +++ b/src/Fields/set!.jl @@ -34,7 +34,7 @@ const AbstractGPUField = AbstractField{X, Y, Z, <:GPU} where {X, Y, Z} """ Set the GPU field `u` to the array or function `v`. """ function set!(u::AbstractGPUField, v::Union{Array, Function}) - cpu_grid = adapt_structure(CPU(), u.grid) + cpu_grid = on_architecture(CPU(), u.grid) v_field = similar(u, cpu_grid) set!(v_field, v) set!(u, v_field) From ab684c57a27a07cb34ab239b0d3a0eb60d324a0c Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 11 Jan 2022 21:46:35 -0700 Subject: [PATCH 065/140] Minor fixes --- src/Distributed/distributed_grids.jl | 10 +++++----- src/Fields/set!.jl | 3 ++- test/test_distributed_poisson_solvers.jl | 8 ++++---- 3 files changed, 11 insertions(+), 10 deletions(-) diff --git a/src/Distributed/distributed_grids.jl b/src/Distributed/distributed_grids.jl index 379657f340..9ca85aeadd 100644 --- a/src/Distributed/distributed_grids.jl +++ b/src/Distributed/distributed_grids.jl @@ -48,11 +48,11 @@ function RectilinearGrid(arch::MultiArch, FT = Float64; TX, TY, TZ, size, halo, x, y, z = validate_rectilinear_grid_args(topology, size, halo, FT, extent, x, y, z) - Nx, Ny, Nz = size + @show Nx, Ny, Nz = size Hx, Hy, Hz = halo ri, rj, rk = arch.local_index - Rx, Ry, Rz = arch.ranks + @show Rx, Ry, Rz = arch.ranks # Make sure we can put an integer number of grid points in each rank. # Will generalize in the future. @@ -95,7 +95,7 @@ function LatitudeLongitudeGrid(arch::MultiArch, Nλ, Nφ, Nz, Hλ, Hφ, Hz, latitude, longitude, topo = validate_lat_lon_grid_args(latitude, longitude, size, halo) - i, j, k = arch.local_index + i, j, k = arch.local_index Rx, Ry, Rz = arch.ranks # Make sure we can put an integer number of grid points in each rank. @@ -160,12 +160,12 @@ function reconstruct_global_grid(grid::RectilinearGrid) 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}(arch, + return RectilinearGrid{TX, TY, TZ}(child_arch, Nx, Ny, Nz, Hx, Hy, Hz, Lx, Ly, Lz, Δxᶠᵃᵃ, Δxᶜᵃᵃ, xᶠᵃᵃ, xᶜᵃᵃ, - Δyᵃᶜᵃ, Δyᵃᶠᵃ, yᵃᶠᵃ, yᵃᶜᵃ, + Δyᵃᶠᵃ, Δyᵃᶜᵃ, yᵃᶠᵃ, yᵃᶜᵃ, Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ) end diff --git a/src/Fields/set!.jl b/src/Fields/set!.jl index ad22c8122c..6ca4292dbb 100644 --- a/src/Fields/set!.jl +++ b/src/Fields/set!.jl @@ -3,6 +3,7 @@ using CUDAKernels using KernelAbstractions: @kernel, @index using Adapt: adapt_structure +using Oceananigans.Grids: on_architecture using Oceananigans.Architectures: device, GPU, CPU, AbstractMultiArchitecture using Oceananigans.Utils: work_layout @@ -34,7 +35,7 @@ const AbstractGPUField = AbstractField{X, Y, Z, <:GPU} where {X, Y, Z} """ Set the GPU field `u` to the array or function `v`. """ function set!(u::AbstractGPUField, v::Union{Array, Function}) - cpu_grid = adapt_structure(CPU(), u.grid) + cpu_grid = on_architecture(CPU(), u.grid) v_field = similar(u, cpu_grid) set!(v_field, v) set!(u, v_field) diff --git a/test/test_distributed_poisson_solvers.jl b/test/test_distributed_poisson_solvers.jl index 400f12f407..607e114e18 100644 --- a/test/test_distributed_poisson_solvers.jl +++ b/test/test_distributed_poisson_solvers.jl @@ -34,10 +34,10 @@ function divergence_free_poisson_solution_triply_periodic(grid_points, ranks) arch = MultiArch(CPU(), ranks=ranks, topology = topo) local_grid = RectilinearGrid(arch, topology=topo, size=grid_points, extent=(1, 2, 3)) - full_grid = reconstruct_global_grid(local_grid) - solver = DistributedFFTBasedPoissonSolver(arch, full_grid, local_grid) + global_grid = reconstruct_global_grid(local_grid) + solver = DistributedFFTBasedPoissonSolver(global_grid, local_grid) - R = random_divergent_source_term(child_architecture(arch), local_grid) + R = random_divergent_source_term(local_grid) first(solver.storage) .= R ϕc = first(solver.storage) @@ -46,7 +46,7 @@ function divergence_free_poisson_solution_triply_periodic(grid_points, ranks) p_bcs = FieldBoundaryConditions(local_grid, (Center, Center, Center)) p_bcs = inject_halo_communication_boundary_conditions(p_bcs, arch.local_rank, arch.connectivity) - ϕ = CenterField(child_architecture(arch), local_grid, p_bcs) # "pressure" + ϕ = CenterField(child_architecture(arch), local_grid, p_bcs) # "pressure" ∇²ϕ = CenterField(child_architecture(arch), local_grid, p_bcs) interior(ϕ) .= real(first(solver.storage)) From 93cffd2295acf538ff5ee33b4eae8bb38fdda574 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 06:46:15 -0700 Subject: [PATCH 066/140] Get arch from local_grid in FFT-based PoissonSolver --- src/Distributed/distributed_fft_based_poisson_solver.jl | 6 ++++-- src/Distributed/distributed_grids.jl | 4 ++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/Distributed/distributed_fft_based_poisson_solver.jl b/src/Distributed/distributed_fft_based_poisson_solver.jl index 900dcc4b66..3048b9c1e7 100644 --- a/src/Distributed/distributed_fft_based_poisson_solver.jl +++ b/src/Distributed/distributed_fft_based_poisson_solver.jl @@ -19,7 +19,7 @@ architecture(solver::DistributedFFTBasedPoissonSolver) = function DistributedFFTBasedPoissonSolver(global_grid, local_grid) - arch = architecture(global_grid) + arch = architecture(local_grid) topo = (TX, TY, TZ) = topology(global_grid) λx = poisson_eigenvalues(global_grid.Nx, global_grid.Lx, 1, TX()) @@ -37,7 +37,9 @@ function DistributedFFTBasedPoissonSolver(global_grid, local_grid) perm_Nx = global_grid.Nx ÷ Ry - λx = λx[(I-1)*perm_Nx+1:I*perm_Nx, :, :] + i₁ = (I-1) * perm_Nx + 1 + i₂ = I * perm_Nx + λx = λx[i₁:i₂, :, :] eigenvalues = (; λx, λy, λz) diff --git a/src/Distributed/distributed_grids.jl b/src/Distributed/distributed_grids.jl index 9ca85aeadd..5a9f54a2e7 100644 --- a/src/Distributed/distributed_grids.jl +++ b/src/Distributed/distributed_grids.jl @@ -48,11 +48,11 @@ function RectilinearGrid(arch::MultiArch, FT = Float64; TX, TY, TZ, size, halo, x, y, z = validate_rectilinear_grid_args(topology, size, halo, FT, extent, x, y, z) - @show Nx, Ny, Nz = size + Nx, Ny, Nz = size Hx, Hy, Hz = halo ri, rj, rk = arch.local_index - @show Rx, Ry, Rz = arch.ranks + Rx, Ry, Rz = arch.ranks # Make sure we can put an integer number of grid points in each rank. # Will generalize in the future. From 8bae80031beab45bc4198084d4fe4471151146ae Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 06:48:38 -0700 Subject: [PATCH 067/140] Only remake grid if halos are too _small_ --- src/Models/ShallowWaterModels/shallow_water_model.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/Models/ShallowWaterModels/shallow_water_model.jl b/src/Models/ShallowWaterModels/shallow_water_model.jl index e11d9e0b8e..9203df3cf9 100644 --- a/src/Models/ShallowWaterModels/shallow_water_model.jl +++ b/src/Models/ShallowWaterModels/shallow_water_model.jl @@ -111,7 +111,8 @@ function ShallowWaterModel(; "when constructing `grid`." Hx, Hy, Hz = inflate_halo_size(grid.Hx, grid.Hy, 0, topology(grid), advection, closure) - grid.Hx, grid.Hy, grid.Hz != (Hx, Hy, 0) && (grid = with_halo((Hx, Hy, 0), grid)) + any((grid.Hx, grid.Hy, grid.Hz) .< (Hx, Hy, 0)) && # halos are too small, remake grid + (grid = with_halo((Hx, Hy, 0), grid)) prognostic_field_names = (:uh, :vh, :h, tracers...) default_boundary_conditions = NamedTuple{prognostic_field_names}(Tuple(FieldBoundaryConditions() From e12883f3f2be34e8fb77f31cf72473650b4ec8cf Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 06:57:43 -0700 Subject: [PATCH 068/140] Brings examples up to date with current Field syntax --- examples/convecting_plankton.jl | 11 +++++------ examples/eady_turbulence.jl | 10 +++++----- examples/horizontal_convection.jl | 16 ++++++++-------- examples/kelvin_helmholtz_instability.jl | 8 ++++---- examples/langmuir_turbulence.jl | 10 +++++----- examples/shallow_water_Bickley_jet.jl | 6 +++--- examples/two_dimensional_turbulence.jl | 10 +++++----- 7 files changed, 35 insertions(+), 36 deletions(-) diff --git a/examples/convecting_plankton.jl b/examples/convecting_plankton.jl index 421cf3fd17..857e75dd1e 100644 --- a/examples/convecting_plankton.jl +++ b/examples/convecting_plankton.jl @@ -29,7 +29,7 @@ # and grazing by zooplankton. # * How to set time-dependent boundary conditions. # * How to use the `TimeStepWizard` to adapt the simulation time-step. -# * How to use `AveragedField` to diagnose spatial averages of model fields. +# * How to use `Average` and `Field` to diagnose spatial averages of model fields. # # ## Install dependencies # @@ -175,7 +175,7 @@ simulation.callbacks[:progress] = Callback(progress, IterationInterval(20)) # and a basic `JLD2OutputWriter` that writes velocities and both # the two-dimensional and horizontally-averaged plankton concentration, -averaged_plankton = AveragedField(model.tracers.P, dims=(1, 2)) +averaged_plankton = Field(Average(model.tracers.P, dims=(1, 2))) outputs = (w = model.velocities.w, plankton = model.tracers.P, @@ -190,10 +190,9 @@ simulation.output_writers[:simple_output] = # !!! info "Using multiple output writers" # Because each output writer is associated with a single output `schedule`, # it often makes sense to use _different_ output writers for different types of output. -# For example, reduced fields like `AveragedField` usually consume less disk space than -# two- or three-dimensional fields, and can thus be output more frequently without -# blowing up your hard drive. An arbitrary number of output writers may be added to -# `simulation.output_writers`. +# For example, smaller outputs that consume less disk space may be written more +# frequently without threatening the capacity of your hard drive. +# An arbitrary number of output writers may be added to `simulation.output_writers`. # # The simulation is set up. Let there be plankton: diff --git a/examples/eady_turbulence.jl b/examples/eady_turbulence.jl index d263f3ea77..90f95864e1 100644 --- a/examples/eady_turbulence.jl +++ b/examples/eady_turbulence.jl @@ -6,7 +6,7 @@ # * How to use a tuple of turbulence closures # * How to use hyperdiffusivity # * How to implement background velocity and tracer distributions -# * How to use `ComputedField`s for output +# * How to build `Field`s that compute output # # ## Install dependencies # @@ -275,16 +275,16 @@ simulation.callbacks[:progress] = Callback(progress, IterationInterval(10)) # ### Output # # To visualize the baroclinic turbulence ensuing in the Eady problem, -# we use `ComputedField`s to diagnose and output vertical vorticity and divergence. -# Note that `ComputedField`s take "AbstractOperations" on `Field`s as input: +# we use computed `Field`s to diagnose and output vertical vorticity and divergence. +# Note that computed `Field`s take "AbstractOperations" on `Field`s as input: u, v, w = model.velocities # unpack velocity `Field`s ## Vertical vorticity [s⁻¹] -ζ = ComputedField(∂x(v) - ∂y(u)) +ζ = Field(∂x(v) - ∂y(u)) ## Horizontal divergence, or ∂x(u) + ∂y(v) [s⁻¹] -δ = ComputedField(-∂z(w)) +δ = Field(-∂z(w)) # With the vertical vorticity, `ζ`, and the horizontal divergence, `δ` in hand, # we create a `JLD2OutputWriter` that saves `ζ` and `δ` and add them to diff --git a/examples/horizontal_convection.jl b/examples/horizontal_convection.jl index 3fdce4b693..815db2db3f 100644 --- a/examples/horizontal_convection.jl +++ b/examples/horizontal_convection.jl @@ -4,7 +4,7 @@ # # This example demonstrates: # -# * How to use `ComputedField`s for output. +# * How to use computed `Field`s for output. # * How to post-process saved output using `FieldTimeSeries`. # # ## Install dependencies @@ -128,18 +128,18 @@ simulation.callbacks[:progress] = Callback(progress, IterationInterval(50)) # ### Output # -# We use `ComputedField`s to diagnose and output the total flow speed, the vorticity, ``\zeta``, -# and the buoyancy, ``b``. Note that `ComputedField`s take "AbstractOperations" on `Field`s as +# We use computed `Field`s to diagnose and output the total flow speed, the vorticity, ``\zeta``, +# and the buoyancy, ``b``. Note that computed `Field`s take "AbstractOperations" on `Field`s as # input: u, v, w = model.velocities # unpack velocity `Field`s b = model.tracers.b # unpack buoyancy `Field` ## total flow speed -s = ComputedField(sqrt(u^2 + w^2)) +s = Field(sqrt(u^2 + w^2)) ## y-component of vorticity -ζ = ComputedField(∂z(u) - ∂x(w)) +ζ = Field(∂z(u) - ∂x(w)) outputs = (s = s, b = b, ζ = ζ) nothing # hide @@ -209,7 +209,7 @@ anim = @animate for i in 1:length(times) ζ_snapshot = interior(ζ_timeseries[i])[:, 1, :] b = b_timeseries[i] - χ = ComputedField(κ * (∂x(b)^2 + ∂z(b)^2)) + χ = Field(κ * (∂x(b)^2 + ∂z(b)^2)) compute!(χ) b_snapshot = interior(b)[:, 1, :] @@ -313,8 +313,8 @@ nothing # hide grid = b_timeseries.grid -∫ⱽ_s² = ReducedField(Nothing, Nothing, Nothing, CPU(), grid, dims=(1, 2, 3)) -∫ⱽ_mod²_∇b = ReducedField(Nothing, Nothing, Nothing, CPU(), grid, dims=(1, 2, 3)) +∫ⱽ_s² = Field{Nothing, Nothing, Nothing}(grid) +∫ⱽ_mod²_∇b = Field{Nothing, Nothing, Nothing}(grid) # We recover the time from the saved `FieldTimeSeries` and construct two empty arrays to store # the volume-averaged kinetic energy and the instantaneous Nusselt number, diff --git a/examples/kelvin_helmholtz_instability.jl b/examples/kelvin_helmholtz_instability.jl index 3003232ecc..06bc043cf5 100644 --- a/examples/kelvin_helmholtz_instability.jl +++ b/examples/kelvin_helmholtz_instability.jl @@ -280,7 +280,7 @@ nothing # hide u, v, w = model.velocities b = model.tracers.b -perturbation_vorticity = ComputedField(∂z(u) - ∂x(w)) +perturbation_vorticity = Field(∂z(u) - ∂x(w)) xF, yF, zF = nodes(perturbation_vorticity) @@ -332,7 +332,7 @@ nothing # hide using Random, Statistics -mean_perturbation_kinetic_energy = AveragedField(1/2 * (u^2 + w^2), dims=(1, 2, 3)) +mean_perturbation_kinetic_energy = Field(Average(1/2 * (u^2 + w^2))) noise(x, y, z) = randn() @@ -377,9 +377,9 @@ rescale!(simulation.model, mean_perturbation_kinetic_energy, target_kinetic_ener # buoyancy (perturbation + basic state). It'll be also neat to plot the kinetic energy time-series # and confirm it grows with the estimated growth rate. -total_vorticity = ComputedField(∂z(u) + ∂z(model.background_fields.velocities.u) - ∂x(w)) +total_vorticity = Field(∂z(u) + ∂z(model.background_fields.velocities.u) - ∂x(w)) -total_b = ComputedField(b + model.background_fields.tracers.b) +total_b = Field(b + model.background_fields.tracers.b) simulation.output_writers[:vorticity] = JLD2OutputWriter(model, (ω=perturbation_vorticity, Ω=total_vorticity, b=b, B=total_b, KE=mean_perturbation_kinetic_energy), diff --git a/examples/langmuir_turbulence.jl b/examples/langmuir_turbulence.jl index 3636974c14..472e4578d3 100644 --- a/examples/langmuir_turbulence.jl +++ b/examples/langmuir_turbulence.jl @@ -225,12 +225,12 @@ simulation.output_writers[:fields] = u, v, w = model.velocities -U = AveragedField(u, dims=(1, 2)) -V = AveragedField(v, dims=(1, 2)) -B = AveragedField(model.tracers.b, dims=(1, 2)) +U = Field(Average(u, dims=(1, 2))) +V = Field(Average(v, dims=(1, 2))) +B = Field(Average(model.tracers.b, dims=(1, 2))) -wu = AveragedField(w * u, dims=(1, 2)) -wv = AveragedField(w * v, dims=(1, 2)) +wu = Field(Average(w * u, dims=(1, 2))) +wv = Field(Average(w * v, dims=(1, 2))) simulation.output_writers[:averages] = JLD2OutputWriter(model, (u=U, v=V, b=B, wu=wu, wv=wv), diff --git a/examples/shallow_water_Bickley_jet.jl b/examples/shallow_water_Bickley_jet.jl index 9f8a4e2e21..9b6cd3ca6b 100644 --- a/examples/shallow_water_Bickley_jet.jl +++ b/examples/shallow_water_Bickley_jet.jl @@ -103,15 +103,15 @@ u = uh / h v = vh / h ## Build and compute mean vorticity discretely -ω = ComputedField(∂x(v) - ∂y(u)) +ω = Field(∂x(v) - ∂y(u)) compute!(ω) ## Copy mean vorticity to a new field -ωⁱ = Field(Face, Face, Nothing, model.architecture, model.grid) +ωⁱ = Field{Face, Face, Nothing}(model.grid) ωⁱ .= ω ## Use this new field to compute the perturbation vorticity -ω′ = ComputedField(ω - ωⁱ) +ω′ = Field(ω - ωⁱ) # and finally set the "true" initial condition with noise, diff --git a/examples/two_dimensional_turbulence.jl b/examples/two_dimensional_turbulence.jl index 68817c23e7..9db1dd3fc6 100644 --- a/examples/two_dimensional_turbulence.jl +++ b/examples/two_dimensional_turbulence.jl @@ -5,7 +5,7 @@ # # * How to run a model with no tracers and no buoyancy model. # * How to use `AbstractOperations`. -# * How to use `ComputedField`s to generate output. +# * How to use computed `Field`s to generate output. # ## Install dependencies # @@ -58,7 +58,7 @@ set!(model, u=uᵢ, v=vᵢ) # the `NamedTuple` model.velocities: u, v, w = model.velocities -# Next we create two objects called `ComputedField`s that calculate +# Next we create two `Field`s that calculate # _(i)_ vorticity that measures the rate at which the fluid rotates # and is defined as # @@ -68,7 +68,7 @@ u, v, w = model.velocities ω = ∂x(v) - ∂y(u) -ω_field = ComputedField(ω) +ω_field = Field(ω) # We also calculate _(ii)_ the _speed_ of the flow, # @@ -78,9 +78,9 @@ u, v, w = model.velocities s = sqrt(u^2 + v^2) -s_field = ComputedField(s) +s_field = Field(s) -# We'll pass these `ComputedField`s to an output writer below to calculate and output them during the simulation. +# We'll pass these `Field`s to an output writer below to calculate and output them during the simulation. simulation = Simulation(model, Δt=0.2, stop_time=50) From 1c5cce679312c5795035bf332bd5658b6f79e5fd Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Wed, 12 Jan 2022 12:33:06 -0500 Subject: [PATCH 069/140] fixed bug in grid --- src/Grids/rectilinear_grid.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl index 0dc9c6f4c3..8f4b752628 100644 --- a/src/Grids/rectilinear_grid.jl +++ b/src/Grids/rectilinear_grid.jl @@ -296,7 +296,7 @@ function RectilinearGrid(architecture::AbstractArchitecture = CPU(), Hx, Hy, Hz, FT(Lx), FT(Ly), FT(Lz), Δxᶠᵃᵃ, Δxᶜᵃᵃ, xᶠᵃᵃ, xᶜᵃᵃ, - Δyᵃᶜᵃ, Δyᵃᶠᵃ, yᵃᶠᵃ, yᵃᶜᵃ, + Δyᵃᶠᵃ, Δyᵃᶜᵃ, yᵃᶠᵃ, yᵃᶜᵃ, Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ) end From 4a5ee20bd19f0f308c1635e822ae2883e0cec4f2 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 11:07:01 -0700 Subject: [PATCH 070/140] Updates docstring output --- src/Grids/rectilinear_grid.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl index 0dc9c6f4c3..7983daeb8d 100644 --- a/src/Grids/rectilinear_grid.jl +++ b/src/Grids/rectilinear_grid.jl @@ -268,7 +268,7 @@ RectilinearGrid{Float64, Periodic, Bounded, Bounded} size (Nx, Ny, Nz): (32, 30, 24) halo (Hx, Hy, Hz): (1, 1, 1) spacing in x: Regular, with spacing 6.25 - spacing in y: Stretched, with spacing min=0.2739052315863262, max=5.22642316338267 + spacing in y: Stretched, with spacing min=0.2739052315863333, max=5.226423163382681 spacing in z: Stretched, with spacing min=0.6826950100338962, max=1.8309085743885056 ``` """ From c4e04d36742a35b056f7a309ca1c3bc00db13cb5 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 13:09:02 -0500 Subject: [PATCH 071/140] May need to relax tolerances in field reduction tests --- test/test_field.jl | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/test/test_field.jl b/test/test_field.jl index 82fe271c7a..7927df4593 100644 --- a/test/test_field.jl +++ b/test/test_field.jl @@ -68,8 +68,9 @@ function run_field_reduction_tests(FT, arch) for (ϕ, ϕ_vals) in zip(ϕs, ϕs_vals) - ε = eps(maximum(ϕ_vals)) + ε = eps(eltype(ϕ_vals)) + # @show ϕ .- ϕ_vals @test all(isapprox.(ϕ, ϕ_vals, atol=ε)) # if this isn't true, reduction tests can't pass # Important to make sure no CUDA scalar operations occur! @@ -280,10 +281,10 @@ end xw, yw, zw = nodes(w) xc, yc, zc = nodes(c) - @test u[1, 2, 3] == f(xu[1], yu[2], zu[3]) - @test v[1, 2, 3] == f(xv[1], yv[2], zv[3]) - @test w[1, 2, 3] == f(xw[1], yw[2], zw[3]) - @test c[1, 2, 3] == f(xc[1], yc[2], zc[3]) + @test u[1, 2, 3] ≈ f(xu[1], yu[2], zu[3]) + @test v[1, 2, 3] ≈ f(xv[1], yv[2], zv[3]) + @test w[1, 2, 3] ≈ f(xw[1], yw[2], zw[3]) + @test c[1, 2, 3] ≈ f(xc[1], yc[2], zc[3]) end end From 07c2f96d277dcd597f53311026fed785436afe65 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 14:54:34 -0500 Subject: [PATCH 072/140] Fusses with tolerance for field reduction tests --- test/test_field.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/test_field.jl b/test/test_field.jl index 7927df4593..f3faffb2f3 100644 --- a/test/test_field.jl +++ b/test/test_field.jl @@ -68,9 +68,9 @@ function run_field_reduction_tests(FT, arch) for (ϕ, ϕ_vals) in zip(ϕs, ϕs_vals) - ε = eps(eltype(ϕ_vals)) + ε = eps(eltype(ϕ_vals)) * 10 * maximum(maximum.(ϕs_vals)) + @info " Testing field reductions with tolerance $ε..." - # @show ϕ .- ϕ_vals @test all(isapprox.(ϕ, ϕ_vals, atol=ε)) # if this isn't true, reduction tests can't pass # Important to make sure no CUDA scalar operations occur! From a72b13ca4b87b547f1b5ce79f3e09f3c3bfcbaff Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 14:57:41 -0500 Subject: [PATCH 073/140] Update show_coordinate --- src/Grids/grid_generation.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/Grids/grid_generation.jl b/src/Grids/grid_generation.jl index b48fe4eaeb..0b8451282b 100644 --- a/src/Grids/grid_generation.jl +++ b/src/Grids/grid_generation.jl @@ -5,7 +5,8 @@ @inline show_coordinate(Δ::Number, T) = "Regular, with spacing $Δ" @inline show_coordinate(Δ::Number, ::Type{Flat}) = "Flattened" -@inline show_coordinate(Δ::AbstractVector, T) = "Stretched, with spacing min=$(minimum(parent(Δ))), max=$(maximum(parent(Δ)))" +@inline show_coordinate(Δ::AbstractVector, T) = @sprintf("Stretched, with spacing min=%.6f, max=%.6f", + minimum(parent(Δ)), maximum(parent(Δ))) get_domain_extent(coord, N) = (coord[1], coord[2]) get_domain_extent(coord::Function, N) = (coord(1), coord(N+1)) From 4f5e6af209585a4b0b0805202dc364ce2c2d552c Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 15:07:57 -0500 Subject: [PATCH 074/140] Fixes NonhydrostaticModel comparison --- test/test_nonhydrostatic_models.jl | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/test/test_nonhydrostatic_models.jl b/test/test_nonhydrostatic_models.jl index f3eaf50915..79b15c3a39 100644 --- a/test/test_nonhydrostatic_models.jl +++ b/test/test_nonhydrostatic_models.jl @@ -115,19 +115,21 @@ include("dependencies_for_runtests.jl") xF, yF, zF = nodes((Face, Face, Face), model.grid; reshape=true) # Form solution arrays - u_answer = u₀.(xF, yC, zC) - v_answer = v₀.(xC, yF, zC) - w_answer = w₀.(xC, yC, zF) - T_answer = T₀.(xC, yC, zC) - S_answer = S₀.(xC, yC, zC) + u_answer = u₀.(xF, yC, zC) |> Array + v_answer = v₀.(xC, yF, zC) |> Array + w_answer = w₀.(xC, yC, zF) |> Array + T_answer = T₀.(xC, yC, zC) |> Array + S_answer = S₀.(xC, yC, zC) |> Array Nx, Ny, Nz = size(model.grid) - u_cpu = XFaceField(grid) - v_cpu = YFaceField(grid) - w_cpu = ZFaceField(grid) - T_cpu = CenterField(grid) - S_cpu = CenterField(grid) + cpu_grid = on_architecture(CPU(), grid) + + u_cpu = XFaceField(cpu_grid) + v_cpu = YFaceField(cpu_grid) + w_cpu = ZFaceField(cpu_grid) + T_cpu = CenterField(cpu_grid) + S_cpu = CenterField(cpu_grid) set!(u_cpu, u) set!(v_cpu, v) From fff0b6b1a849a9e1af421204966ba919ea2afcb6 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 13:17:43 -0700 Subject: [PATCH 075/140] Move show_coordinate to grid_utils and import Printf --- src/Grids/grid_generation.jl | 5 ----- src/Grids/grid_utils.jl | 7 +++++++ 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/Grids/grid_generation.jl b/src/Grids/grid_generation.jl index 0b8451282b..ed9ffeb379 100644 --- a/src/Grids/grid_generation.jl +++ b/src/Grids/grid_generation.jl @@ -3,11 +3,6 @@ @inline adapt_if_vector(to, var) = var @inline adapt_if_vector(to, var::AbstractArray) = Adapt.adapt(to, var) -@inline show_coordinate(Δ::Number, T) = "Regular, with spacing $Δ" -@inline show_coordinate(Δ::Number, ::Type{Flat}) = "Flattened" -@inline show_coordinate(Δ::AbstractVector, T) = @sprintf("Stretched, with spacing min=%.6f, max=%.6f", - minimum(parent(Δ)), maximum(parent(Δ))) - get_domain_extent(coord, N) = (coord[1], coord[2]) get_domain_extent(coord::Function, N) = (coord(1), coord(N+1)) get_domain_extent(coord::AbstractVector, N) = (coord[1], coord[N+1]) diff --git a/src/Grids/grid_utils.jl b/src/Grids/grid_utils.jl index f2cc458712..a2520cfb72 100644 --- a/src/Grids/grid_utils.jl +++ b/src/Grids/grid_utils.jl @@ -1,4 +1,5 @@ using CUDA +using Printf ##### ##### Convenience functions @@ -351,3 +352,9 @@ end ##### struct ZDirection end + +@inline show_coordinate(Δ::Number, T) = "Regular, with spacing $Δ" +@inline show_coordinate(Δ::Number, ::Type{Flat}) = "Flattened" +@inline show_coordinate(Δ::AbstractVector, T) = @sprintf("Stretched, with spacing min=%.6f, max=%.6f", + minimum(parent(Δ)), maximum(parent(Δ))) + From d14325fefb3b1e0135c9d9b5f5aabe711d2675f0 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 13:36:44 -0700 Subject: [PATCH 076/140] Fix doctest for grid --- docs/src/model_setup/grids.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/src/model_setup/grids.md b/docs/src/model_setup/grids.md index b4b6dd6422..0212ef8480 100644 --- a/docs/src/model_setup/grids.md +++ b/docs/src/model_setup/grids.md @@ -115,8 +115,8 @@ RectilinearGrid{Float64, Periodic, Bounded, Bounded} size (Nx, Ny, Nz): (64, 64, 32) halo (Hx, Hy, Hz): (1, 1, 1) spacing in x: Regular, with spacing 156.25 - spacing in y: Stretched, with spacing min=6.022718974138115, max=245.33837163709035 - spacing in z: Stretched, with spacing min=2.407636663901485, max=49.008570164780394 + spacing in y: Stretched, with spacing min=6.022719, max=245.338372 + spacing in z: Stretched, with spacing min=2.407637, max=49.008570 ``` ```@setup 1 From c36f9481e0fab67a857823707d606f1e946f369a Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 13:37:56 -0700 Subject: [PATCH 077/140] Fix docstring in rectilinear_grid --- src/Grids/rectilinear_grid.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl index 53310e914b..5bb6ed61e9 100644 --- a/src/Grids/rectilinear_grid.jl +++ b/src/Grids/rectilinear_grid.jl @@ -236,7 +236,7 @@ RectilinearGrid{Float64, Periodic, Periodic, Bounded} halo (Hx, Hy, Hz): (1, 1, 1) spacing in x: Regular, with spacing 2.0 spacing in y: Regular, with spacing 2.0 - spacing in z: Stretched, with spacing min=0.6826950100338962, max=1.8309085743885056 + spacing in z: Stretched, with spacing min=0.682695, max=1.830909 ``` * A three-dimensional grid with regular spacing in x, cell interfaces that are closely spaced @@ -268,8 +268,8 @@ RectilinearGrid{Float64, Periodic, Bounded, Bounded} size (Nx, Ny, Nz): (32, 30, 24) halo (Hx, Hy, Hz): (1, 1, 1) spacing in x: Regular, with spacing 6.25 - spacing in y: Stretched, with spacing min=0.2739052315863333, max=5.226423163382681 - spacing in z: Stretched, with spacing min=0.6826950100338962, max=1.8309085743885056 + spacing in y: Stretched, with spacing min=0.273905, max=5.226423 + spacing in z: Stretched, with spacing min=0.682695, max=1.830909 ``` """ function RectilinearGrid(architecture::AbstractArchitecture = CPU(), From 7e3e4b00dcf40fe5cf8812f324cbed5c9b10c0a8 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 14:33:37 -0700 Subject: [PATCH 078/140] set! only Field and use child_arch in distributed field type param --- src/Architectures.jl | 9 +-------- src/Distributed/distributed_fields.jl | 4 ++-- src/Fields/set!.jl | 16 ++++++++-------- 3 files changed, 11 insertions(+), 18 deletions(-) diff --git a/src/Architectures.jl b/src/Architectures.jl index 8b0295da15..06ed67988e 100644 --- a/src/Architectures.jl +++ b/src/Architectures.jl @@ -57,14 +57,7 @@ arch_array(::CPU, A::CuArray) = Array(A) arch_array(::GPU, A::Array) = CuArray(A) arch_array(::GPU, A::CuArray) = A -const OffsetCPUArray = OffsetArray{FT, N, <:Array} where {FT, N} -const OffsetGPUArray = OffsetArray{FT, N, <:CuArray} where {FT, N} - -Adapt.adapt_structure(::CPU, a::OffsetCPUArray) = a -Adapt.adapt_structure(::GPU, a::OffsetGPUArray) = a - -Adapt.adapt_structure(::GPU, a::OffsetCPUArray) = OffsetArray(CuArray(a.parent), a.offsets...) -Adapt.adapt_structure(::CPU, a::OffsetGPUArray) = OffsetArray(Array(a.parent), a.offsets...) +arch_array(arch, a::OffsetArray) = OffsetArray(arch_array(arch, a.parent), a.offsets...) device_event(arch) = Event(device(arch)) diff --git a/src/Distributed/distributed_fields.jl b/src/Distributed/distributed_fields.jl index 4b732ede6c..149ea59ac6 100644 --- a/src/Distributed/distributed_fields.jl +++ b/src/Distributed/distributed_fields.jl @@ -4,10 +4,10 @@ import Oceananigans.BoundaryConditions: fill_halo_regions! function Field((LX, LY, LZ)::Tuple, grid::DistributedGrid, data, old_bcs, op, status) arch = architecture(grid) new_bcs = inject_halo_communication_boundary_conditions(old_bcs, arch.local_rank, arch.connectivity) - return Field{LX, LY, LZ}(grid, data, new_bcs, op, status) + return Field{LX, LY, LZ}(grid, data, new_bcs, op, status, child_architecture(arch)) end -const DistributedField = Field{<:Any, <:Any, <:Any, <:Any, <:MultiArch} +const DistributedField = Field{<:Any, <:Any, <:Any, <:Any, <:Any, <:DistributdGrid} fill_halo_regions!(field::DistributedField, arch, args...; kwargs...) = fill_halo_regions!(field.data, diff --git a/src/Fields/set!.jl b/src/Fields/set!.jl index 6ca4292dbb..0ccdd30c36 100644 --- a/src/Fields/set!.jl +++ b/src/Fields/set!.jl @@ -18,10 +18,10 @@ end set!(u::AbstractField, v) = u .= v # fallback # Niceties -const AbstractCPUField = AbstractField{X, Y, Z, <:CPU} where {X, Y, Z} +const CPUField = Field{LX, LY, LZ, O, <:CPU} where {LX, LY, LZ, O} """ Set the CPU field `u` data to the function `f(x, y, z)`. """ -function set!(u::AbstractCPUField, f::Function) +function set!(u::CPUField, f::Function) f_field = FunctionField(location(u), f, u.grid) u .= f_field return nothing @@ -31,10 +31,10 @@ end ##### set! for fields on the GPU ##### -const AbstractGPUField = AbstractField{X, Y, Z, <:GPU} where {X, Y, Z} +const GPUField = Field{LX, LY, LZ, O, <:GPU} where {LX, LY, LZ, O} """ Set the GPU field `u` to the array or function `v`. """ -function set!(u::AbstractGPUField, v::Union{Array, Function}) +function set!(u::GPUField, v::Union{Array, Function}) cpu_grid = on_architecture(CPU(), u.grid) v_field = similar(u, cpu_grid) set!(v_field, v) @@ -47,10 +47,10 @@ end ##### """ Set the CPU field `u` data to the GPU field data of `v`. """ -set!(u::AbstractCPUField, v::AbstractGPUField) = copyto!(parent(u), parent(v)) +set!(u::CPUField, v::GPUField) = copyto!(parent(u), parent(v)) """ Set the GPU field `u` data to the CPU field data of `v`. """ -set!(u::AbstractGPUField, v::AbstractCPUField) = copyto!(parent(u), parent(v)) +set!(u::GPUField, v::CPUField) = copyto!(parent(u), parent(v)) -set!(u::AbstractCPUField, v::AbstractCPUField) = parent(u) .= parent(v) -set!(u::AbstractGPUField, v::AbstractGPUField) = parent(u) .= parent(v) +set!(u::CPUField, v::CPUField) = parent(u) .= parent(v) +set!(u::GPUField, v::GPUField) = parent(u) .= parent(v) From 041939ff867790535163009ecc3bd6ccebf8c6ba Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 14:33:56 -0700 Subject: [PATCH 079/140] on_architecture for conformal cubed sphere grid --- src/Grids/conformal_cubed_sphere_face_grid.jl | 148 +++++++++++------- 1 file changed, 94 insertions(+), 54 deletions(-) diff --git a/src/Grids/conformal_cubed_sphere_face_grid.jl b/src/Grids/conformal_cubed_sphere_face_grid.jl index 93ca6c6d6c..aa60ec3b18 100644 --- a/src/Grids/conformal_cubed_sphere_face_grid.jl +++ b/src/Grids/conformal_cubed_sphere_face_grid.jl @@ -7,37 +7,57 @@ using Adapt: adapt_structure using Oceananigans struct ConformalCubedSphereFaceGrid{FT, TX, TY, TZ, A, R, Arch} <: AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ, Arch} - architecture :: Arch - Nx :: Int - Ny :: Int - Nz :: Int - Hx :: Int - Hy :: Int - Hz :: Int - λᶜᶜᵃ :: A - λᶠᶜᵃ :: A - λᶜᶠᵃ :: A - λᶠᶠᵃ :: A - φᶜᶜᵃ :: A - φᶠᶜᵃ :: A - φᶜᶠᵃ :: A - φᶠᶠᵃ :: A - zᵃᵃᶜ :: R - zᵃᵃᶠ :: R - Δxᶜᶜᵃ :: A - Δxᶠᶜᵃ :: A - Δxᶜᶠᵃ :: A - Δxᶠᶠᵃ :: A - Δyᶜᶜᵃ :: A - Δyᶜᶠᵃ :: A - Δyᶠᶜᵃ :: A - Δyᶠᶠᵃ :: A - Δz :: FT - Azᶜᶜᵃ :: A - Azᶠᶜᵃ :: A - Azᶜᶠᵃ :: A - Azᶠᶠᵃ :: A - radius :: FT + architecture :: Arch + Nx :: Int + Ny :: Int + Nz :: Int + Hx :: Int + Hy :: Int + Hz :: Int + λᶜᶜᵃ :: A + λᶠᶜᵃ :: A + λᶜᶠᵃ :: A + λᶠᶠᵃ :: A + φᶜᶜᵃ :: A + φᶠᶜᵃ :: A + φᶜᶠᵃ :: A + φᶠᶠᵃ :: A + zᵃᵃᶜ :: R + zᵃᵃᶠ :: R + Δxᶜᶜᵃ :: A + Δxᶠᶜᵃ :: A + Δxᶜᶠᵃ :: A + Δxᶠᶠᵃ :: A + Δyᶜᶜᵃ :: A + Δyᶜᶠᵃ :: A + Δyᶠᶜᵃ :: A + Δyᶠᶠᵃ :: A + Δz :: FT + Azᶜᶜᵃ :: A + Azᶠᶜᵃ :: A + Azᶜᶠᵃ :: A + Azᶠᶠᵃ :: A + radius :: FT + + function ConformalCubedSphereFaceGrid{TX, TY, TZ}(architecture, + Nx, Ny, Nz, + Hx, Hy, Hz, + λᶜᶜᵃ :: A, λᶠᶜᵃ :: A, λᶜᶠᵃ :: A, λᶠᶠᵃ :: A, + φᶜᶜᵃ :: A, φᶠᶜᵃ :: A, φᶜᶠᵃ :: A, φᶠᶠᵃ :: A, zᵃᵃᶜ :: R, zᵃᵃᶠ :: R, + Δxᶜᶜᵃ :: A, Δxᶠᶜᵃ :: A, Δxᶜᶠᵃ :: A, Δxᶠᶠᵃ :: A, + Δyᶜᶜᵃ :: A, Δyᶜᶠᵃ :: A, Δyᶠᶜᵃ :: A, Δyᶠᶠᵃ :: A, Δz :: FT, + Azᶜᶜᵃ :: A, Azᶠᶜᵃ :: A, Azᶜᶠᵃ :: A, Azᶠᶠᵃ :: A, + radius :: FT) where {TX, TY, TZ, FT, A, R, Arch} + + return new{FT, TX, TY, TZ, A, R, Arch}(architecture + Nx, Ny, Nz, + Hx, Hy, Hz, + λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ, + φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ, zᵃᵃᶜ, zᵃᵃᶠ, + Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, + Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ, Δz, + Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, radius) + end end function ConformalCubedSphereFaceGrid(arch::AbstractArchitecture = CPU(), @@ -247,16 +267,14 @@ function ConformalCubedSphereFaceGrid(filepath::AbstractString, architecture = C φᶠᶜᵃ = offset_data(zeros(FT, architecture, Txᶠᶜ, Tyᶠᶜ), loc_fc, topology, N, H) φᶜᶠᵃ = offset_data(zeros(FT, architecture, Txᶜᶠ, Tyᶜᶠ), loc_cf, topology, N, H) - Arch = typeof(architecture) - - return ConformalCubedSphereFaceGrid{FT, TX, TY, TZ, typeof(λᶜᶜᵃ), typeof(zᵃᵃᶠ), Arch}( - architecture, Nξ, Nη, Nz, Hx, Hy, Hz, + return ConformalCubedSphereFaceGrid{TX, TY, TZ}(architecture, Nξ, Nη, Nz, Hx, Hy, Hz, λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ, φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ, zᵃᵃᶜ, zᵃᵃᶠ, Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ, Δz, Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, radius) end -function _adapt_structure(to, grid::ConformalCubedSphereFaceGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} +function on_architecture(arch, grid::ConformalCubedSphereFaceGrid) + horizontal_coordinates = (:λᶜᶜᵃ, :λᶠᶜᵃ, :λᶜᶠᵃ, @@ -266,8 +284,6 @@ function _adapt_structure(to, grid::ConformalCubedSphereFaceGrid{FT, TX, TY, TZ} :φᶜᶠᵃ, :φᶠᶠᵃ) - horizontal_coordinate_data = Tuple(adapt_structure(to, getproperty(grid, name)) for name in horizontal_coordinates) - horizontal_grid_spacings = (:Δxᶜᶜᵃ, :Δxᶠᶜᵃ, :Δxᶜᶠᵃ, @@ -277,32 +293,56 @@ function _adapt_structure(to, grid::ConformalCubedSphereFaceGrid{FT, TX, TY, TZ} :Δyᶠᶜᵃ, :Δyᶠᶠᵃ) - horizontal_grid_spacing_data = Tuple(adapt_structure(to, getproperty(grid, name)) for name in horizontal_grid_spacings) - horizontal_areas = (:Azᶜᶜᵃ, :Azᶠᶜᵃ, :Azᶜᶠᵃ, :Azᶠᶠᵃ) - horizontal_area_data = Tuple(adapt_structure(to, getproperty(grid, name)) for name in horizontal_areas) + horizontal_grid_spacing_data = Tuple(arch_array(arch, getproperty(grid, name)) for name in horizontal_grid_spacings) + horizontal_coordinate_data = Tuple(arch_array(arch, getproperty(grid, name)) for name in horizontal_coordinates) + horizontal_area_data = Tuple(arch_array(arch, getproperty(grid, name)) for name in horizontal_areas) - zᵃᵃᶜ = adapt_structure(to, grid.zᵃᵃᶜ) - zᵃᵃᶠ = adapt_structure(to, grid.zᵃᵃᶠ) + zᵃᵃᶜ = arch_array(arch, grid.zᵃᵃᶜ) + zᵃᵃᶠ = arch_array(arch, grid.zᵃᵃᶠ) - A = typeof(horizontal_coordinate_data[1]) - R = typeof(grid.zᵃᵃᶜ) - - Arch = typeof(architecture) + TX, TY, TZ = topology(grid) - return ConformalCubedSphereFaceGrid{FT, TX, TY, TZ, A, R, Arch}(architecture, grid.Nx, grid.Ny, grid.Nz, grid.Hx, grid.Hy, grid.Hz, - horizontal_coordinate_data..., zᵃᵃᶜ, zᵃᵃᶠ, - horizontal_grid_spacing_data..., grid.Δz, - horizontal_area_data..., grid.radius) + return ConformalCubedSphereFaceGrid{TX, TY, TZ}(architecture, grid.Nx, grid.Ny, grid.Nz, grid.Hx, grid.Hy, grid.Hz, + horizontal_coordinate_data..., zᵃᵃᶜ, zᵃᵃᶠ, + horizontal_grid_spacing_data..., grid.Δz, + horizontal_area_data..., grid.radius) end -Adapt.adapt_structure(::CPU, grid::ConformalCubedSphereFaceGrid) = _adapt_structure(CPU(), grid) -Adapt.adapt_structure(to, grid::ConformalCubedSphereFaceGrid) = _adapt_structure(to, grid) - +function Adapt.adapt_structure(to, grid::ConformalCubedSphereFaceGrid) + TX, TY, TZ = topology(grid) + return ConformalCubedSphereFaceGrid{TX, TY, TZ}(nothing, + grid.Nx, grid.Ny, grid.Nz, + grid.Hx, grid.Hy, grid.Hz, + adapt(to, λᶜᶜᵃ), + adapt(to, λᶠᶜᵃ), + adapt(to, λᶜᶠᵃ), + adapt(to, λᶠᶠᵃ), + adapt(to, φᶜᶜᵃ), + adapt(to, φᶠᶜᵃ), + adapt(to, φᶜᶠᵃ), + adapt(to, φᶠᶠᵃ), + adapt(to, zᵃᵃᶜ), + adapt(to, zᵃᵃᶠ), + adapt(to, Δxᶜᶜᵃ), + adapt(to, Δxᶠᶜᵃ), + adapt(to, Δxᶜᶠᵃ), + adapt(to, Δxᶠᶠᵃ), + adapt(to, Δyᶜᶜᵃ), + adapt(to, Δyᶜᶠᵃ), + adapt(to, Δyᶠᶜᵃ), + adapt(to, Δyᶠᶠᵃ), + Δz, + adapt(to, Azᶜᶜᵃ), + adapt(to, Azᶠᶜᵃ), + adapt(to, Azᶜᶠᵃ), + adapt(to, Azᶠᶠᵃ), + radius) +end function Base.show(io::IO, g::ConformalCubedSphereFaceGrid{FT}) where FT print(io, "ConformalCubedSphereFaceGrid{$FT}\n", From d43721359e32908ee68e1093cf8d8a8d129f89f9 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 15:11:50 -0700 Subject: [PATCH 080/140] Fix typo --- src/Distributed/distributed_fields.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Distributed/distributed_fields.jl b/src/Distributed/distributed_fields.jl index 149ea59ac6..7265cf9c65 100644 --- a/src/Distributed/distributed_fields.jl +++ b/src/Distributed/distributed_fields.jl @@ -7,7 +7,7 @@ function Field((LX, LY, LZ)::Tuple, grid::DistributedGrid, data, old_bcs, op, st return Field{LX, LY, LZ}(grid, data, new_bcs, op, status, child_architecture(arch)) end -const DistributedField = Field{<:Any, <:Any, <:Any, <:Any, <:Any, <:DistributdGrid} +const DistributedField = Field{<:Any, <:Any, <:Any, <:Any, <:Any, <:DistributedGrid} fill_halo_regions!(field::DistributedField, arch, args...; kwargs...) = fill_halo_regions!(field.data, From 158e1baccca87c65e6ebcc31fdf9ec84a87c6774 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 15:11:58 -0700 Subject: [PATCH 081/140] Bugfixes in conformal cubed sphere grid --- src/Grids/conformal_cubed_sphere_face_grid.jl | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/Grids/conformal_cubed_sphere_face_grid.jl b/src/Grids/conformal_cubed_sphere_face_grid.jl index aa60ec3b18..71a21c8f75 100644 --- a/src/Grids/conformal_cubed_sphere_face_grid.jl +++ b/src/Grids/conformal_cubed_sphere_face_grid.jl @@ -39,7 +39,7 @@ struct ConformalCubedSphereFaceGrid{FT, TX, TY, TZ, A, R, Arch} <: AbstractHoriz Azᶠᶠᵃ :: A radius :: FT - function ConformalCubedSphereFaceGrid{TX, TY, TZ}(architecture, + function ConformalCubedSphereFaceGrid{TX, TY, TZ}(architecture::Arch, Nx, Ny, Nz, Hx, Hy, Hz, λᶜᶜᵃ :: A, λᶠᶜᵃ :: A, λᶜᶠᵃ :: A, λᶠᶠᵃ :: A, @@ -49,7 +49,7 @@ struct ConformalCubedSphereFaceGrid{FT, TX, TY, TZ, A, R, Arch} <: AbstractHoriz Azᶜᶜᵃ :: A, Azᶠᶜᵃ :: A, Azᶜᶠᵃ :: A, Azᶠᶠᵃ :: A, radius :: FT) where {TX, TY, TZ, FT, A, R, Arch} - return new{FT, TX, TY, TZ, A, R, Arch}(architecture + return new{FT, TX, TY, TZ, A, R, Arch}(architecture, Nx, Ny, Nz, Hx, Hy, Hz, λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ, @@ -171,11 +171,12 @@ function ConformalCubedSphereFaceGrid(arch::AbstractArchitecture = CPU(), Azᶜᶠᵃ = OffsetArray(zeros(Nξ + 2Hx, Nη + 2Hy + 1), -Hx, -Hy) Azᶠᶠᵃ = OffsetArray(zeros(Nξ + 2Hx + 1, Nη + 2Hy + 1), -Hx, -Hy) - return ConformalCubedSphereFaceGrid{FT, TX, TY, TZ, typeof(λᶜᶜᵃ), typeof(zᵃᵃᶠ), Arch}(arch, - Nξ, Nη, Nz, Hx, Hy, Hz, - λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ, φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ, zᵃᵃᶜ, zᵃᵃᶠ, - Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ, Δz, - Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, radius) + return ConformalCubedSphereFaceGrid{TX, TY, TZ}(arch, Nξ, Nη, Nz, Hx, Hy, Hz, + λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ, + φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ, zᵃᵃᶜ, zᵃᵃᶠ, + Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, + Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ, Δz, + Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, radius) end # architecture = CPU() default, assuming that a DataType positional arg From 24e3c17cb31368c9179ca985ecb5c3cdae647b7d Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 20:31:41 -0500 Subject: [PATCH 082/140] Fixes FieldTimeSeries tests --- test/test_output_readers.jl | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/test/test_output_readers.jl b/test/test_output_readers.jl index 29e6321bc3..88894b15f6 100644 --- a/test/test_output_readers.jl +++ b/test/test_output_readers.jl @@ -1,5 +1,7 @@ include("dependencies_for_runtests.jl") +using Oceananigans.OutputReaders: ViewField + function generate_some_interesting_simulation_data(Nx, Ny, Nz; architecture=CPU()) grid = RectilinearGrid(architecture, size=(Nx, Ny, Nz), extent=(64, 64, 32)) @@ -78,12 +80,12 @@ end # This behavior ensures that set! works # but perhaps should be changed in the future - @test parent(u3[1]) isa SubArray - @test parent(v3[1]) isa SubArray - @test parent(w3[1]) isa SubArray - @test parent(T3[1]) isa SubArray - @test parent(b3[1]) isa SubArray - @test parent(ζ3[1]) isa SubArray + @test size(parent(u3[1])) == size(parent(u3))[1:3] + @test size(parent(v3[1])) == size(parent(v3))[1:3] + @test size(parent(w3[1])) == size(parent(w3))[1:3] + @test size(parent(T3[1])) == size(parent(T3))[1:3] + @test size(parent(b3[1])) == size(parent(b3))[1:3] + @test size(parent(ζ3[1])) == size(parent(ζ3))[1:3] @test location(u3) == (Face, Center, Center) @test location(v3) == (Center, Face, Center) @@ -197,6 +199,6 @@ end end end - rm("test_3d_output_with_halos.jld2") - rm("test_1d_output_with_halos.jld2") + rm(filepath3d) + rm(filepath1d) end From 557a40d00a612aa80aedcc87a47987e455f922f8 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 20:37:45 -0500 Subject: [PATCH 083/140] Allow scalar indexing in on_architecture --- src/Grids/grid_generation.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Grids/grid_generation.jl b/src/Grids/grid_generation.jl index ed9ffeb379..e018145548 100644 --- a/src/Grids/grid_generation.jl +++ b/src/Grids/grid_generation.jl @@ -5,7 +5,7 @@ get_domain_extent(coord, N) = (coord[1], coord[2]) get_domain_extent(coord::Function, N) = (coord(1), coord(N+1)) -get_domain_extent(coord::AbstractVector, N) = (coord[1], coord[N+1]) +get_domain_extent(coord::AbstractVector, N) = CUDA.@allowscalar (coord[1], coord[N+1]) get_coord_face(coord::Nothing, i) = 1 get_coord_face(coord::Function, i) = coord(i) From 4f33243615b27839b492ad2f9eebe3d1c9d98013 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 20:57:32 -0500 Subject: [PATCH 084/140] Adds on_architecture for ImmersedBoundaryGrid --- src/ImmersedBoundaries/ImmersedBoundaries.jl | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/ImmersedBoundaries/ImmersedBoundaries.jl b/src/ImmersedBoundaries/ImmersedBoundaries.jl index cf43af1423..55bffcaf4b 100644 --- a/src/ImmersedBoundaries/ImmersedBoundaries.jl +++ b/src/ImmersedBoundaries/ImmersedBoundaries.jl @@ -40,6 +40,7 @@ import Oceananigans.Utils: cell_advection_timescale import Oceananigans.Grids: with_halo, architecture import Oceananigans.Coriolis: φᶠᶠᵃ import Oceananigans.Grids: with_halo, xnode, ynode, znode, all_x_nodes, all_y_nodes, all_z_nodes +import Oceananigans.Grids: on_architecture import Oceananigans.Advection: _advective_momentum_flux_Uu, @@ -136,6 +137,16 @@ all_x_nodes(loc, ibg::ImmersedBoundaryGrid) = all_x_nodes(loc, ibg.grid) all_y_nodes(loc, ibg::ImmersedBoundaryGrid) = all_y_nodes(loc, ibg.grid) all_z_nodes(loc, ibg::ImmersedBoundaryGrid) = all_z_nodes(loc, ibg.grid) +function on_architecture(arch, ibg::ImmersedBoundaryGrid) + underlying_grid = on_architecture(arch, ibg.grid) + + immersed_boundary = ibg.immersed_boundary isa AbstractArray ? + arch_array(arch, ibg.immersed_boundary) : + ibg.immersed_boundary + + return ImmersedBoundaryGrid(grid, immersed_boundary) +end + include("immersed_grid_metrics.jl") include("grid_fitted_immersed_boundaries.jl") include("conditional_fluxes.jl") From 0b95ed76edbd60e051bf6abd852c60feea38d9d3 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 21:17:46 -0500 Subject: [PATCH 085/140] Adds array_type for MultiArch --- src/Distributed/multi_architectures.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/Distributed/multi_architectures.jl b/src/Distributed/multi_architectures.jl index 824d38501f..32fc72962f 100644 --- a/src/Distributed/multi_architectures.jl +++ b/src/Distributed/multi_architectures.jl @@ -1,7 +1,7 @@ using Oceananigans.Architectures using Oceananigans.Grids: topology, validate_tupled_argument -import Oceananigans.Architectures: device, device_event, arch_array +import Oceananigans.Architectures: device, device_event, arch_array, array_type import Oceananigans.Grids: zeros struct MultiArch{A, R, I, ρ, C, γ} <: AbstractMultiArchitecture @@ -56,6 +56,7 @@ device(arch::AbstractMultiArchitecture) = device(child_architecture(arch) device_event(arch::AbstractMultiArchitecture) = device_event(child_architecture(arch)) arch_array(arch::AbstractMultiArchitecture, A) = arch_array(child_architecture(arch), A) zeros(FT, arch::MultiArch, N...) = zeros(FT, child_architecture(arch), N...) +array_type(arch::MultiArch) = array_type(child_architecture(arch)) ##### ##### Converting between index and MPI rank taking k as the fast index From be6994e130d9b06a87b29f67f399f7aeb35ebede Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 21:17:52 -0500 Subject: [PATCH 086/140] Cosmetics --- src/ImmersedBoundaries/grid_fitted_immersed_boundaries.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ImmersedBoundaries/grid_fitted_immersed_boundaries.jl b/src/ImmersedBoundaries/grid_fitted_immersed_boundaries.jl index 94e918c339..ceadd9d502 100644 --- a/src/ImmersedBoundaries/grid_fitted_immersed_boundaries.jl +++ b/src/ImmersedBoundaries/grid_fitted_immersed_boundaries.jl @@ -59,11 +59,11 @@ const GFBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:GridFit zero(eltype(ibg.grid)), Δzᵃᵃᶜ(i, j, k, ibg.grid)) -@inline Δzᶠᶜᶜ(i, j, k, ibg::GFBIBG) = ifelse(solid_node(Face(), Center(), Center(), i , j, k, ibg), +@inline Δzᶠᶜᶜ(i, j, k, ibg::GFBIBG) = ifelse(solid_node(Face(), Center(), Center(), i, j, k, ibg), zero(eltype(ibg)), Δzᵃᵃᶜ(i, j, k, ibg.grid)) -@inline Δzᶜᶠᶜ(i, j, k, ibg::GFBIBG) = ifelse(solid_node(Center(), Face(), Center(), i , j, k, ibg), +@inline Δzᶜᶠᶜ(i, j, k, ibg::GFBIBG) = ifelse(solid_node(Center(), Face(), Center(), i, j, k, ibg), zero(eltype(ibg)), Δzᵃᵃᶜ(i, j, k, ibg.grid)) From 0282c7a4744c18ea149b6835d21ac1d6b98d2e64 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 13 Jan 2022 14:42:55 +1100 Subject: [PATCH 087/140] bump minor release --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index ce6c9f0070..64c62e951d 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "Oceananigans" uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09" -version = "0.67.1" +version = "0.68.0" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" From 843fe21c25755c016c8aeb418ef840b6fd0a1958 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 13 Jan 2022 14:55:08 +1100 Subject: [PATCH 088/140] CenterField infers arch from grid --- test/test_distributed_poisson_solvers.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/test_distributed_poisson_solvers.jl b/test/test_distributed_poisson_solvers.jl index 607e114e18..13a5683259 100644 --- a/test/test_distributed_poisson_solvers.jl +++ b/test/test_distributed_poisson_solvers.jl @@ -46,8 +46,8 @@ function divergence_free_poisson_solution_triply_periodic(grid_points, ranks) p_bcs = FieldBoundaryConditions(local_grid, (Center, Center, Center)) p_bcs = inject_halo_communication_boundary_conditions(p_bcs, arch.local_rank, arch.connectivity) - ϕ = CenterField(child_architecture(arch), local_grid, p_bcs) # "pressure" - ∇²ϕ = CenterField(child_architecture(arch), local_grid, p_bcs) + ϕ = CenterField(local_grid, p_bcs) # "pressure" + ∇²ϕ = CenterField(local_grid, p_bcs) interior(ϕ) .= real(first(solver.storage)) compute_∇²!(∇²ϕ, ϕ, arch, local_grid) From 6ec01fe1ae809ae20de76297579bc63b373f9519 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 13 Jan 2022 15:00:27 +1100 Subject: [PATCH 089/140] avoid ambiguity in `on_architecture` method --- src/Grids/grid_utils.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Grids/grid_utils.jl b/src/Grids/grid_utils.jl index a2520cfb72..1413f1c53f 100644 --- a/src/Grids/grid_utils.jl +++ b/src/Grids/grid_utils.jl @@ -5,7 +5,7 @@ using Printf ##### Convenience functions ##### -on_architecture(::Nothing, grid) = grid +on_architecture(::Nothing, grid::AbstractGrid) = grid Base.length(::Type{Face}, topo, N) = N Base.length(::Type{Face}, ::Type{Bounded}, N) = N+1 From 0c53f78eec672f00e6d535174b81c65514925c03 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 00:27:59 -0500 Subject: [PATCH 090/140] Bugfix in on_architecture for immersed boundary girds --- src/ImmersedBoundaries/ImmersedBoundaries.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ImmersedBoundaries/ImmersedBoundaries.jl b/src/ImmersedBoundaries/ImmersedBoundaries.jl index 55bffcaf4b..e7b706dc29 100644 --- a/src/ImmersedBoundaries/ImmersedBoundaries.jl +++ b/src/ImmersedBoundaries/ImmersedBoundaries.jl @@ -144,7 +144,7 @@ function on_architecture(arch, ibg::ImmersedBoundaryGrid) arch_array(arch, ibg.immersed_boundary) : ibg.immersed_boundary - return ImmersedBoundaryGrid(grid, immersed_boundary) + return ImmersedBoundaryGrid(underlying_grid, immersed_boundary) end include("immersed_grid_metrics.jl") From fcc745466b8b667fc43585a674713a83773397fb Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 00:29:08 -0500 Subject: [PATCH 091/140] Bugfix in adapt_structure for conformal cubed sphere grid --- src/Grids/conformal_cubed_sphere_face_grid.jl | 48 +++++++++---------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/src/Grids/conformal_cubed_sphere_face_grid.jl b/src/Grids/conformal_cubed_sphere_face_grid.jl index 71a21c8f75..f1e6024388 100644 --- a/src/Grids/conformal_cubed_sphere_face_grid.jl +++ b/src/Grids/conformal_cubed_sphere_face_grid.jl @@ -319,30 +319,30 @@ function Adapt.adapt_structure(to, grid::ConformalCubedSphereFaceGrid) return ConformalCubedSphereFaceGrid{TX, TY, TZ}(nothing, grid.Nx, grid.Ny, grid.Nz, grid.Hx, grid.Hy, grid.Hz, - adapt(to, λᶜᶜᵃ), - adapt(to, λᶠᶜᵃ), - adapt(to, λᶜᶠᵃ), - adapt(to, λᶠᶠᵃ), - adapt(to, φᶜᶜᵃ), - adapt(to, φᶠᶜᵃ), - adapt(to, φᶜᶠᵃ), - adapt(to, φᶠᶠᵃ), - adapt(to, zᵃᵃᶜ), - adapt(to, zᵃᵃᶠ), - adapt(to, Δxᶜᶜᵃ), - adapt(to, Δxᶠᶜᵃ), - adapt(to, Δxᶜᶠᵃ), - adapt(to, Δxᶠᶠᵃ), - adapt(to, Δyᶜᶜᵃ), - adapt(to, Δyᶜᶠᵃ), - adapt(to, Δyᶠᶜᵃ), - adapt(to, Δyᶠᶠᵃ), - Δz, - adapt(to, Azᶜᶜᵃ), - adapt(to, Azᶠᶜᵃ), - adapt(to, Azᶜᶠᵃ), - adapt(to, Azᶠᶠᵃ), - radius) + adapt(to, grid.λᶜᶜᵃ), + adapt(to, grid.λᶠᶜᵃ), + adapt(to, grid.λᶜᶠᵃ), + adapt(to, grid.λᶠᶠᵃ), + adapt(to, grid.φᶜᶜᵃ), + adapt(to, grid.φᶠᶜᵃ), + adapt(to, grid.φᶜᶠᵃ), + adapt(to, grid.φᶠᶠᵃ), + adapt(to, grid.zᵃᵃᶜ), + adapt(to, grid.zᵃᵃᶠ), + adapt(to, grid.Δxᶜᶜᵃ), + adapt(to, grid.Δxᶠᶜᵃ), + adapt(to, grid.Δxᶜᶠᵃ), + adapt(to, grid.Δxᶠᶠᵃ), + adapt(to, grid.Δyᶜᶜᵃ), + adapt(to, grid.Δyᶜᶠᵃ), + adapt(to, grid.Δyᶠᶜᵃ), + adapt(to, grid.Δyᶠᶠᵃ), + grid.Δz, + adapt(to, grid.Azᶜᶜᵃ), + adapt(to, grid.Azᶠᶜᵃ), + adapt(to, grid.Azᶜᶠᵃ), + adapt(to, grid.Azᶠᶠᵃ), + grid.radius) end function Base.show(io::IO, g::ConformalCubedSphereFaceGrid{FT}) where FT From 9b5fc07c37361f76c2cf905aefabdc353f3182e6 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 12 Jan 2022 22:57:47 -0700 Subject: [PATCH 092/140] Bug fix in distributed poisson solvers test --- test/test_distributed_poisson_solvers.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/test_distributed_poisson_solvers.jl b/test/test_distributed_poisson_solvers.jl index 13a5683259..3adb6b34ad 100644 --- a/test/test_distributed_poisson_solvers.jl +++ b/test/test_distributed_poisson_solvers.jl @@ -46,8 +46,8 @@ function divergence_free_poisson_solution_triply_periodic(grid_points, ranks) p_bcs = FieldBoundaryConditions(local_grid, (Center, Center, Center)) p_bcs = inject_halo_communication_boundary_conditions(p_bcs, arch.local_rank, arch.connectivity) - ϕ = CenterField(local_grid, p_bcs) # "pressure" - ∇²ϕ = CenterField(local_grid, p_bcs) + ϕ = CenterField(local_grid, boundary_conditions=p_bcs) # "pressure" + ∇²ϕ = CenterField(local_grid, boundary_conditions=p_bcs) interior(ϕ) .= real(first(solver.storage)) compute_∇²!(∇²ϕ, ϕ, arch, local_grid) From 4bf77eb72f5fde595f20ba5e0e03c42738efe168 Mon Sep 17 00:00:00 2001 From: "Gregory L. Wagner" Date: Wed, 12 Jan 2022 22:58:09 -0700 Subject: [PATCH 093/140] Update src/AbstractOperations/kernel_function_operation.jl Co-authored-by: Navid C. Constantinou --- src/AbstractOperations/kernel_function_operation.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/AbstractOperations/kernel_function_operation.jl b/src/AbstractOperations/kernel_function_operation.jl index 0327b64813..145c1ff00f 100644 --- a/src/AbstractOperations/kernel_function_operation.jl +++ b/src/AbstractOperations/kernel_function_operation.jl @@ -57,8 +57,8 @@ u, v, w = model.velocities ``` """ function KernelFunctionOperation{LX, LY, LZ}(kernel_function, grid; - computed_dependencies = (), - parameters = nothing) where {LX, LY, LZ} + computed_dependencies = (), + parameters = nothing) where {LX, LY, LZ} return KernelFunctionOperation{LX, LY, LZ}(kernel_function, computed_dependencies, parameters, grid) end From fb86cdde2cb9404ce3d1cbe2526afb79ae67aa6f Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 07:40:43 -0700 Subject: [PATCH 094/140] Eliminate on_architecture(::Nothing, grid) --- src/Grids/grid_utils.jl | 2 -- src/OutputReaders/field_time_series.jl | 5 +++++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/src/Grids/grid_utils.jl b/src/Grids/grid_utils.jl index 1413f1c53f..7e2f88b696 100644 --- a/src/Grids/grid_utils.jl +++ b/src/Grids/grid_utils.jl @@ -5,8 +5,6 @@ using Printf ##### Convenience functions ##### -on_architecture(::Nothing, grid::AbstractGrid) = grid - Base.length(::Type{Face}, topo, N) = N Base.length(::Type{Face}, ::Type{Bounded}, N) = N+1 Base.length(::Type{Center}, topo, N) = N diff --git a/src/OutputReaders/field_time_series.jl b/src/OutputReaders/field_time_series.jl index a8736e171d..614769e384 100644 --- a/src/OutputReaders/field_time_series.jl +++ b/src/OutputReaders/field_time_series.jl @@ -99,6 +99,11 @@ function FieldTimeSeries(path, name, backend::InMemory; isnothing(times) && (times = [file["timeseries/t/$i"] for i in iterations]) isnothing(location) && (location = file["timeseries/$name/serialized/location"]) + # Default to CPU if neither architecture nor grid is specified + architecture = isnothing(architecture) ? + (isnothing(grid) ? CPU() : Architectures.architecture(grid)) : + architecture + if isnothing(grid) grid = on_architecture(architecture, file["serialized/grid"]) end From cf419f1080a7c0eacaef338a65758ad8e0d50ea2 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 07:42:15 -0700 Subject: [PATCH 095/140] Add arch_array for abstract range --- src/Architectures.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/Architectures.jl b/src/Architectures.jl index 06ed67988e..4844acb684 100644 --- a/src/Architectures.jl +++ b/src/Architectures.jl @@ -57,6 +57,7 @@ arch_array(::CPU, A::CuArray) = Array(A) arch_array(::GPU, A::Array) = CuArray(A) arch_array(::GPU, A::CuArray) = A +arch_array(arch, A::AbstractRange) = A arch_array(arch, a::OffsetArray) = OffsetArray(arch_array(arch, a.parent), a.offsets...) device_event(arch) = Event(device(arch)) From 9fdb82b1f730fe8a14b0b38ceefe6d67181a6dd4 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 10:08:07 -0700 Subject: [PATCH 096/140] Remove Architecture as type parameter in Field and AbstractField plus massive cleanup --- src/AbstractOperations/AbstractOperations.jl | 10 +- src/AbstractOperations/binary_operations.jl | 55 +++----- src/AbstractOperations/multiary_operations.jl | 61 +++------ src/AbstractOperations/unary_operations.jl | 34 ++--- src/Architectures.jl | 12 +- src/CubedSpheres/cubed_sphere_faces.jl | 24 ++-- src/Distributed/multi_architectures.jl | 11 +- src/Fields/Fields.jl | 30 +++-- src/Fields/abstract_field.jl | 59 ++------ src/Fields/averaged_field.jl | 120 ----------------- src/Fields/background_fields.jl | 16 ++- src/Fields/field.jl | 53 +++++--- src/Fields/field_slicer.jl | 2 + src/Fields/field_tuples.jl | 43 ++++++ src/Fields/function_field.jl | 32 ++--- src/Fields/kernel_computed_field.jl | 127 ------------------ src/Fields/set!.jl | 23 +--- src/Fields/show_fields.jl | 8 +- src/Fields/tracer_names.jl | 12 -- src/Fields/validate_field_tuple_grid.jl | 19 --- src/Fields/zero_field.jl | 6 +- src/OutputReaders/field_time_series.jl | 5 +- 22 files changed, 235 insertions(+), 527 deletions(-) delete mode 100644 src/Fields/averaged_field.jl delete mode 100644 src/Fields/kernel_computed_field.jl delete mode 100644 src/Fields/tracer_names.jl delete mode 100644 src/Fields/validate_field_tuple_grid.jl diff --git a/src/AbstractOperations/AbstractOperations.jl b/src/AbstractOperations/AbstractOperations.jl index 19c70961d8..44bff8fa74 100644 --- a/src/AbstractOperations/AbstractOperations.jl +++ b/src/AbstractOperations/AbstractOperations.jl @@ -29,17 +29,15 @@ import Oceananigans.Fields: data, compute_at! ##### Basic functionality ##### -using Oceananigans.Fields: AbstractOperation -# abstract type AbstractOperation{X, Y, Z, A, G, T} <: AbstractField{X, Y, Z, A, G, T, 3} end +abstract type AbstractOperation{LX, LY, LZ, G, T} <: AbstractField{LX, LY, LZ, G, T, 3} end -const AF = AbstractField - -# We (informally) require that all field-like objects define `parent`: -Base.parent(op::AbstractOperation) = op +const AF = AbstractField # used in unary_operations.jl, binary_operations.jl, etc # We have no halos to fill fill_halo_regions!(::AbstractOperation, args...; kwargs...) = nothing +architecture(a::AbstractOperation) = architecture(a.grid) + # AbstractOperation macros add their associated functions to this list const operators = Set() diff --git a/src/AbstractOperations/binary_operations.jl b/src/AbstractOperations/binary_operations.jl index d96071ef0d..75b2acf03f 100644 --- a/src/AbstractOperations/binary_operations.jl +++ b/src/AbstractOperations/binary_operations.jl @@ -1,23 +1,22 @@ const binary_operators = Set() -struct BinaryOperation{X, Y, Z, O, A, B, IA, IB, R, G, T} <: AbstractOperation{X, Y, Z, R, G, T} - op :: O - a :: A - b :: B - ▶a :: IA - ▶b :: IB - architecture :: R - grid :: G +struct BinaryOperation{LX, LY, LZ, O, A, B, IA, IB, G, T} <: AbstractOperation{LX, LY, LZ, G, T} + op :: O + a :: A + b :: B + ▶a :: IA + ▶b :: IB + grid :: G @doc """ - BinaryOperation{X, Y, Z}(op, a, b, ▶a, ▶b, arch, grid) + BinaryOperation{LX, LY, LZ}(op, a, b, ▶a, ▶b, grid) Returns an abstract representation of the binary operation `op(▶a(a), ▶b(b))`. - on `grid` and `arch`itecture, where `▶a` and `▶b` interpolate `a` and `b` to (X, Y, Z). + on `grid`, where `▶a` and `▶b` interpolate `a` and `b` to (LX, LY, LZ). """ - function BinaryOperation{X, Y, Z}(op::O, a::A, b::B, ▶a::IA, ▶b::IB, arch::R, grid::G) where {X, Y, Z, O, A, B, IA, IB, R, G} + function BinaryOperation{LX, LY, LZ}(op::O, a::A, b::B, ▶a::IA, ▶b::IB, grid::G) where {LX, LY, LZ, O, A, B, IA, IB, G} T = eltype(grid) - return new{X, Y, Z, O, A, B, IA, IB, R, G, T}(op, a, b, ▶a, ▶b, arch, grid) + return new{LX, LY, LZ, O, A, B, IA, IB, R, G, T}(op, a, b, ▶a, ▶b, grid) end end @@ -35,8 +34,7 @@ end function _binary_operation(Lc, op, a, b, La, Lb, grid) ▶a = interpolation_operator(La, Lc) ▶b = interpolation_operator(Lb, Lc) - arch = architecture(a, b) - return BinaryOperation{Lc[1], Lc[2], Lc[3]}(op, a, b, ▶a, ▶b, arch, grid) + return BinaryOperation{Lc[1], Lc[2], Lc[3]}(op, a, b, ▶a, ▶b, grid) end const ConcreteLocationType = Union{Type{Face}, Type{Center}} @@ -187,25 +185,6 @@ macro binary(ops...) return expr end -##### -##### Architecture inference for BinaryOperation -##### - -architecture(β::BinaryOperation) = β.architecture - -function architecture(a, b) - arch_a = architecture(a) - arch_b = architecture(b) - - arch_a === arch_b && return arch_a - isnothing(arch_a) && return arch_b - isnothing(arch_b) && return arch_a - - throw(ArgumentError("Operation involves fields on architectures $arch_a and $arch_b")) - - return nothing -end - ##### ##### Nested computations ##### @@ -221,7 +200,11 @@ end ##### "Adapt `BinaryOperation` to work on the GPU via CUDAnative and CUDAdrv." -Adapt.adapt_structure(to, binary::BinaryOperation{X, Y, Z}) where {X, Y, Z} = - BinaryOperation{X, Y, Z}(Adapt.adapt(to, binary.op), Adapt.adapt(to, binary.a), Adapt.adapt(to, binary.b), - Adapt.adapt(to, binary.▶a), Adapt.adapt(to, binary.▶b), nothing, Adapt.adapt(to, binary.grid)) +Adapt.adapt_structure(to, binary::BinaryOperation{LX, LY, LZ}) where {LX, LY, LZ} = + BinaryOperation{LX, LY, LZ}(Adapt.adapt(to, binary.op), + Adapt.adapt(to, binary.a), + Adapt.adapt(to, binary.b), + Adapt.adapt(to, binary.▶a), + Adapt.adapt(to, binary.▶b), + Adapt.adapt(to, binary.grid)) diff --git a/src/AbstractOperations/multiary_operations.jl b/src/AbstractOperations/multiary_operations.jl index d07995d87e..7924475682 100644 --- a/src/AbstractOperations/multiary_operations.jl +++ b/src/AbstractOperations/multiary_operations.jl @@ -1,30 +1,28 @@ const multiary_operators = Set() -struct MultiaryOperation{X, Y, Z, N, O, A, I, R, G, T} <: AbstractOperation{X, Y, Z, R, G, T} - op :: O - args :: A - ▶ :: I - architecture :: R - grid :: G - - function MultiaryOperation{X, Y, Z}(op::O, args::A, ▶::I, arch::R, grid::G) where {X, Y, Z, O, A, I, R, G} +struct MultiaryOperation{LX, LY, LZ, N, O, A, I, G, T} <: AbstractOperation{LX, LY, LZ, G, T} + op :: O + args :: A + ▶ :: I + grid :: G + + function MultiaryOperation{LX, LY, LZ}(op::O, args::A, ▶::I, grid::G) where {LX, LY, LZ, O, A, I, G} T = eltype(grid) N = length(args) - return new{X, Y, Z, N, O, A, I, R, G, T}(op, args, ▶, arch, grid) + return new{LX, LY, LZ, N, O, A, I, R, G, T}(op, args, ▶, grid) end end -@inline Base.getindex(Π::MultiaryOperation{X, Y, Z, N}, i, j, k) where {X, Y, Z, N} = +@inline Base.getindex(Π::MultiaryOperation{LX, LY, LZ, N}, i, j, k) where {LX, LY, LZ, N} = Π.op(ntuple(γ -> Π.▶[γ](i, j, k, Π.grid, Π.args[γ]), Val(N))...) ##### ##### MultiaryOperation construction ##### -function _multiary_operation(L, op, args, Largs, grid) where {X, Y, Z} +function _multiary_operation(L, op, args, Largs, grid) where {LX, LY, LZ} ▶ = Tuple(interpolation_operator(La, L) for La in Largs) - arch = architecture(args...) - return MultiaryOperation{L[1], L[2], L[3]}(op, Tuple(a for a in args), ▶, arch, grid) + return MultiaryOperation{L[1], L[2], L[3]}(op, Tuple(a for a in args), ▶, grid) end # Recompute location of multiary operation @@ -131,36 +129,15 @@ macro multiary(ops...) return expr end -##### -##### Architecture inference for MultiaryOperation -##### - -architecture(Π::MultiaryOperation) = Π.architecture - -function architecture(a, b, c, d...) - - archs = [] - - push!(archs, architecture(a, b)) - push!(archs, architecture(a, c)) - append!(archs, [architecture(a, di) for di in d]) - - for arch in archs - if !(arch === nothing) - return arch - end - end - - return nothing -end - ##### ##### Nested computations ##### function compute_at!(Π::MultiaryOperation, time) - c = Tuple(compute_at!(a, time) for a in Π.args) - return nothing + for a in Π.args + compute_at!(a, time) + end + return Π end ##### @@ -168,7 +145,9 @@ end ##### "Adapt `MultiaryOperation` to work on the GPU via CUDAnative and CUDAdrv." -Adapt.adapt_structure(to, multiary::MultiaryOperation{X, Y, Z}) where {X, Y, Z} = - MultiaryOperation{X, Y, Z}(Adapt.adapt(to, multiary.op), Adapt.adapt(to, multiary.args), - Adapt.adapt(to, multiary.▶), nothing, Adapt.adapt(to, multiary.grid)) +Adapt.adapt_structure(to, multiary::MultiaryOperation{LX, LY, LZ}) where {LX, LY, LZ} = + MultiaryOperation{LX, LY, LZ}(Adapt.adapt(to, multiary.op), + Adapt.adapt(to, multiary.args), + Adapt.adapt(to, multiary.▶), + Adapt.adapt(to, multiary.grid)) diff --git a/src/AbstractOperations/unary_operations.jl b/src/AbstractOperations/unary_operations.jl index 3360f99a28..c3426ad099 100644 --- a/src/AbstractOperations/unary_operations.jl +++ b/src/AbstractOperations/unary_operations.jl @@ -1,21 +1,20 @@ const unary_operators = Set() -struct UnaryOperation{X, Y, Z, O, A, I, R, G, T} <: AbstractOperation{X, Y, Z, R, G, T} - op :: O - arg :: A - ▶ :: I - architecture :: R - grid :: G +struct UnaryOperation{LX, LY, LZ, O, A, I, G, T} <: AbstractOperation{LX, LY, LZ, G, T} + op :: O + arg :: A + ▶ :: I + grid :: G @doc """ - UnaryOperation{X, Y, Z}(op, arg, ▶, grid) + UnaryOperation{LX, LY, LZ}(op, arg, ▶, grid) Returns an abstract `UnaryOperation` representing the action of `op` on `arg`, and subsequent interpolation by `▶` on `grid`. """ - function UnaryOperation{X, Y, Z}(op::O, arg::A, ▶::I, arch::R, grid::G) where {X, Y, Z, O, A, I, R, G} + function UnaryOperation{LX, LY, LZ}(op::O, arg::A, ▶::I, grid::G) where {LX, LY, LZ, O, A, I, G} T = eltype(grid) - return new{X, Y, Z, O, A, I, R, G, T}(op, arg, ▶, arch, grid) + return new{LX, LY, LZ, O, A, I, G, T}(op, arg, ▶, grid) end end @@ -29,8 +28,7 @@ end result from `Larg` to `L`.""" function _unary_operation(L, operator, arg, Larg, grid) ▶ = interpolation_operator(Larg, L) - arch = architecture(arg) - return UnaryOperation{L[1], L[2], L[3]}(operator, arg, ▶, arch, grid) + return UnaryOperation{L[1], L[2], L[3]}(operator, arg, ▶, grid) end # Recompute location of unary operation @@ -113,12 +111,6 @@ macro unary(ops...) return expr end -##### -##### Architecture inference for UnaryOperation -##### - -architecture(υ::UnaryOperation) = υ.architecture - ##### ##### Nested computations ##### @@ -130,7 +122,9 @@ compute_at!(υ::UnaryOperation, time) = compute_at!(υ.arg, time) ##### "Adapt `UnaryOperation` to work on the GPU via CUDAnative and CUDAdrv." -Adapt.adapt_structure(to, unary::UnaryOperation{X, Y, Z}) where {X, Y, Z} = - UnaryOperation{X, Y, Z}(Adapt.adapt(to, unary.op), Adapt.adapt(to, unary.arg), - Adapt.adapt(to, unary.▶), nothing, Adapt.adapt(to, unary.grid)) +Adapt.adapt_structure(to, unary::UnaryOperation{LX, LY, LZ}) where {LX, LY, LZ} = + UnaryOperation{LX, LY, LZ}(Adapt.adapt(to, unary.op), + Adapt.adapt(to, unary.arg), + Adapt.adapt(to, unary.▶), + Adapt.adapt(to, unary.grid)) diff --git a/src/Architectures.jl b/src/Architectures.jl index 4844acb684..929e38d082 100644 --- a/src/Architectures.jl +++ b/src/Architectures.jl @@ -39,7 +39,9 @@ Run Oceananigans on a single NVIDIA CUDA GPU. """ struct GPU <: AbstractArchitecture end -## All extension of these methods is done in Distributed.jl +##### +##### These methods are extended in Distributed.jl +##### device(::CPU) = KernelAbstractions.CPU() device(::GPU) = CUDAKernels.CUDADevice() @@ -49,6 +51,14 @@ architecture(::Number) = nothing architecture(::Array) = CPU() architecture(::CuArray) = GPU() +""" + child_architecture(arch) + +Return `arch`itecture of child processes. +On single-process, non-distributed systems, return `arch`. +""" +child_architecture(arch) = arch + array_type(::CPU) = Array array_type(::GPU) = CuArray diff --git a/src/CubedSpheres/cubed_sphere_faces.jl b/src/CubedSpheres/cubed_sphere_faces.jl index 1938861b1a..9087c80465 100644 --- a/src/CubedSpheres/cubed_sphere_faces.jl +++ b/src/CubedSpheres/cubed_sphere_faces.jl @@ -35,11 +35,11 @@ const CubedSphereData = CubedSphereFaces{<:OffsetArray} const ImmersedConformalCubedSphereFaceGrid = ImmersedBoundaryGrid{FT, TX, TY, TZ, <:ConformalCubedSphereFaceGrid} where {FT, TX, TY, TZ} # CubedSphereFaceField: -const NonImmersedCubedSphereFaceField = AbstractField{X, Y, Z, A, <:ConformalCubedSphereFaceGrid} where {X, Y, Z, A} -const ImmersedCubedSphereFaceField = AbstractField{X, Y, Z, A, <:ImmersedConformalCubedSphereFaceGrid} where {X, Y, Z, A} +const NonImmersedCubedSphereFaceField = AbstractField{LX, LY, LZ, <:ConformalCubedSphereFaceGrid} where {LX, LY, LZ, A} +const ImmersedCubedSphereFaceField = AbstractField{LX, LY, LZ, <:ImmersedConformalCubedSphereFaceGrid} where {LX, LY, LZ, A} -const CubedSphereFaceField = Union{NonImmersedCubedSphereFaceField{X, Y, Z, A}, - ImmersedCubedSphereFaceField{X, Y, Z, A}} where {X, Y, Z, A} +const CubedSphereFaceField = Union{NonImmersedCubedSphereFaceField{LX, LY, LZ, A}, + ImmersedCubedSphereFaceField{LX, LY, LZ, A}} where {LX, LY, LZ, A} # CubedSphereField @@ -48,18 +48,18 @@ const CubedSphereField{LX, LY, LZ, A} = Union{Field{LX, LY, LZ, <:Nothing, A, <:ConformalCubedSphereGrid}, Field{LX, LY, LZ, <:AbstractOperation, A, <:ConformalCubedSphereGrid}} -const CubedSphereAbstractField{LX, LY, LZ, A} = AbstractField{LX, LY, LZ, A, <:ConformalCubedSphereGrid} +const CubedSphereAbstractField{LX, LY, LZ} = AbstractField{LX, LY, LZ, <:ConformalCubedSphereGrid} -const AbstractCubedSphereField{LX, LY, LZ, A} = - Union{CubedSphereAbstractField{LX, LY, LZ, A}, - CubedSphereField{LX, LY, LZ, A}} +const AbstractCubedSphereField{LX, LY, LZ} = + Union{CubedSphereAbstractField{LX, LY, LZ}, + CubedSphereField{LX, LY, LZ}} ##### ##### new data ##### -function new_data(FT, grid::ConformalCubedSphereGrid, (X, Y, Z)) - faces = Tuple(new_data(FT, face_grid, (X, Y, Z)) for face_grid in grid.faces) +function new_data(FT, grid::ConformalCubedSphereGrid, (LX, LY, LZ)) + faces = Tuple(new_data(FT, face_grid, (LX, LY, LZ)) for face_grid in grid.faces) return CubedSphereFaces{typeof(faces[1]), typeof(faces)}(faces) end @@ -67,11 +67,11 @@ end ##### FieldBoundaryConditions ##### -function FieldBoundaryConditions(grid::ConformalCubedSphereGrid, (X, Y, Z); user_defined_bcs...) +function FieldBoundaryConditions(grid::ConformalCubedSphereGrid, (LX, LY, LZ); user_defined_bcs...) faces = Tuple( inject_cubed_sphere_exchange_boundary_conditions( - FieldBoundaryConditions(face_grid, (X, Y, Z); user_defined_bcs...), + FieldBoundaryConditions(face_grid, (LX, LY, LZ); user_defined_bcs...), face_index, grid.face_connectivity ) diff --git a/src/Distributed/multi_architectures.jl b/src/Distributed/multi_architectures.jl index 32fc72962f..f591e32c68 100644 --- a/src/Distributed/multi_architectures.jl +++ b/src/Distributed/multi_architectures.jl @@ -1,7 +1,7 @@ using Oceananigans.Architectures using Oceananigans.Grids: topology, validate_tupled_argument -import Oceananigans.Architectures: device, device_event, arch_array, array_type +import Oceananigans.Architectures: device, device_event, arch_array, array_type, child_architecture import Oceananigans.Grids: zeros struct MultiArch{A, R, I, ρ, C, γ} <: AbstractMultiArchitecture @@ -46,12 +46,11 @@ function MultiArch(child_architecture = CPU(); topology = (Periodic, Periodic, P return MultiArch{A, R, I, ρ, C, γ}(child_architecture, local_rank, local_index, ranks, local_connectivity, communicator) end -child_architecture(arch::MultiArch) = arch.child_architecture -child_architecture(::CPU) = CPU() -child_architecture(::GPU) = GPU() - -# Extending architecture specific methods +##### +##### All the architectures +##### +child_architecture(arch::MultiArch) = arch.child_architecture device(arch::AbstractMultiArchitecture) = device(child_architecture(arch)) device_event(arch::AbstractMultiArchitecture) = device_event(child_architecture(arch)) arch_array(arch::AbstractMultiArchitecture, A) = arch_array(child_architecture(arch), A) diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl index 9bd990d1e4..c617803eff 100644 --- a/src/Fields/Fields.jl +++ b/src/Fields/Fields.jl @@ -1,7 +1,7 @@ module Fields export Face, Center -export AbstractField, Field, Average, Integral, Reduction +export AbstractField, Field, Average, Integral, Reduction, field export CenterField, XFaceField, YFaceField, ZFaceField export BackgroundField export interior, data, xnode, ynode, znode, location @@ -14,16 +14,11 @@ using Oceananigans.Grids using Oceananigans.BoundaryConditions include("abstract_field.jl") -include("field.jl") -include("field_reductions.jl") include("zero_field.jl") -include("averaged_field.jl") -include("kernel_computed_field.jl") include("function_field.jl") +include("field.jl") +include("field_reductions.jl") include("regridding_fields.jl") -include("set!.jl") -include("tracer_names.jl") -include("validate_field_tuple_grid.jl") include("field_tuples.jl") include("background_fields.jl") include("interpolate.jl") @@ -31,4 +26,23 @@ include("field_slicer.jl") include("show_fields.jl") include("broadcasting_abstract_fields.jl") +""" + field(loc, a, grid) + +Build a field from `a` at `loc` and on `grid`. +""" +function field(loc, a::Array, grid) + f = Field(loc, grid) + f .= a + return f end + +field(loc, a::Function, grid) = FunctionField(loc, a, grid) + +function field(loc, f::Field, grid) + loc === location(f) && grid === f.grid && return f + error("Cannot construct field at $loc and on $grid from $f") +end + +include("set!.jl") + diff --git a/src/Fields/abstract_field.jl b/src/Fields/abstract_field.jl index b80cdc735a..62b0fec133 100644 --- a/src/Fields/abstract_field.jl +++ b/src/Fields/abstract_field.jl @@ -20,62 +20,44 @@ const ArchOrNothing = Union{AbstractArchitecture, Nothing} const GridOrNothing = Union{AbstractGrid, Nothing} """ - AbstractField{X, Y, Z, A, G, T, N} + AbstractField{LX, LY, LZ, G, T, N} -Abstract supertype for fields located at `(X, Y, Z)` on architecture `A` +Abstract supertype for fields located at `(LX, LY, LZ)` and defined on a grid `G` with eltype `T` and `N` dimensions. """ -abstract type AbstractField{X, Y, Z, A <: ArchOrNothing, G <: GridOrNothing, T, N} <: AbstractArray{T, N} end -abstract type AbstractOperation{X, Y, Z, A, G, T} <: AbstractField{X, Y, Z, A, G, T, 3} end +abstract type AbstractField{LX, LY, LZ, G <: GridOrNothing, T, N} <: AbstractArray{T, N} end Base.IndexStyle(::AbstractField) = IndexCartesian() -function validate_field_data(loc, data, grid) - Tx, Ty, Tz = total_size(loc, grid) - - if size(data) != (Tx, Ty, Tz) - LX, LY, LZ = loc - e = "Cannot construct field at ($LX, $LY, $LZ) with size(data)=$(size(data)). " * - "`data` must have size ($Tx, $Ty, $Tz)." - throw(ArgumentError(e)) - end - - return nothing -end - # Endpoint for recursive `datatuple` function: +data(a) = a @inline datatuple(obj::AbstractField) = data(obj) - ##### ##### AbstractField functionality ##### -@inline location(a) = (Nothing, Nothing, Nothing) - "Returns the location `(LX, LY, LZ)` of an `AbstractField{LX, LY, LZ}`." @inline location(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = (LX, LY, LZ) # note no instantiation -@inline location(f, i) = location(f)[i] - -@inline instantiated_location(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = (LX(), LY(), LZ()) +@inline location(f::AbstractField, i) = location(f)[i] +instantiated_location(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = (LX(), LY(), LZ()) "Returns the architecture of on which `f` is defined." architecture(f::AbstractField) = architecture(f.grid) "Returns the topology of a fields' `grid`." -@inline topology(f, args...) = topology(f.grid, args...) +topology(f::AbstractField, args...) = topology(f.grid, args...) """ size(f::AbstractField) -Returns the size of an `AbstractField{X, Y, Z}` located at `X, Y, Z`. +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) - -"Returns the length of a field's `data`." -@inline Base.length(f::AbstractField) = prod(size(f)) +Base.length(f::AbstractField) = prod(size(f)) +Base.parent(f::AbstractField) = f """ total_size(field::AbstractField) @@ -85,26 +67,7 @@ both interior points and halo points. """ total_size(f::AbstractField) = total_size(location(f), f.grid) -##### -##### Accessing wrapped arrays -##### - -"Returns `f.data` for `f::Field` or `f` for `f::AbstractArray." -data(a) = nothing # fallback -cpudata(a) = data(a) - -cpudata(f::AbstractField{X, Y, Z, <:GPU}) where {X, Y, Z} = - offset_data(Array(parent(f)), f.grid, location(f)) - -"Returns `f.data.parent` for `f::Field`." -@inline Base.parent(f::AbstractField) = parent(data(f)) - -@inline interior(f::AbstractField) = f - -@inline interior_copy(f::AbstractField{X, Y, Z}) where {X, Y, Z} = - parent(f)[interior_parent_indices(X, topology(f, 1), f.grid.Nx, f.grid.Hx), - interior_parent_indices(Y, topology(f, 2), f.grid.Ny, f.grid.Hy), - interior_parent_indices(Z, topology(f, 3), f.grid.Nz, f.grid.Hz)] +interior(f::AbstractField) = f ##### ##### Coordinates of fields diff --git a/src/Fields/averaged_field.jl b/src/Fields/averaged_field.jl deleted file mode 100644 index 5ee0b8f6cf..0000000000 --- a/src/Fields/averaged_field.jl +++ /dev/null @@ -1,120 +0,0 @@ -using Adapt -using Statistics -using Oceananigans.Grids -using Oceananigans.Grids: interior_parent_indices - -struct AveragedField{X, Y, Z, S, A, D, G, T, N, O} <: AbstractField{X, Y, Z, A, G, T, N} - data :: D - architecture :: A - grid :: G - dims :: NTuple{N, Int} - operand :: O - status :: S - - function AveragedField{X, Y, Z}(data::D, arch::A, grid::G, dims, operand::O; - recompute_safely=true) where {X, Y, Z, D, A, G, O} - - dims = validate_reduced_dims(dims) - validate_reduced_locations(X, Y, Z, dims) - validate_field_data(X, Y, Z, data, grid) - - status = recompute_safely ? nothing : FieldStatus(0.0) - - S = typeof(status) - N = length(dims) - T = eltype(grid) - - return new{X, Y, Z, S, A, D, G, T, N, O}(data, arch, grid, dims, operand, status) - end - - function AveragedField{X, Y, Z}(data::D, arch::A, grid::G, dims, operand::O, status::S) where {X, Y, Z, D, A, G, O, S} - return new{X, Y, Z, S, A, D, G, eltype(grid), length(dims), O}(data, arch, grid, dims, operand, status) - end -end - -""" - AveragedField(operand::AbstractField; dims, data=nothing, recompute_safely=false) - -Returns an AveragedField averaged over `dims`. `dims` is a tuple of integers indicating -spatial dimensions; in a Cartesian coordinate system, `1=x, `2=y`, and `3=z`. - -Arguments -========= - -- `dims`: Tuple of integers specifying the dimensions to average `operand`. - A single integer is also accepted for averaging over a single dimension. - -- `data`: An `OffsetArray` for storing averaged data. - Useful if carefully managing memory allocation. - If unspecified, `data` is created by `Oceananigans.Grids.new_data`. - -- `recompute_safely`: A boolean that's relevant only if the `AveragedField` is used - within another computation. If `recompute_safely=false`, - `AveragedField` will *not* be recomputed before computing any dependent - computations if `AveragedField.status` is consistent with the current state of the simulation. - If `recompute_safely=true`, `AveragedField` is always recomputed - before performing a dependent computation. - -Examples -======= - -```julia -julia> using Oceananigans - -julia> grid = RectilinearGrid(size=(2, 2, 2), x=(0, 1), y=(0, 1), z=(0, 1)); - -julia> c = CenterField(CPU(), grid); - -julia> C_xy = AveragedField(c, dims=(1, 2)) # average over x, y -AveragedField over dims=(1, 2) located at (⋅, ⋅, Center) of Field located at (Center, Center, Center) -├── data: OffsetArrays.OffsetArray{Float64, 3, Array{Float64, 3}}, size: (1, 1, 2) -├── grid: RectilinearGrid{Float64, Periodic, Periodic, Bounded}(Nx=2, Ny=2, Nz=2) -├── dims: (1, 2) -├── operand: Field located at (Center, Center, Center) -└── status: time=0.0 - -julia> C_z = AveragedField(c, dims=3) # averaged over z -AveragedField over dims=(3,) located at (Center, Center, ⋅) of Field located at (Center, Center, Center) -├── data: OffsetArrays.OffsetArray{Float64, 3, Array{Float64, 3}}, size: (2, 2, 1) -├── grid: RectilinearGrid{Float64, Periodic, Periodic, Bounded}(Nx=2, Ny=2, Nz=2) -├── dims: (3,) -├── operand: Field located at (Center, Center, Center) -└── status: time=0.0 -``` -""" -function AveragedField(operand::AbstractField; dims, data=nothing, recompute_safely=true) - - arch = architecture(operand) - loc = reduced_location(location(operand), dims=dims) - grid = operand.grid - - if isnothing(data) - data = new_data(arch, grid, loc) - recompute_safely = false - end - - return AveragedField{loc[1], loc[2], loc[3]}(data, arch, grid, dims, operand, - recompute_safely=recompute_safely) -end - -""" - compute!(avg::AveragedField, time=nothing) - -Compute the average of `avg.operand` and store the result in `avg.data`. -""" -function compute!(avg::AveragedField, time=nothing) - compute_at!(avg.operand, time) - mean!(avg, avg.operand) - return nothing -end - -compute_at!(avg::AveragedField{X, Y, Z, <:FieldStatus}, time) where {X, Y, Z} = - conditional_compute!(avg, time) - -##### -##### Adapt -##### - -Adapt.adapt_structure(to, averaged_field::AveragedField{X, Y, Z}) where {X, Y, Z} = - AveragedField{X, Y, Z}(Adapt.adapt(to, averaged_field.data), nothing, - nothing, averaged_field.dims, nothing, nothing) diff --git a/src/Fields/background_fields.jl b/src/Fields/background_fields.jl index a4c73e1f3f..a4136b25eb 100644 --- a/src/Fields/background_fields.jl +++ b/src/Fields/background_fields.jl @@ -1,3 +1,5 @@ +# TODO: This code belongs in the Models module + function BackgroundVelocityFields(bg, grid, clock) u = :u ∈ keys(bg) ? regularize_background_field(Face, Center, Center, bg[:u], grid, clock) : ZeroField() v = :v ∈ keys(bg) ? regularize_background_field(Center, Face, Center, bg[:v], grid, clock) : ZeroField() @@ -56,15 +58,15 @@ func(x, y, z, t, parameters) """ BackgroundField(func; parameters=nothing) = BackgroundField(func, parameters) -regularize_background_field(X, Y, Z, f::BackgroundField{<:Function}, grid, clock) = - FunctionField{X, Y, Z}(f.func, grid; clock=clock, parameters=f.parameters) +regularize_background_field(LX, LY, LZ, f::BackgroundField{<:Function}, grid, clock) = + FunctionField{LX, LY, LZ}(f.func, grid; clock=clock, parameters=f.parameters) -regularize_background_field(X, Y, Z, func::Function, grid, clock) = - FunctionField{X, Y, Z}(func, grid; clock=clock) +regularize_background_field(LX, LY, LZ, func::Function, grid, clock) = + FunctionField{LX, LY, LZ}(func, grid; clock=clock) -function regularize_background_field(X, Y, Z, field::AbstractField, grid, clock) - if location(field) != (X, Y, Z) - throw(ArgumentError("Cannot use field at $(location(field)) as a background field at $((X, Y, Z))")) +function regularize_background_field(LX, LY, LZ, field::AbstractField, grid, clock) + if location(field) != (LX, LY, LZ) + throw(ArgumentError("Cannot use field at $(location(field)) as a background field at $((LX, LY, LZ))")) end return field diff --git a/src/Fields/field.jl b/src/Fields/field.jl index e7b6cbaa49..8ad90cc49d 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -6,7 +6,7 @@ using Base: @propagate_inbounds import Oceananigans.BoundaryConditions: fill_halo_regions! -struct Field{LX, LY, LZ, O, A, G, T, D, B, S} <: AbstractField{LX, LY, LZ, A, G, T, 3} +struct Field{LX, LY, LZ, O, G, T, D, B, S} <: AbstractField{LX, LY, LZ, G, T, 3} grid :: G data :: D boundary_conditions :: B @@ -14,14 +14,25 @@ struct Field{LX, LY, LZ, O, A, G, T, D, B, S} <: AbstractField{LX, LY, LZ, A, G, status :: S # Inner constructor that does not validate _anything_! - function Field{LX, LY, LZ}(grid::G, data::D, bcs::B, op::O, status::S, - arch::A=architecture(grid)) where {LX, LY, LZ, - G, D, B, O, S, A} + function Field{LX, LY, LZ}(grid::G, data::D, bcs::B, op::O, status::S) where {LX, LY, LZ, G, D, B, O, S} T = eltype(data) - return new{LX, LY, LZ, O, A, G, T, D, B, S}(grid, data, bcs, op, status) + return new{LX, LY, LZ, O, G, T, D, B, S}(grid, data, bcs, op, status) end end +function validate_field_data(loc, data, grid) + Tx, Ty, Tz = total_size(loc, grid) + + if size(data) != (Tx, Ty, Tz) + LX, LY, LZ = loc + e = "Cannot construct field at ($LX, $LY, $LZ) with size(data)=$(size(data)). " * + "`data` must have size ($Tx, $Ty, $Tz)." + throw(ArgumentError(e)) + end + + return nothing +end + # Common outer constructor for all field flavors that validates data and boundary conditions function Field(loc::Tuple, grid::AbstractGrid, data, bcs, op, status) validate_field_data(loc, data, grid) @@ -106,18 +117,27 @@ function interior(f::Field) return view(parent(f), ii, jj, kk) end +interior_copy(f::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = + parent(f)[interior_parent_indices(X, topology(f, 1), f.grid.Nx, f.grid.Hx), + interior_parent_indices(Y, topology(f, 2), f.grid.Ny, f.grid.Hy), + interior_parent_indices(Z, topology(f, 3), f.grid.Nz, f.grid.Hz)] + # Don't use axes(f) to checkbounds; use axes(f.data) -Base.checkbounds(f::Field, I...) = Base.checkbounds(data(f), I...) -@propagate_inbounds Base.getindex(f::Field, inds...) = getindex(data(f), inds...) +Base.checkbounds(f::Field, I...) = Base.checkbounds(f.data, I...) + +@propagate_inbounds Base.getindex(f::Field, inds...) = getindex(f.data, inds...) @propagate_inbounds Base.getindex(f::Field, i::Int) = parent(f)[i] -@propagate_inbounds Base.setindex!(f::Field, val, i, j, k) = setindex!(data(f), val, i, j, k) -@propagate_inbounds Base.lastindex(f::Field) = lastindex(data(f)) -@propagate_inbounds Base.lastindex(f::Field, dim) = lastindex(data(f), dim) -Base.fill!(f::Field, val) = fill!(parent(f), val) +@propagate_inbounds Base.setindex!(f::Field, val, i, j, k) = setindex!(f.data, val, i, j, k) +@propagate_inbounds Base.lastindex(f::Field) = lastindex(f.data) +@propagate_inbounds Base.lastindex(f::Field, dim) = lastindex(f.data, dim) +Base.fill!(f::Field, val) = fill!(parent(f), val) Base.isapprox(ϕ::Field, ψ::Field; kw...) = isapprox(interior(ϕ), interior(ψ); kw...) +Base.parent(f::Field) = parent(f.data) +Adapt.adapt_structure(to, f::Field) = Adapt.adapt(to, f.data) -Adapt.adapt_structure(to, field::Field) = Adapt.adapt(to, field.data) +data(f::Field) = f.data +cpudata(a::Field) = arch_array(CPU(), a.data) ##### ##### Special constructors for tracers and velocity fields @@ -179,10 +199,7 @@ macro compute(def) return expr end -##### -##### Conditional computation -##### - +# Computation "status" for avoiding unnecessary recomputation mutable struct FieldStatus{T} time :: T end @@ -301,10 +318,10 @@ end ##### Field reductions ##### -# Risky to use these without tests. Docs would also be nice. +# TODO: needs test Statistics.dot(a::Field, b::Field) = mapreduce((x, y) -> x * y, +, interior(a), interior(b)) -# TODO: In-place allocations with function mappings need to be fixed in Julia Base... +# TODO: in-place allocations with function mappings need to be fixed in Julia Base... const SumReduction = typeof(Base.sum!) const ProdReduction = typeof(Base.prod!) const MaximumReduction = typeof(Base.maximum!) diff --git a/src/Fields/field_slicer.jl b/src/Fields/field_slicer.jl index f9a249084a..d0eb3eca12 100644 --- a/src/Fields/field_slicer.jl +++ b/src/Fields/field_slicer.jl @@ -2,6 +2,8 @@ import Oceananigans: short_show using Oceananigans.Grids: Face, Bounded, topology using Oceananigans.Fields: location +# TODO: this code probably belongs in the OutputWriters module + """ struct FieldSlicer{I, J, K, W} diff --git a/src/Fields/field_tuples.jl b/src/Fields/field_tuples.jl index 5c4016009c..c1f24c0ae6 100644 --- a/src/Fields/field_tuples.jl +++ b/src/Fields/field_tuples.jl @@ -1,5 +1,48 @@ using Oceananigans.BoundaryConditions: FieldBoundaryConditions, regularize_field_boundary_conditions +# TODO: This code belongs in the Models module + +##### +##### Tracer names +##### + +"Returns true if the first three elements of `names` are `(:u, :v, :w)`." +has_velocities(names) = :u == names[1] && :v == names[2] && :w == names[3] + +# Tuples of length 0-2 cannot contain velocity fields +has_velocities(::Tuple{}) = false +has_velocities(::Tuple{X}) where X = false +has_velocities(::Tuple{X, Y}) where {X, Y} = false + +tracernames(::Nothing) = () +tracernames(name::Symbol) = tuple(name) +tracernames(names::NTuple{N, Symbol}) where N = has_velocities(names) ? names[4:end] : names +tracernames(::NamedTuple{names}) where names = tracernames(names) + +##### +##### Validation +##### + +validate_field_grid(grid, field) = grid === field.grid + +validate_field_grid(grid, field_tuple::NamedTuple) = + all(validate_field_grid(grid, field) for field in field_tuple) + +""" + validate_field_tuple_grid(tuple_name, field_tuple, arch, grid, bcs) + +Validates the grids associated with grids in the (possibly nested) `field_tuple`, +and returns `field_tuple` if validation succeeds. +""" +function validate_field_tuple_grid(tuple_name, field_tuple, grid) + + all(validate_field_grid(grid, field) for field in field_tuple) || + throw(ArgumentError("Model grid and $tuple_name grid are not identical! " * + "Check that the grid used to construct $tuple_name has the correct halo size.")) + + return nothing +end + ##### ##### Velocity fields tuples ##### diff --git a/src/Fields/function_field.jl b/src/Fields/function_field.jl index c21ea9d904..70320b05ed 100644 --- a/src/Fields/function_field.jl +++ b/src/Fields/function_field.jl @@ -1,37 +1,36 @@ - -struct FunctionField{X, Y, Z, C, P, F, G, T} <: AbstractField{X, Y, Z, Nothing, G, T, 3} +struct FunctionField{LX, LY, LZ, C, P, F, G, T} <: AbstractField{LX, LY, LZ, G, T, 3} func :: F grid :: G clock :: C parameters :: P @doc """ - FunctionField{X, Y, Z}(func, grid; clock=nothing, parameters=nothing) where {X, Y, Z} + FunctionField{LX, LY, LZ}(func, grid; clock=nothing, parameters=nothing) where {LX, LY, LZ} - Returns a `FunctionField` on `grid` and at location `X, Y, Z`. + Returns a `FunctionField` on `grid` and at location `LX, LY, LZ`. If `clock` is not specified, then `func` must be a function with signature `func(x, y, z)`. If clock is specified, `func` must be a function with signature `func(x, y, z, t)`, where `t` is internally determined from `clock.time`. - A `FunctionField` will return the result of `func(x, y, z [, t])` at `X, Y, Z` on + A `FunctionField` will return the result of `func(x, y, z [, t])` at `LX, LY, LZ` on `grid` when indexed at `i, j, k`. """ - function FunctionField{X, Y, Z}(func::F, grid::G; clock::C=nothing, parameters::P=nothing) where {X, Y, Z, F, G, C, P} + function FunctionField{LX, LY, LZ}(func::F, grid::G; clock::C=nothing, parameters::P=nothing) where {LX, LY, LZ, F, G, C, P} T = eltype(grid) - return new{X, Y, Z, C, P, F, G, T}(func, grid, clock, parameters) + return new{LX, LY, LZ, C, P, F, G, T}(func, grid, clock, parameters) end @doc """ - FunctionField{X, Y, Z}(func::FunctionField, grid; clock) where {X, Y, Z} + FunctionField{LX, LY, LZ}(func::FunctionField, grid; clock) where {LX, LY, LZ} - Adds `clock` to an existing `FunctionField` and relocates it to `(X, Y, Z)` on `grid`. + Adds `clock` to an existing `FunctionField` and relocates it to `(LX, LY, LZ)` on `grid`. """ - function FunctionField{X, Y, Z}(f::FunctionField, grid::G; clock::C=nothing) where {X, Y, Z, G, C} + function FunctionField{LX, LY, LZ}(f::FunctionField, grid::G; clock::C=nothing) where {LX, LY, LZ, G, C} P = typeof(f.parameters) T = eltype(grid) F = typeof(f.func) - return new{X, Y, Z, C, P, F, G, T}(f.func, grid, clock, f.parameters) + return new{LX, LY, LZ, C, P, F, G, T}(f.func, grid, clock, f.parameters) end end @@ -42,15 +41,11 @@ fieldify(L, a::Function, grid) = FunctionField(L, a, grid) """ FunctionField(L::Tuple, func, grid) -Returns a stationary `FunctionField` on `grid` and at location `L = (X, Y, Z)`, +Returns a stationary `FunctionField` on `grid` and at location `L = (LX, LY, LZ)`, where `func` is callable with signature `func(x, y, z)`. """ FunctionField(L::Tuple, func, grid) = FunctionField{L[1], L[2], L[3]}(func, grid) -# Ordinary functions needed for fields -architecture(f::FunctionField) = nothing -Base.parent(f::FunctionField) = f - # Various possibilities for calling FunctionField.func: @inline call_func(clock, parameters, func, x, y, z) = func(x, y, z, clock.time, parameters) @inline call_func(::Nothing, parameters, func, x, y, z) = func(x, y, z, parameters) @@ -66,8 +61,8 @@ Base.parent(f::FunctionField) = f @inline (f::FunctionField)(x...) = call_func(f.clock, f.parameters, f.func, x...) -Adapt.adapt_structure(to, f::FunctionField{X, Y, Z}) where {X, Y, Z} = - FunctionField{X, Y, Z}(Adapt.adapt(to, f.func), +Adapt.adapt_structure(to, f::FunctionField{LX, LY, LZ}) where {LX, LY, LZ} = + FunctionField{LX, LY, LZ}(Adapt.adapt(to, f.func), Adapt.adapt(to, f.grid), clock = Adapt.adapt(to, f.clock), parameters = Adapt.adapt(to, f.parameters)) @@ -78,3 +73,4 @@ Base.show(io::IO, field::FunctionField) = "├── grid: $(short_show(field.grid))\n", "├── clock: $(short_show(field.clock))\n", "└── parameters: $(field.parameters)") + diff --git a/src/Fields/kernel_computed_field.jl b/src/Fields/kernel_computed_field.jl deleted file mode 100644 index 96b55f43e0..0000000000 --- a/src/Fields/kernel_computed_field.jl +++ /dev/null @@ -1,127 +0,0 @@ -using Oceananigans: AbstractModel -using Oceananigans.Grids -using Oceananigans.Utils: tupleit - -struct KernelComputedField{X, Y, Z, A, S, D, G, T, K, B, F, P} <: AbstractField{X, Y, Z, A, G, T, 3} - data :: D - architecture :: A - grid :: G - kernel :: K - boundary_conditions :: B - computed_dependencies :: F - parameters :: P - status :: S - - function KernelComputedField{X, Y, Z}(kernel::K, arch::A, grid::G; - boundary_conditions::B = FieldBoundaryConditions(grid, (X, Y, Z)), - computed_dependencies = (), - parameters::P = nothing, - data::D = new_data(arch, grid, (X, Y, Z)), - recompute_safely = true) where {X, Y, Z, A, D, B, G, K, P} - - computed_dependencies = tupleit(computed_dependencies) - - # Use FieldStatus if we want to avoid always recomputing - status = recompute_safely ? nothing : FieldStatus(0.0) - - S = typeof(status) - F = typeof(computed_dependencies) - T = eltype(grid) - - return new{X, Y, Z, A, S, D, G, T, K, B, F, P}( - data, arch, grid, kernel, boundary_conditions, computed_dependencies, parameters, status) - end -end - -""" - KernelComputedField(X, Y, Z, kernel, model; - boundary_conditions = FieldBoundaryConditions(grid, (X, Y, Z)), - computed_dependencies = (), - parameters = nothing, - data = nothing, - recompute_safely = true) - -Builds a `KernelComputedField` at `X, Y, Z` computed with `kernel` and `model.architecture` and `model.grid`, with `boundary_conditions`. - -`computed_dependencies` are an iterable of `AbstractField`s or other objects on which `compute!` is called prior to launching `kernel`. - -`data` is a three-dimensional `OffsetArray` of scratch space where the kernel computation is stored. - -If `data=nothing` (the default) then additional memory will be allocated to store the `data` of `KernelComputedField`. - -If `isnothing(parameters)`, `kernel` is launched with the function signature - -`kernel(data, grid, computed_dependencies...)` - -Otherwise, `kernel` is launched with the function signature - -`kernel(data, grid, computed_dependencies..., parameters)` - -`recompute_safely` (default: `true`) determines whether the `KernelComputedField` is "recomputed" if embedded in the expression -tree of another operation. - - If `recompute_safely=true`, the `KernelComputedField` is always recomputed. - - If `recompute_safely=false`, the `KernelComputedField` will not be recomputed if its status is up-to-date. - -Example -======= - -```julia -using KernelAbstractions: @index, @kernel -using Oceananigans.Fields: AveragedField, KernelComputedField, compute! -using Oceananigans.Grids: Center, Face - -@inline ψ′²(i, j, k, grid, ψ, Ψ) = @inbounds (ψ[i, j, k] - Ψ[i, j, k])^2 -@inline ψ′²(i, j, k, grid, ψ, Ψ::Number) = @inbounds (ψ[i, j, k] - Ψ)^2 - -@kernel function compute_variance!(var, grid, ϕ, Φ) - i, j, k = @index(Global, NTuple) - - @inbounds var[i, j, k] = ψ′²(i, j, k, grid, ϕ, Φ) -end - -u, v, w = model.velocities - -U = AveragedField(u, dims=(1, 2)) -V = AveragedField(v, dims=(1, 2)) - -u′² = KernelComputedField(Face, Center, Center, compute_variance!, model; computed_dependencies=(u, U)) -v′² = KernelComputedField(Center, Face, Center, compute_variance!, model; computed_dependencies=(v, V)) -w′² = KernelComputedField(Center, Center, Face, compute_variance!, model; computed_dependencies=(w, 0)) - -compute!(u′²) -compute!(v′²) -compute!(w′²) -``` -""" -KernelComputedField(X, Y, Z, kernel, model::AbstractModel; kwargs...) = - KernelComputedField{X, Y, Z}(kernel, model.architecture, model.grid; kwargs...) - -KernelComputedField(X, Y, Z, kernel, arch::AbstractArchitecture, grid::AbstractGrid; kwargs...) = - KernelComputedField{X, Y, Z}(kernel, arch, grid; kwargs...) - -function compute!(kcf::KernelComputedField{X, Y, Z}) where {X, Y, Z} - - for dependency in kcf.computed_dependencies - compute!(dependency) - end - - arch = architecture(kcf) - - args = isnothing(kcf.parameters) ? - tuple(kcf.data, kcf.grid, kcf.computed_dependencies...) : - tuple(kcf.data, kcf.grid, kcf.computed_dependencies..., kcf.parameters) - - event = launch!(arch, kcf.grid, :xyz, kcf.kernel, args...; - location=(X, Y, Z), include_right_boundaries=true) - - wait(device(arch), event) - - fill_halo_regions!(kcf, arch) - - return nothing -end - -compute!(field::KernelComputedField{X, Y, Z, <:FieldStatus}, time) where {X, Y, Z} = - conditional_compute!(field, time) - -Adapt.adapt_structure(to, kcf::KernelComputedField) = Adapt.adapt(to, kcf.data) diff --git a/src/Fields/set!.jl b/src/Fields/set!.jl index 0ccdd30c36..9cea051225 100644 --- a/src/Fields/set!.jl +++ b/src/Fields/set!.jl @@ -21,9 +21,9 @@ set!(u::AbstractField, v) = u .= v # fallback const CPUField = Field{LX, LY, LZ, O, <:CPU} where {LX, LY, LZ, O} """ Set the CPU field `u` data to the function `f(x, y, z)`. """ -function set!(u::CPUField, f::Function) - f_field = FunctionField(location(u), f, u.grid) - u .= f_field +function set!(u::Field, f::Union{Array, Function}) + f_field = field(location(u), f, u.grid) + set!(u, f_field) return nothing end @@ -35,22 +35,7 @@ const GPUField = Field{LX, LY, LZ, O, <:GPU} where {LX, LY, LZ, O} """ Set the GPU field `u` to the array or function `v`. """ function set!(u::GPUField, v::Union{Array, Function}) - cpu_grid = on_architecture(CPU(), u.grid) - v_field = similar(u, cpu_grid) - set!(v_field, v) - set!(u, v_field) return nothing end -##### -##### Setting fields to fields -##### - -""" Set the CPU field `u` data to the GPU field data of `v`. """ -set!(u::CPUField, v::GPUField) = copyto!(parent(u), parent(v)) - -""" Set the GPU field `u` data to the CPU field data of `v`. """ -set!(u::GPUField, v::CPUField) = copyto!(parent(u), parent(v)) - -set!(u::CPUField, v::CPUField) = parent(u) .= parent(v) -set!(u::GPUField, v::GPUField) = parent(u) .= parent(v) +set!(u::Field, v::Field) = copyto!(parent(u), parent(v)) diff --git a/src/Fields/show_fields.jl b/src/Fields/show_fields.jl index 9963e953b6..aa681e467d 100644 --- a/src/Fields/show_fields.jl +++ b/src/Fields/show_fields.jl @@ -4,18 +4,18 @@ location_str(::Type{Face}) = "Face" location_str(::Type{Center}) = "Center" location_str(::Type{Nothing}) = "⋅" -show_location(X, Y, Z) = "($(location_str(X)), $(location_str(Y)), $(location_str(Z)))" +show_location(LX, LY, LZ) = "($(location_str(X)), $(location_str(Y)), $(location_str(Z)))" -show_location(field::AbstractField{X, Y, Z}) where {X, Y, Z} = show_location(X, Y, Z) +show_location(field::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = show_location(LX, LY, LZ) short_show(m::Missing) = "$m" short_show(field::AbstractField) = string(typeof(field).name.wrapper, " located at ", show_location(field)) short_show(field::AveragedField) = string("AveragedField over dims=$(field.dims) located at ", show_location(field), " of ", short_show(field.operand)) -Base.show(io::IO, field::AbstractField{X, Y, Z, A}) where {X, Y, Z, A} = +Base.show(io::IO, field::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = print(io, "$(short_show(field))\n", - "├── architecture: $A\n", + "├── architecture: $(architecture(field))\n", "└── grid: $(short_show(field.grid))") Base.show(io::IO, field::Field) = diff --git a/src/Fields/tracer_names.jl b/src/Fields/tracer_names.jl deleted file mode 100644 index fa18d53d04..0000000000 --- a/src/Fields/tracer_names.jl +++ /dev/null @@ -1,12 +0,0 @@ -"Returns true if the first three elements of `names` are `(:u, :v, :w)`." -has_velocities(names) = :u == names[1] && :v == names[2] && :w == names[3] - -# Tuples of length 0-2 cannot contain velocity fields -has_velocities(::Tuple{}) = false -has_velocities(::Tuple{X}) where X = false -has_velocities(::Tuple{X, Y}) where {X, Y} = false - -tracernames(::Nothing) = () -tracernames(name::Symbol) = tuple(name) -tracernames(names::NTuple{N, Symbol}) where N = has_velocities(names) ? names[4:end] : names -tracernames(::NamedTuple{names}) where names = tracernames(names) diff --git a/src/Fields/validate_field_tuple_grid.jl b/src/Fields/validate_field_tuple_grid.jl deleted file mode 100644 index 7ab1ece469..0000000000 --- a/src/Fields/validate_field_tuple_grid.jl +++ /dev/null @@ -1,19 +0,0 @@ -validate_field_grid(grid, field) = grid === field.grid - -validate_field_grid(grid, field_tuple::NamedTuple) = - all(validate_field_grid(grid, field) for field in field_tuple) - -""" - validate_field_tuple_grid(tuple_name, field_tuple, arch, grid, bcs) - -Validates the grids associated with grids in the (possibly nested) `field_tuple`, -and returns `field_tuple` if validation succeeds. -""" -function validate_field_tuple_grid(tuple_name, field_tuple, grid) - - all(validate_field_grid(grid, field) for field in field_tuple) || - throw(ArgumentError("Model grid and $tuple_name grid are not identical! " * - "Check that the grid used to construct $tuple_name has the correct halo size.")) - - return nothing -end diff --git a/src/Fields/zero_field.jl b/src/Fields/zero_field.jl index 7b903c9925..834fc32adc 100644 --- a/src/Fields/zero_field.jl +++ b/src/Fields/zero_field.jl @@ -1,4 +1,6 @@ -struct ZeroField <: AbstractField{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, 3} end +struct ZeroField{N} <: AbstractField{Nothing, Nothing, Nothing, Nothing, Nothing, N} end -@inline Base.getindex(::ZeroField, i, j, k) = 0 +ZeroField() = ZeroField{3}() # default + +@inline Base.getindex(::ZeroField, ind...) = 0 diff --git a/src/OutputReaders/field_time_series.jl b/src/OutputReaders/field_time_series.jl index 614769e384..3d8b714278 100644 --- a/src/OutputReaders/field_time_series.jl +++ b/src/OutputReaders/field_time_series.jl @@ -14,7 +14,7 @@ import Oceananigans: short_show import Oceananigans.Fields: Field, set!, interior import Oceananigans.Architectures: architecture -struct FieldTimeSeries{LX, LY, LZ, K, A, T, D, G, B, χ} <: AbstractField{LX, LY, LZ, A, G, T, 4} +struct FieldTimeSeries{LX, LY, LZ, K, T, D, G, B, χ} <: AbstractField{LX, LY, LZ, G, T, 4} data :: D grid :: G boundary_conditions :: B @@ -22,8 +22,7 @@ struct FieldTimeSeries{LX, LY, LZ, K, A, T, D, G, B, χ} <: AbstractField{LX, LY function FieldTimeSeries{LX, LY, LZ, K}(data::D, grid::G, bcs::B, times::χ) where {LX, LY, LZ, K, D, G, B, χ} T = eltype(grid) - A = typeof(architecture(grid)) - return new{LX, LY, LZ, K, A, T, D, G, B, χ}(data, grid, bcs, times) + return new{LX, LY, LZ, K, T, D, G, B, χ}(data, grid, bcs, times) end end From ca77537cfd01f5aee73b1f61a027a8467333226e Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 10:26:34 -0700 Subject: [PATCH 097/140] End the Fields module --- src/Fields/Fields.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl index c617803eff..eb9182b3a6 100644 --- a/src/Fields/Fields.jl +++ b/src/Fields/Fields.jl @@ -46,3 +46,4 @@ end include("set!.jl") +end # module From 966db508bc5dc89f2193fb866f47020b68a45a46 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 10:43:35 -0700 Subject: [PATCH 098/140] Reorganize show and short_show for ReducedComputedFields --- .../metric_field_reductions.jl | 9 +++++++++ src/Fields/field_reductions.jl | 17 +++++++++++++++++ src/Fields/show_fields.jl | 10 +--------- 3 files changed, 27 insertions(+), 9 deletions(-) diff --git a/src/AbstractOperations/metric_field_reductions.jl b/src/AbstractOperations/metric_field_reductions.jl index 0005aaa055..65a18f761b 100644 --- a/src/AbstractOperations/metric_field_reductions.jl +++ b/src/AbstractOperations/metric_field_reductions.jl @@ -1,6 +1,7 @@ using Statistics: mean!, sum! import Oceananigans.Fields: Reduction +import Oceananigans: short_show ##### ##### Metric inference @@ -51,3 +52,11 @@ end Integral(field::AbstractField; dims=:) = Reduction(Integral(), field; dims) const IntegratedField = Field{<:Any, <:Any, <:Any, <:Reduction{<:Integral}} + +##### +##### show +##### + +short_show(r::Reduction{<:Average}) = string("Average of ", short_show(r.operand), " over dims ", r.dims) +short_show(r::Reduction{<:Integral}) = string("Integral of ", short_show(r.operand), " over dims ", r.dims) + diff --git a/src/Fields/field_reductions.jl b/src/Fields/field_reductions.jl index c7025c4f7e..b343f33ca3 100644 --- a/src/Fields/field_reductions.jl +++ b/src/Fields/field_reductions.jl @@ -57,3 +57,20 @@ function compute!(field::ReducedComputedField, time=nothing) return nothing end +##### +##### show +##### + +Base.show(io::IO, field::ReducedComputedField) = + print(io, "$(short_show(field))\n", + "├── data: $(typeof(field.data)), size: $(size(field))\n", + "├── grid: $(short_show(field.grid))\n", + "├── dims: $(field.dims)\n", + "├── operand: $(short_show(field.operand))\n", + "└── status: ", show_status(field.status)) + +short_show(field::ReducedComputedField) = string("Field at ", show_location(field), " via ", short_show(field.operand)) + +short_show(r::Reduction) = string(typeof(r.reduce!), " of ", short_show(r.operand), + " over dims ", r.dims) + diff --git a/src/Fields/show_fields.jl b/src/Fields/show_fields.jl index aa681e467d..4e7152bc45 100644 --- a/src/Fields/show_fields.jl +++ b/src/Fields/show_fields.jl @@ -11,7 +11,6 @@ show_location(field::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = show_locati short_show(m::Missing) = "$m" short_show(field::AbstractField) = string(typeof(field).name.wrapper, " located at ", show_location(field)) -short_show(field::AveragedField) = string("AveragedField over dims=$(field.dims) located at ", show_location(field), " of ", short_show(field.operand)) Base.show(io::IO, field::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = print(io, "$(short_show(field))\n", @@ -27,16 +26,9 @@ Base.show(io::IO, field::Field) = show_status(::Nothing) = "nothing" show_status(status) = "time=$(status.time)" -Base.show(io::IO, field::AveragedField) = - print(io, "$(short_show(field))\n", - "├── data: $(typeof(field.data)), size: $(size(field))\n", - "├── grid: $(short_show(field.grid))\n", - "├── dims: $(field.dims)\n", - "├── operand: $(short_show(field.operand))\n", - "└── status: ", show_status(field.status)) - Base.show(io::IO, field::ZeroField) = print(io, "ZeroField") short_show(array::OffsetArray{T, D, A}) where {T, D, A} = string("OffsetArray{$T, $D, $A}") Base.show(io::IO, ::MIME"text/plain", f::AbstractField) = show(io, f) + From c77cdcbf6224a40de1cf07ff15c20dddda689dbd Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 11:04:36 -0700 Subject: [PATCH 099/140] Numerous bugfixes --- src/AbstractOperations/binary_operations.jl | 2 +- src/AbstractOperations/derivatives.jl | 49 +++++++++---------- src/AbstractOperations/grid_metrics.jl | 9 ++-- .../kernel_function_operation.jl | 5 +- src/AbstractOperations/multiary_operations.jl | 2 +- src/Fields/Fields.jl | 2 + src/Fields/field.jl | 2 - src/Fields/field_slicer.jl | 1 - src/Fields/show_fields.jl | 2 - 9 files changed, 31 insertions(+), 43 deletions(-) diff --git a/src/AbstractOperations/binary_operations.jl b/src/AbstractOperations/binary_operations.jl index 75b2acf03f..03e942a362 100644 --- a/src/AbstractOperations/binary_operations.jl +++ b/src/AbstractOperations/binary_operations.jl @@ -16,7 +16,7 @@ struct BinaryOperation{LX, LY, LZ, O, A, B, IA, IB, G, T} <: AbstractOperation{L """ function BinaryOperation{LX, LY, LZ}(op::O, a::A, b::B, ▶a::IA, ▶b::IB, grid::G) where {LX, LY, LZ, O, A, B, IA, IB, G} T = eltype(grid) - return new{LX, LY, LZ, O, A, B, IA, IB, R, G, T}(op, a, b, ▶a, ▶b, grid) + return new{LX, LY, LZ, O, A, B, IA, IB, G, T}(op, a, b, ▶a, ▶b, grid) end end diff --git a/src/AbstractOperations/derivatives.jl b/src/AbstractOperations/derivatives.jl index 5cbfb02efa..8eee210e86 100644 --- a/src/AbstractOperations/derivatives.jl +++ b/src/AbstractOperations/derivatives.jl @@ -1,23 +1,22 @@ using Oceananigans.Operators: interpolation_code -struct Derivative{X, Y, Z, D, A, I, AD, R, G, T} <: AbstractOperation{X, Y, Z, R, G, T} +struct Derivative{LX, LY, LZ, D, A, I, AD, G, T} <: AbstractOperation{LX, LY, LZ, G, T} ∂ :: D arg :: A ▶ :: I abstract_∂ :: AD - architecture :: R grid :: G @doc """ - Derivative{X, Y, Z}(∂, arg, ▶, grid) + Derivative{LX, LY, LZ}(∂, arg, ▶, grid) Returns an abstract representation of the derivative `∂` on `arg`, and subsequent interpolation by `▶` on `grid`. """ - function Derivative{X, Y, Z}(∂::D, arg::A, ▶::I, abstract_∂::AD, - arch::R, grid::G) where {X, Y, Z, D, A, I, AD, R, G} + function Derivative{LX, LY, LZ}(∂::D, arg::A, ▶::I, abstract_∂::AD, + grid::G) where {LX, LY, LZ, D, A, I, AD, G} T = eltype(grid) - return new{X, Y, Z, D, A, I, AD, R, G, T}(∂, arg, ▶, abstract_∂, arch, grid) + return new{LX, LY, LZ, D, A, I, AD, G, T}(∂, arg, ▶, abstract_∂, grid) end end @@ -29,10 +28,9 @@ end """Create a derivative operator `∂` acting on `arg` at `L∂`, followed by interpolation to `L` on `grid`.""" -function _derivative(L, ∂, arg, L∂, abstract_∂, grid) where {X, Y, Z} +function _derivative(L, ∂, arg, L∂, abstract_∂, grid) where {LX, LY, LZ} ▶ = interpolation_operator(L∂, L) - arch = architecture(arg) - return Derivative{L[1], L[2], L[3]}(∂, arg, ▶, abstract_∂, arch, grid) + return Derivative{L[1], L[2], L[3]}(∂, arg, ▶, abstract_∂, grid) end # Recompute location of derivative @@ -62,8 +60,8 @@ push!(operators, derivative_operators...) Return an abstract representation of an x-derivative acting on field `a` followed by interpolation to `L`, where `L` is a 3-tuple of `Face`s and `Center`s. """ -∂x(L::Tuple, arg::AF{X, Y, Z}) where {X, Y, Z} = - _derivative(L, ∂x(X, Y, Z), arg, (flip(X), Y, Z), ∂x, arg.grid) +∂x(L::Tuple, arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = + _derivative(L, ∂x(LX, LY, LZ), arg, (flip(X), Y, Z), ∂x, arg.grid) """ ∂y(L::Tuple, a::AbstractField) @@ -71,8 +69,8 @@ by interpolation to `L`, where `L` is a 3-tuple of `Face`s and `Center`s. Return an abstract representation of a y-derivative acting on field `a` followed by interpolation to `L`, where `L` is a 3-tuple of `Face`s and `Center`s. """ -∂y(L::Tuple, arg::AF{X, Y, Z}) where {X, Y, Z} = - _derivative(L, ∂y(X, Y, Z), arg, (X, flip(Y), Z), ∂y, arg.grid) +∂y(L::Tuple, arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = + _derivative(L, ∂y(LX, LY, LZ), arg, (X, flip(Y), Z), ∂y, arg.grid) """ ∂z(L::Tuple, a::AbstractField) @@ -80,8 +78,8 @@ by interpolation to `L`, where `L` is a 3-tuple of `Face`s and `Center`s. Return an abstract representation of a z-derivative acting on field `a` followed by interpolation to `L`, where `L` is a 3-tuple of `Face`s and `Center`s. """ -∂z(L::Tuple, arg::AF{X, Y, Z}) where {X, Y, Z} = - _derivative(L, ∂z(X, Y, Z), arg, (X, Y, flip(Z)), ∂z, arg.grid) +∂z(L::Tuple, arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = + _derivative(L, ∂z(LX, LY, LZ), arg, (X, Y, flip(Z)), ∂z, arg.grid) # Defaults """ @@ -89,26 +87,20 @@ by interpolation to `L`, where `L` is a 3-tuple of `Face`s and `Center`s. Return an abstract representation of a x-derivative acting on field `a`. """ -∂x(arg::AF{X, Y, Z}) where {X, Y, Z} = ∂x((flip(X), Y, Z), arg) +∂x(arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = ∂x((flip(X), Y, Z), arg) """ ∂y(a::AbstractField) Return an abstract representation of a y-derivative acting on field `a`. """ -∂y(arg::AF{X, Y, Z}) where {X, Y, Z} = ∂y((X, flip(Y), Z), arg) +∂y(arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = ∂y((X, flip(Y), Z), arg) """ ∂z(a::AbstractField) Return an abstract representation of a z-derivative acting on field `a`. """ -∂z(arg::AF{X, Y, Z}) where {X, Y, Z} = ∂z((X, Y, flip(Z)), arg) - -##### -##### Architecture inference for derivatives -##### - -architecture(∂::Derivative) = ∂.architecture +∂z(arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = ∂z((X, Y, flip(Z)), arg) ##### ##### Nested computations @@ -123,7 +115,10 @@ compute_at!(∂::Derivative, time) = compute_at!(∂.arg, time) ##### "Adapt `Derivative` to work on the GPU via CUDAnative and CUDAdrv." -Adapt.adapt_structure(to, deriv::Derivative{X, Y, Z}) where {X, Y, Z} = - Derivative{X, Y, Z}(Adapt.adapt(to, deriv.∂), Adapt.adapt(to, deriv.arg), - Adapt.adapt(to, deriv.▶), nothing, nothing, Adapt.adapt(to, deriv.grid)) +Adapt.adapt_structure(to, deriv::Derivative{LX, LY, LZ}) where {LX, LY, LZ} = + Derivative{LX, LY, LZ}(Adapt.adapt(to, deriv.∂), + Adapt.adapt(to, deriv.arg), + Adapt.adapt(to, deriv.▶), + nothing, + Adapt.adapt(to, deriv.grid)) diff --git a/src/AbstractOperations/grid_metrics.jl b/src/AbstractOperations/grid_metrics.jl index 589031f1e3..62afe8fbfa 100644 --- a/src/AbstractOperations/grid_metrics.jl +++ b/src/AbstractOperations/grid_metrics.jl @@ -140,16 +140,13 @@ function metric_function(loc, metric::AbstractGridMetric) return eval(metric_function_symbol) end -struct GridMetricOperation{X, Y, Z, A, G, T, M} <: AbstractOperation{X, Y, Z, A, G, T} +struct GridMetricOperation{LX, LY, LZ, G, T, M} <: AbstractOperation{LX, LY, LZ, G, T} metric :: M grid :: G - architecture :: A - function GridMetricOperation{X, Y, Z}(metric::M, grid::G) where {X, Y, Z, M, G} - arch = grid.architecture - A = typeof(arch) + function GridMetricOperation{LX, LY, LZ}(metric::M, grid::G) where {LX, LY, LZ, M, G} T = eltype(grid) - return new{X, Y, Z, A, G, T, M}(metric, grid, arch) + return new{LX, LY, LZ, G, T, M}(metric, grid) end end diff --git a/src/AbstractOperations/kernel_function_operation.jl b/src/AbstractOperations/kernel_function_operation.jl index 145c1ff00f..bd95d0a93a 100644 --- a/src/AbstractOperations/kernel_function_operation.jl +++ b/src/AbstractOperations/kernel_function_operation.jl @@ -1,4 +1,4 @@ -struct KernelFunctionOperation{LX, LY, LZ, P, A, G, T, K, D} <: AbstractOperation{LX, LY, LZ, A, G, T} +struct KernelFunctionOperation{LX, LY, LZ, P, G, T, K, D} <: AbstractOperation{LX, LY, LZ, G, T} kernel_function :: K computed_dependencies :: D parameters :: P @@ -7,8 +7,7 @@ struct KernelFunctionOperation{LX, LY, LZ, P, A, G, T, K, D} <: AbstractOperatio function KernelFunctionOperation{LX, LY, LZ}(kernel_function::K, computed_dependencies::D, parameters::P, grid::G) where {LX, LY, LZ, K, G, D, P} T = eltype(grid) - A = typeof(architecture(grid)) - return new{LX, LY, LZ, P, A, G, T, K, D}(kernel_function, computed_dependencies, parameters, grid) + return new{LX, LY, LZ, P, G, T, K, D}(kernel_function, computed_dependencies, parameters, grid) end end diff --git a/src/AbstractOperations/multiary_operations.jl b/src/AbstractOperations/multiary_operations.jl index 7924475682..544705d783 100644 --- a/src/AbstractOperations/multiary_operations.jl +++ b/src/AbstractOperations/multiary_operations.jl @@ -9,7 +9,7 @@ struct MultiaryOperation{LX, LY, LZ, N, O, A, I, G, T} <: AbstractOperation{LX, function MultiaryOperation{LX, LY, LZ}(op::O, args::A, ▶::I, grid::G) where {LX, LY, LZ, O, A, I, G} T = eltype(grid) N = length(args) - return new{LX, LY, LZ, N, O, A, I, R, G, T}(op, args, ▶, grid) + return new{LX, LY, LZ, N, O, A, I, G, T}(op, args, ▶, grid) end end diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl index eb9182b3a6..6849549345 100644 --- a/src/Fields/Fields.jl +++ b/src/Fields/Fields.jl @@ -13,6 +13,8 @@ using Oceananigans.Architectures using Oceananigans.Grids using Oceananigans.BoundaryConditions +import Oceananigans: short_show + include("abstract_field.jl") include("zero_field.jl") include("function_field.jl") diff --git a/src/Fields/field.jl b/src/Fields/field.jl index 8ad90cc49d..6462a7ea5b 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -101,8 +101,6 @@ function Base.similar(f::Field, grid=f.grid) deepcopy(f.status)) end -data(f::Field) = f.data - # Fallback: cannot infer boundary conditions. boundary_conditions(field) = nothing boundary_conditions(f::Field) = f.boundary_conditions diff --git a/src/Fields/field_slicer.jl b/src/Fields/field_slicer.jl index d0eb3eca12..2313c142d3 100644 --- a/src/Fields/field_slicer.jl +++ b/src/Fields/field_slicer.jl @@ -1,4 +1,3 @@ -import Oceananigans: short_show using Oceananigans.Grids: Face, Bounded, topology using Oceananigans.Fields: location diff --git a/src/Fields/show_fields.jl b/src/Fields/show_fields.jl index 4e7152bc45..8d5884c236 100644 --- a/src/Fields/show_fields.jl +++ b/src/Fields/show_fields.jl @@ -1,5 +1,3 @@ -import Oceananigans: short_show - location_str(::Type{Face}) = "Face" location_str(::Type{Center}) = "Center" location_str(::Type{Nothing}) = "⋅" From 83732b90984711c33b79a54aed72ce8623754451 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 11:16:04 -0700 Subject: [PATCH 100/140] Bugfix for derivatives --- src/AbstractOperations/derivatives.jl | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/AbstractOperations/derivatives.jl b/src/AbstractOperations/derivatives.jl index 8eee210e86..e673f986c5 100644 --- a/src/AbstractOperations/derivatives.jl +++ b/src/AbstractOperations/derivatives.jl @@ -61,7 +61,7 @@ Return an abstract representation of an x-derivative acting on field `a` followe by interpolation to `L`, where `L` is a 3-tuple of `Face`s and `Center`s. """ ∂x(L::Tuple, arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = - _derivative(L, ∂x(LX, LY, LZ), arg, (flip(X), Y, Z), ∂x, arg.grid) + _derivative(L, ∂x(LX, LY, LZ), arg, (flip(LX), LY, LZ), ∂x, arg.grid) """ ∂y(L::Tuple, a::AbstractField) @@ -70,7 +70,7 @@ Return an abstract representation of a y-derivative acting on field `a` followed by interpolation to `L`, where `L` is a 3-tuple of `Face`s and `Center`s. """ ∂y(L::Tuple, arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = - _derivative(L, ∂y(LX, LY, LZ), arg, (X, flip(Y), Z), ∂y, arg.grid) + _derivative(L, ∂y(LX, LY, LZ), arg, (LX, flip(LY), LZ), ∂y, arg.grid) """ ∂z(L::Tuple, a::AbstractField) @@ -79,7 +79,7 @@ Return an abstract representation of a z-derivative acting on field `a` followed by interpolation to `L`, where `L` is a 3-tuple of `Face`s and `Center`s. """ ∂z(L::Tuple, arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = - _derivative(L, ∂z(LX, LY, LZ), arg, (X, Y, flip(Z)), ∂z, arg.grid) + _derivative(L, ∂z(LX, LY, LZ), arg, (LX, LY, flip(LZ)), ∂z, arg.grid) # Defaults """ @@ -87,20 +87,20 @@ by interpolation to `L`, where `L` is a 3-tuple of `Face`s and `Center`s. Return an abstract representation of a x-derivative acting on field `a`. """ -∂x(arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = ∂x((flip(X), Y, Z), arg) +∂x(arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = ∂x((flip(LX), LY, LZ), arg) """ ∂y(a::AbstractField) Return an abstract representation of a y-derivative acting on field `a`. """ -∂y(arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = ∂y((X, flip(Y), Z), arg) +∂y(arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = ∂y((LX, flip(LY), LZ), arg) """ ∂z(a::AbstractField) Return an abstract representation of a z-derivative acting on field `a`. """ -∂z(arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = ∂z((X, Y, flip(Z)), arg) +∂z(arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = ∂z((LX, LY, flip(LZ)), arg) ##### ##### Nested computations From b0d5677ffae98cb505e93b6aba70fcd8ee8a7835 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 11:17:43 -0700 Subject: [PATCH 101/140] Bugfix in CubedSphere --- src/CubedSpheres/cubed_sphere_faces.jl | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/CubedSpheres/cubed_sphere_faces.jl b/src/CubedSpheres/cubed_sphere_faces.jl index 9087c80465..2d17c97ebc 100644 --- a/src/CubedSpheres/cubed_sphere_faces.jl +++ b/src/CubedSpheres/cubed_sphere_faces.jl @@ -35,18 +35,18 @@ const CubedSphereData = CubedSphereFaces{<:OffsetArray} const ImmersedConformalCubedSphereFaceGrid = ImmersedBoundaryGrid{FT, TX, TY, TZ, <:ConformalCubedSphereFaceGrid} where {FT, TX, TY, TZ} # CubedSphereFaceField: -const NonImmersedCubedSphereFaceField = AbstractField{LX, LY, LZ, <:ConformalCubedSphereFaceGrid} where {LX, LY, LZ, A} -const ImmersedCubedSphereFaceField = AbstractField{LX, LY, LZ, <:ImmersedConformalCubedSphereFaceGrid} where {LX, LY, LZ, A} +const NonImmersedCubedSphereFaceField = AbstractField{LX, LY, LZ, <:ConformalCubedSphereFaceGrid} where {LX, LY, LZ} +const ImmersedCubedSphereFaceField = AbstractField{LX, LY, LZ, <:ImmersedConformalCubedSphereFaceGrid} where {LX, LY, LZ} -const CubedSphereFaceField = Union{NonImmersedCubedSphereFaceField{LX, LY, LZ, A}, - ImmersedCubedSphereFaceField{LX, LY, LZ, A}} where {LX, LY, LZ, A} +const CubedSphereFaceField = Union{NonImmersedCubedSphereFaceField{LX, LY, LZ}, + ImmersedCubedSphereFaceField{LX, LY, LZ}} where {LX, LY, LZ} # CubedSphereField # Flavors of CubedSphereField -const CubedSphereField{LX, LY, LZ, A} = - Union{Field{LX, LY, LZ, <:Nothing, A, <:ConformalCubedSphereGrid}, - Field{LX, LY, LZ, <:AbstractOperation, A, <:ConformalCubedSphereGrid}} +const CubedSphereField{LX, LY, LZ} = + Union{Field{LX, LY, LZ, <:Nothing, <:ConformalCubedSphereGrid}, + Field{LX, LY, LZ, <:AbstractOperation, <:ConformalCubedSphereGrid}} const CubedSphereAbstractField{LX, LY, LZ} = AbstractField{LX, LY, LZ, <:ConformalCubedSphereGrid} From 4ac3e9228740ea6512c6223d09da95ac9a20960e Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 11:23:19 -0700 Subject: [PATCH 102/140] Dont import KernelComputedField into CubedSpheres --- src/CubedSpheres/cubed_sphere_faces.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/CubedSpheres/cubed_sphere_faces.jl b/src/CubedSpheres/cubed_sphere_faces.jl index 2d17c97ebc..cd2c623899 100644 --- a/src/CubedSpheres/cubed_sphere_faces.jl +++ b/src/CubedSpheres/cubed_sphere_faces.jl @@ -8,7 +8,7 @@ using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid import Base: getindex, size, show, minimum, maximum import Statistics: mean -import Oceananigans.Fields: AbstractField, Field, minimum, maximum, mean, location, short_show, KernelComputedField +import Oceananigans.Fields: AbstractField, Field, minimum, maximum, mean, location, short_show import Oceananigans.Grids: new_data import Oceananigans.BoundaryConditions: FieldBoundaryConditions From 7b171d3f83c506b032d7e87115ab4e06ba1cca04 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 11:27:00 -0700 Subject: [PATCH 103/140] Simplify CubedSphere set! --- src/CubedSpheres/CubedSpheres.jl | 1 - src/CubedSpheres/cubed_sphere_faces.jl | 9 +++++++- src/CubedSpheres/cubed_sphere_set!.jl | 31 -------------------------- 3 files changed, 8 insertions(+), 33 deletions(-) delete mode 100644 src/CubedSpheres/cubed_sphere_set!.jl diff --git a/src/CubedSpheres/CubedSpheres.jl b/src/CubedSpheres/CubedSpheres.jl index 122cf4dee7..a01d78ed40 100644 --- a/src/CubedSpheres/CubedSpheres.jl +++ b/src/CubedSpheres/CubedSpheres.jl @@ -6,7 +6,6 @@ include("cubed_sphere_utils.jl") include("conformal_cubed_sphere_grid.jl") include("cubed_sphere_exchange_bcs.jl") include("cubed_sphere_faces.jl") -include("cubed_sphere_set!.jl") include("cubed_sphere_halo_filling.jl") include("cubed_sphere_kernel_launching.jl") include("immersed_conformal_cubed_sphere_grid.jl") diff --git a/src/CubedSpheres/cubed_sphere_faces.jl b/src/CubedSpheres/cubed_sphere_faces.jl index cd2c623899..92f0997c62 100644 --- a/src/CubedSpheres/cubed_sphere_faces.jl +++ b/src/CubedSpheres/cubed_sphere_faces.jl @@ -8,7 +8,7 @@ using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid import Base: getindex, size, show, minimum, maximum import Statistics: mean -import Oceananigans.Fields: AbstractField, Field, minimum, maximum, mean, location, short_show +import Oceananigans.Fields: AbstractField, Field, minimum, maximum, mean, location, short_show, set! import Oceananigans.Grids: new_data import Oceananigans.BoundaryConditions: FieldBoundaryConditions @@ -109,6 +109,13 @@ Base.size(data::CubedSphereData) = (size(data.faces[1])..., length(data.faces)) faces(field::AbstractCubedSphereField) = Tuple(get_face(field, face_index) for face_index in 1:length(field.data.faces)) +function set!(u::CubedSphereField, v) + for face = 1:length(u.grid.faces) + set!(get_face(u, face), get_face(v, face)) + end + return nothing +end + minimum(field::AbstractCubedSphereField; dims=:) = minimum(minimum(face_field; dims) for face_field in faces(field)) maximum(field::AbstractCubedSphereField; dims=:) = maximum(maximum(face_field; dims) for face_field in faces(field)) mean(field::AbstractCubedSphereField; dims=:) = mean(mean(face_field; dims) for face_field in faces(field)) diff --git a/src/CubedSpheres/cubed_sphere_set!.jl b/src/CubedSpheres/cubed_sphere_set!.jl deleted file mode 100644 index 8dd94f7206..0000000000 --- a/src/CubedSpheres/cubed_sphere_set!.jl +++ /dev/null @@ -1,31 +0,0 @@ -using Oceananigans.Architectures: CPU -using Oceananigans.Fields: AbstractField - -import Oceananigans.Fields: set! - -const CubedSphereCPUField = CubedSphereField{LX, LY, LZ, <:CPU} where {LX, LY, LZ} -const CubedSphereGPUField = CubedSphereField{LX, LY, LZ, <:GPU} where {LX, LY, LZ} - -# We need to define the function once for CPU fields then again for GPU fields to avoid the method -# ambiguity with Fields.set!. - -function set!(u::CubedSphereCPUField, v::CubedSphereCPUField) - for (u_face, v_face) in zip(faces(u), faces(v)) - @. u_face.data.parent = v_face.data.parent - end - return nothing -end - -function set!(u::CubedSphereGPUField, v::CubedSphereGPUField) - for (u_face, v_face) in zip(faces(u), faces(v)) - @. u_face.data.parent = v_face.data.parent - end - return nothing -end - -set!(field::CubedSphereCPUField, f::Function) = [set_face_field!(field_face, f) for field_face in faces(field)] -set!(field::CubedSphereCPUField, f::Number) = [set_face_field!(field_face, f) for field_face in faces(field)] -set!(field::CubedSphereGPUField, f::Function) = [set_face_field!(field_face, f) for field_face in faces(field)] -set!(field::CubedSphereGPUField, f::Number) = [set_face_field!(field_face, f) for field_face in faces(field)] - -set_face_field!(field, a) = set!(field, a) From 6c10263632f7668ef903d5cfe17d601f3ea4995d Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 11:53:34 -0700 Subject: [PATCH 104/140] Adds back fallback for location --- src/Fields/abstract_field.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/Fields/abstract_field.jl b/src/Fields/abstract_field.jl index 62b0fec133..69eba4c665 100644 --- a/src/Fields/abstract_field.jl +++ b/src/Fields/abstract_field.jl @@ -38,6 +38,7 @@ data(a) = a ##### "Returns the location `(LX, LY, LZ)` of an `AbstractField{LX, LY, LZ}`." +@inline location(a) = (Nothing, Nothing, Nothing) # used in AbstractOperations for location inference @inline location(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = (LX, LY, LZ) # note no instantiation @inline location(f::AbstractField, i) = location(f)[i] instantiated_location(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = (LX(), LY(), LZ()) From 5a47c2920b48ae6bb1bc7466caab56b074c8d441 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 12:17:26 -0700 Subject: [PATCH 105/140] Resolve ambiguities for setting fields on the cubed sphere --- src/CubedSpheres/cubed_sphere_faces.jl | 15 ++++++++++++++- src/Fields/set!.jl | 11 ----------- test/test_cubed_spheres.jl | 4 ++-- 3 files changed, 16 insertions(+), 14 deletions(-) diff --git a/src/CubedSpheres/cubed_sphere_faces.jl b/src/CubedSpheres/cubed_sphere_faces.jl index 92f0997c62..e2720d4ef2 100644 --- a/src/CubedSpheres/cubed_sphere_faces.jl +++ b/src/CubedSpheres/cubed_sphere_faces.jl @@ -109,13 +109,26 @@ Base.size(data::CubedSphereData) = (size(data.faces[1])..., length(data.faces)) faces(field::AbstractCubedSphereField) = Tuple(get_face(field, face_index) for face_index in 1:length(field.data.faces)) -function set!(u::CubedSphereField, v) +##### +##### set! +##### + +function cubed_sphere_set!(u, v) for face = 1:length(u.grid.faces) set!(get_face(u, face), get_face(v, face)) end return nothing end +# Resolve ambiguities +set!(u::CubedSphereField, v) = cubed_sphere_set!(u, v) +set!(u::CubedSphereField, v::Union{Function, Array}) = cubed_sphere_set!(u, v) +set!(u::CubedSphereField, v::CubedSphereField) = cubed_sphere_set!(u, v) + +##### +##### Random utils +##### + minimum(field::AbstractCubedSphereField; dims=:) = minimum(minimum(face_field; dims) for face_field in faces(field)) maximum(field::AbstractCubedSphereField; dims=:) = maximum(maximum(face_field; dims) for face_field in faces(field)) mean(field::AbstractCubedSphereField; dims=:) = mean(mean(face_field; dims) for face_field in faces(field)) diff --git a/src/Fields/set!.jl b/src/Fields/set!.jl index 9cea051225..db7c581e3d 100644 --- a/src/Fields/set!.jl +++ b/src/Fields/set!.jl @@ -27,15 +27,4 @@ function set!(u::Field, f::Union{Array, Function}) return nothing end -##### -##### set! for fields on the GPU -##### - -const GPUField = Field{LX, LY, LZ, O, <:GPU} where {LX, LY, LZ, O} - -""" Set the GPU field `u` to the array or function `v`. """ -function set!(u::GPUField, v::Union{Array, Function}) - return nothing -end - set!(u::Field, v::Field) = copyto!(parent(u), parent(v)) diff --git a/test/test_cubed_spheres.jl b/test/test_cubed_spheres.jl index d1de5641e8..712317b280 100644 --- a/test/test_cubed_spheres.jl +++ b/test/test_cubed_spheres.jl @@ -68,8 +68,8 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels: VerticalVorticityField @test try time_step!(model, 1); true; catch; false; end end - @testset "KernelComputedField on ConformalCubedSphereGrid [$(typeof(arch))]" begin - @info "Testing KernelComputedField on a ConformalCubedSphereGrid [$(typeof(arch))]..." + @testset "VerticalVorticityField on ConformalCubedSphereGrid [$(typeof(arch))]" begin + @info "Testing VerticalVorticityField on a ConformalCubedSphereGrid [$(typeof(arch))]..." ζ = VerticalVorticityField(model) @test ζ isa Field From 89eb639393690f47a19f58c6687e36f35e1acdf6 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 12:29:53 -0700 Subject: [PATCH 106/140] Update DistributedField constructor --- src/Distributed/distributed_fields.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/Distributed/distributed_fields.jl b/src/Distributed/distributed_fields.jl index 7265cf9c65..d66e245b6d 100644 --- a/src/Distributed/distributed_fields.jl +++ b/src/Distributed/distributed_fields.jl @@ -4,10 +4,10 @@ import Oceananigans.BoundaryConditions: fill_halo_regions! function Field((LX, LY, LZ)::Tuple, grid::DistributedGrid, data, old_bcs, op, status) arch = architecture(grid) new_bcs = inject_halo_communication_boundary_conditions(old_bcs, arch.local_rank, arch.connectivity) - return Field{LX, LY, LZ}(grid, data, new_bcs, op, status, child_architecture(arch)) + return Field{LX, LY, LZ}(grid, data, new_bcs, op, status) end -const DistributedField = Field{<:Any, <:Any, <:Any, <:Any, <:Any, <:DistributedGrid} +const DistributedField = Field{<:Any, <:Any, <:Any, <:Any, <:DistributedGrid} fill_halo_regions!(field::DistributedField, arch, args...; kwargs...) = fill_halo_regions!(field.data, From 3cd5e3404291ed4f72c2bf9ad68d435dba03f6a1 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 12:31:35 -0700 Subject: [PATCH 107/140] Fixes show_location --- src/Fields/show_fields.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Fields/show_fields.jl b/src/Fields/show_fields.jl index 8d5884c236..b9e88b8855 100644 --- a/src/Fields/show_fields.jl +++ b/src/Fields/show_fields.jl @@ -2,7 +2,7 @@ location_str(::Type{Face}) = "Face" location_str(::Type{Center}) = "Center" location_str(::Type{Nothing}) = "⋅" -show_location(LX, LY, LZ) = "($(location_str(X)), $(location_str(Y)), $(location_str(Z)))" +show_location(LX, LY, LZ) = "($(location_str(LX)), $(location_str(LY)), $(location_str(LZ)))" show_location(field::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = show_location(LX, LY, LZ) From 95d40c519175cf86b80606fa7fa963801c539283 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 12:50:07 -0700 Subject: [PATCH 108/140] Spruse up info formatting for cubed sphere tests --- test/test_cubed_spheres.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/test/test_cubed_spheres.jl b/test/test_cubed_spheres.jl index 712317b280..5c69155f31 100644 --- a/test/test_cubed_spheres.jl +++ b/test/test_cubed_spheres.jl @@ -16,13 +16,13 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels: VerticalVorticityField for arch in archs - @info "Constructing a ConformalCubedSphereGrid from file [$(typeof(arch))]..." + @info " Constructing a ConformalCubedSphereGrid from file [$(typeof(arch))]..." # Prototype grid and model for subsequent tests cs32_filepath = datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2" grid = ConformalCubedSphereGrid(cs32_filepath, arch, Nz=1, z=(-1, 0)) - @info "Constructing a HydrostaticFreeSurfaceModel on a ConformalCubedSphereGrid [$(typeof(arch))]..." + @info " Constructing a HydrostaticFreeSurfaceModel on a ConformalCubedSphereGrid [$(typeof(arch))]..." free_surface = ExplicitFreeSurface(gravitational_acceleration=0.1) model = HydrostaticFreeSurfaceModel(; grid, free_surface, @@ -37,7 +37,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels: VerticalVorticityField end @testset "CubedSphereData and CubedSphereFields [$(typeof(arch))]" begin - @info "Testing CubedSphereData and CubedSphereFields [$(typeof(arch))]..." + @info " Testing CubedSphereData and CubedSphereFields [$(typeof(arch))]..." c = model.tracers.c η = model.free_surface.η @@ -63,13 +63,13 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels: VerticalVorticityField end @testset "Time stepping a HydrostaticFreeSurfaceModel [$(typeof(arch))]" begin - @info "Time-stepping HydrostaticFreeSurfaceModel on a ConformalCubedSphereGrid [$(typeof(arch))]..." + @info " Time-stepping HydrostaticFreeSurfaceModel on a ConformalCubedSphereGrid [$(typeof(arch))]..." time_step!(model, 1) @test try time_step!(model, 1); true; catch; false; end end @testset "VerticalVorticityField on ConformalCubedSphereGrid [$(typeof(arch))]" begin - @info "Testing VerticalVorticityField on a ConformalCubedSphereGrid [$(typeof(arch))]..." + @info " Testing VerticalVorticityField on a ConformalCubedSphereGrid [$(typeof(arch))]..." ζ = VerticalVorticityField(model) @test ζ isa Field From 4512c642e1e87c216acc210a3bbbdf60475ff357 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 12:50:57 -0700 Subject: [PATCH 109/140] Fix adapt_structure for reduced fields --- src/Fields/field.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/Fields/field.jl b/src/Fields/field.jl index 6462a7ea5b..f438703f7e 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -308,8 +308,9 @@ function Adapt.adapt_structure(to, reduced_field::ReducedField) LX, LY, LZ = location(reduced_field) return Field{LX, LY, LZ}(nothing, adapt(to, reduced_field.data), - nothing, nothing, nothing, - adapt(to, architecture(reduced_field))) + nothing, + nothing, + nothing) end ##### From cfa3d032abb7d3c0c03ad6ce8d2dc8c34b88aafa Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 12:52:01 -0700 Subject: [PATCH 110/140] Fix ViewField alias --- src/OutputReaders/field_time_series.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/OutputReaders/field_time_series.jl b/src/OutputReaders/field_time_series.jl index 3d8b714278..796167dea8 100644 --- a/src/OutputReaders/field_time_series.jl +++ b/src/OutputReaders/field_time_series.jl @@ -201,7 +201,7 @@ end # is there a better way? # FieldTimeSeries[i] returns ViewField -const ViewField = Field{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:SubArray} +const ViewField = Field{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:SubArray} using OffsetArrays: IdOffsetRange From 00137231bb2e3effd139744739cb22da1964c112 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 13:33:10 -0700 Subject: [PATCH 111/140] Bugfix in interior_copy --- src/Fields/field.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/Fields/field.jl b/src/Fields/field.jl index f438703f7e..729ffbc0fb 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -116,9 +116,9 @@ function interior(f::Field) end interior_copy(f::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = - parent(f)[interior_parent_indices(X, topology(f, 1), f.grid.Nx, f.grid.Hx), - interior_parent_indices(Y, topology(f, 2), f.grid.Ny, f.grid.Hy), - interior_parent_indices(Z, topology(f, 3), f.grid.Nz, f.grid.Hz)] + parent(f)[interior_parent_indices(LX, topology(f, 1), f.grid.Nx, f.grid.Hx), + interior_parent_indices(LY, topology(f, 2), f.grid.Ny, f.grid.Hy), + interior_parent_indices(LZ, topology(f, 3), f.grid.Nz, f.grid.Hz)] # Don't use axes(f) to checkbounds; use axes(f.data) Base.checkbounds(f::Field, I...) = Base.checkbounds(f.data, I...) From 731d5cac3da58dc9ba2fe2f853a53ece49328e30 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 15:16:23 -0700 Subject: [PATCH 112/140] Fix field for gpu grids --- src/Fields/Fields.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl index 6849549345..486c08a4b1 100644 --- a/src/Fields/Fields.jl +++ b/src/Fields/Fields.jl @@ -33,8 +33,9 @@ include("broadcasting_abstract_fields.jl") Build a field from `a` at `loc` and on `grid`. """ -function field(loc, a::Array, grid) +function field(loc, a::AbstractArray, grid) f = Field(loc, grid) + a = arch_array(architecture(grid), a) f .= a return f end From 363ba06f4cde5f146d932530e3b0c988e8450d00 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 15:45:24 -0700 Subject: [PATCH 113/140] Add allowscalar for time_discretization determination --- .../HydrostaticFreeSurfaceModels/single_column_model_mode.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl index 71d063155c..da1bb395c2 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl @@ -91,7 +91,7 @@ validate_size(TX, TY, TZ, e::ColumnEnsembleSize) = tuple(e.ensemble[1], e.ensemb validate_halo(TX, TY, TZ, e::ColumnEnsembleSize) = tuple(0, 0, e.Hz) @inline function time_discretization(closure_array::AbstractArray) - first_closure = first(closure_array) # assumes all closures have same time-discretization + first_closure = CUDA.@allowscalar first(closure_array) # assumes all closures have same time-discretization return time_discretization(first_closure) end From c8a833ef832c621bd87ede26bad422c1b2c8d5fa Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 15:54:45 -0700 Subject: [PATCH 114/140] import CUDA into single_column_model_mode.jl --- .../HydrostaticFreeSurfaceModels/single_column_model_mode.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl index da1bb395c2..d106844f1e 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl @@ -1,5 +1,6 @@ using KernelAbstractions: NoneEvent using OffsetArrays: OffsetArray +using CUDA using Oceananigans.Operators: Δzᵃᵃᶜ using Oceananigans.BoundaryConditions: left_gradient, right_gradient, linearly_extrapolate, FBC, VBC, GBC From 188fc4e44267b0c99ef2ae31a7489a09bb8d9baf Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 16:05:15 -0700 Subject: [PATCH 115/140] Bugfix in interior for FieldTimeSeries --- src/OutputReaders/field_time_series.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/OutputReaders/field_time_series.jl b/src/OutputReaders/field_time_series.jl index 796167dea8..22e2589e91 100644 --- a/src/OutputReaders/field_time_series.jl +++ b/src/OutputReaders/field_time_series.jl @@ -214,9 +214,9 @@ Base.parent(vf::ViewField) = view(parent(parent(vf.data)), parent_indices.(vf.da "Returns a view of `f` that excludes halo points." @inline interior(f::FieldTimeSeries{LX, LY, LZ}) where {LX, LY, LZ} = view(parent(f.data), - interior_parent_indices(X, topology(f, 1), f.grid.Nx, f.grid.Hx), - interior_parent_indices(Y, topology(f, 2), f.grid.Ny, f.grid.Hy), - interior_parent_indices(Z, topology(f, 3), f.grid.Nz, f.grid.Hz), + interior_parent_indices(LX, topology(f, 1), f.grid.Nx, f.grid.Hx), + interior_parent_indices(LY, topology(f, 2), f.grid.Ny, f.grid.Hy), + interior_parent_indices(LZ, topology(f, 3), f.grid.Nz, f.grid.Hz), :) ##### From d439132da2808737b7a921ab0e765d80f1876437 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 16:15:51 -0700 Subject: [PATCH 116/140] Ensure that set! for functions is always on cpu --- src/Fields/set!.jl | 29 ++++++++++++++++++++++------- 1 file changed, 22 insertions(+), 7 deletions(-) diff --git a/src/Fields/set!.jl b/src/Fields/set!.jl index db7c581e3d..812ac0d0e7 100644 --- a/src/Fields/set!.jl +++ b/src/Fields/set!.jl @@ -15,15 +15,30 @@ function set!(Φ::NamedTuple; kwargs...) return nothing end -set!(u::AbstractField, v) = u .= v # fallback +set!(u::Field, v) = u .= v # fallback -# Niceties -const CPUField = Field{LX, LY, LZ, O, <:CPU} where {LX, LY, LZ, O} - -""" Set the CPU field `u` data to the function `f(x, y, z)`. """ function set!(u::Field, f::Union{Array, Function}) - f_field = field(location(u), f, u.grid) - set!(u, f_field) + if architecture(u) isa GPU + cpu_grid = on_architecture(CPU(), u.grid) + u_cpu = Field(location(u), cpu_grid) + f_field = field(location(u), f, cpu_grid) + set!(u_cpu, f_field) + set!(u, u_cpu) + elseif architecture(u) isa CPU + f_field = field(location(u), f, u.grid) + set!(u, f_field) + end + + return nothing +end + +function set!(u::Field, f::CuArray) + if architecture(u) isa CPU + f_cpu = arch_array(CPU(), f) + u .= f_cpu + else + u .= f + end return nothing end From 5335b8411291294cabcaed61d98c0cb8281c6f2b Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 16:17:52 -0700 Subject: [PATCH 117/140] Set for Array right --- src/Fields/set!.jl | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/src/Fields/set!.jl b/src/Fields/set!.jl index 812ac0d0e7..6f7fcdcd96 100644 --- a/src/Fields/set!.jl +++ b/src/Fields/set!.jl @@ -17,7 +17,7 @@ end set!(u::Field, v) = u .= v # fallback -function set!(u::Field, f::Union{Array, Function}) +function set!(u::Field, f::Function) if architecture(u) isa GPU cpu_grid = on_architecture(CPU(), u.grid) u_cpu = Field(location(u), cpu_grid) @@ -32,13 +32,9 @@ function set!(u::Field, f::Union{Array, Function}) return nothing end -function set!(u::Field, f::CuArray) - if architecture(u) isa CPU - f_cpu = arch_array(CPU(), f) - u .= f_cpu - else - u .= f - end +function set!(u::Field, f::AbstractArray) + f = arch_array(architecture(u), f) + u .= f return nothing end From 91e32c53d3c5f3bff840209f290a20dd3db98dc1 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 16:20:36 -0700 Subject: [PATCH 118/140] Respect architecture when grid is specified --- src/OutputReaders/field_time_series.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/OutputReaders/field_time_series.jl b/src/OutputReaders/field_time_series.jl index 22e2589e91..c7c14ab6ce 100644 --- a/src/OutputReaders/field_time_series.jl +++ b/src/OutputReaders/field_time_series.jl @@ -104,9 +104,11 @@ function FieldTimeSeries(path, name, backend::InMemory; architecture if isnothing(grid) - grid = on_architecture(architecture, file["serialized/grid"]) + grid = file["serialized/grid"] end + grid = on_architecture(architecture, grid) + if boundary_conditions isa UnspecifiedBoundaryConditions boundary_conditions = file["timeseries/$name/serialized/boundary_conditions"] end From 341abdf0c50e0ef4ee9fdd3fcddce70518ebed55 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 16:22:05 -0700 Subject: [PATCH 119/140] Restrict set! for arrays to specific kinds of arrays --- src/Fields/set!.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Fields/set!.jl b/src/Fields/set!.jl index 6f7fcdcd96..477cbb0e4d 100644 --- a/src/Fields/set!.jl +++ b/src/Fields/set!.jl @@ -32,7 +32,7 @@ function set!(u::Field, f::Function) return nothing end -function set!(u::Field, f::AbstractArray) +function set!(u::Field, f::Union{Array, CuArray, OffsetArray}) f = arch_array(architecture(u), f) u .= f return nothing From a10a8c531b9a51d0384b9141ab97d56864bec566 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 18:47:10 -0500 Subject: [PATCH 120/140] Fix method resolution in cubed sphere set --- src/CubedSpheres/cubed_sphere_faces.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/CubedSpheres/cubed_sphere_faces.jl b/src/CubedSpheres/cubed_sphere_faces.jl index e2720d4ef2..a2026635d7 100644 --- a/src/CubedSpheres/cubed_sphere_faces.jl +++ b/src/CubedSpheres/cubed_sphere_faces.jl @@ -122,7 +122,8 @@ end # Resolve ambiguities set!(u::CubedSphereField, v) = cubed_sphere_set!(u, v) -set!(u::CubedSphereField, v::Union{Function, Array}) = cubed_sphere_set!(u, v) +set!(u::CubedSphereField, v::Function) = cubed_sphere_set!(u, v) +set!(u::CubedSphereField, v::Union{Array, CuArray, OffsetArray}) = cubed_sphere_set!(u, v) set!(u::CubedSphereField, v::CubedSphereField) = cubed_sphere_set!(u, v) ##### From a4754ac22a59345187eafe6ecbacad11cfd95c48 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 19:17:34 -0500 Subject: [PATCH 121/140] Distinguish between CPU and GPU cases in field-field setting --- src/Fields/set!.jl | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/Fields/set!.jl b/src/Fields/set!.jl index 477cbb0e4d..fc6ebb9f49 100644 --- a/src/Fields/set!.jl +++ b/src/Fields/set!.jl @@ -38,4 +38,12 @@ function set!(u::Field, f::Union{Array, CuArray, OffsetArray}) return nothing end -set!(u::Field, v::Field) = copyto!(parent(u), parent(v)) +function set!(u::Field, v::Field) + if architecture(u) === architecture(v) + parent(u) .= parent(v) + else + copyto!(parent(u), parent(v)) + end + return nothing +end + From 154c0138007bd6b5435b1e09887320270db05aae Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 19:29:45 -0500 Subject: [PATCH 122/140] Import CUDA into cubed_sphere_faces --- src/CubedSpheres/cubed_sphere_faces.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/CubedSpheres/cubed_sphere_faces.jl b/src/CubedSpheres/cubed_sphere_faces.jl index a2026635d7..5e78d4a577 100644 --- a/src/CubedSpheres/cubed_sphere_faces.jl +++ b/src/CubedSpheres/cubed_sphere_faces.jl @@ -1,4 +1,5 @@ using Statistics +using CUDA using Oceananigans.Architectures using Oceananigans.AbstractOperations: AbstractOperation From 48add2096f6fa61292b210fcc182b80e776586d5 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 17:42:42 -0700 Subject: [PATCH 123/140] Fix on_architecture for non-conversions --- src/Grids/conformal_cubed_sphere_face_grid.jl | 82 ++++++++++--------- src/Grids/latitude_longitude_grid.jl | 44 +++++----- src/Grids/rectilinear_grid.jl | 34 ++++---- .../single_column_model_mode.jl | 17 +++- 4 files changed, 105 insertions(+), 72 deletions(-) diff --git a/src/Grids/conformal_cubed_sphere_face_grid.jl b/src/Grids/conformal_cubed_sphere_face_grid.jl index f1e6024388..a82af5c512 100644 --- a/src/Grids/conformal_cubed_sphere_face_grid.jl +++ b/src/Grids/conformal_cubed_sphere_face_grid.jl @@ -275,43 +275,51 @@ function ConformalCubedSphereFaceGrid(filepath::AbstractString, architecture = C end function on_architecture(arch, grid::ConformalCubedSphereFaceGrid) - - horizontal_coordinates = (:λᶜᶜᵃ, - :λᶠᶜᵃ, - :λᶜᶠᵃ, - :λᶠᶠᵃ, - :φᶜᶜᵃ, - :φᶠᶜᵃ, - :φᶜᶠᵃ, - :φᶠᶠᵃ) - - horizontal_grid_spacings = (:Δxᶜᶜᵃ, - :Δxᶠᶜᵃ, - :Δxᶜᶠᵃ, - :Δxᶠᶠᵃ, - :Δyᶜᶜᵃ, - :Δyᶜᶠᵃ, - :Δyᶠᶜᵃ, - :Δyᶠᶠᵃ) - - horizontal_areas = (:Azᶜᶜᵃ, - :Azᶠᶜᵃ, - :Azᶜᶠᵃ, - :Azᶠᶠᵃ) - - horizontal_grid_spacing_data = Tuple(arch_array(arch, getproperty(grid, name)) for name in horizontal_grid_spacings) - horizontal_coordinate_data = Tuple(arch_array(arch, getproperty(grid, name)) for name in horizontal_coordinates) - horizontal_area_data = Tuple(arch_array(arch, getproperty(grid, name)) for name in horizontal_areas) - - zᵃᵃᶜ = arch_array(arch, grid.zᵃᵃᶜ) - zᵃᵃᶠ = arch_array(arch, grid.zᵃᵃᶠ) - - TX, TY, TZ = topology(grid) - - return ConformalCubedSphereFaceGrid{TX, TY, TZ}(architecture, grid.Nx, grid.Ny, grid.Nz, grid.Hx, grid.Hy, grid.Hz, - horizontal_coordinate_data..., zᵃᵃᶜ, zᵃᵃᶠ, - horizontal_grid_spacing_data..., grid.Δz, - horizontal_area_data..., grid.radius) + if arch === architecture(grid) + return grid + else + + horizontal_coordinates = (:λᶜᶜᵃ, + :λᶠᶜᵃ, + :λᶜᶠᵃ, + :λᶠᶠᵃ, + :φᶜᶜᵃ, + :φᶠᶜᵃ, + :φᶜᶠᵃ, + :φᶠᶠᵃ) + + horizontal_grid_spacings = (:Δxᶜᶜᵃ, + :Δxᶠᶜᵃ, + :Δxᶜᶠᵃ, + :Δxᶠᶠᵃ, + :Δyᶜᶜᵃ, + :Δyᶜᶠᵃ, + :Δyᶠᶜᵃ, + :Δyᶠᶠᵃ) + + horizontal_areas = (:Azᶜᶜᵃ, + :Azᶠᶜᵃ, + :Azᶜᶠᵃ, + :Azᶠᶠᵃ) + + horizontal_grid_spacing_data = Tuple(arch_array(arch, getproperty(grid, name)) for name in horizontal_grid_spacings) + horizontal_coordinate_data = Tuple(arch_array(arch, getproperty(grid, name)) for name in horizontal_coordinates) + horizontal_area_data = Tuple(arch_array(arch, getproperty(grid, name)) for name in horizontal_areas) + + zᵃᵃᶜ = arch_array(arch, grid.zᵃᵃᶜ) + zᵃᵃᶠ = arch_array(arch, grid.zᵃᵃᶠ) + + TX, TY, TZ = topology(grid) + + new_grid = ConformalCubedSphereFaceGrid{TX, TY, TZ}(architecture, + grid.Nx, grid.Ny, grid.Nz, + grid.Hx, grid.Hy, grid.Hz, + horizontal_coordinate_data..., zᵃᵃᶜ, zᵃᵃᶠ, + horizontal_grid_spacing_data..., grid.Δz, + horizontal_area_data..., grid.radius) + + return new_grid + end end function Adapt.adapt_structure(to, grid::ConformalCubedSphereFaceGrid) diff --git a/src/Grids/latitude_longitude_grid.jl b/src/Grids/latitude_longitude_grid.jl index 7c56847808..b80dc8f1d4 100644 --- a/src/Grids/latitude_longitude_grid.jl +++ b/src/Grids/latitude_longitude_grid.jl @@ -227,25 +227,31 @@ all_z_nodes(::Type{Face}, grid::LatitudeLongitudeGrid) = grid.zᵃᵃᶠ @inline cpu_face_constructor_z(grid::ZRegLatLonGrid) = z_domain(grid) function on_architecture(new_arch, old_grid::LatitudeLongitudeGrid) - - size = (old_grid.Nx, old_grid.Ny, old_grid.Nz) - topo = topology(old_grid) - - x = cpu_face_constructor_x(old_grid) - y = cpu_face_constructor_y(old_grid) - z = cpu_face_constructor_z(old_grid) - - # Remove elements of size and new_halo in Flat directions as expected by grid - # constructor - size = pop_flat_elements(size, topo) - halo = pop_flat_elements(halo_size(old_grid), topo) - - return LatitudeLongitudeGrid(new_arch, eltype(old_grid); - size = size, - longitude = x, - latitude = y, - z = z, - halo = halo) + if new_arch === architecture(old_grid) + return old_grid + else + + size = (old_grid.Nx, old_grid.Ny, old_grid.Nz) + topo = topology(old_grid) + + x = cpu_face_constructor_x(old_grid) + y = cpu_face_constructor_y(old_grid) + z = cpu_face_constructor_z(old_grid) + + # Remove elements of size and new_halo in Flat directions as expected by grid + # constructor + size = pop_flat_elements(size, topo) + halo = pop_flat_elements(halo_size(old_grid), topo) + + new_grid LatitudeLongitudeGrid(new_arch, eltype(old_grid); + size = size, + longitude = x, + latitude = y, + z = z, + halo = halo) + + return new_grid + end end function Adapt.adapt_structure(to, grid::LatitudeLongitudeGrid) diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl index 5bb6ed61e9..e130de95a6 100644 --- a/src/Grids/rectilinear_grid.jl +++ b/src/Grids/rectilinear_grid.jl @@ -414,26 +414,30 @@ function with_halo(new_halo, old_grid::RectilinearGrid) end function on_architecture(new_arch, old_grid::RectilinearGrid) + if new_arch === architecture(old_grid) + return old_grid + else - size = (old_grid.Nx, old_grid.Ny, old_grid.Nz) - topo = topology(old_grid) + size = (old_grid.Nx, old_grid.Ny, old_grid.Nz) + topo = topology(old_grid) - x = cpu_face_constructor_x(old_grid) - y = cpu_face_constructor_y(old_grid) - z = cpu_face_constructor_z(old_grid) + x = cpu_face_constructor_x(old_grid) + y = cpu_face_constructor_y(old_grid) + z = cpu_face_constructor_z(old_grid) - # Remove elements of size and new_halo in Flat directions as expected by grid - # constructor - size = pop_flat_elements(size, topo) - halo = pop_flat_elements(halo_size(old_grid), topo) + # Remove elements of size and new_halo in Flat directions as expected by grid + # constructor + size = pop_flat_elements(size, topo) + halo = pop_flat_elements(halo_size(old_grid), topo) - new_grid = RectilinearGrid(new_arch, eltype(old_grid); - size = size, - x = x, y = y, z = z, - topology = topo, - halo = halo) + new_grid = RectilinearGrid(new_arch, eltype(old_grid); + size = size, + x = x, y = y, z = z, + topology = topo, + halo = halo) - return new_grid + return new_grid + end end ##### diff --git a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl index d106844f1e..b855dd8af9 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl @@ -12,7 +12,7 @@ using Oceananigans.TurbulenceClosures: AbstractTurbulenceClosure using Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivities: _top_tke_flux, CATKEVDArray import Oceananigans.Utils: launch! -import Oceananigans.Grids: validate_size, validate_halo +import Oceananigans.Grids: validate_size, validate_halo, on_architecture import Oceananigans.BoundaryConditions: fill_halo_regions! import Oceananigans.TurbulenceClosures: time_discretization, calculate_diffusivities! import Oceananigans.TurbulenceClosures: ∂ⱼ_τ₁ⱼ, ∂ⱼ_τ₂ⱼ, ∂ⱼ_τ₃ⱼ, ∇_dot_qᶜ @@ -91,6 +91,21 @@ ColumnEnsembleSize(; Nz, ensemble=(0, 0), Hz=1) = ColumnEnsembleSize(ensemble, N validate_size(TX, TY, TZ, e::ColumnEnsembleSize) = tuple(e.ensemble[1], e.ensemble[2], e.Nz) validate_halo(TX, TY, TZ, e::ColumnEnsembleSize) = tuple(0, 0, e.Hz) +function on_architecture(new_arch, old_grid::SingleColumnGrid) + if new_arch === architecture(old_grid) + return old_grid + else + size = ColumnEnsembleSize(; Nz=old_grid.Nz, ensemble=(old_grid.Nx, old_grid.Ny)) + halo = ColumnEnsembleSize(nothing, nothing, old_grid.Hz) + z = cpu_face_constructor_z(old_grid) + + new_grid = RectilinearGrid(new_arch, eltype(old_grid); size, z, halo, + topology = (Flat, Flat, Bounded)) + + return new_grid + end +end + @inline function time_discretization(closure_array::AbstractArray) first_closure = CUDA.@allowscalar first(closure_array) # assumes all closures have same time-discretization return time_discretization(first_closure) From 3b31b0b9f2aa34330e7a06ff1aee0b697fa5cbc6 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 19:53:38 -0500 Subject: [PATCH 124/140] Refactor on_architecture --- src/Architectures.jl | 1 + src/Grids/conformal_cubed_sphere_face_grid.jl | 145 +++++++++--------- src/Grids/latitude_longitude_grid.jl | 42 ++--- src/Grids/rectilinear_grid.jl | 31 ++-- .../single_column_model_mode.jl | 17 +- 5 files changed, 99 insertions(+), 137 deletions(-) diff --git a/src/Architectures.jl b/src/Architectures.jl index 929e38d082..d0ae98f1c3 100644 --- a/src/Architectures.jl +++ b/src/Architectures.jl @@ -69,6 +69,7 @@ arch_array(::GPU, A::CuArray) = A arch_array(arch, A::AbstractRange) = A arch_array(arch, a::OffsetArray) = OffsetArray(arch_array(arch, a.parent), a.offsets...) +arch_array(arch, ::Nothing) = nothing device_event(arch) = Event(device(arch)) diff --git a/src/Grids/conformal_cubed_sphere_face_grid.jl b/src/Grids/conformal_cubed_sphere_face_grid.jl index a82af5c512..9c8cc91a52 100644 --- a/src/Grids/conformal_cubed_sphere_face_grid.jl +++ b/src/Grids/conformal_cubed_sphere_face_grid.jl @@ -8,36 +8,36 @@ using Oceananigans struct ConformalCubedSphereFaceGrid{FT, TX, TY, TZ, A, R, Arch} <: AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ, Arch} architecture :: Arch - Nx :: Int - Ny :: Int - Nz :: Int - Hx :: Int - Hy :: Int - Hz :: Int - λᶜᶜᵃ :: A - λᶠᶜᵃ :: A - λᶜᶠᵃ :: A - λᶠᶠᵃ :: A - φᶜᶜᵃ :: A - φᶠᶜᵃ :: A - φᶜᶠᵃ :: A - φᶠᶠᵃ :: A - zᵃᵃᶜ :: R - zᵃᵃᶠ :: R - Δxᶜᶜᵃ :: A - Δxᶠᶜᵃ :: A - Δxᶜᶠᵃ :: A - Δxᶠᶠᵃ :: A - Δyᶜᶜᵃ :: A - Δyᶜᶠᵃ :: A - Δyᶠᶜᵃ :: A - Δyᶠᶠᵃ :: A - Δz :: FT - Azᶜᶜᵃ :: A - Azᶠᶜᵃ :: A - Azᶜᶠᵃ :: A - Azᶠᶠᵃ :: A - radius :: FT + Nx :: Int + Ny :: Int + Nz :: Int + Hx :: Int + Hy :: Int + Hz :: Int + λᶜᶜᵃ :: A + λᶠᶜᵃ :: A + λᶜᶠᵃ :: A + λᶠᶠᵃ :: A + φᶜᶜᵃ :: A + φᶠᶜᵃ :: A + φᶜᶠᵃ :: A + φᶠᶠᵃ :: A + zᵃᵃᶜ :: R + zᵃᵃᶠ :: R + Δxᶜᶜᵃ :: A + Δxᶠᶜᵃ :: A + Δxᶜᶠᵃ :: A + Δxᶠᶠᵃ :: A + Δyᶜᶜᵃ :: A + Δyᶜᶠᵃ :: A + Δyᶠᶜᵃ :: A + Δyᶠᶠᵃ :: A + Δz :: FT + Azᶜᶜᵃ :: A + Azᶠᶜᵃ :: A + Azᶜᶠᵃ :: A + Azᶠᶠᵃ :: A + radius :: FT function ConformalCubedSphereFaceGrid{TX, TY, TZ}(architecture::Arch, Nx, Ny, Nz, @@ -275,51 +275,46 @@ function ConformalCubedSphereFaceGrid(filepath::AbstractString, architecture = C end function on_architecture(arch, grid::ConformalCubedSphereFaceGrid) - if arch === architecture(grid) - return grid - else - - horizontal_coordinates = (:λᶜᶜᵃ, - :λᶠᶜᵃ, - :λᶜᶠᵃ, - :λᶠᶠᵃ, - :φᶜᶜᵃ, - :φᶠᶜᵃ, - :φᶜᶠᵃ, - :φᶠᶠᵃ) - - horizontal_grid_spacings = (:Δxᶜᶜᵃ, - :Δxᶠᶜᵃ, - :Δxᶜᶠᵃ, - :Δxᶠᶠᵃ, - :Δyᶜᶜᵃ, - :Δyᶜᶠᵃ, - :Δyᶠᶜᵃ, - :Δyᶠᶠᵃ) - - horizontal_areas = (:Azᶜᶜᵃ, - :Azᶠᶜᵃ, - :Azᶜᶠᵃ, - :Azᶠᶠᵃ) - - horizontal_grid_spacing_data = Tuple(arch_array(arch, getproperty(grid, name)) for name in horizontal_grid_spacings) - horizontal_coordinate_data = Tuple(arch_array(arch, getproperty(grid, name)) for name in horizontal_coordinates) - horizontal_area_data = Tuple(arch_array(arch, getproperty(grid, name)) for name in horizontal_areas) - - zᵃᵃᶜ = arch_array(arch, grid.zᵃᵃᶜ) - zᵃᵃᶠ = arch_array(arch, grid.zᵃᵃᶠ) - - TX, TY, TZ = topology(grid) - - new_grid = ConformalCubedSphereFaceGrid{TX, TY, TZ}(architecture, - grid.Nx, grid.Ny, grid.Nz, - grid.Hx, grid.Hy, grid.Hz, - horizontal_coordinate_data..., zᵃᵃᶜ, zᵃᵃᶠ, - horizontal_grid_spacing_data..., grid.Δz, - horizontal_area_data..., grid.radius) - - return new_grid - end + horizontal_coordinates = (:λᶜᶜᵃ, + :λᶠᶜᵃ, + :λᶜᶠᵃ, + :λᶠᶠᵃ, + :φᶜᶜᵃ, + :φᶠᶜᵃ, + :φᶜᶠᵃ, + :φᶠᶠᵃ) + + horizontal_grid_spacings = (:Δxᶜᶜᵃ, + :Δxᶠᶜᵃ, + :Δxᶜᶠᵃ, + :Δxᶠᶠᵃ, + :Δyᶜᶜᵃ, + :Δyᶜᶠᵃ, + :Δyᶠᶜᵃ, + :Δyᶠᶠᵃ) + + horizontal_areas = (:Azᶜᶜᵃ, + :Azᶠᶜᵃ, + :Azᶜᶠᵃ, + :Azᶠᶠᵃ) + + horizontal_grid_spacing_data = Tuple(arch_array(arch, getproperty(grid, name)) for name in horizontal_grid_spacings) + horizontal_coordinate_data = Tuple(arch_array(arch, getproperty(grid, name)) for name in horizontal_coordinates) + horizontal_area_data = Tuple(arch_array(arch, getproperty(grid, name)) for name in horizontal_areas) + + zᵃᵃᶜ = arch_array(arch, grid.zᵃᵃᶜ) + zᵃᵃᶠ = arch_array(arch, grid.zᵃᵃᶠ) + + TX, TY, TZ = topology(grid) + + new_grid = ConformalCubedSphereFaceGrid{TX, TY, TZ}(architecture, + grid.Nx, grid.Ny, grid.Nz, + grid.Hx, grid.Hy, grid.Hz, + horizontal_coordinate_data..., zᵃᵃᶜ, zᵃᵃᶠ, + horizontal_grid_spacing_data..., grid.Δz, + horizontal_area_data..., grid.radius) + + return new_grid end function Adapt.adapt_structure(to, grid::ConformalCubedSphereFaceGrid) diff --git a/src/Grids/latitude_longitude_grid.jl b/src/Grids/latitude_longitude_grid.jl index b80dc8f1d4..c84247d138 100644 --- a/src/Grids/latitude_longitude_grid.jl +++ b/src/Grids/latitude_longitude_grid.jl @@ -227,31 +227,23 @@ all_z_nodes(::Type{Face}, grid::LatitudeLongitudeGrid) = grid.zᵃᵃᶠ @inline cpu_face_constructor_z(grid::ZRegLatLonGrid) = z_domain(grid) function on_architecture(new_arch, old_grid::LatitudeLongitudeGrid) - if new_arch === architecture(old_grid) - return old_grid - else - - size = (old_grid.Nx, old_grid.Ny, old_grid.Nz) - topo = topology(old_grid) - - x = cpu_face_constructor_x(old_grid) - y = cpu_face_constructor_y(old_grid) - z = cpu_face_constructor_z(old_grid) - - # Remove elements of size and new_halo in Flat directions as expected by grid - # constructor - size = pop_flat_elements(size, topo) - halo = pop_flat_elements(halo_size(old_grid), topo) - - new_grid LatitudeLongitudeGrid(new_arch, eltype(old_grid); - size = size, - longitude = x, - latitude = y, - z = z, - halo = halo) - - return new_grid - end + old_properties = (old_grid.Δλᶠᵃᵃ, old_grid.Δλᶜᵃᵃ, old_grid.λᶠᵃᵃ, old_grid.λᶜᵃᵃ, + old_grid.Δφᵃᶠᵃ, old_grid.Δφᵃᶜᵃ, old_grid.φᵃᶠᵃ, old_grid.φᵃᶜᵃ, + old_grid.Δzᵃᵃᶠ, old_grid.Δzᵃᵃᶜ, old_grid.zᵃᵃᶠ, old_grid.zᵃᵃᶜ, + old_grid.Δxᶠᶜ, old_grid.Δxᶜᶠ, old_grid.Δxᶠᶠ, old_grid.Δxᶜᶜ, + old_grid.Δyᶠᶜ, old_grid.Δyᶜᶠ, + old_grid.Azᶠᶜ, old_grid.Azᶜᶠ, old_grid.Azᶠᶠ, old_grid.Azᶜᶜ) + + new_properties = arch_array.(new_arch, old_properties) + + TX, TY, TZ = topology(old_grid) + + return LatitudeLongitudeGrid{TX, TY, TZ}(new_arch, + old_grid.Nλ, old_grid.Nφ, old_grid.Nz, + old_grid.Hλ, old_grid.Hφ, old_grid.Hz, + old_grid.Lx, old_grid.Ly, old_grid.Lz, + new_properties... + old_grid.radius) end function Adapt.adapt_structure(to, grid::LatitudeLongitudeGrid) diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl index e130de95a6..4439dde2b8 100644 --- a/src/Grids/rectilinear_grid.jl +++ b/src/Grids/rectilinear_grid.jl @@ -414,30 +414,19 @@ function with_halo(new_halo, old_grid::RectilinearGrid) end function on_architecture(new_arch, old_grid::RectilinearGrid) - if new_arch === architecture(old_grid) - return old_grid - else - - size = (old_grid.Nx, old_grid.Ny, old_grid.Nz) - topo = topology(old_grid) + old_properties = (old_grid.Δxᶠᵃᵃ, old_grid.Δxᶜᵃᵃ, old_grid.xᶠᵃᵃ, old_grid.xᶜᵃᵃ, + old_grid.Δyᵃᶠᵃ, old_grid.Δyᵃᶜᵃ, old_grid.yᵃᶠᵃ, old_grid.yᵃᶜᵃ, + old_grid.Δzᵃᵃᶠ, old_grid.Δzᵃᵃᶜ, old_grid.zᵃᵃᶠ, old_grid.zᵃᵃᶜ) - x = cpu_face_constructor_x(old_grid) - y = cpu_face_constructor_y(old_grid) - z = cpu_face_constructor_z(old_grid) + new_properties = arch_array.(new_arch, old_properties) - # Remove elements of size and new_halo in Flat directions as expected by grid - # constructor - size = pop_flat_elements(size, topo) - halo = pop_flat_elements(halo_size(old_grid), topo) + TX, TY, TZ = topology(old_grid) - new_grid = RectilinearGrid(new_arch, eltype(old_grid); - size = size, - x = x, y = y, z = z, - topology = topo, - halo = halo) - - return new_grid - end + return RectilinearGrid{TX, TY, TZ}(new_arch, + old_grid.Nx, old_grid.Ny, old_grid.Nz, + old_grid.Hx, old_grid.Hy, old_grid.Hz, + old_grid.Lx, old_grid.Ly, old_grid.Lz, + new_properties...) end ##### diff --git a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl index b855dd8af9..d106844f1e 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl @@ -12,7 +12,7 @@ using Oceananigans.TurbulenceClosures: AbstractTurbulenceClosure using Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivities: _top_tke_flux, CATKEVDArray import Oceananigans.Utils: launch! -import Oceananigans.Grids: validate_size, validate_halo, on_architecture +import Oceananigans.Grids: validate_size, validate_halo import Oceananigans.BoundaryConditions: fill_halo_regions! import Oceananigans.TurbulenceClosures: time_discretization, calculate_diffusivities! import Oceananigans.TurbulenceClosures: ∂ⱼ_τ₁ⱼ, ∂ⱼ_τ₂ⱼ, ∂ⱼ_τ₃ⱼ, ∇_dot_qᶜ @@ -91,21 +91,6 @@ ColumnEnsembleSize(; Nz, ensemble=(0, 0), Hz=1) = ColumnEnsembleSize(ensemble, N validate_size(TX, TY, TZ, e::ColumnEnsembleSize) = tuple(e.ensemble[1], e.ensemble[2], e.Nz) validate_halo(TX, TY, TZ, e::ColumnEnsembleSize) = tuple(0, 0, e.Hz) -function on_architecture(new_arch, old_grid::SingleColumnGrid) - if new_arch === architecture(old_grid) - return old_grid - else - size = ColumnEnsembleSize(; Nz=old_grid.Nz, ensemble=(old_grid.Nx, old_grid.Ny)) - halo = ColumnEnsembleSize(nothing, nothing, old_grid.Hz) - z = cpu_face_constructor_z(old_grid) - - new_grid = RectilinearGrid(new_arch, eltype(old_grid); size, z, halo, - topology = (Flat, Flat, Bounded)) - - return new_grid - end -end - @inline function time_discretization(closure_array::AbstractArray) first_closure = CUDA.@allowscalar first(closure_array) # assumes all closures have same time-discretization return time_discretization(first_closure) From 2a5197d4324898d694c4d223602b9e0292521c94 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 19:55:04 -0500 Subject: [PATCH 125/140] Missing comma --- src/Grids/latitude_longitude_grid.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Grids/latitude_longitude_grid.jl b/src/Grids/latitude_longitude_grid.jl index c84247d138..5a08ee0a81 100644 --- a/src/Grids/latitude_longitude_grid.jl +++ b/src/Grids/latitude_longitude_grid.jl @@ -242,7 +242,7 @@ function on_architecture(new_arch, old_grid::LatitudeLongitudeGrid) old_grid.Nλ, old_grid.Nφ, old_grid.Nz, old_grid.Hλ, old_grid.Hφ, old_grid.Hz, old_grid.Lx, old_grid.Ly, old_grid.Lz, - new_properties... + new_properties..., old_grid.radius) end From cae064f0b3c25ebec1299ce3329b9fdd42c84037 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 20:12:52 -0500 Subject: [PATCH 126/140] Write loop in on_architecture explicitly --- src/Grids/latitude_longitude_grid.jl | 2 +- src/Grids/rectilinear_grid.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/Grids/latitude_longitude_grid.jl b/src/Grids/latitude_longitude_grid.jl index 5a08ee0a81..9fb088d6bb 100644 --- a/src/Grids/latitude_longitude_grid.jl +++ b/src/Grids/latitude_longitude_grid.jl @@ -234,7 +234,7 @@ function on_architecture(new_arch, old_grid::LatitudeLongitudeGrid) old_grid.Δyᶠᶜ, old_grid.Δyᶜᶠ, old_grid.Azᶠᶜ, old_grid.Azᶜᶠ, old_grid.Azᶠᶠ, old_grid.Azᶜᶜ) - new_properties = arch_array.(new_arch, old_properties) + new_properties = Tuple(arch_array(new_arch, p) for p in old_properties) TX, TY, TZ = topology(old_grid) diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl index 4439dde2b8..a0214a1388 100644 --- a/src/Grids/rectilinear_grid.jl +++ b/src/Grids/rectilinear_grid.jl @@ -418,7 +418,7 @@ function on_architecture(new_arch, old_grid::RectilinearGrid) old_grid.Δyᵃᶠᵃ, old_grid.Δyᵃᶜᵃ, old_grid.yᵃᶠᵃ, old_grid.yᵃᶜᵃ, old_grid.Δzᵃᵃᶠ, old_grid.Δzᵃᵃᶜ, old_grid.zᵃᵃᶠ, old_grid.zᵃᵃᶜ) - new_properties = arch_array.(new_arch, old_properties) + new_properties = Tuple(arch_array(new_arch, p) for p in old_properties) TX, TY, TZ = topology(old_grid) From e89900ea15f9465f544c0d91774e92f5dce46547 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 20:14:58 -0500 Subject: [PATCH 127/140] Bugfix in on_architecture for conformal cubed sphere grid --- src/Grids/conformal_cubed_sphere_face_grid.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Grids/conformal_cubed_sphere_face_grid.jl b/src/Grids/conformal_cubed_sphere_face_grid.jl index 9c8cc91a52..cbc3ba53ff 100644 --- a/src/Grids/conformal_cubed_sphere_face_grid.jl +++ b/src/Grids/conformal_cubed_sphere_face_grid.jl @@ -307,7 +307,7 @@ function on_architecture(arch, grid::ConformalCubedSphereFaceGrid) TX, TY, TZ = topology(grid) - new_grid = ConformalCubedSphereFaceGrid{TX, TY, TZ}(architecture, + new_grid = ConformalCubedSphereFaceGrid{TX, TY, TZ}(arch, grid.Nx, grid.Ny, grid.Nz, grid.Hx, grid.Hy, grid.Hz, horizontal_coordinate_data..., zᵃᵃᶜ, zᵃᵃᶠ, From 3a391262c9939b9fcf94c573a1ad6f586aa38160 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 20:16:43 -0500 Subject: [PATCH 128/140] Add numbers to the list of architecture-agnostic objs --- src/Architectures.jl | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/Architectures.jl b/src/Architectures.jl index d0ae98f1c3..54f38d509e 100644 --- a/src/Architectures.jl +++ b/src/Architectures.jl @@ -62,14 +62,15 @@ child_architecture(arch) = arch array_type(::CPU) = Array array_type(::GPU) = CuArray -arch_array(::CPU, A::Array) = A -arch_array(::CPU, A::CuArray) = Array(A) -arch_array(::GPU, A::Array) = CuArray(A) -arch_array(::GPU, A::CuArray) = A +arch_array(::CPU, a::Array) = a +arch_array(::CPU, a::CuArray) = Array(a) +arch_array(::GPU, a::Array) = CuArray(a) +arch_array(::GPU, a::CuArray) = a -arch_array(arch, A::AbstractRange) = A +arch_array(arch, a::AbstractRange) = a arch_array(arch, a::OffsetArray) = OffsetArray(arch_array(arch, a.parent), a.offsets...) arch_array(arch, ::Nothing) = nothing +arch_array(arch, a::Number) = a device_event(arch) = Event(device(arch)) From e9f4e51e51e2b71d27cd5868902a03f256c2f242 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 21:41:08 -0500 Subject: [PATCH 129/140] Slightly reorganize computed field test --- test/test_computed_field.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/test_computed_field.jl b/test/test_computed_field.jl index 4ff1a62f99..8d18dc501d 100644 --- a/test/test_computed_field.jl +++ b/test/test_computed_field.jl @@ -507,7 +507,8 @@ for arch in archs @test try compute!(Field(tke_ccc )); true; catch; false; end computed_tke = Field(tke_ccc) - @test (compute!(computed_tke); all(interior(computed_tke)[2:3, 2:3, 2:3] .== 9/2)) + compute!(computed_tke) + @test all(interior(computed_tke)[2:3, 2:3, 2:3] .== 9/2) end @testset "Computations with Fields [$(typeof(arch))]" begin From 925653f76d22817f659264ccbcf138dc9236c357 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 22:59:17 -0500 Subject: [PATCH 130/140] Add functionality to set! for cpu-gpu transfers of views --- src/Fields/set!.jl | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/Fields/set!.jl b/src/Fields/set!.jl index fc6ebb9f49..6383ecb5dc 100644 --- a/src/Fields/set!.jl +++ b/src/Fields/set!.jl @@ -41,8 +41,15 @@ end function set!(u::Field, v::Field) if architecture(u) === architecture(v) parent(u) .= parent(v) - else - copyto!(parent(u), parent(v)) + elseif parent(v) isa SubArray + u_parent = parent(u) + v_parent = parent(v) + # If u_parent is a view, we have to convert to an Array. + # If u_parent or v_parent is on the GPU, we don't expect + # SubArray. + u_parent isa SubArray && (u_parent = Array(u_parent)) + v_parent isa SubArray && (v_parent = Array(v_parent)) + copyto!(u_parent, v_parent) end return nothing end From d619484c93b23cda4847e132674a8f80ce9703e3 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 13 Jan 2022 23:35:57 -0500 Subject: [PATCH 131/140] Bugfix on_architecture for lat-lon grid --- src/Grids/latitude_longitude_grid.jl | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/Grids/latitude_longitude_grid.jl b/src/Grids/latitude_longitude_grid.jl index 9fb088d6bb..84832629a6 100644 --- a/src/Grids/latitude_longitude_grid.jl +++ b/src/Grids/latitude_longitude_grid.jl @@ -227,12 +227,12 @@ all_z_nodes(::Type{Face}, grid::LatitudeLongitudeGrid) = grid.zᵃᵃᶠ @inline cpu_face_constructor_z(grid::ZRegLatLonGrid) = z_domain(grid) function on_architecture(new_arch, old_grid::LatitudeLongitudeGrid) - old_properties = (old_grid.Δλᶠᵃᵃ, old_grid.Δλᶜᵃᵃ, old_grid.λᶠᵃᵃ, old_grid.λᶜᵃᵃ, - old_grid.Δφᵃᶠᵃ, old_grid.Δφᵃᶜᵃ, old_grid.φᵃᶠᵃ, old_grid.φᵃᶜᵃ, - old_grid.Δzᵃᵃᶠ, old_grid.Δzᵃᵃᶜ, old_grid.zᵃᵃᶠ, old_grid.zᵃᵃᶜ, - old_grid.Δxᶠᶜ, old_grid.Δxᶜᶠ, old_grid.Δxᶠᶠ, old_grid.Δxᶜᶜ, - old_grid.Δyᶠᶜ, old_grid.Δyᶜᶠ, - old_grid.Azᶠᶜ, old_grid.Azᶜᶠ, old_grid.Azᶠᶠ, old_grid.Azᶜᶜ) + old_properties = (old_grid.Δλᶠᵃᵃ, old_grid.Δλᶜᵃᵃ, old_grid.λᶠᵃᵃ, old_grid.λᶜᵃᵃ, + old_grid.Δφᵃᶠᵃ, old_grid.Δφᵃᶜᵃ, old_grid.φᵃᶠᵃ, old_grid.φᵃᶜᵃ, + old_grid.Δzᵃᵃᶠ, old_grid.Δzᵃᵃᶜ, old_grid.zᵃᵃᶠ, old_grid.zᵃᵃᶜ, + old_grid.Δxᶠᶜᵃ, old_grid.Δxᶜᶠᵃ, old_grid.Δxᶠᶠᵃ, old_grid.Δxᶜᶜᵃ, + old_grid.Δyᶠᶜᵃ, old_grid.Δyᶜᶠᵃ, + old_grid.Azᶠᶜᵃ, old_grid.Azᶜᶠᵃ, old_grid.Azᶠᶠᵃ, old_grid.Azᶜᶜᵃ) new_properties = Tuple(arch_array(new_arch, p) for p in old_properties) From b98fb590eea7320c82e0461ff39e627eced11332 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 14 Jan 2022 00:59:01 -0500 Subject: [PATCH 132/140] Regular grid shortcut plus fix on_architecture --- src/Grids/latitude_longitude_grid.jl | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/src/Grids/latitude_longitude_grid.jl b/src/Grids/latitude_longitude_grid.jl index 84832629a6..f1b24d2ea1 100644 --- a/src/Grids/latitude_longitude_grid.jl +++ b/src/Grids/latitude_longitude_grid.jl @@ -75,10 +75,12 @@ struct LatitudeLongitudeGrid{FT, TX, TY, TZ, M, MY, FX, FY, FZ, VX, VY, VZ, Arch end const XRegLatLonGrid = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number} -const YRegLatLonGrid = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number} -const ZRegLatLonGrid = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number} +const YRegLatLonGrid = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number} +const ZRegLatLonGrid = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number} const HRegLatLonGrid = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number, <:Number} +regular_dimensions(::ZRegLatLonGrid) = tuple(3) + """ latitude, longitude and z can be a 2-tuple that specifies the end of the domain (see RegularRectilinearDomain) @@ -87,14 +89,14 @@ or an array or function that specifies the faces (see VerticallyStretchedRectili """ function LatitudeLongitudeGrid(architecture::AbstractArchitecture = CPU(), - FT::DataType = Float64; + FT::DataType = Float64; precompute_metrics = false, size, latitude, longitude, - z, - radius=R_Earth, - halo=(1, 1, 1)) + z, + radius = R_Earth, + halo = (1, 1, 1)) Nλ, Nφ, Nz, Hλ, Hφ, Hz, latitude, longitude, topo = validate_lat_lon_grid_args(latitude, longitude, size, halo) @@ -239,8 +241,8 @@ function on_architecture(new_arch, old_grid::LatitudeLongitudeGrid) TX, TY, TZ = topology(old_grid) return LatitudeLongitudeGrid{TX, TY, TZ}(new_arch, - old_grid.Nλ, old_grid.Nφ, old_grid.Nz, - old_grid.Hλ, old_grid.Hφ, old_grid.Hz, + old_grid.Nx, old_grid.Ny, old_grid.Nz, + old_grid.Hx, old_grid.Hy, old_grid.Hz, old_grid.Lx, old_grid.Ly, old_grid.Lz, new_properties..., old_grid.radius) From 7701b42145c891e919f0f3bd49b34159215a9bac Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 14 Jan 2022 00:59:14 -0500 Subject: [PATCH 133/140] Shortcut for regular grids --- .../metric_field_reductions.jl | 38 +++++++++++++++---- src/Grids/grid_utils.jl | 2 + src/Grids/rectilinear_grid.jl | 21 +++++++--- 3 files changed, 48 insertions(+), 13 deletions(-) diff --git a/src/AbstractOperations/metric_field_reductions.jl b/src/AbstractOperations/metric_field_reductions.jl index 65a18f761b..af674ff836 100644 --- a/src/AbstractOperations/metric_field_reductions.jl +++ b/src/AbstractOperations/metric_field_reductions.jl @@ -1,5 +1,7 @@ using Statistics: mean!, sum! +using Oceananigans.Utils: tupleit +using Oceananigans.Grids: regular_dimensions import Oceananigans.Fields: Reduction import Oceananigans: short_show @@ -22,25 +24,41 @@ reduction_grid_metric(dims) = dims === tuple(1) ? Δx : ##### Metric reductions ##### -"""docstring...""" struct Average end function Reduction(avg::Average, field::AbstractField; dims) dx = reduction_grid_metric(dims) + dims = tupleit(dims) - # Compute "size" (length, area, or volume) of averaging region - metric = GridMetricOperation(location(field), dx, field.grid) - L = sum(metric; dims) - L⁻¹_field_dx = field * dx / L + if dims === regular_dimensions(field.grid) # shortcut! + return Reduction(mean!, field; dims) + else + # Compute "size" (length, area, or volume) of averaging region + metric = GridMetricOperation(location(field), dx, field.grid) + L = sum(metric; dims) - return Reduction(sum!, L⁻¹_field_dx; dims) + # Construct summand of the Average + L⁻¹_field_dx = field * dx / L + + return Reduction(sum!, L⁻¹_field_dx; dims) + end end +""" + Average(field; dims=:) + +Return `Reduction` representing a spatial average of `field` over `dims`. + +Over regularly-spaced dimensions this is equivalent to a numerical `mean!`. + +Over dimensions of variable spacing, `field` is multipled by the +appropriate grid length, area or volume, and divided by the total +spatial extent of the interval. +""" Average(field::AbstractField; dims=:) = Reduction(Average(), field; dims) const AveragedField = Field{<:Any, <:Any, <:Any, <:Reduction{<:Average}} -"""docstring...""" struct Integral end function Reduction(int::Integral, field::AbstractField; dims) @@ -48,7 +66,11 @@ function Reduction(int::Integral, field::AbstractField; dims) return Reduction(sum!, field * dx; dims) end -# Convenience +""" + Integral(field; dims=:) + +Return a `Reduction` representing a spatial integral of `field` over `dims`. +""" Integral(field::AbstractField; dims=:) = Reduction(Integral(), field; dims) const IntegratedField = Field{<:Any, <:Any, <:Any, <:Reduction{<:Integral}} diff --git a/src/Grids/grid_utils.jl b/src/Grids/grid_utils.jl index 7e2f88b696..3093c98412 100644 --- a/src/Grids/grid_utils.jl +++ b/src/Grids/grid_utils.jl @@ -132,6 +132,8 @@ Return 1, which is the 'length' of a field along a reduced dimension. @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) = () + ##### ##### << Indexing >> ##### diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl index a0214a1388..09741bfb7a 100644 --- a/src/Grids/rectilinear_grid.jl +++ b/src/Grids/rectilinear_grid.jl @@ -53,11 +53,22 @@ struct RectilinearGrid{FT, TX, TY, TZ, FX, FY, FZ, VX, VY, VZ, Arch} <: Abstract end end -const XRegRectilinearGrid = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number} -const YRegRectilinearGrid = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Number} -const ZRegRectilinearGrid = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number} -const HRegRectilinearGrid = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number, <:Number} -const RegRectilinearGrid = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number, <:Number, <:Number} +const XRegRectilinearGrid = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number} +const YRegRectilinearGrid = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Number} +const ZRegRectilinearGrid = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number} +const HRegRectilinearGrid = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number, <:Number} +const XYRegRectilinearGrid = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number, <:Number} +const XZRegRectilinearGrid = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number, <:Any, <:Number} +const YZRegRectilinearGrid = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Number, <:Number} +const RegRectilinearGrid = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number, <:Number, <:Number} + +regular_dimensions(::XRegRectilinearGrid) = tuple(1) +regular_dimensions(::YRegRectilinearGrid) = tuple(2) +regular_dimensions(::ZRegRectilinearGrid) = tuple(3) +regular_dimensions(::XYRegRectilinearGrid) = (1, 2) +regular_dimensions(::XZRegRectilinearGrid) = (1, 3) +regular_dimensions(::YZRegRectilinearGrid) = (2, 3) +regular_dimensions(::RegRectilinearGrid) = (1, 2, 3) """ RectilinearGrid([architecture = CPU(), FT = Float64]; From adb2e8b5c4750008ab3bd1788044803103d26ecd Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 14 Jan 2022 00:59:35 -0500 Subject: [PATCH 134/140] Add abstol to some netcdf time averaging tests --- test/test_netcdf_output_writer.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_netcdf_output_writer.jl b/test/test_netcdf_output_writer.jl index c1e72c0af6..4d2fac4fb9 100644 --- a/test/test_netcdf_output_writer.jl +++ b/test/test_netcdf_output_writer.jl @@ -608,7 +608,7 @@ function test_netcdf_time_averaging(arch) for (n, t) in enumerate(single_ds["time"][2:end]) averaging_times = [t - n*Δt for n in 0:stride:window_size-1 if t - n*Δt >= 0] - @test all(isapprox.(single_ds["c1"][:, n+1], c̄1(averaging_times), rtol=rtol)) + @test all(isapprox.(single_ds["c1"][:, n+1], c̄1(averaging_times), rtol=rtol, atol=rtol)) end close(single_ds) From cb615bef794f501a4818437bb63cc7b21a2642d6 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 14 Jan 2022 08:02:43 -0700 Subject: [PATCH 135/140] Add Colon to list of valid dims --- src/AbstractOperations/metric_field_reductions.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/AbstractOperations/metric_field_reductions.jl b/src/AbstractOperations/metric_field_reductions.jl index af674ff836..943fb77875 100644 --- a/src/AbstractOperations/metric_field_reductions.jl +++ b/src/AbstractOperations/metric_field_reductions.jl @@ -18,6 +18,7 @@ reduction_grid_metric(dims) = dims === tuple(1) ? Δx : dims === (1, 3) ? Ay : dims === (2, 3) ? Ax : dims === (1, 2, 3) ? volume : + dims isa Colon ? volume : throw(ArgumentError("Cannot determine grid metric for reducing over dims = $dims")) ##### From 7ba6eeff5945fdef9a45684d16d3a78bff13fc7c Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 14 Jan 2022 12:05:44 -0500 Subject: [PATCH 136/140] Minor updates and bugfix for Average with dims=: --- .../metric_field_reductions.jl | 4 ++-- src/AbstractOperations/multiary_operations.jl | 2 +- .../field_boundary_conditions.jl | 18 ++++++++++-------- src/Fields/abstract_field.jl | 14 +++++++++----- src/Fields/field.jl | 10 ++++++++-- src/Fields/function_field.jl | 4 ++-- src/Grids/rectilinear_grid.jl | 6 +++--- src/OutputReaders/field_time_series.jl | 2 +- .../leith_enstrophy_diffusivity.jl | 2 +- ...an_large_eddy_simulation_regression_test.jl | 2 ++ test/test_regression.jl | 13 ++++++------- 11 files changed, 45 insertions(+), 32 deletions(-) diff --git a/src/AbstractOperations/metric_field_reductions.jl b/src/AbstractOperations/metric_field_reductions.jl index 943fb77875..ae79a13125 100644 --- a/src/AbstractOperations/metric_field_reductions.jl +++ b/src/AbstractOperations/metric_field_reductions.jl @@ -18,7 +18,6 @@ reduction_grid_metric(dims) = dims === tuple(1) ? Δx : dims === (1, 3) ? Ay : dims === (2, 3) ? Ax : dims === (1, 2, 3) ? volume : - dims isa Colon ? volume : throw(ArgumentError("Cannot determine grid metric for reducing over dims = $dims")) ##### @@ -28,8 +27,8 @@ reduction_grid_metric(dims) = dims === tuple(1) ? Δx : struct Average end function Reduction(avg::Average, field::AbstractField; dims) + dims = dims isa Colon ? (1, 2, 3) : tupleit(dims) dx = reduction_grid_metric(dims) - dims = tupleit(dims) if dims === regular_dimensions(field.grid) # shortcut! return Reduction(mean!, field; dims) @@ -63,6 +62,7 @@ const AveragedField = Field{<:Any, <:Any, <:Any, <:Reduction{<:Average}} struct Integral end function Reduction(int::Integral, field::AbstractField; dims) + dims = dims isa Colon ? (1, 2, 3) : tupleit(dims) dx = reduction_grid_metric(dims) return Reduction(sum!, field * dx; dims) end diff --git a/src/AbstractOperations/multiary_operations.jl b/src/AbstractOperations/multiary_operations.jl index 544705d783..443c7759bf 100644 --- a/src/AbstractOperations/multiary_operations.jl +++ b/src/AbstractOperations/multiary_operations.jl @@ -41,7 +41,7 @@ function define_multiary_operator(op) grid = Oceananigans.AbstractOperations.validate_grid(args...) # Convert any functions to FunctionFields - args = Tuple(Oceananigans.Fields.fieldify(Lop, a, grid) for a in args) + args = Tuple(Oceananigans.Fields.fieldify_function(Lop, a, grid) for a in args) Largs = Tuple(Oceananigans.Fields.location(a) for a in args) return Oceananigans.AbstractOperations._multiary_operation(Lop, $op, args, Largs, grid) diff --git a/src/BoundaryConditions/field_boundary_conditions.jl b/src/BoundaryConditions/field_boundary_conditions.jl index c1c37b2ffc..926e5beff7 100644 --- a/src/BoundaryConditions/field_boundary_conditions.jl +++ b/src/BoundaryConditions/field_boundary_conditions.jl @@ -126,9 +126,7 @@ boundary conditions for prognostic model field boundary conditions. Currently, there is no support `ContinuousBoundaryFunction` for immersed boundary conditions. """ -function regularize_field_boundary_conditions(bcs::FieldBoundaryConditions, grid::AbstractGrid, field_name, - prognostic_field_names=nothing) - +function regularize_field_boundary_conditions(bcs::FieldBoundaryConditions, grid::AbstractGrid, field_name::Symbol, prognostic_field_names=nothing) topo = topology(grid) loc = assumed_field_location(field_name) @@ -147,13 +145,17 @@ function regularize_field_boundary_conditions(bcs::FieldBoundaryConditions, grid return FieldBoundaryConditions(west, east, south, north, bottom, top, immersed) end -regularize_field_boundary_conditions(boundary_conditions::NamedTuple, grid, prognostic_field_names) = +# For nested NamedTuples of boundary conditions (eg diffusivity boundary conditions) +regularize_field_boundary_conditions(boundary_conditions::NamedTuple, grid::AbstractGrid, group_name::Symbol, prognostic_field_names=nothing) = NamedTuple(field_name => regularize_field_boundary_conditions(field_bcs, grid, field_name, prognostic_field_names) for (field_name, field_bcs) in pairs(boundary_conditions)) -# For nested NamedTuples of boundary conditions (eg diffusivity boundary conditions) -regularize_field_boundary_conditions(boundary_conditions::NamedTuple, grid, ::Symbol, prognostic_field_names) = +regularize_field_boundary_conditions(::Missing, grid::AbstractGrid, field_name::Symbol, prognostic_field_names=nothing) = missing + +##### +##### Outer interface for model constructors +##### + +regularize_field_boundary_conditions(boundary_conditions::NamedTuple, grid::AbstractGrid, prognostic_field_names::Tuple) = NamedTuple(field_name => regularize_field_boundary_conditions(field_bcs, grid, field_name, prognostic_field_names) for (field_name, field_bcs) in pairs(boundary_conditions)) - -regularize_field_boundary_conditions(::Missing, grid, field_name, prognostic_field_names=nothing) = missing diff --git a/src/Fields/abstract_field.jl b/src/Fields/abstract_field.jl index 69eba4c665..db26f315ef 100644 --- a/src/Fields/abstract_field.jl +++ b/src/Fields/abstract_field.jl @@ -24,6 +24,8 @@ const GridOrNothing = Union{AbstractGrid, Nothing} Abstract supertype for fields located at `(LX, LY, LZ)` and defined on a grid `G` with eltype `T` and `N` dimensions. + +Note: we need the parameter `T` to subtype AbstractArray. """ abstract type AbstractField{LX, LY, LZ, G <: GridOrNothing, T, N} <: AbstractArray{T, N} end @@ -41,13 +43,13 @@ data(a) = a @inline location(a) = (Nothing, Nothing, Nothing) # used in AbstractOperations for location inference @inline location(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = (LX, LY, LZ) # note no instantiation @inline location(f::AbstractField, i) = location(f)[i] -instantiated_location(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = (LX(), LY(), LZ()) +@inline instantiated_location(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = (LX(), LY(), LZ()) "Returns the architecture of on which `f` is defined." architecture(f::AbstractField) = architecture(f.grid) "Returns the topology of a fields' `grid`." -topology(f::AbstractField, args...) = topology(f.grid, args...) +@inline topology(f::AbstractField, args...) = topology(f.grid, args...) """ size(f::AbstractField) @@ -78,6 +80,10 @@ interior(f::AbstractField) = f @propagate_inbounds ynode(j, ψ::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = ynode(LY(), j, ψ.grid) @propagate_inbounds znode(k, ψ::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = znode(LZ(), k, ψ.grid) +@propagate_inbounds xnode(i, j, k, ψ::AbstractField) = xnode(instantiated_location(ψ)..., i, j, k, ψ.grid) +@propagate_inbounds ynode(i, j, k, ψ::AbstractField) = ynode(instantiated_location(ψ)..., i, j, k, ψ.grid) +@propagate_inbounds znode(i, j, k, ψ::AbstractField) = znode(instantiated_location(ψ)..., i, j, k, ψ.grid) + xnodes(ψ::AbstractField) = xnodes(location(ψ, 1), ψ.grid) ynodes(ψ::AbstractField) = ynodes(location(ψ, 2), ψ.grid) znodes(ψ::AbstractField) = znodes(location(ψ, 3), ψ.grid) @@ -94,9 +100,7 @@ for f in (:+, :-) end function Statistics.norm(a::AbstractField) - arch = architecture(a) - grid = a.grid - r = zeros(arch, grid, 1) + r = zeros(a.grid, 1) Base.mapreducedim!(x -> x * x, +, r, a) return CUDA.@allowscalar sqrt(r[1]) end diff --git a/src/Fields/field.jl b/src/Fields/field.jl index 729ffbc0fb..cad72c0b84 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -339,7 +339,13 @@ initialize_reduced_field!(::MinimumReduction, f, r::ReducedField, c) = Base.mapf filltype(f, grid) = eltype(grid) filltype(::Union{AllReduction, AnyReduction}, grid) = Bool -reduced_location(loc; dims) = Tuple(i ∈ dims ? Nothing : loc[i] for i in 1:3) +function reduced_location(loc; dims) + if dims isa Colon + return (Nothing, Nothing, Nothing) + else + return Tuple(i ∈ dims ? Nothing : loc[i] for i in 1:3) + end +end # Allocating and in-place reductions for reduction in (:sum, :maximum, :minimum, :all, :any) @@ -357,7 +363,7 @@ for reduction in (:sum, :maximum, :minimum, :all, :any) # Allocating function Base.$(reduction)(f::Function, c::AbstractField; dims=:) - if dims === (:) + if dims isa Colon r = zeros(architecture(c), c.grid, 1, 1, 1) Base.$(reduction!)(f, r, c) return CUDA.@allowscalar r[1, 1, 1] diff --git a/src/Fields/function_field.jl b/src/Fields/function_field.jl index 70320b05ed..fb9d294937 100644 --- a/src/Fields/function_field.jl +++ b/src/Fields/function_field.jl @@ -35,8 +35,8 @@ struct FunctionField{LX, LY, LZ, C, P, F, G, T} <: AbstractField{LX, LY, LZ, G, end """Return `a`, or convert `a` to `FunctionField` if `a::Function`""" -fieldify(L, a, grid) = a -fieldify(L, a::Function, grid) = FunctionField(L, a, grid) +fieldify_function(L, a, grid) = a +fieldify_function(L, a::Function, grid) = FunctionField(L, a, grid) """ FunctionField(L::Tuple, func, grid) diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl index 09741bfb7a..a74b5effbf 100644 --- a/src/Grids/rectilinear_grid.jl +++ b/src/Grids/rectilinear_grid.jl @@ -62,9 +62,9 @@ const XZRegRectilinearGrid = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Numbe const YZRegRectilinearGrid = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Number, <:Number} const RegRectilinearGrid = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number, <:Number, <:Number} -regular_dimensions(::XRegRectilinearGrid) = tuple(1) -regular_dimensions(::YRegRectilinearGrid) = tuple(2) -regular_dimensions(::ZRegRectilinearGrid) = tuple(3) +regular_dimensions(::XRegRectilinearGrid) = tuple(1) +regular_dimensions(::YRegRectilinearGrid) = tuple(2) +regular_dimensions(::ZRegRectilinearGrid) = tuple(3) regular_dimensions(::XYRegRectilinearGrid) = (1, 2) regular_dimensions(::XZRegRectilinearGrid) = (1, 3) regular_dimensions(::YZRegRectilinearGrid) = (2, 3) diff --git a/src/OutputReaders/field_time_series.jl b/src/OutputReaders/field_time_series.jl index c7c14ab6ce..271c73ca2a 100644 --- a/src/OutputReaders/field_time_series.jl +++ b/src/OutputReaders/field_time_series.jl @@ -21,7 +21,7 @@ struct FieldTimeSeries{LX, LY, LZ, K, T, D, G, B, χ} <: AbstractField{LX, LY, L times :: χ function FieldTimeSeries{LX, LY, LZ, K}(data::D, grid::G, bcs::B, times::χ) where {LX, LY, LZ, K, D, G, B, χ} - T = eltype(grid) + T = eltype(data) return new{LX, LY, LZ, K, T, D, G, B, χ}(data, grid, bcs, times) end end diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/leith_enstrophy_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/leith_enstrophy_diffusivity.jl index c766ce64b4..e064383c21 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/leith_enstrophy_diffusivity.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/leith_enstrophy_diffusivity.jl @@ -75,7 +75,7 @@ end const ArrayOrField = Union{AbstractArray, AbstractField} @inline ψ²(i, j, k, grid, ψ::Function, args...) = ψ(i, j, k, grid, args...)^2 -@inline ψ²(i, j, k, grid, ψ::ArrayOrField, args...) = ψ[i, j, k]^2 +@inline ψ²(i, j, k, grid, ψ::ArrayOrField, args...) = @inbounds ψ[i, j, k]^2 @inline function abs²_∇h_wz(i, j, k, grid, w) wxz² = ℑxᶜᵃᵃ(i, j, k, grid, ψ², ∂xᶠᵃᵃ, ∂zᵃᵃᶜ, w) diff --git a/test/regression_tests/ocean_large_eddy_simulation_regression_test.jl b/test/regression_tests/ocean_large_eddy_simulation_regression_test.jl index 02617b678e..af3111ff2d 100644 --- a/test/regression_tests/ocean_large_eddy_simulation_regression_test.jl +++ b/test/regression_tests/ocean_large_eddy_simulation_regression_test.jl @@ -37,6 +37,8 @@ function run_ocean_large_eddy_simulation_regression_test(arch, grid_type, closur boundary_conditions = (u=u_bcs, T=T_bcs, S=S_bcs) ) + @show model.diffusivity_fields + # We will manually change the stop_iteration as needed. simulation = Simulation(model, Δt=Δt, stop_iteration=0) diff --git a/test/test_regression.jl b/test/test_regression.jl index a73046a270..06e7dc5c5e 100644 --- a/test/test_regression.jl +++ b/test/test_regression.jl @@ -1,11 +1,6 @@ -using Oceananigans.Grids: topology, XRegLatLonGrid, YRegLatLonGrid, ZRegLatLonGrid -using CUDA - -include("utils_for_runtests.jl") -include("data_dependencies.jl") - -archs = test_architectures() +include("dependencies_for_runtests.jl") +using Oceananigans.Grids: topology, XRegLatLonGrid, YRegLatLonGrid, ZRegLatLonGrid function show_hydrostatic_test(grid, free_surface, comp) @@ -71,6 +66,7 @@ include("regression_tests/hydrostatic_free_turbulence_regression_test.jl") for arch in archs for grid_type in [:regular, :vertically_unstretched] + #= @testset "Thermal bubble [$(typeof(arch)), $grid_type grid]" begin @info " Testing thermal bubble regression [$(typeof(arch)), $grid_type grid]" run_thermal_bubble_regression_test(arch, grid_type) @@ -80,6 +76,7 @@ include("regression_tests/hydrostatic_free_turbulence_regression_test.jl") @info " Testing Rayleigh–Bénard tracer regression [$(typeof(arch)), $grid_type grid]" run_rayleigh_benard_regression_test(arch, grid_type) end + =# for closure in (AnisotropicMinimumDissipation(ν=1.05e-6, κ=1.46e-7), SmagorinskyLilly(C=0.23, Cb=1, Pr=1, ν=1.05e-6, κ=1.46e-7)) closurename = string(typeof(closure).name.wrapper) @@ -92,6 +89,7 @@ include("regression_tests/hydrostatic_free_turbulence_regression_test.jl") # Hydrostatic regression test + #= longitude = ((-180, 180), collect(-180:2:180), (-160, 160), collect(-160:2:160)) latitude = ((-60, 60), collect(-60:2:60)) zcoord = ((-90, 0) , collect(-90:30:0)) @@ -127,5 +125,6 @@ include("regression_tests/hydrostatic_free_turbulence_regression_test.jl") end end end + =# end end From ee945bc1c31b371c6e30c28ac0359222c2fbda45 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 14 Jan 2022 12:08:00 -0700 Subject: [PATCH 137/140] Fix show for ReducedComputedField --- src/Fields/field_reductions.jl | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/src/Fields/field_reductions.jl b/src/Fields/field_reductions.jl index b343f33ca3..a186369919 100644 --- a/src/Fields/field_reductions.jl +++ b/src/Fields/field_reductions.jl @@ -65,12 +65,9 @@ Base.show(io::IO, field::ReducedComputedField) = print(io, "$(short_show(field))\n", "├── data: $(typeof(field.data)), size: $(size(field))\n", "├── grid: $(short_show(field.grid))\n", - "├── dims: $(field.dims)\n", "├── operand: $(short_show(field.operand))\n", "└── status: ", show_status(field.status)) -short_show(field::ReducedComputedField) = string("Field at ", show_location(field), " via ", short_show(field.operand)) - -short_show(r::Reduction) = string(typeof(r.reduce!), " of ", short_show(r.operand), - " over dims ", r.dims) - +short_show(r::Reduction) = string(r.reduce!, + " over dims ", r.dims, + " of ", short_show(r.operand)) From 2b6a1194c5a4e372098b847aa104eb858fe82ade Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 14 Jan 2022 12:56:53 -0700 Subject: [PATCH 138/140] Disambiguate bc regularization for CubedSphere --- src/CubedSpheres/CubedSpheres.jl | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/CubedSpheres/CubedSpheres.jl b/src/CubedSpheres/CubedSpheres.jl index a01d78ed40..d3c7243711 100644 --- a/src/CubedSpheres/CubedSpheres.jl +++ b/src/CubedSpheres/CubedSpheres.jl @@ -35,7 +35,10 @@ validate_vertical_velocity_boundary_conditions(w::AbstractCubedSphereField) = import Oceananigans.BoundaryConditions: regularize_field_boundary_conditions -function regularize_field_boundary_conditions(bcs::CubedSphereFaces, grid, field_name, prognostic_field_names) +function regularize_field_boundary_conditions(bcs::CubedSphereFaces, + grid::AbstractGrid, + field_name::Symbol, + prognostic_field_names) faces = Tuple(regularize_field_boundary_conditions(face_bcs, face_grid, field_name, prognostic_field_names) for (face_bcs, face_grid) in zip(bcs.faces, grid.faces)) @@ -43,7 +46,10 @@ function regularize_field_boundary_conditions(bcs::CubedSphereFaces, grid, field return CubedSphereFaces{typeof(faces[1]), typeof(faces)}(faces) end -function regularize_field_boundary_conditions(bcs::FieldBoundaryConditions, grid::ConformalCubedSphereGrid, field_name, prognostic_field_names) +function regularize_field_boundary_conditions(bcs::FieldBoundaryConditions, + grid::ConformalCubedSphereGrid, + field_name::Symbol, + prognostic_field_names) faces = Tuple( inject_cubed_sphere_exchange_boundary_conditions( From 59717c5e3338b4babfbbf20d70c8201c93e80ce8 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 14 Jan 2022 15:11:51 -0500 Subject: [PATCH 139/140] Fixes bug in set! --- src/Fields/set!.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/Fields/set!.jl b/src/Fields/set!.jl index 6383ecb5dc..cd690c2f9a 100644 --- a/src/Fields/set!.jl +++ b/src/Fields/set!.jl @@ -41,7 +41,7 @@ end function set!(u::Field, v::Field) if architecture(u) === architecture(v) parent(u) .= parent(v) - elseif parent(v) isa SubArray + else u_parent = parent(u) v_parent = parent(v) # If u_parent is a view, we have to convert to an Array. @@ -51,6 +51,7 @@ function set!(u::Field, v::Field) v_parent isa SubArray && (v_parent = Array(v_parent)) copyto!(u_parent, v_parent) end + return nothing end From 24f7e8c22c897364746b9eedd584bb8393ede423 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 14 Jan 2022 15:44:15 -0500 Subject: [PATCH 140/140] Restore all regression tests --- test/test_regression.jl | 4 ---- 1 file changed, 4 deletions(-) diff --git a/test/test_regression.jl b/test/test_regression.jl index 06e7dc5c5e..859c086125 100644 --- a/test/test_regression.jl +++ b/test/test_regression.jl @@ -66,7 +66,6 @@ include("regression_tests/hydrostatic_free_turbulence_regression_test.jl") for arch in archs for grid_type in [:regular, :vertically_unstretched] - #= @testset "Thermal bubble [$(typeof(arch)), $grid_type grid]" begin @info " Testing thermal bubble regression [$(typeof(arch)), $grid_type grid]" run_thermal_bubble_regression_test(arch, grid_type) @@ -76,7 +75,6 @@ include("regression_tests/hydrostatic_free_turbulence_regression_test.jl") @info " Testing Rayleigh–Bénard tracer regression [$(typeof(arch)), $grid_type grid]" run_rayleigh_benard_regression_test(arch, grid_type) end - =# for closure in (AnisotropicMinimumDissipation(ν=1.05e-6, κ=1.46e-7), SmagorinskyLilly(C=0.23, Cb=1, Pr=1, ν=1.05e-6, κ=1.46e-7)) closurename = string(typeof(closure).name.wrapper) @@ -89,7 +87,6 @@ include("regression_tests/hydrostatic_free_turbulence_regression_test.jl") # Hydrostatic regression test - #= longitude = ((-180, 180), collect(-180:2:180), (-160, 160), collect(-160:2:160)) latitude = ((-60, 60), collect(-60:2:60)) zcoord = ((-90, 0) , collect(-90:30:0)) @@ -125,6 +122,5 @@ include("regression_tests/hydrostatic_free_turbulence_regression_test.jl") end end end - =# end end