module Matrix.QR.Householder (
   leastSquares,
   decompose, solve, det, detAbs,
   Reflection, reflectMatrix, reflectVector,
   Upper, matrixFromUpper, solveUpper, detUpper,
   ) where

import Matrix.Matrix (mv_mult, m_trans, getRow, getColumn, inner, outer)
import Matrix.Vector (sub, scale, norm)
import DSP.Basic (toMaybe)

import qualified Data.List as List
import Data.Array
         (Array, Ix, bounds, elems, range, rangeSize,
          accum, accumArray, assocs, ixmap, listArray, (!), (//), )


decompose :: (Ix i, Enum i, Ix j, Enum j, RealFloat a) =>
      Array (i,j) a -- ^ A
   -> ([Reflection i a], Upper i j a) -- ^ QR(A)
decompose :: forall i j a.
(Ix i, Enum i, Ix j, Enum j, RealFloat a) =>
Array (i, j) a -> ([Reflection i a], Upper i j a)
decompose Array (i, j) a
a =
   (\([Reflection i a]
qs,[Array j a]
rows) -> ([Reflection i a]
qs, forall i j a. ((i, j), (i, j)) -> [Array j a] -> Upper i j a
Upper (forall i e. Array i e -> (i, i)
bounds Array (i, j) a
a) [Array 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
      (\Array (i, j) a
a0 ->
         let bnds :: ((i, j), (i, j))
bnds@((i
m0,j
_), (i, j)
_) = forall i e. Array i e -> (i, i)
bounds Array (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
$
               let (Reflection i a
q,Array (i, j) a
a1) = forall i j a.
(Ix i, Ix j, RealFloat a) =>
Array (i, j) a -> (Reflection i a, Array (i, j) a)
step Array (i, j) a
a0
               in  ((Reflection i a
q, forall i j e. (Ix i, Ix j) => i -> Array (i, j) e -> Array j e
getRow i
m0 Array (i, j) a
a1), forall i j e.
(Ix i, Enum i, Ix j, Enum j) =>
Array (i, j) e -> Array (i, j) e
submatrix Array (i, j) a
a1))
    forall a b. (a -> b) -> a -> b
$ Array (i, j) a
a

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

step ::
   (Ix i, Ix j, RealFloat a) =>
   Array (i,j) a -> (Reflection i a, Array (i,j) a)
step :: forall i j a.
(Ix i, Ix j, RealFloat a) =>
Array (i, j) a -> (Reflection i a, Array (i, j) a)
step Array (i, j) a
a =
   let (i
m0,j
n0) = forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall i e. Array i e -> (i, i)
bounds Array (i, j) a
a
       z :: Array i a
z = forall i j e. (Ix i, Ix j) => j -> Array (i, j) e -> Array i e
getColumn j
n0 Array (i, j) a
a
       sign :: a -> a
sign a
x = if a
xforall a. Ord a => a -> a -> Bool
<a
0 then -a
1 else a
1
       q :: Reflection i a
q = forall i a. (Ix i, Floating a) => Array i a -> Reflection i a
reflection forall a b. (a -> b) -> a -> b
$ forall i e a.
Ix i =>
(e -> a -> e) -> Array i e -> [(i, a)] -> Array i e
accum forall a. Num a => a -> a -> a
(+) Array i a
z [(i
m0, forall {a} {a}. (Ord a, Num a, Num a) => a -> a
sign(Array i a
zforall i e. Ix i => Array i e -> i -> e
!i
m0) forall a. Num a => a -> a -> a
* forall i a. (Ix i, Floating a) => Array i a -> a
norm Array i a
z)]
   in  (Reflection i a
q, forall i j a.
(Ix i, Ix j, Num a) =>
Reflection i a -> Array (i, j) a -> Array (i, j) a
reflectMatrix Reflection i a
q Array (i, j) a
a)

{-
Submatrices with only Ix constrained indices would not work,
because we cannot reduce a two-dimensional array by only one element.
-}
submatrix :: (Ix i, Enum i, Ix j, Enum j) => Array (i,j) e -> Array (i,j) e
submatrix :: forall i j e.
(Ix i, Enum i, Ix j, Enum j) =>
Array (i, j) e -> Array (i, j) e
submatrix Array (i, j) e
a =
   let ((i
m0,j
n0), (i
m1,j
n1)) = forall i e. Array i e -> (i, i)
bounds Array (i, j) e
a
   in  forall i j e.
(Ix i, Ix j) =>
(i, i) -> (i -> j) -> Array j e -> Array i e
ixmap ((forall a. Enum a => a -> a
succ i
m0, forall a. Enum a => a -> a
succ j
n0), (i
m1,j
n1)) forall a. a -> a
id Array (i, j) e
a


data Upper i j a = Upper ((i,j), (i,j)) [Array j a]

matrixFromUpper :: (Ix i, Ix j, Num a) => Upper i j a -> Array (i,j) a
matrixFromUpper :: forall i j a. (Ix i, Ix j, Num a) => Upper i j a -> Array (i, j) a
matrixFromUpper (Upper bnds :: ((i, j), (i, j))
bnds@((i
m0,j
_n0), (i
m1,j
_n1)) [Array j a]
rows) =
   forall i e a.
Ix i =>
(e -> a -> e) -> e -> (i, i) -> [(i, a)] -> Array i e
accumArray (forall a b. a -> b -> a
const forall a. a -> a
id) a
0 ((i, j), (i, j))
bnds forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall a b. (a -> b) -> a -> b
$
   forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\i
k -> forall a b. (a -> b) -> [a] -> [b]
map (\(j
j,a
a) -> ((i
k,j
j),a
a)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall i e. Ix i => Array i e -> [(i, e)]
assocs) (forall a. Ix a => (a, a) -> [a]
range (i
m0,i
m1)) [Array j a]
rows


newtype Reflection i a = Reflection (Array i a)

reflection :: (Ix i, Floating a) => Array i a -> Reflection i a
reflection :: forall i a. (Ix i, Floating a) => Array i a -> Reflection i a
reflection Array i a
v =
   let normv :: a
normv = forall i a. (Ix i, Floating a) => Array i a -> a
norm Array i a
v
   in  forall i a. Array i a -> Reflection i a
Reflection forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. Fractional a => a -> a -> a
/ ((a
1forall a. Num a => a -> a -> a
-forall a. Num a => a -> a
signum a
normv) forall a. Num a => a -> a -> a
+ a
normv)) Array i a
v

reflectMatrixFull ::
   (Ix i, Ix j, Num a) => Reflection i a -> Array (i,j) a -> Array (i,j) a
reflectMatrixFull :: forall i j a.
(Ix i, Ix j, Num a) =>
Reflection i a -> Array (i, j) a -> Array (i, j) a
reflectMatrixFull (Reflection Array i a
v) Array (i, j) a
a =
   forall i a. (Ix i, Num a) => Array i a -> Array i a -> Array i a
sub Array (i, j) a
a forall a b. (a -> b) -> a -> b
$ forall i a. (Ix i, Num a) => a -> Array i a -> Array i a
scale a
2 forall a b. (a -> b) -> a -> b
$ forall i j a.
(Ix i, Ix j, Num a) =>
Array i a -> Array j a -> Array (i, j) a
outer Array i a
v forall a b. (a -> b) -> a -> b
$ forall i j a.
(Ix i, Ix j, Num a) =>
Array (i, j) a -> Array j a -> Array i a
mv_mult (forall i j a.
(Ix i, Ix j, Num a) =>
Array (i, j) a -> Array (j, i) a
m_trans Array (i, j) a
a) Array i a
v

reflectMatrix ::
   (Ix i, Ix j, Num a) => Reflection i a -> Array (i,j) a -> Array (i,j) a
reflectMatrix :: forall i j a.
(Ix i, Ix j, Num a) =>
Reflection i a -> Array (i, j) a -> Array (i, j) a
reflectMatrix q :: Reflection i a
q@(Reflection Array i a
v) Array (i, j) a
a =
   let (i
k0,i
k1) = forall i e. Array i e -> (i, i)
bounds Array i a
v
       ((i
m0,j
n0), (i
m1,j
n1)) = forall i e. Array i e -> (i, i)
bounds Array (i, j) a
a
       bnds :: ((i, j), (i, j))
bnds = ((i
k0,j
n0),(i
k1,j
n1))
   in  case (forall a. Ord a => a -> a -> Ordering
compare i
k0 i
m0, forall a. Ord a => a -> a -> Ordering
compare i
k1 i
m1) of
         (Ordering
EQ,Ordering
EQ) -> forall i j a.
(Ix i, Ix j, Num a) =>
Reflection i a -> Array (i, j) a -> Array (i, j) a
reflectMatrixFull Reflection i a
q Array (i, j) a
a
         (Ordering
LT,Ordering
_) -> forall a. HasCallStack => [Char] -> a
error [Char]
"reflectMatrix: lower reflection dimension too small"
         (Ordering
_,Ordering
GT) -> forall a. HasCallStack => [Char] -> a
error [Char]
"reflectMatrix: upper reflection dimension too big"
         (Ordering, Ordering)
_ -> forall i a. Ix i => Array i a -> Array i a -> Array i a
replaceSubArray Array (i, j) a
a forall a b. (a -> b) -> a -> b
$ forall i j a.
(Ix i, Ix j, Num a) =>
Reflection i a -> Array (i, j) a -> Array (i, j) a
reflectMatrixFull Reflection i a
q forall a b. (a -> b) -> a -> b
$ forall i a. Ix i => (i, i) -> Array i a -> Array i a
subArray ((i, j), (i, j))
bnds Array (i, j) a
a


reflectVectorFull :: (Ix i, Num a) => Reflection i a -> Array i a -> Array i a
reflectVectorFull :: forall i a.
(Ix i, Num a) =>
Reflection i a -> Array i a -> Array i a
reflectVectorFull (Reflection Array i a
v) Array i a
a = forall i a. (Ix i, Num a) => Array i a -> Array i a -> Array i a
sub Array i a
a forall a b. (a -> b) -> a -> b
$ forall i a. (Ix i, Num a) => a -> Array i a -> Array i a
scale (a
2 forall a. Num a => a -> a -> a
* forall i a. (Ix i, Num a) => Array i a -> Array i a -> a
inner Array i a
v Array i a
a) Array i a
v

reflectVector :: (Ix i, Num a) => Reflection i a -> Array i a -> Array i a
reflectVector :: forall i a.
(Ix i, Num a) =>
Reflection i a -> Array i a -> Array i a
reflectVector q :: Reflection i a
q@(Reflection Array i a
v) Array i a
a =
   let bnds :: (i, i)
bnds@(i
k0,i
k1) = forall i e. Array i e -> (i, i)
bounds Array i a
v
       (i
m0,i
m1) = forall i e. Array i e -> (i, i)
bounds Array i a
a
   in  case (forall a. Ord a => a -> a -> Ordering
compare i
k0 i
m0, forall a. Ord a => a -> a -> Ordering
compare i
k1 i
m1) of
         (Ordering
EQ,Ordering
EQ) -> forall i a.
(Ix i, Num a) =>
Reflection i a -> Array i a -> Array i a
reflectVectorFull Reflection i a
q Array i a
a
         (Ordering
LT,Ordering
_) -> forall a. HasCallStack => [Char] -> a
error [Char]
"reflectVector: lower reflection dimension too small"
         (Ordering
_,Ordering
GT) -> forall a. HasCallStack => [Char] -> a
error [Char]
"reflectVector: upper reflection dimension too big"
         (Ordering, Ordering)
_ -> forall i a. Ix i => Array i a -> Array i a -> Array i a
replaceSubArray Array i a
a forall a b. (a -> b) -> a -> b
$ forall i a.
(Ix i, Num a) =>
Reflection i a -> Array i a -> Array i a
reflectVectorFull Reflection i a
q forall a b. (a -> b) -> a -> b
$ forall i a. Ix i => (i, i) -> Array i a -> Array i a
subArray (i, i)
bnds Array i a
a


subArray :: (Ix i) => (i,i) -> Array i a -> Array i a
subArray :: forall i a. Ix i => (i, i) -> Array i a -> Array i a
subArray (i, i)
bnds = forall i j e.
(Ix i, Ix j) =>
(i, i) -> (i -> j) -> Array j e -> Array i e
ixmap (i, i)
bnds forall a. a -> a
id

replaceSubArray :: (Ix i) => Array i a -> Array i a -> Array i a
replaceSubArray :: forall i a. Ix i => Array i a -> Array i a -> Array i a
replaceSubArray Array i a
x Array i a
y = Array i a
x forall i e. Ix i => Array i e -> [(i, e)] -> Array i e
// forall i e. Ix i => Array i e -> [(i, e)]
assocs Array i a
y


{- |
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)) [Array j a]
rs0) 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) -> [e] -> Array i e
listArray (j
n0,j
n1) forall a b. (a -> b) -> a -> b
$
         forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
            (\(Array j a
r,a
bi) [a]
xs ->
               let (a
a:[a]
as) = forall i e. Array i e -> [e]
elems Array j a
r
               in  (a
bi forall a. Num a => a -> a -> a
- forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall a. Num a => a -> a -> a
(*) [a]
as [a]
xs)) forall a. Fractional a => a -> a -> a
/ a
a forall a. a -> [a] -> [a]
: [a]
xs)
            []
            (forall a b. [a] -> [b] -> [(a, b)]
zip [Array j a]
rs0 (forall i e. Array i e -> [e]
elems Array i a
b))
     else forall a. HasCallStack => [Char] -> a
error [Char]
"solveUpper: vertical bounds mismatch"

solve ::
   (Ix i, Ix j, Fractional a) =>
   ([Reflection i a], Upper i j a) -> Array i a -> Array j a
solve :: forall i j a.
(Ix i, Ix j, Fractional a) =>
([Reflection i a], Upper i j a) -> Array i a -> Array j a
solve ([Reflection 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) =>
Reflection i a -> Array i a -> Array i a
reflectVector) Array i a
b [Reflection i a]
qs

{- |
Solve an 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) =>
   Array (i,j) a -> Array i a -> Array j a
leastSquares :: forall i j a.
(Ix i, Enum i, Ix j, Enum j, RealFloat a) =>
Array (i, j) a -> Array i a -> Array j a
leastSquares = forall i j a.
(Ix i, Ix j, Fractional a) =>
([Reflection 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) =>
Array (i, j) a -> ([Reflection 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)) [Array j a]
rs) =
   if forall a. Ix a => (a, a) -> Int
rangeSize (j
n0,j
n1) forall a. Eq a => a -> a -> Bool
== forall (t :: * -> *) a. Foldable t => t a -> Int
length [Array j a]
rs
     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. [a] -> a
head forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall i e. Array i e -> [e]
elems) [Array j a]
rs
     else a
0

{- |
Only sensible for square matrices,
but the function does not check whether the matrix is square.

For non-square matrices the sign is not of much use.
It depends on the implementation.
It is not uniquely defined by requiring
that the determinant of the orthogonal transformation has positive sign.
-}
det :: (Ix i, Enum i, Ix j, Enum j, RealFloat a) => Array (i,j) a -> a
det :: forall i j a.
(Ix i, Enum i, Ix j, Enum j, RealFloat a) =>
Array (i, j) a -> a
det Array (i, j) a
a =
   let ([Reflection i a]
qs,Upper i j a
u) = forall i j a.
(Ix i, Enum i, Ix j, Enum j, RealFloat a) =>
Array (i, j) a -> ([Reflection i a], Upper i j a)
decompose Array (i, j) a
a
   in  (if forall a. Integral a => a -> Bool
even (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Reflection i a]
qs) then a
1 else -a
1) forall a. Num a => a -> a -> a
* forall i j a. (Ix i, Ix j, Fractional a) => Upper i j a -> a
detUpper Upper i j a
u

{- |
Absolute value of the determinant.
This is also sound for non-square matrices.
-}
detAbs :: (Ix i, Enum i, Ix j, Enum j, RealFloat a) => Array (i,j) a -> a
detAbs :: forall i j a.
(Ix i, Enum i, Ix j, Enum j, RealFloat a) =>
Array (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, 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) =>
Array (i, j) a -> ([Reflection i a], Upper i j a)
decompose