{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Based on https://github.com/jpivarski/python-numpy-mini-course/blob/evaluated/6-compilers.ipynb ."
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [
{
"data": {
"application/javascript": [
"\n",
" require(['notebook/js/codecell'], function(cc) {\n",
" cc.CodeCell.options_default.highlight_modes['magic_text/x-c++src'] =\n",
" {reg: [/^\\s*%%pybind11/]};\n",
" });\n",
" "
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"\n",
" \n",
" "
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"%load_ext ipybind"
]
},
{
"cell_type": "code",
"execution_count": 35,
"metadata": {},
"outputs": [],
"source": [
"import numpy\n",
"import numba"
]
},
{
"cell_type": "code",
"execution_count": 32,
"metadata": {},
"outputs": [],
"source": [
"%%pybind11\n",
"\n",
"#include \n",
"#include \n",
"#include \n",
"\n",
"py::array_t quick(int height, int width, int maxiterations) {\n",
" \n",
" py::array_t fractal({height, width});\n",
" \n",
" auto fractal_uc = fractal.mutable_unchecked<2>();\n",
" \n",
" for (int h = 0; h < height; h++) {\n",
" for (int w = 0; w < width; w++) {\n",
" \n",
" std::complex ci{\n",
" double(h-1)/height - 1,\n",
" 1.5 * (double(w-1)/width - 1)};\n",
" \n",
" std::complex z = ci;\n",
" fractal_uc(h,w) = maxiterations;\n",
" for (int i = 0; i < maxiterations; i++) {\n",
" z = z * z + ci;\n",
" if (std::abs(z) > 2) {\n",
" fractal_uc(h, w) = i;\n",
" break;\n",
" }\n",
" }\n",
" }\n",
" }\n",
" \n",
" return fractal;\n",
"}\n",
"\n",
"PYBIND11_MODULE(py11fractal, m) {\n",
" m.def(\"quick\", quick);\n",
"}\n"
]
},
{
"cell_type": "code",
"execution_count": 34,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 7.45 s, sys: 230 ms, total: 7.68 s\n",
"Wall time: 7.71 s\n"
]
},
{
"data": {
"text/plain": [
"array([[ 0, 0, 0, ..., 20, 20, 20],\n",
" [ 0, 0, 0, ..., 20, 20, 20],\n",
" [ 0, 0, 0, ..., 20, 20, 20],\n",
" ...,\n",
" [ 0, 0, 0, ..., 20, 20, 20],\n",
" [ 0, 0, 0, ..., 20, 20, 20],\n",
" [ 0, 0, 0, ..., 20, 20, 20]], dtype=int32)"
]
},
"execution_count": 34,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"%%time\n",
"quick(8000, 12000)"
]
},
{
"cell_type": "code",
"execution_count": 59,
"metadata": {},
"outputs": [],
"source": [
"@numba.vectorize\n",
"def as_ufunc(c, maxiterations):\n",
" z = c\n",
" for i in range(maxiterations):\n",
" z = z**2 + c\n",
" if abs(z) > 2:\n",
" return i\n",
" return maxiterations\n",
"\n",
"def run_numba_2(height, width, maxiterations=20):\n",
" y, x = numpy.ogrid[-1:0:height*1j, -1.5:0:width*1j]\n",
" c = x + y*1j\n",
" return as_ufunc(c, maxiterations)"
]
},
{
"cell_type": "code",
"execution_count": 61,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 7.27 s, sys: 1.13 s, total: 8.4 s\n",
"Wall time: 8.5 s\n"
]
},
{
"data": {
"text/plain": [
"array([[ 0, 0, 0, ..., 10, 10, 20],\n",
" [ 0, 0, 0, ..., 9, 10, 10],\n",
" [ 0, 0, 0, ..., 9, 9, 9],\n",
" ...,\n",
" [20, 20, 20, ..., 20, 20, 20],\n",
" [20, 20, 20, ..., 20, 20, 20],\n",
" [20, 20, 20, ..., 20, 20, 20]])"
]
},
"execution_count": 61,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"%%time\n",
"run_numba_2(8000, 12000)"
]
},
{
"cell_type": "code",
"execution_count": 57,
"metadata": {},
"outputs": [],
"source": [
"@numba.njit\n",
"def run_numba(height, width, maxiterations):\n",
" fractal = numpy.empty((height, width), dtype=numpy.int32)\n",
" for h in range(height):\n",
" for w in range(width):\n",
" c = ((h-1)/height - 1) + 1.5j*((w-1)/width - 1)\n",
" z = c\n",
" fractal[h, w] = maxiterations\n",
" \n",
" for i in range(maxiterations):\n",
" z = z**2 + c\n",
" if abs(z) > 2:\n",
" fractal[h, w] = i\n",
" break\n",
" return fractal\n"
]
},
{
"cell_type": "code",
"execution_count": 62,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 5.92 s, sys: 177 ms, total: 6.1 s\n",
"Wall time: 6.1 s\n"
]
},
{
"data": {
"text/plain": [
"array([[ 0, 0, 0, ..., 20, 20, 20],\n",
" [ 0, 0, 0, ..., 20, 20, 20],\n",
" [ 0, 0, 0, ..., 20, 20, 20],\n",
" ...,\n",
" [ 0, 0, 0, ..., 20, 20, 20],\n",
" [ 0, 0, 0, ..., 20, 20, 20],\n",
" [ 0, 0, 0, ..., 20, 20, 20]], dtype=int32)"
]
},
"execution_count": 62,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"%%time\n",
"run_numba(8000, 12000, 20)"
]
},
{
"cell_type": "code",
"execution_count": 65,
"metadata": {},
"outputs": [],
"source": [
"@numba.njit(parallel=True)\n",
"def run_numba_p(height, width, maxiterations):\n",
" fractal = numpy.empty((height, width), dtype=numpy.int32)\n",
" for h in numba.prange(height):\n",
" for w in range(width):\n",
" c = ((h-1)/height - 1) + 1.5j*((w-1)/width - 1)\n",
" z = c\n",
" fractal[h, w] = maxiterations\n",
" \n",
" for i in range(maxiterations):\n",
" z = z**2 + c\n",
" if abs(z) > 2:\n",
" fractal[h, w] = i\n",
" break\n",
" return fractal"
]
},
{
"cell_type": "code",
"execution_count": 66,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 9.57 s, sys: 299 ms, total: 9.87 s\n",
"Wall time: 3.07 s\n"
]
},
{
"data": {
"text/plain": [
"array([[ 0, 0, 0, ..., 20, 20, 20],\n",
" [ 0, 0, 0, ..., 20, 20, 20],\n",
" [ 0, 0, 0, ..., 20, 20, 20],\n",
" ...,\n",
" [ 0, 0, 0, ..., 20, 20, 20],\n",
" [ 0, 0, 0, ..., 20, 20, 20],\n",
" [ 0, 0, 0, ..., 20, 20, 20]], dtype=int32)"
]
},
"execution_count": 66,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"%%time\n",
"run_numba_p(8000, 12000, 20)"
]
},
{
"cell_type": "code",
"execution_count": 67,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 9.02 s, sys: 326 ms, total: 9.35 s\n",
"Wall time: 2.86 s\n"
]
},
{
"data": {
"text/plain": [
"array([[ 0, 0, 0, ..., 20, 20, 20],\n",
" [ 0, 0, 0, ..., 20, 20, 20],\n",
" [ 0, 0, 0, ..., 20, 20, 20],\n",
" ...,\n",
" [ 0, 0, 0, ..., 20, 20, 20],\n",
" [ 0, 0, 0, ..., 20, 20, 20],\n",
" [ 0, 0, 0, ..., 20, 20, 20]], dtype=int32)"
]
},
"execution_count": 67,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"%%time\n",
"run_numba_p(8000, 12000, 20)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Anaconda 3.7",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.0"
}
},
"nbformat": 4,
"nbformat_minor": 2
}