VF getindex func + test

CliMA · Sep 1, 2023 · cc4e273 · cc4e273
1 parent 25942f9
commit cc4e273
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/src/Fields/fieldvector.jl b/src/Fields/fieldvector.jl
@@ -22,7 +22,12 @@ struct FieldVector{T, M} <: BlockArrays.AbstractBlockVector{T}
     values::M
 end
 FieldVector{T}(values::M) where {T, M} = FieldVector{T, M}(values)
-
+Base.getindex(fv::Fields.FieldVector, colidx::Fields.ColumnIndex) = (;
+    zip(
+        propertynames(fv),
+        map(x -> getproperty(fv, x)[colidx], propertynames(fv)),
+    )...
+)
 
 """
     Fields.ScalarWrapper(val) <: AbstractArray{T,0}

diff --git a/test/Fields/field.jl b/test/Fields/field.jl
@@ -267,6 +267,20 @@ end
     @test Y.k.z === 3.0
 end
 
+@testset "FieldVector getindex" begin
+    space = spectral_space_2D()
+
+    # arbitrary column index in spectral_space_2D
+    colidx = Fields.ColumnIndex((1, 1), 1)
+
+    y1 = ones(space)
+    y2 = ones(space) .* 2
+    Y = Fields.FieldVector(y1 = y1, y2 = y2)
+
+    @test parent(Y[colidx][1])[1] == 1.0
+    @test parent(Y[colidx][2])[1] == 2.0
+end
+
 @testset "FieldVector array_type" begin
     device = ClimaComms.device()
     context = ClimaComms.SingletonCommsContext(device)