graph problem

gistfile1.json

{
 "metadata": {
  "name": "",
  "signature": "sha256:754cf4f73927f25d89981617f6523263163497ca44833bd7d4afb979a950a727"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "This is an example of high level code that generates a cvxpy problem. The idea is you're creating a single commodity flow problem by specifying the objectives and constraints on the nodes and edges. These objectives and constraints are then automatically gathered together into a problem.\n",
      "\n",
      "Each variable is a scalar, so the parsing/matrix stuffing is quite slow."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%config InlineBackend.figure_format = 'retina'\n",
      "from cvxpy import *\n",
      "import matplotlib.pyplot as plt\n",
      "\n",
      "class Node(object):\n",
      "    def __init__(self, source, cost):\n",
      "        self.source = source\n",
      "        self.cost = cost(self.source)\n",
      "        self.edge_flows = []\n",
      "\n",
      "    def net_flow(self):\n",
      "        \"\"\"The net flow on the node.\"\"\"\n",
      "        return sum(self.edge_flows) + self.source\n",
      "\n",
      "class Edge(object):\n",
      "    def __init__(self, cost):\n",
      "        self.flow = Variable()\n",
      "        self.cost = cost(self.flow)\n",
      "\n",
      "    def connect(self, in_node, out_node):\n",
      "        \"\"\"Connects two nodes via the edge.\"\"\"\n",
      "        in_node.edge_flows.append(-self.flow)\n",
      "        out_node.edge_flows.append(self.flow)\n",
      "\n",
      "import networkx as nx\n",
      "import matplotlib.pyplot as plt\n",
      "import random\n",
      "import numpy as np\n",
      "random.seed(1)\n",
      "m = 20\n",
      "n = 20\n",
      "G = nx.grid_2d_graph(m, n)\n",
      "# print G.edges()\n",
      "#G = nx.DiGraph(G)\n",
      "\n",
      "source = (0, 0)\n",
      "sink = (m-1, n-1)\n",
      "nodes = {k:Node(0, lambda x: 0) for k in G.nodes()}\n",
      "nodes[source] = Node(1, lambda x: 0)\n",
      "nodes[sink] = Node(-1, lambda x: 0)\n",
      "edges = {k:Edge(lambda x: random.random()*(abs(x) + square(x))) for k in G.edges()}\n",
      "for key, edge in edges.items():\n",
      "    edge.connect(nodes[key[0]], nodes[key[1]])\n",
      "    \n",
      "cost = sum([obj.cost for obj in nodes.values() + edges.values()])\n",
      "constraints = [node.net_flow() == 0 for node in nodes.values()]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 3
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%%timeit\n",
      "prob = Problem(Minimize(cost), constraints)\n",
      "prob.get_problem_data(ECOS)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "1 loops, best of 3: 5.93 s per loop\n"
       ]
      }
     ],
     "prompt_number": 4
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# Plot the result.\n",
      "# Looks best with 7 by 7 grid.\n",
      "prob = Problem(Minimize(cost), constraints)\n",
      "prob.solve()\n",
      "pos = dict(zip(G,G)) # dictionary of node names->positions\n",
      "colors = [abs(edges[k].flow).value for k in G.edges()]#range((n-1)*m + (m-1)*n)\n",
      "plt.figure(figsize=(6,4))\n",
      "nodes = nx.draw_networkx_nodes(G,pos,node_color='k', with_labels=False, node_size=100)\n",
      "edges = nx.draw_networkx_edges(G,pos,edge_color=colors,width=4,\n",
      "                               edge_cmap=plt.cm.Blues, arrows=True)\n",
      "plt.colorbar(edges, label='Edge flow')\n",
      "plt.axis('off')\n",
      "%matplotlib inline\n",
      "plt.show()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    }
   ],
   "metadata": {}
  }
 ]
}