NumPy Extensions

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

Documentation

API Reference

Installation

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

Examples

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]])