Merge #1332

1332: allow broadcasting over horizontal spaces r=simonbyrne a=simonbyrne

Fixes #1331.

- [x] Code follows the [style guidelines](https://clima.github.io/ClimateMachine.jl/latest/DevDocs/CodeStyle/) OR N/A.
- [x] Unit tests are included OR N/A.
- [x] Code is exercised in an integration test OR N/A.
- [x] Documentation has been added/updated OR N/A.


Co-authored-by: Simon Byrne <[email protected]>

src/Fields/broadcast.jl

@@ -192,6 +192,10 @@ end
                parent(Spaces.local_geometry_data(space2))
            )
             warn_mismatched_spaces(typeof(space1), typeof(space2))
+        elseif Spaces.issubspace(space2, space1)
+            return space1
+        elseif Spaces.issubspace(space1, space2)
+            return space2
         else
             error_mismatched_spaces(typeof(space1), typeof(space2))
         end
@@ -229,6 +233,9 @@ end
                parent(Spaces.local_geometry_data(space2))
            )
             warn_mismatched_spaces(typeof(space1), typeof(space2))
+        elseif Spaces.issubspace(space2, space1) ||
+               Spaces.issubspace(space1, space2)
+            nothing
         else
             error_mismatched_spaces(typeof(space1), typeof(space2))
         end

src/Spaces/Spaces.jl

@@ -26,6 +26,8 @@ using StaticArrays, ForwardDiff, LinearAlgebra, UnPack, Adapt
 
 abstract type AbstractSpace end
 
+issubspace(::AbstractSpace, ::AbstractSpace) = false
+
 undertype(space::AbstractSpace) =
     Geometry.undertype(eltype(local_geometry_data(space)))
 

src/Spaces/extruded.jl

@@ -32,6 +32,21 @@ struct ExtrudedFiniteDifferenceSpace{
     face_ghost_geometry::LGG
 end
 
+function issubspace(
+    hspace::AbstractSpectralElementSpace,
+    extruded_space::ExtrudedFiniteDifferenceSpace,
+)
+    if hspace === extruded_space.horizontal_space
+        return true
+    end
+    # TODO: improve level handling
+    return Spaces.topology(hspace) ===
+           Spaces.topology(extruded_space.horizontal_space) &&
+           quadrature_style(hspace) ===
+           quadrature_style(extruded_space.horizontal_space)
+end
+
+
 Adapt.adapt_structure(to, space::ExtrudedFiniteDifferenceSpace) =
     ExtrudedFiniteDifferenceSpace(
         space.staggering,

test/Fields/field.jl

@@ -430,21 +430,13 @@ end
         ᶜx_surf = Spaces.level(Y.x, TU.fc_index(1, space))
 
         # Still need to define broadcast rules for surface planes with 3D domains
-        if nameof(typeof(space)) == :ExtrudedFiniteDifferenceSpace
-            @test_broken begin
-                try
-                    domain_surface_bc!(Y.x, ᶜz_surf)
-                    true
-                catch
-                    false
-                end
-            end
-        else
-            domain_surface_bc!(Y.x, ᶜz_surf, ᶜx_surf)
-        end
+        domain_surface_bc!(Y.x, ᶜz_surf, ᶜx_surf)
+
         # Skip spaces incompatible with Fields.bycolumn:
         TU.bycolumnable(space) || continue
-        column_surface_bc!(Y.x, ᶜz_surf, ᶜx_surf)
+        Yc = TU.FieldFromNamedTuple(space, (; x = FT(1)))
+        column_surface_bc!(Yc.x, ᶜz_surf, ᶜx_surf)
+        @test Y.x == Yc.x
         nothing
     end
     nothing