changeset 1437:e3a80ef08d09

Implement size(g,d) for EquidistantGrid and TensorGrid
author Jonatan Werpers <jonatan@werpers.com>
date Wed, 22 Nov 2023 17:34:35 +0100
parents 69c9e6eae686
children 66ad232e45a1 28bd0ea7ee92
files src/Grids/equidistant_grid.jl src/Grids/tensor_grid.jl test/Grids/equidistant_grid_test.jl test/Grids/tensor_grid_test.jl
diffstat 4 files changed, 15 insertions(+), 0 deletions(-) [+]
line wrap: on
line diff
diff -r 69c9e6eae686 -r e3a80ef08d09 src/Grids/equidistant_grid.jl
--- a/src/Grids/equidistant_grid.jl	Fri Aug 25 08:49:07 2023 +0200
+++ b/src/Grids/equidistant_grid.jl	Wed Nov 22 17:34:35 2023 +0100
@@ -29,6 +29,7 @@
 Base.IteratorSize(::Type{<:EquidistantGrid}) = Base.HasShape{1}()
 Base.length(g::EquidistantGrid) = length(g.points)
 Base.size(g::EquidistantGrid) = size(g.points)
+Base.size(g::EquidistantGrid, d) = size(g.points)[d]
 
 
 """
diff -r 69c9e6eae686 -r e3a80ef08d09 src/Grids/tensor_grid.jl
--- a/src/Grids/tensor_grid.jl	Fri Aug 25 08:49:07 2023 +0200
+++ b/src/Grids/tensor_grid.jl	Wed Nov 22 17:34:35 2023 +0100
@@ -46,6 +46,7 @@
 Base.eltype(::Type{<:TensorGrid{T}}) where T = T
 Base.length(g::TensorGrid) = sum(length, g.grids)
 Base.size(g::TensorGrid) = LazyTensors.concatenate_tuples(size.(g.grids)...)
+Base.size(g::TensorGrid, d) = size(g)[d]
 
 
 refine(g::TensorGrid, r::Int) = mapreduce(g->refine(g,r), TensorGrid, g.grids)
diff -r 69c9e6eae686 -r e3a80ef08d09 test/Grids/equidistant_grid_test.jl
--- a/test/Grids/equidistant_grid_test.jl	Fri Aug 25 08:49:07 2023 +0200
+++ b/test/Grids/equidistant_grid_test.jl	Wed Nov 22 17:34:35 2023 +0100
@@ -31,6 +31,8 @@
         @test size(EquidistantGrid(0:10)) == (11,)
         @test size(EquidistantGrid(0:0.1:10)) == (101,)
 
+        @test size(EquidistantGrid(0:0.1:10),1) == 101
+
         @test collect(EquidistantGrid(0:0.1:0.5)) == [0.0, 0.1, 0.2, 0.3, 0.4, 0.5]
 
         @test Base.IteratorSize(EquidistantGrid{Float64, StepRange{Float64}}) == Base.HasShape{1}()
@@ -107,8 +109,13 @@
     @testset "Base" begin
         @test eltype(equidistant_grid(4,0.0,1.0)) == Float64
         @test eltype(equidistant_grid((4,3),(0,0),(1,3))) <: AbstractVector{Float64}
+
         @test size(equidistant_grid(4,0.0,1.0)) == (4,)
         @test size(equidistant_grid((5,3), (0.0,0.0), (2.0,1.0))) == (5,3)
+
+        @test size(equidistant_grid((5,3), (0.0,0.0), (2.0,1.0)),1) == 5
+        @test size(equidistant_grid((5,3), (0.0,0.0), (2.0,1.0)),2) == 3
+
         @test ndims(equidistant_grid(4,0.0,1.0)) == 1
         @test ndims(equidistant_grid((5,3), (0.0,0.0), (2.0,1.0))) == 2
     end
diff -r 69c9e6eae686 -r e3a80ef08d09 test/Grids/tensor_grid_test.jl
--- a/test/Grids/tensor_grid_test.jl	Fri Aug 25 08:49:07 2023 +0200
+++ b/test/Grids/tensor_grid_test.jl	Wed Nov 22 17:34:35 2023 +0100
@@ -91,6 +91,12 @@
         @test size(TensorGrid(g₁, g₄)) == (11,)
         @test size(TensorGrid(g₁, g₄, g₂)) == (11,6)
 
+        @test size(TensorGrid(g₁, g₂, g₃),1) == 11
+        @test size(TensorGrid(g₁, g₂, g₃),2) == 6
+        @test size(TensorGrid(g₁, g₂, g₃),3) == 10
+        @test size(TensorGrid(g₁, g₄, g₂),1) == 11
+        @test size(TensorGrid(g₁, g₄, g₂),2) == 6
+
         @test Base.IteratorSize(TensorGrid(g₁, g₂)) == Base.HasShape{2}()
         @test Base.IteratorSize(TensorGrid(g₁, g₃)) == Base.HasShape{2}()
         @test Base.IteratorSize(TensorGrid(g₁, g₂, g₃)) == Base.HasShape{3}()