1# This file is part of h5py, a Python interface to the HDF5 library.
2#
3# http://www.h5py.org
4#
5# Copyright 2008-2013 Andrew Collette and contributors
6#
7# License:  Standard 3-clause BSD; see "license.txt" for full license terms
8#           and contributor agreement.
9
10"""
11    Demonstrates how to use h5py with the multiprocessing module.
12
13    This module implements a simple multi-process program to generate
14    Mandelbrot set images.  It uses a process pool to do the computations,
15    and a single process to save the results to file.
16
17    Importantly, only one process actually reads/writes the HDF5 file.
18    Remember that when a process is fork()ed, the child inherits the HDF5
19    state from its parent, which can be dangerous if you already have a file
20    open.  Trying to interact with the same file on disk from multiple
21    processes results in undefined behavior.
22
23    If matplotlib is available, the program will read from the HDF5 file and
24    display an image of the fractal in a window.  To re-run the calculation,
25    delete the file "mandelbrot.hdf5".
26
27"""
28
29import numpy as np
30import multiprocessing as mp
31import h5py
32
33# === Parameters for Mandelbrot calculation ===================================
34
35NX = 512
36NY = 512
37ESCAPE = 1000
38
39XSTART = -0.16070135 - 5e-8
40YSTART =  1.0375665 -5e-8
41XEXTENT = 1.0E-7
42YEXTENT = 1.0E-7
43
44xincr = XEXTENT*1.0/NX
45yincr = YEXTENT*1.0/NY
46
47# === Functions to compute set ================================================
48
49def compute_escape(pos):
50    """ Compute the number of steps required to escape from a point on the
51    complex plane """
52    z = 0+0j;
53    for i in range(ESCAPE):
54        z = z**2 + pos
55        if abs(z) > 2:
56            break
57    return i
58
59def compute_row(xpos):
60    """ Compute a 1-D array containing escape step counts for each y-position.
61    """
62    a = np.ndarray((NY,), dtype='i')
63    for y in range(NY):
64        pos = complex(XSTART,YSTART) + complex(xpos, y*yincr)
65        a[y] = compute_escape(pos)
66    return a
67
68# === Functions to run process pool & visualize ===============================
69
70def run_calculation():
71    """ Begin multi-process calculation, and save to file """
72
73    print("Creating %d-process pool" % mp.cpu_count())
74
75    pool = mp.Pool(mp.cpu_count())
76
77    f = h5py.File('mandelbrot.hdf5','w')
78
79    print("Creating output dataset with shape %s x %s" % (NX, NY))
80
81    dset = f.create_dataset('mandelbrot', (NX,NY), 'i')
82    dset.attrs['XSTART'] = XSTART
83    dset.attrs['YSTART'] = YSTART
84    dset.attrs['XEXTENT'] = XEXTENT
85    dset.attrs['YEXTENT'] = YEXTENT
86
87    result = pool.imap(compute_row, (x*xincr for x in range(NX)))
88
89    for idx, arr in enumerate(result):
90        if idx%25 == 0: print("Recording row %s" % idx)
91        dset[idx] = arr
92
93    print("Closing HDF5 file")
94
95    f.close()
96
97    print("Shutting down process pool")
98
99    pool.close()
100    pool.join()
101
102def visualize_file():
103    """ Open the HDF5 file and display the result """
104    try:
105        import pylab as p
106    except ImportError:
107        print("Whoops! Matplotlib is required to view the fractal.")
108        raise
109
110    f = h5py.File('mandelbrot.hdf5','r')
111    dset = f['mandelbrot']
112    a = dset[...]
113    p.imshow(a.transpose())
114
115    print("Displaying fractal. Close window to exit program.")
116    try:
117        p.show()
118    finally:
119        f.close()
120
121if __name__ == '__main__':
122    if not h5py.is_hdf5('mandelbrot.hdf5'):
123        run_calculation()
124    else:
125        print('Fractal found in "mandelbrot.hdf5". Delete file to re-run calculation.')
126    visualize_file()
127