DistributedEnsemble.ipynb

{
 "metadata": {
  "name": "",
  "signature": "sha256:9675a1b79143559c04d3f24bb6d1c67ab1c984aebeea677b94a05601aca1846e"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "This notebook shows how to generate ensembles of trajectories, and then in subsequent steps perform statistical analysis on the generated data.  "
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import IPython.parallel\n",
      "import sys\n",
      "from pyurdme import pyurdme\n",
      "import time\n",
      "import pyurdme\n",
      "import numpy\n",
      "import math\n",
      "\n",
      "# This is to disable annoying messages from FeniCS/UFL/FFC\n",
      "import logging\n",
      "logging.disable(logging.DEBUG)\n",
      "\n",
      "rc = IPython.parallel.Client()\n",
      "# We have to use dill instead of pickle to serialize dynamically defined classes\n",
      "rc[:].use_dill()\n",
      "rv = rc.load_balanced_view()\n",
      "\n",
      "# We are going to need to know the ids of the engines at the engines later on\n",
      "dv = rc[:] \n",
      "dv.scatter('engine_id', rc.ids, flatten=True)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 1,
       "text": [
        "<AsyncResult: scatter>"
       ]
      }
     ],
     "prompt_number": 1
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%%px --local\n",
      "import pyurdme\n",
      "class SimpleDiffusion(pyurdme.URDMEModel):\n",
      "    \"\"\" This simple Hello World model consists of a single\n",
      "        species diffusing on a unit square. The initial \n",
      "        condition is a delta function on the voxel closest \n",
      "        to the center of the grid. \n",
      "    \"\"\"\n",
      "    \n",
      "    def __init__(self):\n",
      "        \n",
      "        pyurdme.URDMEModel.__init__(self,name=\"simple_diffusion\")\n",
      "\n",
      "        A = pyurdme.Species(name=\"A\",diffusion_constant=0.01,dimension=2)\n",
      "\n",
      "        self.add_species([A])\n",
      "\n",
      "        # A unit square\n",
      "        self.mesh = pyurdme.URDMEMesh.generate_unit_square_mesh(40,40)\n",
      "                \n",
      "        # Place the A molecules in the voxel nearest the center of the square\n",
      "        self.set_initial_condition_place_near({A:1000},point=[0.5,0.5])\n",
      "\n",
      "        self.timespan(numpy.linspace(0,1,1000))"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 3
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%%px --local\n",
      "import pyurdme\n",
      "import numpy\n",
      "def add(a, b=None):\n",
      "    if b==None:\n",
      "        return a\n",
      "    return a+b"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 4
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def run_ensemble(model,nt,s):\n",
      "    \"\"\" Generates an ensemble consisting of number_of_trajectories realizations. \n",
      "        Returns a list of result objects. \"\"\"    \n",
      "    \n",
      "    import pyurdme\n",
      "    from pyurdme.nsmsolver import NSMSolver\n",
      "    import sys\n",
      "    import uuid\n",
      "    from molnsutil import PersistentStorage, LocalStorage\n",
      "    \n",
      "    ps  = PersistentStorage()\n",
      "\n",
      "    # Run the solver \n",
      "    solver = NSMSolver(model)\n",
      "    \n",
      "    filenames = []\n",
      "    for i in range(nt):\n",
      "        try:\n",
      "            result = solver.run(seed=s*engine_id*nt+i)\n",
      "            filename = str(uuid.uuid1())\n",
      "            st=ps.put(filename, result)\n",
      "            filenames.append(filename)\n",
      "        except:\n",
      "            pass\n",
      "\n",
      "    ps.provider.connection.close()\n",
      "    return filenames    "
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 21
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%%px --local\n",
      "def map_and_reduce(results, mapper, reducer):\n",
      "    \"Reduces a list of results by applying the map function 'mapper'.  \"\n",
      "    \n",
      "    import dill\n",
      "    import numpy\n",
      "    from molnsutil import PersistentStorage, LocalStorage\n",
      "    ps = PersistentStorage()\n",
      "    ls = LocalStorage()\n",
      "    \n",
      "    num_processed=0\n",
      "    res = None\n",
      "    for i,filename in enumerate(results):\n",
      "        try:\n",
      "            try:\n",
      "                result = ls.get(filename)\n",
      "            except:\n",
      "                result = ps.get(filename)\n",
      "                ls.put(filename, result)\n",
      "                \n",
      "            mapres = mapper(result)\n",
      "            res = reducer(mapres, res)\n",
      "            num_processed +=1\n",
      "        except Exception as e:\n",
      "            raise\n",
      "        \n",
      "    ps.provider.connection.close()\n",
      "    return {'result':res, 'num_sucessful':num_processed, 'num_failed':len(results)-num_processed}"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "class DistributedEnsemble():\n",
      "    \"\"\" A distributed ensemble. \"\"\"\n",
      "\n",
      "    def __init__(self, name=None, model_class=None, model=None, client=None, number_of_realizations=1, persistent=False):\n",
      "        \"\"\" hjhkjhjk \"\"\"\n",
      "        self.model_class = model_class\n",
      "        self.number_of_realizations = number_of_realizations\n",
      "        self.persistent = persistent\n",
      "        \n",
      "        # A chunk list\n",
      "        self.results = [] \n",
      "        # We can build an index that maps filenames to the engines where the file is and \n",
      "        # guide scheduling. \n",
      "        self.file_index = {}\n",
      "        \n",
      "        self.update_client(client)\n",
      "     \n",
      "    def update_client(self, client=None):\n",
      "        if client is None:\n",
      "            self.c = IPython.parallel.Client()\n",
      "        else: \n",
      "            self.c = client\n",
      "        self.c[:].use_dill()\n",
      "        self.dv = self.c[:]\n",
      "        self.lv = self.c.load_balanced_view()\n",
      "        \n",
      "    def rebalance_chunk_list(self):\n",
      "        \"\"\" It seems like it can be necessary to be able to rebalance the chunk list if\n",
      "            the number of engines change. Like if you suddenly have more engines than chunks, you \n",
      "            want to create more chunks. \"\"\"\n",
      "        \n",
      "    def add_realizations(self, number_of_realizations=1, chunk_size=1, blocking=True):\n",
      "        \"\"\" Add a number of realizations to the ensemble. \"\"\"\n",
      "        model = self.model_class()\n",
      "        #num_chunks=number_of_realizations\n",
      "        num_chunks = int(number_of_realizations/chunk_size)\n",
      "        chunks = [chunk_size]*(num_chunks-1)\n",
      "        chunks.append(number_of_realizations-chunk_size*(num_chunks-1))\n",
      "        #num_chunks = number_of_realizations\n",
      "        #num_chunks=self._determine_chunk_size()\n",
      "        #if nt < num_chunks:\n",
      "        #    num_chunks=nt\n",
      "        #else:\n",
      "        #    nt = int(number_of_realizations/num_chunks)\n",
      "        #reminder = number_of_trajectories-nt\n",
      "        \n",
      "        results  = self.lv.map_async(run_ensemble,[model]*num_chunks,chunks,range(num_chunks))\n",
      "        results.wait()\n",
      "        for r in results:\n",
      "            self.results.append(r)\n",
      "        return {'wall_time':results.wall_time, 'results': results}\n",
      "    \n",
      "    def _determine_chunk_size(self):\n",
      "        \"\"\" Determine a good chunk size in some clever way. \"\"\"\n",
      "        num_chunks = len(self.c.ids())\n",
      "        return num_chunks\n",
      "\n",
      "    def save():\n",
      "        \"\"\" Save the data in the object store. \"\"\"\n",
      "        \n",
      "    def _clear_cache(self):\n",
      "        \"\"\" Remove all cached result objects on the engines. \"\"\"\n",
      "        \n",
      "    def delete_realizations(self):\n",
      "        \"\"\" Delete realizations from the object store. \"\"\"\n",
      "        \n",
      "    \n",
      "    def mean(self, mapper=None, number_of_realizations=None, blocking=True, interactive=False):\n",
      "        \"\"\" Compute the mean of the function g(X) based on number_of_realizations realizations\n",
      "            in the ensemble. It has to make sense to say g(result1)+g(result2). \"\"\"\n",
      "        \n",
      "        num_chunks = len(self.c.ids)\n",
      "        \n",
      "        # Now map the postprocessing routine using the view that matches the file locations on the engines. \n",
      "        pr = self.c[:].map_async(aggregate, self.results, [mapper]*len(self.results),[add]*len(self.results))\n",
      "        #pr.wait()\n",
      "        res = {}\n",
      "        num_sucessful=0\n",
      "        for i,p in enumerate(pr):\n",
      "            try:\n",
      "                if i==0:\n",
      "                    meanx = p['result']/p['num_sucessful']\n",
      "                else:\n",
      "                    meanx = meanx+p['result']/p['num_sucessful']\n",
      "                num_sucessful+=1\n",
      "                if interactive:\n",
      "                    print meanx/num_sucessful\n",
      "            except Exception as e:\n",
      "                raise                \n",
      "        #pr.wait()\n",
      "        # This does not work, we can't just assume that the result is a numeric array. What can we assume?\n",
      "        # Can we assume that it makes sense to say pr[0]+pr[1] etc. (i.e. that the '+' operator is defined for the \n",
      "        # result types?)\n",
      "        res['mean'] = meanx/num_sucessful\n",
      "        res['wall_time']=pr.wall_time\n",
      "        #res['variance'] = \n",
      "        #res['confidence_interval'] = \n",
      "        return res\n",
      "        \n",
      "   \n",
      "    def mean_variance(self, g=None, number_of_realizations=None):\n",
      "        \"\"\" Compute the variance (second order central moment) of the function g(X) based on number_of_realizations realizations\n",
      "            in the ensemble. \"\"\"\n",
      "\n",
      "    def moment(self, g=None, order=1, number_of_realizations=None):\n",
      "        \"\"\" Compute the moment of order 'order' of g(X), using number_of_realizations\n",
      "            realizations in the ensemble. \"\"\"\n",
      "\n",
      "    def histogram_density(self, g=None, number_of_realizations=None):\n",
      "        \"\"\" Estimate the probability density function of g(X) based on number_of_realizations realizations\n",
      "            in the ensemble. \"\"\"\n"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 5
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "ensemble = DistributedEnsemble(model_class=SimpleDiffusion, name='testensemble')\n",
      "res  = ensemble.add_realizations(number_of_realizations=1000, chunk_size=10)\n",
      "print \"Time to compute:\", res[\"wall_time\"]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "Time to compute: 52.598316\n"
       ]
      }
     ],
     "prompt_number": 46
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "res  = ensemble.add_realizations(number_of_realizations=1000, chunk_size=10)\n",
      "print \"Time to compute:\", res[\"wall_time\"]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "Time to compute: 51.496302\n"
       ]
      }
     ],
     "prompt_number": 24
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "print ensemble.results[0]\n",
      "print res['wall_time']"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "['8c1b77a8-430f-11e4-acb0-fa163e422853', '8c64c30e-430f-11e4-acb0-fa163e422853', '8ca405b4-430f-11e4-acb0-fa163e422853', '8ce1336c-430f-11e4-acb0-fa163e422853', '8d1e7218-430f-11e4-acb0-fa163e422853', '8d5a8118-430f-11e4-acb0-fa163e422853', '8d974d64-430f-11e4-acb0-fa163e422853', '8dd3be8e-430f-11e4-acb0-fa163e422853', '8e05296a-430f-11e4-acb0-fa163e422853', '8e3d97dc-430f-11e4-acb0-fa163e422853']\n",
        "57.315042\n"
       ]
      }
     ],
     "prompt_number": 18
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%%px --local\n",
      "import pyurdme\n",
      "import numpy\n",
      "def g(result):\n",
      "    \"\"\" Extract the values of A at the endpoint. \"\"\"\n",
      "    #A = numpy.sum(result.get_species(\"A\",-1))\n",
      "    A = result.get_species(\"A\",-1)[400]\n",
      "    return A"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 30
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "ensemble.update_client()\n",
      "res = ensemble.mean(mapper=g)\n",
      "print \"Mean of g:\", res['result']\n",
      "print \"Time to compute:\", res['wall_time']"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "Mean of g: 0.106\n",
        "Time to compute: 14.470887\n"
       ]
      }
     ],
     "prompt_number": 55
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# Store the list of lists of filesnames in the object store\n",
      "from molnsutil import PersistentStorage\n",
      "ps = PersistentStorage()\n",
      "files = []\n",
      "for r in results:\n",
      "    files.append(r)\n",
      "ensemble_data = {'files':files,'num_chunks':num_chunks, 'chunk_size':nt}\n",
      "ps.put('ensemble_data',ensemble_data)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# Clean up on the engines\n",
      "def cleanup(results):\n",
      "    from molnsutil import PersistentStorage, LocalStorage\n",
      "    ps = PersistentStorage()\n",
      "    ls = LocalStorage()\n",
      "    #ps.delete(results)\n",
      "    for filename in results:\n",
      "        try: \n",
      "            ls.delete(filename)\n",
      "        except:\n",
      "            pass\n",
      "        \n",
      "        try:\n",
      "            ps.delete(filename)\n",
      "        except:\n",
      "            pass\n",
      "res = rv.map_async(cleanup, results)\n",
      "res.wait()\n",
      "print \"Time to delete all the results on the engines:\", res.wall_time"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    }
   ],
   "metadata": {}
  }
 ]
}