add ability to extract the resulting clustering from infomap, might be useful

cbits/wrapper.cpp

@@ -5,6 +5,10 @@ typedef infomap::InfomapWrapper IW;
 extern "C" {
   struct network;
 
+  struct cluster_node {
+    unsigned int node, cluster_id;
+  };
+
   network* network_create(const char* cfgstr) {
     IW* iw = new infomap::InfomapWrapper(cfgstr);
     return reinterpret_cast<network*>(iw);
@@ -13,15 +17,28 @@ extern "C" {
     delete reinterpret_cast<IW*>(n);
   }
 
-  void add_link(network* n, unsigned int i, unsigned int j) {
-    reinterpret_cast<IW*>(n)->addLink(i, j);
+  void add_link(network* n, unsigned int i, unsigned int j, double w) {
+    reinterpret_cast<IW*>(n)->addLink(i, j, w);
   }
 
   void run(network* n) {
     reinterpret_cast<IW*>(n)->run();
   }
 
-  double getCodelength(network* n) {
+  double get_codelength(network* n) {
     return reinterpret_cast<IW*>(n)->codelength();
   }
+
+  unsigned int get_num_nodes(network* n) {
+    return reinterpret_cast<IW*>(n)->network().numNodes();
+  }
+
+  void write_modules(network* n, cluster_node* nodes) {
+    auto modules = reinterpret_cast<IW*>(n)->getModules();
+    int i = 0;
+    for(auto node : modules) {
+      nodes[i] = {node.first, node.second};
+      i++;
+    }
+  }
 };

hsinfomap.cabal

@@ -26,11 +26,13 @@ extra-source-files: CHANGELOG.md
 library
     -- Modules included in this executable, other than Main.
     exposed-modules:
+      Data.Graph.Infomap
       Data.Graph.Infomap.Internal
 
     -- LANGUAGE extensions used by modules in this package.
     -- other-extensions:
-    build-depends:    base >=4.13
+    build-depends:    base >=4.13,
+                      containers
     hs-source-dirs: src
     include-dirs: cbits/infomap
     cxx-options: -std=c++14 -O3

src/Data/Graph/Infomap.hs

+module Data.Graph.Infomap where
+
+import Data.Graph.Infomap.Internal
+import Data.Map
+import Foreign.C.Types
+
+import qualified Data.Map as Map
+
+infomap :: NetworkCfg -> Map (Int, Int) Double -> IO [CNode]
+infomap cfg edges = withNetwork cfg $ \n -> do
+  addLinks n edges
+  run n
+  getModules n
+
+  where addLinks n = Map.foldrWithKey (go n) (return ())
+        go n (i, j) w act = addLink n (fromIntegral i)
+                                      (fromIntegral j)
+                                      (CDouble w)
+                         >> act
+
+testHigh = infomap "--silent --two-level -N2" testmap
+  where testmap = Map.fromList
+          [ ((0, 1), 1.0)
+          , ((0, 2), 1.0)
+          , ((0, 3), 1.0)
+          , ((1, 0), 1.0)
+          , ((1, 2), 1.0)
+          , ((2, 1), 1.0)
+          , ((2, 0), 1.0)
+          , ((3, 0), 1.0)
+          , ((3, 4), 1.0)
+          , ((3, 5), 1.0)
+          , ((4, 3), 1.0)
+          , ((4, 5), 1.0)
+          , ((5, 4), 1.0)
+          , ((5, 3), 1.0)
+          ]

src/Data/Graph/Infomap/Internal.hs

@@ -3,10 +3,29 @@ module Data.Graph.Infomap.Internal where
 import Control.Exception
 import Foreign.C.String
 import Foreign.C.Types
+import Foreign.Marshal.Array
 import Foreign.Ptr
+import Foreign.Storable
 
 newtype Network = Network (Ptr Network)
 
+-- node id, cluster id
+data CNode = CNode
+  { nodeId :: CUInt
+  , clusId :: CUInt
+  } deriving Show
+
+instance Storable CNode where
+  sizeOf _ = 2 * sizeOf (undefined :: CUInt)
+  alignment _ = alignment (undefined :: CUInt)
+  peek ptr = do
+    i <- peek (castPtr ptr)
+    c <- peekByteOff ptr (sizeOf i)
+    return (CNode i c)
+  poke ptr (CNode i c) = do
+    poke (castPtr ptr) i
+    pokeByteOff ptr (sizeOf i) c
+
 foreign import ccall unsafe "network_create" createNetwork
   :: CString -> IO Network
 
@@ -14,36 +33,53 @@ foreign import ccall unsafe "network_destroy" destroyNetwork
   :: Network -> IO ()
 
 foreign import ccall unsafe "add_link" addLink
-  :: Network -> CUInt -> CUInt -> IO ()
+  :: Network -> CUInt -> CUInt -> CDouble -> IO ()
 
 foreign import ccall unsafe "run" run
   :: Network -> IO ()
 
-foreign import ccall unsafe "getCodelength" getCodelength
+foreign import ccall unsafe "get_codelength" getCodelength
   :: Network -> IO CDouble
 
+foreign import ccall unsafe "get_num_nodes" getNumNodes
+  :: Network -> IO CUInt
+
+foreign import ccall unsafe "write_modules" writeModules
+  :: Network -> Ptr CNode -> IO ()
+
 type NetworkCfg = String
 
+getModules :: Network -> IO [CNode]
+getModules n = do
+  len <- fromIntegral <$> getNumNodes n
+  allocaArray len $ \ptr -> do
+    writeModules n ptr
+    peekArray len ptr
+
 withNetwork :: NetworkCfg -> (Network -> IO a) -> IO a
 withNetwork cfgstr f = withCString cfgstr $ \cfgcstr ->
   bracket (createNetwork cfgcstr) destroyNetwork f
 
 test :: IO ()
-test = withNetwork "--two-level -N2" $ \n -> do
-  addLink n 0 1
-  addLink n 0 2
-  addLink n 0 3
-  addLink n 1 0
-  addLink n 1 2
-  addLink n 2 1
-  addLink n 2 0
-  addLink n 3 0
-  addLink n 3 4
-  addLink n 3 5
-  addLink n 4 3
-  addLink n 4 5
-  addLink n 5 4
-  addLink n 5 3
+test = withNetwork "--silent --two-level -N2" $ \n -> do
+  addLink n 0 1 1.0
+  addLink n 0 2 1.0
+  addLink n 0 3 1.0
+  addLink n 1 0 1.0
+  addLink n 1 2 1.0
+  addLink n 2 1 1.0
+  addLink n 2 0 1.0
+  addLink n 3 0 1.0
+  addLink n 3 4 1.0
+  addLink n 3 5 1.0
+  addLink n 4 3 1.0
+  addLink n 4 5 1.0
+  addLink n 5 4 1.0
+  addLink n 5 3 1.0
+  sz <- getNumNodes n
+  putStrLn $ "#nodes: " ++ show sz
   run n
   codelen <- getCodelength n
   putStrLn $ "codelength: " ++ show codelen
+  modules <- getModules n
+  putStrLn $ "modules: " ++ show modules