Welcome, guest | Sign In | My Account | Store | Cart

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

Python, 30 lines
 1
 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)

Created by John Schulman on Mon, 19 Apr 2010 (MIT)
Python recipes (4591)
John Schulman's recipes (1)

Required Modules

Other Information and Tasks