Mercurial > repos > public > sbplib_julia
changeset 1349:42ecd4b3e215 refactor/grids
Add support for cartesian indices on TensorGrid
author | Jonatan Werpers <jonatan@werpers.com> |
---|---|
date | Fri, 19 May 2023 14:53:49 +0200 |
parents | 279179677040 |
children | 1e2adeb5593d |
files | src/Grids/tensor_grid.jl test/Grids/tensor_grid_test.jl |
diffstat | 2 files changed, 18 insertions(+), 0 deletions(-) [+] |
line wrap: on
line diff
--- a/src/Grids/tensor_grid.jl Fri May 19 14:31:16 2023 +0200 +++ b/src/Grids/tensor_grid.jl Fri May 19 14:53:49 2023 +0200 @@ -26,6 +26,8 @@ return vcat(ps...) end +Base.getindex(g::TensorGrid, I::CartesianIndex) = g[Tuple(I)...] + function Base.eachindex(g::TensorGrid) szs = LazyTensors.concatenate_tuples(size.(g.grids)...) return CartesianIndices(szs)
--- a/test/Grids/tensor_grid_test.jl Fri May 19 14:31:16 2023 +0200 +++ b/test/Grids/tensor_grid_test.jl Fri May 19 14:53:49 2023 +0200 @@ -36,6 +36,22 @@ @test TensorGrid(g₁, g₄, g₂)[3,2] == [0.2, 1., 2., 2.2] end + @testset "cartesian indexing" begin + cases = [ + (TensorGrid(g₁, g₂), (1,1) ), + (TensorGrid(g₁, g₂), (3,5) ), + (TensorGrid(g₁, g₂), (10,6) ), + (TensorGrid(g₁, g₃), (1,1) ), + (TensorGrid(g₁, g₂, g₃), (3,4,5)), + (TensorGrid(g₁, g₄), (3) ), + (TensorGrid(g₁, g₄, g₂), (3,2) ), + ] + + @testset "i = $is" for (g, is) ∈ cases + @test g[CartesianIndex(is...)] == g[is...] + end + end + @testset "eachindex" begin @test eachindex(TensorGrid(g₁, g₂)) == CartesianIndices((11,6)) @test eachindex(TensorGrid(g₁, g₃)) == CartesianIndices((11,10))