What does (numpy) __array_wrap__ do?

Question

I am diving into the SciPy LinAlg module for the first time, and I saw this function:

def _makearray(a):
    new = asarray(a)
    wrap = getattr(a, "__array_prepare__", new.__array_wrap__)
    return new, wrap

What does __array_wrap__ do exactly? I found the documentation, but I don't understand this explanation:

 At the end of every ufunc, this method is called on the input object with the
 highest array priority, or the output object if one was specified. The ufunc-
 computed array is passed in and whatever is returned is passed to the user. 
 Subclasses inherit a default implementation of this method, which transforms the
 array into a new instance of the object’s class. Subclasses may opt to use this
 method to transform the output array into an instance of the subclass and update
 metadata before returning the array to the user.

Does this just mean it reconverts the output of whatever function back into an array since it was likely broken up into something else for element-by-element processing? Relatedly, regardless of the explanation, what would it mean to get this wrap as an object? What would you do with it?

I am looking at the code for numpy.linalg.inv...what is wrap doing here?

    **a, wrap = _makearray(a)**
    _assertRankAtLeast2(a)
    _assertNdSquareness(a)
    t, result_t = _commonType(a)

    if a.shape[-1] == 0:
        # The inner array is 0x0, the ufunc cannot handle this case
        **return wrap(empty_like(a, dtype=result_t))**

    signature = 'D->D' if isComplexType(t) else 'd->d'
    extobj = get_linalg_error_extobj(_raise_linalgerror_singular)
    ainv = _umath_linalg.inv(a, signature=signature, extobj=extobj)
    return wrap(ainv.astype(result_t))

Answer 1

np.ma.masked_array.__array_wrap__ is an example of a array subclass that updates the metadata (the mask).

File:        /usr/lib/python3/dist-packages/numpy/ma/core.py
Definition:  np.ma.masked_array.__array_wrap__(self, obj, context=None)
Source:
    def __array_wrap__(self, obj, context=None):
        """
        Special hook for ufuncs.
        Wraps the numpy array and sets the mask according to context.
        """

np.matrix.__array_wrap__ appears to inherit the ndarray version. My guess it's because matrix, while a subclass, does not have metadata that needs updating.

Generally the idea with a hook, is that it's a function that is called deep within the normal processing. The default method might not do anything. But it's a way that the subclass can take special action. The class developer writes hooks like this so that the class user will not have to worry about those details. With a __...__ name it isn't part of the public interface - though Python lets us peak under the curtain.

An example of wrapping, i.e returning an array with the same class as the inputs is:

In [659]: np.cumsum(np.arange(10))
Out[659]: array([ 0,  1,  3,  6, 10, 15, 21, 28, 36, 45], dtype=int32)

In [660]: np.cumsum(np.matrix(np.arange(10)))
Out[660]: matrix([[ 0,  1,  3,  6, 10, 15, 21, 28, 36, 45]], dtype=int32

In [665]: np.cumsum(np.ma.masked_array(np.arange(10)))
Out[665]: 
masked_array(data = [ 0  1  3  6 10 15 21 28 36 45],
             mask = False,
       fill_value = 999999)

The returned values are all the same, but the array subclass varies, depending on the input class.

---

cumsum might not be the best example. Masked arrays have their own version of cumsum, one that treats the masked values as 0:

In [679]: m=np.ma.masked_array(np.arange(10),np.arange(10)%2)

In [680]: m
Out[680]: 
masked_array(data = [0 -- 2 -- 4 -- 6 -- 8 --],
             mask = [False  True False  True False  True False  True False  True],
       fill_value = 999999)

In [681]: np.cumsum(m)
Out[681]: 
masked_array(data = [0 -- 2 -- 6 -- 12 -- 20 --],
             mask = [False  True False  True False  True False  True False  True],
       fill_value = 999999)

add.accumulate is similar to cumsum, but doesn't have a special masked version:

In [682]: np.add.accumulate(np.arange(10))
Out[682]: array([ 0,  1,  3,  6, 10, 15, 21, 28, 36, 45], dtype=int32)

In [683]: np.add.accumulate(m)
Out[683]: 
masked_array(data = [ 0  1  3  6 10 15 21 28 36 45],
             mask = False,
       fill_value = 999999)

This last is a masked array, but the mask is the default False, and the masked values were included in the sum.