{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -Wno-redundant-constraints #-}

-- | Arrays with shape information and computations at a type-level.
module Harpie.Fixed.Generic
  ( -- * Usage
    -- $usage

    -- * Fixed Arrays
    Array (..),
    unsafeArray,
    validate,
    safeArray,
    array,
    unsafeModifyShape,
    unsafeModifyVector,

    -- * Dimensions
    Dim,
    pattern Dim,
    Dims,
    pattern Dims,

    -- * Conversion
    FromVector (..),
    toDynamic,
    with,

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

    -- * Indexing
    index,
    unsafeIndex,
    (!),
    (!?),
    tabulate,
    unsafeTabulate,
    backpermute,
    unsafeBackpermute,

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

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

    -- * Vector-level helpers
    fmapA,
    sumA,
    traverseA,

    -- * Element-level functions
    zipWith,
    modify,
    imap,

    -- * Function generalisers
    rowWise,
    colWise,

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

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

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

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

    -- * Search
    find,
    findNoOverlap,
    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
    telecasts,
    transmit,

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

    -- * Shape specializations
    Vector,
    vector,
    vector',
    iota,
    Matrix,

    -- * Math
    uniform,
  )
where

import Data.Bool
import Data.Foldable hiding (find, length, minimum)
import Data.List qualified as List
import Data.Maybe
import Data.Vector.Generic qualified as VG
import Data.Vector.Unboxed qualified as VU
import Fcf hiding (type (&&), type (+), type (++), type (-))
import Fcf qualified
import Fcf.Data.List
import GHC.Generics
import GHC.TypeNats
import Harpie.Array.Generic qualified as A
import Harpie.Shape hiding (asScalar, asSingleton, concatenate, range, rank, reorder, rerank, rotate, size, squeeze)
import Harpie.Shape qualified as S
import Harpie.Sort
import NumHask.Algebra.Action (AdditiveAction (..), DivisiveAction (..), MultiplicativeAction (..), SubtractiveAction (..))
import NumHask.Algebra.Additive (Additive (..), Subtractive (..))
import NumHask.Algebra.Additive qualified as Add
import NumHask.Algebra.Lattice (JoinSemiLattice (..), MeetSemiLattice (..))
import NumHask.Algebra.Metric (Epsilon (..))
import NumHask.Algebra.Multiplicative (Divisive (..), Multiplicative (..))
import NumHask.Algebra.Multiplicative qualified as Mult
import NumHask.Data.Integral (FromInteger (..))
import NumHask.Data.Rational (FromRational (..))
import Prettyprinter hiding (dot, fill)
import System.Random hiding (uniform)
import System.Random.Stateful hiding (uniform)
import Unsafe.Coerce
import Prelude as P hiding (cycle, drop, fromInteger, fromRational, length, negate, repeat, sequence, take, zipWith, (*), (+), (-), (/))

-- $setup
--
-- >>> :m -Prelude
-- >>> :set -XDataKinds
-- >>> :set -Wno-type-defaults
-- >>> :set -Wno-name-shadowing
-- >>> import Prelude hiding (cycle, repeat, take, drop, zipWith, length)
-- >>> import Harpie.Fixed.Generic
-- >>> import Harpie.Shape qualified as S
-- >>> import Harpie.Shape (SNats, Fin (..), Fins (..))
-- >>> import GHC.TypeNats
-- >>> import Data.Vector qualified as Vec
-- >>> import Data.List qualified as List
-- >>> import Prettyprinter hiding (dot,fill)
-- >>> import Data.Functor.Rep hiding (index, tabulate)
-- >>> s = array @Vec.Vector @'[] @Int [1]
-- >>> let v = range @Vec.Vector @'[3]
-- >>> let m = range @Vec.Vector @[2,3]
-- >>> a = range @Vec.Vector @[2,3,4]

-- $usage
--
-- >>> :set -XDataKinds
--
-- 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.Fixed@ 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.Fixed.Generic as F
-- >>> 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)
--
-- The 'Representable' class from [@adjunctions@](https://hackage.haskell.org/package/adjunctions) is used heavily by the module.
--
-- >>> import Data.Functor.Rep hiding (index, tabulate)
--
-- An important base accounting of 'Array' shape is the singleton types 'SNat' (a type-level 'Natural' or 'Nat') from [GHC.TypeNats](https://hackage.haskell.org/package/base/docs/GHC-TypeNats.html) in base.
--
-- >>> import GHC.TypeNats
--
-- The [first-class-families](https://hackage.haskell.org/package/first-class-families) library was used to code most of type-level constraint logic.
--
-- >>> import Fcf qualified
--
-- Examples of arrays:
--
-- An array with no dimensions (a scalar).
--
-- >>> s = array @Vec.Vector @'[] @Int [1]
-- >>> s
-- [1]
-- >>> shape s
-- []
-- >>> pretty s
-- 1
--
-- A single-dimension array (a vector).
--
-- >>> let v = range @Vec.Vector @'[3]
-- >>> pretty v
-- [0,1,2]
--
-- A two-dimensional array (a matrix).
--
-- >>> let m = range @Vec.Vector @[2,3]
-- >>> pretty m
-- [[0,1,2],
--  [3,4,5]]
--
-- An n-dimensional array (n should be finite).
--
-- >>> a = range @Vec.Vector @[2,3,4]
-- >>> 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]
-- >>> 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]]]
--
-- Conversion to a dynamic, value-level shaped 'Harpie.Array.Array'
--
-- >>> toDynamic 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]

-- | A hyperrectangular (or multidimensional) array with a type-level shape.
--
-- >>> array @Vec.Vector @[2,3,4] @Int [1..24]
-- [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24]
-- >>> array [1..24] :: Array Vec.Vector '[2,3,4] Int
-- [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 (array @Vec.Vector @[2,3,4] @Int [1..24])
-- [[[1,2,3,4],
--   [5,6,7,8],
--   [9,10,11,12]],
--  [[13,14,15,16],
--   [17,18,19,20],
--   [21,22,23,24]]]
--
-- >>> array [1,2,3] :: Array Vec.Vector '[2,2] Int
-- *** Exception: Shape Mismatch
-- ...
--
-- In many situations, the use of  [TypeApplication](https://ghc.gitlab.haskell.org/ghc/doc/users_guide/exts/type_applications.html) can lead to a clean coding style.
--
-- >>> array @Vec.Vector @[2,3] @Int [1..6]
-- [1,2,3,4,5,6]
--
-- The main computational entry and exit points are often via 'index' and 'tabulate' with arrays indexed by 'Fins':
--
-- >>> index a (S.UnsafeFins [1,2,3])
-- 23
--
-- >>> :t tabulate id :: Array Vec.Vector [2,3] (Fins [2,3])
-- tabulate id :: Array Vec.Vector [2,3] (Fins [2,3])
--   :: Array Vec.Vector [2, 3] (Fins [2, 3])
-- >>> pretty (tabulate id :: Array Vec.Vector [2,3] (Fins [2,3]))
-- [[[0,0],[0,1],[0,2]],
--  [[1,0],[1,1],[1,2]]]
type role Array nominal nominal nominal

newtype Array v (s :: [Nat]) a = Array (v a)
  deriving stock ((forall a b. (a -> b) -> Array v s a -> Array v s b)
-> (forall a b. a -> Array v s b -> Array v s a)
-> Functor (Array v s)
forall a b. a -> Array v s b -> Array v s a
forall a b. (a -> b) -> Array v s a -> Array v s b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
forall (v :: * -> *) (s :: [Nat]) a b.
Functor v =>
a -> Array v s b -> Array v s a
forall (v :: * -> *) (s :: [Nat]) a b.
Functor v =>
(a -> b) -> Array v s a -> Array v s b
$cfmap :: forall (v :: * -> *) (s :: [Nat]) a b.
Functor v =>
(a -> b) -> Array v s a -> Array v s b
fmap :: forall a b. (a -> b) -> Array v s a -> Array v s b
$c<$ :: forall (v :: * -> *) (s :: [Nat]) a b.
Functor v =>
a -> Array v s b -> Array v s a
<$ :: forall a b. a -> Array v s b -> Array v s a
Functor, (forall m. Monoid m => Array v s m -> m)
-> (forall m a. Monoid m => (a -> m) -> Array v s a -> m)
-> (forall m a. Monoid m => (a -> m) -> Array v s a -> m)
-> (forall a b. (a -> b -> b) -> b -> Array v s a -> b)
-> (forall a b. (a -> b -> b) -> b -> Array v s a -> b)
-> (forall b a. (b -> a -> b) -> b -> Array v s a -> b)
-> (forall b a. (b -> a -> b) -> b -> Array v s a -> b)
-> (forall a. (a -> a -> a) -> Array v s a -> a)
-> (forall a. (a -> a -> a) -> Array v s a -> a)
-> (forall a. Array v s a -> [a])
-> (forall a. Array v s a -> Bool)
-> (forall a. Array v s a -> Int)
-> (forall a. Eq a => a -> Array v s a -> Bool)
-> (forall a. Ord a => Array v s a -> a)
-> (forall a. Ord a => Array v s a -> a)
-> (forall a. Num a => Array v s a -> a)
-> (forall a. Num a => Array v s a -> a)
-> Foldable (Array v s)
forall a. Eq a => a -> Array v s a -> Bool
forall a. Num a => Array v s a -> a
forall a. Ord a => Array v s a -> a
forall m. Monoid m => Array v s m -> m
forall a. Array v s a -> Bool
forall a. Array v s a -> Int
forall a. Array v s a -> [a]
forall a. (a -> a -> a) -> Array v s a -> a
forall m a. Monoid m => (a -> m) -> Array v s a -> m
forall b a. (b -> a -> b) -> b -> Array v s a -> b
forall a b. (a -> b -> b) -> b -> Array v s a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
forall (v :: * -> *) (s :: [Nat]) a.
(Foldable v, Eq a) =>
a -> Array v s a -> Bool
forall (v :: * -> *) (s :: [Nat]) a.
(Foldable v, Num a) =>
Array v s a -> a
forall (v :: * -> *) (s :: [Nat]) a.
(Foldable v, Ord a) =>
Array v s a -> a
forall (v :: * -> *) (s :: [Nat]) m.
(Foldable v, Monoid m) =>
Array v s m -> m
forall (v :: * -> *) (s :: [Nat]) a.
Foldable v =>
Array v s a -> Bool
forall (v :: * -> *) (s :: [Nat]) a.
Foldable v =>
Array v s a -> Int
forall (v :: * -> *) (s :: [Nat]) a.
Foldable v =>
Array v s a -> [a]
forall (v :: * -> *) (s :: [Nat]) a.
Foldable v =>
(a -> a -> a) -> Array v s a -> a
forall (v :: * -> *) (s :: [Nat]) m a.
(Foldable v, Monoid m) =>
(a -> m) -> Array v s a -> m
forall (v :: * -> *) (s :: [Nat]) b a.
Foldable v =>
(b -> a -> b) -> b -> Array v s a -> b
forall (v :: * -> *) (s :: [Nat]) a b.
Foldable v =>
(a -> b -> b) -> b -> Array v s a -> b
$cfold :: forall (v :: * -> *) (s :: [Nat]) m.
(Foldable v, Monoid m) =>
Array v s m -> m
fold :: forall m. Monoid m => Array v s m -> m
$cfoldMap :: forall (v :: * -> *) (s :: [Nat]) m a.
(Foldable v, Monoid m) =>
(a -> m) -> Array v s a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> Array v s a -> m
$cfoldMap' :: forall (v :: * -> *) (s :: [Nat]) m a.
(Foldable v, Monoid m) =>
(a -> m) -> Array v s a -> m
foldMap' :: forall m a. Monoid m => (a -> m) -> Array v s a -> m
$cfoldr :: forall (v :: * -> *) (s :: [Nat]) a b.
Foldable v =>
(a -> b -> b) -> b -> Array v s a -> b
foldr :: forall a b. (a -> b -> b) -> b -> Array v s a -> b
$cfoldr' :: forall (v :: * -> *) (s :: [Nat]) a b.
Foldable v =>
(a -> b -> b) -> b -> Array v s a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> Array v s a -> b
$cfoldl :: forall (v :: * -> *) (s :: [Nat]) b a.
Foldable v =>
(b -> a -> b) -> b -> Array v s a -> b
foldl :: forall b a. (b -> a -> b) -> b -> Array v s a -> b
$cfoldl' :: forall (v :: * -> *) (s :: [Nat]) b a.
Foldable v =>
(b -> a -> b) -> b -> Array v s a -> b
foldl' :: forall b a. (b -> a -> b) -> b -> Array v s a -> b
$cfoldr1 :: forall (v :: * -> *) (s :: [Nat]) a.
Foldable v =>
(a -> a -> a) -> Array v s a -> a
foldr1 :: forall a. (a -> a -> a) -> Array v s a -> a
$cfoldl1 :: forall (v :: * -> *) (s :: [Nat]) a.
Foldable v =>
(a -> a -> a) -> Array v s a -> a
foldl1 :: forall a. (a -> a -> a) -> Array v s a -> a
$ctoList :: forall (v :: * -> *) (s :: [Nat]) a.
Foldable v =>
Array v s a -> [a]
toList :: forall a. Array v s a -> [a]
$cnull :: forall (v :: * -> *) (s :: [Nat]) a.
Foldable v =>
Array v s a -> Bool
null :: forall a. Array v s a -> Bool
$clength :: forall (v :: * -> *) (s :: [Nat]) a.
Foldable v =>
Array v s a -> Int
length :: forall a. Array v s a -> Int
$celem :: forall (v :: * -> *) (s :: [Nat]) a.
(Foldable v, Eq a) =>
a -> Array v s a -> Bool
elem :: forall a. Eq a => a -> Array v s a -> Bool
$cmaximum :: forall (v :: * -> *) (s :: [Nat]) a.
(Foldable v, Ord a) =>
Array v s a -> a
maximum :: forall a. Ord a => Array v s a -> a
$cminimum :: forall (v :: * -> *) (s :: [Nat]) a.
(Foldable v, Ord a) =>
Array v s a -> a
minimum :: forall a. Ord a => Array v s a -> a
$csum :: forall (v :: * -> *) (s :: [Nat]) a.
(Foldable v, Num a) =>
Array v s a -> a
sum :: forall a. Num a => Array v s a -> a
$cproduct :: forall (v :: * -> *) (s :: [Nat]) a.
(Foldable v, Num a) =>
Array v s a -> a
product :: forall a. Num a => Array v s a -> a
Foldable, (forall x. Array v s a -> Rep (Array v s a) x)
-> (forall x. Rep (Array v s a) x -> Array v s a)
-> Generic (Array v s a)
forall x. Rep (Array v s a) x -> Array v s a
forall x. Array v s a -> Rep (Array v s a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall k (v :: k -> *) (s :: [Nat]) (a :: k) x.
Rep (Array v s a) x -> Array v s a
forall k (v :: k -> *) (s :: [Nat]) (a :: k) x.
Array v s a -> Rep (Array v s a) x
$cfrom :: forall k (v :: k -> *) (s :: [Nat]) (a :: k) x.
Array v s a -> Rep (Array v s a) x
from :: forall x. Array v s a -> Rep (Array v s a) x
$cto :: forall k (v :: k -> *) (s :: [Nat]) (a :: k) x.
Rep (Array v s a) x -> Array v s a
to :: forall x. Rep (Array v s a) x -> Array v s a
Generic, Functor (Array v s)
Foldable (Array v s)
(Functor (Array v s), Foldable (Array v s)) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> Array v s a -> f (Array v s b))
-> (forall (f :: * -> *) a.
    Applicative f =>
    Array v s (f a) -> f (Array v s a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> Array v s a -> m (Array v s b))
-> (forall (m :: * -> *) a.
    Monad m =>
    Array v s (m a) -> m (Array v s a))
-> Traversable (Array v s)
forall (t :: * -> *).
(Functor t, Foldable t) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (v :: * -> *) (s :: [Nat]).
Traversable v =>
Functor (Array v s)
forall (v :: * -> *) (s :: [Nat]).
Traversable v =>
Foldable (Array v s)
forall (v :: * -> *) (s :: [Nat]) (m :: * -> *) a.
(Traversable v, Monad m) =>
Array v s (m a) -> m (Array v s a)
forall (v :: * -> *) (s :: [Nat]) (f :: * -> *) a.
(Traversable v, Applicative f) =>
Array v s (f a) -> f (Array v s a)
forall (v :: * -> *) (s :: [Nat]) (m :: * -> *) a b.
(Traversable v, Monad m) =>
(a -> m b) -> Array v s a -> m (Array v s b)
forall (v :: * -> *) (s :: [Nat]) (f :: * -> *) a b.
(Traversable v, Applicative f) =>
(a -> f b) -> Array v s a -> f (Array v s b)
forall (m :: * -> *) a.
Monad m =>
Array v s (m a) -> m (Array v s a)
forall (f :: * -> *) a.
Applicative f =>
Array v s (f a) -> f (Array v s a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Array v s a -> m (Array v s b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Array v s a -> f (Array v s b)
$ctraverse :: forall (v :: * -> *) (s :: [Nat]) (f :: * -> *) a b.
(Traversable v, Applicative f) =>
(a -> f b) -> Array v s a -> f (Array v s b)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Array v s a -> f (Array v s b)
$csequenceA :: forall (v :: * -> *) (s :: [Nat]) (f :: * -> *) a.
(Traversable v, Applicative f) =>
Array v s (f a) -> f (Array v s a)
sequenceA :: forall (f :: * -> *) a.
Applicative f =>
Array v s (f a) -> f (Array v s a)
$cmapM :: forall (v :: * -> *) (s :: [Nat]) (m :: * -> *) a b.
(Traversable v, Monad m) =>
(a -> m b) -> Array v s a -> m (Array v s b)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Array v s a -> m (Array v s b)
$csequence :: forall (v :: * -> *) (s :: [Nat]) (m :: * -> *) a.
(Traversable v, Monad m) =>
Array v s (m a) -> m (Array v s a)
sequence :: forall (m :: * -> *) a.
Monad m =>
Array v s (m a) -> m (Array v s a)
Traversable)
  deriving newtype (Array v s a -> Array v s a -> Bool
(Array v s a -> Array v s a -> Bool)
-> (Array v s a -> Array v s a -> Bool) -> Eq (Array v s a)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Eq (v a) =>
Array v s a -> Array v s a -> Bool
$c== :: forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Eq (v a) =>
Array v s a -> Array v s a -> Bool
== :: Array v s a -> Array v s a -> Bool
$c/= :: forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Eq (v a) =>
Array v s a -> Array v s a -> Bool
/= :: Array v s a -> Array v s a -> Bool
Eq, Eq (Array v s a)
Eq (Array v s a) =>
(Array v s a -> Array v s a -> Ordering)
-> (Array v s a -> Array v s a -> Bool)
-> (Array v s a -> Array v s a -> Bool)
-> (Array v s a -> Array v s a -> Bool)
-> (Array v s a -> Array v s a -> Bool)
-> (Array v s a -> Array v s a -> Array v s a)
-> (Array v s a -> Array v s a -> Array v s a)
-> Ord (Array v s a)
Array v s a -> Array v s a -> Bool
Array v s a -> Array v s a -> Ordering
Array v s a -> Array v s a -> Array v s a
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Ord (v a) =>
Eq (Array v s a)
forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Ord (v a) =>
Array v s a -> Array v s a -> Bool
forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Ord (v a) =>
Array v s a -> Array v s a -> Ordering
forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Ord (v a) =>
Array v s a -> Array v s a -> Array v s a
$ccompare :: forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Ord (v a) =>
Array v s a -> Array v s a -> Ordering
compare :: Array v s a -> Array v s a -> Ordering
$c< :: forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Ord (v a) =>
Array v s a -> Array v s a -> Bool
< :: Array v s a -> Array v s a -> Bool
$c<= :: forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Ord (v a) =>
Array v s a -> Array v s a -> Bool
<= :: Array v s a -> Array v s a -> Bool
$c> :: forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Ord (v a) =>
Array v s a -> Array v s a -> Bool
> :: Array v s a -> Array v s a -> Bool
$c>= :: forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Ord (v a) =>
Array v s a -> Array v s a -> Bool
>= :: Array v s a -> Array v s a -> Bool
$cmax :: forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Ord (v a) =>
Array v s a -> Array v s a -> Array v s a
max :: Array v s a -> Array v s a -> Array v s a
$cmin :: forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Ord (v a) =>
Array v s a -> Array v s a -> Array v s a
min :: Array v s a -> Array v s a -> Array v s a
Ord, Int -> Array v s a -> ShowS
[Array v s a] -> ShowS
Array v s a -> String
(Int -> Array v s a -> ShowS)
-> (Array v s a -> String)
-> ([Array v s a] -> ShowS)
-> Show (Array v s a)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Show (v a) =>
Int -> Array v s a -> ShowS
forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Show (v a) =>
[Array v s a] -> ShowS
forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Show (v a) =>
Array v s a -> String
$cshowsPrec :: forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Show (v a) =>
Int -> Array v s a -> ShowS
showsPrec :: Int -> Array v s a -> ShowS
$cshow :: forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Show (v a) =>
Array v s a -> String
show :: Array v s a -> String
$cshowList :: forall k (v :: k -> *) (s :: [Nat]) (a :: k).
Show (v a) =>
[Array v s a] -> ShowS
showList :: [Array v s a] -> ShowS
Show)

instance
  ( Additive a,
    KnownNats s,
    VG.Vector v a
  ) =>
  Additive (Array v s a)
  where
  + :: Array v s a -> Array v s a -> Array v s a
(+) = (a -> a -> a) -> Array v s a -> Array v s a -> Array v s a
forall (s :: [Nat]) (v :: * -> *) a b c.
(KnownNats s, Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v s a -> Array v s b -> Array v s c
zipWith a -> a -> a
forall a. Additive a => a -> a -> a
(+)
  zero :: Array v s a
zero = a -> Array v s a
forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
a -> Array v s a
konst a
forall a. Additive a => a
zero

instance
  ( Subtractive a,
    KnownNats s,
    VG.Vector v a
  ) =>
  Subtractive (Array v s a)
  where
  negate :: Array v s a -> Array v s a
negate = (a -> a) -> Array v s a -> Array v s a
forall (v :: * -> *) (s :: [Nat]) a b.
(KnownNats s, Vector v a, Vector v b) =>
(a -> b) -> Array v s a -> Array v s b
fmapA a -> a
forall a. Subtractive a => a -> a
negate

instance
  ( Multiplicative a,
    KnownNats s,
    VG.Vector v a
  ) =>
  MultiplicativeAction (Array v s a)
  where
  type Scalar (Array v s a) = a
  |* :: Array v s a -> Scalar (Array v s a) -> Array v s a
(|*) Array v s a
r Scalar (Array v s a)
s = (a -> a) -> Array v s a -> Array v s a
forall (v :: * -> *) (s :: [Nat]) a b.
(KnownNats s, Vector v a, Vector v b) =>
(a -> b) -> Array v s a -> Array v s b
fmapA (a
Scalar (Array v s a)
s a -> a -> a
forall a. Multiplicative a => a -> a -> a
*) Array v s a
r

instance
  ( Additive a,
    KnownNats s,
    VG.Vector v a
  ) =>
  AdditiveAction (Array v s a)
  where
  type AdditiveScalar (Array v s a) = a
  |+ :: Array v s a -> AdditiveScalar (Array v s a) -> Array v s a
(|+) Array v s a
r AdditiveScalar (Array v s a)
s = (a -> a) -> Array v s a -> Array v s a
forall (v :: * -> *) (s :: [Nat]) a b.
(KnownNats s, Vector v a, Vector v b) =>
(a -> b) -> Array v s a -> Array v s b
fmapA (a
AdditiveScalar (Array v s a)
s a -> a -> a
forall a. Additive a => a -> a -> a
+) Array v s a
r

instance
  ( Subtractive a,
    KnownNats s,
    VG.Vector v a
  ) =>
  SubtractiveAction (Array v s a)
  where
  |- :: Array v s a -> AdditiveScalar (Array v s a) -> Array v s a
(|-) Array v s a
r AdditiveScalar (Array v s a)
s = (a -> a) -> Array v s a -> Array v s a
forall (v :: * -> *) (s :: [Nat]) a b.
(KnownNats s, Vector v a, Vector v b) =>
(a -> b) -> Array v s a -> Array v s b
fmapA (\a
x -> a
x a -> a -> a
forall a. Subtractive a => a -> a -> a
- a
AdditiveScalar (Array v s a)
s) Array v s a
r

instance
  ( Divisive a,
    KnownNats s,
    VG.Vector v a
  ) =>
  DivisiveAction (Array v s a)
  where
  |/ :: Array v s a -> Scalar (Array v s a) -> Array v s a
(|/) Array v s a
r Scalar (Array v s a)
s = (Scalar (Array v s a) -> a)
-> Array v s (Scalar (Array v s a)) -> Array v s a
forall (v :: * -> *) (s :: [Nat]) a b.
(KnownNats s, Vector v a, Vector v b) =>
(a -> b) -> Array v s a -> Array v s b
fmapA (Scalar (Array v s a)
-> Scalar (Array v s a) -> Scalar (Array v s a)
forall a. Divisive a => a -> a -> a
/ Scalar (Array v s a)
s) Array v s a
Array v s (Scalar (Array v s a))
r

instance
  ( KnownNats s,
    JoinSemiLattice a,
    VG.Vector v a,
    Eq (v a)
  ) =>
  JoinSemiLattice (Array v s a)
  where
  \/ :: Array v s a -> Array v s a -> Array v s a
(\/) = (a -> a -> a) -> Array v s a -> Array v s a -> Array v s a
forall (s :: [Nat]) (v :: * -> *) a b c.
(KnownNats s, Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v s a -> Array v s b -> Array v s c
zipWith a -> a -> a
forall a. JoinSemiLattice a => a -> a -> a
(\/)

instance
  ( KnownNats s,
    MeetSemiLattice a,
    VG.Vector v a,
    Eq (v a)
  ) =>
  MeetSemiLattice (Array v s a)
  where
  /\ :: Array v s a -> Array v s a -> Array v s a
(/\) = (a -> a -> a) -> Array v s a -> Array v s a -> Array v s a
forall (s :: [Nat]) (v :: * -> *) a b c.
(KnownNats s, Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v s a -> Array v s b -> Array v s c
zipWith a -> a -> a
forall a. MeetSemiLattice a => a -> a -> a
(/\)

instance
  ( KnownNats s,
    Subtractive a,
    Epsilon a,
    VG.Vector v a,
    Eq (v a)
  ) =>
  Epsilon (Array v s a)
  where
  epsilon :: Array v s a
epsilon = a -> Array v s a
forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
a -> Array v s a
konst a
forall a. Epsilon a => a
epsilon

instance
  ( FromInteger a,
    VG.Vector v a
  ) =>
  FromInteger (Array v ('[] :: [Nat]) a)
  where
  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. FromInteger a => Integer -> a
fromInteger Integer
x)

instance
  ( FromRational a,
    VG.Vector v a
  ) =>
  FromRational (Array v ('[] :: [Nat]) a)
  where
  fromRational :: Rational -> Array v '[] a
fromRational Rational
x = a -> Array v '[] a
forall (v :: * -> *) a. Vector v a => a -> Array v '[] a
toScalar (Rational -> a
forall a. FromRational a => Rational -> a
fromRational Rational
x)

instance (KnownNats s, Show a, Show (v a), VG.Vector v a, VG.Vector v (A.Array v a)) => Pretty (Array v s a) where
  pretty :: forall ann. Array v s a -> Doc ann
pretty = 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 s a -> Array v a) -> Array v s a -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array v s a -> Array v a
forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
Array v s a -> Array v a
toDynamic

-- | Tabulate an array from a function on 'Fins'.
tabulate :: forall v s a. (KnownNats s, VG.Vector v a) => (Fins s -> a) -> Array v s a
tabulate :: forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
(Fins s -> a) -> Array v s a
tabulate Fins s -> a
f =
  v a -> Array v s a
forall {k} (v :: k -> *) (s :: [Nat]) (a :: k). v a -> Array v s a
Array (v a -> Array v s a) -> v a -> Array v s a
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> a) -> v a
forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a
VG.generate (Vector Int -> Int
S.size ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s)) (Fins s -> a
f (Fins s -> a) -> (Int -> Fins s) -> Int -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> Fins s
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins ([Int] -> Fins s) -> (Int -> [Int]) -> Int -> Fins s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Vector Int -> [Int]) -> (Int -> Vector Int) -> Int -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Int -> Vector Int
S.shapen ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s))
  where
    s :: [Int]
s = forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @s
{-# INLINE tabulate #-}

-- | Index into an array using 'Fins'.
index :: forall v s a. (KnownNats s, VG.Vector v a) => Array v s a -> Fins s -> a
index :: forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index (Array v a
v) Fins s
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.flatten ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s) ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList (Fins s -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins Fins s
i)))
  where
    s :: [Int]
s = forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @s
{-# INLINE index #-}

-- | Conversion to and from a `V.Vector`
--
-- Note that conversion of an 'Array' to a vector drops shape information, so that:
--
-- > vectorAs . asVector == id
-- > asVector . vectorAs == 'flat'
--
-- >>> asVector (range @Vec.Vector @[2,3]) :: Vec.Vector Int
-- [0,1,2,3,4,5]
--
-- > vectorAs (VG.fromList [0..5]) :: Array Vec.Vector [2,3] Int
-- [0,1,2,3,4,5]
class (VG.Vector v a) => FromVector t v a | t -> a where
  asVector :: t -> v a
  vectorAs :: v a -> t

instance {-# INCOHERENT #-} (VG.Vector v a) => FromVector (v a) v a where
  asVector :: v a -> v a
asVector = v a -> v a
forall a. a -> a
id
  vectorAs :: v a -> v a
vectorAs = v a -> v a
forall a. a -> a
id

instance {-# INCOHERENT #-} (VG.Vector v a) => FromVector [a] v a where
  asVector :: [a] -> v a
asVector = [a] -> v a
forall (v :: * -> *) a. Vector v a => [a] -> v a
VG.fromList
  vectorAs :: v a -> [a]
vectorAs = v a -> [a]
forall (v :: * -> *) a. Vector v a => v a -> [a]
VG.toList

instance {-# INCOHERENT #-} (VG.Vector v a) => FromVector (Array v s a) v a where
  asVector :: Array v s a -> v a
asVector (Array v a
v) = v a
v
  vectorAs :: v a -> Array v s a
vectorAs v a
v = v a -> Array v s a
forall {k} (v :: k -> *) (s :: [Nat]) (a :: k). v a -> Array v s a
Array v a
v

-- | Construct an array without shape validation.
--
-- >>> unsafeArray [0..4] :: Array Vec.Vector [2,3] Int
-- [0,1,2,3,4]
unsafeArray :: (KnownNats s, FromVector t v a, VG.Vector v a) => t -> Array v s a
unsafeArray :: forall (s :: [Nat]) t (v :: * -> *) a.
(KnownNats s, FromVector t v a, Vector v a) =>
t -> Array v s a
unsafeArray (t -> v a
forall t (v :: * -> *) a. FromVector t v a => t -> v a
asVector -> v a
v) = v a -> Array v s a
forall {k} (v :: k -> *) (s :: [Nat]) (a :: k). v a -> Array v s a
Array v a
v

-- | Validate the size and shape of an array.
--
-- >>> validate (unsafeArray [0..4] :: Array Vec.Vector [2,3] Int)
-- False
validate :: (KnownNats s, VG.Vector v a) => Array v s a -> Bool
validate :: forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
Array v s a -> Bool
validate a :: Array v s a
a@(Array v a
v) = Array v s a -> Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Int
size Array v s 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 v a
v

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

-- | Construct an Array, throwing an exception on a bad shape.
--
-- >>> array [0..22] :: Array Vec.Vector [2,3,4] Int
-- *** Exception: Shape Mismatch
-- ...
array :: forall v s a t. (KnownNats s, FromVector t v a, VG.Vector v a) => t -> Array v s a
array :: forall (v :: * -> *) (s :: [Nat]) a t.
(KnownNats s, FromVector t v a, Vector v a) =>
t -> Array v s a
array t
v =
  Array v s a -> Maybe (Array v s a) -> Array v s a
forall a. a -> Maybe a -> a
fromMaybe (String -> Array v s a
forall a. HasCallStack => String -> a
error String
"Shape Mismatch") (t -> Maybe (Array v s a)
forall (s :: [Nat]) t (v :: * -> *) a.
(KnownNats s, FromVector t v a, Vector v a) =>
t -> Maybe (Array v s a)
safeArray t
v)

-- | Unsafely modify an array shape.
--
-- >>> pretty (unsafeModifyShape @Vec.Vector @[3,2] (array @Vec.Vector @[2,3] @Int [0..5]))
-- [[0,1],
--  [2,3],
--  [4,5]]
unsafeModifyShape :: forall v s' s a. (KnownNats s, KnownNats s', VG.Vector v a) => Array v s a -> Array v s' a
unsafeModifyShape :: forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
Array v s a -> Array v s' a
unsafeModifyShape (Array v a
v) = v a -> Array v s' a
forall {k} (v :: k -> *) (s :: [Nat]) (a :: k). v a -> Array v s a
Array v a
v

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

-- Helpers that operate directly on the underlying vector, avoiding
-- Functor/Foldable/Traversable constraints on the vector type.

fmapA :: forall v s a b. (KnownNats s, VG.Vector v a, VG.Vector v b) => (a -> b) -> Array v s a -> Array v s b
fmapA :: forall (v :: * -> *) (s :: [Nat]) a b.
(KnownNats s, Vector v a, Vector v b) =>
(a -> b) -> Array v s a -> Array v s b
fmapA a -> b
f = (v a -> v b) -> Array v s a -> Array v s b
forall (v :: * -> *) (s :: [Nat]) a b.
(KnownNats s, Vector v a, Vector v b) =>
(v a -> v b) -> Array v s a -> Array v s 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)

toListA :: forall v s a. (VG.Vector v a) => Array v s a -> [a]
toListA :: forall (v :: * -> *) (s :: [Nat]) a.
Vector v a =>
Array v s a -> [a]
toListA (Array v a
v) = v a -> [a]
forall (v :: * -> *) a. Vector v a => v a -> [a]
VG.toList v a
v

sumA :: forall v s a. (Add.Additive a, VG.Vector v a) => Array v s a -> a
sumA :: forall (v :: * -> *) (s :: [Nat]) a.
(Additive a, Vector v a) =>
Array v s a -> a
sumA (Array v a
v) = (a -> a -> a) -> a -> v a -> a
forall (v :: * -> *) a b.
Vector v a =>
(a -> b -> b) -> b -> v a -> b
VG.foldr a -> a -> a
forall a. Additive a => a -> a -> a
(Add.+) a
forall a. Additive a => a
Add.zero v a
v

orA :: forall v s. (VG.Vector v Bool) => Array v s Bool -> Bool
orA :: forall (v :: * -> *) (s :: [Nat]).
Vector v Bool =>
Array v s Bool -> Bool
orA (Array v Bool
v) = v Bool -> Bool
forall (v :: * -> *). Vector v Bool => v Bool -> Bool
VG.or v Bool
v

traverseA :: forall f v s a b. (Applicative f, KnownNats s, VG.Vector v a, VG.Vector v b, VG.Vector v (f b)) => (a -> f b) -> Array v s a -> f (Array v s b)
traverseA :: forall (f :: * -> *) (v :: * -> *) (s :: [Nat]) a b.
(Applicative f, KnownNats s, Vector v a, Vector v b,
 Vector v (f b)) =>
(a -> f b) -> Array v s a -> f (Array v s b)
traverseA a -> f b
f (Array v a
v) = v b -> Array v s b
forall {k} (v :: k -> *) (s :: [Nat]) (a :: k). v a -> Array v s a
Array (v b -> Array v s b) -> ([b] -> v b) -> [b] -> Array v s 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 s b) -> f [b] -> f (Array v s 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 v a
v))

-- | Representation of an index into a shape (a type-level [Nat]). 'Dim @0' is commonly thought of as the row of an array.
type Dim = SNat

-- | Pattern synonym for a 'Dim'
pattern Dim :: () => (KnownNat n) => SNat n
pattern $bDim :: forall (n :: Nat). KnownNat n => SNat n
$mDim :: forall {r} {n :: Nat}.
SNat n -> (KnownNat n => r) -> ((# #) -> r) -> r
Dim = SNat

{-# COMPLETE Dim #-}

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

-- | Pattern synonym for a 'Dims'
pattern Dims :: () => (KnownNats ns) => SNats ns
pattern $bDims :: forall (ns :: [Nat]). KnownNats ns => SNats ns
$mDims :: forall {r} {ns :: [Nat]}.
SNats ns -> (KnownNats ns => r) -> ((# #) -> r) -> r
Dims = SNats

{-# COMPLETE Dims #-}

-- | Convert to a dynamic array with shape at the value level.
--
-- >>> toDynamic 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]
toDynamic :: (KnownNats s, VG.Vector v a) => Array v s a -> A.Array v a
toDynamic :: forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
Array v s a -> Array v a
toDynamic a :: Array v s a
a@(Array v a
v) = [Int] -> v a -> Array v a
forall {k} (v :: k -> *) (a :: k). [Int] -> v a -> Array v a
A.arrayV (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v s a -> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape Array v s a
a)) v a
v

-- | Use a dynamic array in a fixed context.
--
-- >>> import Harpie.Array.Generic qualified as A
-- >>> with (A.range @Vec.Vector [2,3,4]) show
-- "[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23]"
--
-- This doesn't work for anything more complex where KnownNats need to be type computed:
--
-- > :t with (A.range [2,3,4]) (pretty . F.takes (Dims @'[0]) (S.SNats @'[1]))
-- ...
--     • Could not deduce ‘S.KnownNats (Fcf.Data.List.Drop_ 1 s)’
-- ...
with ::
  forall v a r.
  (VG.Vector v a) =>
  A.Array v a ->
  (forall s. (KnownNats s) => Array v s a -> r) ->
  r
with :: forall (v :: * -> *) a r.
Vector v a =>
Array v a
-> (forall (s :: [Nat]). KnownNats s => Array v s a -> r) -> r
with Array v a
d forall (s :: [Nat]). KnownNats s => Array v s a -> r
f =
  [Nat] -> (forall (s :: [Nat]). SNats s -> r) -> r
forall r. [Nat] -> (forall (s :: [Nat]). SNats s -> r) -> r
withSomeSNats (Int -> Nat
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Nat) -> [Int] -> [Nat]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f 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
A.shape Array v a
d)) ((forall (s :: [Nat]). SNats s -> r) -> r)
-> (forall (s :: [Nat]). SNats s -> r) -> r
forall a b. (a -> b) -> a -> b
$ \(SNats s
SNats :: SNats s) -> SNats s -> (KnownNats s => r) -> r
forall (ns :: [Nat]) r. SNats ns -> (KnownNats ns => r) -> r
withKnownNats (forall (ns :: [Nat]). KnownNats ns => SNats ns
SNats @s) (Array v s a -> r
forall (s :: [Nat]). KnownNats s => Array v s a -> r
f (forall (v :: * -> *) (s :: [Nat]) a t.
(KnownNats s, FromVector t v a, Vector v a) =>
t -> Array v s a
array @v @s (Array v a -> v a
forall {k} (v :: k -> *) (a :: k). Array v a -> v a
A.asVector Array v a
d)))

-- | Get shape of an Array as a value.
--
-- >>> shape a
-- [2,3,4]
shape :: forall v a s. (KnownNats s, VG.Vector v a) => Array v s a -> VU.Vector Int
shape :: forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape Array v s a
_ = [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @s)
{-# INLINE shape #-}

-- | Get rank of an Array as a value.
--
-- >>> rank a
-- 3
rank :: forall v a s. (KnownNats s, VG.Vector v a) => Array v s a -> Int
rank :: forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Int
rank = Vector Int -> Int
S.rank (Vector Int -> Int)
-> (Array v s a -> Vector Int) -> Array v s a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array v s a -> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape
{-# INLINE rank #-}

-- | Get size of an Array as a value.
--
-- >>> size a
-- 24
size :: forall v a s. (KnownNats s, VG.Vector v a) => Array v s a -> Int
size :: forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Int
size = Vector Int -> Int
S.size (Vector Int -> Int)
-> (Array v s a -> Vector Int) -> Array v s a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array v s a -> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape
{-# INLINE size #-}

-- | 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 @Vec.Vector 0)
-- 1
length :: (KnownNats s, VG.Vector v a) => Array v s a -> Int
length :: forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
Array v s a -> Int
length Array v s a
a = case Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v s a -> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape Array v s a
a) of
  [] -> Int
1
  (Int
x : [Int]
_) -> Int
x

-- | Is the Array empty (has zero number of elements).
--
-- >>> isNull (array [] :: Array Vec.Vector [2,0] ())
-- True
-- >>> isNull (array [4] :: Array Vec.Vector '[] Int)
-- False
isNull :: (KnownNats s, VG.Vector v a) => Array v s a -> Bool
isNull :: forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
Array v s a -> Bool
isNull = (Int
0 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==) (Int -> Bool) -> (Array v s a -> Int) -> Array v s a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array v s a -> Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Int
size

-- | Extract an element at an index, unsafely.
--
-- >>> unsafeIndex a [1,2,3]
-- 23
unsafeIndex :: (KnownNats s, VG.Vector v a) => Array v s a -> [Int] -> a
unsafeIndex :: forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
Array v s a -> [Int] -> a
unsafeIndex Array v s a
a [Int]
xs = Array v s a -> Fins s -> a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v s a
a ([Int] -> Fins s
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins [Int]
xs)

-- | Extract an element at an index, unsafely.
--
-- >>> a ! [1,2,3]
-- 23
(!) :: (KnownNats s, VG.Vector v a) => Array v s a -> [Int] -> a
! :: forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
Array v s a -> [Int] -> a
(!) Array v s a
a [Int]
xs = Array v s a -> Fins s -> a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v s a
a ([Int] -> Fins s
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins [Int]
xs)

infixl 9 !

-- | Extract an element at an index, safely.
--
-- >>> a !? [1,2,3]
-- Just 23
-- >>> a !? [2,3,1]
-- Nothing
(!?) :: (KnownNats s, VG.Vector v a) => Array v s a -> [Int] -> Maybe a
!? :: forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
Array v s a -> [Int] -> Maybe a
(!?) Array v s a
a [Int]
xs = Array v s a -> Fins s -> a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v s a
a (Fins s -> a) -> Maybe (Fins s) -> Maybe a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int] -> Maybe (Fins s)
forall (s :: [Nat]). KnownNats s => [Int] -> Maybe (Fins s)
safeFins [Int]
xs

infixl 9 !?

-- | Tabulate unsafely.
--
-- >>> :t tabulate @Vec.Vector @[2,3] id
-- tabulate @Vec.Vector @[2,3] id
--   :: Array Vec.Vector [2, 3] (Fins [2, 3])
-- >>> :t unsafeTabulate @[2,3] id :: Array Vec.Vector [2,3] [Int]
-- unsafeTabulate @[2,3] id :: Array Vec.Vector [2,3] [Int]
--   :: Array Vec.Vector [2, 3] [Int]
-- >>> pretty (unsafeTabulate @[2,3] id :: Array Vec.Vector [2,3] [Int])
-- [[[0,0],[0,1],[0,2]],
--  [[1,0],[1,1],[1,2]]]
unsafeTabulate :: (KnownNats s, VG.Vector v a) => ([Int] -> a) -> Array v s a
unsafeTabulate :: forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
([Int] -> a) -> Array v s a
unsafeTabulate [Int] -> a
f = (Fins s -> a) -> Array v s a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
(Fins s -> a) -> Array v s a
tabulate ([Int] -> a
f ([Int] -> a) -> (Fins s -> [Int]) -> Fins s -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Fins s -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins)

-- | @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 (backpermute f' a) == backpermute (f . f') a
--
-- Many functions in this module are examples of backpermute usage.
--
-- >>> pretty $ backpermute @Vec.Vector @[4,3,2] (UnsafeFins . List.reverse . fromFins) 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 :: forall v s' s a. (KnownNats s, KnownNats s', VG.Vector v a) => (Fins s' -> Fins s) -> Array v s a -> Array v s' a
backpermute :: forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
(Fins s' -> Fins s) -> Array v s a -> Array v s' a
backpermute Fins s' -> Fins s
f Array v s a
a = (Fins s' -> a) -> Array v s' a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
(Fins s -> a) -> Array v s a
tabulate (Array v s a -> Fins s -> a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v s a
a (Fins s -> a) -> (Fins s' -> Fins s) -> Fins s' -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Fins s' -> Fins s
f)
{-# INLINEABLE backpermute #-}

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

-- | Unsafe backpermute
--
-- >>> pretty $ unsafeBackpermute @Vec.Vector @[4,3,2] List.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]]]
unsafeBackpermute :: forall v s' s a. (KnownNats s, KnownNats s', VG.Vector v a) => ([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute :: forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute [Int] -> [Int]
f Array v s a
a = (Fins s' -> a) -> Array v s' a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
(Fins s -> a) -> Array v s a
tabulate (Array v s a -> Fins s -> a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v s a
a (Fins s -> a) -> (Fins s' -> Fins s) -> Fins s' -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> Fins s
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins ([Int] -> Fins s) -> (Fins s' -> [Int]) -> Fins s' -> Fins s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [Int]
f ([Int] -> [Int]) -> (Fins s' -> [Int]) -> Fins s' -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Fins s' -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins)

{- RULES
   "unsafeBackpermute/unsafeBackpermute" forall f f' (a :: forall a. Array a)). unsafeBackpermute f (unsafeBackpermute f' a) == unsafeBackpermute (f . f') a
-}

-- | Unwrap a scalar.
--
-- >>> s = array @Vec.Vector @'[] @Int [3]
-- >>> :t fromScalar s
-- fromScalar s :: Int
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 -> Fins '[] -> a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v '[] a
a ([Int] -> Fins '[]
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins [])

-- | Wrap a scalar.
--
-- >>> :t toScalar @Vec.Vector @Int 2
-- toScalar @Vec.Vector @Int 2 :: Array Vec.Vector '[] Int
toScalar :: (VG.Vector v a) => a -> Array v '[] a
toScalar :: forall (v :: * -> *) a. Vector v a => a -> Array v '[] a
toScalar a
a = v a -> Array v '[] a
forall {k} (v :: k -> *) (s :: [Nat]) (a :: k). v a -> Array v s a
Array (a -> v a
forall (v :: * -> *) a. Vector v a => a -> v a
VG.singleton a
a)

-- | Is an array a scalar?
--
-- >>> isScalar (toScalar @Vec.Vector (2::Int))
-- True
isScalar :: (KnownNats s, VG.Vector v a) => Array v s a -> Bool
isScalar :: forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
Array v s a -> Bool
isScalar Array v s a
a = Array v s a -> Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Int
rank Array v s 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 @Vec.Vector 4)
-- [4]
asSingleton :: (KnownNats s, KnownNats s', s' ~ Eval (AsSingleton s), VG.Vector v a) => Array v s a -> Array v s' a
asSingleton :: forall (s :: [Nat]) (s' :: [Nat]) (v :: * -> *) a.
(KnownNats s, KnownNats s', s' ~ Eval (AsSingleton s),
 Vector v a) =>
Array v s a -> Array v s' a
asSingleton = Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
Array v s a -> Array v s' a
unsafeModifyShape

-- | Convert an array with shape [1] to being a scalar (Do nothing if not a shape [1] array).
--
-- >>> pretty (asScalar (singleton @Vec.Vector 3))
-- 3
asScalar :: (KnownNats s, KnownNats s', s' ~ Eval (AsScalar s), VG.Vector v a) => Array v s a -> Array v s' a
asScalar :: forall (s :: [Nat]) (s' :: [Nat]) (v :: * -> *) a.
(KnownNats s, KnownNats s', s' ~ Eval (AsScalar s), Vector v a) =>
Array v s a -> Array v s' a
asScalar = Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
Array v s a -> Array v s' a
unsafeModifyShape

-- | An array with no elements.
--
-- >>> toDynamic (empty @Vec.Vector @Int)
-- UnsafeArray [0] []
empty :: (VG.Vector v a) => Array v '[0] a
empty :: forall (v :: * -> *) a. Vector v a => Array v '[0] a
empty = [a] -> Array v '[0] a
forall (v :: * -> *) (s :: [Nat]) a t.
(KnownNats s, FromVector t v a, Vector v a) =>
t -> Array v s a
array []

-- | An enumeration of row-major or [lexicographic](https://en.wikipedia.org/wiki/Lexicographic_order) order.
--
-- >>> pretty (range :: Array Vec.Vector [2,3] Int)
-- [[0,1,2],
--  [3,4,5]]
range :: forall v s. (KnownNats s, VG.Vector v Int) => Array v s Int
range :: forall (v :: * -> *) (s :: [Nat]).
(KnownNats s, Vector v Int) =>
Array v s Int
range = (Fins s -> Int) -> Array v s Int
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
(Fins s -> a) -> Array v s a
tabulate (Vector Int -> Vector Int -> Int
S.flatten ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @s)) (Vector Int -> Int) -> (Fins s -> Vector Int) -> Fins s -> 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) -> (Fins s -> [Int]) -> Fins s -> Vector Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Fins s -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins)

-- | An enumeration of col-major or [colexicographic](https://en.wikipedia.org/wiki/Lexicographic_order) order.
--
-- >>> pretty (corange @Vec.Vector @[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 :: forall v s. (KnownNats s, VG.Vector v Int) => Array v s Int
corange :: forall (v :: * -> *) (s :: [Nat]).
(KnownNats s, Vector v Int) =>
Array v s Int
corange = (Fins s -> Int) -> Array v s Int
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
(Fins s -> a) -> Array v s a
tabulate (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 (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @s))) (Vector Int -> Int) -> (Fins s -> Vector Int) -> Fins s -> 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) -> (Fins s -> [Int]) -> Fins s -> Vector Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [Int]
forall a. [a] -> [a]
List.reverse ([Int] -> [Int]) -> (Fins s -> [Int]) -> Fins s -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Fins s -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins)

-- | Indices of an array shape.
--
-- >>> pretty (indices @[3,3] :: Array Vec.Vector [3,3] [Int])
-- [[[0,0],[0,1],[0,2]],
--  [[1,0],[1,1],[1,2]],
--  [[2,0],[2,1],[2,2]]]
indices :: (KnownNats s, VG.Vector v [Int]) => Array v s [Int]
indices :: forall (s :: [Nat]) (v :: * -> *).
(KnownNats s, Vector v [Int]) =>
Array v s [Int]
indices = (Fins s -> [Int]) -> Array v s [Int]
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
(Fins s -> a) -> Array v s a
tabulate Fins s -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins

-- | The identity array.
--
-- >>> pretty (ident @[3,3] :: Array Vec.Vector [3,3] Int)
-- [[1,0,0],
--  [0,1,0],
--  [0,0,1]]
ident :: (KnownNats s, Add.Additive a, Mult.Multiplicative a, VG.Vector v a) => Array v s a
ident :: forall (s :: [Nat]) a (v :: * -> *).
(KnownNats s, Additive a, Multiplicative a, Vector v a) =>
Array v s a
ident = (Fins s -> a) -> Array v s a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
(Fins s -> a) -> Array v s a
tabulate (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) -> (Fins s -> Bool) -> Fins s -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Bool
S.isDiag (Vector Int -> Bool) -> (Fins s -> Vector Int) -> Fins s -> Bool
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) -> (Fins s -> [Int]) -> Fins s -> Vector Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Fins s -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins)

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

-- | Create an array of shape [1].
--
-- >>> pretty $ singleton @Vec.Vector 1
-- [1]
singleton :: forall v a. (VG.Vector v a) => a -> Array v '[1] a
singleton :: forall (v :: * -> *) a. Vector v a => a -> Array v '[1] a
singleton a
a = v a -> Array v '[1] a
forall {k} (v :: k -> *) (s :: [Nat]) (a :: k). v a -> Array v s a
Array (forall (v :: * -> *) a. Vector v a => a -> v a
VG.singleton @v a
a)

-- | Extract the diagonal of an array.
--
-- >>> pretty $ diag (ident @[3,3] :: Array Vec.Vector [3,3] Int)
-- [1,1,1]
diag :: forall v s' a s. (KnownNats s, KnownNats s', s' ~ Eval (MinDim s), VG.Vector v a) => Array v s a -> Array v s' a
diag :: forall (v :: * -> *) (s' :: [Nat]) a (s :: [Nat]).
(KnownNats s, KnownNats s', s' ~ Eval (MinDim s), Vector v a) =>
Array v s a -> Array v s' a
diag Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute (Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate (Array v s a -> Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Int
rank Array v s a
a) (Int -> [Int]) -> ([Int] -> Int) -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Int] -> Int
S.getDimL Int
0) Array v s a
a

-- | Expand an array to form a diagonal array
--
-- >>> pretty $ undiag (range @Vec.Vector @'[3])
-- [[0,0,0],
--  [0,1,0],
--  [0,0,2]]
undiag :: forall v s' a s. (KnownNats s, KnownNats s', s' ~ Eval ((++) s s), Add.Additive a, VG.Vector v a) => Array v s a -> Array v s' a
undiag :: forall (v :: * -> *) (s' :: [Nat]) a (s :: [Nat]).
(KnownNats s, KnownNats s', s' ~ Eval (s ++ s), Additive a,
 Vector v a) =>
Array v s a -> Array v s' a
undiag Array v s a
a = (Fins s' -> a) -> Array v s' a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
(Fins s -> a) -> Array v s a
tabulate (\Fins s'
xs -> a -> a -> Bool -> a
forall a. a -> a -> Bool -> a
bool a
forall a. Additive a => a
Add.zero (Array v s a -> Fins s -> a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v s a
a ([Int] -> Fins s
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins ([Int] -> Fins s) -> [Int] -> Fins s
forall a b. (a -> b) -> a -> b
$ Int -> [Int]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> [Int]) -> Int -> [Int]
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> Int
S.getDimL Int
0 (Fins s' -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins Fins s'
xs))) ([Int] -> Bool
forall a. Eq a => [a] -> Bool
S.isDiagL (Fins s' -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins Fins s'
xs)))

-- | Zip two arrays at an element level.
--
-- >>> zipWith (-) v v
-- [0,0,0]
zipWith :: (KnownNats s, VG.Vector v a, VG.Vector v b, VG.Vector v c) => (a -> b -> c) -> Array v s a -> Array v s b -> Array v s c
zipWith :: forall (s :: [Nat]) (v :: * -> *) a b c.
(KnownNats s, Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v s a -> Array v s b -> Array v s c
zipWith a -> b -> c
f (Array v a
a) (Array v b
b) = v c -> Array v s c
forall {k} (v :: k -> *) (s :: [Nat]) (a :: k). v a -> Array v s a
Array ((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
a v b
b)

-- | Modify a single value at an index.
--
-- >>> pretty $ modify (S.UnsafeFins [0,0]) (const 100) (range @Vec.Vector @[3,2])
-- [[100,1],
--  [2,3],
--  [4,5]]
modify :: (KnownNats s, VG.Vector v a) => Fins s -> (a -> a) -> Array v s a -> Array v s a
modify :: forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
Fins s -> (a -> a) -> Array v s a -> Array v s a
modify Fins s
ds a -> a
f Array v s a
a = (Fins s -> a) -> Array v s a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
(Fins s -> a) -> Array v s a
tabulate (\Fins s
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 (Fins s
s Fins s -> Fins s -> Bool
forall a. Eq a => a -> a -> Bool
== Fins s
ds) (Array v s a -> Fins s -> a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v s a
a Fins s
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 :: (KnownNats s, VG.Vector v a, VG.Vector v b, VG.Vector v [Int]) => ([Int] -> a -> b) -> Array v s a -> Array v s b
imap :: forall (s :: [Nat]) (v :: * -> *) a b.
(KnownNats s, Vector v a, Vector v b, Vector v [Int]) =>
([Int] -> a -> b) -> Array v s a -> Array v s b
imap [Int] -> a -> b
f Array v s a
a = ([Int] -> a -> b) -> Array v s [Int] -> Array v s a -> Array v s b
forall (s :: [Nat]) (v :: * -> *) a b c.
(KnownNats s, Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v s a -> Array v s b -> Array v s c
zipWith [Int] -> a -> b
f Array v s [Int]
forall (s :: [Nat]) (v :: * -> *).
(KnownNats s, Vector v [Int]) =>
Array v s [Int]
indices Array v s 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..rank xs - 1] xs
--
-- >>> toDynamic $ rowWise indexesT (S.SNats @[1,0]) a
-- UnsafeArray [4] [12,13,14,15]
rowWise :: forall v a ds s s' xs proxy. (KnownNats s, KnownNats ds, ds ~ Eval (DimsOf xs), VG.Vector v a) => (Dims ds -> proxy xs -> Array v s a -> Array v s' a) -> proxy xs -> Array v s a -> Array v s' a
rowWise :: forall (v :: * -> *) a (ds :: [Nat]) (s :: [Nat]) (s' :: [Nat])
       (xs :: [Nat]) (proxy :: [Nat] -> *).
(KnownNats s, KnownNats ds, ds ~ Eval (DimsOf xs), Vector v a) =>
(Dims ds -> proxy xs -> Array v s a -> Array v s' a)
-> proxy xs -> Array v s a -> Array v s' a
rowWise Dims ds -> proxy xs -> Array v s a -> Array v s' a
f proxy xs
xs Array v s a
a = Dims ds -> proxy xs -> Array v s a -> Array v s' a
f (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) proxy xs
xs Array v s 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
--
-- >>> toDynamic $ colWise indexesT (S.SNats @[1,0]) a
-- UnsafeArray [2] [1,13]
colWise :: forall v a ds s s' xs proxy. (KnownNats s, KnownNats ds, ds ~ Eval (EndDimsOf xs s), VG.Vector v a) => (Dims ds -> proxy xs -> Array v s a -> Array v s' a) -> proxy xs -> Array v s a -> Array v s' a
colWise :: forall (v :: * -> *) a (ds :: [Nat]) (s :: [Nat]) (s' :: [Nat])
       (xs :: [Nat]) (proxy :: [Nat] -> *).
(KnownNats s, KnownNats ds, ds ~ Eval (EndDimsOf xs s),
 Vector v a) =>
(Dims ds -> proxy xs -> Array v s a -> Array v s' a)
-> proxy xs -> Array v s a -> Array v s' a
colWise Dims ds -> proxy xs -> Array v s a -> Array v s' a
f proxy xs
xs Array v s a
a = Dims ds -> proxy xs -> Array v s a -> Array v s' a
f (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) proxy xs
xs Array v s a
a

-- | Take the top-most elements across the specified dimension.
--
-- >>> pretty $ take (Dim @2) (SNat @1) a
-- [[[0],
--   [4],
--   [8]],
--  [[12],
--   [16],
--   [20]]]
take :: forall v d t s s' a. (KnownNats s, KnownNats s', s' ~ Eval (TakeDim d t s), VG.Vector v a) => Dim d -> SNat t -> Array v s a -> Array v s' a
take :: forall (v :: * -> *) (d :: Nat) (t :: Nat) (s :: [Nat])
       (s' :: [Nat]) a.
(KnownNats s, KnownNats s', s' ~ Eval (TakeDim d t s),
 Vector v a) =>
Dim d -> SNat t -> Array v s a -> Array v s' a
take Dim d
_ SNat t
_ Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute [Int] -> [Int]
forall a. a -> a
id Array v s a
a

-- | Take the bottom-most elements across the specified dimension.
--
-- >>> pretty $ takeB (Dim @2) (SNat @1) a
-- [[[3],
--   [7],
--   [11]],
--  [[15],
--   [19],
--   [23]]]
takeB :: forall v s s' a d t. (KnownNats s, KnownNats s', s' ~ Eval (TakeDim d t s), VG.Vector v a) => Dim d -> SNat t -> Array v s a -> Array v s' a
takeB :: forall (v :: * -> *) (s :: [Nat]) (s' :: [Nat]) a (d :: Nat)
       (t :: Nat).
(KnownNats s, KnownNats s', s' ~ Eval (TakeDim d t s),
 Vector v a) =>
Dim d -> SNat t -> Array v s a -> Array v s' a
takeB SNat d
Dim SNat t
SNat Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute (Int -> (Int -> Int) -> [Int] -> [Int]
S.modifyDimL (forall (n :: Nat). KnownNat n => Int
valueOf @d) (\Int
x -> Int
x Int -> Int -> Int
forall a. Additive a => a -> a -> a
+ Int -> [Int] -> Int
S.getDimL (forall (n :: Nat). KnownNat n => Int
valueOf @d) (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v s a -> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape Array v s a
a)) Int -> Int -> Int
forall a. Subtractive a => a -> a -> a
- (forall (n :: Nat). KnownNat n => Int
valueOf @t))) Array v s a
a

-- | Drop the top-most elements across the specified dimension.
--
-- >>> pretty $ drop (Dim @2) (SNat @1) a
-- [[[1,2,3],
--   [5,6,7],
--   [9,10,11]],
--  [[13,14,15],
--   [17,18,19],
--   [21,22,23]]]
drop :: forall v s s' a d t. (KnownNats s, KnownNats s', Eval (DropDim d t s) ~ s', VG.Vector v a) => Dim d -> SNat t -> Array v s a -> Array v s' a
drop :: forall (v :: * -> *) (s :: [Nat]) (s' :: [Nat]) a (d :: Nat)
       (t :: Nat).
(KnownNats s, KnownNats s', Eval (DropDim d t s) ~ s',
 Vector v a) =>
Dim d -> SNat t -> Array v s a -> Array v s' a
drop SNat d
Dim SNat t
SNat Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute (Int -> (Int -> Int) -> [Int] -> [Int]
S.modifyDimL (forall (n :: Nat). KnownNat n => Int
valueOf @d) (\Int
x -> Int
x Int -> Int -> Int
forall a. Additive a => a -> a -> a
+ forall (n :: Nat). KnownNat n => Int
valueOf @t)) Array v s a
a

-- | Drop the bottom-most elements across the specified dimension.
--
-- >>> pretty $ dropB (Dim @2) (SNat @1) a
-- [[[0,1,2],
--   [4,5,6],
--   [8,9,10]],
--  [[12,13,14],
--   [16,17,18],
--   [20,21,22]]]
dropB :: forall v s s' a d t. (KnownNats s, KnownNats s', Eval (DropDim d t s) ~ s', VG.Vector v a) => Dim d -> SNat t -> Array v s a -> Array v s' a
dropB :: forall (v :: * -> *) (s :: [Nat]) (s' :: [Nat]) a (d :: Nat)
       (t :: Nat).
(KnownNats s, KnownNats s', Eval (DropDim d t s) ~ s',
 Vector v a) =>
Dim d -> SNat t -> Array v s a -> Array v s' a
dropB Dim d
_ SNat t
_ Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute [Int] -> [Int]
forall a. a -> a
id Array v s a
a

-- ---------------------------------------------------------------------------
-- Matrix block functions
-- ---------------------------------------------------------------------------

-- | Select an index along a dimension.
--
-- >>> let s = select (Dim @2) (S.fin @4 3) a
-- >>> pretty s
-- [[3,7,11],
--  [15,19,23]]
select :: forall v d a p s s'. (KnownNats s, KnownNats s', s' ~ Eval (DeleteDim d s), p ~ Eval (GetDim d s), VG.Vector v a) => Dim d -> Fin p -> Array v s a -> Array v s' a
select :: forall (v :: * -> *) (d :: Nat) a (p :: Nat) (s :: [Nat])
       (s' :: [Nat]).
(KnownNats s, KnownNats s', s' ~ Eval (DeleteDim d s),
 p ~ Eval (GetDim d s), Vector v a) =>
Dim d -> Fin p -> Array v s a -> Array v s' a
select SNat d
Dim Fin p
p Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute (Int -> Int -> [Int] -> [Int]
S.insertDimL (forall (n :: Nat). KnownNat n => Int
valueOf @d) (Fin p -> Int
forall {k} (s :: k). Fin s -> Int
fromFin Fin p
p)) Array v s a
a

-- | Insert along a dimension at a position.
--
-- >>> pretty $ insert (Dim @2) (UnsafeFin 0) a (konst @[2,3] 0 :: Array Vec.Vector [2,3] Int)
-- [[[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]]]
-- >>> toDynamic $ insert (Dim @0) (UnsafeFin 0) (toScalar @Vec.Vector 1) (toScalar @Vec.Vector 2)
-- UnsafeArray [2] [2,1]
insert :: forall v s' s si d p a. (KnownNats s, KnownNats si, KnownNats s', s' ~ Eval (IncAt d s), p ~ Eval (GetDim d s), True ~ Eval (InsertOk d s si), VG.Vector v a) => Dim d -> Fin p -> Array v s a -> Array v si a -> Array v s' a
insert :: forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) (si :: [Nat])
       (d :: Nat) (p :: Nat) a.
(KnownNats s, KnownNats si, KnownNats s', s' ~ Eval (IncAt d s),
 p ~ Eval (GetDim d s), 'True ~ Eval (InsertOk d s si),
 Vector v a) =>
Dim d -> Fin p -> Array v s a -> Array v si a -> Array v s' a
insert SNat d
Dim Fin p
i Array v s a
a Array v si a
b = (Fins s' -> a) -> Array v s' a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
(Fins s -> a) -> Array v s a
tabulate Fins s' -> a
go
  where
    go :: Fins s' -> a
go Fins s'
s
      | Int -> [Int] -> Int
S.getDimL Int
d [Int]
s' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Fin p -> Int
forall {k} (s :: k). Fin s -> Int
fromFin Fin p
i = Array v si a -> Fins si -> a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v si a
b ([Int] -> Fins si
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins (Int -> [Int] -> [Int]
S.deleteDimL Int
d [Int]
s'))
      | Int -> [Int] -> Int
S.getDimL Int
d [Int]
s' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Fin p -> Int
forall {k} (s :: k). Fin s -> Int
fromFin Fin p
i = Array v s a -> Fins s -> a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v s a
a ([Int] -> Fins s
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins [Int]
s')
      | Bool
otherwise = Array v s a -> Fins s -> a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v s a
a ([Int] -> Fins s
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins (Int -> [Int] -> [Int]
S.decAtL Int
d [Int]
s'))
      where
        s' :: [Int]
s' = Fins s' -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins Fins s'
s
    d :: Int
d = forall (n :: Nat). KnownNat n => Int
valueOf @d

-- | Delete along a dimension at a position.
--
-- >>> pretty $ delete (Dim @2) (UnsafeFin 3) a
-- [[[0,1,2],
--   [4,5,6],
--   [8,9,10]],
--  [[12,13,14],
--   [16,17,18],
--   [20,21,22]]]
delete :: forall v d s s' p a. (KnownNats s, KnownNats s', s' ~ Eval (DecAt d s), p ~ 1 + Eval (GetDim d s), VG.Vector v a) => Dim d -> Fin p -> Array v s a -> Array v s' a
delete :: forall (v :: * -> *) (d :: Nat) (s :: [Nat]) (s' :: [Nat])
       (p :: Nat) a.
(KnownNats s, KnownNats s', s' ~ Eval (DecAt d s),
 p ~ (1 + Eval (GetDim d s)), Vector v a) =>
Dim d -> Fin p -> Array v s a -> Array v s' a
delete SNat d
Dim Fin p
p Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute (\[Int]
s -> [Int] -> [Int] -> Bool -> [Int]
forall a. a -> a -> Bool -> a
bool (Int -> [Int] -> [Int]
S.incAtL Int
d [Int]
s) [Int]
s (Int -> [Int] -> Int
S.getDimL Int
d [Int]
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Fin p -> Int
forall {k} (s :: k). Fin s -> Int
fromFin Fin p
p)) Array v s a
a
  where
    d :: Int
d = forall (n :: Nat). KnownNat n => Int
valueOf @d

-- | Insert along a dimension at the end.
--
-- >>> pretty $ append (Dim @2) a (konst @[2,3] 0 :: Array Vec.Vector [2,3] Int)
-- [[[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 :: forall v a d s si s'. (KnownNats s, KnownNats si, KnownNats s', s' ~ Eval (IncAt d s), True ~ Eval (InsertOk d s si), VG.Vector v a) => Dim d -> Array v s a -> Array v si a -> Array v s' a
append :: forall (v :: * -> *) a (d :: Nat) (s :: [Nat]) (si :: [Nat])
       (s' :: [Nat]).
(KnownNats s, KnownNats si, KnownNats s', s' ~ Eval (IncAt d s),
 'True ~ Eval (InsertOk d s si), Vector v a) =>
Dim d -> Array v s a -> Array v si a -> Array v s' a
append (Dim d
Dim :: Dim d) = Dim d
-> Fin
     (If
        (Eval (TyEqImpl d 0 && TyEqImpl s '[]))
        1
        (Eval (FromMaybe (TypeError ...) (GetIndexImpl d s))))
-> Array v s a
-> Array v si a
-> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) (si :: [Nat])
       (d :: Nat) (p :: Nat) a.
(KnownNats s, KnownNats si, KnownNats s', s' ~ Eval (IncAt d s),
 p ~ Eval (GetDim d s), 'True ~ Eval (InsertOk d s si),
 Vector v a) =>
Dim d -> Fin p -> Array v s a -> Array v si a -> Array v s' a
insert (forall (n :: Nat). KnownNat n => SNat n
Dim @d) (Int
-> Fin
     (If
        (Eval (TyEqImpl d 0 && TyEqImpl s '[]))
        1
        (Eval (FromMaybe (TypeError ...) (GetIndexImpl d s))))
forall {k} (s :: k). Int -> Fin s
UnsafeFin (Int -> [Int] -> Int
S.getDimL (forall (n :: Nat). KnownNat n => Int
valueOf @d) (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @s)))

-- | Insert along a dimension at the beginning.
--
-- >>> pretty $ prepend (Dim @2) (konst @[2,3] 0 :: Array Vec.Vector [2,3] Int) 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 :: forall v a d s si s'. (KnownNats s, KnownNats si, KnownNats s', s' ~ Eval (IncAt d s), True ~ Eval (InsertOk d s si), VG.Vector v a) => Dim d -> Array v si a -> Array v s a -> Array v s' a
prepend :: forall (v :: * -> *) a (d :: Nat) (s :: [Nat]) (si :: [Nat])
       (s' :: [Nat]).
(KnownNats s, KnownNats si, KnownNats s', s' ~ Eval (IncAt d s),
 'True ~ Eval (InsertOk d s si), Vector v a) =>
Dim d -> Array v si a -> Array v s a -> Array v s' a
prepend Dim d
d Array v si a
a Array v s a
b = Dim d
-> Fin
     (If
        (Eval (TyEqImpl d 0 && TyEqImpl s '[]))
        1
        (Eval (FromMaybe (TypeError ...) (GetIndexImpl d s))))
-> Array v s a
-> Array v si a
-> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) (si :: [Nat])
       (d :: Nat) (p :: Nat) a.
(KnownNats s, KnownNats si, KnownNats s', s' ~ Eval (IncAt d s),
 p ~ Eval (GetDim d s), 'True ~ Eval (InsertOk d s si),
 Vector v a) =>
Dim d -> Fin p -> Array v s a -> Array v si a -> Array v s' a
insert Dim d
d (Int
-> Fin
     (If
        (Eval (TyEqImpl d 0 && TyEqImpl s '[]))
        1
        (Eval (FromMaybe (TypeError ...) (GetIndexImpl d s))))
forall {k} (s :: k). Int -> Fin s
UnsafeFin Int
0) Array v s a
b Array v si a
a

-- | Concatenate along a dimension.
--
-- >>> shape $ concatenate (Dim @1) a a
-- [2,6,4]
-- >>> toDynamic $ concatenate (Dim @0) (toScalar @Vec.Vector 1) (toScalar @Vec.Vector 2)
-- UnsafeArray [2] [1,2]
-- >>> toDynamic $ concatenate (Dim @0) (array @Vec.Vector @'[1] [0]) (array @Vec.Vector @'[3] [1..3])
-- UnsafeArray [4] [0,1,2,3]
concatenate :: forall v a s0 s1 d s. (KnownNats s0, KnownNats s1, KnownNats s, Eval (Concatenate d s0 s1) ~ s, VG.Vector v a) => Dim d -> Array v s0 a -> Array v s1 a -> Array v s a
concatenate :: forall (v :: * -> *) a (s0 :: [Nat]) (s1 :: [Nat]) (d :: Nat)
       (s :: [Nat]).
(KnownNats s0, KnownNats s1, KnownNats s,
 Eval (Concatenate d s0 s1) ~ s, Vector v a) =>
Dim d -> Array v s0 a -> Array v s1 a -> Array v s a
concatenate SNat d
Dim Array v s0 a
a0 Array v s1 a
a1 = (Fins s -> a) -> Array v s a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
(Fins s -> a) -> Array v s a
tabulate ([Int] -> a
go ([Int] -> a) -> (Fins s -> [Int]) -> Fins s -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Fins s -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins)
  where
    go :: [Int] -> a
go [Int]
s =
      a -> a -> Bool -> a
forall a. a -> a -> Bool -> a
bool
        (Array v s0 a -> Fins s0 -> a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v s0 a
a0 ([Int] -> Fins s0
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins [Int]
s))
        ( Array v s1 a -> Fins s1 -> a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index
            Array v s1 a
a1
            ( [Int] -> Fins s1
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins ([Int] -> Fins s1) -> [Int] -> Fins s1
forall a b. (a -> b) -> a -> b
$
                Int -> Int -> [Int] -> [Int]
S.insertDimL
                  Int
d'
                  (Int -> [Int] -> Int
S.getDimL Int
d' [Int]
s Int -> Int -> Int
forall a. Subtractive a => a -> a -> a
- Int -> [Int] -> Int
S.getDimL Int
d' [Int]
ds0)
                  (Int -> [Int] -> [Int]
S.deleteDimL Int
d' [Int]
s)
            )
        )
        (Int -> [Int] -> Int
S.getDimL Int
d' [Int]
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int -> [Int] -> Int
S.getDimL Int
d' [Int]
ds0)
    ds0 :: [Int]
ds0 = Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v s0 a -> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape Array v s0 a
a0)
    d' :: Int
d' = forall (n :: Nat). KnownNat n => Int
valueOf @d

-- | Combine two arrays as a new dimension of a new array.
--
-- >>> pretty $ couple (Dim @0) (array @Vec.Vector @'[3] [1,2,3]) (array @Vec.Vector @'[3] @Int [4,5,6])
-- [[1,2,3],
--  [4,5,6]]
-- >>> couple (Dim @0) (toScalar @Vec.Vector @Int 0) (toScalar @Vec.Vector 1)
-- [0,1]
couple :: forall v d a s s' se. (KnownNat d, KnownNats s, KnownNats s', KnownNats se, s' ~ Eval (Concatenate d se se), se ~ Eval (InsertDim d 1 s), VG.Vector v a) => Dim d -> Array v s a -> Array v s a -> Array v s' a
couple :: forall (v :: * -> *) (d :: Nat) a (s :: [Nat]) (s' :: [Nat])
       (se :: [Nat]).
(KnownNat d, KnownNats s, KnownNats s', KnownNats se,
 s' ~ Eval (Concatenate d se se), se ~ Eval (InsertDim d 1 s),
 Vector v a) =>
Dim d -> Array v s a -> Array v s a -> Array v s' a
couple Dim d
d Array v s a
a Array v s a
a' = Dim d -> Array v se a -> Array v se a -> Array v s' a
forall (v :: * -> *) a (s0 :: [Nat]) (s1 :: [Nat]) (d :: Nat)
       (s :: [Nat]).
(KnownNats s0, KnownNats s1, KnownNats s,
 Eval (Concatenate d s0 s1) ~ s, Vector v a) =>
Dim d -> Array v s0 a -> Array v s1 a -> Array v s a
concatenate Dim d
d (Dim d -> Array v s a -> Array v se a
forall (s :: [Nat]) (s' :: [Nat]) (d :: Nat) (v :: * -> *) a.
(KnownNats s, KnownNats s', s' ~ Eval (InsertDim d 1 s),
 Vector v a) =>
Dim d -> Array v s a -> Array v s' a
elongate Dim d
d Array v s a
a) (Dim d -> Array v s a -> Array v se a
forall (s :: [Nat]) (s' :: [Nat]) (d :: Nat) (v :: * -> *) a.
(KnownNats s, KnownNats s', s' ~ Eval (InsertDim d 1 s),
 Vector v a) =>
Dim d -> Array v s a -> Array v s' a
elongate Dim d
d Array v s a
a')

-- | Slice along a dimension with the supplied offset & length.
--
-- >>> pretty $ slice (Dim @2) (SNat @1) (SNat @2) a
-- [[[1,2],
--   [5,6],
--   [9,10]],
--  [[13,14],
--   [17,18],
--   [21,22]]]
slice :: forall v a d off l s s'. (KnownNats s, KnownNats s', s' ~ Eval (SetDim d l s), Eval (SliceOk d off l s) ~ True, VG.Vector v a) => Dim d -> SNat off -> SNat l -> Array v s a -> Array v s' a
slice :: forall (v :: * -> *) a (d :: Nat) (off :: Nat) (l :: Nat)
       (s :: [Nat]) (s' :: [Nat]).
(KnownNats s, KnownNats s', s' ~ Eval (SetDim d l s),
 Eval (SliceOk d off l s) ~ 'True, Vector v a) =>
Dim d -> SNat off -> SNat l -> Array v s a -> Array v s' a
slice SNat d
Dim SNat off
SNat SNat l
_ Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute (Int -> (Int -> Int) -> [Int] -> [Int]
S.modifyDimL (forall (n :: Nat). KnownNat n => Int
valueOf @d) (Int -> Int -> Int
forall a. Additive a => a -> a -> a
+ (forall (n :: Nat). KnownNat n => Int
valueOf @off))) Array v s a
a

-- | Rotate an array along a dimension.
--
-- >>> pretty $ rotate (Dim @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 :: forall v d s a. (KnownNats s, VG.Vector v a) => Dim d -> Int -> Array v s a -> Array v s a
rotate :: forall (v :: * -> *) (d :: Nat) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Dim d -> Int -> Array v s a -> Array v s a
rotate SNat d
Dim Int
r Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute (Int -> Int -> [Int] -> [Int] -> [Int]
S.rotateIndexL (forall (n :: Nat). KnownNat n => Int
valueOf @d) Int
r (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v s a -> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape Array v s a
a))) Array v s a
a

-- * multi-dimensional operators

-- | Across the specified dimensions, takes the top-most elements.
--
-- >>> pretty $ takes (Dims @[0,1]) (S.SNats @[1,2]) a
-- [[[0,1,2,3],
--   [4,5,6,7]]]
takes :: forall v ds xs s' s a. (KnownNats s, KnownNats s', s' ~ Eval (SetDims ds xs s), VG.Vector v a) => Dims ds -> SNats xs -> Array v s a -> Array v s' a
takes :: forall (v :: * -> *) (ds :: [Nat]) (xs :: [Nat]) (s' :: [Nat])
       (s :: [Nat]) a.
(KnownNats s, KnownNats s', s' ~ Eval (SetDims ds xs s),
 Vector v a) =>
Dims ds -> SNats xs -> Array v s a -> Array v s' a
takes Dims ds
_ SNats xs
_ Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute [Int] -> [Int]
forall a. a -> a
id Array v s a
a

-- | Across the specified dimensions, takes the bottom-most elements.
--
-- >>> pretty (takeBs (Dims @[0,1]) (S.SNats @[1,2]) a)
-- [[[16,17,18,19],
--   [20,21,22,23]]]
takeBs :: forall v s' s a ds xs. (KnownNats s, KnownNats s', KnownNats ds, KnownNats xs, s' ~ Eval (SetDims ds xs s), VG.Vector v a) => Dims ds -> SNats xs -> Array v s a -> Array v s' a
takeBs :: forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a (ds :: [Nat])
       (xs :: [Nat]).
(KnownNats s, KnownNats s', KnownNats ds, KnownNats xs,
 s' ~ Eval (SetDims ds xs s), Vector v a) =>
Dims ds -> SNats xs -> Array v s a -> Array v s' a
takeBs Dims ds
_ SNats xs
_ Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
List.zipWith Int -> Int -> Int
forall a. Additive a => a -> a -> a
(+) [Int]
start) Array v s a
a
  where
    start :: [Int]
start = (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
List.zipWith (-) (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v s a -> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape Array v s a
a)) ([Int] -> [Int] -> [Int] -> [Int]
S.setDimsL (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @ds) (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @xs) (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v s a -> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape Array v s a
a)))

-- | Across the specified dimensions, drops the top-most elements.
--
-- >>> pretty $ drops (Dims @[0,2]) (S.SNats @[1,3]) a
-- [[[15],
--   [19],
--   [23]]]
drops :: forall v ds xs s' s a. (KnownNats s, KnownNats s', KnownNats ds, KnownNats xs, s' ~ Eval (DropDims ds xs s), VG.Vector v a) => Dims ds -> SNats xs -> Array v s a -> Array v s' a
drops :: forall (v :: * -> *) (ds :: [Nat]) (xs :: [Nat]) (s' :: [Nat])
       (s :: [Nat]) a.
(KnownNats s, KnownNats s', KnownNats ds, KnownNats xs,
 s' ~ Eval (DropDims ds xs s), Vector v a) =>
Dims ds -> SNats xs -> Array v s a -> Array v s' a
drops Dims ds
_ SNats xs
_ Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
List.zipWith Int -> Int -> Int
forall a. Additive a => a -> a -> a
(+) [Int]
start) Array v s a
a
  where
    start :: [Int]
start = (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
List.zipWith (-) (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @s) (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @s')

-- | Across the specified dimensions, drops the bottom-most elements.
--
-- >>> pretty $ dropBs (Dims @[0,2]) (S.SNats @[1,3]) a
-- [[[0],
--   [4],
--   [8]]]
dropBs :: forall v s' s ds xs a. (KnownNats s, KnownNats s', KnownNats ds, KnownNats xs, s' ~ Eval (DropDims ds xs s), VG.Vector v a) => Dims ds -> SNats xs -> Array v s a -> Array v s' a
dropBs :: forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) (ds :: [Nat])
       (xs :: [Nat]) a.
(KnownNats s, KnownNats s', KnownNats ds, KnownNats xs,
 s' ~ Eval (DropDims ds xs s), Vector v a) =>
Dims ds -> SNats xs -> Array v s a -> Array v s' a
dropBs Dims ds
_ SNats xs
_ Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute [Int] -> [Int]
forall a. a -> a
id Array v s a
a

-- | Select by dimensions and indexes.
--
-- >>> pretty $ indexes (Dims @[0,1]) (S.UnsafeFins [1,1]) a
-- [16,17,18,19]
-- >>> pretty $ indexes (Dims @'[1]) (S.fins @'[3] [1]) (range @Vec.Vector @[2,3])
-- [1,4]
indexes :: forall v s' s ds xs a. (KnownNats s, KnownNats s', s' ~ Eval (DeleteDims ds s), xs ~ Eval (GetDims ds s), VG.Vector v a) => Dims ds -> Fins xs -> Array v s a -> Array v s' a
indexes :: forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) (ds :: [Nat])
       (xs :: [Nat]) a.
(KnownNats s, KnownNats s', s' ~ Eval (DeleteDims ds s),
 xs ~ Eval (GetDims ds s), Vector v a) =>
Dims ds -> Fins xs -> Array v s a -> Array v s' a
indexes SNats ds
Dims Fins xs
xs Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute ([Int] -> [Int] -> [Int] -> [Int]
S.insertDimsL (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @ds) (Fins xs -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins Fins xs
xs)) Array v s a
a

-- | Select by dimensions and indexes, supplying indexes as a type.
--
-- >>> pretty $ indexesT (Dims @[0,1]) (S.SNats @[1,1]) a
-- [16,17,18,19]
indexesT :: forall v ds xs s s' a. (KnownNats s, KnownNats ds, KnownNats xs, KnownNats s', s' ~ Eval (DeleteDims ds s), True ~ Eval (IsFins xs =<< GetDims ds s), VG.Vector v a) => Dims ds -> SNats xs -> Array v s a -> Array v s' a
indexesT :: forall (v :: * -> *) (ds :: [Nat]) (xs :: [Nat]) (s :: [Nat])
       (s' :: [Nat]) a.
(KnownNats s, KnownNats ds, KnownNats xs, KnownNats s',
 s' ~ Eval (DeleteDims ds s),
 'True ~ Eval (IsFins xs =<< GetDims ds s), Vector v a) =>
Dims ds -> SNats xs -> Array v s a -> Array v s' a
indexesT Dims ds
ds SNats xs
_ Array v s a
a = Dims ds
-> Fins (Eval (Map (Flip GetDim s) ds))
-> Array v s a
-> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) (ds :: [Nat])
       (xs :: [Nat]) a.
(KnownNats s, KnownNats s', s' ~ Eval (DeleteDims ds s),
 xs ~ Eval (GetDims ds s), Vector v a) =>
Dims ds -> Fins xs -> Array v s a -> Array v s' a
indexes Dims ds
ds ([Int] -> Fins (Eval (Map (Flip GetDim s) ds))
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins ([Int] -> Fins (Eval (Map (Flip GetDim s) ds)))
-> [Int] -> Fins (Eval (Map (Flip GetDim s) ds))
forall a b. (a -> b) -> a -> b
$ forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @xs) Array v s a
a

-- | Slice along dimensions with the supplied offsets and lengths.
--
-- >>> pretty $ slices (Dims @'[2]) (S.SNats @'[1]) (S.SNats @'[2]) a
-- [[[1,2],
--   [5,6],
--   [9,10]],
--  [[13,14],
--   [17,18],
--   [21,22]]]
slices :: forall v a ds ls offs s s'. (KnownNats s, KnownNats s', KnownNats ds, KnownNats ls, KnownNats offs, Eval (SlicesOk ds offs ls s) ~ True, Eval (SetDims ds ls s) ~ s', VG.Vector v a) => Dims ds -> SNats offs -> SNats ls -> Array v s a -> Array v s' a
slices :: forall (v :: * -> *) a (ds :: [Nat]) (ls :: [Nat]) (offs :: [Nat])
       (s :: [Nat]) (s' :: [Nat]).
(KnownNats s, KnownNats s', KnownNats ds, KnownNats ls,
 KnownNats offs, Eval (SlicesOk ds offs ls s) ~ 'True,
 Eval (SetDims ds ls s) ~ s', Vector v a) =>
Dims ds -> SNats offs -> SNats ls -> Array v s a -> Array v s' a
slices Dims ds
_ SNats offs
_ SNats ls
_ Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
List.zipWith Int -> Int -> Int
forall a. Additive a => a -> a -> a
(+) [Int]
o) Array v s a
a
  where
    o :: [Int]
o = [Int] -> [Int] -> [Int] -> [Int]
S.setDimsL (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @ds) (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @offs) (Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate (Array v s a -> Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Int
rank Array v s a
a) Int
0)

-- | Select the first element along the supplied dimensions.
--
-- >>> pretty $ heads (Dims @[0,2]) a
-- [0,4,8]
heads :: forall v a ds s s'. (KnownNats s, KnownNats s', KnownNats ds, s' ~ Eval (DeleteDims ds s), VG.Vector v a) => Dims ds -> Array v s a -> Array v s' a
heads :: forall (v :: * -> *) a (ds :: [Nat]) (s :: [Nat]) (s' :: [Nat]).
(KnownNats s, KnownNats s', KnownNats ds,
 s' ~ Eval (DeleteDims ds s), Vector v a) =>
Dims ds -> Array v s a -> Array v s' a
heads Dims ds
ds Array v s a
a = Dims ds
-> Fins (Eval (Map (Flip GetDim s) ds))
-> Array v s a
-> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) (ds :: [Nat])
       (xs :: [Nat]) a.
(KnownNats s, KnownNats s', s' ~ Eval (DeleteDims ds s),
 xs ~ Eval (GetDims ds s), Vector v a) =>
Dims ds -> Fins xs -> Array v s a -> Array v s' a
indexes Dims ds
ds ([Int] -> Fins (Eval (Map (Flip GetDim s) ds))
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins ([Int] -> Fins (Eval (Map (Flip GetDim s) ds)))
-> [Int] -> Fins (Eval (Map (Flip GetDim s) ds))
forall a b. (a -> b) -> a -> b
$ Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate (forall (s :: [Nat]). KnownNats s => Int
rankOf @ds) Int
0) Array v s a
a

-- | Select the last element along the supplied dimensions.
--
-- >>> pretty $ lasts (Dims @[0,2]) a
-- [15,19,23]
lasts :: forall v ds s s' a. (KnownNats s, KnownNats ds, KnownNats s', s' ~ Eval (DeleteDims ds s), VG.Vector v a) => Dims ds -> Array v s a -> Array v s' a
lasts :: forall (v :: * -> *) (ds :: [Nat]) (s :: [Nat]) (s' :: [Nat]) a.
(KnownNats s, KnownNats ds, KnownNats s',
 s' ~ Eval (DeleteDims ds s), Vector v a) =>
Dims ds -> Array v s a -> Array v s' a
lasts Dims ds
ds Array v s a
a = Dims ds
-> Fins (Eval (Map (Flip GetDim s) ds))
-> Array v s a
-> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) (ds :: [Nat])
       (xs :: [Nat]) a.
(KnownNats s, KnownNats s', s' ~ Eval (DeleteDims ds s),
 xs ~ Eval (GetDims ds s), Vector v a) =>
Dims ds -> Fins xs -> Array v s a -> Array v s' a
indexes Dims ds
ds ([Int] -> Fins (Eval (Map (Flip GetDim s) ds))
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins [Int]
lastds) Array v s a
a
  where
    lastds :: [Int]
lastds = (\Int
i -> Int -> [Int] -> Int
S.getDimL Int
i (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v s a -> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape Array v s a
a)) Int -> Int -> Int
forall a. Subtractive a => a -> a -> a
- Int
1) (Int -> Int) -> [Int] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @ds)

-- | Select the tail elements along the supplied dimensions.
--
-- >>> pretty $ tails (Dims @[0,2]) a
-- [[[13,14,15],
--   [17,18,19],
--   [21,22,23]]]
tails :: forall v ds os s s' a ls. (KnownNats s, KnownNats ds, KnownNats s', KnownNats ls, KnownNats os, Eval (SlicesOk ds os ls s) ~ True, os ~ Eval (Replicate (Eval (Rank ds)) 1), ls ~ Eval (GetLastPositions ds s), s' ~ Eval (SetDims ds ls s), VG.Vector v a) => Dims ds -> Array v s a -> Array v s' a
tails :: forall (v :: * -> *) (ds :: [Nat]) (os :: [Nat]) (s :: [Nat])
       (s' :: [Nat]) a (ls :: [Nat]).
(KnownNats s, KnownNats ds, KnownNats s', KnownNats ls,
 KnownNats os, Eval (SlicesOk ds os ls s) ~ 'True,
 os ~ Eval (Replicate (Eval (Rank ds)) 1),
 ls ~ Eval (GetLastPositions ds s), s' ~ Eval (SetDims ds ls s),
 Vector v a) =>
Dims ds -> Array v s a -> Array v s' a
tails Dims ds
ds Array v s a
a = Dims ds -> SNats os -> SNats ls -> Array v s a -> Array v s' a
forall (v :: * -> *) a (ds :: [Nat]) (ls :: [Nat]) (offs :: [Nat])
       (s :: [Nat]) (s' :: [Nat]).
(KnownNats s, KnownNats s', KnownNats ds, KnownNats ls,
 KnownNats offs, Eval (SlicesOk ds offs ls s) ~ 'True,
 Eval (SetDims ds ls s) ~ s', Vector v a) =>
Dims ds -> SNats offs -> SNats ls -> Array v s a -> Array v s' a
slices Dims ds
ds (forall (ns :: [Nat]). KnownNats ns => SNats ns
SNats @os) (forall (ns :: [Nat]). KnownNats ns => SNats ns
SNats @ls) Array v s a
a

-- | Select the init elements along the supplied dimensions.
--
-- >>> pretty $ inits (Dims @[0,2]) a
-- [[[0,1,2],
--   [4,5,6],
--   [8,9,10]]]
inits :: forall v ds os s s' a ls. (KnownNats s, KnownNats ds, KnownNats s', KnownNats ls, KnownNats os, Eval (SlicesOk ds os ls s) ~ True, os ~ Eval (Replicate (Eval (Rank ds)) 0), ls ~ Eval (GetLastPositions ds s), s' ~ Eval (SetDims ds ls s), VG.Vector v a) => Dims ds -> Array v s a -> Array v s' a
inits :: forall (v :: * -> *) (ds :: [Nat]) (os :: [Nat]) (s :: [Nat])
       (s' :: [Nat]) a (ls :: [Nat]).
(KnownNats s, KnownNats ds, KnownNats s', KnownNats ls,
 KnownNats os, Eval (SlicesOk ds os ls s) ~ 'True,
 os ~ Eval (Replicate (Eval (Rank ds)) 0),
 ls ~ Eval (GetLastPositions ds s), s' ~ Eval (SetDims ds ls s),
 Vector v a) =>
Dims ds -> Array v s a -> Array v s' a
inits Dims ds
ds Array v s a
a = Dims ds -> SNats os -> SNats ls -> Array v s a -> Array v s' a
forall (v :: * -> *) a (ds :: [Nat]) (ls :: [Nat]) (offs :: [Nat])
       (s :: [Nat]) (s' :: [Nat]).
(KnownNats s, KnownNats s', KnownNats ds, KnownNats ls,
 KnownNats offs, Eval (SlicesOk ds offs ls s) ~ 'True,
 Eval (SetDims ds ls s) ~ s', Vector v a) =>
Dims ds -> SNats offs -> SNats ls -> Array v s a -> Array v s' a
slices Dims ds
ds (forall (ns :: [Nat]). KnownNats ns => SNats ns
SNats @os) (forall (ns :: [Nat]). KnownNats ns => SNats ns
SNats @ls) Array v s a
a

-- | Extracts specified dimensions to an outer layer.
--
-- >>> :t extracts (Dims @'[0]) (range @Vec.Vector @[2,3,4])
-- extracts (Dims @'[0]) (range @Vec.Vector @[2,3,4])
--   :: Array Vec.Vector '[2] (Array Vec.Vector [3, 4] Int)
extracts :: forall v ds st si so a. (KnownNats st, KnownNats ds, KnownNats si, KnownNats so, si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st), VG.Vector v a, VG.Vector v (Array v si a)) => Dims ds -> Array v st a -> Array v so (Array v si a)
extracts :: forall (v :: * -> *) (ds :: [Nat]) (st :: [Nat]) (si :: [Nat])
       (so :: [Nat]) a.
(KnownNats st, KnownNats ds, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st),
 Vector v a, Vector v (Array v si a)) =>
Dims ds -> Array v st a -> Array v so (Array v si a)
extracts Dims ds
ds Array v st a
a = (Fins so -> Array v si a) -> Array v so (Array v si a)
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
(Fins s -> a) -> Array v s a
tabulate (\Fins so
s -> Dims ds -> Fins so -> Array v st a -> Array v si a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) (ds :: [Nat])
       (xs :: [Nat]) a.
(KnownNats s, KnownNats s', s' ~ Eval (DeleteDims ds s),
 xs ~ Eval (GetDims ds s), Vector v a) =>
Dims ds -> Fins xs -> Array v s a -> Array v s' a
indexes Dims ds
ds Fins so
s Array v st a
a)

-- | Reduce along specified dimensions, using the supplied fold.
--
-- >>> pretty $ reduces (Dims @'[0]) sum a
-- [66,210]
-- >>> pretty $ reduces (Dims @[0,2]) sum a
-- [[12,15,18,21],
--  [48,51,54,57]]
reduces :: forall v ds st si so a b. (KnownNats st, KnownNats ds, KnownNats si, KnownNats so, si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st), VG.Vector v a, VG.Vector v b, VG.Vector v (Array v si a)) => Dims ds -> (Array v si a -> b) -> Array v st a -> Array v so b
reduces :: forall (v :: * -> *) (ds :: [Nat]) (st :: [Nat]) (si :: [Nat])
       (so :: [Nat]) a b.
(KnownNats st, KnownNats ds, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st),
 Vector v a, Vector v b, Vector v (Array v si a)) =>
Dims ds -> (Array v si a -> b) -> Array v st a -> Array v so b
reduces Dims ds
ds Array v si a -> b
f Array v st a
a = (Array v si a -> b) -> Array v so (Array v si a) -> Array v so b
forall (v :: * -> *) (s :: [Nat]) a b.
(KnownNats s, Vector v a, Vector v b) =>
(a -> b) -> Array v s a -> Array v s b
fmapA Array v si a -> b
f (Dims ds -> Array v st a -> Array v so (Array v si a)
forall (v :: * -> *) (ds :: [Nat]) (st :: [Nat]) (si :: [Nat])
       (so :: [Nat]) a.
(KnownNats st, KnownNats ds, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st),
 Vector v a, Vector v (Array v si a)) =>
Dims ds -> Array v st a -> Array v so (Array v si a)
extracts Dims ds
ds Array v st a
a)

-- | Join inner and outer dimension layers by supplied dimensions.
--
-- >>> let e = extracts (Dims @[1,0]) a
-- >>> let j = joins (Dims @[1,0]) e
-- >>> a == j
-- True
joins :: forall v a ds si so st. (KnownNats ds, KnownNats st, KnownNats si, KnownNats so, Eval (InsertDims ds so si) ~ st, VG.Vector v (Array v si a), VG.Vector v a) => Dims ds -> Array v so (Array v si a) -> Array v st a
joins :: forall (v :: * -> *) a (ds :: [Nat]) (si :: [Nat]) (so :: [Nat])
       (st :: [Nat]).
(KnownNats ds, KnownNats st, KnownNats si, KnownNats so,
 Eval (InsertDims ds so si) ~ st, Vector v (Array v si a),
 Vector v a) =>
Dims ds -> Array v so (Array v si a) -> Array v st a
joins Dims ds
_ Array v so (Array v si a)
a = (Fins st -> a) -> Array v st a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
(Fins s -> a) -> Array v s a
tabulate Fins st -> a
go
  where
    go :: Fins st -> a
go Fins st
s = Array v si a -> Fins si -> a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index (Array v so (Array v si a) -> Fins so -> Array v si a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v so (Array v si a)
a ([Int] -> Fins so
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins ([Int] -> Fins so) -> [Int] -> Fins so
forall a b. (a -> b) -> a -> b
$ [Int] -> [Int] -> [Int]
S.getDimsL (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @ds) (Fins st -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins Fins st
s))) ([Int] -> Fins si
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins ([Int] -> Fins si) -> [Int] -> Fins si
forall a b. (a -> b) -> a -> b
$ [Int] -> [Int] -> [Int]
S.deleteDimsL (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @ds) (Fins st -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins Fins st
s))

-- | Join inner and outer dimension layers in outer dimension order.
--
-- >>> a == join (extracts (Dims @[0,1]) a)
-- True
join :: forall v a si so st ds. (KnownNats st, KnownNats si, KnownNats so, KnownNats ds, ds ~ Eval (DimsOf so), st ~ Eval (InsertDims ds so si), VG.Vector v (Array v si a), VG.Vector v a) => Array v so (Array v si a) -> Array v st a
join :: forall (v :: * -> *) a (si :: [Nat]) (so :: [Nat]) (st :: [Nat])
       (ds :: [Nat]).
(KnownNats st, KnownNats si, KnownNats so, KnownNats ds,
 ds ~ Eval (DimsOf so), st ~ Eval (InsertDims ds so si),
 Vector v (Array v si a), Vector v a) =>
Array v so (Array v si a) -> Array v st a
join Array v so (Array v si a)
a = Dims ds -> Array v so (Array v si a) -> Array v st a
forall (v :: * -> *) a (ds :: [Nat]) (si :: [Nat]) (so :: [Nat])
       (st :: [Nat]).
(KnownNats ds, KnownNats st, KnownNats si, KnownNats so,
 Eval (InsertDims ds so si) ~ st, Vector v (Array v si a),
 Vector v a) =>
Dims ds -> Array v so (Array v si a) -> Array v st a
joins (forall (ns :: [Nat]). KnownNats ns => SNats ns
SNats @ds) Array v so (Array v si a)
a

-- | Traverse along specified dimensions.
--
-- >>> traverses (Dims @'[1]) print (range @Vec.Vector @[2,3])
-- 0
-- 3
-- 1
-- 4
-- 2
-- 5
-- [(),(),(),(),(),()]
traverses :: (Applicative f, KnownNats s, KnownNats si, KnownNats so, si ~ Eval (GetDims ds s), so ~ Eval (DeleteDims ds s), s ~ Eval (InsertDims ds si so), VG.Vector v a, VG.Vector v b, VG.Vector v (Array v so a), VG.Vector v (Array v so b), VG.Vector v (f b), VG.Vector v (f (Array v so b))) => Dims ds -> (a -> f b) -> Array v s a -> f (Array v s b)
traverses :: forall (f :: * -> *) (s :: [Nat]) (si :: [Nat]) (so :: [Nat])
       (ds :: [Nat]) (v :: * -> *) a b.
(Applicative f, KnownNats s, KnownNats si, KnownNats so,
 si ~ Eval (GetDims ds s), so ~ Eval (DeleteDims ds s),
 s ~ Eval (InsertDims ds si so), Vector v a, Vector v b,
 Vector v (Array v so a), Vector v (Array v so b), Vector v (f b),
 Vector v (f (Array v so b))) =>
Dims ds -> (a -> f b) -> Array v s a -> f (Array v s b)
traverses (Dims ds
Dims :: Dims ds) a -> f b
f Array v s a
a = Dims ds -> Array v si (Array v so b) -> Array v s b
forall (v :: * -> *) a (ds :: [Nat]) (si :: [Nat]) (so :: [Nat])
       (st :: [Nat]).
(KnownNats ds, KnownNats st, KnownNats si, KnownNats so,
 Eval (InsertDims ds so si) ~ st, Vector v (Array v si a),
 Vector v a) =>
Dims ds -> Array v so (Array v si a) -> Array v st a
joins (forall (ns :: [Nat]). KnownNats ns => SNats ns
SNats @ds) (Array v si (Array v so b) -> Array v s b)
-> f (Array v si (Array v so b)) -> f (Array v s b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Array v so a -> f (Array v so b))
-> Array v si (Array v so a) -> f (Array v si (Array v so b))
forall (f :: * -> *) (v :: * -> *) (s :: [Nat]) a b.
(Applicative f, KnownNats s, Vector v a, Vector v b,
 Vector v (f b)) =>
(a -> f b) -> Array v s a -> f (Array v s b)
traverseA ((a -> f b) -> Array v so a -> f (Array v so b)
forall (f :: * -> *) (v :: * -> *) (s :: [Nat]) a b.
(Applicative f, KnownNats s, Vector v a, Vector v b,
 Vector v (f b)) =>
(a -> f b) -> Array v s a -> f (Array v s b)
traverseA a -> f b
f) (Dims ds -> Array v s a -> Array v si (Array v so a)
forall (v :: * -> *) (ds :: [Nat]) (st :: [Nat]) (si :: [Nat])
       (so :: [Nat]) a.
(KnownNats st, KnownNats ds, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st),
 Vector v a, Vector v (Array v si a)) =>
Dims ds -> Array v st a -> Array v so (Array v si a)
extracts (Dims ds
forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims :: Dims ds) Array v s a
a)

-- | Maps a function along specified dimensions.
--
-- >>> pretty $ maps (Dims @'[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 :: forall v ds s s' si si' so a b. (KnownNats s, KnownNats s', KnownNats si, KnownNats si', KnownNats so, si ~ Eval (DeleteDims ds s), so ~ Eval (GetDims ds s), s' ~ Eval (InsertDims ds so si'), s ~ Eval (InsertDims ds so si), VG.Vector v a, VG.Vector v b, VG.Vector v (Array v si a), VG.Vector v (Array v si' b)) => Dims ds -> (Array v si a -> Array v si' b) -> Array v s a -> Array v s' b
maps :: forall (v :: * -> *) (ds :: [Nat]) (s :: [Nat]) (s' :: [Nat])
       (si :: [Nat]) (si' :: [Nat]) (so :: [Nat]) a b.
(KnownNats s, KnownNats s', KnownNats si, KnownNats si',
 KnownNats so, si ~ Eval (DeleteDims ds s),
 so ~ Eval (GetDims ds s), s' ~ Eval (InsertDims ds so si'),
 s ~ Eval (InsertDims ds so si), Vector v a, Vector v b,
 Vector v (Array v si a), Vector v (Array v si' b)) =>
Dims ds
-> (Array v si a -> Array v si' b) -> Array v s a -> Array v s' b
maps SNats ds
SNats Array v si a -> Array v si' b
f Array v s a
a = SNats ds -> Array v so (Array v si' b) -> Array v s' b
forall (v :: * -> *) a (ds :: [Nat]) (si :: [Nat]) (so :: [Nat])
       (st :: [Nat]).
(KnownNats ds, KnownNats st, KnownNats si, KnownNats so,
 Eval (InsertDims ds so si) ~ st, Vector v (Array v si a),
 Vector v a) =>
Dims ds -> Array v so (Array v si a) -> Array v st a
joins (forall (ns :: [Nat]). KnownNats ns => SNats ns
SNats @ds) ((Array v si a -> Array v si' b)
-> Array v so (Array v si a) -> Array v so (Array v si' b)
forall (v :: * -> *) (s :: [Nat]) a b.
(KnownNats s, Vector v a, Vector v b) =>
(a -> b) -> Array v s a -> Array v s b
fmapA Array v si a -> Array v si' b
f (SNats ds -> Array v s a -> Array v so (Array v si a)
forall (v :: * -> *) (ds :: [Nat]) (st :: [Nat]) (si :: [Nat])
       (so :: [Nat]) a.
(KnownNats st, KnownNats ds, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st),
 Vector v a, Vector v (Array v si a)) =>
Dims ds -> Array v st a -> Array v so (Array v si a)
extracts (forall (ns :: [Nat]). KnownNats ns => SNats ns
SNats @ds) Array v s a
a))

-- | Filters along specified dimensions (which are flattened as a dynamic array).
--
-- >>> pretty $ filters (Dims @[0,1]) (any ((==0) . (`mod` 7))) a
-- [[0,1,2,3],[4,5,6,7],[12,13,14,15],[20,21,22,23]]
filters :: forall v ds si so a. (KnownNats si, KnownNats so, si ~ Eval (DeleteDims ds so), KnownNats (Eval (GetDims ds so)), VG.Vector v a, VG.Vector v (Array v si a)) => Dims ds -> (Array v si a -> Bool) -> Array v so a -> A.Array v (Array v si a)
filters :: forall (v :: * -> *) (ds :: [Nat]) (si :: [Nat]) (so :: [Nat]) a.
(KnownNats si, KnownNats so, si ~ Eval (DeleteDims ds so),
 KnownNats (Eval (GetDims ds so)), Vector v a,
 Vector v (Array v si a)) =>
Dims ds
-> (Array v si a -> Bool) -> Array v so a -> Array v (Array v si a)
filters SNats ds
Dims Array v si a -> Bool
p Array v so a
a =
  let v' :: v (Array v si a)
v' = (Array v si a -> Bool) -> v (Array v si a) -> v (Array v si a)
forall (v :: * -> *) a. Vector v a => (a -> Bool) -> v a -> v a
VG.filter Array v si a -> Bool
p (Array v (Eval (Map (Flip GetDim so) ds)) (Array v si a)
-> v (Array v si a)
forall t (v :: * -> *) a. FromVector t v a => t -> v a
asVector (SNats ds
-> Array v so a
-> Array v (Eval (Map (Flip GetDim so) ds)) (Array v si a)
forall (v :: * -> *) (ds :: [Nat]) (st :: [Nat]) (si :: [Nat])
       (so :: [Nat]) a.
(KnownNats st, KnownNats ds, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st),
 Vector v a, Vector v (Array v si a)) =>
Dims ds -> Array v st a -> Array v so (Array v si a)
extracts (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) Array v so a
a))
   in [Int] -> v (Array v si a) -> Array v (Array v si a)
forall {k} (v :: k -> *) (a :: k). [Int] -> v a -> Array v a
A.arrayV [v (Array v si a) -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length v (Array v si a)
v'] v (Array v si a)
v'

-- | Zips two arrays with a function along specified dimensions.
--
-- >>> pretty $ zips (Dims @[0,1]) (zipWith (,)) a (reverses (Dims @'[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 :: forall v ds s s' si si' so a b c. (KnownNats s, KnownNats s', KnownNats si, KnownNats si', KnownNats so, si ~ Eval (DeleteDims ds s), so ~ Eval (GetDims ds s), s' ~ Eval (InsertDims ds so si'), s ~ Eval (InsertDims ds so si), VG.Vector v a, VG.Vector v b, VG.Vector v c, VG.Vector v (Array v si a), VG.Vector v (Array v si b), VG.Vector v (Array v si' c)) => Dims ds -> (Array v si a -> Array v si b -> Array v si' c) -> Array v s a -> Array v s b -> Array v s' c
zips :: forall (v :: * -> *) (ds :: [Nat]) (s :: [Nat]) (s' :: [Nat])
       (si :: [Nat]) (si' :: [Nat]) (so :: [Nat]) a b c.
(KnownNats s, KnownNats s', KnownNats si, KnownNats si',
 KnownNats so, si ~ Eval (DeleteDims ds s),
 so ~ Eval (GetDims ds s), s' ~ Eval (InsertDims ds so si'),
 s ~ Eval (InsertDims ds so si), Vector v a, Vector v b, Vector v c,
 Vector v (Array v si a), Vector v (Array v si b),
 Vector v (Array v si' c)) =>
Dims ds
-> (Array v si a -> Array v si b -> Array v si' c)
-> Array v s a
-> Array v s b
-> Array v s' c
zips SNats ds
SNats Array v si a -> Array v si b -> Array v si' c
f Array v s a
a Array v s b
b = SNats ds -> Array v so (Array v si' c) -> Array v s' c
forall (v :: * -> *) a (ds :: [Nat]) (si :: [Nat]) (so :: [Nat])
       (st :: [Nat]).
(KnownNats ds, KnownNats st, KnownNats si, KnownNats so,
 Eval (InsertDims ds so si) ~ st, Vector v (Array v si a),
 Vector v a) =>
Dims ds -> Array v so (Array v si a) -> Array v st a
joins (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) ((Array v si a -> Array v si b -> Array v si' c)
-> Array v so (Array v si a)
-> Array v so (Array v si b)
-> Array v so (Array v si' c)
forall (s :: [Nat]) (v :: * -> *) a b c.
(KnownNats s, Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v s a -> Array v s b -> Array v s c
zipWith Array v si a -> Array v si b -> Array v si' c
f (SNats ds -> Array v s a -> Array v so (Array v si a)
forall (v :: * -> *) (ds :: [Nat]) (st :: [Nat]) (si :: [Nat])
       (so :: [Nat]) a.
(KnownNats st, KnownNats ds, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st),
 Vector v a, Vector v (Array v si a)) =>
Dims ds -> Array v st a -> Array v so (Array v si a)
extracts (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) Array v s a
a) (SNats ds -> Array v s b -> Array v so (Array v si b)
forall (v :: * -> *) (ds :: [Nat]) (st :: [Nat]) (si :: [Nat])
       (so :: [Nat]) a.
(KnownNats st, KnownNats ds, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st),
 Vector v a, Vector v (Array v si a)) =>
Dims ds -> Array v st a -> Array v so (Array v si a)
extracts (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) Array v s b
b))

-- | Modify using the supplied function along dimensions and positions.
--
-- >>> pretty $ modifies (fmap (100+)) (Dims @'[2]) (S.UnsafeFins [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 :: forall v a si s ds so. (KnownNats s, KnownNats si, KnownNats so, si ~ Eval (DeleteDims ds s), so ~ Eval (GetDims ds s), s ~ Eval (InsertDims ds so si), VG.Vector v a, VG.Vector v (Array v si a)) => (Array v si a -> Array v si a) -> Dims ds -> Fins so -> Array v s a -> Array v s a
modifies :: forall (v :: * -> *) a (si :: [Nat]) (s :: [Nat]) (ds :: [Nat])
       (so :: [Nat]).
(KnownNats s, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds s), so ~ Eval (GetDims ds s),
 s ~ Eval (InsertDims ds so si), Vector v a,
 Vector v (Array v si a)) =>
(Array v si a -> Array v si a)
-> Dims ds -> Fins so -> Array v s a -> Array v s a
modifies Array v si a -> Array v si a
f SNats ds
SNats Fins so
ps Array v s a
a = SNats ds -> Array v so (Array v si a) -> Array v s a
forall (v :: * -> *) a (ds :: [Nat]) (si :: [Nat]) (so :: [Nat])
       (st :: [Nat]).
(KnownNats ds, KnownNats st, KnownNats si, KnownNats so,
 Eval (InsertDims ds so si) ~ st, Vector v (Array v si a),
 Vector v a) =>
Dims ds -> Array v so (Array v si a) -> Array v st a
joins (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) (Array v so (Array v si a) -> Array v s a)
-> Array v so (Array v si a) -> Array v s a
forall a b. (a -> b) -> a -> b
$ Fins so
-> (Array v si a -> Array v si a)
-> Array v so (Array v si a)
-> Array v so (Array v si a)
forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
Fins s -> (a -> a) -> Array v s a -> Array v s a
modify Fins so
ps Array v si a -> Array v si a
f (SNats ds -> Array v s a -> Array v so (Array v si a)
forall (v :: * -> *) (ds :: [Nat]) (st :: [Nat]) (si :: [Nat])
       (so :: [Nat]) a.
(KnownNats st, KnownNats ds, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st),
 Vector v a, Vector v (Array v si a)) =>
Dims ds -> Array v st a -> Array v so (Array v si a)
extracts (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) Array v s a
a)

-- | Apply a binary function between successive slices, across dimensions and lags.
--
-- >>> pretty $ diffs (Dims @'[1]) (S.SNats @'[1]) (zipWith (-)) a
-- [[[4,4,4,4],
--   [4,4,4,4]],
--  [[4,4,4,4],
--   [4,4,4,4]]]
diffs :: forall v a b ds ls si si' st st' so postDrop. (KnownNats ls, KnownNats si, KnownNats si', KnownNats st, KnownNats st', KnownNats so, KnownNats postDrop, si ~ Eval (DeleteDims ds postDrop), so ~ Eval (GetDims ds postDrop), st' ~ Eval (InsertDims ds so si'), postDrop ~ Eval (InsertDims ds so si), postDrop ~ Eval (DropDims ds ls st), VG.Vector v a, VG.Vector v b, VG.Vector v (Array v si a), VG.Vector v (Array v si' b)) => Dims ds -> SNats ls -> (Array v si a -> Array v si a -> Array v si' b) -> Array v st a -> Array v st' b
diffs :: forall (v :: * -> *) a b (ds :: [Nat]) (ls :: [Nat]) (si :: [Nat])
       (si' :: [Nat]) (st :: [Nat]) (st' :: [Nat]) (so :: [Nat])
       (postDrop :: [Nat]).
(KnownNats ls, KnownNats si, KnownNats si', KnownNats st,
 KnownNats st', KnownNats so, KnownNats postDrop,
 si ~ Eval (DeleteDims ds postDrop),
 so ~ Eval (GetDims ds postDrop), st' ~ Eval (InsertDims ds so si'),
 postDrop ~ Eval (InsertDims ds so si),
 postDrop ~ Eval (DropDims ds ls st), Vector v a, Vector v b,
 Vector v (Array v si a), Vector v (Array v si' b)) =>
Dims ds
-> SNats ls
-> (Array v si a -> Array v si a -> Array v si' b)
-> Array v st a
-> Array v st' b
diffs SNats ds
SNats SNats ls
xs Array v si a -> Array v si a -> Array v si' b
f Array v st a
a = SNats ds
-> (Array v si a -> Array v si a -> Array v si' b)
-> Array v postDrop a
-> Array v postDrop a
-> Array v st' b
forall (v :: * -> *) (ds :: [Nat]) (s :: [Nat]) (s' :: [Nat])
       (si :: [Nat]) (si' :: [Nat]) (so :: [Nat]) a b c.
(KnownNats s, KnownNats s', KnownNats si, KnownNats si',
 KnownNats so, si ~ Eval (DeleteDims ds s),
 so ~ Eval (GetDims ds s), s' ~ Eval (InsertDims ds so si'),
 s ~ Eval (InsertDims ds so si), Vector v a, Vector v b, Vector v c,
 Vector v (Array v si a), Vector v (Array v si b),
 Vector v (Array v si' c)) =>
Dims ds
-> (Array v si a -> Array v si b -> Array v si' c)
-> Array v s a
-> Array v s b
-> Array v s' c
zips (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) Array v si a -> Array v si a -> Array v si' b
f (SNats ds -> SNats ls -> Array v st a -> Array v postDrop a
forall (v :: * -> *) (ds :: [Nat]) (xs :: [Nat]) (s' :: [Nat])
       (s :: [Nat]) a.
(KnownNats s, KnownNats s', KnownNats ds, KnownNats xs,
 s' ~ Eval (DropDims ds xs s), Vector v a) =>
Dims ds -> SNats xs -> Array v s a -> Array v s' a
drops (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) SNats ls
xs Array v st a
a) (SNats ds -> SNats ls -> Array v st a -> Array v postDrop a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) (ds :: [Nat])
       (xs :: [Nat]) a.
(KnownNats s, KnownNats s', KnownNats ds, KnownNats xs,
 s' ~ Eval (DropDims ds xs s), Vector v a) =>
Dims ds -> SNats xs -> Array v s a -> Array v s' a
dropBs (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) SNats ls
xs Array v st 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 [1,2,3] :: Array Vec.Vector '[3] Int
-- >>> 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] :: Array Vec.Vector [2,2] [Int]
-- >>> 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 :: forall v sc sa sb a b c. (KnownNats sa, KnownNats sb, KnownNats sc, sc ~ Eval ((++) sa sb), VG.Vector v a, VG.Vector v b, VG.Vector v c) => (a -> b -> c) -> Array v sa a -> Array v sb b -> Array v sc c
expand :: forall (v :: * -> *) (sc :: [Nat]) (sa :: [Nat]) (sb :: [Nat]) a b
       c.
(KnownNats sa, KnownNats sb, KnownNats sc, sc ~ Eval (sa ++ sb),
 Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v sa a -> Array v sb b -> Array v sc c
expand a -> b -> c
f Array v sa a
a Array v sb b
b = (Fins sc -> c) -> Array v sc c
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
(Fins s -> a) -> Array v s a
tabulate (\Fins sc
i -> a -> b -> c
f (Array v sa a -> Fins sa -> a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v sa a
a ([Int] -> Fins sa
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins ([Int] -> Fins sa) -> [Int] -> Fins sa
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
List.take Int
r (Fins sc -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins Fins sc
i))) (Array v sb b -> Fins sb -> b
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v sb b
b ([Int] -> Fins sb
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins ([Int] -> Fins sb) -> [Int] -> Fins sb
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
List.drop Int
r (Fins sc -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins Fins sc
i))))
  where
    r :: Int
r = Array v sa a -> Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Int
rank Array v sa a
a

-- | Like expand, but permutes the first array first, rather than the second.
--
-- >>> pretty $ expand (,) v (fmap (+3) v)
-- [[(0,3),(0,4),(0,5)],
--  [(1,3),(1,4),(1,5)],
--  [(2,3),(2,4),(2,5)]]
--
-- >>> pretty $ coexpand (,) v (fmap (+3) v)
-- [[(0,3),(1,3),(2,3)],
--  [(0,4),(1,4),(2,4)],
--  [(0,5),(1,5),(2,5)]]
coexpand :: forall v sc sa sb a b c. (KnownNats sa, KnownNats sb, KnownNats sc, sc ~ Eval ((++) sb sa), VG.Vector v a, VG.Vector v b, VG.Vector v c) => (a -> b -> c) -> Array v sa a -> Array v sb b -> Array v sc c
coexpand :: forall (v :: * -> *) (sc :: [Nat]) (sa :: [Nat]) (sb :: [Nat]) a b
       c.
(KnownNats sa, KnownNats sb, KnownNats sc, sc ~ Eval (sb ++ sa),
 Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v sa a -> Array v sb b -> Array v sc c
coexpand a -> b -> c
f Array v sa a
a Array v sb b
b = (Fins sc -> c) -> Array v sc c
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
(Fins s -> a) -> Array v s a
tabulate (\Fins sc
i -> a -> b -> c
f (Array v sa a -> Fins sa -> a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v sa a
a ([Int] -> Fins sa
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins ([Int] -> Fins sa) -> [Int] -> Fins sa
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
List.drop Int
rb (Fins sc -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins Fins sc
i))) (Array v sb b -> Fins sb -> b
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v sb b
b ([Int] -> Fins sb
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins ([Int] -> Fins sb) -> [Int] -> Fins sb
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
List.take Int
rb (Fins sc -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins Fins sc
i))))
  where
    rb :: Int
rb = Array v sb b -> Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Int
rank Array v sb 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 (Dims @[1,2]) sum (expand (*) m (transpose m))
-- [[5,14],
--  [14,50]]
contract :: forall v a b s ss se s' ds ds'. (KnownNats se, se ~ Eval (DeleteDims ds' s), KnownNats ds', KnownNats s, KnownNats ss, KnownNats s', s' ~ Eval (GetDims ds' s), ss ~ Eval (MinDim se), ds' ~ Eval (ExceptDims ds s), VG.Vector v a, VG.Vector v b, VG.Vector v (Array v se a)) => Dims ds -> (Array v ss a -> b) -> Array v s a -> Array v s' b
contract :: forall (v :: * -> *) a b (s :: [Nat]) (ss :: [Nat]) (se :: [Nat])
       (s' :: [Nat]) (ds :: [Nat]) (ds' :: [Nat]).
(KnownNats se, se ~ Eval (DeleteDims ds' s), KnownNats ds',
 KnownNats s, KnownNats ss, KnownNats s', s' ~ Eval (GetDims ds' s),
 ss ~ Eval (MinDim se), ds' ~ Eval (ExceptDims ds s), Vector v a,
 Vector v b, Vector v (Array v se a)) =>
Dims ds -> (Array v ss a -> b) -> Array v s a -> Array v s' b
contract SNats ds
SNats Array v ss a -> b
f Array v s a
a = (Array v se a -> b) -> Array v s' (Array v se a) -> Array v s' b
forall (v :: * -> *) (s :: [Nat]) a b.
(KnownNats s, Vector v a, Vector v b) =>
(a -> b) -> Array v s a -> Array v s b
fmapA (Array v ss a -> b
f (Array v ss a -> b)
-> (Array v se a -> Array v ss a) -> Array v se a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array v se a -> Array v ss a
forall (v :: * -> *) (s' :: [Nat]) a (s :: [Nat]).
(KnownNats s, KnownNats s', s' ~ Eval (MinDim s), Vector v a) =>
Array v s a -> Array v s' a
diag) (Dims ds' -> Array v s a -> Array v s' (Array v se a)
forall (v :: * -> *) (ds :: [Nat]) (st :: [Nat]) (si :: [Nat])
       (so :: [Nat]) a.
(KnownNats st, KnownNats ds, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st),
 Vector v a, Vector v (Array v si a)) =>
Dims ds -> Array v st a -> Array v so (Array v si a)
extracts (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds') Array v s a
a)

-- | Expand two arrays and then contract the result using the supplied matching dimensions.
--
-- >>> pretty $ prod (Dims @'[1]) (Dims @'[0]) sum (*) (range @Vec.Vector @[2,3]) (range @Vec.Vector @[3,2])
-- [[10,13],
--  [28,40]]
--
-- With full laziness, this computation would be equivalent to:
--
-- > f . diag <$> extracts (Dims @ds') (expand g a b)
prod :: forall v a b c d s0 s1 so0 so1 si st ds0 ds1. (KnownNats so0, KnownNats so1, KnownNats si, KnownNats s0, KnownNats s1, KnownNats st, KnownNats ds0, KnownNats ds1, so0 ~ Eval (DeleteDims ds0 s0), so1 ~ Eval (DeleteDims ds1 s1), si ~ Eval (GetDims ds0 s0), si ~ Eval (GetDims ds1 s1), st ~ Eval ((++) so0 so1), VG.Vector v c, VG.Vector v a, VG.Vector v b, VG.Vector v d) => Dims ds0 -> Dims ds1 -> (Array v si c -> d) -> (a -> b -> c) -> Array v s0 a -> Array v s1 b -> Array v st d
prod :: forall (v :: * -> *) a b c d (s0 :: [Nat]) (s1 :: [Nat])
       (so0 :: [Nat]) (so1 :: [Nat]) (si :: [Nat]) (st :: [Nat])
       (ds0 :: [Nat]) (ds1 :: [Nat]).
(KnownNats so0, KnownNats so1, KnownNats si, KnownNats s0,
 KnownNats s1, KnownNats st, KnownNats ds0, KnownNats ds1,
 so0 ~ Eval (DeleteDims ds0 s0), so1 ~ Eval (DeleteDims ds1 s1),
 si ~ Eval (GetDims ds0 s0), si ~ Eval (GetDims ds1 s1),
 st ~ Eval (so0 ++ so1), Vector v c, Vector v a, Vector v b,
 Vector v d) =>
Dims ds0
-> Dims ds1
-> (Array v si c -> d)
-> (a -> b -> c)
-> Array v s0 a
-> Array v s1 b
-> Array v st d
prod SNats ds0
SNats SNats ds1
SNats Array v si c -> d
g a -> b -> c
f Array v s0 a
a Array v s1 b
b = ([Int] -> d) -> Array v st d
forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
([Int] -> a) -> Array v s a
unsafeTabulate (\[Int]
so -> Array v si c -> d
g (Array v si c -> d) -> Array v si c -> d
forall a b. (a -> b) -> a -> b
$ ([Int] -> c) -> Array v si c
forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
([Int] -> a) -> Array v s a
unsafeTabulate (\[Int]
si -> a -> b -> c
f (Array v s0 a -> [Int] -> a
forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
Array v s a -> [Int] -> a
unsafeIndex Array v s0 a
a ([Int] -> [Int] -> [Int] -> [Int]
S.insertDimsL (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @ds0) [Int]
si (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
List.take Int
sp [Int]
so))) (Array v s1 b -> [Int] -> b
forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
Array v s a -> [Int] -> a
unsafeIndex Array v s1 b
b ([Int] -> [Int] -> [Int] -> [Int]
S.insertDimsL (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @ds1) [Int]
si (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
List.drop Int
sp [Int]
so)))))
  where
    sp :: Int
sp = Array v s0 a -> Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Int
rank Array v s0 a
a Int -> Int -> Int
forall a. Subtractive a => a -> a -> a
- forall (s :: [Nat]). KnownNats s => Int
rankOf @ds0

-- | 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 sum (*) m (transpose m)
-- [[5,14],
--  [14,50]]
--
-- inner product
--
-- >>> pretty $ dot sum (*) v v
-- 5
--
-- matrix-vector multiplication
-- Note that an Array with shape [3] is neither a row vector nor column vector.
--
-- >>> pretty $ dot sum (*) v (transpose m)
-- [5,14]
--
-- >>> pretty $ dot sum (*) m v
-- [5,14]
dot :: forall v a b c d ds0 ds1 s0 s1 so0 so1 st si. (KnownNats s0, KnownNats s1, KnownNats ds0, KnownNats ds1, KnownNats so0, KnownNats so1, KnownNats st, KnownNats si, so0 ~ Eval (DeleteDims ds0 s0), so1 ~ Eval (DeleteDims ds1 s1), si ~ Eval (GetDims ds0 s0), si ~ Eval (GetDims ds1 s1), st ~ Eval ((++) so0 so1), ds0 ~ '[Eval ((Fcf.-) (Eval (Rank s0)) 1)], ds1 ~ '[0], VG.Vector v c, VG.Vector v a, VG.Vector v b, VG.Vector v d) => (Array v si c -> d) -> (a -> b -> c) -> Array v s0 a -> Array v s1 b -> Array v st d
dot :: forall (v :: * -> *) a b c d (ds0 :: [Nat]) (ds1 :: [Nat])
       (s0 :: [Nat]) (s1 :: [Nat]) (so0 :: [Nat]) (so1 :: [Nat])
       (st :: [Nat]) (si :: [Nat]).
(KnownNats s0, KnownNats s1, KnownNats ds0, KnownNats ds1,
 KnownNats so0, KnownNats so1, KnownNats st, KnownNats si,
 so0 ~ Eval (DeleteDims ds0 s0), so1 ~ Eval (DeleteDims ds1 s1),
 si ~ Eval (GetDims ds0 s0), si ~ Eval (GetDims ds1 s1),
 st ~ Eval (so0 ++ so1), ds0 ~ '[Eval (Eval (Rank s0) - 1)],
 ds1 ~ '[0], Vector v c, Vector v a, Vector v b, Vector v d) =>
(Array v si c -> d)
-> (a -> b -> c) -> Array v s0 a -> Array v s1 b -> Array v st d
dot Array v si c -> d
f a -> b -> c
g Array v s0 a
a Array v s1 b
b = Dims ds0
-> Dims ds1
-> (Array v si c -> d)
-> (a -> b -> c)
-> Array v s0 a
-> Array v s1 b
-> Array v st d
forall (v :: * -> *) a b c d (s0 :: [Nat]) (s1 :: [Nat])
       (so0 :: [Nat]) (so1 :: [Nat]) (si :: [Nat]) (st :: [Nat])
       (ds0 :: [Nat]) (ds1 :: [Nat]).
(KnownNats so0, KnownNats so1, KnownNats si, KnownNats s0,
 KnownNats s1, KnownNats st, KnownNats ds0, KnownNats ds1,
 so0 ~ Eval (DeleteDims ds0 s0), so1 ~ Eval (DeleteDims ds1 s1),
 si ~ Eval (GetDims ds0 s0), si ~ Eval (GetDims ds1 s1),
 st ~ Eval (so0 ++ so1), Vector v c, Vector v a, Vector v b,
 Vector v d) =>
Dims ds0
-> Dims ds1
-> (Array v si c -> d)
-> (a -> b -> c)
-> Array v s0 a
-> Array v s1 b
-> Array v st d
prod (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds0) (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds1) Array v si c -> d
f a -> b -> c
g Array v s0 a
a Array v s1 b
b

-- | Array 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 :: forall v a ds0 ds1 s0 s1 so0 so1 st si. (Add.Additive a, Mult.Multiplicative a, KnownNats s0, KnownNats s1, KnownNats ds0, KnownNats ds1, KnownNats so0, KnownNats so1, KnownNats st, KnownNats si, so0 ~ Eval (DeleteDims ds0 s0), so1 ~ Eval (DeleteDims ds1 s1), si ~ Eval (GetDims ds0 s0), si ~ Eval (GetDims ds1 s1), st ~ Eval ((++) so0 so1), ds0 ~ '[Eval ((Fcf.-) (Eval (Rank s0)) 1)], ds1 ~ '[0], VG.Vector v a) => Array v s0 a -> Array v s1 a -> Array v st a
mult :: forall (v :: * -> *) a (ds0 :: [Nat]) (ds1 :: [Nat]) (s0 :: [Nat])
       (s1 :: [Nat]) (so0 :: [Nat]) (so1 :: [Nat]) (st :: [Nat])
       (si :: [Nat]).
(Additive a, Multiplicative a, KnownNats s0, KnownNats s1,
 KnownNats ds0, KnownNats ds1, KnownNats so0, KnownNats so1,
 KnownNats st, KnownNats si, so0 ~ Eval (DeleteDims ds0 s0),
 so1 ~ Eval (DeleteDims ds1 s1), si ~ Eval (GetDims ds0 s0),
 si ~ Eval (GetDims ds1 s1), st ~ Eval (so0 ++ so1),
 ds0 ~ '[Eval (Eval (Rank s0) - 1)], ds1 ~ '[0], Vector v a) =>
Array v s0 a -> Array v s1 a -> Array v st a
mult = (Array v (Eval (Map (Flip GetDim s0) '[Eval (Length s0) - 1])) a
 -> a)
-> (a -> a -> a) -> Array v s0 a -> Array v s1 a -> Array v st a
forall (v :: * -> *) a b c d (ds0 :: [Nat]) (ds1 :: [Nat])
       (s0 :: [Nat]) (s1 :: [Nat]) (so0 :: [Nat]) (so1 :: [Nat])
       (st :: [Nat]) (si :: [Nat]).
(KnownNats s0, KnownNats s1, KnownNats ds0, KnownNats ds1,
 KnownNats so0, KnownNats so1, KnownNats st, KnownNats si,
 so0 ~ Eval (DeleteDims ds0 s0), so1 ~ Eval (DeleteDims ds1 s1),
 si ~ Eval (GetDims ds0 s0), si ~ Eval (GetDims ds1 s1),
 st ~ Eval (so0 ++ so1), ds0 ~ '[Eval (Eval (Rank s0) - 1)],
 ds1 ~ '[0], Vector v c, Vector v a, Vector v b, Vector v d) =>
(Array v si c -> d)
-> (a -> b -> c) -> Array v s0 a -> Array v s1 b -> Array v st d
dot Array v (Eval (Map (Flip GetDim s0) '[Eval (Length s0) - 1])) a
-> a
forall (v :: * -> *) (s :: [Nat]) a.
(Additive a, Vector v a) =>
Array v s a -> a
sumA a -> a -> a
forall a. Multiplicative a => a -> a -> a
(Mult.*)

-- | @windows xs@ are xs-sized windows of an array
--
-- >>> shape $ windows (Dims @[2,2]) (range @Vec.Vector @[4,3,2])
-- [3,2,2,2,2]
windows :: forall v w s ws a. (KnownNats s, KnownNats ws, ws ~ Eval (ExpandWindows w s), VG.Vector v a) => SNats w -> Array v s a -> Array v ws a
windows :: forall (v :: * -> *) (w :: [Nat]) (s :: [Nat]) (ws :: [Nat]) a.
(KnownNats s, KnownNats ws, ws ~ Eval (ExpandWindows w s),
 Vector v a) =>
SNats w -> Array v s a -> Array v ws a
windows SNats w
SNats Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v ws a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute (Int -> [Int] -> [Int]
S.indexWindowsL (forall (s :: [Nat]). KnownNats s => Int
rankOf @w)) Array v s a
a

-- | Find the starting positions of occurences of one array in another.
--
-- >>> a = cycle @Vec.Vector @[4,4] (range @Vec.Vector @'[3])
-- >>> i = array @Vec.Vector @[2,2] [1,2,2,0]
-- >>> pretty $ find i a
-- [[False,True,False],
--  [True,False,False],
--  [False,False,True]]
find :: forall v s' si s a r i' re ws. (Eq a, KnownNats si, KnownNats s, KnownNats s', KnownNats re, KnownNats i', KnownNat r, KnownNats ws, ws ~ Eval (ExpandWindows i' s), r ~ Eval (Rank s), i' ~ Eval (Rerank r si), re ~ Eval (DimWindows ws s), i' ~ Eval (DeleteDims re ws), s' ~ Eval (GetDims re ws), VG.Vector v a, VG.Vector v Bool, Eq (v a), VG.Vector v (Array v i' a)) => Array v si a -> Array v s a -> Array v s' Bool
find :: forall (v :: * -> *) (s' :: [Nat]) (si :: [Nat]) (s :: [Nat]) a
       (r :: Nat) (i' :: [Nat]) (re :: [Nat]) (ws :: [Nat]).
(Eq a, KnownNats si, KnownNats s, KnownNats s', KnownNats re,
 KnownNats i', KnownNat r, KnownNats ws,
 ws ~ Eval (ExpandWindows i' s), r ~ Eval (Rank s),
 i' ~ Eval (Rerank r si), re ~ Eval (DimWindows ws s),
 i' ~ Eval (DeleteDims re ws), s' ~ Eval (GetDims re ws),
 Vector v a, Vector v Bool, Eq (v a), Vector v (Array v i' a)) =>
Array v si a -> Array v s a -> Array v s' Bool
find Array v si a
i Array v s a
a = Array v s' Bool
xs
  where
    i' :: Array v i' a
i' = SNat r -> Array v si a -> Array v i' a
forall (v :: * -> *) (r :: Nat) (s :: [Nat]) (s' :: [Nat]) a.
(KnownNats s, KnownNats s', s' ~ Eval (Rerank r s), Vector v a) =>
SNat r -> Array v s a -> Array v s' a
rerank (forall (n :: Nat). KnownNat n => SNat n
SNat @r) Array v si a
i
    ws :: Array v ws a
ws = SNats i' -> Array v s a -> Array v ws a
forall (v :: * -> *) (w :: [Nat]) (s :: [Nat]) (ws :: [Nat]) a.
(KnownNats s, KnownNats ws, ws ~ Eval (ExpandWindows w s),
 Vector v a) =>
SNats w -> Array v s a -> Array v ws a
windows (forall (ns :: [Nat]). KnownNats ns => SNats ns
SNats @i') Array v s a
a
    xs :: Array v s' Bool
xs = (Array v i' a -> Bool)
-> Array v s' (Array v i' a) -> Array v s' Bool
forall (v :: * -> *) (s :: [Nat]) a b.
(KnownNats s, Vector v a, Vector v b) =>
(a -> b) -> Array v s a -> Array v s b
fmapA (Array v i' a -> Array v i' a -> Bool
forall a. Eq a => a -> a -> Bool
== Array v i' a
i') (Dims re -> Array v ws a -> Array v s' (Array v i' a)
forall (v :: * -> *) (ds :: [Nat]) (st :: [Nat]) (si :: [Nat])
       (so :: [Nat]) a.
(KnownNats st, KnownNats ds, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st),
 Vector v a, Vector v (Array v si a)) =>
Dims ds -> Array v st a -> Array v so (Array v si a)
extracts (forall (ns :: [Nat]). KnownNats ns => SNats ns
SNats @re) Array v ws 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 Vec.Vector [5,5] Int
-- >>> i = konst @[2,2] 1 :: Array Vec.Vector [2,2] Int
-- >>> pretty $ findNoOverlap i a
-- [[True,False,True,False],
--  [False,False,False,False],
--  [True,False,True,False],
--  [False,False,False,False]]
findNoOverlap :: forall v s' si s a r i' re ws. (Eq a, KnownNats si, KnownNats s, KnownNats s', KnownNats re, KnownNats i', KnownNat r, KnownNats ws, ws ~ Eval (ExpandWindows i' s), r ~ Eval (Rank s), i' ~ Eval (Rerank r si), re ~ Eval (DimWindows ws s), i' ~ Eval (DeleteDims re ws), s' ~ Eval (GetDims re ws), VG.Vector v a, VG.Vector v Bool, Eq (v a), VG.Vector v (Array v i' a), VG.Vector v [Int]) => Array v si a -> Array v s a -> Array v s' Bool
findNoOverlap :: forall (v :: * -> *) (s' :: [Nat]) (si :: [Nat]) (s :: [Nat]) a
       (r :: Nat) (i' :: [Nat]) (re :: [Nat]) (ws :: [Nat]).
(Eq a, KnownNats si, KnownNats s, KnownNats s', KnownNats re,
 KnownNats i', KnownNat r, KnownNats ws,
 ws ~ Eval (ExpandWindows i' s), r ~ Eval (Rank s),
 i' ~ Eval (Rerank r si), re ~ Eval (DimWindows ws s),
 i' ~ Eval (DeleteDims re ws), s' ~ Eval (GetDims re ws),
 Vector v a, Vector v Bool, Eq (v a), Vector v (Array v i' a),
 Vector v [Int]) =>
Array v si a -> Array v s a -> Array v s' Bool
findNoOverlap Array v si a
i Array v s a
a = Array v s' Bool
r
  where
    f :: Array v s' Bool
f = Array v si a -> Array v s a -> Array v s' Bool
forall (v :: * -> *) (s' :: [Nat]) (si :: [Nat]) (s :: [Nat]) a
       (r :: Nat) (i' :: [Nat]) (re :: [Nat]) (ws :: [Nat]).
(Eq a, KnownNats si, KnownNats s, KnownNats s', KnownNats re,
 KnownNats i', KnownNat r, KnownNats ws,
 ws ~ Eval (ExpandWindows i' s), r ~ Eval (Rank s),
 i' ~ Eval (Rerank r si), re ~ Eval (DimWindows ws s),
 i' ~ Eval (DeleteDims re ws), s' ~ Eval (GetDims re ws),
 Vector v a, Vector v Bool, Eq (v a), Vector v (Array v i' a)) =>
Array v si a -> Array v s a -> Array v s' Bool
find Array v si a
i Array v s 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
$ v [Int] -> [[Int]]
forall (v :: * -> *) a. Vector v a => v a -> [a]
VG.toList (Array v [Int] -> v [Int]
forall {k} (v :: k -> *) (a :: k). Array v a -> v a
A.asVector ([Int] -> ([Int] -> [Int]) -> Array v [Int]
forall (v :: * -> *) a.
Vector v a =>
[Int] -> ([Int] -> a) -> Array v a
A.tabulate ((\Int
x -> Int
2 Int -> Int -> Int
forall a. Multiplicative a => a -> a -> a
* Int
x Int -> Int -> Int
forall a. Subtractive a => a -> a -> a
- Int
1) (Int -> Int) -> [Int] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
sh) (\[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. Subtractive a => a -> a -> a
- Int
x0 Int -> Int -> Int
forall a. Additive a => a -> a -> a
+ Int
1) [Int]
s [Int]
sh) :: A.Array v [Int]))
    go :: Array v s' Bool -> [Int] -> Bool
go Array v s' Bool
r' [Int]
s = Array v s' Bool -> Fins s' -> Bool
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v s' Bool
f ([Int] -> Fins s'
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins [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 s' Bool -> Fins s' -> Bool
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array v s' Bool
r' (Fins s' -> Bool) -> ([Int] -> Fins s') -> [Int] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> Fins s'
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins) (([Int] -> Bool) -> [[Int]] -> [[Int]]
forall a. (a -> Bool) -> [a] -> [a]
List.filter (\[Int]
x -> [Int] -> [Int] -> Bool
S.isFinsL [Int]
x (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v s' Bool -> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape Array v s' 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. Additive a => a -> a -> a
(+) [Int]
s) ([Int] -> [[Int]]
cl (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v si a -> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape Array v si a
i)))))
    r :: Array v s' Bool
r = ([Int] -> Bool) -> Array v s' Bool
forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
([Int] -> a) -> Array v s a
unsafeTabulate (Array v s' Bool -> [Int] -> Bool
go Array v s' Bool
r)

-- | Check if the first array is a prefix of the second.
--
-- >>> isPrefixOf (array @Vec.Vector @[2,2] [0,1,4,5]) a
-- True
isPrefixOf :: forall v s' s r a. (Eq a, KnownNats s, KnownNats s', KnownNat r, KnownNats (Eval (Rerank r s)), True ~ Eval (IsSubset s' s), r ~ Eval (Rank s'), VG.Vector v a, Eq (v a)) => Array v s' a -> Array v s a -> Bool
isPrefixOf :: forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) (r :: Nat) a.
(Eq a, KnownNats s, KnownNats s', KnownNat r,
 KnownNats (Eval (Rerank r s)), 'True ~ Eval (IsSubset s' s),
 r ~ Eval (Rank s'), Vector v a, Eq (v a)) =>
Array v s' a -> Array v s a -> Bool
isPrefixOf Array v s' a
p Array v s a
a = Array v s' a
p Array v s' a -> Array v s' a -> Bool
forall a. Eq a => a -> a -> Bool
== Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) (r :: Nat) a.
(KnownNats s, KnownNats s', KnownNat r,
 KnownNats (Eval (Rerank r s)), 'True ~ Eval (IsSubset s' s),
 r ~ Eval (Rank s'), Vector v a) =>
Array v s a -> Array v s' a
cut Array v s a
a

-- | Check if the first array is a suffix of the second.
--
-- >>> isSuffixOf (array @Vec.Vector @[2,2] [18,19,22,23]) a
-- True
isSuffixOf :: forall v s' s r a. (Eq a, KnownNats s, KnownNats s', KnownNat r, KnownNats (Eval (Rerank r s)), r ~ Eval (Rank s'), True ~ Eval (IsSubset s' s), VG.Vector v a, Eq (v a)) => Array v s' a -> Array v s a -> Bool
isSuffixOf :: forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) (r :: Nat) a.
(Eq a, KnownNats s, KnownNats s', KnownNat r,
 KnownNats (Eval (Rerank r s)), r ~ Eval (Rank s'),
 'True ~ Eval (IsSubset s' s), Vector v a, Eq (v a)) =>
Array v s' a -> Array v s a -> Bool
isSuffixOf Array v s' a
p Array v s a
a = Array v s' a
p Array v s' a -> Array v s' a -> Bool
forall a. Eq a => a -> a -> Bool
== Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a (r :: Nat).
(KnownNats s, KnownNats s', KnownNat r,
 KnownNats (Eval (Rerank r s)), r ~ Eval (Rank s'),
 'True ~ Eval (IsSubset s' s), Vector v a) =>
Array v s a -> Array v s' a
cutSuffix Array v s a
a

-- | Check if the first array is an infix of the second.
--
-- >>> isInfixOf (array @Vec.Vector @[2,2] [18,19,22,23]) a
-- True
isInfixOf :: forall v s' si s a r i' re ws. (Eq a, KnownNats si, KnownNats s, KnownNats s', KnownNats re, KnownNats i', KnownNat r, KnownNats ws, ws ~ Eval (ExpandWindows i' s), r ~ Eval (Rank s), i' ~ Eval (Rerank r si), re ~ Eval (DimWindows ws s), i' ~ Eval (DeleteDims re ws), s' ~ Eval (GetDims re ws), VG.Vector v a, VG.Vector v Bool, Eq (v a), VG.Vector v (Array v i' a)) => Array v si a -> Array v s a -> Bool
isInfixOf :: forall (v :: * -> *) (s' :: [Nat]) (si :: [Nat]) (s :: [Nat]) a
       (r :: Nat) (i' :: [Nat]) (re :: [Nat]) (ws :: [Nat]).
(Eq a, KnownNats si, KnownNats s, KnownNats s', KnownNats re,
 KnownNats i', KnownNat r, KnownNats ws,
 ws ~ Eval (ExpandWindows i' s), r ~ Eval (Rank s),
 i' ~ Eval (Rerank r si), re ~ Eval (DimWindows ws s),
 i' ~ Eval (DeleteDims re ws), s' ~ Eval (GetDims re ws),
 Vector v a, Vector v Bool, Eq (v a), Vector v (Array v i' a)) =>
Array v si a -> Array v s a -> Bool
isInfixOf Array v si a
p Array v s a
a = Array v s' Bool -> Bool
forall (v :: * -> *) (s :: [Nat]).
Vector v Bool =>
Array v s Bool -> Bool
orA (Array v s' Bool -> Bool) -> Array v s' Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Array v si a -> Array v s a -> Array v s' Bool
forall (v :: * -> *) (s' :: [Nat]) (si :: [Nat]) (s :: [Nat]) a
       (r :: Nat) (i' :: [Nat]) (re :: [Nat]) (ws :: [Nat]).
(Eq a, KnownNats si, KnownNats s, KnownNats s', KnownNats re,
 KnownNats i', KnownNat r, KnownNats ws,
 ws ~ Eval (ExpandWindows i' s), r ~ Eval (Rank s),
 i' ~ Eval (Rerank r si), re ~ Eval (DimWindows ws s),
 i' ~ Eval (DeleteDims re ws), s' ~ Eval (GetDims re ws),
 Vector v a, Vector v Bool, Eq (v a), Vector v (Array v i' a)) =>
Array v si a -> Array v s a -> Array v s' Bool
find Array v si a
p Array v s a
a

-- | 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 @Vec.Vector @'[3] 0 (array @Vec.Vector @'[0] [])
-- [0,0,0]
-- >>> pretty $ fill @Vec.Vector @'[3] 0 (array @Vec.Vector @'[4] [1..4])
-- [1,2,3]
fill :: forall v s' a s. (KnownNats s, KnownNats s', VG.Vector v a, Semigroup (v a)) => a -> Array v s a -> Array v s' a
fill :: forall (v :: * -> *) (s' :: [Nat]) a (s :: [Nat]).
(KnownNats s, KnownNats s', Vector v a, Semigroup (v a)) =>
a -> Array v s a -> Array v s' a
fill a
x (Array v a
v) = v a -> Array v s' a
forall {k} (v :: k -> *) (s :: [Nat]) (a :: k). v a -> Array v s a
Array (Int -> v a -> v a
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
VG.take (Vector Int -> Int
S.size ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @s'))) (v a
v v a -> v a -> v a
forall a. Semigroup a => a -> a -> a
<> Int -> a -> v a
forall (v :: * -> *) a. Vector v a => Int -> a -> v a
VG.replicate (Vector Int -> Int
S.size ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @s')) Int -> Int -> Int
forall a. Subtractive 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.
--
-- >>> toDynamic $ cut @Vec.Vector @'[2] (array @Vec.Vector @'[4] @Int [0..3])
-- UnsafeArray [2] [0,1]
cut :: forall v s' s r a. (KnownNats s, KnownNats s', KnownNat r, KnownNats (Eval (Rerank r s)), True ~ Eval (IsSubset s' s), r ~ Eval (Rank s'), VG.Vector v a) => Array v s a -> Array v s' a
cut :: forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) (r :: Nat) a.
(KnownNats s, KnownNats s', KnownNat r,
 KnownNats (Eval (Rerank r s)), 'True ~ Eval (IsSubset s' s),
 r ~ Eval (Rank s'), Vector v a) =>
Array v s a -> Array v s' a
cut Array v s a
a = ([Int] -> [Int])
-> Array
     v
     (If
        (Eval
           (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
        (Eval
           (UnfoldrCase
              (NumIter 1)
              (If
                 (Eval
                    (Not
                       (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
                 ('Just '(1, (r - Eval (Length s)) - 1))
                 'Nothing))
         <> s)
        (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
           : Drop_ ((Eval (Length s) - r) + 1) s))
     a
-> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute [Int] -> [Int]
forall a. a -> a
id (SNat r
-> Array v s a
-> Array
     v
     (If
        (Eval
           (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
        (Eval
           (UnfoldrCase
              (NumIter 1)
              (If
                 (Eval
                    (Not
                       (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
                 ('Just '(1, (r - Eval (Length s)) - 1))
                 'Nothing))
         <> s)
        (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
           : Drop_ ((Eval (Length s) - r) + 1) s))
     a
forall (v :: * -> *) (r :: Nat) (s :: [Nat]) (s' :: [Nat]) a.
(KnownNats s, KnownNats s', s' ~ Eval (Rerank r s), Vector v a) =>
SNat r -> Array v s a -> Array v s' a
rerank (forall (n :: Nat). KnownNat n => SNat n
SNat @r) Array v s 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.
--
-- >>> toDynamic $ cutSuffix @Vec.Vector @[2,2] a
-- UnsafeArray [2,2] [18,19,22,23]
cutSuffix :: forall v s' s a r. (KnownNats s, KnownNats s', KnownNat r, KnownNats (Eval (Rerank r s)), r ~ Eval (Rank s'), True ~ Eval (IsSubset s' s), VG.Vector v a) => Array v s a -> Array v s' a
cutSuffix :: forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a (r :: Nat).
(KnownNats s, KnownNats s', KnownNat r,
 KnownNats (Eval (Rerank r s)), r ~ Eval (Rank s'),
 'True ~ Eval (IsSubset s' s), Vector v a) =>
Array v s a -> Array v s' a
cutSuffix Array v s a
a = ([Int] -> [Int])
-> Array
     v
     (If
        (Eval
           (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
        (Eval
           (UnfoldrCase
              (NumIter 1)
              (If
                 (Eval
                    (Not
                       (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
                 ('Just '(1, (r - Eval (Length s)) - 1))
                 'Nothing))
         <> s)
        (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
           : Drop_ ((Eval (Length s) - r) + 1) s))
     a
-> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
List.zipWith Int -> Int -> Int
forall a. Additive a => a -> a -> a
(+) [Int]
diffDim) Array
  v
  (If
     (Eval
        (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
     (Eval
        (UnfoldrCase
           (NumIter 1)
           (If
              (Eval
                 (Not
                    (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
              ('Just '(1, (r - Eval (Length s)) - 1))
              'Nothing))
      <> s)
     (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
        : Drop_ ((Eval (Length s) - r) + 1) s))
  a
a'
  where
    a' :: Array
  v
  (If
     (Eval
        (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
     (Eval
        (UnfoldrCase
           (NumIter 1)
           (If
              (Eval
                 (Not
                    (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
              ('Just '(1, (r - Eval (Length s)) - 1))
              'Nothing))
      <> s)
     (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
        : Drop_ ((Eval (Length s) - r) + 1) s))
  a
a' = SNat r
-> Array v s a
-> Array
     v
     (If
        (Eval
           (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
        (Eval
           (UnfoldrCase
              (NumIter 1)
              (If
                 (Eval
                    (Not
                       (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
                 ('Just '(1, (r - Eval (Length s)) - 1))
                 'Nothing))
         <> s)
        (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
           : Drop_ ((Eval (Length s) - r) + 1) s))
     a
forall (v :: * -> *) (r :: Nat) (s :: [Nat]) (s' :: [Nat]) a.
(KnownNats s, KnownNats s', s' ~ Eval (Rerank r s), Vector v a) =>
SNat r -> Array v s a -> Array v s' a
rerank (forall (n :: Nat). KnownNat n => SNat n
SNat @r) Array v s a
a
    diffDim :: [Int]
diffDim = (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
List.zipWith (-) (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array
  v
  (If
     (Eval
        (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
     (Eval
        (UnfoldrCase
           (NumIter 1)
           (If
              (Eval
                 (Not
                    (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
              ('Just '(1, (r - Eval (Length s)) - 1))
              'Nothing))
      <> s)
     (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
        : Drop_ ((Eval (Length s) - r) + 1) s))
  a
-> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape Array
  v
  (If
     (Eval
        (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
     (Eval
        (UnfoldrCase
           (NumIter 1)
           (If
              (Eval
                 (Not
                    (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
              ('Just '(1, (r - Eval (Length s)) - 1))
              'Nothing))
      <> s)
     (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
        : Drop_ ((Eval (Length s) - r) + 1) s))
  a
a')) (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @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.
--
-- >>> toDynamic $ pad @Vec.Vector @'[5] 0 (array @Vec.Vector @'[4] @Int [0..3])
-- UnsafeArray [5] [0,1,2,3,0]
pad :: forall v s' a s r. (KnownNats s, KnownNats s', KnownNat r, KnownNats (Eval (Rerank r s)), r ~ Eval (Rank s'), VG.Vector v a) => a -> Array v s a -> Array v s' a
pad :: forall (v :: * -> *) (s' :: [Nat]) a (s :: [Nat]) (r :: Nat).
(KnownNats s, KnownNats s', KnownNat r,
 KnownNats (Eval (Rerank r s)), r ~ Eval (Rank s'), Vector v a) =>
a -> Array v s a -> Array v s' a
pad a
d Array v s a
a = (Fins s' -> a) -> Array v s' a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
(Fins s -> a) -> Array v s a
tabulate (\Fins s'
s -> a -> a -> Bool -> a
forall a. a -> a -> Bool -> a
bool a
d (Array
  v
  (If
     (Eval
        (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
     (Eval
        (UnfoldrCase
           (NumIter 1)
           (If
              (Eval
                 (Not
                    (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
              ('Just '(1, (r - Eval (Length s)) - 1))
              'Nothing))
      <> s)
     (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
        : Drop_ ((Eval (Length s) - r) + 1) s))
  a
-> Fins
     (If
        (Eval
           (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
        (Eval
           (UnfoldrCase
              (NumIter 1)
              (If
                 (Eval
                    (Not
                       (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
                 ('Just '(1, (r - Eval (Length s)) - 1))
                 'Nothing))
         <> s)
        (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
           : Drop_ ((Eval (Length s) - r) + 1) s))
-> a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array
  v
  (If
     (Eval
        (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
     (Eval
        (UnfoldrCase
           (NumIter 1)
           (If
              (Eval
                 (Not
                    (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
              ('Just '(1, (r - Eval (Length s)) - 1))
              'Nothing))
      <> s)
     (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
        : Drop_ ((Eval (Length s) - r) + 1) s))
  a
a' (Fins s'
-> Fins
     (If
        (Eval
           (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
        (Eval
           (UnfoldrCase
              (NumIter 1)
              (If
                 (Eval
                    (Not
                       (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
                 ('Just '(1, (r - Eval (Length s)) - 1))
                 'Nothing))
         <> s)
        (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
           : Drop_ ((Eval (Length s) - r) + 1) s))
forall a b. a -> b
unsafeCoerce Fins s'
s)) (Fins s' -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins Fins s'
s [Int] -> [Int] -> Bool
`S.isFinsL` Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array
  v
  (If
     (Eval
        (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
     (Eval
        (UnfoldrCase
           (NumIter 1)
           (If
              (Eval
                 (Not
                    (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
              ('Just '(1, (r - Eval (Length s)) - 1))
              'Nothing))
      <> s)
     (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
        : Drop_ ((Eval (Length s) - r) + 1) s))
  a
-> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape Array
  v
  (If
     (Eval
        (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
     (Eval
        (UnfoldrCase
           (NumIter 1)
           (If
              (Eval
                 (Not
                    (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
              ('Just '(1, (r - Eval (Length s)) - 1))
              'Nothing))
      <> s)
     (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
        : Drop_ ((Eval (Length s) - r) + 1) s))
  a
a')))
  where
    a' :: Array
  v
  (If
     (Eval
        (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
     (Eval
        (UnfoldrCase
           (NumIter 1)
           (If
              (Eval
                 (Not
                    (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
              ('Just '(1, (r - Eval (Length s)) - 1))
              'Nothing))
      <> s)
     (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
        : Drop_ ((Eval (Length s) - r) + 1) s))
  a
a' = SNat r
-> Array v s a
-> Array
     v
     (If
        (Eval
           (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
        (Eval
           (UnfoldrCase
              (NumIter 1)
              (If
                 (Eval
                    (Not
                       (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
                 ('Just '(1, (r - Eval (Length s)) - 1))
                 'Nothing))
         <> s)
        (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
           : Drop_ ((Eval (Length s) - r) + 1) s))
     a
forall (v :: * -> *) (r :: Nat) (s :: [Nat]) (s' :: [Nat]) a.
(KnownNats s, KnownNats s', s' ~ Eval (Rerank r s), Vector v a) =>
SNat r -> Array v s a -> Array v s' a
rerank (forall (n :: Nat). KnownNat n => SNat n
SNat @r) Array v s a
a

-- | Left pad an array to form a new shape, supplying a default value for elements outside the shape of the old array.
--
-- >>> toDynamic $ lpad @Vec.Vector @'[5] 0 (array @Vec.Vector @'[4] [0..3])
-- UnsafeArray [5] [0,0,1,2,3]
-- >>> pretty $ lpad @Vec.Vector @[3,3] 0 (range @Vec.Vector @[2,2])
-- [[0,0,0],
--  [0,0,1],
--  [0,2,3]]
lpad :: forall v s' a s r. (KnownNats s, KnownNats s', KnownNat r, KnownNats (Eval (Rerank r s)), r ~ Eval (Rank s'), VG.Vector v a) => a -> Array v s a -> Array v s' a
lpad :: forall (v :: * -> *) (s' :: [Nat]) a (s :: [Nat]) (r :: Nat).
(KnownNats s, KnownNats s', KnownNat r,
 KnownNats (Eval (Rerank r s)), r ~ Eval (Rank s'), Vector v a) =>
a -> Array v s a -> Array v s' a
lpad a
d Array v s a
a = (Fins s' -> a) -> Array v s' a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
(Fins s -> a) -> Array v s a
tabulate (\Fins s'
s -> a -> a -> Bool -> a
forall a. a -> a -> Bool -> a
bool a
d (Array
  v
  (If
     (Eval
        (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
     (Eval
        (UnfoldrCase
           (NumIter 1)
           (If
              (Eval
                 (Not
                    (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
              ('Just '(1, (r - Eval (Length s)) - 1))
              'Nothing))
      <> s)
     (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
        : Drop_ ((Eval (Length s) - r) + 1) s))
  a
-> Fins
     (If
        (Eval
           (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
        (Eval
           (UnfoldrCase
              (NumIter 1)
              (If
                 (Eval
                    (Not
                       (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
                 ('Just '(1, (r - Eval (Length s)) - 1))
                 'Nothing))
         <> s)
        (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
           : Drop_ ((Eval (Length s) - r) + 1) s))
-> a
forall (v :: * -> *) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Array v s a -> Fins s -> a
index Array
  v
  (If
     (Eval
        (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
     (Eval
        (UnfoldrCase
           (NumIter 1)
           (If
              (Eval
                 (Not
                    (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
              ('Just '(1, (r - Eval (Length s)) - 1))
              'Nothing))
      <> s)
     (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
        : Drop_ ((Eval (Length s) - r) + 1) s))
  a
a' ([Int]
-> Fins
     (If
        (Eval
           (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
        (Eval
           (UnfoldrCase
              (NumIter 1)
              (If
                 (Eval
                    (Not
                       (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
                 ('Just '(1, (r - Eval (Length s)) - 1))
                 'Nothing))
         <> s)
        (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
           : Drop_ ((Eval (Length s) - r) + 1) s))
forall {k} (s :: k). [Int] -> Fins s
UnsafeFins ([Int]
 -> Fins
      (If
         (Eval
            (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
         (Eval
            (UnfoldrCase
               (NumIter 1)
               (If
                  (Eval
                     (Not
                        (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
                  ('Just '(1, (r - Eval (Length s)) - 1))
                  'Nothing))
          <> s)
         (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
            : Drop_ ((Eval (Length s) - r) + 1) s)))
-> [Int]
-> Fins
     (If
        (Eval
           (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
        (Eval
           (UnfoldrCase
              (NumIter 1)
              (If
                 (Eval
                    (Not
                       (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
                 ('Just '(1, (r - Eval (Length s)) - 1))
                 'Nothing))
         <> s)
        (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
           : Drop_ ((Eval (Length s) - r) + 1) s))
forall a b. (a -> b) -> a -> b
$ Fins s' -> [Int]
olds Fins s'
s)) (Fins s' -> [Int]
olds Fins s'
s [Int] -> [Int] -> Bool
`S.isFinsL` Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array
  v
  (If
     (Eval
        (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
     (Eval
        (UnfoldrCase
           (NumIter 1)
           (If
              (Eval
                 (Not
                    (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
              ('Just '(1, (r - Eval (Length s)) - 1))
              'Nothing))
      <> s)
     (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
        : Drop_ ((Eval (Length s) - r) + 1) s))
  a
-> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape Array
  v
  (If
     (Eval
        (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
     (Eval
        (UnfoldrCase
           (NumIter 1)
           (If
              (Eval
                 (Not
                    (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
              ('Just '(1, (r - Eval (Length s)) - 1))
              'Nothing))
      <> s)
     (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
        : Drop_ ((Eval (Length s) - r) + 1) s))
  a
a')))
  where
    a' :: Array
  v
  (If
     (Eval
        (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
     (Eval
        (UnfoldrCase
           (NumIter 1)
           (If
              (Eval
                 (Not
                    (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
              ('Just '(1, (r - Eval (Length s)) - 1))
              'Nothing))
      <> s)
     (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
        : Drop_ ((Eval (Length s) - r) + 1) s))
  a
a' = SNat r
-> Array v s a
-> Array
     v
     (If
        (Eval
           (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
        (Eval
           (UnfoldrCase
              (NumIter 1)
              (If
                 (Eval
                    (Not
                       (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
                 ('Just '(1, (r - Eval (Length s)) - 1))
                 'Nothing))
         <> s)
        (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
           : Drop_ ((Eval (Length s) - r) + 1) s))
     a
forall (v :: * -> *) (r :: Nat) (s :: [Nat]) (s' :: [Nat]) a.
(KnownNats s, KnownNats s', s' ~ Eval (Rerank r s), Vector v a) =>
SNat r -> Array v s a -> Array v s' a
rerank (forall (n :: Nat). KnownNat n => SNat n
SNat @r) Array v s a
a
    gap :: [Int]
gap = (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
List.zipWith (-) (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @s') (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array
  v
  (If
     (Eval
        (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
     (Eval
        (UnfoldrCase
           (NumIter 1)
           (If
              (Eval
                 (Not
                    (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
              ('Just '(1, (r - Eval (Length s)) - 1))
              'Nothing))
      <> s)
     (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
        : Drop_ ((Eval (Length s) - r) + 1) s))
  a
-> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape Array
  v
  (If
     (Eval
        (Not (OrdCond (CmpNat r (Eval (Length s))) 'True 'True 'False)))
     (Eval
        (UnfoldrCase
           (NumIter 1)
           (If
              (Eval
                 (Not
                    (OrdCond (CmpNat (r - Eval (Length s)) 0) 'True 'True 'False)))
              ('Just '(1, (r - Eval (Length s)) - 1))
              'Nothing))
      <> s)
     (Eval (Foldr (*) 1 (Take_ ((Eval (Length s) - r) + 1) s))
        : Drop_ ((Eval (Length s) - r) + 1) s))
  a
a'))
    olds :: Fins s' -> [Int]
olds Fins s'
s = (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
List.zipWith (-) (Fins s' -> [Int]
forall {k} (s :: k). Fins s -> [Int]
fromFins Fins s'
s) [Int]
gap

-- | Reshape an array (with the same number of elements).
--
-- >>> pretty $ reshape @Vec.Vector @[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 :: forall v s' s a. (Eval (Size s) ~ Eval (Size s'), KnownNats s, KnownNats s', VG.Vector v a) => Array v s a -> Array v s' a
reshape :: forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(Eval (Size s) ~ Eval (Size s'), KnownNats s, KnownNats s',
 Vector v a) =>
Array v s a -> Array v s' a
reshape = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Vector Int -> [Int]) -> ([Int] -> Vector Int) -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Int -> Vector Int
S.shapen ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
s) (Int -> Vector Int) -> ([Int] -> Int) -> [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') (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)
  where
    s :: [Int]
s = forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @s
    s' :: [Int]
s' = forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @s'

-- | Make an Array single dimensional.
--
-- >>> pretty $ flat (range @Vec.Vector @[2,2])
-- [0,1,2,3]
-- >>> pretty (flat $ toScalar @Vec.Vector 0)
-- [0]
flat :: forall v s' s a. (KnownNats s, KnownNats s', s' ~ '[Eval (Size s)], VG.Vector v a) => Array v s a -> Array v s' a
flat :: forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', s' ~ '[Eval (Size s)], Vector v a) =>
Array v s a -> Array v s' a
flat Array v s a
a = Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
Array v s a -> Array v s' a
unsafeModifyShape Array v s a
a

-- | Reshape an array, repeating the original array. The shape of the array should be a suffix of the new shape.
--
-- >>> pretty $ repeat @Vec.Vector @[2,2,2] (array @Vec.Vector @'[2] [1,2])
-- [[[1,2],
--   [1,2]],
--  [[1,2],
--   [1,2]]]
--
-- > repeat ds (toScalar @Vec.Vector x) == konst ds x
repeat :: forall v s' s a. (KnownNats s, KnownNats s', Eval (IsPrefixOf s s') ~ True, VG.Vector v a) => Array v s a -> Array v s' a
repeat :: forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Eval (IsPrefixOf s s') ~ 'True,
 Vector v a) =>
Array v s a -> Array v s' a
repeat Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
List.drop ([Int] -> Int
forall a. [a] -> Int
S.rankL (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @s') Int -> Int -> Int
forall a. Subtractive a => a -> a -> a
- Array v s a -> Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Int
rank Array v s a
a)) Array v s a
a

-- | Reshape an array, cycling through the elements without regard to the original shape.
--
-- >>> pretty $ cycle @Vec.Vector @[2,2,2] (array @Vec.Vector @'[3] [1,2,3])
-- [[[1,2],
--   [3,1]],
--  [[2,3],
--   [1,2]]]
cycle :: forall v s' s a. (KnownNats s, KnownNats s', VG.Vector v a) => Array v s a -> Array v s' a
cycle :: forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
Array v s a -> Array v s' a
cycle Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Vector Int -> [Int]) -> ([Int] -> Vector Int) -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Int -> Vector Int
S.shapen (Array v s a -> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape Array v s a
a) (Int -> Vector Int) -> ([Int] -> Int) -> [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 s a -> Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Int
size Array v s a
a) (Int -> Int) -> ([Int] -> Int) -> [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 (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @s')) (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) Array v s 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 (SNat @4) a)
-- [1,2,3,4]
-- >>> shape (rerank (SNat @2) a)
-- [6,4]
--
-- > flat == rerank 1
rerank :: forall v r s s' a. (KnownNats s, KnownNats s', s' ~ Eval (Rerank r s), VG.Vector v a) => SNat r -> Array v s a -> Array v s' a
rerank :: forall (v :: * -> *) (r :: Nat) (s :: [Nat]) (s' :: [Nat]) a.
(KnownNats s, KnownNats s', s' ~ Eval (Rerank r s), Vector v a) =>
SNat r -> Array v s a -> Array v s' a
rerank SNat r
_ Array v s a
a = Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
Array v s a -> Array v s' a
unsafeModifyShape Array v s a
a

-- | Change the order of dimensions.
--
-- >>> pretty $ reorder (Dims @[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 :: forall v ds s s' a. (KnownNats s, KnownNats s', s' ~ Eval (Reorder s ds), VG.Vector v a) => SNats ds -> Array v s a -> Array v s' a
reorder :: forall (v :: * -> *) (ds :: [Nat]) (s :: [Nat]) (s' :: [Nat]) a.
(KnownNats s, KnownNats s', s' ~ Eval (Reorder s ds),
 Vector v a) =>
SNats ds -> Array v s a -> Array v s' a
reorder SNats ds
SNats Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute (\[Int]
s -> [Int] -> [Int] -> [Int] -> [Int]
S.insertDimsL (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @ds) [Int]
s []) Array v s a
a

-- | Remove single dimensions.
--
-- >>> let sq = array [1..24] :: Array Vec.Vector '[2,1,3,4,1] Int
-- >>> shape $ squeeze sq
-- [2,3,4]
--
-- >>> shape $ squeeze (singleton @Vec.Vector 0)
-- []
squeeze :: forall v s t a. (KnownNats s, KnownNats t, t ~ Eval (Squeeze s), VG.Vector v a) => Array v s a -> Array v t a
squeeze :: forall (v :: * -> *) (s :: [Nat]) (t :: [Nat]) a.
(KnownNats s, KnownNats t, t ~ Eval (Squeeze s), Vector v a) =>
Array v s a -> Array v t a
squeeze = Array v s a -> Array v t a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
Array v s a -> Array v s' a
unsafeModifyShape

-- | Insert a single dimension at the supplied position.
--
-- >>> shape $ elongate (SNat @1) a
-- [2,1,3,4]
-- >>> toDynamic $ elongate (SNat @0) (toScalar @Vec.Vector 1)
-- UnsafeArray [1] [1]
elongate :: (KnownNats s, KnownNats s', s' ~ Eval (InsertDim d 1 s), VG.Vector v a) => Dim d -> Array v s a -> Array v s' a
elongate :: forall (s :: [Nat]) (s' :: [Nat]) (d :: Nat) (v :: * -> *) a.
(KnownNats s, KnownNats s', s' ~ Eval (InsertDim d 1 s),
 Vector v a) =>
Dim d -> Array v s a -> Array v s' a
elongate Dim d
_ Array v s a
a = Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
Array v s a -> Array v s' a
unsafeModifyShape Array v s a
a

-- | Reverse indices eg transposes the element A/ijk/ to A/kji/.
--
-- >>> (transpose a) ! [1,0,0] == a ! [0,0,1]
-- True
-- >>> pretty $ transpose (array @Vec.Vector @[2,2,2] [1..8])
-- [[[1,5],
--   [3,7]],
--  [[2,6],
--   [4,8]]]
transpose :: forall v a s s'. (KnownNats s, KnownNats s', s' ~ Eval (Reverse s), VG.Vector v a) => Array v s a -> Array v s' a
transpose :: forall (v :: * -> *) a (s :: [Nat]) (s' :: [Nat]).
(KnownNats s, KnownNats s', s' ~ Eval (Reverse s), Vector v a) =>
Array v s a -> Array v s' a
transpose Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute [Int] -> [Int]
forall a. [a] -> [a]
List.reverse Array v s a
a

-- | Inflate (or replicate) an array by inserting a new dimension given a supplied dimension and size.
--
-- >>> pretty $ inflate (SNat @0) (SNat @2) (array @Vec.Vector @'[3] [0,1,2])
-- [[0,1,2],
--  [0,1,2]]
inflate :: forall v s' s d x a. (KnownNats s, KnownNats s', s' ~ Eval (InsertDim d x s), VG.Vector v a) => Dim d -> SNat x -> Array v s a -> Array v s' a
inflate :: forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) (d :: Nat)
       (x :: Nat) a.
(KnownNats s, KnownNats s', s' ~ Eval (InsertDim d x s),
 Vector v a) =>
Dim d -> SNat x -> Array v s a -> Array v s' a
inflate SNat d
SNat SNat x
_ Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute (Int -> [Int] -> [Int]
S.deleteDimL (forall (n :: Nat). KnownNat n => Int
valueOf @d)) Array v s a
a

-- | Intercalate an array along dimensions.
--
-- >>> pretty $ intercalate (SNat @2) (konst @[2,3] 0 :: Array Vec.Vector [2,3] Int) 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 :: forall v d ds n n' s si st a. (KnownNats s, KnownNats si, KnownNats st, KnownNats ds, KnownNat n, KnownNat n', ds ~ '[d], si ~ Eval (DeleteDim d s), n ~ Eval (GetDim d s), n' ~ Eval ((Fcf.-) (Eval ((Fcf.+) n n)) 1), st ~ Eval (InsertDim d n' si), VG.Vector v a, VG.Vector v (Array v si a)) => Dim d -> Array v si a -> Array v s a -> Array v st a
intercalate :: forall (v :: * -> *) (d :: Nat) (ds :: [Nat]) (n :: Nat)
       (n' :: Nat) (s :: [Nat]) (si :: [Nat]) (st :: [Nat]) a.
(KnownNats s, KnownNats si, KnownNats st, KnownNats ds, KnownNat n,
 KnownNat n', ds ~ '[d], si ~ Eval (DeleteDim d s),
 n ~ Eval (GetDim d s), n' ~ Eval (Eval (n + n) - 1),
 st ~ Eval (InsertDim d n' si), Vector v a,
 Vector v (Array v si a)) =>
Dim d -> Array v si a -> Array v s a -> Array v st a
intercalate SNat d
SNat Array v si a
i Array v s a
a =
  Dims ds -> Array v '[n'] (Array v si a) -> Array v st a
forall (v :: * -> *) a (ds :: [Nat]) (si :: [Nat]) (so :: [Nat])
       (st :: [Nat]).
(KnownNats ds, KnownNats st, KnownNats si, KnownNats so,
 Eval (InsertDims ds so si) ~ st, Vector v (Array v si a),
 Vector v a) =>
Dims ds -> Array v so (Array v si a) -> Array v st a
joins
    (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds)
    ( forall (v :: * -> *) (n :: Nat) a t.
(FromVector t v a, KnownNat n, Vector v a) =>
t -> Array v '[n] a
vector @v @n'
        ( Array v si a -> [Array v si a] -> [Array v si a]
forall a. a -> [a] -> [a]
List.intersperse
            Array v si a
i
            (Array v '[n] (Array v si a) -> [Array v si a]
forall (v :: * -> *) (s :: [Nat]) a.
Vector v a =>
Array v s a -> [a]
toListA (Dims ds -> Array v s a -> Array v '[n] (Array v si a)
forall (v :: * -> *) (ds :: [Nat]) (st :: [Nat]) (si :: [Nat])
       (so :: [Nat]) a.
(KnownNats st, KnownNats ds, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st),
 Vector v a, Vector v (Array v si a)) =>
Dims ds -> Array v st a -> Array v so (Array v si a)
extracts (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) Array v s a
a))
        )
    )

-- | Intersperse an element along dimensions.
--
-- >>> pretty $ intersperse (SNat @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 :: forall v d ds n n' s si st a. (KnownNats s, KnownNats si, KnownNats st, KnownNats ds, KnownNat n, KnownNat n', ds ~ '[d], si ~ Eval (DeleteDim d s), n ~ Eval (GetDim d s), n' ~ n + n - 1, st ~ Eval (InsertDim d n' si), VG.Vector v a, VG.Vector v (Array v si a)) => Dim d -> a -> Array v s a -> Array v st a
intersperse :: forall (v :: * -> *) (d :: Nat) (ds :: [Nat]) (n :: Nat)
       (n' :: Nat) (s :: [Nat]) (si :: [Nat]) (st :: [Nat]) a.
(KnownNats s, KnownNats si, KnownNats st, KnownNats ds, KnownNat n,
 KnownNat n', ds ~ '[d], si ~ Eval (DeleteDim d s),
 n ~ Eval (GetDim d s), n' ~ ((n + n) - 1),
 st ~ Eval (InsertDim d n' si), Vector v a,
 Vector v (Array v si a)) =>
Dim d -> a -> Array v s a -> Array v st a
intersperse (SNat d
SNat :: SNat d) a
x Array v s a
a = SNat d -> Array v si a -> Array v s a -> Array v st a
forall (v :: * -> *) (d :: Nat) (ds :: [Nat]) (n :: Nat)
       (n' :: Nat) (s :: [Nat]) (si :: [Nat]) (st :: [Nat]) a.
(KnownNats s, KnownNats si, KnownNats st, KnownNats ds, KnownNat n,
 KnownNat n', ds ~ '[d], si ~ Eval (DeleteDim d s),
 n ~ Eval (GetDim d s), n' ~ Eval (Eval (n + n) - 1),
 st ~ Eval (InsertDim d n' si), Vector v a,
 Vector v (Array v si a)) =>
Dim d -> Array v si a -> Array v s a -> Array v st a
intercalate (forall (n :: Nat). KnownNat n => SNat n
SNat @d) (forall (s :: [Nat]) (v :: * -> *) a.
(KnownNats s, Vector v a) =>
a -> Array v s a
konst @si a
x) Array v s a
a

-- | Concatenate dimensions, creating a new dimension at the supplied postion.
--
-- >>> pretty $ concats (Dims @[0,1]) (SNat @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 :: forall v s s' newd ds a. (KnownNats s, KnownNats s', s' ~ Eval (ConcatDims ds newd s), VG.Vector v a) => Dims ds -> SNat newd -> Array v s a -> Array v s' a
concats :: forall (v :: * -> *) (s :: [Nat]) (s' :: [Nat]) (newd :: Nat)
       (ds :: [Nat]) a.
(KnownNats s, KnownNats s', s' ~ Eval (ConcatDims ds newd s),
 Vector v a) =>
Dims ds -> SNat newd -> Array v s a -> Array v s' a
concats SNats ds
SNats SNat newd
SNat Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s' a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute ([Int] -> Int -> [Int] -> [Int] -> [Int]
S.unconcatDimsIndexL [Int]
ds Int
n (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v s a -> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape Array v s a
a))) Array v s a
a
  where
    n :: Int
n = forall (n :: Nat). KnownNat n => Int
valueOf @newd
    ds :: [Int]
ds = forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @ds

-- | Reverses element order along specified dimensions.
--
-- >>> pretty $ reverses (Dims @[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 :: forall v ds s a. (KnownNats s, VG.Vector v a) => Dims ds -> Array v s a -> Array v s a
reverses :: forall (v :: * -> *) (ds :: [Nat]) (s :: [Nat]) a.
(KnownNats s, Vector v a) =>
Dims ds -> Array v s a -> Array v s a
reverses SNats ds
SNats Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute ([Int] -> [Int] -> [Int] -> [Int]
S.reverseIndexL (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @ds) (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array v s a -> Vector Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Vector Int
shape Array v s a
a))) Array v s a
a

-- | Rotate an array by/along dimensions & offsets.
--
-- >>> pretty $ rotates (Dims @'[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 :: forall v a ds s. (KnownNats s, True ~ Eval (IsDims ds s), VG.Vector v a) => Dims ds -> [Int] -> Array v s a -> Array v s a
rotates :: forall (v :: * -> *) a (ds :: [Nat]) (s :: [Nat]).
(KnownNats s, 'True ~ Eval (IsDims ds s), Vector v a) =>
Dims ds -> [Int] -> Array v s a -> Array v s a
rotates SNats ds
SNats [Int]
rs Array v s a
a = ([Int] -> [Int]) -> Array v s a -> Array v s a
forall (v :: * -> *) (s' :: [Nat]) (s :: [Nat]) a.
(KnownNats s, KnownNats s', Vector v a) =>
([Int] -> [Int]) -> Array v s a -> Array v s' a
unsafeBackpermute ([Int] -> [Int] -> [Int] -> [Int] -> [Int]
S.rotatesIndexL (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @ds) [Int]
rs (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @s)) Array v s a
a

-- | Sort an array along the supplied dimensions.
--
-- >>> pretty $ sorts (Dims @'[0]) (array @Vec.Vector @[2,2] [2,3,1,4])
-- [[1,4],
--  [2,3]]
-- >>> pretty $ sorts (Dims @'[1]) (array @Vec.Vector @[2,2] [2,3,1,4])
-- [[2,3],
--  [1,4]]
-- >>> pretty $ sorts (Dims @[0,1]) (array @Vec.Vector @[2,2] [2,3,1,4])
-- [[1,2],
--  [3,4]]
sorts :: forall v ds s a si so. (Ord a, KnownNats s, KnownNats si, KnownNats so, si ~ Eval (DeleteDims ds s), so ~ Eval (GetDims ds s), s ~ Eval (InsertDims ds so si), VG.Vector v a, Ord (v a), VG.Vector v (Array v si a)) => Dims ds -> Array v s a -> Array v s a
sorts :: forall (v :: * -> *) (ds :: [Nat]) (s :: [Nat]) a (si :: [Nat])
       (so :: [Nat]).
(Ord a, KnownNats s, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds s), so ~ Eval (GetDims ds s),
 s ~ Eval (InsertDims ds so si), Vector v a, Ord (v a),
 Vector v (Array v si a)) =>
Dims ds -> Array v s a -> Array v s a
sorts SNats ds
SNats Array v s a
a = SNats ds -> Array v so (Array v si a) -> Array v s a
forall (v :: * -> *) a (ds :: [Nat]) (si :: [Nat]) (so :: [Nat])
       (st :: [Nat]).
(KnownNats ds, KnownNats st, KnownNats si, KnownNats so,
 Eval (InsertDims ds so si) ~ st, Vector v (Array v si a),
 Vector v a) =>
Dims ds -> Array v so (Array v si a) -> Array v st a
joins (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) (Array v so (Array v si a) -> Array v s a)
-> Array v so (Array v si a) -> Array v s a
forall a b. (a -> b) -> a -> b
$ (v (Array v si a) -> v (Array v si a))
-> Array v so (Array v si a) -> Array v so (Array v si a)
forall (v :: * -> *) (s :: [Nat]) a b.
(KnownNats s, Vector v a, Vector v b) =>
(v a -> v b) -> Array v s a -> Array v s b
unsafeModifyVector (Vector (Array v si a) -> v (Array v si a)
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert (Vector (Array v si a) -> v (Array v si a))
-> (v (Array v si a) -> Vector (Array v si a))
-> v (Array v si a)
-> v (Array v si a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Array v si a) -> Vector (Array v si a)
forall a. Ord a => Vector a -> Vector a
sortV (Vector (Array v si a) -> Vector (Array v si a))
-> (v (Array v si a) -> Vector (Array v si a))
-> v (Array v si a)
-> Vector (Array v si a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v (Array v si a) -> Vector (Array v si a)
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert) (SNats ds -> Array v s a -> Array v so (Array v si a)
forall (v :: * -> *) (ds :: [Nat]) (st :: [Nat]) (si :: [Nat])
       (so :: [Nat]) a.
(KnownNats st, KnownNats ds, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st),
 Vector v a, Vector v (Array v si a)) =>
Dims ds -> Array v st a -> Array v so (Array v si a)
extracts (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) Array v s a
a)

-- | The indices into the array if it were sorted by a comparison function along the dimensions supplied.
--
-- >>> import Data.Ord (Down (..))
-- >>> toDynamic $ sortsBy (Dims @'[0]) (fmap Down) (array @Vec.Vector @[2,2] [2,3,1,4])
-- UnsafeArray [2,2] [2,3,1,4]
sortsBy :: forall v ds s a b si so. (Ord b, KnownNats s, KnownNats si, KnownNats so, si ~ Eval (DeleteDims ds s), so ~ Eval (GetDims ds s), s ~ Eval (InsertDims ds so si), VG.Vector v a, VG.Vector v (Array v si a), Ord (v b)) => Dims ds -> (Array v si a -> Array v si b) -> Array v s a -> Array v s a
sortsBy :: forall (v :: * -> *) (ds :: [Nat]) (s :: [Nat]) a b (si :: [Nat])
       (so :: [Nat]).
(Ord b, KnownNats s, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds s), so ~ Eval (GetDims ds s),
 s ~ Eval (InsertDims ds so si), Vector v a,
 Vector v (Array v si a), Ord (v b)) =>
Dims ds
-> (Array v si a -> Array v si b) -> Array v s a -> Array v s a
sortsBy SNats ds
SNats Array v si a -> Array v si b
c Array v s a
a = SNats ds -> Array v so (Array v si a) -> Array v s a
forall (v :: * -> *) a (ds :: [Nat]) (si :: [Nat]) (so :: [Nat])
       (st :: [Nat]).
(KnownNats ds, KnownNats st, KnownNats si, KnownNats so,
 Eval (InsertDims ds so si) ~ st, Vector v (Array v si a),
 Vector v a) =>
Dims ds -> Array v so (Array v si a) -> Array v st a
joins (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) (Array v so (Array v si a) -> Array v s a)
-> Array v so (Array v si a) -> Array v s a
forall a b. (a -> b) -> a -> b
$ (v (Array v si a) -> v (Array v si a))
-> Array v so (Array v si a) -> Array v so (Array v si a)
forall (v :: * -> *) (s :: [Nat]) a b.
(KnownNats s, Vector v a, Vector v b) =>
(v a -> v b) -> Array v s a -> Array v s b
unsafeModifyVector (Vector (Array v si a) -> v (Array v si a)
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert (Vector (Array v si a) -> v (Array v si a))
-> (v (Array v si a) -> Vector (Array v si a))
-> v (Array v si a)
-> v (Array v si a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Array v si a -> Array v si b)
-> Vector (Array v si a) -> Vector (Array v si a)
forall b a. Ord b => (a -> b) -> Vector a -> Vector a
sortByV Array v si a -> Array v si b
c (Vector (Array v si a) -> Vector (Array v si a))
-> (v (Array v si a) -> Vector (Array v si a))
-> v (Array v si a)
-> Vector (Array v si a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v (Array v si a) -> Vector (Array v si a)
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert) (SNats ds -> Array v s a -> Array v so (Array v si a)
forall (v :: * -> *) (ds :: [Nat]) (st :: [Nat]) (si :: [Nat])
       (so :: [Nat]) a.
(KnownNats st, KnownNats ds, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st),
 Vector v a, Vector v (Array v si a)) =>
Dims ds -> Array v st a -> Array v so (Array v si a)
extracts (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) Array v s a
a)

-- | The indices into the array if it were sorted along the dimensions supplied.
--
-- >>> orders (Dims @'[0]) (array @Vec.Vector @[2,2] [2,3,1,4])
-- [1,0]
orders :: forall v ds s a si so. (Ord a, KnownNats s, KnownNats si, KnownNats so, si ~ Eval (DeleteDims ds s), so ~ Eval (GetDims ds s), s ~ Eval (InsertDims ds so si), VG.Vector v a, VG.Vector v Int, Ord (v a), VG.Vector v (Array v si a)) => Dims ds -> Array v s a -> Array v so Int
orders :: forall (v :: * -> *) (ds :: [Nat]) (s :: [Nat]) a (si :: [Nat])
       (so :: [Nat]).
(Ord a, KnownNats s, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds s), so ~ Eval (GetDims ds s),
 s ~ Eval (InsertDims ds so si), Vector v a, Vector v Int,
 Ord (v a), Vector v (Array v si a)) =>
Dims ds -> Array v s a -> Array v so Int
orders SNats ds
SNats Array v s a
a = (v (Array v si a) -> v Int)
-> Array v so (Array v si a) -> Array v so Int
forall (v :: * -> *) (s :: [Nat]) a b.
(KnownNats s, Vector v a, Vector v b) =>
(v a -> v b) -> Array v s a -> Array v s b
unsafeModifyVector (Vector Int -> v Int
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert (Vector Int -> v Int)
-> (v (Array v si a) -> Vector Int) -> v (Array v si a) -> v Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Array v si a) -> Vector Int
forall a. Ord a => Vector a -> Vector Int
orderV (Vector (Array v si a) -> Vector Int)
-> (v (Array v si a) -> Vector (Array v si a))
-> v (Array v si a)
-> Vector Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v (Array v si a) -> Vector (Array v si a)
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert) (SNats ds -> Array v s a -> Array v so (Array v si a)
forall (v :: * -> *) (ds :: [Nat]) (st :: [Nat]) (si :: [Nat])
       (so :: [Nat]) a.
(KnownNats st, KnownNats ds, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st),
 Vector v a, Vector v (Array v si a)) =>
Dims ds -> Array v st a -> Array v so (Array v si a)
extracts (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) Array v s a
a)

-- | The indices into the array if it were sorted by a comparison function along the dimensions supplied.
--
-- >>> import Data.Ord (Down (..))
-- >>> ordersBy (Dims @'[0]) (fmap Down) (array @Vec.Vector @[2,2] [2,3,1,4])
-- [0,1]
ordersBy :: forall v ds s a b si so. (Ord b, KnownNats s, KnownNats si, KnownNats so, si ~ Eval (DeleteDims ds s), so ~ Eval (GetDims ds s), s ~ Eval (InsertDims ds so si), VG.Vector v a, VG.Vector v Int, VG.Vector v (Array v si a), Ord (v b)) => Dims ds -> (Array v si a -> Array v si b) -> Array v s a -> Array v so Int
ordersBy :: forall (v :: * -> *) (ds :: [Nat]) (s :: [Nat]) a b (si :: [Nat])
       (so :: [Nat]).
(Ord b, KnownNats s, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds s), so ~ Eval (GetDims ds s),
 s ~ Eval (InsertDims ds so si), Vector v a, Vector v Int,
 Vector v (Array v si a), Ord (v b)) =>
Dims ds
-> (Array v si a -> Array v si b) -> Array v s a -> Array v so Int
ordersBy SNats ds
SNats Array v si a -> Array v si b
c Array v s a
a = (v (Array v si a) -> v Int)
-> Array v so (Array v si a) -> Array v so Int
forall (v :: * -> *) (s :: [Nat]) a b.
(KnownNats s, Vector v a, Vector v b) =>
(v a -> v b) -> Array v s a -> Array v s b
unsafeModifyVector (Vector Int -> v Int
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert (Vector Int -> v Int)
-> (v (Array v si a) -> Vector Int) -> v (Array v si a) -> v Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Array v si a -> Array v si b)
-> Vector (Array v si a) -> Vector Int
forall b a. Ord b => (a -> b) -> Vector a -> Vector Int
orderByV Array v si a -> Array v si b
c (Vector (Array v si a) -> Vector Int)
-> (v (Array v si a) -> Vector (Array v si a))
-> v (Array v si a)
-> Vector Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v (Array v si a) -> Vector (Array v si a)
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert) (SNats ds -> Array v s a -> Array v so (Array v si a)
forall (v :: * -> *) (ds :: [Nat]) (st :: [Nat]) (si :: [Nat])
       (so :: [Nat]) a.
(KnownNats st, KnownNats ds, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st),
 Vector v a, Vector v (Array v si a)) =>
Dims ds -> Array v st a -> Array v so (Array v si a)
extracts (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) Array v s a
a)

-- | Apply a binary array function to two arrays with matching shapes across the supplied (matching) dimensions.
--
-- >>> a = array @Vec.Vector @[2,3] [0..5]
-- >>> b = array @Vec.Vector @'[3] [6..8]
-- >>> pretty $ telecasts (Dims @'[1]) (Dims @'[0]) (concatenate (SNat @0)) a b
-- [[0,3,6],
--  [1,4,7],
--  [2,5,8]]
telecasts :: forall v sa sb sc sia sib sic ma mb a b c soa sob ds. (KnownNats sa, KnownNats sb, KnownNats sc, KnownNats sia, KnownNats sib, KnownNats sic, KnownNats soa, KnownNats sob, KnownNats ds, ds ~ Eval (DimsOf soa), sia ~ Eval (DeleteDims ma sa), sib ~ Eval (DeleteDims mb sb), soa ~ Eval (GetDims ma sa), sob ~ Eval (GetDims mb sb), soa ~ sob, sc ~ Eval (InsertDims ds soa sic), VG.Vector v a, VG.Vector v b, VG.Vector v c, VG.Vector v (Array v sia a), VG.Vector v (Array v sib b), VG.Vector v (Array v sic c)) => SNats ma -> SNats mb -> (Array v sia a -> Array v sib b -> Array v sic c) -> Array v sa a -> Array v sb b -> Array v sc c
telecasts :: forall (v :: * -> *) (sa :: [Nat]) (sb :: [Nat]) (sc :: [Nat])
       (sia :: [Nat]) (sib :: [Nat]) (sic :: [Nat]) (ma :: [Nat])
       (mb :: [Nat]) a b c (soa :: [Nat]) (sob :: [Nat]) (ds :: [Nat]).
(KnownNats sa, KnownNats sb, KnownNats sc, KnownNats sia,
 KnownNats sib, KnownNats sic, KnownNats soa, KnownNats sob,
 KnownNats ds, ds ~ Eval (DimsOf soa),
 sia ~ Eval (DeleteDims ma sa), sib ~ Eval (DeleteDims mb sb),
 soa ~ Eval (GetDims ma sa), sob ~ Eval (GetDims mb sb), soa ~ sob,
 sc ~ Eval (InsertDims ds soa sic), Vector v a, Vector v b,
 Vector v c, Vector v (Array v sia a), Vector v (Array v sib b),
 Vector v (Array v sic c)) =>
SNats ma
-> SNats mb
-> (Array v sia a -> Array v sib b -> Array v sic c)
-> Array v sa a
-> Array v sb b
-> Array v sc c
telecasts SNats ma
SNats SNats mb
SNats Array v sia a -> Array v sib b -> Array v sic c
f Array v sa a
a Array v sb b
b = Array v sob (Array v sic c) -> Array v sc c
forall (v :: * -> *) a (si :: [Nat]) (so :: [Nat]) (st :: [Nat])
       (ds :: [Nat]).
(KnownNats st, KnownNats si, KnownNats so, KnownNats ds,
 ds ~ Eval (DimsOf so), st ~ Eval (InsertDims ds so si),
 Vector v (Array v si a), Vector v a) =>
Array v so (Array v si a) -> Array v st a
join ((Array v sia a -> Array v sib b -> Array v sic c)
-> Array v sob (Array v sia a)
-> Array v sob (Array v sib b)
-> Array v sob (Array v sic c)
forall (s :: [Nat]) (v :: * -> *) a b c.
(KnownNats s, Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v s a -> Array v s b -> Array v s c
zipWith Array v sia a -> Array v sib b -> Array v sic c
f (SNats ma -> Array v sa a -> Array v sob (Array v sia a)
forall (v :: * -> *) (ds :: [Nat]) (st :: [Nat]) (si :: [Nat])
       (so :: [Nat]) a.
(KnownNats st, KnownNats ds, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st),
 Vector v a, Vector v (Array v si a)) =>
Dims ds -> Array v st a -> Array v so (Array v si a)
extracts (forall (ns :: [Nat]). KnownNats ns => SNats ns
SNats @ma) Array v sa a
a) (SNats mb -> Array v sb b -> Array v sob (Array v sib b)
forall (v :: * -> *) (ds :: [Nat]) (st :: [Nat]) (si :: [Nat])
       (so :: [Nat]) a.
(KnownNats st, KnownNats ds, KnownNats si, KnownNats so,
 si ~ Eval (DeleteDims ds st), so ~ Eval (GetDims ds st),
 Vector v a, Vector v (Array v si a)) =>
Dims ds -> Array v st a -> Array v so (Array v si a)
extracts (forall (ns :: [Nat]). KnownNats ns => SNats ns
SNats @mb) Array v sb b
b))

-- | Apply a binary array function to two arrays where the shape of the first array is a prefix of the second array.
--
-- >>> a = array @Vec.Vector @[2,3] [0..5]
-- >>> pretty $ transmit (zipWith (+)) (toScalar @Vec.Vector 1) a
-- [[1,2,3],
--  [4,5,6]]
transmit :: forall v sa sb sc a b c ds sib sic sob. (KnownNats sa, KnownNats sb, KnownNats sc, KnownNats ds, KnownNats sib, KnownNats sic, KnownNats sob, ds ~ Eval (EnumFromTo (Eval (Rank sa)) (Eval (Rank sb) - 1)), sib ~ Eval (DeleteDims ds sb), sob ~ Eval (GetDims ds sb), sb ~ Eval (InsertDims ds sob sib), sc ~ Eval (InsertDims ds sob sic), True ~ Eval (IsPrefixOf sa sb), VG.Vector v a, VG.Vector v b, VG.Vector v c, VG.Vector v (Array v sib b), VG.Vector v (Array v sic c)) => (Array v sa a -> Array v sib b -> Array v sic c) -> Array v sa a -> Array v sb b -> Array v sc c
transmit :: forall (v :: * -> *) (sa :: [Nat]) (sb :: [Nat]) (sc :: [Nat]) a b
       c (ds :: [Nat]) (sib :: [Nat]) (sic :: [Nat]) (sob :: [Nat]).
(KnownNats sa, KnownNats sb, KnownNats sc, KnownNats ds,
 KnownNats sib, KnownNats sic, KnownNats sob,
 ds ~ Eval (EnumFromTo (Eval (Rank sa)) (Eval (Rank sb) - 1)),
 sib ~ Eval (DeleteDims ds sb), sob ~ Eval (GetDims ds sb),
 sb ~ Eval (InsertDims ds sob sib),
 sc ~ Eval (InsertDims ds sob sic), 'True ~ Eval (IsPrefixOf sa sb),
 Vector v a, Vector v b, Vector v c, Vector v (Array v sib b),
 Vector v (Array v sic c)) =>
(Array v sa a -> Array v sib b -> Array v sic c)
-> Array v sa a -> Array v sb b -> Array v sc c
transmit Array v sa a -> Array v sib b -> Array v sic c
f Array v sa a
a Array v sb b
b = Dims ds
-> (Array v sib b -> Array v sic c) -> Array v sb b -> Array v sc c
forall (v :: * -> *) (ds :: [Nat]) (s :: [Nat]) (s' :: [Nat])
       (si :: [Nat]) (si' :: [Nat]) (so :: [Nat]) a b.
(KnownNats s, KnownNats s', KnownNats si, KnownNats si',
 KnownNats so, si ~ Eval (DeleteDims ds s),
 so ~ Eval (GetDims ds s), s' ~ Eval (InsertDims ds so si'),
 s ~ Eval (InsertDims ds so si), Vector v a, Vector v b,
 Vector v (Array v si a), Vector v (Array v si' b)) =>
Dims ds
-> (Array v si a -> Array v si' b) -> Array v s a -> Array v s' b
maps (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) (Array v sa a -> Array v sib b -> Array v sic c
f Array v sa a
a) Array v sb b
b

-- | A one-dimensional array.
type Vector v s a = Array v '[s] a

-- | Create a one-dimensional array.
--
-- >>> pretty $ vector @Vec.Vector @3 @Int [2,3,4]
-- [2,3,4]
vector :: forall v n a t. (FromVector t v a, KnownNat n, VG.Vector v a) => t -> Array v '[n] a
vector :: forall (v :: * -> *) (n :: Nat) a t.
(FromVector t v a, KnownNat n, Vector v a) =>
t -> Array v '[n] a
vector t
xs = t -> Array v '[n] a
forall (v :: * -> *) (s :: [Nat]) a t.
(KnownNats s, FromVector t v a, Vector v a) =>
t -> Array v s a
array t
xs

-- | vector with an explicit SNat rather than a KnownNat constraint.
--
-- >>> pretty $ vector' @Vec.Vector @Int (SNat @3) [2,3,4]
-- [2,3,4]
vector' :: forall v a n t. (FromVector t v a, VG.Vector v a) => SNat n -> t -> Array v '[n] a
vector' :: forall (v :: * -> *) a (n :: Nat) t.
(FromVector t v a, Vector v a) =>
SNat n -> t -> Array v '[n] a
vector' SNat n
n t
xs = SNat n -> (KnownNat n => Array v '[n] a) -> Array v '[n] a
forall (n :: Nat) r. SNat n -> (KnownNat n => r) -> r
withKnownNat SNat n
n (t -> Array v '[n] a
forall (v :: * -> *) (n :: Nat) a t.
(FromVector t v a, KnownNat n, Vector v a) =>
t -> Array v '[n] a
vector t
xs)

-- | Vector specialisation of 'range'
--
-- >>> toDynamic $ iota @Vec.Vector @5
-- UnsafeArray [5] [0,1,2,3,4]
iota :: forall v n. (KnownNat n, VG.Vector v Int) => Vector v n Int
iota :: forall (v :: * -> *) (n :: Nat).
(KnownNat n, Vector v Int) =>
Vector v n Int
iota = Array v '[n] Int
forall (v :: * -> *) (s :: [Nat]).
(KnownNats s, Vector v Int) =>
Array v s Int
range

-- | A two-dimensional array.
type Matrix v m n a = Array v '[m, n] a

-- * row (first dimension) specializations

-- | Add a new row
--
-- >>> pretty $ cons (array @Vec.Vector @'[2] [0,1]) (array @Vec.Vector @[2,2] [2,3,4,5])
-- [[0,1],
--  [2,3],
--  [4,5]]
cons :: forall v st s sh a. (KnownNats st, KnownNats s, KnownNats sh, True ~ Eval (InsertOk 0 st sh), s ~ Eval (IncAt 0 st), sh ~ Eval (DeleteDim 0 st), VG.Vector v a) => Array v sh a -> Array v st a -> Array v s a
cons :: forall (v :: * -> *) (st :: [Nat]) (s :: [Nat]) (sh :: [Nat]) a.
(KnownNats st, KnownNats s, KnownNats sh,
 'True ~ Eval (InsertOk 0 st sh), s ~ Eval (IncAt 0 st),
 sh ~ Eval (DeleteDim 0 st), Vector v a) =>
Array v sh a -> Array v st a -> Array v s a
cons =
  Dim 0 -> Array v sh a -> Array v st a -> Array v s a
forall (v :: * -> *) a (d :: Nat) (s :: [Nat]) (si :: [Nat])
       (s' :: [Nat]).
(KnownNats s, KnownNats si, KnownNats s', s' ~ Eval (IncAt d s),
 'True ~ Eval (InsertOk d s si), Vector v a) =>
Dim d -> Array v si a -> Array v s a -> Array v s' a
prepend (forall (n :: Nat). KnownNat n => SNat n
SNat @0)

-- | Add a new row at the end
--
-- >>> pretty $ snoc (array @Vec.Vector @[2,2] [0,1,2,3]) (array @Vec.Vector @'[2] [4,5])
-- [[0,1],
--  [2,3],
--  [4,5]]
snoc :: forall v si s sl a. (KnownNats si, KnownNats s, KnownNats sl, True ~ Eval (InsertOk 0 si sl), s ~ Eval (IncAt 0 si), sl ~ Eval (DeleteDim 0 si), VG.Vector v a) => Array v si a -> Array v sl a -> Array v s a
snoc :: forall (v :: * -> *) (si :: [Nat]) (s :: [Nat]) (sl :: [Nat]) a.
(KnownNats si, KnownNats s, KnownNats sl,
 'True ~ Eval (InsertOk 0 si sl), s ~ Eval (IncAt 0 si),
 sl ~ Eval (DeleteDim 0 si), Vector v a) =>
Array v si a -> Array v sl a -> Array v s a
snoc = Dim 0 -> Array v si a -> Array v sl a -> Array v s a
forall (v :: * -> *) a (d :: Nat) (s :: [Nat]) (si :: [Nat])
       (s' :: [Nat]).
(KnownNats s, KnownNats si, KnownNats s', s' ~ Eval (IncAt d s),
 'True ~ Eval (InsertOk d s si), Vector v a) =>
Dim d -> Array v s a -> Array v si a -> Array v s' a
append (forall (n :: Nat). KnownNat n => SNat n
SNat @0)

-- | split an array into the first row and the remaining rows.
--
-- >>> import Data.Bifunctor (bimap)
-- >>> bimap toDynamic toDynamic $ uncons (array @Vec.Vector @[3,2] [0..5])
-- (UnsafeArray [2] [0,1],UnsafeArray [2,2] [2,3,4,5])
uncons :: forall v a s sh st ls os ds. (KnownNats s, KnownNats sh, KnownNats st, ds ~ '[0], sh ~ Eval (DeleteDims ds s), KnownNats ls, KnownNats os, os ~ Eval (Replicate (Eval (Rank ds)) 1), ls ~ Eval (GetLastPositions ds s), Eval (SlicesOk ds os ls s) ~ True, st ~ Eval (SetDims ds ls s), VG.Vector v a) => Array v s a -> (Array v sh a, Array v st a)
uncons :: forall (v :: * -> *) a (s :: [Nat]) (sh :: [Nat]) (st :: [Nat])
       (ls :: [Nat]) (os :: [Nat]) (ds :: [Nat]).
(KnownNats s, KnownNats sh, KnownNats st, ds ~ '[0],
 sh ~ Eval (DeleteDims ds s), KnownNats ls, KnownNats os,
 os ~ Eval (Replicate (Eval (Rank ds)) 1),
 ls ~ Eval (GetLastPositions ds s),
 Eval (SlicesOk ds os ls s) ~ 'True, st ~ Eval (SetDims ds ls s),
 Vector v a) =>
Array v s a -> (Array v sh a, Array v st a)
uncons Array v s a
a = (Dims ds -> Array v s a -> Array v sh a
forall (v :: * -> *) a (ds :: [Nat]) (s :: [Nat]) (s' :: [Nat]).
(KnownNats s, KnownNats s', KnownNats ds,
 s' ~ Eval (DeleteDims ds s), Vector v a) =>
Dims ds -> Array v s a -> Array v s' a
heads (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) Array v s a
a, Dims ds -> Array v s a -> Array v st a
forall (v :: * -> *) (ds :: [Nat]) (os :: [Nat]) (s :: [Nat])
       (s' :: [Nat]) a (ls :: [Nat]).
(KnownNats s, KnownNats ds, KnownNats s', KnownNats ls,
 KnownNats os, Eval (SlicesOk ds os ls s) ~ 'True,
 os ~ Eval (Replicate (Eval (Rank ds)) 1),
 ls ~ Eval (GetLastPositions ds s), s' ~ Eval (SetDims ds ls s),
 Vector v a) =>
Dims ds -> Array v s a -> Array v s' a
tails (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) Array v s a
a)

-- | split an array into the initial rows and the last row.
--
-- >>> import Data.Bifunctor (bimap)
-- >>> bimap toDynamic toDynamic $ unsnoc (array @Vec.Vector @[3,2] [0..5])
-- (UnsafeArray [2,2] [0,1,2,3],UnsafeArray [2] [4,5])
unsnoc :: forall v ds os s a ls si sl. (KnownNats s, KnownNats ds, KnownNats si, KnownNats ls, KnownNats os, KnownNats sl, ds ~ '[0], Eval (SlicesOk ds os ls s) ~ True, os ~ Eval (Replicate (Eval (Rank ds)) 0), ls ~ Eval (GetLastPositions ds s), si ~ Eval (SetDims ds ls s), sl ~ Eval (DeleteDims ds s), VG.Vector v a) => Array v s a -> (Array v si a, Array v sl a)
unsnoc :: forall (v :: * -> *) (ds :: [Nat]) (os :: [Nat]) (s :: [Nat]) a
       (ls :: [Nat]) (si :: [Nat]) (sl :: [Nat]).
(KnownNats s, KnownNats ds, KnownNats si, KnownNats ls,
 KnownNats os, KnownNats sl, ds ~ '[0],
 Eval (SlicesOk ds os ls s) ~ 'True,
 os ~ Eval (Replicate (Eval (Rank ds)) 0),
 ls ~ Eval (GetLastPositions ds s), si ~ Eval (SetDims ds ls s),
 sl ~ Eval (DeleteDims ds s), Vector v a) =>
Array v s a -> (Array v si a, Array v sl a)
unsnoc Array v s a
a = (Dims ds -> Array v s a -> Array v si a
forall (v :: * -> *) (ds :: [Nat]) (os :: [Nat]) (s :: [Nat])
       (s' :: [Nat]) a (ls :: [Nat]).
(KnownNats s, KnownNats ds, KnownNats s', KnownNats ls,
 KnownNats os, Eval (SlicesOk ds os ls s) ~ 'True,
 os ~ Eval (Replicate (Eval (Rank ds)) 0),
 ls ~ Eval (GetLastPositions ds s), s' ~ Eval (SetDims ds ls s),
 Vector v a) =>
Dims ds -> Array v s a -> Array v s' a
inits (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) Array v s a
a, Dims ds -> Array v s a -> Array v sl a
forall (v :: * -> *) (ds :: [Nat]) (s :: [Nat]) (s' :: [Nat]) a.
(KnownNats s, KnownNats ds, KnownNats s',
 s' ~ Eval (DeleteDims ds s), Vector v a) =>
Dims ds -> Array v s a -> Array v s' a
lasts (forall (ns :: [Nat]). KnownNats ns => SNats ns
Dims @ds) Array v s a
a)

-- | Convenience pattern for row extraction and consolidation at the beginning of an Array.
--
-- >>> (x:<xs) = array @Vec.Vector @'[4] [0..3]
-- >>> toDynamic x
-- UnsafeArray [] [0]
-- >>> toDynamic xs
-- UnsafeArray [3] [1,2,3]
-- >>> toDynamic (x:<xs)
-- UnsafeArray [4] [0,1,2,3]
pattern (:<) :: forall v s sh st a os ls ds. (KnownNats s, KnownNats sh, KnownNats st, True ~ Eval (InsertOk 0 st sh), s ~ Eval (IncAt 0 st), ds ~ '[0], sh ~ Eval (DeleteDims ds s), KnownNats ls, KnownNats os, Eval (SlicesOk ds os ls s) ~ True, os ~ Eval (Replicate (Eval (Rank ds)) 1), ls ~ Eval (GetLastPositions ds s), st ~ Eval (SetDims ds ls s), VG.Vector v a) => Array v sh a -> Array v st a -> Array v s a
pattern x $b:< :: forall (v :: * -> *) (s :: [Nat]) (sh :: [Nat]) (st :: [Nat]) a
       (os :: [Nat]) (ls :: [Nat]) (ds :: [Nat]).
(KnownNats s, KnownNats sh, KnownNats st,
 'True ~ Eval (InsertOk 0 st sh), s ~ Eval (IncAt 0 st), ds ~ '[0],
 sh ~ Eval (DeleteDims ds s), KnownNats ls, KnownNats os,
 Eval (SlicesOk ds os ls s) ~ 'True,
 os ~ Eval (Replicate (Eval (Rank ds)) 1),
 ls ~ Eval (GetLastPositions ds s), st ~ Eval (SetDims ds ls s),
 Vector v a) =>
Array v sh a -> Array v st a -> Array v s a
$m:< :: forall {r} {v :: * -> *} {s :: [Nat]} {sh :: [Nat]} {st :: [Nat]}
       {a} {os :: [Nat]} {ls :: [Nat]} {ds :: [Nat]}.
(KnownNats s, KnownNats sh, KnownNats st,
 'True ~ Eval (InsertOk 0 st sh), s ~ Eval (IncAt 0 st), ds ~ '[0],
 sh ~ Eval (DeleteDims ds s), KnownNats ls, KnownNats os,
 Eval (SlicesOk ds os ls s) ~ 'True,
 os ~ Eval (Replicate (Eval (Rank ds)) 1),
 ls ~ Eval (GetLastPositions ds s), st ~ Eval (SetDims ds ls s),
 Vector v a) =>
Array v s a
-> (Array v sh a -> Array v st a -> r) -> ((# #) -> r) -> r
:< xs <- (uncons -> (x, xs))
  where
    Array v sh a
x :< Array v st a
xs = Array v sh a -> Array v st a -> Array v s a
forall (v :: * -> *) (st :: [Nat]) (s :: [Nat]) (sh :: [Nat]) a.
(KnownNats st, KnownNats s, KnownNats sh,
 'True ~ Eval (InsertOk 0 st sh), s ~ Eval (IncAt 0 st),
 sh ~ Eval (DeleteDim 0 st), Vector v a) =>
Array v sh a -> Array v st a -> Array v s a
cons Array v sh a
x Array v st a
xs

infix 5 :<

{-# COMPLETE (:<) :: Array #-}

-- | Convenience pattern for row extraction and consolidation at the end of an Array.
--
-- >>> (xs:>x) = array @Vec.Vector @'[4] [0..3]
-- >>> toDynamic x
-- UnsafeArray [] [3]
-- >>> toDynamic xs
-- UnsafeArray [3] [0,1,2]
-- >>> toDynamic (xs:>x)
-- UnsafeArray [4] [0,1,2,3]
pattern (:>) :: forall v si sl s a ds ls os. (KnownNats si, KnownNats sl, KnownNats s, True ~ Eval (InsertOk 0 si sl), s ~ Eval (IncAt 0 si), KnownNats ds, KnownNats ls, KnownNats os, sl ~ Eval (DeleteDim 0 si), ds ~ '[0], Eval (SlicesOk ds os ls s) ~ True, os ~ Eval (Replicate (Eval (Rank ds)) 0), ls ~ Eval (GetLastPositions ds s), si ~ Eval (SetDims ds ls s), sl ~ Eval (DeleteDims ds s), VG.Vector v a) => Array v si a -> Array v sl a -> Array v s a
pattern xs $b:> :: forall (v :: * -> *) (si :: [Nat]) (sl :: [Nat]) (s :: [Nat]) a
       (ds :: [Nat]) (ls :: [Nat]) (os :: [Nat]).
(KnownNats si, KnownNats sl, KnownNats s,
 'True ~ Eval (InsertOk 0 si sl), s ~ Eval (IncAt 0 si),
 KnownNats ds, KnownNats ls, KnownNats os,
 sl ~ Eval (DeleteDim 0 si), ds ~ '[0],
 Eval (SlicesOk ds os ls s) ~ 'True,
 os ~ Eval (Replicate (Eval (Rank ds)) 0),
 ls ~ Eval (GetLastPositions ds s), si ~ Eval (SetDims ds ls s),
 sl ~ Eval (DeleteDims ds s), Vector v a) =>
Array v si a -> Array v sl a -> Array v s a
$m:> :: forall {r} {v :: * -> *} {si :: [Nat]} {sl :: [Nat]} {s :: [Nat]}
       {a} {ds :: [Nat]} {ls :: [Nat]} {os :: [Nat]}.
(KnownNats si, KnownNats sl, KnownNats s,
 'True ~ Eval (InsertOk 0 si sl), s ~ Eval (IncAt 0 si),
 KnownNats ds, KnownNats ls, KnownNats os,
 sl ~ Eval (DeleteDim 0 si), ds ~ '[0],
 Eval (SlicesOk ds os ls s) ~ 'True,
 os ~ Eval (Replicate (Eval (Rank ds)) 0),
 ls ~ Eval (GetLastPositions ds s), si ~ Eval (SetDims ds ls s),
 sl ~ Eval (DeleteDims ds s), Vector v a) =>
Array v s a
-> (Array v si a -> Array v sl a -> r) -> ((# #) -> r) -> r
:> x <- (unsnoc -> (xs, x))
  where
    Array v si a
xs :> Array v sl a
x = Array v si a -> Array v sl a -> Array v s a
forall (v :: * -> *) (si :: [Nat]) (s :: [Nat]) (sl :: [Nat]) a.
(KnownNats si, KnownNats s, KnownNats sl,
 'True ~ Eval (InsertOk 0 si sl), s ~ Eval (IncAt 0 si),
 sl ~ Eval (DeleteDim 0 si), Vector v a) =>
Array v si a -> Array v sl a -> Array v s a
snoc Array v si a
xs Array v sl a
x

infix 5 :>

{-# COMPLETE (:>) :: Array #-}

-- | Generate an array of uniform random variates between a range.
--
-- >>> import System.Random.Stateful hiding (uniform)
-- >>> g <- newIOGenM (mkStdGen 42)
-- >>> u <- uniform @Vec.Vector @[2,3,4] @Int g (0,9)
-- >>> 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 :: forall v s a g m. (StatefulGen g m, UniformRange a, KnownNats s, VG.Vector v a) => g -> (a, a) -> m (Array v s a)
uniform :: forall (v :: * -> *) (s :: [Nat]) a g (m :: * -> *).
(StatefulGen g m, UniformRange a, KnownNats s, Vector v a) =>
g -> (a, a) -> m (Array v s a)
uniform g
g (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 (forall (s :: [Nat]). KnownNats s => [Int]
valuesOf @s))) ((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) :: m (v a)
  pure $ array @v @s v