Skip to content

Commit

Permalink
the tutorials are all done now
Browse files Browse the repository at this point in the history
  • Loading branch information
@jpivarski
jpivarski committed Dec 20, 2024
1 parent e91d929 commit d8efb14
Showing 1 changed file with 176 additions and 4 deletions.
180 changes: 176 additions & 4 deletions docs/src/numba.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,19 +10,191 @@
},
{
"cell_type": "markdown",
"id": "26a7599a-5b33-4f6d-85d2-788295abe176",
"id": "03064225-ea6b-4f6b-a5bf-08a47b62350e",
"metadata": {},
"source": [
"First, [install](../index.md#installation) and import Vector."
"First, [install](../index.md#installation) and import Vector and [Numba](https://numba.pydata.org/)."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "2ccf22f1-fc21-4afe-8e46-baef45a1e85b",
"id": "df8bea35-cf3c-4f87-8e3d-8ffe06264811",
"metadata": {},
"outputs": [],
"source": []
"source": [
"import vector\n",
"import numba as nb\n",
"import numpy as np"
]
},
{
"cell_type": "markdown",
"id": "7fff9620-227b-4e5e-a211-4daf1fb96298",
"metadata": {},
"source": [
"Numba is a just-in-time (JIT) compiler for a mathematically relevant subset of NumPy and Python. It allows you to write fast code without leaving the Python environment. The drawback of Numba is that it can only compile code blocks involving objects and functions that it recognizes.\n",
"\n",
"The Vector library includes extensions to inform Numba about [vector objects](object.md), ~~[arrays of vectors](numpy.md)~~, and [arrays of Awkward Arrays](awkward.md). At the time of writing, the implementation of vector NumPy arrays is incomplete (see issue [#43](https://github.com/scikit-hep/vector/issues/43)).\n",
"\n",
"Consider the following function:"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "0e7f7c5c-2ce3-4f55-9159-363910fac6b9",
"metadata": {},
"outputs": [],
"source": [
"@nb.njit\n",
"def compute_mass(v1, v2):\n",
" return (v1 + v2).mass"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "16d0fe04-fc8d-4d1f-bb68-0f869c1d2bdb",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"8.0"
]
},
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"compute_mass(vector.obj(px=1, py=2, pz=3, E=4), vector.obj(px=-1, py=-2, pz=-3, E=4))"
]
},
{
"cell_type": "markdown",
"id": "69c36f6d-3b54-4db8-9890-9245d64fefe7",
"metadata": {},
"source": [
"When the two `MomentumObject4D` objects are passed as arguments, Numba recognizes them and replaces the Python objects with low-level structs. When it compiles the function, it recognizes `+` as the 4D `add` function and recognizes `.mass` as the `tau` component of the result.\n",
"\n",
"Although this demonstrates that Numba can manipulate vector objects, there is no performance advantage (and a likely disadvantage) to compiling a calculation on just a few vectors. The advantage comes when many vectors are involved, in arrays."
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "63f614a2-969d-465d-8d74-a44f6f77eab4",
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<pre>[[],\n",
" [{x: 1.96, y: -0.654, z: -0.443, t: 8.88}, ..., {x: -0.0976, y: 0.877, ...}],\n",
" [{x: -0.543, y: 1.82, z: 0.177, t: 10.1}],\n",
" [{x: 0.0705, y: 0.816, z: -0.607, t: 10.1}, ..., {x: -0.756, y: -2.23, ...}],\n",
" [{x: -1.61, y: -1.14, z: 0.312, t: 10.2}, ..., {x: -0.878, y: -0.803, ...}],\n",
" [{x: -1.4, y: -1.09, z: 1.13, t: 9.03}],\n",
" [{x: 0.634, y: -0.0388, z: -0.44, t: 10.5}],\n",
" [{x: -0.767, y: 0.0841, z: -0.344, t: 9.34}],\n",
" [{x: -1.89, y: 0.295, z: -1.95, t: 10.1}],\n",
" [{x: -1.85, y: -0.832, z: -0.816, t: 8.2}],\n",
" ...,\n",
" [{x: 0.496, y: 0.791, z: 1.67, t: 10.8}],\n",
" [{x: -1.09, y: 1.42, z: -0.299, t: 11.4}],\n",
" [{x: -0.885, y: -0.537, z: -0.0475, t: 9.95}, {x: -0.615, y: 1.19, ...}],\n",
" [{x: -0.503, y: -3.06, z: 0.893, t: 9.61}, ..., {x: 0.263, y: -0.493, ...}],\n",
" [{x: 2.16, y: -1.11, z: 0.696, t: 11.2}, {x: 0.824, y: 0.604, ...}],\n",
" [],\n",
" [{x: 1.55, y: 0.497, z: -0.764, t: 10.3}, {x: 0.311, y: 0.545, ...}],\n",
" [],\n",
" [{x: 0.0877, y: 0.334, z: -1.5, t: 10.5}, {x: 1.26, y: 1.94, ...}]]\n",
"---------------------------------------------------------------------------------------------\n",
"backend: cpu\n",
"nbytes: 3.0 kB\n",
"type: 50 * var * Momentum4D[\n",
" x: float64,\n",
" y: float64,\n",
" z: float64,\n",
" t: float64\n",
"]</pre>"
],
"text/plain": [
"<MomentumArray4D [[], ..., [{x: 0.0877, ...}, ...]] type='50 * var * Moment...'>"
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# This is still not a large number. You want millions.\n",
"array = vector.Array(\n",
" [\n",
" [\n",
" dict(\n",
" {x: np.random.normal(0, 1) for x in (\"px\", \"py\", \"pz\")},\n",
" E=np.random.normal(10, 1),\n",
" )\n",
" for inner in range(np.random.poisson(1.5))\n",
" ]\n",
" for outer in range(50)\n",
" ]\n",
")\n",
"array"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "f6edd6c9-a423-4662-bcae-c685686633c8",
"metadata": {},
"outputs": [],
"source": [
"@nb.njit\n",
"def compute_masses(array):\n",
" out = np.empty(len(array), np.float64)\n",
" for i, event in enumerate(array):\n",
" total = vector.obj(px=0.0, py=0.0, pz=0.0, E=0.0)\n",
" for vec in event:\n",
" total = total + vec\n",
" out[i] = total.mass\n",
" return out"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "fa854373-aeba-4274-8ab8-c1c85d5a931b",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"array([ 0. , 29.59334472, 9.96825092, 29.98088102, 30.022318 ,\n",
" 8.77904697, 10.52123699, 9.30633562, 9.69685471, 7.90724598,\n",
" 8.22083106, 12.10506017, 9.14678916, 42.27902654, 10.08573923,\n",
" 9.98477624, 0. , 22.79438116, 29.12919935, 32.07848403,\n",
" 29.6325778 , 0. , 8.11447927, 39.07959905, 30.54027295,\n",
" 18.39609413, 21.15982485, 18.62455179, 7.75621961, 19.52526457,\n",
" 19.52907948, 20.40038164, 8.84074954, 9.18655937, 30.22076618,\n",
" 8.47538375, 10.42253874, 17.84485932, 0. , 41.55064913,\n",
" 0. , 10.60246245, 11.24522316, 19.32603825, 31.28879051,\n",
" 20.49852241, 0. , 20.21304572, 0. , 20.76651039])"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"compute_masses(array)"
]
}
],
"metadata": {
Expand Down

0 comments on commit d8efb14

Please sign in to comment.