Bisection.hs

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

module Math.Root.Isolation.Bisection
  ( -- * The bisection method for root isolation
    Bisection
  , bisection

    -- ** Configuration options
  , BisectionCoordPicker
  , BisectionOptions(..), defaultBisectionOptions

    -- *** Helper code for picking which dimension to bisect
  , CoordBisectionData(..)
  , spread, normVI, maxVI
  , sortOnArgNE
  )
  where

-- base
import Data.Kind
  ( Type )
import Data.Foldable
 ( toList )
import Data.Functor
  ( (<&>) )
import qualified Data.List.NonEmpty as NE
  ( NonEmpty(..), cons, filter
  , head, nonEmpty, singleton, sort
  )
import Data.Semigroup
  ( Arg(..), Dual(..) )
import GHC.TypeNats
  ( Nat, KnownNat
  , type (<=)
  )

-- brush-strokes
import Math.Algebra.Dual
  ( D )
import Math.Interval
import Math.Linear
import Math.Module
  ( Module(..) )
import Math.Monomial
  ( MonomialBasis(..), linearMonomial, zeroMonomial )
import qualified Math.Ring as Ring
import Math.Root.Isolation.Core

--------------------------------------------------------------------------------
-- Bisection

-- | The bisection algorithm; see 'bisection'.
data Bisection
instance BoxCt n d => RootIsolationAlgorithm Bisection n d where
  type instance StepDescription Bisection = ( String, Double )
  type instance RootIsolationAlgorithmOptions Bisection n d = BisectionOptions n d
  rootIsolationAlgorithm
    ( BisectionOptions { canHaveSols, fallbackBisectionCoord } )
    thisRoundHist prevRoundsHist eqs box = do
      let ( boxes, whatBis ) =
            bisection
              ( canHaveSols eqs )
              ( fallbackBisectionCoord thisRoundHist prevRoundsHist eqs )
              box
      return ( whatBis, boxes )
  {-# 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 bisection method.
type BisectionOptions :: Nat -> Nat -> Type
data BisectionOptions n d =
  BisectionOptions
  { -- | Custom function to check whether the given box might contain solutions to the
    -- given equations.
    --
    -- If you always return @True@, then we will always bisect along
    -- the dimension picked by the 'fallbackBisectionCoord' function.
    --
    -- NB: only return 'False' if non-existence of solutions is guaranteed
    -- (otherwise, the root isolation algorithm might not be consistent).
    canHaveSols :: !( ( 𝕀ℝ n -> D 1 n ( 𝕀ℝ d ) ) -> Box n -> Bool )
    -- | Heuristic to choose which coordinate dimension to bisect.
    --
    -- It's only a fallback, as we prefer to bisect along coordinate dimensions
    -- that minimise the number of sub-boxes created.
  , fallbackBisectionCoord :: !( BisectionCoordPicker n d )
  }

-- | A function to choose along which coordination dimension we should bisect.
type BisectionCoordPicker n d
  =  [ ( RootIsolationStep, Box n ) ]
  -> BoxHistory n
  -> ( 𝕀ℝ n -> D 1 n ( 𝕀ℝ d ) )
  -> forall r. Show r => ( NE.NonEmpty ( Fin n, r ) -> ( r, String ) )

-- | Default options for the bisection method.
defaultBisectionOptions
  :: forall n d
  .  ( 1 <= n, BoxCt n d, Show ( 𝕀ℝ d ) )
  => Double -> Double
  -> Box n -> BisectionOptions n d
defaultBisectionOptions minWidth _ε_eq box =
  BisectionOptions
    { canHaveSols =
        \ eqs box' ->
          -- box(0)-consistency
          let iRange' :: Box d
              iRange' = eqs box' `monIndex` zeroMonomial
          in all ( unT ( origin @Double ) ∈ ) ( coordinates iRange' )

          -- box(1)-consistency
          --let box1Options = Box1Options _ε_eq ( toList $ universe @n ) ( toList $ universe @d )
          --in not $ null $ makeBox1Consistent _minWidth box1Options eqs box'

          -- box(2)-consistency
          --let box2Options = Box2Options _ε_eq 0.001 ( toList $ universe @n ) ( toList $ universe @d )
          --    box'' = makeBox2Consistent _minWidth box2Options eqs box'
          --    iRange'' :: Box d
          --    iRange'' = eqs box'' `monIndex` zeroMonomial
          --in unT ( origin @Double ) ∈ iRange''
    , fallbackBisectionCoord =
        \ _thisRoundHist _prevRoundsHist eqs possibleCoordChoices ->
          let datPerCoord =
                possibleCoordChoices <&> \ ( i, r ) ->
                  CoordBisectionData
                    { coordIndex = i
                    , coordInterval = box `index` i
                    , coordJacobianColumn = eqs box `monIndex` ( linearMonomial i )
                    , coordBisectionData = r
                    }

              -- First, check if the largest dimension is over 20 times larger
              -- than the smallest dimension; if so bisect that dimension.
          in case sortOnArgNE ( width . coordInterval ) datPerCoord of
                Arg _ d NE.:| [] ->
                  ( coordBisectionData d, "" )
                ( Arg w0 _ NE.:| ds ) ->
                  let Arg w1 d1 = last ds
                  in if w1 >= 20 * w0
                     then ( coordBisectionData d1, "tooWide" )
                     -- Otherwise, pick the coordination dimension with maximum spread
                     -- (spread = width * Jacobian column norm).
                     else
                      let isTooSmall ( Arg ( Dual w ) _ ) = w < minWidth
                      in case NE.filter ( not . isTooSmall ) $ sortOnArgNE ( Dual . spread ) datPerCoord of
                        [] -> ( coordBisectionData d1, "tooWide'" )
                        Arg _ d : _ -> ( coordBisectionData d, "spread" )
              -- TODO: pick a dimension that previous Newton steps did not
              -- manage to narrow well?
    }
{-# INLINEABLE defaultBisectionOptions #-}

sortOnArgNE :: Ord b => ( a -> b ) -> NE.NonEmpty a -> NE.NonEmpty ( Arg b a )
sortOnArgNE f = NE.sort . fmap ( \ a -> Arg ( f a ) a )
{-# INLINEABLE sortOnArgNE #-}

-- | A utility datatype that is useful in implementing bisection dimension picker
-- functions ('fallbackBisectionCoord').
type CoordBisectionData :: Nat -> Nat -> Type -> Type
data CoordBisectionData n d r =
  CoordBisectionData
    { coordIndex :: !( Fin n )
    , coordInterval :: !𝕀
    , coordJacobianColumn :: !( 𝕀ℝ d )
    , coordBisectionData :: !r
    }
deriving stock instance ( Show ( 𝕀ℝ d ), Show r )
                     => Show ( CoordBisectionData n d r )

spread :: ( BoxCt n d, Representable 𝕀 ( 𝕀ℝ d ) )
       => CoordBisectionData n d r -> Double
spread ( CoordBisectionData { coordInterval = cd, coordJacobianColumn = j_cd } )
  = width cd * normVI j_cd
{-# INLINEABLE spread #-}

normVI :: ( Applicative ( Vec d ), Representable 𝕀 ( 𝕀ℝ d ) ) => 𝕀ℝ d -> Double
normVI box =
  sqrt $ sum ( nm1 <$> coordinates box )
    where
      nm1 :: 𝕀 -> Double
      nm1 ( 𝕀 lo hi ) = max ( abs lo ) ( abs hi ) Ring.^ 2
{-# INLINEABLE normVI #-}

maxVI :: ( Applicative ( Vec d ), Representable 𝕀 ( 𝕀ℝ d ) ) => 𝕀ℝ d -> Double
maxVI box =
  maximum ( maxAbs <$> coordinates box )
    where
      maxAbs :: 𝕀 -> Double
      maxAbs ( 𝕀 lo hi ) = max ( abs lo ) ( abs hi )
{-# INLINEABLE maxVI #-}

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

-- | Bisect the given box.
--
-- (The difficult part lies in determining along which coordinate
-- dimension to bisect.)
bisection
  :: forall n
  .  ( 1 <= n, KnownNat n, Show ( 𝕀ℝ n ), Representable 𝕀 ( 𝕀ℝ n ) )
  => ( Box n -> Bool )
     -- ^ how to check whether a box contains solutions
  -> ( forall r. Show r => NE.NonEmpty ( Fin n, r ) -> ( r, String ) )
      -- ^ heuristic bisection coordinate picker
  -> Box n
      -- ^ the box to bisect
  -> ( [ Box n ], ( String, Double ) )
bisection canHaveSols pickFallbackBisCoord box =
  case NE.nonEmpty solsList of
    Nothing ->
      -- We discarded dimensions along which bisection was useless
      -- (because the interval was canonical in that dimension).
      -- If there are no dimensions left, then don't do any bisection.
      -- (TODO: we shouldn't really ever get here.)
      ( [ box ], ( "noBis", 0 / 0 ) )
    Just solsNE ->
      case findFewestSols solsNE of
        -- If there is a coordinate for which bisection results in no solutions,
        -- or in fewer sub-boxes with solutions than any other coordinate choice,
        -- pick that coordinate for bisection.
        Arg nbSols ( ( i, ( mid, subBoxesWithSols ) ) NE.:| [] ) ->
          ( subBoxesWithSols, ( "cd = " ++ show i ++ "(#subs=" ++ show nbSols ++ ")", mid ) )
        -- Otherwise, fall back to the provided heuristic.
        Arg _nbSols is ->
          let ( ( mid, subBoxesWithSols ), why ) = pickFallbackBisCoord is
          in ( subBoxesWithSols, ( why, mid ) )
  where
    solsList =
      [ Arg ( fromIntegral $ length subBoxesWithSols ) ( i, ( mid, subBoxesWithSols ) )
      | i <- toList $ universe @n
      , let ( mid, subBoxes ) = bisectInCoord box i
            subBoxesWithSols = NE.filter canHaveSols subBoxes
        -- discard coordinate dimensions in which the box is a singleton
      , length subBoxes >= 2 || null subBoxesWithSols
      ]
{-# INLINEABLE bisection #-}

-- | Bisect a box in the given coordinate dimension.
bisectInCoord
  :: Representable 𝕀 ( 𝕀ℝ n )
  => Box n -> Fin n -> ( Double, NE.NonEmpty ( Box n ) )
bisectInCoord box i =
  let z = box `index` i
      zs' = bisect z
  in ( sup ( NE.head zs' )
     , fmap ( \ z' -> set i z' box ) zs' )
{-# INLINEABLE bisectInCoord #-}

-- | Return the elements with the least argument.
--
-- NB: this function shortcuts as soon as it finds an element with argument 0.
findFewestSols :: forall a. NE.NonEmpty ( Arg Word a ) -> Arg Word ( NE.NonEmpty a )
findFewestSols ( Arg nbSols arg NE.:| args )
  | nbSols == 0
  = Arg 0 $ NE.singleton arg
  | otherwise
  = go nbSols ( NE.singleton arg ) args
  where
    go :: Word -> NE.NonEmpty a -> [ Arg Word a ] -> Arg Word ( NE.NonEmpty a )
    go bestNbSolsSoFar bestSoFar [] = Arg bestNbSolsSoFar bestSoFar
    go bestNbSolsSoFar bestSoFar ( ( Arg nbSols' arg' ) : args' )
      | nbSols' == 0
      = Arg 0 $ NE.singleton arg'
      | otherwise
      = case compare nbSols' bestNbSolsSoFar of
          LT -> go nbSols' ( NE.singleton arg' ) args'
          GT -> go bestNbSolsSoFar bestSoFar args'
          EQ -> go bestNbSolsSoFar ( arg' `NE.cons` bestSoFar ) args'