Implement diag function in Massiv and roundtrip it against accelerate

src/Gargantext/Core/LinearAlgebra.hs

@@ -20,6 +20,7 @@ module Gargantext.Core.LinearAlgebra (
   , createIndices
 
   -- * Operations on matrixes
+  , diag
   , termDivNan
   ) where
 
@@ -31,7 +32,7 @@ import Data.Set qualified as S
 import Data.Set (Set)
 import Prelude
 import Data.Massiv.Array qualified as A
-import Data.Massiv.Array (D, Matrix)
+import Data.Massiv.Array (D, Matrix, Vector)
 
 newtype Index = Index { _Index :: Int }
   deriving newtype (Eq, Show, Ord, Num, Enum)
@@ -44,6 +45,16 @@ createIndices = set2indices . map2set
 
     set2indices :: Ord t => Set t -> Bimap Index t
     set2indices s = foldr (uncurry Bimap.insert) Bimap.empty (zip [0..] $ S.toList s)
+
+-- | Computes the diagnonal matrix of the input one.
+diag :: Num e => Matrix D e -> Vector D e
+diag matrix =
+  let (A.Sz2 _rows cols) = A.size matrix
+      newSize = A.Sz1 cols
+  in A.backpermute' newSize (\j -> j A.:. j) matrix
+
+
+
 -- | `distributional m` returns the distributional distance between terms each
 -- pair of terms as a matrix.  The argument m is the matrix $[n_{ij}]_{i,j}$
 -- where $n_{ij}$ is the coocccurrence between term $i$ and term $j$.
@@ -81,21 +92,19 @@ createIndices = set2indices . map2set
 -- /IMPORTANT/: As this function computes the diagonal matrix in order to carry on the computation
 -- the input has to be a square matrix, or this function will fail at runtime.
 {-
-distributional :: A.Matrix r Int -> A.Matrix r Double
-distributional m' = interpret $ result
+distributional :: Matrix D Int -> Matrix D Double
+distributional m' = result
  where
     m = map A.fromIntegral $ use m'
     n = dim m'
 
-    diag_m = diag m
+    diag_m = A.diag m
 
     d_1 = replicate (constant (Z :. n :. All)) diag_m
     d_2 = replicate (constant (Z :. All :. n)) diag_m
 
-    mi    = (.*) (termDivNan m d_1) (termDivNan m d_2)
+    mi    = (A..*) (termDivNan m d_1) (termDivNan m d_2)
 
-    -- w = (.-) mi d_mi
-    
     -- The matrix permutations is taken care of below by directly replicating
     -- the matrix mi, making the matrix w unneccessary and saving one step.
     w_1 = replicate (constant (Z :. All :. n :. All)) mi
@@ -104,7 +113,7 @@ distributional m' = interpret $ result
 
     -- The matrix ii = [r_{i,j,k}]_{i,j,k} has r_(i,j,k) = 0 if k = i OR k = j 
     -- and r_(i,j,k) = 1 otherwise (i.e. k /= i AND k /= j). 
-    ii = generate (constant (Z :. n :. n :. n)) 
+    ii = generate (constant (Z :. n :. n :. n))
         (lift1 (\( i A.:. j A.:. k) -> cond ((&&) ((/=) k i) ((/=) k j)) 1 0))
 
     z_1 = sum ((A..*) w' ii)

test/Test/Core/LinearAlgebra.hs

@@ -2,6 +2,7 @@
 {-# LANGUAGE TypeOperators #-}
 {-# LANGUAGE MultiWayIf #-}
 {-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 
 module Test.Core.LinearAlgebra where
 
@@ -70,13 +71,19 @@ accelerate2MassivMatrix m =
   let (Z :. _r :. c) = A.arrayShape m
   in Massiv.delay $ Massiv.fromLists' @Massiv.U Massiv.Par $ Split.chunksOf c (A.toList m)
 
-massiv2AccelerateMatrix :: Elt a => Massiv.Matrix Massiv.D a -> Matrix a
+massiv2AccelerateMatrix :: (Elt a, Massiv.Source r a) => Massiv.Matrix r a -> Matrix a
 massiv2AccelerateMatrix m =
   let m' = Massiv.toLists2 m
       r  = Prelude.length m'
       c  = maybe 0 Prelude.length (headMay m')
   in A.fromList (Z :. r :. c) (mconcat m')
 
+massiv2AccelerateVector :: (Massiv.Source r a, Elt a) => Massiv.Vector r a -> Vector a
+massiv2AccelerateVector m =
+  let m' = Massiv.toList m
+      r  = Prelude.length m'
+  in A.fromList (Z :. r) m'
+
 --
 -- Main test runner
 --
@@ -87,6 +94,7 @@ tests :: TestTree
 tests = testGroup "LinearAlgebra" [
       testProperty "createIndices roundtrip" (compareImplementations (LA.createIndices @Int @Int) Legacy.createIndices mapCreateIndices)
     , testProperty "termDivNan"     compareTermDivNan
+    , testProperty "diag"           compareDiag
     , testGroup "distributional" [
         testProperty "2x2" (compareDistributional twoByTwo)
       , testProperty "roundtrips" (compareImplementations' @(SquareMatrix Int) (Legacy.distributionalWith Naive.run . _SquareMatrix)
@@ -107,6 +115,12 @@ compareTermDivNan i1 i2
         accelerate = Naive.run (Legacy.termDivNan (use i1) (use i2))
     in accelerate === massiv2AccelerateMatrix massiv
 
+compareDiag :: SquareMatrix Int -> Property
+compareDiag (SquareMatrix i1)
+  = let massiv     = Massiv.computeAs Massiv.U $ LA.diag (accelerate2MassivMatrix i1)
+        accelerate =  Naive.run (Legacy.diag (use i1))
+    in accelerate === massiv2AccelerateVector massiv
+
 compareDistributional :: Matrix Int
                       -> Property
 compareDistributional i1