|
| 1 | +#!/usr/bin/env python |
| 2 | + |
| 3 | +from argparse import ArgumentParser |
| 4 | +from collections import Counter |
| 5 | +import multiprocessing as mp |
| 6 | +from multiprocessing.managers import SharedMemoryManager |
| 7 | +import numpy as np |
| 8 | +import os |
| 9 | +import sys |
| 10 | + |
| 11 | + |
| 12 | +def init_z(nr_points): |
| 13 | + x = np.linspace(-1.8, 1.8, nr_points) |
| 14 | + y = np.linspace(-1.8j, 1.8j, nr_points) |
| 15 | + X, Y = np.meshgrid(x, np.flip(y)) |
| 16 | + return (X + Y).ravel() |
| 17 | + |
| 18 | + |
| 19 | +def compute_partial_julia(args): |
| 20 | + z_shmem, n_shmem, idx_begin, idx_end, max_iters, max_norm = args |
| 21 | + z_size = np.dtype(np.complex).itemsize |
| 22 | + z_array = np.ndarray((idx_end - idx_begin, ), dtype=np.complex, |
| 23 | + buffer=z_shmem.buf[z_size*idx_begin:z_size*idx_end]) |
| 24 | + n_size = np.dtype(np.int32).itemsize |
| 25 | + n = np.ndarray((idx_end - idx_begin, ), dtype=np.int32, |
| 26 | + buffer=n_shmem.buf[n_size*idx_begin:n_size*idx_end]) |
| 27 | + for i, z in enumerate(z_array): |
| 28 | + while (n[i] <= max_iters and np.abs(z) <= max_norm): |
| 29 | + z = z**2 - 0.622772 + 0.42193j |
| 30 | + n[i] += 1 |
| 31 | + return os.getpid() |
| 32 | + |
| 33 | + |
| 34 | +def compute_julia(nr_points=100, pool_size=2, work_size=15, verbose=False, |
| 35 | + max_iters=255, max_norm=2.0): |
| 36 | + size = nr_points**2 |
| 37 | + z_size = np.dtype(np.complex).itemsize |
| 38 | + n_size = np.dtype(np.int32).itemsize |
| 39 | + with SharedMemoryManager() as shmem_mgr: |
| 40 | + with mp.Pool(pool_size) as pool: |
| 41 | + z_shmem = shmem_mgr.SharedMemory(size=z_size*size) |
| 42 | + z_buf = np.ndarray((size, ), dtype=np.complex, buffer=z_shmem.buf) |
| 43 | + z_buf[:] = init_z(nr_points) |
| 44 | + n_shmem = shmem_mgr.SharedMemory(size=n_size*size) |
| 45 | + n_buf = np.ndarray((size, ), dtype=np.int32, buffer=n_shmem.buf) |
| 46 | + n_buf[:] = np.zeros((size, ), dtype=np.int32) |
| 47 | + args = [(z_shmem, n_shmem, i*work_size, min(z_buf.size, (i + 1)*work_size), |
| 48 | + max_iters, max_norm) |
| 49 | + for i in range(int(np.ceil(z_buf.size/work_size)))] |
| 50 | + if verbose: |
| 51 | + print(args, file=sys.stderr) |
| 52 | + pid_counter = Counter() |
| 53 | + for pid in pool.imap_unordered(compute_partial_julia, args): |
| 54 | + pid_counter[pid] += 1 |
| 55 | + if verbose: |
| 56 | + print(pid_counter, file=sys.stderr) |
| 57 | + return n_buf.copy().reshape(nr_points, nr_points) |
| 58 | + |
| 59 | + |
| 60 | +def main(): |
| 61 | + arg_parser = ArgumentParser(description='compute pi') |
| 62 | + arg_parser.add_argument('--pool_size', type=int, default=2, help='pool size') |
| 63 | + arg_parser.add_argument('--work_size', type=int, default=10, |
| 64 | + help='number of points per work item') |
| 65 | + arg_parser.add_argument('--nr_points', type=int, default=10, |
| 66 | + help='size of the image n x n') |
| 67 | + arg_parser.add_argument('--verbose', action='store_true', help='verbose output') |
| 68 | + options = arg_parser.parse_args() |
| 69 | + constructor = None |
| 70 | + n = compute_julia(nr_points=options.nr_points, pool_size=options.pool_size, |
| 71 | + work_size=options.work_size, verbose=options.verbose) |
| 72 | + np.savetxt(sys.stdout, n, fmt='%3d') |
| 73 | + return 0 |
| 74 | + |
| 75 | +if __name__ == '__main__': |
| 76 | + status = main() |
| 77 | + sys.exit(status) |
0 commit comments