safe-tensor-0.1.0.0: src/Math/Tensor/Safe.hs
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE NoStarIsType #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE StandaloneDeriving #-}
-----------------------------------------------------------------------------
{-|
Module : Math.Tensor.Safe
Description : Dependently typed tensor algebra.
Copyright : (c) Nils Alex, 2020
License : MIT
Maintainer : nils.alex@fau.de
Stability : experimental
Dependently typed tensor algebra.
-}
-----------------------------------------------------------------------------
module Math.Tensor.Safe
( -- * The Tensor GADT
Tensor(..)
-- * Generic Rank of a Tensor
-- |A vector space is the product of a label and a dimension.
, VSpace(..)
-- |The generic tensor rank is a list of vector spaces with label, dimension and
-- associated index list.
, GRank
-- |The rank of a tensor is a generic rank specialized to 'Symbol' and 'Nat'
, Rank
-- * Length-typed assocs lists
-- |Type-level naturals used internally.
, N(..)
, -- |Length-typed vector used internally.
Vec(..)
, vecFromListUnsafe
, -- * Conversion from and to lists
fromList
, fromList'
, toList
, -- * Tensor algebra
(&+), (&-), (&*), removeZeros
, -- * Contraction
contract
, -- * Transpositions
transpose
, transposeMult
, -- * Relabelling
relabel
) where
import Math.Tensor.Safe.TH
import Math.Tensor.Safe.Proofs
import Math.Tensor.Safe.Vector
import Data.Kind (Type)
import Data.Constraint (Dict(Dict), (:-)(Sub))
import Data.Singletons
( Sing
, SingI (sing)
, withSingI, fromSing
)
import Data.Singletons.Prelude
( SBool (STrue, SFalse)
, SList (SNil)
, SMaybe (SJust)
, SOrdering (SLT, SEQ, SGT)
, STuple2 (STuple2)
, Tail
, sFst, sSnd, sHead, sTail
, sCompare, (%==)
)
import Data.Singletons.Prelude.Maybe
( IsJust
, sIsJust
)
import Data.Singletons.Decide
( Decision (Proved, Disproved)
, (:~:) (Refl)
, (%~)
)
import Data.Singletons.TypeLits (Nat, Symbol)
import Data.Maybe (catMaybes)
import Data.Bifunctor (first,second)
import Data.List (foldl',groupBy,sortBy)
-- |The 'Tensor' type is parameterized by its generalized 'Rank' @r@ and holds
-- arbitrary values @v@.
data Tensor :: Rank -> Type -> Type where
ZeroTensor :: forall (r :: Rank) v . Sane r ~ 'True => Tensor r v -- ^
-- A tensor of any sane rank type can be zero.
Scalar :: forall v. !v -> Tensor '[] v -- ^
-- A tensor of empty rank is a scalar holding some value.
Tensor :: forall (r :: Rank) (r' :: Rank) v.
(Sane r ~ 'True, Tail' r ~ r') =>
[(Int, Tensor r' v)] -> Tensor r v -- ^
-- A non-zero tensor of sane non-empty rank is represented as an assocs list of
-- component-value pairs. The keys must be unique and in ascending order.
-- The values are tensors of the next-lower rank.
deriving instance Eq v => Eq (Tensor r v)
deriving instance Show v => Show (Tensor r v)
instance Functor (Tensor r) where
fmap _ ZeroTensor = ZeroTensor
fmap f (Scalar s) = Scalar $ f s
fmap f (Tensor ms) = Tensor $ fmap (fmap (fmap f)) ms
-- |Union of assocs lists with a merging function if a component is present in both lists
-- and two functions to treat components only present in either list.
unionWith :: (a -> b -> c) -> (a -> c) -> (b -> c) -> [(Int, a)] -> [(Int, b)] -> [(Int, c)]
unionWith _ _ f [] ys = fmap (fmap f) ys
unionWith _ f _ xs [] = fmap (fmap f) xs
unionWith f g h xs@((ix,vx):xs') ys@((iy,vy):ys') =
case ix `compare` iy of
LT -> (ix, g vx) : unionWith f g h xs' ys
EQ -> (ix, f vx vy) : unionWith f g h xs' ys'
GT -> (iy, h vy) : unionWith f g h xs ys'
-- |Given a 'Num' and 'Eq' instance, remove all zero values from the tensor,
-- eventually replacing a zero @Scalar@ or an empty @Tensor@ with @ZeroTensor@.
removeZeros :: (Num v, Eq v) => Tensor r v -> Tensor r v
removeZeros ZeroTensor = ZeroTensor
removeZeros (Scalar s) = if s == 0 then ZeroTensor else Scalar s
removeZeros (Tensor ms) =
case ms' of
[] -> ZeroTensor
_ -> Tensor ms'
where
ms' = filter
(\(_, t) ->
case t of
ZeroTensor -> False
_ -> True) $
fmap (fmap removeZeros) ms
-- |Tensor addition. Ranks of summands and sum coincide.
-- Zero values are removed from the result.
(&+) :: forall (r :: Rank) (r' :: Rank) v.
((r ~ r'), Num v, Eq v) =>
Tensor r v -> Tensor r' v -> Tensor r v
(&+) ZeroTensor t = t
(&+) t ZeroTensor = t
(&+) (Scalar s) (Scalar s') =
if s'' == 0 then ZeroTensor else Scalar s''
where
s'' = s + s'
(&+) (Tensor xs) (Tensor xs') = removeZeros $ Tensor xs''
where
xs'' = unionWith (&+) id id xs xs'
(&+) _ _ = error "Cannot add scalar and tensor! Should have been caught by the type system!"
infixl 6 &+
-- |Tensor subtraction. Ranks of operands and difference coincide.
-- Zero values are removed from the result.
(&-) :: forall (r :: Rank) (r' :: Rank) v.
((r ~ r'), Num v, Eq v) =>
Tensor r v -> Tensor r' v -> Tensor r v
(&-) t1 t2 = t1 &+ fmap negate t2
infixl 6 &-
-- |Tensor multiplication, ranks of factors passed explicitly as singletons.
mult :: forall (r :: Rank) (r' :: Rank) (r'' :: Rank) v.
(Num v, 'Just r'' ~ MergeR r r') =>
Sing r -> Sing r' -> Tensor r v -> Tensor r' v -> Tensor r'' v
mult _ _ (Scalar s) (Scalar t) = Scalar (s*t)
mult sr sr' (Scalar s) (Tensor ms) =
case saneTail'Proof sr' of
Sub Dict -> Tensor $ fmap (fmap (mult sr (sTail' sr') (Scalar s))) ms
mult sr sr' (Tensor ms) (Scalar s) =
case saneTail'Proof sr of
Sub Dict -> Tensor $ fmap (fmap (\t -> mult (sTail' sr) sr' t (Scalar s))) ms
mult sr sr' (Tensor ms) (Tensor ms') =
let sh = sHead' sr
sh' = sHead' sr'
st = sTail' sr
st' = sTail' sr'
in case saneMergeRProof sr sr' of
Sub Dict ->
case sh of
STuple2 sv si ->
case sh' of
STuple2 sv' si' ->
case sCompare sv sv' of
SLT -> case proofMergeLT sr sr' of
Sub Dict ->
case saneTail'Proof sr of
Sub Dict -> Tensor $ fmap (fmap (\t -> mult st sr' t (Tensor ms'))) ms
SGT -> case proofMergeGT sr sr' of
Sub Dict ->
case saneTail'Proof sr' of
Sub Dict -> Tensor $ fmap (fmap (mult sr st' (Tensor ms))) ms'
SEQ -> case proofMergeIxNotEQ sr sr' of
Sub Dict ->
case sIxCompare si si' of
SLT -> case proofMergeIxLT sr sr' of
Sub Dict ->
case saneTail'Proof sr of
Sub Dict -> Tensor $ fmap (fmap (\t -> mult st sr' t (Tensor ms'))) ms
SGT -> case proofMergeIxGT sr sr' of
Sub Dict ->
case saneTail'Proof sr' of
Sub Dict -> Tensor $ fmap (fmap (mult sr st' (Tensor ms))) ms'
mult sr sr' ZeroTensor ZeroTensor =
case saneMergeRProof sr sr' of
Sub Dict -> ZeroTensor
mult sr sr' ZeroTensor (Scalar _) =
case saneMergeRProof sr sr' of
Sub Dict -> ZeroTensor
mult sr sr' ZeroTensor (Tensor _) =
case saneMergeRProof sr sr' of
Sub Dict -> ZeroTensor
mult sr sr' (Scalar _) ZeroTensor =
case saneMergeRProof sr sr' of
Sub Dict -> ZeroTensor
mult sr sr' (Tensor _) ZeroTensor =
case saneMergeRProof sr sr' of
Sub Dict -> ZeroTensor
-- |Tensor multiplication. Ranks of factors must not overlap. The Product
-- rank is the merged rank of the factors.
(&*) :: forall (r :: Rank) (r' :: Rank) (r'' :: Rank) v.
(Num v, 'Just r'' ~ MergeR r r', SingI r, SingI r') =>
Tensor r v -> Tensor r' v -> Tensor r'' v
(&*) = mult (sing :: Sing r) (sing :: Sing r')
infixl 7 &*
contract' :: forall (r :: Rank) (r' :: Rank) v.
(r' ~ ContractR r, Num v, Eq v)
=> Sing r -> Tensor r v -> Tensor r' v
contract' sr t = case sContractR sr %~ sr of
Proved Refl -> t
Disproved _ -> contract'' sr t
contract'' :: forall (r :: Rank) (r' :: Rank) v.
(r' ~ ContractR r, Num v, Eq v)
=> Sing r -> Tensor r v -> Tensor r' v
contract'' sr ZeroTensor =
case saneContractProof sr of
Sub Dict -> ZeroTensor
contract'' _ (Scalar v) = Scalar v
contract'' sr (Tensor ms) =
case sTail' sr of
SNil ->
case singletonContractProof sr of
Sub Dict -> Tensor ms
st ->
case saneContractProof sr of
Sub Dict ->
let st' = sTail' st
sh = sHead' sr
sv = sFst sh
si = sSnd sh
sh' = sHead' st
sv' = sFst sh'
si' = sSnd sh'
in case sv %== sv' of
SFalse ->
case contractTailDiffVProof sr of
Sub Dict -> removeZeros $ Tensor $ fmap (fmap (contract'' st)) ms
STrue -> case si of
SICon sa -> case si' of
SICov sb -> case sa %== sb of
STrue ->
let ms' = fmap (\(i, v) -> case v of
Tensor vs ->
case filter (\(i', _) -> i == i') vs of
[] -> Nothing
[(_, v')] -> Just v'
_ -> error "duplicate key in tensor assoc list") ms
ms'' = catMaybes ms' :: [Tensor (Tail' (Tail' r)) v]
in case saneTail'Proof sr of
Sub Dict ->
case saneTail'Proof st of
Sub Dict ->
case contractTailSameVSameIProof sr of
Sub Dict -> contract' st' $ foldl' (&+) ZeroTensor ms''
SFalse ->
case contractTailSameVDiffIProof sr of
Sub Dict -> removeZeros $ Tensor $ fmap (fmap (contract'' st)) ms
SICon _ ->
case contractTailSameVNoCovProof sr of
Sub Dict -> removeZeros $ Tensor $ fmap (fmap (contract'' st)) ms
SICov _ ->
case contractTailSameVNoConProof sr of
Sub Dict -> removeZeros $ Tensor $ fmap (fmap (contract'' st)) ms
-- |Tensor contraction. Contracting a tensor is the identity function on non-contractible tensors.
-- Otherwise, the result is the contracted tensor with the contracted labels removed from the rank.
contract :: forall (r :: Rank) (r' :: Rank) v.
(r' ~ ContractR r, SingI r, Num v, Eq v)
=> Tensor r v -> Tensor r' v
contract = contract' (sing :: Sing r)
-- |Tensor transposition. Given a vector space and two index labels, the result is a tensor with
-- the corresponding entries swapped. Only possible if the indices are part of the rank. The
-- rank remains untouched.
transpose :: forall (vs :: VSpace Symbol Nat) (a :: Ix Symbol) (b :: Ix Symbol) (r :: Rank) v.
(CanTranspose vs a b r ~ 'True, SingI r) =>
Sing vs -> Sing a -> Sing b -> Tensor r v -> Tensor r v
transpose _ _ _ ZeroTensor = ZeroTensor
transpose _ _ _ (Scalar _) = error "This is not possible, might yet have to convince the type system."
transpose v a b t@(Tensor ms) =
case a `sCompare` b of
SEQ -> t
SGT -> case sCanTranspose v b a (sing :: Sing r) %~ STrue of
Proved Refl -> transpose v b a t
SLT ->
let sr = sing :: Sing r
sh = sHead' sr
sv = sFst sh
si = sSnd sh
st = sTail' sr
in withSingI st $
case sv %~ v of
Proved Refl -> case si %~ a of
Proved Refl -> let sr' = sRemoveUntil b sr
in withSingI sr' $
case sSane sr' %~ STrue of
Proved Refl ->
let tl = toTListUntil b t
tl' = fmap (\(i:is, val) -> (last is : (init is ++ [i]),val)) tl
tl'' = sortBy (\(i,_) (i',_) -> i `compare` i') tl'
in fromTList tl''
Disproved _ -> case sCanTranspose v a b st of
STrue -> Tensor $ fmap (fmap (transpose v a b)) ms
Disproved _ -> case sCanTranspose v a b st of
STrue -> Tensor $ fmap (fmap (transpose v a b)) ms
-- |Transposition of multiple labels. Given a vector space and a list of transpositions, the
-- result is a tensor with the corresponding entries swapped. Only possible if the indices are
-- part of the rank. The rank remains untouched.
transposeMult :: forall (vs :: VSpace Symbol Nat) (tl :: TransList Symbol) (r :: Rank) v.
(IsJust (Transpositions vs tl r) ~ 'True, SingI r) =>
Sing vs -> Sing tl -> Tensor r v -> Tensor r v
transposeMult _ _ ZeroTensor = ZeroTensor
transposeMult sv stl tens@(Tensor ms) =
let sr = sing :: Sing r
sh = sHead' sr
st = sTail' sr
sr' = sTail sr
sts = sTranspositions sv stl sr
in case sv %~ sFst sh of
Proved Refl ->
case sSane sr' %~ STrue of
Proved Refl ->
case sts of
SJust sts' ->
withSingI (sFst (sHead sr)) $
withSingI sr' $
let sn = sLengthIL (sSnd (sHead sr))
n = fromSing sn
ts = fromSing sts'
ts' = go ts $ take' n 0
xs = toTListWhile tens
xs' = fmap (first (transposeIndices ts')) xs
xs'' = sortBy (\(i,_) (i',_) -> i `compare` i') xs'
in fromTList xs''
Disproved _ ->
withSingI st $
case sIsJust (sTranspositions sv stl st) %~ STrue of
Proved Refl -> Tensor $ fmap (fmap (transposeMult sv stl)) ms
where
take' :: N -> Int -> [Int]
take' Z i = [i]
take' (S n) i = i : take' n (i+1)
transposeIndices :: [Int] -> [Int] -> [Int]
transposeIndices ts' is = fmap snd $
sortBy (\(i,_) (i',_) -> i `compare` i') $
zip ts' is
go :: [(N,N)] -> [Int] -> [Int]
go [] is = is
go ((s,t):ts) (i:is) =
case s' `compare` i of
EQ -> t' : go ts is
GT -> i : go ((s,t):ts) is
LT -> error $ "illegal permutation" <> show ((s,t):ts) <> "\t" <> show (i:is)
where
s' = toInt s
t' = toInt t
go _ [] = error "cannot transpose elements of empty list"
-- |Tensor relabelling. Given a vector space and a list of relabellings, the result is a tensor
-- with the resulting rank after relabelling. Only possible if labels to be renamed are part of
-- the rank and if uniqueness of labels after relabelling is preserved.
relabel :: forall (vs :: VSpace Symbol Nat) (rl :: RelabelList Symbol) (r1 :: Rank) (r2 :: Rank) v.
(RelabelR vs rl r1 ~ 'Just r2, Sane r2 ~ 'True, SingI r1, SingI r2) =>
Sing vs -> Sing rl -> Tensor r1 v -> Tensor r2 v
relabel _ _ ZeroTensor = ZeroTensor
relabel sv srl tens@(Tensor ms) =
let sr1 = sing :: Sing r1
sr2 = sing :: Sing r2
sh = sHead' sr1
sr1' = sTail' sr1
sr2' = sTail' sr2
sr1'' = sTail sr1
sts = sRelabelTranspositions srl (sSnd (sHead sr1))
in case sv %~ sFst sh of
Proved Refl ->
case sSane sr1'' %~ STrue of
Proved Refl ->
case sts of
SJust sts' ->
withSingI (sFst (sHead sr1)) $
withSingI sr1'' $
let sn = sLengthIL (sSnd (sHead sr1))
n = fromSing sn
ts = fromSing sts'
ts' = go ts $ take' n 0
xs = toTListWhile tens
xs' = fmap (first (transposeIndices ts')) xs
xs'' = sortBy (\(i,_) (i',_) -> i `compare` i') xs'
in fromTList xs''
Disproved _ ->
case sRelabelR sv srl sr1' %~ SJust sr2' of
Proved Refl ->
case sSane sr2' %~ STrue of
Proved Refl -> withSingI sr1' $ withSingI sr2' $ Tensor $ fmap (fmap (relabel sv srl)) ms
where
take' :: N -> Int -> [Int]
take' Z i = [i]
take' (S n) i = i : take' n (i+1)
transposeIndices :: [Int] -> [Int] -> [Int]
transposeIndices ts' is = fmap snd $
sortBy (\(i,_) (i',_) -> i `compare` i') $
zip ts' is
go :: [(N,N)] -> [Int] -> [Int]
go [] is = is
go ((s,t):ts) (i:is) =
case s' `compare` i of
EQ -> t' : go ts is
GT -> i : go ((s,t):ts) is
LT -> error $ "illegal permutation" <> show ((s,t):ts) <> "\t" <> show (i:is)
where
s' = toInt s
t' = toInt t
go _ [] = error "cannot transpose elements of empty list"
-- |Get assocs list from tensor. Keys are length-indexed vectors of indices.
toList :: forall r v n.
(SingI r, SingI n, LengthR r ~ n) =>
Tensor r v -> [(Vec n Int, v)]
toList ZeroTensor = []
toList (Scalar s) = [(VNil, s)]
toList (Tensor ms) =
let st = sTail' (sing :: Sing r)
sn = sing :: Sing n
sm = sLengthR st
in case st of
SNil ->
case sn of
SS SZ -> fmap (\(i, Scalar s) -> (VCons i VNil, s)) ms
_ ->
case sn of
SS sm' ->
withSingI sm' $
case sm %~ sm' of
Proved Refl ->
concatMap (\(i, v) -> case v of Tensor _ -> fmap (first (VCons i)) (withSingI st $ toList v)) ms
fromList' :: forall r v n.
(Sane r ~ 'True, LengthR r ~ n) =>
Sing r -> [(Vec n Int, v)] -> Tensor r v
fromList' _ [] = ZeroTensor
fromList' sr xs =
let sn = sLengthR sr
st = sTail' sr
sm = sLengthR st
in case sn of
SZ ->
case sr %~ SNil of
Proved Refl -> Scalar $ snd (head xs)
SS sm' ->
withSingI sm' $
case sm %~ sm' of
Proved Refl ->
withSingI st $
case sSane st %~ STrue of
Proved Refl ->
case fmap (\(i `VCons` is,v) -> (i,(is ,v))) xs of
xs' -> Tensor $ fmap (fromList' st) <$> myGroup xs'
where
myGroup :: Eq k => [(k,a)] -> [(k, [a])]
myGroup ys =
let ys' = groupBy (\(i,_) (i',_) -> i == i') ys
in fmap (\x -> (fst $ head x, fmap snd x)) ys'
-- |Construct 'Tensor' from assocs list. Keys are length-indexed vectors of indices.
fromList :: forall r v n.
(SingI r, Sane r ~ 'True, LengthR r ~ n) =>
[(Vec n Int, v)] -> Tensor r v
fromList =
let sr = sing :: Sing r
in fromList' sr
-- |Decompose tensor into assocs list with keys being lists of indices for the first vector space
-- and values being the tensors with lower rank for the remaining vector spaces.
toTListWhile :: forall r v.
(SingI r, Sane r ~ 'True) =>
Tensor r v -> [([Int], Tensor (Tail r) v)]
toTListWhile (Tensor ms) =
let sr = sing :: Sing r
st = sTail' sr
in case st %~ sTail sr of
Proved Refl -> fmap (first pure) ms
Disproved _ ->
case sSane st %~ STrue of
Proved Refl ->
case sTail sr %~ sTail st of
Proved Refl ->
withSingI st $
withSingI (sFst (sHead st)) $
let ms' = fmap (second toTListWhile) ms
in concatMap (\(i, xs) -> fmap (first (i :)) xs) ms'
-- |Decompose tensor into assocs list with keys being lists of indices up to and including the
-- desired label, and values being tensors of corresponding lower rank.
toTListUntil :: forall (a :: Ix Symbol) r r' v.
(SingI r, SingI r', RemoveUntil a r ~ r', Sane r ~ 'True, Sane r' ~ 'True) =>
Sing a -> Tensor r v -> [([Int], Tensor r' v)]
toTListUntil sa (Tensor ms) =
let sr = sing :: Sing r
st = sTail' sr
sh = sHead' sr
in case sSnd sh %~ sa of
Proved Refl -> withSingI st $
case st %~ (sing :: Sing r') of
Proved Refl -> fmap (first pure) ms
Disproved _ ->
withSingI st $
case sSane st %~ STrue of
Proved Refl ->
case sRemoveUntil sa st %~ (sing :: Sing r') of
Proved Refl ->
let ms' = fmap (second (toTListUntil sa)) ms
in concatMap (\(i, xs) -> fmap (first (i :)) xs) ms'
-- |Construct tensor from assocs list. Keys are lists of indices, values are
-- tensors of lower rank. Used internally for tensor algebra.
fromTList :: forall r r' v.(Sane r ~ 'True, Sane r' ~ 'True, SingI r, SingI r') =>
[([Int], Tensor r v)] -> Tensor r' v
fromTList [] = ZeroTensor
fromTList xs@((i0,t0):ys)
| null i0 = if null ys
then case (sing :: Sing r) %~ (sing :: Sing r') of
Proved Refl -> t0
else error $ "illegal assocs in fromTList : " ++ show (fmap fst xs)
| otherwise =
let sr' = sing :: Sing r'
st' = sTail' sr'
in withSingI st' $
case sSane st' of
STrue -> Tensor $ fmap (fmap fromTList) xs'''
where
xs' = fmap (\(i:is,v) -> (i,(is,v))) xs
xs'' = groupBy (\(i,_) (i',_) -> i == i') xs'
xs''' = fmap (\x -> (fst $ head x, map snd x)) xs''