toysolver-0.5.0: src/ToySolver/SAT/MessagePassing/SurveyPropagation.hs
{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE ScopedTypeVariables, BangPatterns, TypeFamilies #-}
-----------------------------------------------------------------------------
-- |
-- Module : ToySolver.SAT.MessagePassing.SurveyPropagation
-- Copyright : (c) Masahiro Sakai 2016
-- License : BSD-style
--
-- Maintainer : masahiro.sakai@gmail.com
-- Stability : provisional
-- Portability : non-portable (ScopedTypeVariables, BangPatterns, TypeFamilies)
--
-- References:
--
-- * Alfredo Braunstein, Marc Mézard and Riccardo Zecchina.
-- Survey Propagation: An Algorithm for Satisfiability,
-- <http://arxiv.org/abs/cs/0212002>
--
-- * Corrie Scalisi. Visualizing Survey Propagation in 3-SAT Factor Graphs,
-- <http://classes.soe.ucsc.edu/cmps290c/Winter06/proj/corriereport.pdf>.
--
-----------------------------------------------------------------------------
module ToySolver.SAT.MessagePassing.SurveyPropagation
(
-- * The Solver type
Solver
, newSolver
, deleteSolver
-- * Problem information
, getNVars
, getNConstraints
-- * Parameters
, getTolerance
, setTolerance
, getIterationLimit
, setIterationLimit
, getNThreads
, setNThreads
-- * Computing marginal distributions
, initializeRandom
, initializeRandomDirichlet
, propagate
, getVarProb
-- * Solving
, fixLit
, unfixLit
-- * Debugging
, printInfo
) where
import Control.Applicative
import Control.Concurrent
import Control.Concurrent.STM
import Control.Exception
import Control.Loop
import Control.Monad
import qualified Data.Array.IArray as A
import qualified Data.IntMap as IntMap
import qualified Data.IntSet as IntSet
import Data.IORef
import Data.Maybe (fromJust)
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM
import Data.Vector.Generic ((!))
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable as VGM
import qualified Numeric.Log as L
import qualified System.Random.MWC as Rand
import qualified System.Random.MWC.Distributions as Rand
import qualified ToySolver.SAT.Types as SAT
infixr 8 ^*
(^*) :: Num a => L.Log a -> a -> L.Log a
L.Exp a ^* b = L.Exp (a*b)
comp :: (RealFloat a, L.Precise a) => L.Log a -> L.Log a
comp (L.Exp a) = L.Exp $ L.log1p $ max (-1) $ negate (exp a)
type ClauseIndex = Int
type EdgeIndex = Int
data Solver
= Solver
{ svVarEdges :: !(V.Vector (VU.Vector EdgeIndex))
, svVarProbT :: !(VUM.IOVector (L.Log Double))
, svVarProbF :: !(VUM.IOVector (L.Log Double))
, svVarFixed :: !(VM.IOVector (Maybe Bool))
, svClauseEdges :: !(V.Vector (VU.Vector EdgeIndex))
, svClauseWeight :: !(VU.Vector Double)
, svEdgeLit :: !(VU.Vector SAT.Lit) -- i
, svEdgeClause :: !(VU.Vector ClauseIndex) -- a
, svEdgeSurvey :: !(VUM.IOVector (L.Log Double)) -- η_{a → i}
, svEdgeProbU :: !(VUM.IOVector (L.Log Double)) -- Π^u_{i → a} / (Π^u_{i → a} + Π^s_{i → a} + Π^0_{i → a})
, svTolRef :: !(IORef Double)
, svIterLimRef :: !(IORef (Maybe Int))
, svNThreadsRef :: !(IORef Int)
}
newSolver :: Int -> [(Double, SAT.Clause)] -> IO Solver
newSolver nv clauses = do
let num_clauses = length clauses
num_edges = sum [length c | (_,c) <- clauses]
varEdgesRef <- newIORef IntMap.empty
clauseEdgesM <- VGM.new num_clauses
(edgeLitM :: VUM.IOVector SAT.Lit) <- VGM.new num_edges
(edgeClauseM :: VUM.IOVector ClauseIndex) <- VGM.new num_edges
ref <- newIORef 0
forM_ (zip [0..] clauses) $ \(i,(_,c)) -> do
es <- forM c $ \lit -> do
e <- readIORef ref
modifyIORef' ref (+1)
modifyIORef' varEdgesRef (IntMap.insertWith IntSet.union (abs lit) (IntSet.singleton e))
VGM.unsafeWrite edgeLitM e lit
VGM.unsafeWrite edgeClauseM e i
return e
VGM.unsafeWrite clauseEdgesM i (VG.fromList es)
varEdges <- readIORef varEdgesRef
clauseEdges <- VG.unsafeFreeze clauseEdgesM
edgeLit <- VG.unsafeFreeze edgeLitM
edgeClause <- VG.unsafeFreeze edgeClauseM
-- Initialize all surveys with non-zero values.
-- If we initialize to zero, following trivial solution exists:
--
-- η_{a→i} = 0 for all i and a.
--
-- Π^0_{i→a} = 1, Π^u_{i→a} = Π^s_{i→a} = 0 for all i and a,
--
-- \^{Π}^{0}_i = 1, \^{Π}^{+}_i = \^{Π}^{-}_i = 0
--
edgeSurvey <- VGM.replicate num_edges 0.5
edgeProbU <- VGM.new num_edges
varFixed <- VGM.replicate nv Nothing
varProbT <- VGM.new nv
varProbF <- VGM.new nv
tolRef <- newIORef 0.01
maxIterRef <- newIORef (Just 1000)
nthreadsRef <- newIORef 1
let solver = Solver
{ svVarEdges = VG.generate nv $ \i ->
case IntMap.lookup (i+1) varEdges of
Nothing -> VG.empty
Just es -> VG.fromListN (IntSet.size es) (IntSet.toList es)
, svVarProbT = varProbT
, svVarProbF = varProbF
, svVarFixed = varFixed
, svClauseEdges = clauseEdges
, svClauseWeight = VG.fromListN (VG.length clauseEdges) (map fst clauses)
, svEdgeLit = edgeLit
, svEdgeClause = edgeClause
, svEdgeSurvey = edgeSurvey
, svEdgeProbU = edgeProbU
, svTolRef = tolRef
, svIterLimRef = maxIterRef
, svNThreadsRef = nthreadsRef
}
return solver
deleteSolver :: Solver -> IO ()
deleteSolver solver = return ()
initializeRandom :: Solver -> Rand.GenIO -> IO ()
initializeRandom solver gen = do
VG.forM_ (svClauseEdges solver) $ \es -> do
case VG.length es of
0 -> return ()
1 -> VGM.unsafeWrite (svEdgeSurvey solver) (es ! 0) 1
n -> do
let p :: Double
p = 1 / fromIntegral n
VG.forM_ es $ \e -> do
d <- Rand.uniformR (p*0.5, p*1.5) gen
VGM.unsafeWrite (svEdgeSurvey solver) e (realToFrac d)
initializeRandomDirichlet :: Solver -> Rand.GenIO -> IO ()
initializeRandomDirichlet solver gen = do
VG.forM_ (svClauseEdges solver) $ \es -> do
case VG.length es of
0 -> return ()
1 -> VGM.unsafeWrite (svEdgeSurvey solver) (es ! 0) 1
len -> do
(ps :: V.Vector Double) <- Rand.dirichlet (VG.replicate len 1) gen
numLoop 0 (len-1) $ \i -> do
VGM.unsafeWrite (svEdgeSurvey solver) (es ! i) (realToFrac (ps ! i))
-- | number of variables of the problem.
getNVars :: Solver -> IO Int
getNVars solver = return $ VG.length (svVarEdges solver)
-- | number of constraints of the problem.
getNConstraints :: Solver -> IO Int
getNConstraints solver = return $ VG.length (svClauseEdges solver)
-- | number of edges of the factor graph
getNEdges :: Solver -> IO Int
getNEdges solver = return $ VG.length (svEdgeLit solver)
getTolerance :: Solver -> IO Double
getTolerance solver = readIORef (svTolRef solver)
setTolerance :: Solver -> Double -> IO ()
setTolerance solver !tol = writeIORef (svTolRef solver) tol
getIterationLimit :: Solver -> IO (Maybe Int)
getIterationLimit solver = readIORef (svIterLimRef solver)
setIterationLimit :: Solver -> Maybe Int -> IO ()
setIterationLimit solver val = writeIORef (svIterLimRef solver) val
getNThreads :: Solver -> IO Int
getNThreads solver = readIORef (svNThreadsRef solver)
setNThreads :: Solver -> Int -> IO ()
setNThreads solver val = writeIORef (svNThreadsRef solver) val
propagate :: Solver -> IO Bool
propagate solver = do
nthreads <- getNThreads solver
if nthreads > 1 then
propagateMT solver nthreads
else
propagateST solver
propagateST :: Solver -> IO Bool
propagateST solver = do
tol <- getTolerance solver
lim <- getIterationLimit solver
nv <- getNVars solver
nc <- getNConstraints solver
let max_v_len = VG.maximum $ VG.map VG.length $ svVarEdges solver
max_c_len = VG.maximum $ VG.map VG.length $ svClauseEdges solver
tmp <- VGM.new (max (max_v_len * 2) max_c_len)
let loop !i
| Just l <- lim, i >= l = return False
| otherwise = do
numLoop 1 nv $ \v -> updateEdgeProb solver v tmp
let f maxDelta c = max maxDelta <$> updateEdgeSurvey solver c tmp
delta <- foldM f 0 [0 .. nc-1]
if delta <= tol then do
numLoop 1 nv $ \v -> computeVarProb solver v
return True
else
loop (i+1)
loop 0
data WorkerCommand
= WCUpdateEdgeProb
| WCUpdateSurvey
| WCComputeVarProb
| WCTerminate
propagateMT :: Solver -> Int -> IO Bool
propagateMT solver nthreads = do
tol <- getTolerance solver
lim <- getIterationLimit solver
nv <- getNVars solver
nc <- getNConstraints solver
mask $ \restore -> do
ex <- newEmptyTMVarIO
let wait :: STM a -> IO a
wait m = join $ atomically $ liftM return m `orElse` liftM throwIO (takeTMVar ex)
workers <- do
let mV = (nv + nthreads - 1) `div` nthreads
mC = (nc + nthreads - 1) `div` nthreads
forM [0..nthreads-1] $ \i -> do
let lbV = mV * i + 1 -- inclusive
ubV = min (lbV + mV) (nv + 1) -- exclusive
lbC = mC * i -- exclusive
ubC = min (lbC + mC) nc -- exclusive
let max_v_len = VG.maximum $ VG.map VG.length $ VG.slice (lbV - 1) (ubV - lbV) (svVarEdges solver)
max_c_len = VG.maximum $ VG.map VG.length $ VG.slice lbC (ubC - lbC) (svClauseEdges solver)
tmp <- VGM.new (max (max_v_len*2) max_c_len)
reqVar <- newEmptyMVar
respVar <- newEmptyTMVarIO
respVar2 <- newEmptyTMVarIO
th <- forkIO $ do
let loop = do
cmd <- takeMVar reqVar
case cmd of
WCTerminate -> return ()
WCUpdateEdgeProb -> do
numLoop lbV (ubV-1) $ \v -> updateEdgeProb solver v tmp
atomically $ putTMVar respVar ()
loop
WCUpdateSurvey -> do
let f maxDelta c = max maxDelta <$> updateEdgeSurvey solver c tmp
delta <- foldM f 0 [lbC .. ubC-1]
atomically $ putTMVar respVar2 delta
loop
WCComputeVarProb -> do
numLoop lbV (ubV-1) $ \v -> computeVarProb solver v
atomically $ putTMVar respVar ()
loop
restore loop `catch` \(e :: SomeException) -> atomically (tryPutTMVar ex e >> return ())
return (th, reqVar, respVar, respVar2)
let loop !i
| Just l <- lim, i >= l = return False
| otherwise = do
mapM_ (\(_,reqVar,_,_) -> putMVar reqVar WCUpdateEdgeProb) workers
mapM_ (\(_,_,respVar,_) -> wait (takeTMVar respVar)) workers
mapM_ (\(_,reqVar,_,_) -> putMVar reqVar WCUpdateSurvey) workers
delta <- foldM (\delta (_,_,_,respVar2) -> max delta <$> wait (takeTMVar respVar2)) 0 workers
if delta <= tol then do
mapM_ (\(_,reqVar,_,_) -> putMVar reqVar WCComputeVarProb) workers
mapM_ (\(_,_,respVar,_) -> wait (takeTMVar respVar)) workers
mapM_ (\(_,reqVar,_,_) -> putMVar reqVar WCTerminate) workers
return True
else
loop (i+1)
ret <- try $ restore $ loop 0
case ret of
Right b -> return b
Left (e :: SomeException) -> do
mapM_ (\(th,_,_,_) -> killThread th) workers
throwIO e
-- tmp1 must have at least @VG.length (svVarEdges solver ! (v - 1)) * 2@ elements
updateEdgeProb :: Solver -> SAT.Var -> VUM.IOVector (L.Log Double) -> IO ()
updateEdgeProb solver v tmp = do
let i = v - 1
edges = svVarEdges solver ! i
m <- VGM.unsafeRead (svVarFixed solver) i
case m of
Just val -> do
VG.forM_ edges $ \e -> do
let lit = svEdgeLit solver ! e
flag = (lit > 0) == val
VGM.unsafeWrite (svEdgeProbU solver) e (if flag then 0 else 1)
Nothing -> do
let f !k !val1_pre !val2_pre
| k >= VG.length edges = return ()
| otherwise = do
let e = edges ! k
a = svEdgeClause solver ! e
VGM.unsafeWrite tmp (k*2) val1_pre
VGM.unsafeWrite tmp (k*2+1) val2_pre
eta_ai <- VGM.unsafeRead (svEdgeSurvey solver) e -- η_{a→i}
let w = svClauseWeight solver ! a
lit2 = svEdgeLit solver ! e
val1_pre' = if lit2 > 0 then val1_pre * comp eta_ai ^* w else val1_pre
val2_pre' = if lit2 > 0 then val2_pre else val2_pre * comp eta_ai ^* w
f (k+1) val1_pre' val2_pre'
f 0 1 1
-- tmp ! (k*2) == Π_{a∈edges[0..k-1], a∈V^{+}(i)} (1 - eta_ai)^{w_i}
-- tmp ! (k*2+1) == Π_{a∈edges[0..k-1], a∈V^{-}(i)} (1 - eta_ai)^{w_i}
let g !k !val1_post !val2_post
| k < 0 = return ()
| otherwise = do
let e = edges ! k
a = svEdgeClause solver ! e
lit2 = svEdgeLit solver ! e
val1_pre <- VGM.unsafeRead tmp (k*2)
val2_pre <- VGM.unsafeRead tmp (k*2+1)
let val1 = val1_pre * val1_post -- val1 == Π_{b∈edges, b∈V^{+}(i), a≠b} (1 - eta_bi)^{w_i}
val2 = val2_pre * val2_post -- val2 == Π_{b∈edges, b∈V^{-}(i), a≠b} (1 - eta_bi)^{w_i}
eta_ai <- VGM.unsafeRead (svEdgeSurvey solver) e -- η_{a→i}
let w = svClauseWeight solver ! a
val1_post' = if lit2 > 0 then val1_post * comp eta_ai ^* w else val1_post
val2_post' = if lit2 > 0 then val2_post else val2_post * comp eta_ai ^* w
let pi_0 = val1 * val2 -- Π^0_{i→a}
pi_u = if lit2 > 0 then comp val2 * val1 else comp val1 * val2 -- Π^u_{i→a}
pi_s = if lit2 > 0 then comp val1 * val2 else comp val2 * val1 -- Π^s_{i→a}
VGM.unsafeWrite (svEdgeProbU solver) e (pi_u / L.sum [pi_0, pi_u, pi_s])
g (k-1) val1_post' val2_post'
g (VG.length edges - 1) 1 1
-- tmp must have at least @VG.length (svClauseEdges solver ! a)@ elements
updateEdgeSurvey :: Solver -> ClauseIndex -> VUM.IOVector (L.Log Double) -> IO Double
updateEdgeSurvey solver a tmp = do
let edges = svClauseEdges solver ! a
let f !k !p_pre
| k >= VG.length edges = return ()
| otherwise = do
let e = edges ! k
VGM.unsafeWrite tmp k p_pre
p <- VGM.unsafeRead (svEdgeProbU solver) e
-- p is the probability of lit being false, if the edge does not exist.
f (k+1) (p_pre * p)
let g !k !p_post !maxDelta
| k < 0 = return maxDelta
| otherwise = do
let e = edges ! k
-- p_post == Π_{e∈edges[k+1..]} p_e
p_pre <- VGM.unsafeRead tmp k -- Π_{e∈edges[0..k-1]} p_e
p <- VGM.unsafeRead (svEdgeProbU solver) e
eta_ai <- VGM.unsafeRead (svEdgeSurvey solver) e
let eta_ai' = p_pre * p_post -- Π_{e∈edges[0,..,k-1,k+1,..]} p_e
VGM.unsafeWrite (svEdgeSurvey solver) e eta_ai'
let delta = abs (realToFrac eta_ai' - realToFrac eta_ai)
g (k-1) (p_post * p) (max delta maxDelta)
f 0 1
-- tmp ! k == Π_{e∈edges[0..k-1]} p_e
g (VG.length edges - 1) 1 0
computeVarProb :: Solver -> SAT.Var -> IO ()
computeVarProb solver v = do
let i = v - 1
f (val1,val2) e = do
let lit = svEdgeLit solver ! e
a = svEdgeClause solver ! e
w = svClauseWeight solver ! a
eta_ai <- VGM.unsafeRead (svEdgeSurvey solver) e
let val1' = if lit > 0 then val1 * comp eta_ai ^* w else val1
val2' = if lit < 0 then val2 * comp eta_ai ^* w else val2
return (val1',val2')
(val1,val2) <- VG.foldM' f (1,1) (svVarEdges solver ! i)
let p0 = val1 * val2 -- \^{Π}^{0}_i
pp = comp val1 * val2 -- \^{Π}^{+}_i
pn = comp val2 * val1 -- \^{Π}^{-}_i
let wp = pp / (pp + pn + p0)
wn = pn / (pp + pn + p0)
VGM.unsafeWrite (svVarProbT solver) i wp -- W^{(+)}_i
VGM.unsafeWrite (svVarProbF solver) i wn -- W^{(-)}_i
-- | Get the marginal probability of the variable to be @True@, @False@ and unspecified respectively.
getVarProb :: Solver -> SAT.Var -> IO (Double, Double, Double)
getVarProb solver v = do
pt <- realToFrac <$> VGM.unsafeRead (svVarProbT solver) (v - 1)
pf <- realToFrac <$> VGM.unsafeRead (svVarProbF solver) (v - 1)
return (pt, pf, 1 - (pt + pf))
fixLit :: Solver -> SAT.Lit -> IO ()
fixLit solver lit = do
VGM.unsafeWrite (svVarFixed solver) (abs lit - 1) (if lit > 0 then Just True else Just False)
unfixLit :: Solver -> SAT.Lit -> IO ()
unfixLit solver lit = do
VGM.unsafeWrite (svVarFixed solver) (abs lit - 1) Nothing
printInfo :: Solver -> IO ()
printInfo solver = do
(surveys :: VU.Vector (L.Log Double)) <- VG.freeze (svEdgeSurvey solver)
(u :: VU.Vector (L.Log Double)) <- VG.freeze (svEdgeProbU solver)
let xs = [(clause, lit, eta, u ! e) | (e, eta) <- zip [0..] (VG.toList surveys), let lit = svEdgeLit solver ! e, let clause = svEdgeClause solver ! e]
putStrLn $ "edges: " ++ show xs
(pt :: VU.Vector (L.Log Double)) <- VG.freeze (svVarProbT solver)
(pf :: VU.Vector (L.Log Double)) <- VG.freeze (svVarProbF solver)
nv <- getNVars solver
let xs2 = [(v, realToFrac (pt ! i) :: Double, realToFrac (pf ! i) :: Double, realToFrac (pt ! i) - realToFrac (pf ! i) :: Double) | v <- [1..nv], let i = v - 1]
putStrLn $ "vars: " ++ show xs2