Possible to have logical operations (e.g. `ndarray.__eq__`) not return `bool` in NumPy?

问题

Question

What's the most elegant way to have NumPy not always force np.equal / np.ndarray.__eq__ return arrays of type bool, especially when dtype=object?

Problem

Project issue: https://github.com/RobotLocomotion/drake/issues/8315

We're presently using NumPy in a project, and we have symbolic scalars that we're using in the arrays. We define __eq__ to return a formula - here's a simple class that kinda reflects it:

import numpy as np

class Custom(object):
    def __init__(self, value):
        self.value = value

    def __eq__(self, rhs):
        return "eq({}, {})".format(self, rhs)

    def __str__(self):
        return repr(self.value)

    def __repr__(self):
        return repr(self.value)

It appears that ndarray.__eq__ is always converted to a boolean type, even if dtype=object, because it converts using something like __nonzero__:

>>> # Scalar case works.
>>> a = Custom('a')
>>> print(a == a)
eq('a', 'a')
>>> # Not what we want.
>>> av = np.array([a, a])
>>> print(np.equal(av, av))
[ True  True]

While something like sympy can certainly handle symbolic matrices, I'm hesitant to start mixing array / linear algebra interfaces. And of course, we could go the route of machine learning packages and implement our own (or possibly use theirs), such as TensorFlow.

Blind Hacking

Granted, we can make our own NumPy UFunc, which propagates the type we want:

>>> # Try custom ufunc:
>>> generic_equal = np.frompyfunc(lambda a, b: a == b, 2, 1)
>>> print(generic_equal(av, av))
["eq('a', 'a')" "eq('a', 'a')"]

We can apparently also override ndarray.__eq__ using set_numeric_ops (but that doesn't affect np.equal), discovered thanks to this post https://stackoverflow.com/a/45602324/7829525:

>>> # Try replacing.
>>> np.set_numeric_ops(equal=generic_equal)
>>> print(av == av)
["eq('a', 'a')" "eq('a', 'a')"]
>>> print(np.equal(av, av))
[ True  True]

And then we can monkey-patch np.equal:

>>> # Now replace original ufunc.
>>> np.equal = generic_equal
>>> print(np.equal(av, av))
["eq('a', 'a')" "eq('a', 'a')"]

This kinda works, except it completely destroys the original np.equal UFunc, and gives me the heeby-jeebies that we have to mutate all of numpy to try and carve out functionality for one datatype -- it'd be nice to extend the existing np.equal to maybe teach it something else.

Next Steps

From looking at the C API, there's the function PyUFunc_RegisterLoopForType which can add to an existing UFunc, but I do not see an easy-to-access Python version for at least prototype.

A related fact is we're using pybind11, which has tons of convenient mechanisms for handling Eigen matrices and interfacing them with NumPy. We made a patch to allow for non-POD types to be converted to np.array(..., dtype=object) types (hence this whole question); however, this prevents the same memory block from being accessed, so we cannot use the Eigen::Ref<> casting that pybind11/eigen.h offers: https://github.com/RobotLocomotion/drake/issues/8116

I've been tinkering with the apparent test module in the numpy source for custom dtypes: https://github.com/numpy/numpy/blob/4092a9e160cc247a4a45724579a0c829733688ca/numpy/core/src/umath/test_rational.c.src#L1192 which I came across thanks to this post: https://github.com/numpy/numpy/issues/5719

From that, it looks like you can willy-nilly register your own UFuncs for a custom DType, so I'm going to try and see if that also solves this problem.

来源：https://stackoverflow.com/questions/49205075/possible-to-have-logical-operations-e-g-ndarray-eq-not-return-bool-in

标签

python

numpy

logical-operators

symbolic-math