Implement np.array/np.asarray

[shenker]

Let's say I have a function like the following where doing the calculation involves doing a list comprehension which generates a list, and I want to turn it back into a numpy array before returning. The units of each element of the list comprehension will have the same units as some input x: if x is a normal Python number/numpy array, then I can turn the list into an array with a simple np.array(lst) or np.asarray(lst). However, if x has units, this fails.

My goal is to avoid introducing any pint-dependent code in do_calculation (e.g., let's say I have a big library of Python functions that calculate various physical formulae, part of the beauty of pint is if I pass in Quantities as inputs to these, the result is automatically carries units, but the library itself has no knowledge or dependency on pint; this list comprehension issue is the only place hangup).

Ideally, np.array/np.asarray would dispatch to pint.Quantities.from_sequence when any member of the input iterable is a pint.Quantity, but I'm not sure that NEP18 provides for that kind of thing.

Example:

import numpy as np
import pint
u = pint._DEFAULT_REGISTRY
def do_calculation(x, N):
    lst = [x * i for i in range(N)]
    ary = np.array(lst)
    ary2 = ary * 2
    return ary2
do_calculation(3, 10) # works
do_calculation(3 * u.meter, 10) # fails with ValueError: setting an array element with a sequence.

Is this something that will eventually be addressed by improvements in NEP18 support? Is there another best practice I should be following?

[keewis]

unfortunately, array creation functions (such as numpy.array or numpy.asarray, but also numpy.arange / numpy.linspace) are not covered by NEP18, so we can't override those functions using __array_function__.

In your example, though, I think you can (and should) work around that by using numpy functionality:

In [2]: def do_calculation(x, N):
   ...:     ary = x * np.arange(N)
   ...:     ary2 = ary * 2
   ...:     return ary2
   ...:
   ...: do_calculation(3, 10), do_calculation(3 * u.meter, 10)
Out[2]:
(array([ 0,  6, 12, 18, 24, 30, 36, 42, 48, 54]),
 array([ 0,  6, 12, 18, 24, 30, 36, 42, 48, 54]) <Unit('meter')>)

[jthielen]

In addition to @keewis's great answer, it might be worth mentioning that np.array and np.asarray are always expected to return a numpy.ndarray (they are the functions for explicitly converting to ndarray), so they should not be expected to return something that isn't like a pint.Quantity.

[mocquin]

It should be noted that NEP 35 seems to open the possibility to override the array creation routine.
From the abstract :
We propose the introduction of a new keyword argument like= to all array creation functions to address one of the shortcomings of __array_function__, as described by NEP 18 1. The like= keyword argument will create an instance of the argument’s type, enabling direct creation of non-NumPy arrays. The target array type must implement the __array_function__ protocol