An extension library for NumPy that implements common array operations not present in NumPy.
npext.fill_na(...)
npext.drop_na(...)
npext.rolling(...)
npext.expanding(...)
npext.rolling_apply(...)
npext.expanding_apply(...)
# etc
Regular installation:
pip install numpy_ext
For development:
git clone https://github.com/3jane/numpy_ext.git
cd numpy_ext
pip install -e .[dev] # note: make sure you are using pip>=20
Here are few common examples of how the library is used. The rest is available in the documentation.
- Apply a function to a rolling window over the provided array
import numpy as np
import numpy_ext as npext
a = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
window = 3
npext.rolling_apply(np.sum, window, a)
> array([nan, nan, 3., 6., 9., 12., 15., 18., 21., 24.])
- Same as the above, but with a custom function, two input arrays and parallel computation using
joblib
:
def func(array_first, array_second, param):
return (np.min(array_first) + np.sum(array_second)) * param
a = np.array([0, 1, 2, 3])
b = np.array([3, 2, 1, 0])
npext.rolling_apply(func, 2, a, b, n_jobs=2, param=-1)
> array([nan, -5., -4., -3.])
- Same as the first example, but using rolling function:
a = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
window = 3
rolls = npext.rolling(a, window, as_array=True)
np.sum(rolls, axis=1)
> array([nan, nan, 3., 6., 9., 12., 15., 18., 21., 24.])
- Apply a function with multiple output to a rolling window over the provided array, with no nans prepend
res = npext.rolling_apply(
lambda x: (max(x), min(x)),
3,
np.array([1, 2, 5, 1, 6, 4, 0]),
prepend_nans=False,
)
> array([[5, 1],
[5, 1],
[6, 1],
[6, 1],
[6, 0]])
The software is distributed under MIT license.