Isolation.hs

{-# LANGUAGE ScopedTypeVariables #-}

module Math.Root.Isolation
  (
    -- * Main entry point: execute root isolation strategies
    isolateRootsIn
  , Box
  , DoneBoxes(..)

    -- * General API for root isolation algorithms
  , BoxCt, BoxHistory
  , RootIsolationAlgorithm(..)
  , RootIsolationAlgorithmWithOptions(..)

    -- ** Inspecting history
  , RootIsolationStep(..)

    -- ** Visualising history
  , RootIsolationTree(..)
  , showRootIsolationTree

    -- * Configuration of the root isolation methods
  , RootIsolationOptions(..), defaultRootIsolationOptions
  , defaultRootIsolationAlgorithms

    -- * The bisection method
  , Bisection
  , BisectionOptions(..), BisectionCoordPicker
  , defaultBisectionOptions

    -- * The interval Newton method
  , NewtonOptions(..)
  , defaultNewtonOptions
    -- ** Options for the Gauss–Seidel step
  , GaussSeidelOptions(..), Preconditioner(..)
  , GaussSeidelUpdateMethod(..)

    -- * Box-consistency methods

    -- ** The @box(1)@-consistency algorithm
  , Box1
  , Box1Options(..)
  , defaultBox1Options

    -- *** Narrowing operators for the @box(1)@-consistency algorithm
  , NarrowingMethod(..)
  , narrowingMethods
  , AdaptiveShavingOptions(..)
  , defaultAdaptiveShavingOptions

    -- ** The @box(2)@-consistency algorithm
  , Box2
  , Box2Options(..)
  , defaultBox2Options
  )
  where

-- base
import Data.Kind
  ( Type )
import Data.List.NonEmpty
  ( NonEmpty )
import qualified Data.List.NonEmpty as NE
  ( NonEmpty(..), last, singleton )
import GHC.TypeNats
  ( Nat )

-- parallel
--import Control.Parallel.Strategies
--  ( Strategy, using, parTraversable, rpar )

-- transformers
import Control.Monad.Trans.Writer.CPS
  ( Writer, runWriter, tell )

-- brush-strokes
import Math.Algebra.Dual
  ( D )
import Math.Interval
import Math.Linear
import Math.Module
  ( Module(..) )
import Math.Monomial
  ( MonomialBasis(..), zeroMonomial )

import Math.Root.Isolation.Bisection
import Math.Root.Isolation.Core
import Math.Root.Isolation.Narrowing
import Math.Root.Isolation.Newton
import Math.Root.Isolation.Newton.GaussSeidel

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

-- | Options for the root isolation methods in 'isolateRootsIn'.
type RootIsolationOptions :: Nat -> Nat -> Type
newtype RootIsolationOptions n d
  = RootIsolationOptions
  { -- | Given a box and its history, return a round of root isolation strategies
    -- to run, in sequence, on this box.
    --
    -- A return value of @Left whyStop@ indicates stopping any further iterations
    -- on this box, e.g. because it is too small.
   rootIsolationAlgorithms
     :: BoxHistory n
     -> Box n
     -> Either String ( NE.NonEmpty ( RootIsolationAlgorithmWithOptions n d ) )
  }

defaultRootIsolationOptions :: forall n d. BoxCt n d => RootIsolationOptions n d
defaultRootIsolationOptions =
  RootIsolationOptions
    { rootIsolationAlgorithms = \ hist ->
        defaultRootIsolationAlgorithms minWidth ε_eq
          ( defaultNewtonOptions @n @d hist ) hist
    }
  where
    minWidth = 1e-7
    ε_eq = 5e-3
{-# INLINEABLE defaultRootIsolationOptions #-}

defaultRootIsolationAlgorithms
  :: forall n d
  .  BoxCt n d
  => Double -- ^ minimum width of boxes (don't bisect further)
  -> Double -- ^ threshold for progress
  -> NewtonOptions n d -- ^ options for Newton's method
  -> BoxHistory n
  -> Box n
  -> Either String ( NE.NonEmpty ( RootIsolationAlgorithmWithOptions n d ) )
defaultRootIsolationAlgorithms minWidth ε_eq newtonOptions history box =
  case history of
    lastRoundBoxes : _
      -- If, in the last round of strategies, we didn't try bisection...
      | all ( \case { IsolationStep @Bisection _ -> False; _ -> True } . fst ) lastRoundBoxes
      , ( _step, lastRoundFirstBox ) <- NE.last lastRoundBoxes
      -- ...and the last round didn't sufficiently reduce the size of the box...
      , not $ box `sufficientlySmallerThan` lastRoundFirstBox
      -- ...then try bisecting the box...
      -- ...unless the box is already too small, in which case we give up.
      -> if verySmall
         then Left $ "widths <= " ++ show minWidth
         else Right $ NE.singleton $ AlgoWithOptions @Bisection _bisOptions
    -- Otherwise, do a normal round.
    -- Currently: we try an interval Gauss–Seidel.
    -- (box(1)- and box(2)-consistency don't seem to help when using
    -- the complete interval union Gauss–Seidel step)
    _ -> Right $ AlgoWithOptions @Newton newtonOptions
            NE.:| []

  where
    verySmall = and $ ( \ cd -> width cd <= minWidth ) <$> coordinates box

    _bisOptions    = defaultBisectionOptions @n @d minWidth ε_eq box
    _box1Options   = defaultBox1Options @n @d ( minWidth * 100 ) ε_eq
    _box2Options   = ( defaultBox2Options @n @d minWidth ε_eq ) { box2LambdaMin = 0.001 }

    -- Did we reduce the box width by at least ε_eq
    -- in at least one of the coordinates?
    sufficientlySmallerThan :: Box n -> Box n -> Bool
    b1 `sufficientlySmallerThan` b2 =
      or $
        ( \ cd1 cd2 -> width cd2 - width cd1 >= ε_eq )
          <$> coordinates b1
          <*> coordinates b2
{-# INLINEABLE defaultRootIsolationAlgorithms #-}

--------------------------------------------------------------------------------
-- Main driver

-- | Use various methods using interval arithmetic in order to solve a system of
-- non-linear equations.
--
-- Currently supported algorithms:
--
--  - interval Newton method with Gauss–Seidel step for inversion
--    of the interval Jacobian,
--  - coordinate-wise Newton method (@box(1)@-consistency algorithm),
--  - @box(2)@-consistency algorithm,
--  - bisection.
--
-- Returns @(tree, DoneBoxes sols dunno)@ where @tree@ is the full tree
-- (useful for debugging), sols are boxes that contain a unique solution
-- (and to which Newton's method will converge starting from anywhere inside
-- the box), and @dunno@ are boxes that we have given up processing (usually
-- because they are too small); they may or may not enclose solutions.
isolateRootsIn
  :: forall n d
  .  BoxCt n d
  => RootIsolationOptions n d
      -- ^ configuration (which algorithms to use, and with what parameters)
  -> ( 𝕀ℝ n -> D 1 n ( 𝕀ℝ d ) )
      -- ^ equations to solve
  -> Box n
      -- ^ initial search domain
  -> ( [ ( Box n, RootIsolationTree ( Box n ) ) ], DoneBoxes n )
isolateRootsIn ( RootIsolationOptions { rootIsolationAlgorithms } )
  eqs initBox = runWriter $ map ( initBox , ) <$> go [] initBox

  where

    go :: BoxHistory n
       -> Box n -- box to work on
       -> Writer ( DoneBoxes n )
                 [ RootIsolationTree ( Box n ) ]
    go history cand
      | -- Check the range of the equations contains zero.
        not $ unT ( origin :: T ( ℝ d ) ) ∈∈ iRange
        -- Box doesn't contain a solution: discard it.
      = return [ RootIsolationLeaf "rangeTest" cand ]
      | otherwise
      = case rootIsolationAlgorithms history cand of
          Right strats -> doStrategies history strats cand
          Left whyStop -> do
            -- We are giving up on this box (e.g. because it is too small).
            tell $ noDoneBoxes { doneGiveUpBoxes = [ ( cand, whyStop ) ] }
            return [ RootIsolationLeaf whyStop cand ]
      where
        iRange :: Box d
        iRange = eqs cand `monIndex` zeroMonomial

    -- Run a round of root isolation strategies, then recur.
    doStrategies
      :: BoxHistory n
      -> NonEmpty ( RootIsolationAlgorithmWithOptions n d )
      -> Box n
      -> Writer ( DoneBoxes n )
                [ RootIsolationTree ( Box n ) ]
    doStrategies prevRoundsHist = do_strats []
      where
        do_strats :: [ ( RootIsolationStep, Box n ) ]
                  -> NE.NonEmpty ( RootIsolationAlgorithmWithOptions n d )
                  -> Box n
                  -> Writer ( DoneBoxes n )
                            [ RootIsolationTree ( Box n )]
        do_strats thisRoundHist ( algo NE.:| algos ) box = do
          -- Run one strategy in the round.
          ( step, boxes ) <- doStrategy algo thisRoundHist prevRoundsHist eqs box
          case algos of
            -- If there are other algorithms to run in this round, run them next.
            nextAlgo : otherAlgos ->
              let thisRoundHist' = ( step, box ) : thisRoundHist
              in recur step ( do_strats thisRoundHist' ( nextAlgo NE.:| otherAlgos ) ) boxes
            -- Otherwise, we have done one full round of strategies.
            -- Recur back to the top (calling 'go').
            [] ->
              let thisRoundHist' = ( step, box ) NE.:| thisRoundHist
                  history = thisRoundHist' : prevRoundsHist
              in recur step ( go history ) boxes

    recur :: RootIsolationStep
          -> ( Box n -> Writer ( DoneBoxes n ) [ RootIsolationTree ( Box n ) ] )
          -> [ Box n ]
          -> Writer ( DoneBoxes n ) [ RootIsolationTree ( Box n ) ]
    recur step r cands = do
      let rest :: [ ( [ ( Box n, RootIsolationTree ( Box n ) ) ], DoneBoxes n ) ]
          rest =
            map ( \ c -> runWriter $ do { trees <- r c; return [ (c, tree) | tree <- trees ] } ) cands
          ( rest2, dones ) = unzip rest --( rest `using` myParTraversable rpar )
      tell ( mconcat dones )
      return [ RootIsolationStep step $ concat rest2 ]
{-# INLINEABLE isolateRootsIn #-}

--myParTraversable :: Strategy a -> Strategy [a]
--myParTraversable _ [] = return []
--myParTraversable strat (a:as) = (a:) <$> parTraversable strat as

-- | Execute a root isolation strategy, replacing the input box with
-- (hopefully smaller) output boxes.
doStrategy
  :: forall n d
  .  RootIsolationAlgorithmWithOptions n d
  -> [ ( RootIsolationStep, Box n ) ]
  -> BoxHistory n
  -> ( 𝕀ℝ n -> D 1 n ( 𝕀ℝ d ) )
  -> Box n
  -> Writer ( DoneBoxes n )
            ( RootIsolationStep, [ Box n ] )
doStrategy algo thisRoundHist prevRoundsHist eqs cand =
  case algo of
    AlgoWithOptions @algo options -> do
      ( step, res ) <- rootIsolationAlgorithm @algo options thisRoundHist prevRoundsHist eqs cand
      return ( SomeRootIsolationStep @algo step, res )