-- SPDX-License-Identifier: MPL-2.0

-- | Private checked reference linear algebra shared by scientific algorithms.
module HMath.Signal.Internal (finite, require, checked, vector, solve, mm, mv, tr, eye, dot, norm, chunks, replace) where

import Data.List (maximumBy, transpose)
import Data.Ord (comparing)
import HMath.Error (MathError (..))

finite :: Double -> Bool
finite :: Double -> Bool
finite Double
x = Bool -> Bool
not (Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
x Bool -> Bool -> Bool
|| Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
x)
require :: Bool -> String -> Either MathError ()
require :: Bool -> String -> Either MathError ()
require Bool
p String
s = if Bool
p then () -> Either MathError ()
forall a b. b -> Either a b
Right () else MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
s)
checked :: Double -> Either MathError Double
checked :: Double -> Either MathError Double
checked Double
x = if Double -> Bool
finite Double
x then Double -> Either MathError Double
forall a b. b -> Either a b
Right Double
x else MathError -> Either MathError Double
forall a b. a -> Either a b
Left MathError
NonFiniteResult
vector :: [Double] -> Either MathError [Double]
vector :: [Double] -> Either MathError [Double]
vector = (Double -> Either MathError Double)
-> [Double] -> Either MathError [Double]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Double -> Either MathError Double
checked
replace :: Int -> a -> [a] -> [a]
replace :: forall a. Int -> a -> [a] -> [a]
replace Int
i a
x [a]
xs = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
i [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a
x] [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
xs
chunks :: Int -> [a] -> [[a]]
chunks :: forall a. Int -> [a] -> [[a]]
chunks Int
n [a]
xs | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 Bool -> Bool -> Bool
|| [a] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
xs = [] | Bool
otherwise = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
n [a]
xs [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
chunks Int
n (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
n [a]
xs)
dot :: [Double] -> [Double] -> Double
dot :: [Double] -> [Double] -> Double
dot [Double]
a [Double]
b = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
a [Double]
b)
norm :: [Double] -> Double
norm :: [Double] -> Double
norm [Double]
xs =
    let s :: Double
s = [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (Double
0 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
forall a. Num a => a -> a
abs [Double]
xs)
     in if Double
s Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else Double
s Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sqrt ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [(Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
s) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) | Double
x <- [Double]
xs])
tr :: [[a]] -> [[a]]
tr :: forall a. [[a]] -> [[a]]
tr = [[a]] -> [[a]]
forall a. [[a]] -> [[a]]
transpose
mv :: [[Double]] -> [Double] -> [Double]
mv :: [[Double]] -> [Double] -> [Double]
mv [[Double]]
a [Double]
x = ([Double] -> Double) -> [[Double]] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map ([Double] -> [Double] -> Double
`dot` [Double]
x) [[Double]]
a
mm :: [[Double]] -> [[Double]] -> [[Double]]
mm :: [[Double]] -> [[Double]] -> [[Double]]
mm [[Double]]
a [[Double]]
b = [[[Double] -> [Double] -> Double
dot [Double]
r [Double]
c | [Double]
c <- [[Double]] -> [[Double]]
forall a. [[a]] -> [[a]]
tr [[Double]]
b] | [Double]
r <- [[Double]]
a]
eye :: Int -> [[Double]]
eye :: Int -> [[Double]]
eye Int
n = [[if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j then Double
1 else Double
0 | Int
j <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]

{- | Partial-pivoted elimination. Explicit square/finite contracts; singular
systems fail rather than yielding an arbitrary least-squares solution.
-}
solve :: [[Double]] -> [Double] -> Either MathError [Double]
solve :: [[Double]] -> [Double] -> Either MathError [Double]
solve [[Double]]
a [Double]
b = do
    let n :: Int
n = [[Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double]]
a
    Bool -> String -> Either MathError ()
require
        (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
2048 Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
b Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n Bool -> Bool -> Bool
&& ([Double] -> Bool) -> [[Double]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n) (Int -> Bool) -> ([Double] -> Int) -> [Double] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) [[Double]]
a)
        String
"invalid square reference solve (maximum 2048)"
    _ <- [Double] -> Either MathError [Double]
vector ([[Double]] -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Double]]
a [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
b)
    upper <- foldl step (Right (zipWith (\[Double]
r Double
x -> [Double]
r [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double
x]) a b)) [0 .. n - 1]
    let back Int
i [Double]
acc
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = [Double] -> Either MathError [Double]
vector [Double]
acc
            | Bool
otherwise =
                let row :: [Double]
row = [[Double]]
upper [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i; x :: Double
x = ([Double] -> Double
forall a. HasCallStack => [a] -> a
last [Double]
row Double -> Double -> Double
forall a. Num a => a -> a -> a
- [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
n [Double]
row)) [Double]
acc)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ ([Double]
row [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)
                 in Double -> Either MathError Double
checked Double
x Either MathError Double
-> Either MathError [Double] -> Either MathError [Double]
forall a b.
Either MathError a -> Either MathError b -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> [Double] -> Either MathError [Double]
back (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Double
x Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: [Double]
acc)
    back (n - 1) []
  where
    step :: Either MathError [[Double]] -> Int -> Either MathError [[Double]]
step Either MathError [[Double]]
result Int
k = do
        rows <- Either MathError [[Double]]
result
        let n = [[Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double]]
rows
            p = (Int -> Int -> Ordering) -> [Int] -> Int
forall (t :: * -> *) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
maximumBy ((Int -> Double) -> Int -> Int -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (Double -> Double
forall a. Num a => a -> a
abs (Double -> Double) -> (Int -> Double) -> Int -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
k) ([Double] -> Double) -> (Int -> [Double]) -> Int -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([[Double]]
rows [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!!))) [Int
k .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
            pivot = [[Double]]
rows [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
p [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
k
        if pivot == 0
            then Left (SingularMatrix k)
            else do
                let swapped = Int -> [Double] -> [[Double]] -> [[Double]]
forall a. Int -> a -> [a] -> [a]
replace Int
p ([[Double]]
rows [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
k) (Int -> [Double] -> [[Double]] -> [[Double]]
forall a. Int -> a -> [a] -> [a]
replace Int
k ([[Double]]
rows [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
p) [[Double]]
rows)
                    row = [[Double]]
swapped [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
k
                    next =
                        [if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
k then [Double]
r else (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-) [Double]
r ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Num a => a -> a -> a
* ([Double]
r [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
k Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
pivot)) [Double]
row) | (Int
i, [Double]
r) <- [Int] -> [[Double]] -> [(Int, [Double])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [[Double]]
swapped]
                _ <- vector (concat next)
                Right next