Solve.hs

{-# LANGUAGE ScopedTypeVariables #-}

{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}

module Math.Linear.Solve
  ( linearSolve, fromEigen, toEigen, withEigenSem )
  where

-- base
import Control.Concurrent.QSem
import Control.Exception
  ( evaluate )
import Data.Maybe
  ( fromJust )
import GHC.TypeNats
  ( KnownNat )
import System.IO.Unsafe
  ( unsafePerformIO )

-- deepseq
import Control.DeepSeq

-- eigen
import qualified Eigen.Matrix as Eigen
  ( Matrix, toList, fromList, generate, unsafeCoeff )
import qualified Eigen.Solver.LA as Eigen
  ( Decomposition(..), solve )

-- brush-strokes
import Math.Linear

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

linearSolve :: Mat22 -> T ( ℝ 2 ) -> T ( ℝ 2 )
linearSolve ( Mat22 a b c d ) ( V2 p q ) = V2 u v
  where
    [[u],[v]] = withEigenSem
              $ Eigen.toList
              $ Eigen.solve Eigen.JacobiSVD
                  ( fromJust ( Eigen.fromList [[a,b],[c,d]] )
                    :: Eigen.Matrix 2 2 Double )
                  ( fromJust ( Eigen.fromList [[p],[q]] )
                    :: Eigen.Matrix 2 1 Double )

toEigen :: forall n d
        . ( KnownNat d, Representable Double ( ℝ n ) )
        => Vec d ( ℝ n ) -> Eigen.Matrix n d Double
toEigen cols =
  Eigen.generate $ \ r c ->
    ( cols ! Fin ( fromIntegral c + 1 ) ) `index` ( Fin ( fromIntegral r + 1 ) )
{-# INLINEABLE toEigen #-}

fromEigen :: forall n d
          .  ( KnownNat d, Representable Double ( ℝ n ) )
          => Eigen.Matrix n d Double -> Vec d ( ℝ n )
fromEigen mat =
  fmap
    ( \ ( Fin c ) ->
      tabulate $ \ ( Fin r ) ->
        Eigen.unsafeCoeff ( fromIntegral r - 1 ) ( fromIntegral c - 1 ) mat
    )
    ( universe @d )
{-# INLINEABLE fromEigen #-}

-- | Semaphore to ensure we don't ever call Eigen concurrently, which
-- seems to cause random crashes.
eigenSem :: QSem
eigenSem = unsafePerformIO $ newQSem 1
{-# NOINLINE eigenSem #-}

-- | Take a lock on the 'eigenSem' semaphore, to ensure we don't call Eigen
-- concurrently (which seems to cause random crashes).
withEigenSem :: NFData a => a -> a
withEigenSem v = unsafePerformIO $ do
  waitQSem eigenSem
  res <- evaluate v
  signalQSem eigenSem
  return res