Fix for #2358

Author: ray-chew

src/Grids/level.jl

@@ -1,18 +1,19 @@
 struct LevelGrid{
     G <: AbstractExtrudedFiniteDifferenceGrid,
     L <: Union{Int, PlusHalf{Int}},
+    Q <: Quadratures.QuadratureStyle,
 } <: AbstractGrid
     full_grid::G
     level::L
+    quadrature_style::Q
 end
 
-quadrature_style(levelgrid::LevelGrid) =
-    quadrature_style(levelgrid.full_grid.horizontal_grid)
+quadrature_style(levelgrid::LevelGrid) = levelgrid.quadrature_style
 
 level(
     grid::AbstractExtrudedFiniteDifferenceGrid,
     level::Union{Int, PlusHalf{Int}},
-) = LevelGrid(grid, level)
+) = LevelGrid(grid, level, quadrature_style(grid))
 
 topology(levelgrid::LevelGrid) = topology(levelgrid.full_grid)
 
@@ -21,22 +22,28 @@ topology(levelgrid::LevelGrid) = topology(levelgrid.full_grid)
 # need to extract the weights at a particular level.
 dss_weights(levelgrid::LevelGrid, _) = dss_weights(levelgrid.full_grid, nothing)
 
-local_geometry_type(::Type{LevelGrid{G, L}}) where {G, L} =
+local_geometry_type(::Type{LevelGrid{G, L, Q}}) where {G, L, Q} =
     local_geometry_type(G)
-local_geometry_data(levelgrid::LevelGrid{<:Any, Int}, ::Nothing) = level(
+
+local_geometry_data(levelgrid::LevelGrid{<:Any, Int, <:Any}, ::Nothing) = level(
     local_geometry_data(levelgrid.full_grid, CellCenter()),
     levelgrid.level,
 )
-local_geometry_data(levelgrid::LevelGrid{<:Any, PlusHalf{Int}}, ::Nothing) =
-    level(
-        local_geometry_data(levelgrid.full_grid, CellFace()),
-        levelgrid.level + half,
-    )
-global_geometry(levlgrid::LevelGrid) = global_geometry(levlgrid.full_grid)
+local_geometry_data(
+    levelgrid::LevelGrid{<:Any, PlusHalf{Int}, <:Any},
+    ::Nothing,
+) = level(
+    local_geometry_data(levelgrid.full_grid, CellFace()),
+    levelgrid.level + half,
+)
+global_geometry(levelgrid::LevelGrid) = global_geometry(levelgrid.full_grid)
 
 ## GPU compatibility
-Adapt.adapt_structure(to, grid::LevelGrid) =
-    LevelGrid(Adapt.adapt(to, grid.full_grid), grid.level)
+Adapt.adapt_structure(to, grid::LevelGrid) = LevelGrid(
+    Adapt.adapt(to, grid.full_grid),
+    grid.level,
+    quadrature_style(grid),
+)
 
 ## aliases
 const LevelCubedSphereSpectralElementGrid2D =

src/InputOutput/readers.jl

@@ -522,7 +522,24 @@ function read_grid_new(reader, name)
         else
             level = attrs(group)["level_half"] + half
         end
-        return Grids.LevelGrid(full_grid, level)
+
+        # Check if quadrature attributes exist in current group, otherwise use horizontal_grid
+        if haskey(attrs(group), "quadrature_num_points") &&
+           haskey(attrs(group), "quadrature_type")
+            group = reader.file["grids/horizontal_grid"]
+            npts = attrs(group)["quadrature_num_points"]
+            quadrature_style =
+                _scan_quadrature_style(attrs(group)["quadrature_type"], npts)
+        else
+            horizontal_group = reader.file["grids/horizontal_grid"]
+            npts = attrs(horizontal_group)["quadrature_num_points"]
+            quadrature_style = _scan_quadrature_style(
+                attrs(horizontal_group)["quadrature_type"],
+                npts,
+            )
+        end
+
+        return Grids.LevelGrid(full_grid, level, quadrature_style)
     else
         error("Unsupported grid type $type")
     end

test/Fields/unit_field.jl

@@ -689,6 +689,31 @@ end
     end
 end
 
+@testset "Levels of nonlocal Fields and nonlocal Field broadcasts" begin
+    FT = Float64
+    gradh = Operators.Gradient()
+    # Todo: Make this work over all spaces; currently broken for everything else.
+    for space in (
+        TU.CenterExtrudedFiniteDifferenceSpace(FT),
+        TU.FaceExtrudedFiniteDifferenceSpace(FT),
+    )
+        TU.levelable(space) || continue
+        field = fill((; x = FT(1)), space)
+
+        op_on_level_of_field = gradh.(Fields.Field(
+            Spaces.level(Fields.field_values(field.x), 1),
+            Spaces.level(space, TU.fc_index(1, space)),
+        ))
+
+        @test op_on_level_of_field ==
+              (Spaces.level(gradh.(field.x), TU.fc_index(1, space)))
+
+        @test_broken op_on_level_of_field == Base.materialize(
+            (Spaces.level(lazy.(gradh.(field.x)), TU.fc_index(1, space))),
+        )
+    end
+end
+
 @testset "Columns of Fields and Field broadcasts" begin
     FT = Float64
     for space in TU.all_spaces(FT)
@@ -736,7 +761,7 @@ end
                 is_cuda && space isa OneSlabIndexSpace
             @test slab_of_field == Base.materialize(
                 Spaces.slab(lazy.(identity.(field)), indices...),
-            ) broken = is_cuda
+            ) broken = is_cuda && space isa OneSlabIndexSpace
             # TODO: Figure out why some of these tests are broken on GPUs.
         end
     end

test/InputOutput/unit_hybrid3dcubedsphere_topography.jl

@@ -13,6 +13,11 @@ using ClimaCore:
     InputOutput,
     Grids
 
+@isdefined(TU) || include(
+    joinpath(pkgdir(ClimaCore), "test", "TestUtilities", "TestUtilities.jl"),
+);
+import .TestUtilities as TU;
+
 using ClimaComms
 const comms_ctx = ClimaComms.context(ClimaComms.CPUSingleThreaded())
 pid, nprocs = ClimaComms.init(comms_ctx)
@@ -54,7 +59,7 @@ end
                 z_max / 8 .* (
                     cosd.(Fields.coordinate_field(h_space).lat) .+
                     cosd.(Fields.coordinate_field(h_space).long) .+ 1
-                )
+                ),
             )
 
         z_mesh = Meshes.IntervalMesh(z_domain, nelems = z_elem)
@@ -90,3 +95,50 @@ end
         end
     end
 end
+
+
+@testset "HDF5 restart test for a Named Tuple of Levels of a 3D hybrid cubed sphere for deep" begin
+    # This I/O is used for the computation of the topographic drag
+    FT = Float32
+
+    for space in (
+        TU.CenterExtrudedFiniteDifferenceSpace(FT, context = comms_ctx),
+        TU.FaceExtrudedFiniteDifferenceSpace(FT, context = comms_ctx),
+    )
+        TU.levelable(space) || continue
+        field = fill((; x = FT(1)), space)
+
+        level_of_field = Fields.Field(
+            Spaces.level(Fields.field_values(field.x), 1),
+            Spaces.level(space, TU.fc_index(1, space)),
+        )
+
+        fake_drag = fill(
+            (;
+                t11 = FT(0.0),
+                t12 = FT(0.0),
+                t21 = FT(0.0),
+                t22 = FT(0.0),
+                hmin = FT(0.0),
+                hmax = FT(0.0),
+            ),
+            axes(level_of_field),
+        )
+
+        # write field vector to hdf5 file
+        InputOutput.HDF5Writer(filename, comms_ctx) do writer
+            InputOutput.write!(writer, fake_drag, "fake_drag")
+        end
+
+        InputOutput.HDF5Reader(filename, comms_ctx) do reader
+            restart_fake_drag = InputOutput.read_field(reader, "fake_drag") # read fieldvector from hdf5 file
+
+            # The underlying space is of a different instance, so we cannot use == to check for equivalence.
+            # Instead, we make sure that the values and types are the same.
+            @test typeof(restart_fake_drag) == typeof(fake_drag)
+            @test maximum(
+                abs.(parent(fake_drag.t21) .- parent(restart_fake_drag.t21)),
+            ) == 0.0f0
+        end
+    end
+end