-- SPDX-License-Identifier: MPL-2.0

{- | Checked finite complex arithmetic over "Data.Complex". Angles are radians,
positive rotations are counterclockwise, and logarithms/powers use the principal
branch (argument in (-pi,pi], with IEEE signed-zero behavior inherited from
"Data.Complex"). Poles, zero reciprocals, and nonfinite results are errors.
Matrices use the existing row-major @M.Matrix@; FFTs use the existing negative
forward-exponent convention. The reference matrix budget is 256 per dimension.
-}
module HMath.Complex (
    Complex,
    complex,
    realPart,
    imagPart,
    magnitude,
    phase,
    conjugate,
    reciprocal,
    add,
    subtractComplex,
    multiply,
    divide,
    power,
    Elementary (..),
    elementary,
    fromPolar,
    toPolar,
    rotate,
    complexMatrix,
    multiplyMatrices,
    conjugateTranspose,
    complexFFT,
    complexIFFT,
) where

import Data.Complex (Complex ((:+)), conjugate, imagPart, magnitude, phase, realPart)
import qualified Data.Complex as C
import HMath.Error (MathError (..))
import qualified HMath.Fourier as F
import qualified HMath.Matrix as M

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)
valid :: Complex Double -> Bool
valid :: Complex Double -> Bool
valid Complex Double
z = Double -> Bool
finite (Complex Double -> Double
forall a. Complex a -> a
realPart Complex Double
z) Bool -> Bool -> Bool
&& Double -> Bool
finite (Complex Double -> Double
forall a. Complex a -> a
imagPart Complex Double
z)
checked :: Complex Double -> Either MathError (Complex Double)
checked :: Complex Double -> Either MathError (Complex Double)
checked Complex Double
z = if Complex Double -> Bool
valid Complex Double
z then Complex Double -> Either MathError (Complex Double)
forall a b. b -> Either a b
Right Complex Double
z else MathError -> Either MathError (Complex Double)
forall a b. a -> Either a b
Left MathError
NonFiniteResult
inputs :: [Complex Double] -> Either MathError ()
inputs :: [Complex Double] -> Either MathError ()
inputs [Complex Double]
zs = if (Complex Double -> Bool) -> [Complex Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Complex Double -> Bool
valid [Complex Double]
zs 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
DomainError String
"nonfinite complex input")

-- | Construct a finite number from real and imaginary parts.
complex :: Double -> Double -> Either MathError (Complex Double)
complex :: Double -> Double -> Either MathError (Complex Double)
complex Double
x Double
y = Complex Double -> Either MathError (Complex Double)
checked (Double
x Double -> Double -> Complex Double
forall a. a -> a -> Complex a
:+ Double
y)

-- | Reciprocal using scaled division to avoid squaring a large norm.
reciprocal :: Complex Double -> Either MathError (Complex Double)
reciprocal :: Complex Double -> Either MathError (Complex Double)
reciprocal Complex Double
z = Complex Double
-> Complex Double -> Either MathError (Complex Double)
divide (Double
1 Double -> Double -> Complex Double
forall a. a -> a -> Complex a
:+ Double
0) Complex Double
z

-- | Checked sum.
add :: Complex Double -> Complex Double -> Either MathError (Complex Double)
add :: Complex Double
-> Complex Double -> Either MathError (Complex Double)
add Complex Double
a Complex Double
b = [Complex Double] -> Either MathError ()
inputs [Complex Double
a, Complex Double
b] Either MathError ()
-> Either MathError (Complex Double)
-> Either MathError (Complex 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
>> Complex Double -> Either MathError (Complex Double)
checked (Complex Double
a Complex Double -> Complex Double -> Complex Double
forall a. Num a => a -> a -> a
+ Complex Double
b)

-- | Checked difference.
subtractComplex :: Complex Double -> Complex Double -> Either MathError (Complex Double)
subtractComplex :: Complex Double
-> Complex Double -> Either MathError (Complex Double)
subtractComplex Complex Double
a Complex Double
b = [Complex Double] -> Either MathError ()
inputs [Complex Double
a, Complex Double
b] Either MathError ()
-> Either MathError (Complex Double)
-> Either MathError (Complex 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
>> Complex Double -> Either MathError (Complex Double)
checked (Complex Double
a Complex Double -> Complex Double -> Complex Double
forall a. Num a => a -> a -> a
- Complex Double
b)

-- | Checked product.
multiply :: Complex Double -> Complex Double -> Either MathError (Complex Double)
multiply :: Complex Double
-> Complex Double -> Either MathError (Complex Double)
multiply Complex Double
a Complex Double
b = [Complex Double] -> Either MathError ()
inputs [Complex Double
a, Complex Double
b] Either MathError ()
-> Either MathError (Complex Double)
-> Either MathError (Complex 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
>> Complex Double -> Either MathError (Complex Double)
checked (Complex Double
a Complex Double -> Complex Double -> Complex Double
forall a. Num a => a -> a -> a
* Complex Double
b)

-- | Smith scaled division. A zero divisor is rejected.
divide :: Complex Double -> Complex Double -> Either MathError (Complex Double)
divide :: Complex Double
-> Complex Double -> Either MathError (Complex Double)
divide Complex Double
a Complex Double
b = do
    [Complex Double] -> Either MathError ()
inputs [Complex Double
a, Complex Double
b]
    if Complex Double
b Complex Double -> Complex Double -> Bool
forall a. Eq a => a -> a -> Bool
== Complex Double
0
        then MathError -> Either MathError (Complex Double)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"zero complex divisor")
        else
            let ar :: Double
ar = Complex Double -> Double
forall a. Complex a -> a
realPart Complex Double
a
                ai :: Double
ai = Complex Double -> Double
forall a. Complex a -> a
imagPart Complex Double
a
                br :: Double
br = Complex Double -> Double
forall a. Complex a -> a
realPart Complex Double
b
                bi :: Double
bi = Complex Double -> Double
forall a. Complex a -> a
imagPart Complex Double
b
                -- Normalize both numerator and denominator before forming products.
                s :: Double
s = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max (Double -> Double
forall a. Num a => a -> a
abs Double
br) (Double -> Double
forall a. Num a => a -> a
abs Double
bi)
                u :: Double
u = Double
br Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
s
                v :: Double
v = Double
bi Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
s
                den :: Double
den = Double
u Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
u Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
v
             in Complex Double -> Either MathError (Complex Double)
checked (((Double
ar Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
s Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
u Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
ai Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
s Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
v) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
den) Double -> Double -> Complex Double
forall a. a -> a -> Complex a
:+ ((Double
ai Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
s Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
u Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
ar Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
s Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
v) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
den))

-- | Principal complex power. 0^0 is 1; other zero-base exponents must be positive real.
power :: Complex Double -> Complex Double -> Either MathError (Complex Double)
power :: Complex Double
-> Complex Double -> Either MathError (Complex Double)
power Complex Double
a Complex Double
b = do
    [Complex Double] -> Either MathError ()
inputs [Complex Double
a, Complex Double
b]
    if Complex Double
a Complex Double -> Complex Double -> Bool
forall a. Eq a => a -> a -> Bool
== Complex Double
0
        then
            if Complex Double
b Complex Double -> Complex Double -> Bool
forall a. Eq a => a -> a -> Bool
== Complex Double
0
                then Complex Double -> Either MathError (Complex Double)
forall a b. b -> Either a b
Right Complex Double
1
                else
                    if Complex Double -> Double
forall a. Complex a -> a
imagPart Complex Double
b Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 Bool -> Bool -> Bool
&& Complex Double -> Double
forall a. Complex a -> a
realPart Complex Double
b Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0
                        then Complex Double -> Either MathError (Complex Double)
forall a b. b -> Either a b
Right Complex Double
0
                        else MathError -> Either MathError (Complex Double)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"undefined zero complex power")
        else Complex Double -> Either MathError (Complex Double)
checked (Complex Double
a Complex Double -> Complex Double -> Complex Double
forall a. Floating a => a -> a -> a
** Complex Double
b)

-- | Elementary functions with principal inverse-function branches.
data Elementary
    = Exp
    | Log
    | Sqrt
    | Sin
    | Cos
    | Tan
    | Sinh
    | Cosh
    | Tanh
    | Asin
    | Acos
    | Atan
    | Asinh
    | Acosh
    | Atanh
    deriving (Elementary -> Elementary -> Bool
(Elementary -> Elementary -> Bool)
-> (Elementary -> Elementary -> Bool) -> Eq Elementary
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Elementary -> Elementary -> Bool
== :: Elementary -> Elementary -> Bool
$c/= :: Elementary -> Elementary -> Bool
/= :: Elementary -> Elementary -> Bool
Eq, Int -> Elementary -> ShowS
[Elementary] -> ShowS
Elementary -> String
(Int -> Elementary -> ShowS)
-> (Elementary -> String)
-> ([Elementary] -> ShowS)
-> Show Elementary
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Elementary -> ShowS
showsPrec :: Int -> Elementary -> ShowS
$cshow :: Elementary -> String
show :: Elementary -> String
$cshowList :: [Elementary] -> ShowS
showList :: [Elementary] -> ShowS
Show)

-- | Apply an elementary function, rejecting poles and nonfinite results.
elementary :: Elementary -> Complex Double -> Either MathError (Complex Double)
elementary :: Elementary -> Complex Double -> Either MathError (Complex Double)
elementary Elementary
op Complex Double
z = do
    [Complex Double] -> Either MathError ()
inputs [Complex Double
z]
    if Elementary
op Elementary -> Elementary -> Bool
forall a. Eq a => a -> a -> Bool
== Elementary
Log Bool -> Bool -> Bool
&& Complex Double
z Complex Double -> Complex Double -> Bool
forall a. Eq a => a -> a -> Bool
== Complex Double
0
        then MathError -> Either MathError (Complex Double)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"logarithm of zero")
        else Complex Double -> Either MathError (Complex Double)
checked (Complex Double -> Either MathError (Complex Double))
-> Complex Double -> Either MathError (Complex Double)
forall a b. (a -> b) -> a -> b
$ case Elementary
op of
            Elementary
Exp -> Complex Double -> Complex Double
forall a. Floating a => a -> a
exp Complex Double
z
            Elementary
Log -> Complex Double -> Complex Double
forall a. Floating a => a -> a
log Complex Double
z
            Elementary
Sqrt -> Complex Double -> Complex Double
forall a. Floating a => a -> a
sqrt Complex Double
z
            Elementary
Sin -> Complex Double -> Complex Double
forall a. Floating a => a -> a
sin Complex Double
z
            Elementary
Cos -> Complex Double -> Complex Double
forall a. Floating a => a -> a
cos Complex Double
z
            Elementary
Tan -> Complex Double -> Complex Double
forall a. Floating a => a -> a
tan Complex Double
z
            Elementary
Sinh -> Complex Double -> Complex Double
forall a. Floating a => a -> a
sinh Complex Double
z
            Elementary
Cosh -> Complex Double -> Complex Double
forall a. Floating a => a -> a
cosh Complex Double
z
            Elementary
Tanh -> Complex Double -> Complex Double
forall a. Floating a => a -> a
tanh Complex Double
z
            Elementary
Asin -> Complex Double -> Complex Double
forall a. Floating a => a -> a
asin Complex Double
z
            Elementary
Acos -> Complex Double -> Complex Double
forall a. Floating a => a -> a
acos Complex Double
z
            Elementary
Atan -> Complex Double -> Complex Double
forall a. Floating a => a -> a
atan Complex Double
z
            Elementary
Asinh -> Complex Double -> Complex Double
forall a. Floating a => a -> a
asinh Complex Double
z
            Elementary
Acosh -> Complex Double -> Complex Double
forall a. Floating a => a -> a
acosh Complex Double
z
            Elementary
Atanh -> Complex Double -> Complex Double
forall a. Floating a => a -> a
atanh Complex Double
z

-- | Polar construction with nonnegative finite radius.
fromPolar :: Double -> Double -> Either MathError (Complex Double)
fromPolar :: Double -> Double -> Either MathError (Complex Double)
fromPolar Double
radius Double
angle
    | Bool -> Bool
not (Double -> Bool
finite Double
radius Bool -> Bool -> Bool
&& Double -> Bool
finite Double
angle) Bool -> Bool -> Bool
|| Double
radius Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 = MathError -> Either MathError (Complex Double)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"invalid polar coordinates")
    | Bool
otherwise = Complex Double -> Either MathError (Complex Double)
checked (Double -> Double -> Complex Double
forall a. Floating a => a -> a -> Complex a
C.mkPolar Double
radius Double
angle)

-- | Radius and principal phase; the phase of zero is defined by Data.Complex.
toPolar :: Complex Double -> Either MathError (Double, Double)
toPolar :: Complex Double -> Either MathError (Double, Double)
toPolar Complex Double
z = [Complex Double] -> Either MathError ()
inputs [Complex Double
z] Either MathError ()
-> Either MathError (Double, Double)
-> Either MathError (Double, 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
>> if Double -> Bool
finite (Complex Double -> Double
forall a. RealFloat a => Complex a -> a
magnitude Complex Double
z) then (Double, Double) -> Either MathError (Double, Double)
forall a b. b -> Either a b
Right (Complex Double -> Double
forall a. RealFloat a => Complex a -> a
magnitude Complex Double
z, Complex Double -> Double
forall a. RealFloat a => Complex a -> a
phase Complex Double
z) else MathError -> Either MathError (Double, Double)
forall a b. a -> Either a b
Left MathError
NonFiniteResult

-- | Counterclockwise rotation in the complex plane.
rotate :: Double -> Complex Double -> Either MathError (Complex Double)
rotate :: Double -> Complex Double -> Either MathError (Complex Double)
rotate Double
angle Complex Double
z = Double -> Double -> Either MathError (Complex Double)
fromPolar Double
1 Double
angle Either MathError (Complex Double)
-> (Complex Double -> Either MathError (Complex Double))
-> Either MathError (Complex Double)
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
>>= Complex Double
-> Complex Double -> Either MathError (Complex Double)
multiply Complex Double
z

-- | Checked finite complex matrix, at most 256 rows and columns.
complexMatrix :: Int -> Int -> [Complex Double] -> Either MathError (M.Matrix (Complex Double))
complexMatrix :: Int
-> Int
-> [Complex Double]
-> Either MathError (Matrix (Complex Double))
complexMatrix Int
rows Int
cols [Complex Double]
xs
    | Int
rows Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
256 Bool -> Bool -> Bool
|| Int
cols Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
256 = MathError -> Either MathError (Matrix (Complex Double))
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"complex matrix exceeds reference budget")
    | Bool
otherwise = [Complex Double] -> Either MathError ()
inputs [Complex Double]
xs Either MathError ()
-> Either MathError (Matrix (Complex Double))
-> Either MathError (Matrix (Complex 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
-> Int
-> [Complex Double]
-> Either MathError (Matrix (Complex Double))
forall a. Int -> Int -> [a] -> Either MathError (Matrix a)
M.fromListChecked Int
rows Int
cols [Complex Double]
xs

-- | Matrix multiplication with finite-result validation.
multiplyMatrices ::
    M.Matrix (Complex Double) ->
    M.Matrix (Complex Double) ->
    Either MathError (M.Matrix (Complex Double))
multiplyMatrices :: Matrix (Complex Double)
-> Matrix (Complex Double)
-> Either MathError (Matrix (Complex Double))
multiplyMatrices Matrix (Complex Double)
a Matrix (Complex Double)
b = do
    let (Int
r, Int
k) = Matrix (Complex Double) -> (Int, Int)
forall a. Matrix a -> (Int, Int)
M.shape Matrix (Complex Double)
a; (Int
_, Int
c) = Matrix (Complex Double) -> (Int, Int)
forall a. Matrix a -> (Int, Int)
M.shape Matrix (Complex Double)
b
    _ <- Int
-> Int
-> [Complex Double]
-> Either MathError (Matrix (Complex Double))
complexMatrix Int
r Int
k (Matrix (Complex Double) -> [Complex Double]
forall a. Matrix a -> [a]
M.toList Matrix (Complex Double)
a)
    let (br, bc) = M.shape b
    _ <- complexMatrix br bc (M.toList b)
    out <- M.multiplyChecked a b
    complexMatrix r c (M.toList out)

-- | Hermitian transpose with finite-input validation.
conjugateTranspose :: M.Matrix (Complex Double) -> Either MathError (M.Matrix (Complex Double))
conjugateTranspose :: Matrix (Complex Double)
-> Either MathError (Matrix (Complex Double))
conjugateTranspose Matrix (Complex Double)
a =
    let t :: Matrix (Complex Double)
t = Matrix (Complex Double) -> Matrix (Complex Double)
forall a. Matrix a -> Matrix a
M.transposeMatrix Matrix (Complex Double)
a; (Int
r, Int
c) = Matrix (Complex Double) -> (Int, Int)
forall a. Matrix a -> (Int, Int)
M.shape Matrix (Complex Double)
t
     in Int
-> Int
-> [Complex Double]
-> Either MathError (Matrix (Complex Double))
complexMatrix Int
r Int
c ((Complex Double -> Complex Double)
-> [Complex Double] -> [Complex Double]
forall a b. (a -> b) -> [a] -> [b]
map Complex Double -> Complex Double
forall a. Num a => Complex a -> Complex a
conjugate (Matrix (Complex Double) -> [Complex Double]
forall a. Matrix a -> [a]
M.toList Matrix (Complex Double)
t))

-- | Existing checked mixed-radix/Bluestein forward FFT.
complexFFT :: [Complex Double] -> Either MathError [Complex Double]
complexFFT :: [Complex Double] -> Either MathError [Complex Double]
complexFFT = [Complex Double] -> Either MathError [Complex Double]
F.fft

-- | Existing checked inverse FFT, normalized by sample count.
complexIFFT :: [Complex Double] -> Either MathError [Complex Double]
complexIFFT :: [Complex Double] -> Either MathError [Complex Double]
complexIFFT = [Complex Double] -> Either MathError [Complex Double]
F.ifft