linear-1.23.3: src/Linear/V.hs
{-# LANGUAGE CPP #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE EmptyDataDecls #-}
{-# LANGUAGE MultiParamTypeClasses, FlexibleContexts, FlexibleInstances, UndecidableInstances #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE DeriveGeneric #-}
-----------------------------------------------------------------------------
-- |
-- Copyright : (C) 2012-2015 Edward Kmett
-- License : BSD-style (see the file LICENSE)
--
-- Maintainer : Edward Kmett <ekmett@gmail.com>
-- Stability : experimental
-- Portability : non-portable
--
-- n-D Vectors
----------------------------------------------------------------------------
module Linear.V
( V(V,toVector)
#ifdef MIN_VERSION_template_haskell
, int
#endif
, dim
, Dim(..)
, reifyDim
, reifyVector
, reifyDimNat
, reifyVectorNat
, fromVector
, Finite(..)
, _V, _V'
) where
import Control.Applicative
import Control.DeepSeq (NFData)
import Control.Monad
import Control.Monad.Fix
import Control.Monad.Trans.State
import Control.Monad.Zip
import Control.Lens as Lens
import Data.Binary as Binary
import Data.Bytes.Serial
import Data.Complex
import Data.Data
import Data.Distributive
import Data.Foldable as Foldable
import qualified Data.Foldable.WithIndex as WithIndex
import Data.Functor.Bind
import Data.Functor.Classes
import Data.Functor.Rep as Rep
import qualified Data.Functor.WithIndex as WithIndex
import Data.Hashable
import Data.Hashable.Lifted
import Data.Kind
import Data.Reflection as R
import Data.Serialize as Cereal
import qualified Data.Traversable.WithIndex as WithIndex
import qualified Data.Vector as V
import Data.Vector (Vector)
import Data.Vector.Fusion.Util (Box(..))
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Generic.Mutable as M
import Foreign.Ptr
import Foreign.Storable
import GHC.TypeLits
import GHC.Generics (Generic, Generic1)
#if !(MIN_VERSION_reflection(1,3,0)) && defined(MIN_VERSION_template_haskell)
import Language.Haskell.TH
#endif
import Linear.Epsilon
import Linear.Metric
import Linear.Vector
import Prelude as P
#if !(MIN_VERSION_base(4,11,0))
import Data.Semigroup
#endif
import System.Random (Random(..))
import System.Random.Stateful (Uniform(..), UniformRange(..))
class Dim n where
reflectDim :: p n -> Int
type role V nominal representational
class Finite v where
type Size (v :: Type -> Type) :: Nat -- this should allow kind k, for Reifies k Int
toV :: v a -> V (Size v) a
default toV :: Foldable v => v a -> V (Size v) a
toV = V . V.fromList . Foldable.toList
fromV :: V (Size v) a -> v a
instance Finite Complex where
type Size Complex = 2
toV (a :+ b) = V (V.fromListN 2 [a, b])
fromV (V v) = (v V.! 0) :+ (v V.! 1)
_V :: (Finite u, Finite v) => Iso (V (Size u) a) (V (Size v) b) (u a) (v b)
_V = iso fromV toV
_V' :: Finite v => Iso (V (Size v) a) (V (Size v) b) (v a) (v b)
_V' = iso fromV toV
instance Finite (V (n :: Nat)) where
type Size (V n) = n
toV = id
fromV = id
newtype V n a = V { toVector :: V.Vector a } deriving (Eq,Ord,Show,Read,NFData
,Generic,Generic1
)
dim :: forall n a. Dim n => V n a -> Int
dim _ = reflectDim (Proxy :: Proxy n)
{-# INLINE dim #-}
instance KnownNat n => Dim (n :: Nat) where
reflectDim = fromInteger . natVal
{-# INLINE reflectDim #-}
instance (Dim n, Random a) => Random (V n a) where
random = runState (V <$> V.replicateM (reflectDim (Proxy :: Proxy n)) (state random))
randomR (V ls,V hs) = runState (V <$> V.zipWithM (\l h -> state $ randomR (l,h)) ls hs)
instance (Dim n, Uniform a) => Uniform (V n a) where
uniformM g = V <$> V.replicateM (reflectDim (Proxy :: Proxy n)) (uniformM g)
instance (Dim n, UniformRange a) => UniformRange (V n a) where
uniformRM (V ls, V hs) g = V <$> V.zipWithM (\l h -> uniformRM (l, h) g) ls hs
#if (MIN_VERSION_random(1,3,0))
isInRange (V ls, V hs) (V xs) = V.and $ V.zipWith3 (\l h x -> isInRange (l, h) x) ls hs xs
#endif
data ReifiedDim (s :: Type)
retagDim :: (Proxy s -> a) -> proxy (ReifiedDim s) -> a
retagDim f _ = f Proxy
{-# INLINE retagDim #-}
instance Reifies s Int => Dim (ReifiedDim s) where
reflectDim = retagDim reflect
{-# INLINE reflectDim #-}
reifyDimNat :: Int -> (forall (n :: Nat). KnownNat n => Proxy n -> r) -> r
reifyDimNat i f = R.reifyNat (fromIntegral i) f
{-# INLINE reifyDimNat #-}
reifyVectorNat :: forall a r. Vector a -> (forall (n :: Nat). KnownNat n => V n a -> r) -> r
reifyVectorNat v f = reifyNat (fromIntegral $ V.length v) $ \(Proxy :: Proxy n) -> f (V v :: V n a)
{-# INLINE reifyVectorNat #-}
reifyDim :: Int -> (forall (n :: Type). Dim n => Proxy n -> r) -> r
reifyDim i f = R.reify i (go f) where
go :: (Proxy (ReifiedDim n) -> a) -> proxy n -> a
go g _ = g Proxy
{-# INLINE reifyDim #-}
reifyVector :: forall a r. Vector a -> (forall (n :: Type). Dim n => V n a -> r) -> r
reifyVector v f = reifyDim (V.length v) $ \(Proxy :: Proxy n) -> f (V v :: V n a)
{-# INLINE reifyVector #-}
instance Dim n => Dim (V n a) where
reflectDim _ = reflectDim (Proxy :: Proxy n)
{-# INLINE reflectDim #-}
instance (Dim n, Semigroup a) => Semigroup (V n a) where
(<>) = liftA2 (<>)
instance (Dim n, Monoid a) => Monoid (V n a) where
mempty = pure mempty
#if !(MIN_VERSION_base(4,11,0))
mappend = liftA2 mappend
#endif
instance Functor (V n) where
fmap f (V as) = V (fmap f as)
{-# INLINE fmap #-}
instance WithIndex.FunctorWithIndex Int (V n) where
imap f (V as) = V (Lens.imap f as)
{-# INLINE imap #-}
instance Foldable (V n) where
fold (V as) = fold as
{-# INLINE fold #-}
foldMap f (V as) = Foldable.foldMap f as
{-# INLINE foldMap #-}
foldr f z (V as) = V.foldr f z as
{-# INLINE foldr #-}
foldl f z (V as) = V.foldl f z as
{-# INLINE foldl #-}
foldr' f z (V as) = V.foldr' f z as
{-# INLINE foldr' #-}
foldl' f z (V as) = V.foldl' f z as
{-# INLINE foldl' #-}
foldr1 f (V as) = V.foldr1 f as
{-# INLINE foldr1 #-}
foldl1 f (V as) = V.foldl1 f as
{-# INLINE foldl1 #-}
length (V as) = V.length as
{-# INLINE length #-}
null (V as) = V.null as
{-# INLINE null #-}
toList (V as) = V.toList as
{-# INLINE toList #-}
elem a (V as) = V.elem a as
{-# INLINE elem #-}
maximum (V as) = V.maximum as
{-# INLINE maximum #-}
minimum (V as) = V.minimum as
{-# INLINE minimum #-}
sum (V as) = V.sum as
{-# INLINE sum #-}
product (V as) = V.product as
{-# INLINE product #-}
instance WithIndex.FoldableWithIndex Int (V n) where
ifoldMap f (V as) = ifoldMap f as
{-# INLINE ifoldMap #-}
instance Traversable (V n) where
traverse f (V as) = V <$> traverse f as
{-# INLINE traverse #-}
instance WithIndex.TraversableWithIndex Int (V n) where
itraverse f (V as) = V <$> itraverse f as
{-# INLINE itraverse #-}
#if !MIN_VERSION_lens(5,0,0)
instance Lens.FunctorWithIndex Int (V n) where imap = WithIndex.imap
instance Lens.FoldableWithIndex Int (V n) where ifoldMap = WithIndex.ifoldMap
instance Lens.TraversableWithIndex Int (V n) where itraverse = WithIndex.itraverse
#endif
instance Apply (V n) where
V as <.> V bs = V (V.zipWith id as bs)
{-# INLINE (<.>) #-}
instance Dim n => Applicative (V n) where
pure = V . V.replicate (reflectDim (Proxy :: Proxy n))
{-# INLINE pure #-}
V as <*> V bs = V (V.zipWith id as bs)
{-# INLINE (<*>) #-}
instance Bind (V n) where
V as >>- f = V $ V.generate (V.length as) $ \i ->
toVector (f (as `V.unsafeIndex` i)) `V.unsafeIndex` i
{-# INLINE (>>-) #-}
instance Dim n => Monad (V n) where
#if !(MIN_VERSION_base(4,11,0))
return = V . V.replicate (reflectDim (Proxy :: Proxy n))
{-# INLINE return #-}
#endif
V as >>= f = V $ V.generate (reflectDim (Proxy :: Proxy n)) $ \i ->
toVector (f (as `V.unsafeIndex` i)) `V.unsafeIndex` i
{-# INLINE (>>=) #-}
instance Dim n => Additive (V n) where
zero = pure 0
{-# INLINE zero #-}
liftU2 f (V as) (V bs) = V (V.zipWith f as bs)
{-# INLINE liftU2 #-}
liftI2 f (V as) (V bs) = V (V.zipWith f as bs)
{-# INLINE liftI2 #-}
instance (Dim n, Num a) => Num (V n a) where
V as + V bs = V $ V.zipWith (+) as bs
{-# INLINE (+) #-}
V as - V bs = V $ V.zipWith (-) as bs
{-# INLINE (-) #-}
V as * V bs = V $ V.zipWith (*) as bs
{-# INLINE (*) #-}
negate = fmap negate
{-# INLINE negate #-}
abs = fmap abs
{-# INLINE abs #-}
signum = fmap signum
{-# INLINE signum #-}
fromInteger = pure . fromInteger
{-# INLINE fromInteger #-}
instance (Dim n, Fractional a) => Fractional (V n a) where
recip = fmap recip
{-# INLINE recip #-}
V as / V bs = V $ V.zipWith (/) as bs
{-# INLINE (/) #-}
fromRational = pure . fromRational
{-# INLINE fromRational #-}
instance (Dim n, Floating a) => Floating (V n a) where
pi = pure pi
{-# INLINE pi #-}
exp = fmap exp
{-# INLINE exp #-}
sqrt = fmap sqrt
{-# INLINE sqrt #-}
log = fmap log
{-# INLINE log #-}
V as ** V bs = V $ V.zipWith (**) as bs
{-# INLINE (**) #-}
logBase (V as) (V bs) = V $ V.zipWith logBase as bs
{-# INLINE logBase #-}
sin = fmap sin
{-# INLINE sin #-}
tan = fmap tan
{-# INLINE tan #-}
cos = fmap cos
{-# INLINE cos #-}
asin = fmap asin
{-# INLINE asin #-}
atan = fmap atan
{-# INLINE atan #-}
acos = fmap acos
{-# INLINE acos #-}
sinh = fmap sinh
{-# INLINE sinh #-}
tanh = fmap tanh
{-# INLINE tanh #-}
cosh = fmap cosh
{-# INLINE cosh #-}
asinh = fmap asinh
{-# INLINE asinh #-}
atanh = fmap atanh
{-# INLINE atanh #-}
acosh = fmap acosh
{-# INLINE acosh #-}
instance Dim n => Distributive (V n) where
distribute f = V $ V.generate (reflectDim (Proxy :: Proxy n)) $ \i -> fmap (\(V v) -> V.unsafeIndex v i) f
{-# INLINE distribute #-}
instance Hashable a => Hashable (V n a) where
hashWithSalt s0 (V v) =
V.foldl' (\s a -> s `hashWithSalt` a) s0 v
`hashWithSalt` V.length v
instance Dim n => Hashable1 (V n) where
liftHashWithSalt h s0 (V v) =
V.foldl' (\s a -> h s a) s0 v
`hashWithSalt` V.length v
{-# INLINE liftHashWithSalt #-}
instance (Dim n, Storable a) => Storable (V n a) where
sizeOf _ = reflectDim (Proxy :: Proxy n) * sizeOf (undefined:: a)
{-# INLINE sizeOf #-}
alignment _ = alignment (undefined :: a)
{-# INLINE alignment #-}
poke ptr (V xs) = Foldable.forM_ [0..reflectDim (Proxy :: Proxy n)-1] $ \i ->
pokeElemOff ptr' i (V.unsafeIndex xs i)
where ptr' = castPtr ptr
{-# INLINE poke #-}
peek ptr = V <$> V.generateM (reflectDim (Proxy :: Proxy n)) (peekElemOff ptr')
where ptr' = castPtr ptr
{-# INLINE peek #-}
instance (Dim n, Epsilon a) => Epsilon (V n a) where
nearZero = nearZero . quadrance
{-# INLINE nearZero #-}
instance Dim n => Metric (V n) where
dot (V a) (V b) = V.sum $ V.zipWith (*) a b
{-# INLINE dot #-}
-- TODO: instance (Dim n, Ix a) => Ix (V n a)
fromVector :: forall n a. Dim n => Vector a -> Maybe (V n a)
fromVector v
| V.length v == reflectDim (Proxy :: Proxy n) = Just (V v)
| otherwise = Nothing
#if !(MIN_VERSION_reflection(1,3,0)) && defined(MIN_VERSION_template_haskell)
data Z -- 0
data D (n :: *) -- 2n
data SD (n :: *) -- 2n+1
data PD (n :: *) -- 2n-1
instance Reifies Z Int where
reflect _ = 0
{-# INLINE reflect #-}
retagD :: (Proxy n -> a) -> proxy (D n) -> a
retagD f _ = f Proxy
{-# INLINE retagD #-}
retagSD :: (Proxy n -> a) -> proxy (SD n) -> a
retagSD f _ = f Proxy
{-# INLINE retagSD #-}
retagPD :: (Proxy n -> a) -> proxy (PD n) -> a
retagPD f _ = f Proxy
{-# INLINE retagPD #-}
instance Reifies n Int => Reifies (D n) Int where
reflect = (\n -> n+n) <$> retagD reflect
{-# INLINE reflect #-}
instance Reifies n Int => Reifies (SD n) Int where
reflect = (\n -> n+n+1) <$> retagSD reflect
{-# INLINE reflect #-}
instance Reifies n Int => Reifies (PD n) Int where
reflect = (\n -> n+n-1) <$> retagPD reflect
{-# INLINE reflect #-}
-- | This can be used to generate a template haskell splice for a type level version of a given 'int'.
--
-- This does not use GHC TypeLits, instead it generates a numeric type by hand similar to the ones used
-- in the \"Functional Pearl: Implicit Dimurations\" paper by Oleg Kiselyov and Chung-Chieh Shan.
int :: Int -> TypeQ
int n = case quotRem n 2 of
(0, 0) -> conT ''Z
(q,-1) -> conT ''PD `appT` int q
(q, 0) -> conT ''D `appT` int q
(q, 1) -> conT ''SD `appT` int q
_ -> error "ghc is bad at math"
#endif
instance Dim n => Representable (V n) where
type Rep (V n) = Int
tabulate = V . V.generate (reflectDim (Proxy :: Proxy n))
{-# INLINE tabulate #-}
index (V xs) i = xs V.! i
{-# INLINE index #-}
type instance Index (V n a) = Int
type instance IxValue (V n a) = a
instance Ixed (V n a) where
ix i f v@(V as)
| i < 0 || i >= V.length as = pure v
| otherwise = vLens i f v
{-# INLINE ix #-}
instance Dim n => MonadZip (V n) where
mzip (V as) (V bs) = V $ V.zip as bs
mzipWith f (V as) (V bs) = V $ V.zipWith f as bs
instance Dim n => MonadFix (V n) where
mfix f = tabulate $ \r -> let a = Rep.index (f a) r in a
instance Each (V n a) (V n b) a b where
each = traverse
{-# INLINE each #-}
instance (Bounded a, Dim n) => Bounded (V n a) where
minBound = pure minBound
{-# INLINE minBound #-}
maxBound = pure maxBound
{-# INLINE maxBound #-}
vConstr :: Constr
vConstr = mkConstr vDataType "variadic" [] Prefix
{-# NOINLINE vConstr #-}
vDataType :: DataType
vDataType = mkDataType "Linear.V.V" [vConstr]
{-# NOINLINE vDataType #-}
instance (Typeable (V n), Typeable (V n a), Dim n, Data a) => Data (V n a) where
gfoldl f z (V as) = z (V . V.fromList) `f` V.toList as
toConstr _ = vConstr
gunfold k z c = case constrIndex c of
1 -> k (z (V . V.fromList))
_ -> error "gunfold"
dataTypeOf _ = vDataType
dataCast1 f = gcast1 f
instance Dim n => Serial1 (V n) where
serializeWith = traverse_
deserializeWith f = sequenceA $ pure f
instance (Dim n, Serial a) => Serial (V n a) where
serialize = traverse_ serialize
deserialize = sequenceA $ pure deserialize
instance (Dim n, Binary a) => Binary (V n a) where
put = serializeWith Binary.put
get = deserializeWith Binary.get
instance (Dim n, Serialize a) => Serialize (V n a) where
put = serializeWith Cereal.put
get = deserializeWith Cereal.get
instance Eq1 (V n) where
liftEq f0 (V as0) (V bs0) = go f0 (V.toList as0) (V.toList bs0) where
go _ [] [] = True
go f (a:as) (b:bs) = f a b && go f as bs
go _ _ _ = False
instance Ord1 (V n) where
liftCompare f0 (V as0) (V bs0) = go f0 (V.toList as0) (V.toList bs0) where
go f (a:as) (b:bs) = f a b `mappend` go f as bs
go _ [] [] = EQ
go _ _ [] = GT
go _ [] _ = LT
instance Show1 (V n) where
liftShowsPrec _ g d (V as) = showParen (d > 10) $ showString "V " . g (V.toList as)
instance Dim n => Read1 (V n) where
liftReadsPrec _ g d = readParen (d > 10) $ \r ->
[ (V (V.fromList as), r2)
| ("V",r1) <- lex r
, (as, r2) <- g r1
, P.length as == reflectDim (Proxy :: Proxy n)
]
data instance U.Vector (V n a) = V_VN {-# UNPACK #-} !Int !(U.Vector a)
data instance U.MVector s (V n a) = MV_VN {-# UNPACK #-} !Int !(U.MVector s a)
instance (Dim n, U.Unbox a) => U.Unbox (V n a)
instance (Dim n, U.Unbox a) => M.MVector U.MVector (V n a) where
{-# INLINE basicLength #-}
{-# INLINE basicUnsafeSlice #-}
{-# INLINE basicOverlaps #-}
{-# INLINE basicUnsafeNew #-}
{-# INLINE basicUnsafeRead #-}
{-# INLINE basicUnsafeWrite #-}
basicLength (MV_VN n _) = n
basicUnsafeSlice m n (MV_VN _ v) = MV_VN n (M.basicUnsafeSlice (d*m) (d*n) v)
where d = reflectDim (Proxy :: Proxy n)
basicOverlaps (MV_VN _ v) (MV_VN _ u) = M.basicOverlaps v u
basicUnsafeNew n = liftM (MV_VN n) (M.basicUnsafeNew (d*n))
where d = reflectDim (Proxy :: Proxy n)
basicUnsafeRead (MV_VN _ v) i =
liftM V $ V.generateM d (\j -> M.basicUnsafeRead v (d*i+j))
where d = reflectDim (Proxy :: Proxy n)
basicUnsafeWrite (MV_VN _ v0) i (V vn0) = let d0 = V.length vn0 in go v0 vn0 d0 (d0*i) 0
where
go v vn d o j
| j >= d = return ()
| otherwise = do
a <- liftBox $ G.basicUnsafeIndexM vn j
M.basicUnsafeWrite v o a
go v vn d (o+1) (j+1)
basicInitialize (MV_VN _ v) = M.basicInitialize v
{-# INLINE basicInitialize #-}
liftBox :: Monad m => Box a -> m a
liftBox (Box a) = return a
{-# INLINE liftBox #-}
instance (Dim n, U.Unbox a) => G.Vector U.Vector (V n a) where
{-# INLINE basicUnsafeFreeze #-}
{-# INLINE basicUnsafeThaw #-}
{-# INLINE basicLength #-}
{-# INLINE basicUnsafeSlice #-}
{-# INLINE basicUnsafeIndexM #-}
basicUnsafeFreeze (MV_VN n v) = liftM ( V_VN n) (G.basicUnsafeFreeze v)
basicUnsafeThaw ( V_VN n v) = liftM (MV_VN n) (G.basicUnsafeThaw v)
basicLength ( V_VN n _) = n
basicUnsafeSlice m n (V_VN _ v) = V_VN n (G.basicUnsafeSlice (d*m) (d*n) v)
where d = reflectDim (Proxy :: Proxy n)
basicUnsafeIndexM (V_VN _ v) i =
liftM V $ V.generateM d (\j -> G.basicUnsafeIndexM v (d*i+j))
where d = reflectDim (Proxy :: Proxy n)
vLens :: Int -> Lens' (V n a) a
vLens i = \f (V v) -> f (v V.! i) <&> \a -> V (v V.// [(i, a)])
{-# INLINE vLens #-}
instance ( 1 <= n) => Field1 (V n a) (V n a) a a where _1 = vLens 0
instance ( 2 <= n) => Field2 (V n a) (V n a) a a where _2 = vLens 1
instance ( 3 <= n) => Field3 (V n a) (V n a) a a where _3 = vLens 2
instance ( 4 <= n) => Field4 (V n a) (V n a) a a where _4 = vLens 3
instance ( 5 <= n) => Field5 (V n a) (V n a) a a where _5 = vLens 4
instance ( 6 <= n) => Field6 (V n a) (V n a) a a where _6 = vLens 5
instance ( 7 <= n) => Field7 (V n a) (V n a) a a where _7 = vLens 6
instance ( 8 <= n) => Field8 (V n a) (V n a) a a where _8 = vLens 7
instance ( 9 <= n) => Field9 (V n a) (V n a) a a where _9 = vLens 8
instance (10 <= n) => Field10 (V n a) (V n a) a a where _10 = vLens 9
instance (11 <= n) => Field11 (V n a) (V n a) a a where _11 = vLens 10
instance (12 <= n) => Field12 (V n a) (V n a) a a where _12 = vLens 11
instance (13 <= n) => Field13 (V n a) (V n a) a a where _13 = vLens 12
instance (14 <= n) => Field14 (V n a) (V n a) a a where _14 = vLens 13
instance (15 <= n) => Field15 (V n a) (V n a) a a where _15 = vLens 14
instance (16 <= n) => Field16 (V n a) (V n a) a a where _16 = vLens 15
instance (17 <= n) => Field17 (V n a) (V n a) a a where _17 = vLens 16
instance (18 <= n) => Field18 (V n a) (V n a) a a where _18 = vLens 17
instance (19 <= n) => Field19 (V n a) (V n a) a a where _19 = vLens 18