module Data.Graph.Generators.Random.WattsStrogatz (
wattsStrogatzGraph,
wattsStrogatzGraph',
wattsStrogatzContext,
selectWithProbability
)
where
import System.Random.MWC
import Control.Monad
import Data.Graph.Generators
import Control.Applicative ((<$>))
import qualified Data.Map as Map
import qualified Data.Set as Set
wattsStrogatzContext :: GenIO
-> Int
-> [Int]
-> Double
-> IO GraphContext
wattsStrogatzContext gen n allNodes p = do
let endpoints = selectWithProbability gen p allNodes
inEdges <- endpoints
outEdges <- endpoints
return $ GraphContext inEdges n outEdges
wattsStrogatzGraph :: GenIO
-> Int
-> Int
-> Double
-> IO GraphInfo
wattsStrogatzGraph gen n k p = do
let allNodes = [0..n-1]
let insert m (i, js) = Map.insert i js m
let initialEdges = foldl (insert) (Map.empty) [ (i, forward_neighbors i) | i <- allNodes ]
allEdges <- rewrites (return initialEdges) $
[ (i, j) |
(i,s) <- Map.toList initialEdges,
j <- Set.toList s ]
return $ GraphInfo n (delineate allEdges)
where
k' = fromInteger.toInteger $ k
forward_neighbors :: Int -> Set.Set Int
forward_neighbors i = foldr (Set.insert) Set.empty $
fmap (`mod` n)
[i+1..i+k']
rewrites :: IO (Map.Map Int (Set.Set Int))
-> [(Int, Int)]
-> IO (Map.Map Int (Set.Set Int))
rewrites ioedges [] = ioedges
rewrites ioedges (t:tuples) = do
r <- uniform gen :: IO Double
edges <- ioedges :: IO (Map.Map Int (Set.Set Int))
if (r > p)
then rewrites (return edges) tuples
else do es <- rewrite t edges
rewrites (return es) tuples
rewrite (i, j1) edges = do
r <- uniform gen :: IO Double
let j2 = floor $ r*((fromInteger.toInteger) n)
if (((member (i, j2) edges) || (member (j2, i) edges)) || (i == j2))
then rewrite (i, j1) edges
else return $ swap (i, j1) (i, j2) edges
delineate :: Map.Map a (Set.Set b) -> [(a, b)]
delineate m = [ (i, j) |
(i, js) <- Map.toList m,
j <- Set.toList js ]
member :: (Int, Int) -> Map.Map Int (Set.Set Int) -> Bool
member (i, j) m = Set.member j (Map.findWithDefault Set.empty i m)
swap :: (Ord a, Ord b)
=> (a, b) -> (a, b)
-> Map.Map a (Set.Set b)
-> Map.Map a (Set.Set b)
swap (k1, v1) (k2, v2) m =
let set1 = Map.findWithDefault Set.empty k1 m in
let set2 = Set.insert v2 $ Set.delete v1 set1 in
let m2 = Map.insert k2 set2 $ Map.delete k1 m in m2
wattsStrogatzGraph' :: Int
-> Int
-> Double
-> IO GraphInfo
wattsStrogatzGraph' n k p =
withSystemRandom . asGenIO $ \gen -> wattsStrogatzGraph gen n k p
selectWithProbability :: GenIO
-> Double
-> [a]
-> IO [a]
selectWithProbability _ _ [] = return []
selectWithProbability gen p (x:xs) = do
r <- uniform gen :: IO Double
let v = [ x | r <= p ]
liftM2 (++) (return v) $ selectWithProbability gen p xs