Mojo vs Python: Comparison with Examples - Part III: Speed Test with Complex Numbers

As I promised in the previous article, we’re going to take a look at our first speed comparison between Mojo and Python.

Speed Comparison using Complex Numbers

You can find the complete code in my GitHub repository dedicated to this series of articles.

In this speed comparison I use complex numbers. Inspired by a HackerRank task, I built a class representing a complex number in Mojo as well as in Python, then in a loop of 10 million iterations I perform some mathematical operations with complex numbers and get an integer result from each iteration. The final result is the sum of all these iteration results.

In my very first article about Mojo I already did almost the same speed comparison, however, that time developers were provided only a JupyterHub-based playground to “play on” and Mojo was quite in its early stages. So, I thought it would be interesting to repeat it properly.

The Speed Test

Python:

from timeit import timeit


class Complex:

	def __init__(self, real, imaginary):
		self.real = real
		self.imaginary = imaginary

	def __add__(self, other):
		return Complex(self.real + other.real, self.imaginary + other.imaginary)

	def __sub__(self, other):
		return Complex(self.real - other.real, self.imaginary - other.imaginary)

	def __mul__(self, other):
		r = (self.real * other.real) + (self.imaginary * other.imaginary * -1)
		i = (self.real * other.imaginary) + (self.imaginary * other.real)
		return Complex(r, i)


def fn_speed_test():
    result = 0
    for i in range(10_000_000):
        x = Complex(i + 91, i)
        y = Complex(i, i - 2)
        sum = x + y
        sub = x - y
        mul = x * y
        result += (sum.real + sub.real + mul.real)
    print(result)


if __name__ == '__main__':
    execution_time_ns = timeit(fn_speed_test, number=1) * 1e9
    print("Execution time: {} nanoseconds".format(int(execution_time_ns)))

4750001345000000
Execution time: 8830614791 nanoseconds

Mojo:

from time import time_function


struct Complex:
	var real: Int
	var imaginary: Int
  
	fn __init__(inout self, real: Int, imaginary: Int):
		self.real = real
		self.imaginary = imaginary
  
	fn __add__(self, other: Complex) -> Complex:
		return Complex(self.real + other.real, self.imaginary + other.imaginary)

	fn __sub__(self, other: Complex) -> Complex:
		return Complex(self.real - other.real, self.imaginary - other.imaginary)

	fn __mul__(self, other: Complex) -> Complex:
		var r: Int = (self.real * other.real) + 
		             (self.imaginary * other.imaginary * -1)
		var i: Int = (self.real * other.imaginary) + (self.imaginary * other.real)
		return Complex(r, i)
  
  
fn fn_speed_test() capturing -> None:
    var result: Int = 0
    var x: Complex
    var y: Complex
    var sum: Complex
    var sub: Complex
    var mul: Complex
    for i in range(10_000_000):
        x = Complex(i + 91, i)
        y = Complex(i, i - 2)
        sum = x + y
        sub = x - y
        mul = x * y
        result += (sum.real + sub.real + mul.real)
    print(result)


fn main():
    var execution_time_ns: Int = time_function[fn_speed_test]()
    print("Execution time: " + str(execution_time_ns) + " nanoseconds")

4750001345000000
Execution time: 41000 nanoseconds

The result is mind-blowing. Mojo is around 200K times faster than Python. Well, at least in this example, but even if there are more balanced ones, it’s simply unbelievable.

We are leveraging Mojo’s static and memory-safe struct. Structs are similar to classes in Python, but for instance you cannot add methods to them at runtime. In exchange you get higher performance and safety. In a struct you can have fn or def methods as well, but all fields must be declared with var. In our example we declare methods with fn.

In a future article we will take a closer look at structs, for now I’d just like to point out the great similarity between Python’s and Mojo’s magic methods such as __init__(), __add__(), __str__() and others. Another great example of Mojo’s Pythonic style.

This article ends here, but I’m already excited about further speed comparisons. I already have some ideas for the next article, but let it be a surprise. Good luck on your journey with Mojo and Happy Coding! 🔥🔥🔥