{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fno-warn-incomplete-uni-patterns #-}

-- | Arrays with shape information and computations at a value-level.
--
-- This is the generic core: the backing vector family is a parameter @v@.
-- Concrete modules ('Harpie.Array', 'Harpie.Array.Storable',
-- 'Harpie.Array.Unboxed') provide type aliases for the common vector families.
module Harpie.Array.Generic
  ( -- * Usage
    -- $usage

    -- * Harpie Arrays
    Array (..),
    array,
    (><),
    validate,
    safeArray,
    unsafeArray,
    unsafeArrayL,
    unsafeModifyShape,
    unsafeModifyVector,

    -- * Dimensions
    Dim,
    Dims,

    -- * Conversion
    asVector,
    arrayV,

    -- * Shape Access
    shape,
    rank,
    size,
    length,
    isNull,

    -- * Indexing
    index,
    (!),
    (!?),
    tabulate,
    backpermute,

    -- * Scalars
    fromScalar,
    toScalar,
    isScalar,
    asSingleton,
    asScalar,

    -- * Array v Creation
    empty,
    range,
    corange,
    indices,
    ident,
    konst,
    singleton,
    diag,
    undiag,

    -- * Element-level functions
    fmapA,
    foldrA,
    toListA,
    zipWith,
    zipWithSafe,
    modify,
    imap,

    -- * Function generalisers
    rowWise,
    colWise,
    dimsWise,

    -- * Single-dimension functions
    take,
    drop,
    select,
    insert,
    delete,
    append,
    prepend,
    concatenate,
    couple,
    slice,
    rotate,

    -- * Multi-dimension functions
    takes,
    drops,
    indexes,
    slices,
    heads,
    lasts,
    tails,
    inits,

    -- * Function application
    extracts,
    reduces,
    joins,
    joinsSafe,
    join,
    joinSafe,
    traverses,
    maps,
    filters,
    zips,
    zipsSafe,
    modifies,
    diffs,

    -- * Array v expansion & contraction
    expand,
    coexpand,
    contract,
    prod,
    dot,
    mult,
    windows,

    -- * Search
    find,
    findNoOverlap,
    findIndices,
    isPrefixOf,
    isSuffixOf,
    isInfixOf,

    -- * Shape manipulations
    fill,
    cut,
    cutSuffix,
    pad,
    lpad,
    reshape,
    flat,
    repeat,
    cycle,
    rerank,
    reorder,
    squeeze,
    elongate,
    transpose,
    inflate,
    intercalate,
    intersperse,
    concats,
    reverses,
    rotates,

    -- * Sorting
    sorts,
    sortsBy,
    orders,
    ordersBy,

    -- * Transmission
    transmit,
    transmitSafe,
    transmitOp,
    telecasts,
    telecastsSafe,

    -- * Row specializations
    pattern (:<),
    cons,
    uncons,
    pattern (:>),
    snoc,
    unsnoc,

    -- * Shape specializations
    iota,

    -- * Math
    uniform,
  )
where

import Control.Monad hiding (join)
import Data.Bool
import Data.Foldable hiding (find, length, minimum)
import Data.Function
import Data.List qualified as List
import Data.Ord (comparing)
import Data.Vector.Algorithms.Intro (sortBy)
import Data.Vector.Generic qualified as VG
import Data.Vector.Unboxed qualified as VU
import GHC.Generics
import Harpie.Shape hiding (asScalar, asSingleton, concatenate, range, rank, reorder, rerank, rotate, size, squeeze)
import Harpie.Shape qualified as S
import NumHask.Algebra.Additive qualified as Add
import NumHask.Algebra.Multiplicative qualified as Mult
import Prettyprinter hiding (dot, fill)
import System.Random hiding (uniform)
import System.Random.Stateful hiding (uniform)
import Prelude as P hiding (cycle, drop, length, repeat, take, zip, zipWith)

-- $setup
-- >>> :m -Prelude
-- >>> import Prelude hiding (take, drop, zipWith, length, cycle, repeat)
-- >>> import Harpie.Array.Generic
-- >>> import Harpie.Shape qualified as S
-- >>> import Data.Vector (Vector)
-- >>> import Data.Vector.Generic qualified as VG
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> import Prettyprinter hiding (dot, fill)
-- >>> import System.Random.Stateful (StatefulGen, UniformRange)
-- >>> import Data.List qualified as List
-- >>> let sumA = foldrA (+) 0
-- >>> let anyA p = foldrA (\x acc -> p x || acc) False
-- >>> let range = Harpie.Array.Generic.range :: [Int] -> Array Vector Int
-- >>> let array = Harpie.Array.Generic.array :: forall a. [Int] -> [a] -> Array Vector a
-- >>> let (><) = (Harpie.Array.Generic.><) :: forall a. [Int] -> [a] -> Array Vector a
-- >>> let toScalar = Harpie.Array.Generic.toScalar :: forall a. a -> Array Vector a
-- >>> let empty = Harpie.Array.Generic.empty :: Array Vector Int
-- >>> let corange = Harpie.Array.Generic.corange :: [Int] -> Array Vector Int
-- >>> let indices = Harpie.Array.Generic.indices :: [Int] -> Array Vector [Int]
-- >>> let ident = Harpie.Array.Generic.ident :: [Int] -> Array Vector Int
-- >>> let konst = Harpie.Array.Generic.konst :: forall a. [Int] -> a -> Array Vector a
-- >>> let singleton = Harpie.Array.Generic.singleton :: forall a. a -> Array Vector a
-- >>> let iota = Harpie.Array.Generic.iota :: Int -> Array Vector Int
-- >>> let uniform = Harpie.Array.Generic.uniform :: forall a g m. (StatefulGen g m, UniformRange a) => g -> [Int] -> (a, a) -> m (Array Vector a)
-- >>> let asArray xs = array [Prelude.length xs] xs
-- >>> let s = 1 :: Array Vector Int
-- >>> s
-- UnsafeArray [] [1]
-- >>> pretty s
-- 1
-- >>> let v = range [3] :: Array Vector Int
-- >>> v
-- UnsafeArray [3] [0,1,2]
-- >>> let m = range [2,3] :: Array Vector Int
-- >>> pretty m
-- [[0,1,2],
--  [3,4,5]]
-- >>> let a = range [2,3,4] :: Array Vector Int
-- >>> a
-- UnsafeArray [2,3,4] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23]
-- >>> pretty a
-- [[[0,1,2,3],
--   [4,5,6,7],
--   [8,9,10,11]],
--  [[12,13,14,15],
--   [16,17,18,19],
--   [20,21,22,23]]]

-- $usage
--
-- Several names used in @harpie@ conflict with [Prelude](https://hackage.haskell.org/package/base/docs/Prelude.html):
--
-- >>> import Prelude hiding (cycle, repeat, take, drop, zipWith, length)
--
-- In general, 'Array' functionality is contained in @Harpie.Array@ and shape  functionality is contained in @Harpie.Shape@. These two modules also have name clashes and at least one needs to be qualified:
--
-- >>> import Harpie.Array.Generic as A
-- >>> import Harpie.Shape qualified as S
--
-- [@prettyprinter@](https://hackage.haskell.org/package/prettyprinter) is used to prettily render arrays to better visualise shape.
--
-- >>> import Prettyprinter hiding (dot,fill)
--
-- Examples of arrays:
--
-- An array with no dimensions (a scalar).
--
-- >>> s = 1 :: Array Vector Int
-- >>> s
-- UnsafeArray [] [1]
-- >>> shape s
-- []
-- >>> pretty s
-- 1
--
-- A single-dimension array (a vector).
--
-- >>> let v = range [3]
-- >>> pretty v
-- [0,1,2]
--
-- A two-dimensional array (a matrix).
--
-- >>> let m = range [2,3]
-- >>> pretty m
-- [[0,1,2],
--  [3,4,5]]
--
-- An n-dimensional array (n should be finite).
--
-- >>> a = range [2,3,4]
-- >>> a
-- UnsafeArray [2,3,4] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23]
-- >>> pretty a
-- [[[0,1,2,3],
--   [4,5,6,7],
--   [8,9,10,11]],
--  [[12,13,14,15],
--   [16,17,18,19],
--   [20,21,22,23]]]

-- | A hyperrectangular (or multidimensional) array with a value-level shape.
--
-- >>> let a = array [2,3,4] [1..24] :: Array Vector Int
-- >>> a
-- UnsafeArray [2,3,4] [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24]
--
-- >>> pretty a
-- [[[1,2,3,4],
--   [5,6,7,8],
--   [9,10,11,12]],
--  [[13,14,15,16],
--   [17,18,19,20],
--   [21,22,23,24]]]
data Array v a = UnsafeArray !(VU.Vector Int) !(VU.Vector Int) !(v a)
  deriving stock ((forall x. Array v a -> Rep (Array v a) x)
-> (forall x. Rep (Array v a) x -> Array v a)
-> Generic (Array v a)
forall x. Rep (Array v a) x -> Array v a
forall x. Array v a -> Rep (Array v a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall k (v :: k -> *) (a :: k) x. Rep (Array v a) x -> Array v a
forall k (v :: k -> *) (a :: k) x. Array v a -> Rep (Array v a) x
$cfrom :: forall k (v :: k -> *) (a :: k) x. Array v a -> Rep (Array v a) x
from :: forall x. Array v a -> Rep (Array v a) x
$cto :: forall k (v :: k -> *) (a :: k) x. Rep (Array v a) x -> Array v a
to :: forall x. Rep (Array v a) x -> Array v a
Generic)

type role Array representational nominal

instance (Eq (v a)) => Eq (Array v a) where
  (UnsafeArray Vector Int
s Vector Int
_ v a
v) == :: Array v a -> Array v a -> Bool
== (UnsafeArray Vector Int
s' Vector Int
_ v a
v') = Vector Int
s Vector Int -> Vector Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector Int
s' Bool -> Bool -> Bool
&& v a
v v a -> v a -> Bool
forall a. Eq a => a -> a -> Bool
== v a
v'

instance (Ord (v a)) => Ord (Array v a) where
  compare :: Array v a -> Array v a -> Ordering
compare (UnsafeArray Vector Int
s Vector Int
_ v a
v) (UnsafeArray Vector Int
s' Vector Int
_ v a
v') = Vector Int -> Vector Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Vector Int
s Vector Int
s' Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> v a -> v a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare v a
v v a
v'

instance (Show (v a)) => Show (Array v a) where
  showsPrec :: Int -> Array v a -> ShowS
showsPrec Int
p (UnsafeArray Vector Int
s Vector Int
_ v a
v) =
    Bool -> ShowS -> ShowS
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10) (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$
      String -> ShowS
showString String
"UnsafeArray " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> ShowS
forall a. Show a => a -> ShowS
shows (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList Vector Int
s) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showString String
" " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v a -> ShowS
forall a. Show a => a -> ShowS
shows v a
v

-- | Internal smart constructor: precomputes strides for O(1) zero-allocation indexing.
unsafeArray :: VU.Vector Int -> v a -> Array v a
unsafeArray :: forall {k} (v :: k -> *) (a :: k). Vector Int -> v a -> Array v a
unsafeArray Vector Int
s v a
v = Vector Int -> Vector Int -> v a -> Array v a
forall {k} (v :: k -> *) (a :: k).
Vector Int -> Vector Int -> v a -> Array v a
UnsafeArray Vector Int
s (Int -> Vector Int -> Vector Int
forall a. Unbox a => Int -> Vector a -> Vector a
VU.drop Int
1 ((Int -> Int -> Int) -> Int -> Vector Int -> Vector Int
forall a b.
(Unbox a, Unbox b) =>
(a -> b -> b) -> b -> Vector a -> Vector b
VU.scanr Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) Int
1 Vector Int
s)) v a
v
{-# INLINE unsafeArray #-}

-- | Internal smart constructor from a list shape.
unsafeArrayL :: [Int] -> v a -> Array v a
unsafeArrayL :: forall {k} (v :: k -> *) (a :: k). [Int] -> v a -> Array v a
unsafeArrayL [Int]
s = Vector Int -> v a -> Array v a
forall {k} (v :: k -> *) (a :: k). Vector Int -> v a -> Array v a
unsafeArray ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s)
{-# INLINE unsafeArrayL #-}

-- | Functor-like map with an explicit vector constraint.
fmapA :: (VG.Vector v a, VG.Vector v b) => (a -> b) -> Array v a -> Array v b
fmapA :: forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> Array v a -> Array v b
fmapA a -> b
f = (v a -> v b) -> Array v a -> Array v b
forall {k} {k} (v :: k -> *) (a :: k) (v' :: k -> *) (b :: k).
(v a -> v' b) -> Array v a -> Array v' b
unsafeModifyVector ((a -> b) -> v a -> v b
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
VG.map a -> b
f)

-- | Foldable-like fold with an explicit vector constraint.
foldrA :: (VG.Vector v a) => (a -> b -> b) -> b -> Array v a -> b
foldrA :: forall (v :: * -> *) a b.
Vector v a =>
(a -> b -> b) -> b -> Array v a -> b
foldrA a -> b -> b
f b
x0 Array v a
a = (a -> b -> b) -> b -> v a -> b
forall (v :: * -> *) a b.
Vector v a =>
(a -> b -> b) -> b -> v a -> b
VG.foldr a -> b -> b
f b
x0 (Array v a -> v a
forall {k} (v :: k -> *) (a :: k). Array v a -> v a
asVector Array v a
a)

-- | Convert an array to a list.
toListA :: (VG.Vector v a) => Array v a -> [a]
toListA :: forall (v :: * -> *) a. Vector v a => Array v a -> [a]
toListA = v a -> [a]
forall (v :: * -> *) a. Vector v a => v a -> [a]
VG.toList (v a -> [a]) -> (Array v a -> v a) -> Array v a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array v a -> v a
forall {k} (v :: k -> *) (a :: k). Array v a -> v a
asVector

-- | Traversable-like traversal through a list intermediary.
traverseA :: (Applicative f, VG.Vector v a, VG.Vector v b, VG.Vector v (f b)) => (a -> f b) -> Array v a -> f (Array v b)
traverseA :: forall (f :: * -> *) (v :: * -> *) a b.
(Applicative f, Vector v a, Vector v b, Vector v (f b)) =>
(a -> f b) -> Array v a -> f (Array v b)
traverseA a -> f b
f Array v a
a = Vector Int -> v b -> Array v b
forall {k} (v :: k -> *) (a :: k). Vector Int -> v a -> Array v a
unsafeArray (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a) (v b -> Array v b) -> ([b] -> v b) -> [b] -> Array v b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [b] -> v b
forall (v :: * -> *) a. Vector v a => [a] -> v a
VG.fromList ([b] -> Array v b) -> f [b] -> f (Array v b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [f b] -> f [b]
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
forall (f :: * -> *) a. Applicative f => [f a] -> f [a]
sequenceA (v (f b) -> [f b]
forall (v :: * -> *) a. Vector v a => v a -> [a]
VG.toList ((a -> f b) -> v a -> v (f b)
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
VG.map a -> f b
f (Array v a -> v a
forall {k} (v :: k -> *) (a :: k). Array v a -> v a
asVector Array v a
a)))

instance (Show a, Show (v a), VG.Vector v a, VG.Vector v (Array v a)) => Pretty (Array v a) where
  pretty :: forall ann. Array v a -> Doc ann
pretty a :: Array v a
a@(UnsafeArray Vector Int
_ Vector Int
_ v a
v) = case Array v a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank Array v a
a of
    Int
0 -> a -> Doc ann
forall a ann. Show a => a -> Doc ann
viaShow (v a -> a
forall (v :: * -> *) a. Vector v a => v a -> a
VG.head v a
v)
    Int
1 -> v a -> Doc ann
forall a ann. Show a => a -> Doc ann
viaShow v a
v
    Int
_ ->
      String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty String
"["
        Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
indent
          Int
0
          ( [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
vsep
              ( Doc ann -> [Doc ann] -> [Doc ann]
forall ann. Doc ann -> [Doc ann] -> [Doc ann]
punctuate Doc ann
forall ann. Doc ann
comma ([Doc ann] -> [Doc ann]) -> [Doc ann] -> [Doc ann]
forall a b. (a -> b) -> a -> b
$
                  Array v a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Array v a -> Doc ann
pretty
                    (Array v a -> Doc ann) -> [Array v a] -> [Doc ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Array v (Array v a) -> [Array v a]
forall (v :: * -> *) a. Vector v a => Array v a -> [a]
toListA ([Int] -> Array v a -> Array v (Array v a)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int
0] Array v a
a)
              )
          )
        Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty String
"]"

-- * conversions

instance (Num a, VG.Vector v a) => Num (Array v a) where
  + :: Array v a -> Array v a -> Array v a
(+) = (a -> a -> a) -> Array v a -> Array v a -> Array v a
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v a -> Array v b -> Array v c
zipWith a -> a -> a
forall a. Num a => a -> a -> a
(+)
  (-) = (a -> a -> a) -> Array v a -> Array v a -> Array v a
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v a -> Array v b -> Array v c
zipWith (-)
  * :: Array v a -> Array v a -> Array v a
(*) = String -> Array v a -> Array v a -> Array v a
forall a. HasCallStack => String -> a
error String
"multiplication not defined"
  abs :: Array v a -> Array v a
abs = (a -> a) -> Array v a -> Array v a
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> Array v a -> Array v b
fmapA a -> a
forall a. Num a => a -> a
abs
  signum :: Array v a -> Array v a
signum = (a -> a) -> Array v a -> Array v a
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> Array v a -> Array v b
fmapA a -> a
forall a. Num a => a -> a
signum
  fromInteger :: Integer -> Array v a
fromInteger Integer
x = a -> Array v a
forall (v :: * -> *) a. Vector v a => a -> Array v a
toScalar (Integer -> a
forall a. Num a => Integer -> a
fromInteger Integer
x)

-- | Convert an array to its underlying vector. Drops shape information.
--
-- > asVector . arrayV [size a] . asVector == a
--
-- >>> asVector (range [2,3])
-- [0,1,2,3,4,5]
asVector :: Array v a -> v a
asVector :: forall {k} (v :: k -> *) (a :: k). Array v a -> v a
asVector (UnsafeArray Vector Int
_ Vector Int
_ v a
v) = v a
v

-- | Construct an array from a shape and a list of elements.
--
-- >>> array [2,3] [0..5]
-- UnsafeArray [2,3] [0,1,2,3,4,5]
array :: (VG.Vector v a) => [Int] -> [a] -> Array v a
array :: forall (v :: * -> *) a. Vector v a => [Int] -> [a] -> Array v a
array [Int]
s [a]
xs = [Int] -> v a -> Array v a
forall {k} (v :: k -> *) (a :: k). [Int] -> v a -> Array v a
unsafeArrayL [Int]
s ([a] -> v a
forall (v :: * -> *) a. Vector v a => [a] -> v a
VG.fromList [a]
xs)

-- | Construct an array from a shape and an already-correctly-sized vector.
arrayV :: [Int] -> v a -> Array v a
arrayV :: forall {k} (v :: k -> *) (a :: k). [Int] -> v a -> Array v a
arrayV = [Int] -> v a -> Array v a
forall {k} (v :: k -> *) (a :: k). [Int] -> v a -> Array v a
unsafeArrayL

-- | Construct an array from a shape and a value without any shape validation.
--
-- >>> array [2,3] [0..5]
-- UnsafeArray [2,3] [0,1,2,3,4,5]
infixl 4 ><

-- | Construct an Array.
--
-- >>> pretty $ [2,3] >< [0..5]
-- [[0,1,2],
--  [3,4,5]]
(><) :: (VG.Vector v a) => [Int] -> [a] -> Array v a
>< :: forall (v :: * -> *) a. Vector v a => [Int] -> [a] -> Array v a
(><) = [Int] -> [a] -> Array v a
forall (v :: * -> *) a. Vector v a => [Int] -> [a] -> Array v a
array

-- | Validate the size and shape of an array.
--
-- >>> validate (array [2,3,4] [1..23] :: Array Vector Int)
-- False
validate :: (VG.Vector v a) => Array v a -> Bool
validate :: forall (v :: * -> *) a. Vector v a => Array v a -> Bool
validate Array v a
a = Array v a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
size Array v a
a Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length (Array v a -> v a
forall {k} (v :: k -> *) (a :: k). Array v a -> v a
asVector Array v a
a)

-- | Construct an Array, checking shape.
--
-- >>> safeArray [2,3,4] [0..23] == Just a
-- True
safeArray :: (VG.Vector v a) => [Int] -> [a] -> Maybe (Array v a)
safeArray :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> [a] -> Maybe (Array v a)
safeArray [Int]
s [a]
xs =
  Maybe (Array v a) -> Maybe (Array v a) -> Bool -> Maybe (Array v a)
forall a. a -> a -> Bool -> a
bool Maybe (Array v a)
forall a. Maybe a
Nothing (Array v a -> Maybe (Array v a)
forall a. a -> Maybe a
Just Array v a
a) (Array v a -> Bool
forall (v :: * -> *) a. Vector v a => Array v a -> Bool
validate Array v a
a)
  where
    a :: Array v a
a = [Int] -> [a] -> Array v a
forall (v :: * -> *) a. Vector v a => [Int] -> [a] -> Array v a
array [Int]
s [a]
xs

-- | Unsafely modify an array shape.
--
-- >>> unsafeModifyShape (VU.map (+1)) (array [2,3] [0..5])
-- UnsafeArray [3,4] [0,1,2,3,4,5]
unsafeModifyShape :: (VU.Vector Int -> VU.Vector Int) -> Array v a -> Array v a
unsafeModifyShape :: forall {k} (v :: k -> *) (a :: k).
(Vector Int -> Vector Int) -> Array v a -> Array v a
unsafeModifyShape Vector Int -> Vector Int
f (UnsafeArray Vector Int
s Vector Int
_ v a
v) = Vector Int -> v a -> Array v a
forall {k} (v :: k -> *) (a :: k). Vector Int -> v a -> Array v a
unsafeArray (Vector Int -> Vector Int
f Vector Int
s) v a
v

-- | Unsafely modify an array vector.
--
-- >>> unsafeModifyVector (VG.map (+1)) (array [2,3] [0..5])
-- UnsafeArray [2,3] [1,2,3,4,5,6]
unsafeModifyVector :: (v a -> v' b) -> Array v a -> Array v' b
unsafeModifyVector :: forall {k} {k} (v :: k -> *) (a :: k) (v' :: k -> *) (b :: k).
(v a -> v' b) -> Array v a -> Array v' b
unsafeModifyVector v a -> v' b
f (UnsafeArray Vector Int
s Vector Int
_ v a
v) = Vector Int -> v' b -> Array v' b
forall {k} (v :: k -> *) (a :: k). Vector Int -> v a -> Array v a
unsafeArray Vector Int
s (v a -> v' b
f v a
v)

-- | Representation of an index into a shape (an [Int]). The index is a dimension of the shape.
type Dim = Int

-- | Representation of indexes into a shape (an [Int]). The indexes are dimensions of the shape.
type Dims = [Int]

-- | shape of an Array
--
-- >>> shape a
-- [2,3,4]
shape :: Array v a -> VU.Vector Int
shape :: forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape (UnsafeArray Vector Int
s Vector Int
_ v a
_) = Vector Int
s

-- | rank of an Array
--
-- >>> rank a
-- 3
rank :: Array v a -> Int
rank :: forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank = Vector Int -> Int
S.rank (Vector Int -> Int)
-> (Array v a -> Vector Int) -> Array v a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape

-- | size of an Array, which is the total number of elements, if the Array Vector is valid.
--
-- >>> size a
-- 24
size :: Array v a -> Int
size :: forall {k} (v :: k -> *) (a :: k). Array v a -> Int
size = Vector Int -> Int
S.size (Vector Int -> Int)
-> (Array v a -> Vector Int) -> Array v a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape

-- | Number of rows (first dimension size) in an Array. As a convention, a scalar value is still a single row.
--
-- >>> length a
-- 2
-- >>> length (toScalar 0)
-- 1
length :: Array v a -> Int
length :: forall {k} (v :: k -> *) (a :: k). Array v a -> Int
length Array v a
a = Int -> Int -> Bool -> Int
forall a. a -> a -> Bool -> a
bool (Vector Int -> Int
forall a. Unbox a => Vector a -> a
VU.head (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a)) Int
1 (Vector Int -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a))

-- | Is the Array Vector empty (has zero number of elements).
--
-- >>> isNull ([2,0] >< [] :: Array Vector ())
-- True
-- >>> isNull ([] >< [4] :: Array Vector Int)
-- False
isNull :: Array v a -> Bool
isNull :: forall {k} (v :: k -> *) (a :: k). Array v a -> Bool
isNull = (Int
0 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==) (Int -> Bool) -> (Array v a -> Int) -> Array v a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array v a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
size

-- | Extract an element at an index, unsafely.
--
-- >>> index a [1,2,3]
-- 23
indexV :: (VG.Vector v a) => Array v a -> VU.Vector Int -> a
indexV :: forall (v :: * -> *) a. Vector v a => Array v a -> Vector Int -> a
indexV (UnsafeArray Vector Int
_ Vector Int
strides v a
v) Vector Int
i = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
VG.unsafeIndex v a
v (Vector Int -> Vector Int -> Int
S.flattenStrides Vector Int
strides Vector Int
i)
{-# NOINLINE indexV #-}

-- | Extract an element at an index, unsafely.
--
-- >>> index a [1,2,3]
-- 23
index :: (VG.Vector v a) => Array v a -> [Int] -> a
index :: forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v a
a = Array v a -> Vector Int -> a
forall (v :: * -> *) a. Vector v a => Array v a -> Vector Int -> a
indexV Array v a
a (Vector Int -> a) -> ([Int] -> Vector Int) -> [Int] -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList
{-# INLINE index #-}

infixl 9 !

-- | Extract an element at an index, unsafely.
--
-- >>> a ! [1,2,3]
-- 23
(!) :: (VG.Vector v a) => Array v a -> [Int] -> a
! :: forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
(!) = Array v a -> [Int] -> a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index

-- | Extract an element at an index, safely.
--
-- >>> a !? [1,2,3]
-- Just 23
-- >>> a !? [2,3,1]
-- Nothing
(!?) :: (VG.Vector v a) => Array v a -> [Int] -> Maybe a
!? :: forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> Maybe a
(!?) Array v a
a [Int]
xs = Maybe a -> Maybe a -> Bool -> Maybe a
forall a. a -> a -> Bool -> a
bool Maybe a
forall a. Maybe a
Nothing (a -> Maybe a
forall a. a -> Maybe a
Just (Array v a
a Array v a -> [Int] -> a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
! [Int]
xs)) ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
xs Vector Int -> Vector Int -> Bool
`S.isFins` Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a)

-- | Tabulate an array supplying a shape and a vector tabulation function.
tabulateV :: (VG.Vector v a) => VU.Vector Int -> (VU.Vector Int -> a) -> Array v a
tabulateV :: forall (v :: * -> *) a.
Vector v a =>
Vector Int -> (Vector Int -> a) -> Array v a
tabulateV Vector Int
ds Vector Int -> a
f =
  let strs :: Vector Int
strs = Vector Int -> Vector Int
S.stridesOf Vector Int
ds
   in Vector Int -> Vector Int -> v a -> Array v a
forall {k} (v :: k -> *) (a :: k).
Vector Int -> Vector Int -> v a -> Array v a
UnsafeArray Vector Int
ds Vector Int
strs (Int -> (Int -> a) -> v a
forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a
VG.generate (Vector Int -> Int
S.size Vector Int
ds) (Vector Int -> a
f (Vector Int -> a) -> (Int -> Vector Int) -> Int -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Int -> Vector Int
S.shapenStrides Vector Int
strs))

-- | Tabulate an array supplying a shape and a tabulation function.
--
-- >>> tabulate [2,3,4] (S.flatten (VU.fromList [2,3,4]) . VU.fromList) == a
-- True
tabulate :: (VG.Vector v a) => [Int] -> ([Int] -> a) -> Array v a
tabulate :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate [Int]
ds [Int] -> a
f = Vector Int -> (Vector Int -> a) -> Array v a
forall (v :: * -> *) a.
Vector v a =>
Vector Int -> (Vector Int -> a) -> Array v a
tabulateV ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) ([Int] -> a
f ([Int] -> a) -> (Vector Int -> [Int]) -> Vector Int -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList)

-- | @backpermute@ is a tabulation where the contents of an array do not need to be accessed, and is thus a fulcrum for leveraging laziness and fusion via the rule:
--
-- > backpermute f g (backpermute f' g' a) == backpermute (f . f') (g . g') a
--
-- Many functions in this module are examples of backpermute usage.
--
-- >>> pretty $ backpermute VU.reverse VU.reverse a
-- [[[0,12],
--   [4,16],
--   [8,20]],
--  [[1,13],
--   [5,17],
--   [9,21]],
--  [[2,14],
--   [6,18],
--   [10,22]],
--  [[3,15],
--   [7,19],
--   [11,23]]]
backpermute :: (VG.Vector v a) => (VU.Vector Int -> VU.Vector Int) -> (VU.Vector Int -> VU.Vector Int) -> Array v a -> Array v a
backpermute :: forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute Vector Int -> Vector Int
f Vector Int -> Vector Int
g Array v a
a = Vector Int -> (Vector Int -> a) -> Array v a
forall (v :: * -> *) a.
Vector v a =>
Vector Int -> (Vector Int -> a) -> Array v a
tabulateV (Vector Int -> Vector Int
f (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a)) (Array v a -> Vector Int -> a
forall (v :: * -> *) a. Vector v a => Array v a -> Vector Int -> a
indexV Array v a
a (Vector Int -> a) -> (Vector Int -> Vector Int) -> Vector Int -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Vector Int
g)
{-# INLINEABLE backpermute #-}

{- RULES
   "backpermute/backpermute" forall f f' g g' (a :: forall a. Array v a)). backpermute f g (backpermute f' g' a) == backpermute (f . f') (g . g') a

-}

-- | Unwrap a scalar.
--
-- >>> let s = array [] [3] :: Array Vector Int
-- >>> fromScalar s
-- 3
fromScalar :: (VG.Vector v a) => Array v a -> a
fromScalar :: forall (v :: * -> *) a. Vector v a => Array v a -> a
fromScalar Array v a
a = Array v a -> [Int] -> a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v a
a ([] :: [Int])

-- | Wrap a scalar.
--
-- >>> :t toScalar 2
-- toScalar 2 :: Num a => Array Vector a
toScalar :: (VG.Vector v a) => a -> Array v a
toScalar :: forall (v :: * -> *) a. Vector v a => a -> Array v a
toScalar a
a = [Int] -> ([Int] -> a) -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate [] (a -> [Int] -> a
forall a b. a -> b -> a
const a
a)

-- | Is an array a Scalar?
--
-- >>> isScalar (toScalar (2::Int))
-- True
isScalar :: Array v a -> Bool
isScalar :: forall {k} (v :: k -> *) (a :: k). Array v a -> Bool
isScalar Array v a
a = Array v a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank Array v a
a Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0

-- | Convert a scalar to being a dimensioned array. Do nothing if not a scalar.
--
-- >>> asSingleton (toScalar 4)
-- UnsafeArray [1] [4]
asSingleton :: Array v a -> Array v a
asSingleton :: forall {k} (v :: k -> *) (a :: k). Array v a -> Array v a
asSingleton = (Vector Int -> Vector Int) -> Array v a -> Array v a
forall {k} (v :: k -> *) (a :: k).
(Vector Int -> Vector Int) -> Array v a -> Array v a
unsafeModifyShape Vector Int -> Vector Int
S.asSingleton

-- | Convert an array with shape [1] to being a scalar (Do nothing if not a shape [1] array).
--
-- >>> asScalar (singleton 3)
-- UnsafeArray [] [3]
asScalar :: Array v a -> Array v a
asScalar :: forall {k} (v :: k -> *) (a :: k). Array v a -> Array v a
asScalar = (Vector Int -> Vector Int) -> Array v a -> Array v a
forall {k} (v :: k -> *) (a :: k).
(Vector Int -> Vector Int) -> Array v a -> Array v a
unsafeModifyShape Vector Int -> Vector Int
S.asScalar

-- * Creation

-- | An array with no elements.
--
-- >>> empty
-- UnsafeArray [0] []
empty :: (VG.Vector v a) => Array v a
empty :: forall (v :: * -> *) a. Vector v a => Array v a
empty = [Int] -> [a] -> Array v a
forall (v :: * -> *) a. Vector v a => [Int] -> [a] -> Array v a
array [Int
0] []

-- | An enumeration of row-major or [lexicographic](https://en.wikipedia.org/wiki/Lexicographic_order) order.
--
-- >>> pretty $ range [2,3]
-- [[0,1,2],
--  [3,4,5]]
range :: (VG.Vector v Int) => [Int] -> Array v Int
range :: forall (v :: * -> *). Vector v Int => [Int] -> Array v Int
range [Int]
xs = [Int] -> ([Int] -> Int) -> Array v Int
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate [Int]
xs (Vector Int -> Vector Int -> Int
S.flatten ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
xs) (Vector Int -> Int) -> ([Int] -> Vector Int) -> [Int] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList)

-- | An enumeration of col-major or [colexicographic](https://en.wikipedia.org/wiki/Lexicographic_order) order.
--
-- >>> pretty (corange [2,3,4])
-- [[[0,6,12,18],
--   [2,8,14,20],
--   [4,10,16,22]],
--  [[1,7,13,19],
--   [3,9,15,21],
--   [5,11,17,23]]]
corange :: (VG.Vector v Int) => [Int] -> Array v Int
corange :: forall (v :: * -> *). Vector v Int => [Int] -> Array v Int
corange [Int]
xs = [Int] -> ([Int] -> Int) -> Array v Int
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate [Int]
xs (Vector Int -> Vector Int -> Int
S.flatten ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList ([Int] -> [Int]
forall a. [a] -> [a]
List.reverse [Int]
xs)) (Vector Int -> Int) -> ([Int] -> Vector Int) -> [Int] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList ([Int] -> Vector Int) -> ([Int] -> [Int]) -> [Int] -> Vector Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [Int]
forall a. [a] -> [a]
List.reverse)

-- | Indices of an array shape.
--
-- >>> pretty $ indices [3,3]
-- [[[0,0],[0,1],[0,2]],
--  [[1,0],[1,1],[1,2]],
--  [[2,0],[2,1],[2,2]]]
indices :: (VG.Vector v [Int]) => [Int] -> Array v [Int]
indices :: forall (v :: * -> *). Vector v [Int] => [Int] -> Array v [Int]
indices [Int]
ds = [Int] -> ([Int] -> [Int]) -> Array v [Int]
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate [Int]
ds [Int] -> [Int]
forall a. a -> a
id

-- | The identity array.
--
-- >>> pretty $ ident [3,3]
-- [[1,0,0],
--  [0,1,0],
--  [0,0,1]]
ident :: (Add.Additive a, Mult.Multiplicative a, VG.Vector v a) => [Int] -> Array v a
ident :: forall a (v :: * -> *).
(Additive a, Multiplicative a, Vector v a) =>
[Int] -> Array v a
ident [Int]
ds = [Int] -> ([Int] -> a) -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate [Int]
ds (a -> a -> Bool -> a
forall a. a -> a -> Bool -> a
bool a
forall a. Additive a => a
Add.zero a
forall a. Multiplicative a => a
Mult.one (Bool -> a) -> ([Int] -> Bool) -> [Int] -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Bool
isDiag (Vector Int -> Bool) -> ([Int] -> Vector Int) -> [Int] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList)

-- | Create an array composed of a single value.
--
-- >>> pretty $ konst [3,2] 1
-- [[1,1],
--  [1,1],
--  [1,1]]
konst :: (VG.Vector v a) => [Int] -> a -> Array v a
konst :: forall (v :: * -> *) a. Vector v a => [Int] -> a -> Array v a
konst [Int]
ds a
a = [Int] -> ([Int] -> a) -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate [Int]
ds (a -> [Int] -> a
forall a b. a -> b -> a
const a
a)

-- | Create an array of shape [1].
--
-- >>> pretty $ singleton 1
-- [1]
-- >>> singleton 3 == toScalar 3
-- False
--
-- >>> asVector (singleton 3) == asVector (toScalar 3)
-- True
singleton :: (VG.Vector v a) => a -> Array v a
singleton :: forall (v :: * -> *) a. Vector v a => a -> Array v a
singleton a
a = [Int] -> v a -> Array v a
forall {k} (v :: k -> *) (a :: k). [Int] -> v a -> Array v a
unsafeArrayL [Int
1] (a -> v a
forall (v :: * -> *) a. Vector v a => a -> v a
VG.singleton a
a)

-- | Extract the diagonal of an array.
--
-- >>> pretty $ diag (ident [3,3])
-- [1,1,1]
diag ::
  (VG.Vector v a) =>
  Array v a ->
  Array v a
diag :: forall (v :: * -> *) a. Vector v a => Array v a -> Array v a
diag Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute Vector Int -> Vector Int
S.minDim (Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate (Array v a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank Array v a
a) (Int -> Vector Int)
-> (Vector Int -> Int) -> Vector Int -> Vector Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Vector Int -> Int
S.getDim Int
0) Array v a
a

-- | Expand the array to form a diagonal array.
--
-- >>> pretty $ undiag (range [3])
-- [[0,0,0],
--  [0,1,0],
--  [0,0,2]]
undiag ::
  (VG.Vector v a) =>
  (Add.Additive a) =>
  Array v a ->
  Array v a
undiag :: forall (v :: * -> *) a.
(Vector v a, Additive a) =>
Array v a -> Array v a
undiag Array v a
a = [Int] -> ([Int] -> a) -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a Vector Int -> Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector a -> Vector a
VU.++ Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a)) (\[Int]
xs -> a -> a -> Bool -> a
forall a. a -> a -> Bool -> a
bool a
forall a. Additive a => a
Add.zero (Array v a -> [Int] -> a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v a
a (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
List.take (Array v a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank Array v a
a) [Int]
xs)) (Vector Int -> Bool
isDiag ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
xs)))

-- | Zip two arrays at an element level.
--
-- >>> zipWith (-) v v
-- UnsafeArray [3] [0,0,0]
zipWith :: (VG.Vector v a, VG.Vector v b, VG.Vector v c) => (a -> b -> c) -> Array v a -> Array v b -> Array v c
zipWith :: forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v a -> Array v b -> Array v c
zipWith a -> b -> c
f (UnsafeArray Vector Int
s Vector Int
_ v a
v) (UnsafeArray Vector Int
_ Vector Int
_ v b
v') = Vector Int -> v c -> Array v c
forall {k} (v :: k -> *) (a :: k). Vector Int -> v a -> Array v a
unsafeArray Vector Int
s ((a -> b -> c) -> v a -> v b -> v c
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> v a -> v b -> v c
VG.zipWith a -> b -> c
f v a
v v b
v')

-- | Zip two arrays at an element level, checking for shape consistency.
--
-- >>> zipWithSafe (-) (range [3]) (range [4])
-- Nothing
zipWithSafe :: (VG.Vector v a, VG.Vector v b, VG.Vector v c) => (a -> b -> c) -> Array v a -> Array v b -> Maybe (Array v c)
zipWithSafe :: forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v a -> Array v b -> Maybe (Array v c)
zipWithSafe a -> b -> c
f (UnsafeArray Vector Int
s Vector Int
_ v a
v) (UnsafeArray Vector Int
s' Vector Int
_ v b
v') = Maybe (Array v c) -> Maybe (Array v c) -> Bool -> Maybe (Array v c)
forall a. a -> a -> Bool -> a
bool Maybe (Array v c)
forall a. Maybe a
Nothing (Array v c -> Maybe (Array v c)
forall a. a -> Maybe a
Just (Array v c -> Maybe (Array v c)) -> Array v c -> Maybe (Array v c)
forall a b. (a -> b) -> a -> b
$ Vector Int -> v c -> Array v c
forall {k} (v :: k -> *) (a :: k). Vector Int -> v a -> Array v a
unsafeArray Vector Int
s ((a -> b -> c) -> v a -> v b -> v c
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> v a -> v b -> v c
VG.zipWith a -> b -> c
f v a
v v b
v')) (Vector Int
s Vector Int -> Vector Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector Int
s')

-- | Modify a single value at an index.
--
-- >>> pretty $ modify [0,0] (const 100) (range [3,2])
-- [[100,1],
--  [2,3],
--  [4,5]]
modify :: (VG.Vector v a) => [Int] -> (a -> a) -> Array v a -> Array v a
modify :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> (a -> a) -> Array v a -> Array v a
modify [Int]
ds a -> a
f Array v a
a = [Int] -> ([Int] -> a) -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a)) (\[Int]
s -> (a -> a) -> (a -> a) -> Bool -> a -> a
forall a. a -> a -> Bool -> a
bool a -> a
forall a. a -> a
id a -> a
f ([Int]
s [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
== [Int]
ds) (Array v a -> [Int] -> a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v a
a [Int]
s))

-- | Maps an index function at element-level.
--
-- >>> pretty $ imap (\xs x -> x - sum xs) a
-- [[[0,0,0,0],
--   [3,3,3,3],
--   [6,6,6,6]],
--  [[11,11,11,11],
--   [14,14,14,14],
--   [17,17,17,17]]]
imap ::
  (VG.Vector v a, VG.Vector v b, VG.Vector v [Int]) =>
  ([Int] -> a -> b) ->
  Array v a ->
  Array v b
imap :: forall (v :: * -> *) a b.
(Vector v a, Vector v b, Vector v [Int]) =>
([Int] -> a -> b) -> Array v a -> Array v b
imap [Int] -> a -> b
f Array v a
a = ([Int] -> a -> b) -> Array v [Int] -> Array v a -> Array v b
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v a -> Array v b -> Array v c
zipWith [Int] -> a -> b
f ([Int] -> Array v [Int]
forall (v :: * -> *). Vector v [Int] => [Int] -> Array v [Int]
indices (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a))) Array v a
a

-- | With a function that takes dimensions and (type-level) parameters, apply the parameters to the initial dimensions. ie
--
-- > rowWise f xs = f [0..] xs
--
-- >>> rowWise indexes [1,0] a
-- UnsafeArray [4] [12,13,14,15]
rowWise :: (Dims -> [x] -> Array v a -> Array v a) -> [x] -> Array v a -> Array v a
rowWise :: forall {k} x (v :: k -> *) (a :: k).
([Int] -> [x] -> Array v a -> Array v a)
-> [x] -> Array v a -> Array v a
rowWise [Int] -> [x] -> Array v a -> Array v a
f [x]
xs Array v a
a = [Int] -> [x] -> Array v a -> Array v a
f [Int
0 .. ([x] -> Int
forall a. [a] -> Int
S.rankL [x]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)] [x]
xs Array v a
a

-- | With a function that takes dimensions and (type-level) parameters, apply the parameters to the the last dimensions. ie
--
-- > colWise f xs = f (List.reverse [0 .. (rank a - 1)]) xs
--
-- >>> colWise indexes [1,0] a
-- UnsafeArray [2] [1,13]
colWise :: (Dims -> [x] -> Array v a -> Array v a) -> [x] -> Array v a -> Array v a
colWise :: forall {k} x (v :: k -> *) (a :: k).
([Int] -> [x] -> Array v a -> Array v a)
-> [x] -> Array v a -> Array v a
colWise [Int] -> [x] -> Array v a -> Array v a
f [x]
xs Array v a
a = [Int] -> [x] -> Array v a -> Array v a
f ([Int] -> [Int]
forall a. [a] -> [a]
List.reverse [(Array v a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank Array v a
a Int -> Int -> Int
forall a. Num a => a -> a -> a
- [x] -> Int
forall a. [a] -> Int
S.rankL [x]
xs) .. (Array v a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank Array v a
a Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)]) [x]
xs Array v a
a

-- | With a function that takes a dimension and a parameter, fold dimensions and parameters using the function.
--
-- >>> dimsWise take [0,2] [1,2] a
-- UnsafeArray [1,3,2] [0,1,4,5,8,9]
dimsWise :: (Dim -> x -> Array v a -> Array v a) -> Dims -> [x] -> Array v a -> Array v a
dimsWise :: forall {k} x (v :: k -> *) (a :: k).
(Int -> x -> Array v a -> Array v a)
-> [Int] -> [x] -> Array v a -> Array v a
dimsWise Int -> x -> Array v a -> Array v a
f [Int]
ds [x]
xs Array v a
a = (Array v a -> (Int, x) -> Array v a)
-> Array v a -> [(Int, x)] -> Array v a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Array v a
a' (Int
d, x
x) -> Int -> x -> Array v a -> Array v a
f Int
d x
x Array v a
a') Array v a
a ([Int] -> [x] -> [(Int, x)]
forall a b. [a] -> [b] -> [(a, b)]
List.zip [Int]
ds [x]
xs)

-- | Take the top-most elements across the specified dimension. Negative values take the bottom-most. No index check is performed.
--
-- > take d x == takes [(d,x)]
--
-- >>> pretty $ take 2 1 a
-- [[[0],
--   [4],
--   [8]],
--  [[12],
--   [16],
--   [20]]]
-- >>> pretty $ take 2 (-1) a
-- [[[3],
--   [7],
--   [11]],
--  [[15],
--   [19],
--   [23]]]
take ::
  (VG.Vector v a) =>
  Dim ->
  Int ->
  Array v a ->
  Array v a
take :: forall (v :: * -> *) a.
Vector v a =>
Int -> Int -> Array v a -> Array v a
take Int
d Int
t Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute Vector Int -> Vector Int
dsNew (Int -> (Int -> Int) -> Vector Int -> Vector Int
modifyDim Int
d (\Int
x -> Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int -> Bool -> Int
forall a. a -> a -> Bool -> a
bool Int
0 (Int -> Vector Int -> Int
getDim Int
d (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
t) (Int
t Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0))) Array v a
a
  where
    dsNew :: Vector Int -> Vector Int
dsNew = Int -> Int -> Vector Int -> Vector Int
takeDim Int
d (Int -> Int
forall a. Num a => a -> a
abs Int
t)

-- | Drop the top-most elements across the specified dimension. Negative values take the bottom-most.
--
-- >>> pretty $ drop 2 1 a
-- [[[1,2,3],
--   [5,6,7],
--   [9,10,11]],
--  [[13,14,15],
--   [17,18,19],
--   [21,22,23]]]
-- >>> pretty $ drop 2 (-1) a
-- [[[0,1,2],
--   [4,5,6],
--   [8,9,10]],
--  [[12,13,14],
--   [16,17,18],
--   [20,21,22]]]
drop ::
  (VG.Vector v a) =>
  Dim ->
  Int ->
  Array v a ->
  Array v a
drop :: forall (v :: * -> *) a.
Vector v a =>
Int -> Int -> Array v a -> Array v a
drop Int
d Int
t Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute Vector Int -> Vector Int
dsNew (Int -> (Int -> Int) -> Vector Int -> Vector Int
modifyDim Int
d (\Int
x -> Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int -> Bool -> Int
forall a. a -> a -> Bool -> a
bool Int
t Int
0 (Int
t Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0))) Array v a
a
  where
    dsNew :: Vector Int -> Vector Int
dsNew = Int -> Int -> Vector Int -> Vector Int
dropDim Int
d (Int -> Int
forall a. Num a => a -> a
abs Int
t)

-- | Select an index along a dimension.
--
-- >>> let s = select 2 3 a
-- >>> pretty s
-- [[3,7,11],
--  [15,19,23]]
select ::
  (VG.Vector v a) =>
  Dim ->
  Int ->
  Array v a ->
  Array v a
select :: forall (v :: * -> *) a.
Vector v a =>
Int -> Int -> Array v a -> Array v a
select Int
d Int
x Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute (Int -> Vector Int -> Vector Int
deleteDim Int
d) (Int -> Int -> Vector Int -> Vector Int
insertDim Int
d Int
x) Array v a
a

-- | Insert along a dimension at a position.
--
-- >>> pretty $ insert 2 0 a (konst [2,3] 0)
-- [[[0,0,1,2,3],
--   [0,4,5,6,7],
--   [0,8,9,10,11]],
--  [[0,12,13,14,15],
--   [0,16,17,18,19],
--   [0,20,21,22,23]]]
-- >>> insert 0 0 (toScalar 1) (toScalar 2)
-- UnsafeArray [2] [2,1]
insert ::
  (VG.Vector v a) =>
  Dim ->
  Int ->
  Array v a ->
  Array v a ->
  Array v a
insert :: forall (v :: * -> *) a.
Vector v a =>
Int -> Int -> Array v a -> Array v a -> Array v a
insert Int
d Int
i Array v a
a Array v a
b = [Int] -> ([Int] -> a) -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Int -> Vector Int -> Vector Int
S.incAt Int
d (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a))) [Int] -> a
go
  where
    go :: [Int] -> a
go [Int]
s =
      let s' :: Vector Int
s' = [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s
       in case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Vector Int -> Int
S.getDim Int
d Vector Int
s') Int
i of
            Ordering
EQ -> Array v a -> [Int] -> a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v a
b (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Int -> Vector Int -> Vector Int
S.deleteDim Int
d Vector Int
s'))
            Ordering
LT -> Array v a -> [Int] -> a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v a
a [Int]
s
            Ordering
GT -> Array v a -> [Int] -> a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v a
a (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Int -> Vector Int -> Vector Int
S.decAt Int
d Vector Int
s'))

-- | Delete along a dimension at a position.
--
-- >>> pretty $ delete 2 0 a
-- [[[1,2,3],
--   [5,6,7],
--   [9,10,11]],
--  [[13,14,15],
--   [17,18,19],
--   [21,22,23]]]
delete ::
  (VG.Vector v a) =>
  Dim ->
  Int ->
  Array v a ->
  Array v a
delete :: forall (v :: * -> *) a.
Vector v a =>
Int -> Int -> Array v a -> Array v a
delete Int
d Int
i Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute (Int -> Vector Int -> Vector Int
decAt Int
d) (\Vector Int
s -> Vector Int -> Vector Int -> Bool -> Vector Int
forall a. a -> a -> Bool -> a
bool (Int -> Vector Int -> Vector Int
incAt Int
d Vector Int
s) Vector Int
s (Int -> Vector Int -> Int
getDim Int
d Vector Int
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
i)) Array v a
a

-- | Insert along a dimension at the end.
--
-- >>> pretty $ append 2 a (konst [2,3] 0)
-- [[[0,1,2,3,0],
--   [4,5,6,7,0],
--   [8,9,10,11,0]],
--  [[12,13,14,15,0],
--   [16,17,18,19,0],
--   [20,21,22,23,0]]]
append ::
  (VG.Vector v a) =>
  Dim ->
  Array v a ->
  Array v a ->
  Array v a
append :: forall (v :: * -> *) a.
Vector v a =>
Int -> Array v a -> Array v a -> Array v a
append Int
d Array v a
a Array v a
b = Int -> Int -> Array v a -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
Int -> Int -> Array v a -> Array v a -> Array v a
insert Int
d (Int -> Vector Int -> Int
getDim Int
d (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a)) Array v a
a Array v a
b

-- | Insert along a dimension at the beginning.
--
-- >>> pretty $ prepend 2 (konst [2,3] 0) a
-- [[[0,0,1,2,3],
--   [0,4,5,6,7],
--   [0,8,9,10,11]],
--  [[0,12,13,14,15],
--   [0,16,17,18,19],
--   [0,20,21,22,23]]]
prepend ::
  (VG.Vector v a) =>
  Dim ->
  Array v a ->
  Array v a ->
  Array v a
prepend :: forall (v :: * -> *) a.
Vector v a =>
Int -> Array v a -> Array v a -> Array v a
prepend Int
d Array v a
a Array v a
b = Int -> Int -> Array v a -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
Int -> Int -> Array v a -> Array v a -> Array v a
insert Int
d Int
0 Array v a
b Array v a
a

-- | Concatenate along a dimension.
--
-- >>> shape $ concatenate 1 a a
-- [2,6,4]
-- >>> concatenate 0 (toScalar 1) (toScalar 2)
-- UnsafeArray [2] [1,2]
-- >>> concatenate 0 (toScalar 0) (asArray [1..3])
-- UnsafeArray [4] [0,1,2,3]
concatenate ::
  (VG.Vector v a) =>
  Dim ->
  Array v a ->
  Array v a ->
  Array v a
concatenate :: forall (v :: * -> *) a.
Vector v a =>
Int -> Array v a -> Array v a -> Array v a
concatenate Int
d Array v a
a0 Array v a
a1 = [Int] -> ([Int] -> a) -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Int -> Vector Int -> Vector Int -> Vector Int
S.concatenate Int
d (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a0) (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a1))) [Int] -> a
go
  where
    go :: [Int] -> a
go [Int]
s =
      let s' :: Vector Int
s' = [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s
       in a -> a -> Bool -> a
forall a. a -> a -> Bool -> a
bool
            (Array v a -> [Int] -> a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v a
a0 [Int]
s)
            ( Array v a -> [Int] -> a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index
                Array v a
a1
                ( Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList
                    ( Int -> Int -> Vector Int -> Vector Int
S.insertDim
                        Int
d
                        (Int -> Vector Int -> Int
S.getDim Int
d Vector Int
s' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Vector Int -> Int
S.getDim Int
d Vector Int
ds0)
                        (Int -> Vector Int -> Vector Int
S.deleteDim Int
d Vector Int
s')
                    )
                )
            )
            (Int -> Vector Int -> Int
S.getDim Int
d Vector Int
s' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int -> Vector Int -> Int
S.getDim Int
d Vector Int
ds0)
    ds0 :: Vector Int
ds0 = Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a0

-- | Combine two arrays as a new dimension of a new array.
--
-- >>> pretty $ couple 0 (asArray [1,2,3]) (asArray [4,5,6::Int])
-- [[1,2,3],
--  [4,5,6]]
couple :: (VG.Vector v a) => Int -> Array v a -> Array v a -> Array v a
couple :: forall (v :: * -> *) a.
Vector v a =>
Int -> Array v a -> Array v a -> Array v a
couple Int
d Array v a
a Array v a
a' = Int -> Array v a -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
Int -> Array v a -> Array v a -> Array v a
concatenate Int
d (Int -> Array v a -> Array v a
forall {k} (v :: k -> *) (a :: k). Int -> Array v a -> Array v a
elongate Int
d Array v a
a) (Int -> Array v a -> Array v a
forall {k} (v :: k -> *) (a :: k). Int -> Array v a -> Array v a
elongate Int
d Array v a
a')

-- | Slice along a dimension with the supplied offset & length.
--
-- >>> let s = slice 2 1 2 a
-- >>> pretty s
-- [[[1,2],
--   [5,6],
--   [9,10]],
--  [[13,14],
--   [17,18],
--   [21,22]]]
slice ::
  (VG.Vector v a) =>
  Dim ->
  Int ->
  Int ->
  Array v a ->
  Array v a
slice :: forall (v :: * -> *) a.
Vector v a =>
Int -> Int -> Int -> Array v a -> Array v a
slice Int
d Int
o Int
l Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute (Int -> Int -> Vector Int -> Vector Int
setDim Int
d Int
l) (Int -> (Int -> Int) -> Vector Int -> Vector Int
modifyDim Int
d (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
o)) Array v a
a

-- | Rotate an array along a dimension.
--
-- >>> pretty $ rotate 1 2 a
-- [[[8,9,10,11],
--   [0,1,2,3],
--   [4,5,6,7]],
--  [[20,21,22,23],
--   [12,13,14,15],
--   [16,17,18,19]]]
rotate ::
  (VG.Vector v a) =>
  Dim ->
  Int ->
  Array v a ->
  Array v a
rotate :: forall (v :: * -> *) a.
Vector v a =>
Int -> Int -> Array v a -> Array v a
rotate Int
d Int
r Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute Vector Int -> Vector Int
forall a. a -> a
id (Int -> Int -> Vector Int -> Vector Int -> Vector Int
rotateIndex Int
d Int
r (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a)) Array v a
a

-- * multi-dimension operators

-- | Takes the top-most elements across the supplied dimension,n tuples. Negative values take the bottom-most.
--
-- > takes == dimsWise take
--
-- >>> pretty $ takes [0,2] [1,-3] a
-- [[[1,2,3],
--   [5,6,7],
--   [9,10,11]]]
takes ::
  (VG.Vector v a) =>
  Dims ->
  [Int] ->
  Array v a ->
  Array v a
takes :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> [Int] -> Array v a -> Array v a
takes [Int]
ds [Int]
xs Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute (Vector Int -> Vector Int -> Vector Int
forall a b. a -> b -> a
const Vector Int
dsNew) ((Int -> Int -> Int) -> Vector Int -> Vector Int -> Vector Int
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
VU.zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Vector Int
start) Array v a
a
  where
    dsNew :: Vector Int
dsNew = Vector Int -> Vector Int -> Vector Int -> Vector Int
S.setDims ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
xsAbs) (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a)
    start :: Vector Int
start = (Int -> Int -> Int) -> Vector Int -> Vector Int -> Vector Int
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
VU.zipWith (\Int
x Int
s -> Int -> Int -> Bool -> Int
forall a. a -> a -> Bool -> a
bool Int
0 (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x) (Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0)) (Vector Int -> Vector Int -> Vector Int -> Vector Int
S.setDims ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
xs) (Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate (Array v a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank Array v a
a) Int
0)) (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a)
    xsAbs :: [Int]
xsAbs = (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Int -> Int
forall a. Num a => a -> a
abs [Int]
xs

-- | Drops the top-most elements. Negative values drop the bottom-most.
--
-- >>> pretty $ drops [0,1,2] [1,2,-3] a
-- [[[20]]]
drops ::
  (VG.Vector v a) =>
  Dims ->
  [Int] ->
  Array v a ->
  Array v a
drops :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> [Int] -> Array v a -> Array v a
drops [Int]
ds [Int]
xs Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute (Vector Int -> Vector Int -> Vector Int
forall a b. a -> b -> a
const Vector Int
dsNew) ((Int -> Int -> Int) -> Vector Int -> Vector Int -> Vector Int
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
VU.zipWith (\Int
d' Int
s' -> Int -> Int -> Bool -> Int
forall a. a -> a -> Bool -> a
bool (Int
d' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
s') Int
s' (Int
d' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0)) Vector Int
xsNew) Array v a
a
  where
    dsNew :: Vector Int
dsNew = Vector Int -> Vector Int -> Vector Int -> Vector Int
S.dropDims ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
xsAbs) (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a)
    xsNew :: Vector Int
xsNew = Vector Int -> Vector Int -> Vector Int -> Vector Int
S.setDims ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
xs) (Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate (Array v a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank Array v a
a) Int
0)
    xsAbs :: [Int]
xsAbs = (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Int -> Int
forall a. Num a => a -> a
abs [Int]
xs

-- | Select by dimensions and indexes.
--
-- >>> let s = indexes [0,1] [1,1] a
-- >>> pretty s
-- [16,17,18,19]
indexes :: (VG.Vector v a) => Dims -> [Int] -> Array v a -> Array v a
indexes :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> [Int] -> Array v a -> Array v a
indexes [Int]
ds [Int]
xs Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute (Vector Int -> Vector Int -> Vector Int
forall a b. a -> b -> a
const (Vector Int -> Vector Int -> Vector Int
S.deleteDims ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a))) (Vector Int -> Vector Int -> Vector Int -> Vector Int
S.insertDims ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
xs)) Array v a
a

-- | Slice along dimensions with the supplied offsets and lengths.
--
-- >>> let s = slices [2,0] [1,1] [2,1] a
-- >>> pretty s
-- [[[13,14],
--   [17,18],
--   [21,22]]]
slices :: (VG.Vector v a) => Dims -> [Int] -> [Int] -> Array v a -> Array v a
slices :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> [Int] -> [Int] -> Array v a -> Array v a
slices [Int]
ds [Int]
os [Int]
ls Array v a
a = (Int -> (Int, Int) -> Array v a -> Array v a)
-> [Int] -> [(Int, Int)] -> Array v a -> Array v a
forall {k} x (v :: k -> *) (a :: k).
(Int -> x -> Array v a -> Array v a)
-> [Int] -> [x] -> Array v a -> Array v a
dimsWise (\Int
d (Int
o, Int
l) -> Int -> Int -> Int -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
Int -> Int -> Int -> Array v a -> Array v a
slice Int
d Int
o Int
l) [Int]
ds ([Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
List.zip [Int]
os [Int]
ls) Array v a
a

-- | Select the first element along the supplied dimensions.
--
-- >>> pretty $ heads [0,2] a
-- [0,4,8]
heads :: (VG.Vector v a) => Dims -> Array v a -> Array v a
heads :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
heads [Int]
ds Array v a
a = [Int] -> [Int] -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> [Int] -> Array v a -> Array v a
indexes [Int]
ds (Int -> Int -> [Int]
forall a. Int -> a -> [a]
List.replicate ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
List.length [Int]
ds) Int
0) Array v a
a

-- | Select the last element along the supplied dimensions.
--
-- >>> pretty $ lasts [0,2] a
-- [15,19,23]
lasts :: (VG.Vector v a) => Dims -> Array v a -> Array v a
lasts :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
lasts [Int]
ds Array v a
a = [Int] -> [Int] -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> [Int] -> Array v a -> Array v a
indexes [Int]
ds [Int]
lastds Array v a
a
  where
    lastds :: [Int]
lastds = (\Int
i -> Int -> Vector Int -> Int
S.getDim Int
i (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int -> Int) -> [Int] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
ds

-- | Select the tail elements along the supplied dimensions.
--
-- >>> pretty $ tails [0,2] a
-- [[[13,14,15],
--   [17,18,19],
--   [21,22,23]]]
tails :: (VG.Vector v a) => Dims -> Array v a -> Array v a
tails :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
tails [Int]
ds Array v a
a = [Int] -> [Int] -> [Int] -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> [Int] -> [Int] -> Array v a -> Array v a
slices [Int]
ds [Int]
os (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList Vector Int
ls) Array v a
a
  where
    os :: [Int]
os = Int -> Int -> [Int]
forall a. Int -> a -> [a]
List.replicate ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
List.length [Int]
ds) Int
1
    ls :: Vector Int
ls = Vector Int -> Vector Int -> Vector Int
S.getLastPositions ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a)

-- | Select the init elements along the supplied dimensions.
--
-- >>> pretty $ inits [0,2] a
-- [[[0,1,2],
--   [4,5,6],
--   [8,9,10]]]
inits :: (VG.Vector v a) => Dims -> Array v a -> Array v a
inits :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
inits [Int]
ds Array v a
a = [Int] -> [Int] -> [Int] -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> [Int] -> [Int] -> Array v a -> Array v a
slices [Int]
ds [Int]
os (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList Vector Int
ls) Array v a
a
  where
    os :: [Int]
os = Int -> Int -> [Int]
forall a. Int -> a -> [a]
List.replicate ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
List.length [Int]
ds) Int
0
    ls :: Vector Int
ls = Vector Int -> Vector Int -> Vector Int
S.getLastPositions ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a)

-- | Extracts dimensions to an outer layer.
--
-- >>> pretty $ fmapA shape (extracts [0] a)
-- [[3,4],[3,4]]
extracts ::
  (VG.Vector v (Array v a), VG.Vector v a) =>
  Dims ->
  Array v a ->
  Array v (Array v a)
extracts :: forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int]
ds Array v a
a = [Int] -> ([Int] -> Array v a) -> Array v (Array v a)
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Vector Int -> Vector Int -> Vector Int
S.getDims ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a))) [Int] -> Array v a
go
  where
    go :: [Int] -> Array v a
go [Int]
s = [Int] -> [Int] -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> [Int] -> Array v a -> Array v a
indexes [Int]
ds [Int]
s Array v a
a

-- | Reduce along specified dimensions, using the supplied fold.
--
-- >>> pretty $ reduces [0] sumA a
-- [66,210]
-- >>> pretty $ reduces [0,2] sumA a
-- [[12,15,18,21],
--  [48,51,54,57]]
reduces ::
  (VG.Vector v a, VG.Vector v b, VG.Vector v (Array v a)) =>
  Dims ->
  (Array v a -> b) ->
  Array v a ->
  Array v b
reduces :: forall (v :: * -> *) a b.
(Vector v a, Vector v b, Vector v (Array v a)) =>
[Int] -> (Array v a -> b) -> Array v a -> Array v b
reduces [Int]
ds Array v a -> b
f Array v a
a = (Array v a -> b) -> Array v (Array v a) -> Array v b
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> Array v a -> Array v b
fmapA Array v a -> b
f ([Int] -> Array v a -> Array v (Array v a)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int]
ds Array v a
a)

-- | Join inner and outer dimension layers by supplied dimensions. No checks on shape.
--
-- >>> let e = extracts [1,0] a
-- >>> let j = joins [1,0] e
-- >>> a == j
-- True
joins ::
  (VG.Vector v (Array v a), VG.Vector v a) =>
  Dims ->
  Array v (Array v a) ->
  Array v a
joins :: forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v (Array v a) -> Array v a
joins [Int]
ds Array v (Array v a)
a = [Int] -> ([Int] -> a) -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Vector Int -> Vector Int -> Vector Int -> Vector Int
S.insertDims ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) Vector Int
so Vector Int
si)) [Int] -> a
go
  where
    go :: [Int] -> a
go [Int]
s = Array v a -> [Int] -> a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index (Array v (Array v a) -> [Int] -> Array v a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v (Array v a)
a (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Vector Int -> Vector Int -> Vector Int
S.getDims ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s)))) (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Vector Int -> Vector Int -> Vector Int
S.deleteDims ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s)))
    so :: Vector Int
so = Array v (Array v a) -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v (Array v a)
a
    si :: Vector Int
si = Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape (Array v (Array v a) -> [Int] -> Array v a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v (Array v a)
a (Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate (Array v (Array v a) -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank Array v (Array v a)
a) Int
0))

-- | Join inner and outer dimension layers by supplied dimensions. Check inner layer shape.
--
-- >>> let e = extracts [1,0] a
-- >>> (Just j) = joinsSafe [1,0] e
-- >>> a == j
-- True
joinsSafe ::
  (VG.Vector v (Array v a), VG.Vector v a, VG.Vector v (VU.Vector Int)) =>
  Dims ->
  Array v (Array v a) ->
  Maybe (Array v a)
joinsSafe :: forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a, Vector v (Vector Int)) =>
[Int] -> Array v (Array v a) -> Maybe (Array v a)
joinsSafe [Int]
ds Array v (Array v a)
a =
  Maybe (Array v a) -> Maybe (Array v a) -> Bool -> Maybe (Array v a)
forall a. a -> a -> Bool -> a
bool
    Maybe (Array v a)
forall a. Maybe a
Nothing
    (Array v a -> Maybe (Array v a)
forall a. a -> Maybe a
Just (Array v a -> Maybe (Array v a)) -> Array v a -> Maybe (Array v a)
forall a b. (a -> b) -> a -> b
$ [Int] -> Array v (Array v a) -> Array v a
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v (Array v a) -> Array v a
joins [Int]
ds Array v (Array v a)
a)
    (Array v (Vector Int) -> Bool
forall a (v :: * -> *). (Eq a, Vector v a) => Array v a -> Bool
allEqual ((Array v a -> Vector Int)
-> Array v (Array v a) -> Array v (Vector Int)
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> Array v a -> Array v b
fmapA Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v (Array v a)
a))

-- | Join inner and outer dimension layers in outer dimension order.
--
-- >>> a == join (extracts [0,1] a)
-- True
join ::
  (VG.Vector v (Array v a), VG.Vector v a) =>
  Array v (Array v a) ->
  Array v a
join :: forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
Array v (Array v a) -> Array v a
join Array v (Array v a)
a = [Int] -> Array v (Array v a) -> Array v a
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v (Array v a) -> Array v a
joins (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Vector Int -> Vector Int
S.dimsOf (Array v (Array v a) -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v (Array v a)
a))) Array v (Array v a)
a

-- | Join inner and outer dimension layers in outer dimension order, checking for consistent inner dimension shape.
--
-- >>> joinSafe (extracts [0,1] a)
-- Just (UnsafeArray [2,3,4] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23])
joinSafe ::
  (VG.Vector v (Array v a), VG.Vector v a, VG.Vector v (VU.Vector Int)) =>
  Array v (Array v a) ->
  Maybe (Array v a)
joinSafe :: forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a, Vector v (Vector Int)) =>
Array v (Array v a) -> Maybe (Array v a)
joinSafe Array v (Array v a)
a =
  Maybe (Array v a) -> Maybe (Array v a) -> Bool -> Maybe (Array v a)
forall a. a -> a -> Bool -> a
bool
    Maybe (Array v a)
forall a. Maybe a
Nothing
    (Array v a -> Maybe (Array v a)
forall a. a -> Maybe a
Just (Array v a -> Maybe (Array v a)) -> Array v a -> Maybe (Array v a)
forall a b. (a -> b) -> a -> b
$ Array v (Array v a) -> Array v a
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
Array v (Array v a) -> Array v a
join Array v (Array v a)
a)
    (Array v (Vector Int) -> Bool
forall a (v :: * -> *). (Eq a, Vector v a) => Array v a -> Bool
allEqual ((Array v a -> Vector Int)
-> Array v (Array v a) -> Array v (Vector Int)
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> Array v a -> Array v b
fmapA Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v (Array v a)
a))

-- | Satisfy a predicate across all elements
allEqual :: (Eq a, VG.Vector v a) => Array v a -> Bool
allEqual :: forall a (v :: * -> *). (Eq a, Vector v a) => Array v a -> Bool
allEqual Array v a
a = case v a -> [a]
forall (v :: * -> *) a. Vector v a => v a -> [a]
VG.toList (Array v a -> v a
forall {k} (v :: k -> *) (a :: k). Array v a -> v a
asVector Array v a
a) of
  [] -> Bool
True
  (a
x : [a]
xs) -> (a -> Bool) -> [a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
x) [a]
xs

-- | Traverse along specified dimensions.
--
-- traverses [1] print (range [2,3])
-- 0
-- 3
-- 1
-- 4
-- 2
-- 5
-- UnsafeArray [2,3] [(),(),(),(),(),()]
traverses ::
  (VG.Vector v a, VG.Vector v b, VG.Vector v (Array v a), VG.Vector v (Array v b), VG.Vector v (f b), VG.Vector v (f (Array v b))) =>
  (Applicative f) =>
  Dims ->
  (a -> f b) ->
  Array v a ->
  f (Array v b)
traverses :: forall (v :: * -> *) a b (f :: * -> *).
(Vector v a, Vector v b, Vector v (Array v a),
 Vector v (Array v b), Vector v (f b), Vector v (f (Array v b)),
 Applicative f) =>
[Int] -> (a -> f b) -> Array v a -> f (Array v b)
traverses [Int]
ds a -> f b
f Array v a
a = [Int] -> Array v (Array v b) -> Array v b
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v (Array v a) -> Array v a
joins [Int]
ds (Array v (Array v b) -> Array v b)
-> f (Array v (Array v b)) -> f (Array v b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Array v a -> f (Array v b))
-> Array v (Array v a) -> f (Array v (Array v b))
forall (f :: * -> *) (v :: * -> *) a b.
(Applicative f, Vector v a, Vector v b, Vector v (f b)) =>
(a -> f b) -> Array v a -> f (Array v b)
traverseA ((a -> f b) -> Array v a -> f (Array v b)
forall (f :: * -> *) (v :: * -> *) a b.
(Applicative f, Vector v a, Vector v b, Vector v (f b)) =>
(a -> f b) -> Array v a -> f (Array v b)
traverseA a -> f b
f) ([Int] -> Array v a -> Array v (Array v a)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int]
ds Array v a
a)

-- | Maps a function along specified dimensions.
--
-- >>> pretty $ maps [1] transpose a
-- [[[0,12],
--   [4,16],
--   [8,20]],
--  [[1,13],
--   [5,17],
--   [9,21]],
--  [[2,14],
--   [6,18],
--   [10,22]],
--  [[3,15],
--   [7,19],
--   [11,23]]]
maps ::
  (VG.Vector v a, VG.Vector v b, VG.Vector v (Array v a), VG.Vector v (Array v b)) =>
  Dims ->
  (Array v a -> Array v b) ->
  Array v a ->
  Array v b
maps :: forall (v :: * -> *) a b.
(Vector v a, Vector v b, Vector v (Array v a),
 Vector v (Array v b)) =>
[Int] -> (Array v a -> Array v b) -> Array v a -> Array v b
maps [Int]
ds Array v a -> Array v b
f Array v a
a = [Int] -> Array v (Array v b) -> Array v b
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v (Array v a) -> Array v a
joins [Int]
ds ((Array v a -> Array v b)
-> Array v (Array v a) -> Array v (Array v b)
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> Array v a -> Array v b
fmapA Array v a -> Array v b
f ([Int] -> Array v a -> Array v (Array v a)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int]
ds Array v a
a))

-- | Filters along specified dimensions (which are flattened).
--
-- >>> pretty $ filters [0,1] (anyA ((==0) . (`mod` 7))) a
-- [[0,1,2,3],
--  [4,5,6,7],
--  [12,13,14,15],
--  [20,21,22,23]]
filters ::
  (VG.Vector v a, VG.Vector v (Array v a)) =>
  Dims ->
  (Array v a -> Bool) ->
  Array v a ->
  Array v a
filters :: forall (v :: * -> *) a.
(Vector v a, Vector v (Array v a)) =>
[Int] -> (Array v a -> Bool) -> Array v a -> Array v a
filters [Int]
ds Array v a -> Bool
p Array v a
a = Array v (Array v a) -> Array v a
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
Array v (Array v a) -> Array v a
join ([Int] -> v (Array v a) -> Array v (Array v a)
forall {k} (v :: k -> *) (a :: k). [Int] -> v a -> Array v a
unsafeArrayL [v (Array v a) -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length v (Array v a)
v'] v (Array v a)
v')
  where
    v' :: v (Array v a)
v' = (Array v a -> Bool) -> v (Array v a) -> v (Array v a)
forall (v :: * -> *) a. Vector v a => (a -> Bool) -> v a -> v a
VG.filter Array v a -> Bool
p (Array v (Array v a) -> v (Array v a)
forall {k} (v :: k -> *) (a :: k). Array v a -> v a
asVector ([Int] -> Array v a -> Array v (Array v a)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int]
ds Array v a
a))

-- | Zips two arrays with a function along specified dimensions.
--
-- >>> pretty $ zips [0,1] (zipWith (,)) a (reverses [0] a)
-- [[[(0,12),(1,13),(2,14),(3,15)],
--   [(4,16),(5,17),(6,18),(7,19)],
--   [(8,20),(9,21),(10,22),(11,23)]],
--  [[(12,0),(13,1),(14,2),(15,3)],
--   [(16,4),(17,5),(18,6),(19,7)],
--   [(20,8),(21,9),(22,10),(23,11)]]]
zips ::
  (VG.Vector v a, VG.Vector v b, VG.Vector v c, VG.Vector v (Array v a), VG.Vector v (Array v b), VG.Vector v (Array v c)) =>
  Dims ->
  (Array v a -> Array v b -> Array v c) ->
  Array v a ->
  Array v b ->
  Array v c
zips :: forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c, Vector v (Array v a),
 Vector v (Array v b), Vector v (Array v c)) =>
[Int]
-> (Array v a -> Array v b -> Array v c)
-> Array v a
-> Array v b
-> Array v c
zips [Int]
ds Array v a -> Array v b -> Array v c
f Array v a
a Array v b
b = [Int] -> Array v (Array v c) -> Array v c
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v (Array v a) -> Array v a
joins [Int]
ds ((Array v a -> Array v b -> Array v c)
-> Array v (Array v a)
-> Array v (Array v b)
-> Array v (Array v c)
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v a -> Array v b -> Array v c
zipWith Array v a -> Array v b -> Array v c
f ([Int] -> Array v a -> Array v (Array v a)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int]
ds Array v a
a) ([Int] -> Array v b -> Array v (Array v b)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int]
ds Array v b
b))

-- | Zips two arrays with a function along specified dimensions, checking shapes.
--
-- >>> zipsSafe [0] (zipWith (,)) (asArray [1::Int]) (asArray [1,2::Int])
-- Nothing
zipsSafe ::
  (VG.Vector v a, VG.Vector v b, VG.Vector v c, VG.Vector v (Array v a), VG.Vector v (Array v b), VG.Vector v (Array v c)) =>
  Dims ->
  (Array v a -> Array v b -> Array v c) ->
  Array v a ->
  Array v b ->
  Maybe (Array v c)
zipsSafe :: forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c, Vector v (Array v a),
 Vector v (Array v b), Vector v (Array v c)) =>
[Int]
-> (Array v a -> Array v b -> Array v c)
-> Array v a
-> Array v b
-> Maybe (Array v c)
zipsSafe [Int]
ds Array v a -> Array v b -> Array v c
f Array v a
a Array v b
b =
  Maybe (Array v c) -> Maybe (Array v c) -> Bool -> Maybe (Array v c)
forall a. a -> a -> Bool -> a
bool
    (Array v c -> Maybe (Array v c)
forall a. a -> Maybe a
Just (Array v c -> Maybe (Array v c)) -> Array v c -> Maybe (Array v c)
forall a b. (a -> b) -> a -> b
$ [Int] -> Array v (Array v c) -> Array v c
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v (Array v a) -> Array v a
joins [Int]
ds ((Array v a -> Array v b -> Array v c)
-> Array v (Array v a)
-> Array v (Array v b)
-> Array v (Array v c)
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v a -> Array v b -> Array v c
zipWith Array v a -> Array v b -> Array v c
f ([Int] -> Array v a -> Array v (Array v a)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int]
ds Array v a
a) ([Int] -> Array v b -> Array v (Array v b)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int]
ds Array v b
b)))
    Maybe (Array v c)
forall a. Maybe a
Nothing
    (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a Vector Int -> Vector Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Array v b -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v b
b)

-- | Modify using the supplied function along dimensions & positions.
--
-- >>> pretty $ modifies (fmapA (100+)) [2] [0] a
-- [[[100,1,2,3],
--   [104,5,6,7],
--   [108,9,10,11]],
--  [[112,13,14,15],
--   [116,17,18,19],
--   [120,21,22,23]]]
modifies ::
  (VG.Vector v a, VG.Vector v (Array v a)) =>
  (Array v a -> Array v a) ->
  Dims ->
  [Int] ->
  Array v a ->
  Array v a
modifies :: forall (v :: * -> *) a.
(Vector v a, Vector v (Array v a)) =>
(Array v a -> Array v a)
-> [Int] -> [Int] -> Array v a -> Array v a
modifies Array v a -> Array v a
f [Int]
ds [Int]
ps Array v a
a = [Int] -> Array v (Array v a) -> Array v a
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v (Array v a) -> Array v a
joins [Int]
ds (Array v (Array v a) -> Array v a)
-> Array v (Array v a) -> Array v a
forall a b. (a -> b) -> a -> b
$ [Int]
-> (Array v a -> Array v a)
-> Array v (Array v a)
-> Array v (Array v a)
forall (v :: * -> *) a.
Vector v a =>
[Int] -> (a -> a) -> Array v a -> Array v a
modify [Int]
ps Array v a -> Array v a
f ([Int] -> Array v a -> Array v (Array v a)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int]
ds Array v a
a)

-- | Apply a binary function between successive slices, across dimensions and lags.
--
-- >>> pretty $ diffs [1] [1] (zipWith (-)) a
-- [[[4,4,4,4],
--   [4,4,4,4]],
--  [[4,4,4,4],
--   [4,4,4,4]]]
diffs :: (VG.Vector v a, VG.Vector v b, VG.Vector v (Array v a), VG.Vector v (Array v b)) => Dims -> [Int] -> (Array v a -> Array v a -> Array v b) -> Array v a -> Array v b
diffs :: forall (v :: * -> *) a b.
(Vector v a, Vector v b, Vector v (Array v a),
 Vector v (Array v b)) =>
[Int]
-> [Int]
-> (Array v a -> Array v a -> Array v b)
-> Array v a
-> Array v b
diffs [Int]
ds [Int]
xs Array v a -> Array v a -> Array v b
f Array v a
a = [Int]
-> (Array v a -> Array v a -> Array v b)
-> Array v a
-> Array v a
-> Array v b
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c, Vector v (Array v a),
 Vector v (Array v b), Vector v (Array v c)) =>
[Int]
-> (Array v a -> Array v b -> Array v c)
-> Array v a
-> Array v b
-> Array v c
zips [Int]
ds Array v a -> Array v a -> Array v b
f ([Int] -> [Int] -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> [Int] -> Array v a -> Array v a
drops [Int]
ds [Int]
xs Array v a
a) ([Int] -> [Int] -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> [Int] -> Array v a -> Array v a
drops [Int]
ds ((Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Int -> Int
forall a. Num a => a -> a
P.negate [Int]
xs) Array v a
a)

-- | Product two arrays using the supplied binary function.
--
-- For context, if the function is multiply, and the arrays are tensors,
-- then this can be interpreted as a [tensor product](https://en.wikipedia.org/wiki/Tensor_product).
-- The concept of a tensor product is a dense crossroad, and a complete treatment is elsewhere.  To quote the wiki article:
--
-- ... the tensor product can be extended to other categories of mathematical objects in addition to vector spaces, such as to matrices, tensors, algebras, topological vector spaces, and modules. In each such case the tensor product is characterized by a similar universal property: it is the freest bilinear operation. The general concept of a "tensor product" is captured by monoidal categories; that is, the class of all things that have a tensor product is a monoidal category.
--
-- >>> x = array [3] [1,2,3]
-- >>> pretty $ expand (*) x x
-- [[1,2,3],
--  [2,4,6],
--  [3,6,9]]
--
-- Alternatively, expand can be understood as representing the permutation of element pairs of two arrays, so like the Applicative List instance.
--
-- >>> i2 = indices [2,2]
-- >>> pretty $ expand (,) i2 i2
-- [[[[([0,0],[0,0]),([0,0],[0,1])],
--    [([0,0],[1,0]),([0,0],[1,1])]],
--   [[([0,1],[0,0]),([0,1],[0,1])],
--    [([0,1],[1,0]),([0,1],[1,1])]]],
--  [[[([1,0],[0,0]),([1,0],[0,1])],
--    [([1,0],[1,0]),([1,0],[1,1])]],
--   [[([1,1],[0,0]),([1,1],[0,1])],
--    [([1,1],[1,0]),([1,1],[1,1])]]]]
expand ::
  (VG.Vector v a, VG.Vector v b, VG.Vector v c) =>
  (a -> b -> c) ->
  Array v a ->
  Array v b ->
  Array v c
expand :: forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v a -> Array v b -> Array v c
expand a -> b -> c
f Array v a
a Array v b
b = [Int] -> ([Int] -> c) -> Array v c
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a Vector Int -> Vector Int -> Vector Int
forall a. Semigroup a => a -> a -> a
<> Array v b -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v b
b)) (\[Int]
i -> a -> b -> c
f (Array v a -> [Int] -> a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v a
a (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
List.take Int
r [Int]
i)) (Array v b -> [Int] -> b
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v b
b (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
List.drop Int
r [Int]
i)))
  where
    r :: Int
r = Array v a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank Array v a
a

-- | Like expand, but permutes the first array first, rather than the second.
--
-- >>> pretty $ expand (,) v (fmapA (+3) v)
-- [[(0,3),(0,4),(0,5)],
--  [(1,3),(1,4),(1,5)],
--  [(2,3),(2,4),(2,5)]]
--
-- >>> pretty $ coexpand (,) v (fmapA (+3) v)
-- [[(0,3),(1,3),(2,3)],
--  [(0,4),(1,4),(2,4)],
--  [(0,5),(1,5),(2,5)]]
--
-- The output shape is @shape b ++ shape a@, so the first array's axes occupy
-- the suffix of the product shape rather than the prefix.  Equivalently, it is
-- @expand@ followed by the block-swap permutation that exchanges the two
-- operand shapes.
coexpand ::
  (VG.Vector v a, VG.Vector v b, VG.Vector v c) =>
  (a -> b -> c) ->
  Array v a ->
  Array v b ->
  Array v c
coexpand :: forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v a -> Array v b -> Array v c
coexpand a -> b -> c
f Array v a
a Array v b
b = [Int] -> ([Int] -> c) -> Array v c
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v b -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v b
b Vector Int -> Vector Int -> Vector Int
forall a. Semigroup a => a -> a -> a
<> Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a)) (\[Int]
i -> a -> b -> c
f (Array v a -> [Int] -> a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v a
a (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
List.drop Int
rb [Int]
i)) (Array v b -> [Int] -> b
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v b
b (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
List.take Int
rb [Int]
i)))
  where
    rb :: Int
rb = Array v b -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank Array v b
b

-- | Contract an array by applying the supplied (folding) function on diagonal elements of the dimensions.
--
-- This generalises a tensor contraction by allowing the number of contracting diagonals to be other than 2.
--
--
-- >>> pretty $ contract [1,2] sumA (expand (*) m (transpose m))
-- [[5,14],
--  [14,50]]
contract ::
  (VG.Vector v a, VG.Vector v b, VG.Vector v (Array v a)) =>
  Dims ->
  (Array v a -> b) ->
  Array v a ->
  Array v b
contract :: forall (v :: * -> *) a b.
(Vector v a, Vector v b, Vector v (Array v a)) =>
[Int] -> (Array v a -> b) -> Array v a -> Array v b
contract [Int]
ds Array v a -> b
f Array v a
a = (Array v a -> b) -> Array v (Array v a) -> Array v b
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> Array v a -> Array v b
fmapA (Array v a -> b
f (Array v a -> b) -> (Array v a -> Array v a) -> Array v a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array v a -> Array v a
forall (v :: * -> *) a. Vector v a => Array v a -> Array v a
diag) ([Int] -> Array v a -> Array v (Array v a)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Vector Int -> Vector Int -> Vector Int
S.exceptDims ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a))) Array v a
a)

-- | Product two arrays using the supplied function and then contract the result using the supplied matching dimensions and function.
--
-- >>> pretty $ prod [1] [0] sumA (*) (range [2,3]) (range [3,2])
-- [[10,13],
--  [28,40]]
--
-- With full laziness, this computation would be equivalent to:
--
-- > f . diag <$> extracts ds' (expand g a b)
prod ::
  (VG.Vector v a, VG.Vector v b, VG.Vector v c, VG.Vector v d) =>
  Dims ->
  Dims ->
  (Array v c -> d) ->
  (a -> b -> c) ->
  Array v a ->
  Array v b ->
  Array v d
prod :: forall (v :: * -> *) a b c d.
(Vector v a, Vector v b, Vector v c, Vector v d) =>
[Int]
-> [Int]
-> (Array v c -> d)
-> (a -> b -> c)
-> Array v a
-> Array v b
-> Array v d
prod [Int]
ds0 [Int]
ds1 Array v c -> d
g a -> b -> c
f Array v a
a Array v b
b =
  Vector Int -> (Vector Int -> d) -> Array v d
forall (v :: * -> *) a.
Vector v a =>
Vector Int -> (Vector Int -> a) -> Array v a
tabulateV
    (Vector Int -> Vector Int -> Vector Int
S.deleteDims Vector Int
ds0V (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a) Vector Int -> Vector Int -> Vector Int
forall a. Semigroup a => a -> a -> a
<> Vector Int -> Vector Int -> Vector Int
S.deleteDims Vector Int
ds1V (Array v b -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v b
b))
    ( \Vector Int
so ->
        Array v c -> d
g (Array v c -> d) -> Array v c -> d
forall a b. (a -> b) -> a -> b
$
          Vector Int -> (Vector Int -> c) -> Array v c
forall (v :: * -> *) a.
Vector v a =>
Vector Int -> (Vector Int -> a) -> Array v a
tabulateV
            (Vector Int -> Vector Int -> Vector Int
S.getDims Vector Int
ds0V (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a))
            ( \Vector Int
si ->
                a -> b -> c
f
                  (Array v a -> Vector Int -> a
forall (v :: * -> *) a. Vector v a => Array v a -> Vector Int -> a
indexV Array v a
a (Vector Int -> Vector Int -> Vector Int -> Vector Int
S.insertDims Vector Int
ds0V Vector Int
si (Int -> Vector Int -> Vector Int
forall a. Unbox a => Int -> Vector a -> Vector a
VU.take Int
sp Vector Int
so)))
                  (Array v b -> Vector Int -> b
forall (v :: * -> *) a. Vector v a => Array v a -> Vector Int -> a
indexV Array v b
b (Vector Int -> Vector Int -> Vector Int -> Vector Int
S.insertDims Vector Int
ds1V Vector Int
si (Int -> Vector Int -> Vector Int
forall a. Unbox a => Int -> Vector a -> Vector a
VU.drop Int
sp Vector Int
so)))
            )
    )
  where
    ds0V :: Vector Int
ds0V = [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds0
    ds1V :: Vector Int
ds1V = [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds1
    sp :: Int
sp = Array v a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank Array v a
a Int -> Int -> Int
forall a. Num a => a -> a -> a
- Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
ds0V

-- | A generalisation of a dot operation, which is a multiplicative expansion of two arrays and sum contraction along the middle two dimensions.
--
-- matrix multiplication
--
-- >>> pretty $ dot sumA (*) m (transpose m)
-- [[5,14],
--  [14,50]]
--
-- inner product
--
-- >>> pretty $ dot sumA (*) v v
-- 5
--
-- matrix-vector multiplication
-- Note that an Array Vector with shape [3] is neither a row vector nor column vector.
--
-- >>> pretty $ dot sumA (*) v (transpose m)
-- [5,14]
--
-- >>> pretty $ dot sumA (*) m v
-- [5,14]
dot ::
  (VG.Vector v a, VG.Vector v b, VG.Vector v c, VG.Vector v d, VG.Vector v (Array v c)) =>
  (Array v c -> d) ->
  (a -> b -> c) ->
  Array v a ->
  Array v b ->
  Array v d
dot :: forall (v :: * -> *) a b c d.
(Vector v a, Vector v b, Vector v c, Vector v d,
 Vector v (Array v c)) =>
(Array v c -> d)
-> (a -> b -> c) -> Array v a -> Array v b -> Array v d
dot Array v c -> d
f a -> b -> c
g Array v a
a Array v b
b = [Int] -> (Array v c -> d) -> Array v c -> Array v d
forall (v :: * -> *) a b.
(Vector v a, Vector v b, Vector v (Array v a)) =>
[Int] -> (Array v a -> b) -> Array v a -> Array v b
contract [Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
r] Array v c -> d
f ((a -> b -> c) -> Array v a -> Array v b -> Array v c
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v a -> Array v b -> Array v c
expand a -> b -> c
g Array v a
a Array v b
b)
  where
    r :: Int
r = Array v a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank Array v a
a

-- | Array Vector multiplication.
--
-- matrix multiplication
--
-- >>> pretty $ mult m (transpose m)
-- [[5,14],
--  [14,50]]
--
-- inner product
--
-- >>> pretty $ mult v v
-- 5
--
-- matrix-vector multiplication
--
-- >>> pretty $ mult v (transpose m)
-- [5,14]
--
-- >>> pretty $ mult m v
-- [5,14]
mult ::
  (VG.Vector v a) =>
  (Add.Additive a, Mult.Multiplicative a) =>
  Array v a ->
  Array v a ->
  Array v a
mult :: forall (v :: * -> *) a.
(Vector v a, Additive a, Multiplicative a) =>
Array v a -> Array v a -> Array v a
mult Array v a
a Array v a
b = [Int]
-> [Int]
-> (Array v a -> a)
-> (a -> a -> a)
-> Array v a
-> Array v a
-> Array v a
forall (v :: * -> *) a b c d.
(Vector v a, Vector v b, Vector v c, Vector v d) =>
[Int]
-> [Int]
-> (Array v c -> d)
-> (a -> b -> c)
-> Array v a
-> Array v b
-> Array v d
prod [Array v a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank Array v a
a Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] [Int
0] ((a -> a -> a) -> a -> Array v a -> a
forall (v :: * -> *) a b.
Vector v a =>
(a -> b -> b) -> b -> Array v a -> b
foldrA a -> a -> a
forall a. Additive a => a -> a -> a
(Add.+) a
forall a. Additive a => a
Add.zero) a -> a -> a
forall a. Multiplicative a => a -> a -> a
(Mult.*) Array v a
a Array v a
b

-- | @windows xs@ are xs-sized windows of an array
--
-- >>> shape $ windows [2,2] (range [4,3,2])
-- [3,2,2,2,2]
windows :: (VG.Vector v a) => [Int] -> Array v a -> Array v a
windows :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
windows [Int]
xs Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute (Vector Int -> Vector Int -> Vector Int
S.expandWindows ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
xs)) (Int -> Vector Int -> Vector Int
S.indexWindows ([Int] -> Int
forall a. [a] -> Int
S.rankL [Int]
xs)) Array v a
a

-- | Find the starting positions of occurences of one array in another.
--
-- >>> a = cycle [4,4] (range [3]) :: Array Vector Int
-- >>> i = array [2,2] [1,2,2,0] :: Array Vector Int
-- >>> pretty $ find i a
-- [[False,True,False],
--  [True,False,False],
--  [False,False,True]]
find :: (Eq (v a), VG.Vector v Bool, VG.Vector v a, VG.Vector v (Array v a)) => Array v a -> Array v a -> Array v Bool
find :: forall (v :: * -> *) a.
(Eq (v a), Vector v Bool, Vector v a, Vector v (Array v a)) =>
Array v a -> Array v a -> Array v Bool
find Array v a
i Array v a
a = Array v Bool
xs
  where
    i' :: Array v a
i' = Int -> Array v a -> Array v a
forall {k} (v :: k -> *) (a :: k). Int -> Array v a -> Array v a
rerank (Array v a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank Array v a
a) Array v a
i
    ws :: Array v a
ws = [Int] -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
windows (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
i')) Array v a
a
    xs :: Array v Bool
xs = (Array v a -> Bool) -> Array v (Array v a) -> Array v Bool
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> Array v a -> Array v b
fmapA (Array v a -> Array v a -> Bool
forall a. Eq a => a -> a -> Bool
== Array v a
i') ([Int] -> Array v a -> Array v (Array v a)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Vector Int -> Vector Int -> Vector Int
S.dimWindows (Vector Int -> Vector Int -> Vector Int
S.expandWindows (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
i') (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a)) (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a))) Array v a
ws)

-- | Find the ending positions of one array in another except where the array overlaps with another copy.
--
-- >>> a = konst [5,5] 1 :: Array Vector Int
-- >>> i = konst [2,2] 1 :: Array Vector Int
-- >>> pretty $ findNoOverlap i a
-- [[True,False,True,False],
--  [False,False,False,False],
--  [True,False,True,False],
--  [False,False,False,False]]
findNoOverlap :: (Eq (v a), VG.Vector v Bool, VG.Vector v a, VG.Vector v (Array v a)) => Array v a -> Array v a -> Array v Bool
findNoOverlap :: forall (v :: * -> *) a.
(Eq (v a), Vector v Bool, Vector v a, Vector v (Array v a)) =>
Array v a -> Array v a -> Array v Bool
findNoOverlap Array v a
i Array v a
a = Array v Bool
r
  where
    f :: Array v Bool
f = Array v a -> Array v a -> Array v Bool
forall (v :: * -> *) a.
(Eq (v a), Vector v Bool, Vector v a, Vector v (Array v a)) =>
Array v a -> Array v a -> Array v Bool
find Array v a
i Array v a
a

    cl :: [Int] -> [[Int]]
    cl :: [Int] -> [[Int]]
cl [Int]
sh =
      ([Int] -> Bool) -> [[Int]] -> [[Int]]
forall a. (a -> Bool) -> [a] -> [a]
List.filter (Bool -> Bool
P.not (Bool -> Bool) -> ([Int] -> Bool) -> [Int] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) ([Int] -> Bool) -> ([Int] -> [Int]) -> [Int] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [Int]
forall a. HasCallStack => [a] -> [a]
List.init) ([[Int]] -> [[Int]]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> a -> b
$
        ([Int] -> Bool) -> [[Int]] -> [[Int]]
forall a. (a -> Bool) -> [a] -> [a]
List.filter (Bool -> Bool
P.not (Bool -> Bool) -> ([Int] -> Bool) -> [Int] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0)) ([[Int]] -> [[Int]]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> a -> b
$
          ([Int] -> [Int]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\[Int]
s -> (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
List.zipWith (\Int
x Int
x0 -> Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
x0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [Int]
s [Int]
sh) ((Int -> [Int]) -> [Int] -> [[Int]]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (\Int
x -> [Int
0 .. (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2)]) [Int]
sh)
    go :: Array v Bool -> [Int] -> Bool
go Array v Bool
r' [Int]
s = Array v Bool -> [Int] -> Bool
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v Bool
f [Int]
s Bool -> Bool -> Bool
&& Bool -> Bool
not (([Int] -> Bool) -> [[Int]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Array v Bool -> [Int] -> Bool
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v Bool
r') (([Int] -> Bool) -> [[Int]] -> [[Int]]
forall a. (a -> Bool) -> [a] -> [a]
List.filter (\[Int]
x -> Vector Int -> Vector Int -> Bool
S.isFins ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
x) (Array v Bool -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v Bool
f)) ([[Int]] -> [[Int]]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> a -> b
$ ([Int] -> [Int]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
List.zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) [Int]
s) ([Int] -> [[Int]]
cl (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
i)))))
    r :: Array v Bool
r = [Int] -> ([Int] -> Bool) -> Array v Bool
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v Bool -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v Bool
f)) (Array v Bool -> [Int] -> Bool
go Array v Bool
r)

-- | Find the indices of the starting location of one array in another.
--
-- >>> b = cycle [4,4] (range [3]) :: Array Vector Int
-- >>> i = array [2,2] [1,2,2,0] :: Array Vector Int
-- >>> pretty $ findIndices i b
-- [[0,1],[1,0],[2,2]]
findIndices :: (Eq (v a), VG.Vector v [Int], VG.Vector v a, VG.Vector v Bool, VG.Vector v ([Int], Bool), VG.Vector v (Array v a)) => Array v a -> Array v a -> Array v [Int]
findIndices :: forall (v :: * -> *) a.
(Eq (v a), Vector v [Int], Vector v a, Vector v Bool,
 Vector v ([Int], Bool), Vector v (Array v a)) =>
Array v a -> Array v a -> Array v [Int]
findIndices Array v a
i Array v a
a = [Int] -> v [Int] -> Array v [Int]
forall {k} (v :: k -> *) (a :: k). [Int] -> v a -> Array v a
unsafeArrayL [v ([Int], Bool) -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length v ([Int], Bool)
v'] ((([Int], Bool) -> [Int]) -> v ([Int], Bool) -> v [Int]
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
VG.map ([Int], Bool) -> [Int]
forall a b. (a, b) -> a
fst v ([Int], Bool)
v')
  where
    v' :: v ([Int], Bool)
v' = (([Int], Bool) -> Bool) -> v ([Int], Bool) -> v ([Int], Bool)
forall (v :: * -> *) a. Vector v a => (a -> Bool) -> v a -> v a
VG.filter ([Int], Bool) -> Bool
forall a b. (a, b) -> b
snd (Array v ([Int], Bool) -> v ([Int], Bool)
forall {k} (v :: k -> *) (a :: k). Array v a -> v a
asVector (([Int] -> Bool -> ([Int], Bool))
-> Array v Bool -> Array v ([Int], Bool)
forall (v :: * -> *) a b.
(Vector v a, Vector v b, Vector v [Int]) =>
([Int] -> a -> b) -> Array v a -> Array v b
imap (,) Array v Bool
b))
    b :: Array v Bool
b = Array v a -> Array v a -> Array v Bool
forall (v :: * -> *) a.
(Eq (v a), Vector v Bool, Vector v a, Vector v (Array v a)) =>
Array v a -> Array v a -> Array v Bool
find Array v a
i Array v a
a

-- | Check if the first array is a prefix of the second
--
-- >>> isPrefixOf (array [2,2] [0,1,4,5]) a
-- True
isPrefixOf :: (Eq (v a), VG.Vector v a) => Array v a -> Array v a -> Bool
isPrefixOf :: forall (v :: * -> *) a.
(Eq (v a), Vector v a) =>
Array v a -> Array v a -> Bool
isPrefixOf Array v a
p Array v a
a = Array v a
p Array v a -> Array v a -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
cut (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
p)) Array v a
a

-- | Check if the first array is a suffix of the second
--
-- >>> isSuffixOf (array [2,2] [18,19,22,23]) a
-- True
isSuffixOf :: (Eq (v a), VG.Vector v a) => Array v a -> Array v a -> Bool
isSuffixOf :: forall (v :: * -> *) a.
(Eq (v a), Vector v a) =>
Array v a -> Array v a -> Bool
isSuffixOf Array v a
p Array v a
a = Array v a
p Array v a -> Array v a -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
cutSuffix (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
p)) Array v a
a

-- | Check if the first array is an infix of the second
--
-- >>> isInfixOf (array [2,2] [18,19,22,23]) a
-- True
isInfixOf :: (Eq (v a), VG.Vector v a, VG.Vector v Bool, VG.Vector v (Array v a)) => Array v a -> Array v a -> Bool
isInfixOf :: forall (v :: * -> *) a.
(Eq (v a), Vector v a, Vector v Bool, Vector v (Array v a)) =>
Array v a -> Array v a -> Bool
isInfixOf Array v a
p Array v a
a = (Bool -> Bool -> Bool) -> Bool -> Array v Bool -> Bool
forall (v :: * -> *) a b.
Vector v a =>
(a -> b -> b) -> b -> Array v a -> b
foldrA Bool -> Bool -> Bool
(||) Bool
False (Array v a -> Array v a -> Array v Bool
forall (v :: * -> *) a.
(Eq (v a), Vector v Bool, Vector v a, Vector v (Array v a)) =>
Array v a -> Array v a -> Array v Bool
find Array v a
p Array v a
a)

-- * shape manipulation

-- | Fill an array with the supplied value without regard to the original shape or cut the array values to match array size.
--
-- > validate (def x a) == True
--
-- >>> pretty $ fill 0 (array [3] [])
-- [0,0,0]
-- >>> pretty $ fill 0 (array [3] [1..4])
-- [1,2,3]
fill :: (VG.Vector v a) => a -> Array v a -> Array v a
fill :: forall (v :: * -> *) a. Vector v a => a -> Array v a -> Array v a
fill a
x (UnsafeArray Vector Int
s Vector Int
_ v a
v) = Vector Int -> v a -> Array v a
forall {k} (v :: k -> *) (a :: k). Vector Int -> v a -> Array v a
unsafeArray Vector Int
s (Int -> v a -> v a
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
VG.take (Vector Int -> Int
S.size Vector Int
s) (v a
v v a -> v a -> v a
forall (v :: * -> *) a. Vector v a => v a -> v a -> v a
VG.++ Int -> a -> v a
forall (v :: * -> *) a. Vector v a => Int -> a -> v a
VG.replicate (Vector Int -> Int
S.size Vector Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length v a
v) a
x))

-- | Cut an array to form a new (smaller) shape. Errors if the new shape is larger. The old array is reranked to the rank of the new shape first.
--
-- >>> cut [2] (array [4] [0..3] :: Array Vector Int)
-- UnsafeArray [2] [0,1]
cut ::
  (VG.Vector v a) =>
  [Int] ->
  Array v a ->
  Array v a
cut :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
cut [Int]
s' Array v a
a = Array v a -> Array v a -> Bool -> Array v a
forall a. a -> a -> Bool -> a
bool (String -> Array v a
forall a. HasCallStack => String -> a
error String
"bad cut") ([Int] -> ([Int] -> a) -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate [Int]
s' (Array v a -> [Int] -> a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v a
a')) (Vector Int -> Vector Int -> Bool
S.isSubset ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s') (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a))
  where
    a' :: Array v a
a' = Int -> Array v a -> Array v a
forall {k} (v :: k -> *) (a :: k). Int -> Array v a -> Array v a
rerank ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
List.length [Int]
s') Array v a
a

-- | Cut an array to form a new (smaller) shape, using suffix elements. Errors if the new shape is larger. The old array is reranked to the rank of the new shape first.
--
-- >>> cutSuffix [2,2] a
-- UnsafeArray [2,2] [18,19,22,23]
cutSuffix ::
  (VG.Vector v a) =>
  [Int] ->
  Array v a ->
  Array v a
cutSuffix :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
cutSuffix [Int]
s' Array v a
a = Array v a -> Array v a -> Bool -> Array v a
forall a. a -> a -> Bool -> a
bool (String -> Array v a
forall a. HasCallStack => String -> a
error String
"bad cut") ([Int] -> ([Int] -> a) -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate [Int]
s' (Array v a -> [Int] -> a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v a
a' ([Int] -> a) -> ([Int] -> [Int]) -> [Int] -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
List.zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) [Int]
diffDim)) (Vector Int -> Vector Int -> Bool
S.isSubset ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s') (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a))
  where
    a' :: Array v a
a' = Int -> Array v a -> Array v a
forall {k} (v :: k -> *) (a :: k). Int -> Array v a -> Array v a
rerank ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
List.length [Int]
s') Array v a
a
    diffDim :: [Int]
diffDim = Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList ((Int -> Int -> Int) -> Vector Int -> Vector Int -> Vector Int
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
VU.zipWith (-) (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a') ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s'))

-- | Pad an array to form a new shape, supplying a default value for elements outside the shape of the old array. The old array is reranked to the rank of the new shape first.
--
-- >>> pad 0 [5] (array [4] [0..3] :: Array Vector Int)
-- UnsafeArray [5] [0,1,2,3,0]
pad ::
  (VG.Vector v a) =>
  a ->
  [Int] ->
  Array v a ->
  Array v a
pad :: forall (v :: * -> *) a.
Vector v a =>
a -> [Int] -> Array v a -> Array v a
pad a
d [Int]
s' Array v a
a = [Int] -> ([Int] -> a) -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate [Int]
s' (\[Int]
s -> a -> a -> Bool -> a
forall a. a -> a -> Bool -> a
bool a
d (Array v a -> [Int] -> a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v a
a' [Int]
s) ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s Vector Int -> Vector Int -> Bool
`S.isFins` Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a'))
  where
    a' :: Array v a
a' = Int -> Array v a -> Array v a
forall {k} (v :: k -> *) (a :: k). Int -> Array v a -> Array v a
rerank ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
List.length [Int]
s') Array v a
a

-- | Left pad an array to form a new shape, supplying a default value for elements outside the shape of the old array.
--
-- >>> lpad 0 [5] (array [4] [0..3] :: Array Vector Int)
-- UnsafeArray [5] [0,0,1,2,3]
-- >>> pretty $ lpad 0 [3,3] (range [2,2] :: Array Vector Int)
-- [[0,0,0],
--  [0,0,1],
--  [0,2,3]]
lpad ::
  (VG.Vector v a) =>
  a ->
  [Int] ->
  Array v a ->
  Array v a
lpad :: forall (v :: * -> *) a.
Vector v a =>
a -> [Int] -> Array v a -> Array v a
lpad a
d [Int]
s' Array v a
a = [Int] -> ([Int] -> a) -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
tabulate [Int]
s' (\[Int]
s -> a -> a -> Bool -> a
forall a. a -> a -> Bool -> a
bool a
d (Array v a -> [Int] -> a
forall (v :: * -> *) a. Vector v a => Array v a -> [Int] -> a
index Array v a
a' ([Int] -> [Int]
olds [Int]
s)) ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList ([Int] -> [Int]
olds [Int]
s) Vector Int -> Vector Int -> Bool
`S.isFins` Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a'))
  where
    a' :: Array v a
a' = Int -> Array v a -> Array v a
forall {k} (v :: k -> *) (a :: k). Int -> Array v a -> Array v a
rerank ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
List.length [Int]
s') Array v a
a
    gap :: [Int]
gap = Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList ((Int -> Int -> Int) -> Vector Int -> Vector Int -> Vector Int
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
VU.zipWith (-) ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s') (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a'))
    olds :: [Int] -> [Int]
olds [Int]
s = (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
List.zipWith (-) [Int]
s [Int]
gap

-- | Reshape an array (with the same or less number of elements).
--
-- >>> pretty $ reshape [4,3,2] a
-- [[[0,1],
--   [2,3],
--   [4,5]],
--  [[6,7],
--   [8,9],
--   [10,11]],
--  [[12,13],
--   [14,15],
--   [16,17]],
--  [[18,19],
--   [20,21],
--   [22,23]]]
reshape ::
  (VG.Vector v a) =>
  [Int] ->
  Array v a ->
  Array v a
reshape :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
reshape [Int]
s Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute (Vector Int -> Vector Int -> Vector Int
forall a b. a -> b -> a
const ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s)) (Vector Int -> Int -> Vector Int
S.shapen (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a) (Int -> Vector Int)
-> (Vector Int -> Int) -> Vector Int -> Vector Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Vector Int -> Int
S.flatten ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s)) Array v a
a

-- | Make an Array Vector single dimensional.
--
-- >>> pretty $ flat (range [2,2])
-- [0,1,2,3]
-- >>> pretty (flat $ toScalar 0)
-- [0]
flat :: Array v a -> Array v a
flat :: forall {k} (v :: k -> *) (a :: k). Array v a -> Array v a
flat Array v a
a = (Vector Int -> Vector Int) -> Array v a -> Array v a
forall {k} (v :: k -> *) (a :: k).
(Vector Int -> Vector Int) -> Array v a -> Array v a
unsafeModifyShape (Int -> Vector Int
forall a. Unbox a => a -> Vector a
VU.singleton (Int -> Vector Int)
-> (Vector Int -> Int) -> Vector Int -> Vector Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Int
S.size) Array v a
a

-- | Reshape an array, repeating the original array. The shape of the array should be a suffix of the new shape.
--
-- >>> pretty $ repeat [2,2,2] (array [2] [1,2])
-- [[[1,2],
--   [1,2]],
--  [[1,2],
--   [1,2]]]
--
-- > repeat ds (toScalar x) == konst ds x
repeat ::
  (VG.Vector v a) =>
  [Int] ->
  Array v a ->
  Array v a
repeat :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
repeat [Int]
s Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute (Vector Int -> Vector Int -> Vector Int
forall a b. a -> b -> a
const ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s)) (Int -> Vector Int -> Vector Int
forall a. Unbox a => Int -> Vector a -> Vector a
VU.drop ([Int] -> Int
forall a. [a] -> Int
S.rankL [Int]
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Array v a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank Array v a
a)) Array v a
a

-- | Reshape an array, cycling through the elements without regard to the original shape.
--
-- >>> pretty $ cycle [2,2,2] (array [3] [1,2,3])
-- [[[1,2],
--   [3,1]],
--  [[2,3],
--   [1,2]]]
cycle ::
  (VG.Vector v a) =>
  [Int] ->
  Array v a ->
  Array v a
cycle :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
cycle [Int]
s Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute (Vector Int -> Vector Int -> Vector Int
forall a b. a -> b -> a
const ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s)) (Vector Int -> Int -> Vector Int
S.shapen (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a) (Int -> Vector Int)
-> (Vector Int -> Int) -> Vector Int -> Vector Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Array v a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
size Array v a
a) (Int -> Int) -> (Vector Int -> Int) -> Vector Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Vector Int -> Int
S.flatten ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s)) Array v a
a

-- | Change rank by adding new dimensions at the front, if the new rank is greater, or combining dimensions (from left to right) into rows, if the new rank is lower.
--
-- >>> shape (rerank 4 a)
-- [1,2,3,4]
-- >>> shape (rerank 2 a)
-- [6,4]
--
-- > flat == rerank 1
rerank :: Int -> Array v a -> Array v a
rerank :: forall {k} (v :: k -> *) (a :: k). Int -> Array v a -> Array v a
rerank Int
r Array v a
a = (Vector Int -> Vector Int) -> Array v a -> Array v a
forall {k} (v :: k -> *) (a :: k).
(Vector Int -> Vector Int) -> Array v a -> Array v a
unsafeModifyShape (Int -> Vector Int -> Vector Int
S.rerank Int
r) Array v a
a

-- | Change the order of dimensions.
--
-- >>> pretty $ reorder [2,0,1] a
-- [[[0,4,8],
--   [12,16,20]],
--  [[1,5,9],
--   [13,17,21]],
--  [[2,6,10],
--   [14,18,22]],
--  [[3,7,11],
--   [15,19,23]]]
reorder ::
  (VG.Vector v a) =>
  Dims ->
  Array v a ->
  Array v a
reorder :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
reorder [Int]
ds Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute (Vector Int -> Vector Int -> Vector Int
`S.reorder` [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) (\Vector Int
s -> Vector Int -> Vector Int -> Vector Int -> Vector Int
S.insertDims ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) Vector Int
s Vector Int
forall a. Unbox a => Vector a
VU.empty) Array v a
a

-- | Remove single dimensions.
--
-- >>> let sq = array [2,1,3,4,1] [1..24] :: Array Vector Int
-- >>> shape $ squeeze sq
-- [2,3,4]
--
-- >>> shape $ squeeze (singleton 0)
-- []
squeeze ::
  Array v a ->
  Array v a
squeeze :: forall {k} (v :: k -> *) (a :: k). Array v a -> Array v a
squeeze Array v a
a = (Vector Int -> Vector Int) -> Array v a -> Array v a
forall {k} (v :: k -> *) (a :: k).
(Vector Int -> Vector Int) -> Array v a -> Array v a
unsafeModifyShape Vector Int -> Vector Int
S.squeeze Array v a
a

-- | Insert a single dimension at the supplied position.
--
-- >>> shape $ elongate 1 a
-- [2,1,3,4]
-- >>> elongate 0 (toScalar 1)
-- UnsafeArray [1] [1]
elongate ::
  Dim ->
  Array v a ->
  Array v a
elongate :: forall {k} (v :: k -> *) (a :: k). Int -> Array v a -> Array v a
elongate Int
d Array v a
a = (Vector Int -> Vector Int) -> Array v a -> Array v a
forall {k} (v :: k -> *) (a :: k).
(Vector Int -> Vector Int) -> Array v a -> Array v a
unsafeModifyShape (Int -> Int -> Vector Int -> Vector Int
insertDim Int
d Int
1) Array v a
a

-- | Reverse indices eg transposes the element A/ijk/ to A/kji/.
--
-- >>> index (transpose a) [1,0,0] == index a [0,0,1]
-- True
-- >>> pretty $ transpose (array [2,2,2] [1..8])
-- [[[1,5],
--   [3,7]],
--  [[2,6],
--   [4,8]]]
transpose :: (VG.Vector v a) => Array v a -> Array v a
transpose :: forall (v :: * -> *) a. Vector v a => Array v a -> Array v a
transpose Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector a
VU.reverse Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector a
VU.reverse Array v a
a

-- | Inflate an array by inserting a new dimension given a supplied dimension and size.
--
-- alt name: replicate
--
-- >>> pretty $ inflate 0 2 (array [3] [0,1,2])
-- [[0,1,2],
--  [0,1,2]]
inflate ::
  (VG.Vector v a) =>
  Dim ->
  Int ->
  Array v a ->
  Array v a
inflate :: forall (v :: * -> *) a.
Vector v a =>
Int -> Int -> Array v a -> Array v a
inflate Int
d Int
n Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute (Int -> Int -> Vector Int -> Vector Int
insertDim Int
d Int
n) (Int -> Vector Int -> Vector Int
deleteDim Int
d) Array v a
a

-- | Intercalate an array along dimensions.
--
-- >>> pretty $ intercalate 2 (konst [2,3] 0) a
-- [[[0,0,1,0,2,0,3],
--   [4,0,5,0,6,0,7],
--   [8,0,9,0,10,0,11]],
--  [[12,0,13,0,14,0,15],
--   [16,0,17,0,18,0,19],
--   [20,0,21,0,22,0,23]]]
intercalate :: (VG.Vector v a, VG.Vector v (Array v a)) => Dim -> Array v a -> Array v a -> Array v a
intercalate :: forall (v :: * -> *) a.
(Vector v a, Vector v (Array v a)) =>
Int -> Array v a -> Array v a -> Array v a
intercalate Int
d Array v a
i Array v a
a = [Int] -> Array v (Array v a) -> Array v a
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v (Array v a) -> Array v a
joins [Int
d] (Array v (Array v a) -> Array v a)
-> Array v (Array v a) -> Array v a
forall a b. (a -> b) -> a -> b
$ [Int] -> v (Array v a) -> Array v (Array v a)
forall {k} (v :: k -> *) (a :: k). [Int] -> v a -> Array v a
unsafeArrayL [[Array v a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
List.length [Array v a]
xs] ([Array v a] -> v (Array v a)
forall (v :: * -> *) a. Vector v a => [a] -> v a
VG.fromList [Array v a]
xs)
  where
    xs :: [Array v a]
xs = Array v a -> [Array v a] -> [Array v a]
forall a. a -> [a] -> [a]
List.intersperse Array v a
i (v (Array v a) -> [Array v a]
forall (v :: * -> *) a. Vector v a => v a -> [a]
VG.toList (Array v (Array v a) -> v (Array v a)
forall {k} (v :: k -> *) (a :: k). Array v a -> v a
asVector ([Int] -> Array v a -> Array v (Array v a)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int
d] Array v a
a)))

-- | Intersperse an element along dimensions.
--
-- >>> pretty $ intersperse 2 0 a
-- [[[0,0,1,0,2,0,3],
--   [4,0,5,0,6,0,7],
--   [8,0,9,0,10,0,11]],
--  [[12,0,13,0,14,0,15],
--   [16,0,17,0,18,0,19],
--   [20,0,21,0,22,0,23]]]
intersperse :: (VG.Vector v a, VG.Vector v (Array v a)) => Dim -> a -> Array v a -> Array v a
intersperse :: forall (v :: * -> *) a.
(Vector v a, Vector v (Array v a)) =>
Int -> a -> Array v a -> Array v a
intersperse Int
d a
i Array v a
a = Int -> Array v a -> Array v a -> Array v a
forall (v :: * -> *) a.
(Vector v a, Vector v (Array v a)) =>
Int -> Array v a -> Array v a -> Array v a
intercalate Int
d ([Int] -> a -> Array v a
forall (v :: * -> *) a. Vector v a => [Int] -> a -> Array v a
konst (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Int -> Vector Int -> Vector Int
S.deleteDim Int
d (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a))) a
i) Array v a
a

-- | Concatenate and replace dimensions, creating a new dimension at the supplied postion.
--
-- >>> pretty $ concats [0,1] 1 a
-- [[0,4,8,12,16,20],
--  [1,5,9,13,17,21],
--  [2,6,10,14,18,22],
--  [3,7,11,15,19,23]]
concats ::
  (VG.Vector v a) =>
  Dims ->
  Int ->
  Array v a ->
  Array v a
concats :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> Int -> Array v a -> Array v a
concats [Int]
ds Int
n Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute (Vector Int -> Int -> Vector Int -> Vector Int
S.concatDims ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) Int
n) (Vector Int -> Int -> Vector Int -> Vector Int -> Vector Int
S.unconcatDimsIndex ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) Int
n (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a)) Array v a
a

-- | Reverses element order along specified dimensions.
--
-- >>> pretty $ reverses [0,1] a
-- [[[20,21,22,23],
--   [16,17,18,19],
--   [12,13,14,15]],
--  [[8,9,10,11],
--   [4,5,6,7],
--   [0,1,2,3]]]
reverses ::
  (VG.Vector v a) =>
  Dims ->
  Array v a ->
  Array v a
reverses :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
reverses [Int]
ds Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute Vector Int -> Vector Int
forall a. a -> a
id (Vector Int -> Vector Int -> Vector Int -> Vector Int
S.reverseIndex ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a)) Array v a
a

-- | Rotate an array by/along dimensions & offsets.
--
-- >>> pretty $ rotates [1] [2] a
-- [[[8,9,10,11],
--   [0,1,2,3],
--   [4,5,6,7]],
--  [[20,21,22,23],
--   [12,13,14,15],
--   [16,17,18,19]]]
rotates ::
  (VG.Vector v a) =>
  Dims ->
  [Int] ->
  Array v a ->
  Array v a
rotates :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> [Int] -> Array v a -> Array v a
rotates [Int]
ds [Int]
rs Array v a
a = (Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
(Vector Int -> Vector Int)
-> (Vector Int -> Vector Int) -> Array v a -> Array v a
backpermute Vector Int -> Vector Int
forall a. a -> a
id (Vector Int -> Vector Int -> Vector Int -> Vector Int -> Vector Int
S.rotatesIndex ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds) ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
rs) (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a)) Array v a
a

-- | Generic vector sort helpers.
sortG :: (Ord a, VG.Vector v a, VG.Vector v Int) => v a -> v a
sortG :: forall a (v :: * -> *).
(Ord a, Vector v a, Vector v Int) =>
v a -> v a
sortG v a
a = v a -> v Int -> v a
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Vector v Int) =>
v a -> v Int -> v a
VG.backpermute v a
a (v a -> v Int
forall a (v :: * -> *).
(Ord a, Vector v a, Vector v Int) =>
v a -> v Int
orderG v a
a)

sortByG :: (Ord b, VG.Vector v a, VG.Vector v Int) => (a -> b) -> v a -> v a
sortByG :: forall b (v :: * -> *) a.
(Ord b, Vector v a, Vector v Int) =>
(a -> b) -> v a -> v a
sortByG a -> b
c v a
a = v a -> v Int -> v a
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Vector v Int) =>
v a -> v Int -> v a
VG.backpermute v a
a ((a -> b) -> v a -> v Int
forall b (v :: * -> *) a.
(Ord b, Vector v a, Vector v Int) =>
(a -> b) -> v a -> v Int
orderByG a -> b
c v a
a)

orderG :: (Ord a, VG.Vector v a, VG.Vector v Int) => v a -> v Int
orderG :: forall a (v :: * -> *).
(Ord a, Vector v a, Vector v Int) =>
v a -> v Int
orderG v a
a = (forall s. Mutable v s Int -> ST s ()) -> v Int -> v Int
forall (v :: * -> *) a.
Vector v a =>
(forall s. Mutable v s a -> ST s ()) -> v a -> v a
VG.modify ((Int -> Int -> Ordering)
-> Mutable v (PrimState (ST s)) Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> m ()
sortBy Int -> Int -> Ordering
comp) v Int
init0
  where
    comp :: Int -> Int -> Ordering
comp = (Int -> a) -> Int -> Int -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing ((Int -> a) -> Int -> Int -> Ordering)
-> (Int -> a) -> Int -> Int -> Ordering
forall a b. (a -> b) -> a -> b
$ v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
VG.unsafeIndex v a
a
    init0 :: v Int
init0 = Int -> (Int -> Int) -> v Int
forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a
VG.generate (v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length v a
a) Int -> Int
forall a. a -> a
id

orderByG :: (Ord b, VG.Vector v a, VG.Vector v Int) => (a -> b) -> v a -> v Int
orderByG :: forall b (v :: * -> *) a.
(Ord b, Vector v a, Vector v Int) =>
(a -> b) -> v a -> v Int
orderByG a -> b
c v a
a = (forall s. Mutable v s Int -> ST s ()) -> v Int -> v Int
forall (v :: * -> *) a.
Vector v a =>
(forall s. Mutable v s a -> ST s ()) -> v a -> v a
VG.modify ((Int -> Int -> Ordering)
-> Mutable v (PrimState (ST s)) Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> m ()
sortBy Int -> Int -> Ordering
comp) v Int
init0
  where
    comp :: Int -> Int -> Ordering
comp = (Int -> b) -> Int -> Int -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing ((Int -> b) -> Int -> Int -> Ordering)
-> (Int -> b) -> Int -> Int -> Ordering
forall a b. (a -> b) -> a -> b
$ a -> b
c (a -> b) -> (Int -> a) -> Int -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
VG.unsafeIndex v a
a
    init0 :: v Int
init0 = Int -> (Int -> Int) -> v Int
forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a
VG.generate (v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length v a
a) Int -> Int
forall a. a -> a
id

-- * sorting

-- | Sort an array along the supplied dimensions.
--
-- >>> sorts [0] (array [2,2] [2,3,1,4])
-- UnsafeArray [2,2] [1,4,2,3]
-- >>> sorts [1] (array [2,2] [2,3,1,4])
-- UnsafeArray [2,2] [2,3,1,4]
-- >>> sorts [0,1] (array [2,2] [2,3,1,4])
-- UnsafeArray [2,2] [1,2,3,4]
sorts :: (Ord (v a), VG.Vector v a, VG.Vector v Int, VG.Vector v (Array v a)) => Dims -> Array v a -> Array v a
sorts :: forall (v :: * -> *) a.
(Ord (v a), Vector v a, Vector v Int, Vector v (Array v a)) =>
[Int] -> Array v a -> Array v a
sorts [Int]
ds Array v a
a = [Int] -> Array v (Array v a) -> Array v a
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v (Array v a) -> Array v a
joins [Int]
ds (Array v (Array v a) -> Array v a)
-> Array v (Array v a) -> Array v a
forall a b. (a -> b) -> a -> b
$ (v (Array v a) -> v (Array v a))
-> Array v (Array v a) -> Array v (Array v a)
forall {k} {k} (v :: k -> *) (a :: k) (v' :: k -> *) (b :: k).
(v a -> v' b) -> Array v a -> Array v' b
unsafeModifyVector v (Array v a) -> v (Array v a)
forall a (v :: * -> *).
(Ord a, Vector v a, Vector v Int) =>
v a -> v a
sortG ([Int] -> Array v a -> Array v (Array v a)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int]
ds Array v a
a)

-- | The indices into the array if it were sorted by a comparison function along the dimensions supplied.
--
-- >>> import Data.Ord (Down (..))
-- >>> sortsBy [0] (fmapA Down) (array [2,2] [2,3,1,4])
-- UnsafeArray [2,2] [2,3,1,4]
sortsBy :: (Ord (v b), VG.Vector v a, VG.Vector v Int, VG.Vector v (Array v a)) => Dims -> (Array v a -> Array v b) -> Array v a -> Array v a
sortsBy :: forall (v :: * -> *) b a.
(Ord (v b), Vector v a, Vector v Int, Vector v (Array v a)) =>
[Int] -> (Array v a -> Array v b) -> Array v a -> Array v a
sortsBy [Int]
ds Array v a -> Array v b
c Array v a
a = [Int] -> Array v (Array v a) -> Array v a
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v (Array v a) -> Array v a
joins [Int]
ds (Array v (Array v a) -> Array v a)
-> Array v (Array v a) -> Array v a
forall a b. (a -> b) -> a -> b
$ (v (Array v a) -> v (Array v a))
-> Array v (Array v a) -> Array v (Array v a)
forall {k} {k} (v :: k -> *) (a :: k) (v' :: k -> *) (b :: k).
(v a -> v' b) -> Array v a -> Array v' b
unsafeModifyVector ((Array v a -> Array v b) -> v (Array v a) -> v (Array v a)
forall b (v :: * -> *) a.
(Ord b, Vector v a, Vector v Int) =>
(a -> b) -> v a -> v a
sortByG Array v a -> Array v b
c) ([Int] -> Array v a -> Array v (Array v a)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int]
ds Array v a
a)

-- | The indices into the array if it were sorted along the dimensions supplied.
--
-- >>> orders [0] (array [2,2] [2,3,1,4])
-- UnsafeArray [2] [1,0]
orders :: (Ord (v a), VG.Vector v Int, VG.Vector v a, VG.Vector v (Array v a)) => Dims -> Array v a -> Array v Int
orders :: forall (v :: * -> *) a.
(Ord (v a), Vector v Int, Vector v a, Vector v (Array v a)) =>
[Int] -> Array v a -> Array v Int
orders [Int]
ds Array v a
a = [Int] -> v Int -> Array v Int
forall {k} (v :: k -> *) (a :: k). [Int] -> v a -> Array v a
unsafeArrayL [v Int -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length v Int
v'] v Int
v'
  where
    v' :: v Int
v' = v (Array v a) -> v Int
forall a (v :: * -> *).
(Ord a, Vector v a, Vector v Int) =>
v a -> v Int
orderG (Array v (Array v a) -> v (Array v a)
forall {k} (v :: k -> *) (a :: k). Array v a -> v a
asVector ([Int] -> Array v a -> Array v (Array v a)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int]
ds Array v a
a))

-- | The indices into the array if it were sorted by a comparison function along the dimensions supplied.
--
-- >>> import Data.Ord (Down (..))
-- >>> ordersBy [0] (fmapA Down) (array [2,2] [2,3,1,4])
-- UnsafeArray [2] [0,1]
ordersBy :: (Ord (v b), VG.Vector v Int, VG.Vector v a, VG.Vector v (Array v a)) => Dims -> (Array v a -> Array v b) -> Array v a -> Array v Int
ordersBy :: forall (v :: * -> *) b a.
(Ord (v b), Vector v Int, Vector v a, Vector v (Array v a)) =>
[Int] -> (Array v a -> Array v b) -> Array v a -> Array v Int
ordersBy [Int]
ds Array v a -> Array v b
c Array v a
a = [Int] -> v Int -> Array v Int
forall {k} (v :: k -> *) (a :: k). [Int] -> v a -> Array v a
unsafeArrayL [v Int -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length v Int
v'] v Int
v'
  where
    v' :: v Int
v' = (Array v a -> Array v b) -> v (Array v a) -> v Int
forall b (v :: * -> *) a.
(Ord b, Vector v a, Vector v Int) =>
(a -> b) -> v a -> v Int
orderByG Array v a -> Array v b
c (Array v (Array v a) -> v (Array v a)
forall {k} (v :: k -> *) (a :: k). Array v a -> v a
asVector ([Int] -> Array v a -> Array v (Array v a)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int]
ds Array v a
a))

-- * transmission

-- | Apply a binary array function to two arrays with matching shapes across the supplied dimensions. No check on shapes.
--
-- >>> a = array [2,3] [0..5]
-- >>> b = array [3] [0..2]
-- >>> pretty $ telecasts [1] [0] (concatenate 0) a b
-- [[0,1,2],
--  [3,4,5],
--  [0,1,2]]
telecasts :: (VG.Vector v a, VG.Vector v b, VG.Vector v c, VG.Vector v (Array v a), VG.Vector v (Array v b), VG.Vector v (Array v c)) => Dims -> Dims -> (Array v a -> Array v b -> Array v c) -> Array v a -> Array v b -> Array v c
telecasts :: forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c, Vector v (Array v a),
 Vector v (Array v b), Vector v (Array v c)) =>
[Int]
-> [Int]
-> (Array v a -> Array v b -> Array v c)
-> Array v a
-> Array v b
-> Array v c
telecasts [Int]
dsa [Int]
dsb Array v a -> Array v b -> Array v c
f Array v a
a Array v b
b = (Array v a -> Array v b -> Array v c)
-> Array v (Array v a)
-> Array v (Array v b)
-> Array v (Array v c)
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v a -> Array v b -> Array v c
zipWith Array v a -> Array v b -> Array v c
f ([Int] -> Array v a -> Array v (Array v a)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int]
dsa Array v a
a) ([Int] -> Array v b -> Array v (Array v b)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int]
dsb Array v b
b) Array v (Array v c)
-> (Array v (Array v c) -> Array v c) -> Array v c
forall a b. a -> (a -> b) -> b
& [Int] -> Array v (Array v c) -> Array v c
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v (Array v a) -> Array v a
joins [Int]
dsa

-- | Apply a binary array function to two arrays with matching shapes across the supplied dimensions. Checks shape.
--
-- >>> a = array [2,3] [0..5]
-- >>> b = array [1] [1]
-- >>> telecastsSafe [0] [0] (zipWith (+)) a b
-- Nothing
telecastsSafe :: (VG.Vector v a, VG.Vector v b, VG.Vector v c, VG.Vector v (Array v a), VG.Vector v (Array v b), VG.Vector v (Array v c)) => Dims -> Dims -> (Array v a -> Array v b -> Array v c) -> Array v a -> Array v b -> Maybe (Array v c)
telecastsSafe :: forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c, Vector v (Array v a),
 Vector v (Array v b), Vector v (Array v c)) =>
[Int]
-> [Int]
-> (Array v a -> Array v b -> Array v c)
-> Array v a
-> Array v b
-> Maybe (Array v c)
telecastsSafe [Int]
dsa [Int]
dsb Array v a -> Array v b -> Array v c
f Array v a
a Array v b
b =
  Maybe (Array v c) -> Maybe (Array v c) -> Bool -> Maybe (Array v c)
forall a. a -> a -> Bool -> a
bool
    (Array v c -> Maybe (Array v c)
forall a. a -> Maybe a
Just (Array v c -> Maybe (Array v c)) -> Array v c -> Maybe (Array v c)
forall a b. (a -> b) -> a -> b
$ [Int]
-> [Int]
-> (Array v a -> Array v b -> Array v c)
-> Array v a
-> Array v b
-> Array v c
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c, Vector v (Array v a),
 Vector v (Array v b), Vector v (Array v c)) =>
[Int]
-> [Int]
-> (Array v a -> Array v b -> Array v c)
-> Array v a
-> Array v b
-> Array v c
telecasts [Int]
dsa [Int]
dsb Array v a -> Array v b -> Array v c
f Array v a
a Array v b
b)
    Maybe (Array v c)
forall a. Maybe a
Nothing
    (Array v (Array v a) -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape ([Int] -> Array v a -> Array v (Array v a)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int]
dsa Array v a
a) Vector Int -> Vector Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Array v (Array v b) -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape ([Int] -> Array v b -> Array v (Array v b)
forall (v :: * -> *) a.
(Vector v (Array v a), Vector v a) =>
[Int] -> Array v a -> Array v (Array v a)
extracts [Int]
dsb Array v b
b))

-- | Apply a binary array function to two arrays where the shape of the first array is a prefix of the second array. No checks on shape.
--
-- >>> a = array [2,3] [0..5]
-- >>> pretty $ transmit (zipWith (+)) (toScalar 1) a
-- [[1,2,3],
--  [4,5,6]]
transmit :: (VG.Vector v b, VG.Vector v c, VG.Vector v (Array v b), VG.Vector v (Array v c)) => (Array v a -> Array v b -> Array v c) -> Array v a -> Array v b -> Array v c
transmit :: forall (v :: * -> *) b c a.
(Vector v b, Vector v c, Vector v (Array v b),
 Vector v (Array v c)) =>
(Array v a -> Array v b -> Array v c)
-> Array v a -> Array v b -> Array v c
transmit Array v a -> Array v b -> Array v c
f Array v a
a Array v b
b = [Int] -> (Array v b -> Array v c) -> Array v b -> Array v c
forall (v :: * -> *) a b.
(Vector v a, Vector v b, Vector v (Array v a),
 Vector v (Array v b)) =>
[Int] -> (Array v a -> Array v b) -> Array v a -> Array v b
maps [Int]
ds (Array v a -> Array v b -> Array v c
f Array v a
a) Array v b
b
  where
    ds :: [Int]
ds = [(Array v a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank Array v a
a) .. (Array v b -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
rank Array v b
b Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)]

-- | Apply a binary array function to two arrays where the shape of the first array is a prefix of the second array. Checks shape.
--
-- >>> a = array [2,3] [0..5]
-- >>> transmitSafe (zipWith (+)) (array [3] [1,2,3]) a
-- Nothing
transmitSafe :: (VG.Vector v b, VG.Vector v c, VG.Vector v (Array v b), VG.Vector v (Array v c)) => (Array v a -> Array v b -> Array v c) -> Array v a -> Array v b -> Maybe (Array v c)
transmitSafe :: forall (v :: * -> *) b c a.
(Vector v b, Vector v c, Vector v (Array v b),
 Vector v (Array v c)) =>
(Array v a -> Array v b -> Array v c)
-> Array v a -> Array v b -> Maybe (Array v c)
transmitSafe Array v a -> Array v b -> Array v c
f Array v a
a Array v b
b = Maybe (Array v c) -> Maybe (Array v c) -> Bool -> Maybe (Array v c)
forall a. a -> a -> Bool -> a
bool Maybe (Array v c)
forall a. Maybe a
Nothing (Array v c -> Maybe (Array v c)
forall a. a -> Maybe a
Just (Array v c -> Maybe (Array v c)) -> Array v c -> Maybe (Array v c)
forall a b. (a -> b) -> a -> b
$ (Array v a -> Array v b -> Array v c)
-> Array v a -> Array v b -> Array v c
forall (v :: * -> *) b c a.
(Vector v b, Vector v c, Vector v (Array v b),
 Vector v (Array v c)) =>
(Array v a -> Array v b -> Array v c)
-> Array v a -> Array v b -> Array v c
transmit Array v a -> Array v b -> Array v c
f Array v a
a Array v b
b) (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a) [Int] -> [Int] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`List.isPrefixOf` Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v b -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v b
b))

-- | Transmit an operation if the first array is a prefix of the second or vice versa.
--
-- >>> pretty $ transmitOp (*) a (asArray [1,2])
-- [[[0,1,2,3],
--   [4,5,6,7],
--   [8,9,10,11]],
--  [[24,26,28,30],
--   [32,34,36,38],
--   [40,42,44,46]]]
transmitOp :: (VG.Vector v a, VG.Vector v b, VG.Vector v c, VG.Vector v (Array v a), VG.Vector v (Array v b), VG.Vector v (Array v c)) => (a -> b -> c) -> Array v a -> Array v b -> Array v c
transmitOp :: forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c, Vector v (Array v a),
 Vector v (Array v b), Vector v (Array v c)) =>
(a -> b -> c) -> Array v a -> Array v b -> Array v c
transmitOp a -> b -> c
f Array v a
a Array v b
b
  | Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a Vector Int -> Vector Int -> Bool
forall a. Eq a => a -> a -> Bool
== Array v b -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v b
b = (a -> b -> c) -> Array v a -> Array v b -> Array v c
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v a -> Array v b -> Array v c
zipWith a -> b -> c
f Array v a
a Array v b
b
  | Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a) [Int] -> [Int] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`List.isPrefixOf` Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v b -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v b
b) = (Array v a -> Array v b -> Array v c)
-> Array v a -> Array v b -> Array v c
forall (v :: * -> *) b c a.
(Vector v b, Vector v c, Vector v (Array v b),
 Vector v (Array v c)) =>
(Array v a -> Array v b -> Array v c)
-> Array v a -> Array v b -> Array v c
transmit ((a -> b -> c) -> Array v a -> Array v b -> Array v c
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v a -> Array v b -> Array v c
zipWith a -> b -> c
f) Array v a
a Array v b
b
  | Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v b -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v b
b) [Int] -> [Int] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`List.isPrefixOf` Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v a -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
shape Array v a
a) = (Array v b -> Array v a -> Array v c)
-> Array v b -> Array v a -> Array v c
forall (v :: * -> *) b c a.
(Vector v b, Vector v c, Vector v (Array v b),
 Vector v (Array v c)) =>
(Array v a -> Array v b -> Array v c)
-> Array v a -> Array v b -> Array v c
transmit ((b -> a -> c) -> Array v b -> Array v a -> Array v c
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v a -> Array v b -> Array v c
zipWith ((a -> b -> c) -> b -> a -> c
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> b -> c
f)) Array v b
b Array v a
a
  | Bool
otherwise = String -> Array v c
forall a. HasCallStack => String -> a
error String
"bad shapes"

-- | Vector specialisation of 'range'
--
-- >>> iota 5
-- UnsafeArray [5] [0,1,2,3,4]
iota :: (VG.Vector v Int) => Int -> Array v Int
iota :: forall (v :: * -> *). Vector v Int => Int -> Array v Int
iota Int
n = [Int] -> Array v Int
forall (v :: * -> *). Vector v Int => [Int] -> Array v Int
range [Int
n]

-- * row (first dimension) specializations

-- | Add a new row
--
-- >>> pretty $ cons (array [2] [0,1]) (array [2,2] [2,3,4,5])
-- [[0,1],
--  [2,3],
--  [4,5]]
cons :: (VG.Vector v a) => Array v a -> Array v a -> Array v a
cons :: forall (v :: * -> *) a.
Vector v a =>
Array v a -> Array v a -> Array v a
cons = Int -> Array v a -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
Int -> Array v a -> Array v a -> Array v a
prepend Int
0

-- | split an array into the first row and the remaining rows.
--
-- >>> uncons (array [3,2] [0..5])
-- (UnsafeArray [2] [0,1],UnsafeArray [2,2] [2,3,4,5])
uncons :: (VG.Vector v a) => Array v a -> (Array v a, Array v a)
uncons :: forall (v :: * -> *) a.
Vector v a =>
Array v a -> (Array v a, Array v a)
uncons Array v a
a = ([Int] -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
heads [Int
0] Array v a
a', [Int] -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
tails [Int
0] Array v a
a')
  where
    a' :: Array v a
a' = Array v a -> Array v a
forall {k} (v :: k -> *) (a :: k). Array v a -> Array v a
asSingleton Array v a
a

-- | Convenience pattern for row extraction and consolidation at the beginning of an Array.
--
-- >>> (x:<xs) = array [4] [0..3]
-- >>> x
-- UnsafeArray [] [0]
-- >>> xs
-- UnsafeArray [3] [1,2,3]
-- >>> (x:<xs)
-- UnsafeArray [4] [0,1,2,3]
pattern (:<) :: (VG.Vector v a) => Array v a -> Array v a -> Array v a
pattern x $b:< :: forall (v :: * -> *) a.
Vector v a =>
Array v a -> Array v a -> Array v a
$m:< :: forall {r} {v :: * -> *} {a}.
Vector v a =>
Array v a -> (Array v a -> Array v a -> r) -> ((# #) -> r) -> r
:< xs <- (uncons -> (x, xs))
  where
    Array v a
x :< Array v a
xs = Array v a -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
Array v a -> Array v a -> Array v a
cons Array v a
x Array v a
xs

infix 5 :<

-- | Add a new row at the end
--
-- >>> pretty $ snoc (array [2,2] [0,1,2,3]) (array [2] [4,5])
-- [[0,1],
--  [2,3],
--  [4,5]]
snoc :: (VG.Vector v a) => Array v a -> Array v a -> Array v a
snoc :: forall (v :: * -> *) a.
Vector v a =>
Array v a -> Array v a -> Array v a
snoc = Int -> Array v a -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
Int -> Array v a -> Array v a -> Array v a
append Int
0

-- | split an array into the initial rows and the last row.
--
-- >>> unsnoc (array [3,2] [0..5])
-- (UnsafeArray [2,2] [0,1,2,3],UnsafeArray [2] [4,5])
unsnoc :: (VG.Vector v a) => Array v a -> (Array v a, Array v a)
unsnoc :: forall (v :: * -> *) a.
Vector v a =>
Array v a -> (Array v a, Array v a)
unsnoc Array v a
a = ([Int] -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
inits [Int
0] Array v a
a', [Int] -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
[Int] -> Array v a -> Array v a
lasts [Int
0] Array v a
a')
  where
    a' :: Array v a
a' = Array v a -> Array v a
forall {k} (v :: k -> *) (a :: k). Array v a -> Array v a
asSingleton Array v a
a

-- | Convenience pattern for row extraction and consolidation at the end of an Array.
--
-- >>> (xs:>x) = array [4] [0..3]
-- >>> x
-- UnsafeArray [] [3]
-- >>> xs
-- UnsafeArray [3] [0,1,2]
-- >>> (xs:>x)
-- UnsafeArray [4] [0,1,2,3]
pattern (:>) :: (VG.Vector v a) => Array v a -> Array v a -> Array v a
pattern xs $b:> :: forall (v :: * -> *) a.
Vector v a =>
Array v a -> Array v a -> Array v a
$m:> :: forall {r} {v :: * -> *} {a}.
Vector v a =>
Array v a -> (Array v a -> Array v a -> r) -> ((# #) -> r) -> r
:> x <- (unsnoc -> (xs, x))
  where
    Array v a
xs :> Array v a
x = Array v a -> Array v a -> Array v a
forall (v :: * -> *) a.
Vector v a =>
Array v a -> Array v a -> Array v a
snoc Array v a
xs Array v a
x

infix 5 :>

-- * Math

-- | Generate an array of uniform random variates between a range.
--
-- >>> import System.Random.Stateful hiding (uniform)
-- >>> g <- newIOGenM (mkStdGen 42)
-- >>> u <- uniform g [2,3,4] (0,9 :: Int)
-- >>> pretty u
-- [[[0,7,0,2],
--   [1,7,4,2],
--   [5,9,8,2]],
--  [[9,8,1,0],
--   [2,2,8,2],
--   [2,8,0,6]]]
uniform :: (StatefulGen g m, UniformRange a, VG.Vector v a) => g -> [Int] -> (a, a) -> m (Array v a)
uniform :: forall g (m :: * -> *) a (v :: * -> *).
(StatefulGen g m, UniformRange a, Vector v a) =>
g -> [Int] -> (a, a) -> m (Array v a)
uniform g
g [Int]
ds (a, a)
r = do
  v <- Int -> m a -> m (v a)
forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a) =>
Int -> m a -> m (v a)
VG.replicateM (Vector Int -> Int
S.size ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
ds)) ((a, a) -> g -> m a
forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *). StatefulGen g m => (a, a) -> g -> m a
uniformRM (a, a)
r g
g)
  pure $ unsafeArray (VU.fromList ds) v