Mercurial > repos > public > sbplib_julia

changeset 1266:a4ddae8b5d49 refactor/grids

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Add tests for TensorGrid and make them pass
author Jonatan Werpers <jonatan@werpers.com>
date Fri, 24 Feb 2023 21:42:28 +0100
parents 9c9ea2900250
children 30729cba1095
files src/Grids/Grids.jl src/Grids/grid.jl src/Grids/tensor_grid.jl test/Grids/tensor_grid_test.jl
diffstat 4 files changed, 178 insertions(+), 30 deletions(-) [+]
line wrap: on
line diff
--- a/src/Grids/Grids.jl	Fri Feb 24 20:47:56 2023 +0100
+++ b/src/Grids/Grids.jl	Fri Feb 24 21:42:28 2023 +0100
@@ -2,6 +2,7 @@
 
 using Sbplib.RegionIndices
 using Sbplib.LazyTensors
+using StaticArrays
 
 # Grid
 export Grid
@@ -11,14 +12,14 @@
 export TensorGrid
 export ZeroDimGrid
 
+export TensorGridBoundary
+
 export eval_on
 export getcomponent
 
 # BoundaryIdentifier
 export BoundaryIdentifier
-export CartesianBoundary
-export dim
-export region
+
 
 # EquidistantGrid
 export EquidistantGrid
--- a/src/Grids/grid.jl	Fri Feb 24 20:47:56 2023 +0100
+++ b/src/Grids/grid.jl	Fri Feb 24 21:42:28 2023 +0100
@@ -33,6 +33,7 @@
 # TODO
 """
 function boundary_grid end
+# TBD Can we implement a version here that accepts multiple ids and grouped boundaries? Maybe we need multiblock stuff?
 
 
 # TODO: Make sure that all grids implement all of the above.
--- a/src/Grids/tensor_grid.jl	Fri Feb 24 20:47:56 2023 +0100
+++ b/src/Grids/tensor_grid.jl	Fri Feb 24 21:42:28 2023 +0100
@@ -1,24 +1,21 @@
+"""
+    TensorGrid{T,D} <: Grid{T,D}
+
+* Only supports HasShape grids at the moment
+"""
 struct TensorGrid{T,D,GT<:NTuple{N,Grid} where N} <: Grid{T,D}
     grids::GT
 
     function TensorGrid(gs...)
-        T = eltype(gs[1]) # All gs should have the same T
-        D = sum(ndims,gs)
+        # T = combined_coordinate_vector_type(eltype.(gs)...)
+        T = mapreduce(eltype, combined_coordinate_vector_type, gs)
+        D = sum(ndims, gs)
 
         return new{T,D,typeof(gs)}(gs)
     end
 end
 
 # Indexing interface
-# TODO
-# Iteration interface
-# TODO
-
-
-function Base.size(g::TensorGrid)
-    return LazyTensors.concatenate_tuples(size.(g.grids)...)
-end
-
 function Base.getindex(g::TensorGrid, I...)
     szs = ndims.(g.grids)
 
@@ -28,18 +25,32 @@
     return vcat(ps...)
 end
 
-IndexStyle(::TensorGrid) = IndexCartesian()
-
 function Base.eachindex(g::TensorGrid)
     szs = LazyTensors.concatenate_tuples(size.(g.grids)...)
     return CartesianIndices(szs)
 end
 
+# Iteration interface
+Base.iterate(g::TensorGrid) = iterate(Iterators.product(g.grids...)) |> _iterate_combine_coords
+Base.iterate(g::TensorGrid, state) = iterate(Iterators.product(g.grids...), state) |> _iterate_combine_coords
+_iterate_combine_coords(::Nothing) = nothing
+_iterate_combine_coords((next,state)) = combine_coordinates(next...), state
 
-struct TensorBoundary{N, BID<:BoundaryIdentifier} <: BoundaryIdentifier end
-grid_id(::TensorBoundary{N, BID}) where {N, BID} = N
-boundary_id(::TensorBoundary{N, BID}) where {N, BID} = BID()
+Base.IteratorSize(::Type{<:TensorGrid{<:Any, D}}) where D = Base.HasShape{D}()
+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)...)
+
 
+refine(g::TensorGrid, r::Int) = mapreduce(g->refine(g,r), TensorGrid, g.grids)
+coarsen(g::TensorGrid, r::Int) = mapreduce(g->coarsen(g,r), TensorGrid, g.grids)
+
+"""
+# TODO:
+"""
+struct TensorGridBoundary{N, BID} <: BoundaryIdentifier end
+grid_id(::TensorGridBoundary{N, BID}) where {N, BID} = N
+boundary_id(::TensorGridBoundary{N, BID}) where {N, BID} = BID()
 
 """
     boundary_identifiers(::TensorGrid)
@@ -47,21 +58,44 @@
 Returns a tuple containing the boundary identifiers for the grid.
 """
 function boundary_identifiers(g::TensorGrid)
-    n = length(g.grids)
     per_grid = map(eachindex(g.grids)) do i
-        return map(bid -> TensorBoundary{i, bid}(), boundary_identifiers(g.grids[i]))
+        return map(bid -> TensorGridBoundary{i, typeof(bid)}(), boundary_identifiers(g.grids[i]))
     end
     return LazyTensors.concatenate_tuples(per_grid...)
 end
 
 
 """
-    boundary_grid(grid::TensorGrid, id::TensorBoundary)
+    boundary_grid(grid::TensorGrid, id::TensorGridBoundary)
 
 The grid for the boundary specified by `id`.
 """
-function boundary_grid(g::TensorGrid, bid::TensorBoundary)
+function boundary_grid(g::TensorGrid, bid::TensorGridBoundary)
     local_boundary_grid = boundary_grid(g.grids[grid_id(bid)], boundary_id(bid))
     new_grids = Base.setindex(g.grids, local_boundary_grid, grid_id(bid))
     return TensorGrid(new_grids...)
 end
+
+
+function combined_coordinate_vector_type(coordinate_types...)
+    coord_length(::Type{<:Number}) = 1
+    coord_length(T::Type{<:SVector}) = length(T)
+
+    coord_type(T::Type{<:Number}) = T
+    coord_type(T::Type{<:SVector}) = eltype(T)
+
+
+    combined_coord_length = mapreduce(coord_length, +, coordinate_types)
+    combined_coord_type = mapreduce(coord_type, promote_type, coordinate_types)
+
+    if combined_coord_length == 1
+        return combined_coord_type
+    else
+        return SVector{combined_coord_length, combined_coord_type}
+    end
+end
+
+function combine_coordinates(coords...)
+    return mapreduce(SVector, vcat, coords)
+    # return SVector(coords...)
+end
--- a/test/Grids/tensor_grid_test.jl	Fri Feb 24 20:47:56 2023 +0100
+++ b/test/Grids/tensor_grid_test.jl	Fri Feb 24 21:42:28 2023 +0100
@@ -1,14 +1,126 @@
+using Test
 using Sbplib.Grids
-using Test
+using StaticArrays
+using Sbplib.RegionIndices
 
 @testset "TensorGrid" begin
-    @test_broken false
+    g₁ = EquidistantGrid(range(0,1,length=11))
+    g₂ = EquidistantGrid(range(2,3,length=6))
+    g₃ = EquidistantGrid(1:10)
+    g₄ = ZeroDimGrid(@SVector[1,2])
+
+    @test TensorGrid(g₁, g₂) isa TensorGrid
+    @test TensorGrid(g₁, g₂) isa Grid{SVector{2,Float64}, 2}
+    @test TensorGrid(g₃, g₃) isa Grid{SVector{2,Int}, 2}
+    @test TensorGrid(g₁, g₂, g₃) isa Grid{SVector{3,Float64}, 3}
+    @test TensorGrid(g₁, g₄) isa Grid{SVector{3,Float64}, 1}
+    @test TensorGrid(g₁, g₄, g₂) isa Grid{SVector{4,Float64}, 2}
+
+    @testset "Indexing Interface" begin
+        @test TensorGrid(g₁, g₂)[1,1] isa SVector{2,Float64}
+        @test TensorGrid(g₁, g₂)[1,1] == [0.0,2.0]
+        @test TensorGrid(g₁, g₂)[3,5] == [0.2,2.8]
+        @test TensorGrid(g₁, g₂)[10,6] == [0.9,3.0]
+
+        @test TensorGrid(g₁, g₃)[1,1] isa SVector{2,Float64}
+        @test TensorGrid(g₁, g₃)[1,1] == [0.0,1.0]
+
+        @test TensorGrid(g₁, g₂, g₃)[3,4,5] isa SVector{3,Float64}
+        @test TensorGrid(g₁, g₂, g₃)[3,4,5] == [0.2, 2.6, 5.0]
 
+        @test TensorGrid(g₁, g₄)[3] isa SVector{3,Float64}
+        @test TensorGrid(g₁, g₄)[3] == [0.2, 1., 2.]
+
+        @test TensorGrid(g₁, g₄, g₂)[3,2] isa SVector{4,Float64}
+        @test TensorGrid(g₁, g₄, g₂)[3,2] == [0.2, 1., 2., 2.2]
+
+        @test eachindex(TensorGrid(g₁, g₂)) == CartesianIndices((11,6))
+        @test eachindex(TensorGrid(g₁, g₃)) == CartesianIndices((11,10))
+        @test eachindex(TensorGrid(g₁, g₂, g₃)) == CartesianIndices((11,6,10))
+        @test eachindex(TensorGrid(g₁, g₄)) == CartesianIndices((11,))
+        @test eachindex(TensorGrid(g₁, g₄, g₂)) == CartesianIndices((11,6))
+    end
+
+    @testset "Iterator interface" begin
+        @test eltype(TensorGrid(g₁, g₂)) == SVector{2,Float64}
+        @test eltype(TensorGrid(g₁, g₃)) == SVector{2,Float64}
+        @test eltype(TensorGrid(g₁, g₂, g₃)) == SVector{3,Float64}
+        @test eltype(TensorGrid(g₁, g₄)) == SVector{3,Float64}
+        @test eltype(TensorGrid(g₁, g₄, g₂)) == SVector{4,Float64}
+
+        @test size(TensorGrid(g₁, g₂)) == (11,6)
+        @test size(TensorGrid(g₁, g₃)) == (11,10)
+        @test size(TensorGrid(g₁, g₂, g₃)) == (11,6,10)
+        @test size(TensorGrid(g₁, g₄)) == (11,)
+        @test size(TensorGrid(g₁, g₄, g₂)) == (11,6)
 
-    @testset "restrict" begin
-        @test_broken restrict(g, 1:2) == nothing
-        @test_broken restrict(g, 2:3) == nothing
-        @test_broken restrict(g, [1,3]) == nothing
-        @test_broken restrict(g, [2,1]) == nothing
+        @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}()
+        @test Base.IteratorSize(TensorGrid(g₁, g₄)) == Base.HasShape{1}()
+        @test Base.IteratorSize(TensorGrid(g₁, g₄, g₂)) == Base.HasShape{2}()
+
+        @test iterate(TensorGrid(g₁, g₂))[1] isa SVector{2,Float64}
+        @test iterate(TensorGrid(g₁, g₃))[1] isa SVector{2,Float64}
+        @test iterate(TensorGrid(g₁, g₂, g₃))[1] isa SVector{3,Float64}
+        @test iterate(TensorGrid(g₁, g₄))[1] isa SVector{3,Float64}
+        @test iterate(TensorGrid(g₁, g₄, g₂))[1] isa SVector{4,Float64}
+
+        @test collect(TensorGrid(g₁, g₂)) == [@SVector[x,y] for x ∈ range(0,1,length=11), y ∈ range(2,3,length=6)]
+        @test collect(TensorGrid(g₁, g₃)) == [@SVector[x,y] for x ∈ range(0,1,length=11), y ∈ 1:10]
+        @test collect(TensorGrid(g₁, g₂, g₃)) == [@SVector[x,y,z] for x ∈ range(0,1,length=11), y ∈ range(2,3,length=6), z ∈ 1:10]
+        @test collect(TensorGrid(g₁, g₄)) == [@SVector[x,1,2] for x ∈ range(0,1,length=11)]
+        @test collect(TensorGrid(g₁, g₄, g₂)) == [@SVector[x,1,2,y] for x ∈ range(0,1,length=11), y ∈ range(2,3,length=6)]
+    end
+
+    @testset "refine" begin
+        g1(n) = EquidistantGrid(range(0,1,length=n))
+        g2(n) = EquidistantGrid(range(2,3,length=n))
+
+        @test refine(TensorGrid(g1(11), g2(6)),1) == TensorGrid(g1(11), g2(6))
+        @test refine(TensorGrid(g1(11), g2(6)),2) == TensorGrid(g1(21), g2(11))
+        @test refine(TensorGrid(g1(11), g2(6)),3) == TensorGrid(g1(31), g2(16))
+        @test_broken refine(TensorGrid(g1(11), g₄), 1) == TensorGrid(g1(11), g₄)
+        @test_broken refine(TensorGrid(g1(11), g₄), 2) == TensorGrid(g1(21), g₄)
+    end
+
+    @testset "coarsen" begin
+        g1(n) = EquidistantGrid(range(0,1,length=n))
+        g2(n) = EquidistantGrid(range(2,3,length=n))
+
+        @test coarsen(TensorGrid(g1(11), g2(6)),1) == TensorGrid(g1(11), g2(6))
+        @test coarsen(TensorGrid(g1(21), g2(11)),2) == TensorGrid(g1(11), g2(6))
+        @test coarsen(TensorGrid(g1(31), g2(16)),3) == TensorGrid(g1(11), g2(6))
+        @test_broken coarsen(TensorGrid(g1(11), g₄), 1) == TensorGrid(g1(11), g₄)
+        @test_broken coarsen(TensorGrid(g1(21), g₄), 2) == TensorGrid(g1(11), g₄)
+    end
+
+    @testset "boundary_identifiers" begin
+        @test boundary_identifiers(TensorGrid(g₁, g₂)) == map((n,id)->TensorGridBoundary{n,id}(), (1,1,2,2), (Lower,Upper,Lower,Upper))
+        @test_broken boundary_identifiers(TensorGrid(g₁, g₄)) == (TensorGridBoundary{1,Lower}(),TensorGridBoundary{1,Upper}())
+    end
+
+    @testset "boundary_grid" begin
+        @test boundary_grid(TensorGrid(g₁, g₂), TensorGridBoundary{1, Upper}()) == TensorGrid(ZeroDimGrid(g₁[end]), g₂)
+        @test boundary_grid(TensorGrid(g₁, g₄), TensorGridBoundary{1, Upper}()) == TensorGrid(ZeroDimGrid(g₁[end]), g₄)
     end
 end
+
+@testset "combined_coordinate_vector_type" begin
+    @test Grids.combined_coordinate_vector_type(Float64) == Float64
+    @test Grids.combined_coordinate_vector_type(Float64, Int) == SVector{2,Float64}
+    @test Grids.combined_coordinate_vector_type(Float32, Int16, Int32) == SVector{3,Float32}
+
+    @test Grids.combined_coordinate_vector_type(SVector{2,Float64}) == SVector{2,Float64}
+    @test Grids.combined_coordinate_vector_type(SVector{2,Float64}, SVector{1,Float64}) == SVector{3,Float64}
+    @test Grids.combined_coordinate_vector_type(SVector{2,Float64}, SVector{1,Int}, SVector{3, Float32}) == SVector{6,Float64}
+end
+
+@testset "combine_coordinates" begin
+    @test Grids.combine_coordinates(1,2,3) isa SVector{3, Int}
+    @test Grids.combine_coordinates(1,2,3) == [1,2,3]
+    @test Grids.combine_coordinates(1,2.,3) isa SVector{3, Float64}
+    @test Grids.combine_coordinates(1,2.,3) == [1,2,3]
+    @test Grids.combine_coordinates(1,@SVector[2.,3]) isa SVector{3, Float64}
+    @test Grids.combine_coordinates(1,@SVector[2.,3]) == [1,2,3]
+end