Add fieldvector unit tests

Author: charleskawczynski

src/Fields/fieldvector.jl

@@ -469,36 +469,54 @@ end
 
 
 # Recursively compare contents of similar fieldvectors
-_rcompare(pass, x::T, y::T) where {T <: Field} =
-    pass && _rcompare(pass, field_values(x), field_values(y))
-_rcompare(pass, x::T, y::T) where {T <: DataLayouts.AbstractData} =
+_rcompare(pass, x::T, y::T; strict) where {T <: Field} =
+    pass && _rcompare(pass, field_values(x), field_values(y); strict)
+_rcompare(pass, x::T, y::T; strict) where {T <: DataLayouts.AbstractData} =
     pass && (parent(x) == parent(y))
-_rcompare(pass, x::T, y::T) where {T} = pass && (x == y)
+_rcompare(pass, x::T, y::T; strict) where {T} = pass && (x == y)
 
-function _rcompare(pass, x::T, y::T) where {T <: Union{FieldVector, NamedTuple}}
+_rcompare(pass, x::NamedTuple, y::NamedTuple; strict) =
+    _rcompare_nt(pass, x, y; strict)
+_rcompare(pass, x::FieldVector, y::FieldVector; strict) =
+    _rcompare_nt(pass, x, y; strict)
+
+function _rcompare_nt(pass, x, y; strict)
+    length(propertynames(x)) ≠ length(propertynames(y)) && return false
+    if strict
+        typeof(x) == typeof(y) || return false
+    end
     for pn in propertynames(x)
-        pass &= _rcompare(pass, getproperty(x, pn), getproperty(y, pn))
+        pass &= _rcompare(pass, getproperty(x, pn), getproperty(y, pn); strict)
     end
     return pass
 end
 
 """
-    rcompare(x::T, y::T) where {T <: Union{FieldVector, NamedTuple}}
+    rcompare(x::T, y::T; strict = true) where {T <: Union{FieldVector, NamedTuple}}
 
 Recursively compare given fieldvectors via `==`.
 Returns `true` if `x == y` recursively.
 
 FieldVectors with different types are considered different.
 """
-rcompare(x::T, y::T) where {T <: Union{FieldVector, NamedTuple}} =
-    _rcompare(true, x, y)
+rcompare(
+    x::T,
+    y::T;
+    strict = true,
+) where {T <: Union{FieldVector, NamedTuple}} = _rcompare(true, x, y; strict)
 
-rcompare(x::T, y::T) where {T <: FieldVector} = _rcompare(true, x, y)
+rcompare(x::T, y::T; strict = true) where {T <: FieldVector} =
+    _rcompare(true, x, y; strict)
 
-rcompare(x::T, y::T) where {T <: NamedTuple} = _rcompare(true, x, y)
+rcompare(x::T, y::T; strict = true) where {T <: NamedTuple} =
+    _rcompare(true, x, y; strict)
 
 # FieldVectors with different types are always different
-rcompare(x::FieldVector, y::FieldVector) = false
+rcompare(x::FieldVector, y::FieldVector; strict::Bool = true) =
+    strict ? false : _rcompare(true, x, y; strict)
+
+rcompare(x::NamedTuple, y::NamedTuple; strict::Bool = true) =
+    strict ? false : _rcompare(true, x, y; strict)
 
 # Define == to call rcompare for two fieldvectors
-Base.:(==)(x::FieldVector, y::FieldVector) = rcompare(x, y)
+Base.:(==)(x::FieldVector, y::FieldVector) = rcompare(x, y; strict = true)

test/Fields/unit_field.jl

@@ -11,6 +11,7 @@ ClimaComms.@import_required_backends
 using OrderedCollections
 using StaticArrays, IntervalSets
 import ClimaCore
+import ClimaCore.InputOutput
 import ClimaCore.Utilities: PlusHalf
 import ClimaCore.DataLayouts
 import ClimaCore.DataLayouts: IJFH
@@ -330,6 +331,67 @@ end
     @test occursin("==================== Difference found:", s)
 end
 
+@testset "FieldVector restart + broadcasting" begin
+    FT = Float32
+    root_path = "bucket_restart"
+    rm(root_path; force = true, recursive = true)
+    mkpath(root_path)
+
+    radius = FT(100)
+    depth = FT(3.5)
+    nelements = (1, 10)
+    npolynomial = 1
+    dz_tuple = nothing
+    context = ClimaComms.context()
+
+    vertdomain = Domains.IntervalDomain(
+        Geometry.ZPoint(FT(-depth)),
+        Geometry.ZPoint(FT(0));
+        boundary_names = (:bottom, :top),
+    )
+    vertmesh = Meshes.IntervalMesh(vertdomain; nelems = nelements[2])
+    device = ClimaComms.device()
+    vert_center_space = Spaces.CenterFiniteDifferenceSpace(device, vertmesh)
+    horzdomain = Domains.SphereDomain(radius)
+    horzmesh = Meshes.EquiangularCubedSphere(horzdomain, nelements[1])
+    horztopology = Topologies.Topology2D(context, horzmesh)
+    quad = Spaces.Quadratures.GLL{npolynomial + 1}()
+    space = Spaces.SpectralElementSpace2D(horztopology, quad)
+
+    vars = (;
+        bucket = (;
+            W = Fields.Field(FT, space),
+            T = Fields.Field(FT, space),
+            Ws = Fields.Field(FT, space),
+            σS = Fields.Field(FT, space),
+        )
+    )
+
+    Y = Fields.FieldVector(; vars...)
+    Y.bucket.T .= 280.0
+    Y.bucket.W .= 0.05
+    Y.bucket.Ws .= 0.0
+    Y.bucket.σS .= 0.08
+
+    output_file = joinpath(root_path, "day0.0.hdf5")
+    hdfwriter = InputOutput.HDF5Writer(output_file, context)
+    InputOutput.write_attributes!(
+        hdfwriter,
+        "/",
+        Dict("time" => 0.0, "hash" => "1234"),
+    )
+    InputOutput.write!(hdfwriter, Y, "Y")
+    Base.close(hdfwriter)
+
+    restart_file = output_file
+    Y_restart = Fields.FieldVector(; vars...)
+    hdfreader = InputOutput.HDF5Reader(restart_file, context)
+    Y_restart = InputOutput.read_field(hdfreader, "Y")
+    close(hdfreader)
+    Y .= Y_restart
+    @test Fields.rcompare(Y, Y_restart; strict = false)
+end
+
 # https://github.com/CliMA/ClimaCore.jl/issues/1465
 @testset "Diagonal FieldVector broadcast expressions" begin
     FT = Float64