"""Demonstrate the usage of numpy.memmap with joblib.Parallel

This example shows how to preallocate data in memmap arrays both for input and
output of the parallel worker processes.

Sample output for this program::

    [Worker 93486] Sum for row 0 is -1599.756454
    [Worker 93487] Sum for row 1 is -243.253165
    [Worker 93488] Sum for row 3 is 610.201883
    [Worker 93489] Sum for row 2 is 187.982005
    [Worker 93489] Sum for row 7 is 326.381617
    [Worker 93486] Sum for row 4 is 137.324438
    [Worker 93489] Sum for row 8 is -198.225809
    [Worker 93487] Sum for row 5 is -1062.852066
    [Worker 93488] Sum for row 6 is 1666.334107
    [Worker 93486] Sum for row 9 is -463.711714
    Expected sums computed in the parent process:
    [-1599.75645426  -243.25316471   187.98200458   610.20188337   137.32443803
     -1062.85206633  1666.33410715   326.38161713  -198.22580876  -463.71171369]
    Actual sums computed by the worker processes:
    [-1599.75645426  -243.25316471   187.98200458   610.20188337   137.32443803
     -1062.85206633  1666.33410715   326.38161713  -198.22580876  -463.71171369]

"""
import tempfile
import gc
import shutil
import os
import numpy as np

from joblib import Parallel, delayed
from joblib import load, dump


def sum_row(input, output, i):
    """Compute the sum of a row in input and store it in output"""
    sum_ = input[i, :].sum()
    print "[Worker %d] Sum for row %d is %f" % (os.getpid(), i, sum_)
    output[i] = sum_

if __name__ == "__main__":
    rng = np.random.RandomState(42)
    folder = tempfile.mkdtemp()
    samples_name = os.path.join(folder, 'samples')
    sums_name = os.path.join(folder, 'sums')
    try:
        # Generate some data and an allocate an output buffer
        samples = rng.normal(size=(10, int(1e6)))
        sums = np.memmap(sums_name, dtype=samples.dtype,
                         shape=samples.shape[0], mode='w+')

        # Dump the input data to disk to free the memory
        dump(samples, samples_name)
        samples = load(samples_name, mmap_mode='r')

        # Make sure that the original arrays are no longer in memory before
        # forking
        gc.collect()

        # Fork the worker processes to perform computation concurrently
        Parallel(n_jobs=4)(delayed(sum_row)(samples, sums, i)
                           for i in xrange(samples.shape[0]))

        # Compare the results from the output buffer with the ground truth
        print "Expected sums computed in the parent process:"
        expected_result = samples.sum(axis=1)
        print expected_result

        print "Actual sums computed by the worker processes:"
        print sums

        assert np.allclose(expected_result, sums)
    finally:
        shutil.rmtree(folder)

@@ -1391,17 +1391,17 @@
    print
- 
+(
 %22%5BWorker
@@ -1451,16 +1451,17 @@
 i, sum_)
+)
 %0A    out
@@ -2337,17 +2337,16 @@
 or i in 
-x
 range(sa
@@ -2451,17 +2451,17 @@
    print
- 
+(
 %22Expecte
@@ -2499,16 +2499,17 @@
 rocess:%22
+)
 %0A       
@@ -2560,17 +2560,17 @@
    print
- 
+(
 expected
@@ -2576,16 +2576,17 @@
 d_result
+)
 %0A%0A      
@@ -2592,17 +2592,17 @@
    print
- 
+(
 %22Actual 
@@ -2640,16 +2640,17 @@
 cesses:%22
+)
 %0A       
@@ -2655,21 +2655,22 @@
    print
- 
+(
 sums
+)
 %0A%0A      

Make example work with Python 3