"lanes"-like functions from Rust ndarray

[goertzenator]

massiv has the slicing functions outerSlices, innerSlices, and withinSlices that return Array D Ix1 (Array D (Lower ix) e). I would like to have functions that provide the opposite: Array D (Lower ix) (Array D Ix1 e).

Some means to stack them back together would be useful, along with maps and zips that operate on Ix1 "lanes" in an Array.

Rust ndarray has functions like this named rows(), columns(), and lanes().

My current need is to operate on inner Ix1 "lanes" to multiply them by a Ix1 window function. I have been able to do this by the following sequence:

1. resize n dim array to a 2d array
2. call innerSlices
3. map the resulting Ix1 arrays
4. stack the result back together with stackInnerSlices stackOuterSlices
5. resize the array back to its original dimension
   I think this approach is only good for the inner dimension, not outer or within dimensions.

If there is a simpler approach using existing massiv functions I would be happy to hear it.

[goertzenator]

Here is my mapInner implementation. I'm sure there are much smarter ways to do this, but I just wanted to put this here as a reference.

-- Map over 1d inner arrays.
mapInner ::
  forall r r' ix e e'.
  (Index ix, Index (Lower ix), Source r' e', Source r e) =>
  (Array r Ix1 e -> Array r' Ix1 e') ->
  Array r ix e ->
  Array DL ix e'
mapInner f = unlanesInner . map f . lanesInner

-- Produce array of 1d arrays from the innermost dimension.
lanesInner ::
  forall r ix e.
  ( Index (Lower ix)
  , Index ix
  , Source r e
  ) =>
  Array r ix e ->
  Array D (Lower ix) (Array r Ix1 e)
lanesInner arr = resize' lowerSize sliced
 where
  (lowerSize, innerSize) = unsnocSz (size arr)
  size2d = Sz2 (totalElem lowerSize) (unSz innerSize)
  (outer2d, _) = unconsSz size2d
  sliced = makeArray Seq outer2d (arr2d !>)
  arr2d = resize' (Sz2 (totalElem lowerSize) (unSz innerSize)) arr

-- Combine array of 1d arrays into a larger array.
unlanesInner ::
  forall ix r r' e.
  ( Index ix
  , Index (Lower ix)
  , Source r' e
  , Source r (Array r' Ix1 e)
  ) =>
  Array r (Lower ix) (Array r' Ix1 e) ->
  Array DL ix e
unlanesInner arr = resize' destSize (throwEither . stackOuterSlicesM @r' @Ix2 . toList $ arr)
 where
  dimSize = maybe 0 (unSz . size) (headMay . toList $ arr)
  destSize = size arr `snocSz` Sz1 dimSize

[goertzenator]

I started tinkering on a PR yesterday, and while doing so came up with another idea. Consider the existing outer slice operator that takes an Int (Ix1) and produces an array with index type one lower than that of the input array:

(!>)
  :: forall r ix e
   . (HasCallStack, Index ix, Index (Lower ix), Source r e)
  => Array r ix e
  -> Int
  -> Array r (Lower ix) e

Now what I was originally talking about was to have an alternative slice where you would specify an index of Lower ix and get an array of Ix1:

(alt_!>)
  :: forall r ix e
   . (HasCallStack, Index ix, Index (Lower ix), Source r e)
  => Array r ix e
  -> Lower ix
  -> Array r Ix1 e

The idea is, can we extend the original slice operator to take a generic index. Something like:

(!>)
  :: forall r ix ix2 e
   . (HasCallStack, Index ix, Index ix2, Index (Subtracted ix ix2), Source r e)
  => Array r ix e
  -> ix2
  -> Array r (Subtracted ix ix2) e

This would cover the previous two cases and everything in between (ie, Slice a 4d array with a 2d index and get a 2d result).

I am willing to do some tinkering on this but have virtually no type programming experience. If a type programming ninja reads this and knows this kind of thing can't be done, please tell me!

EDIT: I am tinkering on this and it is coming along really well. I hope to post a PR soon.