Support np.vectorize and np.frompyfunc

[ngoldbaum]

• unyt version: master
• Python version: 3.7
• Operating System: Ubuntu 18.04

Description

Right now using a function wrapped with np.vectorize or np.frompyfunc raises an error.

This was originally reported against yt as yt-project/yt#2465.

What I Did

Running this script:

import unyt as u
import numpy as np


def myfunc(a, b):
    return a + b


vmyfunc = np.vectorize(myfunc)

print(vmyfunc([1, 2, 3]*u.g, [1, 2, 3]*u.kg))

Produces the following traceback:

Traceback (most recent call last):
  File "test.py", line 11, in <module>
    print(vmyfunc([1, 2, 3]*u.g, [1, 2, 3]*u.kg))
  File "/home/goldbaum/.pyenv/versions/3.7.3/lib/python3.7/site-packages/numpy/lib/function_base.py", line 2091, in __call__
    return self._vectorize_call(func=func, args=vargs)
  File "/home/goldbaum/.pyenv/versions/3.7.3/lib/python3.7/site-packages/numpy/lib/function_base.py", line 2167, in _vectorize_call
    outputs = ufunc(*inputs)
  File "/home/goldbaum/Documents/unyt/unyt/array.py", line 1692, in __array_ufunc__
    unit_operator = self._ufunc_registry[ufunc]
KeyError: <ufunc '? (vectorized)'>

Similarly:

import unyt as u
import numpy as np


def myfunc(a, b):
    return a + b


vmyfunc = np.frompyfunc(myfunc, 2, 1)

print(vmyfunc([1, 2, 3]*u.g, [1, 2, 3]*u.kg))

Traceback (most recent call last):
  File "test.py", line 11, in <module>
    print(vmyfunc([1, 2, 3]*u.g, [1, 2, 3]*u.kg))
  File "/home/goldbaum/Documents/unyt/unyt/array.py", line 1692, in __array_ufunc__
    unit_operator = self._ufunc_registry[ufunc]
KeyError: <ufunc '? (vectorized)'>

For what it's worth, astropy produces exactly the same error:

Traceback (most recent call last):
  File "test.py", line 11, in <module>
    print(vmyfunc([1, 2, 3]*u.g, [1, 2, 3]*u.kg))
  File "/home/goldbaum/.pyenv/versions/3.7.3/lib/python3.7/site-packages/numpy/lib/function_base.py", line 2091, in __call__
    return self._vectorize_call(func=func, args=vargs)
  File "/home/goldbaum/.pyenv/versions/3.7.3/lib/python3.7/site-packages/numpy/lib/function_base.py", line 2167, in _vectorize_call
    outputs = ufunc(*inputs)
  File "/home/goldbaum/.pyenv/versions/3.7.3/lib/python3.7/site-packages/astropy/units/quantity.py", line 446, in __array_ufunc__
    converters, unit = converters_and_unit(function, method, *inputs)
  File "/home/goldbaum/.pyenv/versions/3.7.3/lib/python3.7/site-packages/astropy/units/quantity_helper/converters.py", line 157, in converters_and_unit
    ufunc_helper = UFUNC_HELPERS[function]
  File "/home/goldbaum/.pyenv/versions/3.7.3/lib/python3.7/site-packages/astropy/units/quantity_helper/converters.py", line 89, in __missing__
    .format(ufunc.__name__))
TypeError: unknown ufunc ? (vectorized).  If you believe this ufunc should be supported, please raise an issue on https://github.com/astropy/astropy

[ngoldbaum]

I think the most straightforward way to support this would be to cast the arguments to ndarray and just ignore the units. That's more useful than completely erroring out. We could probably also print a warning saying that the units are being ignored.

I don't think we can support ufuncs generated via frompyfunc without degrading to ignoring the units because we'd have no ability to introspect and determine what the units should be.