{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE Safe #-}

{- | 
Module      :  Physics.Learn.RootFinding
Copyright   :  (c) Scott N. Walck 2012-2017
License     :  BSD3 (see LICENSE)
Maintainer  :  Scott N. Walck <walck@lvc.edu>
Stability   :  experimental

Functions for approximately solving equations like f(x) = 0.
These functions proceed by assuming that f is continuous,
and that a root is bracketed.  A bracket around a root consists
of numbers a, b such that f(a) f(b) <= 0.  Since the product
changes sign, there must be an x with a < x < b such that f(x) = 0.
-}

module Physics.Learn.RootFinding
    ( findRoots
    , findRootsN
    , findRoot
    , bracketRoot
    , bracketRootStep
    )
    where

-- | Given an initial bracketing of a root
--   (an interval (a,b) for which f(a) f(b) <= 0),
--   produce a bracket (c,d) for which |c-d| < desired accuracy.
bracketRoot :: (Ord a, Fractional a) =>
               a         -- ^ desired accuracy
            -> (a -> a)  -- ^ function
            -> (a,a)     -- ^ initial bracket
            -> (a,a)     -- ^ final bracket
bracketRoot :: forall a.
(Ord a, Fractional a) =>
a -> (a -> a) -> (a, a) -> (a, a)
bracketRoot a
dx a -> a
f (a
a,a
b)
    = let fa :: a
fa = a -> a
f a
a
          fb :: a
fb = a -> a
f a
b
          bRoot :: ((a, a), (a, a)) -> (a, a)
bRoot ((a
c,a
fc),(a
d,a
fd)) = let m :: a
m = (a
c forall a. Num a => a -> a -> a
+ a
d) forall a. Fractional a => a -> a -> a
/ a
2
                                      fm :: a
fm = a -> a
f a
m
                                  in if forall a. Num a => a -> a
abs (a
c forall a. Num a => a -> a -> a
- a
d) forall a. Ord a => a -> a -> Bool
<  a
dx
                                     then (a
c,a
d)
                                     else if a
fc forall a. Num a => a -> a -> a
* a
fm forall a. Ord a => a -> a -> Bool
<= a
0
                                          then ((a, a), (a, a)) -> (a, a)
bRoot ((a
c,a
fc),(a
m,a
fm))
                                          else ((a, a), (a, a)) -> (a, a)
bRoot ((a
m,a
fm),(a
d,a
fd))
      in if a
fa forall a. Num a => a -> a -> a
* a
fb forall a. Ord a => a -> a -> Bool
> a
0
         then forall a. HasCallStack => [Char] -> a
error [Char]
"bracketRoot:  initial interval is not a bracket"
         else ((a, a), (a, a)) -> (a, a)
bRoot ((a
a,a
fa),(a
b,a
fb))

-- | Given a bracketed root, return a half-width bracket.
bracketRootStep :: (Ord a, Fractional a) =>
                   (a -> a)       -- ^ function
                -> ((a,a),(a,a))  -- ^ original bracket
                -> ((a,a),(a,a))  -- ^ new bracket
bracketRootStep :: forall a.
(Ord a, Fractional a) =>
(a -> a) -> ((a, a), (a, a)) -> ((a, a), (a, a))
bracketRootStep a -> a
f ((a
a,a
fa),(a
b,a
fb))
    = let m :: a
m = (a
a forall a. Num a => a -> a -> a
+ a
b) forall a. Fractional a => a -> a -> a
/ a
2
          fm :: a
fm = a -> a
f a
m
      in if a
fa forall a. Num a => a -> a -> a
* a
fm forall a. Ord a => a -> a -> Bool
<= a
0
         then ((a
a,a
fa),(a
m,a
fm))
         else ((a
m,a
fm),(a
b,a
fb))

findRootMachinePrecision :: (Double -> Double)
                         -> ((Double,Double),(Double,Double))
                         -> Double
findRootMachinePrecision :: (Double -> Double)
-> ((Double, Double), (Double, Double)) -> Double
findRootMachinePrecision Double -> Double
f ((Double
c,Double
fc),(Double
d,Double
fd))
    = let m :: Double
m = (Double
c forall a. Num a => a -> a -> a
+ Double
d) forall a. Fractional a => a -> a -> a
/ Double
2
          fm :: Double
fm = Double -> Double
f Double
m
      in if Double
fc forall a. Eq a => a -> a -> Bool
== Double
0
         then Double
c
         else if Double
fd forall a. Eq a => a -> a -> Bool
== Double
0
              then Double
d
              else if Double
c forall a. Eq a => a -> a -> Bool
== Double
m
                   then Double
c
                   else if Double
m forall a. Eq a => a -> a -> Bool
== Double
d
                        then Double
d
                        else if Double
fc forall a. Num a => a -> a -> a
* Double
fm forall a. Ord a => a -> a -> Bool
<= Double
0
                             then (Double -> Double)
-> ((Double, Double), (Double, Double)) -> Double
findRootMachinePrecision Double -> Double
f ((Double
c,Double
fc),(Double
m,Double
fm))
                             else (Double -> Double)
-> ((Double, Double), (Double, Double)) -> Double
findRootMachinePrecision Double -> Double
f ((Double
m,Double
fm),(Double
d,Double
fd))

-- | Find a single root in a bracketed region.
--   The algorithm continues until it exhausts the
--   precision of a 'Double'.  This could cause the function to hang.
findRoot :: (Double -> Double)  -- ^ function
         -> (Double,Double)     -- ^ initial bracket
         -> Double              -- ^ approximate root
findRoot :: (Double -> Double) -> (Double, Double) -> Double
findRoot Double -> Double
f (Double
a,Double
b)
    = let fa :: Double
fa = Double -> Double
f Double
a
          fb :: Double
fb = Double -> Double
f Double
b
      in if Double
fa forall a. Num a => a -> a -> a
* Double
fb forall a. Ord a => a -> a -> Bool
> Double
0
         then forall a. HasCallStack => [Char] -> a
error [Char]
"bracketRoot:  initial interval is not a bracket"
         else (Double -> Double)
-> ((Double, Double), (Double, Double)) -> Double
findRootMachinePrecision Double -> Double
f ((Double
a,Double
fa),(Double
b,Double
fb))

-- | Find a list of roots for a function over a given range.
--   First parameter is the initial number of intervals to
--   use to find the roots.  If roots are closely spaced,
--   this number of intervals may need to be large.
findRootsN :: Int                 -- ^ initial number of intervals to use
           -> (Double -> Double)  -- ^ function
           -> (Double,Double)     -- ^ range over which to search
           -> [Double]            -- ^ list of roots
findRootsN :: Int -> (Double -> Double) -> (Double, Double) -> [Double]
findRootsN Int
n Double -> Double
f (Double
a,Double
b)
    = let dx :: Double
dx = (Double
b forall a. Num a => a -> a -> a
- Double
a) forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
          xs :: [Double]
xs = [Double
a,Double
aforall a. Num a => a -> a -> a
+Double
dx..Double
b]
      in forall a b. (a -> b) -> [a] -> [b]
map ((Double -> Double)
-> ((Double, Double), (Double, Double)) -> Double
findRootMachinePrecision Double -> Double
f) [((Double
x0,Double
fx0),(Double
x1,Double
fx1)) | (Double
x0,Double
x1) <- forall a b. [a] -> [b] -> [(a, b)]
zip [Double]
xs (forall a. [a] -> [a]
tail [Double]
xs), let fx0 :: Double
fx0 = Double -> Double
f Double
x0, let fx1 :: Double
fx1 = Double -> Double
f Double
x1, Double
fx0 forall a. Num a => a -> a -> a
* Double
fx1 forall a. Ord a => a -> a -> Bool
<= Double
0]

-- | Find a list of roots for a function over a given range.
--   There are no guarantees that all roots will be found.
--   Uses 'findRootsN' with 1000 intervals.
findRoots :: (Double -> Double)  -- ^ function
          -> (Double,Double)     -- ^ range over which to search
          -> [Double]            -- ^ list of roots
findRoots :: (Double -> Double) -> (Double, Double) -> [Double]
findRoots = Int -> (Double -> Double) -> (Double, Double) -> [Double]
findRootsN Int
1000