brush-strokes-0.1.0.0: src/lib/Math/Roots.hs

{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE ScopedTypeVariables   #-}

{-# OPTIONS_GHC -Wno-name-shadowing #-}

module Math.Roots
  ( -- * Quadratic solver
    solveQuadratic

    -- * Laguerre method
  , roots, realRoots
  , laguerre

    -- * Newton–Raphson
  , newtonRaphson

    -- * Modified Halley's method M2
  , halleyM2

    -- * N-dimensional Newton's method
  , newton
  )
  where

-- base
import Control.Monad
  ( unless, when )
import Control.Monad.ST
  ( ST, runST )
import Data.Complex
  ( Complex(..), conjugate, magnitude, realPart, imagPart )
import Data.Maybe
  ( mapMaybe )
import GHC.TypeNats
  ( KnownNat )

-- deepseq
import Control.DeepSeq
  ( NFData, force )

-- eigen
import qualified Eigen.Matrix as Eigen
  ( Matrix )
import qualified Eigen.Solver.LA as Eigen
  ( Decomposition(..), solve )

-- fp-ieee
import Numeric.Floating.IEEE.NaN
  ( RealFloatNaN(copySign) )

-- primitive
import Control.Monad.Primitive
  ( PrimMonad(PrimState) )
import Data.Primitive.PrimArray
  ( PrimArray, MutablePrimArray
  , primArrayFromList
  , getSizeofMutablePrimArray, sizeofPrimArray
  , unsafeThawPrimArray
  , shrinkMutablePrimArray, readPrimArray, writePrimArray
  )
import Data.Primitive.Types
  ( Prim )

-- rounded-hw
import Numeric.Rounded.Hardware.Internal
  ( fusedMultiplyAdd )

-- brush-strokes
import Math.Algebra.Dual
import Math.Epsilon
  ( epsilon, nearZero )
import Math.Linear
import Math.Linear.Solve
  ( fromEigen, toEigen, withEigenSem )
import Math.Module
import Math.Monomial

--------------------------------------------------------------------------------

-- | Real solutions to a quadratic equation.
--
-- Coefficients are given in order of increasing degree.
--
-- Implementation taken from https://pbr-book.org/4ed/Utilities/Mathematical_Infrastructure#Quadratic.
solveQuadratic
  :: forall a. RealFloatNaN a
  => a -- ^ constant coefficient
  -> a -- ^ linear coefficient
  -> a -- ^ quadratic coefficient
  -> [ a ]
solveQuadratic c b a
  | isNaN a || isNaN b || isNaN c
  || isInfinite a || isInfinite b || isInfinite c
  = []
  | otherwise
  = case (a == 0, c == 0) of
      -- First, handle all cases in which a or c is zero.

      -- bx = 0
      (True , True )
        | b == 0
        -> [ 0, 0.5, 1 ] -- convention
        | otherwise
        -> [ 0 ]
      -- bx + c = 0
      (True , False)
        | b == 0
        -> [ ]
        | otherwise
        -> [ -c / b ]
      -- ax² + bx = 0
      (False, True )
        | b == 0
        -> [ 0 ]
        | signum a == signum b
        -> [ -b / a, 0 ]
        | otherwise
        -> [ 0, -b / a ]
      -- General case: ax² + bx + c = 0
      (False, False)
        | discr < 0
        -> []
        | otherwise
        -> let rootDiscr = sqrt discr
               q = -0.5 * ( b + copySign rootDiscr b )
               x1 = q / a
               x2 = c / q
            in if x1 > x2
               then [ x2, x1 ]
               else [ x1, x2 ]
        where discr = discriminant a b c
{-# INLINEABLE solveQuadratic #-}
{-# SPECIALISE solveQuadratic @Float  #-}
{-# SPECIALISE solveQuadratic @Double #-}
-- TODO: implement the version from the paper
--   "The Ins and Outs of Solving Quadratic Equations with Floating-Point Arithmetic"
-- which is even more robust.

-- | Kahan's method for computing the discriminant \( b^2 - 4ac \),
-- using a fused multiply-add operation to avoid cancellation in the naive
-- formula (if \( b^2 \) and \( 4ac \) are close).
--
-- From "The Ins and Outs of Solving Quadratic Equations with Floating-Point Arithmetic",
-- (Frédéric Goualard, 2023).
discriminant :: RealFloat a => a -> a -> a -> a
discriminant a b c
  -- b² and 4ac are different enough that b² - 4ac gives a good answer
  | 3 * abs d_naive >= b² - m4ac
  = d_naive
  | otherwise
  = let dp = fma b b -b²
        dq = fma ( 4 * a ) c m4ac
    in d_naive + ( dp - dq )
  where
    b² = b * b
    m4ac = -4 * a * c
    d_naive = b² + m4ac
    fma = fusedMultiplyAdd
{-# INLINEABLE discriminant #-}
{-# SPECIALISE discriminant @Float  #-}
{-# SPECIALISE discriminant @Double #-}

--------------------------------------------------------------------------------
-- Root finding using Laguerre's method

-- | Find real roots of a polynomial with real coefficients.
--
-- Coefficients are given in order of decreasing degree, e.g.:
--  x² + 7 is given by [ 1, 0, 7 ].
realRoots :: forall a. ( RealFloat a, Prim a, NFData a ) => Int -> [ a ] -> [ a ]
realRoots maxIters coeffs = mapMaybe isReal ( roots epsilon maxIters coeffs )
  where
    isReal :: Complex a -> Maybe a
    isReal ( a :+ b )
      | abs b < epsilon = Just a
      | otherwise       = Nothing

-- | Compute all roots of a polynomial with real coefficients using Laguerre's method and (forward) deflation.
--
-- Polynomial coefficients are given in order of descending degree (e.g. constant coefficient last).
--
-- N.B. The forward deflation process is only guaranteed to be numerically stable
-- if Laguerre's method finds roots in increasing order of magnitude
-- (which this function attempts to do).
roots :: forall a. ( RealFloat a, Prim a, NFData a ) => a -> Int -> [ a ] -> [ Complex a ]
roots eps maxIters coeffs = runST do
  let
    coeffPrimArray :: PrimArray a
    coeffPrimArray = primArrayFromList coeffs
    sz :: Int
    !sz = sizeofPrimArray coeffPrimArray
    n :: Int
    !n = sz - 1
  ( p :: MutablePrimArray s a ) <- unsafeThawPrimArray coeffPrimArray
  let
    go :: Int -> [ Complex a ] -> ST s [ Complex a ]
    go 0 rs = return rs
    go d rs = do
      -- Start at 0, attempting to find the root with smallest absolute value.
      -- This improves numerical stability of the forward deflation process.
      !r <- force <$> laguerre eps maxIters p 0
      if d == 1
      then pure ( r : rs )
      else
        -- real root
        if abs ( imagPart r ) < epsilon
        then do
          deflate ( realPart r ) p
          go ( d - 1 ) ( r : rs )
        else do
          deflateConjugatePair r p
          go ( d - 2 ) ( r : conjugate r : rs )
  go n []

-- | Forward deflation of a polynomial by a root: factors out the root.
--
-- Mutates the input polynomial.
deflate :: forall a m s. ( Num a, Prim a, PrimMonad m, s ~ PrimState m ) => a -> MutablePrimArray s a -> m ()
deflate r p = do
  deg <- subtract 1 <$> getSizeofMutablePrimArray p
  when ( deg >= 2 ) do
    shrinkMutablePrimArray p deg
    let
      go :: a -> Int -> m ()
      go !b !i = unless ( i >= deg ) do
        ai <- readPrimArray p i
        let
          bi :: a
          !bi = ai + r * b
        writePrimArray p i bi
        go bi ( i + 1 )
    a0 <- readPrimArray p 0
    go a0 1

-- | Forward deflation of a polynomial with real coefficients by a pair of complex-conjugate roots.
--
-- Mutates the input polynomial.
deflateConjugatePair :: forall a m s. ( Num a, Prim a, PrimMonad m, s ~ PrimState m ) => Complex a -> MutablePrimArray s a -> m ()
deflateConjugatePair ( x :+ y ) p = do
  deg <- subtract 1 <$> getSizeofMutablePrimArray p
  when ( deg >= 3 ) do
    shrinkMutablePrimArray p ( deg - 1 )
    let
      c1, c2 :: a
      !c1 = 2 * x
      !c2 = x * x + y * y
    a0 <- readPrimArray p 0
    a1 <- readPrimArray p 1
    let
      b1 :: a
      !b1 = a1 + c1 * a0
    writePrimArray p 1 b1
    let
      go :: a -> a -> Int -> m ()
      go !b !b' !i = unless ( i >= deg - 1 ) do
        ai <- readPrimArray p i
        let
          bi :: a
          !bi = ai + c1 * b - c2 * b'
        writePrimArray p i bi
        go bi b ( i + 1 )
    go b1 a0 2

-- | Laguerre's method.
--
-- Does not perform any mutation.
laguerre
  :: forall a m s
  .  ( RealFloat a, Prim a, PrimMonad m, s ~ PrimState m )
  => a                    -- ^ error tolerance
  -> Int                  -- ^ max number of iterations
  -> MutablePrimArray s a -- ^ polynomial
  -> Complex a            -- ^ initial point
  -> m ( Complex a )
laguerre eps maxIters p x0 = do
  let
    go :: Int -> Complex a -> m ( Complex a )
    go !stepNo !x = do
      x' <- laguerreStep stepNo p x
      if stepNo >= maxIters || magnitude ( x' - x ) < eps
      then pure x'
      else go ( stepNo + 1 ) x'
  go 1 x0

-- | Take a single step in Laguerre's method.
--
-- Does not perform any mutation.
laguerreStep
  :: forall a m s
  .  ( RealFloat a, Prim a, PrimMonad m, s ~ PrimState m )
  => Int -- ^ step number
  -> MutablePrimArray s a -- ^ polynomial
  -> Complex a            -- ^ initial point
  -> m ( Complex a )
laguerreStep stepNo p x = do
  sz  <- getSizeofMutablePrimArray p
  let
    n :: a
    !n = fromIntegral $ sz - 1
  ( px, p'x, p''x ) <- evalDerivatives p x
  if magnitude px == 0
  then pure x
  else do
    let
      !g     = p'x / px
      !g²    = g * g
      !h     = g² - p''x / px
      !m     = multiplicityEstimator n ( g² / h )
      !delta = sqrt $ ( ( n - m ) / m ) *: ( n *: h - g² )
      !gp    = g + delta
      !gm    = g - delta
      !denom
        | sqMag gm > sqMag gp
        = gm
        | otherwise
        = gp
    pure $
      if sqMag denom == 0
      then x + ( 0.001 :+ 0.001 ) -- random perturbation to escape saddle point
      else x - ( n :+ 0 ) / denom

  where
    (*:) :: a -> Complex a -> Complex a
    r *: (u :+ v) = ( r * u ) :+ ( r * v )

    sqMag :: Complex a -> a
    sqMag (u :+ v) = u * u + v * v

    multiplicityEstimator :: a -> Complex a -> a
    multiplicityEstimator n est

      -- Burn-in period: unmodified Laguerre for the first few steps
      | stepNo < 5

      -- Don't use an estimator if it is NaN/Infinity
      || isNaN realEst || isInfinite realEst
      || isNaN imagEst || isInfinite imagEst

      -- Only use an estimate if it is near an integer in [1..n]
      || sqMag ( est - ( rounded :+ 0 ) ) > 0.01
      || rounded >= n + 0.5
      || rounded <= 1
      = 1

      | otherwise
      = rounded
      where
        realEst, imagEst, rounded :: a
        realEst = realPart est
        imagEst = imagPart est
        rounded = fromIntegral ( round realEst :: Int )

-- | Evaluate @p(x)@, @p'(x)@ and @p''(x)@, for a polynomial with real
-- coefficients @p@, at a complex number.
--
-- Does not perform any mutation.
evalDerivatives
  :: forall a m s
  . ( RealFloat a, Prim a, PrimMonad m, s ~ PrimState m )
  => MutablePrimArray s a
  -> Complex a
  -> m ( Complex a, Complex a, Complex a )
evalDerivatives p x = do
  n <- getSizeofMutablePrimArray p
  a0 <- readPrimArray p 0
  let
    -- Generalised Horner's method
    go :: Int -> Complex a -> Complex a -> Complex a -> m ( Complex a, Complex a, Complex a )
    go !i !px !dpx !ddpx
      | i >= n
      = pure (px, dpx, ddpx * 2) -- Multiply second derivative by 2 at the end
      | otherwise
      = do
          ai <- readPrimArray p i
          let
            !ddpx' = ddpx * x + dpx
            !dpx'  = dpx * x + px
            !px'   = px * x + (ai :+ 0)
          go (i + 1) px' dpx' ddpx'
  go 1 (a0 :+ 0) 0 0

{-# SPECIALISE maxRealFloat @Float  #-}
{-# SPECIALISE maxRealFloat @Double #-}
maxRealFloat :: forall r. RealFloat r => r
maxRealFloat = encodeFloat m n
  where
    !b = floatRadix @r 0
    !e = floatDigits @r 0
    !(_, !e') = floatRange @r 0
    !m = b ^ e - 1
    !n = e' - e

{-
for some reason GHC isn't able to optimise maxRealFloat @Double into

maxDouble :: Double
maxDouble = D# 1.7976931348623157e308##
-}

--------------------------------------------------------------------------------
-- Newton–Raphson

-- | Newton–Raphson root-finding with Armijo backtracking line search and
-- quadratic interpolation fallback.
newtonRaphson
  :: forall r
  .  ( RealFloat r, Show r )
  => Word                 -- ^ maximum number of iterations
  -> Int                  -- ^ desired digits of precision
  -> ( r -> (# r, r #) )  -- ^ function and its derivative
  -> r                    -- ^ initial guess
  -> Maybe r
newtonRaphson maxIters digits f x0 =
  let (# f_x0, f'_x0 #) = f x0
  in go 0 x0 f_x0 f'_x0
  where
    !factor = encodeFloat 1 ( 1 - digits )

    go :: Word -> r -> r -> r -> Maybe r
    go !iters !x !f_x !f'_x
      | f_x == 0
      = Just x
      | iters >= maxIters || f'_x == 0 -- trapped at a local extremum
      = Nothing
      | otherwise
      = case newtonRaphsonStep factor f x f_x f'_x of
          (# done | #) -> done
          (# | (# x', fx', f'x' #) #) ->
            go ( iters + 1 ) x' fx' f'x'
{-# SPECIALISE newtonRaphson @Float  #-}
{-# SPECIALISE newtonRaphson @Double #-}
{-# INLINEABLE newtonRaphson #-}

newtonRaphsonStep
  :: ( RealFloat r, Show r )
  => r
  -> ( r -> (# r, r #) )
  -> r -> r -> r
  -> (# Maybe r | (# r, r, r #) #)
newtonRaphsonStep !factor f !x !f_x !f'_x = lineSearch 1
  where
    -- Armijo "sufficient decrease" parameter
    !c1 = 1e-4

    !g_0    = f_x * f_x
    !p     = -f_x / f'_x
    -- Directional derivative of f(x)^2 along the Newton step
    !slope = -2 * g_0

    -- Backtracking line search
    lineSearch !α

      -- Termination criteria:
      --  1. Close enough to a solution.
      | abs δ <= max epsilon ( abs ( x' * factor ) )
      = (# Just x' | #)
      --  2. α got too small, the algorithm has stalled
      | α < 1e-4
      = (# Nothing | #)

      -- Otherwise, either accept step or shrink α.
      | otherwise
      = let
          !(# f_x', f'_x' #) = f x'
          !g_α = f_x' * f_x'
        in
          -- Armijo condition for accepting step
          if g_α <= g_0 + c1 * α * slope
          then (# | (# x', f_x', f'_x' #) #)
          else
            -- Quadratic interpolation fallback
            let !α_tmp = ( g_0 * α * α ) / ( g_α - g_0 + 2 * g_0 * α )
                -- Clamp α to [0.1 * α, 0.5 * α] to prevent numerical instability
                !α' = max ( 0.1 * α ) ( min ( 0.5 * α ) α_tmp )
            in lineSearch α'
       where
        !δ = α * p
        !x' = x + δ
{-# SPECIALISE newtonRaphsonStep @Float  #-}
{-# SPECIALISE newtonRaphsonStep @Double #-}
{-# INLINEABLE newtonRaphsonStep #-}

--------------------------------------------------------------------------------
-- Modified Halley method

-- | Take a single step in the M2 modified Halley method.
--
-- Taken from @Some variants of Halley’s method with memory and their applications for solving several chemical problems@
-- by A. Cordero, H. Ramos & J.R. Torregrosa, J Math Chem 58, 751–774 (2020).
--
-- @https://doi.org/10.1007/s10910-020-01108-3@
halleyM2Step
  :: ( RealFloat a, Show a )
  => (# a, (# a, a #) #) -> (# a, (# a, a #) #) -> a
halleyM2Step (# x_nm1, (# f_nm1, f'_nm1 #) #) (# x_n, (# f_n, f'_n #) #)
  | nearZero dx
  = if nearZero f'_n
    then x_n -- Stalled: Derivative is zero, cannot take a Newton step
    else x_n - ( f_n / f'_n ) -- simple Newton step fallback
  | nearZero num && nearZero denom
  = 0.001 * signum num * signum denom
  | nearZero num
  = num
  | otherwise
  = num / denom
  where
    !u = f_n * f_nm1 * (f_n - f_nm1)
    !dx = x_n - x_nm1
    !g1 = f_nm1 ^ ( 2 :: Int ) * f'_n
    !g2 = f_n ^ ( 2 :: Int ) * f'_nm1
    !num = (x_n + x_nm1) * u - dx * ( g1 * x_n + g2 * x_nm1)
    !denom = 2 * u - dx * ( g1 + g2 )

{-# SPECIALISE halleyM2Step @Float  #-}
{-# SPECIALISE halleyM2Step @Double #-}
{-# INLINEABLE halleyM2Step #-}

{-# SPECIALISE halleyM2 @Float  #-}
{-# SPECIALISE halleyM2 @Double #-}
{-# INLINEABLE halleyM2 #-}
halleyM2
  :: forall r
  .  ( RealFloat r, Show r )
  => Word                -- ^ maximum number of iterations
  -> Int                 -- ^ desired digits of precision
  -> ( r -> (# r, r #) ) -- ^ function and its derivative
  -> r                   -- ^ initial guess
  -> Maybe r
halleyM2 maxIters digits f x0 =
  let
    y0 = (# x0, df_x0 #)
    !df_x0@(# f_x0, f'_x0 #) = f x0
    !factor = encodeFloat 1 ( 1 - digits )

    -- Start off with a Newton step to get the second initial x value.
  in case newtonRaphsonStep factor f x0 f_x0 f'_x0 of
       -- Newton converged immediately on the first step
       (# Just x1 | #) -> Just x1

       -- Newton stalled immediately (e.g., zero derivative)
       (# Nothing | #) -> Nothing

       -- Newton took a successful step, we have x1. Proceed with Halley M2.
       (# | (# x1, f_x1, f'_x1 #) #) ->
         let y1 = (# x1, (# f_x1, f'_x1 #) #)
         in go 0 y0 y1
  where
    !factor = encodeFloat 1 ( 1 - digits )

    go :: Word -> (# r, (# r, r #) #) -> (# r, (# r, r #) #) -> Maybe r
    go i y_nm1 y_n@(# x_n, _ #) =
      let !x_np1 = halleyM2Step y_nm1 y_n
      in
        if
          | i >= maxIters
          || abs ( x_np1 - x_n ) <= abs ( x_n * factor )
          -> Just x_np1
          | otherwise
          -> go (i+1) y_n (# x_np1, f x_np1 #)

--------------------------------------------------------------------------------

-- | @n@-dimensional Newton's method.
newton
  :: forall n d
  .  ( KnownNat n, KnownNat d
     , Show ( ℝ n ), Show ( ℝ d )
     , NFData ( ℝ n )
     , Module Double ( T ( ℝ n ) )
     , MonomialBasis ( D 1 n ), Deg ( D 1 n ) ~ 1, Vars ( D 1 n ) ~ n
     , Representable Double ( ℝ d )
     , Representable Double ( ℝ n )
     )
  => ( ℝ n -> D 1 n ( ℝ d ) )
      -- ^ equations
  -> ℝ n
     -- ^ initial guess
  -> ℝ n
newton f x0 = go 0 x0
  where
    maxIters :: Int
    maxIters = 40
    relError, absError :: Double
    relError = 1e-9
    absError = 1e-11

    go :: Int -> ℝ n -> ℝ n
    go !i !x
      | i >= maxIters
      = x
      | norm ( unT δ ) < absError + relError * norm x
      = x'
      | otherwise
      = go ( i + 1 ) x'
      where
        ( _norm_fx, δ ) = newtonStep x
        x' = unT $ T x ^-^ δ

    norm :: forall i. Representable Double ( ℝ i ) => ℝ i -> Double
    norm x =
      sqrt $ sum ( ( \ x -> x * x ) <$> coordinates x )
    {-# INLINEABLE norm #-}

    newtonStep :: ℝ n -> ( Double, T ( ℝ n ) )
    newtonStep x =
      let df_x = f x
          f_x :: ℝ d
          f_x = df_x `monIndex` zeroMonomial
          j_x :: Vec n ( ℝ d )
          j_x = fmap ( \ i -> df_x `monIndex` linearMonomial i ) ( universe @n )
          δ :: ℝ n
          !δ = withEigenSem $
            ( \ ( Vec cols ) -> case cols of [ !c ] -> c; _ -> error "internal error in Newton's method" ) $
              fromEigen $
                Eigen.solve Eigen.JacobiSVD
                  ( toEigen j_x :: Eigen.Matrix d n Double )
                  ( toEigen ( Vec [ f_x ] ) :: Eigen.Matrix d 1 Double )
      in ( norm f_x, T δ )

{-# INLINEABLE newton #-}
{-# SPECIALISE newton @1 @1 #-}
{-# SPECIALISE newton @1 @2 #-}
{-# SPECIALISE newton @1 @3 #-}
{-# SPECIALISE newton @2 @2 #-}
{-# SPECIALISE newton @2 @3 #-}

-- TODO: implement the W4 method?