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