One of our example classes is Matrix
, a 2 by 2 matrix.
You can perform operations such as add or multiply on matrices. The most obvious way to do this would be to define an add
function:
p = Matrix(1, 2, 3, 4) q = Matrix(5, 6, 7, 8) r = p.add(q)
That is ok, but it would be nice if we could write:
p = Matrix(1, 2, 3, 4) q = Matrix(5, 6, 7, 8) r = p + q
Well we can! In fact we can override all the arithmetic operators if we wish. In this section we will look at + and *, but the technique can apply to any operator.
Just to recap the basics of matrix algebra, the sum of two matrices:
$$\begin{pmatrix}a & b\\c & d\end{pmatrix} + \begin{pmatrix}e & f\\g & h\end{pmatrix}$$
is
$$\begin{pmatrix}a + e & b + f\\c + g & d + h\end{pmatrix}$$
You can override the addition operator for a class by providing an __add__
method:
class Matrix: def __init__(self, a, b, c, d): self.data = [a, b, c, d] def __add__(self, other): if isinstance(other, Matrix): return Matrix(self.data[0] + other.data[0], self.data[1] + other.data[1], self.data[2] + other.data[2], self.data[3] + other.data[3]) else: return NotImplemented
The __add__
method accepts a parameter other
. We first check other
is a Matrix
. If it is, the method creates a brand new Matrix
object whose elements formed by adding the elements of other
to self
.
If other
is not a Matrix
, our code doesn't know how to handle it. In that case we must return NotImplemented
. Python can then decide what to do (this is covered in more detail below).
Here is how this is used:
p = Matrix(1, 2, 3, 4) q = Matrix(5, 6, 7, 8) r = p + q print(r)
We create two matrices, p
and q
. We then perform the calculation p + q
. This calls the __add__
method on the first object p
, passing the second object r
as the other
parameter.
The __add__
function returns a Matrix
that is the result of the addition, and this gets assigned to r
. The result is printed:
[6, 8][10, 12]
Matrix multiplication is a more interesting case, because you can multiply a matrix by another matrix, or alternatively you can multiply it by a scalar (ie an ordinary number).
The product of two matrices:
$$\begin{pmatrix}a & b\\c & d\end{pmatrix} . \begin{pmatrix}e & f\\g & h\end{pmatrix}$$
is
$$\begin{pmatrix}a.e + b.g & a.f + b.h\\c.e + d.g & c.f + d.h\end{pmatrix}$$
You can also multiply a matrix by a scalar (an ordinary number n
):
$$\begin{pmatrix}a & b\\c & d\end{pmatrix} . n$$
giving:
$$\begin{pmatrix}a.n & b.n\\c.n & d.n\end{pmatrix}$$
Here is a version of the Matrix
class with an implementation of __mul__
class Matrix: def __init__(self, a, b, c, d): self.a = a self.b = b self.c = c self.d = d def __mul__(self, other): if isinstance(other, (int, float)): return Matrix(self.data[0] * other, self.data[1] * other, self.data[2] * other, self.data[3] * other) elif isinstance(other, Matrix): return Matrix(self.data[0] * other.data[0] + self.data[1] * other.data[1], self.data[0] * other.data[1] + self.data[1] * other.data[3], self.data[2] * other.data[0] + self.data[3] * other.data[1], self.data[2] * other.data[1] + self.data[3] * other.data[3]) else: return NotImplemented
If you look at __mul__
, you will see that the first thing we do is to check if other
is a scalar. We do this by checking if it is an instance of int
or float
(you could also check complex
if you wanted the Matrix class to support complex number, but we won't bother in this example).
If the value is a number, we execute the code for the scalar multiplication equation above.
If the value is not a scalar, we check if it is a Matrix
, and execute the code for the matrix multiplication equation above.
If the value is neither a number nor a Matrix
we return NotImplemented
.
Here is an example:
p = Matrix(1, 2, 3, 4) q = Matrix(5, 6, 7, 8) print(p*2) print(p*q)
This prints
[2, 4][6, 8] [17, 22][39, 50]
as expected.
What if we try to do this:
print(2*p)
Unfortunately our existing code doesn't quite cope with this situation. We get an error:
TypeError: unsupported operand type(s) for *: 'int' and 'Matrix'
So what has happened here? We are trying to multiply 2*p
:
int
.int.__mul__
passing in the second value p
which is a Matrix
.int
is a built-in type, its __mul__
function knows nothing of out Matrix
type, so it returns NotImplemented
.You might think that Python would give an error at this point, but actually it tries one last thing:
p
has a __rmul__
method. If not it gives an error.p.__rmul__
passing in the first value 2.p.__rmul__
can handle an integer type, the value will be calculated.p.__rmul__
returns NotImplemented
and Python gives an error.So, we can handle this extra case by implementing __rmul__
for our Matrix
class:
def __rmul__(self, other): if isinstance(other, (int, float)): return Matrix(self.data[0] * other, self.data[1] * other, self.data[2] * other, self.data[3] * other) else: return NotImplemented
In this case, self
is the second operand p
, and other
is the first operand 2. This is because __rmul__
reverses the arguments.
Since other
is an int
, our code executes and creates the correct result. In this case, the code for handing numbers is identical in __mul__
and __rmul__
because for matrices p*2
and 2*p
are the same. That won't be true for all data types and all operators of course.
Notice that if both operands are of type Matrix, the case will always be handled by __mul__
, so there is no need to handle that case in __rmul__
. This is generally true for all data types and operators.
What if we try something crazy like:
print(p*'abc')
Our __mul__
code checks the type of the other
value. It isn't a number, it isn't a Matrix
, so we return NotImplemented
.
Python will then check if str
has an __rmul__
method. It does, but it can't handle our Matrix
type so again it returns NotImplemented
.
Python gives an error.
There is an additional case to consider, the in place operators such as +=
and *=
. They are covered in the next section
Here is a summary of the available numerical operators:
Method | Symbol |
---|---|
__add__ | + |
__sub__ | - |
__mul__ | * |
__matmul__ | @ |
__truediv__ | / |
__floordiv__ | // |
__mod__ | % |
__divmod__ | |
__pow__ | ** |
__lshift__ | << |
__rshift__ | >> |
__and__ | & |
__xor__ | ^ |
__or__ | | |
If you found this article useful, you might be interested in the book Functional Programming in Python, or other books, by the same author.
<<PrevCopyright (c) Axlesoft Ltd 2020