npx 0.0.10

Some useful extensions for NumPy

npx

NumPy and SciPy are large libraries used everywhere in scientific computing. That's why breaking backwards-compatibility comes as 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.

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

dot

npx.dot(a, b)

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

npx.solve(A, b)

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 minumum 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:

Corresponding report: https://github.com/numpy/numpy/issues/11156.

unique_rows

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.

Corresponding report: https://github.com/numpy/numpy/issues/11136.

isin_rows

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.

SciPy Krylov methods

sol, info = npx.cg(A, b, tol=1.0e-10)
sol, info = npx.minres(A, b, tol=1.0e-10)
sol, info = npx.gmres(A, b, tol=1.0e-10)

sol is the solution of the linear system A @ x = b (or None if no convergence), and info contains some useful data, e.g., info.resnorms. The methods are wrappers around SciPy's iterative solvers.

SciPy minimization

def f(x):
    return (x ** 2 - 2) ** 2


x0 = 1.5
out = npx.minimize(f, x0)

In SciPy, the result from a minimization out.x will always have shape (n,), no matter the input vector. npx changes this to respect the input vector shape.

Corresponding report: https://github.com/scipy/scipy/issues/13869.

License

npx is published under the MIT license.