Improvements and testing for sub2ind

examples/demo_sub2ind.jl

@@ -0,0 +1,15 @@
+using Comodo
+
+A = rand(30)
+B = rand(5,6)
+C = rand(3,5,2)
+
+ind = [1,2,3,4,8,12,30]
+
+IJK_A = [[1],[2],[3],[4],[8],[12],[30]]
+IJK_B = [[1, 1], [2, 1], [3, 1], [4, 1], [3, 2], [2, 3], [5, 6]]
+IJK_C = [[1, 1, 1], [2, 1, 1], [3, 1, 1], [1, 2, 1], [2, 3, 1], [3, 4, 1], [3, 5, 2]]
+
+all(sub2ind(size(A),IJK_A)==ind)
+all(sub2ind(size(B),IJK_B)==ind)
+all(sub2ind(size(C),IJK_C)==ind)

src/functions.jl

@@ -598,7 +598,7 @@ function ind2sub(siz::Union{Tuple{Vararg{Int64, N}}, Vector{Int64}},ind::Union{T
 end
 
 # ind2sub helper function to parse just a single linear index and produce a single subscript index set 
-function ind2sub_(ind::Integer,numDim::Int64,k::Union{Int64,Vector{Int64},Tuple{Vararg{Int64, N}}}) where N    
+function ind2sub_(ind::Integer,numDim::Int64,k::Union{Int64,Vector{Int64},Tuple{Vararg{Int64, N}}}) where N
     a = Vector{Int64}(undef,numDim) # Initialise a
     for q ∈ numDim:-1:1   # For all dimensions     
         if isone(q) # First 1st dimension
@@ -620,32 +620,41 @@ end
 Converts the subscript indices in `A`, for a matrix/array with size `siz`, to 
 the equivalent linear indices.  
 """
-function sub2ind(siz,A)
-
+function sub2ind(siz::Union{Tuple{Vararg{Int64, N}}, Vector{Int64}},A::Vector{Vector{Int64}})  where N    
     numDim = length(siz)
-    k = cumprod([siz[i] for i ∈ eachindex(siz)],dims=1)
-
-    ind = Vector{Int64}(undef,length(A))
-    for i ∈ eachindex(A)        
-        ind[i] = round(Int64,A[i][1]);
-        for q=2:numDim
-            ind[i] += (A[i][q].-1).*k[q-1]
-        end        
+    if numDim==1
+        ind = [A[q][1] for q ∈ eachindex(A)]
+    else
+        k = cumprod([siz[i] for i ∈ eachindex(siz)],dims=1)
+        ind = Vector{Int64}(undef,length(A))
+        for i ∈ eachindex(A)        
+            a = A[i]
+            if length(a)==numDim
+                if any(a.>siz) || any(a.<1)
+                    error("BoundsError: Indices in A[$i] exceed bounds implied by size provided")
+                end    
+            else                
+                error("Incorrect number of indices in  A[$i], size implies number of indices should be $numDim")
+            end            
+            ind[i] = sub2ind_(a,numDim,k)
+        end
     end
-
     return ind
 end
 
-function sub2ind(siz,A::Vector{Int64})
-
-    numDim = length(siz)
-    k = cumprod([siz[i] for i ∈ eachindex(siz)],dims=1)
-
-    ind = A[1];
-    for q=2:numDim
-        ind += (A[q].-1).*k[q-1]
-    end        
+function sub2ind(siz::Union{Tuple{Vararg{Int64, N}}, Vector{Int64}},A::Vector{Int64})  where N    
+    return sub2ind(siz,[A])
+end
 
+function sub2ind_(a::Union{Tuple{Vararg{Int64, N}}, Vector{Int64}},numDim::Int64,k::Union{Int64,Vector{Int64},Tuple{Vararg{Int64, N}}})  where N
+    if numDim==1 
+        ind = a[1]
+    else        
+        ind = a[1]
+        for i=2:numDim
+            ind += (a[i].-1).*k[i-1]
+        end 
+    end
     return ind
 end
 

test/runtests.jl

@@ -53,6 +53,36 @@ end
     end
 end
 
+@testset "sub2ind" verbose = true begin
+    ind = [1,2,3,4,8,12,30]
+    A = rand(30)
+    IJK_A = [[1],[2],[3],[4],[8],[12],[30]]
+    B = rand(5,6) 
+    IJK_B = [[1, 1], [2, 1], [3, 1], [4, 1], [3, 2], [2, 3], [5, 6]]
+    C = rand(3,5,2)
+    IJK_C = [[1, 1, 1], [2, 1, 1], [3, 1, 1], [1, 2, 1], [2, 3, 1], [3, 4, 1], [3, 5, 2]]
+
+    @testset "1D i.e. Vector" begin        
+        @test sub2ind(size(A),IJK_A)==ind
+    end
+
+    @testset "2D i.e. 2D Matrix" begin    
+        @test sub2ind(size(B),IJK_B)==ind
+    end
+
+    @testset "3D i.e. 3D matrix" begin        
+        @test sub2ind(size(C),IJK_C)==ind
+    end
+
+    @testset "Vector specifying indices 1D" begin        
+        @test [sub2ind(size(A),[i])[1] for i ∈ ind]==ind # IJK = ind for 1D case
+    end
+
+    @testset "Vector specifying size" begin        
+        @test sub2ind(collect(size(C)),IJK_C)==ind
+    end
+end
+
 @testset "elements2indices" verbose = true begin
     @testset "Tri. faces" begin
         F = Vector{TriangleFace{Int64}}(undef, 3)