module Matrix.QR.Givens  (
   leastSquares,
   decompose, solve, det, detAbs,
   Rotation, rotateVector,
   Upper, solveUpper, detUpper,
   ) where

import qualified Matrix.Sparse as Sparse
import DSP.Basic (toMaybe, (^!))

import Control.Monad (mfilter)

import qualified Data.Foldable as Fold
import qualified Data.List as List
import qualified Data.Map as Map
import Data.Map (Map)
import Data.Array
         (Array, Ix, array, bounds, elems, rangeSize, range, (!), (//), )


data Rotation i a = Rotation (i,i) (a,a)
   deriving Int -> Rotation i a -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall i a. (Show i, Show a) => Int -> Rotation i a -> ShowS
forall i a. (Show i, Show a) => [Rotation i a] -> ShowS
forall i a. (Show i, Show a) => Rotation i a -> String
showList :: [Rotation i a] -> ShowS
$cshowList :: forall i a. (Show i, Show a) => [Rotation i a] -> ShowS
show :: Rotation i a -> String
$cshow :: forall i a. (Show i, Show a) => Rotation i a -> String
showsPrec :: Int -> Rotation i a -> ShowS
$cshowsPrec :: forall i a. (Show i, Show a) => Int -> Rotation i a -> ShowS
Show

data Upper i j a = Upper ((i,j), (i,j)) (Map i (Map j a))
   deriving Int -> Upper i j a -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall i j a.
(Show i, Show j, Show a) =>
Int -> Upper i j a -> ShowS
forall i j a. (Show i, Show j, Show a) => [Upper i j a] -> ShowS
forall i j a. (Show i, Show j, Show a) => Upper i j a -> String
showList :: [Upper i j a] -> ShowS
$cshowList :: forall i j a. (Show i, Show j, Show a) => [Upper i j a] -> ShowS
show :: Upper i j a -> String
$cshow :: forall i j a. (Show i, Show j, Show a) => Upper i j a -> String
showsPrec :: Int -> Upper i j a -> ShowS
$cshowsPrec :: forall i j a.
(Show i, Show j, Show a) =>
Int -> Upper i j a -> ShowS
Show

{- |
The decomposition routine is pretty simple.
It does not try to minimize fill-up by a clever ordering of rotations.
However, for banded matrices it will work as expected.
-}
decompose :: (Ix i, Enum i, Ix j, Enum j, RealFloat a) =>
      Sparse.Matrix i j a -- ^ A
   -> ([Rotation i a], Upper i j a) -- ^ QR(A)
decompose :: forall i j a.
(Ix i, Enum i, Ix j, Enum j, RealFloat a) =>
Matrix i j a -> ([Rotation i a], Upper i j a)
decompose Matrix i j a
a =
   (\([[Rotation i a]]
qs,[(i, Map j a)]
rows) -> (forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Rotation i a]]
qs, forall i j a. ((i, j), (i, j)) -> Map i (Map j a) -> Upper i j a
Upper (forall i j a. Matrix i j a -> ((i, j), (i, j))
Sparse.bounds Matrix i j a
a) forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(i, Map j a)]
rows)) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
   forall a b. [(a, b)] -> ([a], [b])
unzip forall b c a. (b -> c) -> (a -> b) -> a -> c
.
   forall b a. (b -> Maybe (a, b)) -> b -> [a]
List.unfoldr
      (\Matrix i j a
a0 ->
         let bnds :: ((i, j), (i, j))
bnds@((i
m0,j
_), (i, j)
_) = forall i j a. Matrix i j a -> ((i, j), (i, j))
Sparse.bounds Matrix i j a
a0
             ([Rotation i a]
q,Matrix i j a
a1) = forall i j a.
(Ord i, Ord j, RealFloat a) =>
Matrix i j a -> ([Rotation i a], Matrix i j a)
step Matrix i j a
a0
         in  forall a. Bool -> a -> Maybe a
toMaybe (Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall i. Ix i => (i, i) -> Bool
emptyRange ((i, j), (i, j))
bnds) forall a b. (a -> b) -> a -> b
$
               (([Rotation i a]
q, (i
m0, forall i j a. (Ord i, Ord j) => i -> Matrix i j a -> Map j a
Sparse.getRow i
m0 Matrix i j a
a1)), forall i j a.
(Ord i, Enum i, Ord j, Enum j) =>
Matrix i j a -> Matrix i j a
submatrix Matrix i j a
a1))
    forall a b. (a -> b) -> a -> b
$ Matrix i j a
a

-- cf. QR.Householder
emptyRange :: (Ix i) => (i,i) -> Bool
emptyRange :: forall i. Ix i => (i, i) -> Bool
emptyRange = forall (t :: * -> *) a. Foldable t => t a -> Bool
null forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ix a => (a, a) -> [a]
range

-- | assumes that the first column is empty except the top-most element
submatrix ::
   (Ord i, Enum i, Ord j, Enum j) => Sparse.Matrix i j a -> Sparse.Matrix i j a
submatrix :: forall i j a.
(Ord i, Enum i, Ord j, Enum j) =>
Matrix i j a -> Matrix i j a
submatrix Matrix i j a
a =
   let ((i
m0,j
n0), (i, j)
mn1) = forall i j a. Matrix i j a -> ((i, j), (i, j))
Sparse.bounds Matrix i j a
a
   in  forall i j a.
(Ord i, Ord j) =>
((i, j), (i, j)) -> Map i (Map j a) -> Matrix i j a
Sparse.fromRows ((forall a. Enum a => a -> a
succ i
m0, forall a. Enum a => a -> a
succ j
n0), (i, j)
mn1) forall a b. (a -> b) -> a -> b
$
       forall k a. Ord k => k -> Map k a -> Map k a
Map.delete i
m0 forall a b. (a -> b) -> a -> b
$ forall i j a. (Ord i, Ord j) => Matrix i j a -> Map i (Map j a)
Sparse.toRows Matrix i j a
a

step ::
   (Ord i, Ord j, RealFloat a) =>
   Sparse.Matrix i j a -> ([Rotation i a], Sparse.Matrix i j a)
step :: forall i j a.
(Ord i, Ord j, RealFloat a) =>
Matrix i j a -> ([Rotation i a], Matrix i j a)
step Matrix i j a
a =
   let bnds :: ((i, j), (i, j))
bnds@((i
m0,j
n0), (i, j)
_) = forall i j a. Matrix i j a -> ((i, j), (i, j))
Sparse.bounds Matrix i j a
a
       rows :: Map i (Map j a)
rows = forall i j a. (Ord i, Ord j) => Matrix i j a -> Map i (Map j a)
Sparse.toRows Matrix i j a
a
       topRow :: Map j a
topRow = forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault forall k a. Map k a
Map.empty i
m0 Map i (Map j a)
rows
   in  (\((a
xi,Map j a
xrem), Map i (Maybe (Rotation i a), Map j a)
finalRows) ->
         (forall k a. Map k a -> [a]
Map.elems forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> Maybe b) -> Map k a -> Map k b
Map.mapMaybe forall a b. (a, b) -> a
fst Map i (Maybe (Rotation i a), Map j a)
finalRows,
          forall i j a.
(Ord i, Ord j) =>
((i, j), (i, j)) -> Map i (Map j a) -> Matrix i j a
Sparse.fromRows ((i, j), (i, j))
bnds forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert i
m0 (forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert j
n0 a
xi Map j a
xrem) forall a b. (a -> b) -> a -> b
$
            forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> b
snd Map i (Maybe (Rotation i a), Map j a)
finalRows)) forall a b. (a -> b) -> a -> b
$
       forall a k b c.
(a -> k -> b -> (a, c)) -> a -> Map k b -> (a, Map k c)
Map.mapAccumWithKey
         (\(a
xi,Map j a
xrem) i
mi Map j a
yrow ->
            let yrem :: Map j a
yrem = forall k a. Ord k => k -> Map k a -> Map k a
Map.delete j
n0 Map j a
yrow
                jrot :: (a, a) -> Maybe (Rotation i a)
jrot = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall i a. (i, i) -> (a, a) -> Rotation i a
Rotation (i
m0,i
mi)
            in  case forall (m :: * -> *) a. MonadPlus m => (a -> Bool) -> m a -> m a
mfilter (a
0forall a. Eq a => a -> a -> Bool
/=) forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup j
n0 Map j a
yrow of
                  Maybe a
Nothing -> ((a
xi,Map j a
xrem),(forall a. Maybe a
Nothing,Map j a
yrem))
                  Just a
yi ->
                     if a
xiforall a. Eq a => a -> a -> Bool
==a
0
                        then ((a
yi,Map j a
yrem), (forall {a}. (a, a) -> Maybe (Rotation i a)
jrot (a
0,-a
1), forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. Num a => a -> a
negate Map j a
xrem))
                        else
                           let rot :: (a, a)
rot = forall a. RealFloat a => (a, a) -> (a, a)
rotationFromXY (a
xi,a
yi)
                               (Map j a
rx,Map j a
ry) = forall j a.
(Ord j, Num a) =>
(a, a) -> (Map j a, Map j a) -> (Map j a, Map j a)
rotateRows (a, a)
rot (Map j a
xrem, Map j a
yrem)
                           in  ((forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall a. Num a => (a, a) -> (a, a) -> (a, a)
rotateSingle (a, a)
rot (a
xi,a
yi), Map j a
rx),
                                (forall {a}. (a, a) -> Maybe (Rotation i a)
jrot (a, a)
rot, Map j a
ry)))
         (forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault a
0 j
n0 Map j a
topRow, forall k a. Ord k => k -> Map k a -> Map k a
Map.delete j
n0 Map j a
topRow)
         (forall k a. Ord k => k -> Map k a -> Map k a
Map.delete i
m0 Map i (Map j a)
rows)

-- | The argument must not be (0,0).
rotationFromXY :: (RealFloat a) => (a,a) -> (a,a)
rotationFromXY :: forall a. RealFloat a => (a, a) -> (a, a)
rotationFromXY (a
x,a
y) =
   if forall a. Num a => a -> a
abs a
x forall a. Ord a => a -> a -> Bool
> forall a. Num a => a -> a
abs a
y
     then let q :: a
q = a
yforall a. Fractional a => a -> a -> a
/a
x; k :: a
k = forall a. Floating a => a -> a
recipNorm a
q in (a
k,-a
qforall a. Num a => a -> a -> a
*a
k)
     else let q :: a
q = a
xforall a. Fractional a => a -> a -> a
/a
y; k :: a
k = forall a. Floating a => a -> a
recipNorm a
q in (a
qforall a. Num a => a -> a -> a
*a
k,-a
k)

recipNorm :: Floating a => a -> a
recipNorm :: forall a. Floating a => a -> a
recipNorm a
q = forall a. Fractional a => a -> a
recip forall a b. (a -> b) -> a -> b
$ forall a. Floating a => a -> a
sqrt (a
1forall a. Num a => a -> a -> a
+a
qforall a. Num a => a -> Int -> a
^!Int
2)

rotateSingle :: (Num a) => (a,a) -> (a,a) -> (a,a)
rotateSingle :: forall a. Num a => (a, a) -> (a, a) -> (a, a)
rotateSingle (a
c,a
s) (a
x,a
y) = (a
cforall a. Num a => a -> a -> a
*a
xforall a. Num a => a -> a -> a
-a
sforall a. Num a => a -> a -> a
*a
y, a
sforall a. Num a => a -> a -> a
*a
xforall a. Num a => a -> a -> a
+a
cforall a. Num a => a -> a -> a
*a
y)

rotateRows ::
   (Ord j, Num a) => (a,a) -> (Map j a, Map j a) -> (Map j a, Map j a)
rotateRows :: forall j a.
(Ord j, Num a) =>
(a, a) -> (Map j a, Map j a) -> (Map j a, Map j a)
rotateRows (a
c,a
s) (Map j a
xs,Map j a
ys) =
   let rs :: Map j (a, a)
rs =
         forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
Map.intersectionWith (forall a b c. ((a, b) -> c) -> a -> b -> c
curry forall a b. (a -> b) -> a -> b
$ forall a. Num a => (a, a) -> (a, a) -> (a, a)
rotateSingle (a
c,a
s)) Map j a
xs Map j a
ys
         forall k a. Ord k => Map k a -> Map k a -> Map k a
`Map.union`
         forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\a
x -> ( a
cforall a. Num a => a -> a -> a
*a
x, a
sforall a. Num a => a -> a -> a
*a
x)) (forall k a b. Ord k => Map k a -> Map k b -> Map k a
Map.difference Map j a
xs Map j a
ys)
         forall k a. Ord k => Map k a -> Map k a -> Map k a
`Map.union`
         forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\a
y -> (-a
sforall a. Num a => a -> a -> a
*a
y, a
cforall a. Num a => a -> a -> a
*a
y)) (forall k a b. Ord k => Map k a -> Map k b -> Map k a
Map.difference Map j a
ys Map j a
xs)

   in  (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst Map j (a, a)
rs, forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> b
snd Map j (a, a)
rs)

rotateVector :: (Ix i, Num a) => Rotation i a -> Array i a -> Array i a
rotateVector :: forall i a. (Ix i, Num a) => Rotation i a -> Array i a -> Array i a
rotateVector (Rotation (i
i0,i
i1) (a, a)
cs) Array i a
x =
   let (a
y0,a
y1) = forall a. Num a => (a, a) -> (a, a) -> (a, a)
rotateSingle (a, a)
cs (Array i a
xforall i e. Ix i => Array i e -> i -> e
!i
i0,Array i a
xforall i e. Ix i => Array i e -> i -> e
!i
i1)
   in  Array i a
x forall i e. Ix i => Array i e -> [(i, e)] -> Array i e
// [(i
i0,a
y0),(i
i1,a
y1)]


{- |
Assumes that 'Upper' matrix is at least as high as wide
and that it has full rank.
-}
solveUpper ::
   (Ix i, Ix j, Fractional a) => Upper i j a -> Array i a -> Array j a
solveUpper :: forall i j a.
(Ix i, Ix j, Fractional a) =>
Upper i j a -> Array i a -> Array j a
solveUpper (Upper ((i
m0,j
n0), (i
m1,j
n1)) Map i (Map j a)
rows0) Array i a
b =
   if forall i e. Array i e -> (i, i)
bounds Array i a
b forall a. Eq a => a -> a -> Bool
== (i
m0,i
m1)
     then
         forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (j
n0,j
n1) forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
Map.toList forall a b. (a -> b) -> a -> b
$
         forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
            (\(Map j a
row,a
bi) Map j a
xs ->
               let ((j
j,a
a),Map j a
as) = forall k a. Map k a -> ((k, a), Map k a)
Map.deleteFindMin Map j a
row
               in  forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert j
j
                     ((a
bi forall a. Num a => a -> a -> a
- forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
Fold.sum (forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
Map.intersectionWith forall a. Num a => a -> a -> a
(*) Map j a
as Map j a
xs)) forall a. Fractional a => a -> a -> a
/ a
a) Map j a
xs)
            forall k a. Map k a
Map.empty
            (forall a b. [a] -> [b] -> [(a, b)]
zip (forall k a. Map k a -> [a]
Map.elems Map i (Map j a)
rows0) (forall i e. Array i e -> [e]
elems Array i a
b))
     else forall a. HasCallStack => String -> a
error String
"solveUpper: vertical bounds mismatch"

solve ::
   (Ix i, Ix j, Fractional a) =>
   ([Rotation i a], Upper i j a) -> Array i a -> Array j a
solve :: forall i j a.
(Ix i, Ix j, Fractional a) =>
([Rotation i a], Upper i j a) -> Array i a -> Array j a
solve ([Rotation i a]
qs, Upper i j a
u) Array i a
b = forall i j a.
(Ix i, Ix j, Fractional a) =>
Upper i j a -> Array i a -> Array j a
solveUpper Upper i j a
u forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall i a. (Ix i, Num a) => Rotation i a -> Array i a -> Array i a
rotateVector) Array i a
b [Rotation i a]
qs

{- |
Solve a sparse overconstrained linear problem, i.e. minimize @||Ax-b||@.
@A@ must have dimensions @m x n@ with @m>=n@ and it must have full-rank.
None of these conditions is checked.
-}
leastSquares ::
   (Ix i, Enum i, Ix j, Enum j, RealFloat a) =>
   Sparse.Matrix i j a -> Array i a -> Array j a
leastSquares :: forall i j a.
(Ix i, Enum i, Ix j, Enum j, RealFloat a) =>
Matrix i j a -> Array i a -> Array j a
leastSquares = forall i j a.
(Ix i, Ix j, Fractional a) =>
([Rotation i a], Upper i j a) -> Array i a -> Array j a
solve forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall i j a.
(Ix i, Enum i, Ix j, Enum j, RealFloat a) =>
Matrix i j a -> ([Rotation i a], Upper i j a)
decompose


detUpper ::
   (Ix i, Ix j, Fractional a) => Upper i j a -> a
detUpper :: forall i j a. (Ix i, Ix j, Fractional a) => Upper i j a -> a
detUpper (Upper ((i
_m0,j
n0), (i
_m1,j
n1)) Map i (Map j a)
rows) =
   if forall a. Ix a => (a, a) -> Int
rangeSize (j
n0,j
n1) forall a. Eq a => a -> a -> Bool
== forall k a. Map k a -> Int
Map.size Map i (Map j a)
rows
     then forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> (k, a)
Map.findMin) forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [a]
Map.elems Map i (Map j a)
rows
     else a
0

{- |
Only sensible for square matrices,
but the function does not check whether the matrix is square.
-}
det :: (Ix i, Enum i, Ix j, Enum j, RealFloat a) => Sparse.Matrix i j a -> a
det :: forall i j a.
(Ix i, Enum i, Ix j, Enum j, RealFloat a) =>
Matrix i j a -> a
det = forall i j a. (Ix i, Ix j, Fractional a) => Upper i j a -> a
detUpper forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall i j a.
(Ix i, Enum i, Ix j, Enum j, RealFloat a) =>
Matrix i j a -> ([Rotation i a], Upper i j a)
decompose

{- |
Absolute value of the determinant.
This is also sound for non-square matrices.
-}
detAbs :: (Ix i, Enum i, Ix j, Enum j, RealFloat a) => Sparse.Matrix i j a -> a
detAbs :: forall i j a.
(Ix i, Enum i, Ix j, Enum j, RealFloat a) =>
Matrix i j a -> a
detAbs = forall a. Num a => a -> a
abs forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall i j a.
(Ix i, Enum i, Ix j, Enum j, RealFloat a) =>
Matrix i j a -> a
det