-- |A module defining a recursive multi-dimensional tensor data type and -- arithmetic functions for it. module Tensor ( -- * Data types Tensor(Scalar, Vector), TensorAlgebra, -- * Basic functions tenVerify, isScalar, order, tenDims, -- * Folds, maps etc. tenFold, tenMap, tenZipWith, -- * Creation and export of tensors list2tensor, tensor2list, dim2tensor, dim3tensor, -- * Special tensors kronecker, levicivita, -- * Arithmetic scalarmul, tenAdd, tenBinOp, fullcontract, externalprod, contract, selfcontract, -- * Index rearrangement transpose, swapind, revind, reindex, slice ) where import Text.ParserCombinators.ReadP -- -************************************************************************* -- Data types -- -************************************************************************* -- |The Tensor data type is defined recursively as either a Scalar or an array -- of sub-tensors called Vector. Equal size of all sub-tensors in a Vector is -- assumed by most functions below. Use the tenVerify function to verify -- correctness if you build tensors using constructors. data Tensor n = Scalar n | Vector [Tensor n] -- |Function type for folds. The first function processes Scalars, the second -- combines the results from sub-tensors. type TensorAlgebra n r = (n -> r, [r] -> r) -- -************************************************************************* -- Class instances -- -************************************************************************* instance (Eq n) => Eq (Tensor n) where (==) a b = tenDims a == tenDims b && tenZipWith (uncurry (==), and) a b (/=) a b = tenDims a /= tenDims b || tenZipWith (uncurry (/=), or) a b instance Functor Tensor where fmap = tenMap instance (Show n) => Show (Tensor n) where show (Scalar n) = show n show (Vector l) = "[" ++ (sl l) ++ "]" where sl [] = "" sl (x:xs) | length xs == 0 = show x | otherwise = show x ++ ',' : sl xs instance (Read n) => Read (Tensor n) where readsPrec _ = readP_to_S \$ readS_to_P readsScalar +++ readS_to_P readsVector -- |Parse a scalar and wrap it into a Scalar constructor readsScalar :: Read n => String -> [(Tensor n, String)] -- = ReadS (Tensor n) readsScalar = map (\(n,s) -> (Scalar n,s)) . reads -- |Parse a vector or higher-order tensor. The reads function used will -- recurse via the list read function to the Tensor read function again. -- Because of this, the tensor consistency check is quite wasteful. It would -- suffice to call tenVerify once on the final result, but AFAIK there is no -- hook in the Read class that allows that. readsVector :: Read n => String -> [(Tensor n, String)] -- = ReadS (Tensor n) readsVector = filter (tenVerify . fst) . map (\(l,s) -> (Vector l,s)) . reads -- -************************************************************************* -- Basic tensor functions -- -************************************************************************* -- |Verify correctness of Tensor data. To be used if you build tensors -- yourself from constructors. tenVerify :: Tensor n -> Bool tenVerify = decide . tenFold minmaxsizes where decide = all (uncurry (==)) -- Min/max size for every tensor index. minmaxsizes :: TensorAlgebra n [(Int, Int)] minmaxsizes = (const [], merge) merge [] = [(-1, 0)] -- catch empty Vector merge (x:xs) = (dup \$ length (x:xs)) : foldl (mmZipWith mm) x xs dup x = (x, x) mm (a,b) (c,d) = (min a c, max b d) -- Special version of zipWith which causes minmaxsize to fail for -- inconsistent tensor order. mmZipWith :: (a -> b -> (Int, Int)) -> [a] -> [b] -> [(Int, Int)] mmZipWith _ [] [] = [] mmZipWith f (x:xs) [] = (-1, 0) : mmZipWith f xs [] mmZipWith f [] (y:ys) = (-1, 0) : mmZipWith f [] ys mmZipWith f (x:xs) (y:ys) = f x y : mmZipWith f xs ys -- |Concise check for Scalar isScalar :: Tensor n -> Bool isScalar (Scalar _) = True isScalar (Vector _) = False -- |Tensor order order :: Tensor n -> Int order (Scalar _) = 0 order (Vector (h:r)) = 1 + (order h) order (Vector []) = error "order: hit empty Vector" -- |Tensor dimensions, starting with outer Vector = first index tenDims :: Tensor n -> [Int] tenDims (Scalar _) = [] tenDims (Vector l) = length l : tenDims (head l) -- -************************************************************************* -- Folds, maps etc. -- -************************************************************************* -- |Fold for the Tensor data type. tenFold :: TensorAlgebra n r -> Tensor n -> r tenFold (fs, _) (Scalar s) = fs s tenFold f@(_, fv) (Vector v) = fv \$ map (tenFold f) v -- |Apply a function to all elements of the tensor tenMap :: (a -> b) -> Tensor a -> Tensor b tenMap f = tenFold (Scalar . f, Vector) -- |zipWith for tensors with the same order and dimensions tenZipWith :: TensorAlgebra (a, b) c -> Tensor a -> Tensor b -> c tenZipWith (fs, _) (Scalar a) (Scalar b) = fs (a, b) tenZipWith _ (Vector _) (Scalar _) = error "tenZipWith: order of first argument exceeds that of second" tenZipWith _ (Scalar _) (Vector _) = error "tenZipWith: order of second argument exceeds that of first" tenZipWith f@(_, fv) (Vector a) (Vector b) | null a || null b = error "tenZipWith: hit empty vector" | length a /= length b = error "tenZipWith: mismatching dimensions" | otherwise = fv \$ zipWith (tenZipWith f) a b -- The following functions are not exported. They are used internally for -- rearranging indices and contractions. -- |map for Vectors vecMap :: (Tensor a -> Tensor b) -> Tensor a -> Tensor b vecMap f (Vector l) = Vector \$ map f l -- |map for n-th order Vectors of Vectors. Used for operations on the n-th -- index, possibly after rearrangement. nVecMap :: Int -> (Tensor a -> Tensor b) -> Tensor a -> Tensor b nVecMap n = (!! n) . iterate vecMap -- |zipWith for Vectors vecZipWith :: (Tensor a -> Tensor b -> Tensor c) -> Tensor a -> Tensor b -> Tensor c vecZipWith f (Vector l1) (Vector l2) = Vector \$ zipWith f l1 l2 -- |(zipWith id) for Vectors, for completeness vecZid :: [Tensor a -> Tensor b] -> Tensor a -> Tensor b vecZid fv (Vector l) = Vector \$ zipWith id fv l -- |zipWith for n-th order Vectors of Vectors. Used for contraction of two -- tensors with the n-th index, possibly after rearrangement. nVecZipWith :: Int -> (Tensor a -> Tensor b -> Tensor c) -> Tensor a -> Tensor b -> Tensor c nVecZipWith n = (!! n) . iterate vecZipWith -- |Pull a function returning a pair out by one Vector level. vecUnzipWith :: (Tensor a -> (Tensor b, Tensor c)) -> Tensor a -> (Tensor b, Tensor c) vecUnzipWith f (Vector l) = (Vector \$ fst fl, Vector \$ snd fl) where fl = unzip \$ map f l -- |Pull a function returning a pair out by n Vector levels. nVecUnzipWith :: Int -> (Tensor a -> (Tensor b, Tensor c)) -> Tensor a -> (Tensor b, Tensor c) nVecUnzipWith n = (!! n) . iterate vecUnzipWith -- -************************************************************************* -- Creation and export of tensors -- -************************************************************************* -- |Build Tensor from list of dimensions and list of elements. -- The last tensor index = innermost Vector element changes fastest. list2tensor :: [Int] -> [n] -> Tensor n list2tensor s d | any (<= 0) s = error "list2vector: zero or negative dimension" | head sc /= length d = error "list2vector: size mismatch" | otherwise = l2t (tail sc) d where sc = scanr (*) 1 s l2t :: [Int] -> [n] -> Tensor n l2t [] (x:xs) = Scalar x l2t (c:cs) (x:xs) = Vector \$ map (l2t cs) \$ split c (x:xs) split :: Int -> [n] -> [[n]] split _ [] = [] split n l = take n l : split n (drop n l) -- |Turn tensor into pair of list of dimensions and list of elements tensor2list :: Tensor n -> ([Int], [n]) tensor2list = tenFold t2l where t2l :: TensorAlgebra n ([Int], [n]) t2l = (scalar2l, vectorl2l) scalar2l x = ([], [x]) vectorl2l [] = error "tensor2list: hit empty Vector" vectorl2l l = (length l : fst (head l), concat \$ map snd l) -- |Build 2D tensor from nested lists. Haskell's strong typing forbids -- generalising this. In order to guard against higher-order nested lists -- being passed, this function is restricted to Num's. dim2tensor :: Num n => [[n]] -> Tensor n dim2tensor l | not \$ tenVerify result = error "dim2tensor: inconsistent tensor dimension or empty sub-tensor" | otherwise = result where result = Vector \$ map (Vector . map Scalar) l -- |Build 3D tensor from nested lists. Haskell's strong typing forbids -- generalising this. In order to guard against higher-order nested lists -- being passed, this function is restricted to Num's. dim3tensor :: Num n => [[[n]]] -> Tensor n dim3tensor l | not \$ tenVerify result = error "dim3tensor: inconsistent tensor dimension or empty sub-tensor" | otherwise = result where result = Vector \$ map (Vector . map (Vector . map Scalar)) l -- -************************************************************************* -- Special tensors -- -************************************************************************* -- |Unit matrix or Kronecker symbol in given dimension kronecker :: Num n => Int -> Tensor n kronecker d | d <= 0 = error "kronecker: dimension must be positive" | otherwise = Vector \$ kr 0 where kr e | e == d = [] | otherwise = Vector (map (eq2one e) [0..d-1]) : kr (e+1) eq2one x y | x == y = Scalar 1 | otherwise = Scalar 0 -- |Levi-Civita or fully antisymmetric tensor in d dimensions (and of order d) levicivita :: Num n => Int -> Tensor n levicivita d | d < 2 = error "levicivita: order must be >=2" | otherwise = lcbuilder d [] 0 where -- Build a slice of a Levi-Civita tensor. The index list (second argument) -- is reversed. The third argument is the number of inversions of the -- index list. lcbuilder :: Num n => Int -> [Int] -> Int -> Tensor n lcbuilder d i inv | length i == d = Scalar \$ (-1) ^ inv | otherwise = Vector \$ buildloop d i inv [1..d] buildloop :: Num n => Int -> [Int] -> Int -> [Int] -> [Tensor n] buildloop _ _ _ [] = [] buildloop d i inv (x:xs) = buildsub d i inv x : buildloop d i inv xs buildsub :: Num n => Int -> [Int] -> Int -> Int -> Tensor n buildsub d i inv x | elem x i = lczeros d (d - (length i) - 1) | otherwise = lcbuilder d (x:i) (inv + addinv i x) addinv :: [Int] -> Int -> Int addinv i x = length \$ filter (x<) i lczeros :: Num n => Int -> Int -> Tensor n lczeros _ 0 = Scalar 0 lczeros d o = Vector (replicate d \$ lczeros d \$ o-1) -- -************************************************************************* -- Arithmetic -- -************************************************************************* -- |Multiplication with a (non-wrapped) scalar scalarmul :: Num n => n -> Tensor n -> Tensor n scalarmul n = tenFold (Scalar . (n*), Vector) -- |Add two equally-sized tensors element by element tenAdd :: Num n => Tensor n -> Tensor n -> Tensor n tenAdd = tenZipWith (Scalar . uncurry (+), Vector) -- |Arbitrary element-wise binary operator tenBinOp :: (a -> b -> c) -> Tensor a -> Tensor b -> Tensor c tenBinOp f = tenZipWith (Scalar . uncurry f, Vector) -- |Full contraction with respect to all indices of two tensors with equal -- dimensions. The result is a scalar not wrapped in a Scalar constructor. fullcontract :: Num n => Tensor n -> Tensor n -> n fullcontract = tenZipWith (uncurry (*), sum) -- |External product (no contractions) externalprod :: Num n => Tensor n -> Tensor n -> Tensor n externalprod t1 t2 = tenFold (flip scalarmul t2, Vector) t1 -- |General contraction of two tensors. The first argument contains pairs of -- one-based indices to be contracted. The contraction is performed by -- shifting indices to be contracted to the end, then using fullcontract on -- corresponding inner tensors of the two tensor arguments. contract :: Num n => [(Int,Int)] -> Tensor n -> Tensor n -> Tensor n contract inds x y = nVecMap (ox-li) (\xs -> nVecMap (oy-li) (Scalar . fullcontract xs) y') x' where inds' = unzip \$ sortind \$ checkind inds x' = pushall (fst inds') x y' = pushall (snd inds') y ox = order x oy = order y li = length inds checkind :: [(Int,Int)] -> [(Int,Int)] checkind [] = [] checkind ((i0,i1):it) | i0 < 1 || i0 > ox = error \$ "contract: index " ++ (show i0) ++ " out of range for first operand" | i1 < 1 || i1 > ox = error \$ "contract: index " ++ (show i1) ++ " out of range for second operand" | otherwise = (i0, i1) : checkind it sortind :: [(Int,Int)] -> [(Int,Int)] sortind [] = [] sortind (ph:pt) = inspair ph \$ sortind pt inspair :: (Int,Int) -> [(Int,Int)] -> [(Int,Int)] inspair p [] = [p] inspair p l@(lh:lt) | cmpindpair p lh = p : l | otherwise = lh : inspair p lt cmpindpair :: (Int,Int) -> (Int,Int) -> Bool cmpindpair (i0,i1) (i2,i3) | i0 == i2 = error \$ "contract: duplicate index " ++ (show i0) ++ " for first operand" | i1 == i3 = error \$ "contract: duplicate index " ++ (show i1) ++ " for second operand" | otherwise = i0+i1 > i2+i3 -- |Contract two indices of one tensor. The indices are one-based. selfcontract :: Num n => Int -> Int -> Tensor n -> Tensor n selfcontract a b t | i0 < 0 = error \$ "selfcontract: index number " ++ (show (i0+1)) ++ " out of range" | i1 >= order t = error \$ "selfcontract: index number " ++ (show (i1+1)) ++ " out of range" | otherwise = nVecMap i0 (secon (i1-i0)) t where i0 = (min a b) - 1 i1 = (max a b) - 1 secon 1 v | vlength (vhead v) /= vlength v = error "selfcontract: dimension mismatch" | otherwise = sctrace v secon b v = vecMap (secon (b-1)) \$ transpose v sctrace v | vlength v == 1 = vhead \$ vhead v | otherwise = tenAdd (vhead \$ vhead v) (sctrace \$ vecMap vtail \$ vtail v) -- -************************************************************************* -- Index rearrangement -- -************************************************************************* -- This is the most "interesting" section, as Haskell's hierarchical structures -- do not lend themselves well to this kind of thing. -- |Exchange the first two indices transpose :: Tensor n -> Tensor n transpose v | vnull \$ vhead v = Vector [] | otherwise = vcons (fst headtail) (transpose \$ snd headtail) where headtail = vecUnzipWith vheadtail v -- |Swap two indices. As for all exported functions, these are one-based. swapind :: Int -> Int -> Tensor n -> Tensor n swapind a b t | a == b = t | i0 < 0 = error \$ "swapind: index "++(show \$ i0+1)++" out of range" | i1 >= order t = error \$ "swapind: index "++(show \$ i1+1)++" out of range" | otherwise = nVecMap i0 (swind (i1-i0)) t where i0 = (min a b) - 1 i1 = (max a b) - 1 swind 1 v = transpose v swind b v = pullindex (b-1) \$ pushindex b v -- |Reverse order of indices revind :: Tensor n -> Tensor n revind s@(Scalar _) = s revind v@(Vector vl) = rimerge \$ map revind vl where rimerge :: [Tensor n] -> Tensor n rimerge l | isScalar \$ head l = Vector l | vnull \$ head l = Vector [] | otherwise = vcons (rimerge \$ map vhead l) (rimerge \$ map vtail l) -- |Rearrange indices. The first argument contains a list of index numbers -- that are moved to the front / outside in the given order. The remaining -- indices, if any, are left in their existing order at the back. reindex :: [Int] -> Tensor n -> Tensor n reindex il t = ristep il ord t where ord = order t ristep :: [Int] -> Int -> Tensor n -> Tensor n ristep [] _ t = t ristep il _ (Scalar _) = error "reindex: index list exceeds tensor order" ristep il@(ih:it) o t | ih < 1 || ih > o = error \$ "reindex: index "++(show ih)++" out of range" | otherwise = vecMap (ristep (rmind il) (o-1)) \$ pullindex (ih-1) t -- Remove the first index number from the list, and decrement those that -- exceed it so the list can be used on the sub-tensors. rmind :: [Int] -> [Int] rmind (ih:it) = map (conddec ih) it conddec i x | x < i = x | x == i = error \$ "reindex: duplicate index " ++ (show i) | otherwise = x-1 -- |Slice of a tensor for specific values of some indices. The first argument -- is a list of indices and their values, both one-based. slice :: [(Int,Int)] -> Tensor n -> Tensor n slice [] t = t slice inds t | head il < 1 || il !! (length il - 1) > order t = error "slice: index number out of range" | otherwise = slrec il gl vl t where indsvals = unzip \$ inssort (\(x,y) (z,w) -> x < z) inds il = fst indsvals gl = inds2gaps \$ fst indsvals vl = snd indsvals -- Recurse over indices to evaluate slrec :: [Int] -> [Int] -> [Int] -> Tensor n -> Tensor n slrec [] [] [] t = t slrec (i:is) (g:gs) (v:vs) t = nVecMap g (slrec is gs vs . sleval i v) t -- Evaluate one index by returning the corresponding list element of the -- Vector that is the last argument sleval :: Int -> Int -> Tensor n -> Tensor n sleval i v (Vector l) | v < 1 || v > length l = error \$ "slice: value " ++ (show v) ++ " for index " ++ (show i) ++ " out of range" | otherwise = l !! (v-1) -- The following are used only internally and not exported -- |Shift the topmost index downward by n places. pushindex :: Int -> Tensor n -> Tensor n pushindex n v@(Vector (vh:vt)) | null vt = nVecMap n (\x -> Vector [x]) vh | otherwise = nVecZipWith n vcons vh (pushindex n \$ Vector vt) -- |Shift nth index upward to the top. n is zero-based. pullindex :: Int -> Tensor n -> Tensor n pullindex n v@(Vector _) | vnull \$ (!! n) \$ iterate vhead v = Vector [] | otherwise = vcons (fst headtail) (pullindex n \$ snd headtail) where headtail = nVecUnzipWith n vheadtail v -- |Push some indices to the back in the given order pushall :: [Int] -> Tensor n -> Tensor n pushall [] t = t pushall inds t = pastep gl pl t where ord = order t gl = inds2gaps \$ fst indspushs pl = snd indspushs -- Pair of ascending list of indices and list of push depths, generated -- from index list indspushs :: ([Int],[Int]) indspushs = unzip \$ map fixpush \$ zip [0..] \$ reorder2orig inds -- Permute index numbers into original order after zipping with 0-based -- destination order reorder2orig = inssort (\(x,y) (z,w) -> y < w) . zip [0..] -- Generate pair of original index and push depth from its relative -- position before and after the pushall fixpush :: (Int,(Int,Int)) -> (Int,Int) fixpush (newi, (oldi, ind)) = (ind, ord-ind+(min oldi newi)) -- Recursively descend into the nested tensor by the amount given by the -- gaps between the index numbers, and push them back by the amount -- determined by fixpush pastep :: [Int] -> [Int] -> Tensor n -> Tensor n pastep [] [] t = t pastep (g:gs) (p:ps) t = nVecMap g (pastep gs ps . pushindex p) t -- |Insertion sort, used by pushall and slice to order the indices to work on. -- The first argument is a comparison function which must return true iff its -- first argument is to be sorted before its second. inssort :: (t -> t -> Bool) -> [t] -> [t] inssort _ [] = [] inssort cmp (xh:xt) = isinsert cmp xh \$ inssort cmp xt where isinsert :: (t -> t -> Bool) -> t -> [t] -> [t] isinsert cmp x [] = [x] isinsert cmp x xl@(xh:xt) | cmp x xh = x:xl | otherwise = xh : isinsert cmp x xt -- |Compute the number of indices to skip between indices to do something to. -- Used by slice and pushall. The result has the same number of entries as the -- argument; the first is the first entry of the argument minus one, giving the -- number of initial indices to ignore. inds2gaps :: [Int] -> [Int] inds2gaps [] = [] inds2gaps l@(h:t) = (h-1) : zipWith (\x y -> x-y-1) t l -- Small auxiliary functions - standard list functions for Vectors -- It would be nice to generalise the following auxiliary functions, but bad -- for debugging because errors would always occur in vecDo. vecDo :: ([Tensor a] -> b) -> Tensor a -> b vecDo f (Vector l) = f l vcons :: Tensor n -> Tensor n -> Tensor n vcons a (Vector b) = Vector (a:b) -- vcons a = vecDo (Vector . (a:)) vnull :: Tensor n -> Bool vnull (Vector v) = null v -- vnull = vecDo null vlength :: Tensor n -> Int vlength (Vector v) = length v -- vlength = vecDo length vhead :: Tensor n -> Tensor n vhead (Vector (vh:vt)) = vh -- vhead = vecDo head vtail :: Tensor n -> Tensor n vtail (Vector (vh:vt)) = Vector vt -- vtail = vecDo (Vector . tail) vheadtail :: Tensor n -> (Tensor n, Tensor n) vheadtail (Vector (vh:vt)) = (vh, Vector vt)