changeset 1428:a936b414283a feature/grids/curvilinear

Merge bugfix/grids/complete_interface_impl
author Jonatan Werpers <jonatan@werpers.com>
date Thu, 24 Aug 2023 08:50:10 +0200
parents 26e168924cf1 (current diff) 18e21601da2d (diff)
children e87f3465b770
files
diffstat 5 files changed, 95 insertions(+), 0 deletions(-) [+]
line wrap: on
line diff
--- a/TODO.md	Wed Aug 23 15:51:32 2023 +0200
+++ b/TODO.md	Thu Aug 24 08:50:10 2023 +0200
@@ -35,6 +35,7 @@
  - [ ] Multiblock grids
  - [ ] Periodic grids
  - [ ] Grids with modified boundary closures
+ - [ ] Support indexing with `:`.
 
 
 ### Benchmarks
--- a/src/Grids/equidistant_grid.jl	Wed Aug 23 15:51:32 2023 +0200
+++ b/src/Grids/equidistant_grid.jl	Thu Aug 24 08:50:10 2023 +0200
@@ -20,6 +20,8 @@
 Base.firstindex(g::EquidistantGrid) = firstindex(g.points)
 Base.lastindex(g::EquidistantGrid) = lastindex(g.points)
 
+Base.axes(g::EquidistantGrid, d) = axes(g.points, d)
+
 # Iteration interface
 Base.iterate(g::EquidistantGrid) = iterate(g.points)
 Base.iterate(g::EquidistantGrid, state) = iterate(g.points, state)
--- a/src/Grids/tensor_grid.jl	Wed Aug 23 15:51:32 2023 +0200
+++ b/src/Grids/tensor_grid.jl	Thu Aug 24 08:50:10 2023 +0200
@@ -31,6 +31,11 @@
     return CartesianIndices(szs)
 end
 
+function Base.axes(g::TensorGrid, d)
+    i, ld = grid_and_local_dim_index(ndims.(g.grids), d)
+    return axes(g.grids[i], ld)
+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
@@ -95,3 +100,28 @@
 function combine_coordinates(coords...)
     return mapreduce(SVector, vcat, coords)
 end
+
+"""
+   grid_and_local_dim_index(nds, d)
+
+Given a tuple of number of dimensions `nds`, and a global dimension index `d`,
+calculate which grid index, and local dimension, `d` corresponds to.
+
+`nds` would come from broadcasting `ndims` on the grids tuple of a
+`TensorGrid`. If you are interested in a dimension `d` of a tensor grid `g`
+```julia
+gi, ldi = grid_and_local_dim_index(ndims.(g.grids), d)
+```
+tells you which grid it belongs to (`gi`) and wich index it is at within that
+grid (`ldi`).
+"""
+function grid_and_local_dim_index(nds, d)
+    I = findfirst(>=(d), cumsum(nds))
+
+    if I == 1
+        return (1, d)
+    else
+        return (I, d-cumsum(nds)[I-1])
+    end
+    # TBD: Is there a cleaner way to compute this?
+end
--- a/test/Grids/equidistant_grid_test.jl	Wed Aug 23 15:51:32 2023 +0200
+++ b/test/Grids/equidistant_grid_test.jl	Thu Aug 24 08:50:10 2023 +0200
@@ -19,6 +19,9 @@
         @test g[end] == 10.0
 
         @test all(eachindex(g) .== 1:101)
+
+        @test firstindex(g) == 1
+        @test lastindex(g) == 101
     end
 
     @testset "Iterator interface" begin
@@ -35,6 +38,10 @@
 
     @testset "Base" begin
         @test ndims(EquidistantGrid(0:10)) == 1
+
+        g = EquidistantGrid(0:0.1:10)
+        @test axes(g,1) == 1:101
+        @test axes(g) == (1:101,)
     end
 
     @testset "spacing" begin
--- a/test/Grids/tensor_grid_test.jl	Wed Aug 23 15:51:32 2023 +0200
+++ b/test/Grids/tensor_grid_test.jl	Thu Aug 24 08:50:10 2023 +0200
@@ -34,6 +34,11 @@
 
             @test TensorGrid(g₁, g₄, g₂)[3,2] isa SVector{4,Float64}
             @test TensorGrid(g₁, g₄, g₂)[3,2] == [0.2, 1., 2., 2.2]
+
+            g = TensorGrid(g₁, g₂)
+            @test g[begin, begin] == g[1,1]
+            @test g[begin, end] == g[1,6]
+            @test g[end, end] == g[11,6]
         end
 
         @testset "cartesian indexing" begin
@@ -59,6 +64,18 @@
             @test eachindex(TensorGrid(g₁, g₄)) == CartesianIndices((11,))
             @test eachindex(TensorGrid(g₁, g₄, g₂)) == CartesianIndices((11,6))
         end
+
+        @testset "firstindex" begin
+            @test firstindex(TensorGrid(g₁, g₂, g₃), 1) == 1
+            @test firstindex(TensorGrid(g₁, g₂, g₃), 2) == 1
+            @test firstindex(TensorGrid(g₁, g₂, g₃), 3) == 1
+        end
+
+        @testset "lastindex" begin
+            @test lastindex(TensorGrid(g₁, g₂, g₃), 1) == 11
+            @test lastindex(TensorGrid(g₁, g₂, g₃), 2) == 6
+            @test lastindex(TensorGrid(g₁, g₂, g₃), 3) == 10
+        end
     end
 
     @testset "Iterator interface" begin
@@ -93,6 +110,16 @@
         @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 "Base" begin
+        g₁ = EquidistantGrid(range(0,1,length=11))
+        g₂ = EquidistantGrid(range(2,3,length=6))
+        g = TensorGrid(g₁, g₂)
+
+        @test axes(g, 1) == 1:11
+        @test axes(g, 2) == 1:6
+        @test axes(g) == (1:11,1:6)
+    end
+
     @testset "refine" begin
         g1(n) = EquidistantGrid(range(0,1,length=n))
         g2(n) = EquidistantGrid(range(2,3,length=n))
@@ -144,3 +171,31 @@
     @test Grids.combine_coordinates(1,@SVector[2.,3]) isa SVector{3, Float64}
     @test Grids.combine_coordinates(1,@SVector[2.,3]) == [1,2,3]
 end
+
+@testset "grid_and_local_dim_index" begin
+    cases = [
+        ((1,), 1) => (1,1),
+
+        ((1,1), 1) => (1,1),
+        ((1,1), 2) => (2,1),
+
+        ((1,2), 1) => (1,1),
+        ((1,2), 2) => (2,1),
+        ((1,2), 3) => (2,2),
+
+        ((2,1), 1) => (1,1),
+        ((2,1), 2) => (1,2),
+        ((2,1), 3) => (2,1),
+
+        ((2,1,3), 1) => (1,1),
+        ((2,1,3), 2) => (1,2),
+        ((2,1,3), 3) => (2,1),
+        ((2,1,3), 4) => (3,1),
+        ((2,1,3), 5) => (3,2),
+        ((2,1,3), 6) => (3,3),
+    ]
+
+    @testset "grid_and_local_dim_index$args" for (args, expected) ∈ cases
+        @test Grids.grid_and_local_dim_index(args...) == expected
+    end
+end