MFE: added a Numba Version

benchmarks/magnetic_field_extrapolation/python_numba/magnetic_field_extrapolation.py

@@ -0,0 +1,184 @@
+# This Python file uses the following encoding: utf-8
+from benchpress import util
+import numpy as np
+import math
+
+#import numba
+
+def window(B,a=0.37):
+    assert (len(B.shape) == 2)
+    assert (B.shape[0] == B.shape[1])
+    n = B.shape[0]
+    wl = np.ones_like(B[0])
+    b = int(np.ceil((a * (n-1) / 2)))
+    wl[:b]  =  0.5 * (1 + np.cos(math.pi*(2 * np.arange(b) / (a * (n-1)) - 1)))
+    wl[-b:] =  0.5 * (1 + np.cos(math.pi*(2 * np.arange(b-1,-1,-1) / (a * (n-1)) - 1)))
+    wl *= wl
+    w = np.sqrt(wl+wl[:,None])
+    return B*w
+
+
+def calcB_vec(B_x0, alpha=0.0,
+          x_min = 0.0, x_max = 0.25,
+          y_min = 0.0, y_max = 1.0,
+          z_min = 0.0, z_max = 1.0):
+
+    n = len(B_x0)
+    x = np.linspace(x_min,x_max,num=n, endpoint=False).astype(B_x0.dtype,copy=False)
+    y = np.linspace(y_min,y_max,num=n).astype(B_x0.dtype,copy=False)
+    z = np.linspace(z_min,z_max,num=n).astype(B_x0.dtype,copy=False)
+    u = np.arange(n,dtype=B_x0.dtype)
+
+    #Making C
+    C = 4.0 / (n-1.0)**2 * np.sum(np.sum((B_x0 * np.sin(math.pi/y_max * u * y[:,None])[:,:,None])[:,None] * np.sin(math.pi/z_max * u * z[:,None])[:,None],-1),-1)
+    l = np.pi**2 * ((u**2 / y_max)[:,None] + (u**2 / z_max))
+    l[0,0] = 1.0
+    r = np.sqrt(l - alpha**2)
+
+    # Calculating B
+    sincos = np.sin(math.pi/y_max * u * y[:, None])[:, None, :, None] * (u * np.cos(math.pi/z_max * u * z[:,None]))[None, :, None, :]
+    cossin = (u * np.cos(math.pi/y_max * u * y[:,None]))[:, None, :, None] * np.sin(math.pi/z_max * u * z[:,None])[None, :, None, :]
+    temp_x = C * np.sin(math.pi/y_max * u * y[:,None])[:, None, :, None] * np.sin(math.pi/z_max * u * z[:,None])[None, :, None, :]
+    temp_y = C / l * (alpha * math.pi / z_max * sincos - r * math.pi / y_max * cossin)
+    temp_z = C / l * (alpha * math.pi / y_max * cossin + r * math.pi / z_max * sincos)
+    exprx = np.exp((-r * x[:, None, None]))
+
+    Bx = np.sum(np.sum(temp_x * exprx[:,None,None],-1),-1)
+    By = np.sum(np.sum(temp_y * exprx[:,None,None],-1),-1)
+    Bz = np.sum(np.sum(temp_z * exprx[:,None,None],-1),-1)
+    return (Bx, By, Bz)
+
+
+def calcB_loop(B_x0, alpha=0.0,
+          x_min = 0.0, x_max = 0.25,
+          y_min = 0.0, y_max = 1.0,
+          z_min = 0.0, z_max = 1.0):
+
+    n = len(B_x0)
+
+    x = np.linspace(x_min,x_max,num=n,endpoint=False).astype(B_x0.dtype,copy=False)
+    y = np.linspace(y_min,y_max,num=n).astype(B_x0.dtype,copy=False)
+    z = np.linspace(z_min,z_max,num=n).astype(B_x0.dtype,copy=False)
+    u = np.arange(n,dtype=B_x0.dtype)
+
+    # Making C
+    C = 4.0 / (n-1.0)**2 * np.sum(np.sum((B_x0 * np.sin(math.pi/y_max * u * y[:,None])[:,:,None])[:,None] * np.sin(math.pi/z_max * u * z[:,None])[:,None],-1),-1)
+    l = np.pi**2 * ((u**2 / y_max)[:,None] + (u**2 / z_max))
+    l[0,0] = 1.0
+    r = np.sqrt(l - alpha**2)
+
+    # Calculating B
+    Bx = np.empty((n,n,n),dtype=B_x0.dtype)
+    By = np.empty((n,n,n),dtype=B_x0.dtype)
+    Bz = np.empty((n,n,n),dtype=B_x0.dtype)
+    for i in range(n):
+        for j in range(n):
+            Bx[:, i, j] = 0
+            By[:, i, j] = 0
+            Bz[:, i, j] = 0
+            temp_x = np.empty((n, n), dtype=B_x0.dtype)
+            temp_y = np.empty((n, n), dtype=B_x0.dtype)
+            temp_z = np.empty((n, n), dtype=B_x0.dtype)
+            exprx = np.empty((n, n, n), dtype=B_x0.dtype)
+            for k in range(n):
+                for m in range(n):
+                    sincos = np.sin(np.pi * u[k] * y[i] / y_max) * (u[m] * np.cos(np.pi * u[m] * z[j] / z_max))
+                    cossin = (u[k] * np.cos(np.pi * u[k] * y[i] / y_max)) * (np.sin(np.pi * u[m] * z[j] / z_max))
+                    temp_x[k,m] = C[k,m] * (np.sin(np.pi * u[k] * y[i] / y_max) * (np.sin(np.pi * u[m] * z[j] / z_max)))
+                    temp_y[k,m] = C[k,m] / l[k,m] * (alpha * np.pi / z_max * sincos - r[k,m] * np.pi / y_max * cossin)
+                    temp_z[k,m] = C[k,m] / l[k,m] * (alpha * np.pi / y_max * cossin + r[k,m] * np.pi / z_max * sincos)
+                    for q in range(n):
+                        exprx[k,m,q] = np.exp(-r[m,q] * x[k])
+            for k in range(n):
+                for m in range(n):
+                    for q in range(n):
+                        Bx[k, i, j] += temp_x[m, q] * exprx[k, m, q]
+                        By[k, i, j] += temp_y[m, q] * exprx[k, m, q]
+                        Bz[k, i, j] += temp_z[m, q] * exprx[k, m, q]
+    return (Bx, By, Bz)
+
+
+def cal_B(B_x0, Bx, By, Bz, u, y, y_max, z, z_max, alpha, C, l, r, x):
+    n = B_x0.shape[0]
+    for i in range(n):
+        for j in range(n):
+            Bx[:, i, j] = 0
+            By[:, i, j] = 0
+            Bz[:, i, j] = 0
+            temp_x = np.empty((n, n), dtype=B_x0.dtype)
+            temp_y = np.empty((n, n), dtype=B_x0.dtype)
+            temp_z = np.empty((n, n), dtype=B_x0.dtype)
+            exprx = np.empty((n, n, n), dtype=B_x0.dtype)
+            for k in range(n):
+                for m in range(n):
+                    sincos = np.sin(np.pi * u[k] * y[i] / y_max) * (u[m] * np.cos(np.pi * u[m] * z[j] / z_max))
+                    cossin = (u[k] * np.cos(np.pi * u[k] * y[i] / y_max)) * (np.sin(np.pi * u[m] * z[j] / z_max))
+                    temp_x[k,m] = C[k,m] * (np.sin(np.pi * u[k] * y[i] / y_max) * (np.sin(np.pi * u[m] * z[j] / z_max)))
+                    temp_y[k,m] = C[k,m] / l[k,m] * (alpha * np.pi / z_max * sincos - r[k,m] * np.pi / y_max * cossin)
+                    temp_z[k,m] = C[k,m] / l[k,m] * (alpha * np.pi / y_max * cossin + r[k,m] * np.pi / z_max * sincos)
+                    for q in range(n):
+                        exprx[k,m,q] = np.exp(-r[m,q] * x[k])
+            for k in range(n):
+                for m in range(n):
+                    for q in range(n):
+                        Bx[k, i, j] += temp_x[m, q] * exprx[k, m, q]
+                        By[k, i, j] += temp_y[m, q] * exprx[k, m, q]
+                        Bz[k, i, j] += temp_z[m, q] * exprx[k, m, q]
+
+
+def calcB_numba(B_x0, alpha=0.0,
+          x_min = 0.0, x_max = 0.25,
+          y_min = 0.0, y_max = 1.0,
+          z_min = 0.0, z_max = 1.0):
+
+    n = len(B_x0)
+
+    x = np.linspace(x_min,x_max,num=n,endpoint=False).astype(B_x0.dtype,copy=False)
+    y = np.linspace(y_min,y_max,num=n).astype(B_x0.dtype,copy=False)
+    z = np.linspace(z_min,z_max,num=n).astype(B_x0.dtype,copy=False)
+    u = np.arange(n,dtype=B_x0.dtype)
+
+    # Making C
+    C = 4.0 / (n-1.0)**2 * np.sum(np.sum((B_x0 * np.sin(math.pi/y_max * u * y[:,None])[:,:,None])[:,None] * np.sin(math.pi/z_max * u * z[:,None])[:,None],-1),-1)
+    l = np.pi**2 * ((u**2 / y_max)[:,None] + (u**2 / z_max))
+    l[0,0] = 1.0
+    r = np.sqrt(l - alpha**2)
+
+    # Calculating B
+    Bx = np.empty((n,n,n),dtype=B_x0.dtype)
+    By = np.empty((n,n,n),dtype=B_x0.dtype)
+    Bz = np.empty((n,n,n),dtype=B_x0.dtype)
+    cal_B(B_x0, Bx, By, Bz, u, y, y_max, z, z_max, alpha, C, l, r, x)
+    return (Bx, By, Bz)
+
+
+def main():
+
+    B = util.Benchmark()
+
+    if B.inputfn is None:
+        B_x0 = B.random_array((B.size[0],B.size[1]), dtype=B.dtype)
+    else:
+        inputfn = B.inputfn if B.inputfn else '../idl_input-float64_512*512.npz'
+        sd = { 512:1, 256:2, 128:4, 64:8, 32:16, 16:32, 8:64}
+        try:
+            h = sd[B.size[0]]
+            w = sd[B.size[1]]
+        except KeyError:
+            raise ValueError('Only valid sizes are: '+str(sd.keys()))
+        B_x0 = B.load_array(inputfn, 'input', dtype=B.dtype)[::h,::w]
+
+    B.start()
+    for _ in range(B.size[2]):
+        Rx, Ry, Rz = calcB_loop(window(B_x0.copy()))
+        Rx2, Ry2, Rz2 = calcB_numba(window(B_x0.copy()))
+        assert np.allclose(Rx, Rx2)
+    B.stop()
+    B.pprint()
+
+    if B.outputfn:
+        R = Rx+Ry+Rz
+        B.tofile(B.outputfn, {'res': R, 'res_x': Rx, 'res_y': Ry, 'res_z': Rz})
+
+if __name__ == '__main__':
+    main()

benchmarks/magnetic_field_extrapolation/python_numba/readme.rst

@@ -0,0 +1 @@
+Magnetic Field Extrapolation

benchmarks/magnetic_field_extrapolation/python_numpy/magnetic_field_extrapolation.py

@@ -16,78 +16,38 @@ def window(B,a=0.37):
     w = np.sqrt(wl+wl[:,None])
     return B*w
 
-def calcB(B, alpha=1.0,
+
+def calcB_vec(B_x0, alpha=0.0,
           x_min = 0.0, x_max = 0.25,
           y_min = 0.0, y_max = 1.0,
           z_min = 0.0, z_max = 1.0):
 
-    n = len(B)
-    x = np.linspace(x_min,x_max,num=n,endpoint=False)
-    y = np.linspace(y_min,y_max,num=n)
-    z = np.linspace(z_min,z_max,num=n)
-    u = np.arange(n,dtype=B.dtype)
-
-
-    uy = math.pi/y_max * u * y[:,None]
-    uz = math.pi/z_max * u * z[:,None]
-
-    sinuy = np.sin(uy)
-    sinuz = np.sin(uz)
+    n = len(B_x0)
+    x = np.linspace(x_min,x_max,num=n, endpoint=False).astype(B_x0.dtype,copy=False)
+    y = np.linspace(y_min,y_max,num=n).astype(B_x0.dtype,copy=False)
+    z = np.linspace(z_min,z_max,num=n).astype(B_x0.dtype,copy=False)
+    u = np.arange(n,dtype=B_x0.dtype)
 
-#    C =  4.0 / (n-1.0)**2 * np.sum(np.sum((B * sinuy[:,:,None])[:,None] * sinuz[:,None],-1),-1)
-    C = np.empty((n,n,n,n))
-    C[:] = sinuy[:,None,:,None]
-    C *= B
-    C *= sinuz[:,None]
-    C = 4.0 / (n-1.0)**2 * np.sum(np.sum(C,-1),-1)
-
-    l = math.pi**2 * ((u**2 / y_max)[:,None] + (u**2 / z_max))
+    # Making C
+    C = 4.0 / (n-1.0)**2 * np.sum(np.sum((B_x0 * np.sin(math.pi/y_max * u * y[:,None])[:,:,None])[:,None] * np.sin(math.pi/z_max * u * z[:,None])[:,None],-1),-1)
+    l = np.pi**2 * ((u**2 / y_max)[:,None] + (u**2 / z_max))
     l[0,0] = 1.0
     r = np.sqrt(l - alpha**2)
 
-    C_l = C / l
-
-    d5 = (n,n,n,n,n)
-
-#    sincos = sinuy[:,None,:,None] * (u * np.cos(uz))[None,:,None,:]
-    sincos = np.empty(d5)
-    sincos[:] = sinuy[:,None,:,None]
-    sincos *= (u * np.cos(math.pi/z_max * u * z[:,None]))[None,:,None,:]
-
-#    cossin = (u * np.cos(uy))[:,None,:,None] * sinuz[None,:,None,:]
-    cossin = np.empty(d5)
-    cossin[:] = sinuz[None,:,None,:]
-    cossin *= (u * np.cos(math.pi/y_max * u * y[:,None]))[:,None,:,None]
-
-#    exprx = np.exp((-r * x[:,None,None]))[:,None,None]
-    exprx = np.empty(d5)
-    exprx[:] = -x[:,None,None,None,None]
-    exprx *= r
-    exprx = np.exp(exprx)
-
-#    temp_x = C * sinuy[:,None,:,None] * sinuz[None,:,None,:]
-    temp_x = np.empty(d5)
-    temp_x[:] = sinuy[:,None,:,None]
-    temp_x *= sinuz[None,:,None,:]
-    temp_x *= C
-
-#    temp_y = C / l * (alpha * math.pi / z_max * sincos - r * math.pi / y_max * cossin)
-    temp_y = -(math.pi/y_max) * cossin
-    temp_y *= r
-    temp_y += (alpha*math.pi/z_max) * sincos
-    temp_y *= C_l
-
-#    temp_z = C / l * (alpha * math.pi / y_max * cossin + r * math.pi / z_max * sincos)
-    temp_z = (math.pi/z_max) * sincos
-    temp_z *= r
-    temp_z += (alpha*math.pi/y_max) * cossin
-    temp_z *= C_l
-
-    Bx = np.sum(np.sum(temp_x * exprx,-1),-1)
-    By = np.sum(np.sum(temp_y * exprx,-1),-1)
-    Bz = np.sum(np.sum(temp_z * exprx,-1),-1)
+    # Calculating B
+    sincos = np.sin(math.pi/y_max * u * y[:, None])[:, None, :, None] * (u * np.cos(math.pi/z_max * u * z[:,None]))[None, :, None, :]
+    cossin = (u * np.cos(math.pi/y_max * u * y[:,None]))[:, None, :, None] * np.sin(math.pi/z_max * u * z[:,None])[None, :, None, :]
+    temp_x = C * np.sin(math.pi/y_max * u * y[:,None])[:, None, :, None] * np.sin(math.pi/z_max * u * z[:,None])[None, :, None, :]
+    temp_y = C / l * (alpha * math.pi / z_max * sincos - r * math.pi / y_max * cossin)
+    temp_z = C / l * (alpha * math.pi / y_max * cossin + r * math.pi / z_max * sincos)
+    exprx = np.exp((-r * x[:, None, None]))
+
+    Bx = np.sum(np.sum(temp_x * exprx[:,None,None],-1),-1)
+    By = np.sum(np.sum(temp_y * exprx[:,None,None],-1),-1)
+    Bz = np.sum(np.sum(temp_z * exprx[:,None,None],-1),-1)
     return (Bx, By, Bz)
 
+
 def main():
 
     B = util.Benchmark()