Some useful extensions for NumPy


Keywords
numpy, python
License
Other
Install
pip install npx==0.1.2

Documentation

npx

PyPi Version PyPI pyversions GitHub stars Downloads

gh-actions codecov Code style: black

NumPy is a large library used everywhere in scientific computing. That's why breaking backwards-compatibility comes at a significant cost and is almost always avoided, even if the API of some methods is arguably lacking. This package provides drop-in wrappers "fixing" those.

scipyx does the same for SciPy.

If you have a fix for a NumPy method that can't go upstream for some reason, feel free to PR here.

dot

import npx
import numpy as np

a = np.random.rand(3, 4, 5)
b = np.random.rand(5, 2, 2)

out = npx.dot(a, b)
# out.shape == (3, 4, 2, 2)

Forms the dot product between the last axis of a and the first axis of b.

(Not the second-last axis of b as numpy.dot(a, b).)

np.solve

import npx
import numpy as np

A = np.random.rand(3, 3)
b = np.random.rand(3, 10, 4)

out = npx.solve(A, b)
# out.shape == (3, 10, 4)

Solves a linear equation system with a matrix of shape (n, n) and an array of shape (n, ...). The output has the same shape as the second argument.

sum_at/add_at

npx.sum_at(a, idx, minlength=0)
npx.add_at(out, idx, a)

Returns an array with entries of a summed up at indices idx with a minimum length of minlength. idx can have any shape as long as it's matching a. The output shape is (minlength,...).

The numpy equivalent numpy.add.at is much slower:

memory usage

Relevant issue reports:

unique

import npx

a = [0.1, 0.15, 0.7]
a_unique = npx.unique(a, tol=2.0e-1)

assert all(a_unique == [0.1, 0.7])

npx's unique() works just like NumPy's, except that it provides a parameter tol (default 0.0) which allows the user to set a tolerance. The real line is essentially partitioned into bins of size tol and at most one representative of each bin is returned.

unique_rows

import npx
import numpy as np

a = np.random.randint(0, 5, size=(100, 2))

npx.unique_rows(a, return_inverse=False, return_counts=False)

Returns the unique rows of the integer array a. The numpy alternative np.unique(a, axis=0) is slow.

Relevant issue reports:

isin_rows

import npx
import numpy as np

a = [[0, 1], [0, 2]]
b = np.random.randint(0, 5, size=(100, 2))

npx.isin_rows(a, b)

Returns a boolean array of length len(a) specifying if the rows a[k] appear in b. Similar to NumPy's own np.isin which only works for scalars.

mean

import npx

a = [1.0, 2.0, 5.0]
npx.mean(a, p=3)

Returns the generalized mean of a given list. Handles the cases +-np.inf (max/min) and0 (geometric mean) correctly. Also does well for large p.

Relevant NumPy issues:

License

This software is published under the BSD-3-Clause license.