Merge branch 'main' into ncc-ss/distributed-reduction

Author: simone-silvestri

src/BoundaryConditions/boundary_condition.jl

@@ -109,18 +109,15 @@ DistributedCommunicationBoundaryCondition(val; kwargs...) = BoundaryCondition(Di
 #     * additional arguments to `fill_halo_regions` enter `getbc` after the `grid` argument:
 #           so `fill_halo_regions!(c, clock, fields)` translates to `getbc(bc, i, j, grid, clock, fields)`, etc.
 
-@inline getbc(bc, args...) = bc.condition(args...) # fallback!
+@inline getbc(bc::BoundaryCondition, args...) = getbc(bc.condition, args...) # unwrap
+@inline getbc(condition, args...) = condition(args...) # fallback!
 
-@inline getbc(::BC{<:Open, Nothing}, ::Integer, ::Integer, grid::AbstractGrid, args...) = zero(grid)
-@inline getbc(::BC{<:Flux, Nothing}, ::Integer, ::Integer, grid::AbstractGrid, args...) = zero(grid)
-@inline getbc(::Nothing,             ::Integer, ::Integer, grid::AbstractGrid, args...) = zero(grid)
+@inline getbc(::Nothing, ::Integer, ::Integer, grid::AbstractGrid, args...) = zero(grid)
 
-@inline getbc(bc::BC{<:Any, <:Number}, args...) = bc.condition
-@inline getbc(bc::BC{<:Any, <:AbstractArray}, i::Integer, j::Integer, grid::AbstractGrid, args...) = @inbounds bc.condition[i, j]
-
-# Support for Ref boundary conditions
 const NumberRef = Base.RefValue{<:Number}
-@inline getbc(bc::BC{<:Any, <:NumberRef}, args...) = bc.condition[]
+@inline getbc(condition::NumberRef, args...) = condition[]
+@inline getbc(condition::Number, args...) = condition
+@inline getbc(condition::AbstractArray, i::Integer, j::Integer, grid::AbstractGrid, args...) = @inbounds condition[i, j]
 
 #####
 ##### Validation with topology

src/BoundaryConditions/continuous_boundary_function.jl

@@ -116,15 +116,14 @@ end
 @inline z_boundary_node(i, j, k, grid::YFlatGrid,  ℓx, ℓy) = tuple(ξnode(i, j, k, grid, ℓx, nothing, Face()))
 @inline z_boundary_node(i, j, k, grid::XYFlatGrid, ℓx, ℓy) = tuple()
 
-const XBoundaryFunction{LY, LZ, S} = BoundaryCondition{<:Any, <:ContinuousBoundaryFunction{Nothing, LY, LZ, S}} where {LY, LZ, S}
-const YBoundaryFunction{LX, LZ, S} = BoundaryCondition{<:Any, <:ContinuousBoundaryFunction{LX, Nothing, LZ, S}} where {LX, LZ, S}
-const ZBoundaryFunction{LX, LY, S} = BoundaryCondition{<:Any, <:ContinuousBoundaryFunction{LX, LY, Nothing, S}} where {LX, LY, S}
+const XBoundaryFunction{LY, LZ, S} = ContinuousBoundaryFunction{Nothing, LY, LZ, S} where {LY, LZ, S}
+const YBoundaryFunction{LX, LZ, S} = ContinuousBoundaryFunction{LX, Nothing, LZ, S} where {LX, LZ, S}
+const ZBoundaryFunction{LX, LY, S} = ContinuousBoundaryFunction{LX, LY, Nothing, S} where {LX, LY, S}
 
 # Return ContinuousBoundaryFunction on east or west boundaries.
-@inline function getbc(bc::XBoundaryFunction{LY, LZ, S}, j::Integer, k::Integer,
+@inline function getbc(cbf::XBoundaryFunction{LY, LZ, S}, j::Integer, k::Integer,
                        grid::AbstractGrid, clock, model_fields, args...) where {LY, LZ, S}
 
-    cbf = bc.condition
     i, i′ = domain_boundary_indices(S(), grid.Nx)
     args = user_function_arguments(i, j, k, grid, model_fields, cbf.parameters, cbf)
     X = x_boundary_node(i′, j, k, grid, LY(), LZ())
@@ -133,10 +132,9 @@ const ZBoundaryFunction{LX, LY, S} = BoundaryCondition{<:Any, <:ContinuousBounda
 end
 
 # Return ContinuousBoundaryFunction on south or north boundaries.
-@inline function getbc(bc::YBoundaryFunction{LX, LZ, S}, i::Integer, k::Integer,
+@inline function getbc(cbf::YBoundaryFunction{LX, LZ, S}, i::Integer, k::Integer,
                        grid::AbstractGrid, clock, model_fields, args...) where {LX, LZ, S}
 
-    cbf = bc.condition
     j, j′ = domain_boundary_indices(S(), grid.Ny)
     args = user_function_arguments(i, j, k, grid, model_fields, cbf.parameters, cbf)
     X = y_boundary_node(i, j′, k, grid, LX(), LZ())
@@ -145,10 +143,9 @@ end
 end
 
 # Return ContinuousBoundaryFunction on bottom or top boundaries.
-@inline function getbc(bc::ZBoundaryFunction{LX, LY, S}, i::Integer, j::Integer,
+@inline function getbc(cbf::ZBoundaryFunction{LX, LY, S}, i::Integer, j::Integer,
                        grid::AbstractGrid, clock, model_fields, args...) where {LX, LY, S}
 
-    cbf = bc.condition
     k, k′ = domain_boundary_indices(S(), grid.Nz)
     args = user_function_arguments(i, j, k, grid, model_fields, cbf.parameters, cbf)
     X = z_boundary_node(i, j, k′, grid, LX(), LY())
@@ -161,10 +158,9 @@ end
 #####
 
 # Return ContinuousBoundaryFunction on the east or west interface of a cell adjacent to an immersed boundary
-@inline function getbc(bc::XBoundaryFunction{LY, LZ, S}, i::Integer, j::Integer, k::Integer,
+@inline function getbc(cbf::XBoundaryFunction{LY, LZ, S}, i::Integer, j::Integer, k::Integer,
                        grid::AbstractGrid, clock, model_fields, args...) where {LY, LZ, S}
 
-    cbf = bc.condition
     i′ = cell_boundary_index(S(), i)
     args = user_function_arguments(i, j, k, grid, model_fields, cbf.parameters, cbf)
     X = node(i′, j, k, grid, Face(), LY(), LZ())
@@ -173,10 +169,9 @@ end
 end
 
 # Return ContinuousBoundaryFunction on the south or north interface of a cell adjacent to an immersed boundary
-@inline function getbc(bc::YBoundaryFunction{LX, LZ, S}, i::Integer, j::Integer, k::Integer,
+@inline function getbc(cbf::YBoundaryFunction{LX, LZ, S}, i::Integer, j::Integer, k::Integer,
                        grid::AbstractGrid, clock, model_fields, args...) where {LX, LZ, S}
 
-    cbf = bc.condition
     j′ = cell_boundary_index(S(), j)
     args = user_function_arguments(i, j, k, grid, model_fields, cbf.parameters, cbf)
     X = node(i, j′, k, grid, LX(), Face(), LZ())
@@ -185,10 +180,9 @@ end
 end
 
 # Return ContinuousBoundaryFunction on the bottom or top interface of a cell adjacent to an immersed boundary
-@inline function getbc(bc::ZBoundaryFunction{LX, LY, S}, i::Integer, j::Integer, k::Integer,
+@inline function getbc(cbf::ZBoundaryFunction{LX, LY, S}, i::Integer, j::Integer, k::Integer,
                        grid::AbstractGrid, clock, model_fields, args...) where {LX, LY, S}
 
-    cbf = bc.condition
     k′ = cell_boundary_index(S(), k)
     args = user_function_arguments(i, j, k, grid, model_fields, cbf.parameters, cbf)
     X = node(i, j, k′, grid, LX(), LY(), Face())

src/BoundaryConditions/discrete_boundary_function.jl

@@ -32,21 +32,19 @@ struct DiscreteBoundaryFunction{P, F}
 end
 
 const UnparameterizedDBF = DiscreteBoundaryFunction{<:Nothing}
-const UnparameterizedDBFBC = BoundaryCondition{<:Any, <:UnparameterizedDBF}
-const DBFBC = BoundaryCondition{<:Any, <:DiscreteBoundaryFunction}
 
-@inline getbc(bc::UnparameterizedDBFBC, i::Integer, j::Integer, grid::AbstractGrid, clock, model_fields, args...) =
-    bc.condition.func(i, j, grid, clock, model_fields)
+@inline getbc(condition::UnparameterizedDBF, i::Integer, j::Integer, grid::AbstractGrid, clock, model_fields, args...) =
+    condition.func(i, j, grid, clock, model_fields)
 
-@inline getbc(bc::DBFBC, i::Integer, j::Integer, grid::AbstractGrid, clock, model_fields, args...) =
-    bc.condition.func(i, j, grid, clock, model_fields, bc.condition.parameters)
+@inline getbc(condition::DiscreteBoundaryFunction, i::Integer, j::Integer, grid::AbstractGrid, clock, model_fields, args...) =
+    condition.func(i, j, grid, clock, model_fields, condition.parameters)
 
 # 3D function for immersed boundary conditions
-@inline getbc(bc::UnparameterizedDBFBC, i::Integer, j::Integer, k::Integer, grid::AbstractGrid, clock, model_fields, args...) =
-    bc.condition.func(i, j, k, grid, clock, model_fields)
+@inline getbc(condition::UnparameterizedDBF, i::Integer, j::Integer, k::Integer, grid::AbstractGrid, clock, model_fields, args...) =
+    condition.func(i, j, k, grid, clock, model_fields)
 
-@inline getbc(bc::DBFBC, i::Integer, j::Integer, k::Integer, grid::AbstractGrid, clock, model_fields, args...) =
-    bc.condition.func(i, j, k, grid, clock, model_fields, bc.condition.parameters)
+@inline getbc(condition::DiscreteBoundaryFunction, i::Integer, j::Integer, k::Integer, grid::AbstractGrid, clock, model_fields, args...) =
+    condition.func(i, j, k, grid, clock, model_fields, condition.parameters)
 
 # Don't re-convert DiscreteBoundaryFunctions passed to BoundaryCondition constructor
 BoundaryCondition(Classification::DataType, condition::DiscreteBoundaryFunction) = BoundaryCondition(Classification(), condition)

src/BoundaryConditions/perturbation_advection_open_boundary_matching_scheme.jl

@@ -116,24 +116,19 @@ end
 @inline function step_right_boundary!(bc::FPAOBC, l, m, boundary_indices, boundary_adjacent_indices,
                                       grid, u, clock, model_fields, ΔX)
     Δt = clock.last_stage_Δt
-
     Δt = ifelse(isinf(Δt), 0, Δt)
 
     ūⁿ⁺¹ = getbc(bc, l, m, grid, clock, model_fields)
-
     uᵢⁿ     = @inbounds getindex(u, boundary_indices...)
     uᵢ₋₁ⁿ⁺¹ = @inbounds getindex(u, boundary_adjacent_indices...)
 
     U = max(0, min(1, Δt / ΔX * ūⁿ⁺¹))
-
     pa = bc.classification.matching_scheme
 
     τ = ifelse(ūⁿ⁺¹ >= 0, pa.outflow_timescale, pa.inflow_timescale)
-
     τ̃ = Δt / τ
 
     uᵢⁿ⁺¹ = uᵢⁿ + U * (uᵢ₋₁ⁿ⁺¹ - ūⁿ⁺¹) + (ūⁿ⁺¹ - uᵢⁿ) * τ̃
-
     @inbounds setindex!(u, uᵢⁿ⁺¹, boundary_indices...)
 
     return nothing
@@ -142,24 +137,19 @@ end
 @inline function step_left_boundary!(bc::FPAOBC, l, m, boundary_indices, boundary_adjacent_indices, boundary_secret_storage_indices,
                                      grid, u, clock, model_fields, ΔX)
     Δt = clock.last_stage_Δt
-
     Δt = ifelse(isinf(Δt), 0, Δt)
 
     ūⁿ⁺¹ = getbc(bc, l, m, grid, clock, model_fields)
-
     uᵢⁿ     = @inbounds getindex(u, boundary_secret_storage_indices...)
     uᵢ₋₁ⁿ⁺¹ = @inbounds getindex(u, boundary_adjacent_indices...)
 
     U = min(0, max(-1, Δt / ΔX * ūⁿ⁺¹))
-
     pa = bc.classification.matching_scheme
 
     τ = ifelse(ūⁿ⁺¹ <= 0, pa.outflow_timescale, pa.inflow_timescale)
-
     τ̃ = Δt / τ
 
     u₁ⁿ⁺¹ = uᵢⁿ - U * (uᵢ₋₁ⁿ⁺¹ - ūⁿ⁺¹) + (ūⁿ⁺¹ - uᵢⁿ) * τ̃
-
     @inbounds setindex!(u, u₁ⁿ⁺¹, boundary_indices...)
     @inbounds setindex!(u, u₁ⁿ⁺¹, boundary_secret_storage_indices...)
 
@@ -168,7 +158,6 @@ end
 
 @inline function _fill_east_halo!(j, k, grid, u, bc::PAOBC, ::Tuple{Face, Any, Any}, clock, model_fields)
     i = grid.Nx + 1
-
     boundary_indices = (i, j, k)
     boundary_adjacent_indices = (i-1, j, k)
 
@@ -193,7 +182,6 @@ end
 
 @inline function _fill_north_halo!(i, k, grid, u, bc::PAOBC, ::Tuple{Any, Face, Any}, clock, model_fields)
     j = grid.Ny + 1
-
     boundary_indices = (i, j, k)
     boundary_adjacent_indices = (i, j-1, k)
 
@@ -218,7 +206,6 @@ end
 
 @inline function _fill_top_halo!(i, j, grid, u, bc::PAOBC, ::Tuple{Any, Any, Face}, clock, model_fields)
     k = grid.Nz + 1
-
     boundary_indices = (i, j, k)
     boundary_adjacent_indices = (i, j, k-1)
 

src/BoundaryConditions/polar_boundary_condition.jl

@@ -31,7 +31,7 @@ maybe_polar_boundary_condition(grid, side, ::Type{Center},  LZ) = PolarValueBoun
 maybe_polar_boundary_condition(grid, side, ::Type{Face},    LZ) = PolarOpenBoundaryCondition(grid, side, LZ)
 
 # Just a column
-@inline getbc(pv::BC{<:Any, <:PolarValue}, i, k, args...) = @inbounds pv.condition.data[1, 1, k]
+@inline getbc(pv::PolarValue, i, k, args...) = @inbounds pv.data[1, 1, k]
 
 @kernel function _average_pole_value!(data, c, j, grid, loc)
     i′, j′, k = @index(Global, NTuple)

src/Fields/Fields.jl

@@ -30,7 +30,6 @@ include("field_indices.jl")
 include("scans.jl")
 include("regridding_fields.jl")
 include("field_tuples.jl")
-include("background_fields.jl")
 include("interpolate.jl")
 include("show_fields.jl")
 include("broadcasting_abstract_fields.jl")

src/Fields/background_fields.jl

@@ -544,20 +544,15 @@ const ReducedField = Union{XReducedField,
 @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)
 
-const XFieldBC = BoundaryCondition{<:Any, XReducedField}
-const YFieldBC = BoundaryCondition{<:Any, YReducedField}
-const ZFieldBC = BoundaryCondition{<:Any, ZReducedField}
-
 # Boundary conditions reduced in one direction --- drop boundary-normal index
-@inline getbc(bc::XFieldBC, j::Integer, k::Integer, grid::AbstractGrid, args...) = @inbounds bc.condition[1, j, k]
-@inline getbc(bc::YFieldBC, i::Integer, k::Integer, grid::AbstractGrid, args...) = @inbounds bc.condition[i, 1, k]
-@inline getbc(bc::ZFieldBC, i::Integer, j::Integer, grid::AbstractGrid, args...) = @inbounds bc.condition[i, j, 1]
+@inline getbc(condition::XReducedField, j::Integer, k::Integer, grid::AbstractGrid, args...) = @inbounds condition[1, j, k]
+@inline getbc(condition::YReducedField, i::Integer, k::Integer, grid::AbstractGrid, args...) = @inbounds condition[i, 1, k]
+@inline getbc(condition::ZReducedField, i::Integer, j::Integer, grid::AbstractGrid, args...) = @inbounds condition[i, j, 1]
 
 # Boundary conditions reduced in two directions are ambiguous, so that's hard...
 
 # 0D boundary conditions --- easy case
-const XYZFieldBC = BoundaryCondition{<:Any, XYZReducedField}
-@inline getbc(bc::XYZFieldBC, ::Integer, ::Integer, ::AbstractGrid, args...) = @inbounds bc.condition[1, 1, 1]
+@inline getbc(condition::XYZReducedField, ::Integer, ::Integer, ::AbstractGrid, args...) = @inbounds condition[1, 1, 1]
 
 # Preserve location when adapting fields reduced on one or more dimensions
 function Adapt.adapt_structure(to, reduced_field::ReducedField)

src/Fields/field.jl

@@ -1,7 +1,7 @@
 module Models
 
 export
-    NonhydrostaticModel,
+    NonhydrostaticModel, BackgroundField, BackgroundFields,
     ShallowWaterModel, ConservativeFormulation, VectorInvariantFormulation,
     HydrostaticFreeSurfaceModel, ZStar, ZCoordinate,
     ExplicitFreeSurface, ImplicitFreeSurface, SplitExplicitFreeSurface,
@@ -101,7 +101,7 @@ include("HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl")
 include("ShallowWaterModels/ShallowWaterModels.jl")
 include("LagrangianParticleTracking/LagrangianParticleTracking.jl")
 
-using .NonhydrostaticModels: NonhydrostaticModel, PressureField
+using .NonhydrostaticModels: NonhydrostaticModel, PressureField, BackgroundField, BackgroundFields
 
 using .HydrostaticFreeSurfaceModels:
     HydrostaticFreeSurfaceModel,

src/Models/Models.jl

@@ -1,6 +1,6 @@
 module NonhydrostaticModels
 
-export NonhydrostaticModel
+export NonhydrostaticModel, BackgroundField, BackgroundFields
 
 using DocStringExtensions
 
@@ -65,6 +65,7 @@ nonhydrostatic_pressure_solver(grid) = nonhydrostatic_pressure_solver(architectu
 ##### NonhydrostaticModel definition
 #####
 
+include("background_fields.jl")
 include("nonhydrostatic_model.jl")
 include("pressure_field.jl")
 include("show_nonhydrostatic_model.jl")