如何获取两个不同的numpy.ndarray子类的__matmul__以返回特定的子类？

Question

I have two np.ndarray subclasses. Tuple @ Matrix returns a Tuple , but Matrix @ Tuple returns a Matrix . How might I have it return a Tuple instead?

import numpy as np

class Tuple(np.ndarray):
    def __new__(cls, input_array, info=None):
        return np.asarray(input_array).view(cls)

class Matrix(np.ndarray):
    def __new__(cls, input_array, info=None):
        return np.asarray(input_array).view(cls)

def scaling(x, y, z):
    m = Matrix(np.identity(4))
    m[0, 0] = x
    m[1, 1] = y
    m[2, 2] = z
    return m

example:

>>> Tuple([1,2,3,4]) @ scaling(2,2,2)
Tuple([2., 4., 6., 4.])

>>> scaling(2,2,2) @ Tuple([1,2,3,4])
Matrix([2., 4., 6., 4.])   # XXXX I'd like this to be a Tuple

PS: Matrix @ Matrix should return Matrix

Answer 1

You can overload the __matmul__ method to return a Tuple - and if you want to be a Tuple if any of the variables is a Tuple and a Matrix otherwise, I think this'll work:

class Matrix(np.ndarray):
    def __new__(cls, input_array, info=None):
        return np.asarray(input_array).view(cls)

    def __matmul__(m1, m2):
         return (m2.T @ m1.T).T if isinstance(m2, Tuple) else np.matmul(m1, m2)

Answer 2

I made a mistake in copying from the np.matrix example.

class Tuple(np.ndarray): 
    __array_priority__ = 10 
    def __new__(cls, input_array, info=None): 
        return np.asarray(input_array).view(cls) 
class Matrix(np.ndarray):
    __array_priority__ = 5.0 
    def __new__(cls, input_array, info=None): 
        return np.asarray(input_array).view(cls)

In [2]: def scaling(x, y, z):  
   ...:      ...:     m = Matrix(np.identity(4))  
   ...:      ...:     m[0, 0] = x  
   ...:      ...:     m[1, 1] = y  
   ...:      ...:     m[2, 2] = z  
   ...:      ...:     return m  
   ...:                                                                                                                                  
In [3]: Tuple([1,2,3,4]) @ scaling(2,2,2)                                                                                                
Out[3]: Tuple([2., 4., 6., 4.])
In [4]: scaling(2,2,2) @ Tuple([1,2,3,4])                                                                                                
Out[4]: Tuple([2., 4., 6., 4.])

===

Taking a clue from the np.matrix definition: numpy.matrixlib.defmatrix.py

Add a __array_priority__ attribute:

In [382]: class Tuple(np.ndarray): 
     ...:     def __new__(cls, input_array, info=None): 
     ...:         __array_priority = 10 
     ...:         return np.asarray(input_array).view(cls) 
     ...:  
     ...: class Matrix(np.ndarray): 
     ...:     def __new__(cls, input_array, info=None): 
     ...:         __array_priority = 5 
     ...:         return np.asarray(input_array).view(cls) 
     ...:                                                                                            
In [383]:                                                                                            
In [383]: def scaling(x, y, z): 
     ...:     m = Matrix(np.identity(4)) 
     ...:     m[0, 0] = x 
     ...:     m[1, 1] = y 
     ...:     m[2, 2] = z 
     ...:     return m 
     ...:                                                                                            
In [384]: Tuple([1,2,3,4]) @ scaling(2,2,2)                                                          
Out[384]: Tuple([2., 4., 6., 4.])
In [385]: scaling(2,2,2) @ Tuple([1,2,3,4])                                                          
Out[385]: Matrix([2., 4., 6., 4.])

Answer 3

One way of solving this is by implementing a custom __matmul__ in Matrix and __rmatmul__ in Tuple :

import numpy as np

class Tuple(np.ndarray):
    def __new__(cls, input_array, info=None):
        return np.asarray(input_array).view(cls)

    def __rmatmul__(self, other):
        return super().__matmul__(other)

class Matrix(np.ndarray):
    def __new__(cls, input_array, info=None):
        return np.asarray(input_array).view(cls)

    def __matmul__(self, other):
        if not isinstance(other, Matrix):
            return NotImplemented
        return super().__matmul__(other)

def scaling(x, y, z):
    m = Matrix(np.identity(4))
    m[0, 0] = x
    m[1, 1] = y
    m[2, 2] = z
    return m

scaling(2,2,2) @ scaling(2,2,2)
# Matrix([[4., 0., 0., 0.],
#         [0., 4., 0., 0.],
#         [0., 0., 4., 0.],
#         [0., 0., 0., 1.]])
Tuple([1,2,3,4]) @ scaling(2,2,2)
# Tuple([2., 4., 6., 4.])
scaling(2,2,2) @ Tuple([1,2,3,4])
# Tuple([2., 4., 6., 4.])

Answer 4

Just overload __matmul__ of Matrix class to return tuple instead

class Matrix(np.ndarray):
    def __new__(cls, input_array, info=None):
        return np.asarray(input_array).view(cls)

    def __matmul__(self, other):
        return other @ self