-- SPDX-License-Identifier: MPL-2.0

{- | Column-pivoted Householder QR. @A[:, permutation] = Q R@, with square Q.
Least squares requires m >= n and full column rank. Ridge uses augmented QR,
avoiding formation of the normal equations. Relative rank tolerance defaults
to max(m,n) times machine epsilon; callers may select an explicit tolerance.
-}
module HMath.Matrix.QR (
    QR,
    qrDecompose,
    qrDecomposeWithTolerance,
    qrQ,
    qrR,
    qrPermutation,
    qrRank,
    applyQ,
    applyQT,
    leastSquares,
    leastSquaresMatrix,
    solveQR,
    solveMatrixQR,
    ridgeSolve,
    ridgeSolveMatrix,
) where

import Control.Monad (foldM)
import Data.List (maximumBy, transpose)
import Data.Ord (comparing)
import HMath.Error
import HMath.Matrix
import HMath.Matrix.Internal

{- | Opaque column-pivoted Householder factors satisfying A[:,p]=Q*R, with numerical-rank
estimate. Construct with qrDecompose.
-}
data QR a = QR (Matrix a) (Matrix a) [Int] Int deriving (QR a -> QR a -> Bool
(QR a -> QR a -> Bool) -> (QR a -> QR a -> Bool) -> Eq (QR a)
forall a. Eq a => QR a -> QR a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => QR a -> QR a -> Bool
== :: QR a -> QR a -> Bool
$c/= :: forall a. Eq a => QR a -> QR a -> Bool
/= :: QR a -> QR a -> Bool
Eq, Int -> QR a -> ShowS
[QR a] -> ShowS
QR a -> String
(Int -> QR a -> ShowS)
-> (QR a -> String) -> ([QR a] -> ShowS) -> Show (QR a)
forall a. Show a => Int -> QR a -> ShowS
forall a. Show a => [QR a] -> ShowS
forall a. Show a => QR a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> QR a -> ShowS
showsPrec :: Int -> QR a -> ShowS
$cshow :: forall a. Show a => QR a -> String
show :: QR a -> String
$cshowList :: forall a. Show a => [QR a] -> ShowS
showList :: [QR a] -> ShowS
Show)

-- | Extract the square m-by-m orthogonal Q factor from a QR decomposition.
qrQ :: QR a -> Matrix a
qrQ :: forall a. QR a -> Matrix a
qrQ (QR Matrix a
q Matrix a
_ [Int]
_ Int
_) = Matrix a
q

{- | Extract the m-by-n upper trapezoidal R factor. Column ordering corresponds to
qrPermutation.
-}
qrR :: QR a -> Matrix a
qrR :: forall a. QR a -> Matrix a
qrR (QR Matrix a
_ Matrix a
r [Int]
_ Int
_) = Matrix a
r

-- | Return zero-based original column indices in pivoted order, satisfying A[:,p]=Q*R.
qrPermutation :: QR a -> [Int]
qrPermutation :: forall a. QR a -> [Int]
qrPermutation (QR Matrix a
_ Matrix a
_ [Int]
p Int
_) = [Int]
p

{- | Return estimated numerical rank selected by the factorization's relative diagonal-pivot
tolerance.
-}
qrRank :: QR a -> Int
qrRank :: forall a. QR a -> Int
qrRank (QR Matrix a
_ Matrix a
_ [Int]
_ Int
r) = Int
r

{- | Column-pivoted Householder QR of a finite m-by-n matrix using relative rank tolerance
max(m,n)*machine epsilon. Return factors satisfying A[:,permutation]=Q*R. Numeric/input
failures return Left MathError; rank deficiency is recorded rather than rejected at
factorization.
-}
qrDecompose :: (RealFloat a) => Matrix a -> Either MathError (QR a)
qrDecompose :: forall a. RealFloat a => Matrix a -> Either MathError (QR a)
qrDecompose Matrix a
a = a -> Matrix a -> Either MathError (QR a)
forall a. RealFloat a => a -> Matrix a -> Either MathError (QR a)
qrDecomposeWithTolerance (Matrix a -> a
forall a. RealFloat a => Matrix a -> a
defaultTolerance Matrix a
a) Matrix a
a

{- | Given finite nonnegative relative rank tolerance then finite matrix, compute
column-pivoted Householder QR. Reject invalid tolerance/non-finite arithmetic. Zero
tolerance counts every nonzero diagonal pivot as significant.
-}
qrDecomposeWithTolerance :: (RealFloat a) => a -> Matrix a -> Either MathError (QR a)
qrDecomposeWithTolerance :: forall a. RealFloat a => a -> Matrix a -> Either MathError (QR a)
qrDecomposeWithTolerance a
tol Matrix a
a
    | 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 (QR a)
forall a b. a -> Either a b
Left MathError
InvalidTolerance
    | Bool
otherwise = do
        Matrix a -> Either MathError ()
forall a. RealFloat a => Matrix a -> Either MathError ()
finiteMatrix Matrix a
a
        q0 <- Int -> Either MathError (Matrix a)
forall a. Num a => Int -> Either MathError (Matrix a)
identityChecked Int
m
        (q, r, p) <- foldM step (entries q0, entries a, [0 .. n - 1]) [0 .. min m n - 1]
        qm <- fromListChecked m m (concat q)
        rm <- fromListChecked m n (concat r)
        finiteMatrix qm
        finiteMatrix rm
        let diagonals = [a -> a
forall a. Num a => a -> a
abs ([[a]]
r [[a]] -> Int -> [a]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
i) | Int
i <- [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]]
            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]
diagonals)
            rank = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a
x | a
x <- [a]
diagonals, a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
0, a
tol a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 Bool -> Bool -> Bool
|| a
x a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
scale a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
tol]
        Right (QR qm rm p rank)
  where
    (Int
m, Int
n) = Matrix a -> (Int, Int)
forall a. Matrix a -> (Int, Int)
shape Matrix a
a
    step :: ([[a]], [[a]], [a]) -> Int -> Either MathError ([[a]], [[a]], [a])
step ([[a]]
q, [[a]]
r, [a]
p) Int
k = do
        let candidates :: [(Int, a)]
candidates = [(Int
j, [a] -> a
forall a. RealFloat a => [a] -> a
stableNorm [[[a]]
r [[a]] -> Int -> [a]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
j | Int
i <- [Int
k .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]) | Int
j <- [Int
k .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
            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, a)]
candidates)
            rs :: [[a]]
rs = ([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> [a] -> [a]
forall a. Int -> Int -> [a] -> [a]
swap Int
k Int
pivot) [[a]]
r
            x :: [a]
x = [[[a]]
rs [[a]] -> Int -> [a]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
k | Int
i <- [Int
k .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
            s :: a
s = [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]
x)
        if a
s a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0
            then ([[a]], [[a]], [a]) -> Either MathError ([[a]], [[a]], [a])
forall a b. b -> Either a b
Right ([[a]]
q, [[a]]
rs, Int -> Int -> [a] -> [a]
forall a. Int -> Int -> [a] -> [a]
swap Int
k Int
pivot [a]
p)
            else do
                let u :: [a]
u = (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
s) [a]
x
                    alpha :: a
alpha = if [a]
u [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]
u) else [a] -> a
forall a. RealFloat a => [a] -> a
stableNorm [a]
u
                    v0 :: [a]
v0 = Int -> a -> [a] -> [a]
forall a. Int -> a -> [a] -> [a]
replace Int
0 ([a]
u [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
0 a -> a -> a
forall a. Num a => a -> a -> a
- a
alpha) [a]
u
                    v :: [a]
v = (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]
v0) [a]
v0
                    transform :: [a] -> [a]
transform [a]
xs = let t :: a
t = 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
k [a]
xs) in Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
k [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
z a
w -> a
z a -> a -> a
forall a. Num a => a -> a -> a
- a
t a -> a -> a
forall a. Num a => a -> a -> a
* a
w) (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
k [a]
xs) [a]
v
                    rnew :: [[a]]
rnew = [[a]] -> [[a]]
forall a. [[a]] -> [[a]]
transpose (([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map [a] -> [a]
transform ([[a]] -> [[a]]
forall a. [[a]] -> [[a]]
transpose [[a]]
rs))
                    qnew :: [[a]]
qnew = ([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map [a] -> [a]
transform [[a]]
q
                    rclean :: [[a]]
rclean = [[if 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 then a
0 else [[a]]
rnew [[a]] -> Int -> [a]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
j | Int
j <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] | Int
i <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
                [a] -> Either MathError ()
forall a. RealFloat a => [a] -> Either MathError ()
finiteValues ([[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[a]]
qnew [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[a]]
rclean)
                ([[a]], [[a]], [a]) -> Either MathError ([[a]], [[a]], [a])
forall a b. b -> Either a b
Right ([[a]]
qnew, [[a]]
rclean, Int -> Int -> [a] -> [a]
forall a. Int -> Int -> [a] -> [a]
swap Int
k Int
pivot [a]
p)

{- | Given QR factorization then a matrix with m rows, return Q times that matrix.
Dimension/size errors return Left MathError; generic Num arithmetic is unchanged.
-}
applyQ :: (Num a) => QR a -> Matrix a -> Either MathError (Matrix a)
applyQ :: forall a. Num a => QR a -> Matrix a -> Either MathError (Matrix a)
applyQ QR a
f = Matrix a -> Matrix a -> Either MathError (Matrix a)
forall a.
Num a =>
Matrix a -> Matrix a -> Either MathError (Matrix a)
multiplyChecked (QR a -> Matrix a
forall a. QR a -> Matrix a
qrQ QR a
f)

{- | Given QR factorization then a matrix with m rows, return transpose(Q) times it. Checked
matrix dimension errors propagate.
-}
applyQT :: (Num a) => QR a -> Matrix a -> Either MathError (Matrix a)
applyQT :: forall a. Num a => QR a -> Matrix a -> Either MathError (Matrix a)
applyQT QR a
f = Matrix a -> Matrix a -> Either MathError (Matrix a)
forall a.
Num a =>
Matrix a -> Matrix a -> Either MathError (Matrix a)
multiplyChecked (Matrix a -> Matrix a
forall a. Matrix a -> Matrix a
transposeMatrix (QR a -> Matrix a
forall a. QR a -> Matrix a
qrQ QR a
f))

{- | Given QR factors and an m-by-k RHS, return n-by-k least-squares coefficients, undoing
column pivoting. Require m>=n and full column rank; shape/rank/finiteness failures
return Left MathError.
-}
solveMatrixQR :: (RealFloat a) => QR a -> Matrix a -> Either MathError (Matrix a)
solveMatrixQR :: forall a.
RealFloat a =>
QR a -> Matrix a -> Either MathError (Matrix a)
solveMatrixQR QR a
f Matrix a
b
    | Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n = MathError -> Either MathError (Matrix a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"leastSquares: underdetermined system")
    | (Int, Int) -> Int
forall a b. (a, b) -> a
fst (Matrix a -> (Int, Int)
forall a. Matrix a -> (Int, Int)
shape Matrix a
b) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
m = MathError -> Either MathError (Matrix a)
forall a b. a -> Either a b
Left (Int -> Int -> MathError
RightHandSideMismatch Int
m ((Int, Int) -> Int
forall a b. (a, b) -> a
fst (Matrix a -> (Int, Int)
forall a. Matrix a -> (Int, Int)
shape Matrix a
b)))
    | QR a -> Int
forall a. QR a -> Int
qrRank QR a
f Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n = MathError -> Either MathError (Matrix a)
forall a b. a -> Either a b
Left (Int -> MathError
RankDeficient (QR a -> Int
forall a. QR a -> Int
qrRank QR a
f))
    | Bool
otherwise = do
        Matrix a -> Either MathError ()
forall a. RealFloat a => Matrix a -> Either MathError ()
finiteMatrix Matrix a
b
        if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
|| (Int, Int) -> Int
forall a b. (a, b) -> b
snd (Matrix a -> (Int, Int)
forall a. Matrix a -> (Int, Int)
shape Matrix a
b) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
            then Int -> Int -> [a] -> Either MathError (Matrix a)
forall a. Int -> Int -> [a] -> Either MathError (Matrix a)
fromListChecked Int
n ((Int, Int) -> Int
forall a b. (a, b) -> b
snd (Matrix a -> (Int, Int)
forall a. Matrix a -> (Int, Int)
shape Matrix a
b)) []
            else do
                y <- QR a -> Matrix a -> Either MathError (Matrix a)
forall a. Num a => QR a -> Matrix a -> Either MathError (Matrix a)
applyQT QR a
f Matrix a
b
                let solveColumn Int
j = ([a] -> Int -> Either MathError [a])
-> [a] -> [Int] -> Either MathError [a]
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (Matrix a -> Int -> [a] -> Int -> Either MathError [a]
back Matrix a
y Int
j) [] ([Int] -> [Int]
forall a. [a] -> [a]
reverse [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1])
                    back Matrix a
ymat Int
j [a]
xs Int
i = do
                        let value :: a
value = (Matrix a -> Int -> Int -> a
forall a. Matrix a -> Int -> Int -> a
get Matrix a
ymat Int
i Int
j 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 [Matrix a -> Int -> Int -> a
forall a. Matrix a -> Int -> Int -> a
get Matrix a
r Int
i Int
k a -> a -> a
forall a. Num a => a -> a -> a
* ([a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) | Int
k <- [Int
i 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]]) a -> a -> a
forall a. Fractional a => a -> a -> a
/ Matrix a -> Int -> Int -> a
forall a. Matrix a -> Int -> Int -> a
get Matrix a
r Int
i Int
i
                        [a] -> Either MathError ()
forall a. RealFloat a => [a] -> Either MathError ()
finiteValues [a
value]
                        [a] -> Either MathError [a]
forall a b. b -> Either a b
Right (a
value a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
xs)
                columns <- traverse solveColumn [0 .. snd (shape b) - 1]
                build
                    n
                    (snd (shape b))
                    ( \Int
i Int
j ->
                        let pos :: Int
pos = case Int -> [(Int, Int)] -> Maybe Int
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Int
i ([Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip (QR a -> [Int]
forall a. QR a -> [Int]
qrPermutation QR a
f) [Int
0 ..]) of Just Int
k -> Int
k; Maybe Int
Nothing -> Int
i
                         in [[a]]
columns [[a]] -> Int -> [a]
forall a. HasCallStack => [a] -> Int -> a
!! Int
j [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
pos
                    )
  where
    r :: Matrix a
r = QR a -> Matrix a
forall a. QR a -> Matrix a
qrR QR a
f; (Int
m, Int
n) = Matrix a -> (Int, Int)
forall a. Matrix a -> (Int, Int)
shape Matrix a
r

{- | Given full-column-rank QR factors of an m-by-n matrix then m-element right-hand side,
return n least-squares coefficients in original column order. Reject underdetermined
systems, incompatible RHS, deficient rank, or non-finite input/result.
-}
solveQR :: (RealFloat a) => QR a -> [a] -> Either MathError [a]
solveQR :: forall a. RealFloat a => QR a -> [a] -> Either MathError [a]
solveQR QR a
f [a]
b = do
    rhs <- Int -> Int -> [a] -> Either MathError (Matrix a)
forall a. Int -> Int -> [a] -> Either MathError (Matrix a)
fromListChecked ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
b) Int
1 [a]
b
    toList <$> solveMatrixQR f rhs

{- | Given finite m-by-n design matrix then m-element RHS, factor with pivoted QR and return
least-squares coefficients. Require m>=n and full column rank; decomposition/solve
errors propagate.
-}
leastSquares :: (RealFloat a) => Matrix a -> [a] -> Either MathError [a]
leastSquares :: forall a. RealFloat a => Matrix a -> [a] -> Either MathError [a]
leastSquares Matrix a
a [a]
b = Matrix a -> Either MathError (QR a)
forall a. RealFloat a => Matrix a -> Either MathError (QR a)
qrDecompose Matrix a
a Either MathError (QR a)
-> (QR 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
>>= (QR a -> [a] -> Either MathError [a]
forall a. RealFloat a => QR a -> [a] -> Either MathError [a]
`solveQR` [a]
b)

{- | Given finite m-by-n design and m-by-k RHS matrices, return n-by-k least-squares
solution via QR. Underdetermined/rank-deficient or invalid input returns Left
MathError.
-}
leastSquaresMatrix :: (RealFloat a) => Matrix a -> Matrix a -> Either MathError (Matrix a)
leastSquaresMatrix :: forall a.
RealFloat a =>
Matrix a -> Matrix a -> Either MathError (Matrix a)
leastSquaresMatrix Matrix a
a Matrix a
b = Matrix a -> Either MathError (QR a)
forall a. RealFloat a => Matrix a -> Either MathError (QR a)
qrDecompose Matrix a
a Either MathError (QR a)
-> (QR a -> Either MathError (Matrix a))
-> Either MathError (Matrix 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
>>= (QR a -> Matrix a -> Either MathError (Matrix a)
forall a.
RealFloat a =>
QR a -> Matrix a -> Either MathError (Matrix a)
`solveMatrixQR` Matrix a
b)

{- | Given finite nonnegative lambda, design matrix, then RHS list, minimize squared
residual plus lambda times squared coefficient norm. Use augmented QR, not normal
equations. Zero lambda delegates to leastSquares; dimension/numeric errors propagate.
-}
ridgeSolve :: (RealFloat a) => a -> Matrix a -> [a] -> Either MathError [a]
ridgeSolve :: forall a.
RealFloat a =>
a -> Matrix a -> [a] -> Either MathError [a]
ridgeSolve a
lambda Matrix a
a [a]
b = do
    rhs <- Int -> Int -> [a] -> Either MathError (Matrix a)
forall a. Int -> Int -> [a] -> Either MathError (Matrix a)
fromListChecked ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
b) Int
1 [a]
b
    toList <$> ridgeSolveMatrix lambda a rhs

{- | Given lambda, m-by-n design, then m-by-k RHS, return n-by-k ridge coefficients via
augmented QR. Reject negative/non-finite lambda, RHS mismatch, shape overflow, or
numerical failure.
-}
ridgeSolveMatrix :: (RealFloat a) => a -> Matrix a -> Matrix a -> Either MathError (Matrix a)
ridgeSolveMatrix :: forall a.
RealFloat a =>
a -> Matrix a -> Matrix a -> Either MathError (Matrix a)
ridgeSolveMatrix a
lambda Matrix a
a Matrix a
b
    | a
lambda 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
lambda Bool -> Bool -> Bool
|| a -> Bool
forall a. RealFloat a => a -> Bool
isInfinite a
lambda =
        MathError -> Either MathError (Matrix a)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"ridgeSolve: finite nonnegative regularization required")
    | (Int, Int) -> Int
forall a b. (a, b) -> a
fst (Matrix a -> (Int, Int)
forall a. Matrix a -> (Int, Int)
shape Matrix a
b) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
m = MathError -> Either MathError (Matrix a)
forall a b. a -> Either a b
Left (Int -> Int -> MathError
RightHandSideMismatch Int
m ((Int, Int) -> Int
forall a b. (a, b) -> a
fst (Matrix a -> (Int, Int)
forall a. Matrix a -> (Int, Int)
shape Matrix a
b)))
    | a
lambda a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 = Matrix a -> Matrix a -> Either MathError (Matrix a)
forall a.
RealFloat a =>
Matrix a -> Matrix a -> Either MathError (Matrix a)
leastSquaresMatrix Matrix a
a Matrix a
b
    | Bool
otherwise = do
        Matrix a -> Either MathError ()
forall a. RealFloat a => Matrix a -> Either MathError ()
finiteMatrix Matrix a
a
        Matrix a -> Either MathError ()
forall a. RealFloat a => Matrix a -> Either MathError ()
finiteMatrix Matrix a
b
        if Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
m Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
n 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 (Matrix a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"ridgeSolve: shape overflow")
            else do
                augmented <-
                    Int -> Int -> (Int -> Int -> a) -> Either MathError (Matrix a)
forall a.
Int -> Int -> (Int -> Int -> a) -> Either MathError (Matrix a)
build (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) Int
n (\Int
i Int
j -> if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
m then Matrix a -> Int -> Int -> a
forall a. Matrix a -> Int -> Int -> a
get Matrix a
a Int
i Int
j else if Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
m Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j then a -> a
forall a. Floating a => a -> a
sqrt a
lambda else a
0)
                rhs <- build (m + n) (snd (shape b)) (\Int
i Int
j -> if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
m then Matrix a -> Int -> Int -> a
forall a. Matrix a -> Int -> Int -> a
get Matrix a
b Int
i Int
j else a
0)
                leastSquaresMatrix augmented rhs
  where
    (Int
m, Int
n) = Matrix a -> (Int, Int)
forall a. Matrix a -> (Int, Int)
shape Matrix a
a