changeset 1388:c0208286234e bugfix/grids/complete_interface_impl

Add `grid_and_local_dim_index`
author Jonatan Werpers <jonatan@werpers.com>
date Tue, 15 Aug 2023 22:11:20 +0200
parents 3d6425c36d32
children d2219cc8316b
files src/Grids/tensor_grid.jl test/Grids/tensor_grid_test.jl
diffstat 2 files changed, 45 insertions(+), 0 deletions(-) [+]
line wrap: on
line diff
--- a/src/Grids/tensor_grid.jl	Tue Aug 15 21:20:29 2023 +0200
+++ b/src/Grids/tensor_grid.jl	Tue Aug 15 22:11:20 2023 +0200
@@ -97,3 +97,20 @@
 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 `d`,
+calculate which grid index and local dimension `d` corresponds to.
+"""
+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/tensor_grid_test.jl	Tue Aug 15 21:20:29 2023 +0200
+++ b/test/Grids/tensor_grid_test.jl	Tue Aug 15 22:11:20 2023 +0200
@@ -156,3 +156,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