1.17。示例

原文: http://numba.pydata.org/numba-doc/latest/user/examples.html

1.17.1。 Mandelbrot

  1. #! /usr/bin/env python
  2. # -*- coding: utf-8 -*-
  3. from __future__ import print_function, division, absolute_import
  4. from timeit import default_timer as timer
  5. from matplotlib.pylab import imshow, jet, show, ion
  6. import numpy as np
  7. from numba import jit
  8. @jit
  9. def mandel(x, y, max_iters):
  10. """
  11. Given the real and imaginary parts of a complex number,
  12. determine if it is a candidate for membership in the Mandelbrot
  13. set given a fixed number of iterations.
  14. """
  15. i = 0
  16. c = complex(x,y)
  17. z = 0.0j
  18. for i in range(max_iters):
  19. z = z*z + c
  20. if (z.real*z.real + z.imag*z.imag) >= 4:
  21. return i
  22. return 255
  23. @jit
  24. def create_fractal(min_x, max_x, min_y, max_y, image, iters):
  25. height = image.shape[0]
  26. width = image.shape[1]
  27. pixel_size_x = (max_x - min_x) / width
  28. pixel_size_y = (max_y - min_y) / height
  29. for x in range(width):
  30. real = min_x + x * pixel_size_x
  31. for y in range(height):
  32. imag = min_y + y * pixel_size_y
  33. color = mandel(real, imag, iters)
  34. image[y, x] = color
  35. return image
  36. image = np.zeros((500 * 2, 750 * 2), dtype=np.uint8)
  37. s = timer()
  38. create_fractal(-2.0, 1.0, -1.0, 1.0, image, 20)
  39. e = timer()
  40. print(e - s)
  41. imshow(image)
  42. #jet()
  43. #ion()
  44. show()

1.17.2。移动平均线

  1. #!/usr/bin/env python
  2. """
  3. A moving average function using @guvectorize.
  4. """
  5. import numpy as np
  6. from numba import guvectorize
  7. @guvectorize(['void(float64[:], intp[:], float64[:])'], '(n),()->(n)')
  8. def move_mean(a, window_arr, out):
  9. window_width = window_arr[0]
  10. asum = 0.0
  11. count = 0
  12. for i in range(window_width):
  13. asum += a[i]
  14. count += 1
  15. out[i] = asum / count
  16. for i in range(window_width, len(a)):
  17. asum += a[i] - a[i - window_width]
  18. out[i] = asum / count
  19. arr = np.arange(20, dtype=np.float64).reshape(2, 10)
  20. print(arr)
  21. print(move_mean(arr, 3))

1.17.3。多线程

下面的代码展示了使用 nogil 功能时潜在的性能提升。例如,在 4 核机器上,我打印出以下结果:

  1. numpy (1 thread) 145 ms
  2. numba (1 thread) 128 ms
  3. numba (4 threads) 35 ms

注意

在 Python 3 下,您可以使用标准的 concurrent.futures 模块,而不是手工生成线程和调度任务。

  1. #!/usr/bin/env python
  2. from __future__ import print_function, division, absolute_import
  3. import math
  4. import threading
  5. from timeit import repeat
  6. import numpy as np
  7. from numba import jit
  8. nthreads = 4
  9. size = 10**6
  10. def func_np(a, b):
  11. """
  12. Control function using Numpy.
  13. """
  14. return np.exp(2.1 * a + 3.2 * b)
  15. @jit('void(double[:], double[:], double[:])', nopython=True, nogil=True)
  16. def inner_func_nb(result, a, b):
  17. """
  18. Function under test.
  19. """
  20. for i in range(len(result)):
  21. result[i] = math.exp(2.1 * a[i] + 3.2 * b[i])
  22. def timefunc(correct, s, func, *args, **kwargs):
  23. """
  24. Benchmark *func* and print out its runtime.
  25. """
  26. print(s.ljust(20), end=" ")
  27. # Make sure the function is compiled before we start the benchmark
  28. res = func(*args, **kwargs)
  29. if correct is not None:
  30. assert np.allclose(res, correct), (res, correct)
  31. # time it
  32. print('{:>5.0f} ms'.format(min(repeat(lambda: func(*args, **kwargs),
  33. number=5, repeat=2)) * 1000))
  34. return res
  35. def make_singlethread(inner_func):
  36. """
  37. Run the given function inside a single thread.
  38. """
  39. def func(*args):
  40. length = len(args[0])
  41. result = np.empty(length, dtype=np.float64)
  42. inner_func(result, *args)
  43. return result
  44. return func
  45. def make_multithread(inner_func, numthreads):
  46. """
  47. Run the given function inside *numthreads* threads, splitting its
  48. arguments into equal-sized chunks.
  49. """
  50. def func_mt(*args):
  51. length = len(args[0])
  52. result = np.empty(length, dtype=np.float64)
  53. args = (result,) + args
  54. chunklen = (length + numthreads - 1) // numthreads
  55. # Create argument tuples for each input chunk
  56. chunks = [[arg[i * chunklen:(i + 1) * chunklen] for arg in args]
  57. for i in range(numthreads)]
  58. # Spawn one thread per chunk
  59. threads = [threading.Thread(target=inner_func, args=chunk)
  60. for chunk in chunks]
  61. for thread in threads:
  62. thread.start()
  63. for thread in threads:
  64. thread.join()
  65. return result
  66. return func_mt
  67. func_nb = make_singlethread(inner_func_nb)
  68. func_nb_mt = make_multithread(inner_func_nb, nthreads)
  69. a = np.random.rand(size)
  70. b = np.random.rand(size)
  71. correct = timefunc(None, "numpy (1 thread)", func_np, a, b)
  72. timefunc(correct, "numba (1 thread)", func_nb, a, b)
  73. timefunc(correct, "numba (%d threads)" % nthreads, func_nb_mt, a, b)