aboutsummaryrefslogtreecommitdiff
path: root/src/Data/Array/Mixed
diff options
context:
space:
mode:
Diffstat (limited to 'src/Data/Array/Mixed')
-rw-r--r--src/Data/Array/Mixed/Permutation.hs5
-rw-r--r--src/Data/Array/Mixed/XArray.hs331
2 files changed, 336 insertions, 0 deletions
diff --git a/src/Data/Array/Mixed/Permutation.hs b/src/Data/Array/Mixed/Permutation.hs
index 83a5ee4..6ff3bdc 100644
--- a/src/Data/Array/Mixed/Permutation.hs
+++ b/src/Data/Array/Mixed/Permutation.hs
@@ -61,6 +61,11 @@ permToList (x `PCons` l) = TN.fromSNat x : permToList l
permToList' :: Perm list -> [Int]
permToList' = map fromIntegral . permToList
+-- | Utility class for generating permutations from type class information.
+class KnownPerm l where makePerm :: Perm l
+instance KnownPerm '[] where makePerm = PNil
+instance (KnownNat n, KnownPerm l) => KnownPerm (n : l) where makePerm = natSing `PCons` makePerm
+
-- ** Applying permutations
diff --git a/src/Data/Array/Mixed/XArray.hs b/src/Data/Array/Mixed/XArray.hs
new file mode 100644
index 0000000..cc0a6a5
--- /dev/null
+++ b/src/Data/Array/Mixed/XArray.hs
@@ -0,0 +1,331 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE ImportQualifiedPost #-}
+{-# LANGUAGE NoStarIsType #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE StandaloneKindSignatures #-}
+{-# LANGUAGE StrictData #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeOperators #-}
+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
+module Data.Array.Mixed.XArray where
+
+import Control.DeepSeq (NFData(..))
+import Data.Array.Ranked qualified as ORB
+import Data.Array.RankedS qualified as S
+import Data.Coerce
+import Data.Kind
+import Data.Proxy
+import Data.Type.Equality
+import Data.Type.Ord
+import Data.Vector.Storable qualified as VS
+import Foreign.Storable (Storable)
+import GHC.Generics (Generic)
+import GHC.TypeLits
+
+import Data.Array.Mixed.Internal.Arith
+import Data.Array.Mixed.Lemmas
+import Data.Array.Mixed.Permutation
+import Data.Array.Mixed.Shape
+import Data.Array.Mixed.Types
+
+
+type XArray :: [Maybe Nat] -> Type -> Type
+newtype XArray sh a = XArray (S.Array (Rank sh) a)
+ deriving (Show, Eq, Generic)
+
+-- | Only on scalars, because lexicographical ordering is strange on multi-dimensional arrays.
+deriving instance (Ord a, Storable a) => Ord (XArray '[] a)
+
+instance NFData a => NFData (XArray sh a)
+
+
+shape :: forall sh a. StaticShX sh -> XArray sh a -> IShX sh
+shape = \ssh (XArray arr) -> go ssh (S.shapeL arr)
+ where
+ go :: StaticShX sh' -> [Int] -> IShX sh'
+ go ZKX [] = ZSX
+ go (n :!% ssh) (i : l) = fromSMayNat (\_ -> SUnknown i) SKnown n :$% go ssh l
+ go _ _ = error "Invalid shapeL"
+
+fromVector :: forall sh a. Storable a => IShX sh -> VS.Vector a -> XArray sh a
+fromVector sh v
+ | Dict <- lemKnownNatRank sh
+ = XArray (S.fromVector (shxToList sh) v)
+
+toVector :: Storable a => XArray sh a -> VS.Vector a
+toVector (XArray arr) = S.toVector arr
+
+scalar :: Storable a => a -> XArray '[] a
+scalar = XArray . S.scalar
+
+-- | Will throw if the array does not have the casted-to shape.
+cast :: forall sh1 sh2 sh' a. Rank sh1 ~ Rank sh2
+ => StaticShX sh1 -> IShX sh2 -> StaticShX sh'
+ -> XArray (sh1 ++ sh') a -> XArray (sh2 ++ sh') a
+cast ssh1 sh2 ssh' (XArray arr)
+ | Refl <- lemRankApp ssh1 ssh'
+ , Refl <- lemRankApp (ssxFromShape sh2) ssh'
+ = let arrsh :: IShX sh1
+ (arrsh, _) = shxSplitApp (Proxy @sh') ssh1 (shape (ssxAppend ssh1 ssh') (XArray arr))
+ in if shxToList arrsh == shxToList sh2
+ then XArray arr
+ else error $ "Data.Array.Mixed.cast: Cannot cast (" ++ show arrsh ++ ") to (" ++ show sh2 ++ ")"
+
+unScalar :: Storable a => XArray '[] a -> a
+unScalar (XArray a) = S.unScalar a
+
+replicate :: forall sh sh' a. Storable a => IShX sh -> StaticShX sh' -> XArray sh' a -> XArray (sh ++ sh') a
+replicate sh ssh' (XArray arr)
+ | Dict <- lemKnownNatRankSSX ssh'
+ , Dict <- lemKnownNatRankSSX (ssxAppend (ssxFromShape sh) ssh')
+ , Refl <- lemRankApp (ssxFromShape sh) ssh'
+ = XArray (S.stretch (shxToList sh ++ S.shapeL arr) $
+ S.reshape (map (const 1) (shxToList sh) ++ S.shapeL arr) $
+ arr)
+
+replicateScal :: forall sh a. Storable a => IShX sh -> a -> XArray sh a
+replicateScal sh x
+ | Dict <- lemKnownNatRank sh
+ = XArray (S.constant (shxToList sh) x)
+
+generate :: Storable a => IShX sh -> (IIxX sh -> a) -> XArray sh a
+generate sh f = fromVector sh $ VS.generate (shxSize sh) (f . ixxFromLinear sh)
+
+-- generateM :: (Monad m, Storable a) => IShX sh -> (IIxX sh -> m a) -> m (XArray sh a)
+-- generateM sh f | Dict <- lemKnownNatRank sh =
+-- XArray . S.fromVector (shxShapeL sh)
+-- <$> VS.generateM (shxSize sh) (f . ixxFromLinear sh)
+
+indexPartial :: Storable a => XArray (sh ++ sh') a -> IIxX sh -> XArray sh' a
+indexPartial (XArray arr) ZIX = XArray arr
+indexPartial (XArray arr) (i :.% idx) = indexPartial (XArray (S.index arr i)) idx
+
+index :: forall sh a. Storable a => XArray sh a -> IIxX sh -> a
+index xarr i
+ | Refl <- lemAppNil @sh
+ = let XArray arr' = indexPartial xarr i :: XArray '[] a
+ in S.unScalar arr'
+
+append :: forall n m sh a. Storable a
+ => StaticShX sh -> XArray (n : sh) a -> XArray (m : sh) a -> XArray (AddMaybe n m : sh) a
+append ssh (XArray a) (XArray b)
+ | Dict <- lemKnownNatRankSSX ssh
+ = XArray (S.append a b)
+
+-- | If the prefix of the shape of the input array (@sh@) is empty (i.e.
+-- contains a zero), then there is no way to deduce the full shape of the output
+-- array (more precisely, the @sh2@ part): that could only come from calling
+-- @f@, and there are no subarrays to call @f@ on. @orthotope@ errors out in
+-- this case; we choose to fill the shape with zeros wherever we cannot deduce
+-- what it should be.
+--
+-- For example, if:
+--
+-- @
+-- arr :: XArray '[Just 3, Just 0, Just 4, Just 2, Nothing] Int -- of shape [3, 0, 4, 2, 21]
+-- f :: XArray '[Just 2, Nothing] Int -> XArray '[Just 5, Nothing, Just 17] Float
+-- @
+--
+-- then:
+--
+-- @
+-- rerank _ _ _ f arr :: XArray '[Just 3, Just 0, Just 4, Just 5, Nothing, Just 17] Float
+-- @
+--
+-- and this result will have shape @[3, 0, 4, 5, 0, 17]@. Note the second @0@
+-- in this shape: we don't know if @f@ intended to return an array with shape 0
+-- here (it probably didn't), but there is no better number to put here absent
+-- a subarray of the input to pass to @f@.
+--
+-- In this particular case the fact that @sh@ is empty was evident from the
+-- type-level information, but the same situation occurs when @sh@ consists of
+-- @Nothing@s, and some of those happen to be zero at runtime.
+rerank :: forall sh sh1 sh2 a b.
+ (Storable a, Storable b)
+ => StaticShX sh -> StaticShX sh1 -> StaticShX sh2
+ -> (XArray sh1 a -> XArray sh2 b)
+ -> XArray (sh ++ sh1) a -> XArray (sh ++ sh2) b
+rerank ssh ssh1 ssh2 f xarr@(XArray arr)
+ | Dict <- lemKnownNatRankSSX (ssxAppend ssh ssh2)
+ = let (sh, _) = shxSplitApp (Proxy @sh1) ssh (shape (ssxAppend ssh ssh1) xarr)
+ in if any (== 0) (shxToList sh)
+ then XArray (S.fromList (shxToList (shxAppend sh (shxCompleteZeros ssh2))) [])
+ else case () of
+ () | Dict <- lemKnownNatRankSSX ssh
+ , Dict <- lemKnownNatRankSSX ssh2
+ , Refl <- lemRankApp ssh ssh1
+ , Refl <- lemRankApp ssh ssh2
+ -> XArray (S.rerank @(Rank sh) @(Rank sh1) @(Rank sh2)
+ (\a -> let XArray r = f (XArray a) in r)
+ arr)
+
+rerankTop :: forall sh1 sh2 sh a b.
+ (Storable a, Storable b)
+ => StaticShX sh1 -> StaticShX sh2 -> StaticShX sh
+ -> (XArray sh1 a -> XArray sh2 b)
+ -> XArray (sh1 ++ sh) a -> XArray (sh2 ++ sh) b
+rerankTop ssh1 ssh2 ssh f = transpose2 ssh ssh2 . rerank ssh ssh1 ssh2 f . transpose2 ssh1 ssh
+
+-- | The caveat about empty arrays at @rerank@ applies here too.
+rerank2 :: forall sh sh1 sh2 a b c.
+ (Storable a, Storable b, Storable c)
+ => StaticShX sh -> StaticShX sh1 -> StaticShX sh2
+ -> (XArray sh1 a -> XArray sh1 b -> XArray sh2 c)
+ -> XArray (sh ++ sh1) a -> XArray (sh ++ sh1) b -> XArray (sh ++ sh2) c
+rerank2 ssh ssh1 ssh2 f xarr1@(XArray arr1) (XArray arr2)
+ | Dict <- lemKnownNatRankSSX (ssxAppend ssh ssh2)
+ = let (sh, _) = shxSplitApp (Proxy @sh1) ssh (shape (ssxAppend ssh ssh1) xarr1)
+ in if any (== 0) (shxToList sh)
+ then XArray (S.fromList (shxToList (shxAppend sh (shxCompleteZeros ssh2))) [])
+ else case () of
+ () | Dict <- lemKnownNatRankSSX ssh
+ , Dict <- lemKnownNatRankSSX ssh2
+ , Refl <- lemRankApp ssh ssh1
+ , Refl <- lemRankApp ssh ssh2
+ -> XArray (S.rerank2 @(Rank sh) @(Rank sh1) @(Rank sh2)
+ (\a b -> let XArray r = f (XArray a) (XArray b) in r)
+ arr1 arr2)
+
+-- | The list argument gives indices into the original dimension list.
+transpose :: forall is sh a. (IsPermutation is, Rank is <= Rank sh)
+ => StaticShX sh
+ -> Perm is
+ -> XArray sh a
+ -> XArray (PermutePrefix is sh) a
+transpose ssh perm (XArray arr)
+ | Dict <- lemKnownNatRankSSX ssh
+ , Refl <- lemRankApp (ssxPermute perm (ssxTakeLen perm ssh)) (ssxDropLen perm ssh)
+ , Refl <- lemRankPermute (Proxy @(TakeLen is sh)) perm
+ , Refl <- lemRankDropLen ssh perm
+ = XArray (S.transpose (permToList' perm) arr)
+
+-- | The list argument gives indices into the original dimension list.
+--
+-- The permutation (the list) must have length <= @n@. If it is longer, this
+-- function throws.
+transposeUntyped :: forall n sh a.
+ SNat n -> StaticShX sh -> [Int]
+ -> XArray (Replicate n Nothing ++ sh) a -> XArray (Replicate n Nothing ++ sh) a
+transposeUntyped sn ssh perm (XArray arr)
+ | length perm <= fromSNat' sn
+ , Dict <- lemKnownNatRankSSX (ssxAppend (ssxReplicate sn) ssh)
+ = XArray (S.transpose perm arr)
+ | otherwise
+ = error "Data.Array.Mixed.transposeUntyped: Permutation longer than length of unshaped prefix of shape type"
+
+transpose2 :: forall sh1 sh2 a.
+ StaticShX sh1 -> StaticShX sh2
+ -> XArray (sh1 ++ sh2) a -> XArray (sh2 ++ sh1) a
+transpose2 ssh1 ssh2 (XArray arr)
+ | Refl <- lemRankApp ssh1 ssh2
+ , Refl <- lemRankApp ssh2 ssh1
+ , Dict <- lemKnownNatRankSSX (ssxAppend ssh1 ssh2)
+ , Dict <- lemKnownNatRankSSX (ssxAppend ssh2 ssh1)
+ , Refl <- lemRankAppComm ssh1 ssh2
+ , let n1 = ssxLength ssh1
+ = XArray (S.transpose (ssxIotaFrom n1 ssh2 ++ ssxIotaFrom 0 ssh1) arr)
+
+sumFull :: (Storable a, NumElt a) => XArray sh a -> a
+sumFull (XArray arr) =
+ S.unScalar $
+ numEltSum1Inner (SNat @0) $
+ S.fromVector [product (S.shapeL arr)] $
+ S.toVector arr
+
+sumInner :: forall sh sh' a. (Storable a, NumElt a)
+ => StaticShX sh -> StaticShX sh' -> XArray (sh ++ sh') a -> XArray sh a
+sumInner ssh ssh' arr
+ | Refl <- lemAppNil @sh
+ = let (_, sh') = shxSplitApp (Proxy @sh') ssh (shape (ssxAppend ssh ssh') arr)
+ sh'F = shxFlatten sh' :$% ZSX
+ ssh'F = ssxFromShape sh'F
+
+ go :: XArray (sh ++ '[Flatten sh']) a -> XArray sh a
+ go (XArray arr')
+ | Refl <- lemRankApp ssh ssh'F
+ , let sn = listxLengthSNat (let StaticShX l = ssh in l)
+ = XArray (numEltSum1Inner sn arr')
+
+ in go $
+ transpose2 ssh'F ssh $
+ reshapePartial ssh' ssh sh'F $
+ transpose2 ssh ssh' $
+ arr
+
+sumOuter :: forall sh sh' a. (Storable a, NumElt a)
+ => StaticShX sh -> StaticShX sh' -> XArray (sh ++ sh') a -> XArray sh' a
+sumOuter ssh ssh' arr
+ | Refl <- lemAppNil @sh
+ = let (sh, _) = shxSplitApp (Proxy @sh') ssh (shape (ssxAppend ssh ssh') arr)
+ shF = shxFlatten sh :$% ZSX
+ in sumInner ssh' (ssxFromShape shF) $
+ transpose2 (ssxFromShape shF) ssh' $
+ reshapePartial ssh ssh' shF $
+ arr
+
+fromListOuter :: forall n sh a. Storable a
+ => StaticShX (n : sh) -> [XArray sh a] -> XArray (n : sh) a
+fromListOuter ssh l
+ | Dict <- lemKnownNatRankSSX ssh
+ = case ssh of
+ SKnown m :!% _ | fromSNat' m /= length l ->
+ error $ "Data.Array.Mixed.fromListOuter: length of list (" ++ show (length l) ++ ")" ++
+ "does not match the type (" ++ show (fromSNat' m) ++ ")"
+ _ -> XArray (S.ravel (ORB.fromList [length l] (coerce @[XArray sh a] @[S.Array (Rank sh) a] l)))
+
+toListOuter :: Storable a => XArray (n : sh) a -> [XArray sh a]
+toListOuter (XArray arr) =
+ case S.shapeL arr of
+ 0 : _ -> []
+ _ -> coerce (ORB.toList (S.unravel arr))
+
+fromList1 :: Storable a => StaticShX '[n] -> [a] -> XArray '[n] a
+fromList1 ssh l =
+ let n = length l
+ in case ssh of
+ SKnown m :!% _ | fromSNat' m /= n ->
+ error $ "Data.Array.Mixed.fromList1: length of list (" ++ show n ++ ")" ++
+ "does not match the type (" ++ show (fromSNat' m) ++ ")"
+ _ -> XArray (S.fromVector [n] (VS.fromListN n l))
+
+toList1 :: Storable a => XArray '[n] a -> [a]
+toList1 (XArray arr) = S.toList arr
+
+-- | Throws if the given shape is not, in fact, empty.
+empty :: forall sh a. Storable a => IShX sh -> XArray sh a
+empty sh
+ | Dict <- lemKnownNatRank sh
+ = XArray (S.constant (shxToList sh)
+ (error "Data.Array.Mixed.empty: shape was not empty"))
+
+slice :: SNat i -> SNat n -> XArray (Just (i + n + k) : sh) a -> XArray (Just n : sh) a
+slice i n (XArray arr) = XArray (S.slice [(fromSNat' i, fromSNat' n)] arr)
+
+sliceU :: Int -> Int -> XArray (Nothing : sh) a -> XArray (Nothing : sh) a
+sliceU i n (XArray arr) = XArray (S.slice [(i, n)] arr)
+
+rev1 :: XArray (n : sh) a -> XArray (n : sh) a
+rev1 (XArray arr) = XArray (S.rev [0] arr)
+
+-- | Throws if the given array and the target shape do not have the same number of elements.
+reshape :: forall sh1 sh2 a. Storable a => StaticShX sh1 -> IShX sh2 -> XArray sh1 a -> XArray sh2 a
+reshape ssh1 sh2 (XArray arr)
+ | Dict <- lemKnownNatRankSSX ssh1
+ , Dict <- lemKnownNatRank sh2
+ = XArray (S.reshape (shxToList sh2) arr)
+
+-- | Throws if the given array and the target shape do not have the same number of elements.
+reshapePartial :: forall sh1 sh2 sh' a. Storable a => StaticShX sh1 -> StaticShX sh' -> IShX sh2 -> XArray (sh1 ++ sh') a -> XArray (sh2 ++ sh') a
+reshapePartial ssh1 ssh' sh2 (XArray arr)
+ | Dict <- lemKnownNatRankSSX (ssxAppend ssh1 ssh')
+ , Dict <- lemKnownNatRankSSX (ssxAppend (ssxFromShape sh2) ssh')
+ = XArray (S.reshape (shxToList sh2 ++ drop (ssxLength ssh1) (S.shapeL arr)) arr)
+
+-- this was benchmarked to be (slightly) faster than S.iota, S.generate and S.fromVector(VS.enumFromTo).
+iota :: (Enum a, Storable a) => SNat n -> XArray '[Just n] a
+iota sn = XArray (S.fromVector [fromSNat' sn] (VS.fromListN (fromSNat' sn) [toEnum 0 .. toEnum (fromSNat' sn - 1)]))