1.17。示例
原文: http://numba.pydata.org/numba-doc/latest/user/examples.html
1.17.1。 Mandelbrot
#! /usr/bin/env python# -*- coding: utf-8 -*-from __future__ import print_function, division, absolute_importfrom timeit import default_timer as timerfrom matplotlib.pylab import imshow, jet, show, ionimport numpy as npfrom numba import jit@jitdef mandel(x, y, max_iters):"""Given the real and imaginary parts of a complex number,determine if it is a candidate for membership in the Mandelbrotset given a fixed number of iterations."""i = 0c = complex(x,y)z = 0.0jfor i in range(max_iters):z = z*z + cif (z.real*z.real + z.imag*z.imag) >= 4:return ireturn 255@jitdef create_fractal(min_x, max_x, min_y, max_y, image, iters):height = image.shape[0]width = image.shape[1]pixel_size_x = (max_x - min_x) / widthpixel_size_y = (max_y - min_y) / heightfor x in range(width):real = min_x + x * pixel_size_xfor y in range(height):imag = min_y + y * pixel_size_ycolor = mandel(real, imag, iters)image[y, x] = colorreturn imageimage = np.zeros((500 * 2, 750 * 2), dtype=np.uint8)s = timer()create_fractal(-2.0, 1.0, -1.0, 1.0, image, 20)e = timer()print(e - s)imshow(image)#jet()#ion()show()
1.17.2。移动平均线
#!/usr/bin/env python"""A moving average function using @guvectorize."""import numpy as npfrom numba import guvectorize@guvectorize(['void(float64[:], intp[:], float64[:])'], '(n),()->(n)')def move_mean(a, window_arr, out):window_width = window_arr[0]asum = 0.0count = 0for i in range(window_width):asum += a[i]count += 1out[i] = asum / countfor i in range(window_width, len(a)):asum += a[i] - a[i - window_width]out[i] = asum / countarr = np.arange(20, dtype=np.float64).reshape(2, 10)print(arr)print(move_mean(arr, 3))
1.17.3。多线程
下面的代码展示了使用 nogil 功能时潜在的性能提升。例如,在 4 核机器上,我打印出以下结果:
numpy (1 thread) 145 msnumba (1 thread) 128 msnumba (4 threads) 35 ms
注意
在 Python 3 下,您可以使用标准的 concurrent.futures 模块,而不是手工生成线程和调度任务。
#!/usr/bin/env pythonfrom __future__ import print_function, division, absolute_importimport mathimport threadingfrom timeit import repeatimport numpy as npfrom numba import jitnthreads = 4size = 10**6def func_np(a, b):"""Control function using Numpy."""return np.exp(2.1 * a + 3.2 * b)@jit('void(double[:], double[:], double[:])', nopython=True, nogil=True)def inner_func_nb(result, a, b):"""Function under test."""for i in range(len(result)):result[i] = math.exp(2.1 * a[i] + 3.2 * b[i])def timefunc(correct, s, func, *args, **kwargs):"""Benchmark *func* and print out its runtime."""print(s.ljust(20), end=" ")# Make sure the function is compiled before we start the benchmarkres = func(*args, **kwargs)if correct is not None:assert np.allclose(res, correct), (res, correct)# time itprint('{:>5.0f} ms'.format(min(repeat(lambda: func(*args, **kwargs),number=5, repeat=2)) * 1000))return resdef make_singlethread(inner_func):"""Run the given function inside a single thread."""def func(*args):length = len(args[0])result = np.empty(length, dtype=np.float64)inner_func(result, *args)return resultreturn funcdef make_multithread(inner_func, numthreads):"""Run the given function inside *numthreads* threads, splitting itsarguments into equal-sized chunks."""def func_mt(*args):length = len(args[0])result = np.empty(length, dtype=np.float64)args = (result,) + argschunklen = (length + numthreads - 1) // numthreads# Create argument tuples for each input chunkchunks = [[arg[i * chunklen:(i + 1) * chunklen] for arg in args]for i in range(numthreads)]# Spawn one thread per chunkthreads = [threading.Thread(target=inner_func, args=chunk)for chunk in chunks]for thread in threads:thread.start()for thread in threads:thread.join()return resultreturn func_mtfunc_nb = make_singlethread(inner_func_nb)func_nb_mt = make_multithread(inner_func_nb, nthreads)a = np.random.rand(size)b = np.random.rand(size)correct = timefunc(None, "numpy (1 thread)", func_np, a, b)timefunc(correct, "numba (1 thread)", func_nb, a, b)timefunc(correct, "numba (%d threads)" % nthreads, func_nb_mt, a, b)
