Pad Center fields on BoundedTopology so that all Fields have identical horizontal size (#4144)

* Pad Center fields on BoundedTopology to even out all fields

* Change abstraction

* Better hack

* sharded grid

* add sharding length

* no need for this

* Update ext/OceananigansReactantExt/Grids/sharded_grids.jl

* using instantiate

* remove the show

* Update src/Grids/grid_utils.jl

* Extend total_size

* Update src/Grids/new_data.jl

* sharding_total_length -> reactant_total_length (#4329)

* Fix tests now that padding differs between reactant and vanilla

---------

Co-authored-by: Simone Silvestri <[email protected]>
Co-authored-by: Mosè Giordano <[email protected]>

Author: 3 people

ext/OceananigansReactantExt/Grids/Grids.jl

@@ -3,17 +3,19 @@ module Grids
 export constant_with_arch
 
 using Reactant
+using OffsetArrays
 
 using Oceananigans
 using Oceananigans: Distributed
 using Oceananigans.Architectures: ReactantState, CPU
 using Oceananigans.Grids: AbstractGrid, AbstractUnderlyingGrid, StaticVerticalDiscretization, MutableVerticalDiscretization
-using Oceananigans.Grids: Center, Face, RightConnected, LeftConnected, Periodic, Bounded, Flat
+using Oceananigans.Grids: Center, Face, RightConnected, LeftConnected, Periodic, Bounded, Flat, BoundedTopology
 using Oceananigans.Fields: Field
 using Oceananigans.ImmersedBoundaries: GridFittedBottom, AbstractImmersedBoundary
 
 import ..OceananigansReactantExt: deconcretize
 import Oceananigans.Grids: LatitudeLongitudeGrid, RectilinearGrid, OrthogonalSphericalShellGrid
+import Oceananigans.Grids: total_length, offset_data
 import Oceananigans.OrthogonalSphericalShellGrids: RotatedLatitudeLongitudeGrid, TripolarGrid
 import Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, materialize_immersed_boundary
 
@@ -39,5 +41,53 @@ const ShardedGrid{FT, TX, TY, TZ} = AbstractGrid{FT, TX, TY, TZ, <:ShardedDistri
 include("serial_grids.jl")
 include("sharded_grids.jl")
 
+function total_size(grid::ReactantGrid, loc, indices)
+    sz = size(grid)
+    halo_sz = halo_size(grid)
+    topo = topology(grid)
+    return reactant_total_size(loc, topo, sz, halo_sz, indices)
+end
+
+function reactant_total_size(loc, topo, sz, halo_sz, indices=default_indices(Val(length(loc))))
+    D = length(loc)
+    return Tuple(reactant_total_length(instantiate(loc[d]), instantiate(topo[d]), sz[d], halo_sz[d], indices[d]) for d = 1:D)
+end
+
+reactant_total_length(loc, topo, N, H, ::Colon) = reactant_total_length(loc, topo, N, H)
+reactant_total_length(loc, topo, N, H, ind::AbstractUnitRange) = min(reactant_total_length(loc, topo, N, H), length(ind))
+reactant_total_length(loc, topo, N, H) = Oceananigans.Grids.total_length(loc, topo, N, H)
+reactant_total_length(::Face, ::BoundedTopology, N, H=0) = N + 2H
+
+reactant_offset_indices(loc, topo, N, H=0) = 1 - H : N + H
+reactant_offset_indices(::Nothing, topo, N, H=0) = 1:1
+reactant_offset_indices(ℓ,         topo, N, H, ::Colon) = reactant_offset_indices(ℓ, topo, N, H)
+reactant_offset_indices(ℓ,         topo, N, H, r::AbstractUnitRange) = r
+reactant_offset_indices(::Nothing, topo, N, H, ::AbstractUnitRange) = 1:1
+
+function Oceananigans.Grids.new_data(FT::DataType, arch::Union{ReactantState, ShardedDistributed},
+        loc, topo, sz, halo_sz, indices=default_indices(length(loc)))
+
+    Tsz = reactant_total_size(loc, topo, sz, halo_sz, indices)
+    underlying_data = zeros(arch, FT, Tsz...)
+    indices = validate_indices(indices, loc, topo, sz, halo_sz)
+
+    return offset_data(underlying_data, loc, topo, sz, halo_sz, indices)
+end
+
+# The type parameter for indices helps / encourages the compiler to fully type infer `offset_data`
+function offset_data(underlying_data::ConcreteRArray, loc, topo, N, H, indices::T=default_indices(length(loc))) where T
+    loc = map(instantiate, loc)
+    topo = map(instantiate, topo)
+    ii = map(reactant_offset_indices, loc, topo, N, H, indices)
+    # Add extra indices for arrays of higher dimension than loc, topo, etc.
+    # Use the "`ntuple` trick" so the compiler can infer the type of `extra_ii`
+    extra_ii = ntuple(Val(ndims(underlying_data)-length(ii))) do i
+        Base.@_inline_meta
+        axes(underlying_data, i+length(ii))
+    end
+
+    return OffsetArray(underlying_data, ii..., extra_ii...)
+end
+
 end # module
 

ext/OceananigansReactantExt/Grids/sharded_grids.jl

@@ -2,8 +2,9 @@ using Oceananigans.Architectures: architecture
 using Oceananigans.Grids: AbstractGrid
 using Oceananigans.OrthogonalSphericalShellGrids
 using Oceananigans.Grids: R_Earth, validate_lat_lon_grid_args, generate_coordinate, with_precomputed_metrics, validate_rectilinear_grid_args
+using Oceananigans.Grids: default_indices, validate_indices, offset_data, instantiate, halo_size, topology
 
-import Oceananigans.Grids: zeros, StaticVerticalDiscretization
+import Oceananigans.Grids: zeros, StaticVerticalDiscretization, total_size
 import Oceananigans.Architectures: child_architecture
 
 import Oceananigans.DistributedComputations: 
@@ -214,64 +215,38 @@ function TripolarGrid(arch::ShardedDistributed,
 
     # Needed for partitial array assembly
     # device_to_array_slices = Reactant.sharding_to_array_slices(sharding, global_size)
-    irange = Colon()
-    jrange = Colon()
     FT = eltype(global_grid)
 
-    # Partitioning the Coordinates
-    λᶠᶠᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :λᶠᶠᵃ, irange, jrange)
-    φᶠᶠᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :φᶠᶠᵃ, irange, jrange)
-    λᶠᶜᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :λᶠᶜᵃ, irange, jrange)
-    φᶠᶜᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :φᶠᶜᵃ, irange, jrange)
-    λᶜᶠᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :λᶜᶠᵃ, irange, jrange)
-    φᶜᶠᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :φᶜᶠᵃ, irange, jrange)
-    λᶜᶜᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :λᶜᶜᵃ, irange, jrange)
-    φᶜᶜᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :φᶜᶜᵃ, irange, jrange)
-
-    # # Partitioning the Metrics
-    Δxᶜᶜᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :Δxᶜᶜᵃ, irange, jrange)
-    Δxᶠᶜᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :Δxᶠᶜᵃ, irange, jrange)
-    Δxᶜᶠᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :Δxᶜᶠᵃ, irange, jrange)
-    Δxᶠᶠᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :Δxᶠᶠᵃ, irange, jrange)
-    Δyᶜᶜᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :Δyᶜᶜᵃ, irange, jrange)
-    Δyᶠᶜᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :Δyᶠᶜᵃ, irange, jrange)
-    Δyᶜᶠᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :Δyᶜᶠᵃ, irange, jrange)
-    Δyᶠᶠᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :Δyᶠᶠᵃ, irange, jrange)
-    Azᶜᶜᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :Azᶜᶜᵃ, irange, jrange)
-    Azᶠᶜᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :Azᶠᶜᵃ, irange, jrange)
-    Azᶜᶠᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :Azᶜᶠᵃ, irange, jrange)
-    Azᶠᶠᵃ = OrthogonalSphericalShellGrids.partition_tripolar_metric(global_grid, :Azᶠᶠᵃ, irange, jrange)
-
     # Copying the z coordinate to all the devices: we pass a NamedSharding of `nothing`s
     # (a NamedSharding of nothings represents a copy to all devices)
     # ``1'' here is the maximum number of dimensions of the fields of ``z''
     replicate = Sharding.NamedSharding(arch.connectivity, ntuple(Returns(nothing), 1)) 
 
-    grid = OrthogonalSphericalShellGrid{Periodic,RightConnected,Bounded}(arch,
+    grid = OrthogonalSphericalShellGrid{Periodic, RightConnected, Bounded}(arch,
         global_size...,
         halo...,
         convert(FT, global_grid.Lz),
-        Reactant.to_rarray(λᶜᶜᵃ; sharding),
-        Reactant.to_rarray(λᶠᶜᵃ; sharding),
-        Reactant.to_rarray(λᶜᶠᵃ; sharding),
-        Reactant.to_rarray(λᶠᶠᵃ; sharding),
-        Reactant.to_rarray(φᶜᶜᵃ; sharding),
-        Reactant.to_rarray(φᶠᶜᵃ; sharding),
-        Reactant.to_rarray(φᶜᶠᵃ; sharding),
-        Reactant.to_rarray(φᶠᶠᵃ; sharding),
+        Reactant.to_rarray(global_grid.λᶜᶜᵃ; sharding),
+        Reactant.to_rarray(global_grid.λᶠᶜᵃ; sharding),
+        Reactant.to_rarray(global_grid.λᶜᶠᵃ; sharding),
+        Reactant.to_rarray(global_grid.λᶠᶠᵃ; sharding),
+        Reactant.to_rarray(global_grid.φᶜᶜᵃ; sharding),
+        Reactant.to_rarray(global_grid.φᶠᶜᵃ; sharding),
+        Reactant.to_rarray(global_grid.φᶜᶠᵃ; sharding),
+        Reactant.to_rarray(global_grid.φᶠᶠᵃ; sharding),
         sharded_z_direction(global_grid.z; sharding=replicate), # Replicating on all devices
-        Reactant.to_rarray(Δxᶜᶜᵃ; sharding),
-        Reactant.to_rarray(Δxᶠᶜᵃ; sharding),
-        Reactant.to_rarray(Δxᶜᶠᵃ; sharding),
-        Reactant.to_rarray(Δxᶠᶠᵃ; sharding),
-        Reactant.to_rarray(Δyᶜᶜᵃ; sharding),
-        Reactant.to_rarray(Δyᶠᶜᵃ; sharding),
-        Reactant.to_rarray(Δyᶜᶠᵃ; sharding),
-        Reactant.to_rarray(Δyᶠᶠᵃ; sharding),
-        Reactant.to_rarray(Azᶜᶜᵃ; sharding),
-        Reactant.to_rarray(Azᶠᶜᵃ; sharding),
-        Reactant.to_rarray(Azᶜᶠᵃ; sharding),
-        Reactant.to_rarray(Azᶠᶠᵃ; sharding),
+        Reactant.to_rarray(global_grid.Δxᶜᶜᵃ; sharding),
+        Reactant.to_rarray(global_grid.Δxᶠᶜᵃ; sharding),
+        Reactant.to_rarray(global_grid.Δxᶜᶠᵃ; sharding),
+        Reactant.to_rarray(global_grid.Δxᶠᶠᵃ; sharding),
+        Reactant.to_rarray(global_grid.Δyᶜᶜᵃ; sharding),
+        Reactant.to_rarray(global_grid.Δyᶠᶜᵃ; sharding),
+        Reactant.to_rarray(global_grid.Δyᶜᶠᵃ; sharding),
+        Reactant.to_rarray(global_grid.Δyᶠᶠᵃ; sharding),
+        Reactant.to_rarray(global_grid.Azᶜᶜᵃ; sharding),
+        Reactant.to_rarray(global_grid.Azᶠᶜᵃ; sharding),
+        Reactant.to_rarray(global_grid.Azᶜᶠᵃ; sharding),
+        Reactant.to_rarray(global_grid.Azᶠᶠᵃ; sharding),
         convert(FT, global_grid.radius),
         global_grid.conformal_mapping)
 

src/Grids/abstract_grid.jl

@@ -42,7 +42,7 @@ const XYZFlatGrid = AbstractGrid{<:Any, Flat, Flat, Flat}
 isrectilinear(grid) = false
 
 # Fallback
-@inline  get_active_column_map(::AbstractGrid) = nothing
+@inline get_active_column_map(::AbstractGrid) = nothing
 @inline get_active_cells_map(::AbstractGrid, any_map_type) = nothing
 
 """

src/Grids/vertical_discretization.jl

@@ -58,7 +58,7 @@ MutableVerticalDiscretization(r_faces) = MutableVerticalDiscretization(r_faces,
 ####
 
 const RegularStaticVerticalDiscretization  = StaticVerticalDiscretization{<:Any, <:Any, <:Number}
-const RegularMutableVerticalDiscretization = MutableVerticalDiscretization{<:Any,  <:Any, <:Number}
+const RegularMutableVerticalDiscretization = MutableVerticalDiscretization{<:Any, <:Any, <:Number}
 
 const RegularVerticalCoordinate = Union{RegularStaticVerticalDiscretization, RegularMutableVerticalDiscretization}
 

test/reactant_test_utils.jl

@@ -60,23 +60,29 @@ function test_reactant_model_correctness(GridType, ModelType, grid_kw, model_kw;
     u, v, w = model.velocities
     ru, rv, rw = r_model.velocities
 
-    # They will not be equal because r_model halos are not
-    # filled during set!
-    @test !(parent(u) ≈ parent(ru))
-    @test !(parent(v) ≈ parent(rv))
-    @test !(parent(w) ≈ parent(rw))
-
+    # Note that r_model halos are not filled during set!
+    # It's complicated to test this currently because the halo
+    # regions have different paddings, so we don't do it.
+    
     Oceananigans.TimeSteppers.update_state!(r_model)
 
     # Test that fields were set correctly
     @info "  After setting an initial condition:"
-    @show maximum(abs.(parent(u) .- parent(ru)))
-    @show maximum(abs.(parent(v) .- parent(rv)))
-    @show maximum(abs.(parent(w) .- parent(rw)))
+    rui = Array(interior(ru)) 
+    rvi = Array(interior(rv)) 
+    rwi = Array(interior(rw)) 
+
+    ui = Array(interior(u)) 
+    vi = Array(interior(v)) 
+    wi = Array(interior(w)) 
+
+    @show maximum(abs.(ui .- rui))
+    @show maximum(abs.(vi .- rvi))
+    @show maximum(abs.(wi .- rwi))
 
-    @test parent(u) ≈ parent(ru)
-    @test parent(v) ≈ parent(rv)
-    @test parent(w) ≈ parent(rw)
+    @test ui ≈ rui
+    @test vi ≈ rvi
+    @test wi ≈ rwi
 
     # Deduce a stable time-step
     Δx = minimum_xspacing(grid)
@@ -116,17 +122,25 @@ function test_reactant_model_correctness(GridType, ModelType, grid_kw, model_kw;
     @test iteration(r_simulation) == iteration(simulation)
     @test time(r_simulation) == time(simulation)
 
-    @show maximum(abs, parent(u))
-    @show maximum(abs, parent(v))
-    @show maximum(abs, parent(w))
+    rui = Array(interior(ru)) 
+    rvi = Array(interior(rv)) 
+    rwi = Array(interior(rw)) 
+
+    ui = Array(interior(u)) 
+    vi = Array(interior(v)) 
+    wi = Array(interior(w)) 
+
+    @show maximum(abs, ui)
+    @show maximum(abs, vi)
+    @show maximum(abs, wi)
 
-    @show maximum(abs.(parent(u) .- parent(ru)))
-    @show maximum(abs.(parent(v) .- parent(rv)))
-    @show maximum(abs.(parent(w) .- parent(rw)))
+    @show maximum(abs.(ui .- rui))
+    @show maximum(abs.(vi .- rvi))
+    @show maximum(abs.(wi .- rwi))
 
-    @test parent(u) ≈ parent(ru)
-    @test parent(v) ≈ parent(rv)
-    @test parent(w) ≈ parent(rw)
+    @test ui ≈ rui
+    @test vi ≈ rvi
+    @test wi ≈ rwi
 
     # Running a few more time-steps works too:
     r_simulation.stop_iteration += 2