GaussSeidel.hs

{-# LANGUAGE ScopedTypeVariables #-}

-- | The Gauss–Seidel method for solving systems of
-- interval linear equations
module Math.Root.Isolation.Newton.GaussSeidel
  ( -- *  Gauss–Seidel step
    gaussSeidelUpdate
    -- ** Options for the Gauss–Seidel method
  , GaussSeidelOptions(..)
  , Preconditioner(..), GaussSeidelUpdateMethod(..)
  , defaultGaussSeidelOptions
  , precondition
  )
  where

-- base
import Prelude hiding ( unzip )
import Data.Coerce
  ( coerce )
import Data.Kind
  ( Type )
import Data.Foldable
  ( toList )
import Data.Proxy
  ( Proxy(..) )
import Data.Type.Ord
  ( OrderingI(..) )
import GHC.TypeNats
  ( Nat, KnownNat, type (<=), cmpNat )

-- eigen
import qualified Eigen.Matrix as Eigen
  ( inverse )

-- brush-strokes
import Math.Interval
import Math.Linear
import Math.Linear.Solve
  ( fromEigen, toEigen )
import Math.Root.Isolation.Core
import Math.Root.Isolation.Utils

--------------------------------------------------------------------------------
-- Gauss–Seidel

-- | Options for the Gauss–Seidel method.
type GaussSeidelOptions :: Nat -> Nat -> Type
data GaussSeidelOptions n d =
  GaussSeidelOptions
    { -- | Which preconditioner to use?
      gsPreconditioner :: !Preconditioner
      -- | Function that projects over the equations we will consider
      -- (the identity for a well-determined problem, or a projection for
      -- an overdetermined system).
    , gsPickEqs :: !( 𝕀ℝ d -> 𝕀ℝ n )
      -- | Whether to use a partial or a complete Gauss–Seidel update
    , gsUpdate :: !GaussSeidelUpdateMethod
    }
instance Show ( GaussSeidelOptions n d ) where
  show ( GaussSeidelOptions { gsPreconditioner, gsUpdate } ) =
    unlines
      [ "GaussSeidelOptions"
      , "  preconditioner: " ++ show gsPreconditioner
      , "  update: " ++ show gsUpdate
      ]

-- | Whether to use a partial or a complete Gauss–Seidel update.
data GaussSeidelUpdateMethod
  = GS_Partial
  | GS_Complete
  deriving stock Show

-- | Default options for the Gauss–Seidel step.
defaultGaussSeidelOptions
  :: forall n d
  .  ( KnownNat n, KnownNat d
     , 1 <= n, 1 <= d, n <= d
     , Representable Double ( ℝ n ), Representable 𝕀 ( 𝕀ℝ n )
     , Representable Double ( ℝ d ), Representable 𝕀 ( 𝕀ℝ d )
     )
  => BoxHistory n
  -> GaussSeidelOptions n d
defaultGaussSeidelOptions history =
  GaussSeidelOptions
    { gsPreconditioner = InverseMidpoint
    , gsPickEqs =
        case cmpNat @n @d Proxy Proxy of
          EQI -> id
          LTI ->
            -- If there are more equations (d) than variables (n),
            -- pick a size n subset of the variables
            -- (go through all combinations cyclically).
            let choices :: [ Vec n ( Fin d ) ]
                choices = choose @d @n
                choice :: Vec n ( Fin d )
                choice = choices !! ( length history `mod` length choices )
            in \ u -> tabulate \ i -> index u ( choice ! i )
    , gsUpdate = GS_Complete
    }
{-# INLINEABLE defaultGaussSeidelOptions #-}

-- | Preconditioner to use with the interval Gauss–Seidel method.
data Preconditioner
  = NoPreconditioning
  | InverseMidpoint
  deriving stock Show

-- | A partial or complete Gauss–Seidel step for the equation \( A X = B \),
-- refining the initial guess box for \( X \) into up to \( 2^n \) (disjoint) new boxes.
gaussSeidelUpdate
  :: forall n
  .  ( Representable Double ( ℝ n )
     , Representable 𝕀 ( 𝕀ℝ n )
     , Eq ( ℝ n )
     , Show ( 𝕀ℝ n )
     )
  => GaussSeidelUpdateMethod -- ^ which step method to use
  -> Vec n ( 𝕀ℝ n )    -- ^ columns of \( A \)
  -> 𝕀ℝ n              -- ^ \( B \)
  -> T ( 𝕀ℝ n )        -- ^ initial box \( X \)
  -> [ ( T ( 𝕀ℝ n ), Bool ) ]
gaussSeidelUpdate upd as b x =
  case upd of
    GS_Partial  -> gaussSeidelStep          as b x
    GS_Complete -> gaussSeidelStep_Complete as b x
{-# INLINEABLE gaussSeidelUpdate #-}

-- | Take one interval Gauss–Seidel step for the equation \( A X = B \),
-- refining the initial guess box for \( X \) into up to \( 2^n \) (disjoint) new boxes.
--
-- The boolean indicates whether the Gauss–Seidel step was a contraction.
gaussSeidelStep
  :: forall n
  .  ( Representable Double ( ℝ n )
     , Representable 𝕀 ( 𝕀ℝ n )
     , Eq ( ℝ n )
     )
  => Vec n ( 𝕀ℝ n ) -- ^ columns of \( A \)
  -> 𝕀ℝ n           -- ^ \( B \)
  -> T ( 𝕀ℝ n )     -- ^ initial box \( X \)
  -> [ ( T ( 𝕀ℝ n ), Bool ) ]
gaussSeidelStep as b ( T x0 ) = coerce $
  forEachCoord @n ( x0, True ) $ \ i ( x, contraction ) -> do
    -- For each i, we have an equation: sum_j a_ij * x_j = b_i
    --
    -- Re-arrange this with x_i on the left to get an iteration:
    --   x_i' = ( b_i - sum { j /= i } a_ij * x_j ) / a_ii
    --
    -- We perform each iteration in turn (for i = 1, ..., n),
    -- **using the latest updated value of each x_j** in each iteration.
    let s = b `index` i - sum [ ( as ! j ) `index` i * x `index` j
                              | j <- toList ( universe @n ), j /= i ]
        x_i = x `index` i
        a_ii = ( as ! i ) `index` i
    -- Take a shortcut before performing the division if possible.
    if | not $ 0 ∈ ( s - a_ii * x_i )
       -- No solutions: don't bother performing a division.
       -> [ ]
       | 0 ∈ s && 0 ∈ a_ii
       -- The division would produce [-oo,+oo]: don't do anything.
       -> [ ( x, False ) ]
       -- Otherwise, perform the division.
       | otherwise
       -> do
           x_i'0 <- s ⊘ a_ii
           ( x_i', sub_i ) <- x_i'0 ∩ x_i
           return $ ( set i x_i' x, sub_i && contraction )
{-# INLINEABLE gaussSeidelStep #-}

-- | The complete interval-union Gauss–Seidel step.
--
-- Algorithm 2 from:
--   "Using interval unions to solve linear systems of equations with uncertainties"
--    (Montanher, Domes, Schichl, Neumaier) (2017)
gaussSeidelStep_Complete
  :: forall n
  .  ( Representable Double ( ℝ n )
     , Representable 𝕀 ( 𝕀ℝ n )
     , Show ( 𝕀ℝ n )
     )
  => Vec n ( 𝕀ℝ n ) -- ^ columns of \( A \)
  -> 𝕀ℝ n           -- ^ \( B \)
  -> T ( 𝕀ℝ n )     -- ^ initial box \( X \)
  -> [ ( T ( 𝕀ℝ n ), Bool ) ]
gaussSeidelStep_Complete as b ( T x0 ) = coerce $ do
  ( x', subs ) <-
    forEachCoord @n ( x0, pure False ) $ \ i ( x, contractions ) -> do
      let s = b `index` i - sum [ ( as ! k ) `index` i * x `index` k
                                | k <- toList ( universe @n ) ]
      ( x', subs ) <- fromComponents \ j -> do
        let x_j = x `index` j
            a_ij = ( as ! j ) `index` i
            s_j = s ⊖ ( -a_ij * x_j )
              -- NB: the paper is missing the minus sign here.
              -- In the definition of s we subtract all terms including
              -- the a_ij * x_j term, so we need to add it back.

        -- Shortcut division if possible (see gaussSeidelStep for commentary).
        -- NB: there is a typo in the paper here, it should NOT be
        --
        --  0 ∈ ( s_j - a_ii * x_j )
        --
        -- because we are dividing by a_ij, not a_ii.
        if | not $ 0 ∈ ( s_j - a_ij * x_j )
           -> [ ]
           | 0 ∈ s_j && 0 ∈ a_ij
           -> [ ( x_j, False ) ]
           | otherwise
           -> do
               x_j'0 <- s_j ⊘ a_ij
               ( x_j', sub_j ) <- x_j'0 ∩ x_j
               return $ ( x_j', sub_j )
      return ( x', (||) <$> contractions <*> subs )
  return ( x', and subs )
{-# INLINEABLE gaussSeidelStep_Complete #-}

-- | Pre-condition the system \( AX = B \).
precondition
  :: forall n
  .  ( KnownNat n
     , Representable Double ( ℝ n )
     , Representable 𝕀 ( 𝕀ℝ n )
     )
  => Preconditioner -- ^ pre-conditioning method to use
  -> Vec n ( 𝕀ℝ n ) -- ^ columns of \( A \)
  -> 𝕀ℝ n           -- ^ \( B \)
  -> ( Vec n ( 𝕀ℝ n ), 𝕀ℝ n )
precondition meth as b =
  case meth of
    NoPreconditioning
      -> ( as, b )
    InverseMidpoint
      | let mat = toEigen $ fmap boxMidpoint as
      , let precond = Eigen.inverse mat
            doPrecond = matMulVec ( fromEigen precond )
      -- TODO: avoid this when condition number is small?
      -> ( fmap doPrecond as, doPrecond b )
{-# INLINEABLE precondition #-}