Newton.hs

{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE UndecidableInstances #-}

module Math.Root.Isolation.Newton
  ( -- * The interval Newton method,
    -- with Gauss–Seidel step or explicit linear programming
    Newton
  , intervalNewton

    -- ** Configuration options
  , NewtonOptions(..)
  , defaultNewtonOptions
  )
  where

-- base
import Prelude hiding ( unzip )
import Control.Arrow
  ( first )
import Data.Bifunctor
  ( Bifunctor(bimap) )
import Data.Kind
  ( Type )
import Data.List
  ( partition )
import GHC.TypeNats
  ( Nat, KnownNat, type (<=) )

-- deepseq
import Control.DeepSeq
  ( force )

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

-- brush-strokes
import Math.Algebra.Dual
  ( D )
import Math.Interval
import Math.Linear
import Math.Module
  ( Module(..) )
import Math.Monomial
  ( MonomialBasis(..), linearMonomial, zeroMonomial )
import Math.Root.Isolation.Core
import Math.Root.Isolation.Newton.GaussSeidel
import Math.Root.Isolation.Newton.LP
import Math.Root.Isolation.Utils

--------------------------------------------------------------------------------
-- Interval Newton

-- | The interval Newton method; see 'intervalNewton'.
data Newton
instance (Show (NewtonOptions n d), BoxCt n d) => RootIsolationAlgorithm Newton n d where
  type instance StepDescription Newton = (String, TimeInterval)
  type instance RootIsolationAlgorithmOptions Newton n d = NewtonOptions n d
  rootIsolationAlgorithm opts _thisRoundHist _prevRoundsHist eqs box =
    first ( "Newton " ++ show opts, ) <$> intervalNewton @n @d opts eqs box
  {-# INLINEABLE rootIsolationAlgorithm #-}
  {-# SPECIALISE rootIsolationAlgorithm @2 @3 #-}
    -- NB: including this to be safe. The specialiser seems to sometimes
    -- be able to generate this specialisation on its own, and sometimes not.

-- | Options for the interval Newton method.
type NewtonOptions :: Nat -> Nat -> Type
data NewtonOptions n d where
  -- | Use the Gauss–Seidel method to solve linear systems.
  NewtonGaussSeidel
    :: GaussSeidelOptions n d -> NewtonOptions n d
  -- | Use linear programming to solve linear systems (2 dimensions only).
  NewtonLP
    :: NewtonOptions 2 d

deriving stock instance Show ( NewtonOptions n d )

-- | Default options for the interval Newton method.
defaultNewtonOptions
  :: forall n d
  .  ( KnownNat n, KnownNat d
     , 1 <= n, 1 <= d, n <= d
     , Representable Double ( ℝ n )
     , Representable Double ( ℝ d )
     , Representable 𝕀 ( 𝕀ℝ n )
     , Representable 𝕀 ( 𝕀ℝ d )
     )
  => BoxHistory n
  -> NewtonOptions n d
defaultNewtonOptions history =
  NewtonGaussSeidel $ defaultGaussSeidelOptions history
{-# INLINEABLE defaultNewtonOptions #-}

-- | Interval Newton method with Gauss–Seidel step.
intervalNewton
  :: forall n d
  .  BoxCt n d
  => NewtonOptions n d
  -> ( 𝕀ℝ n -> D 1 n ( 𝕀ℝ d ) )
      -- ^ equations
  -> 𝕀ℝ n
     -- ^ box
  -> Writer ( DoneBoxes n ) ( TimeInterval, [ 𝕀ℝ n ] )
intervalNewton opts eqs x = case opts of
  NewtonGaussSeidel
    ( GaussSeidelOptions
      { gsPreconditioner = precondMeth
      , gsPickEqs = pickEqs
      , gsUpdate
      } ) ->
    let x_mid = point $ boxMidpoint x
        f :: 𝕀ℝ n -> 𝕀ℝ n
        f = \ x_0 -> pickEqs $ eqs x_0 `monIndex` zeroMonomial
        f'_x :: Vec n ( 𝕀ℝ n )
        f'_x = fmap ( \ i -> pickEqs $ eqs x `monIndex` linearMonomial i ) ( universe @n )

        -- Interval Newton method: take one Gauss–Seidel step
        -- for the system of equations f'(x) ( x - x_mid ) = - f(x_mid).
        minus_f_x_mid = unT $ -1 *^ T ( boxMidpoint $ f x_mid )

        -- Precondition the above linear system into A ( x - x_mid ) = B.
        !( !a, !b ) = force $
           precondition precondMeth
            f'_x ( point minus_f_x_mid )
        !x'_0 = force ( T x ^-^ T x_mid )

        -- NB: we have to change coordinates, putting the midpoint of the box
        -- at the origin, in order to take a Gauss–Seidel step.
        ( x's, dt ) =
          timeInterval
            ( gaussSeidelUpdate gsUpdate a b )
            x'_0
        gsGuesses = map ( first ( \ x' -> unT $ x' ^+^ T x_mid ) ) x's
        ( done, todo ) = bimap ( map fst ) ( map fst )
                       $ partition snd gsGuesses
    in -- If the Gauss–Seidel step was a contraction, then the box
       -- contains a unique solution (by the Banach fixed point theorem).
       --
       -- These boxes can thus be directly added to the solution set:
       -- Newton's method is guaranteed to converge to the unique solution.
      do tell $ noDoneBoxes { doneSolBoxes = done }
         return ( dt, todo )
  NewtonLP ->
    -- TODO: reduce duplication with the above.
    let x_mid = point $ boxMidpoint x
        f :: 𝕀ℝ 2 -> 𝕀ℝ d
        f = \ x_0 -> eqs x_0 `monIndex` zeroMonomial
        f'_x :: Vec 2 ( 𝕀ℝ d )
        f'_x = fmap ( \ i -> eqs x `monIndex` linearMonomial i ) ( universe @2 )

        minus_f_x_mid = unT $ -1 *^ T ( boxMidpoint $ f x_mid )
        !( !a, !b ) = force ( f'_x, point minus_f_x_mid )
        !x'_0 = force ( T x ^-^ T x_mid )
        ( x's, dt ) =
          timeInterval
            ( solveIntervalLinearEquations a b ) x'_0
        lpGuesses = map ( first ( \ x' -> unT $ x' ^+^ T x_mid ) ) x's
        ( done, todo ) = bimap ( map fst ) ( map fst )
                       $ partition snd lpGuesses
    in do tell $ noDoneBoxes { doneSolBoxes = done }
          return ( dt, todo )
{-# INLINEABLE intervalNewton #-}
{-# SPECIALISE intervalNewton @2 @3 #-}

{-

mbDeg = topologicalDegree 0.005 f x
det = case f'_x of
        Vec [ c1, c2 ] ->
          let a_11 = c1 `index` Fin 1
              a_12 = c2 `index` Fin 1
              a_21 = c1 `index` Fin 2
              a_22 = c2 `index` Fin 2
          in a_11 * a_22 - a_12 * a_21
        _ -> error "TODO: just testing n=2 here"

if | not $ 0 ∈ det
   , mbDeg == Just 0
   -> return []
     -- If the Jacobian is invertible over the box, then the topological
     -- degree tells us exactly how many solutions there are in the box.
-}