VF getindex func + test

move func to indices.jl prior format
CliMA · Sep 2, 2023 · ace210f · ace210f
1 parent 25942f9
commit ace210f
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 1 deletion.
diff --git a/src/Fields/indices.jl b/src/Fields/indices.jl
@@ -17,7 +17,13 @@ end
 
 Base.@propagate_inbounds Base.getindex(field::Field, colidx::ColumnIndex) =
     column(field, colidx)
-
+Base.@propagate_inbounds Base.getindex(fv::FieldVector, colidx::ColumnIndex) =
+    (;
+        zip(
+            propertynames(fv),
+            map(x -> getproperty(fv, x)[colidx], propertynames(fv)),
+        )...
+    )
 Base.@propagate_inbounds function column(
     field::SpectralElementField1D,
     colidx::ColumnIndex{1},

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)