Add restart support for masked spaces (#2212)

charleskawczynski · web-flow · commit 8fb42cf13046 · 2025-04-12T09:47:44.000-04:00
Remove unused variable

Try MPI-safe mktempfile
diff --git a/src/InputOutput/readers.jl b/src/InputOutput/readers.jl
@@ -414,6 +414,49 @@ function read_grid(reader, name)
     end
 end
 
+"""
+    read_data_layout(dataset, topology)
+
+Read a datalayout from a `dataset`, with a given `topology`.
+
+This should cooperate with datasets written by `write!` for datalayouts.
+"""
+function read_data_layout(dataset, topology)
+    ArrayType = ClimaComms.array_type(topology)
+    data_layout = HDF5.read_attribute(dataset, "type")
+    has_horizontal = occursin('I', data_layout)
+    DataLayout = _scan_data_layout(data_layout)
+    array = HDF5.read(dataset)
+    has_horizontal &&
+        (h_dim = DataLayouts.h_dim(DataLayouts.singleton(DataLayout)))
+    if topology isa Topologies.Topology2D
+        nd = ndims(array)
+        localidx = ntuple(d -> d == h_dim ? topology.local_elem_gidx : (:), nd)
+        data = ArrayType(array[localidx...])
+    else
+        data = ArrayType(read(array))
+    end
+    has_horizontal && (Nij = size(data, findfirst("I", data_layout)[1]))
+    # For when `Nh` is added back to the type space
+    #     Nhd = Nh_dim(data_layout)
+    #     Nht = Nhd == -1 ? () : (size(data, Nhd),)
+    ElType = read_type(HDF5.read_attribute(dataset, "data_eltype"))
+    if data_layout in ("VIJFH", "VIFH")
+        Nv = size(data, 1)
+        # values = DataLayout{ElType, Nv, Nij, Nht...}(data) # when Nh is in type-domain
+        values = DataLayout{ElType, Nv, Nij}(data)
+    elseif data_layout in ("VF",)
+        Nv = size(data, 1)
+        values = DataLayout{ElType, Nv}(data)
+    elseif data_layout in ("DataF",)
+        values = DataLayout{ElType}(data)
+    else
+        # values = DataLayout{ElType, Nij, Nht...}(data) # when Nh is in type-domain
+        values = DataLayout{ElType, Nij}(data)
+    end
+    return values
+end
+
 function read_grid_new(reader, name)
     group = reader.file["grids/$name"]
     type = attrs(group)["type"]
@@ -426,11 +469,21 @@ function read_grid_new(reader, name)
         if type == "SpectralElementGrid1D"
             return Grids.SpectralElementGrid1D(topology, quadrature_style)
         else
-            return Grids.SpectralElementGrid2D(
+            enable_mask = haskey(attrs(group), "grid_mask")
+            grid = Grids.SpectralElementGrid2D(
                 topology,
                 quadrature_style;
                 enable_bubble,
+                enable_mask,
             )
+            if enable_mask
+                mask_type = keys(reader.file["grid_mask"])[1]
+                @assert mask_type == "IJHMask"
+                ds_is_active = reader.file["grid_mask"]["IJHMask"]["is_active"]
+                is_active = read_data_layout(ds_is_active, topology)
+                Grids.set_mask!(grid, is_active)
+            end
+            return grid
         end
     elseif type == "FiniteDifferenceGrid"
         topology = read_topology(reader, attrs(group)["topology"])
diff --git a/src/InputOutput/writers.jl b/src/InputOutput/writers.jl
@@ -321,6 +321,7 @@ end
 
 # Grids
 #
+defaultname(::DataLayouts.IJHMask) = "IJHMask"
 defaultname(::Grids.SpectralElementGrid1D) = "horizontal_grid"
 defaultname(::Grids.SpectralElementGrid2D) = "horizontal_grid"
 defaultname(::Grids.ExtrudedFiniteDifferenceGrid) =
@@ -374,6 +375,13 @@ function write_new!(
     )
     write_attribute(group, "bubble", grid.enable_bubble ? "true" : "false")
     write_attribute(group, "topology", write!(writer, Spaces.topology(grid)))
+    if !(grid.mask isa DataLayouts.NoMask)
+        write_attribute(
+            group,
+            "grid_mask",
+            write!(writer, grid.mask, Spaces.topology(grid)),
+        )
+    end
     return name
 end
 
@@ -450,6 +458,17 @@ function write_new!(
     return name
 end
 
+function write!(
+    writer::HDF5Writer,
+    mask::DataLayouts.IJHMask,
+    topology::Topologies.AbstractTopology,
+    name::AbstractString = defaultname(mask),
+)
+    group = create_group(writer.file, "grid_mask/$name")
+    write!(writer, group, mask.is_active, "is_active", topology)
+    return name
+end
+
 # write fields
 """
     write!(writer::HDF5Writer, field::Fields.Field, name::AbstractString)
@@ -511,6 +530,109 @@ function write!(
     )
 end
 
+"""
+    write!(
+        writer::HDF5Writer,
+        values::DataLayouts.AbstractData,
+        name::AbstractString,
+        topology::Topologies.AbstractTopology,
+    )
+
+Write an object of type `AbstractData` and name `name` to the HDF5 file.
+
+The `values` should belong to a `Field` whose `space`'s topology is
+`topology(axes(field))`.
+"""
+function write!(
+    writer::HDF5Writer,
+    group,
+    values::DataLayouts.AbstractData,
+    name::AbstractString,
+    topology::Topologies.AbstractTopology,
+)
+    if topology isa Topologies.Topology2D &&
+       !(writer.context isa ClimaComms.SingletonCommsContext)
+        nelems = Topologies.nelems(topology)
+        (; local_elem_gidx) = topology
+        _write_mpi!(group, values, name; nelems, local_elem_gidx)
+    else
+        _write!(group, values, name)
+    end
+end
+
+function write_plain_array!(group, array::AbstractArray, name::AbstractString)
+    nd = ndims(array)
+    dims = size(array)
+    localidx = ntuple(d -> (:), nd)
+    dataset =
+        create_dataset(group, name, datatype(eltype(array)), dataspace(dims))
+    dataset[localidx...] = array
+    return dataset
+end
+
+"""
+    _write_mpi!(
+        writer::HDF5Writer,
+        data::DataLayouts.AbstractData,
+        name::AbstractString,
+        nelems,
+        local_elem_gidx
+    )
+
+This is an internal method, meant to be used for writing data layouts to the
+HDF5 file.
+
+This method should be used for distributed datalayouts.
+"""
+function _write_mpi!(
+    group,
+    values::DataLayouts.AbstractData,
+    name::AbstractString;
+    nelems,
+    local_elem_gidx,
+)
+    h_dim = DataLayouts.h_dim(DataLayouts.singleton(values))
+    array = parent(values)
+    nd = ndims(array)
+    dims = ntuple(d -> d == h_dim ? nelems : size(array, d), nd)
+    localidx = ntuple(d -> d == h_dim ? local_elem_gidx : (:), nd)
+    dataset = create_dataset(
+        group,
+        "data/$name",
+        datatype(eltype(array)),
+        dataspace(dims);
+        dxpl_mpio = :collective,
+    )
+    dataset[localidx...] = array
+    write_attribute(dataset, "array", array)
+    write_attribute(dataset, "data_layout", string(nameof(typeof(values))))
+    write_attribute(dataset, "data_eltype", string(eltype(values)))
+    return name
+end
+
+"""
+    _write!(
+        writer::HDF5Writer,
+        data::DataLayouts.AbstractData,
+        name::AbstractString,
+    )
+
+This is an internal method, meant to be used for writing data layouts to the
+HDF5 file.
+
+This method should be used when this is not a distributed datalayout.
+"""
+function _write!(group, values::DataLayouts.AbstractData, name::AbstractString;)
+    h_dim = DataLayouts.h_dim(DataLayouts.singleton(values))
+    array = parent(values)
+    dataset = write_plain_array!(group, array, name)
+    write_attribute(dataset, "array", array)
+    write_attribute(dataset, "type", string(nameof(typeof(values))))
+    write_attribute(dataset, "data_eltype", string(eltype(values)))
+    return name
+end
+
+
 """
     write!(
         writer::HDF5Writer,
diff --git a/test/InputOutput/unit_spectralelement2d.jl b/test/InputOutput/unit_spectralelement2d.jl
@@ -1,8 +1,13 @@
+#=
+julia --project=.buildkite
+using Revise; include("test/InputOutput/unit_spectralelement2d.jl")
+=#
 using Test
 using ClimaComms
 ClimaComms.@import_required_backends
 using LinearAlgebra
 import ClimaCore
+import ClimaCore.Utilities.Cache
 import ClimaCore:
     Domains,
     Fields,
@@ -41,16 +46,7 @@ function init_state(local_geometry, p)
     return (ρ = ρ, u = u, ρθ = ρ * θ)
 end
 
-@testset "restart test for 2D spectral element simulations" begin
-    parameters = (
-        ϵ = 0.1,  # perturbation size for initial condition
-        l = 0.5, # Gaussian width
-        k = 0.5, # Sinusoidal wavenumber
-        ρ₀ = 1.0, # reference density
-        c = 2,
-        g = 10,
-        D₄ = 1e-4, # hyperdiffusion coefficient
-    )
+function init_space(context; enable_bubble, enable_mask)
     domain = Domains.RectangleDomain(
         Domains.IntervalDomain(
             Geometry.XPoint(-2π),
@@ -67,31 +63,92 @@ end
     Nq = 4
     quad = Quadratures.GLL{Nq}()
     mesh = Meshes.RectilinearMesh(domain, n1, n2)
+    grid_topology = Topologies.Topology2D(context, mesh)
+
+    return Spaces.SpectralElementSpace2D(
+        grid_topology,
+        quad;
+        enable_bubble,
+        enable_mask,
+    )
+end
+
+# function mktempfile(f)
+#     mktempdir() do dir
+#         cd(dir) do
+#             f(tempname(dir))
+#         end
+#     end
+# end
+
+function mktempfile(f, context)
+    filename =
+        ClimaComms.iamroot(context) ? tempname(pwd(); cleanup = true) : ""
+    filename = ClimaComms.bcast(context, filename)
+    f(filename)
+end
+
+@testset "restart test for 2D spectral element simulations" begin
     device = ClimaComms.device()
     @info "Using device" device
-    context = ClimaComms.SingletonCommsContext(device)
-    grid_topology = Topologies.Topology2D(context, mesh)
+    context = ClimaComms.context(device)
+    parameters = (
+        ϵ = 0.1,  # perturbation size for initial condition
+        l = 0.5, # Gaussian width
+        k = 0.5, # Sinusoidal wavenumber
+        ρ₀ = 1.0, # reference density
+        c = 2,
+        g = 10,
+        D₄ = 1e-4, # hyperdiffusion coefficient
+    )
+    space = init_space(context; enable_bubble = true, enable_mask = false)
+    y0 = init_state.(Fields.local_geometry_field(space), Ref(parameters))
+    Y = Fields.FieldVector(y0 = y0)
+
+    # write field vector to hdf5 file
+    mktempfile(context) do filename
+        InputOutput.HDF5Writer(filename, context) do writer
+            InputOutput.write!(writer, "Y" => Y) # write field vector from hdf5 file
+        end
+        Cache.clean_cache!()
+        InputOutput.HDF5Reader(filename, context) do reader
+            restart_Y = InputOutput.read_field(reader, "Y") # read fieldvector from hdf5 file
+            @test restart_Y == Y # test if restart is exact
+            @test axes(restart_Y) == axes(Y) # test if restart is exact for space
+        end
+    end
 
-    for enable_bubble in (true, false)
-        space =
-            Spaces.SpectralElementSpace2D(grid_topology, quad; enable_bubble)
+    # Test with masks
+    space = init_space(context; enable_bubble = true, enable_mask = true)
+    y0 = init_state.(Fields.local_geometry_field(space), Ref(parameters))
+    Y = Fields.FieldVector(y0 = y0)
 
-        y0 = init_state.(Fields.local_geometry_field(space), Ref(parameters))
-        Y = Fields.FieldVector(y0 = y0)
+    Spaces.set_mask!(space) do coords
+        rand() > 0.5
+    end
 
-        # write field vector to hdf5 file
-        filename = tempname(pwd())
+    # write field vector to hdf5 file
+    mktempfile(context) do filename
         InputOutput.HDF5Writer(filename, context) do writer
             InputOutput.write!(writer, "Y" => Y) # write field vector from hdf5 file
         end
+
         InputOutput.HDF5Reader(filename, context) do reader
+            # We need to clean the cache so that the next read of space
+            # does not use a pointer to the cached one.
+            Cache.clean_cache!()
             restart_Y = InputOutput.read_field(reader, "Y") # read fieldvector from hdf5 file
-            ClimaComms.allowscalar(device) do
-                @test restart_Y == Y # test if restart is exact
-            end
-            ClimaComms.allowscalar(device) do
-                @test axes(restart_Y) == axes(Y) # test if restart is exact for space
-            end
+
+            is_active_restart =
+                parent(Spaces.get_mask(axes(restart_Y.y0)).is_active)
+            is_active = parent(Spaces.get_mask(axes(Y.y0)).is_active)
+            @test is_active == is_active_restart
+
+            # Test that we're not doing trivial pointer comparisons
+            @test !(axes(Y.y0) === axes(restart_Y.y0))
+            @test restart_Y == Y
+            @test typeof(axes(restart_Y.y0)) == typeof(axes(Y.y0))
+            @test axes(restart_Y) == axes(Y) # test if restart is exact for space
         end
     end
 end
