Test that output is saved in cubed sphere simulations

src/Models/HydrostaticFreeSurfaceModels/set_hydrostatic_free_surface_model.jl

@@ -62,7 +62,7 @@ model.velocities.u
             throw(ArgumentError("name $fldname not found in model.velocities, model.tracers, or model.free_surface"))
         end
 
-        @apply_regionally set!(ϕ, value)
+        set!(ϕ, value)
     end
 
     # initialize!(model)

src/MultiRegion/multi_region_boundary_conditions.jl

@@ -13,7 +13,7 @@ using Oceananigans.BoundaryConditions:
 
 import Oceananigans.BoundaryConditions: fill_halo_regions!, fill_halo_event!
 
-@inline bc_str(::MultiRegionObject) = "MultiRegion Boundary Conditions"
+@inline bc_str(::MultiRegionObject) = "MultiRegionObject"
 
 @inline function fill_halo_regions!(fields::NamedTuple, grid::ConformalCubedSphereGridOfSomeKind, args...; kwargs...)
     u = haskey(fields, :u) ? fields.u : nothing

src/MultiRegion/multi_region_field.jl

@@ -1,6 +1,6 @@
 using Oceananigans.AbstractOperations: AbstractOperation, compute_computed_field!
 using Oceananigans.BoundaryConditions: default_auxiliary_bc
-using Oceananigans.Fields: FunctionField, data_summary, AbstractField, instantiated_location
+using Oceananigans.Fields: FunctionField, data_summary, AbstractField, instantiated_location, set_to_function!
 using Oceananigans.Operators: assumed_field_location
 using Oceananigans.OutputWriters: output_indices
 
@@ -9,8 +9,17 @@ using Base: @propagate_inbounds
 import Oceananigans.BoundaryConditions: regularize_field_boundary_conditions
 import Oceananigans.Diagnostics: hasnan
 import Oceananigans.DistributedComputations: reconstruct_global_field, CommunicationBuffers
-import Oceananigans.Fields: set!, compute!, compute_at!, interior, communication_buffers,
-                            validate_field_data, validate_boundary_conditions, validate_indices
+import Oceananigans.Fields: set_to_array!,
+                            set_to_field!,
+                            compute!, 
+                            compute_at!, 
+                            interior, 
+                            communication_buffers,
+                            validate_field_data, 
+                            validate_boundary_conditions, 
+                            validate_indices,
+                            set!
+
 import Oceananigans.Grids: xnodes, ynodes
 
 import Base: fill!, axes
@@ -111,15 +120,20 @@ function reconstruct_global_indices(indices, p::YPartition, N)
     return (idx_x, idx_y, idx_z)
 end
 
-## Functions applied regionally
-set!(mrf::MultiRegionField, v)  = apply_regionally!(set!,  mrf, v)
-fill!(mrf::MultiRegionField, v) = apply_regionally!(fill!, mrf, v)
+# Set fields regionally
+set_to_array!(mrf::MultiRegionField, a) = apply_regionally!(set_to_array!, mrf, a)
+set_to_field!(mrf::MultiRegionField, v) = apply_regionally!(set_to_field!, mrf, v)
 
-set!(mrf::MultiRegionField, a::Number)  = apply_regionally!(set!,  mrf, a)
+# Exporting some set! methods
+# set! a function field can be safely done regionally
+set!(mrf::MultiRegionField, f::FunctionField) = apply_regionally!(set!, mrf, f)
+set!(mrf::MultiRegionField, f::Function)      = set_to_function!(mrf, f)
+set!(mrf::MultiRegionField, mro::MultiRegionObject) = apply_regionally!(set!, mrf, mro)
+
+# Fill fields regionally
+fill!(mrf::MultiRegionField, v) = apply_regionally!(fill!, mrf, v)
 fill!(mrf::MultiRegionField, a::Number) = apply_regionally!(fill!, mrf, a)
 
-set!(mrf::MultiRegionField, f::Function) = apply_regionally!(set!, mrf, f)
-set!(u::MultiRegionField, v::MultiRegionField) = apply_regionally!(set!, u, v)
 compute!(mrf::GriddedMultiRegionField, time=nothing) = apply_regionally!(compute!, mrf, time)
 
 # Disambiguation (same as computed_field.jl:64)
@@ -216,3 +230,23 @@ ynodes(ψ::AbstractField{<:Any, <:Any, <:Any, <:OrthogonalSphericalShellGrid}) =
 # Convenience
 @propagate_inbounds Base.getindex(mrf::MultiRegionField, r::Int) = getregion(mrf, r)
 @propagate_inbounds Base.lastindex(mrf::MultiRegionField) = lastindex(mrf.grid)
+
+import Base: ==
+
+function ==(a::MultiRegionField, b::MultiRegionField)
+    if regions(a) == regions(b)
+        return all(a[r] == b[r] for r in regions(a))
+    else
+        return false
+    end
+end
+
+import Base: isapprox    
+
+function isapprox(a::MultiRegionField, b::MultiRegionField; kw...)
+    if regions(a) == regions(b)
+        return all(isapprox(a[r], b[r]; kw...) for r in regions(a))
+    else
+        return false
+    end
+end

src/OrthogonalSphericalShellGrids/conformal_cubed_sphere_panel.jl

@@ -34,7 +34,8 @@ function Adapt.adapt_structure(to, conformal_mapping::CubedSphereConformalMappin
         adapt(to, conformal_mapping.ξᶠᵃᵃ),
         adapt(to, conformal_mapping.ηᵃᶠᵃ),
         adapt(to, conformal_mapping.ξᶜᵃᵃ),
-        adapt(to, conformal_mapping.ηᵃᶜᵃ))
+        adapt(to, conformal_mapping.ηᵃᶜᵃ)
+    )
 end
 
 const ConformalCubedSpherePanelGrid{FT, TX, TY, TZ, CZ, CC, FC, CF, FF, Arch} =

src/Utils/multi_region_transformation.jl

@@ -14,6 +14,12 @@ struct MultiRegionObject{R}
     regional_objects :: R
 end
 
+function Base.summary(mo::MultiRegionObject)
+    obj = summary(first(mo.regional_objects))
+    N = length(mo.regional_objects)
+    return "MultiRegionObject with $N × $obj"
+end
+
 #####
 ##### Convenience structs
 #####

test/runtests.jl

@@ -232,7 +232,6 @@ CUDA.allowscalar() do
         end
     end
 
-
     # Tests for Enzyme extension
     if group == :enzyme || group == :all
         @testset "Enzyme extension tests" begin