module Multilinear.Generic.MultiCore (
Tensor(..),
(!), isScalar, isSimple, isFiniteTensor,
tensorIndex, _standardize, _mergeScalars,
_map, _contractedIndices, _elemByElem, zipT,
(.+), (.-), (.*), (+.), (-.), (*.),
Multilinear.Generic.MultiCore.map,
Multilinear.Generic.MultiCore.filter,
Multilinear.Generic.MultiCore.filterIndex,
Multilinear.Generic.MultiCore.zipWith
) where
import Control.DeepSeq
import qualified Control.Parallel.Strategies as Parallel
import Data.Foldable
import Data.List
import Data.Maybe
import qualified Data.Set as Set
import qualified Data.Vector as Boxed
import qualified Data.Vector.Unboxed as Unboxed
import GHC.Generics
import Multilinear.Class as Multilinear
import qualified Multilinear.Index as Index
import qualified Multilinear.Index.Finite as Finite
incompatibleTypes :: String
incompatibleTypes = "Incompatible tensor types!"
scalarIndices :: String
scalarIndices = "Scalar has no indices!"
indexNotFound :: String
indexNotFound = "This tensor has not such index!"
tensorOfScalars :: String
tensorOfScalars = "Tensor construction error! Vector of scalars"
data Tensor a where
Scalar :: {
scalarVal :: a
} -> Tensor a
SimpleFinite :: {
tensorFiniteIndex :: Finite.Index,
tensorScalars :: Unboxed.Vector a
} -> Tensor a
FiniteTensor :: {
tensorFiniteIndex :: Finite.Index,
tensorsFinite :: Boxed.Vector (Tensor a)
} -> Tensor a
deriving (Eq, Generic)
{-# INLINE isScalar #-}
isScalar :: Unboxed.Unbox a => Tensor a -> Bool
isScalar x = case x of
Scalar _ -> True
_ -> False
{-# INLINE isSimple #-}
isSimple :: Unboxed.Unbox a => Tensor a -> Bool
isSimple x = case x of
SimpleFinite _ _ -> True
_ -> False
{-# INLINE isFiniteTensor #-}
isFiniteTensor :: Unboxed.Unbox a => Tensor a -> Bool
isFiniteTensor x = case x of
FiniteTensor _ _ -> True
_ -> False
{-# INLINE tensorIndex #-}
tensorIndex :: Unboxed.Unbox a => Tensor a -> Index.TIndex
tensorIndex x = case x of
Scalar _ -> error scalarIndices
SimpleFinite i _ -> Index.toTIndex i
FiniteTensor i _ -> Index.toTIndex i
{-# INLINE firstTensor #-}
firstTensor :: Unboxed.Unbox a => Tensor a -> Tensor a
firstTensor x = case x of
FiniteTensor _ ts -> Boxed.head ts
_ -> x
{-# INLINE (!) #-}
(!) :: Unboxed.Unbox a => Tensor a
-> Int
-> Tensor a
t ! i = case t of
Scalar _ -> error scalarIndices
SimpleFinite ind ts -> Scalar $ ts Unboxed.! (i `mod` Finite.indexSize ind)
FiniteTensor ind ts -> ts Boxed.! (i `mod` Finite.indexSize ind)
instance NFData a => NFData (Tensor a)
{-# INLINE _standardize #-}
_standardize :: (Unboxed.Unbox a) => Tensor a -> Tensor a
_standardize tens = foldl' (<<<|) tens $ Index.indexName <$> (Index.isContravariant `Prelude.filter` indices tens)
instance (
Unboxed.Unbox a, Show a
) => Show (Tensor a) where
show = show' . _standardize
where
show' x = case x of
Scalar v -> show v
SimpleFinite index ts -> show index ++ "S: " ++ case index of
Finite.Contravariant _ _ -> "\n" ++ tail (Unboxed.foldl' (\string e -> string ++ "\n |" ++ show e) "" ts)
_ -> "[" ++ tail (Unboxed.foldl' (\string e -> string ++ "," ++ show e) "" ts) ++ "]"
FiniteTensor index ts -> show index ++ "T: " ++ case index of
Finite.Contravariant _ _ -> "\n" ++ tail (Boxed.foldl' (\string e -> string ++ "\n |" ++ show e) "" ts)
_ -> "[" ++ tail (Boxed.foldl' (\string e -> string ++ "," ++ show e) "" ts) ++ "]"
{-# INLINE _mergeScalars #-}
_mergeScalars :: Unboxed.Unbox a => Tensor a -> Tensor a
_mergeScalars x = case x of
(FiniteTensor index1 ts1) -> case ts1 Boxed.! 0 of
Scalar _ -> SimpleFinite index1 $ Unboxed.generate (Boxed.length ts1) (\i -> scalarVal (ts1 Boxed.! i))
_ -> FiniteTensor index1 $ _mergeScalars <$> ts1
_ -> x
_map :: (
Unboxed.Unbox a, Unboxed.Unbox b, NFData b
) => (a -> b)
-> Tensor a
-> Tensor b
_map f x = case x of
Scalar v -> Scalar $ f v
SimpleFinite index ts -> SimpleFinite index (f `Unboxed.map` ts)
FiniteTensor index ts ->
let len = Boxed.length ts
lts = Boxed.toList $ _map f <$> ts
ltsp = lts `Parallel.using` Parallel.parListChunk (len `div` 8) Parallel.rdeepseq
in FiniteTensor index $ Boxed.fromList ltsp
_transpose :: Boxed.Vector (Boxed.Vector a)
-> Boxed.Vector (Boxed.Vector a)
_transpose v =
let outerS = Boxed.length v
innerS = Boxed.length $ v Boxed.! 0
in Boxed.generate innerS (\i -> Boxed.generate outerS $ \j -> v Boxed.! j Boxed.! i)
_contractedIndices ::
Tensor Double
-> Tensor Double
-> Set.Set String
_contractedIndices t1 t2 =
let iContravariantNames1 = Set.fromList $ Index.indexName <$> (Index.isContravariant `Prelude.filter` indices t1)
iCovariantNames1 = Set.fromList $ Index.indexName <$> (Index.isCovariant `Prelude.filter` indices t1)
iContravariantNames2 = Set.fromList $ Index.indexName <$> (Index.isContravariant `Prelude.filter` indices t2)
iCovariantNames2 = Set.fromList $ Index.indexName <$> (Index.isCovariant `Prelude.filter` indices t2)
in
Set.intersection iCovariantNames1 iContravariantNames2 `Set.union`
Set.intersection iContravariantNames1 iCovariantNames2
{-# INLINE _elemByElem' #-}
_elemByElem' :: (Num a, Unboxed.Unbox a, NFData a)
=> Tensor a
-> Tensor a
-> (a -> a -> a)
-> (Tensor a -> Tensor a -> Tensor a)
-> Tensor a
_elemByElem' (Scalar x1) (Scalar x2) f _ = Scalar $ f x1 x2
_elemByElem' (Scalar x) t f _ = (x `f`) `Multilinear.Generic.MultiCore.map` t
_elemByElem' t (Scalar x) f _ = (`f` x) `Multilinear.Generic.MultiCore.map` t
_elemByElem' t1@(SimpleFinite index1 v1) t2@(SimpleFinite index2 _) f op
| Index.indexName index1 == Index.indexName index2 = op t1 t2
| otherwise = FiniteTensor index1 $ Boxed.generate (Unboxed.length v1)
(\i -> (\x -> f x `Multilinear.Generic.MultiCore.map` t2) (v1 Unboxed.! i))
_elemByElem' t1@(FiniteTensor index1 v1) t2@(FiniteTensor index2 v2) f op
| Index.indexName index1 == Index.indexName index2 = op t1 t2
| Index.indexName index1 `Data.List.elem` indicesNames t2 =
let len2 = Finite.indexSize index2
rl = Boxed.toList $ (\x -> _elemByElem' t1 x f op) <$> v2
rlp = rl `Parallel.using` Parallel.parListChunk (len2 `div` 8) Parallel.rdeepseq
in FiniteTensor index2 $ Boxed.fromList rlp
| otherwise =
let len1 = Finite.indexSize index1
rl = Boxed.toList $ (\x -> _elemByElem' x t2 f op) <$> v1
rlp = rl `Parallel.using` Parallel.parListChunk (len1 `div` 8) Parallel.rdeepseq
in FiniteTensor index1 $ Boxed.fromList rlp
_elemByElem' t1@(SimpleFinite index1 _) t2@(FiniteTensor index2 v2) f op
| Index.indexName index1 == Index.indexName index2 = op t1 t2
| otherwise =
let len2 = Finite.indexSize index2
rl = Boxed.toList $ (\x -> _elemByElem' t1 x f op) <$> v2
rlp = rl `Parallel.using` Parallel.parListChunk (len2 `div` 8) Parallel.rdeepseq
in FiniteTensor index2 $ Boxed.fromList rlp
_elemByElem' t1@(FiniteTensor index1 v1) t2@(SimpleFinite index2 _) f op
| Index.indexName index1 == Index.indexName index2 = op t1 t2
| otherwise =
let len1 = Finite.indexSize index1
rl = Boxed.toList $ (\x -> _elemByElem' x t2 f op) <$> v1
rlp = rl `Parallel.using` Parallel.parListChunk (len1 `div` 8) Parallel.rdeepseq
in FiniteTensor index1 $ Boxed.fromList rlp
{-# INLINE _elemByElem #-}
_elemByElem :: (Num a, Unboxed.Unbox a, NFData a)
=> Tensor a
-> Tensor a
-> (a -> a -> a)
-> (Tensor a -> Tensor a -> Tensor a)
-> Tensor a
_elemByElem t1 t2 f op =
let commonIndices =
if indices t1 /= indices t2 then
Data.List.filter (`Data.List.elem` indicesNames t2) $ indicesNames t1
else []
t1' = foldl' (|>>>) t1 commonIndices
t2' = foldl' (|>>>) t2 commonIndices
in _mergeScalars $ _elemByElem' t1' t2' f op
{-# INLINE zipT #-}
zipT :: (
Num a, NFData a, Unboxed.Unbox a
) => (a -> a -> a)
-> Tensor a
-> Tensor a
-> Tensor a
zipT f t1@(SimpleFinite index1 v1) t2@(SimpleFinite index2 v2) =
if index1 == index2 then
SimpleFinite index1 $ Unboxed.zipWith f v1 v2
else dot t1 t2
zipT f t1@(FiniteTensor index1 v1) t2@(FiniteTensor index2 v2) =
if index1 == index2 then let
l1l = Boxed.length v1
l3 = Boxed.toList (Boxed.zipWith (zipT f) v1 v2) `Parallel.using` Parallel.parListChunk (l1l `div` 8) Parallel.rdeepseq
in FiniteTensor index1 $ Boxed.fromList l3
else dot t1 t2
zipT _ _ _ = error $ "zipT: " ++ scalarIndices
{-# INLINE dot #-}
dot :: (Num a, Unboxed.Unbox a, NFData a)
=> Tensor a
-> Tensor a
-> Tensor a
dot (SimpleFinite i1@(Finite.Covariant count1 _) ts1') (SimpleFinite i2@(Finite.Contravariant count2 _) ts2')
| count1 == count2 =
Scalar $ Unboxed.sum $ Unboxed.zipWith (*) ts1' ts2'
| otherwise = contractionErr "simple-simple" (Index.toTIndex i1) (Index.toTIndex i2)
dot (SimpleFinite i1@(Finite.Contravariant count1 _) ts1') (SimpleFinite i2@(Finite.Covariant count2 _) ts2')
| count1 == count2 =
Scalar $ Unboxed.sum $ Unboxed.zipWith (*) ts1' ts2'
| otherwise = contractionErr "simple-simple" (Index.toTIndex i1) (Index.toTIndex i2)
dot t1@(SimpleFinite _ _) t2@(SimpleFinite _ _) = zipT (*) t1 t2
dot (FiniteTensor i1@(Finite.Covariant count1 _) ts1') (FiniteTensor i2@(Finite.Contravariant count2 _) ts2')
| count1 == count2 =
let zipList = Boxed.toList $ Boxed.zipWith (*) ts1' ts2'
zipListPar = zipList `Parallel.using` Parallel.parListChunk (count1 `div` 8) Parallel.rdeepseq
in Boxed.sum $ Boxed.fromList zipListPar
| otherwise = contractionErr "finite-finite" (Index.toTIndex i1) (Index.toTIndex i2)
dot (FiniteTensor i1@(Finite.Contravariant count1 _) ts1') (FiniteTensor i2@(Finite.Covariant count2 _) ts2')
| count1 == count2 =
let zipList = Boxed.toList $ Boxed.zipWith (*) ts1' ts2'
zipListPar = zipList `Parallel.using` Parallel.parListChunk (count1 `div` 8) Parallel.rdeepseq
in Boxed.sum $ Boxed.fromList zipListPar
| otherwise = contractionErr "finite-finite" (Index.toTIndex i1) (Index.toTIndex i2)
dot t1@(FiniteTensor _ _) t2@(FiniteTensor _ _) = zipT (*) t1 t2
dot t1' t2' = contractionErr "other" (tensorIndex t1') (tensorIndex t2')
{-# INLINE contractionErr #-}
contractionErr :: String
-> Index.TIndex
-> Index.TIndex
-> Tensor a
contractionErr variant i1' i2' = error $
"Tensor product: " ++ variant ++ " - " ++ incompatibleTypes ++
" - index1 is " ++ show i1' ++
" and index2 is " ++ show i2'
{-# INLINE zipErr #-}
zipErr :: String
-> Index.TIndex
-> Index.TIndex
-> Tensor a
zipErr variant i1' i2' = error $
"zipT: " ++ variant ++ " - " ++ incompatibleTypes ++
" - index1 is " ++ show i1' ++
" and index2 is " ++ show i2'
instance (Unboxed.Unbox a, Num a, NFData a) => Num (Tensor a) where
{-# INLINE (+) #-}
t1 + t2 = _elemByElem t1 t2 (+) $ zipT (+)
{-# INLINE (-) #-}
t1 - t2 = _elemByElem t1 t2 (-) $ zipT (-)
{-# INLINE (*) #-}
t1 * t2 = _elemByElem t1 t2 (*) dot
{-# INLINE abs #-}
abs t = abs `Multilinear.Generic.MultiCore.map` t
{-# INLINE signum #-}
signum t = signum `Multilinear.Generic.MultiCore.map` t
{-# INLINE fromInteger #-}
fromInteger x = Scalar $ fromInteger x
instance (Unboxed.Unbox a, Fractional a, NFData a) => Fractional (Tensor a) where
{-# INLINE (/) #-}
_ / _ = error "TODO"
{-# INLINE fromRational #-}
fromRational x = Scalar $ fromRational x
instance (Unboxed.Unbox a, Floating a, NFData a) => Floating (Tensor a) where
{-# INLINE pi #-}
pi = Scalar pi
{-# INLINE exp #-}
exp t = exp `Multilinear.Generic.MultiCore.map` t
{-# INLINE log #-}
log t = log `Multilinear.Generic.MultiCore.map` t
{-# INLINE sin #-}
sin t = sin `Multilinear.Generic.MultiCore.map` t
{-# INLINE cos #-}
cos t = cos `Multilinear.Generic.MultiCore.map` t
{-# INLINE asin #-}
asin t = asin `Multilinear.Generic.MultiCore.map` t
{-# INLINE acos #-}
acos t = acos `Multilinear.Generic.MultiCore.map` t
{-# INLINE atan #-}
atan t = atan `Multilinear.Generic.MultiCore.map` t
{-# INLINE sinh #-}
sinh t = sinh `Multilinear.Generic.MultiCore.map` t
{-# INLINE cosh #-}
cosh t = cosh `Multilinear.Generic.MultiCore.map` t
{-# INLINE asinh #-}
asinh t = acosh `Multilinear.Generic.MultiCore.map` t
{-# INLINE acosh #-}
acosh t = acosh `Multilinear.Generic.MultiCore.map` t
{-# INLINE atanh #-}
atanh t = atanh `Multilinear.Generic.MultiCore.map` t
instance (Unboxed.Unbox a) => Multilinear Tensor a where
fromIndices [u] [] [s] [] f =
SimpleFinite (Finite.Contravariant s [u]) $ Unboxed.generate s $ \x -> f [x] []
fromIndices [] [d] [] [s] f =
SimpleFinite (Finite.Covariant s [d]) $ Unboxed.generate s $ \x -> f [] [x]
fromIndices (u:us) d (s:size) dsize f =
FiniteTensor (Finite.Contravariant s [u]) $ Boxed.generate s (\x -> fromIndices us d size dsize (\uss dss -> f (x:uss) dss) )
fromIndices u (d:ds) usize (s:size) f =
FiniteTensor (Finite.Covariant s [d]) $ Boxed.generate s (\x -> fromIndices u ds usize size (\uss dss -> f uss (x:dss)) )
fromIndices _ _ _ _ _ = error "Indices and its sizes incompatible!"
{-# INLINE el #-}
el (Scalar x) _ = Scalar x
el t1@(SimpleFinite index1 _) (inds,vals) =
let indvals = zip inds vals
val = Data.List.find (\(n,_) -> [n] == Index.indexName index1) indvals
in if isJust val
then t1 ! snd (fromJust val)
else t1
el t1@(FiniteTensor index1 v1) (inds,vals) =
let indvals = zip inds vals
val = Data.List.find (\(n,_) -> [n] == Index.indexName index1) indvals
indvals1 = Data.List.filter (\(n,_) -> [n] /= Index.indexName index1) indvals
inds1 = Data.List.map fst indvals1
vals1 = Data.List.map snd indvals1
in if isJust val
then el (t1 ! snd (fromJust val)) (inds1,vals1)
else FiniteTensor index1 $ (\t -> el t (inds,vals)) <$> v1
{-# INLINE indices #-}
indices x = case x of
Scalar _ -> []
FiniteTensor i ts -> Index.toTIndex i : indices (head $ toList ts)
SimpleFinite i _ -> [Index.toTIndex i]
{-# INLINE order #-}
order x = case x of
Scalar _ -> (0,0)
SimpleFinite index _ -> case index of
Finite.Contravariant _ _ -> (1,0)
Finite.Covariant _ _ -> (0,1)
_ -> let (cnvr, covr) = order $ firstTensor x
in case tensorIndex x of
Index.Contravariant _ _ -> (cnvr+1,covr)
Index.Covariant _ _ -> (cnvr,covr+1)
{-# INLINE size #-}
size t iname = case t of
Scalar _ -> error scalarIndices
SimpleFinite index _ ->
if Index.indexName index == iname
then Finite.indexSize index
else error indexNotFound
FiniteTensor index _ ->
if Index.indexName index == iname
then Finite.indexSize index
else size (firstTensor t) iname
{-# INLINE ($|) #-}
Scalar x $| _ = Scalar x
SimpleFinite (Finite.Contravariant isize _) ts $| (u:_, _) = SimpleFinite (Finite.Contravariant isize [u]) ts
SimpleFinite (Finite.Covariant isize _) ts $| (_, d:_) = SimpleFinite (Finite.Covariant isize [d]) ts
FiniteTensor (Finite.Contravariant isize _) ts $| (u:us, ds) = FiniteTensor (Finite.Contravariant isize [u]) $ ($| (us,ds)) <$> ts
FiniteTensor (Finite.Covariant isize _) ts $| (us, d:ds) = FiniteTensor (Finite.Covariant isize [d]) $ ($| (us,ds)) <$> ts
t $| _ = t
{-# INLINE (/\) #-}
Scalar x /\ _ = Scalar x
FiniteTensor index ts /\ n
| Index.indexName index == n =
FiniteTensor (Finite.Contravariant (Finite.indexSize index) n) $ (/\ n) <$> ts
| otherwise =
FiniteTensor index $ (/\ n) <$> ts
t1@(SimpleFinite index ts) /\ n
| Index.indexName index == n =
SimpleFinite (Finite.Contravariant (Finite.indexSize index) n) ts
| otherwise = t1
{-# INLINE (\/) #-}
Scalar x \/ _ = Scalar x
FiniteTensor index ts \/ n
| Index.indexName index == n =
FiniteTensor (Finite.Covariant (Finite.indexSize index) n) $ (\/ n) <$> ts
| otherwise =
FiniteTensor index $ (\/ n) <$> ts
t1@(SimpleFinite index ts) \/ n
| Index.indexName index == n =
SimpleFinite (Finite.Covariant (Finite.indexSize index) n) ts
| otherwise = t1
{-# INLINE transpose #-}
transpose (Scalar x) = Scalar x
transpose (FiniteTensor (Finite.Covariant count name) ts) =
FiniteTensor (Finite.Contravariant count name) (Multilinear.transpose <$> ts)
transpose (FiniteTensor (Finite.Contravariant count name) ts) =
FiniteTensor (Finite.Covariant count name) (Multilinear.transpose <$> ts)
transpose (SimpleFinite (Finite.Covariant count name) ts) =
SimpleFinite (Finite.Contravariant count name) ts
transpose (SimpleFinite (Finite.Contravariant count name) ts) =
SimpleFinite (Finite.Covariant count name) ts
Scalar x `shiftRight` _ = Scalar x
t1@(SimpleFinite _ _) `shiftRight` _ = t1
t1@(FiniteTensor index1 ts1) `shiftRight` ind
| Data.List.length (indicesNames t1) > 1 && Index.indexName index1 /= ind =
FiniteTensor index1 $ (|>> ind) <$> ts1
| Data.List.length (indicesNames t1) > 1 && Index.indexName index1 == ind =
let index2 = tensorFiniteIndex (ts1 Boxed.! 0)
dane = if isSimple (ts1 Boxed.! 0)
then (\un -> Boxed.generate (Unboxed.length un) (\i -> Scalar $ un Unboxed.! i)) <$>
(tensorScalars <$> ts1)
else tensorsFinite <$> ts1
result = FiniteTensor index2 $ FiniteTensor index1 <$> (_transpose dane)
in _mergeScalars result
| otherwise = t1
{-# INLINE (.+) #-}
(.+) :: (
Unboxed.Unbox a, Num a, NFData a
) => Tensor a
-> a
-> Tensor a
t .+ x = (+x) `Multilinear.Generic.MultiCore.map` t
{-# INLINE (.-) #-}
(.-) :: (
Unboxed.Unbox a, Num a, NFData a
) => Tensor a
-> a
-> Tensor a
t .- x = (\p -> p - x) `Multilinear.Generic.MultiCore.map` t
{-# INLINE (.*) #-}
(.*) :: (
Unboxed.Unbox a, Num a, NFData a
) => Tensor a
-> a
-> Tensor a
t .* x = (*x) `Multilinear.Generic.MultiCore.map` t
{-# INLINE (+.) #-}
(+.) :: (
Unboxed.Unbox a, Num a, NFData a
) => a
-> Tensor a
-> Tensor a
x +. t = (x+) `Multilinear.Generic.MultiCore.map` t
{-# INLINE (-.) #-}
(-.) :: (
Unboxed.Unbox a, Num a, NFData a
) => a
-> Tensor a
-> Tensor a
x -. t = (x-) `Multilinear.Generic.MultiCore.map` t
{-# INLINE (*.) #-}
(*.) :: (
Unboxed.Unbox a, Num a, NFData a
) => a
-> Tensor a
-> Tensor a
x *. t = (x*) `Multilinear.Generic.MultiCore.map` t
{-# INLINE map #-}
map :: (
Unboxed.Unbox a, Unboxed.Unbox b, NFData b
) => (a -> b)
-> Tensor a
-> Tensor b
map f x = case x of
Scalar v -> Scalar $ f v
SimpleFinite index ts -> SimpleFinite index (f `Unboxed.map` ts)
FiniteTensor index ts ->
let len = Boxed.length ts
lts = Boxed.toList $ Multilinear.Generic.MultiCore.map f <$> ts
ltsp = lts `Parallel.using` Parallel.parListChunk (len `div` 8) Parallel.rdeepseq
in FiniteTensor index $ Boxed.fromList ltsp
{-# INLINE filter #-}
filter :: (
Unboxed.Unbox a
) => (String -> Int -> Bool)
-> Tensor a
-> Tensor a
filter _ (Scalar x) = Scalar x
filter f (SimpleFinite index ts) =
let iname = Finite.indexName' index
ts' = (\i _ -> f iname i) `Unboxed.ifilter` ts
in SimpleFinite index { Finite.indexSize = Unboxed.length ts' } ts'
filter f (FiniteTensor index ts) =
let iname = Finite.indexName' index
ts' = Multilinear.Generic.MultiCore.filter f <$> ((\i _ -> f iname i) `Boxed.ifilter` ts)
ts'' =
(\case
(SimpleFinite _ ts) -> not $ Unboxed.null ts
(FiniteTensor _ ts) -> not $ Boxed.null ts
_ -> error $ "Filter: " ++ tensorOfScalars
) `Boxed.filter` ts'
in FiniteTensor index { Finite.indexSize = Boxed.length ts'' } ts''
{-# INLINE filterIndex #-}
filterIndex :: (
Unboxed.Unbox a
) => String
-> (Int -> Bool)
-> Tensor a
-> Tensor a
filterIndex iname f = Multilinear.Generic.MultiCore.filter (\i n -> i /= iname || f n)
{-# INLINE zipWith' #-}
zipWith' :: (
Unboxed.Unbox a, Unboxed.Unbox b, Unboxed.Unbox c, NFData c
) => (a -> b -> c)
-> Tensor a
-> Tensor b
-> Tensor c
zipWith' f (Scalar x1) (Scalar x2) = Scalar $ f x1 x2
zipWith' f t (Scalar x) = (`f` x) `_map` t
zipWith' f (Scalar x) t = (x `f`) `_map` t
zipWith' f (SimpleFinite index1 v1) (SimpleFinite index2 v2) =
if index1 == index2 then
SimpleFinite index1 $ Unboxed.zipWith f v1 v2
else zipErr "simple-simple" (Index.toTIndex index1) (Index.toTIndex index2)
zipWith' f (FiniteTensor index1 v1) (FiniteTensor index2 v2) =
if index1 == index2 then
FiniteTensor index1 $ Boxed.zipWith (Multilinear.Generic.MultiCore.zipWith f) v1 v2
else zipErr "finite-finite" (Index.toTIndex index1) (Index.toTIndex index2)
zipWith' _ _ _ = error "Invalid indices to peroform zip!"
{-# INLINE zipWith #-}
zipWith :: (
Unboxed.Unbox a, Unboxed.Unbox b, Unboxed.Unbox c, NFData c
) => (a -> b -> c)
-> Tensor a
-> Tensor b
-> Tensor c
zipWith f t1 t2 =
let t1' = _standardize t1
t2' = _standardize t2
in zipWith' f t1' t2'