Source

metis-python / metis.py

Full commit
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
"""
METIS for Python
================

Wrapper for the METIS library for partitioning graphs (and other stuff).

This library is unrelated to PyMetis, except that they wrap the same library.
PyMetis is a Boost Python extension, while this library is pure python and will
run under PyPy and interpreters with similarly compatible ctypes libraries.

NetworkX_ is recommended for representing graphs for use with this wrapper,
but it isn't required. Simple adjacency lists are supported as well.

.. _NetworkX: http://networkx.lanl.gov/

The function of primary interest in this module is :func:`part_graph`.

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 :func:`networkx_to_metis` and :func:`adjlist_to_metis` functions.
These are automatically called by :func:`part_graph`, so there is
little need to call them manually.

See the BitBucket repository_  for updates and issue reporting.

.. _repository: http://bitbucket.org/kw/metis-python

Installation
============

It's on PyPI, so installation should be as easy as::

    pip install metis
          -or-
    easy_install metis

METIS itself is not included with this wrapper. Get it here_.

.. _here: http://glaros.dtc.umn.edu/gkhome/views/metis

Note that the shared library is needed, and isn't enabled by default
by the configuration process. Turn it on by issuing::

    make config shared=1

Your operating system's package manager might know about METIS,
but this wrapper was designed for use with METIS 5. Packages with
METIS 4 will not work.

This wrapper uses a few environment variables:

.. envvar:: METIS_DLL

    This wrapper uses Python's ctypes module to interface with the METIS
    shared library. If it is unable to automatically locate the library, you
    may specify the full path to the library file in this environment variable.

.. envvar:: METIS_IDXTYPEWIDTH
.. envvar:: METIS_REALTYPEWIDTH

    The sizes of the :c:type:`idx_t` and :c:type:`real_t` types are not
    easily determinable at runtime, so they can be provided with these 
    environment variables. The default value for each of these (at both compile 
    time and in this library) 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.

Example
=======

    >>> import networkx as nx
    >>> import metis 
    >>> G = metis.example_networkx()
    >>> (edgecuts, parts) = metis.part_graph(G, 3)
    >>> colors = ['red','blue','green']
    >>> for i, p in enumerate(parts):
    ...     G.node[i]['color'] = colors[p]
    ... 
    >>> nx.write_dot(G, 'example.dot') # Requires pydot or pygraphviz

.. graphviz:: example.dot

"""
# Copyright (c) 2012 Ken Watford
#
# Permission is hereby granted, free of charge, to any person 
# obtaining a copy of this software and associated documentation 
# files (the "Software"), to deal in the Software without 
# restriction, including without limitation the rights to use, 
# copy, modify, merge, publish, distribute, sublicense, and/or 
# sell copies of the Software, and to permit persons to whom the 
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be 
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES 
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR 
# ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF 
# CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#
# tl;dr - MIT license.

__version__ = '0.1a'

import ctypes
from ctypes import POINTER as P, byref
import os, sys, operator as op
from warnings import warn
from collections import namedtuple
try:
    import networkx
except ImportError:
    networkx = None

__all__ = ['part_graph', 'networkx_to_metis', 'adjlist_to_metis']

# Sadly, METIS does not currently include any API call to determine
# the correct datatypes. So we either have to guess, let the user tell
# us, try to infer it by checking API behavior on test inputs, or
# look for the header and parse out the preprocessor macros.
# Since we're in a bit of a hurry, for now we'll use the defaults
# and let the user specify if this is wrong with env vars.
IDXTYPEWIDTH  = os.getenv('METIS_IDXTYPEWIDTH', '32')
REALTYPEWIDTH = os.getenv('METIS_REALTYPEWIDTH', '32')

if IDXTYPEWIDTH == '32':
    idx_t = ctypes.c_int32
elif IDXTYPEWIDTH == '64':
    idx_t = ctypes.c_int64
else:
    raise EnvironmentError('Env var METIS_IDXTYPEWIDTH must be "32" or "64"')

if REALTYPEWIDTH == '32':
    real_t = ctypes.c_float
elif REALTYPEWIDTH == '64':
    real_t = ctypes.c_double
else:
    raise EnvironmentError('Env var METIS_REALTYPEWIDTH must be "32" or "64"')


METIS_NOPTIONS = 40

# The _enum and _bitfield base classes come from my PyCL project
# They make enum constants a little more friendly.
class _enum(ctypes.c_int32):
    # Base class for various enums
    def __eq__(self, other):
        if not isinstance(other, self.__class__):
            return False
        else:
            return self.value == other.value
    def __ne__(self, other):
        return not(self == other)
    def __hash__(self):
        return self.value.__hash__()
    def __repr__(self):
        by_value = self.__class__._by_value
        names = []
        if self in by_value:
            return by_value[self]
        else:
            return "UNKNOWN(0x%x)" % self.value
    def shortname(self):        
        return self._by_value_short.get(self, 'unknown')
    @classmethod
    def fromname(cls, name):
        if name in cls._by_name:
            return cls._by_name[name]
        elif name in cls._by_name_short:
            return cls._by_name_short[name]
        else:
            raise KeyError('Unknown name: %s' % name) 
    @classmethod
    def toname(cls, val):
        if val in self._by_value_short:
            return cls._by_value_short[value]
        elif val in cls._by_value:
            return cls._by_value[value]
        else:
            raise KeyError

class _bitfield(ctypes.c_int32):
    # Base class for bitfield values
    # Bitwise operations for combining flags are supported.
    def __or__(self, other):
        assert isinstance(other, self.__class__)
        return self.__class__(self.value | other.value)
    def __and__(self, other):
        assert isinstance(other, self.__class__)
        return self.__class__(self.value & other.value)
    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)
    def __contains__(self, other):
        assert isinstance(other, self.__class__)
        return (self.value & other.value) == other.value
    def __hash__(self):
        return self.value.__hash__()
    def __eq__(self, other):
        if not isinstance(other, self.__class__):
            return False
        else:
            return self.value == other.value
    def __ne__(self, other):
        return not(self == other)
    def __repr__(self):
        by_value = self.__class__._by_value
        names = []
        if self in by_value:
            return by_value[self]
        for val in by_value:
            if val in self:
                names.append(by_value[val])
        if names:
            return " | ".join(names)
        else:
            return "UNKNOWN(0x%x)" % self.value

class rstatus_et(_enum):
    METIS_OK              =  1    #/*!< Returned normally */
    METIS_ERROR_INPUT     = -2   #/*!< Returned due to erroneous inputs and/or options */
    METIS_ERROR_MEMORY    = -3   #/*!< Returned due to insufficient memory */
    METIS_ERROR           = -4   #/*!< Some other errors */

class moptype_et(_enum):
    METIS_OP_PMETIS = 0
    METIS_OP_KMETIS = 1
    METIS_OP_OMETIS = 2

class moptions_et(_enum):
    METIS_OPTION_PTYPE     =  0
    METIS_OPTION_OBJTYPE   =  1
    METIS_OPTION_CTYPE     =  2
    METIS_OPTION_IPTYPE    =  3
    METIS_OPTION_RTYPE     =  4
    METIS_OPTION_DBGLVL    =  5
    METIS_OPTION_NITER     =  6
    METIS_OPTION_NCUTS     =  7
    METIS_OPTION_SEED      =  8
    METIS_OPTION_MINCONN   =  9 
    METIS_OPTION_CONTIG    = 10
    METIS_OPTION_COMPRESS  = 11
    METIS_OPTION_CCORDER   = 12 
    METIS_OPTION_PFACTOR   = 13
    METIS_OPTION_NSEPS     = 14
    METIS_OPTION_UFACTOR   = 15
    METIS_OPTION_NUMBERING = 16
    #/* Used for command-line parameter purposes */
    METIS_OPTION_HELP      = 17
    METIS_OPTION_TPWGTS    = 18 
    METIS_OPTION_NCOMMON   = 19
    METIS_OPTION_NOOUTPUT  = 20
    METIS_OPTION_BALANCE   = 21 
    METIS_OPTION_GTYPE     = 22
    METIS_OPTION_UBVEC     = 23

class mptype_et(_enum):
    METIS_PTYPE_DEFAULT = -1
    METIS_PTYPE_RB   = 0
    METIS_PTYPE_KWAY = 1              

class mgtype_et(_enum):
    METIS_GTYPE_DEFAULT = -1
    METIS_GTYPE_DUAL  = 0
    METIS_GTYPE_NODAL = 1

class mctype_et(_enum):
    METIS_CTYPE_DEFAULT = -1
    METIS_CTYPE_RM   = 0
    METIS_CTYPE_SHEM = 1

class miptype_et(_enum):
    METIS_IPTYPE_DEFAULT = -1
    METIS_IPTYPE_GROW    = 0
    METIS_IPTYPE_RANDOM  = 1
    METIS_IPTYPE_EDGE    = 2
    METIS_IPTYPE_NODE    = 3
    METIS_IPTYPE_METISRB = 4

class mrtype_et(_enum):
    METIS_RTYPE_DEFAULT   = -1
    METIS_RTYPE_FM        = 0
    METIS_RTYPE_GREEDY    = 1
    METIS_RTYPE_SEP2SIDED = 2
    METIS_RTYPE_SEP1SIDED = 3

class mdbglvl_et(_bitfield):
    METIS_DBG_DEFAULT    = -1
    METIS_DBG_INFO       = 1       #/*!< Shows various diagnostic messages */
    METIS_DBG_TIME       = 2       #/*!< Perform timing analysis */
    METIS_DBG_COARSEN    = 4       #/*!< Show the coarsening progress */
    METIS_DBG_REFINE     = 8       #/*!< Show the refinement progress */
    METIS_DBG_IPART      = 16      #/*!< Show info on initial partitioning */
    METIS_DBG_MOVEINFO   = 32      #/*!< Show info on vertex moves during refinement */
    METIS_DBG_SEPINFO    = 64      #/*!< Show info on vertex moves during sep refinement */
    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
    METIS_OBJTYPE_CUT  = 0
    METIS_OBJTYPE_VOL  = 1
    METIS_OBJTYPE_NODE = 2

# For enums and bitfields, do magic. Each type gets a registry of the
# names and values of their defined elements, to support pretty printing.
# Further, each of the class variables (which are defined using ints) is
# upgraded to be a member of the class in question.
# Additionally, each of the constants is copied into the module scope.
for cls in (_enum.__subclasses__() + _bitfield.__subclasses__()):
    if cls.__name__ not in globals() or cls.__name__.startswith('_'):
        # Don't apply this to types that ctypes makes automatically,
        # like the _be classes. Doing so will overwrite the declared
        # constants at global scope, which is really weird.
        continue
    cls._by_name = dict()
    cls._by_value = dict()
    cls._by_name_short = dict()
    cls._by_value_short = dict()
    if not cls.__doc__:
        cls.__doc__ = ""
    for name, value in cls.__dict__.items():
        if isinstance(value, int):
            obj = cls(value)
            setattr(cls, name, obj)
            cls._by_name[name] = obj
            cls._by_value[obj] = name
            shortname = name.split('_')[-1].lower()
            cls._by_name_short[shortname] = obj
            cls._by_value_short[obj] = shortname
            globals()[name] = obj
            cls.__doc__ += """
            .. attribute:: %s
            """ % name
# cleanup
del cls; del name; del value; del obj


# Convert values taken from option array into appropriate enum
_opt_types = {
    METIS_OPTION_PTYPE   : mptype_et,
    METIS_OPTION_OBJTYPE : mobjtype_et,
    METIS_OPTION_CTYPE   : mctype_et,
    METIS_OPTION_GTYPE   : mgtype_et,
    METIS_OPTION_IPTYPE  : miptype_et,
    METIS_OPTION_RTYPE   : mrtype_et,
    METIS_OPTION_DBGLVL  : mdbglvl_et,
    }

class METIS_Options(object):
    """ Represents the 'options' array used to represent all 
    nearly all options that can be given to METIS functions.
    Will be used when extra keyword arguments are are used in wrappers.

    Note that I spent way too much time on this. 

    """
    def __init__(self, options=None, **opts):
        self.array = (idx_t*METIS_NOPTIONS)()
        _METIS_SetDefaultOptions(self.array)
        if options:
            for opt, val in options.keys():
                self[opt] = val
        for opt, val in opts.items():
            self[opt] = val

    def keys(self):
        return moptions_et._by_name_short.keys()

    def __getitem__(self, opt):        
        if isinstance(opt, basestring):
            opt = moptions_et.fromname(opt)
        val = self.array[opt.value]   
        if opt in _opt_types:
            val = _opt_types[opt](val)
            if isinstance(val, _enum):
                val = val.shortname()
        return val

    def __setitem__(self, opt, val):
        if isinstance(opt, basestring):
            opt = moptions_et.fromname(opt)
        if isinstance(val, basestring) and opt in _opt_types:
            val = _opt_types[opt].fromname(val)
        try:
            self.array[opt.value] = val
        except TypeError:
            raise TypeError("Bad type for option %s: %s" %
                (opt, val.__class__.__name__))

    def __repr__(self):
        """ Only show non-default options """
        nondefaults = []
        for opt in self.keys():
            realind = moptions_et.fromname(opt).value
            if self.array[realind] != -1:
                val = self[opt]
                nondefaults.append('%s=%r' % (opt, val))
        return 'METIS_Options(' + ', '.join(nondefaults) + ')'
   


# Attempt to locate and load the appropriate shared library
_dll_filename = os.getenv('METIS_DLL')
if not _dll_filename:
    try:
        from ctypes.util import find_library as _find_library
        _dll_filename = _find_library('metis')
    except ImportError:
        pass
if _dll_filename:
    try:
        _dll = ctypes.cdll.LoadLibrary(_dll_filename)
    except:        
        raise RuntimeError('Could not load METIS dll: %s' % _dll_filename)
else:
    if os.environ.get('READTHEDOCS', None) == 'True':
        # Don't care if we can load the DLL on RTD.
        _dll = None
    else:
        raise RuntimeError('Could not locate METIS dll. Please set the METIS_DLL environment variable to its full path.')

# Wrapping conveniences

def _wrapdll(*argtypes, **kw):
    """
    Decorator used to simplify wrapping METIS functions a bit.

    The positional arguments represent the ctypes argument types the
    C-level function expects, and will be used to do argument type checking.

    If a `res` keyword argument is given, it represents the C-level
    function's expected return type. The default is `rstatus_et`
    
    If an `err` keyword argument is given, it represents an error checker
    that should be run after low-level calls. The `_result_errcheck` and
    `_lastarg_errcheck` functions should be sufficient for most OpenCL
    functions. `_result_errcheck` is the default value.

    The decorated function should have the same name as the underlying
    METIS function, since the function name is used to do the lookup. The
    C-level function pointer will be stored in the decorated function's
    `call` attribute, and should be used by the decorated function to
    perform the actual call(s). The wrapped function is otherwise untouched.

    """
    def dowrap(f):
        if f.__name__.startswith('_'):
            name = f.__name__[1:]
        else:
            name = f.__name__
        if _dll:
            wrapped_func = getattr(_dll, name)        
            wrapped_func.argtypes = argtypes
            res = kw.pop('res', rstatus_et)
            wrapped_func.restype = res
            err = kw.pop('err', _result_errcheck)
            wrapped_func.errcheck = err
            f.call = wrapped_func
        else:
            def nodll(*args, **kw):
                raise NotImplemented, "No METIS DLL"
            f.call = nodll
        return f
    return dowrap

# Translate METIS status messages into Python exceptions
class METIS_Error(Exception): pass
class METIS_MemoryError(METIS_Error, MemoryError): pass
class METIS_InputError(METIS_Error, ValueError): pass
class METIS_OtherError(METIS_Error): pass

def _result_errcheck(result, func, args):
    """
    For use in the errcheck attribute of a ctypes function wrapper.

    Most METIS functions return rstatus_et. This checks it for
    an error code and raises an appropriate exception if it finds one.

    This is the default error checker when using _wrapdll
    """
    if result != METIS_OK:
        if result == METIS_ERROR_INPUT: raise METIS_InputError
        if result == METIS_ERROR_MEMORY: raise METIS_MemoryError
        if result == METIS_ERROR: raise METIS_OtherError
        raise RuntimeError("Error raising error: Bad error.") # lolwut
    return result

# Graph helpers

METIS_Graph = namedtuple('METIS_Graph',
    'nvtxs ncon xadj adjncy vwgt vsize adjwgt')

def networkx_to_metis(G):
    """
    Convert NetworkX graph into something METIS can consume
    The graph may specify weights and sizes using the following
    graph attributes:

    * ``edge_weight_attr``
    * ``node_weight_attr`` (multiple names allowed)
    * ``node_size_attr``

    For example::

        >>> G.edge[0][1]['weight'] = 3
        >>> G.node[0]['quality'] = 5
        >>> G.node[0]['specialness'] = 8
        >>> G.graph['edge_weight_attr'] = 'weight'
        >>> G.graph['node_weight_attr'] = ['quality', 'specialness']

    If node_weight_attr is a list instead of a string, then multiple
    node weight labels can be provided. 

    All weights must be integer values. If an attr label is specified but 
    a node/edge is missing that attribute, it defaults to 1.

    If a graph attribute is not provided, no defaut is used. That is, if
    ``edge_weight_attr`` is not set, then ``'weight'`` is not used as the
    default, and the graph will appear unweighted to METIS.
    """
    n = G.number_of_nodes()
    m = G.number_of_edges()
    nvtxs = idx_t(n)
    
    H = networkx.convert_node_labels_to_integers(G)
    xadj = (idx_t*(n+1))()
    adjncy = (idx_t*(2*m))()

    # Check graph attributes for weight/size labels
    edgew = G.graph.get('edge_weight_attr', None)
    nodew = G.graph.get('node_weight_attr', [])
    nodesz = G.graph.get('node_size_attr', None)

    if edgew:
        adjwgt = (idx_t*(2*m))()
    else:
        adjwgt = None

    if nodew:
        if isinstance(nodew, basestring):
            nodew = [nodew]            
        nc = len(nodew)
        ncon = idx_t(nc)
        vwgt = (idx_t*(n*len(nodew)))()                
    else:
        ncon = idx_t(1)
        vwgt = None

    if nodesz:
        vsize = (idx_t*n)()
    else:
        vsize = None

    # Fill in each array
    xadj[0] = e = 0
    for i in H.node:
        for c,w in enumerate(nodew):
            try:            
                vwgt[i*nc+c] = H.node[i].get(w, 1)
            except TypeError:
                raise TypeError("Node weights must be integers" )

        if nodesz:
            try:
                vsize[i] = H.node[i].get(nodesz, 1)
            except TypeError:
                raise TypeError("Node sizes must be integers")

        for j, attr in H.edge[i].iteritems():
            adjncy[e] = j            
            if edgew:
                try:
                    adjwgt[e] = attr.get(edgew, 1)
                except TypeError:
                    raise TypeError("Edge weights must be integers")
            e += 1
        xadj[i+1] = e
    
    return METIS_Graph(nvtxs, ncon, xadj, adjncy, vwgt, vsize, adjwgt)

def adjlist_to_metis(adjlist, nodew=None, nodesz=None):
    """
    Rudimentary adjacency list converter.
    Primarily of use if you don't have or don't want to use NetworkX.

    :param adjlist: A list of tuples. Each list element represents a node or vertex
      in the graph. Each item in the tuples represents an edge. These items may be 
      single integers representing neighbor index, or they may be an (index, weight)
      tuple if you want weighted edges. Default weight is 1 for missing weights.
      
      The graph must be undirected, and each edge must be represented twice (once for
      each node). The weights should be identical, if provided.
    :param nodew: is a list of node weights, and must be the same size as `adjlist` if given.
      If desired, the elements of `nodew` may be tuples of the same size (>= 1) to provided
      multiple weights for each node. 
    :param nodesz: is a list of node sizes. These are relevant when doing volume-based
      partitioning. 

    Note that all weights and sizes must be non-negative integers.
    """
    n = len(adjlist)
    m2 = sum(map(len, adjlist))    

    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:
        if isinstance(nodew[0]):
            vwgt = (idx_t*n)(*nodew)
        else: # Assume a list of them
            nw = len(nodew[0])
            ncon = idx_t(nw)
            vwgt = (idx_t*(nw*n))(*reduce(op.add,nodew))
    else:
        vwgt = None

    if nodesz:
        vsize = (idx_t*n)(*nodesz)
    else:
        vsize = None

    xadj[0] = e = 0
    for i, adj in enumerate(adjlist):
        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)




### Wrapped METIS functions ###

@_wrapdll(P(idx_t))
def _METIS_SetDefaultOptions(optarray):
    _METIS_SetDefaultOptions.call(optarray)

@_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(nvtxs, ncon, xadj, adjncy, vwgt, vsize, 
                adjwgt, nparts, tpwgts, ubvec, options, objval, part):
    """
    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).

    :param graph: may be a NetworkX graph, an adjacency list, or a :class:`METIS_Graph`
      named tuple. To use the named tuple approach, you'll need to
      read the METIS manual for the meanings of the fields.
  
      See :func:`networkx_to_metis` for help and details on how the
      graph is converted and how node/edge weights and sizes can
      be specified.
  
      See :func:`adjlist_to_metis` for information on the use of adjacency lists.
      The extra ``nodew`` and ``nodesz`` keyword arguments 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.
    :param nparts: The target number of partitions. You might get fewer.
    :param tpwgts: Target partition weights. For each partition, there should
      be one (float) weight for each node constraint. That is, if `nparts` is 3 and 
      each node of the graph has two weights, then tpwgts might look like this::
  
        [(0.5, 0.1), (0.25, 0.8), (0.25, 0.1)]

      This list may be provided flattened. The internal tuples are for convenience.
      The partition weights for each constraint must sum to 1. 
    :param ubvec: The load imalance tolerance for each node constraint. Should be
      a list of floating point values each greater than 1.

    :param recursive: Determines whether the partitioning should be done by 
      direct k-way cuts or by a series of recursive cuts. These correspond to 
      :c:func:`METIS_PartGraphKway` and :c:func:`METIS_PartGraphRecursive` in
      METIS's C API.

    Any additional METIS options may be specified as keyword parameters.

    For k-way clustering, the appropriate options are::

        objtype   = 'cut' or 'vol' 
        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
        minconn   = bool, minimize degree of subdomain graph
        contig    = bool, force contiguous partitions
        seed      = integer, RNG seed
        numbering = 0 (C-style) or 1 (Fortran-style) indices
        dbglvl    = Debug flag bitfield

    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)
    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):
        if isinstance(tpwgts[0], (tuple, list)):
            tpwgts = reduce(op.add, tpwgts)   
        tpwgts = (real_t*len(tpwgts))(*tpwgts)
    if ubvec and not isinstance(ubvec, ctypes.Array):
        ubvec = (real_t*len(ubvect))(*ubvec)

    if tpwgts: assert len(tpgwgts) == nparts * graph.ncon
    if ubvec: assert len(ubvec) == graph.ncon

    nparts_var = idx_t(nparts)

    objval = idx_t()
    partition = (idx_t*graph.nvtxs.value)()

    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)

def example_adjlist():
    return [[1, 2, 3, 4], [0], [0], [0], [0, 5], [4, 6], [13, 5, 7], 
            [8, 6], [9, 10, 11, 12, 7], [8], [8], [8], [8], [14, 6], [13, 15],
            [16, 17, 18, 14], [15], [15], [15]]

def example_networkx():
    G = networkx.Graph()
    G.add_star([0,1,2,3,4])
    G.add_path([4,5,6,7,8])
    G.add_star([8,9,10,11,12])
    G.add_path([6,13,14,15])
    G.add_star([15,16,17,18])
    return G

def test():
    adjlist = example_adjlist()

    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])

    print "Testing recursive cut"
    cuts, parts = part_graph(adjlist, 3, recursive=True, dbglvl=METIS_DBG_ALL)
    assert cuts == 2
    assert set(parts) == set([0,1,2])

    if networkx:
        print "Testing with NetworkX"
        G = example_networkx()
        cuts, parts = part_graph(G, 3)
        assert cuts == 2
        assert set(parts) == set([0,1,2])

    print "METIS appears to be working."    

if __name__ == '__main__':
    test()