February 15, 2024

MojošŸ”„ ā™„ļø Python: Calculating and plotting a Valentineā€™s day ā™„ļø using Mojo and Python

On Valentineā€™s Day yesterday, I wanted to create something special to celebrate my love for Mojo and Python. My search on the interwebs led me to a nifty little equation that plots a heart. The equation is quite simple and Iā€™ll refer to this as the ā€œheart equationā€ through the rest of this blog post:

$f(x) = (x)^{\frac{2}{3}} - 0.9 \sqrt{3.3 - x^2} \sin(a \pi x)$

In this equation if you vary the variable a from 0 to 20 in small increments, you get this cool animation:

The heart equation is fairly easy to implement in Python using NumPy, but in this blog post I'll show you how you can implement it in MojošŸ”„. Mojo is still young and its standard library doesnā€™t yet have a robust NumPy-like data structure for implementing math equations, but with relatively low effort you can implement a custom data structure (Iā€™ll call MojoArray in this blog post) that supports basic vectorized math operations we need, to implement the heart equation. Once we've created our MojoArray data structure, we can implement the equation and plot it using the following code:

Mojo
let np_arr = np.arange(-2,2,0.01) let x = MojoArray.from_numpy(np_arr) let y = (x**2)**(1/3.) - 0.9*((3.3-(x*x)).sqrt())*(15*3.14*x).sin() plt.plot(x.to_numpy(),y.to_numpy(),'r')

Since Mojo offers first-class support for Python interoperability, weā€™ll see how you can do all the computation in MojošŸ”„ and leverage Python and Matplotlib library to plot the animation you see above.

Note: The code for this example is available on GitHub. In this blog post, Iā€™ll only share specific code snippets to explain key implementation details.

What youā€™ll learn:

  1. How to create a custom MojošŸ”„ data structure that lets do vectorized math operations
  2. How to use double underscore aka dunder methods to overload math operators that make writing math equations easy
  3. How to use Mojo standard library Math functions that operate on SIMD type and extend them to work on your data structure
  4. How to leverage Python libraries to import NumPy arrays and plot using Matplotlib

Creating an Array data structure that supports vectorized Math operations

To implement the heart equation, I need a data structure that supports vectorized Math operations. For example, in this equation x is a vector, and I need sqrt, pow, sin, add, mul and div to operate on vectors.

$f(x) = (x)^{\frac{2}{3}} - 0.9 \sqrt{3.3 - x^2} \sin(a \pi x)$

Mojo Math module in the standard library supports common mathematical operations but they operate on SIMD types only. SIMD types allow you perform a single operation on small vectors whoā€™s length depends on the CPU type. For example, to compute sin on a SIMD vector, in Mojo you can run:

Mojo
from math import sin print(sin(SIMD[DType.float64,8].splat(3.141592654)))

Output:

Output
[-4.1020685703066231e-10, -4.1020685703066231e-10, -4.1020685703066231e-10, -4.1020685703066231e-10, -4.1020685703066231e-10, -4.1020685703066231e-10, -4.1020685703066231e-10, -4.1020685703066231e-10]

Which gives you very small numbers as youā€™d expect, since sin of approximate value of pi, is approximately 0. To implement the heart equation, we'll need to work with large vectors, therefore, we need a Mojo struct that offers methods to perform Math operations on large vectors, not just SIMD types.

Mojo standard library offers a Tensor data structure that can store higher order tensors, but it is quite an overkill for this simple example. The Tensor type also doesnā€™t support basic Math operations which we need. With a simple, custom data structure called MojoArray, we can only implement the features we need.

Here is the skeleton of the MojoArray data structure that weā€™ll be using to calculate our heart equation.Ā 

Mojo
struct MojoArray[dtype: DType = DType.float64](Stringable): var _ptr: DTypePointer[dtype] var numel: Int # Initializers fn __init__(inout self, numel: Int): fn __init__(inout self, numel: Int, _ptr: DTypePointer[dtype]): fn __init__(inout self, *data: Scalar[dtype]): fn __copyinit__(inout self, other: Self): fn __getitem__(self, idx: Int) -> Scalar[dtype]: # Math operators fn __neg__(self)->Self: fn __pow__(self, p: Scalar[dtype])->Self: fn __mul__(self, other: Self)->Self: fn __mul__(self, s: Scalar[dtype])->Self: fn __add__(self, s: Scalar[dtype])->Self: fn __add__(self, other: Self)->Self: fn __sub__(self, s: Scalar[dtype])->Self: fn __sub__(self, other: Self)->Self: # Math functions fn sqrt(self)->Self: fn cos(self)->Self: fn sin(self)->Self: fn abs(self)->Self: # Helper functions to get data in and out of NumPy fn from_numpy(np_array: PythonObject) raises->Self: fn to_numpy(self) raises->PythonObject: # Helper functions that extend SIMD math module functions to MojoArrays fn _elemwise_transform[]() fn _elemwise_array_math[]() fn _elemwise_scalar_math[]()

The MojoArray data structure is parameterized on dtype which defaults to float64, which is also the default type of NumPy arrays. MojoArray struct also defines double underscore aka dunder methods that overload common mathematical operations such as addition, subtraction, multiplication to work on two MojoArrays or MojoArray and scalar values.

Letā€™s take a closer look at the implementation details of a few of these. Most of the math operations fall into one of these categories:

Elementwise transforms: _elemwise_transform[]()

This helper function implements Trigonometric functions, and abs. that transform a MojoArray. _elemwise_transform[]() applies the function func on all the elements in the MojoArray.

Mojo
fn _elemwise_transform[func: fn[dtype: DType, width: Int](SIMD[dtype, width])->SIMD[dtype, width]](self) -> Self: alias simd_width: Int = simdwidthof[dtype]() let new_array = Self(self.numel) @parameter fn elemwise_vectorize[simd_width: Int](idx: Int) -> None: new_array._ptr.simd_store[simd_width](idx, func[dtype, simd_width](self._ptr.simd_load[simd_width](idx))) vectorize[simd_width, elemwise_vectorize](self.numel) return new_array

I use _elemwise_transform to implement sqrt, cos, sin and abs:

Mojo
fn sqrt(self)->Self: return self._elemwise_transform[math.sqrt]() fn cos(self)->Self: return self._elemwise_transform[math.cos]() fn sin(self)->Self: return self._elemwise_transform[math.sin]() fn abs(self)->Self: return self._elemwise_transform[math.abs]()

Elementwise vector-scalar operations _elemwise_scalar_math[]():Ā 

This helper function implements multiplication, addition, and subtraction between MojoArrays and scalars. _elemwise_array_math accepts a scalar and applies the function func on all the elements in the MojoArray.

Mojo
fn _elemwise_scalar_math[func: fn[dtype: DType, width: Int](SIMD[dtype, width],SIMD[dtype, width])->SIMD[dtype, width]](self, s: Scalar[dtype]) -> Self: alias simd_width: Int = simdwidthof[dtype]() let new_array = Self(self.numel) @parameter fn elemwise_vectorize[simd_width: Int](idx: Int) -> None: new_array._ptr.simd_store[simd_width](idx, func[dtype, simd_width](self._ptr.simd_load[simd_width](idx), SIMD[dtype, simd_width](s))) vectorize[simd_width, elemwise_vectorize](self.numel) return new_array

We use _elemwise_scalar_math to implement negative, addition, subtraction, and multiplication between a vectors and a scalar:

Mojo
fn __neg__(self)->Self: return self._elemwise_scalar_math[math.mul](-1.0) fn __mul__(self, s: Scalar[dtype])->Self: return self._elemwise_scalar_math[math.mul](s) fn __add__(self, s: Scalar[dtype])->Self: return self._elemwise_scalar_math[math.add](s) fn __sub__(self, s: Scalar[dtype])->Self: return self._elemwise_scalar_math[math.sub](s)

Elementwise vector-vector operations _elemwise_array_math[]()

This helper function implements multiplication, addition, and subtraction between two MojoArrays.Ā 

Mojo
fn _elemwise_array_math[func: fn[dtype: DType, width: Int](SIMD[dtype, width],SIMD[dtype, width])->SIMD[dtype, width]](self, other: Self) -> Self: alias simd_width: Int = simdwidthof[dtype]() let new_array = Self(self.numel) @parameter fn elemwise_vectorize[simd_width: Int](idx: Int) -> None: new_array._ptr.simd_store[simd_width](idx, func[dtype, simd_width](self._ptr.simd_load[simd_width](idx), other._ptr.simd_load[simd_width](idx))) vectorize[simd_width, elemwise_vectorize](self.numel) return new_array

We use _elemwise_array_math to implement addition, subtraction, and multiplication between two vectors

Mojo
fn __mul__(self, other: Self)->Self: return self._elemwise_array_math[math.mul](other) fn __add__(self, other: Self)->Self: return self._elemwise_array_math[math.add](other) fn __sub__(self, other: Self)->Self: return self._elemwise_array_math[math.sub](other)

Implementing the Heart Equation

With our MojoArray data structure in place, weā€™re now ready to implement our heart equation. In the main() function, we perform the following steps:

  1. Load Python modules NumPy and Matplotlib
  2. Create NumPy arrays and use that to instantiate our MojoArray data structure
  3. Implement the heart equation y = (x**2)**(1/3.) - 0.9*((3.3-(x*x)).sqrt())*(a[i]*3.14*x).sin()
  4. Loop over values of a to create the animation I showed at the top of the blog post.
Mojo
def main(): var np = Python.import_module("numpy") var plt = Python.import_module("matplotlib.pyplot") var np_arr = np.arange(-2,2,0.01) var x = MojoArray.from_numpy(np_arr) var fig = plt.figure() var ax = fig.add_subplot() ax.set_xlim([-3,3]) ax.set_ylim([-3,3]) var a = MojoArray.from_numpy(np.linspace(0,20,100)) for i in range(a.numel): var y = (x**2)**(1/3.) - 0.9*((3.3-(x*x)).sqrt())*(a[i]*3.14*x).sin() ax.cla() var title = ax.set_title("Mojo ā¤ļø Python") title.set_fontsize(20) ax.set_axis_off() ax.plot(x.to_numpy(),y.to_numpy(),'r') plt.pause(0.1) plt.draw()

Conclusion

Happy Valentinesā€™ week! Hope you enjoyed reading this blog post on how to implement the heart equation in Mojo. My hope is that you can reuse parts of this example in your own workflows. Download Mojo to run this example and share your feedback with us! Here are some additional resources to get started:

Until next timešŸ”„!

Shashank Prasanna
,
AI Developer Advocate

Shashank Prasanna

AI Developer Advocate

Shashank is an engineer, educator and doodler. He writes and talks about machine learning, specialized machine learning hardware (AI Accelerators) and AI Infrastructure in the cloud. He previously worked at Meta, AWS, NVIDIA, MathWorks (MATLAB) and Oracle in developer relations and marketing, product management, and software development roles and hold an M.S. in electrical engineering.