Overloading Mathematical Operations in Python

Python allows you to implement and customize mathematical operations within a class using a technique called overloading.

What is overloading? Overloading is the ability to define multiple versions of a method or operator with the same name but different parameter types or counts. In Python, overloading is typically achieved using magic methods, which enable you to customize the behavior of standard operators when they're applied to objects of user-defined classes.

In other words, it lets you use operators like +, -, *, etc., and specify how they should behave when used on objects of your class.

Let’s look at a practical example.

The assignment and increment operator += is used to add a value to an existing value stored in a numeric variable.

For example, you might assign the number 1 to a variable called num.

num = 1

Then, use the += operator to add 2 to the current value.

num += 2

The final result is 1+2=3

print(num)

3

The += operator works well with numeric values, but it doesn’t work if you want to add vectors together.

However, you can implement the += operator for a class by defining the special method __iadd__(), which is automatically called when the += operator is used on an instance of that class.

With __iadd__, you can define custom behavior for the += operator in your class.

For instance, let’s create a Vector class that represents a vector in 2D space.

Next, implement the __iadd__() method to add another vector or a pair of coordinates (x, y) to the current vector.

class Vector:
    def __init__(self, x: float = 0.0, y: float = 0.0):
        self.x = x
        self.y = y

    def __iadd__(self, other):
        if isinstance(other, Vector):
          # If other is a Vector, add corresponding coordinates
            self.x += other.x
            self.y += other.y
        elif isinstance(other, (tuple, list)) and len(other) == 2:
            # If other is a tuple or list with two elements, add the x and y values
            self.x += other[0]
            self.y += other[1]
        else:
            raise TypeError("other must be a Vector or a tuple/list of two numbers")

        # Return the current object to support chaining
        return self

    def __repr__(self):
        return f"Vector({self.x}, {self.y})"

It’s crucial to return self at the end of the method to ensure the operation happens in-place, modifying the current instance. This approach is especially useful when working with mutable objects.

Now, define two objects v1 and v2 of type Vector:

v1 = Vector(1, 2)
v2 = Vector(3, 4)

Next, use the += operator to add the second vector to the first:

v1 += v2

Now, print the result:

print(v1)

Vector(4, 6)

The first vector has been updated by adding the second one.

$$ \vec{v}_1 + \vec{v}_2 = \begin{pmatrix} 1 \\ 2 \end{pmatrix} + \begin{pmatrix} 3 \\ 4 \end{pmatrix} = \begin{pmatrix} 4 \\ 6 \end{pmatrix} $$

You can also add a vector and a tuple (x, y):

v1 += (1, 1)
print(v1)

Vector(5, 7)

Or, add a vector and a list (x, y):

v1 += [2, 3]
print(v1)

Vector(7, 10)

However, keep in mind that this operation is not commutative.

The __iadd__() method is only called if the first operand is an instance of the Vector class where you implemented the special method.

For instance, if the first operand is a tuple, the operation will raise an exception:

v3 = (1, 1)
v3 += v1

Traceback (most recent call last):
TypeError: can only concatenate tuple (not "Vector") to tuple

Using this technique, you can implement and customize all mathematical operators.

Here’s a complete list of special methods you can use in Python:

Magic Method Operator
__add__ +
__sub__ -
__mul__ *
__floordiv__ //
__div__ /
__truediv__ /
__mod__ %
__divmod__ divmod(a, b)
__pow__ ** pow(a, power, modulo=None)
__lshift__ <<
__rshift__ >>
__and__ &
__xor__ ^
__or__ |
__neg__ -
__pos__ +
__abs__ abs(a)
__invert__ ~

Commutative Operations

One thing to note is that overloading arithmetic operations is not inherently commutative.

For example, let’s implement the multiplication of a vector by a scalar in the Vector class using the __mul__() method:

class Vector:
    def __init__(self, x, y):
        self.x = x
        self.y = y

    def __mul__(self, scalar):
        if isinstance(scalar, (int, float)):
            return Vector(self.x * scalar, self.y * scalar)

    def __repr__(self):
        return f"Vector({self.x}, {self.y})"

The overloaded __mul__() method allows you to use the * operator to multiply a vector by a scalar value.

v = Vector(1, 2)
print(v * 3)

When you invoke this method in the operation v * 3, Python recognizes that the first operand (v) is a vector and the second operand (3) is a scalar.

As a result, it calls the __mul__() method in the Vector class and returns the correct result:

Vector(3, 6)

However, if you try 3 * v, Python doesn’t recognize a vector as the first operand in the multiplication.

In this case, Python doesn’t call the __mul__() method, leading to an error:

v = Vector(1, 2)
print(3 * v)

Traceback (most recent call last):
TypeError: unsupported operand type(s) for *: 'int' and 'Vector'

To make the multiplication commutative, you need to implement the __rmul__() method in the Vector class.

This method reverses the order of the operands, allowing right-to-left evaluation.

class Vector:
    def __init__(self, x, y):
        self.x = x
        self.y = y

    def __mul__(self, scalar):
        if isinstance(scalar, (int, float)):
            return Vector(self.x * scalar, self.y * scalar)

    def __rmul__(self, scalar):
        # Use the same logic to make the multiplication commutative
        return self.__mul__(scalar)

    def __repr__(self):
        return f"Vector({self.x}, {self.y})"

Now, multiplying a vector by a scalar works whether the vector is the first operand or the second.

v = Vector(1, 2)
print(v * 3)
print(3 * v)

Vector(3, 6)
Vector(3, 6)

Similarly, you can make addition commutative by implementing the __radd__() method.

Note. Keep in mind that reverse methods (e.g., __radd__(), __rmul__(), etc.) exist only for certain mathematical operations. For instance, the __iadd__() method doesn’t have a default reverse method; it must be created manually.




Report a mistake or post a question




FacebookTwitterLinkedinLinkedin