VF getindex func + test

move func to indices.jl

prior format

make output a FieldVector

rev + type stable

Co-authored-by: @charleskawczynski  <[email protected]>

add name test

rm parent idx in test

src/Fields/indices.jl

@@ -17,7 +17,13 @@ end
 
 Base.@propagate_inbounds Base.getindex(field::Field, colidx::ColumnIndex) =
     column(field, colidx)
-
+Base.@propagate_inbounds function Base.getindex(
+    fv::FieldVector{T},
+    colidx::ColumnIndex,
+) where {T}
+    values = map(x -> x[colidx], _values(fv))
+    return FieldVector{T, typeof(values)}(values)
+end
 Base.@propagate_inbounds function column(
     field::SpectralElementField1D,
     colidx::ColumnIndex{1},

test/Fields/field.jl

@@ -1,4 +1,6 @@
 using Test
+using JET
+
 using ClimaComms
 using OrderedCollections
 using StaticArrays, IntervalSets
@@ -267,6 +269,37 @@ end
     @test Y.k.z === 3.0
 end
 
+function call_getcolumn(fv, colidx)
+    fvcol = fv[colidx]
+    nothing
+end
+function call_getproperty(fv)
+    fva = fv.c.a
+    nothing
+end
+@testset "FieldVector getindex" begin
+    cspace = TU.CenterExtrudedFiniteDifferenceSpace(Float32)
+    fspace = Spaces.ExtrudedFiniteDifferenceSpace{Spaces.CellFace}(cspace)
+    c = fill((a = Float32(1), b = Float32(2)), cspace)
+    f = fill((x = Float32(1), y = Float32(2)), fspace)
+    fv = Fields.FieldVector(; c, f)
+    colidx = Fields.ColumnIndex((1, 1), 1) # arbitrary index
+
+    @test all(parent(fv.c.a[colidx]) .== Float32(1))
+    @test all(parent(fv.f.y[colidx]) .== Float32(2))
+    @test propertynames(fv) == propertynames(fv[colidx])
+
+    # JET tests
+    # prerequisite
+    call_getproperty(fv) # compile first
+    @test_opt call_getproperty(fv)
+
+    call_getcolumn(fv, colidx) # compile first
+    @test_opt call_getcolumn(fv, colidx)
+    p = @allocated call_getcolumn(fv, colidx)
+    @test p ≤ 32
+end
+
 @testset "FieldVector array_type" begin
     device = ClimaComms.device()
     context = ClimaComms.SingletonCommsContext(device)