Skip to content

Commit

Permalink
Merge pull request #2050 from CliMA/ck/dl_get_single_fid
Browse files Browse the repository at this point in the history 
Define `getindex_field` and `setindex_field`
  • Loading branch information
@charleskawczynski
charleskawczynski authored Oct 18, 2024
2 parents 90572d1 + 09e48c6 commit 78894fc
Show file tree
Hide file tree
Showing 4 changed files with 226 additions and 0 deletions.
4 changes: 4 additions & 0 deletions .buildkite/pipeline.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,10 @@ steps:
key: unit_data_copyto
command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_copyto.jl"

- label: "Unit: getindex_field"
key: unit_data_getindex_field
command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_getindex_field.jl"

- label: "Unit: mapreduce"
key: unit_data_mapreduce
command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_mapreduce.jl"
Expand Down
51 changes: 51 additions & 0 deletions src/DataLayouts/DataLayouts.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1309,6 +1309,15 @@ type parameters.
@inline to_data_specific(::VIJFHSingleton, I::Tuple) = (I[4], I[1], I[2], 1, I[5])
@inline to_data_specific(::VIFHSingleton, I::Tuple) = (I[4], I[1], 1, I[5])

@inline to_data_specific_field(::DataFSingleton, I::Tuple) = (I[3],)
@inline to_data_specific_field(::VFSingleton, I::Tuple) = (I[4], I[3])
@inline to_data_specific_field(::IFSingleton, I::Tuple) = (I[1], I[3])
@inline to_data_specific_field(::IJFSingleton, I::Tuple) = (I[1], I[2], I[3])
@inline to_data_specific_field(::IJFHSingleton, I::Tuple) = (I[1], I[2], I[3], I[5])
@inline to_data_specific_field(::IFHSingleton, I::Tuple) = (I[1], I[3], I[5])
@inline to_data_specific_field(::VIJFHSingleton, I::Tuple) = (I[4], I[1], I[2], I[3], I[5])
@inline to_data_specific_field(::VIFHSingleton, I::Tuple) = (I[4], I[1], I[3], I[5])

"""
bounds_condition(data::AbstractData, I::Tuple)
Expand Down Expand Up @@ -1482,6 +1491,48 @@ end
)
end

"""
getindex_field(data, ci::CartesianIndex{5})
Returns the value of the data at universal index `ci`,
for the specific field `f` in the `CartesianIndex`.
The universal index order is `CartesianIndex(i, j, f, v, h)`, see
see the notation in [`DataLayouts`](@ref) for more information.
"""
@inline function getindex_field(
data::Union{DataF, IJF, IJFH, IFH, VIJFH, VIFH, VF, IF},
I::CartesianIndex, # universal index
)
@boundscheck bounds_condition(data, I) || throw(BoundsError(data, I))
@inbounds Base.getindex(
parent(data),
CartesianIndex(to_data_specific_field(singleton(data), I.I)),
)
end

"""
setindex_field!(data, val::Real, ci::CartesianIndex{5})
Stores the value `val` of the data at universal index `ci`,
for the specific field `f` in the `CartesianIndex`.
The universal index order is `CartesianIndex(i, j, f, v, h)`, see
see the notation in [`DataLayouts`](@ref) for more information.
"""
@inline function setindex_field!(
data::Union{DataF, IJF, IJFH, IFH, VIJFH, VIFH, VF, IF},
val::Real,
I::CartesianIndex, # universal index
)
@boundscheck bounds_condition(data, I) || throw(BoundsError(data, I))
@inbounds Base.setindex!(
parent(data),
val,
CartesianIndex(to_data_specific_field(singleton(data), I.I)),
)
end

if VERSION ≥ v"1.11.0-beta"
### --------------- Support for multi-dimensional indexing
# TODO: can we remove this? It's not needed for Julia 1.10,
Expand Down
170 changes: 170 additions & 0 deletions test/DataLayouts/unit_getindex_field.jl
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,170 @@
#=
julia --project
using Revise; include(joinpath("test", "DataLayouts", "unit_getindex_field.jl"))
=#
using Test
using ClimaCore.DataLayouts
using ClimaCore.DataLayouts: getindex_field, setindex_field!
using ClimaCore.DataLayouts: to_data_specific_field, singleton
import ClimaCore.Geometry
import ClimaComms
using StaticArrays
ClimaComms.@import_required_backends
import Random
Random.seed!(1234)

universal_axes(data) =
map(size(data)) do i
s =
i == DataLayouts.field_dim(singleton(data)) ?
DataLayouts.ncomponents(data) : i
Base.OneTo(s)
end

universal_field_index(I::CartesianIndex, f) =
CartesianIndex(map(i -> i == 3 ? f : i, I.I))

function test_copyto_float!(data)
Random.seed!(1234)
# Normally we'd use `similar` here, but https://github.com/CliMA/ClimaCore.jl/issues/1803
rand_data = DataLayouts.rebuild(data, similar(parent(data)))
ArrayType = ClimaComms.array_type(ClimaComms.device())
FT = eltype(parent(data))
parent(rand_data) .= ArrayType(rand(FT, DataLayouts.farray_size(data)))
# For a float, getindex and getindex_field return the same thing
for I in CartesianIndices(universal_axes(data))
@test getindex_field(data, I) == getindex(data, I)
end
for I in CartesianIndices(universal_axes(data))
setindex_field!(data, FT(prod(I.I)), I)
end
for I in CartesianIndices(universal_axes(data))
@test getindex_field(data, I) == prod(I.I)
end
end

function test_copyto!(data)
Random.seed!(1234)
# Normally we'd use `similar` here, but https://github.com/CliMA/ClimaCore.jl/issues/1803
rand_data = DataLayouts.rebuild(data, similar(parent(data)))
ArrayType = ClimaComms.array_type(ClimaComms.device())
FT = eltype(parent(data))
parent(rand_data) .= ArrayType(rand(FT, DataLayouts.farray_size(data)))

for I in CartesianIndices(universal_axes(data))
for f in 1:DataLayouts.ncomponents(data)
UFI = universal_field_index(I, f)
DSI = CartesianIndex(to_data_specific_field(singleton(data), UFI.I))
@test getindex_field(data, UFI) == parent(data)[DSI]
end
end

for I in CartesianIndices(universal_axes(data))
for f in 1:DataLayouts.ncomponents(data)
UFI = universal_field_index(I, f)
DSI = CartesianIndex(to_data_specific_field(singleton(data), UFI.I))
val = parent(data)[DSI]
setindex_field!(data, val + 1, UFI)
@test parent(data)[DSI] == val + 1
end
end
end

@testset "copyto! with Nf = 1" begin
device = ClimaComms.device()
ArrayType = ClimaComms.array_type(device)
FT = Float64
S = FT
Nv = 4
Ni = Nij = 3
Nh = 5
Nk = 6
data = DataF{S}(ArrayType{FT}, zeros)
test_copyto_float!(data)
data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh)
test_copyto_float!(data)
data = IFH{S}(ArrayType{FT}, zeros; Ni, Nh)
test_copyto_float!(data)
data = IJF{S}(ArrayType{FT}, zeros; Nij)
test_copyto_float!(data)
data = IF{S}(ArrayType{FT}, zeros; Ni)
test_copyto_float!(data)
data = VF{S}(ArrayType{FT}, zeros; Nv)
test_copyto_float!(data)
data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh)
test_copyto_float!(data)
data = VIFH{S}(ArrayType{FT}, zeros; Nv, Ni, Nh)
test_copyto_float!(data)
# data = DataLayouts.IJKFVH{S}(ArrayType{FT}, zeros; Nij,Nk,Nv,Nh); test_copyto_float!(data) # TODO: test
# data = DataLayouts.IH1JH2{S}(ArrayType{FT}, zeros; Nij); test_copyto_float!(data) # TODO: test
end

@testset "copyto! with Nf > 1" begin
device = ClimaComms.device()
ArrayType = ClimaComms.array_type(device)
FT = Float64
S = Tuple{FT, FT}
Nv = 4
Ni = Nij = 3
Nh = 5
Nk = 6
data = DataF{S}(ArrayType{FT}, zeros)
test_copyto!(data)
data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh)
test_copyto!(data)
data = IFH{S}(ArrayType{FT}, zeros; Ni, Nh)
test_copyto!(data)
data = IJF{S}(ArrayType{FT}, zeros; Nij)
test_copyto!(data)
data = IF{S}(ArrayType{FT}, zeros; Ni)
test_copyto!(data)
data = VF{S}(ArrayType{FT}, zeros; Nv)
test_copyto!(data)
data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh)
test_copyto!(data)
data = VIFH{S}(ArrayType{FT}, zeros; Nv, Ni, Nh)
test_copyto!(data)
# TODO: test this
# data = DataLayouts.IJKFVH{S}(ArrayType{FT}, zeros; Nij,Nk,Nv,Nh); test_copyto!(data) # TODO: test
# data = DataLayouts.IH1JH2{S}(ArrayType{FT}, zeros; Nij); test_copyto!(data) # TODO: test
end

@testset "copyto! views with Nf > 1" begin
device = ClimaComms.device()
ArrayType = ClimaComms.array_type(device)
data_view(data) = DataLayouts.rebuild(
data,
SubArray(
parent(data),
ntuple(
i -> Base.Slice(Base.OneTo(DataLayouts.farray_size(data, i))),
ndims(data),
),
),
)
FT = Float64
S = Tuple{FT, FT}
Nv = 4
Ni = Nij = 3
Nh = 5
Nk = 6
# Rather than using level/slab/column, let's just make views/SubArrays
# directly so that we can easily test all cases:
data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh)
test_copyto!(data_view(data))
data = IFH{S}(ArrayType{FT}, zeros; Ni, Nh)
test_copyto!(data_view(data))
data = IJF{S}(ArrayType{FT}, zeros; Nij)
test_copyto!(data_view(data))
data = IF{S}(ArrayType{FT}, zeros; Ni)
test_copyto!(data_view(data))
data = VF{S}(ArrayType{FT}, zeros; Nv)
test_copyto!(data_view(data))
data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh)
test_copyto!(data_view(data))
data = VIFH{S}(ArrayType{FT}, zeros; Nv, Ni, Nh)
test_copyto!(data_view(data))
# TODO: test this
# data = DataLayouts.IJKFVH{S}(ArrayType{FT}, zeros; Nij,Nk,Nv,Nh); test_copyto!(data) # TODO: test
# data = DataLayouts.IH1JH2{S}(ArrayType{FT}, zeros; Nij); test_copyto!(data) # TODO: test
end
1 change: 1 addition & 0 deletions test/runtests.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ UnitTest("DataLayouts fill" ,"DataLayouts/unit_fill.jl"),
UnitTest("DataLayouts ndims" ,"DataLayouts/unit_ndims.jl"),
UnitTest("DataLayouts array<->data" ,"DataLayouts/unit_data2array.jl"),
UnitTest("DataLayouts get_struct" ,"DataLayouts/unit_struct.jl"),
UnitTest("DataLayouts get/set_index_field" ,"DataLayouts/unit_getindex_field.jl"),
UnitTest("Recursive" ,"RecursiveApply/unit_recursive_apply.jl"),
UnitTest("PlusHalf" ,"Utilities/unit_plushalf.jl"),
UnitTest("DataLayouts 0D" ,"DataLayouts/data0d.jl"),
Expand Down

0 comments on commit 78894fc

Please sign in to comment.