From 43098dec0124af354e2093d4040d1a2acb23a59b Mon Sep 17 00:00:00 2001
From: sheaf <sam.derbyshire@gmail.com>
Date: Mon, 9 Jan 2023 23:59:00 +0100
Subject: [PATCH] Newton-Raphson implementation

---
 src/splines/Math/Roots.hs | 122 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 121 insertions(+), 1 deletion(-)

diff --git a/src/splines/Math/Roots.hs b/src/splines/Math/Roots.hs
index 10f0ac5..c3fd8bb 100644
--- a/src/splines/Math/Roots.hs
+++ b/src/splines/Math/Roots.hs
@@ -1,4 +1,7 @@
-{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE DuplicateRecordFields #-}
+{-# LANGUAGE ScopedTypeVariables   #-}
+
+{-# OPTIONS_GHC -Wno-name-shadowing #-}
 
 module Math.Roots where
 
@@ -274,3 +277,120 @@ derivative p = do
       go ( i + 1 )
   go 0
   pure p'
+
+--------------------------------------------------------------------------------
+
+data NewtonRaphson
+  = NewtonRaphson
+    { f :: Double -> (# Double, Double #)
+    , maxIters     :: Word
+    , min_x, max_x :: Double
+    , digits       :: Int }
+
+-- Newton–Raphson implementation taken from Boost C++ library:
+-- https://github.com/boostorg/math/blob/0dc6a70caa6bbec2b6ae25eede36c430f0ccae13/include/boost/math/tools/roots.hpp#L217
+
+newtonRaphson :: NewtonRaphson
+              -> Double
+              -> Maybe Double
+newtonRaphson ( NewtonRaphson { f, maxIters, min_x, max_x, digits } ) x0 =
+    go $
+      NewtonRaphsonState
+        { f_x_prev = 0
+        , x = x0
+        , δ = maxDouble, δ1 = maxDouble
+        , iters = 0
+        , min_x, max_x
+        , min_f_x = 0, max_f_x = 0 }
+  where
+    factor = encodeFloat 1 ( 1 - digits )
+    go ( NewtonRaphsonState { f_x_prev, x, δ = δ1, δ1 = δ2, iters, min_x, max_x, min_f_x, max_f_x } )
+      = case f x of
+          (# f_x, f'_x #)
+            | f_x == 0
+            -> Just x
+            | δ <-
+               if | f'_x == 0
+                  -> handleZeroDerivative f_x_prev x f_x δ1 min_x max_x
+                  | otherwise
+                  -> f_x / f'_x
+            , (# δ, δ1 #) <-
+                if | abs ( δ * δ ) > abs δ2
+                   , let shift = if δ > 0 then 0.5 * ( x - min_x ) else 0.5 * ( x - max_x )
+                   , δ <- if x /= 0 && ( abs shift > abs x )
+                          then signum δ * abs x * 1.1
+                          else shift
+                   -> (# δ, 3 * δ #)
+                   | otherwise
+                   -> (# δ, δ1 #)
+            , let new_x = x - δ
+            , (# δ, new_x #) <-
+               if | new_x <= min_x
+                  , δ <- 0.5 * ( x - min_x )
+                  -> (# δ, x - δ #)
+                  | new_x >= max_x
+                  , δ <- 0.5 * ( x - max_x )
+                  -> (# δ, x - δ #)
+                  | otherwise
+                  -> (# δ, new_x #)
+            -> if
+                |  abs δ <= abs ( new_x * factor )
+                || new_x == min_x || new_x == max_x
+                -> Just x
+                | iters >= maxIters
+                -> Nothing
+                | (# min_x, max_x, min_f_x, max_f_x #) <-
+                     if δ > 0
+                     then (# min_x, x, min_f_x, f_x #)
+                     else (# x, max_x, f_x, max_f_x #)
+                -> if min_f_x * max_f_x > 0
+                   then Nothing
+                   else
+                    go $ NewtonRaphsonState
+                         { f_x_prev = f_x, x = new_x
+                         , δ, δ1
+                         , iters = iters + 1
+                         , min_x, max_x
+                         , min_f_x, max_f_x
+                         }
+
+    handleZeroDerivative :: Double
+                         -> Double -> Double
+                         -> Double
+                         -> Double -> Double
+                         -> Double
+    handleZeroDerivative f_x_prev x f_x δ min_x max_x
+      -- Handle zero derivative on first iteration.
+      | f_x_prev == 0
+      , x_prev <- if x <= 0.5 * ( min_x + max_x ) then max_x else min_x
+      , (# f_x_prev, _ #) <- f x_prev
+      , δ <- x_prev - x
+      = finish f_x_prev δ
+      | otherwise
+      = finish f_x_prev δ
+
+      where
+        finish f_x_prev δ
+          | signum f_x_prev * signum f_x < 0
+          = if δ < 0 then 0.5 * ( x - min_x ) else 0.5 * ( x - max_x )
+          | otherwise
+          = if δ < 0 then 0.5 * ( x - max_x ) else 0.5 * ( x - min_x )
+
+-- | Loop state for the 'newtonRaphson' function.
+data NewtonRaphsonState =
+  NewtonRaphsonState
+    { f_x_prev         :: !Double
+    , x                :: !Double
+    , δ, δ1            :: !Double
+    , iters            :: !Word
+    , min_x, max_x     :: !Double
+    , min_f_x, max_f_x :: !Double }
+
+maxDouble :: Double
+maxDouble = encodeFloat m n
+  where
+    b = floatRadix ( 0 :: Double )
+    e = floatDigits ( 0 :: Double )
+    (_, e') = floatRange ( 0 :: Double )
+    m = b ^ e - 1
+    n = e' - e