This recipe adapts the infix operator trick from http://code.activestate.com/recipes/384122-infix-operators/ to give the appropriate behavior with numpy arrays, so you can write A *dot* B for np.dot(A,B)
UPDATE A solution to the dot problem was recently added to the numpy trunk: the dot method was added to the ndarray class so you can write a.dot(b). See http://projects.scipy.org/numpy/ticket/1456
Download
Copy to clipboard1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 | import numpy as np
class Infix(np.ndarray):
"""
Creates a new infix operator that correcly acts on numpy arrays and scalars, used as X *op* Y.
The main motivation is to use np.dot as an infix operator for matrix multiplication.
example:
>>> x = np.array([1, 1, 1])
>>> x *dot* x
3
>>> 1 + x *dot* x # Multiplication has higher precedence than addition
4
"""
def __new__(cls, function):
obj = np.ndarray.__new__(cls, 0)
obj.function = function
return obj
def __array_finalize__(self, obj):
if obj is None: return
self.function = getattr(obj, 'function', None)
def __rmul__(self, other):
return Infix(lambda x, self=self, other=other: self.function(other, x))
def __mul__(self, other):
return self.function(other)
def __call__(self, value1, value2):
return self.function(value1, value2)
dot = Infix(np.dot)
outer = Infix(np.outer)
|
In Ferdinand Jamitzky's recipe for the Infix class, the following happens:
>>> dot = Infix(np.dot)
>>> x = np.array([1, 2, 3])
>>> x |dot| x
np.array([1, 1, 1], dtype=object)
We want it to equal
>>> np.dot(x, x)
14
The problem is that the __or__ method of array gets called instead of the __ror__ (rightward or) method of the dot object. So this expression evaluates as x.__or__(dot).__or__(x).
Numpy's ndarray class defines special methods for the infix operators: __or__, __mul__, __div__ etc. They all act element-wise. Therefore, the left operand's method gets called.
This recipe exploits the following rule of evaluation order: in the expression X * Y, if Y's class is a subclass of X's class, then Y's __mul__ method gets called. Thus if we define Infix to be a subclass of ndarray, then it will act on the left operand appropriately (see [3].)
Instead of using the | and <<,>> operators as in FJ's version, this recipe uses * so that the *dot* operator takes precedence over addition in an expression. (You can define other methods like __sub__ or __or__ if you want a different precedence level (see [4].)
References:
[1] infix operators trick: http://code.activestate.com/recipes/384122-infix-operators/
[2] subclassing ndarray: http://docs.scipy.org/doc/numpy/user/basics.subclassing.html
[3] python documentation on coercion rules: http://docs.python.org/release/2.5.2/ref/coercion-rules.html
[4] python documentation on evaluation order: http://docs.python.org/reference/expressions.html#summary
[5] discussion about adding math operators (2008): http://mail.python.org/pipermail/python-dev/2000-August /008194.html
[6] http://www.python.org/dev/peps/pep-0211/ (deferred)
[7] http://www.python.org/dev/peps/pep-0225/ (deferred)
