Add sub2ind and ind2sub functions

These let you convert "subscript" sequences into linear tensor indexes and vice-versa. These are missing from both `Matlab` (where they have the same names) and `numpy` (where they are called `ravel_index` and `unravel_index`).

src/arraymancer/tensor/accessors.nim

@@ -17,15 +17,15 @@ import ./private/p_accessors,
        ./data_structure
 
 proc atContiguousIndex*[T](t: Tensor[T], idx: int): T {.noSideEffect,inline.} =
-  ## Return value of tensor at contiguous index
+  ## Return value of tensor at a contiguous index
   ## i.e. as treat the tensor as flattened
   when T is KnownSupportsCopyMem:
     return t.unsafe_raw_buf[t.getContiguousIndex(idx)]
   else:
     return t.storage.raw_buffer[t.getContiguousIndex(idx)]
 
 proc atContiguousIndex*[T](t: var Tensor[T], idx: int): var T {.noSideEffect,inline.} =
-  ## Return value of tensor at contiguous index (mutable)
+  ## Return value of tensor at a contiguous index (mutable)
   ## i.e. as treat the tensor as flattened
   when T is KnownSupportsCopyMem:
     return t.unsafe_raw_buf[t.getContiguousIndex(idx)]
@@ -359,3 +359,131 @@ iterator enumerateAxis*[T](t: Tensor[T], axis: int): (int, Tensor[T]) {.inline.}
 
 iterator enumerateAxis*[T](t: Tensor[T], axis, offset, size: int): (int, Tensor[T]) {.inline.}=
   enumerateAxisIterator(t, axis, offset, size)
+
+proc sub2ind*(shape: seq[int] | MetaData, sub: seq[int]) : int =
+  ## Convert a subscript sequence into a linear element index given a Tensor shape
+  ##
+  ## Inputs:
+  ##   - A Tensor "shape" (as a sequence of integers or a Tensor.shape)
+  ##   - A Tensor "subscript" as a sequence of integers
+  ##
+  ## Result:
+  ##   - The integer index corresponding to the input subscript given the input shape
+  ##
+  ## Note:
+  ##   - Raises `ValueError` if any of the subscript values equals or exceeds
+  ##     the corresponding shape dimensions.
+  ##
+  ## Examples:
+  ## ```nim
+  ## let shape = [2, 7, 6]
+  ## echo sub2ind(shape, @[0, 0, 3])
+  ## # 3
+  ## echo sub2ind(shape, @[0, 1, 2])
+  ## # 8
+  ## echo sub2ind(shape, @[1, 6, 4])
+  ## # 82
+  ## let t = newTensor[int](2, 7, 6)
+  ## echo sub2ind(t.shape, @[1, 6, 4])
+  ## # 82
+  ## ```
+  result = 0
+  var ax_stride = 1
+  for ax in countdown(sub.high, 0):
+    when compileOption("boundChecks"):
+      if unlikely(sub[ax] >= shape[ax]):
+        raise newException(ValueError,
+          "sub2ind called with an subscript value for axis #" & $ax &
+          " (" & $sub[ax] &
+          ") that is larger than corresponding axis size (" &
+          $shape[ax] & ")")
+    result += sub[ax] * ax_stride
+    ax_stride *= shape[ax]
+
+proc sub2ind*(t: Tensor, sub: seq[int]) : int {. inline .} =
+  ## Convert a Tensor subscript sequence into a linear element index
+  ##
+  ## Inputs:
+  ##   - A Tensor
+  ##   - sub: a Tensor "subscript" as a sequence of integers
+  ##
+  ## Result:
+  ##   - The integer index corresponding to the input subscript given the input shape
+  ##
+  ## Note:
+  ##   - Raises `ValueError` if any of the subscript values equals or exceeds
+  ##     the corresponding shape dimensions.
+  ##
+  ## Examples:
+  ## ```nim
+  ## let t = newTensor[int](2, 7, 6)
+  ## echo sub2ind(t, @[0, 0, 3])
+  ## # 3
+  ## echo sub2ind(t, @[0, 1, 2])
+  ## # 8
+  ## echo sub2ind(t, @[1, 6, 4])
+  ## # 82
+  ## ```
+  sub2ind(t.shape, sub)
+
+proc ind2sub*(shape: seq[int] | MetaData, ind: int) : seq[int] {. noinit .} =
+  ## Convert a linear element index into a subscript sequence given a Tensor shape
+  ##
+  ## Inputs:
+  ##   - A Tensor
+  ##   - An (integer) index to an element in the input Tensor
+  ##
+  ## Result:
+  ##   - The subscript integer sequence corresponding to the input index given the input shape
+  ##
+  ## Note:
+  ##   - Raises `IndexDefect` if any of the index equals or exceeds
+  ##     the tensor size.
+  ##
+  ## Examples:
+  ## ```nim
+  ## let shape = [2, 7, 6]
+  ## echo ind2sub(shape, 3)
+  ## # @[0, 0, 3]
+  ## echo ind2sub(shape, 76)
+  ## # @[1, 5, 4]
+  ## ```
+  when compileOption("boundChecks"):
+    when shape is MetaData:
+      let num_elements = product(shape)
+    else:
+      let num_elements = prod(shape)
+    if ind >= num_elements:
+      raise newException(IndexDefect,
+        "ind2sub called with an index (" & $ind &
+        ") that is larger than the number of elements (" &
+        $num_elements & ")")
+  result = newSeq[int](shape.len)
+  var ax_stride = 1
+  for ax in countdown(shape.len - 1, 0):
+    result[ax] = (ind div ax_stride) mod shape[ax]
+    ax_stride *= shape[ax]
+
+proc ind2sub*(t: Tensor, ind: int) : seq[int] {. noinit, inline .} =
+  ## Convert a Tensor linear element index into a subscript sequence
+  ##
+  ## Inputs:
+  ##   - A Tensor
+  ##   - An (integer) index to an element in the input Tensor
+  ##
+  ## Result:
+  ##   - The subscript integer sequence corresponding to the input index given the input shape
+  ##
+  ## Note:
+  ##   - Raises `IndexDefect` if any of the index equals or exceeds
+  ##     the tensor size.
+  ##
+  ## Examples:
+  ## ```nim
+  ## let t = newTensor[int](2, 7, 6)
+  ## echo ind2sub(t, 3)
+  ## # @[0, 0, 3]
+  ## echo ind2sub(t, 76)
+  ## # @[1, 5, 4]
+  ## ```
+  ind2sub(t.shape, ind)

tests/tensor/test_accessors.nim

@@ -107,5 +107,27 @@ proc main() =
         res += ai + bi
       check: res == 36
 
+    test "sub2ind":
+      let shape = @[2, 7, 6]
+      let t = newTensor[int](shape)
+
+      check: sub2ind(shape, @[0, 0, 3]) == 3
+      check: sub2ind(shape, @[1, 5, 4]) == 76
+      check: sub2ind(t.shape,@[0, 0, 3]) == 3
+      check: sub2ind(t.shape, @[1, 5, 4]) == 76
+      check: sub2ind(t, @[0, 0, 3]) == 3
+      check: sub2ind(t, @[1, 5, 4]) == 76
+
+    test "ind2sub":
+      let shape = @[2, 7, 6]
+      let t = newTensor[int](shape)
+
+      check: ind2sub(shape, 3) == @[0, 0, 3]
+      check: ind2sub(shape, 76) == @[1, 5, 4]
+      check: ind2sub(t.shape, 3) == @[0, 0, 3]
+      check: ind2sub(t.shape, 76) == @[1, 5, 4]
+      check: ind2sub(t, 3) == @[0, 0, 3]
+      check: ind2sub(t, 76) == @[1, 5, 4]
+
 main()
 GC_fullCollect()