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))