Fix reading/writing purely vertical spaces

This previously was not supported
CliMA · Dec 16, 2024 · 7a098cf · 7a098cf
1 parent 340603b
commit 7a098cf
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 3 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -4,6 +4,7 @@ ClimaCore.jl Release Notes
 main
 -------
 
+
  - We've added new convenience constructors for spaces PR [2082](https://github.com/CliMA/ClimaCore.jl/pull/2082). Here are links to the new constructors:
    - [ExtrudedCubedSphereSpace]()
    - [CubedSphereSpace]()
@@ -12,7 +13,12 @@ main
    - [SliceXZSpace]()
    - [RectangleXYSpace]()
 
+### ![][badge-🐛bugfix] Bug fixes
+
+- Fixed writing/reading purely vertical spaces
+
 v0.14.20
+--------
 
  - We've added new convenience constructors for grids PR [1848](https://github.com/CliMA/ClimaCore.jl/pull/1848). Here are links to the new constructors:
    - [ExtrudedCubedSphereGrid](https://github.com/CliMA/ClimaCore.jl/blob/cbb193042fac3b4bef33251fbc0f232427bfe506/src/CommonGrids/CommonGrids.jl#L85-L144)

diff --git a/src/InputOutput/readers.jl b/src/InputOutput/readers.jl
@@ -152,17 +152,19 @@ function _scan_data_layout(layoutstring::AbstractString)
         "IFH",
         "IHF",
         "IF",
+        "VF",
         "VIJFH",
         "VIJHF",
         "VIFH",
         "VIHF",
-    )
+    ) "datalayout is $layoutstring"
     layoutstring == "IJFH" && return DataLayouts.IJFH
     layoutstring == "IJHF" && return DataLayouts.IJHF
     layoutstring == "IJF" && return DataLayouts.IJF
     layoutstring == "IFH" && return DataLayouts.IFH
     layoutstring == "IHF" && return DataLayouts.IHF
     layoutstring == "IF" && return DataLayouts.IF
+    layoutstring == "VF" && return DataLayouts.VF
     layoutstring == "VIJFH" && return DataLayouts.VIJFH
     layoutstring == "VIJHF" && return DataLayouts.VIJHF
     return DataLayouts.VIFH
@@ -482,8 +484,10 @@ function read_field(reader::HDF5Reader, name::AbstractString)
         topology = Spaces.topology(space)
         ArrayType = ClimaComms.array_type(topology)
         data_layout = attrs(obj)["data_layout"]
+        has_horizontal = occursin(data_layout, "I")
         DataLayout = _scan_data_layout(data_layout)
-        h_dim = DataLayouts.h_dim(DataLayouts.singleton(DataLayout))
+        has_horizontal &&
+            (h_dim = DataLayouts.h_dim(DataLayouts.singleton(DataLayout)))
         if topology isa Topologies.Topology2D
             nd = ndims(obj)
             localidx =
@@ -492,7 +496,7 @@ function read_field(reader::HDF5Reader, name::AbstractString)
         else
             data = ArrayType(read(obj))
         end
-        Nij = size(data, findfirst("I", data_layout)[1])
+        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),)
@@ -501,6 +505,9 @@ function read_field(reader::HDF5Reader, name::AbstractString)
             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)
         else
             # values = DataLayout{ElType, Nij, Nht...}(data) # when Nh is in type-domain
             values = DataLayout{ElType, Nij}(data)

diff --git a/test/InputOutput/finitedifference.jl b/test/InputOutput/finitedifference.jl
@@ -0,0 +1,49 @@
+using Test
+import ClimaCore
+
+using ClimaComms
+const comms_ctx = ClimaComms.context(ClimaComms.CPUSingleThreaded())
+pid, nprocs = ClimaComms.init(comms_ctx)
+filename = ClimaComms.bcast(comms_ctx, tempname(pwd()))
+if ClimaComms.iamroot(comms_ctx)
+    @info "Comms context" comms_ctx nprocs filename
+end
+
+@testset "HDF5 restart test for 1d finite difference space" begin
+    FT = Float32
+
+    z_min = FT(0)
+    z_max = FT(30e3)
+    z_elem = 10
+    center_staggering = ClimaCore.Grids.CellCenter()
+    face_staggering = ClimaCore.Grids.CellFace()
+
+    center_space = ClimaCore.CommonSpaces.ColumnSpace(;
+        z_min,
+        z_max,
+        z_elem,
+        staggering = center_staggering,
+    )
+
+    face_space = ClimaCore.CommonSpaces.ColumnSpace(;
+        z_min,
+        z_max,
+        z_elem,
+        staggering = face_staggering,
+    )
+
+    center_field = Fields.local_geometry_field(center_space)
+    face_field = Fields.local_geometry_field(face_space)
+
+    Y = ClimaCore.Fields.FieldVector(; c = center_field, f = face_field)
+
+    # write field vector to hdf5 file
+    writer = ClimaCore.InputOutput.HDF5Writer(filename, comms_ctx)
+    ClimaCore.InputOutput.write!(writer, Y, "Y")
+    close(writer)
+
+    reader = ClimaCore.InputOutput.HDF5Reader(filename, comms_ctx)
+    restart_Y = ClimaCore.InputOutput.read_field(reader, "Y") # read fieldvector from hdf5 file
+    close(reader)
+    @test restart_Y == Y # test if restart is exact
+end