SHOW:
|
|
- or go back to the newest paste.
| 1 | import numpy as np | |
| 2 | import time | |
| 3 | from joblib import Parallel, delayed, parallel_backend | |
| 4 | from extra_fns import * | |
| 5 | ||
| 6 | time.perf_counter() | |
| 7 | nj = 2 | |
| 8 | set_par = True | |
| 9 | split_var = True | |
| 10 | ||
| 11 | # define 3d grid | |
| 12 | nd = 3 | |
| 13 | nx = 250 | |
| 14 | ny = 250 | |
| 15 | nz = 250 | |
| 16 | x = np.linspace(0, 1, nx) | |
| 17 | y = np.linspace(0, 1, ny) | |
| 18 | z = np.linspace(0, 1, nz) | |
| 19 | ||
| 20 | # positions of gaussians in space | |
| 21 | pgrid = np.linspace(0.05, 0.95 , 20) | |
| 22 | Xp, Yp, Zp = np.meshgrid(pgrid,pgrid,pgrid) | |
| 23 | xp = Xp.ravel() | |
| 24 | yp = Yp.ravel() | |
| 25 | zp = Zp.ravel() | |
| 26 | Np = np.size(xp) | |
| 27 | s = np.ones(Np) # intensity of each gaussian | |
| 28 | # compact gaussian representation | |
| 29 | sigma = x[1]-x[0] | |
| 30 | max_dist = sigma*(-2*np.log(10e-3)) | |
| 31 | ||
| 32 | # 3D domain: | |
| 33 | I = np.zeros((ny, nx, nz)) | |
| 34 | dx = x[1] - x[0] | |
| 35 | dy = y[1] - y[0] | |
| 36 | dz = z[1] - z[0] | |
| 37 | ||
| 38 | ||
| 39 | dix = np.ceil(max_dist/dx) | |
| 40 | diy = np.ceil(max_dist/dy) | |
| 41 | diz = np.ceil(max_dist/dz) | |
| 42 | def run_test(set_par, split_var, xp, yp, zp, s): | |
| 43 | def add_loc_gaussian(i): | |
| 44 | ix = round((xp[i] - x[0]) / dx) | |
| 45 | iy = round((yp[i] - y[0]) / dy) | |
| 46 | iz = round((zp[i] - z[0]) / dz) | |
| 47 | iix = np.arange(max(0, ix - dix), min(nx, ix + dix), 1, dtype=int) | |
| 48 | iiy = np.arange(max(0, iy - diy), min(ny, iy + diy), 1, dtype=int) | |
| 49 | iiz = np.arange(max(0, iz - diz), min(nz, iz + diz), 1, dtype=int) | |
| 50 | ddx = dx * iix - xp[i] | |
| 51 | ddy = dy * iiy - yp[i] | |
| 52 | ddz = dz * iiz - zp[i] | |
| 53 | gx = np.exp(-1 / (2 * sigma ** 2) * ddx ** 2) | |
| 54 | gy = np.exp(-1 / (2 * sigma ** 2) * ddy ** 2) | |
| 55 | gz = np.exp(-1 / (2 * sigma ** 2) * ddz ** 2) | |
| 56 | gx = gx[np.newaxis,:, np.newaxis] | |
| 57 | gy = gy[:,np.newaxis, np.newaxis] | |
| 58 | gz = gz[np.newaxis, np.newaxis, :] | |
| 59 | I[np.ix_(iiy, iix, iiz)] += s[i] * gy*gx*gz | |
| 60 | if set_par and split_var: # case 1 | |
| 61 | mp = int(Np/nj) # hard code this test fn for two cores | |
| 62 | xp_list = [xp[:mp],xp[mp:]] | |
| 63 | yp_list = [yp[:mp],yp[mp:]] | |
| 64 | zp_list = [zp[:mp],zp[mp:]] | |
| 65 | sp_list = [s[:mp],s[mp:]] | |
| 66 | def core_loop(j): | |
| 67 | xpt = xp_list[j] | |
| 68 | ypt = yp_list[j] | |
| 69 | zpt = zp_list[j] | |
| 70 | spt = sp_list[j] | |
| 71 | def add_loc_gaussian_s(i): | |
| 72 | ix = round((xpt[i] - x[0]) / dx) | |
| 73 | iy = round((ypt[i] - y[0]) / dy) | |
| 74 | iz = round((zpt[i] - z[0]) / dz) | |
| 75 | iix = np.arange(max(0, ix - dix), min(nx, ix + dix), 1, dtype=int) | |
| 76 | iiy = np.arange(max(0, iy - diy), min(ny, iy + diy), 1, dtype=int) | |
| 77 | iiz = np.arange(max(0, iz - diz), min(nz, iz + diz), 1, dtype=int) | |
| 78 | ddx = dx * iix - xpt[i] | |
| 79 | ddy = dy * iiy - ypt[i] | |
| 80 | ddz = dz * iiz - zpt[i] | |
| 81 | gx = np.exp(-1 / (2 * sigma ** 2) * ddx ** 2) | |
| 82 | gy = np.exp(-1 / (2 * sigma ** 2) * ddy ** 2) | |
| 83 | gz = np.exp(-1 / (2 * sigma ** 2) * ddz ** 2) | |
| 84 | gx = gx[np.newaxis, :, np.newaxis] | |
| 85 | gy = gy[:, np.newaxis, np.newaxis] | |
| 86 | gz = gz[np.newaxis, np.newaxis, :] | |
| 87 | I[np.ix_(iiy, iix, iiz)] += spt[i] * gy * gx * gz | |
| 88 | for i in range(np.size(xpt)): | |
| 89 | add_loc_gaussian_s(i) | |
| 90 | Parallel(n_jobs=2, require='sharedmem')(delayed(core_loop)(i) for i in range(2)) | |
| 91 | elif set_par: # case 2 | |
| 92 | Parallel(n_jobs=nj, require='sharedmem')(delayed(add_loc_gaussian)(i) for i in range(Np)) | |
| 93 | else: # case 3 | |
| 94 | for i in range(0,Np): | |
| 95 | add_loc_gaussian(i) | |
| 96 | ||
| 97 | run_test(set_par, split_var, xp, yp, zp, s) | |
| 98 | print("Time taken: {} s".format(time.perf_counter())) |
