chempartlib2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from chempartlib2 import *\n",
    "from networkit import *\n",
    "from betterRepair import repairPartition\n",
    "import math, sys, subprocess, functools, operator, time, random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def spiralLayout(G, k, rowheight = 10, colwidth = 10):\n",
    "    \"\"\"\n",
    "    Return two lists, of x and y coordinates for a spiral layout of G.\n",
    "\n",
    "    k nodes are put in one row, keywords rowheight and colwidth determine spacing\n",
    "    \"\"\"\n",
    "    n = G.numberOfNodes()\n",
    "    z = G.upperNodeIdBound()\n",
    "    x = [0 for i in range(z)]\n",
    "    y = [0 for i in range(z)]\n",
    "    for i in range(z):\n",
    "        if G.hasNode(i):\n",
    "            if int(i / k) % 2 > 0:\n",
    "                x[i] = colwidth*(k-(i % k)-1)\n",
    "            else:\n",
    "                x[i] = colwidth*(i % k)\n",
    "            \n",
    "            y[i] = rowheight*int(i / k)\n",
    "            \n",
    "            # adapt coordinates for rounded bends\n",
    "            \n",
    "            ydelta = int(rowheight / 4)\n",
    "            xdelta = colwidth*(1-math.cos(math.pi/3))\n",
    "            rightwards = int(i / k) % 2 == 0\n",
    "    \n",
    "            if i % k == k-1:\n",
    "                y[i] += ydelta\n",
    "                x[i] = x[i] - xdelta if rightwards else x[i] + xdelta\n",
    "            if i > 0 and i % k == 0:\n",
    "                y[i] -= ydelta\n",
    "                x[i] = x[i] - xdelta if not rightwards else x[i] + xdelta\n",
    "        \n",
    "    for i in range(z):\n",
    "        x[i] += 1# gephi ignores coordinates with value 0\n",
    "        y[i] += 1\n",
    "    return x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def naivePartition(G, k):\n",
    "    \"\"\"\n",
    "    Chop a new fragment off G every n/k nodes\n",
    "    \"\"\"\n",
    "    n = G.numberOfNodes()\n",
    "    naivePart = partitioning.Partition(n)\n",
    "    naivePart.allToSingletons()\n",
    "    for i in range(n):\n",
    "        naivePart.moveToSubset(int(i/math.ceil(n/k)), i)\n",
    "    naivePart.compact()\n",
    "    return naivePart"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def greedyPartition(G, k, imbalance, isCharged=[]):\n",
    "    \"\"\"\n",
    "    Starting with singleton clusters, greedily merge the heaviest edge as long as it is smaller than sizelimit.\n",
    "    \"\"\"\n",
    "    n = G.numberOfNodes()\n",
    "    if len(isCharged) > 0:\n",
    "        assert(len(isCharged)==n)\n",
    "    else:\n",
    "        isCharged = [False for i in range(n)]\n",
    "    n = G.numberOfNodes()\n",
    "    part = partitioning.Partition(n)\n",
    "    part.allToSingletons()\n",
    "    chargedPartitions = set([part.subsetOf(i) for i in range(n) if isCharged[i]])\n",
    "    sizelimit = int(math.ceil(n / k)*(1+imbalance))\n",
    "    remainingFragments = n\n",
    "\n",
    "    \n",
    "    def getWeight(edge):\n",
    "        return G.weight(edge[0], edge[1])\n",
    "    \n",
    "    sortedEdges = sorted(G.edges(), key=getWeight)\n",
    "    \n",
    "    # merge heaviest edge, as long as allowed\n",
    "    while len(sortedEdges) > 0:\n",
    "        allowed = True\n",
    "        heaviestEdge = sortedEdges.pop()\n",
    "        firstPart = part.subsetOf(heaviestEdge[0])\n",
    "        secondPart = part.subsetOf(heaviestEdge[1])\n",
    "        if firstPart in chargedPartitions and secondPart in chargedPartitions:\n",
    "            allowed = False\n",
    "        sizeMap = part.subsetSizeMap()\n",
    "        if sizeMap[firstPart] + sizeMap[secondPart] > sizelimit:\n",
    "            allowed = False\n",
    "        partSet = {firstPart, secondPart}\n",
    "        for i in range(n-2):\n",
    "            if part[i] in partSet and part[i+2] in partSet and not part[i+1] in partSet:\n",
    "                allowed = False #otherwise, would create single embedded node\n",
    "        if allowed:\n",
    "            part.mergeSubsets(firstPart, secondPart)\n",
    "            remainingFragments -= 1\n",
    "            if firstPart in chargedPartitions or secondPart in chargedPartitions:\n",
    "                chargedPartitions.add(part.subsetOf(heaviestEdge[0]))\n",
    "    \n",
    "    part.compact()\n",
    "    return part"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def mlPartition(G, k, imbalance, isCharged=[]):\n",
    "    \"\"\"\n",
    "    Use a multi-level approach with Fiduccia-Matheyses to partition G.\n",
    "\n",
    "    Subsets have size at most (1+imbalance)*ceil(n/k)\n",
    "    \"\"\"\n",
    "    n = G.numberOfNodes()\n",
    "    if len(isCharged) > 0:\n",
    "        assert(len(isCharged)==n)\n",
    "        if k > 0:\n",
    "            assert(sum(isCharged) <= k)\n",
    "    else:\n",
    "        isCharged = [False for i in range(n)]\n",
    "    \n",
    "    listOfChargedNodes = [i for i in range(n) if isCharged[i]]\n",
    "    greedy = greedyPartition(G, k, imbalance, isCharged)\n",
    "    try:\n",
    "        dynamic = continuousPartition(G, k, imbalance, isCharged)\n",
    "        if partitioning.computeEdgeCut(greedy, G) < partitioning.computeEdgeCut(dynamic, G):\n",
    "            initial = greedy\n",
    "        else:\n",
    "            initial = dynamic\n",
    "    except ValueError:\n",
    "        initial = greedy\n",
    "    # Problem: The single node repair in C++ may create invalid partitions.\n",
    "    # Better to not use it and repair in Python.\n",
    "    mlp = partitioning.MultiLevelPartitioner(G, k, imbalance, False, listOfChargedNodes, False, initial)\n",
    "    mlp.run()\n",
    "    return mlp.getPartition()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def exportToGephi(G, xcoords, ycoords, part):\n",
    "    \"\"\"\n",
    "    Export graph to Gephi, along with coordinates and partition\n",
    "    \"\"\"\n",
    "    client = gephi.streaming.GephiStreamingClient()\n",
    "    client.clearGraph()\n",
    "    client.exportGraph(G)\n",
    "    client.exportNodeValues(G, part, \"partition\")\n",
    "    client.exportNodeValues(G, xcoords, 'x')\n",
    "    client.exportNodeValues(G, [-elem for elem in ycoords], 'y')\n",
    "    client.exportEdgeValues(G, [G.weight(u,v) for u,v in G.edges()], 'Weight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def moveAllowed(G, partition, v, targetBlock, maxBlockSize = 0, isCharged = []):\n",
    "    \"\"\"\n",
    "    Returns False/True, depending on whether it is allowed to move node v to block targetBlock without violating the charge and gap constraints\n",
    "    \"\"\"\n",
    "    \n",
    "    z = G.upperNodeIdBound()\n",
    "    n = G.numberOfNodes()\n",
    "    assert(n <= z)\n",
    "    assert(len(partition) == z)\n",
    "    assert(G.hasNode(v))\n",
    "    assert(len(isCharged) == 0 or len(isCharged) == z)\n",
    "    \n",
    "    if maxBlockSize == 0:\n",
    "        maxBlockSize = n\n",
    "        \n",
    "    targetBlockSize = sum([partition[u] == targetBlock for u in G.nodes()])\n",
    "    \n",
    "    if len(isCharged) == 0:\n",
    "        isCharged = [False for i in range(z)]\n",
    "    \n",
    "    # move would be forbidden if node and target partition are already charged\n",
    "    if isCharged[v] and any([isCharged[u] for u in G.nodes() if partition[u] == targetBlock and u != v]):\n",
    "        return False\n",
    "    \n",
    "    # move would also be forbidden if gap in main chain of size 1 is created\n",
    "    if v-2 >= 0 and G.hasNode(v-2) and partition[v-2] == targetBlock and G.hasNode(v-1) and partition[v-1] != targetBlock:\n",
    "        return False\n",
    "    \n",
    "    if G.hasNode(v+2) and partition[v+2] == targetBlock and G.hasNode(v+1) and partition[v+1] != targetBlock:\n",
    "        return False\n",
    "    \n",
    "    # move is forbidden if targetBlock is already at maximum capacity\n",
    "    if targetBlockSize > maxBlockSize or (targetBlockSize == maxBlockSize and partition[v] != targetBlock):\n",
    "        return False\n",
    "    \n",
    "    # move is allowed!\n",
    "    return True\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def getCutWeights(G, part, v):\n",
    "    \"\"\"\n",
    "    Returns a dictionary with block numbers as keys and cut weights as values\n",
    "    \"\"\"\n",
    "    n = G.numberOfNodes()\n",
    "    assert(len(part) == n)\n",
    "    assert(v < n)\n",
    "    \n",
    "    sums = {}\n",
    "    \n",
    "    for block in set(part):\n",
    "        sums[block] = 0\n",
    "        \n",
    "    for u in G.neighbors(v):\n",
    "        sums[part[v]] += G.weight(u,v)\n",
    "        \n",
    "    return sums"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def repairPartition(G, partition, imbalance = 0.2, isCharged = []):\n",
    "    n = G.numberOfNodes()\n",
    "\n",
    "    fragmentSet = set(partition)\n",
    "    k = len(fragmentSet)\n",
    "\n",
    "    #TODO: what happens if the partition contains large fragment IDs?\n",
    "    fragmentSizes = [0 for f in fragmentSet]\n",
    "    fragmentCharges = [[] for f in fragmentSet]\n",
    "    edgeCuts = [[0 for f in fragmentSet] for v in G.nodes()]\n",
    "\n",
    "    gapsFound = False\n",
    "\n",
    "    for v in G.nodes():\n",
    "        fragmentSizes[partition[v]] += 1\n",
    "        if isCharged[v]:\n",
    "            fragmentCharges[partition[v]].append(v)\n",
    "        if G.hasNode(v+2) and partition[v+2] == partition[v] and G.hasNode(v+1) and partition[v+1] != partition[v]:\n",
    "            gapsFound = True\n",
    "\n",
    "        for u in G.neighbors(v):\n",
    "            edgeCuts[v][partition[u]] += G.weight(v, u)\n",
    "\n",
    "    maxBlockSize = int(math.ceil(n / k)*(1+imbalance))\n",
    "    if max(fragmentSizes) <= maxBlockSize and max([len(group) for group in fragmentCharges]) <= 1 and not gapsFound:\n",
    "        return partition\n",
    "\n",
    "    maxGain = [- float('inf') for v in G.nodes()]\n",
    "    maxTarget = [-1 for v in G.nodes()]\n",
    "    heap = []\n",
    "\n",
    "    def allowed(v, target):\n",
    "        allowed = True\n",
    "        if isCharged[v] and ((len(fragmentCharges[target]) > 0 and v not in fragmentCharges[target]) or len(fragmentCharges[target]) > 1):\n",
    "            allowed = False\n",
    "        if fragmentSizes[target] > maxBlockSize:\n",
    "            allowed = False\n",
    "        if fragmentSizes[target] == maxBlockSize and partition[v] != target:\n",
    "            allowed = False\n",
    "        if G.hasNode(v+2) and part[v+2] == target and part[v+1] != target:\n",
    "            allowed = False\n",
    "        if v > 0 and G.hasNode(v-1) and G.hasNode(v+1) and part[v-1] == part[v+1] and target != part[v-1]:\n",
    "            allowed = False\n",
    "        return allowed\n",
    "\n",
    "\n",
    "    for v in G.nodes():\n",
    "        for neighbor in G.neighbors(v):\n",
    "            target = partition[neighbor]# actually, I only need to iterate over the partitions. This is more for an easier asymptotic running time\n",
    "\n",
    "            if allowed(v, target) and edgeCuts[v][target] - edgeCuts[v][partition[v]] > maxGain[v]:\n",
    "                maxGain[v] = edgeCuts[v][target] - edgeCuts[v][partition[v]]\n",
    "                maxTarget[v] = target\n",
    "\n",
    "        heap.heappush((maxGain[v], v))\n",
    "\n",
    "    visited = [False for v in range(n)]\n",
    "\n",
    "    while len(heap) > 0:\n",
    "        (key, v) = heappop()\n",
    "        fragment = partition[v]\n",
    "\n",
    "        visited[v] = True\n",
    "\n",
    "        def gapAt(v):\n",
    "            return v > 0 and G.hasNode(v) and G.hasNode(v+2) and partition[v] == partition[v+2] and partition[v] != partition[v+1]\n",
    "\n",
    "        # if fragment of v is alright, skip node\n",
    "        if fragmentSizes[fragment] <= maxBlockSize and (not isCharged[v] or len(fragmentCharges[fragment]) <= 1) and not gapAt(v-2) and not gapAt(v-1) and not gapAt(v):\n",
    "            continue\n",
    "\n",
    "        if key == - math.inf:\n",
    "            raise ValueError(\"Partition could not be repaired.\")\n",
    "\n",
    "        # otherwise, move v to best allowed fragment and update data structures\n",
    "        fragmentSizes[partition[v]] -= 1\n",
    "        fragmentSizes[maxTarget[v]] += 1\n",
    "\n",
    "        if isCharged[v]:\n",
    "            fragmentCharges[partition[v]].remove(v)\n",
    "            fragmentCharges[maxTarget[v]].append(v)\n",
    "\n",
    "        for neighbor in G.neighbors(v):\n",
    "            edgeCuts[neighbor][partition[v]] -= G.weight(neighbor, v)\n",
    "            edgeCuts[neighbor][maxTarget[v]] += G.weight(neighbor, v)\n",
    "\n",
    "        partition[v] = maxTarget[v]\n",
    "\n",
    "        # update max gains and queue positions of neighbors\n",
    "        for neighbor in G.neighbors(v):\n",
    "            if visited[neighbor]:\n",
    "                continue\n",
    "\n",
    "            oldKey = maxGain[neighbor]\n",
    "            for target in fragmentSet:\n",
    "                if allowed(neighbor, target) and edgeCuts[neighbor][target] - edgeCuts[neighbor][partition[neighbor]] > maxGain[neighbor]:\n",
    "                    maxGain[neighbor] = edgeCuts[neighbor][target] - edgeCuts[neighbor][partition[neighbor]]\n",
    "                    maxTarget[neighbor] = target\n",
    "\n",
    "            if maxGain[neighbor] != oldKey:\n",
    "                heap.remove((oldKey, neighbor))\n",
    "                heap.heapify()\n",
    "                heap.heappush((maxGain[neighbor], neighbor))\n",
    "\n",
    "    return partition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def kaHiPWrapper(G, k, imbalance = 0.2, pathToKaHiP = '/home/moritzl/Gadgets/KaHIP/deploy/kaffpa', multiple=False):\n",
    "    \"\"\"\n",
    "    Calls KaHiP, an external partitioner.\n",
    "    \"\"\"\n",
    "    \n",
    "    tempFileName = 'tempForKaHiP.graph'\n",
    "    outputFileName = 'tmppartition'+str(k)\n",
    "    n = G.numberOfNodes()\n",
    "    \n",
    "    maxWeight = max([G.weight(u,v) for (u,v) in G.edges()])\n",
    "    \n",
    "    \"\"\"\n",
    "    KaHiP only accepts integer weights, thus we scale and round them.\n",
    "    Weights must be under 1 million, otherwise the METIS graph writer switches to scientific notation,\n",
    "    which confuses KaHiP\n",
    "    \"\"\"\n",
    "    scalingFactor = int((10**6-1)/maxWeight)\n",
    "    \n",
    "    \n",
    "    #copy and scale graph\n",
    "    Gscaled = G.copyNodes()\n",
    "    for (u,v) in G.edges():\n",
    "        Gscaled.addEdge(u,v,int(G.weight(u,v)*scalingFactor))\n",
    "    \n",
    "    # write out temporary file\n",
    "    writeGraph(Gscaled, tempFileName, Format.METIS)\n",
    "    \n",
    "    # call KaHIP\n",
    "    callList = [pathToKaHiP, '--k='+str(k), '--imbalance='+str(int(imbalance*100)), '--preconfiguration=strong']\n",
    "    if multiple:\n",
    "        callList.append('--time_limit=1')\n",
    "    callList.append(tempFileName)\n",
    "    subprocess.call(callList)\n",
    "    \n",
    "    # read in partition\n",
    "    part = community.PartitionReader().read(outputFileName)\n",
    "    \n",
    "    # remove temporary files\n",
    "    subprocess.call(['rm', tempFileName])\n",
    "    subprocess.call(['rm', outputFileName])\n",
    "    \n",
    "    return part"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def partitionValid(G, partition, maxBlockSize = 0, isCharged = []):\n",
    "    z = G.upperNodeIdBound()\n",
    "    n = G.numberOfNodes()\n",
    "    if len(partition) != z:\n",
    "        return False\n",
    "    \n",
    "    if len(isCharged) != 0 and len(isCharged) != z:\n",
    "        return False\n",
    "    \n",
    "    if len(isCharged) == 0:\n",
    "        isCharged = [False for i in range(z)]\n",
    "        \n",
    "    if maxBlockSize == 0:\n",
    "        maxBlockSize = n\n",
    "        \n",
    "    chargedFragments = set()\n",
    "    \n",
    "    fragmentSizes = {}\n",
    "    \n",
    "    for v in range(G.numberOfNodes()):\n",
    "        if not G.hasNode(v):\n",
    "            print(\"Node \", v, \" not in graph.\")\n",
    "            return False\n",
    "            \n",
    "        # partition invalid if two charged nodes in same fragment\n",
    "        if isCharged[v]:\n",
    "            if partition[v] in chargedFragments:\n",
    "                print(\"Node\", v, \" is charged, but fragment\", partition[v], \"already has a charged node.\")\n",
    "                return False\n",
    "            else:\n",
    "                chargedFragments.add(partition[v])\n",
    "    \n",
    "        # partition also invalid if gaps of size 1 exist\n",
    "        if G.hasNode(v+2) and partition[v+2] == partition[v] and G.hasNode(v+1) and partition[v+1] != partition[v]:\n",
    "            print(\"Nodes\", v, \"and\", v+2, \"are in fragment\", partition[v], \"but\", v+1, \"is in fragment\", partition[v+1])\n",
    "            return False\n",
    "    \n",
    "        # partition invalid if fragment is larger than allowed\n",
    "        if not partition[v] in fragmentSizes:\n",
    "            fragmentSizes[partition[v]] = 1\n",
    "        else:\n",
    "            fragmentSizes[partition[v]] += 1\n",
    "        if fragmentSizes[partition[v]] > maxBlockSize:\n",
    "            print(\"Fragment\", partition[v], \"contains\", fragmentSizes[partition[v]], \"nodes, more than\", maxBlockSize)\n",
    "            return False\n",
    "    \n",
    "    # no reason to complain found, partition is valid\n",
    "    return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def comparePartitionQuality(G, k, imbalance, chargedNodes = set(), silent=False):\n",
    "    n = G.numberOfNodes()\n",
    "    \n",
    "    isCharged = [v in chargedNodes for v in range(G.numberOfNodes())]\n",
    "    sizelimit = int(math.ceil(n / k)*(1+imbalance))\n",
    "    if not silent:\n",
    "        print(\"Size limit:\", sizelimit)\n",
    "    result = {}\n",
    "    \n",
    "    before = time.time()\n",
    "    ml = mlPartition(G, k, imbalance, isCharged)\n",
    "    timeML = time.time() - before\n",
    "    if not silent:\n",
    "        print(\"MultiLevel:\", partitioning.computeEdgeCut(ml, G))\n",
    "        print(\"Time:\", timeML)\n",
    "    if not partitionValid(G, ml, sizelimit, isCharged):\n",
    "        ml = repairPartition(G, ml, imbalance, isCharged)\n",
    "        if not silent:\n",
    "            print(\"Repaired Multilevel:\", partitioning.computeEdgeCut(ml, G))\n",
    "            partitionValid(G, ml, sizelimit, isCharged)\n",
    "    if not silent:\n",
    "        print(\"Effective k\", str(ml.numberOfSubsets()))\n",
    "        print()\n",
    "    result['ml'] = partitioning.computeEdgeCut(ml, G)\n",
    "    \n",
    "    before = time.time()\n",
    "    greedy = greedyPartition(G, k, imbalance, isCharged)\n",
    "    timeGreedy = time.time() - before\n",
    "    if not silent:\n",
    "        print(\"Greedy:\", partitioning.computeEdgeCut(greedy, G))\n",
    "        print(\"Time:\", timeGreedy)\n",
    "    if not partitionValid(G, greedy, sizelimit, isCharged):\n",
    "        greedy = repairPartition(G, greedy, imbalance, isCharged)\n",
    "        if not silent:\n",
    "            print(\"Repaired Greedy:\", partitioning.computeEdgeCut(greedy, G))\n",
    "        assert(partitionValid(G, greedy, sizelimit, isCharged))\n",
    "    if not silent:\n",
    "        print(\"Effective k\", str(greedy.numberOfSubsets()))\n",
    "        print()\n",
    "    result['greedy'] = partitioning.computeEdgeCut(greedy, G)\n",
    "    \n",
    "    X = int(n / k)\n",
    "    tolerance = int(math.ceil(n / k)*(1+imbalance)) - X\n",
    "    \n",
    "    try:\n",
    "        before = time.time()\n",
    "        cont = continuousPartition(G, k, imbalance, isCharged)\n",
    "        timeDP = time.time() - before\n",
    "        if not silent:\n",
    "            print(\"Dynamic Programming:\", partitioning.computeEdgeCut(cont, G))\n",
    "            print(\"Time:\", timeDP)\n",
    "        if not partitionValid(G, cont, sizelimit, isCharged):\n",
    "            cont = repairPartition(G, cont, imbalance, isCharged)\n",
    "            if not silent:\n",
    "                print(\"Repaired Dynamic:\", partitioning.computeEdgeCut(cont, G))\n",
    "            assert(partitionValid(G, cont, sizelimit, isCharged))\n",
    "        result['cont'] = partitioning.computeEdgeCut(cont, G)\n",
    "        if not silent:\n",
    "            print(\"Effective k\", str(cont.numberOfSubsets()))\n",
    "    except ValueError as e:\n",
    "        print(e)\n",
    "\n",
    "    print()\n",
    "        \n",
    "    before = time.time()\n",
    "    ka = kaHiPWrapper(G, k, imbalance)\n",
    "    timeKa = time.time() - before\n",
    "    if not silent:\n",
    "        print(\"Raw KaHIP:\", partitioning.computeEdgeCut(ka, G))\n",
    "        print(\"Time:\", timeKa)\n",
    "    if not partitionValid(G, ka, sizelimit, isCharged):\n",
    "        ka = repairPartition(G, ka, imbalance, isCharged)\n",
    "        if not silent:\n",
    "            print(\"Repaired KaHiP:\", partitioning.computeEdgeCut(ka, G))\n",
    "            partitionValid(G, ka, sizelimit, isCharged)\n",
    "    if not silent:\n",
    "        print(\"Effective k\", str(ka.numberOfSubsets()))\n",
    "        print()\n",
    "    result['ka'] = partitioning.computeEdgeCut(ka, G)\n",
    "    result['bestOfFour'] = min([result[key] for key in result])\n",
    "\n",
    "\n",
    "    before = time.time()\n",
    "    naive = naivePartition(G, k)\n",
    "    timeNaive = time.time() - before\n",
    "    if not silent:\n",
    "        print(\"Naive:\", partitioning.computeEdgeCut(naive, G))\n",
    "        print(\"Time:\", timeNaive)\n",
    "    if not partitionValid(G, naive, sizelimit, isCharged):\n",
    "        naive = repairPartition(G, naive, imbalance, isCharged)        \n",
    "        if not silent:\n",
    "            print(\"Repaired Naive:\", partitioning.computeEdgeCut(naive, G))\n",
    "        assert(partitionValid(G, naive, sizelimit, isCharged))\n",
    "    if not silent:\n",
    "        print(\"Effective k\", str(naive.numberOfSubsets()))\n",
    "    result['naive'] = partitioning.computeEdgeCut(naive, G)\n",
    "    result['bestOfFive'] = min([result[key] for key in result])\n",
    "\n",
    "    if not silent:\n",
    "        print(str(result['bestOfFour'] / result['naive']))\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def chargesValid(G, klist, minEpsilon, isCharged):\n",
    "    assert(type(klist) is list)\n",
    "    assert(type(minEpsilon) is float)\n",
    "    assert(type(isCharged) is list)\n",
    "    assert(len(isCharged) == G.numberOfNodes())\n",
    "    for k in klist:\n",
    "        try:\n",
    "            part = continuousPartition(G, k, minEpsilon, isCharged)\n",
    "        except ValueError as e:\n",
    "            return False\n",
    "    return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def readCharges(path):\n",
    "    \"\"\"\n",
    "    Reads file at path, returns a list of charged nodes\n",
    "    \"\"\"\n",
    "    chargedNodes = []\n",
    "    \n",
    "    with open(path, 'r') as f:\n",
    "        for line in f:\n",
    "            chargedNodes.append(int(line)-1)\n",
    "    \n",
    "    return chargedNodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def main(numberOfChargedNodes = 5):\n",
    "    pathPrefix = \"/home/moritzl/ChemPart/\"\n",
    "    graphSuffix = \"_complete.graph\"\n",
    "    chargeSuffix = \"_charges.resid\"\n",
    "    Gnames = [\"ubiquitin\", \"bubble\", \"br\", \"fmo\", \"gfp\"]\n",
    "    #kList = [2**i for i in range(3,6)]\n",
    "    scores = []\n",
    "    epsilon = 0.2\n",
    "    initialTime = time.time()\n",
    "    \n",
    "    for Gname in Gnames:\n",
    "        G = readGraph(pathPrefix + Gname + graphSuffix, Format.METIS)\n",
    "        chargedNodes = []#readCharges(pathPrefix + Gname + chargeSuffix)\n",
    "        n = G.numberOfNodes()\n",
    "        graphScores = []\n",
    "        \n",
    "        if n > 100:\n",
    "            kList = [8,12,16,20,24]\n",
    "        else:\n",
    "            kList = [2,4,6,8]\n",
    "        \n",
    "        print(\"Graph:\", Gname, \"with\", n, \" nodes.\")\n",
    "        print(\"chargedNodes =\", chargedNodes)\n",
    "        for k in kList:\n",
    "            if len(chargedNodes) > k:\n",
    "                continue\n",
    "            print(\"k = \", k)\n",
    "            qualities = comparePartitionQuality(G, k, epsilon, chargedNodes, False)\n",
    "            scores.append(qualities)\n",
    "            graphScores.append(qualities)\n",
    "            print('------------------------------------------------------------------')\n",
    "        graphResults = {}\n",
    "        for score in graphScores:\n",
    "            for key in score:\n",
    "                if not key in graphResults:\n",
    "                    graphResults[key] = []\n",
    "                graphResults[key].append(score[key])\n",
    "                \n",
    "        gMeans = {}\n",
    "        for method in graphResults:\n",
    "            gMeans[method] = functools.reduce(operator.mul, graphResults[method], 1) ** (1/len(graphResults[method]))\n",
    "            print(method, ':', str(gMeans[method]))\n",
    "        print(\"Ratio:\", gMeans['bestOfFour'] / gMeans['naive'])\n",
    "        print('##################################################################')\n",
    "    \n",
    "    results = {}\n",
    "    for score in scores:\n",
    "        for key in score:\n",
    "            if not key in results:\n",
    "                results[key] = []\n",
    "            results[key].append(score[key])\n",
    "            \n",
    "    print(\"Elapsed Time:\", time.time() - initialTime)\n",
    "            \n",
    "    print(\"Geometric Means:\")\n",
    "    for method in results:\n",
    "        print(method, ':', str(functools.reduce(operator.mul, results[method], 1) ** (1/len(results[method]))))\n",
    "        \n",
    "    print(\"Arithmetic Means:\")\n",
    "    for method in results:\n",
    "        print(method, ':', str(sum(results[method]) / len(results[method])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def runAndLogExperiments(runs = 1, charges = False, epsilonList=[0.1,0.2]):\n",
    "    pathPrefix = \"/home/moritzl/ChemPart/\"\n",
    "    graphSuffix = \"_complete.graph\"\n",
    "    chargeSuffix = \"_charges.resid\"\n",
    "    algoList = ['ml', 'greedy', 'ka', 'naive', 'cont']\n",
    "    Gnames = [\"ubiquitin\", \"bubble\", \"br\", \"fmo\", \"gfp\", \"fmo\"]\n",
    "    #kList = [2**i for i in range(1,6)]\n",
    "    maxIterations = 100\n",
    "    \n",
    "    for Gname in Gnames:\n",
    "        G = readGraph(pathPrefix + Gname + graphSuffix, Format.METIS)\n",
    "        potentiallyCharged = readCharges(pathPrefix + Gname + chargeSuffix)\n",
    "        \n",
    "        n = G.numberOfNodes()\n",
    "        kList = []\n",
    "        if n > 100:\n",
    "            kList = [8,12,16,20,24]\n",
    "        else:\n",
    "            kList = [2,4,6,8]\n",
    "        \n",
    "        for run in range(runs):\n",
    "            with open(Gname+'-results-'+str(run)+'.dat', 'w') as f:\n",
    "                f.write('\\t'.join(['k']+[str(e) for e in epsilonList]+['label\\n']))\n",
    "\n",
    "                for k in kList:\n",
    "                    chargedNodes = []\n",
    "                    if charges:\n",
    "                        valid = False\n",
    "                        i = 0\n",
    "\n",
    "                        while not valid and i < maxIterations:\n",
    "                            chargedNodes = random.sample(potentiallyCharged, int(k*0.8))\n",
    "                            isCharged = [v in chargedNodes for v in G.nodes()]\n",
    "                            valid = chargesValid(G, [k], min(epsilonList), isCharged)\n",
    "                            i += 1\n",
    "                        if not valid:\n",
    "                            print(\"No valid charges found after \"+str(i)+\" iterations.\")\n",
    "                            continue\n",
    "\n",
    "                        print(chargedNodes)\n",
    "\n",
    "                    # data format: one line per algorithm and k\n",
    "                    lineDict = {}\n",
    "                    for algo in algoList:\n",
    "                        lineDict[algo] = [str(k)]\n",
    "\n",
    "                    for epsilon in epsilonList:\n",
    "                        qualities = comparePartitionQuality(G, k, epsilon, chargedNodes, True)\n",
    "                        for algo in algoList:\n",
    "                            if algo in qualities:\n",
    "                                lineDict[algo].append(str(qualities[algo]/qualities['naive']))\n",
    "                            else:\n",
    "                                lineDict[algo].append('NA')\n",
    "                        print('Experiments done for k=', k, ', epsilon=', epsilon)\n",
    "\n",
    "                    for i in range(len(algoList)):\n",
    "                        algo = algoList[i]\n",
    "                        lineDict[algo].append(str(i))\n",
    "                        f.write('\\t'.join(lineDict[algo])+'\\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def averageLogs(runs, Gnames = [\"ubiquitin\", \"bubble\", \"br\", \"fmo\", \"gfp\"]):\n",
    "    for Gname in Gnames:\n",
    "\n",
    "        sumEntries = []\n",
    "        numEntries = []\n",
    "\n",
    "        for run in range(runs):\n",
    "            filename = Gname+'-results-'+str(run)+'.dat'\n",
    "            with open(filename, 'r') as f:\n",
    "                f.readline()# remove header data\n",
    "                lineNumber = 0\n",
    "                for line in f:\n",
    "                    lineList = line.split('\\t')\n",
    "\n",
    "                    if len(sumEntries) < lineNumber+1:\n",
    "                        sumEntries.append([0 for field in lineList])\n",
    "                        numEntries.append([0 for field in lineList])\n",
    "\n",
    "                    for i in range(len(sumEntries[lineNumber]), len(lineList)):\n",
    "                        sumEntries[lineNumber].append(0)\n",
    "                        numEntries[lineNumber].append(0)\n",
    "\n",
    "                    for i in range(len(lineList)):\n",
    "                        sumEntries[lineNumber][i] +=  float(lineList[i])\n",
    "                        numEntries[lineNumber][i] += 1\n",
    "\n",
    "                    lineNumber += 1\n",
    "        assert(len(sumEntries) == len(numEntries))\n",
    "\n",
    "        outputname = Gname+'-results-averaged.dat'\n",
    "        with open(outputname, 'w') as f:\n",
    "            for rowIndex in range(len(sumEntries)):\n",
    "                assert(len(sumEntries[rowIndex]) == len(numEntries[rowIndex]))\n",
    "                linelist = [str(sumEntries[rowIndex][colIndex] / numEntries[rowIndex][colIndex]) for colIndex in range(len(sumEntries[rowIndex])) ]\n",
    "                f.write('\\t'.join(linelist)+'\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "setLogLevel('DEBUG')\n",
    "runAndLogExperiments(1, False, [0.3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "averageLogs(20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    # execute only if run as a script\n",
    "    main(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "pathPrefix = \"/home/moritzl/ChemPart/\"\n",
    "Gnames = [\"ubiquitin_complete\", \"bubble_complete\", \"br_complete\", \"fmo_complete\", \"gfp_complete\"]#, \"ubiquitin\", \"bubble\", \"bubble_new\"]\n",
    "\n",
    "kList = [2**i for i in range(1,7)]\n",
    "epsilon = 0.2\n",
    "\n",
    "G = readGraph(pathPrefix + Gnames[0] + \".graph\", Format.METIS)\n",
    "chargedNodes = set()\n",
    "for v in range(1):\n",
    "    chargedNodes.add(G.randomNode())\n",
    "isCharged = [v in chargedNodes for v in range(G.numberOfNodes())]\n",
    "\n",
    "chargesValid(G, kList, epsilon, isCharged)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# compare times for ml partitioning \n",
    "#setLogLevel('DEBUG')\n",
    "pathPrefix = \"/home/moritzl/ChemPart/\"\n",
    "Gname = \"fmo_complete\"\n",
    "k = 8\n",
    "epsilon = 0.2\n",
    "G = readGraph(pathPrefix + Gname + \".graph\", Format.METIS)\n",
    "isCharged = [False for v in G.nodes()]\n",
    "before = time.time()\n",
    "mlp = partitioning.MultiLevelPartitioner(G, k, epsilon, False, [])\n",
    "mlp.run()\n",
    "ml = mlp.getPartition()\n",
    "rawtime = time.time() - before\n",
    "initial = continuousPartition(G, k, epsilon, isCharged)\n",
    "before = time.time()\n",
    "mlp = partitioning.MultiLevelPartitioner(G, k, epsilon, False, [], False, initial)\n",
    "mlp.run()\n",
    "ml = mlp.getPartition()\n",
    "timeWithinit = time.time() - before\n",
    "print(rawtime, timeWithinit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "for k in kList:\n",
    "    ml = mlPartition(G, k, epsilon)\n",
    "    print(\"MultiLevel\", k, \":\", partitioning.computeEdgeCut(ml, G))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ml = mlPartition(G, 2, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "k = 8\n",
    "epsilon = 0.2\n",
    "Gname = Gnames[1]\n",
    "G = readGraph(pathPrefix + Gname + \".graph\", Format.METIS)\n",
    "\n",
    "ml = mlPartition(G, k, epsilon)\n",
    "ka = kaHiPWrapper(G, k, epsilon)\n",
    "ml = repairPartition(G, ml, epsilon)\n",
    "ka = repairPartition(G, ka, epsilon)\n",
    "cont = continuousPartition(G, k, epsilon)\n",
    "naive = naivePartition(G, k)\n",
    "ml.compact()\n",
    "ka.compact()\n",
    "cont.compact()\n",
    "\n",
    "# write partitions\n",
    "community.PartitionWriter().write(ml, 'MultiLevel-k-'+str(k)+'-imbalance-'+str(epsilon)+'-'+Gname+'.part')\n",
    "community.PartitionWriter().write(ka, 'KaHiP-k-'+str(k)+'-imbalance-'+str(epsilon)+'-'+Gname+'.part')\n",
    "community.PartitionWriter().write(cont, 'DP-k-'+str(k)+'-imbalance-'+str(epsilon)+'-'+Gname+'.part')\n",
    "community.PartitionWriter().write(naive, 'Naive-k-'+str(k)+'-'+Gname+'.part')\n",
    "\n",
    "# prepare coordinates for gephi\n",
    "xcoords, ycoords = spiralLayout(G, k)\n",
    "\n",
    "# remove light edges from graph\n",
    "edgelist = G.edges()\n",
    "cutoff = 0.05\n",
    "for a,b in edgelist:\n",
    "    if G.weight(a,b) < cutoff:\n",
    "        G.removeEdge(a,b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(\"ML:\")\n",
    "partitioning.inspectPartitions(ml, G)\n",
    "        \n",
    "#exportToGephi(G, xcoords, ycoords, ml)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(\"Ka:\")\n",
    "partitioning.inspectPartitions(ka, G)\n",
    "\n",
    "#exportToGephi(G, xcoords, ycoords, ka)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "print(\"DP:\")\n",
    "partitioning.inspectPartitions(cont, G)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(\"Na:\")\n",
    "partitioning.inspectPartitions(naive, G)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "k = 8\n",
    "naive = naivePartition(G, k)\n",
    "partitioning.inspectPartitions(naive, G)\n",
    "gain = partitioning.MultiLevelPartitioner.fiducciaMattheysesStep(G, naive, epsilon, set(), set())\n",
    "partitioning.inspectPartitions(naive, G)\n",
    "print(gain)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.4.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}