def __init__(self): lib = load_library('metis') # Try to determine the data types METIS was compiled with with silence(): self._probe_types(lib) # METIS_SetDefaultOptions self.METIS_SetDefaultOptions = lib.METIS_SetDefaultOptions self.METIS_SetDefaultOptions.argtypes = [c_void_p] self.METIS_SetDefaultOptions.errcheck = self._errcheck # METIS_PartGraphKway self.METIS_PartGraphKway = lib.METIS_PartGraphKway self.METIS_PartGraphKway.argtypes = [ POINTER(self.metis_int), POINTER(self.metis_int), c_void_p, c_void_p, c_void_p, c_void_p, c_void_p, POINTER(self.metis_int), c_void_p, c_void_p, c_void_p, POINTER(self.metis_int), c_void_p ] self.METIS_PartGraphKway.errcheck = self._errcheck # METIS_PartGraphRecursive self.METIS_PartGraphRecursive = lib.METIS_PartGraphRecursive self.METIS_PartGraphRecursive.argtypes = [ POINTER(self.metis_int), POINTER(self.metis_int), c_void_p, c_void_p, c_void_p, c_void_p, c_void_p, POINTER(self.metis_int), c_void_p, c_void_p, c_void_p, POINTER(self.metis_int), c_void_p ] self.METIS_PartGraphRecursive.errcheck = self._errcheck
def _partition_graph(self, graph, partwts): w = self._wrappers # Type conversion vtab = np.asanyarray(graph.vtab, dtype=w.metis_int_np) etab = np.asanyarray(graph.etab, dtype=w.metis_int_np) vwts = np.asanyarray(graph.vwts, dtype=w.metis_int_np) ewts = np.asanyarray(graph.ewts, dtype=w.metis_int_np) partwts = np.array(partwts, dtype=w.metis_flt_np) # Normalise the weights partwts /= np.sum(partwts) # Allocate the partition array parts = np.empty(len(vtab) - 1, dtype=w.metis_int_np) # Allocate our options array opts = np.empty(w.METIS_NOPTIONS, dtype=w.metis_int_np) w.METIS_SetDefaultOptions(opts.ctypes) # Process our options for k, v in self.opts.items(): opts[getattr(w, f'METIS_OPTION_{k.upper()}')] = v # Select the partitioning function if self.opts['ptype'] == self.enum_opts['ptype']['rb']: part_graph_fn = w.METIS_PartGraphRecursive else: part_graph_fn = w.METIS_PartGraphKway # Integer parameters nvert, nconst = w.metis_int(len(vtab) - 1), w.metis_int(1) npart, objval = w.metis_int(len(partwts)), w.metis_int() # Partition with silence(): part_graph_fn( nvert, nconst, vtab.ctypes, etab.ctypes, vwts.ctypes, None, ewts.ctypes, npart, partwts.ctypes, None, opts.ctypes, objval, parts.ctypes ) # Check for invalid partition numbers if np.max(parts) >= len(partwts): raise RuntimeError('Invalid partition number from METIS') return parts