Commits

Ken Watford committed 76966c1

Refactor the two part_graph functions into one. Add a test.

  • Participants
  • Parent commits 7567cdd

Comments (0)

Files changed (1)

 is 32, but they may be set to 64 if desired. If the values do not match
 what was used to compile the library, Bad Things(TM) will occur.
 
-The functions of primary interest in this module are:
+The function of primary interest in this module is:
 
-    `metis.part_graph_kway`
-    `metis.part_graph_recur`
+    `metis.part_graph`
 
-They are also the only objects export by "from metis import *".
+It is also the only object currently export by "from metis import *".
 Other objects in the module may be of interest to those looking to 
 mangle their graph datastructures into the required format. Examples
 of this include the `networkx_to_metis` and `adjlist_to_metis` functions.
-These are automatically called by the `part_graph_` functions, so there is
+These are automatically called by `part_graph`, so there is
 little need to call them manually.
 
 See http://bitbucket.org/kw/metis-python for updates and issue reporting.
 except ImportError:
     networkx = None
 
-__all__ = ['part_graph_kway', 'part_graph_recur']
+__all__ = ['part_graph']
 
 # Sadly, METIS does not currently include any API call to determine
 # the correct datatypes. So we either have to guess, let the user tell
     def __and__(self, other):
         assert isinstance(other, self.__class__)
         return self.__class__(self.value & other.value)
-    def __xor__(self):
-        assert isinstance(other, self.__class__)
+    def __xor__(self, other):
+        assert isinstance(other, (int, self.__class__))
         return self.__class__(self.value ^ other.value)
     def __not__(self):
         return self.__class__(~self.value)
     METIS_DBG_CONNINFO   = 128     #/*!< Show info on minimization of subdomain connectivity */
     METIS_DBG_CONTIGINFO = 256     #/*!< Show info on elimination of connected components */ 
     METIS_DBG_MEMORY     = 2048    #/*!< Show info related to wspace allocation */
+    METIS_DBG_ALL        = sum(2**i for i in range(9)+[11])
 
 class mobjtype_et(_enum):
     METIS_OBJTYPE_DEFAULT = -1
     xadj = (idx_t*(n+1))()
     adjncy = (idx_t*m2)()
     adjwgt = (idx_t*m2)()
+    seen_adjwgt = False # Don't use adjwgt unless we've seen any
 
     ncon = idx_t(1)
     if nodew:
         for j in adj:
             try:
                 adjncy[e], adjwgt[e] = j
+                seen_adjwgt = True
             except TypeError:
                 adjncy[e], adjwgt[e] = j, 1
             e += 1        
         xadj[i+1] = e
 
+    if not seen_adjwgt:
+        adjwgt = None
+
     return METIS_Graph(idx_t(n), ncon, xadj, adjncy, vwgt, vsize, adjwgt)
 
 
 @_wrapdll(P(idx_t), P(idx_t), P(idx_t), P(idx_t),
     P(idx_t), P(idx_t), P(idx_t), P(idx_t), P(real_t),
     P(real_t), P(idx_t), P(idx_t), P(idx_t))
-def _METIS_PartGraphKway(graph, nparts=2, 
-    tpwgts=None, ubvec=None, **opts):
+def _METIS_PartGraphKway(nvtxs, ncon, xadj, adjncy, vwgt, vsize, 
+                adjwgt, nparts, tpwgts, ubvec, options, objval, part):
     """
-    Perform k-way partitioning. 
-    Returns a 2-tuple `(part, objval)`, where `part` is a list of
+    Called by `part_graph`
+    """
+    return _METIS_PartGraphKway.call(nvtxs, ncon, xadj, adjncy, vwgt, vsize, 
+                adjwgt, nparts, tpwgts, ubvec, options, objval, part)
+
+@_wrapdll(P(idx_t), P(idx_t), P(idx_t), P(idx_t),
+    P(idx_t), P(idx_t), P(idx_t), P(idx_t), P(real_t),
+    P(real_t), P(idx_t), P(idx_t), P(idx_t))
+def _METIS_PartGraphRecursive(nvtxs, ncon, xadj, adjncy, vwgt, vsize, 
+                adjwgt, nparts, tpwgts, ubvec, options, objval, part):
+    """
+    Called by `part_graph` 
+    """
+    return _METIS_PartGraphRecursive.call(nvtxs, ncon, xadj, adjncy, vwgt, vsize, 
+                adjwgt, nparts, tpwgts, ubvec, options, objval, part)
+
+### End METIS wrappers. ###
+
+def part_graph(graph, nparts=2, 
+    tpwgts=None, ubvec=None, recursive=False, **opts):
+    """
+    Perform graph partitioning using k-way or recursive methods.
+
+    Returns a 2-tuple `(objval, parts)`, where `parts` is a list of
     partition indices corresponding and `objval` is the value of 
     the objective function that was minimized (either the edge cuts
     or the total volume).
     correct type (`metis.real_t`). 
 
     Any additional METIS options may be specified as keyword parameters.
-    For this function, the appropriate options are:
+
+    For k-way clustering, the appropriate options are:
 
        objtype   = 'cut' or 'vol' 
        ctype     = 'rm' or 'shem' 
        numbering = 0 (C-style) or 1 (Fortran-style) indices
        dbglvl    = Debug flag bitfield
 
-    See the METIS manual for specifics. 
+    For recursive clustering, the appropraite options are:
+
+       ctype     = 'rm' or 'shem' 
+       iptype    = 'grow', 'random', 'edge', 'node'
+       rtype     = 'fm', 'greedy', 'sep2sided', 'sep1sided'
+       ncuts     = integer, number of cut attempts (default = 1)
+       niter     = integer, number of iterations (default = 10)
+       ufactor   = integer, maximum load imbalance of (1+x)/1000
+       seed      = integer, RNG seed
+       numbering = 0 (C-style) or 1 (Fortran-style) indices
+       dbglvl    = Debug flag bitfield
+
+    See the METIS manual for specific meaning of each option.
     """
+
     options = METIS_Options(**opts)
     if networkx and isinstance(graph, networkx.Graph):
         graph = networkx_to_metis(graph)
     objval = idx_t()
     partition = (idx_t*graph.nvtxs.value)()
 
-    _METIS_PartGraphKway.call(
-        byref(graph.nvtxs), byref(graph.ncon), graph.xadj,
-        graph.adjncy, graph.vwgt, graph.vsize, graph.adjwgt,
-        byref(nparts_var), tpwgts, ubvec, options.array, 
-        byref(objval), partition)
+    args = (byref(graph.nvtxs), byref(graph.ncon), graph.xadj,
+            graph.adjncy, graph.vwgt, graph.vsize, graph.adjwgt,
+            byref(nparts_var), tpwgts, ubvec, options.array, 
+            byref(objval), partition)
+    if recursive:
+        _METIS_PartGraphRecursive(*args)
+    else:
+        _METIS_PartGraphKway(*args)
+    
+    return objval.value, list(partition)
 
-    return list(partition), objval.value
+def test():
+    adjlist = [[1, 2, 3, 4], [0], [0], [0], [0, 5], [4, 6], [5, 7, 13], 
+                [8, 6], [7], [10, 11, 12], [9], [9], [9], [6, 14], [13, 15],
+                [14, 16, 17, 18], [15], [15], [15]] 
+    print "Testing k-way cut"
+    cuts, parts = part_graph(adjlist, 3, recursive=False, dbglvl=METIS_DBG_ALL)
+    assert cuts == 2
+    assert set(parts) == set([0,1,2])
 
-@_wrapdll(P(idx_t), P(idx_t), P(idx_t), P(idx_t),
-    P(idx_t), P(idx_t), P(idx_t), P(idx_t), P(real_t),
-    P(real_t), P(idx_t), P(idx_t), P(idx_t))
-def _METIS_PartGraphRecursive(graph, nparts=2, 
-    tpwgts=None, ubvec=None, **opts):
-    """
-    Perform recursive partitioning. 
-    Returns a 2-tuple `(part, objval)`, where `part` is a list of
-    partition indices corresponding and `objval` is the value of 
-    the objective function that was minimized (either the edge cuts
-    or the total volume).
+    print "Testing recursive cut"
+    cuts, parts = part_graph(adjlist, 3, recursive=True, dbglvl=METIS_DBG_ALL)
+    assert cuts <= 3
+    assert set(parts) == set([0,1,2])
 
-    `graph` may be a NetworkX graph, an adjacency list, or a METIS_Graph 
-    named tuple. To use the named tuple approach, you'll need to
-    read the METIS manual for the meanings of the fields.
+    if networkx:
+        print "Testing with NetworkX"
+        G = networkx.Graph()
+        G.add_star([0,1,2,3,4])
+        G.add_path([4,5,6,7,8])
+        G.add_star([9,10,11,12])
+        G.add_path([6,13,14,15])
+        G.add_star([15,16,17,18])
+        cuts, parts = part_graph(adjlist, 3)
+        assert cuts == 2
+        assert set(parts) == set([0,1,2])
 
-    See `networkx_to_metis` for help and details on how the
-    graph is converted and how node/edge weights and sizes can
-    be specified.
+    print "METIS appears to be working."
 
-    See `adjlist_to_metis` for information on the use of adjacency lists.
-    The extra `nodew` and `nodesz` options of that function may be given 
-    directly to this function and will be forwarded to the converter. 
-    Alternatively, a dictionary can be provided as `graph` and its items
-    will be passed as keyword arguments.
-
-    `nparts` is the target number of partitions. You might get fewer.
-
-    You probably won't need to specify `tpwgts` and `ubvec`, but if 
-    you do, you'll read the METIS manual to see what they are. Any
-    sequence type with floats may be provided so long as the number of
-    elements is correct. If a ctypes Array is provided, it must be the
-    correct type (`metis.real_t`). 
-
-    Any additional METIS options may be specified as keyword parameters.
-    For this function, the appropriate options are:
-
-       ctype     = 'rm' or 'shem' 
-       iptype    = 'grow', 'random', 'edge', 'node'
-       rtype     = 'fm', 'greedy', 'sep2sided', 'sep1sided'
-       ncuts     = integer, number of cut attempts (default = 1)
-       niter     = integer, number of iterations (default = 10)
-       ufactor   = integer, maximum load imbalance of (1+x)/1000
-       seed      = integer, RNG seed
-       numbering = 0 (C-style) or 1 (Fortran-style) indices
-       dbglvl    = Debug flag bitfield
-
-    See the METIS manual for specifics. 
-    """
-    options = METIS_Options(**opts)
-    if networkx and isinstance(graph, networkx.Graph):
-        graph = networkx_to_metis(graph)
-    elif isinstance(graph, (list,tuple)):
-        nodesz = opts.pop('nodesz', None)
-        nodew  = opts.pop('nodew', None)
-        graph = adjlist_to_metis(graph, nodew, nodesz)
-    elif isinstance(graph, dict):
-        # Check if this has METIS_Graph fields or an adjlist
-        if 'nvtxs' in graph:
-            graph = METIS_Graph(**graph)
-        elif 'adjlist' in graph:
-            graph = adlist_to_metis(**graph)
-
-    if tpwgts and not isinstance(tpwgts, ctypes.Array):
-        tpwgts = (real_t*len(tpwgts))(*tpwgts)
-    if ubvec and not isinstance(ubvec, ctypes.Array):
-        ubvec = (real_t*len(ubvect))(*ubvec)
-    nparts_var = idx_t(nparts)
-
-    objval = idx_t()
-    partition = (idx_t*graph.nvtxs.value)()
-
-    _METIS_PartGraphRecursive.call(
-        byref(graph.nvtxs), byref(graph.ncon), graph.xadj,
-        graph.adjncy, graph.vwgt, graph.vsize, graph.adjwgt,
-        byref(nparts_var), tpwgts, ubvec, options.array, 
-        byref(objval), partition)
-
-    return list(partition), objval.value
-
-part_graph_kway = _METIS_PartGraphKway
-
-part_graph_recur = _METIS_PartGraphRecursive
-
-### End METIS wrappers. ###
-
-