Measurements from telescopes are then represented as an array of pixels that encode the pointing of the instrument at each timestamp and the measurement output.
import pandas as pd
import numba
import numpy as np
For simplicity let's assume we have a sky with 50K pixels:
NPIX = 50000
And we have 50 million measurement from our instrument:
NTIME = int(50 * 1e6)
The pointing of our instrument is an array of pixels, random in our sample case:
pixels = np.random.randint(0, NPIX-1, NTIME)
Our data are also random:
timeline = np.random.randn(NTIME)
One of the most common operations is to sum all of our measurements in a sky map, so the value of each pixel in our sky map will be the sum of each individual measurement.
The easiest way is to use the groupby
operation in pandas
:
timeline_pandas = pd.Series(timeline, index=pixels)
timeline_pandas.head()
11105 -0.498438 16106 -1.103009 44723 0.392057 16687 0.086918 15197 1.946979 dtype: float64
%time m = timeline_pandas.groupby(level=0).sum()
CPU times: user 2.49 s, sys: 604 ms, total: 3.09 s Wall time: 3.1 s
We would like to improve the performance of this operation using numba
, which allows to produce automatically C-speed compiled code from pure python functions.
First we need to develop a pure python version of the code, test it, and then have numba
optimize it:
def groupby_python(index, value, output):
for i in range(index.shape[0]):
output[index[i]] += value[i]
m_python = np.zeros_like(m)
%time groupby_python(pixels, timeline, m_python)
CPU times: user 25.8 s, sys: 5 ms, total: 25.8 s Wall time: 25.8 s
np.testing.assert_allclose(m_python, m)
Pure Python is slower than the pandas
version implemented in cython
.
numba.jit
gets an input function and creates an compiled version with does not depend on slow Python calls, this is enforced by nopython=True
, numba
would throw an error if it would not be possible to run in nopython
mode.
groupby_numba = numba.jit(groupby_python, nopython=True)
m_numba = np.zeros_like(m)
%time groupby_numba(pixels, timeline, m_numba)
CPU times: user 411 ms, sys: 7 ms, total: 418 ms Wall time: 441 ms
np.testing.assert_allclose(m_numba, m)
Performance improvement is about 50x compared to Python and up to 10x compared to Pandas, pretty good!
The exact same result is obtained if we use numba.jit
as a decorator:
@numba.jit(nopython=True)
def groupby_numba(index, value, output):
for i in range(index.shape[0]):
output[index[i]] += value[i]