harpie (empty) → 0.1.0.0
raw patch · 9 files changed
+6809/−0 lines, 9 filesdep +QuickCheckdep +adjunctionsdep +base
Dependencies added: QuickCheck, adjunctions, base, distributive, doctest-parallel, first-class-families, prettyprinter, quickcheck-instances, random, vector, vector-algorithms
Files
- ChangeLog.md +5/−0
- LICENSE +30/−0
- harpie.cabal +101/−0
- readme.md +117/−0
- src/Harpie/Array.hs +2012/−0
- src/Harpie/Fixed.hs +2831/−0
- src/Harpie/Shape.hs +1643/−0
- src/Harpie/Sort.hs +60/−0
- test/doctests.hs +10/−0
+ ChangeLog.md view
@@ -0,0 +1,5 @@+0.1+===++- cleaved from numhask-array-0.11+
+ LICENSE view
@@ -0,0 +1,30 @@+Copyright (c) 2016, Tony Day++All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++ * Redistributions of source code must retain the above copyright+ notice, this list of conditions and the following disclaimer.++ * Redistributions in binary form must reproduce the above+ copyright notice, this list of conditions and the following+ disclaimer in the documentation and/or other materials provided+ with the distribution.++ * Neither the name of Tony Day nor the names of other+ contributors may be used to endorse or promote products derived+ from this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ harpie.cabal view
@@ -0,0 +1,101 @@+cabal-version: 3.0+name: harpie+version: 0.1.0.0+license: BSD-3-Clause+license-file: LICENSE+copyright: Tony Day (c) 2016-2024+category: math+author: Tony Day+maintainer: tonyday567@gmail.com+homepage: https://github.com/tonyday567/harpie#readme+bug-reports: https://github.com/tonyday567/harpie/issues+synopsis: Haskell array programming.+description:+ This package provides Haskell array programming, interface and environment.++ Module names clash with each other and with the Prelude.++ == Usage++ >>> import Harpie.Fixed qualified as F+ >>> import Harpie.Shape qualified as S+ >>> import Harpie.Array qualified as A++ >>> a = F.range @[2,3,4]+ >>> F.shape a+ [2,3,4]+ >>> pretty a+ [[[0,1,2,3],+ [4,5,6,7],+ [8,9,10,11]],+ [[12,13,14,15],+ [16,17,18,19],+ [20,21,22,23]]]++ >>> a = A.range [2,3,4]+ >>> F.shape a+ [2,3,4]+ >>> 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]]]++build-type: Simple+tested-with:+ , GHC == 9.10.1+ , GHC == 9.8.2+ , GHC == 9.6.5+extra-doc-files:+ ChangeLog.md+ readme.md++source-repository head+ type: git+ location: https://github.com/tonyday567/harpie++common ghc-options-stanza+ ghc-options:+ -Wall+ -Wcompat+ -Widentities+ -Wincomplete-record-updates+ -Wincomplete-uni-patterns+ -Wpartial-fields+ -Wredundant-constraints++common ghc2021-stanza+ default-language: GHC2021++library+ import: ghc-options-stanza+ import: ghc2021-stanza+ hs-source-dirs: src+ build-depends:+ , QuickCheck >=2.15.0 && <2.16+ , adjunctions >=4.0 && <5+ , base >=4.14 && <5+ , distributive >=0.4 && <0.7+ , first-class-families >=0.8.1 && <0.9+ , prettyprinter >=1.7 && <1.8+ , quickcheck-instances >=0.3.31 && <0.4+ , random >=1.2 && <1.3+ , vector >=0.12.3 && <0.14+ , vector-algorithms >=0.9.0 && <0.10+ exposed-modules:+ Harpie.Array+ Harpie.Fixed+ Harpie.Shape+ Harpie.Sort++test-suite doctests+ import: ghc2021-stanza+ main-is: doctests.hs+ hs-source-dirs: test+ build-depends:+ , base >=4.14 && <5+ , doctest-parallel >=0.3 && <0.4+ ghc-options: -threaded+ type: exitcode-stdio-1.0
+ readme.md view
@@ -0,0 +1,117 @@+harpie+===++[](https://hackage.haskell.org/package/harpie)+[](https://github.com/tonyday567/harpie/actions) <img src="https://static.wikia.nocookie.net/monster/images/2/28/Harpy.jpg" alt="harpie" width="100"/>++Haskell array programming, interface and environment (harpie).++**harpie** is an array programming library written in Haskell. Features include:++- **Rank polymorphism:** vectors, matrices, tensors and scalars all use the same API and functions and combine naturally, as the math gods intend.++- **Both value- and type- level shapes:** consistent functionality for accessing and transforming shape is provided at both value-level and type-level.++- **Dimensional agnosticism:** the granularity of an array is flexible with the same function able to be applied over rows, columns, over multiple dimensions at once, or at element-level.++- **Pure Haskell:** promoting a higher-kinded, ergonomic style of array programming, in idiomatic Haskell style.++The library is experimental and educational (at least for the authors) and likely to remain so. Collaboration is most welcome.++Usage+===++Naming conventions clash with the prelude and with each other, so importing should be qualified.++``` haskell+import qualified Harpie.Array as A+import qualified Harpie.Shape as S+import qualified Harpie.Fixed as F++-- >>> a = F.range @[2,3,4]+-- >>> F.shape a+-- [2,3,4]+-- >>> pretty a+-- [[[0,1,2,3],+-- [4,5,6,7],+-- [8,9,10,11]],+-- [[12,13,14,15],+-- [16,17,18,19],+-- [20,21,22,23]]]++-- >>> a = A.range [2,3,4]+-- >>> F.shape a+-- [2,3,4]+-- >>> 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]]]+```+++Design notes+===++Haskell utility+---++The library attempts to be idiomatic Haskell and otherwise fit in with the language ecosystem. In particular, boxed Vectors are used as the array container to enforce immutability, permit lazy expression, and allow arbitrary element types.++Consistency of type- and value- level list algorithms.+---++The library is an attempt to provide a consistent approach to array programming whether or not array shape is held and computed at value-level or at type-level. ++The Harpie.Shape module contains common list algorithms for value-level shapes (ie Int list operatorions) and type-level shapes (a type-level Nat list operations) that is as close to the same as possible. They cannot be identical because type and value programming in Haskell are very different languages. The [first-class-families](https://hackage.haskell.org/package/first-class-families) library was used to achieve this.++Is it safe?+---++Harpie.Fixed arrays sit at around level 4.5 in [Justin Le's type safety heirarchy](https://blog.jle.im/entry/levels-of-type-safety-haskell-lists.html). They are designed with **static** type-safety in mind; a run-time error in shape or index computation is a bug. Typed shape information, however, is modelled on GHC.TypeNats with Nat and thus [Nat] being opaque types rather than inductively-structured. This makes compiler KnownNat and KnownNats discovery and proof witnessing problematic.++Instead of dependently-typed approaches, the library leans into switching from fixed to untyped shape representation if shape is not known at runtime. ++Is it fast?+---++Maybe. Does [vector](https://hackage.haskell.org/package/vector) streaming lead to efficient code? If it does then harpie should be able to access this efficiency.++If you have a bunch of continguous data in need of strictly simple array processing then there are plenty of faster array programming options out there, in Haskell and elsewhere. If you need close-to-the-metal performance, try:++- [hmatrix](https://hackage.haskell.org/package/hmatrix) provides bindings to [BLAS](https://en.wikipedia.org/wiki/Basic_Linear_Algebra_Subprograms) and 45 years of Fortran tweaking.+- [hasktorch](https://github.com/hasktorch/hasktorch) or [tensorflow](https://hackage.haskell.org/package/tensorflow) bindings to the might of the Python machince learning community.+- [massive](https://hackage.haskell.org/package/massiv) or [accelerate](https://hackage.haskell.org/package/accelerate) for homegrown Haskell speed.+- [repa](https://hackage.haskell.org/package/repa) for archeology.+- [orthotope](https://hackage.haskell.org/package/orthotope) for a more APL treatment of the problem domain, and for unboxed vector usage. ++backpermute+---++As computational complexity increases, either in array programmimg or in connectivity to other problem domains, harpie performance is (one-day) expected to come into her own. Backpermute fusion is the secret sauce.++A canonical backpermute function is detailed in [Regular, Shape-polymorphic, Parallel Arrays in Haskell](https://benl.ouroborus.net/papers/2010-rarrays/repa-icfp2010.pdf) and would be implemented in harpie as:++> repa_backpermute f a = tabulate (f (shape a)) (index a . f)++The harpie backpermute splits the function application into two parts, the macro-shape change and the pre-indexing operation:++> harpie_backpermute f g a = tabulate (f (shape a)) (index a . g)++This, more general specification, should (one-day) allow more opportunities for the fusion rule to kick in:++> forall f f' g g' (a :: forall a. Array a)). backpermute f g (backpermute f' g' a) == backpermute (f . f') (g . g') a++Type Ugliness+===++Harpies are monstrous beasties, fierce but not conventional or conventionally attractive:++- type errors and information quickly explode in complexity and types can run to hundreds of lines. It's quite ugly.++- Constraints cannot be encapsulated in the function they originate and thus have to be propogated in line with usage. Type-level programming in Haskell is a good idea in search of some composition.+++
+ src/Harpie/Array.hs view
@@ -0,0 +1,2012 @@+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE FunctionalDependencies #-}+{-# LANGUAGE PatternSynonyms #-}+{-# LANGUAGE RoleAnnotations #-}+{-# LANGUAGE ViewPatterns #-}+{-# OPTIONS_GHC -fno-warn-incomplete-uni-patterns #-}++-- | Arrays with shape information and computations at a value-level.+module Harpie.Array+ ( -- * Usage+ -- $usage++ -- * Harpie Arrays+ Array (..),+ array,+ (><),+ validate,+ safeArray,+ unsafeModifyShape,+ unsafeModifyVector,++ -- * Dimensions+ Dim,+ Dims,++ -- * Conversion+ FromVector (..),+ FromArray (..),++ -- * Shape Access+ shape,+ rank,+ size,+ length,+ isNull,++ -- * Indexing+ index,+ (!),+ (!?),+ tabulate,+ backpermute,++ -- * Scalars+ fromScalar,+ toScalar,+ isScalar,+ asSingleton,+ asScalar,++ -- * Array Creation+ empty,+ range,+ corange,+ indices,+ ident,+ konst,+ singleton,+ diag,+ undiag,++ -- * Element-level functions+ zipWith,+ zipWithSafe,+ modify,+ imap,++ -- * Function generalisers+ rowWise,+ colWise,+ dimsWise,++ -- * Single-dimension functions+ take,+ drop,+ select,+ insert,+ delete,+ append,+ prepend,+ concatenate,+ couple,+ slice,+ rotate,++ -- * Multi-dimension functions+ takes,+ drops,+ indexes,+ slices,+ heads,+ lasts,+ tails,+ inits,++ -- * Function application+ extracts,+ reduces,+ joins,+ joinsSafe,+ join,+ joinSafe,+ traverses,+ maps,+ filters,+ zips,+ zipsSafe,+ modifies,+ diffs,++ -- * Array expansion & contraction+ expand,+ coexpand,+ contract,+ prod,+ dot,+ mult,+ windows,++ -- * Search+ find,+ findNoOverlap,+ findIndices,+ isPrefixOf,+ isSuffixOf,+ isInfixOf,++ -- * Shape manipulations+ fill,+ cut,+ cutSuffix,+ pad,+ lpad,+ reshape,+ flat,+ repeat,+ cycle,+ rerank,+ reorder,+ squeeze,+ elongate,+ transpose,+ inflate,+ intercalate,+ intersperse,+ concats,+ reverses,+ rotates,++ -- * Sorting+ sorts,+ sortsBy,+ orders,+ ordersBy,++ -- * Transmission+ transmit,+ transmitSafe,+ transmitOp,+ telecasts,+ telecastsSafe,++ -- * Row specializations+ pattern (:<),+ cons,+ uncons,+ pattern (:>),+ snoc,+ unsnoc,++ -- * Shape specializations+ iota,++ -- * Math+ uniform,+ invtri,+ inverse,+ chol,+ )+where++import Control.Monad hiding (join)+import Data.Bool+import Data.Foldable hiding (find, length, minimum)+import Data.Function+import Data.List qualified as List+import Data.Vector qualified as V+import GHC.Generics+import Harpie.Shape hiding (asScalar, asSingleton, concatenate, range, rank, reorder, rerank, rotate, size, squeeze)+import Harpie.Shape qualified as S+import Harpie.Sort+import Prettyprinter hiding (dot, fill)+import System.Random hiding (uniform)+import System.Random.Stateful hiding (uniform)+import Prelude as P hiding (cycle, drop, length, repeat, take, zip, zipWith)++-- $setup+-- >>> :m -Prelude+-- >>> import Prelude hiding (take, drop, zipWith, length, cycle, repeat)+-- >>> import Harpie.Array as A+-- >>> import Harpie.Shape qualified as S+-- >>> import Data.Vector qualified as V+-- >>> import Prettyprinter hiding (dot, fill)+-- >>> import Data.List qualified as List+-- >>> let s = 1 :: Array Int+-- >>> s+-- UnsafeArray [] [1]+-- >>> pretty s+-- 1+-- >>> let v = range [3]+-- >>> v+-- UnsafeArray [3] [0,1,2]+-- >>> let m = range [2,3]+-- >>> pretty m+-- [[0,1,2],+-- [3,4,5]]+-- >>> let a = range [2,3,4]+-- >>> a+-- UnsafeArray [2,3,4] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23]+-- >>> pretty a+-- [[[0,1,2,3],+-- [4,5,6,7],+-- [8,9,10,11]],+-- [[12,13,14,15],+-- [16,17,18,19],+-- [20,21,22,23]]]++-- $usage+--+-- Several names used in @harpie@ conflict with [Prelude](https://hackage.haskell.org/package/base/docs/Prelude.html):+--+-- >>> import Prelude hiding (cycle, repeat, take, drop, zipWith, length)+--+-- In general, 'Array' functionality is contained in @Harpie.Array@ and shape functionality is contained in @Harpie.Shape@. These two modules also have name clashes and at least one needs to be qualified:+--+-- >>> import Harpie.Array as A+-- >>> import Harpie.Shape qualified as S+--+-- [@prettyprinter@](https://hackage.haskell.org/package/prettyprinter) is used to prettily render arrays to better visualise shape.+--+-- >>> import Prettyprinter hiding (dot,fill)+--+-- Examples of arrays:+--+-- An array with no dimensions (a scalar).+--+-- >>> s = 1 :: Array Int+-- >>> s+-- UnsafeArray [] [1]+-- >>> shape s+-- []+-- >>> pretty s+-- 1+--+-- A single-dimension array (a vector).+--+-- >>> let v = range [3]+-- >>> pretty v+-- [0,1,2]+--+-- A two-dimensional array (a matrix).+--+-- >>> let m = range [2,3]+-- >>> pretty m+-- [[0,1,2],+-- [3,4,5]]+--+-- An n-dimensional array (n should be finite).+--+-- >>> a = range [2,3,4]+-- >>> a+-- UnsafeArray [2,3,4] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23]+-- >>> pretty a+-- [[[0,1,2,3],+-- [4,5,6,7],+-- [8,9,10,11]],+-- [[12,13,14,15],+-- [16,17,18,19],+-- [20,21,22,23]]]++-- | A hyperrectangular (or multidimensional) array with a value-level shape.+--+-- >>> let a = array [2,3,4] [1..24] :: Array Int+-- >>> a+-- UnsafeArray [2,3,4] [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24]+--+-- >>> pretty a+-- [[[1,2,3,4],+-- [5,6,7,8],+-- [9,10,11,12]],+-- [[13,14,15,16],+-- [17,18,19,20],+-- [21,22,23,24]]]+data Array a = UnsafeArray [Int] (V.Vector a)+ deriving stock (Generic)+ deriving stock (Eq, Ord, Show)++type role Array representational++instance Functor Array where+ fmap f = unsafeModifyVector (V.map f)++instance Foldable Array where+ foldr f x0 a = V.foldr f x0 (asVector a)++instance Traversable Array where+ traverse f (UnsafeArray s v) =+ array s <$> traverse f v++instance (Show a) => Pretty (Array a) where+ pretty a@(UnsafeArray _ v) = case rank a of+ 0 -> viaShow (V.head v)+ 1 -> viaShow v+ _ ->+ pretty "["+ <> indent+ 0+ ( vsep+ ( punctuate comma $+ pretty+ <$> toList (extracts [0] a)+ )+ )+ <> pretty "]"++-- * conversions++instance (Num a) => Num (Array a) where+ (+) = zipWith (+)+ (-) = zipWith (-)+ (*) = error "multiplication not defined"+ abs = fmap abs+ signum = fmap signum+ fromInteger x = toScalar (fromInteger x)++-- | 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 [2,3])+-- [0,1,2,3,4,5]+--+-- >>> vectorAs (V.fromList [0..5]) :: Array Int+-- UnsafeArray [6] [0,1,2,3,4,5]+class FromVector t a | t -> a where+ asVector :: t -> V.Vector a+ vectorAs :: V.Vector a -> t++instance FromVector (V.Vector a) a where+ asVector = id+ vectorAs = id++instance FromVector [a] a where+ asVector = V.fromList+ vectorAs = V.toList++instance FromVector (Array a) a where+ asVector (UnsafeArray _ v) = v+ vectorAs v = UnsafeArray [V.length v] v++-- | Conversion to and from an `Array`+--+-- Note that conversion of an 'Array' to a `FromArray` likely drops shape information, so that:+--+-- > arrayAs . asArray == flat+-- > asArray . arrayAs == id+--+-- >>> asArray ([0..5::Int])+-- UnsafeArray [6] [0,1,2,3,4,5]+--+-- >>> arrayAs (range [2,3]) :: [Int]+-- [0,1,2,3,4,5]+class FromArray t a | t -> a where+ asArray :: t -> Array a+ arrayAs :: Array a -> t++instance FromArray (Array a) a where+ asArray = id+ arrayAs = id++instance FromArray [a] a where+ asArray l = UnsafeArray [S.rank l] (V.fromList l)+ arrayAs (UnsafeArray _ v) = V.toList v++instance FromArray (V.Vector a) a where+ asArray v = UnsafeArray [V.length v] v+ arrayAs (UnsafeArray _ v) = v++-- | Construct an array from a shape and a value without any shape validation.+--+-- >>> array [2,3] [0..5]+-- UnsafeArray [2,3] [0,1,2,3,4,5]+array :: (FromVector t a) => [Int] -> t -> Array a+array s (asVector -> v) = UnsafeArray s v++infixl 4 ><++-- | Construct an Array.+--+-- >>> pretty $ [2,3] >< [0..5]+-- [[0,1,2],+-- [3,4,5]]+(><) :: (FromVector t a) => [Int] -> t -> Array a+(><) = array++-- | Validate the size and shape of an array.+--+-- >>> validate (array [2,3,4] [1..23] :: Array Int)+-- False+validate :: Array a -> Bool+validate a = size a == V.length (asVector a)++-- | Construct an Array, checking shape.+--+-- >>> safeArray [2,3,4] [0..23] == Just a+-- True+safeArray :: (FromVector t a) => [Int] -> t -> Maybe (Array a)+safeArray s v =+ bool Nothing (Just a) (validate a)+ where+ a = UnsafeArray s (asVector v)++-- | Unsafely modify an array shape.+--+-- >>> unsafeModifyShape (fmap (+1) :: [Int] -> [Int]) (array [2,3] [0..5])+-- UnsafeArray [3,4] [0,1,2,3,4,5]+unsafeModifyShape :: ([Int] -> [Int]) -> Array a -> Array a+unsafeModifyShape f (UnsafeArray s v) = UnsafeArray (f s) v++-- | Unsafely modify an array vector.+--+-- >>> unsafeModifyVector (V.map (+1)) (array [2,3] [0..5])+-- UnsafeArray [2,3] [1,2,3,4,5,6]+unsafeModifyVector :: (FromVector u a) => (FromVector v b) => (u -> v) -> Array a -> Array b+unsafeModifyVector f (UnsafeArray s v) = UnsafeArray s (asVector (f (vectorAs v)))++-- | Representation of an index into a shape (an [Int]). The index is a dimension of the shape.+type Dim = Int++-- | Representation of indexes into a shape (an [Int]). The indexes are dimensions of the shape.+type Dims = [Int]++-- | shape of an Array+--+-- >>> shape a+-- [2,3,4]+shape :: Array a -> [Int]+shape (UnsafeArray s _) = s++-- | rank of an Array+--+-- >>> rank a+-- 3+rank :: Array a -> Int+rank = List.length . shape++-- | size of an Array, which is the total number of elements, if the Array is valid.+--+-- >>> size a+-- 24+size :: Array a -> Int+size = S.size . shape++-- | Number of rows (first dimension size) in an Array. As a convention, a scalar value is still a single row.+--+-- >>> length a+-- 2+-- >>> length (toScalar 0)+-- 1+length :: Array a -> Int+length a = case shape a of+ [] -> 1+ (x : _) -> x++-- | Is the Array empty (has zero number of elements).+--+-- >>> isNull ([2,0] >< [] :: Array ())+-- True+-- >>> isNull ([] >< [4] :: Array Int)+-- False+isNull :: Array a -> Bool+isNull = (0 ==) . size++-- | Extract an element at an index, unsafely.+--+-- >>> index a [1,2,3]+-- 23+index :: Array a -> [Int] -> a+index (UnsafeArray s v) i = V.unsafeIndex v (flatten s i)++infixl 9 !++-- | Extract an element at an index, unsafely.+--+-- >>> a ! [1,2,3]+-- 23+(!) :: Array a -> [Int] -> a+(!) = index++-- | Extract an element at an index, safely.+--+-- >>> a !? [1,2,3]+-- Just 23+-- >>> a !? [2,3,1]+-- Nothing+(!?) :: Array a -> [Int] -> Maybe a+(!?) a xs = bool Nothing (Just (a ! xs)) (xs `isFins` shape a)++-- | Tabulate an array supplying a shape and a tabulation function.+--+-- >>> tabulate [2,3,4] (S.flatten [2,3,4]) == a+-- True+tabulate :: [Int] -> ([Int] -> a) -> Array a+tabulate ds f =+ UnsafeArray ds (V.generate (V.product (asVector ds)) (f . shapen ds))++-- | @backpermute@ is a tabulation where the contents of an array do not need to be accessed, and is thus a fulcrum for leveraging laziness and fusion via the rule:+--+-- > backpermute f g (backpermute f' g' a) == backpermute (f . f') (g . g') a+--+-- Many functions in this module are examples of backpermute usage.+--+-- >>> pretty $ backpermute List.reverse 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]]]+backpermute :: ([Int] -> [Int]) -> ([Int] -> [Int]) -> Array a -> Array a+backpermute f g a = tabulate (f (shape a)) (index a . g)+{-# INLINEABLE backpermute #-}++{- RULES+ "backpermute/backpermute" forall f f' g g' (a :: forall a. Array a)). backpermute f g (backpermute f' g' a) == backpermute (f . f') (g . g') a++-}++-- | Unwrap a scalar.+--+-- >>> let s = array [] [3] :: Array Int+-- >>> fromScalar s+-- 3+fromScalar :: Array a -> a+fromScalar a = index a ([] :: [Int])++-- | Wrap a scalar.+--+-- >>> :t toScalar 2+-- toScalar 2 :: Num a => Array a+toScalar :: a -> Array a+toScalar a = tabulate [] (const a)++-- | Is an array a Scalar?+--+-- >>> isScalar (toScalar (2::Int))+-- True+isScalar :: Array a -> Bool+isScalar a = rank a == 0++-- | Convert a scalar to being a dimensioned array. Do nothing if not a scalar.+--+-- >>> asSingleton (toScalar 4)+-- UnsafeArray [1] [4]+asSingleton :: Array a -> Array a+asSingleton = unsafeModifyShape S.asSingleton++-- | Convert an array with shape [1] to being a scalar (Do nothing if not a shape [1] array).+--+-- >>> asScalar (singleton 3)+-- UnsafeArray [] [3]+asScalar :: Array a -> Array a+asScalar = unsafeModifyShape S.asScalar++-- * Creation++-- | An array with no elements.+--+-- >>> empty+-- UnsafeArray [0] []+empty :: Array a+empty = array [0] []++-- | An enumeration of row-major or [lexicographic](https://en.wikipedia.org/wiki/Lexicographic_order) order.+--+-- >>> pretty $ range [2,3]+-- [[0,1,2],+-- [3,4,5]]+range :: [Int] -> Array Int+range xs = tabulate xs (flatten xs)++-- | An enumeration of col-major or [colexicographic](https://en.wikipedia.org/wiki/Lexicographic_order) order.+--+-- >>> pretty (corange [2,3,4])+-- [[[0,6,12,18],+-- [2,8,14,20],+-- [4,10,16,22]],+-- [[1,7,13,19],+-- [3,9,15,21],+-- [5,11,17,23]]]+corange :: [Int] -> Array Int+corange xs = tabulate xs (flatten (List.reverse xs) . List.reverse)++-- | Indices of an array shape.+--+-- >>> pretty $ indices [3,3]+-- [[[0,0],[0,1],[0,2]],+-- [[1,0],[1,1],[1,2]],+-- [[2,0],[2,1],[2,2]]]+indices :: [Int] -> Array [Int]+indices ds = tabulate ds id++-- | The identity array.+--+-- >>> pretty $ ident [3,3]+-- [[1,0,0],+-- [0,1,0],+-- [0,0,1]]+ident :: (Num a) => [Int] -> Array a+ident ds = tabulate ds (bool 0 1 . isDiag)++-- | Create an array composed of a single value.+--+-- >>> pretty $ konst [3,2] 1+-- [[1,1],+-- [1,1],+-- [1,1]]+konst :: [Int] -> a -> Array a+konst ds a = tabulate ds (const a)++-- | Create an array of shape [1].+--+-- >>> pretty $ singleton 1+-- [1]+-- >>> singleton 3 == toScalar 3+-- False+--+-- >>> asVector (singleton 3) == asVector (toScalar 3)+-- True+singleton :: a -> Array a+singleton a = UnsafeArray [1] (V.singleton a)++-- | Extract the diagonal of an array.+--+-- >>> pretty $ diag (ident [3,3])+-- [1,1,1]+diag ::+ Array a ->+ Array a+diag a = backpermute minDim (replicate (rank a) . getDim 0) a++-- | Expand the array to form a diagonal array.+--+-- >>> pretty $ undiag (range [3])+-- [[0,0,0],+-- [0,1,0],+-- [0,0,2]]+undiag ::+ (Num a) =>+ Array a ->+ Array a+undiag a = tabulate (shape a <> shape a) (\xs -> bool 0 (index a xs) (isDiag xs))++-- | Zip two arrays at an element level.+--+-- >>> zipWith (-) v v+-- UnsafeArray [3] [0,0,0]+zipWith :: (a -> b -> c) -> Array a -> Array b -> Array c+zipWith f (UnsafeArray s v) (UnsafeArray _ v') = UnsafeArray s (V.zipWith f v v')++-- | Zip two arrays at an element level, checking for shape consistency.+--+-- >>> zipWithSafe (-) (range [3]) (range [4])+-- Nothing+zipWithSafe :: (a -> b -> c) -> Array a -> Array b -> Maybe (Array c)+zipWithSafe f (UnsafeArray s v) (UnsafeArray s' v') = bool Nothing (Just $ UnsafeArray s (V.zipWith f v v')) (s == s')++-- | Modify a single value at an index.+--+-- >>> pretty $ modify [0,0] (const 100) (range [3,2])+-- [[100,1],+-- [2,3],+-- [4,5]]+modify :: [Int] -> (a -> a) -> Array a -> Array a+modify ds f a = tabulate (shape a) (\s -> bool id f (s == ds) (index a 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 ::+ ([Int] -> a -> b) ->+ Array a ->+ Array b+imap f a = zipWith f (indices (shape a)) a++-- | With a function that takes dimensions and (type-level) parameters, apply the parameters to the initial dimensions. ie+--+-- > rowWise f xs = f [0..] xs+--+-- >>> rowWise indexes [1,0] a+-- UnsafeArray [4] [12,13,14,15]+rowWise :: (Dims -> [x] -> Array a -> Array a) -> [x] -> Array a -> Array a+rowWise f xs a = f [0 .. (S.rank xs - 1)] xs a++-- | With a function that takes dimensions and (type-level) parameters, apply the parameters to the the last dimensions. ie+--+-- > colWise f xs = f (List.reverse [0 .. (rank a - 1)]) xs+--+-- >>> colWise indexes [1,0] a+-- UnsafeArray [2] [1,13]+colWise :: (Dims -> [x] -> Array a -> Array a) -> [x] -> Array a -> Array a+colWise f xs a = f (List.reverse [(rank a - S.rank xs) .. (rank a - 1)]) xs a++-- | With a function that takes a dimension and a parameter, fold dimensions and parameters using the function.+--+-- >>> dimsWise take [0,2] [1,2] a+-- UnsafeArray [1,3,2] [0,1,4,5,8,9]+dimsWise :: (Dim -> x -> Array a -> Array a) -> Dims -> [x] -> Array a -> Array a+dimsWise f ds xs a = foldl' (\a' (d, x) -> f d x a') a (List.zip ds xs)++-- | Take the top-most elements across the specified dimension. Negative values take the bottom-most. No index check is performed.+--+-- > take d x == takes [(d,x)]+--+-- >>> pretty $ take 2 1 a+-- [[[0],+-- [4],+-- [8]],+-- [[12],+-- [16],+-- [20]]]+-- >>> pretty $ take 2 (-1) a+-- [[[3],+-- [7],+-- [11]],+-- [[15],+-- [19],+-- [23]]]+take ::+ Dim ->+ Int ->+ Array a ->+ Array a+take d t a = backpermute dsNew (modifyDim d (\x -> x + bool 0 (getDim d (shape a) + t) (t < 0))) a+ where+ dsNew = takeDim d (abs t)++-- | Drop the top-most elements across the specified dimension. Negative values take the bottom-most.+--+-- >>> pretty $ drop 2 1 a+-- [[[1,2,3],+-- [5,6,7],+-- [9,10,11]],+-- [[13,14,15],+-- [17,18,19],+-- [21,22,23]]]+-- >>> pretty $ drop 2 (-1) a+-- [[[0,1,2],+-- [4,5,6],+-- [8,9,10]],+-- [[12,13,14],+-- [16,17,18],+-- [20,21,22]]]+drop ::+ Dim ->+ Int ->+ Array a ->+ Array a+drop d t a = backpermute dsNew (modifyDim d (\x -> x + bool t 0 (t < 0))) a+ where+ dsNew = dropDim d (abs t)++-- | Select an index along a dimension.+--+-- >>> let s = select 2 3 a+-- >>> pretty s+-- [[3,7,11],+-- [15,19,23]]+select ::+ Dim ->+ Int ->+ Array a ->+ Array a+select d x a = backpermute (deleteDim d) (insertDim d x) a++-- | Insert along a dimension at a position.+--+-- >>> pretty $ insert 2 0 a (konst [2,3] 0)+-- [[[0,0,1,2,3],+-- [0,4,5,6,7],+-- [0,8,9,10,11]],+-- [[0,12,13,14,15],+-- [0,16,17,18,19],+-- [0,20,21,22,23]]]+-- >>> insert 0 0 (toScalar 1) (toScalar 2)+-- UnsafeArray [2] [2,1]+insert ::+ Dim ->+ Int ->+ Array a ->+ Array a ->+ Array a+insert d i a b = tabulate (incAt d (shape a)) go+ where+ go s+ | getDim d s == i = index b (deleteDim d s)+ | getDim d s < i = index a s+ | otherwise = index a (decAt d s)++-- | Delete along a dimension at a position.+--+-- >>> pretty $ delete 2 0 a+-- [[[0,1,2],+-- [4,5,6],+-- [8,9,10]],+-- [[12,13,14],+-- [16,17,18],+-- [20,21,22]]]+delete ::+ Dim ->+ Int ->+ Array a ->+ Array a+delete d i a = backpermute (decAt d) (\s -> bool s (incAt d s) (getDim d s < i)) a++-- | Insert along a dimension at the end.+--+-- >>> pretty $ append 2 a (konst [2,3] 0)+-- [[[0,1,2,3,0],+-- [4,5,6,7,0],+-- [8,9,10,11,0]],+-- [[12,13,14,15,0],+-- [16,17,18,19,0],+-- [20,21,22,23,0]]]+append ::+ Dim ->+ Array a ->+ Array a ->+ Array a+append d a b = insert d (getDim d (shape a)) a b++-- | Insert along a dimension at the beginning.+--+-- >>> pretty $ prepend 2 (konst [2,3] 0) a+-- [[[0,0,1,2,3],+-- [0,4,5,6,7],+-- [0,8,9,10,11]],+-- [[0,12,13,14,15],+-- [0,16,17,18,19],+-- [0,20,21,22,23]]]+prepend ::+ Dim ->+ Array a ->+ Array a ->+ Array a+prepend d a b = insert d 0 b a++-- | Concatenate along a dimension.+--+-- >>> shape $ concatenate 1 a a+-- [2,6,4]+-- >>> concatenate 0 (toScalar 1) (toScalar 2)+-- UnsafeArray [2] [1,2]+-- >>> concatenate 0 (toScalar 0) (asArray [1..3])+-- UnsafeArray [4] [0,1,2,3]+concatenate ::+ Dim ->+ Array a ->+ Array a ->+ Array a+concatenate d a0 a1 = tabulate (S.concatenate d (shape a0) (shape a1)) go+ where+ go s =+ bool+ (index a0 s)+ ( index+ a1+ ( insertDim+ d+ (getDim d s - getDim d ds0)+ (deleteDim d s)+ )+ )+ (getDim d s >= getDim d ds0)+ ds0 = shape a0++-- | Combine two arrays as a new dimension of a new array.+--+-- >>> pretty $ couple 0 (asArray [1,2,3]) (asArray [4,5,6::Int])+-- [[1,2,3],+-- [4,5,6]]+couple :: Int -> Array a -> Array a -> Array a+couple d a a' = concatenate d (elongate d a) (elongate d a')++-- | Slice along a dimension with the supplied offset & length.+--+-- >>> let s = slice 2 1 2 a+-- >>> pretty s+-- [[[1,2],+-- [5,6],+-- [9,10]],+-- [[13,14],+-- [17,18],+-- [21,22]]]+slice ::+ Dim ->+ Int ->+ Int ->+ Array a ->+ Array a+slice d o l a = backpermute (setDim d l) (modifyDim d (+ o)) a++-- | Rotate an array along a dimension.+--+-- >>> pretty $ rotate 1 2 a+-- [[[8,9,10,11],+-- [0,1,2,3],+-- [4,5,6,7]],+-- [[20,21,22,23],+-- [12,13,14,15],+-- [16,17,18,19]]]+rotate ::+ Dim ->+ Int ->+ Array a ->+ Array a+rotate d r a = backpermute id (rotateIndex d r (shape a)) a++-- * multi-dimension operators++-- | Takes the top-most elements across the supplied dimension,n tuples. Negative values take the bottom-most.+--+-- > takes == dimsWise take+--+-- >>> pretty $ takes [0,2] [1,-3] a+-- [[[1,2,3],+-- [5,6,7],+-- [9,10,11]]]+takes ::+ Dims ->+ [Int] ->+ Array a ->+ Array a+takes ds xs a = backpermute dsNew (List.zipWith (+) start) a+ where+ dsNew = setDims ds xsAbs+ start = List.zipWith (\x s -> bool 0 (s + x) (x < 0)) (setDims ds xs (replicate (rank a) 0)) (shape a)+ xsAbs = fmap abs xs++-- | Drops the top-most elements. Negative values drop the bottom-most.+--+-- >>> pretty $ drops [0,1,2] [1,2,-3] a+-- [[[20]]]+drops ::+ Dims ->+ [Int] ->+ Array a ->+ Array a+drops ds xs a = backpermute dsNew (List.zipWith (\d' s' -> bool (d' + s') s' (d' < 0)) xsNew) a+ where+ dsNew = dropDims ds xsAbs+ xsNew = setDims ds xs (replicate (rank a) 0)+ xsAbs = fmap abs xs++-- | Select by dimensions and indexes.+--+-- >>> let s = indexes [0,1] [1,1] a+-- >>> pretty s+-- [16,17,18,19]+indexes ::+ Dims ->+ [Int] ->+ Array a ->+ Array a+indexes ds xs a = backpermute (deleteDims ds) (insertDims ds xs) a++-- | Slice along dimensions with the supplied offsets and lengths.+--+-- >>> let s = slices [2,0] [1,1] [2,1] a+-- >>> pretty s+-- [[[13,14],+-- [17,18],+-- [21,22]]]+slices ::+ Dims ->+ [Int] ->+ [Int] ->+ Array a ->+ Array a+slices ds os ls a = dimsWise (\d (o, l) -> slice d o l) ds (List.zip os ls) a++-- | Select the first element along the supplied dimensions.+--+-- >>> pretty $ heads [0,2] a+-- [0,4,8]+heads :: Dims -> Array a -> Array a+heads ds a = indexes ds (List.replicate (S.rank ds) 0) a++-- | Select the last element along the supplied dimensions.+--+-- >>> pretty $ lasts [0,2] a+-- [15,19,23]+lasts :: Dims -> Array a -> Array a+lasts ds a = indexes ds lastds a+ where+ lastds = (\i -> getDim i (shape a) - 1) <$> ds++-- | Select the tail elements along the supplied dimensions.+--+-- >>> pretty $ tails [0,2] a+-- [[[13,14,15],+-- [17,18,19],+-- [21,22,23]]]+tails :: Dims -> Array a -> Array a+tails ds a = slices ds os ls a+ where+ os = List.replicate (S.rank ds) 1+ ls = getLastPositions ds (shape a)++-- | Select the init elements along the supplied dimensions.+--+-- >>> pretty $ inits [0,2] a+-- [[[0,1,2],+-- [4,5,6],+-- [8,9,10]]]+inits :: Dims -> Array a -> Array a+inits ds a = slices ds os ls a+ where+ os = List.replicate (S.rank ds) 0+ ls = getLastPositions ds (shape a)++-- | Extracts dimensions to an outer layer.+--+-- >>> pretty $ shape <$> extracts [0] a+-- [[3,4],[3,4]]+extracts ::+ Dims ->+ Array a ->+ Array (Array a)+extracts ds a = tabulate (getDims ds (shape a)) go+ where+ go s = indexes ds s a++-- | Reduce along specified dimensions, using the supplied fold.+--+-- >>> pretty $ reduces [0] sum a+-- [66,210]+-- >>> pretty $ reduces [0,2] sum a+-- [[12,15,18,21],+-- [48,51,54,57]]+reduces ::+ Dims ->+ (Array a -> b) ->+ Array a ->+ Array b+reduces ds f a = fmap f (extracts ds a)++-- | Join inner and outer dimension layers by supplied dimensions. No checks on shape.+--+-- >>> let e = extracts [1,0] a+-- >>> let j = joins [1,0] e+-- >>> a == j+-- True+joins ::+ Dims ->+ Array (Array a) ->+ Array a+joins ds a = tabulate (insertDims ds so si) go+ where+ go s = index (index a (getDims ds s)) (deleteDims ds s)+ so = shape a+ si = shape (index a (replicate (rank a) 0))++-- | Join inner and outer dimension layers by supplied dimensions. Check inner layer shape.+--+-- >>> let e = extracts [1,0] a+-- >>> (Just j) = joinsSafe [1,0] e+-- >>> a == j+-- True+joinsSafe ::+ Dims ->+ Array (Array a) ->+ Maybe (Array a)+joinsSafe ds a =+ bool+ Nothing+ (Just $ joins ds a)+ (allEqual (fmap shape a))++-- | Join inner and outer dimension layers in outer dimension order.+--+-- >>> a == join (extracts [0,1] a)+-- True+join ::+ Array (Array a) ->+ Array a+join a = joins (S.dimsOf (shape a)) a++-- | Join inner and outer dimension layers in outer dimension order, checking for consistent inner dimension shape.+--+-- >>> joinSafe (extracts [0,1] a)+-- Just (UnsafeArray [2,3,4] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23])+joinSafe ::+ Array (Array a) ->+ Maybe (Array a)+joinSafe a =+ bool+ Nothing+ (Just $ join a)+ (allEqual (fmap shape a))++-- | Satisfy a predicate across all elements+allEqual :: (Eq a) => Array a -> Bool+allEqual a = case arrayAs a of+ [] -> True+ (x : xs) -> all (== x) xs++-- | Traverse along specified dimensions.+--+-- traverses [1] print (range [2,3])+-- 0+-- 3+-- 1+-- 4+-- 2+-- 5+-- UnsafeArray [2,3] [(),(),(),(),(),()]+traverses ::+ (Applicative f) =>+ Dims ->+ (a -> f b) ->+ Array a ->+ f (Array b)+traverses ds f a = joins ds <$> traverse (traverse f) (extracts ds a)++-- | Maps a function along specified dimensions.+--+-- >>> pretty $ maps [1] transpose a+-- [[[0,12],+-- [4,16],+-- [8,20]],+-- [[1,13],+-- [5,17],+-- [9,21]],+-- [[2,14],+-- [6,18],+-- [10,22]],+-- [[3,15],+-- [7,19],+-- [11,23]]]+maps ::+ Dims ->+ (Array a -> Array b) ->+ Array a ->+ Array b+maps ds f a = joins ds (fmap f (extracts ds a))++-- | Filters along specified dimensions (which are flattened).+--+-- >>> pretty $ filters [0,1] (any ((==0) . (`mod` 7))) a+-- [[0,1,2,3],+-- [4,5,6,7],+-- [12,13,14,15],+-- [20,21,22,23]]+filters ::+ Dims ->+ (Array a -> Bool) ->+ Array a ->+ Array a+filters ds p a = join (asArray $ V.filter p $ asVector (extracts ds a))++-- | Zips two arrays with a function along specified dimensions.+--+-- >>> pretty $ zips [0,1] (zipWith (,)) a (reverses [0] a)+-- [[[(0,12),(1,13),(2,14),(3,15)],+-- [(4,16),(5,17),(6,18),(7,19)],+-- [(8,20),(9,21),(10,22),(11,23)]],+-- [[(12,0),(13,1),(14,2),(15,3)],+-- [(16,4),(17,5),(18,6),(19,7)],+-- [(20,8),(21,9),(22,10),(23,11)]]]+zips ::+ Dims ->+ (Array a -> Array b -> Array c) ->+ Array a ->+ Array b ->+ Array c+zips ds f a b = joins ds (zipWith f (extracts ds a) (extracts ds b))++-- | Zips two arrays with a function along specified dimensions, checking shapes.+--+-- >>> zipsSafe [0] (zipWith (,)) (asArray [1::Int]) (asArray [1,2::Int])+-- Nothing+zipsSafe ::+ Dims ->+ (Array a -> Array b -> Array c) ->+ Array a ->+ Array b ->+ Maybe (Array c)+zipsSafe ds f a b =+ bool+ (Just $ joins ds (zipWith f (extracts ds a) (extracts ds b)))+ Nothing+ (shape a /= (shape b :: [Int]))++-- | Modify using the supplied function along dimensions & positions.+--+-- >>> pretty $ modifies (fmap (100+)) [2] [0] a+-- [[[100,1,2,3],+-- [104,5,6,7],+-- [108,9,10,11]],+-- [[112,13,14,15],+-- [116,17,18,19],+-- [120,21,22,23]]]+modifies ::+ (Array a -> Array a) ->+ Dims ->+ [Int] ->+ Array a ->+ Array a+modifies f ds ps a = joins ds $ modify ps f (extracts ds a)++-- | Apply a binary function between successive slices, across dimensions and lags.+--+-- >>> pretty $ diffs [1] [1] (zipWith (-)) a+-- [[[4,4,4,4],+-- [4,4,4,4]],+-- [[4,4,4,4],+-- [4,4,4,4]]]+diffs :: Dims -> [Int] -> (Array a -> Array a -> Array b) -> Array a -> Array b+diffs ds xs f a = zips ds f (drops ds xs a) (drops ds (fmap P.negate xs) a)++-- | Product two arrays using the supplied binary function.+--+-- For context, if the function is multiply, and the arrays are tensors,+-- then this can be interpreted as a [tensor product](https://en.wikipedia.org/wiki/Tensor_product).+-- The concept of a tensor product is a dense crossroad, and a complete treatment is elsewhere. To quote the wiki article:+--+-- ... the tensor product can be extended to other categories of mathematical objects in addition to vector spaces, such as to matrices, tensors, algebras, topological vector spaces, and modules. In each such case the tensor product is characterized by a similar universal property: it is the freest bilinear operation. The general concept of a "tensor product" is captured by monoidal categories; that is, the class of all things that have a tensor product is a monoidal category.+--+-- >>> x = array [3] [1,2,3]+-- >>> pretty $ expand (*) x x+-- [[1,2,3],+-- [2,4,6],+-- [3,6,9]]+--+-- Alternatively, expand can be understood as representing the permutation of element pairs of two arrays, so like the Applicative List instance.+--+-- >>> i2 = indices [2,2]+-- >>> pretty $ expand (,) i2 i2+-- [[[[([0,0],[0,0]),([0,0],[0,1])],+-- [([0,0],[1,0]),([0,0],[1,1])]],+-- [[([0,1],[0,0]),([0,1],[0,1])],+-- [([0,1],[1,0]),([0,1],[1,1])]]],+-- [[[([1,0],[0,0]),([1,0],[0,1])],+-- [([1,0],[1,0]),([1,0],[1,1])]],+-- [[([1,1],[0,0]),([1,1],[0,1])],+-- [([1,1],[1,0]),([1,1],[1,1])]]]]+expand ::+ (a -> b -> c) ->+ Array a ->+ Array b ->+ Array c+expand f a b = tabulate (shape a <> shape b) (\i -> f (index a (List.take r i)) (index b (List.drop r i)))+ where+ r = rank 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 ::+ (a -> b -> c) ->+ Array a ->+ Array b ->+ Array c+coexpand f a b = tabulate (shape a <> shape b) (\i -> f (index a (List.drop r i)) (index b (List.take r i)))+ where+ r = rank a++-- | Contract an array by applying the supplied (folding) function on diagonal elements of the dimensions.+--+-- This generalises a tensor contraction by allowing the number of contracting diagonals to be other than 2.+--+--+-- >>> pretty $ contract [1,2] sum (expand (*) m (transpose m))+-- [[5,14],+-- [14,50]]+contract ::+ Dims ->+ (Array a -> b) ->+ Array a ->+ Array b+contract ds f a = f . diag <$> extracts (exceptDims ds (shape a)) a++-- | Product two arrays using the supplied function and then contract the result using the supplied matching dimensions and function.+--+-- >>> pretty $ prod [1] [0] sum (*) (range [2,3]) (range [3,2])+-- [[10,13],+-- [28,40]]+--+-- With full laziness, this computation would be equivalent to:+--+-- > f . diag <$> extracts ds' (expand g a b)+prod ::+ Dims ->+ Dims ->+ (Array c -> d) ->+ (a -> b -> c) ->+ Array a ->+ Array b ->+ Array d+prod ds0 ds1 g f a b = tabulate (S.deleteDims ds0 (shape a) <> S.deleteDims ds1 (shape b)) (\so -> g $ tabulate (S.getDims ds0 (shape a)) (\si -> f (index a (S.insertDims ds0 si (List.take sp so))) (index b (S.insertDims ds1 si (List.drop sp so)))))+ where+ sp = rank a - S.rank 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 ::+ (Array c -> d) ->+ (a -> b -> c) ->+ Array a ->+ Array b ->+ Array d+dot f g a b = contract [r - 1, r] f (expand g a b)+ where+ r = rank a++-- | 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 ::+ (Num a) =>+ Array a ->+ Array a ->+ Array a+mult = dot sum (*)++-- | @windows xs@ are xs-sized windows of an array+--+-- >>> shape $ windows [2,2] (range [4,3,2])+-- [3,2,2,2,2]+windows :: [Int] -> Array a -> Array a+windows xs a = backpermute (expandWindows xs) (indexWindows (S.rank xs)) a++-- | Find the starting positions of occurences of one array in another.+--+-- >>> a = cycle [4,4] (range [3]) :: Array Int+-- >>> i = array [2,2] [1,2,2,0] :: Array Int+-- >>> pretty $ find i a+-- [[False,True,False],+-- [True,False,False],+-- [False,False,True]]+find :: (Eq a) => Array a -> Array a -> Array Bool+find i a = xs+ where+ i' = rerank (rank a) i+ ws = windows (shape i') a+ xs = fmap (== i') (extracts (dimWindows (expandWindows (shape i') (shape a)) (shape 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 Int+-- >>> i = konst [2,2] 1 :: Array Int+-- >>> pretty $ findNoOverlap i a+-- [[True,False,True,False],+-- [False,False,False,False],+-- [True,False,True,False],+-- [False,False,False,False]]+findNoOverlap :: (Eq a) => Array a -> Array a -> Array Bool+findNoOverlap i a = r+ where+ f :: Array Bool+ f = find i a++ cl :: [Int] -> [[Int]]+ cl sh = List.filter (P.not . any (> 0) . List.init) $ List.filter (P.not . all (>= 0)) $ arrayAs $ tabulate ((\x -> 2 * x - 1) <$> sh) (\s -> List.zipWith (\x x0 -> x - x0 + 1) s sh)+ go r' s = index f s && not (any (index r') (List.filter (\x -> isFins x (shape f)) $ fmap (List.zipWith (+) s) (cl (shape i))))+ r = tabulate (shape f) (go r)++-- | Find the indices of the starting location of one array in another.+--+-- >>> b = cycle [4,4] (range [3]) :: Array Int+-- >>> i = array [2,2] [1,2,2,0] :: Array Int+-- >>> pretty $ findIndices i b+-- [[0,1],[1,0],[2,2]]+findIndices :: (Eq a) => Array a -> Array a -> Array [Int]+findIndices i a = fmap fst $ vectorAs $ V.filter snd $ asVector $ imap (,) b+ where+ b = find i a++-- | Check if the first array is a prefix of the second+--+-- >>> isPrefixOf (array [2,2] [0,1,4,5]) a+-- True+isPrefixOf :: (Eq a) => Array a -> Array a -> Bool+isPrefixOf p a = p == cut (shape p) a++-- | Check if the first array is a suffix of the second+--+-- >>> isSuffixOf (array [2,2] [18,19,22,23]) a+-- True+isSuffixOf :: (Eq a) => Array a -> Array a -> Bool+isSuffixOf p a = p == cutSuffix (shape p) a++-- | Check if the first array is an infix of the second+--+-- >>> isInfixOf (array [2,2] [18,19,22,23]) a+-- True+isInfixOf :: (Eq a) => Array a -> Array a -> Bool+isInfixOf p a = or $ find p a++-- * shape manipulation++-- | Fill an array with the supplied value without regard to the original shape or cut the array values to match array size.+--+-- > validate (def x a) == True+--+-- >>> pretty $ fill 0 (array [3] [])+-- [0,0,0]+-- >>> pretty $ fill 0 (array [3] [1..4])+-- [1,2,3]+fill :: a -> Array a -> Array a+fill x (UnsafeArray s v) = UnsafeArray s (V.take (S.size s) (v <> V.replicate (S.size s - V.length v) x))++-- | Cut an array to form a new (smaller) shape. Errors if the new shape is larger. The old array is reranked to the rank of the new shape first.+--+-- >>> cut [2] (array [4] [0..3] :: Array Int)+-- UnsafeArray [2] [0,1]+cut ::+ [Int] ->+ Array a ->+ Array a+cut s' a = bool (error "bad cut") (tabulate s' (index a')) (isSubset s' (shape a))+ where+ a' = rerank (S.rank s') a++-- | Cut an array to form a new (smaller) shape, using suffix elements. Errors if the new shape is larger. The old array is reranked to the rank of the new shape first.+--+-- >>> cutSuffix [2,2] a+-- UnsafeArray [2,2] [18,19,22,23]+cutSuffix ::+ [Int] ->+ Array a ->+ Array a+cutSuffix s' a = bool (error "bad cut") (tabulate s' (index a' . List.zipWith (+) diffDim)) (isSubset s' (shape a))+ where+ a' = rerank (S.rank s') a+ diffDim = List.zipWith (-) (shape a') s'++-- | Pad an array to form a new shape, supplying a default value for elements outside the shape of the old array. The old array is reranked to the rank of the new shape first.+--+-- >>> pad 0 [5] (array [4] [0..3] :: Array Int)+-- UnsafeArray [5] [0,1,2,3,0]+pad ::+ a ->+ [Int] ->+ Array a ->+ Array a+pad d s' a = tabulate s' (\s -> bool d (index a' s) (s `isFins` shape a'))+ where+ a' = rerank (S.rank s') a++-- | Left pad an array to form a new shape, supplying a default value for elements outside the shape of the old array.+--+-- >>> lpad 0 [5] (array [4] [0..3] :: Array Int)+-- UnsafeArray [5] [0,0,1,2,3]+-- >>> pretty $ lpad 0 [3,3] (range [2,2] :: Array Int)+-- [[0,0,0],+-- [0,0,1],+-- [0,2,3]]+lpad ::+ a ->+ [Int] ->+ Array a ->+ Array a+lpad d s' a = tabulate s' (\s -> bool d (index a' (olds s)) (olds s `S.isFins` shape a'))+ where+ a' = rerank (S.rank s') a+ gap = List.zipWith (-) s' (shape a')+ olds s = List.zipWith (-) s gap++-- | Reshape an array (with the same or less number of elements).+--+-- >>> pretty $ reshape [4,3,2] a+-- [[[0,1],+-- [2,3],+-- [4,5]],+-- [[6,7],+-- [8,9],+-- [10,11]],+-- [[12,13],+-- [14,15],+-- [16,17]],+-- [[18,19],+-- [20,21],+-- [22,23]]]+reshape ::+ [Int] ->+ Array a ->+ Array a+reshape s a = backpermute (const s) (shapen (shape a) . flatten s) a++-- | Make an Array single dimensional.+--+-- >>> pretty $ flat (range [2,2])+-- [0,1,2,3]+-- >>> pretty (flat $ toScalar 0)+-- [0]+flat :: Array a -> Array a+flat a = unsafeModifyShape (pure . S.size) a++-- | Reshape an array, repeating the original array. The shape of the array should be a suffix of the new shape.+--+-- >>> pretty $ repeat [2,2,2] (array [2] [1,2])+-- [[[1,2],+-- [1,2]],+-- [[1,2],+-- [1,2]]]+--+-- > repeat ds (toScalar x) == konst ds x+repeat ::+ [Int] ->+ Array a ->+ Array a+repeat s a = backpermute (const s) (List.drop (S.rank s - rank a)) a++-- | Reshape an array, cycling through the elements without regard to the original shape.+--+-- >>> pretty $ cycle [2,2,2] (array [3] [1,2,3])+-- [[[1,2],+-- [3,1]],+-- [[2,3],+-- [1,2]]]+cycle ::+ [Int] ->+ Array a ->+ Array a+cycle s a = backpermute (const s) (shapen (shape a) . (`mod` size a) . flatten s) a++-- | Change rank by adding new dimensions at the front, if the new rank is greater, or combining dimensions (from left to right) into rows, if the new rank is lower.+--+-- >>> shape (rerank 4 a)+-- [1,2,3,4]+-- >>> shape (rerank 2 a)+-- [6,4]+--+-- > flat == rerank 1+rerank :: Int -> Array a -> Array a+rerank r a = unsafeModifyShape (S.rerank r) a++-- | Change the order of dimensions.+--+-- >>> pretty $ reorder [2,0,1] a+-- [[[0,4,8],+-- [12,16,20]],+-- [[1,5,9],+-- [13,17,21]],+-- [[2,6,10],+-- [14,18,22]],+-- [[3,7,11],+-- [15,19,23]]]+reorder ::+ Dims ->+ Array a ->+ Array a+reorder ds a = backpermute (`S.reorder` ds) (\s -> insertDims ds s []) a++-- | Remove single dimensions.+--+-- >>> let sq = array [2,1,3,4,1] [1..24] :: Array Int+-- >>> shape $ squeeze sq+-- [2,3,4]+--+-- >>> shape $ squeeze (singleton 0)+-- []+squeeze ::+ Array a ->+ Array a+squeeze a = unsafeModifyShape S.squeeze a++-- | Insert a single dimension at the supplied position.+--+-- >>> shape $ elongate 1 a+-- [2,1,3,4]+-- >>> elongate 0 (toScalar 1)+-- UnsafeArray [1] [1]+elongate ::+ Dim ->+ Array a ->+ Array a+elongate d a = unsafeModifyShape (insertDim d 1) a++-- | Reverse indices eg transposes the element A/ijk/ to A/kji/.+--+-- >>> index (transpose a) [1,0,0] == index a [0,0,1]+-- True+-- >>> pretty $ transpose (array [2,2,2] [1..8])+-- [[[1,5],+-- [3,7]],+-- [[2,6],+-- [4,8]]]+transpose :: Array a -> Array a+transpose a = backpermute List.reverse List.reverse a++-- | Inflate an array by inserting a new dimension given a supplied dimension and size.+--+-- alt name: replicate+--+-- >>> pretty $ inflate 0 2 (array [3] [0,1,2])+-- [[0,1,2],+-- [0,1,2]]+inflate ::+ Dim ->+ Int ->+ Array a ->+ Array a+inflate d n a = backpermute (insertDim d n) (deleteDim d) a++-- | Intercalate an array along dimensions.+--+-- >>> pretty $ intercalate 2 (konst [2,3] 0) a+-- [[[0,0,1,0,2,0,3],+-- [4,0,5,0,6,0,7],+-- [8,0,9,0,10,0,11]],+-- [[12,0,13,0,14,0,15],+-- [16,0,17,0,18,0,19],+-- [20,0,21,0,22,0,23]]]+intercalate :: Dim -> Array a -> Array a -> Array a+intercalate d i a = joins [d] $ asArray (List.intersperse i (arrayAs (extracts [d] a)))++-- | Intersperse an element along dimensions.+--+-- >>> pretty $ intersperse 2 0 a+-- [[[0,0,1,0,2,0,3],+-- [4,0,5,0,6,0,7],+-- [8,0,9,0,10,0,11]],+-- [[12,0,13,0,14,0,15],+-- [16,0,17,0,18,0,19],+-- [20,0,21,0,22,0,23]]]+intersperse :: Dim -> a -> Array a -> Array a+intersperse d i a = intercalate d (konst (deleteDim d (shape a)) i) a++-- | Concatenate and replace dimensions, creating a new dimension at the supplied postion.+--+-- >>> pretty $ concats [0,1] 1 a+-- [[0,4,8,12,16,20],+-- [1,5,9,13,17,21],+-- [2,6,10,14,18,22],+-- [3,7,11,15,19,23]]+concats ::+ Dims ->+ Int ->+ Array a ->+ Array a+concats ds n a = backpermute (concatDims ds n) (unconcatDimsIndex ds n (shape a)) a++-- | Reverses element order along specified dimensions.+--+-- >>> pretty $ reverses [0,1] a+-- [[[20,21,22,23],+-- [16,17,18,19],+-- [12,13,14,15]],+-- [[8,9,10,11],+-- [4,5,6,7],+-- [0,1,2,3]]]+reverses ::+ Dims ->+ Array a ->+ Array a+reverses ds a = backpermute id (reverseIndex ds (shape a)) a++-- | Rotate an array by/along dimensions & offsets.+--+-- >>> pretty $ rotates [1] [2] a+-- [[[8,9,10,11],+-- [0,1,2,3],+-- [4,5,6,7]],+-- [[20,21,22,23],+-- [12,13,14,15],+-- [16,17,18,19]]]+rotates ::+ Dims ->+ [Int] ->+ Array a ->+ Array a+rotates ds rs a = backpermute id (rotatesIndex ds rs (shape a)) a++-- * sorting++-- | Sort an array along the supplied dimensions.+--+-- >>> sorts [0] (array [2,2] [2,3,1,4])+-- UnsafeArray [2,2] [1,4,2,3]+-- >>> sorts [1] (array [2,2] [2,3,1,4])+-- UnsafeArray [2,2] [2,3,1,4]+-- >>> sorts [0,1] (array [2,2] [2,3,1,4])+-- UnsafeArray [2,2] [1,2,3,4]+sorts :: (Ord a) => Dims -> Array a -> Array a+sorts ds a = joins ds $ unsafeModifyVector sortV (extracts ds a)++-- | The indices into the array if it were sorted by a comparison function along the dimensions supplied.+--+-- >>> import Data.Ord (Down (..))+-- >>> sortsBy [0] (fmap Down) (array [2,2] [2,3,1,4])+-- UnsafeArray [2,2] [2,3,1,4]+sortsBy :: (Ord b) => Dims -> (Array a -> Array b) -> Array a -> Array a+sortsBy ds c a = joins ds $ unsafeModifyVector (sortByV c) (extracts ds a)++-- | The indices into the array if it were sorted along the dimensions supplied.+--+-- >>> orders [0] (array [2,2] [2,3,1,4])+-- UnsafeArray [2] [1,0]+orders :: (Ord a) => Dims -> Array a -> Array Int+orders ds a = unsafeModifyVector orderV (extracts ds a)++-- | The indices into the array if it were sorted by a comparison function along the dimensions supplied.+--+-- >>> import Data.Ord (Down (..))+-- >>> ordersBy [0] (fmap Down) (array [2,2] [2,3,1,4])+-- UnsafeArray [2] [0,1]+ordersBy :: (Ord b) => Dims -> (Array a -> Array b) -> Array a -> Array Int+ordersBy ds c a = unsafeModifyVector (orderByV c) (extracts ds a)++-- * transmission++-- | Apply a binary array function to two arrays with matching shapes across the supplied dimensions. No check on shapes.+--+-- >>> a = array [2,3] [0..5]+-- >>> b = array [3] [0..2]+-- >>> pretty $ telecasts [1] [0] (concatenate 0) a b+-- [[0,1,2],+-- [3,4,5],+-- [0,1,2]]+telecasts :: Dims -> Dims -> (Array a -> Array b -> Array c) -> Array a -> Array b -> Array c+telecasts dsa dsb f a b = zipWith f (extracts dsa a) (extracts dsb b) & joins dsa++-- | Apply a binary array function to two arrays with matching shapes across the supplied dimensions. Checks shape.+--+-- >>> a = array [2,3] [0..5]+-- >>> b = array [1] [1]+-- >>> telecastsSafe [0] [0] (zipWith (+)) a b+-- Nothing+telecastsSafe :: Dims -> Dims -> (Array a -> Array b -> Array c) -> Array a -> Array b -> Maybe (Array c)+telecastsSafe dsa dsb f a b =+ bool+ (Just $ telecasts dsa dsb f a b)+ Nothing+ (shape (extracts dsa a) /= (shape (extracts dsb b) :: [Int]))++-- | Apply a binary array function to two arrays where the shape of the first array is a prefix of the second array. No checks on shape.+--+-- >>> a = array [2,3] [0..5]+-- >>> pretty $ transmit (zipWith (+)) (toScalar 1) a+-- [[1,2,3],+-- [4,5,6]]+transmit :: (Array a -> Array b -> Array c) -> Array a -> Array b -> Array c+transmit f a b = maps ds (f a) b+ where+ ds = [(rank a) .. (rank b - 1)]++-- | Apply a binary array function to two arrays where the shape of the first array is a prefix of the second array. Checks shape.+--+-- >>> a = array [2,3] [0..5]+-- >>> transmitSafe (zipWith (+)) (array [3] [1,2,3]) a+-- Nothing+transmitSafe :: (Array a -> Array b -> Array c) -> Array a -> Array b -> Maybe (Array c)+transmitSafe f a b = bool Nothing (Just $ transmit f a b) (shape a `List.isPrefixOf` shape b)++-- | Transmit an operation if the first array is a prefix of the second or vice versa.+--+-- >>> pretty $ transmitOp (*) a (asArray [1,2])+-- [[[0,1,2,3],+-- [4,5,6,7],+-- [8,9,10,11]],+-- [[24,26,28,30],+-- [32,34,36,38],+-- [40,42,44,46]]]+transmitOp :: (a -> b -> c) -> Array a -> Array b -> Array c+transmitOp f a b+ | shape a == shape b = zipWith f a b+ | shape a `List.isPrefixOf` shape b = transmit (zipWith f) a b+ | shape b `List.isPrefixOf` shape a = transmit (zipWith (flip f)) b a+ | otherwise = error "bad shapes"++-- | Vector specialisation of 'range'+--+-- >>> iota 5+-- UnsafeArray [5] [0,1,2,3,4]+iota :: Int -> Array Int+iota n = range [n]++-- * row (first dimension) specializations++-- | Add a new row+--+-- >>> pretty $ cons (array [2] [0,1]) (array [2,2] [2,3,4,5])+-- [[0,1],+-- [2,3],+-- [4,5]]+cons :: Array a -> Array a -> Array a+cons = prepend 0++-- | split an array into the first row and the remaining rows.+--+-- >>> uncons (array [3,2] [0..5])+-- (UnsafeArray [2] [0,1],UnsafeArray [2,2] [2,3,4,5])+uncons :: Array a -> (Array a, Array a)+uncons a = (heads [0] a', tails [0] a')+ where+ a' = asSingleton a++-- | Convenience pattern for row extraction and consolidation at the beginning of an Array.+--+-- >>> (x:<xs) = array [4] [0..3]+-- >>> x+-- UnsafeArray [] [0]+-- >>> xs+-- UnsafeArray [3] [1,2,3]+-- >>> (x:<xs)+-- UnsafeArray [4] [0,1,2,3]+pattern (:<) :: Array a -> Array a -> Array a+pattern x :< xs <- (uncons -> (x, xs))+ where+ x :< xs = cons x xs++infix 5 :<++{-# COMPLETE (:<) :: Array #-}++-- | Add a new row at the end+--+-- >>> pretty $ snoc (array [2,2] [0,1,2,3]) (array [2] [4,5])+-- [[0,1],+-- [2,3],+-- [4,5]]+snoc :: Array a -> Array a -> Array a+snoc = append 0++-- | split an array into the initial rows and the last row.+--+-- >>> unsnoc (array [3,2] [0..5])+-- (UnsafeArray [2,2] [0,1,2,3],UnsafeArray [2] [4,5])+unsnoc :: Array a -> (Array a, Array a)+unsnoc a = (inits [0] a', lasts [0] a')+ where+ a' = asSingleton a++-- | Convenience pattern for row extraction and consolidation at the end of an Array.+--+-- >>> (xs:>x) = array [4] [0..3]+-- >>> x+-- UnsafeArray [] [3]+-- >>> xs+-- UnsafeArray [3] [0,1,2]+-- >>> (xs:>x)+-- UnsafeArray [4] [0,1,2,3]+pattern (:>) :: Array a -> Array a -> Array a+pattern xs :> x <- (unsnoc -> (xs, x))+ where+ xs :> x = snoc xs x++infix 5 :>++{-# COMPLETE (:>) :: Array #-}++-- * Math++-- | Generate an array of uniform random variates between a range.+--+-- >>> import System.Random.Stateful hiding (uniform)+-- >>> g <- newIOGenM (mkStdGen 42)+-- >>> u <- uniform g [2,3,4] (0,9 :: Int)+-- >>> pretty u+-- [[[0,7,0,2],+-- [1,7,4,2],+-- [5,9,8,2]],+-- [[9,8,1,0],+-- [2,2,8,2],+-- [2,8,0,6]]]+uniform :: (StatefulGen g m, UniformRange a) => g -> [Int] -> (a, a) -> m (Array a)+uniform g ds r = do+ v <- V.replicateM (S.size ds) (uniformRM r g)+ pure $ UnsafeArray ds v++-- | Inverse of a square matrix.+--+-- >>> e = array [3,3] [4,12,-16,12,37,-43,-16,-43,98] :: Array Double+-- >>> pretty (inverse e)+-- [[49.36111111111111,-13.555555555555554,2.1111111111111107],+-- [-13.555555555555554,3.7777777777777772,-0.5555555555555555],+-- [2.1111111111111107,-0.5555555555555555,0.1111111111111111]]+--+-- > mult (inverse a) a == a+inverse :: (Floating a) => Array a -> Array a+inverse a = mult (invtri (transpose (chol a))) (invtri (chol a))++-- | [Inversion of a Triangular Matrix](https://math.stackexchange.com/questions/1003801/inverse-of-an-invertible-upper-triangular-matrix-of-order-3)+--+-- >>> t = array [3,3] ([1,0,1,0,1,2,0,0,1] :: [Double]) :: Array Double+-- >>> pretty (invtri t)+-- [[1.0,0.0,-1.0],+-- [0.0,1.0,-2.0],+-- [0.0,0.0,1.0]]+-- >>> ident (shape t) == mult t (invtri t)+-- True+invtri :: (Fractional a) => Array a -> Array a+invtri a = i+ where+ ti = undiag (fmap recip (diag a))+ tl = zipWith (-) a (undiag (diag a))+ l = fmap negate (dot sum (*) ti tl)+ pow xs x = foldr ($) (ident (shape xs)) (replicate x (mult xs))+ zero' = konst (shape a) 0+ add = zipWith (+)+ sum' = foldl' add zero'+ i = mult (sum' (fmap (pow l) (range [n]))) ti+ n = S.getDim 0 (shape a)++-- | Cholesky decomposition using the <https://en.wikipedia.org/wiki/Cholesky_decomposition#The_Cholesky_algorithm Cholesky-Crout> algorithm.+--+-- >>> e = array [3,3] [4,12,-16,12,37,-43,-16,-43,98] :: Array Double+-- >>> pretty (chol e)+-- [[2.0,0.0,0.0],+-- [6.0,1.0,0.0],+-- [-8.0,5.0,3.0]]+-- >>> mult (chol e) (transpose (chol e)) == e+-- True+chol :: (Floating a) => Array a -> Array a+chol a =+ let l =+ tabulate+ (shape a)+ ( \[i, j] ->+ bool+ ( 1+ / index l [j, j]+ * ( index a [i, j]+ - sum+ ( (\k -> index l [i, k] * index l [j, k])+ <$> [0 .. (j - 1)]+ )+ )+ )+ ( sqrt+ ( index a [i, i]+ - sum+ ( (\k -> index l [j, k] ^ (2 :: Int))+ <$> [0 .. (j - 1)]+ )+ )+ )+ (i == j)+ )+ in l
+ src/Harpie/Fixed.hs view
@@ -0,0 +1,2831 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE FunctionalDependencies #-}+{-# LANGUAGE PatternSynonyms #-}+{-# LANGUAGE RebindableSyntax #-}+{-# LANGUAGE RoleAnnotations #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UndecidableInstances #-}+{-# LANGUAGE ViewPatterns #-}+{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}+{-# OPTIONS_GHC -Wno-redundant-constraints #-}++-- | Arrays with shape information and computations at a type-level.+module Harpie.Fixed+ ( -- * Usage+ -- $usage++ -- * Fixed Arrays+ Array (..),+ unsafeArray,+ validate,+ safeArray,+ array,+ unsafeModifyShape,+ unsafeModifyVector,++ -- * Dimensions+ Dim,+ pattern Dim,+ Dims,+ pattern Dims,++ -- * Conversion+ FromVector (..),+ toDynamic,+ with,+ SomeArray (..),+ someArray,++ -- * 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,++ -- * 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,+ invtri,+ inverse,+ chol,+ )+where++import Data.Bool+import Data.Distributive (Distributive (..))+import Data.Foldable hiding (find, length, minimum)+import Data.Functor.Classes+import Data.Functor.Rep+import Data.List qualified as List+import Data.Maybe+import Data.Vector qualified as V+import Fcf hiding (type (&&), type (+), type (++), type (-))+import Fcf qualified+import Fcf.Data.List+import GHC.Generics+import GHC.TypeNats+import Harpie.Array 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 Prettyprinter hiding (dot, fill)+import System.Random hiding (uniform)+import System.Random.Stateful hiding (uniform)+import Test.QuickCheck hiding (tabulate, vector)+import Test.QuickCheck.Instances.Natural ()+import Unsafe.Coerce+import Prelude as P hiding (cycle, drop, length, repeat, sequence, take, zipWith)+import Prelude qualified++-- $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 as F+-- >>> import Harpie.Shape qualified as S+-- >>> import Harpie.Shape (SNats (..), Fin (..), Fins (..))+-- >>> import GHC.TypeNats+-- >>> import Data.List qualified as List+-- >>> import Prettyprinter hiding (dot,fill)+-- >>> import Data.Functor.Rep+-- >>> s = 1 :: Array '[] Int+-- >>> s+-- [1]+-- >>> shape s+-- []+-- >>> pretty s+-- 1+-- >>> let v = range @'[3]+-- >>> pretty v+-- [0,1,2]+-- >>> let m = range @[2,3]+-- >>> pretty m+-- [[0,1,2],+-- [3,4,5]]+-- >>> a = range @[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]]]++-- $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 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+--+-- 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 function constraints.+--+-- >>> import Fcf qualified+--+-- Examples of arrays:+--+-- An array with no dimensions (a scalar).+--+-- >>> s = 1 :: Array '[] Int+-- >>> s+-- [1]+-- >>> shape s+-- []+-- >>> pretty s+-- 1+--+-- A single-dimension array (a vector).+--+-- >>> let v = range @'[3]+-- >>> pretty v+-- [0,1,2]+--+-- A two-dimensional array (a matrix).+--+-- >>> let m = range @[2,3]+-- >>> pretty m+-- [[0,1,2],+-- [3,4,5]]+--+-- An n-dimensional array (n should be finite).+--+-- >>> a = range @[2,3,4]+-- >>> a+-- [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 @[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 '[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 @[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 '[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 @[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 [2,3] (Fins [2,3])+-- tabulate id :: Array [2,3] (Fins [2,3])+-- :: Array [2, 3] (Fins [2, 3])+-- >>> pretty (tabulate id :: Array [2,3] (Fins [2,3]))+-- [[[0,0],[0,1],[0,2]],+-- [[1,0],[1,1],[1,2]]]+type role Array nominal representational++newtype Array (s :: [Nat]) a where+ Array :: V.Vector a -> Array s a+ deriving stock (Functor, Foldable, Generic, Traversable)+ deriving newtype (Eq, Eq1, Ord, Ord1, Show, Show1)++instance (Num a, KnownNats s) => Num (Array s a) where+ (+) = zipWith (+)+ (-) = zipWith (-)+ (*) = error "multiplication not defined"+ abs = fmap abs+ signum = fmap signum+ fromInteger x = konst @s (fromInteger x)++instance (KnownNats s, Show a) => Pretty (Array s a) where+ pretty = pretty . toDynamic++instance+ (KnownNats s) =>+ Data.Distributive.Distributive (Array s)+ where+ distribute :: (KnownNats s, Functor f) => f (Array s a) -> Array s (f a)+ distribute = distributeRep+ {-# INLINE distribute #-}++instance+ forall s.+ (KnownNats s) =>+ Representable (Array s)+ where+ type Rep (Array s) = Fins s++ tabulate f =+ Array . V.generate (S.size s) $ (f . UnsafeFins . shapen s)+ where+ s = valuesOf @s+ {-# INLINE tabulate #-}++ index (Array v) i = V.unsafeIndex v (flatten s (fromFins i))+ where+ s = 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 @[2,3])+-- [0,1,2,3,4,5]+--+-- >>> import Data.Vector qualified as V+-- >>> vectorAs (V.fromList [0..5]) :: Array [2,3] Int+-- [0,1,2,3,4,5]+class FromVector t a | t -> a where+ asVector :: t -> V.Vector a+ vectorAs :: V.Vector a -> t++instance FromVector (V.Vector a) a where+ asVector = id+ vectorAs = id++instance FromVector [a] a where+ asVector = V.fromList+ vectorAs = V.toList++instance FromVector (Array s a) a where+ asVector (Array v) = v+ vectorAs v = Array v++-- | Construct an array without shape validation.+--+-- >>> unsafeArray [0..4] :: Array [2,3] Int+-- [0,1,2,3,4]+unsafeArray :: (KnownNats s, FromVector t a) => t -> Array s a+unsafeArray (asVector -> v) = Array v++-- | Validate the size and shape of an array.+--+-- >>> validate (unsafeArray [0..4] :: Array [2,3] Int)+-- False+validate :: (KnownNats s) => Array s a -> Bool+validate a = size a == V.length (asVector a)++-- | Construct an Array, checking shape.+--+-- >>> (safeArray [0..23] :: Maybe (Array [2,3,4] Int)) == Just a+-- True+safeArray :: (KnownNats s, FromVector t a) => t -> Maybe (Array s a)+safeArray v =+ bool Nothing (Just a) (validate a)+ where+ a = unsafeArray v++-- | Construct an Array, checking shape.+--+-- >>> array [0..22] :: Array [2,3,4] Int+-- *** Exception: Shape Mismatch+-- ...+array :: forall s a t. (KnownNats s, FromVector t a) => t -> Array s a+array v =+ fromMaybe (error "Shape Mismatch") (safeArray v)++-- | Unsafely modify an array shape.+--+-- >>> pretty (unsafeModifyShape @[3,2] (array @[2,3] @Int [0..5]))+-- [[0,1],+-- [2,3],+-- [4,5]]+unsafeModifyShape :: forall s' s a. (KnownNats s, KnownNats s') => Array s a -> Array s' a+unsafeModifyShape a = unsafeArray (asVector a)++-- | Unsafely modify an array vector.+--+-- >>> import Data.Vector qualified as V+-- >>> pretty (unsafeModifyVector (V.map (+1)) (array [0..5] :: Array [2,3] Int))+-- [[1,2,3],+-- [4,5,6]]+unsafeModifyVector :: (KnownNats s) => (FromVector u a) => (FromVector v b) => (u -> v) -> Array s a -> Array s b+unsafeModifyVector f a = unsafeArray (asVector (f (vectorAs (asVector a))))++-- | Representation of an index into a shape (a type-level [Nat]). The index is a dimension of the shape.+type Dim = SNat++-- | Pattern synonym for a 'Dim'+pattern Dim :: () => (KnownNat n) => SNat n+pattern 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 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) => Array s a -> A.Array a+toDynamic a = A.array (shape a) (asVector a)++-- | Use a dynamic array in a fixed context.+--+-- >>> import qualified Harpie.Array as A+-- >>> with (A.range [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 a r.+ A.Array a ->+ (forall s. (KnownNats s) => Array s a -> r) ->+ r+with d f =+ withSomeSNats (fromIntegral <$> A.shape d) $ \(SNats :: SNats s) -> withKnownNats (SNats @s) (f (array @s (A.asVector d)))++-- | Sigma type for an 'Array'+--+-- A fixed Array where shape was unknown at runtime.+--+-- The library design encourages the use of value-level shape arrays (in @Harpie.Array@) via 'toDynamic' in preference to dependent-type styles of coding. In particular, no attempt has been made to prove to the compiler that a particular Shape (resulting from any of the supplied functions) exists. Life is short.+--+-- > P.take 4 <$> sample' arbitrary :: IO [SomeArray Int]+-- [SomeArray SNats @'[] [0],SomeArray SNats @'[0] [],SomeArray SNats @[1, 1] [1],SomeArray SNats @[5, 1, 4] [2,1,0,2,-6,0,5,6,-1,-4,0,5,-1,6,4,-6,1,0,3,-1]]+data SomeArray a = forall s. SomeArray (SNats s) (Array s a)++deriving instance (Show a) => Show (SomeArray a)++instance Functor SomeArray where+ fmap f (SomeArray sn a) = SomeArray sn (fmap f a)++instance Foldable SomeArray where+ foldMap f (SomeArray _ a) = foldMap f a++-- | Contruct a SomeArray+someArray :: forall s t a. (FromVector t a) => SNats s -> t -> SomeArray a+someArray s t = SomeArray s (Array (asVector t))++instance (Arbitrary a) => Arbitrary (SomeArray a) where+ arbitrary = do+ s <- arbitrary :: Gen [Small Nat]+ let s' = Prelude.take 3 (getSmall <$> s)+ v <- V.replicateM (product (Prelude.fromIntegral <$> s')) arbitrary+ withSomeSNats s' $ \sn -> pure (someArray sn v)++-- | Get shape of an Array as a value.+--+-- >>> shape a+-- [2,3,4]+shape :: forall a s. (KnownNats s) => Array s a -> [Int]+shape _ = valuesOf @s+{-# INLINE shape #-}++-- | Get rank of an Array as a value.+--+-- >>> rank a+-- 3+rank :: forall a s. (KnownNats s) => Array s a -> Int+rank = S.rank . shape+{-# INLINE rank #-}++-- | Get size of an Array as a value.+--+-- >>> size a+-- 24+size :: forall a s. (KnownNats s) => Array s a -> Int+size = S.size . 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 0)+-- 1+length :: (KnownNats s) => Array s a -> Int+length a = case shape a of+ [] -> 1+ (x : _) -> x++-- | Is the Array empty (has zero number of elements).+--+-- >>> isNull (array [] :: Array [2,0] ())+-- True+-- >>> isNull (array [4] :: Array '[] Int)+-- False+isNull :: (KnownNats s) => Array s a -> Bool+isNull = (0 ==) . size++-- | Extract an element at an index, unsafely.+--+-- >>> unsafeIndex a [1,2,3]+-- 23+unsafeIndex :: (KnownNats s) => Array s a -> [Int] -> a+unsafeIndex a xs = index a (UnsafeFins xs)++-- | Extract an element at an index, unsafely.+--+-- >>> a ! [1,2,3]+-- 23+(!) :: (KnownNats s) => Array s a -> [Int] -> a+(!) a xs = index a (UnsafeFins xs)++infixl 9 !++-- | Extract an element at an index, safely.+--+-- >>> a !? [1,2,3]+-- Just 23+-- >>> a !? [2,3,1]+-- Nothing+(!?) :: (KnownNats s) => Array s a -> [Int] -> Maybe a+(!?) a xs = index a <$> safeFins xs++infixl 9 !?++-- | Tabulate unsafely.+--+-- >>> :t tabulate @(Array [2,3]) id+-- tabulate @(Array [2,3]) id :: Array [2, 3] (Fins [2, 3])+-- >>> :t unsafeTabulate @[2,3] id+-- unsafeTabulate @[2,3] id :: Array [2, 3] [Int]+-- >>> pretty $ unsafeTabulate @[2,3] id+-- [[[0,0],[0,1],[0,2]],+-- [[1,0],[1,1],[1,2]]]+unsafeTabulate :: (KnownNats s) => ([Int] -> a) -> Array s a+unsafeTabulate f = tabulate (f . 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 @[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 s' s a. (KnownNats s, KnownNats s') => (Fins s' -> Fins s) -> Array s a -> Array s' a+backpermute f a = tabulate (index a . 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 @[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 s' s a. (KnownNats s, KnownNats s') => ([Int] -> [Int]) -> Array s a -> Array s' a+unsafeBackpermute f a = tabulate (index a . UnsafeFins . f . 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 @'[] @Int [3]+-- >>> :t fromScalar s+-- fromScalar s :: Int+fromScalar :: Array '[] a -> a+fromScalar a = index a (UnsafeFins [])++-- | Wrap a scalar.+--+-- >>> :t toScalar @Int 2+-- toScalar @Int 2 :: Array '[] Int+toScalar :: a -> Array '[] a+toScalar a = Array (V.singleton a)++-- | Is an array a scalar?+--+-- >>> isScalar (toScalar (2::Int))+-- True+isScalar :: (KnownNats s) => Array s a -> Bool+isScalar a = rank a == 0++-- | Convert a scalar to being a dimensioned array. Do nothing if not a scalar.+--+-- >>> asSingleton (toScalar 4)+-- [4]+asSingleton :: (KnownNats s, KnownNats s', s' ~ Eval (AsSingleton s)) => Array s a -> Array s' a+asSingleton = unsafeModifyShape++-- | Convert an array with shape [1] to being a scalar (Do nothing if not a shape [1] array).+--+-- >>> pretty (asScalar (singleton 3))+-- 3+asScalar :: (KnownNats s, KnownNats s', s' ~ Eval (AsScalar s)) => Array s a -> Array s' a+asScalar = unsafeModifyShape++-- | An array with no elements.+--+-- >>> toDynamic empty+-- UnsafeArray [0] []+empty :: Array '[0] a+empty = array []++-- | An enumeration of row-major or [lexicographic](https://en.wikipedia.org/wiki/Lexicographic_order) order.+--+-- >>> pretty (range :: Array [2,3] Int)+-- [[0,1,2],+-- [3,4,5]]+range :: forall s. (KnownNats s) => Array s Int+range = tabulate (S.flatten (valuesOf @s) . fromFins)++-- | An enumeration of col-major or [colexicographic](https://en.wikipedia.org/wiki/Lexicographic_order) order.+--+-- >>> pretty (corange @[2,3,4])+-- [[[0,6,12,18],+-- [2,8,14,20],+-- [4,10,16,22]],+-- [[1,7,13,19],+-- [3,9,15,21],+-- [5,11,17,23]]]+corange :: forall s. (KnownNats s) => Array s Int+corange = tabulate (S.flatten (List.reverse (valuesOf @s)) . List.reverse . fromFins)++-- | Indices of an array shape.+--+-- >>> pretty $ indices @[3,3]+-- [[[0,0],[0,1],[0,2]],+-- [[1,0],[1,1],[1,2]],+-- [[2,0],[2,1],[2,2]]]+indices :: (KnownNats s) => Array s [Int]+indices = tabulate fromFins++-- | The identity array.+--+-- >>> pretty $ ident @[3,3]+-- [[1,0,0],+-- [0,1,0],+-- [0,0,1]]+ident :: (KnownNats s, Num a) => Array s a+ident = tabulate (bool 0 1 . S.isDiag . fromFins)++-- | Create an array composed of a single value.+--+-- >>> pretty $ konst @[3,2] 1+-- [[1,1],+-- [1,1],+-- [1,1]]+konst :: (KnownNats s) => a -> Array s a+konst a = tabulate (const a)++-- | Create an array of shape [1].+--+-- >>> pretty $ singleton 1+-- [1]+singleton :: a -> Array '[1] a+singleton a = unsafeArray (V.singleton a)++-- | Extract the diagonal of an array.+--+-- >>> pretty $ diag (ident @[3,3])+-- [1,1,1]+diag ::+ forall s' a s.+ ( KnownNats s,+ KnownNats s',+ s' ~ Eval (MinDim s)+ ) =>+ Array s a ->+ Array s' a+diag a = unsafeBackpermute (replicate (rank a) . getDim 0) a++-- | Expand an array to form a diagonal array+--+-- >>> pretty $ undiag (range @'[3])+-- [[0,0,0],+-- [0,1,0],+-- [0,0,2]]+undiag ::+ forall s' a s.+ ( KnownNats s,+ KnownNats s',+ s' ~ Eval ((++) s s),+ Num a+ ) =>+ Array s a ->+ Array s' a+undiag a = tabulate (\xs -> bool 0 (index a (UnsafeFins $ pure $ getDim 0 (fromFins xs))) (isDiag (fromFins xs)))++-- | Zip two arrays at an element level.+--+-- >>> zipWith (-) v v+-- [0,0,0]+zipWith :: (KnownNats s) => (a -> b -> c) -> Array s a -> Array s b -> Array s c+zipWith f (asVector -> a) (asVector -> b) = unsafeArray (V.zipWith f a b)++-- | Modify a single value at an index.+--+-- >>> pretty $ modify (S.UnsafeFins [0,0]) (const 100) (range @[3,2])+-- [[100,1],+-- [2,3],+-- [4,5]]+modify :: (KnownNats s) => Fins s -> (a -> a) -> Array s a -> Array s a+modify ds f a = tabulate (\s -> bool id f (s == ds) (index a 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) =>+ ([Int] -> a -> b) ->+ Array s a ->+ Array s b+imap f a = zipWith f indices 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 a ds s s' xs proxy.+ ( KnownNats s,+ KnownNats ds,+ ds ~ Eval (DimsOf xs)+ ) =>+ (Dims ds -> proxy xs -> Array s a -> Array s' a) ->+ proxy xs ->+ Array s a ->+ Array s' a+rowWise f xs a = f (Dims @ds) xs 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 a ds s s' xs proxy.+ ( KnownNats s,+ KnownNats ds,+ ds ~ Eval (EndDimsOf xs s)+ ) =>+ (Dims ds -> proxy xs -> Array s a -> Array s' a) ->+ proxy xs ->+ Array s a ->+ Array s' a+colWise f xs a = f (Dims @ds) xs 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 d t s s' a.+ ( KnownNats s,+ KnownNats s',+ s' ~ Eval (TakeDim d t s)+ ) =>+ Dim d ->+ SNat t ->+ Array s a ->+ Array s' a+take _ _ a = unsafeBackpermute id 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 s s' a d t.+ ( KnownNats s,+ KnownNats s',+ s' ~ Eval (TakeDim d t s)+ ) =>+ Dim d ->+ SNat t ->+ Array s a ->+ Array s' a+takeB Dim SNat a = unsafeBackpermute (modifyDim (valueOf @d) (\x -> x + getDim (valueOf @d) (shape a) - (valueOf @t))) 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 s s' a d t.+ ( KnownNats s,+ KnownNats s',+ Eval (DropDim d t s) ~ s'+ ) =>+ Dim d ->+ SNat t ->+ Array s a ->+ Array s' a+drop Dim SNat a = unsafeBackpermute (S.modifyDim (valueOf @d) (\x -> x + valueOf @t)) 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 s s' a d t.+ ( KnownNats s,+ KnownNats s',+ Eval (DropDim d t s) ~ s'+ ) =>+ Dim d ->+ SNat t ->+ Array s a ->+ Array s' a+dropB _ _ a = unsafeBackpermute id a++-- | 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 d a p s s'.+ ( KnownNats s,+ KnownNats s',+ s' ~ Eval (DeleteDim d s),+ p ~ Eval (GetDim d s)+ ) =>+ Dim d ->+ Fin p ->+ Array s a ->+ Array s' a+select Dim p a = unsafeBackpermute (S.insertDim (valueOf @d) (fromFin p)) a++-- | Insert along a dimension at a position.+--+-- >>> pretty $ insert (Dim @2) (UnsafeFin 0) a (konst @[2,3] 0)+-- [[[0,0,1,2,3],+-- [0,4,5,6,7],+-- [0,8,9,10,11]],+-- [[0,12,13,14,15],+-- [0,16,17,18,19],+-- [0,20,21,22,23]]]+-- >>> toDynamic $ insert (Dim @0) (UnsafeFin 0) (toScalar 1) (toScalar 2)+-- UnsafeArray [2] [2,1]+insert ::+ forall 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)+ ) =>+ Dim d ->+ Fin p ->+ Array s a ->+ Array si a ->+ Array s' a+insert Dim i a b = tabulate go+ where+ go s+ | getDim d s' == fromFin i = index b (UnsafeFins (deleteDim d s'))+ | getDim d s' < fromFin i = index a (UnsafeFins s')+ | otherwise = index a (UnsafeFins (decAt d s'))+ where+ s' = fromFins s+ d = 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 d s s' p a.+ ( KnownNats s,+ KnownNats s',+ s' ~ Eval (DecAt d s),+ p ~ 1 + Eval (GetDim d s)+ ) =>+ Dim d ->+ Fin p ->+ Array s a ->+ Array s' a+delete Dim p a = unsafeBackpermute (\s -> bool (incAt d s) s (getDim d s < fromFin p)) a+ where+ d = valueOf @d++-- | Insert along a dimension at the end.+--+-- >>> pretty $ append (Dim @2) a (konst @[2,3] 0)+-- [[[0,1,2,3,0],+-- [4,5,6,7,0],+-- [8,9,10,11,0]],+-- [[12,13,14,15,0],+-- [16,17,18,19,0],+-- [20,21,22,23,0]]]+append ::+ forall a d s si s'.+ ( KnownNats s,+ KnownNats si,+ KnownNats s',+ s' ~ Eval (IncAt d s),+ True ~ Eval (InsertOk d s si)+ ) =>+ Dim d ->+ Array s a ->+ Array si a ->+ Array s' a+append (Dim :: Dim d) = insert (Dim @d) (UnsafeFin (getDim (valueOf @d) (valuesOf @s)))++-- | Insert along a dimension at the beginning.+--+-- >>> pretty $ prepend (Dim @2) (konst @[2,3] 0) a+-- [[[0,0,1,2,3],+-- [0,4,5,6,7],+-- [0,8,9,10,11]],+-- [[0,12,13,14,15],+-- [0,16,17,18,19],+-- [0,20,21,22,23]]]+prepend ::+ forall a d s si s'.+ ( KnownNats s,+ KnownNats si,+ KnownNats s',+ s' ~ Eval (IncAt d s),+ True ~ Eval (InsertOk d s si)+ ) =>+ Dim d ->+ Array si a ->+ Array s a ->+ Array s' a+prepend d a b = insert d (UnsafeFin 0) b a++-- | Concatenate along a dimension.+--+-- >>> shape $ concatenate (Dim @1) a a+-- [2,6,4]+-- >>> toDynamic $ concatenate (Dim @0) (toScalar 1) (toScalar 2)+-- UnsafeArray [2] [1,2]+-- >>> toDynamic $ concatenate (Dim @0) (array @'[1] [0]) (array @'[3] [1..3])+-- UnsafeArray [4] [0,1,2,3]+concatenate ::+ forall a s0 s1 d s.+ ( KnownNats s0,+ KnownNats s1,+ KnownNats s,+ Eval (Concatenate d s0 s1) ~ s+ ) =>+ Dim d ->+ Array s0 a ->+ Array s1 a ->+ Array s a+concatenate Dim a0 a1 = tabulate (go . fromFins)+ where+ go s =+ bool+ (index a0 (UnsafeFins s))+ ( index+ a1+ ( UnsafeFins $+ insertDim+ d'+ (getDim d' s - getDim d' ds0)+ (deleteDim d' s)+ )+ )+ (getDim d' s >= getDim d' ds0)+ ds0 = shape a0+ d' = valueOf @d++-- | Combine two arrays as a new dimension of a new array.+--+-- >>> pretty $ couple (Dim @0) (array @'[3] [1,2,3]) (array @'[3] @Int [4,5,6])+-- [[1,2,3],+-- [4,5,6]]+-- >>> couple (Dim @0) (toScalar @Int 0) (toScalar 1)+-- [0,1]+couple ::+ forall 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)+ ) =>+ Dim d ->+ Array s a ->+ Array s a ->+ Array s' a+couple d a a' = concatenate d (elongate d a) (elongate d 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 a d off l s s'.+ ( KnownNats s,+ KnownNats s',+ s' ~ Eval (SetDim d l s),+ Eval (SliceOk d off l s) ~ True+ ) =>+ Dim d ->+ SNat off ->+ SNat l ->+ Array s a ->+ Array s' a+slice Dim SNat _ a = unsafeBackpermute (S.modifyDim (valueOf @d) (+ (valueOf @off))) 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 d s a.+ (KnownNats s) =>+ Dim d ->+ Int ->+ Array s a ->+ Array s a+rotate Dim r a = unsafeBackpermute (rotateIndex (valueOf @d) r (shape 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 ds xs s' s a.+ ( KnownNats s,+ KnownNats s',+ s' ~ Eval (SetDims ds xs s)+ ) =>+ Dims ds ->+ SNats xs ->+ Array s a ->+ Array s' a+takes _ _ a = unsafeBackpermute id a++-- | Across the specified dimesnions, 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 s' s a ds xs.+ ( KnownNats s,+ KnownNats s',+ KnownNats ds,+ KnownNats xs,+ s' ~ Eval (SetDims ds xs s)+ ) =>+ Dims ds ->+ SNats xs ->+ Array s a ->+ Array s' a+takeBs _ _ a = unsafeBackpermute (List.zipWith (+) start) a+ where+ start = List.zipWith (-) (shape a) (S.setDims (valuesOf @ds) (valuesOf @xs) (shape 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 ds xs s' s a.+ ( KnownNats s,+ KnownNats s',+ KnownNats ds,+ KnownNats xs,+ s' ~ Eval (DropDims ds xs s)+ ) =>+ Dims ds ->+ SNats xs ->+ Array s a ->+ Array s' a+drops _ _ a = unsafeBackpermute (List.zipWith (+) start) a+ where+ start = List.zipWith (-) (valuesOf @s) (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 s' s ds xs a.+ ( KnownNats s,+ KnownNats s',+ KnownNats ds,+ KnownNats xs,+ s' ~ Eval (DropDims ds xs s)+ ) =>+ Dims ds ->+ SNats xs ->+ Array s a ->+ Array s' a+dropBs _ _ a = unsafeBackpermute id a++-- | Select by dimensions and indexes.+--+-- >>> pretty $ indexes (Dims @[0,1]) (S.UnsafeFins [1,1]) a+-- [16,17,18,19]+indexes ::+ forall ds s s' xs a.+ ( KnownNats s,+ KnownNats s',+ s' ~ Eval (DeleteDims ds s),+ xs ~ Eval (GetDims ds s)+ ) =>+ Dims ds ->+ Fins xs ->+ Array s a ->+ Array s' a+indexes Dims xs a = unsafeBackpermute (S.insertDims (valuesOf @ds) (fromFins xs)) 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 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)+ ) =>+ Dims ds ->+ SNats xs ->+ Array s a ->+ Array s' a+indexesT ds _ a = indexes ds (UnsafeFins $ valuesOf @xs) 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 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'+ ) =>+ Dims ds ->+ SNats offs ->+ SNats ls ->+ Array s a ->+ Array s' a+slices _ _ _ a = unsafeBackpermute (List.zipWith (+) o) a+ where+ o = S.setDims (valuesOf @ds) (valuesOf @offs) (replicate (rank a) 0)++-- | Select the first element along the supplied dimensions.+--+-- >>> pretty $ heads (Dims @[0,2]) a+-- [0,4,8]+heads ::+ forall a ds s s'.+ ( KnownNats s,+ KnownNats s',+ KnownNats ds,+ s' ~ Eval (DeleteDims ds s)+ ) =>+ Dims ds ->+ Array s a ->+ Array s' a+heads ds a = indexes ds (UnsafeFins $ replicate (rankOf @ds) 0) a++-- | Select the last element along the supplied dimensions.+--+-- >>> pretty $ lasts (Dims @[0,2]) a+-- [15,19,23]+lasts ::+ forall ds s s' a.+ ( KnownNats s,+ KnownNats ds,+ KnownNats s',+ s' ~ Eval (DeleteDims ds s)+ ) =>+ Dims ds ->+ Array s a ->+ Array s' a+lasts ds a = indexes ds (UnsafeFins lastds) a+ where+ lastds = (\i -> getDim i (shape a) - 1) <$> (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 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)+ ) =>+ Dims ds ->+ Array s a ->+ Array s' a+tails ds a = slices ds (SNats @os) (SNats @ls) 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 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)+ ) =>+ Dims ds ->+ Array s a ->+ Array s' a+inits ds a = slices ds (SNats @os) (SNats @ls) a++-- | Extracts specified dimensions to an outer layer.+--+-- >>> :t extracts (Dims @'[0]) (range @[2,3,4])+-- extracts (Dims @'[0]) (range @[2,3,4])+-- :: Array '[2] (Array [3, 4] Int)+extracts ::+ forall ds st si so a.+ ( KnownNats st,+ KnownNats ds,+ KnownNats si,+ KnownNats so,+ si ~ Eval (DeleteDims ds st),+ so ~ Eval (GetDims ds st)+ ) =>+ Dims ds ->+ Array st a ->+ Array so (Array si a)+extracts ds a = tabulate (\s -> indexes ds s 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 ds st si so a b.+ ( KnownNats st,+ KnownNats ds,+ KnownNats si,+ KnownNats so,+ si ~ Eval (DeleteDims ds st),+ so ~ Eval (GetDims ds st)+ ) =>+ Dims ds ->+ (Array si a -> b) ->+ Array st a ->+ Array so b+reduces ds f a = fmap f (extracts ds 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 a ds si so st.+ ( KnownNats ds,+ KnownNats st,+ KnownNats si,+ KnownNats so,+ Eval (InsertDims ds so si) ~ st+ ) =>+ Dims ds ->+ Array so (Array si a) ->+ Array st a+joins _ a = tabulate go+ where+ go s = index (index a (UnsafeFins $ S.getDims (valuesOf @ds) (fromFins s))) (UnsafeFins $ S.deleteDims (valuesOf @ds) (fromFins s))++-- | Join inner and outer dimension layers in outer dimension order.+--+-- >>> a == join (extracts (Dims @[0,1]) a)+-- True+join ::+ forall a si so st ds.+ ( KnownNats st,+ KnownNats si,+ KnownNats so,+ KnownNats ds,+ ds ~ Eval (DimsOf so),+ st ~ Eval (InsertDims ds so si)+ ) =>+ Array so (Array si a) ->+ Array st a+join a = joins (SNats @ds) a++-- | Traverse along specified dimensions.+--+-- >>> traverses (Dims @'[1]) print (range @[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)+ ) =>+ Dims ds ->+ (a -> f b) ->+ Array s a ->+ f (Array s b)+traverses (Dims :: Dims ds) f a = joins (SNats @ds) <$> traverse (traverse f) (extracts (Dims :: Dims ds) 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 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)+ ) =>+ Dims ds ->+ (Array si a -> Array si' b) ->+ Array s a ->+ Array s' b+maps SNats f a = joins (SNats @ds) (fmap f (extracts (SNats @ds) 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 ds si so a.+ ( KnownNats si,+ KnownNats so,+ si ~ Eval (DeleteDims ds so),+ KnownNats (Eval (GetDims ds so))+ ) =>+ Dims ds ->+ (Array si a -> Bool) ->+ Array so a ->+ A.Array (Array si a)+filters Dims p a = A.asArray $ V.filter p $ asVector (extracts (Dims @ds) a)++-- | 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 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)+ ) =>+ Dims ds ->+ (Array si a -> Array si b -> Array si' c) ->+ Array s a ->+ Array s b ->+ Array s' c+zips SNats f a b = joins (Dims @ds) (zipWith f (extracts (Dims @ds) a) (extracts (Dims @ds) 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 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)+ ) =>+ (Array si a -> Array si a) ->+ Dims ds ->+ Fins so ->+ Array s a ->+ Array s a+modifies f SNats ps a = joins (Dims @ds) $ modify ps f (extracts (Dims @ds) 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 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)+ ) =>+ Dims ds ->+ SNats ls ->+ (Array si a -> Array si a -> Array si' b) ->+ Array st a ->+ Array st' b+diffs SNats xs f a = zips (Dims @ds) f (drops (Dims @ds) xs a) (dropBs (Dims @ds) xs 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 '[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]+-- >>> 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 sc sa sb a b c.+ ( KnownNats sa,+ KnownNats sb,+ KnownNats sc,+ sc ~ Eval ((++) sa sb)+ ) =>+ (a -> b -> c) ->+ Array sa a ->+ Array sb b ->+ Array sc c+expand f a b = tabulate (\i -> f (index a (UnsafeFins $ List.take r (fromFins i))) (index b (UnsafeFins $ List.drop r (fromFins i))))+ where+ r = rank 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 sc sa sb a b c.+ ( KnownNats sa,+ KnownNats sb,+ KnownNats sc,+ sc ~ Eval ((++) sa sb)+ ) =>+ (a -> b -> c) ->+ Array sa a ->+ Array sb b ->+ Array sc c+coexpand f a b = tabulate (\i -> f (index a (UnsafeFins $ List.drop r (fromFins i))) (index b (UnsafeFins $ List.take r (fromFins i))))+ where+ r = rank a++-- | 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 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)+ ) =>+ Dims ds ->+ (Array ss a -> b) ->+ Array s a ->+ Array s' b+contract SNats f a = f . diag <$> extracts (Dims @ds') a++-- | Expand two arrays and then contract the result using the supplied matching dimensions.+--+-- >>> pretty $ prod (Dims @'[1]) (Dims @'[0]) sum (*) (range @[2,3]) (range @[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 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)+ ) =>+ Dims ds0 ->+ Dims ds1 ->+ (Array si c -> d) ->+ (a -> b -> c) ->+ Array s0 a ->+ Array s1 b ->+ Array st d+prod SNats SNats g f a b = unsafeTabulate (\so -> g $ unsafeTabulate (\si -> f (unsafeIndex a (S.insertDims (valuesOf @ds0) si (List.take sp so))) (unsafeIndex b (S.insertDims (valuesOf @ds1) si (List.drop sp so)))))+ where+ sp = rank a - 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 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]+ ) =>+ (Array si c -> d) ->+ (a -> b -> c) ->+ Array s0 a ->+ Array s1 b ->+ Array st d+dot f g a b = prod (Dims @ds0) (Dims @ds1) f g a 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 a ds0 ds1 s0 s1 so0 so1 st si.+ ( Num 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]+ ) =>+ Array s0 a ->+ Array s1 a ->+ Array st a+mult = dot sum (*)++-- | @windows xs@ are xs-sized windows of an array+--+-- >>> shape $ windows (Dims @[2,2]) (range @[4,3,2])+-- [3,2,2,2,2]+windows ::+ forall w s ws a.+ ( KnownNats s,+ KnownNats ws,+ ws ~ Eval (ExpandWindows w s)+ ) =>+ SNats w -> Array s a -> Array ws a+windows SNats a = unsafeBackpermute (S.indexWindows (rankOf @w)) a++-- | Find the starting positions of occurences of one array in another.+--+-- >>> a = cycle @[4,4] (range @'[3])+-- >>> i = array @[2,2] [1,2,2,0]+-- >>> pretty $ find i a+-- [[False,True,False],+-- [True,False,False],+-- [False,False,True]]+find ::+ forall 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)+ ) =>+ Array si a -> Array s a -> Array s' Bool+find i a = xs+ where+ i' = rerank (SNat @r) i+ ws = windows (SNats @i') a+ xs = fmap (== i') (extracts (SNats @re) ws)++-- | Find the ending positions of one array in another except where the array overlaps with another copy.+--+-- >>> a = konst @[5,5] @Int 1+-- >>> i = konst @[2,2] @Int 1+-- >>> pretty $ findNoOverlap i a+-- [[True,False,True,False],+-- [False,False,False,False],+-- [True,False,True,False],+-- [False,False,False,False]]+findNoOverlap ::+ forall 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)+ ) =>+ Array si a -> Array s a -> Array s' Bool+findNoOverlap i a = r+ where+ f = find i a++ cl :: [Int] -> [[Int]]+ cl sh = List.filter (P.not . any (> 0) . List.init) $ List.filter (P.not . all (>= 0)) $ A.arrayAs $ A.tabulate ((\x -> 2 * x - 1) <$> sh) (\s -> List.zipWith (\x x0 -> x - x0 + 1) s sh)+ go r' s = index f (UnsafeFins s) && not (any (index r' . UnsafeFins) (List.filter (\x -> isFins x (shape f)) $ fmap (List.zipWith (+) s) (cl (shape i))))+ r = unsafeTabulate (go r)++-- | Check if the first array is a prefix of the second.+--+-- >>> isPrefixOf (array @[2,2] [0,1,4,5]) a+-- True+isPrefixOf ::+ forall 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')+ ) =>+ Array s' a -> Array s a -> Bool+isPrefixOf p a = p == cut a++-- | Check if the first array is a suffix of the second.+--+-- >>> isSuffixOf (array @[2,2] [18,19,22,23]) a+-- True+isSuffixOf ::+ forall 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)+ ) =>+ Array s' a -> Array s a -> Bool+isSuffixOf p a = p == cutSuffix a++-- | Check if the first array is an infix of the second.+--+-- >>> isInfixOf (array @[2,2] [18,19,22,23]) a+-- True+isInfixOf ::+ forall 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)+ ) =>+ Array si a -> Array s a -> Bool+isInfixOf p a = or $ find p 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 @'[3] 0 (array @'[0] [])+-- [0,0,0]+-- >>> pretty $ fill @'[3] 0 (array @'[4] [1..4])+-- [1,2,3]+fill ::+ forall s' a s.+ ( KnownNats s,+ KnownNats s'+ ) =>+ a -> Array s a -> Array s' a+fill x (Array v) = Array (V.take (S.size (valuesOf @s')) (v <> V.replicate (S.size (valuesOf @s') - V.length v) 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 @'[2] (array @'[4] @Int [0..3])+-- UnsafeArray [2] [0,1]+cut ::+ forall s' s r a.+ ( KnownNats s,+ KnownNats s',+ KnownNat r,+ KnownNats (Eval (Rerank r s)),+ True ~ Eval (IsSubset s' s),+ r ~ Eval (Rank s')+ ) =>+ Array s a ->+ Array s' a+cut a = unsafeBackpermute id (rerank (SNat @r) 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 @[2,2] a+-- UnsafeArray [2,2] [18,19,22,23]+cutSuffix ::+ forall s' s a r.+ ( KnownNats s,+ KnownNats s',+ KnownNat r,+ KnownNats (Eval (Rerank r s)),+ r ~ Eval (Rank s'),+ True ~ Eval (IsSubset s' s)+ ) =>+ Array s a ->+ Array s' a+cutSuffix a = unsafeBackpermute (List.zipWith (+) diffDim) a'+ where+ a' = rerank (SNat @r) a+ diffDim = List.zipWith (-) (shape a') (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 @'[5] 0 (array @'[4] @Int [0..3])+-- UnsafeArray [5] [0,1,2,3,0]+pad ::+ forall s' a s r.+ ( KnownNats s,+ KnownNats s',+ KnownNat r,+ KnownNats (Eval (Rerank r s)),+ r ~ Eval (Rank s')+ ) =>+ a ->+ Array s a ->+ Array s' a+pad d a = tabulate (\s -> bool d (index a' (unsafeCoerce s)) (fromFins s `S.isFins` shape a'))+ where+ a' = rerank (SNat @r) 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 @'[5] 0 (array @'[4] [0..3])+-- UnsafeArray [5] [0,0,1,2,3]+-- >>> pretty $ lpad @[3,3] 0 (range @[2,2])+-- [[0,0,0],+-- [0,0,1],+-- [0,2,3]]+lpad ::+ forall s' a s r.+ ( KnownNats s,+ KnownNats s',+ KnownNat r,+ KnownNats (Eval (Rerank r s)),+ r ~ Eval (Rank s')+ ) =>+ a ->+ Array s a ->+ Array s' a+lpad d a = tabulate (\s -> bool d (index a' (UnsafeFins $ olds s)) (olds s `S.isFins` shape a'))+ where+ a' = rerank (SNat @r) a+ gap = List.zipWith (-) (valuesOf @s') (shape a')+ olds s = List.zipWith (-) (fromFins s) gap++-- | Reshape an array (with the same number of elements).+--+-- >>> pretty $ reshape @[4,3,2] a+-- [[[0,1],+-- [2,3],+-- [4,5]],+-- [[6,7],+-- [8,9],+-- [10,11]],+-- [[12,13],+-- [14,15],+-- [16,17]],+-- [[18,19],+-- [20,21],+-- [22,23]]]+reshape ::+ forall s' s a.+ ( Eval (Size s) ~ Eval (Size s'),+ KnownNats s,+ KnownNats s'+ ) =>+ Array s a ->+ Array s' a+reshape = unsafeBackpermute (shapen s . flatten s')+ where+ s = valuesOf @s+ s' = valuesOf @s'++-- | Make an Array single dimensional.+--+-- >>> pretty $ flat (range @[2,2])+-- [0,1,2,3]+-- >>> pretty (flat $ toScalar 0)+-- [0]+flat ::+ forall s' s a.+ ( KnownNats s,+ KnownNats s',+ s' ~ '[Eval (Size s)]+ ) =>+ Array s a ->+ Array s' a+flat a = unsafeModifyShape a++-- | Reshape an array, repeating the original array. The shape of the array should be a suffix of the new shape.+--+-- >>> pretty $ repeat @[2,2,2] (array @'[2] [1,2])+-- [[[1,2],+-- [1,2]],+-- [[1,2],+-- [1,2]]]+--+-- > repeat ds (toScalar x) == konst ds x+repeat ::+ forall s' s a.+ ( KnownNats s,+ KnownNats s',+ Eval (IsPrefixOf s s') ~ True+ ) =>+ Array s a ->+ Array s' a+repeat a = unsafeBackpermute (List.drop (S.rank (valuesOf @s') - rank a)) a++-- | Reshape an array, cycling through the elements without regard to the original shape.+--+-- >>> pretty $ cycle @[2,2,2] (array @'[3] [1,2,3])+-- [[[1,2],+-- [3,1]],+-- [[2,3],+-- [1,2]]]+cycle ::+ forall s' s a.+ ( KnownNats s,+ KnownNats s'+ ) =>+ Array s a ->+ Array s' a+cycle a = unsafeBackpermute (S.shapen (shape a) . (`mod` size a) . S.flatten (valuesOf @s')) 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 r s s' a.+ ( KnownNats s,+ KnownNats s',+ s' ~ Eval (Rerank r s)+ ) =>+ SNat r -> Array s a -> Array s' a+rerank _ a = unsafeModifyShape 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 ds s s' a.+ ( KnownNats s,+ KnownNats s',+ s' ~ Eval (Reorder s ds)+ ) =>+ SNats ds ->+ Array s a ->+ Array s' a+reorder SNats a = unsafeBackpermute (\s -> S.insertDims (valuesOf @ds) s []) a++-- | Remove single dimensions.+--+-- >>> let sq = array [1..24] :: Array '[2,1,3,4,1] Int+-- >>> shape $ squeeze sq+-- [2,3,4]+--+-- >>> shape $ squeeze (singleton 0)+-- []+squeeze ::+ forall s t a.+ ( KnownNats s,+ KnownNats t,+ t ~ Eval (Squeeze s)+ ) =>+ Array s a ->+ Array t a+squeeze = unsafeModifyShape++-- | Insert a single dimension at the supplied position.+--+-- >>> shape $ elongate (SNat @1) a+-- [2,1,3,4]+-- >>> toDynamic $ elongate (SNat @0) (toScalar 1)+-- UnsafeArray [1] [1]+elongate ::+ ( KnownNats s,+ KnownNats s',+ s' ~ Eval (InsertDim d 1 s)+ ) =>+ Dim d ->+ Array s a ->+ Array s' a+elongate _ a = unsafeModifyShape 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 @[2,2,2] [1..8])+-- [[[1,5],+-- [3,7]],+-- [[2,6],+-- [4,8]]]+transpose ::+ forall a s s'. (KnownNats s, KnownNats s', s' ~ Eval (Reverse s)) => Array s a -> Array s' a+transpose a = unsafeBackpermute List.reverse a++-- | Inflate (or replicate) an array by inserting a new dimension given a supplied dimension and size.+--+-- >>> pretty $ inflate (SNat @0) (SNat @2) (array @'[3] [0,1,2])+-- [[0,1,2],+-- [0,1,2]]+inflate ::+ forall s' s d x a.+ ( KnownNats s,+ KnownNats s',+ s' ~ Eval (InsertDim d x s)+ ) =>+ Dim d ->+ SNat x ->+ Array s a ->+ Array s' a+inflate SNat _ a = unsafeBackpermute (S.deleteDim (valueOf @d)) a++-- | Intercalate an array along dimensions.+--+-- >>> pretty $ intercalate (SNat @2) (konst @[2,3] 0) a+-- [[[0,0,1,0,2,0,3],+-- [4,0,5,0,6,0,7],+-- [8,0,9,0,10,0,11]],+-- [[12,0,13,0,14,0,15],+-- [16,0,17,0,18,0,19],+-- [20,0,21,0,22,0,23]]]+intercalate ::+ forall 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)+ ) =>+ Dim d -> Array si a -> Array s a -> Array st a+intercalate SNat i a =+ joins+ (Dims @ds)+ ( vector @n'+ ( List.intersperse+ i+ (toList (extracts (Dims @ds) 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 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)+ ) =>+ Dim d -> a -> Array s a -> Array st a+intersperse (SNat :: SNat d) x a = intercalate (SNat @d) (konst @si x) 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 s s' newd ds a.+ ( KnownNats s,+ KnownNats s',+ s' ~ Eval (ConcatDims ds newd s)+ ) =>+ Dims ds ->+ SNat newd ->+ Array s a ->+ Array s' a+concats SNats SNat a = unsafeBackpermute (unconcatDimsIndex ds n (shape a)) a+ where+ n = valueOf @newd+ ds = 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 ds s a.+ (KnownNats s) =>+ Dims ds ->+ Array s a ->+ Array s a+reverses SNats a = unsafeBackpermute (reverseIndex (valuesOf @ds) (shape 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 a ds s.+ ( KnownNats s,+ True ~ Eval (IsDims ds s)+ ) =>+ Dims ds ->+ [Int] ->+ Array s a ->+ Array s a+rotates SNats rs a = unsafeBackpermute (rotatesIndex (valuesOf @ds) rs (valuesOf @s)) a++-- | Sort an array along the supplied dimensions.+--+-- >>> pretty $ sorts (Dims @'[0]) (array @[2,2] [2,3,1,4])+-- [[1,4],+-- [2,3]]+-- >>> pretty $ sorts (Dims @'[1]) (array @[2,2] [2,3,1,4])+-- [[2,3],+-- [1,4]]+-- >>> pretty $ sorts (Dims @[0,1]) (array @[2,2] [2,3,1,4])+-- [[1,2],+-- [3,4]]+sorts ::+ forall 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)+ ) =>+ Dims ds -> Array s a -> Array s a+sorts SNats a = joins (Dims @ds) $ unsafeModifyVector sortV (extracts (Dims @ds) 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 @[2,2] [2,3,1,4])+-- UnsafeArray [2,2] [2,3,1,4]+sortsBy ::+ forall 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)+ ) =>+ Dims ds -> (Array si a -> Array si b) -> Array s a -> Array s a+sortsBy SNats c a = joins (Dims @ds) $ unsafeModifyVector (sortByV c) (extracts (Dims @ds) a)++-- | The indices into the array if it were sorted along the dimensions supplied.+--+-- >>> orders (Dims @'[0]) (array @[2,2] [2,3,1,4])+-- [1,0]+orders ::+ forall 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)+ ) =>+ Dims ds -> Array s a -> Array so Int+orders SNats a = unsafeModifyVector orderV (extracts (Dims @ds) 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 @[2,2] [2,3,1,4])+-- [0,1]+ordersBy ::+ forall 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)+ ) =>+ Dims ds -> (Array si a -> Array si b) -> Array s a -> Array so Int+ordersBy SNats c a = unsafeModifyVector (orderByV c) (extracts (Dims @ds) a)++-- | Apply a binary array function to two arrays with matching shapes across the supplied (matching) dimensions.+--+-- >>> a = array @[2,3] [0..5]+-- >>> b = array @'[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 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)+ ) =>+ SNats ma -> SNats mb -> (Array sia a -> Array sib b -> Array sic c) -> Array sa a -> Array sb b -> Array sc c+telecasts SNats SNats f a b = join (zipWith f (extracts (SNats @ma) a) (extracts (SNats @mb) 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 @[2,3] [0..5]+-- >>> pretty $ transmit (zipWith (+)) (toScalar 1) a+-- [[1,2,3],+-- [4,5,6]]+transmit ::+ forall 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)+ ) =>+ (Array sa a -> Array sib b -> Array sic c) -> Array sa a -> Array sb b -> Array sc c+transmit f a b = maps (Dims @ds) (f a) b++-- | A one-dimensional array.+type Vector s a = Array '[s] a++-- | Create a one-dimensional array.+--+-- >>> pretty $ vector @3 @Int [2,3,4]+-- [2,3,4]+vector ::+ forall n a t.+ ( FromVector t a,+ KnownNat n+ ) =>+ t ->+ Array '[n] a+vector xs = array xs++-- | vector with an explicit SNat rather than a KnownNat constraint.+--+-- >>> pretty $ vector' @Int (SNat @3) [2,3,4]+-- [2,3,4]+vector' ::+ forall a n t.+ (FromVector t a) =>+ SNat n ->+ t ->+ Array '[n] a+vector' n xs = withKnownNat n (vector xs)++-- | Vector specialisation of 'range'+--+-- >>> toDynamic $ iota @5+-- UnsafeArray [5] [0,1,2,3,4]+iota :: forall n. (KnownNat n) => Vector n Int+iota = range++-- | A two-dimensional array.+type Matrix m n a = Array '[m, n] a++-- * row (first dimension) specializations++-- | Add a new row+--+-- >>> pretty $ cons (array @'[2] [0,1]) (array @[2,2] [2,3,4,5])+-- [[0,1],+-- [2,3],+-- [4,5]]+cons ::+ forall 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)+ ) =>+ Array sh a -> Array st a -> Array s a+cons =+ prepend (SNat @0)++-- | Add a new row at the end+--+-- >>> pretty $ snoc (array @[2,2] [0,1,2,3]) (array @'[2] [4,5])+-- [[0,1],+-- [2,3],+-- [4,5]]+snoc ::+ forall 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)+ ) =>+ Array si a -> Array sl a -> Array s a+snoc = append (SNat @0)++-- | split an array into the first row and the remaining rows.+--+-- >>> import Data.Bifunctor (bimap)+-- >>> bimap toDynamic toDynamic $ uncons (array @[3,2] [0..5])+-- (UnsafeArray [2] [0,1],UnsafeArray [2,2] [2,3,4,5])+uncons ::+ forall 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)+ ) =>+ Array s a -> (Array sh a, Array st a)+uncons a = (heads (Dims @ds) a, tails (Dims @ds) a)++-- | split an array into the initial rows and the last row.+--+-- >>> import Data.Bifunctor (bimap)+-- >>> bimap toDynamic toDynamic $ unsnoc (array @[3,2] [0..5])+-- (UnsafeArray [2,2] [0,1,2,3],UnsafeArray [2] [4,5])+unsnoc ::+ forall 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)+ ) =>+ Array s a -> (Array si a, Array sl a)+unsnoc a = (inits (Dims @ds) a, lasts (Dims @ds) a)++-- | Convenience pattern for row extraction and consolidation at the beginning of an Array.+--+-- >>> (x:<xs) = array @'[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 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)+ ) =>+ Array sh a -> Array st a -> Array s a+pattern x :< xs <- (uncons -> (x, xs))+ where+ x :< xs = cons x xs++infix 5 :<++{-# COMPLETE (:<) :: Array #-}++-- | Convenience pattern for row extraction and consolidation at the end of an Array.+--+-- >>> (xs:>x) = array @'[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 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)+ ) =>+ Array si a -> Array sl a -> Array s a+pattern xs :> x <- (unsnoc -> (xs, x))+ where+ xs :> x = snoc xs 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 @[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 s a g m.+ ( StatefulGen g m,+ UniformRange a,+ KnownNats s+ ) =>+ g -> (a, a) -> m (Array s a)+uniform g r = do+ v <- V.replicateM (S.size (valuesOf @s)) (uniformRM r g)+ pure $ array v++-- | Inverse of a square matrix.+--+-- > A.mult (D.inverse a) a == a+--+-- >>> e = array @[3,3] @Double [4,12,-16,12,37,-43,-16,-43,98]+-- >>> pretty (inverse e)+-- [[49.36111111111111,-13.555555555555554,2.1111111111111107],+-- [-13.555555555555554,3.7777777777777772,-0.5555555555555555],+-- [2.1111111111111107,-0.5555555555555555,0.1111111111111111]]+inverse :: (Eq a, Floating a, KnownNat m) => Matrix m m a -> Matrix m m a+inverse a = mult (invtri (transpose (chol a))) (invtri (chol a))++-- | [Inversion of a Triangular Matrix](https://math.stackexchange.com/questions/1003801/inverse-of-an-invertible-upper-triangular-matrix-of-order-3)+--+-- >>> t = array @[3,3] @Double [1,0,1,0,1,2,0,0,1]+-- >>> pretty (invtri t)+-- [[1.0,0.0,-1.0],+-- [0.0,1.0,-2.0],+-- [0.0,0.0,1.0]]+--+-- >>> ident == mult t (invtri t)+-- True+invtri :: forall a n. (KnownNat n, Floating a, Eq a) => Matrix n n a -> Matrix n n a+invtri a = i+ where+ ti = undiag (fmap recip (diag a))+ tl = zipWith (-) a (undiag (diag a))+ l = fmap negate (dot sum (*) ti tl)+ pow xs x = foldr ($) (ident @[n, n]) (replicate x (mult xs))+ zero' = konst @[n, n] 0+ add = zipWith (+)+ sum' = foldl' add zero'+ i = mult (sum' (fmap (pow l) (range @'[n]))) ti++-- | Cholesky decomposition using the <https://en.wikipedia.org/wiki/Cholesky_decomposition#The_Cholesky_algorithm Cholesky-Crout> algorithm.+--+-- >>> e = array @[3,3] @Double [4,12,-16,12,37,-43,-16,-43,98]+-- >>> pretty (chol e)+-- [[2.0,0.0,0.0],+-- [6.0,1.0,0.0],+-- [-8.0,5.0,3.0]]+-- >>> mult (chol e) (transpose (chol e)) == e+-- True+chol :: (KnownNat m, Floating a) => Matrix m m a -> Matrix m m a+chol a =+ let l =+ unsafeTabulate+ ( \[i, j] ->+ bool+ ( 1+ / unsafeIndex l [j, j]+ * ( unsafeIndex a [i, j]+ - sum+ ( (\k -> unsafeIndex l [i, k] * unsafeIndex l [j, k])+ <$> ([0 .. (j - 1)] :: [Int])+ )+ )+ )+ ( sqrt+ ( unsafeIndex a [i, i]+ - sum+ ( (\k -> unsafeIndex l [j, k] ^ (2 :: Int))+ <$> ([0 .. (j - 1)] :: [Int])+ )+ )+ )+ (i == j)+ )+ in l
+ src/Harpie/Shape.hs view
@@ -0,0 +1,1643 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE PatternSynonyms #-}+{-# LANGUAGE RoleAnnotations #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UndecidableInstances #-}+{-# LANGUAGE ViewPatterns #-}++-- | Functions for manipulating shape. The module tends to supply equivalent functionality at type-level and value-level with functions of the same name (except for capitalization).+module Harpie.Shape+ ( -- * Type-level Nat+ SNat,+ pattern SNat,+ valueOf,++ -- * Type-level [Nat]+ SNats,+ pattern SNats,+ fromSNats,+ KnownNats (..),+ natVals,+ withKnownNats,+ SomeNats,+ someNatVals,+ withSomeSNats,++ -- * Shape+ valuesOf,+ rankOf,+ sizeOf,+ Fin (..),+ fin,+ safeFin,+ Fins (..),+ fins,+ safeFins,++ -- * Shape Operators at value- and type- level.+ rank,+ Rank,+ range,+ Range,+ rerank,+ Rerank,+ dimsOf,+ DimsOf,+ endDimsOf,+ EndDimsOf,+ size,+ Size,+ flatten,+ shapen,+ asSingleton,+ AsSingleton,+ asScalar,+ AsScalar,+ isSubset,+ IsSubset,+ exceptDims,+ ExceptDims,+ reorder,+ Reorder,+ ReorderOk,+ squeeze,+ Squeeze,++ -- * Primitives+ Min,+ Max,+ minimum,+ Minimum,++ -- * Position+ isFin,+ IsFin,+ isFins,+ IsFins,+ isDim,+ IsDim,+ isDims,+ IsDims,+ lastPos,+ LastPos,+ minDim,+ MinDim,++ -- * combinators+ EnumFromTo,+ Foldl',++ -- * single dimension+ GetIndex,+ SetIndex,+ getDim,+ GetDim,+ modifyDim,+ ModifyDim,+ incAt,+ IncAt,+ decAt,+ DecAt,+ setDim,+ SetDim,+ takeDim,+ TakeDim,+ dropDim,+ DropDim,+ deleteDim,+ DeleteDim,+ insertDim,+ InsertDim,+ InsertOk,+ SliceOk,+ SlicesOk,+ concatenate,+ Concatenate,+ ConcatenateOk,++ -- * multiple dimension+ getDims,+ GetDims,+ getLastPositions,+ GetLastPositions,+ modifyDims,+ insertDims,+ InsertDims,+ preDeletePositions,+ PreDeletePositions,+ preInsertPositions,+ PreInsertPositions,+ setDims,+ SetDims,+ deleteDims,+ DeleteDims,+ dropDims,+ DropDims,+ concatDims,+ ConcatDims,++ -- * value-only operations+ unconcatDimsIndex,+ reverseIndex,+ rotate,+ rotateIndex,+ rotatesIndex,+ isDiag,++ -- * windowed+ expandWindows,+ ExpandWindows,+ indexWindows,+ dimWindows,+ DimWindows,++ -- * Fcf re-exports+ Eval,+ type (++),+ )+where++import Data.Bool+import Data.Foldable hiding (minimum)+import Data.Function+import Data.List qualified as List+import Data.Maybe+import Data.Proxy+import Data.Type.Bool hiding (Not)+import Data.Type.Equality+import Data.Type.Ord hiding (Max, Min)+import Fcf hiding (type (&&), type (+), type (++), type (-), type (<), type (>), type (||))+import Fcf qualified+import Fcf.Class.Foldable+import Fcf.Data.List+import GHC.Exts+import GHC.TypeLits (ErrorMessage (..))+import GHC.TypeLits qualified as L+import GHC.TypeNats+import Prelude as P hiding (minimum)++-- $setup+-- >>> :set -XDataKinds+-- >>> :set -XTypeFamilies+-- >>> import Prelude+-- >>> import Fcf+-- >>> import GHC.Exts ()+-- >>> import Harpie.Shape as S++-- | Get the value of a type level Nat.+-- Use with explicit type application+--+-- >>> valueOf @42+-- 42+valueOf :: forall n. (KnownNat n) => Int+valueOf = fromIntegral $ fromSNat (SNat @n)+{-# INLINE valueOf #-}++type role SNats nominal++-- | A value-level witness for a type-level list of natural numbers.+--+-- Obtain an SNats value using:+--+-- - The natsSing method of KnownNats+-- - The SNats pattern+-- - The withSomeSNats function+--+-- >>> :t SNats @[2,3,4]+-- SNats @[2,3,4] :: KnownNats [2, 3, 4] => SNats [2, 3, 4]+-- >>> SNats @[2,3,4]+-- SNats @[2, 3, 4]+newtype SNats (ns :: [Nat]) = UnsafeSNats [Nat]++instance Eq (SNats ns) where+ _ == _ = True++instance Ord (SNats ns) where+ compare _ _ = EQ++-- | Matches GHC printing quirks.+instance Show (SNats ns) where+ show (UnsafeSNats s) = "SNats @" <> bool "" "'" (length s < 2) <> "[" <> mconcat (List.intersperse ", " (show <$> s)) <> "]"++-- | A explicitly bidirectional pattern synonym relating an 'SNats' to a 'KnownNats' constraint.+--+-- As an expression: Constructs an explicit 'SNats' ns value from an implicit 'KnownNats' ns constraint:+--+-- > SNat @n :: KnownNat n => SNat n+--+-- As a pattern: Matches on an explicit SNats n value bringing an implicit KnownNats n constraint into scope:+--+-- > f :: SNats ns -> ..+-- > f SNat = {- KnownNats ns in scope -}+--+-- or, if you need to both bring the KnownNats into scope and reuse the SNats input:+--+-- > f (SNats :: SNats s) = g (SNats @s)+pattern SNats :: forall ns. () => (KnownNats ns) => SNats ns+pattern SNats <- (knownNatsInstance -> KnownNatsInstance)+ where+ SNats = natsSing++{-# COMPLETE SNats #-}++-- | Return the value-level list of naturals in an SNats ns value.+--+-- >>> fromSNats (SNats @[2,3,4])+-- [2,3,4]+fromSNats :: SNats s -> [Nat]+fromSNats (UnsafeSNats s) = s++-- An internal data type that is only used for defining the SNat pattern+-- synonym.+data KnownNatsInstance (ns :: [Nat]) where+ KnownNatsInstance :: (KnownNats ns) => KnownNatsInstance ns++-- An internal function that is only used for defining the SNat pattern+-- synonym.+knownNatsInstance :: SNats ns -> KnownNatsInstance ns+knownNatsInstance dims = withKnownNats dims KnownNatsInstance++-- | Reflect a list of naturals.+--+-- >>> natsSing @'[2]+-- SNats @'[2]+class KnownNats (ns :: [Nat]) where+ natsSing :: SNats ns++instance KnownNats '[] where+ natsSing = UnsafeSNats []++instance (KnownNat n, KnownNats s) => KnownNats (n ': s) where+ natsSing = UnsafeSNats (fromSNat (SNat :: SNat n) : fromSNats (SNats :: SNats s))++-- | Obtain a value-level list of naturals from a type-level proxy+--+-- >>> natVals (SNats @[2,3,4])+-- [2,3,4]+natVals :: forall ns proxy. (KnownNats ns) => proxy ns -> [Nat]+natVals _ = case natsSing :: SNats ns of+ UnsafeSNats xs -> xs++-- | Convert an explicit SNats ns value into an implicit KnownNats ns constraint.+withKnownNats ::+ forall ns rep (r :: TYPE rep).+ SNats ns -> ((KnownNats ns) => r) -> r+withKnownNats = withDict @(KnownNats ns)++-- | Convert a list of naturals into an SNats ns value, where ns is a fresh type-level list of naturals.+withSomeSNats ::+ forall rep (r :: TYPE rep).+ [Nat] -> (forall s. SNats s -> r) -> r+withSomeSNats s k = k (UnsafeSNats s)+{-# NOINLINE withSomeSNats #-}++-- | An unknown type-level list of naturals.+data SomeNats = forall s. (KnownNats s) => SomeNats (Proxy s)++-- | Promote a list of naturals to unknown type-level+someNatVals :: [Nat] -> SomeNats+someNatVals s =+ withSomeSNats+ s+ ( \(sn :: SNats s) ->+ withKnownNats sn (SomeNats @s Proxy)+ )++-- * shape primitives++-- | The value of a 'KnownNats'.+--+-- >>> valuesOf @[2,3,4]+-- [2,3,4]+valuesOf :: forall s. (KnownNats s) => [Int]+valuesOf = fmap fromIntegral (fromSNats (SNats :: SNats s))+{-# INLINE valuesOf #-}++-- | The rank (or length) of a KnownNats.+--+-- >>> rankOf @[2,3,4]+-- 3+rankOf :: forall s. (KnownNats s) => Int+rankOf = length (valuesOf @s)+{-# INLINE rankOf #-}++-- | The size (or product) of a KnownNats.+--+-- >>> sizeOf @[2,3,4]+-- 24+sizeOf :: forall s. (KnownNats s) => Int+sizeOf = product (valuesOf @s)+{-# INLINE sizeOf #-}++-- | Fin most often represents a (finite) zero-based index for a single dimension (of a multi-dimensioned hyper-rectangular array).+type role Fin nominal++newtype Fin s+ = UnsafeFin+ { fromFin :: Int+ }+ deriving stock (Eq, Ord)++instance Show (Fin n) where+ show (UnsafeFin x) = show x++-- | Construct a Fin.+--+-- Errors on out-of-bounds+--+-- >>> fin @2 1+-- 1+--+-- >>> fin @2 2+-- *** Exception: value outside bounds+-- ...+fin :: forall n. (KnownNat n) => Int -> Fin n+fin x = fromMaybe (error "value outside bounds") (safeFin x)++-- | Construct a Fin safely.+--+-- >>> safeFin 1 :: Maybe (Fin 2)+-- Just 1+--+-- >>> safeFin 2 :: Maybe (Fin 2)+-- Nothing+safeFin :: forall n. (KnownNat n) => Int -> Maybe (Fin n)+safeFin x = bool Nothing (Just (UnsafeFin x)) (x >= 0 && x < valueOf @n)++-- | Fins most often represents (finite) indexes for multiple dimensions (of a multi-dimensioned hyper-rectangular array).+type role Fins nominal++newtype Fins s+ = UnsafeFins+ { fromFins :: [Int]+ }+ deriving stock (Eq, Ord, Functor)++instance Show (Fins n) where+ show (UnsafeFins x) = show x++-- | Construct a Fins.+--+-- Errors on out-of-bounds+--+-- >>> fins @[2,3,4] [1,2,3]+-- [1,2,3]+--+-- >>> fins @[2,3,4] [1,2,5]+-- *** Exception: value outside bounds+-- ...+fins :: forall s. (KnownNats s) => [Int] -> Fins s+fins x = fromMaybe (error "value outside bounds") (safeFins x)++-- | Construct a Fins safely.+--+-- >>> safeFins [1,2,3] :: Maybe (Fins [2,3,4])+-- Just [1,2,3]+--+-- >>> safeFins [2] :: Maybe (Fins '[2])+-- Nothing+safeFins :: forall s. (KnownNats s) => [Int] -> Maybe (Fins s)+safeFins xs = bool Nothing (Just (UnsafeFins xs)) (isFins xs (valuesOf @s))++-- | Number of dimensions+--+-- >>> rank @Int [2,3,4]+-- 3+rank :: [a] -> Int+rank = length+{-# INLINE rank #-}++-- | Number of dimensions+--+-- >>> :k! Eval (Rank [2,3,4])+-- ...+-- = 3+data Rank :: [a] -> Exp Natural++type instance+ Eval (Rank xs) =+ Eval (Length xs)++-- | Enumerate a range of rank n+--+-- >>> range 0+-- []+--+-- >>> range 3+-- [0,1,2]+range :: Int -> [Int]+range n = [0 .. (n - 1)]++-- | Enumerate a range of rank n+--+-- >>> :k! Eval (Range 0)+-- ...+-- = '[]+--+-- >>> :k! Eval (Range 3)+-- ...+-- = [0, 1, 2]+data Range :: Nat -> Exp [Nat]++type instance+ Eval (Range x) =+ If (x == 0) '[] (Eval (EnumFromTo 0 (Eval ((Fcf.-) x 1))))++-- | Create a new rank by adding ones to the left, if the new rank is greater, or combining dimensions (from left to right) into rows, if the new rank is lower.+--+-- >>> rerank 4 [2,3,4]+-- [1,2,3,4]+-- >>> rerank 2 [2,3,4]+-- [6,4]+rerank :: Int -> [Int] -> [Int]+rerank r xs =+ replicate (r - r') 1+ <> bool [] [product (take (r' - r + 1) xs)] (r <= r')+ <> drop (r' - r + 1) xs+ where+ r' = rank xs++-- | Create a new rank by adding ones to the left, if the new rank is greater, or combining dimensions (from left to right) into rows, if the new rank is lower.+--+-- >>> :k! Eval (Rerank 4 [2,3,4])+-- ...+-- = [1, 2, 3, 4]+-- >>> :k! Eval (Rerank 2 [2,3,4])+-- ...+-- = [6, 4]+data Rerank :: Nat -> [Nat] -> Exp [Nat]++type instance+ Eval (Rerank r xs) =+ If+ (Eval ((Fcf.>) r (Eval (Rank xs))))+ (Eval (Eval (Replicate (Eval ((Fcf.-) r (Eval (Rank xs)))) 1) ++ xs))+ ( Eval+ ( '[Eval (Size (Eval (Take (Eval ((Fcf.+) (Eval ((Fcf.-) (Eval (Rank xs)) r)) 1)) xs)))]+ ++ Eval (Drop (Eval ((Fcf.+) (Eval ((Fcf.-) (Eval (Rank xs)) r)) 1)) xs)+ )+ )++-- | Enumerate the dimensions of a shape.+--+-- dimsOf [2,3,4]+-- [0,1,2]+dimsOf :: [Int] -> [Int]+dimsOf s = range (rank s)++-- | Enumerate the dimensions of a shape.+--+-- >>> :k! Eval (DimsOf [2,3,4])+-- ...+-- = [0, 1, 2]+data DimsOf :: [Nat] -> Exp [Nat]++type instance+ Eval (DimsOf xs) =+ Eval (Range =<< Rank xs)++-- | Enumerate the final dimensions of a shape.+--+-- >>> endDimsOf [1,0] [2,3,4]+-- [2,1]+endDimsOf :: [Int] -> [Int] -> [Int]+endDimsOf xs s = take (rank xs) (List.reverse (dimsOf s))++-- | Enumerate the final dimensions of a shape.+--+-- >>> :k! Eval (EndDimsOf [1,0] [2,3,4])+-- ...+-- = [2, 1]+data EndDimsOf :: [Nat] -> [Nat] -> Exp [Nat]++type instance+ Eval (EndDimsOf xs s) =+ Eval (LiftM2 Take (Rank xs) (Reverse =<< DimsOf s))++-- | Total number of elements (if the list is the shape of a hyper-rectangular array).+--+-- >>> size [2,3,4]+-- 24+size :: [Int] -> Int+size [] = 1+size [x] = x+size xs = P.product xs+{-# INLINE size #-}++-- | Total number of elements (if the list is the shape of a hyper-rectangular array).+--+-- >>> :k! (Eval (Size [2,3,4]))+-- ...+-- = 24+data Size :: [Nat] -> Exp Nat++type instance Eval (Size xs) = Eval (Foldr (Fcf.*) 1 xs)++-- | Convert from a n-dimensional shape list index to a flat index, which, technically is the lexicographic position of the position in a row-major array.+--+-- >>> flatten [2,3,4] [1,1,1]+-- 17+--+-- >>> flatten [] [1,1,1]+-- 0+flatten :: [Int] -> [Int] -> Int+flatten [] _ = 0+flatten _ [x'] = x'+flatten ns xs = sum $ zipWith (*) xs (drop 1 $ scanr (*) 1 ns)+{-# INLINE flatten #-}++-- | Convert from a flat index to a shape index.+--+-- >>> shapen [2,3,4] 17+-- [1,1,1]+shapen :: [Int] -> Int -> [Int]+shapen [] _ = []+shapen [_] x' = [x']+shapen [_, y] x' = let (i, j) = divMod x' y in [i, j]+shapen ns x =+ fst $+ foldr+ ( \a (acc, r) ->+ let (d, m) = divMod r a+ in (m : acc, d)+ )+ ([], x)+ ns+{-# INLINE shapen #-}++-- | Convert a scalar to a dimensioned shape+--+-- >>> asSingleton []+-- [1]+-- >>> asSingleton [2,3,4]+-- [2,3,4]+asSingleton :: [Int] -> [Int]+asSingleton [] = [1]+asSingleton x = x++-- | Convert a scalar to a dimensioned shape+-- >>> :k! Eval (AsSingleton '[])+-- ...+-- = '[1]+-- >>> :k! Eval (AsSingleton [2,3,4])+-- ...+-- = [2, 3, 4]+data AsSingleton :: [Nat] -> Exp [Nat]++type instance+ Eval (AsSingleton xs) =+ If (xs == '[]) '[1] xs++-- | Convert a (potentially) [1] dimensioned shape to a scalar shape+--+-- >>> asScalar [1]+-- []+-- >>> asScalar [2,3,4]+-- [2,3,4]+asScalar :: [Int] -> [Int]+asScalar [1] = []+asScalar x = x++-- | Convert a (potentially) [1] dimensioned shape to a scalar shape+-- >>> :k! Eval (AsScalar '[1])+-- ...+-- = '[]+-- >>> :k! Eval (AsScalar [2,3,4])+-- ...+-- = [2, 3, 4]+data AsScalar :: [Nat] -> Exp [Nat]++type instance+ Eval (AsScalar xs) =+ If (xs == '[1]) '[] xs++lte :: [Int] -> [Int] -> Bool+lte xs ys =+ and (zipWith (<=) xs ys)+ && rank xs == rank ys++data LTE :: [Nat] -> [Nat] -> Exp Bool++type instance+ Eval (LTE xs ys) =+ Eval+ ( LiftM2+ (Fcf.&&)+ (And =<< ZipWith (Fcf.<=) xs ys)+ (LiftM2 TyEq (Rank xs) (Rank ys))+ )++-- | Check if a shape is a subset (<=) another shape after reranking.+--+-- >>> isSubset [2,3,4] [2,3,4]+-- True+--+-- >>> isSubset [1,2] [2,3,4]+-- True+--+-- >>> isSubset [2,1] [1]+-- False+isSubset :: [Int] -> [Int] -> Bool+isSubset xs ys = lte (rerank (rank ys) xs) ys++-- | Check if a shape is a subset (<=) another shape after reranking.+--+-- >>> :k! Eval (IsSubset [2,3,4] [2,3,4])+-- ...+-- = True+--+-- >>> :k! Eval (IsSubset [1,2] [2,3,4])+-- ...+-- = True+--+-- >>> :k! Eval (IsSubset [2,1] '[1])+-- ...+-- = False+data IsSubset :: [Nat] -> [Nat] -> Exp Bool++type instance+ Eval (IsSubset xs ys) =+ Eval (LTE (Eval (Rerank (Eval (Rank ys)) xs)) ys)++-- | Compute dimensions for a shape other than the supplied dimensions.+--+-- >>> exceptDims [1,2] [2,3,4]+-- [0]+exceptDims :: [Int] -> [Int] -> [Int]+exceptDims ds s = deleteDims ds [0 .. (rank s - 1)]++-- | Compute dimensions for a shape other than the supplied dimensions.+--+-- >>> :k! Eval (ExceptDims [1,2] [2,3,4])+-- ...+-- = '[0]+data ExceptDims :: [Nat] -> [Nat] -> Exp [Nat]++type instance+ Eval (ExceptDims ds s) =+ Eval (DeleteDims ds =<< EnumFromTo 0 (Eval ((Fcf.-) (Eval (Rank s)) 1)))++-- | Reorder the dimensions of shape according to a list of positions.+--+-- >>> reorder [2,3,4] [2,0,1]+-- [4,2,3]+reorder :: [Int] -> [Int] -> [Int]+reorder [] _ = []+reorder _ [] = []+reorder s (d : ds) = getDim d s : reorder s ds++-- | Reorder the dimensions of shape according to a list of positions.+--+-- >>> :k! Eval (Reorder [2,3,4] [2,0,1])+-- ...+-- = [4, 2, 3]+data Reorder :: [Nat] -> [Nat] -> Exp [Nat]++type instance+ Eval (Reorder ds xs) =+ If+ (Eval (ReorderOk ds xs))+ (Eval (Map (Flip GetDim ds) xs))+ (L.TypeError ('Text "Reorder dimension indices out of bounds"))++-- | Test if a Reorder is valid.+--+-- >>> :k! Eval (ReorderOk [2,3,4] [0,1])+-- ...+-- = False+data ReorderOk :: [Nat] -> [Nat] -> Exp Bool++type instance+ Eval (ReorderOk ds xs) =+ Eval (TyEq (Eval (Rank ds)) (Eval (Rank xs)))+ && Eval (And =<< Map (Flip IsFin (Eval (Rank ds))) xs)++-- | remove 1's from a list+--+-- >>> squeeze [0,1,2,3]+-- [0,2,3]+squeeze :: [Int] -> [Int]+squeeze = filter (/= 1)++-- | Remove 1's from a list.+--+-- >>> :k! (Eval (Squeeze [0,1,2,3]))+-- ...+-- = [0, 2, 3]+data Squeeze :: [a] -> Exp [a]++type instance+ Eval (Squeeze xs) =+ Eval (Filter (Not <=< TyEq 1) xs)++-- | minimum of a list+--+-- >>> S.minimum []+-- *** Exception: zero-ranked+-- ...+-- >>> S.minimum [2,3,4]+-- 2+minimum :: [Int] -> Int+minimum [] = error "zero-ranked"+minimum [x] = x+minimum (x : xs) = P.min x (minimum xs)++-- | minimum of a list+--+-- >>> :k! Eval (Minimum '[])+-- ...+-- = (TypeError ...)+--+-- >>> :k! Eval (Minimum [2,3,4])+-- ...+-- = 2+data Minimum :: [a] -> Exp a++type instance Eval (Minimum '[]) = L.TypeError (L.Text "zero ranked")++type instance+ Eval (Minimum (x ': xs)) =+ Eval (Foldr Min x xs)++-- | Minimum of two type values.+--+-- >>> :k! Eval (Min 0 1)+-- ...+-- = 0+data Min :: a -> a -> Exp a++type instance Eval (Min a b) = If (a <? b) a b++-- | Maximum of two type values.+--+-- >>> :k! Eval (Max 0 1)+-- ...+-- = 1+data Max :: a -> a -> Exp a++type instance Eval (Max a b) = If (a >? b) a b++-- | Check if i is a valid Fin (aka in-bounds index of a dimension)+--+-- >>> isFin 0 2+-- True+-- >>> isFin 2 2+-- False+isFin :: Int -> Int -> Bool+isFin i d = 0 <= i && i + 1 <= d++-- | Check if i is a valid Fin (aka in-bounds index of a dimension)+--+-- >>> :k! Eval (IsFin 0 2)+-- ...+-- = True+-- >>> :k! Eval (IsFin 2 2)+-- ...+-- = False+data IsFin :: Nat -> Nat -> Exp Bool++type instance+ Eval (IsFin x d) =+ x <? d++-- | Check if i is a valid Fins (aka in-bounds index of a Shape)+--+-- >>> isFins [0,1] [2,2]+-- True+-- >>> isFins [0,1] [2,1]+-- False+isFins :: [Int] -> [Int] -> Bool+isFins xs ds = length xs == length ds && and (zipWith isFin xs ds)++-- | Check if i is a valid Fins (aka in-bounds index of a Shape)+--+-- >>> :k! Eval (IsFins [0,1] [2,2])+-- ...+-- = True+-- >>> :k! Eval (IsFins [0,1] [2,1])+-- ...+-- = False+data IsFins :: [Nat] -> [Nat] -> Exp Bool++type instance+ Eval (IsFins xs ds) =+ Eval (And (Eval (ZipWith IsFin xs ds)))+ && Eval (LiftM2 TyEq (Rank xs) (Rank ds))++-- | Is a value a valid dimension of a shape.+--+-- >>> isDim 2 [2,3,4]+-- True+-- >>> isDim 0 []+-- True+isDim :: Int -> [Int] -> Bool+isDim d s = isFin d (rank s) || d == 0 && null s++-- | Is a value a valid dimension of a shape.+--+-- >>> :k! Eval (IsDim 2 [2,3,4])+-- ...+-- = True+-- >>> :k! Eval (IsDim 0 '[])+-- ...+-- = True+data IsDim :: Nat -> [Nat] -> Exp Bool++type instance+ Eval (IsDim d s) =+ Eval (IsFin d =<< Rank s)+ || (0 == d && s == '[])++-- | Are values valid dimensions of a shape.+--+-- >>> isDims [2,1] [2,3,4]+-- True+-- >>> isDims [0] []+-- True+isDims :: [Int] -> [Int] -> Bool+isDims ds s = all (`isDim` s) ds++-- | Are values valid dimensions of a shape.+--+-- >>> :k! Eval (IsDims [2,1] [2,3,4])+-- ...+-- = True+-- >>> :k! Eval (IsDims '[0] '[])+-- ...+-- = True+data IsDims :: [Nat] -> [Nat] -> Exp Bool++type instance+ Eval (IsDims ds s) =+ Eval (And =<< Map (Flip IsDim s) ds)++-- | Get the last position of a dimension of a shape.+--+-- >>> lastPos 2 [2,3,4]+-- 3+-- >>> lastPos 0 []+-- 0+lastPos :: Int -> [Int] -> Int+lastPos d s =+ bool (getDim d s - 1) 0 (0 == d && null s)++-- | Get the last position of a dimension of a shape.+--+-- >>> :k! Eval (LastPos 2 [2,3,4])+-- ...+-- = 3+-- >>> :k! Eval (LastPos 0 '[])+-- ...+-- = 0+data LastPos :: Nat -> [Nat] -> Exp Nat++type instance+ Eval (LastPos d s) =+ If+ (0 == d && s == '[])+ 0+ (Eval (GetDim d s) - 1)++-- | Get the minimum dimension as a singleton dimension.+--+-- >>> minDim [2,3,4]+-- [2]+-- >>> minDim []+-- []+minDim :: [Int] -> [Int]+minDim [] = []+minDim s = [minimum s]++-- | Get the minimum dimension as a singleton dimension.+--+-- >>> :k! Eval (MinDim [2,3,4])+-- ...+-- = '[2]+-- >>> :k! Eval (MinDim '[])+-- ...+-- = '[]+data MinDim :: [Nat] -> Exp [Nat]++type instance+ Eval (MinDim s) =+ If+ (s == '[])+ '[]+ '[Eval (Minimum s)]++-- | Enumerate between two Nats+--+-- >>> :k! Eval (EnumFromTo 0 3)+-- ...+-- = [0, 1, 2, 3]+data EnumFromTo :: Nat -> Nat -> Exp [Nat]++type instance Eval (EnumFromTo a b) = Eval (Unfoldr (EnumFromToHelper b) a)++data EnumFromToHelper :: Nat -> Nat -> Exp (Maybe (a, Nat))++type instance+ Eval (EnumFromToHelper b a) =+ If+ (a >? b)+ 'Nothing+ ('Just '(a, a + 1))++-- | Left fold.+--+-- >>> :k! Eval (Foldl' (Fcf.+) 0 [1,2,3])+-- ...+-- = 6+data Foldl' :: (b -> a -> Exp b) -> b -> t a -> Exp b++type instance Eval (Foldl' f y '[]) = y++type instance Eval (Foldl' f y (x ': xs)) = Eval (Foldl' f (Eval (f y x)) xs)++-- | Get an element at a given index.+--+-- >>> :kind! Eval (GetIndex 2 [2,3,4])+-- ...+-- = Just 4+data GetIndex :: Nat -> [a] -> Exp (Maybe a)++type instance Eval (GetIndex d xs) = GetIndexImpl d xs++type family GetIndexImpl (n :: Nat) (xs :: [k]) where+ GetIndexImpl _ '[] = 'Nothing+ GetIndexImpl 0 (x ': _) = 'Just x+ GetIndexImpl n (_ ': xs) = GetIndexImpl (n - 1) xs++-- | Get the dimension of a shape at the supplied index. Error if out-of-bounds.+--+-- >>> getDim 1 [2,3,4]+-- 3+-- >>> getDim 3 [2,3,4]+-- *** Exception: getDim outside bounds+-- ...+-- >>> getDim 0 []+-- 1+getDim :: Int -> [Int] -> Int+getDim 0 [] = 1+getDim i s = fromMaybe (error "getDim outside bounds") (s List.!? i)++-- | Get the dimension of a shape at the supplied index. Error if out-of-bounds or non-computable (usually unknown to the compiler).+--+-- >>> :k! Eval (GetDim 1 [2,3,4])+-- ...+-- = 3+-- >>> :k! Eval (GetDim 3 [2,3,4])+-- ...+-- = (TypeError ...)+-- >>> :k! Eval (GetDim 0 '[])+-- ...+-- = 1+data GetDim :: Nat -> [Nat] -> Exp Nat++type instance+ Eval (GetDim n xs) =+ If+ (Eval (And [Eval (TyEq n 0), Eval (TyEq xs ('[] :: [Nat]))]))+ 1+ (Eval (FromMaybe (L.TypeError (L.Text "GetDim out of bounds or non-computable: " :<>: ShowType n :<>: L.Text " " :<>: ShowType xs)) (Eval (GetIndex n xs))))++-- | modify an index at a specific dimension. Errors if out of bounds.+--+-- >>> modifyDim 0 (+1) [0,1,2]+-- [1,1,2]+-- >>> modifyDim 0 (+1) []+-- [2]+modifyDim :: Int -> (Int -> Int) -> [Int] -> [Int]+modifyDim 0 f [] = [f 1]+modifyDim d f xs =+ getDim d xs+ & f+ & (: drop (d + 1) xs)+ & (take d xs <>)++-- | modify an index at a specific dimension. Errors if out of bounds.+--+-- >>> :k! Eval (ModifyDim 0 ((Fcf.+) 1) [0,1,2])+-- ...+-- = [1, 1, 2]+data ModifyDim :: Nat -> (Nat -> Exp Nat) -> [Nat] -> Exp [Nat]++type instance+ Eval (ModifyDim d f s) =+ Eval (LiftM2 (Fcf.++) (Take d s) (LiftM2 Cons (f =<< GetDim d s) (Drop (d + 1) s)))++-- | Increment the index at a dimension of a shape by 1. Scalars turn into singletons.+--+-- >>> incAt 1 [2,3,4]+-- [2,4,4]+-- >>> incAt 0 []+-- [2]+incAt :: Int -> [Int] -> [Int]+incAt d ds = modifyDim d (+ 1) (asSingleton ds)++-- | Increment the index at a dimension of a shape by 1. Scalars turn into singletons.+--+-- >>> :k! Eval (IncAt 1 [2,3,4])+-- ...+-- = [2, 4, 4]+-- >>> :k! Eval (IncAt 0 '[])+-- ...+-- = '[2]+data IncAt :: Nat -> [Nat] -> Exp [Nat]++type instance+ Eval (IncAt d ds) =+ Eval (ModifyDim d ((Fcf.+) 1) (Eval (AsSingleton ds)))++-- | Decrement the index at a dimension os a shape by 1.+--+-- >>> decAt 1 [2,3,4]+-- [2,2,4]+decAt :: Int -> [Int] -> [Int]+decAt d = modifyDim d (\x -> x - 1)++-- | Decrement the index at a dimension of a shape by 1.+--+-- >>> :k! Eval (DecAt 1 [2,3,4])+-- ...+-- = [2, 2, 4]+data DecAt :: Nat -> [Nat] -> Exp [Nat]++type instance+ Eval (DecAt d ds) =+ Eval (ModifyDim d (Flip (Fcf.-) 1) ds)++-- | replace an index at a specific dimension, or transform a scalar into being 1-dimensional.+--+-- >>> setDim 0 1 [2,3,4]+-- [1,3,4]+-- >>> setDim 0 3 []+-- [3]+setDim :: Int -> Int -> [Int] -> [Int]+setDim d x = modifyDim d (const x)++-- | replace an index at a specific dimension.+--+-- >>> :k! Eval (SetDim 0 1 [2,3,4])+-- ...+-- = [1, 3, 4]+data SetDim :: Nat -> Nat -> [Nat] -> Exp [Nat]++type instance+ Eval (SetDim d x ds) =+ Eval (ModifyDim d (ConstFn x) ds)++data SetDimUncurried :: (Nat, Nat) -> [Nat] -> Exp [Nat]++type instance+ Eval (SetDimUncurried xs ds) =+ Eval (SetDim (Eval (Fst xs)) (Eval (Snd xs)) ds)++-- | Take along a dimension.+--+-- >>> takeDim 0 1 [2,3,4]+-- [1,3,4]+takeDim :: Int -> Int -> [Int] -> [Int]+takeDim d t = modifyDim d (min t)++-- | Take along a dimension.+--+-- >>> :k! Eval (TakeDim 0 1 [2,3,4])+-- ...+-- = [1, 3, 4]+data TakeDim :: Nat -> Nat -> [Nat] -> Exp [Nat]++type instance+ Eval (TakeDim d t s) =+ Eval+ (ModifyDim d (Min t) s)++-- | Drop along a dimension.+--+-- >>> dropDim 2 1 [2,3,4]+-- [2,3,3]+dropDim :: Int -> Int -> [Int] -> [Int]+dropDim d t = modifyDim d (max 0 . (\x -> x - t))++-- | Drop along a dimension.+--+-- >>> :k! Eval (DropDim 2 1 [2,3,4])+-- ...+-- = [2, 3, 3]+data DropDim :: Nat -> Nat -> [Nat] -> Exp [Nat]++type instance+ Eval (DropDim d t s) =+ Eval+ ( ModifyDim+ d+ (Max 0 <=< Flip (Fcf.-) t)+ s+ )++-- | delete the i'th dimension. No effect on a scalar.+--+-- >>> deleteDim 1 [2, 3, 4]+-- [2,4]+-- >>> deleteDim 2 []+-- []+deleteDim :: Int -> [Int] -> [Int]+deleteDim i s = take i s ++ drop (i + 1) s++-- | delete the i'th dimension+--+-- >>> :k! Eval (DeleteDim 1 [2, 3, 4])+-- ...+-- = [2, 4]+-- >>> :k! Eval (DeleteDim 1 '[])+-- ...+-- = '[]+data DeleteDim :: Nat -> [Nat] -> Exp [Nat]++type instance+ Eval (DeleteDim i ds) =+ Eval (LiftM2 (Fcf.++) (Take i ds) (Drop (i + 1) ds))++-- | Insert a new dimension at a position (or at the end if > rank).+--+-- >>> insertDim 1 3 [2,4]+-- [2,3,4]+-- >>> insertDim 0 4 []+-- [4]+insertDim :: Int -> Int -> [Int] -> [Int]+insertDim d i s = take d s ++ (i : drop d s)++-- | Insert a new dimension at a position (or at the end if > rank).+--+-- >>> :k! Eval (InsertDim 1 3 [2,4])+-- ...+-- = [2, 3, 4]+-- >>> :k! Eval (InsertDim 0 4 '[])+-- ...+-- = '[4]+data InsertDim :: Nat -> Nat -> [Nat] -> Exp [Nat]++type instance+ Eval (InsertDim d i ds) =+ Eval (LiftM2 (Fcf.++) (Take d ds) (Cons i =<< Drop d ds))++data InsertDimUncurried :: (Nat, Nat) -> [Nat] -> Exp [Nat]++type instance+ Eval (InsertDimUncurried xs ds) =+ Eval (InsertDim (Eval (Fst xs)) (Eval (Snd xs)) ds)++-- | Is a slice ok constraint.+--+-- >>> :k! Eval (InsertOk 2 [2,3,4] [2,3])+-- ...+-- = True+-- >>> :k! Eval (InsertOk 0 '[] '[])+-- ...+-- = True+data InsertOk :: Nat -> [Nat] -> [Nat] -> Exp Bool++type instance+ Eval (InsertOk d s si) =+ Eval+ ( And+ [ Eval (IsDim d s),+ Eval (TyEq si (Eval (DeleteDim d s)))+ ]+ )++-- | Is a slice ok?+--+-- >>> :k! Eval (SliceOk 1 1 2 [2,3,4])+-- ...+-- = True+data SliceOk :: Nat -> Nat -> Nat -> [Nat] -> Exp Bool++type instance+ Eval (SliceOk d off l s) =+ Eval+ ( And+ [ Eval (IsFin off =<< GetDim d s),+ Eval ((Fcf.<) l =<< GetDim d s),+ Eval ((Fcf.<) (off + l) (Eval (GetDim d s) + 1)),+ Eval (IsDim d s)+ ]+ )++-- | Combine elements of two lists pairwise.+data ZipWith3 :: (a -> b -> c -> Exp d) -> [a] -> [b] -> [c] -> Exp [d]++type instance Eval (ZipWith3 _f '[] _bs _cs) = '[]++type instance Eval (ZipWith3 _f _as '[] _cs) = '[]++type instance Eval (ZipWith3 _f _as _bs '[]) = '[]++type instance+ Eval (ZipWith3 f (a ': as) (b ': bs) (c ': cs)) =+ Eval (f a b c) ': Eval (ZipWith3 f as bs cs)++data SliceOk_ :: [Nat] -> Nat -> Nat -> Nat -> Exp Bool++type instance Eval (SliceOk_ s d off l) = Eval (SliceOk d off l s)++-- | Are slices ok?+--+-- >>> :k! Eval (SlicesOk '[1] '[1] '[2] [2,3,4])+-- ...+-- = True+data SlicesOk :: [Nat] -> [Nat] -> [Nat] -> [Nat] -> Exp Bool++type instance+ Eval (SlicesOk ds offs ls s) =+ Eval (And =<< ZipWith3 (SliceOk_ s) ds offs ls)++-- | concatenate two arrays at dimension i+--+-- Bespoke logic for scalars.+--+-- >>> concatenate 1 [2,3,4] [2,3,4]+-- [2,6,4]+-- >>> concatenate 0 [3] []+-- [4]+-- >>> concatenate 0 [] [3]+-- [4]+-- >>> concatenate 0 [] []+-- [2]+concatenate :: Int -> [Int] -> [Int] -> [Int]+concatenate _ [] [] = [2]+concatenate _ [] [x] = [x + 1]+concatenate _ [x] [] = [x + 1]+concatenate i s0 s1 = take i s0 ++ (getDim i s0 + getDim i s1 : drop (i + 1) s0)++-- | concatenate two arrays at dimension i+--+-- Bespoke logic for scalars.+--+-- >>> :k! Eval (Concatenate 1 [2,3,4] [2,3,4])+-- ...+-- = [2, 6, 4]+-- >>> :k! Eval (Concatenate 0 '[3] '[])+-- ...+-- = '[4]+-- >>> :k! Eval (Concatenate 0 '[] '[3])+-- ...+-- = '[4]+-- >>> :k! Eval (Concatenate 0 '[] '[])+-- ...+-- = '[2]+data Concatenate :: Nat -> [Nat] -> [Nat] -> Exp [Nat]++type instance+ Eval (Concatenate i s0 s1) =+ If+ (Eval (ConcatenateOk i s0 s1))+ (Eval (Eval (Take i s0) ++ (Eval (GetDim i s0) + Eval (GetDim i s1) : Eval (Drop (i + 1) s0))))+ (L.TypeError (L.Text "Concatenate Mis-matched shapes."))++-- | Concatenate is Ok if ranks are the same and the non-indexed portion of the shapes are the same.+data ConcatenateOk :: Nat -> [Nat] -> [Nat] -> Exp Bool++type instance+ Eval (ConcatenateOk i s0 s1) =+ Eval (IsDim i s0)+ && Eval (IsDim i s1)+ && Eval (LiftM2 TyEq (DeleteDim i s0) (DeleteDim i s1))+ && Eval (LiftM2 TyEq (Rank =<< AsSingleton s0) (Rank =<< AsSingleton s1))++-- * multiple dimension manipulations++-- | Get dimensions of a shape.+--+-- >>> getDims [2,0] [2,3,4]+-- [4,2]+-- >>> getDims [2] []+-- []+getDims :: [Int] -> [Int] -> [Int]+getDims _ [] = []+getDims i s = (`getDim` s) <$> i++-- | Get dimensions of a shape.+--+-- >>> :k! Eval (GetDims [2,0] [2,3,4])+-- ...+-- = [4, 2]+-- >>> :k! Eval (GetDims '[2] '[])+-- ...+-- = '[(TypeError ...)]+data GetDims :: [Nat] -> [Nat] -> Exp [Nat]++type instance+ Eval (GetDims xs ds) =+ Eval (Map (Flip GetDim ds) xs)++-- | Get the index of the last position in the selected dimensions of a shape. Errors on a 0-dimension.+--+-- >>> getLastPositions [2,0] [2,3,4]+-- [3,1]+-- >>> getLastPositions [0] [0]+-- [-1]+getLastPositions :: [Int] -> [Int] -> [Int]+getLastPositions ds s =+ fmap (\x -> x - 1) (getDims ds s)++-- | Get the index of the last position in the selected dimensions of a shape. Errors on a 0-dimension.+--+-- >>> :k! Eval (GetLastPositions [2,0] [2,3,4])+-- ...+-- = [3, 1]+-- >>> :k! Eval (GetLastPositions '[0] '[0])+-- ...+-- = '[0 GHC.TypeNats.- 1]+data GetLastPositions :: [Nat] -> [Nat] -> Exp [Nat]++type instance+ Eval (GetLastPositions ds s) =+ Eval (Map (Flip (Fcf.-) 1) (Eval (GetDims ds s)))++-- | modify dimensions of a shape with (separate) functions.+--+-- >>> modifyDims [0,1] [(+1), (+5)] [2,3,4]+-- [3,8,4]+modifyDims :: [Int] -> [Int -> Int] -> [Int] -> [Int]+modifyDims ds fs ns = foldl' (\ns' (d, f) -> modifyDim d f ns') ns (zip ds fs)++-- | Convert a list of positions that reference deletions according to a final shape to 1 that references deletions relative to an initial shape.+--+-- To delete the positions [1,2,5] from a list, for example, you need to delete position 1, (arriving at a 4 element list), then position 1, arriving at a 3 element list, and finally position 3.+--+-- >>> preDeletePositions [1,2,5]+-- [1,1,3]+--+-- >>> preDeletePositions [1,2,0]+-- [1,1,0]+preDeletePositions :: [Int] -> [Int]+preDeletePositions as = reverse (go as [])+ where+ go [] r = r+ go (x : xs) r = go (decPast x <$> xs) (x : r)+ decPast x y = bool (y - 1) y (y < x)++-- | Convert a list of positions that reference deletions according to a final shape to 1 that references deletions relative to an initial shape.+--+-- To delete the positions [1,2,5] from a list, for example, you need to delete position 1, (arriving at a 4 element list), then position 1, arriving at a 3 element list, and finally position 3.+--+-- >>> :k! Eval (PreDeletePositions [1,2,5])+-- ...+-- = [1, 1, 3]+--+-- >>> :k! Eval (PreDeletePositions [1,2,0])+-- ...+-- = [1, 1, 0]+data PreDeletePositions :: [Nat] -> Exp [Nat]++type instance+ Eval (PreDeletePositions xs) =+ Eval (Reverse (Eval (PreDeletePositionsGo xs '[])))++data PreDeletePositionsGo :: [Nat] -> [Nat] -> Exp [Nat]++type instance Eval (PreDeletePositionsGo '[] rs) = rs++type instance+ Eval (PreDeletePositionsGo (x : xs) r) =+ Eval (PreDeletePositionsGo (Eval (Map (DecPast x) xs)) (x : r))++data DecPast :: Nat -> Nat -> Exp Nat++type instance+ Eval (DecPast x d) =+ If (x + 1 <=? d) (d - 1) d++-- | Convert a list of position that reference insertions according to a final shape to 1 that references list insertions relative to an initial shape.+--+-- To insert into positions [1,2,0] from a list, starting from a 2 element list, for example, you need to insert at position 0, (arriving at a 3 element list), then position 1, arriving at a 4 element list, and finally position 0.+--+-- > preInsertPositions == reverse . preDeletePositions . reverse+-- >>> preInsertPositions [1,2,5]+-- [1,2,5]+--+-- >>> preInsertPositions [1,2,0]+-- [0,1,0]+preInsertPositions :: [Int] -> [Int]+preInsertPositions = reverse . preDeletePositions . reverse++-- | Convert a list of position that reference insertions according to a final shape to 1 that references list insertions relative to an initial shape.+--+-- To insert into positions [1,2,0] from a list, starting from a 2 element list, for example, you need to insert at position 0, (arriving at a 3 element list), then position 1, arriving at a 4 element list, and finally position 0.+--+-- > preInsertPositions == reverse . preDeletePositions . reverse+-- >>> :k! Eval (PreInsertPositions [1,2,5])+-- ...+-- = [1, 2, 5]+--+-- >>> :k! Eval (PreInsertPositions [1,2,0])+-- ...+-- = [0, 1, 0]+data PreInsertPositions :: [Nat] -> Exp [Nat]++type instance+ Eval (PreInsertPositions xs) =+ Eval (Reverse =<< (PreDeletePositions =<< Reverse xs))++-- | drop dimensions of a shape according to a list of positions (where position refers to the initial shape)+--+-- >>> deleteDims [1,0] [2, 3, 4]+-- [4]+deleteDims :: [Int] -> [Int] -> [Int]+deleteDims i s = foldl' (flip deleteDim) s (preDeletePositions i)++-- | drop dimensions of a shape according to a list of positions (where position refers to the initial shape)+--+-- >>> :k! Eval (DeleteDims [1,0] [2, 3, 4])+-- ...+-- = '[4]+data DeleteDims :: [Nat] -> [Nat] -> Exp [Nat]++type instance+ Eval (DeleteDims xs ds) =+ Eval (Foldl' (Flip DeleteDim) ds =<< PreDeletePositions xs)++-- | Insert a list of dimensions according to dimensions and positions. Note that the list of positions references the final shape and not the initial shape.+--+-- >>> insertDims [0] [5] []+-- [5]+-- >>> insertDims [1,0] [3,2] [4]+-- [2,3,4]+insertDims :: [Int] -> [Int] -> [Int] -> [Int]+insertDims ds xs s = foldl' (flip (uncurry insertDim)) s ps+ where+ ps = zip (preInsertPositions ds) xs++-- | insert a list of dimensions according to dimension,position tuple lists. Note that the list of positions references the final shape and not the initial shape.+--+-- >>> :k! Eval (InsertDims '[0] '[5] '[])+-- ...+-- = '[5]+-- >>> :k! Eval (InsertDims [1,0] [3,2] '[4])+-- ...+-- = [2, 3, 4]+data InsertDims :: [Nat] -> [Nat] -> [Nat] -> Exp [Nat]++type instance+ Eval (InsertDims ds xs s) =+ Eval (Foldl' (Flip InsertDimUncurried) s =<< Flip Zip xs =<< PreInsertPositions ds)++-- | Set dimensions of a shape.+--+-- >>> setDims [0,1] [1,5] [2,3,4]+-- [1,5,4]+--+-- >>> setDims [0] [3] []+-- [3]+setDims :: [Int] -> [Int] -> [Int] -> [Int]+setDims ds xs ns = foldl' (\ns' (d, x) -> setDim d x ns') ns (zip ds xs)++-- | Set dimensions of a shape.+--+-- >>> :k! Eval (SetDims [0,1] [1,5] [2,3,4])+-- ...+-- = [1, 5, 4]+--+-- >>> :k! Eval (SetDims '[0] '[3] '[])+-- ...+-- = '[3]+data SetDims :: [Nat] -> [Nat] -> [Nat] -> Exp [Nat]++type instance+ Eval (SetDims ds xs ns) =+ Eval (Foldl' (Flip SetDimUncurried) ns =<< Zip ds xs)++-- | Drop a number of elements of a shape along the supplied dimensions.+--+-- >>> dropDims [0,2] [1,3] [2,3,4]+-- [1,3,1]+dropDims :: [Int] -> [Int] -> [Int] -> [Int]+dropDims ds xs s = setDims ds xs' s+ where+ xs' = zipWith (-) (getDims ds s) xs++-- | Drop a number of elements of a shape along the supplied dimensions.+--+-- >>> :k! Eval (DropDims [0,2] [1,3] [2,3,4])+-- ...+-- = [1, 3, 1]+data DropDims :: [Nat] -> [Nat] -> [Nat] -> Exp [Nat]++type instance+ Eval (DropDims ds xs s) =+ Eval (SetDims ds (Eval (ZipWith (Fcf.-) (Eval (GetDims ds s)) xs)) s)++-- | Concatenate and replace dimensions, creating a new dimension at the supplied postion.+--+-- >>> concatDims [0,1] 1 [2,3,4]+-- [4,6]+concatDims :: [Int] -> Int -> [Int] -> [Int]+concatDims ds n s = insertDim n (size $ getDims ds s) (deleteDims ds s)++-- | Drop a number of elements of a shape along the supplied dimensions.+--+-- >>> :k! Eval (ConcatDims [0,1] 1 [2,3,4])+-- ...+-- = [4, 6]+data ConcatDims :: [Nat] -> Nat -> [Nat] -> Exp [Nat]++type instance+ Eval (ConcatDims ds n s) =+ Eval (InsertDim n (Eval (Size (Eval (GetDims ds s)))) (Eval (DeleteDims ds s)))++-- | Unconcatenate and reinsert dimensions for an index.+--+-- >>> unconcatDimsIndex [0,1] 1 [4,6] [2,3]+-- [0,3,2]+unconcatDimsIndex :: [Int] -> Int -> [Int] -> [Int] -> [Int]+unconcatDimsIndex ds n s i = insertDims ds (shapen (getDims ds s) (getDim n i)) (deleteDim n i)++-- | reverse an index along specific dimensions.+--+-- >>> reverseIndex [0] [2,3,4] [0,1,2]+-- [1,1,2]+reverseIndex :: [Int] -> [Int] -> [Int] -> [Int]+reverseIndex ds ns xs = fmap (\(i, x, n) -> bool x (n - 1 - x) (i `elem` ds)) (zip3 [0 ..] xs ns)++-- | rotate a list+--+-- >>> rotate 1 [0..3]+-- [1,2,3,0]+-- >>> rotate (-1) [0..3]+-- [3,0,1,2]+rotate :: Int -> [a] -> [a]+rotate r xs = drop r' xs <> take r' xs+ where+ r' = r `mod` List.length xs++-- | rotate an index along a specific dimension.+--+-- >>> rotateIndex 0 1 [2,3,4] [0,1,2]+-- [1,1,2]+rotateIndex :: Int -> Int -> [Int] -> [Int] -> [Int]+rotateIndex d r s = modifyDim d (\x -> (x + r) `mod` getDim d s)++-- | rotate an index along specific dimensions.+--+-- >>> rotatesIndex [0] [1] [2,3,4] [0,1,2]+-- [1,1,2]+rotatesIndex :: [Int] -> [Int] -> [Int] -> [Int] -> [Int]+rotatesIndex ds rs s xs = foldr (\(d, r) acc -> rotateIndex d r s acc) xs (zip ds rs)++-- | Test whether an index is a diagonal one.+--+-- >>> isDiag [2,2,2]+-- True+-- >>> isDiag [1,2]+-- False+isDiag :: (Eq a) => [a] -> Bool+isDiag [] = True+isDiag [_] = True+isDiag [x, y] = x == y+isDiag (x : y : xs) = x == y && isDiag (y : xs)++-- | Expanded shape of a windowed array+--+-- >>> expandWindows [2,2] [4,3,2]+-- [3,2,2,2,2]+expandWindows :: [Int] -> [Int] -> [Int]+expandWindows ws ds = List.zipWith (\s' x' -> s' - x' + 1) ds ws <> ws <> List.drop (rank ws) ds++-- | Expanded shape of a windowed array+--+-- >>> :k! Eval (ExpandWindows [2,2] [4,3,2])+-- ...+-- = [3, 2, 2, 2, 2]+data ExpandWindows :: [Nat] -> [Nat] -> Exp [Nat]++type instance+ Eval (ExpandWindows ws ds) =+ Eval (Eval (ZipWith (Fcf.-) (Eval (Map ((Fcf.+) 1) ds)) ws) ++ Eval (ws ++ Eval (Drop (Eval (Rank ws)) ds)))++-- | Index into windows of an expanded windowed array, given a rank of the windows.+--+-- >>> indexWindows 2 [0,1,2,1,1]+-- [2,2,1]+indexWindows :: Int -> [Int] -> [Int]+indexWindows r ds = List.zipWith (+) (List.take r ds) (List.take r (List.drop r ds)) <> List.drop (r + r) ds++-- | Dimensions of a windowed array.+--+-- >>> dimWindows [2,2] [2,3,4]+-- [0,1,2]+dimWindows :: [Int] -> [Int] -> [Int]+dimWindows ws s = range (rank s) <> [rank s * 2 .. (rank ws - 1)]++-- | Dimensions of a windowed array.+--+-- >>> :k! Eval (DimWindows [2,2] [4,3,2])+-- ...+-- = [0, 1, 2]+data DimWindows :: [Nat] -> [Nat] -> Exp [Nat]++type instance+ Eval (DimWindows ws s) =+ Eval (Eval (Range =<< Rank s) ++ Eval (EnumFromTo (Eval ((Fcf.*) 2 (Eval (Rank s)))) (Eval (Rank ws) - 1)))
+ src/Harpie/Sort.hs view
@@ -0,0 +1,60 @@+-- | 'Vector' sort routines.+module Harpie.Sort+ ( sortV,+ sortByV,+ orderV,+ orderByV,+ )+where++import Data.Ord+import Data.Vector (Vector, convert, unsafeIndex)+import Data.Vector qualified as V+import Data.Vector.Algorithms.Intro (sortBy)+import Data.Vector.Unboxed (generate, modify)+import Prelude++-- $setup+-- >>> :m -Prelude+-- >>> import Data.Vector qualified as V+-- >>> import Data.Ord (Down (..))+-- >>> import Prelude (Int)+-- >>> :set -XDataKinds+-- >>> :set -XTypeFamilies+-- >>> :set -XFlexibleContexts++-- | return the sorted array+--+-- >>> sortV (V.fromList [3,1,4,2,0,5::Int])+-- [0,1,2,3,4,5]+sortV :: (Ord a) => Vector a -> Vector a+sortV a = V.unsafeBackpermute a (orderV a)++-- | return the array sorted by the comparison function+--+-- >>> sortByV Down (V.fromList [3,1,4,2,0,5::Int])+-- [5,4,3,2,1,0]+sortByV :: (Ord b) => (a -> b) -> Vector a -> Vector a+sortByV c a = V.unsafeBackpermute a (orderByV c a)++-- | returns the indices of the elements in ascending order.+--+-- >>> orderV (V.fromList [0..5::Int])+-- [0,1,2,3,4,5]+orderV :: (Ord a) => Vector a -> Vector Int+orderV a = idx+ where+ idx = convert $ modify (sortBy comp) init0+ comp = comparing $ unsafeIndex a -- comparing function+ init0 = generate (V.length a) id -- [0..size - 1]++-- | returns the indices of the elements in order given a comparison function.+--+-- >>> orderByV Down (V.fromList [0..5::Int])+-- [5,4,3,2,1,0]+orderByV :: (Ord b) => (a -> b) -> Vector a -> Vector Int+orderByV c a = idx+ where+ idx = convert $ modify (sortBy comp) init0+ comp = comparing $ c . unsafeIndex a -- comparing function+ init0 = generate (V.length a) id -- [0..size - 1]
+ test/doctests.hs view
@@ -0,0 +1,10 @@+module Main where++import System.Environment (getArgs)+import Test.DocTest (mainFromCabal)+import Prelude (IO, (=<<))++main :: IO ()+main = mainFromCabal "harpie" =<< getArgs++-- (\a -> toDynamic (F.takeBs (Dims @'[0]) (S.SNats @'[1]) a) == (A.takes [0] [-1] (toDynamic a))) (F.range @[2,3,4])