-- SPDX-License-Identifier: MPL-2.0

{- | Sparse matrices use canonical sorted coordinates, coalesced duplicates and
no explicit zeros. Formats carry public metadata and are validated on import.
CSR/CSC indices within each segment must be strictly increasing. DIA stores
one value per row of a diagonal, with zero padding. BSR allows edge blocks
only with zero padding. ELLPACK uses -1 and zero for unused row-major slots.
-}
module HMath.Sparse (
    Sparse,
    SparseFormat (..),
    sparse,
    sparseShape,
    nonzeros,
    sparseEntries,
    fromFormat,
    toCOO,
    toCSR,
    toCSC,
    toDIA,
    toBSR,
    toELLPACK,
    sparseToDense,
    denseToSparse,
    sparseTranspose,
    sparseAdd,
    sparseSubtract,
    spmv,
    spmm,
    sparseKronecker,
    sparseTriangularSolve,
    sparseDiagonalSolve,
    SparseLU,
    sparseLU,
    sparseLULower,
    sparseLUUpper,
    sparseLUPermutation,
    sparseSolveLU,
    sparseCholesky,
    SparseQR,
    sparseQR,
    sparseQRUpper,
    sparseQRApplyQT,
    sparseGraph,
    sparseStorageBytes,
    SparseTensor,
    sparseTensor,
    sparseTensorEntries,
    sparseTensorShape,
    sparseTensorLookup,
    sparseEigenReference,
    sparseSVDReference,
) where

import Control.Monad (foldM)
import Data.Array (accumArray, elems, listArray, (!))
import Data.List (maximumBy, sortOn)
import qualified Data.Map.Strict as M
import Data.Ord (comparing)
import Foreign.Storable (sizeOf)
import HMath.Error
import HMath.Matrix
import HMath.Matrix.Internal (dotV, finiteValues, stableNorm, swap)
import HMath.Matrix.LAPACK (Eigen, SVD, eigen, svd)

{- | Canonical immutable sparse matrix with summed duplicate coordinates and explicit zeros
removed. Construct using sparse or checked format conversions.
-}
data Sparse a = Sparse Int Int (M.Map (Int, Int) a) deriving (Sparse a -> Sparse a -> Bool
(Sparse a -> Sparse a -> Bool)
-> (Sparse a -> Sparse a -> Bool) -> Eq (Sparse a)
forall a. Eq a => Sparse a -> Sparse a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Sparse a -> Sparse a -> Bool
== :: Sparse a -> Sparse a -> Bool
$c/= :: forall a. Eq a => Sparse a -> Sparse a -> Bool
/= :: Sparse a -> Sparse a -> Bool
Eq, Int -> Sparse a -> ShowS
[Sparse a] -> ShowS
Sparse a -> String
(Int -> Sparse a -> ShowS)
-> (Sparse a -> String) -> ([Sparse a] -> ShowS) -> Show (Sparse a)
forall a. Show a => Int -> Sparse a -> ShowS
forall a. Show a => [Sparse a] -> ShowS
forall a. Show a => Sparse a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Sparse a -> ShowS
showsPrec :: Int -> Sparse a -> ShowS
$cshow :: forall a. Show a => Sparse a -> String
show :: Sparse a -> String
$cshowList :: forall a. Show a => [Sparse a] -> ShowS
showList :: [Sparse a] -> ShowS
Show)

{- | Interchange representation selector carrying COO,CSR,CSC,DIA,BSR,or ELLPACK arrays;
validate through fromFormat.
-}
data SparseFormat a
    = COO Int Int [(Int, Int, a)]
    | CSR Int Int [Int] [Int] [a]
    | CSC Int Int [Int] [Int] [a]
    | DIA Int Int [(Int, [a])]
    | BSR Int Int Int Int [Int] [Int] [[a]]
    | ELLPACK Int Int Int [Int] [a]
    deriving (SparseFormat a -> SparseFormat a -> Bool
(SparseFormat a -> SparseFormat a -> Bool)
-> (SparseFormat a -> SparseFormat a -> Bool)
-> Eq (SparseFormat a)
forall a. Eq a => SparseFormat a -> SparseFormat a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => SparseFormat a -> SparseFormat a -> Bool
== :: SparseFormat a -> SparseFormat a -> Bool
$c/= :: forall a. Eq a => SparseFormat a -> SparseFormat a -> Bool
/= :: SparseFormat a -> SparseFormat a -> Bool
Eq, Int -> SparseFormat a -> ShowS
[SparseFormat a] -> ShowS
SparseFormat a -> String
(Int -> SparseFormat a -> ShowS)
-> (SparseFormat a -> String)
-> ([SparseFormat a] -> ShowS)
-> Show (SparseFormat a)
forall a. Show a => Int -> SparseFormat a -> ShowS
forall a. Show a => [SparseFormat a] -> ShowS
forall a. Show a => SparseFormat a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> SparseFormat a -> ShowS
showsPrec :: Int -> SparseFormat a -> ShowS
$cshow :: forall a. Show a => SparseFormat a -> String
show :: SparseFormat a -> String
$cshowList :: forall a. Show a => [SparseFormat a] -> ShowS
showList :: [SparseFormat a] -> ShowS
Show)

shapeValid :: Int -> Int -> Either MathError ()
shapeValid :: Int -> Int -> Either MathError ()
shapeValid Int
r Int
c
    | Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError ()
forall a b. a -> Either a b
Left ((Int, Int) -> MathError
NegativeDimensions (Int
r, Int
c))
    | Bool
otherwise = () -> Either MathError ()
forall a b. b -> Either a b
Right ()

{- | Given rows,columns,then (row,column,value) entries, validate coordinates, sum
duplicates and remove explicit zeros. Return canonical ordered sparse storage; invalid
shape/indices or size overflow fail.
-}
sparse :: (Eq a, Num a) => Int -> Int -> [(Int, Int, a)] -> Either MathError (Sparse a)
sparse :: forall a.
(Eq a, Num a) =>
Int -> Int -> [(Int, Int, a)] -> Either MathError (Sparse a)
sparse Int
r Int
c [(Int, Int, a)]
xs = do
    Int -> Int -> Either MathError ()
shapeValid Int
r Int
c
    if ((Int, Int, a) -> Bool) -> [(Int, Int, a)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\(Int
i, Int
j, a
_) -> Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
r Bool -> Bool -> Bool
|| Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
c) [(Int, Int, a)]
xs
        then MathError -> Either MathError (Sparse a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"sparse: coordinate out of bounds")
        else Sparse a -> Either MathError (Sparse a)
forall a b. b -> Either a b
Right (Int -> Int -> Map (Int, Int) a -> Sparse a
forall a. Int -> Int -> Map (Int, Int) a -> Sparse a
Sparse Int
r Int
c ((a -> Bool) -> Map (Int, Int) a -> Map (Int, Int) a
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
0) ((a -> a -> a) -> [((Int, Int), a)] -> Map (Int, Int) a
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith a -> a -> a
forall a. Num a => a -> a -> a
(+) [((Int
i, Int
j), a
v) | (Int
i, Int
j, a
v) <- [(Int, Int, a)]
xs])))

-- | Return declared (row count,column count), including empty dimensions.
sparseShape :: Sparse a -> (Int, Int)
sparseShape :: forall a. Sparse a -> (Int, Int)
sparseShape (Sparse Int
r Int
c Map (Int, Int) a
_) = (Int
r, Int
c)

-- | Return canonical nonzero (row,column,value) triples in ascending coordinate order.
sparseEntries :: Sparse a -> [(Int, Int, a)]
sparseEntries :: forall a. Sparse a -> [(Int, Int, a)]
sparseEntries (Sparse Int
_ Int
_ Map (Int, Int) a
xs) = [(Int
i, Int
j, a
v) | ((Int
i, Int
j), a
v) <- Map (Int, Int) a -> [((Int, Int), a)]
forall k a. Map k a -> [(k, a)]
M.toAscList Map (Int, Int) a
xs]

-- | Return stored nonzero count after duplicate coalescing and zero removal.
nonzeros :: Sparse a -> Int
nonzeros :: forall a. Sparse a -> Int
nonzeros (Sparse Int
_ Int
_ Map (Int, Int) a
xs) = Map (Int, Int) a -> Int
forall k a. Map k a -> Int
M.size Map (Int, Int) a
xs

segments :: Int -> Int -> [Int] -> [Int] -> [a] -> Either MathError [[(Int, a)]]
segments :: forall a.
Int
-> Int -> [Int] -> [Int] -> [a] -> Either MathError [[(Int, a)]]
segments Int
outer Int
inner [Int]
ptr [Int]
indices [a]
values
    | Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
outer Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int) = Either MathError [[(Int, a)]]
forall {b}. Either MathError b
bad
    | [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ptr Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
outer Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Bool -> Bool -> Bool
|| [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
indices Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
values = Either MathError [[(Int, a)]]
forall {b}. Either MathError b
bad
    | [Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
ptr Bool -> Bool -> Bool
|| [Int]
ptr [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
0 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0 Bool -> Bool -> Bool
|| [Int] -> Int
forall a. HasCallStack => [a] -> a
last [Int]
ptr Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
values = Either MathError [[(Int, a)]]
forall {b}. Either MathError b
bad
    | ((Int, Int) -> Bool) -> [(Int, Int)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\(Int
a, Int
b) -> Int
a Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
b Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
a Bool -> Bool -> Bool
|| Int
b Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
values) ([Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
ptr (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop Int
1 [Int]
ptr)) = Either MathError [[(Int, a)]]
forall {b}. Either MathError b
bad
    | (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Int
i -> Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
inner) [Int]
indices = Either MathError [[(Int, a)]]
forall {b}. Either MathError b
bad
    | ([(Int, a)] -> Bool) -> [[(Int, a)]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\[(Int, a)]
part -> ((Int, Int) -> Bool) -> [(Int, Int)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Int -> Int -> Bool) -> (Int, Int) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(>=)) ([Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Int, a) -> Int) -> [(Int, a)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, a) -> Int
forall a b. (a, b) -> a
fst [(Int, a)]
part) (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop Int
1 (((Int, a) -> Int) -> [(Int, a)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, a) -> Int
forall a b. (a, b) -> a
fst [(Int, a)]
part)))) [[(Int, a)]]
parts = Either MathError [[(Int, a)]]
forall {b}. Either MathError b
bad
    | Bool
otherwise = [[(Int, a)]] -> Either MathError [[(Int, a)]]
forall a b. b -> Either a b
Right [[(Int, a)]]
parts
  where
    bad :: Either MathError b
bad = MathError -> Either MathError b
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"sparse compressed format: invalid pointers, indices, or ordering")
    parts :: [[(Int, a)]]
parts = [Int -> [(Int, a)] -> [(Int, a)]
forall a. Int -> [a] -> [a]
take ([Int]
ptr [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
- [Int]
ptr [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
i) (Int -> [(Int, a)] -> [(Int, a)]
forall a. Int -> [a] -> [a]
drop ([Int]
ptr [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
i) ([Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
indices [a]
values)) | Int
i <- [Int
0 .. Int
outer Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]

{- | Validate COO/CSR/CSC/DIA/BSR/ELLPACK dimensions, pointers/indices and values, then
construct canonical sparse storage. Invalid layout or coordinates return Left
MathError.
-}
fromFormat :: (Eq a, Num a) => SparseFormat a -> Either MathError (Sparse a)
fromFormat :: forall a.
(Eq a, Num a) =>
SparseFormat a -> Either MathError (Sparse a)
fromFormat (COO Int
r Int
c [(Int, Int, a)]
xs) = Int -> Int -> [(Int, Int, a)] -> Either MathError (Sparse a)
forall a.
(Eq a, Num a) =>
Int -> Int -> [(Int, Int, a)] -> Either MathError (Sparse a)
sparse Int
r Int
c [(Int, Int, a)]
xs
fromFormat (CSR Int
r Int
c [Int]
ptr [Int]
indices [a]
values) = do
    Int -> Int -> Either MathError ()
shapeValid Int
r Int
c
    rows <- Int
-> Int -> [Int] -> [Int] -> [a] -> Either MathError [[(Int, a)]]
forall a.
Int
-> Int -> [Int] -> [Int] -> [a] -> Either MathError [[(Int, a)]]
segments Int
r Int
c [Int]
ptr [Int]
indices [a]
values
    sparse r c [(i, j, v) | (i, row) <- zip [0 ..] rows, (j, v) <- row]
fromFormat (CSC Int
r Int
c [Int]
ptr [Int]
indices [a]
values) = Sparse a -> Sparse a
forall a. Sparse a -> Sparse a
sparseTranspose (Sparse a -> Sparse a)
-> Either MathError (Sparse a) -> Either MathError (Sparse a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SparseFormat a -> Either MathError (Sparse a)
forall a.
(Eq a, Num a) =>
SparseFormat a -> Either MathError (Sparse a)
fromFormat (Int -> Int -> [Int] -> [Int] -> [a] -> SparseFormat a
forall a. Int -> Int -> [Int] -> [Int] -> [a] -> SparseFormat a
CSR Int
c Int
r [Int]
ptr [Int]
indices [a]
values)
fromFormat (DIA Int
r Int
c [(Int, [a])]
diagonals) = do
    Int -> Int -> Either MathError ()
shapeValid Int
r Int
c
    if ((Int, [a]) -> Bool) -> [(Int, [a])] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
r) (Int -> Bool) -> ((Int, [a]) -> Int) -> (Int, [a]) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([a] -> Int) -> ((Int, [a]) -> [a]) -> (Int, [a]) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, [a]) -> [a]
forall a b. (a, b) -> b
snd) [(Int, [a])]
diagonals
        Bool -> Bool -> Bool
|| Map Int () -> Int
forall a. Map Int a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([(Int, ())] -> Map Int ()
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Int
k, ()) | (Int
k, [a]
_) <- [(Int, [a])]
diagonals]) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [(Int, [a])] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, [a])]
diagonals
        Bool -> Bool -> Bool
|| ((Int, [a]) -> Bool) -> [(Int, [a])] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any
            ( \(Int
k, [a]
vs) ->
                ((Int, a) -> Bool) -> [(Int, a)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any
                    (\(Int
i, a
v) -> (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
k Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 Bool -> Bool -> Bool
|| Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
k Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
c) Bool -> Bool -> Bool
&& a
v a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
0)
                    ([Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] [a]
vs)
            )
            [(Int, [a])]
diagonals
        then MathError -> Either MathError (Sparse a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"DIA: invalid lengths, duplicate offsets, or nonzero padding")
        else Int -> Int -> [(Int, Int, a)] -> Either MathError (Sparse a)
forall a.
(Eq a, Num a) =>
Int -> Int -> [(Int, Int, a)] -> Either MathError (Sparse a)
sparse Int
r Int
c [(Int
i, Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k, a
v) | (Int
k, [a]
vs) <- [(Int, [a])]
diagonals, (Int
i, a
v) <- [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [a]
vs, a
v a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
0]
fromFormat (BSR Int
r Int
c Int
br Int
bc [Int]
ptr [Int]
cols [[a]]
blocks) = do
    Int -> Int -> Either MathError ()
shapeValid Int
r Int
c
    if Int
br Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 Bool -> Bool -> Bool
|| Int
bc Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 Bool -> Bool -> Bool
|| Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
br Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
bc Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int)
        then MathError -> Either MathError (Sparse a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"BSR: invalid block size")
        else do
            let nr :: Int
nr = Integer -> Int
forall a. Num a => Integer -> a
fromInteger ((Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
r Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
br Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
br)
                nc :: Int
nc = Integer -> Int
forall a. Num a => Integer -> a
fromInteger ((Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
c Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
bc Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
bc)
            rows <- Int
-> Int
-> [Int]
-> [Int]
-> [[a]]
-> Either MathError [[(Int, [a])]]
forall a.
Int
-> Int -> [Int] -> [Int] -> [a] -> Either MathError [[(Int, a)]]
segments Int
nr Int
nc [Int]
ptr [Int]
cols [[a]]
blocks
            if any ((/= br * bc) . length) blocks
                then Left (InvalidArgument "BSR: incorrect block element count")
                else do
                    let coords =
                            [ (Int
bi Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
br Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i, Int
bj Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
bc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j, [a]
vs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
bc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j))
                            | (Int
bi, [(Int, [a])]
row) <- [Int] -> [[(Int, [a])]] -> [(Int, [(Int, [a])])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [[(Int, [a])]]
rows
                            , (Int
bj, [a]
vs) <- [(Int, [a])]
row
                            , Int
i <- [Int
0 .. Int
br Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                            , Int
j <- [Int
0 .. Int
bc Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                            ]
                    if any (\(Int
i, Int
j, a
v) -> (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
r Bool -> Bool -> Bool
|| Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
c) Bool -> Bool -> Bool
&& a
v a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
0) coords
                        then Left (InvalidArgument "BSR: nonzero edge padding")
                        else sparse r c [(i, j, v) | (i, j, v) <- coords, i < r, j < c]
fromFormat (ELLPACK Int
r Int
c Int
width [Int]
indices [a]
values) = do
    Int -> Int -> Either MathError ()
shapeValid Int
r Int
c
    if Int
width Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
        Bool -> Bool -> Bool
|| Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
r Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
width Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int)
        Bool -> Bool -> Bool
|| [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
indices Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
width
        Bool -> Bool -> Bool
|| [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
values Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
width
        then MathError -> Either MathError (Sparse a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"ELLPACK: invalid width or slot count")
        else do
            let rows :: [[(Int, a)]]
rows = [Int -> [(Int, a)] -> [(Int, a)]
forall a. Int -> [a] -> [a]
take Int
width (Int -> [(Int, a)] -> [(Int, a)]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
width) ([Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
indices [a]
values)) | Int
i <- [Int
0 .. Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
                valid :: [(Int, b)] -> Bool
valid [(Int, b)]
row =
                    let active :: [(Int, b)]
active = ((Int, b) -> Bool) -> [(Int, b)] -> [(Int, b)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= (-Int
1)) (Int -> Bool) -> ((Int, b) -> Int) -> (Int, b) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, b) -> Int
forall a b. (a, b) -> a
fst) [(Int, b)]
row; padding :: [(Int, b)]
padding = Int -> [(Int, b)] -> [(Int, b)]
forall a. Int -> [a] -> [a]
drop ([(Int, b)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, b)]
active) [(Int, b)]
row
                     in ((Int, b) -> Bool) -> [(Int, b)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\(Int
j, b
_) -> Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
c) [(Int, b)]
active
                            Bool -> Bool -> Bool
&& ((Int, b) -> Bool) -> [(Int, b)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\(Int
j, b
v) -> Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== (-Int
1) Bool -> Bool -> Bool
&& b
v b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
0) [(Int, b)]
padding
                            Bool -> Bool -> Bool
&& [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Int -> Int -> Bool) -> [Int] -> [Int] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) (((Int, b) -> Int) -> [(Int, b)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, b) -> Int
forall a b. (a, b) -> a
fst [(Int, b)]
active) (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop Int
1 (((Int, b) -> Int) -> [(Int, b)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, b) -> Int
forall a b. (a, b) -> a
fst [(Int, b)]
active)))
            if Bool -> Bool
not (([(Int, a)] -> Bool) -> [[(Int, a)]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all [(Int, a)] -> Bool
forall {b}. (Eq b, Num b) => [(Int, b)] -> Bool
valid [[(Int, a)]]
rows)
                then MathError -> Either MathError (Sparse a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"ELLPACK: malformed indices or padding")
                else Int -> Int -> [(Int, Int, a)] -> Either MathError (Sparse a)
forall a.
(Eq a, Num a) =>
Int -> Int -> [(Int, Int, a)] -> Either MathError (Sparse a)
sparse Int
r Int
c [(Int
i, Int
j, a
v) | (Int
i, [(Int, a)]
row) <- [Int] -> [[(Int, a)]] -> [(Int, [(Int, a)])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [[(Int, a)]]
rows, (Int
j, a
v) <- [(Int, a)]
row, Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= (-Int
1)]

{- | Return coordinate-list SparseFormat with explicit row/column/value arrays, preserving
canonical nonzeros.
-}
toCOO :: Sparse a -> SparseFormat a
toCOO :: forall a. Sparse a -> SparseFormat a
toCOO Sparse a
a = let (Int
r, Int
c) = Sparse a -> (Int, Int)
forall a. Sparse a -> (Int, Int)
sparseShape Sparse a
a in Int -> Int -> [(Int, Int, a)] -> SparseFormat a
forall a. Int -> Int -> [(Int, Int, a)] -> SparseFormat a
COO Int
r Int
c (Sparse a -> [(Int, Int, a)]
forall a. Sparse a -> [(Int, Int, a)]
sparseEntries Sparse a
a)

{- | Return compressed sparse row format with row offsets, column indices and nonzero
values.
-}
toCSR :: Sparse a -> SparseFormat a
toCSR :: forall a. Sparse a -> SparseFormat a
toCSR Sparse a
a =
    let (Int
r, Int
c) = Sparse a -> (Int, Int)
forall a. Sparse a -> (Int, Int)
sparseShape Sparse a
a
        xs :: [(Int, Int, a)]
xs = Sparse a -> [(Int, Int, a)]
forall a. Sparse a -> [(Int, Int, a)]
sparseEntries Sparse a
a
        counts :: Map Int Int
counts = (Int -> Int -> Int) -> [(Int, Int)] -> Map Int Int
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) [(Int
i, Int
1 :: Int) | (Int
i, Int
_, a
_) <- [(Int, Int, a)]
xs]
     in Int -> Int -> [Int] -> [Int] -> [a] -> SparseFormat a
forall a. Int -> Int -> [Int] -> [Int] -> [a] -> SparseFormat a
CSR
            Int
r
            Int
c
            ((Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 [Int -> Int -> Map Int Int -> Int
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Int
0 Int
i Map Int Int
counts | Int
i <- [Int
0 .. Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]])
            [Int
j | (Int
_, Int
j, a
_) <- [(Int, Int, a)]
xs]
            [a
v | (Int
_, Int
_, a
v) <- [(Int, Int, a)]
xs]

-- | Return compressed sparse column format with column offsets, row indices and values.
toCSC :: Sparse a -> SparseFormat a
toCSC :: forall a. Sparse a -> SparseFormat a
toCSC Sparse a
a = case Sparse a -> SparseFormat a
forall a. Sparse a -> SparseFormat a
toCSR (Sparse a -> Sparse a
forall a. Sparse a -> Sparse a
sparseTranspose Sparse a
a) of CSR Int
c Int
r [Int]
p [Int]
i [a]
v -> Int -> Int -> [Int] -> [Int] -> [a] -> SparseFormat a
forall a. Int -> Int -> [Int] -> [Int] -> [a] -> SparseFormat a
CSC Int
r Int
c [Int]
p [Int]
i [a]
v; SparseFormat a
_ -> String -> SparseFormat a
forall a. HasCallStack => String -> a
error String
"toCSC invariant"

{- | Return diagonal sparse format; absent positions in represented diagonals are padded
with zero.
-}
toDIA :: (Num a) => Sparse a -> SparseFormat a
toDIA :: forall a. Num a => Sparse a -> SparseFormat a
toDIA Sparse a
a =
    let (Int
r, Int
c) = Sparse a -> (Int, Int)
forall a. Sparse a -> (Int, Int)
sparseShape Sparse a
a
        xs :: [(Int, Int, a)]
xs = Sparse a -> [(Int, Int, a)]
forall a. Sparse a -> [(Int, Int, a)]
sparseEntries Sparse a
a
        offsets :: [Int]
offsets = Map Int () -> [Int]
forall k a. Map k a -> [k]
M.keys ([(Int, ())] -> Map Int ()
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i, ()) | (Int
i, Int
j, a
_) <- [(Int, Int, a)]
xs])
        vals :: Map (Int, Int) a
vals = [((Int, Int), a)] -> Map (Int, Int) a
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [((Int
i, Int
j), a
v) | (Int
i, Int
j, a
v) <- [(Int, Int, a)]
xs]
     in Int -> Int -> [(Int, [a])] -> SparseFormat a
forall a. Int -> Int -> [(Int, [a])] -> SparseFormat a
DIA Int
r Int
c [(Int
k, [a -> (Int, Int) -> Map (Int, Int) a -> a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault a
0 (Int
i, Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) Map (Int, Int) a
vals | Int
i <- [Int
0 .. Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]) | Int
k <- [Int]
offsets]

{- | Given positive block rows,block columns,then sparse matrix, build block-sparse-row
format with zero padding within edge blocks. Invalid block dimensions or size overflow
fails.
-}
toBSR :: (Num a) => Int -> Int -> Sparse a -> Either MathError (SparseFormat a)
toBSR :: forall a.
Num a =>
Int -> Int -> Sparse a -> Either MathError (SparseFormat a)
toBSR Int
br Int
bc Sparse a
a
    | Int
br Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 Bool -> Bool -> Bool
|| Int
bc Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 Bool -> Bool -> Bool
|| Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
br Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
bc Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int) =
        MathError -> Either MathError (SparseFormat a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"toBSR: invalid block size")
    | Bool
otherwise = do
        let groups :: Map (Int, Int) [(Int, Int, a)]
groups =
                ([(Int, Int, a)] -> [(Int, Int, a)] -> [(Int, Int, a)])
-> [((Int, Int), [(Int, Int, a)])]
-> Map (Int, Int) [(Int, Int, a)]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith
                    [(Int, Int, a)] -> [(Int, Int, a)] -> [(Int, Int, a)]
forall a. [a] -> [a] -> [a]
(++)
                    [((Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
br, Int
j Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
bc), [(Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
br, Int
j Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
bc, a
v)]) | (Int
i, Int
j, a
v) <- Sparse a -> [(Int, Int, a)]
forall a. Sparse a -> [(Int, Int, a)]
sparseEntries Sparse a
a]
            nr :: Int
nr = Integer -> Int
forall a. Num a => Integer -> a
fromInteger ((Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
r Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
br Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
br)
            counts :: Map Int Int
counts = (Int -> Int -> Int) -> [(Int, Int)] -> Map Int Int
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) [(Int
i, Int
1 :: Int) | ((Int
i, Int
_), [(Int, Int, a)]
_) <- Map (Int, Int) [(Int, Int, a)] -> [((Int, Int), [(Int, Int, a)])]
forall k a. Map k a -> [(k, a)]
M.toAscList Map (Int, Int) [(Int, Int, a)]
groups]
            block :: [(Int, Int, a)] -> [a]
block [(Int, Int, a)]
vs =
                let vals :: Map (Int, Int) a
vals = [((Int, Int), a)] -> Map (Int, Int) a
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [((Int
i, Int
j), a
v) | (Int
i, Int
j, a
v) <- [(Int, Int, a)]
vs]
                 in [a -> (Int, Int) -> Map (Int, Int) a -> a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault a
0 (Int
i, Int
j) Map (Int, Int) a
vals | Int
i <- [Int
0 .. Int
br Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Int
j <- [Int
0 .. Int
bc Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
        SparseFormat a -> Either MathError (SparseFormat a)
forall a b. b -> Either a b
Right
            ( Int
-> Int -> Int -> Int -> [Int] -> [Int] -> [[a]] -> SparseFormat a
forall a.
Int
-> Int -> Int -> Int -> [Int] -> [Int] -> [[a]] -> SparseFormat a
BSR
                Int
r
                Int
c
                Int
br
                Int
bc
                ((Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 [Int -> Int -> Map Int Int -> Int
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Int
0 Int
i Map Int Int
counts | Int
i <- [Int
0 .. Int
nr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]])
                [Int
j | ((Int
_, Int
j), [(Int, Int, a)]
_) <- Map (Int, Int) [(Int, Int, a)] -> [((Int, Int), [(Int, Int, a)])]
forall k a. Map k a -> [(k, a)]
M.toAscList Map (Int, Int) [(Int, Int, a)]
groups]
                [[(Int, Int, a)] -> [a]
forall {a}. Num a => [(Int, Int, a)] -> [a]
block [(Int, Int, a)]
vs | ((Int, Int)
_, [(Int, Int, a)]
vs) <- Map (Int, Int) [(Int, Int, a)] -> [((Int, Int), [(Int, Int, a)])]
forall k a. Map k a -> [(k, a)]
M.toAscList Map (Int, Int) [(Int, Int, a)]
groups]
            )
  where
    (Int
r, Int
c) = Sparse a -> (Int, Int)
forall a. Sparse a -> (Int, Int)
sparseShape Sparse a
a

{- | Return ELLPACK padded row storage, using row width equal to maximum nonzeros per row.
Invalid/overflowing storage dimensions fail.
-}
toELLPACK :: (Num a) => Sparse a -> Either MathError (SparseFormat a)
toELLPACK :: forall a. Num a => Sparse a -> Either MathError (SparseFormat a)
toELLPACK Sparse a
a
    | Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
r Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
width Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int) =
        MathError -> Either MathError (SparseFormat a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"toELLPACK: slot overflow")
    | Bool
otherwise =
        SparseFormat a -> Either MathError (SparseFormat a)
forall a b. b -> Either a b
Right
            ( Int -> Int -> Int -> [Int] -> [a] -> SparseFormat a
forall a. Int -> Int -> Int -> [Int] -> [a] -> SparseFormat a
ELLPACK
                Int
r
                Int
c
                Int
width
                ([[Int]] -> [Int]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [((Int, a) -> Int) -> [(Int, a)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, a) -> Int
forall a b. (a, b) -> a
fst (Int -> [(Int, a)]
row Int
i) [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate (Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
- [(Int, a)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Int -> [(Int, a)]
row Int
i)) (-Int
1) | Int
i <- [Int
0 .. Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]])
                ([[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [((Int, a) -> a) -> [(Int, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (Int, a) -> a
forall a b. (a, b) -> b
snd (Int -> [(Int, a)]
row Int
i) [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Int -> a -> [a]
forall a. Int -> a -> [a]
replicate (Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
- [(Int, a)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Int -> [(Int, a)]
row Int
i)) a
0 | Int
i <- [Int
0 .. Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]])
            )
  where
    (Int
r, Int
c) = Sparse a -> (Int, Int)
forall a. Sparse a -> (Int, Int)
sparseShape Sparse a
a
    rows :: Map Int [(Int, a)]
rows = ([(Int, a)] -> [(Int, a)] -> [(Int, a)])
-> [(Int, [(Int, a)])] -> Map Int [(Int, a)]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith [(Int, a)] -> [(Int, a)] -> [(Int, a)]
forall a. [a] -> [a] -> [a]
(++) [(Int
i, [(Int
j, a
v)]) | (Int
i, Int
j, a
v) <- Sparse a -> [(Int, Int, a)]
forall a. Sparse a -> [(Int, Int, a)]
sparseEntries Sparse a
a]
    row :: Int -> [(Int, a)]
row Int
i = ((Int, a) -> Int) -> [(Int, a)] -> [(Int, a)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Int, a) -> Int
forall a b. (a, b) -> a
fst ([(Int, a)] -> Int -> Map Int [(Int, a)] -> [(Int, a)]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [] Int
i Map Int [(Int, a)]
rows)
    width :: Int
width = [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (Int
0 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: ([(Int, a)] -> Int) -> [[(Int, a)]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map [(Int, a)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Map Int [(Int, a)] -> [[(Int, a)]]
forall k a. Map k a -> [a]
M.elems Map Int [(Int, a)]
rows))

{- | Return a dense Matrix with the same dimensions and zeros for missing entries. Checked
dense element-count overflow can fail.
-}
sparseToDense :: (Num a) => Sparse a -> Either MathError (Matrix a)
sparseToDense :: forall a. Num a => Sparse a -> Either MathError (Matrix a)
sparseToDense (Sparse Int
r Int
c Map (Int, Int) a
xs) =
    Int -> Int -> [a] -> Either MathError (Matrix a)
forall a. Int -> Int -> [a] -> Either MathError (Matrix a)
fromListChecked
        Int
r
        Int
c
        ( if Int
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
|| Int
c Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
            then []
            else [a -> (Int, Int) -> Map (Int, Int) a -> a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault a
0 (Int
i, Int
j) Map (Int, Int) a
xs | Int
i <- [Int
0 .. Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Int
j <- [Int
0 .. Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
        )

{- | Convert a dense matrix to canonical sparse storage, discarding exact zeros and
preserving dimensions.
-}
denseToSparse :: (Eq a, Num a) => Matrix a -> Sparse a
denseToSparse :: forall a. (Eq a, Num a) => Matrix a -> Sparse a
denseToSparse Matrix a
a = Int -> Int -> Map (Int, Int) a -> Sparse a
forall a. Int -> Int -> Map (Int, Int) a -> Sparse a
Sparse Int
r Int
c ([((Int, Int), a)] -> Map (Int, Int) a
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [((Int
k Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
c, Int
k Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
c), a
v) | (Int
k, a
v) <- [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] (Matrix a -> [a]
forall a. Matrix a -> [a]
toList Matrix a
a), a
v a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
0])
  where
    (Int
r, Int
c) = Matrix a -> (Int, Int)
forall a. Matrix a -> (Int, Int)
shape Matrix a
a

{- | Exchange row/column coordinates and dimensions without changing values, returning
canonical sparse order.
-}
sparseTranspose :: Sparse a -> Sparse a
sparseTranspose :: forall a. Sparse a -> Sparse a
sparseTranspose (Sparse Int
r Int
c Map (Int, Int) a
xs) = Int -> Int -> Map (Int, Int) a -> Sparse a
forall a. Int -> Int -> Map (Int, Int) a -> Sparse a
Sparse Int
c Int
r (((Int, Int) -> (Int, Int)) -> Map (Int, Int) a -> Map (Int, Int) a
forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
M.mapKeys (\(Int
i, Int
j) -> (Int
j, Int
i)) Map (Int, Int) a
xs)

{- | Add equal-shaped sparse matrices, coalescing entries and dropping resulting zeros.
Shape mismatch or checked construction errors fail.
-}
sparseAdd :: (Eq a, Num a) => Sparse a -> Sparse a -> Either MathError (Sparse a)
sparseAdd :: forall a.
(Eq a, Num a) =>
Sparse a -> Sparse a -> Either MathError (Sparse a)
sparseAdd (Sparse Int
r Int
c Map (Int, Int) a
a) (Sparse Int
s Int
t Map (Int, Int) a
b)
    | (Int
r, Int
c) (Int, Int) -> (Int, Int) -> Bool
forall a. Eq a => a -> a -> Bool
/= (Int
s, Int
t) = MathError -> Either MathError (Sparse a)
forall a b. a -> Either a b
Left ((Int, Int) -> (Int, Int) -> MathError
ShapeMismatch (Int
r, Int
c) (Int
s, Int
t))
    | Bool
otherwise = Sparse a -> Either MathError (Sparse a)
forall a b. b -> Either a b
Right (Int -> Int -> Map (Int, Int) a -> Sparse a
forall a. Int -> Int -> Map (Int, Int) a -> Sparse a
Sparse Int
r Int
c ((a -> Bool) -> Map (Int, Int) a -> Map (Int, Int) a
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
0) ((a -> a -> a)
-> Map (Int, Int) a -> Map (Int, Int) a -> Map (Int, Int) a
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith a -> a -> a
forall a. Num a => a -> a -> a
(+) Map (Int, Int) a
a Map (Int, Int) a
b)))

{- | Subtract right sparse values from left for equal shapes; coalesce and remove zeros.
Shape mismatch fails.
-}
sparseSubtract :: (Eq a, Num a) => Sparse a -> Sparse a -> Either MathError (Sparse a)
sparseSubtract :: forall a.
(Eq a, Num a) =>
Sparse a -> Sparse a -> Either MathError (Sparse a)
sparseSubtract Sparse a
a (Sparse Int
r Int
c Map (Int, Int) a
b) = Sparse a -> Sparse a -> Either MathError (Sparse a)
forall a.
(Eq a, Num a) =>
Sparse a -> Sparse a -> Either MathError (Sparse a)
sparseAdd Sparse a
a (Int -> Int -> Map (Int, Int) a -> Sparse a
forall a. Int -> Int -> Map (Int, Int) a -> Sparse a
Sparse Int
r Int
c ((a -> a) -> Map (Int, Int) a -> Map (Int, Int) a
forall a b k. (a -> b) -> Map k a -> Map k b
M.map a -> a
forall a. Num a => a -> a
negate Map (Int, Int) a
b))

{- | Given sparse matrix then dense vector, return matrix-vector product. Vector length must
equal column count; mismatch returns Left MathError.
-}
spmv :: (Num a) => Sparse a -> [a] -> Either MathError [a]
spmv :: forall a. Num a => Sparse a -> [a] -> Either MathError [a]
spmv (Sparse Int
r Int
c Map (Int, Int) a
xs) [a]
vector
    | [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
vector Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
c = MathError -> Either MathError [a]
forall a b. a -> Either a b
Left (Int -> Int -> MathError
VectorLengthMismatch Int
c ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
vector))
    | Int
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = [a] -> Either MathError [a]
forall a b. b -> Either a b
Right []
    | Bool
otherwise =
        let input :: Array Int a
input = (Int, Int) -> [a] -> Array Int a
forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (Int
0, Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [a]
vector
            sums :: Array Int a
sums = (a -> a -> a) -> a -> (Int, Int) -> [(Int, a)] -> Array Int a
forall i e a.
Ix i =>
(e -> a -> e) -> e -> (i, i) -> [(i, a)] -> Array i e
accumArray a -> a -> a
forall a. Num a => a -> a -> a
(+) a
0 (Int
0, Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [(Int
i, a
v a -> a -> a
forall a. Num a => a -> a -> a
* (Array Int a
input Array Int a -> Int -> a
forall i e. Ix i => Array i e -> i -> e
! Int
j)) | ((Int
i, Int
j), a
v) <- Map (Int, Int) a -> [((Int, Int), a)]
forall k a. Map k a -> [(k, a)]
M.toAscList Map (Int, Int) a
xs]
         in [a] -> Either MathError [a]
forall a b. b -> Either a b
Right (Array Int a -> [a]
forall i e. Array i e -> [e]
elems Array Int a
sums)

{- | Multiply compatible sparse matrices using sparse row contributions. Inner-dimension
mismatch or output construction errors return Left MathError.
-}
spmm :: (Eq a, Num a) => Sparse a -> Sparse a -> Either MathError (Sparse a)
spmm :: forall a.
(Eq a, Num a) =>
Sparse a -> Sparse a -> Either MathError (Sparse a)
spmm (Sparse Int
r Int
c Map (Int, Int) a
a) (Sparse Int
s Int
t Map (Int, Int) a
b)
    | Int
c Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
s = MathError -> Either MathError (Sparse a)
forall a b. a -> Either a b
Left (Int -> Int -> MathError
InnerDimensionMismatch Int
c Int
s)
    | Bool
otherwise =
        Int -> Int -> [(Int, Int, a)] -> Either MathError (Sparse a)
forall a.
(Eq a, Num a) =>
Int -> Int -> [(Int, Int, a)] -> Either MathError (Sparse a)
sparse Int
r Int
t [(Int
i, Int
j, a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
y) | ((Int
i, Int
k), a
x) <- Map (Int, Int) a -> [((Int, Int), a)]
forall k a. Map k a -> [(k, a)]
M.toAscList Map (Int, Int) a
a, (Int
j, a
y) <- [(Int, a)] -> Int -> Map Int [(Int, a)] -> [(Int, a)]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [] Int
k Map Int [(Int, a)]
rowsB]
  where
    rowsB :: Map Int [(Int, a)]
rowsB = ([(Int, a)] -> [(Int, a)] -> [(Int, a)])
-> [(Int, [(Int, a)])] -> Map Int [(Int, a)]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith [(Int, a)] -> [(Int, a)] -> [(Int, a)]
forall a. [a] -> [a] -> [a]
(++) [(Int
k, [(Int
j, a
v)]) | ((Int
k, Int
j), a
v) <- Map (Int, Int) a -> [((Int, Int), a)]
forall k a. Map k a -> [(k, a)]
M.toAscList Map (Int, Int) a
b]

{- | Return sparse Kronecker product, multiplying row/column dimensions and nonzero pairs.
Dimension/product overflow fails.
-}
sparseKronecker :: (Eq a, Num a) => Sparse a -> Sparse a -> Either MathError (Sparse a)
sparseKronecker :: forall a.
(Eq a, Num a) =>
Sparse a -> Sparse a -> Either MathError (Sparse a)
sparseKronecker Sparse a
a Sparse a
b
    | Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
r Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
s Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int)
        Bool -> Bool -> Bool
|| Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
c Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
t Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int) =
        MathError -> Either MathError (Sparse a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"sparseKronecker: dimension overflow")
    | Bool
otherwise =
        Int -> Int -> [(Int, Int, a)] -> Either MathError (Sparse a)
forall a.
(Eq a, Num a) =>
Int -> Int -> [(Int, Int, a)] -> Either MathError (Sparse a)
sparse
            (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
s)
            (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
t)
            [(Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k, Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l, a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
y) | (Int
i, Int
j, a
x) <- Sparse a -> [(Int, Int, a)]
forall a. Sparse a -> [(Int, Int, a)]
sparseEntries Sparse a
a, (Int
k, Int
l, a
y) <- Sparse a -> [(Int, Int, a)]
forall a. Sparse a -> [(Int, Int, a)]
sparseEntries Sparse a
b]
  where
    (Int
r, Int
c) = Sparse a -> (Int, Int)
forall a. Sparse a -> (Int, Int)
sparseShape Sparse a
a; (Int
s, Int
t) = Sparse a -> (Int, Int)
forall a. Sparse a -> (Int, Int)
sparseShape Sparse a
b

{- | Given lower Bool,square finite sparse matrix,then RHS list, solve by forward/back
substitution. Require correct triangular structure and matching RHS length; zero
diagonal pivots or non-finite arithmetic fail.
-}
sparseTriangularSolve :: (RealFloat a) => Bool -> Sparse a -> [a] -> Either MathError [a]
sparseTriangularSolve :: forall a.
RealFloat a =>
Bool -> Sparse a -> [a] -> Either MathError [a]
sparseTriangularSolve Bool
lower a :: Sparse a
a@(Sparse Int
n Int
m Map (Int, Int) a
values) [a]
b
    | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
m = MathError -> Either MathError [a]
forall a b. a -> Either a b
Left ((Int, Int) -> MathError
NonSquareMatrix (Int
n, Int
m))
    | [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
b Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n = MathError -> Either MathError [a]
forall a b. a -> Either a b
Left (Int -> Int -> MathError
VectorLengthMismatch Int
n ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
b))
    | ((Int, Int, a) -> Bool) -> [(Int, Int, a)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\(Int
i, Int
j, a
_) -> if Bool
lower then Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
i else Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
i) (Sparse a -> [(Int, Int, a)]
forall a. Sparse a -> [(Int, Int, a)]
sparseEntries Sparse a
a) =
        MathError -> Either MathError [a]
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"sparseTriangularSolve: wrong triangular structure")
    | Bool
otherwise = do
        [a] -> Either MathError ()
forall a. RealFloat a => [a] -> Either MathError ()
finiteValues ([a]
b [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Map (Int, Int) a -> [a]
forall k a. Map k a -> [a]
M.elems Map (Int, Int) a
values)
        solved <- (Map Int a -> Int -> Either MathError (Map Int a))
-> Map Int a -> [Int] -> Either MathError (Map Int a)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM Map Int a -> Int -> Either MathError (Map Int a)
step Map Int a
forall k a. Map k a
M.empty (if Bool
lower then [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] else [Int] -> [Int]
forall a. [a] -> [a]
reverse [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1])
        Right [M.findWithDefault 0 i solved | i <- [0 .. n - 1]]
  where
    rows :: Map Int [(Int, a)]
rows = ([(Int, a)] -> [(Int, a)] -> [(Int, a)])
-> [(Int, [(Int, a)])] -> Map Int [(Int, a)]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith [(Int, a)] -> [(Int, a)] -> [(Int, a)]
forall a. [a] -> [a] -> [a]
(++) [(Int
i, [(Int
j, a
v)]) | ((Int
i, Int
j), a
v) <- Map (Int, Int) a -> [((Int, Int), a)]
forall k a. Map k a -> [(k, a)]
M.toAscList Map (Int, Int) a
values]
    step :: Map Int a -> Int -> Either MathError (Map Int a)
step Map Int a
xs Int
i =
        let pivot :: a
pivot = a -> (Int, Int) -> Map (Int, Int) a -> a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault a
0 (Int
i, Int
i) Map (Int, Int) a
values
            v :: a
v =
                ([a]
b [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
i a -> a -> a
forall a. Num a => a -> a -> a
- [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [a
x a -> a -> a
forall a. Num a => a -> a -> a
* a -> Int -> Map Int a -> a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault a
0 Int
j Map Int a
xs | (Int
j, a
x) <- [(Int, a)] -> Int -> Map Int [(Int, a)] -> [(Int, a)]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [] Int
i Map Int [(Int, a)]
rows, Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
i])
                    a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
pivot
         in if a
pivot a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 then MathError -> Either MathError (Map Int a)
forall a b. a -> Either a b
Left (Int -> MathError
SingularMatrix Int
i) else [a] -> Either MathError ()
forall a. RealFloat a => [a] -> Either MathError ()
finiteValues [a
v] Either MathError ()
-> Either MathError (Map Int a) -> Either MathError (Map Int a)
forall a b.
Either MathError a -> Either MathError b -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Map Int a -> Either MathError (Map Int a)
forall a b. b -> Either a b
Right (Int -> a -> Map Int a -> Map Int a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
i a
v Map Int a
xs)

{- | Given diagonal square sparse matrix then RHS, divide by each diagonal pivot.
Off-diagonal entries, shape mismatch, zero pivots or non-finite data fail.
-}
sparseDiagonalSolve :: (RealFloat a) => Sparse a -> [a] -> Either MathError [a]
sparseDiagonalSolve :: forall a. RealFloat a => Sparse a -> [a] -> Either MathError [a]
sparseDiagonalSolve Sparse a
a [a]
b
    | ((Int, Int, a) -> Bool) -> [(Int, Int, a)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\(Int
i, Int
j, a
_) -> Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
j) (Sparse a -> [(Int, Int, a)]
forall a. Sparse a -> [(Int, Int, a)]
sparseEntries Sparse a
a) =
        MathError -> Either MathError [a]
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"sparseDiagonalSolve: diagonal matrix required")
    | Bool
otherwise = Bool -> Sparse a -> [a] -> Either MathError [a]
forall a.
RealFloat a =>
Bool -> Sparse a -> [a] -> Either MathError [a]
sparseTriangularSolve Bool
True Sparse a
a [a]
b

{- | Opaque sparse row-pivoted LU factors with reconstruction permutation; solve through
sparseSolveLU.
-}
data SparseLU a = SparseLU (Sparse a) (Sparse a) [Int] deriving (SparseLU a -> SparseLU a -> Bool
(SparseLU a -> SparseLU a -> Bool)
-> (SparseLU a -> SparseLU a -> Bool) -> Eq (SparseLU a)
forall a. Eq a => SparseLU a -> SparseLU a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => SparseLU a -> SparseLU a -> Bool
== :: SparseLU a -> SparseLU a -> Bool
$c/= :: forall a. Eq a => SparseLU a -> SparseLU a -> Bool
/= :: SparseLU a -> SparseLU a -> Bool
Eq, Int -> SparseLU a -> ShowS
[SparseLU a] -> ShowS
SparseLU a -> String
(Int -> SparseLU a -> ShowS)
-> (SparseLU a -> String)
-> ([SparseLU a] -> ShowS)
-> Show (SparseLU a)
forall a. Show a => Int -> SparseLU a -> ShowS
forall a. Show a => [SparseLU a] -> ShowS
forall a. Show a => SparseLU a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> SparseLU a -> ShowS
showsPrec :: Int -> SparseLU a -> ShowS
$cshow :: forall a. Show a => SparseLU a -> String
show :: SparseLU a -> String
$cshowList :: forall a. Show a => [SparseLU a] -> ShowS
showList :: [SparseLU a] -> ShowS
Show)

-- | Extract unit lower-triangular sparse factor from a SparseLU result.
sparseLULower :: SparseLU a -> Sparse a
sparseLULower :: forall a. SparseLU a -> Sparse a
sparseLULower (SparseLU Sparse a
l Sparse a
_ [Int]
_) = Sparse a
l

-- | Extract upper-triangular sparse factor from a SparseLU result.
sparseLUUpper :: SparseLU a -> Sparse a
sparseLUUpper :: forall a. SparseLU a -> Sparse a
sparseLUUpper (SparseLU Sparse a
_ Sparse a
u [Int]
_) = Sparse a
u

-- | Return LU row permutation in pivoted order, for reconstructing P*A=L*U.
sparseLUPermutation :: SparseLU a -> [Int]
sparseLUPermutation :: forall a. SparseLU a -> [Int]
sparseLUPermutation (SparseLU Sparse a
_ Sparse a
_ [Int]
p) = [Int]
p

{- | Given finite nonnegative relative pivot tolerance then square finite sparse matrix,
compute sparse LU with row pivoting. Invalid tolerance/input, zero/near-singular pivots
or numerical failure returns Left MathError.
-}
sparseLU :: (RealFloat a) => a -> Sparse a -> Either MathError (SparseLU a)
sparseLU :: forall a.
RealFloat a =>
a -> Sparse a -> Either MathError (SparseLU a)
sparseLU a
tol a :: Sparse a
a@(Sparse Int
n Int
m Map (Int, Int) a
values)
    | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
m = MathError -> Either MathError (SparseLU a)
forall a b. a -> Either a b
Left ((Int, Int) -> MathError
NonSquareMatrix (Int
n, Int
m))
    | a
tol a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0 Bool -> Bool -> Bool
|| a -> Bool
forall a. RealFloat a => a -> Bool
isNaN a
tol Bool -> Bool -> Bool
|| a -> Bool
forall a. RealFloat a => a -> Bool
isInfinite a
tol = MathError -> Either MathError (SparseLU a)
forall a b. a -> Either a b
Left MathError
InvalidTolerance
    | Bool
otherwise = do
        [a] -> Either MathError ()
forall a. RealFloat a => [a] -> Either MathError ()
finiteValues (Map (Int, Int) a -> [a]
forall k a. Map k a -> [a]
M.elems Map (Int, Int) a
values)
        (rows, lower, p) <- ((Map Int (Map Int a), Map (Int, Int) a, [Int])
 -> Int
 -> Either MathError (Map Int (Map Int a), Map (Int, Int) a, [Int]))
-> (Map Int (Map Int a), Map (Int, Int) a, [Int])
-> [Int]
-> Either MathError (Map Int (Map Int a), Map (Int, Int) a, [Int])
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (Map Int (Map Int a), Map (Int, Int) a, [Int])
-> Int
-> Either MathError (Map Int (Map Int a), Map (Int, Int) a, [Int])
forall {a}.
(Map Int (Map Int a), Map (Int, Int) a, [a])
-> Int
-> Either MathError (Map Int (Map Int a), Map (Int, Int) a, [a])
step (Map Int (Map Int a)
rowMaps, Map (Int, Int) a
forall k a. Map k a
M.empty, [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]) [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
        l <- sparse n n ([(i, i, 1) | i <- [0 .. n - 1]] ++ [(i, j, v) | ((i, j), v) <- M.toAscList lower])
        u <- sparse n n [(i, j, v) | (i, row) <- M.toAscList rows, (j, v) <- M.toAscList row]
        Right (SparseLU l u p)
  where
    rowMaps :: Map Int (Map Int a)
rowMaps = (Map Int a -> Map Int a -> Map Int a)
-> [(Int, Map Int a)] -> Map Int (Map Int a)
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith Map Int a -> Map Int a -> Map Int a
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union [(Int
i, Int -> a -> Map Int a
forall k a. k -> a -> Map k a
M.singleton Int
j a
v) | (Int
i, Int
j, a
v) <- Sparse a -> [(Int, Int, a)]
forall a. Sparse a -> [(Int, Int, a)]
sparseEntries Sparse a
a]
    scale :: a
scale = [a] -> a
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (a
0 a -> [a] -> [a]
forall a. a -> [a] -> [a]
: (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map a -> a
forall a. Num a => a -> a
abs (Map (Int, Int) a -> [a]
forall k a. Map k a -> [a]
M.elems Map (Int, Int) a
values))
    step :: (Map Int (Map Int a), Map (Int, Int) a, [a])
-> Int
-> Either MathError (Map Int (Map Int a), Map (Int, Int) a, [a])
step (Map Int (Map Int a)
rows, Map (Int, Int) a
l, [a]
p) Int
k = do
        let row :: Int -> Map Int a
row Int
i = Map Int a -> Int -> Map Int (Map Int a) -> Map Int a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Map Int a
forall k a. Map k a
M.empty Int
i Map Int (Map Int a)
rows
            pv :: Int -> a
pv Int
i = a -> Int -> Map Int a -> a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault a
0 Int
k (Int -> Map Int a
row Int
i)
            pivot :: Int
pivot = (Int, a) -> Int
forall a b. (a, b) -> a
fst (((Int, a) -> (Int, a) -> Ordering) -> [(Int, a)] -> (Int, a)
forall (t :: * -> *) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
maximumBy (((Int, a) -> a) -> (Int, a) -> (Int, a) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (Int, a) -> a
forall a b. (a, b) -> b
snd) [(Int
i, a -> a
forall a. Num a => a -> a
abs (Int -> a
pv Int
i)) | Int
i <- [Int
k .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]])
            pivotValue :: a
pivotValue = Int -> a
pv Int
pivot
        if a
pivotValue a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0
            then MathError
-> Either MathError (Map Int (Map Int a), Map (Int, Int) a, [a])
forall a b. a -> Either a b
Left (Int -> MathError
SingularMatrix Int
k)
            else
                if a -> a
forall a. Num a => a -> a
abs a
pivotValue a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
tol a -> a -> a
forall a. Num a => a -> a -> a
* a
scale
                    then MathError
-> Either MathError (Map Int (Map Int a), Map (Int, Int) a, [a])
forall a b. a -> Either a b
Left (Int -> MathError
NearSingularMatrix Int
k)
                    else do
                        let swapped :: Map Int (Map Int a)
swapped = Int -> Map Int a -> Map Int (Map Int a) -> Map Int (Map Int a)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
pivot (Int -> Map Int a
row Int
k) (Int -> Map Int a -> Map Int (Map Int a) -> Map Int (Map Int a)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
k (Int -> Map Int a
row Int
pivot) Map Int (Map Int a)
rows)
                            lsw :: Map (Int, Int) a
lsw = ((Int, Int) -> (Int, Int)) -> Map (Int, Int) a -> Map (Int, Int) a
forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
M.mapKeys (\(Int
i, Int
j) -> (if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
k then Int
pivot else if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
pivot then Int
k else Int
i, Int
j)) Map (Int, Int) a
l
                            pr :: Map Int a
pr = Map Int a -> Int -> Map Int (Map Int a) -> Map Int a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Map Int a
forall k a. Map k a
M.empty Int
k Map Int (Map Int a)
swapped
                            update :: Either MathError (Map a (Map Int a), Map (a, Int) a)
-> a -> Either MathError (Map a (Map Int a), Map (a, Int) a)
update Either MathError (Map a (Map Int a), Map (a, Int) a)
acc a
i = do
                                (rs, ls) <- Either MathError (Map a (Map Int a), Map (a, Int) a)
acc
                                let ri = Map Int a -> a -> Map a (Map Int a) -> Map Int a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Map Int a
forall k a. Map k a
M.empty a
i Map a (Map Int a)
rs
                                    multiplier = a -> Int -> Map Int a -> a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault a
0 Int
k Map Int a
ri a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
pivotValue
                                    next = (a -> Bool) -> Map Int a -> Map Int a
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
0) (Int -> Map Int a -> Map Int a
forall k a. Ord k => k -> Map k a -> Map k a
M.delete Int
k ((a -> a -> a) -> Map Int a -> Map Int a -> Map Int a
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith a -> a -> a
forall a. Num a => a -> a -> a
(+) Map Int a
ri ((a -> a) -> Map Int a -> Map Int a
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (a -> a
forall a. Num a => a -> a
negate (a -> a) -> (a -> a) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a
multiplier a -> a -> a
forall a. Num a => a -> a -> a
*)) Map Int a
pr)))
                                finiteValues (multiplier : M.elems next)
                                Right (M.insert i next rs, if multiplier == 0 then ls else M.insert (i, k) multiplier ls)
                        (nextRows, nextL) <- (Either MathError (Map Int (Map Int a), Map (Int, Int) a)
 -> Int -> Either MathError (Map Int (Map Int a), Map (Int, Int) a))
-> Either MathError (Map Int (Map Int a), Map (Int, Int) a)
-> [Int]
-> Either MathError (Map Int (Map Int a), Map (Int, Int) a)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Either MathError (Map Int (Map Int a), Map (Int, Int) a)
-> Int -> Either MathError (Map Int (Map Int a), Map (Int, Int) a)
forall {a}.
Ord a =>
Either MathError (Map a (Map Int a), Map (a, Int) a)
-> a -> Either MathError (Map a (Map Int a), Map (a, Int) a)
update ((Map Int (Map Int a), Map (Int, Int) a)
-> Either MathError (Map Int (Map Int a), Map (Int, Int) a)
forall a b. b -> Either a b
Right (Map Int (Map Int a)
swapped, Map (Int, Int) a
lsw)) [Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                        Right (nextRows, nextL, swap k pivot p)

{- | Given sparse LU factors then RHS list, apply row permutation and triangular solves.
Dimension mismatch, singularity or non-finite input/result fails.
-}
sparseSolveLU :: (RealFloat a) => SparseLU a -> [a] -> Either MathError [a]
sparseSolveLU :: forall a. RealFloat a => SparseLU a -> [a] -> Either MathError [a]
sparseSolveLU SparseLU a
f [a]
b
    | [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
b Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n = MathError -> Either MathError [a]
forall a b. a -> Either a b
Left (Int -> Int -> MathError
VectorLengthMismatch Int
n ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
b))
    | Bool
otherwise =
        Bool -> Sparse a -> [a] -> Either MathError [a]
forall a.
RealFloat a =>
Bool -> Sparse a -> [a] -> Either MathError [a]
sparseTriangularSolve Bool
True (SparseLU a -> Sparse a
forall a. SparseLU a -> Sparse a
sparseLULower SparseLU a
f) [[a]
b [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
i | Int
i <- SparseLU a -> [Int]
forall a. SparseLU a -> [Int]
sparseLUPermutation SparseLU a
f]
            Either MathError [a]
-> ([a] -> Either MathError [a]) -> Either MathError [a]
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Bool -> Sparse a -> [a] -> Either MathError [a]
forall a.
RealFloat a =>
Bool -> Sparse a -> [a] -> Either MathError [a]
sparseTriangularSolve Bool
False (SparseLU a -> Sparse a
forall a. SparseLU a -> Sparse a
sparseLUUpper SparseLU a
f)
  where
    n :: Int
n = (Int, Int) -> Int
forall a b. (a, b) -> a
fst (Sparse a -> (Int, Int)
forall a. Sparse a -> (Int, Int)
sparseShape (SparseLU a -> Sparse a
forall a. SparseLU a -> Sparse a
sparseLULower SparseLU a
f))

{- | Factor a symmetric positive-definite finite sparse matrix into a lower sparse factor.
Invalid shape/symmetry, nonpositive pivots or numerical failure returns Left MathError.
-}
sparseCholesky :: (RealFloat a) => Sparse a -> Either MathError (Sparse a)
sparseCholesky :: forall a. RealFloat a => Sparse a -> Either MathError (Sparse a)
sparseCholesky a :: Sparse a
a@(Sparse Int
n Int
m Map (Int, Int) a
values)
    | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
m = MathError -> Either MathError (Sparse a)
forall a b. a -> Either a b
Left ((Int, Int) -> MathError
NonSquareMatrix (Int
n, Int
m))
    | ((Int, Int, a) -> Bool) -> [(Int, Int, a)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any
        (\(Int
i, Int
j, a
v) -> a -> a
forall a. Num a => a -> a
abs (a
v a -> a -> a
forall a. Num a => a -> a -> a
- a -> (Int, Int) -> Map (Int, Int) a -> a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault a
0 (Int
j, Int
i) Map (Int, Int) a
values) a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
1e-12 a -> a -> a
forall a. Num a => a -> a -> a
* a -> a -> a
forall a. Ord a => a -> a -> a
max a
1 (a -> a
forall a. Num a => a -> a
abs a
v))
        (Sparse a -> [(Int, Int, a)]
forall a. Sparse a -> [(Int, Int, a)]
sparseEntries Sparse a
a) =
        MathError -> Either MathError (Sparse a)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"sparseCholesky: symmetric matrix required")
    | Bool
otherwise = do
        [a] -> Either MathError ()
forall a. RealFloat a => [a] -> Either MathError ()
finiteValues (Map (Int, Int) a -> [a]
forall k a. Map k a -> [a]
M.elems Map (Int, Int) a
values)
        l <- (Map (Int, Int) a -> Int -> Either MathError (Map (Int, Int) a))
-> Map (Int, Int) a -> [Int] -> Either MathError (Map (Int, Int) a)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM Map (Int, Int) a -> Int -> Either MathError (Map (Int, Int) a)
row Map (Int, Int) a
forall k a. Map k a
M.empty [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
        sparse n n [(i, j, v) | ((i, j), v) <- M.toAscList l]
  where
    row :: Map (Int, Int) a -> Int -> Either MathError (Map (Int, Int) a)
row Map (Int, Int) a
l Int
i = (Map (Int, Int) a -> Int -> Either MathError (Map (Int, Int) a))
-> Map (Int, Int) a -> [Int] -> Either MathError (Map (Int, Int) a)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (Int
-> Map (Int, Int) a -> Int -> Either MathError (Map (Int, Int) a)
entry Int
i) Map (Int, Int) a
l [Int
0 .. Int
i]
    entry :: Int
-> Map (Int, Int) a -> Int -> Either MathError (Map (Int, Int) a)
entry Int
i Map (Int, Int) a
l Int
j = do
        let value :: a
value =
                a -> (Int, Int) -> Map (Int, Int) a -> a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault a
0 (Int
i, Int
j) Map (Int, Int) a
values
                    a -> a -> a
forall a. Num a => a -> a -> a
- [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [a -> (Int, Int) -> Map (Int, Int) a -> a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault a
0 (Int
i, Int
k) Map (Int, Int) a
l a -> a -> a
forall a. Num a => a -> a -> a
* a -> (Int, Int) -> Map (Int, Int) a -> a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault a
0 (Int
j, Int
k) Map (Int, Int) a
l | Int
k <- [Int
0 .. Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
        if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j Bool -> Bool -> Bool
&& a
value a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0
            then MathError -> Either MathError (Map (Int, Int) a)
forall a b. a -> Either a b
Left (Int -> MathError
NotPositiveDefinite Int
i)
            else do
                let x :: a
x = if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j then a -> a
forall a. Floating a => a -> a
sqrt a
value else a
value a -> a -> a
forall a. Fractional a => a -> a -> a
/ a -> (Int, Int) -> Map (Int, Int) a -> a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault a
0 (Int
j, Int
j) Map (Int, Int) a
l
                [a] -> Either MathError ()
forall a. RealFloat a => [a] -> Either MathError ()
finiteValues [a
x]
                Map (Int, Int) a -> Either MathError (Map (Int, Int) a)
forall a b. b -> Either a b
Right (if a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 then Map (Int, Int) a
l else (Int, Int) -> a -> Map (Int, Int) a -> Map (Int, Int) a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (Int
i, Int
j) a
x Map (Int, Int) a
l)

{- | Opaque sparse QR factors with implicit orthogonal transform; use sparseQRUpper and
sparseQRApplyQT.
-}
data SparseQR a = SparseQR (Sparse a) [(Int, [a])] deriving (SparseQR a -> SparseQR a -> Bool
(SparseQR a -> SparseQR a -> Bool)
-> (SparseQR a -> SparseQR a -> Bool) -> Eq (SparseQR a)
forall a. Eq a => SparseQR a -> SparseQR a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => SparseQR a -> SparseQR a -> Bool
== :: SparseQR a -> SparseQR a -> Bool
$c/= :: forall a. Eq a => SparseQR a -> SparseQR a -> Bool
/= :: SparseQR a -> SparseQR a -> Bool
Eq, Int -> SparseQR a -> ShowS
[SparseQR a] -> ShowS
SparseQR a -> String
(Int -> SparseQR a -> ShowS)
-> (SparseQR a -> String)
-> ([SparseQR a] -> ShowS)
-> Show (SparseQR a)
forall a. Show a => Int -> SparseQR a -> ShowS
forall a. Show a => [SparseQR a] -> ShowS
forall a. Show a => SparseQR a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> SparseQR a -> ShowS
showsPrec :: Int -> SparseQR a -> ShowS
$cshow :: forall a. Show a => SparseQR a -> String
show :: SparseQR a -> String
$cshowList :: forall a. Show a => [SparseQR a] -> ShowS
showList :: [SparseQR a] -> ShowS
Show)

-- | Extract upper-triangular/trapezoidal sparse R from a sparse QR factorization.
sparseQRUpper :: SparseQR a -> Sparse a
sparseQRUpper :: forall a. SparseQR a -> Sparse a
sparseQRUpper (SparseQR Sparse a
r [(Int, [a])]
_) = Sparse a
r

{- | Sparse Householder QR retains only the R nonzeros and compact reflectors;
fill is retained, never silently dropped. Reflection vectors are dense tails.
-}

{- | Given finite nonnegative tolerance then sparse matrix, compute its checked sparse QR
factors. Invalid tolerance, input or numerical failure returns Left MathError.
-}
sparseQR :: (RealFloat a) => Sparse a -> Either MathError (SparseQR a)
sparseQR :: forall a. RealFloat a => Sparse a -> Either MathError (SparseQR a)
sparseQR a :: Sparse a
a@(Sparse Int
m Int
n Map (Int, Int) a
values) = do
    [a] -> Either MathError ()
forall a. RealFloat a => [a] -> Either MathError ()
finiteValues (Map (Int, Int) a -> [a]
forall k a. Map k a -> [a]
M.elems Map (Int, Int) a
values)
    (result, reflectors) <- ((Sparse a, [(Int, [a])])
 -> Int -> Either MathError (Sparse a, [(Int, [a])]))
-> (Sparse a, [(Int, [a])])
-> [Int]
-> Either MathError (Sparse a, [(Int, [a])])
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (Sparse a, [(Int, [a])])
-> Int -> Either MathError (Sparse a, [(Int, [a])])
forall {a}.
RealFloat a =>
(Sparse a, [(Int, [a])])
-> Int -> Either MathError (Sparse a, [(Int, [a])])
step (Sparse a
a, []) [Int
0 .. Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
m Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    Right (SparseQR result reflectors)
  where
    step :: (Sparse a, [(Int, [a])])
-> Int -> Either MathError (Sparse a, [(Int, [a])])
step (Sparse Int
r Int
c Map (Int, Int) a
entriesMap, [(Int, [a])]
reflectors) Int
k = do
        let xs :: [a]
xs = [a -> (Int, Int) -> Map (Int, Int) a -> a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault a
0 (Int
i, Int
k) Map (Int, Int) a
entriesMap | Int
i <- [Int
k .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
            scale :: a
scale = [a] -> a
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (a
0 a -> [a] -> [a]
forall a. a -> [a] -> [a]
: (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map a -> a
forall a. Num a => a -> a
abs [a]
xs)
        if a
scale a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0
            then (Sparse a, [(Int, [a])])
-> Either MathError (Sparse a, [(Int, [a])])
forall a b. b -> Either a b
Right (Int -> Int -> Map (Int, Int) a -> Sparse a
forall a. Int -> Int -> Map (Int, Int) a -> Sparse a
Sparse Int
r Int
c Map (Int, Int) a
entriesMap, [(Int, [a])]
reflectors)
            else do
                let scaled :: [a]
scaled = (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
scale) [a]
xs
                    alpha :: a
alpha = if [a]
scaled [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
0 a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
0 then a -> a
forall a. Num a => a -> a
negate ([a] -> a
forall a. RealFloat a => [a] -> a
stableNorm [a]
scaled) else [a] -> a
forall a. RealFloat a => [a] -> a
stableNorm [a]
scaled
                    vector :: [a]
vector = ([a]
scaled [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
0 a -> a -> a
forall a. Num a => a -> a -> a
- a
alpha) a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
1 [a]
scaled
                    unit :: [a]
unit = (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a -> a -> a
forall a. Fractional a => a -> a -> a
/ [a] -> a
forall a. RealFloat a => [a] -> a
stableNorm [a]
vector) [a]
vector
                    columns :: Map Int [(Int, a)]
columns = ([(Int, a)] -> [(Int, a)] -> [(Int, a)])
-> [(Int, [(Int, a)])] -> Map Int [(Int, a)]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith [(Int, a)] -> [(Int, a)] -> [(Int, a)]
forall a. [a] -> [a] -> [a]
(++) [(Int
j, [(Int
i, a
x)]) | ((Int
i, Int
j), a
x) <- Map (Int, Int) a -> [((Int, Int), a)]
forall k a. Map k a -> [(k, a)]
M.toAscList Map (Int, Int) a
entriesMap, Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
k, Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
k]
                    updates :: [((Int, Int), a)]
updates =
                        [ ((Int
i, Int
j), a
old a -> a -> a
forall a. Num a => a -> a -> a
- a
2 a -> a -> a
forall a. Num a => a -> a -> a
* a
coefficient a -> a -> a
forall a. Num a => a -> a -> a
* ([a]
unit [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k)))
                        | (Int
j, [(Int, a)]
col) <- Map Int [(Int, a)] -> [(Int, [(Int, a)])]
forall k a. Map k a -> [(k, a)]
M.toAscList Map Int [(Int, a)]
columns
                        , let vals :: Map Int a
vals = [(Int, a)] -> Map Int a
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Int, a)]
col
                        , let coefficient :: a
coefficient = [a] -> [a] -> a
forall a. Num a => [a] -> [a] -> a
dotV [a]
unit [a -> Int -> Map Int a -> a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault a
0 Int
i Map Int a
vals | Int
i <- [Int
k .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
                        , Int
i <- [Int
k .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                        , let old :: a
old = a -> Int -> Map Int a -> a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault a
0 Int
i Map Int a
vals
                        ]
                    next :: Map (Int, Int) a
next = (a -> Bool) -> Map (Int, Int) a -> Map (Int, Int) a
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
0) (Map (Int, Int) a -> Map (Int, Int) a -> Map (Int, Int) a
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union ([((Int, Int), a)] -> Map (Int, Int) a
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [((Int, Int), a)]
updates) Map (Int, Int) a
entriesMap)
                    cleaned :: Map (Int, Int) a
cleaned = ((Int, Int) -> a -> Bool) -> Map (Int, Int) a -> Map (Int, Int) a
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (\(Int
i, Int
j) a
_ -> Bool -> Bool
not (Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
k Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
k)) Map (Int, Int) a
next
                [a] -> Either MathError ()
forall a. RealFloat a => [a] -> Either MathError ()
finiteValues (Map (Int, Int) a -> [a]
forall k a. Map k a -> [a]
M.elems Map (Int, Int) a
cleaned)
                (Sparse a, [(Int, [a])])
-> Either MathError (Sparse a, [(Int, [a])])
forall a b. b -> Either a b
Right (Int -> Int -> Map (Int, Int) a -> Sparse a
forall a. Int -> Int -> Map (Int, Int) a -> Sparse a
Sparse Int
r Int
c Map (Int, Int) a
cleaned, [(Int, [a])]
reflectors [(Int, [a])] -> [(Int, [a])] -> [(Int, [a])]
forall a. [a] -> [a] -> [a]
++ [(Int
k, [a]
unit)])

{- | Given sparse QR factors then compatible RHS list, apply transpose(Q) without
constructing dense Q. Invalid length or non-finite arithmetic fails.
-}
sparseQRApplyQT :: (RealFloat a) => SparseQR a -> [a] -> Either MathError [a]
sparseQRApplyQT :: forall a. RealFloat a => SparseQR a -> [a] -> Either MathError [a]
sparseQRApplyQT (SparseQR Sparse a
r [(Int, [a])]
reflectors) [a]
b
    | [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
b Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= (Int, Int) -> Int
forall a b. (a, b) -> a
fst (Sparse a -> (Int, Int)
forall a. Sparse a -> (Int, Int)
sparseShape Sparse a
r) = MathError -> Either MathError [a]
forall a b. a -> Either a b
Left (Int -> Int -> MathError
VectorLengthMismatch ((Int, Int) -> Int
forall a b. (a, b) -> a
fst (Sparse a -> (Int, Int)
forall a. Sparse a -> (Int, Int)
sparseShape Sparse a
r)) ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
b))
    | Bool
otherwise = [a] -> Either MathError ()
forall a. RealFloat a => [a] -> Either MathError ()
finiteValues [a]
b Either MathError () -> Either MathError [a] -> Either MathError [a]
forall a b.
Either MathError a -> Either MathError b -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ([a] -> (Int, [a]) -> Either MathError [a])
-> [a] -> [(Int, [a])] -> Either MathError [a]
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM [a] -> (Int, [a]) -> Either MathError [a]
forall {a}.
RealFloat a =>
[a] -> (Int, [a]) -> Either MathError [a]
apply [a]
b [(Int, [a])]
reflectors
  where
    apply :: [a] -> (Int, [a]) -> Either MathError [a]
apply [a]
xs (Int
start, [a]
v) =
        let coefficient :: a
coefficient = a
2 a -> a -> a
forall a. Num a => a -> a -> a
* [a] -> [a] -> a
forall a. Num a => [a] -> [a] -> a
dotV [a]
v (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
start [a]
xs)
            next :: [a]
next = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
start [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ (a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\a
x a
y -> a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
coefficient a -> a -> a
forall a. Num a => a -> a -> a
* a
y) (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
start [a]
xs) [a]
v
         in [a] -> Either MathError ()
forall a. RealFloat a => [a] -> Either MathError ()
finiteValues [a]
next Either MathError () -> Either MathError [a] -> Either MathError [a]
forall a b.
Either MathError a -> Either MathError b -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [a] -> Either MathError [a]
forall a b. b -> Either a b
Right [a]
next

{- | For a square sparse matrix, return one adjacency row per vertex with stored
(column,value) neighbors. Non-square matrix returns Left MathError.
-}
sparseGraph :: Sparse a -> Either MathError [(Int, [(Int, a)])]
sparseGraph :: forall a. Sparse a -> Either MathError [(Int, [(Int, a)])]
sparseGraph Sparse a
a
    | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
m = MathError -> Either MathError [(Int, [(Int, a)])]
forall a b. a -> Either a b
Left ((Int, Int) -> MathError
NonSquareMatrix (Int
n, Int
m))
    | Bool
otherwise =
        let rows :: Map Int [(Int, a)]
rows = ([(Int, a)] -> [(Int, a)] -> [(Int, a)])
-> [(Int, [(Int, a)])] -> Map Int [(Int, a)]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith [(Int, a)] -> [(Int, a)] -> [(Int, a)]
forall a. [a] -> [a] -> [a]
(++) [(Int
i, [(Int
j, a
v)]) | (Int
i, Int
j, a
v) <- Sparse a -> [(Int, Int, a)]
forall a. Sparse a -> [(Int, Int, a)]
sparseEntries Sparse a
a]
         in [(Int, [(Int, a)])] -> Either MathError [(Int, [(Int, a)])]
forall a b. b -> Either a b
Right [(Int
i, ((Int, a) -> Int) -> [(Int, a)] -> [(Int, a)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Int, a) -> Int
forall a b. (a, b) -> a
fst ([(Int, a)] -> Int -> Map Int [(Int, a)] -> [(Int, a)]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [] Int
i Map Int [(Int, a)]
rows)) | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
  where
    (Int
n, Int
m) = Sparse a -> (Int, Int)
forall a. Sparse a -> (Int, Int)
sparseShape Sparse a
a

{- | Logical payload bytes: two machine-width indices and caller-provided scalar
width per nonzero. Runtime Map node overhead is implementation dependent.
-}
sparseStorageBytes :: Int -> Sparse a -> Either MathError Integer
sparseStorageBytes :: forall a. Int -> Sparse a -> Either MathError Integer
sparseStorageBytes Int
scalarBytes Sparse a
a
    | Int
scalarBytes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"sparseStorageBytes: negative scalar width")
    | Bool
otherwise =
        Integer -> Either MathError Integer
forall a b. b -> Either a b
Right
            ( Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Sparse a -> Int
forall a. Sparse a -> Int
nonzeros Sparse a
a) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* (Integer
2 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int -> Int
forall a. Storable a => a -> Int
sizeOf (Int
forall a. HasCallStack => a
undefined :: Int)) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
scalarBytes)
                Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
2 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int -> Int
forall a. Storable a => a -> Int
sizeOf (Int
forall a. HasCallStack => a
undefined :: Int))
            )

{- | Canonical sparse multidimensional entries, constructed by sparseTensor with checked
coordinates and duplicate coalescing.
-}
data SparseTensor a = SparseTensor [Int] (M.Map [Int] a) deriving (SparseTensor a -> SparseTensor a -> Bool
(SparseTensor a -> SparseTensor a -> Bool)
-> (SparseTensor a -> SparseTensor a -> Bool)
-> Eq (SparseTensor a)
forall a. Eq a => SparseTensor a -> SparseTensor a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => SparseTensor a -> SparseTensor a -> Bool
== :: SparseTensor a -> SparseTensor a -> Bool
$c/= :: forall a. Eq a => SparseTensor a -> SparseTensor a -> Bool
/= :: SparseTensor a -> SparseTensor a -> Bool
Eq, Int -> SparseTensor a -> ShowS
[SparseTensor a] -> ShowS
SparseTensor a -> String
(Int -> SparseTensor a -> ShowS)
-> (SparseTensor a -> String)
-> ([SparseTensor a] -> ShowS)
-> Show (SparseTensor a)
forall a. Show a => Int -> SparseTensor a -> ShowS
forall a. Show a => [SparseTensor a] -> ShowS
forall a. Show a => SparseTensor a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> SparseTensor a -> ShowS
showsPrec :: Int -> SparseTensor a -> ShowS
$cshow :: forall a. Show a => SparseTensor a -> String
show :: SparseTensor a -> String
$cshowList :: forall a. Show a => [SparseTensor a] -> ShowS
showList :: [SparseTensor a] -> ShowS
Show)

{- | Given dimensions then (coordinate-list,value) entries, validate rank/bounds, sum
duplicate coordinates and remove zeros. Invalid shape/indices return Left MathError.
-}
sparseTensor :: (Eq a, Num a) => [Int] -> [([Int], a)] -> Either MathError (SparseTensor a)
sparseTensor :: forall a.
(Eq a, Num a) =>
[Int] -> [([Int], a)] -> Either MathError (SparseTensor a)
sparseTensor [Int]
dims [([Int], a)]
values
    | (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) [Int]
dims
        Bool -> Bool -> Bool
|| (([Int], a) -> Bool) -> [([Int], a)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\([Int]
is, a
_) -> [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
is Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims Bool -> Bool -> Bool
|| [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or ((Int -> Int -> Bool) -> [Int] -> [Int] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
i Int
d -> Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
d) [Int]
is [Int]
dims)) [([Int], a)]
values =
        MathError -> Either MathError (SparseTensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"sparseTensor: shape or index invalid")
    | Bool
otherwise = SparseTensor a -> Either MathError (SparseTensor a)
forall a b. b -> Either a b
Right ([Int] -> Map [Int] a -> SparseTensor a
forall a. [Int] -> Map [Int] a -> SparseTensor a
SparseTensor [Int]
dims ((a -> Bool) -> Map [Int] a -> Map [Int] a
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
0) ((a -> a -> a) -> [([Int], a)] -> Map [Int] a
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith a -> a -> a
forall a. Num a => a -> a -> a
(+) [([Int], a)]
values)))

-- | Return axis dimensions of a sparse tensor; [] denotes a scalar.
sparseTensorShape :: SparseTensor a -> [Int]
sparseTensorShape :: forall a. SparseTensor a -> [Int]
sparseTensorShape (SparseTensor [Int]
dims Map [Int] a
_) = [Int]
dims

{- | Return canonical nonzero (coordinate list,value) pairs in lexicographic coordinate
order.
-}
sparseTensorEntries :: SparseTensor a -> [([Int], a)]
sparseTensorEntries :: forall a. SparseTensor a -> [([Int], a)]
sparseTensorEntries (SparseTensor [Int]
_ Map [Int] a
xs) = Map [Int] a -> [([Int], a)]
forall k a. Map k a -> [(k, a)]
M.toAscList Map [Int] a
xs

{- | Given sparse tensor then full zero-based coordinates, return stored value or zero.
Wrong rank or out-of-range indices fail.
-}
sparseTensorLookup :: (Num a) => SparseTensor a -> [Int] -> Either MathError a
sparseTensorLookup :: forall a. Num a => SparseTensor a -> [Int] -> Either MathError a
sparseTensorLookup (SparseTensor [Int]
dims Map [Int] a
xs) [Int]
is
    | [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
is Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims Bool -> Bool -> Bool
|| [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or ((Int -> Int -> Bool) -> [Int] -> [Int] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
i Int
d -> Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
d) [Int]
is [Int]
dims) =
        MathError -> Either MathError a
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"sparseTensorLookup: index out of bounds")
    | Bool
otherwise = a -> Either MathError a
forall a b. b -> Either a b
Right (a -> [Int] -> Map [Int] a -> a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault a
0 [Int]
is Map [Int] a
xs)

{- | Small-problem spectral reference backend. The caller explicitly authorizes
a dense element budget; oversized matrices fail before generating zeros.
This is not a scalable sparse eigensolver/SVD and never silently densifies.
-}
sparseEigenReference :: Integer -> Sparse Double -> Either MathError Eigen
sparseEigenReference :: Integer -> Sparse Double -> Either MathError Eigen
sparseEigenReference Integer
budget Sparse Double
a = Integer -> Sparse Double -> Either MathError (Matrix Double)
referenceDense Integer
budget Sparse Double
a Either MathError (Matrix Double)
-> (Matrix Double -> Either MathError Eigen)
-> Either MathError Eigen
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Matrix Double -> Either MathError Eigen
eigen

{- | Given explicit dense-byte budget then sparse Double matrix, densify only if within
budget and run checked dense SVD. This is a reference fallback, not a sparse iterative
SVD; budget/backend/numeric errors fail.
-}
sparseSVDReference :: Integer -> Sparse Double -> Either MathError SVD
sparseSVDReference :: Integer -> Sparse Double -> Either MathError SVD
sparseSVDReference Integer
budget Sparse Double
a = Integer -> Sparse Double -> Either MathError (Matrix Double)
referenceDense Integer
budget Sparse Double
a Either MathError (Matrix Double)
-> (Matrix Double -> Either MathError SVD) -> Either MathError SVD
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Matrix Double -> Either MathError SVD
svd

referenceDense :: Integer -> Sparse Double -> Either MathError (Matrix Double)
referenceDense :: Integer -> Sparse Double -> Either MathError (Matrix Double)
referenceDense Integer
budget Sparse Double
a
    | Integer
budget Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 Bool -> Bool -> Bool
|| Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
r Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
c Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
budget =
        MathError -> Either MathError (Matrix Double)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"sparse spectral reference: dense element budget exceeded")
    | Bool
otherwise = Sparse Double -> Either MathError (Matrix Double)
forall a. Num a => Sparse a -> Either MathError (Matrix a)
sparseToDense Sparse Double
a
  where
    (Int
r, Int
c) = Sparse Double -> (Int, Int)
forall a. Sparse a -> (Int, Int)
sparseShape Sparse Double
a