1 files changed, 59 insertions, 78 deletions
diff --git a/sem4/hpp/miniproject/multiproc.py b/sem4/hpp/miniproject/multiproc.py
index be39326..96fe788 100644
--- a/sem4/hpp/miniproject/multiproc.py
+++ b/sem4/hpp/miniproject/multiproc.py
@@ -1,107 +1,88 @@
-#!/usr/bin/env python3
-
 import numpy as np
-import matplotlib.pyplot as plt
-import time
 import multiprocessing as mp
-from multiprocessing import RawArray
+
+# Local import
+from optimised import numpylota
 
 # c-mesh limits
 limitre = ( -2, 1 )
 limitim = ( -1.5, 1.5 )
 
+def createmangelctx(pre, pim, T, l, savez):
+    """
+    Calculates the context used when calculating mangelbrot
 
-def worker(gridchunk, s, step, T, l):
-    global rs
-    # Preallocate z array
-    rschunk = rs[s : s + step]
+    :param pre: Number of real numbers in cmatrix
+    :param pim: Number of imaginary numbers in cmatrix
+    :param T: Mangelbrot threshold
+    :param l: Iterations
+    """
+    sre = ( limitre[1] - limitre[0] ) / (pre-1)
 
-    z = np.zeros(rschunk.shape)
-    print(f"rschunk.shape: {rschunk.shape}, gridchunk.shape: {gridchunk.shape}, T: {T}, l: {l}")
+    ctx = (sre, pim, T, l, savez)
+    
+    return ctx
 
-    # Calculate ι for all complex numbers
-    for i in range(l):
-        # This will generate warnings for some of the values rising above T.
-        # Because these values are above T they are not used, thus the warnings
-        # can be ignored
-        z = z*z + gridchunk
+def mangelstep(ctx, re):
+    """
+    Calculates a single mangelbrot row.
+    
+    :param ctx: Context containing information about mangelbrot
+    :param re: Row to calculate
+    """
 
-        # This will generate 1 in all the places 
-        # where z < T and zeros elsewhere
-        below = (np.abs(z) < T)
+    # Unpack context
+    (sre, pim, T, l, savez) = ctx
 
-        # Add this to the result
-        # Because the ones that pass T are 0 
-        # they will stop counting.
-        #
-        # If a specific z never reaches >= T its value in rs will
-        # be l
-        rschunk += below
+    # Create c-mesh row
+    im = np.linspace(limitim[0], limitim[1], pim, dtype=complex)
+    np.add(1j * im, limitre[0] + sre * re, im)
 
-    np.divide(rschunk, l, out=rschunk)
+    # Calculate ι
+    rs = numpyiota(im, T, l, savez)
+    if savez:
+        return rs
+    else:
+        # It takes some time to unpack a list of tuples
+        # So if no savez do not return a tuple
+        return rs[0]
 
-def mangel(pre, pim, T, l, workers):
+def mangel(pre, pim, T, l, savez):
     """
-    Calculate the mangelbrot image with multiple processes
+    Calculate the mangelbrot image
     (pre, pim) discribes the image size. Use T and l to tune the mangelbrot
-
-    This will split the image in horizontal parts and distribute it 
-    between the workers.
-    Because the result array is row major, data will be nicely together if
-    the workers work with rows not columns.
-
-    Pre must be devisible by workers.
-
-    The result is saved in rs. Sorry couln't get numpy references through to the process as arguments
+    This function uses the global variables limitre and limitim to determine
+    the c-mesh range.
 
     :param pre: Number of real numbers used
     :param pim: Number of imaginary numbers
     :param T: Mangelbrot threshold
     :param l: Iterations
-    :param workers: Number of workers. 
+    :param savez: Return z as the second element of returned tuple
     """
 
-    # Used to calculate c-mesh
-    re = np.linspace(limitre[0], limitre[1], pre)
-    im = np.linspace(limitim[0], limitim[1], pim)
-
-    # Calculate c-mesh
-    grid = np.add.outer(re, 1j * im)
-
-    # Calculate the partition variables
-    step = int(pre / workers)
-
-
-    # Loop chunks and start the workers
-    wl = []
-    for s in range(0, pre, step):
-        gridchunk = grid[s : s + step]
-        p = mp.Process(target=worker, args=(gridchunk, s, step, T, l))
-        wl.append(p)
-        p.start()
-
-    # Wait for them to be done
-    for p in wl:
-        p.join()
+    ctx = createmangelctx(pre, pim, T, l, savez)
 
-    return rs
+    # Number of processes
+    procs = mp.cpu_count()
 
-rs = np.full((500, 500), 0.5)
+    pool = mp.Pool(processes=procs)
 
-start = time.time()
-arr = mangel(500, 500, 2, 100, 1)
-end = time.time()
+    args = ((ctx, i) for i in range(pre)) 
+    result = pool.starmap_async(mangelstep, args)
+    pool.close()
+    pool.join()
 
-plt.imshow(arr, cmap=plt.cm.hot, vmin=0, vmax=1)
-plt.savefig("mult.png")
+    result = result.get()
+    if savez:
+        z = [None] * pre
+        rs = [None] * pre
 
-print(f"Took {end - start} seconds")
-"""
-        p = mp.Process(target=worker, args=(rs[s:s + step], grid[s:s + step], T, l))
-        wl.append(p)
-        p.start()
+        # Unzip rs
+        for i, r in enumerate(result):
+            rs[i] = r[0]
+            z[i] = r[1]
 
-    # Wait for them to be done
-    for p in wl:
-        p.join()
-"""
+        return (np.vstack(rs), np.vstack(z))
+    else:
+        return (np.vstack(result), None)