07_threading.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Threading and parallel processing\n",
    "\n",
    "\n",
    "The Python language has built-in support for multi-threading in the\n",
    "[`threading`](https://docs.python.org/3.5/library/threading.html) module, and\n",
    "true parallelism in the\n",
    "[`multiprocessing`](https://docs.python.org/3.5/library/multiprocessing.html)\n",
    "module.  If you want to be impressed, skip straight to the section on\n",
    "[`multiprocessing`](todo).\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "## Threading\n",
    "\n",
    "\n",
    "The [`threading`](https://docs.python.org/3.5/library/threading.html) module\n",
    "provides a traditional multi-threading API that should be familiar to you if\n",
    "you have worked with threads in other languages.\n",
    "\n",
    "\n",
    "Running a task in a separate thread in Python is easy - simply create a\n",
    "`Thread` object, and pass it the function or method that you want it to\n",
    "run. Then call its `start` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import threading\n",
    "\n",
    "def longRunningTask(niters):\n",
    "    for i in range(niters):\n",
    "        if i % 2 == 0: print('Tick')\n",
    "        else:          print('Tock')\n",
    "        time.sleep(0.5)\n",
    "\n",
    "t = threading.Thread(target=longRunningTask, args=(8,))\n",
    "\n",
    "t.start()\n",
    "\n",
    "while t.is_alive():\n",
    "    time.sleep(0.4)\n",
    "    print('Waiting for thread to finish...')\n",
    "print('Finished!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also `join` a thread, which will block execution in the current thread\n",
    "until the thread that has been `join`ed has finished:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = threading.Thread(target=longRunningTask, args=(6, ))\n",
    "t.start()\n",
    "\n",
    "print('Joining thread ...')\n",
    "t.join()\n",
    "print('Finished!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Subclassing `Thread`\n",
    "\n",
    "\n",
    "It is also possible to sub-class the `Thread` class, and override its `run`\n",
    "method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LongRunningThread(threading.Thread):\n",
    "    def __init__(self, niters, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.niters = niters\n",
    "\n",
    "    def run(self):\n",
    "        for i in range(self.niters):\n",
    "            if i % 2 == 0: print('Tick')\n",
    "            else:          print('Tock')\n",
    "            time.sleep(0.5)\n",
    "\n",
    "t = LongRunningThread(6)\n",
    "t.start()\n",
    "t.join()\n",
    "print('Done')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Daemon threads\n",
    "\n",
    "\n",
    "By default, a Python application will not exit until _all_ active threads have\n",
    "finished execution.  If you want to run a task in the background for the\n",
    "duration of your application, you can mark it as a `daemon` thread - when all\n",
    "non-daemon threads in a Python application have finished, all daemon threads\n",
    "will be halted, and the application will exit.\n",
    "\n",
    "\n",
    "You can mark a thread as being a daemon by setting an attribute on it after\n",
    "creation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = threading.Thread(target=longRunningTask)\n",
    "t.daemon = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See the [`Thread`\n",
    "documentation](https://docs.python.org/3.5/library/threading.html#thread-objects)\n",
    "for more details.\n",
    "\n",
    "\n",
    "### Thread synchronisation\n",
    "\n",
    "\n",
    "The `threading` module provides some useful thread-synchronisation primitives\n",
    "- the `Lock`, `RLock` (re-entrant `Lock`), and `Event` classes.  The\n",
    "`threading` module also provides `Condition` and `Semaphore` classes - refer\n",
    "to the [documentation](https://docs.python.org/3.5/library/threading.html) for\n",
    "more details.\n",
    "\n",
    "\n",
    "#### `Lock`\n",
    "\n",
    "\n",
    "The [`Lock`](https://docs.python.org/3.5/library/threading.html#lock-objects)\n",
    "class (and its re-entrant version, the\n",
    "[`RLock`](https://docs.python.org/3.5/library/threading.html#rlock-objects))\n",
    "prevents a block of code from being accessed by more than one thread at a\n",
    "time. For example, if we have multiple threads running this `task` function,\n",
    "their [outputs](https://www.youtube.com/watch?v=F5fUFnfPpYU) will inevitably\n",
    "become intertwined:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def task():\n",
    "    for i in range(5):\n",
    "        print('{} Woozle '.format(i), end='')\n",
    "        time.sleep(0.1)\n",
    "        print('Wuzzle')\n",
    "\n",
    "threads = [threading.Thread(target=task) for i in range(5)]\n",
    "for t in threads:\n",
    "    t.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But if we protect the critical section with a `Lock` object, the output will\n",
    "look more sensible:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lock = threading.Lock()\n",
    "\n",
    "def task():\n",
    "\n",
    "    for i in range(5):\n",
    "        with lock:\n",
    "            print('{} Woozle '.format(i), end='')\n",
    "            time.sleep(0.1)\n",
    "            print('Wuzzle')\n",
    "\n",
    "threads = [threading.Thread(target=task) for i in range(5)]\n",
    "for t in threads:\n",
    "    t.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Instead of using a `Lock` object in a `with` statement, it is also possible\n",
    "> to manually call its `acquire` and `release` methods:\n",
    ">\n",
    ">     def task():\n",
    ">         for i in range(5):\n",
    ">             lock.acquire()\n",
    ">             print('{} Woozle '.format(i), end='')\n",
    ">             time.sleep(0.1)\n",
    ">             print('Wuzzle')\n",
    ">             lock.release()\n",
    "\n",
    "\n",
    "Python does not have any built-in constructs to implement `Lock`-based mutual\n",
    "exclusion across several functions or methods - each function/method must\n",
    "explicitly acquire/release a shared `Lock` instance. However, it is relatively\n",
    "straightforward to implement a decorator which does this for you:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mutex(func, lock):\n",
    "    def wrapper(*args):\n",
    "        with lock:\n",
    "            func(*args)\n",
    "    return wrapper\n",
    "\n",
    "class MyClass(object):\n",
    "\n",
    "    def __init__(self):\n",
    "        lock = threading.Lock()\n",
    "        self.safeFunc1 = mutex(self.safeFunc1, lock)\n",
    "        self.safeFunc2 = mutex(self.safeFunc2, lock)\n",
    "\n",
    "    def safeFunc1(self):\n",
    "        time.sleep(0.1)\n",
    "        print('safeFunc1 start')\n",
    "        time.sleep(0.2)\n",
    "        print('safeFunc1 end')\n",
    "\n",
    "    def safeFunc2(self):\n",
    "        time.sleep(0.1)\n",
    "        print('safeFunc2 start')\n",
    "        time.sleep(0.2)\n",
    "        print('safeFunc2 end')\n",
    "\n",
    "mc = MyClass()\n",
    "\n",
    "f1threads = [threading.Thread(target=mc.safeFunc1) for i in range(4)]\n",
    "f2threads = [threading.Thread(target=mc.safeFunc2) for i in range(4)]\n",
    "\n",
    "for t in f1threads + f2threads:\n",
    "    t.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try removing the `mutex` lock from the two methods in the above code, and see\n",
    "what it does to the output.\n",
    "\n",
    "\n",
    "#### `Event`\n",
    "\n",
    "\n",
    "The\n",
    "[`Event`](https://docs.python.org/3.5/library/threading.html#event-objects)\n",
    "class is essentially a boolean [semaphore][semaphore-wiki]. It can be used to\n",
    "signal events between threads. Threads can `wait` on the event, and be awoken\n",
    "when the event is `set` by another thread:\n",
    "\n",
    "\n",
    "[semaphore-wiki]: https://en.wikipedia.org/wiki/Semaphore_(programming)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "processingFinished = threading.Event()\n",
    "\n",
    "def processData(data):\n",
    "    print('Processing data ...')\n",
    "    time.sleep(2)\n",
    "    print('Result: {}'.format(data.mean()))\n",
    "    processingFinished.set()\n",
    "\n",
    "data = np.random.randint(1, 100, 100)\n",
    "\n",
    "t = threading.Thread(target=processData, args=(data,))\n",
    "t.start()\n",
    "\n",
    "processingFinished.wait()\n",
    "print('Processing finished!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Global Interpreter Lock (GIL)\n",
    "\n",
    "\n",
    "The [_Global Interpreter\n",
    "Lock_](https://docs.python.org/3/c-api/init.html#thread-state-and-the-global-interpreter-lock)\n",
    "is an implementation detail of [CPython](https://github.com/python/cpython)\n",
    "(the official Python interpreter).  The GIL means that a multi-threaded\n",
    "program written in pure Python is not able to take advantage of multiple\n",
    "cores - this essentially means that only one thread may be executing at any\n",
    "point in time.\n",
    "\n",
    "\n",
    "The `threading` module does still have its uses though, as this GIL problem\n",
    "does not affect tasks which involve calls to system or natively compiled\n",
    "libraries (e.g. file and network I/O, Numpy operations, etc.). So you can,\n",
    "for example, perform some expensive processing on a Numpy array in a thread\n",
    "running on one core, whilst having another thread (e.g. user interaction)\n",
    "running on another core.\n",
    "\n",
    "\n",
    "## Multiprocessing\n",
    "\n",
    "\n",
    "For true parallelism, you should check out the\n",
    "[`multiprocessing`](https://docs.python.org/3.5/library/multiprocessing.html)\n",
    "module.\n",
    "\n",
    "\n",
    "The `multiprocessing` module spawns sub-processes, rather than threads, and so\n",
    "is not subject to the GIL constraints that the `threading` module suffers\n",
    "from. It provides two APIs - a \"traditional\" equivalent to that provided by\n",
    "the `threading` module, and a powerful higher-level API.\n",
    "\n",
    "\n",
    "### `threading`-equivalent API\n",
    "\n",
    "\n",
    "The\n",
    "[`Process`](https://docs.python.org/3.5/library/multiprocessing.html#the-process-class)\n",
    "class is the `multiprocessing` equivalent of the\n",
    "[`threading.Thread`](https://docs.python.org/3.5/library/threading.html#thread-objects)\n",
    "class.  `multprocessing` also has equivalents of the [`Lock` and `Event`\n",
    "classes](https://docs.python.org/3.5/library/multiprocessing.html#synchronization-between-processes),\n",
    "and the other synchronisation primitives provided by `threading`.\n",
    "\n",
    "\n",
    "So you can simply replace `threading.Thread` with `multiprocessing.Process`,\n",
    "and you will have true parallelism.\n",
    "\n",
    "\n",
    "Because your \"threads\" are now independent processes, you need to be a little\n",
    "careful about how to share information across them. Fortunately, the\n",
    "`multiprocessing` module provides [`Queue` and `Pipe`\n",
    "classes](https://docs.python.org/3.5/library/multiprocessing.html#exchanging-objects-between-processes)\n",
    "which make it easy to share data across processes.\n",
    "\n",
    "\n",
    "### Higher-level API - the `multiprocessing.Pool`\n",
    "\n",
    "\n",
    "The real advantages of `multiprocessing` lie in its higher level API, centered\n",
    "around the [`Pool`\n",
    "class](https://docs.python.org/3.5/library/multiprocessing.html#using-a-pool-of-workers).\n",
    "\n",
    "\n",
    "Essentially, you create a `Pool` of worker processes - you specify the number\n",
    "of processes when you create the pool.\n",
    "\n",
    "\n",
    "> The best number of processes to use for a `Pool` will depend on the system\n",
    "> you are running on (number of cores), and the tasks you are running (e.g.\n",
    "> I/O bound or CPU bound).\n",
    "\n",
    "\n",
    "Once you have created a `Pool`, you can use its methods to automatically\n",
    "parallelise tasks. The most useful are the `map`, `starmap` and\n",
    "`apply_async` methods.\n",
    "\n",
    "\n",
    "#### `Pool.map`\n",
    "\n",
    "\n",
    "The\n",
    "[`Pool.map`](https://docs.python.org/3.5/library/multiprocessing.html#multiprocessing.pool.Pool.map)\n",
    "method is the multiprocessing equivalent of the built-in\n",
    "[`map`](https://docs.python.org/3.5/library/functions.html#map) function - it\n",
    "is given a function, and a sequence, and it applies the function to each\n",
    "element in the sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import                    time\n",
    "import multiprocessing as mp\n",
    "import numpy           as np\n",
    "\n",
    "def crunchImage(imgfile):\n",
    "\n",
    "    # Load a nifti image, do stuff\n",
    "    # to it. Use your imagination\n",
    "    # to fill in this function.\n",
    "    time.sleep(2)\n",
    "\n",
    "    # numpy's random number generator\n",
    "    # will be initialised in the same\n",
    "    # way in each process, so let's\n",
    "    # re-seed it.\n",
    "    np.random.seed()\n",
    "    result = np.random.randint(1, 100, 1)\n",
    "\n",
    "    print(imgfile, ':', result)\n",
    "\n",
    "    return result\n",
    "\n",
    "imgfiles = ['{:02d}.nii.gz'.format(i) for i in range(20)]\n",
    "\n",
    "p = mp.Pool(processes=16)\n",
    "\n",
    "print('Crunching images...')\n",
    "\n",
    "start   = time.time()\n",
    "results = p.map(crunchImage, imgfiles)\n",
    "end     = time.time()\n",
    "\n",
    "print('Total execution time: {:0.2f} seconds'.format(end - start))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `Pool.map` method only works with functions that accept one argument, such\n",
    "as our `crunchImage` function above. If you have a function which accepts\n",
    "multiple arguments, use the\n",
    "[`Pool.starmap`](https://docs.python.org/3.5/library/multiprocessing.html#multiprocessing.pool.Pool.starmap)\n",
    "method instead:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def crunchImage(imgfile, modality):\n",
    "    time.sleep(2)\n",
    "\n",
    "    np.random.seed()\n",
    "\n",
    "    if modality == 't1':\n",
    "        result = np.random.randint(1, 100, 1)\n",
    "    elif modality == 't2':\n",
    "        result = np.random.randint(100, 200, 1)\n",
    "\n",
    "    print(imgfile, ': ', result)\n",
    "\n",
    "    return result\n",
    "\n",
    "imgfiles   = ['t1_{:02d}.nii.gz'.format(i) for i in range(10)] + \\\n",
    "             ['t2_{:02d}.nii.gz'.format(i) for i in range(10)]\n",
    "modalities = ['t1'] * 10 + ['t2'] * 10\n",
    "\n",
    "pool = mp.Pool(processes=16)\n",
    "\n",
    "args = [(f, m) for f, m in zip(imgfiles, modalities)]\n",
    "\n",
    "print('Crunching images...')\n",
    "\n",
    "start   = time.time()\n",
    "results = pool.starmap(crunchImage, args)\n",
    "end     = time.time()\n",
    "\n",
    "print('Total execution time: {:0.2f} seconds'.format(end - start))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `map` and `starmap` methods also have asynchronous equivalents `map_async`\n",
    "and `starmap_async`, which return immediately. Refer to the\n",
    "[`Pool`](https://docs.python.org/3.5/library/multiprocessing.html#module-multiprocessing.pool)\n",
    "documentation for more details.\n",
    "\n",
    "\n",
    "#### `Pool.apply_async`\n",
    "\n",
    "\n",
    "The\n",
    "[`Pool.apply`](https://docs.python.org/3.5/library/multiprocessing.html#multiprocessing.pool.Pool.apply)\n",
    "method will execute a function on one of the processes, and block until it has\n",
    "finished.  The\n",
    "[`Pool.apply_async`](https://docs.python.org/3.5/library/multiprocessing.html#multiprocessing.pool.Pool.apply_async)\n",
    "method returns immediately, and is thus more suited to asynchronously\n",
    "scheduling multiple jobs to run in parallel.\n",
    "\n",
    "\n",
    "`apply_async` returns an object of type\n",
    "[`AsyncResult`](https://docs.python.org/3.5/library/multiprocessing.html#multiprocessing.pool.AsyncResult).\n",
    "An `AsyncResult` object has `wait` and `get` methods which will block until\n",
    "the job has completed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import                    time\n",
    "import multiprocessing as mp\n",
    "import numpy           as np\n",
    "\n",
    "\n",
    "def linear_registration(src, ref):\n",
    "    time.sleep(1)\n",
    "\n",
    "    return np.eye(4)\n",
    "\n",
    "def nonlinear_registration(src, ref, affine):\n",
    "\n",
    "    time.sleep(3)\n",
    "\n",
    "    # this number represents a non-linear warp\n",
    "    # field - use your imagination people!\n",
    "    np.random.seed()\n",
    "    return np.random.randint(1, 100, 1)\n",
    "\n",
    "t1s = ['{:02d}_t1.nii.gz'.format(i) for i in range(20)]\n",
    "std = 'MNI152_T1_2mm.nii.gz'\n",
    "\n",
    "pool = mp.Pool(processes=16)\n",
    "\n",
    "print('Running structural-to-standard registration '\n",
    "      'on {} subjects...'.format(len(t1s)))\n",
    "\n",
    "# Run linear registration on all the T1s.\n",
    "#\n",
    "# We build a list of AsyncResult objects\n",
    "linresults = [pool.apply_async(linear_registration, (t1, std))\n",
    "              for t1 in t1s]\n",
    "\n",
    "# Then we wait for each job to finish,\n",
    "# and replace its AsyncResult object\n",
    "# with the actual result - an affine\n",
    "# transformation matrix.\n",
    "start = time.time()\n",
    "for i, r in enumerate(linresults):\n",
    "    linresults[i] = r.get()\n",
    "end = time.time()\n",
    "\n",
    "print('Linear registrations completed in '\n",
    "      '{:0.2f} seconds'.format(end - start))\n",
    "\n",
    "# Run non-linear registration on all the T1s,\n",
    "# using the linear registrations to initialise.\n",
    "nlinresults = [pool.apply_async(nonlinear_registration, (t1, std, aff))\n",
    "               for (t1, aff) in zip(t1s, linresults)]\n",
    "\n",
    "# Wait for each non-linear reg to finish,\n",
    "# and store the resulting warp field.\n",
    "start = time.time()\n",
    "for i, r in enumerate(nlinresults):\n",
    "    nlinresults[i] = r.get()\n",
    "end = time.time()\n",
    "\n",
    "print('Non-linear registrations completed in '\n",
    "      '{:0.2f} seconds'.format(end - start))\n",
    "\n",
    "print('Non linear registrations:')\n",
    "for t1, result in zip(t1s, nlinresults):\n",
    "    print(t1, ':', result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sharing data between processes\n",
    "\n",
    "\n",
    "When you use the `Pool.map` method (or any of the other methods we have shown)\n",
    "to run a function on a sequence of items, those items must be copied into the\n",
    "memory of each of the child processes. When the child processes are finished,\n",
    "the data that they return then has to be copied back to the parent process.\n",
    "\n",
    "\n",
    "Any items which you wish to pass to a function that is executed by a `Pool`\n",
    "must be - the built-in\n",
    "[`pickle`](https://docs.python.org/3.5/library/pickle.html) module is used by\n",
    "`multiprocessing` to serialise and de-serialise the data passed into and\n",
    "returned from a child process. The majority of standard Python types (`list`,\n",
    "`dict`, `str` etc), and Numpy arrays can be pickled and unpickled, so you only\n",
    "need to worry about this detail if you are passing objects of a custom type\n",
    "(e.g. instances of classes that you have written, or that are defined in some\n",
    "third-party library).\n",
    "\n",
    "\n",
    "There is obviously some overhead in copying data back and forth between the\n",
    "main process and the worker processes.  For most computationally intensive\n",
    "tasks, this communication overhead is not important - the performance\n",
    "bottleneck is typically going to be the computation time, rather than I/O\n",
    "between the parent and child processes. You may need to spend some time\n",
    "adjusting the way in which you split up your data, and the number of\n",
    "processes, in order to get the best performance.\n",
    "\n",
    "\n",
    "However, if you have determined that copying data between processes is having\n",
    "a substantial impact on your performance, the `multiprocessing` module\n",
    "provides the [`Value`, `Array`, and `RawArray`\n",
    "classes](https://docs.python.org/3.5/library/multiprocessing.html#shared-ctypes-objects),\n",
    "which allow you to share individual values, or arrays of values, respectively.\n",
    "\n",
    "\n",
    "The `Array` and `RawArray` classes essentially wrap a typed pointer (from the\n",
    "built-in [`ctypes`](https://docs.python.org/3.5/library/ctypes.html) module)\n",
    "to a block of memory. We can use the `Array` or `RawArray` class to share a\n",
    "Numpy array between our worker processes. The difference between an `Array`\n",
    "and a `RawArray` is that the former offers synchronised (i.e. process-safe)\n",
    "access to the shared memory. This is necessary if your child processes will be\n",
    "modifying the same parts of your data.\n",
    "\n",
    "\n",
    "Due to the way that shared memory works, in order to share a Numpy array\n",
    "between different processes you need to structure your code so that the\n",
    "array(s) you want to share are accessible at the _module level_. Furthermore,\n",
    "we need to make sure that our input and output arrays are located in shared\n",
    "memory - we can do this via the `Array` or `RawArray`.\n",
    "\n",
    "\n",
    "As an example, let's say we want to parallelise processing of an image by\n",
    "having each worker process perform calculations on a chunk of the image.\n",
    "First, let's define a function which does the calculation on a specified set\n",
    "of image coordinates:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiprocessing as mp\n",
    "import ctypes\n",
    "import numpy as np\n",
    "np.set_printoptions(suppress=True)\n",
    "\n",
    "\n",
    "def process_chunk(shape, idxs):\n",
    "\n",
    "    # Get references to our\n",
    "    # input/output data, and\n",
    "    # create Numpy array views\n",
    "    # into them.\n",
    "    sindata  = process_chunk.input_data\n",
    "    soutdata = process_chunk.output_data\n",
    "    indata   = np.ctypeslib.as_array(sindata) .reshape(shape)\n",
    "    outdata  = np.ctypeslib.as_array(soutdata).reshape(shape)\n",
    "\n",
    "    # Do the calculation on\n",
    "    # the specified voxels\n",
    "    outdata[idxs] = indata[idxs] ** 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Rather than passing the input and output data arrays in as arguments to the\n",
    "`process_chunk` function, we set them as attributes of the `process_chunk`\n",
    "function. This makes the input/output data accessible at the module level,\n",
    "which is required in order to share the data between the main process and the\n",
    "child processes.\n",
    "\n",
    "\n",
    "Now let's define a second function which process an entire image. It does the\n",
    "following:\n",
    "\n",
    "\n",
    "1. Initialises shared memory areas to store the input and output data.\n",
    "2. Copies the input data into shared memory.\n",
    "3. Sets the input and output data as attributes of the `process_chunk` function.\n",
    "4. Creates sets of indices into the input data which, for each worker process,\n",
    "   specifies the portion of the data that it is responsible for.\n",
    "5. Creates a worker pool, and runs the `process_chunk` function for each set\n",
    "   of indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_dataset(data):\n",
    "\n",
    "    nprocs   = 8\n",
    "    origData = data\n",
    "\n",
    "    # Create arrays to store the\n",
    "    # input and output data\n",
    "    sindata  = mp.RawArray(ctypes.c_double, data.size)\n",
    "    soutdata = mp.RawArray(ctypes.c_double, data.size)\n",
    "    data     = np.ctypeslib.as_array(sindata).reshape(data.shape)\n",
    "    outdata  = np.ctypeslib.as_array(soutdata).reshape(data.shape)\n",
    "\n",
    "    # Copy the input data\n",
    "    # into shared memory\n",
    "    data[:]  = origData\n",
    "\n",
    "    # Make the input/output data\n",
    "    # accessible to the process_chunk\n",
    "    # function. This must be done\n",
    "    # *before* the worker pool is created.\n",
    "    process_chunk.input_data  = sindata\n",
    "    process_chunk.output_data = soutdata\n",
    "\n",
    "    # number of boxels to be computed\n",
    "    # by each worker process.\n",
    "    nvox = int(data.size / nprocs)\n",
    "\n",
    "    # Generate coordinates for\n",
    "    # every voxel in the image\n",
    "    xlen, ylen, zlen = data.shape\n",
    "    xs, ys, zs = np.meshgrid(np.arange(xlen),\n",
    "                             np.arange(ylen),\n",
    "                             np.arange(zlen))\n",
    "\n",
    "    xs = xs.flatten()\n",
    "    ys = ys.flatten()\n",
    "    zs = zs.flatten()\n",
    "\n",
    "    # We're going to pass each worker\n",
    "    # process a list of indices, which\n",
    "    # specify the data items which that\n",
    "    # worker process needs to compute.\n",
    "    xs = [xs[nvox * i:nvox * i + nvox] for i in range(nprocs)] + \\\n",
    "         [xs[nvox * nprocs:]]\n",
    "    ys = [ys[nvox * i:nvox * i + nvox] for i in range(nprocs)] + \\\n",
    "         [ys[nvox * nprocs:]]\n",
    "    zs = [zs[nvox * i:nvox * i + nvox] for i in range(nprocs)] + \\\n",
    "         [zs[nvox * nprocs:]]\n",
    "\n",
    "    # Build the argument lists for\n",
    "    # each worker process.\n",
    "    args = [(data.shape, (x, y, z)) for x, y, z in zip(xs, ys, zs)]\n",
    "\n",
    "    # Create a pool of worker\n",
    "    # processes and run the jobs.\n",
    "    pool   = mp.Pool(processes=nprocs)\n",
    "\n",
    "    pool.starmap(process_chunk, args)\n",
    "\n",
    "    return outdata"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can call our `process_data` function just like any other function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = np.array(np.arange(64).reshape((4, 4, 4)), dtype=np.float64)\n",
    "\n",
    "outdata = process_dataset(data)\n",
    "\n",
    "print('Input')\n",
    "print(data)\n",
    "\n",
    "print('Output')\n",
    "print(outdata)"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}