brush-strokes-0.1.0.0: src/lib/Math/Root/Isolation/Newton/LP.hs

{-# LANGUAGE CApiFFI #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE ParallelListComp #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | A linear programming approach for solving systems of
-- interval linear equations.
module Math.Root.Isolation.Newton.LP
  ( solveIntervalLinearEquations )
  where

-- base
import Control.Arrow
  ( (&&&) )
import Data.Coerce
  ( coerce )
import Data.Foldable
  ( toList )
import Foreign.C.Types
  ( CDouble(..), CInt(..), CUInt(..) )
import Foreign.Marshal
  ( allocaArray, peekArray, pokeArray, with, withArray )
import Foreign.Ptr
  ( Ptr, castPtr )
import Foreign.Storable
  ( Storable(..) )
import GHC.TypeNats
  ( KnownNat, natVal' )
import GHC.Exts
  ( proxy# )
import System.IO.Unsafe
  ( unsafePerformIO )

-- brush-strokes
import Math.Interval
import Math.Linear

--------------------------------------------------------------------------------
-- Linear programming approach to solving systems of linear equations
-- (in two variables).

-- | Solve the system of linear equations \( A X = B \)
-- using linear programming.
solveIntervalLinearEquations
  :: ( KnownNat d, Representable 𝕀 ( 𝕀ℝ d ) )
  => Vec 2 ( 𝕀ℝ d )    -- ^ columns of \( A \)
  -> 𝕀ℝ d              -- ^ \( B \)
  -> T ( 𝕀ℝ 2 )        -- ^ initial box \( X \)
  -> [ ( T ( 𝕀ℝ 2 ), Bool ) ]
solveIntervalLinearEquations a b x =
  let !sols = unsafePerformIO $ intervalSystem2D_LP a b x
  in if
      | any hasNaNs sols
      -> [ ( x, False ) ]
      | length sols <= 1
      -> map ( id &&& isSubBox x ) sols
      | otherwise
      -> map ( , False ) sols

-- Assuming the second box is a subset of the second box, returns whether it
-- is in fact a strict subset.
isSubBox :: T ( 𝕀ℝ 2 ) -> T ( 𝕀ℝ 2 ) -> Bool
isSubBox
  ( T ( 𝕀ℝ2 ( 𝕀 x1_min x1_max ) ( 𝕀 y1_min y1_max ) ) )
  ( T ( 𝕀ℝ2 ( 𝕀 x2_min x2_max ) ( 𝕀 y2_min y2_max ) ) )
    =  ( ( x2_min > x1_min && x2_max < x1_max ) || x2_min == x2_max )
    && ( ( y2_min > y1_min && y2_max < y1_max ) || y2_min == y2_max )

hasNaNs :: T (𝕀ℝ 2) -> Bool
hasNaNs ( T ( 𝕀ℝ2 ( 𝕀 x_min x_max ) ( 𝕀 y_min y_max ) ) ) =
  any ( \ x -> isNaN x || isInfinite x ) [ x_min, y_min, x_max, y_max ]

intervalSystem2D_LP
  :: forall d
  .  ( KnownNat d, Representable 𝕀 ( 𝕀ℝ d ) )
  => Vec 2 ( 𝕀ℝ d )
  -> 𝕀ℝ d
  -> T ( 𝕀ℝ 2 )
  -> IO [ T ( 𝕀ℝ 2 ) ]
intervalSystem2D_LP a b x =
  allocaArray 4 \ ptrSolutions ->
    with ( CBox $ unT $ x ) \ ptrBox ->
      withArray ( mkEquationArray a b ) \ ptrEqs -> do
        CInt nbSols <-
          interval_system_2d ptrSolutions ptrBox ptrEqs ( fromIntegral d )
        if nbSols < 0
        then
          error $ unlines
            [ "interval_system_2d returned with exit code " ++ show nbSols
            , "This probably means it was given invalid input." ]
        else
          coerce <$> peekArray ( fromIntegral nbSols ) ptrSolutions
  where
    d = natVal' @d proxy#

mkEquationArray
  :: ( KnownNat d, Representable 𝕀 ( 𝕀ℝ d ) )
  => Vec 2 ( 𝕀ℝ d ) -> 𝕀ℝ d -> [ CEqn ]
mkEquationArray ( Vec [ a_x, a_y ] ) b =
  [ CEqn a_x_i a_y_i b_i
  | a_x_i <- toList $ coordinates a_x
  | a_y_i <- toList $ coordinates a_y
  | b_i <- toList $ coordinates b
  ]
mkEquationArray _ _ = error "impossible"

foreign import ccall "interval_system_2d"
  interval_system_2d :: Ptr CBox -> Ptr CBox -> Ptr CEqn -> CUInt -> IO CInt

data CEqn = CEqn !𝕀 !𝕀 !𝕀
instance Storable CEqn where
  sizeOf _ = 6 * sizeOf @Double undefined
  alignment _ = 4 * alignment @Double undefined
  peek ptr = do
    [ CDouble a_min, CDouble a_max, CDouble b_min, CDouble b_max, CDouble c_min, CDouble c_max ]
      <- peekArray 6 ( castPtr ptr :: Ptr CDouble )
    return $
      CEqn ( 𝕀 a_min a_max ) ( 𝕀 b_min b_max ) ( 𝕀 c_min c_max )
  poke ptr ( CEqn ( 𝕀 a_min a_max ) ( 𝕀 b_min b_max ) ( 𝕀 c_min c_max ) )
    = pokeArray ( castPtr ptr ) [ CDouble a_min, CDouble a_max, CDouble b_min, CDouble b_max, CDouble c_min, CDouble c_max ]


newtype CBox = CBox ( 𝕀ℝ 2 )
instance Storable CBox where
  sizeOf _ = 4 * sizeOf @Double undefined
  alignment _ = 4 * alignment @Double undefined
  peek ptr = do
    [ CDouble x_min, CDouble x_max, CDouble y_min, CDouble y_max ] <- peekArray 4 ( castPtr ptr :: Ptr CDouble )
    return $
      CBox ( 𝕀ℝ2 ( 𝕀 x_min x_max ) ( 𝕀 y_min y_max ) )
  poke ptr ( CBox ( 𝕀ℝ2 ( 𝕀 x_min x_max ) ( 𝕀 y_min y_max ) ) ) =
    pokeArray ( castPtr ptr ) [ CDouble x_min, CDouble x_max, CDouble y_min, CDouble y_max ]