Core.hs

{-# LANGUAGE AllowAmbiguousTypes     #-}
{-# LANGUAGE PolyKinds               #-}
{-# LANGUAGE ScopedTypeVariables     #-}
{-# LANGUAGE UndecidableInstances    #-}
{-# LANGUAGE UndecidableSuperClasses #-}

-- | Core definitions and utilities common to root isolation methods.
module Math.Root.Isolation.Core
  ( -- * Root isolation types
    Box
  , DoneBoxes(..), noDoneBoxes

    -- * General typeclass for root isolation methods
  , BoxCt
  , RootIsolationAlgorithm(..)
  , RootIsolationAlgorithmWithOptions(..)

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

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

    -- * Utility functions
  , pipeFunctionsWhileTrue
  , forEachCoord

    -- * Timing
  , TimeInterval(..), timeInterval
  ) where

-- base
import Data.Foldable
  ( toList )
import Data.Kind
  ( Type, Constraint )
import qualified Data.List.NonEmpty as NE
  ( NonEmpty )
import Data.Type.Equality
  ( (:~~:)(HRefl) )
import Data.Typeable
  ( Typeable, heqT )
import Numeric
  ( showFFloat )
import GHC.Clock
  ( getMonotonicTime )
import GHC.TypeNats
  ( Nat, KnownNat, type (<=) )
import System.IO.Unsafe
  ( unsafePerformIO )

-- containers
import Data.Tree
  ( Tree(..) )

-- deepseq
import Control.DeepSeq
  ( NFData(..), deepseq )

-- transformers
import Control.Monad.Trans.State.Strict as State
  ( State, get, put )
import Control.Monad.Trans.Writer.CPS
  ( Writer )

-- tree-view
import Data.Tree.View
  ( showTree )

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

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

-- | An axis-aligned box in @n@-dimensions.
type Box n = 𝕀ℝ n

-- | Dimension constraints for root isolation in a system of equations:
--
--   - @n@: number of variables
--   - @d@: number of equations
--
-- NB: we require n <= d (no support for under-constrained systems).
--
-- NB: in practice, this constraint should specialise away.
type BoxCt n d =
  ( KnownNat n, KnownNat d
  , 1 <= n, 1 <= d, n <= d

  , Show ( 𝕀ℝ n ), Show ( ℝ n )
  , Show ( 𝕀ℝ d ), Show ( ℝ d )
  , Eq ( ℝ n )
  , Representable Double ( ℝ n )
  , Representable 𝕀 ( 𝕀ℝ n )
  , MonomialBasis ( D 1 n )
  , Deg ( D 1 n ) ~ 1
  , Vars ( D 1 n ) ~ n
  , Module Double ( T ( ℝ n ) )
  , Module 𝕀 ( T ( 𝕀ℝ n ) )
  , NFData ( ℝ n )
  , NFData ( 𝕀ℝ n )

  , Ord ( ℝ d )
  , Module Double ( T ( ℝ d ) )
  , Representable Double ( ℝ d )
  , Representable 𝕀 ( 𝕀ℝ d )
  , NFData ( ℝ d )
  , NFData ( 𝕀ℝ d )
  )
-- | Boxes we are done with and will not continue processing.
data DoneBoxes n =
  DoneBoxes
    { -- | Boxes which definitely contain a unique solution.
      doneSolBoxes    :: ![ Box n ]
      -- | Boxes which may or may not contain solutions,
      -- and that we have stopped processing for some reason.
    , doneGiveUpBoxes :: ![ ( Box n, String ) ]
    }
deriving stock instance Show ( Box n ) => Show ( DoneBoxes n )

instance Semigroup ( DoneBoxes n ) where
  DoneBoxes a1 b1 <> DoneBoxes a2 b2 = DoneBoxes ( a1 <> a2 ) ( b1 <> b2 )
instance Monoid ( DoneBoxes n ) where
  mempty = noDoneBoxes
noDoneBoxes :: DoneBoxes n
noDoneBoxes = DoneBoxes [] []

--------------------------------------------------------------------------------
-- Class for all root isolation algorithms.
--
-- This keeps the implementation open-ended, and allows inspection of
-- other root isolation methods, so that heuristics can look at
-- what happened in previous steps to decide what to do.

-- | Existential wrapper over any root isolation algorithm,
-- with the options necessary to run it.
type RootIsolationAlgorithmWithOptions :: Nat -> Nat -> Type
data RootIsolationAlgorithmWithOptions n d where
  AlgoWithOptions
    :: forall {k :: Type} {n :: Nat} {d :: Nat} ( ty :: k )
    .  RootIsolationAlgorithm ty n d
    => RootIsolationAlgorithmOptions ty n d
    -> RootIsolationAlgorithmWithOptions n d

-- | Type-class for root isolation algorithms.
--
-- This design keeps the set of root isolation algorithms open-ended,
-- while retaining the ability to inspect previous steps (using the
-- 'IsolationStep' pattern).
type RootIsolationAlgorithm :: forall {k :: Type}. k -> Nat -> Nat -> Constraint
class ( Typeable ty, Show ( StepDescription ty ), BoxCt n d )
    => RootIsolationAlgorithm ty n d where
  -- | The type of additional information about an algorithm step.
  --
  -- Only really useful for debugging; gets stored in 'RootIsolationTree's.
  type StepDescription ty
  -- | Configuration options expected by this root isolation method.
  type RootIsolationAlgorithmOptions ty n d = r | r -> ty n d
  -- | Run one step of the root isolation method.
  --
  -- This gets given the equations and a box, and should attempt to
  -- shrink the box in some way, returning smaller boxes.
  --
  -- Should return:
  --
  --  - a description of the step taken (see 'StepDescription'),
  --  - new boxes to process (the return value of type @['Box' n]@),
  --    which can be empty if the algorithm can prove that the input
  --    bix does not contain any solutions;
  --  - (as a writer side-effect) boxes to definitely stop processing; see 'DoneBoxes'.
  rootIsolationAlgorithm
    :: RootIsolationAlgorithmOptions ty n d
        -- ^ options for this root isolation algorithm
    -> [ ( RootIsolationStep, Box n ) ]
        -- ^ history of the current round
    -> BoxHistory n
        -- ^ previous rounds history
    -> ( 𝕀ℝ n -> D 1 n ( 𝕀ℝ d ) )
      -- ^ equations
    -> Box n
       -- ^ box
    -> Writer ( DoneBoxes n ) ( StepDescription ty, [ Box n ] )

-- | Match on an unknown root isolation algorithm step with a known algorithm.
pattern IsolationStep
  :: forall ( ty :: Type )
  .  Typeable ty
  => StepDescription ty
  -> RootIsolationStep
pattern IsolationStep stepDescr
          <- ( rootIsolationAlgorithmStep_maybe @ty -> Just stepDescr )
  -- NB: this pattern could also return @RootIsolationAlgorithm n d ty@ evidence,
  -- but it's simpler to not do so for now.

-- | Helper function used to define the 'IsolationStep' pattern.
--
-- Inspects whether an existential 'RootIsolationStep' packs a step for
-- the given algorithm.
rootIsolationAlgorithmStep_maybe
  :: forall ty
  .  Typeable ty
  => RootIsolationStep -> Maybe ( StepDescription ty )
rootIsolationAlgorithmStep_maybe ( SomeRootIsolationStep @existential descr )
  | Just HRefl <- heqT @existential @ty
  = Just descr
  | otherwise
  = Nothing
{-# INLINEABLE rootIsolationAlgorithmStep_maybe #-}

-- | History for a given box: what was the outcome of previous root isolation
-- methods?
type BoxHistory n = [ NE.NonEmpty ( RootIsolationStep, Box n ) ]

-- | A description of a step taken when isolating roots.
data RootIsolationStep where
  SomeRootIsolationStep
    :: forall step
    . ( Typeable step
      , Show ( StepDescription step )
      )
    => StepDescription step
    -> RootIsolationStep

instance Show RootIsolationStep where
  showsPrec p ( SomeRootIsolationStep stepDescr ) = showsPrec p stepDescr

--------------------------------------------------------------------------------
-- Trees recording steps taken by the algorithm, for visualisation & debugging.

-- | A tree recording the steps taken when isolating roots.
data RootIsolationTree d
  = RootIsolationLeaf String d
  | RootIsolationStep RootIsolationStep [ ( d, RootIsolationTree d ) ]

instance {-# OVERLAPPING #-} ( Representable 𝕀 ( 𝕀ℝ n ), Show ( Box n ) ) => Show ( Box n, RootIsolationTree ( Box n ) ) where
  show ( cand, tree ) = showTree $ showRootIsolationTree cand tree

showRootIsolationTree
  :: ( Representable 𝕀 ( 𝕀ℝ n ), Show ( Box n ) )
  => Box n -> RootIsolationTree ( Box n ) -> Tree String
showRootIsolationTree cand (RootIsolationLeaf why l) = Node (show cand ++ " " ++ showArea (boxArea cand) ++ " " ++ why ++ " " ++ show l) []
showRootIsolationTree cand (RootIsolationStep s ts)
  = Node (show cand ++ " abc " ++ showArea (boxArea cand) ++ " " ++ show s) $ map (\ (c,t) -> showRootIsolationTree c t) ts

{-# INLINEABLE boxArea #-}
boxArea :: Representable 𝕀 ( 𝕀ℝ n ) => Box n -> Double
boxArea box = product ( width <$> coordinates box )

showArea :: Double -> String
showArea area = "(area " ++ showFFloat (Just 6) area "" ++ ")"

--------------------------------------------------------------------------------
-- Utilities for feeding one computation into the next.

-- NB: I tried adding RULES on these functions in order to perform
-- "loop unrolling"-like optimisations, but this does not seem to
-- improve the performance.

-- | Run an effectful computation several times in sequence, piping its output
-- into the next input, once for each coordinate dimension.
forEachCoord :: forall n a m. ( KnownNat n, Monad m ) => a -> ( Fin n -> a -> m a ) -> m a
forEachCoord a0 f = go ( toList $ universe @n ) a0
  where
    go [] a = return a
    go ( i : is ) a = do
      a' <- f i a
      go is a'
{-# INLINEABLE forEachCoord #-}

-- | Apply each function in turn, piping the output of one function into
-- the next.
--
-- Once all the functions have been applied, check whether the Bool is True.
-- If it is, go around again with all the functions; otherwise, stop.
pipeFunctionsWhileTrue :: [ a -> State Bool [ a ] ] -> a -> State Bool [ a ]
pipeFunctionsWhileTrue fns = go fns
  where
    go [] x = do
      doAnotherRound <- State.get
      if doAnotherRound
      then do { State.put False ; go fns x }
      else return [ x ]
    go ( f : fs ) x = do
      xs <- f x
      concat <$> traverse ( go fs ) xs

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

newtype TimeInterval = TimeInterval Double
  deriving newtype Show

timeInterval :: NFData b => ( a -> b ) -> a -> ( b, TimeInterval )
timeInterval f a = unsafePerformIO $ do
  bef <- getMonotonicTime
  let !b = f a
  b `deepseq` return ()
  aft <- getMonotonicTime
  return $ ( b, TimeInterval ( aft - bef ) )