working on tensor contractions

4 files changed, 133 insertions, 35 deletions

diff --git a/examples/pru.hs b/examples/pru.hs
index bddc08f..a935d93 100644
--- a/examples/pru.hs
+++ b/examples/pru.hs
@@ -7,7 +7,7 @@ import Data.Packed.Internal.Tensor
 
 import Complex
 import Numeric(showGFloat)
-import Data.List(transpose,intersperse)
+import Data.List(transpose,intersperse,sort)
 import Foreign.Storable
 
 r >< c = f where
@@ -22,8 +22,6 @@ r >|< c = f where
                                  ++show (dim v) ++"in ("++show r++"><"++show c++")"
         where v = fromList l
 
-
-
 vr = fromList [1..15::Double]
 vc = fromList (map (\x->x :+ (x+1)) [1..15::Double])
 
@@ -62,38 +60,56 @@ main = do
 t = T [(4,(Covariant,"p")),(2,(Covariant,"q")),(3,(Contravariant,"r"))] $ fromList [1..24::Double]
 
 
-findIdx name t = ((d1,d2),m) where
-    (d1,d2) = span (\(_,(_,n)) -> n /=name) (dims t)
-    c = product (map fst d2)
-    m = matrixFromVector RowMajor c (ten t)
 
-putFirstIdx name t = (nd,m')
-    where ((d1,d2),m) = findIdx name t
-          m' = matrixFromVector RowMajor c $ cdat $ trans m
-          nd = d2++d1
-          c = dim (ten t) `div` (fst $ head d2)
+t1 = T [(4,(Covariant,"p")),(4,(Contravariant,"q")),(2,(Covariant,"r"))] $ fromList [1..32::Double]
+t2 = T [(4,(Covariant,"p")),(4,(Contravariant,"q"))] $ fromList [1..16::Double]
 
-part t (name,k) = if k<0 || k>=l
-                    then error $ "part "++show (name,k)++" out of range in "++show t
-                    else T {dims = ds, ten = toRows m !! k}
-    where (d:ds,m) = putFirstIdx name t
-          (l,_) = d
 
-parts t name = map f (toRows m) 
-    where (d:ds,m) = putFirstIdx name t
-          (l,_) = d
-          f t = T {dims=ds, ten=t}
 
-t1 = T [(4,(Covariant,"p")),(4,(Contravariant,"q")),(2,(Covariant,"r"))] $ fromList [1..32::Double]
-t2 = T [(4,(Covariant,"p")),(4,(Contravariant,"q"))] $ fromList [1..16::Double]
 
-contract1 t name1 name2 = map head $ zipWith drop [0..] (map (flip parts name2) (parts t name1))
 
-sumT ls = foldl (zipWith (+)) [0,0..] (map (toList.ten) ls)
 
-on f g = \x y -> f (g x) (g y)
+delta i j | i==j      = 1
+          | otherwise = 0
+
+e i n = fromList [ delta k i | k <- [1..n]]
+
+ident n = fromRows [ e i n | i <- [1..n]]
+
+diag l = reshape c $ fromList $ [ l!!(i-1) * delta k i | k <- [1..c], i <- [1..c]]
+    where c = length l
+
+tensorFromVector idx v = T {dims = [(dim v,idx)], ten = v}
+tensorFromMatrix idxr idxc m = T {dims = [(rows m,idxr),(cols m,idxc)], ten = cdat m}
+
+td = tensorFromMatrix (Contravariant,"i") (Covariant,"j") $ diag [1..4] :: Tensor Double
+
+tn = tensorFromMatrix (Contravariant,"i") (Covariant,"j") $ (2><3) [1..6] :: Tensor Double
+tt = tensorFromMatrix (Contravariant,"i") (Covariant,"j") $ (2><3) [1..6] :: Tensor Double
+
+tq = T [(3,(Covariant,"p")),(2,(Covariant,"q")),(2,(Covariant,"r"))] $ fromList [11 .. 22] :: Tensor Double
+
+r1 = contraction tt "j" tq "p"
+r1' = contraction' tt "j" tq "p"
+
+pru = do
+    mapM_ (putStrLn.shdims.dims.normal) (contractions t1 t2)
+    let t1 = contraction tt "i" tq "q"
+    print $ normal t1
+    print $ foldl part t1 [("j",0),("p'",1),("r'",1)]
+    let t2 = contraction' tt "i" tq "q"
+    print $ normal t2
+    print $ foldl part t2 [("j",0),("p'",1),("r'",1)]
+    let t1 = contraction tq "q" tt "i"
+    print $ normal t1
+    print $ foldl part t1 [("j'",0),("p",1),("r",1)]
+    let t2 = contraction' tq "q" tt "i"
+    print $ normal t2
+    print $ foldl part t2 [("j'",0),("p",1),("r",1)]
+
+
+names t = sort $ map (snd.snd) (dims t)
+
+normal t = tridx (names t) t
 
-contract t1 n1 t2 n2 = T (tail d1++tail d2) (cdat m)
-    where (d1,m1) = putFirstIdx n1 t1
-          (d2,m2) = putFirstIdx n2 t2
-          m = multiply RowMajor (trans m2) m1
\ No newline at end of file
+contractions t1 t2 = [ contraction t1 n1 t2 n2 | n1 <- names t1, n2 <- names t2, compatIdx t1 n1 t2 n2 ]
diff --git a/lib/Data/Packed/Internal/Matrix.hs b/lib/Data/Packed/Internal/Matrix.hs
index a2a70dd..ec6657a 100644
--- a/lib/Data/Packed/Internal/Matrix.hs
+++ b/lib/Data/Packed/Internal/Matrix.hs
@@ -65,6 +65,15 @@ partit :: Int -> [a] -> [[a]]
 partit _ [] = []
 partit n l  = take n l : partit n (drop n l)
 
+-- | obtains the common value of a property of a list
+common :: (Eq a) => (b->a) -> [b] -> Maybe a
+common f = commonval . map f where
+    commonval :: (Eq a) => [a] -> Maybe a
+    commonval [] = Nothing
+    commonval [a] = Just a
+    commonval (a:b:xs) = if a==b then commonval (b:xs) else Nothing
+
+
 toLists m | fortran m = transpose $ partit (rows m) . toList . dat $ m
           | otherwise = partit (cols m) . toList . dat $ m
 
@@ -115,7 +124,7 @@ transdataAux fun c1 d c2 =
         else unsafePerformIO $ do
             v <- createVector (dim d)
             fun r1 c1 (ptr d) r2 c2 (ptr v) // check "transdataAux" [d]
-            putStrLn "---> transdataAux"
+            --putStrLn "---> transdataAux"
             return v
   where r1 = dim d `div` c1
         r2 = dim d `div` c2
@@ -136,6 +145,12 @@ transdata c1 d c2 | isReal baseOf d = scast $ transdataR c1 (scast d) c2
 --{-# RULES "transdataR" transdata=transdataR #-}
 --{-# RULES "transdataC" transdata=transdataC #-}
 
+-- | creates a Matrix from a list of vectors
+fromRows :: Field t => [Vector t] -> Matrix t
+fromRows vs = case common dim vs of
+    Nothing -> error "fromRows applied to [] or to vectors with different sizes"
+    Just c  -> reshape c (join vs)
+
 -- | extracts the rows of a matrix as a list of vectors
 toRows :: Storable t => Matrix t -> [Vector t]
 toRows m = toRows' 0 where
diff --git a/lib/Data/Packed/Internal/Tensor.hs b/lib/Data/Packed/Internal/Tensor.hs
index 960e3c5..b66d6b8 100644
--- a/lib/Data/Packed/Internal/Tensor.hs
+++ b/lib/Data/Packed/Internal/Tensor.hs
@@ -1,4 +1,4 @@
-{-# OPTIONS_GHC -fglasgow-exts -fallow-undecidable-instances #-}
+--{-# OPTIONS_GHC -fglasgow-exts -fallow-undecidable-instances #-}
 -----------------------------------------------------------------------------
 -- |
 -- Module      :  Data.Packed.Internal.Tensor
@@ -19,7 +19,7 @@ import Data.Packed.Internal.Vector
 import Data.Packed.Internal.Matrix
 import Foreign.Storable
 
-data IdxTp = Covariant | Contravariant deriving Show
+data IdxTp = Covariant | Contravariant deriving (Show,Eq)
 
 data Tensor t = T { dims   :: [(Int,(IdxTp,String))]
                   , ten    :: Vector t
@@ -36,5 +36,70 @@ instance (Show a,Storable a) => Show (Tensor a) where
     show T {dims = ds, ten = t} = "("++shdims ds ++") "++ show (toList t)
 
 
-shdims [(n,(t,name))] = name++"["++show n++"]"
-shdims (d:ds) = shdims [d] ++ "><"++ shdims ds
\ No newline at end of file
+shdims [(n,(t,name))] = name ++ sym t ++"["++show n++"]"
+    where sym Covariant     = "_"
+          sym Contravariant = "^"
+shdims (d:ds) = shdims [d] ++ "><"++ shdims ds
+
+
+
+findIdx name t = ((d1,d2),m) where
+    (d1,d2) = span (\(_,(_,n)) -> n /=name) (dims t)
+    c = product (map fst d2)
+    m = matrixFromVector RowMajor c (ten t)
+
+putFirstIdx name t = (nd,m')
+    where ((d1,d2),m) = findIdx name t
+          m' = matrixFromVector RowMajor c $ cdat $ trans m
+          nd = d2++d1
+          c = dim (ten t) `div` (fst $ head d2)
+
+part t (name,k) = if k<0 || k>=l
+                    then error $ "part "++show (name,k)++" out of range in "++show t
+                    else T {dims = ds, ten = toRows m !! k}
+    where (d:ds,m) = putFirstIdx name t
+          (l,_) = d
+
+parts t name = map f (toRows m)
+    where (d:ds,m) = putFirstIdx name t
+          (l,_) = d
+          f t = T {dims=ds, ten=t}
+
+concatRename l1 l2 = l1 ++ map ren l2 where
+    ren (n,(t,s)) = if {- s `elem` fs -} True then (n,(t,s++"'")) else (n,(t,s))
+    fs = map (snd.snd) l1
+
+prod (T d1 v1) (T d2 v2) = T (concatRename d1 d2) (outer v1 v2)
+
+contraction t1 n1 t2 n2 =
+    if compatIdx t1 n1 t2 n2
+        then T (concatRename (tail d1) (tail d2)) (cdat m)
+        else error "wrong contraction'"
+  where (d1,m1) = putFirstIdx n1 t1
+        (d2,m2) = putFirstIdx n2 t2
+        m = multiply RowMajor (trans m1) m2
+
+sumT ls = foldl (zipWith (+)) [0,0..] (map (toList.ten) ls)
+
+contract1 t name1 name2 = T d $ fromList $ sumT y
+    where d = dims (head y)
+          x = (map (flip parts name2) (parts t name1))
+          y = map head $ zipWith drop [0..] x
+
+contraction' t1 n1 t2 n2 =
+    if compatIdx t1 n1 t2 n2
+        then contract1 (prod t1 t2) n1 (n2++"'")
+        else error "wrong contraction'"
+
+tridx [] t = t
+tridx (name:rest) t = T (d:ds) (join ts) where
+    ((_,d:_),_) = findIdx name t
+    ps = map (tridx rest) (parts t name)
+    ts = map ten ps
+    ds = dims (head ps)
+
+compatIdxAux (n1,(t1,_)) (n2, (t2,_)) = t1 /= t2 && n1 == n2
+
+compatIdx t1 n1 t2 n2 = compatIdxAux d1 d2 where
+    d1 = head $ snd $ fst $ findIdx n1 t1
+    d2 = head $ snd $ fst $ findIdx n2 t2
diff --git a/lib/Data/Packed/Internal/Vector.hs b/lib/Data/Packed/Internal/Vector.hs
index 6ed9339..36d5df7 100644
--- a/lib/Data/Packed/Internal/Vector.hs
+++ b/lib/Data/Packed/Internal/Vector.hs
@@ -35,6 +35,8 @@ instance (Storable a, RealFloat a) => Storable (Complex a) where    --
     poke p (a :+ b) = pokeArray (castPtr p) [a,b]                   --
 ----------------------------------------------------------------------
 
+on f g = \x y -> f (g x) (g y)
+
 (//) :: x -> (x -> y) -> y
 infixl 0 //
 (//) = flip ($)