An example of defining a point structure in Mojo

Mojo is a relatively new language (introduced in 2022) designed for parallel computing. It is compatible with SIMD (Single Instruction Multiple Data) architectures and is intentionally interoperable with Python. SIMD allows for the application of vectorized operations, resulting in significant speedups in processing.

Currently, Mojo is available only for Ubuntu and macOS.

In this example, we will define a point geometry and implement it using Mojo.

In this context, using SIMD to store point coordinates is advantageous because it enables the application of vectorized operations to those coordinates.

To define a SIMD variable, it’s necessary to specify both the data type and the number of stored values. Here is an example of initializing a SIMD variable:

var vec = SIMD[DType.int8, 4](1, 2, 3, 4)

In Mojo, you cannot use Float32 or Float64 directly to define the variable type. Instead, you need to use DType.float32 or DType.float64 for floating-point values.

Another requirement for SIMD variables is that the number of elements must be a power of two. Therefore, it is not possible to use SIMD[DType.float64, 3] to store just the x, y, and z coordinates of a point. To meet the minimum size requirement of 4 elements, you can add a time variable (t) of type Float64 to the point coordinates. This avoids wasting memory space.

The point coordinates are stored as follows:

var coords: SIMD[DType.float64, 4]

Similar to Python, a class-like structure in Mojo is initialized using the __init__ method. To indicate that the instance is mutable, the self parameter is marked as inout:

fn __init__(inout self, x: Float64, y: Float64, z: Float64, t: Float64 = 0.0)

Since the values are stored in a SIMD variable, the elements are accessed by their index:

fn x(self) -> Float64:
	return self.coords[0]
    

To calculate the distance between two points, you cannot use the sum() or sqrt() functions directly on a SIMD variable. Instead, you must extract the scalar values by index and perform the scalar operations manually.

The square root function is imported from the math module like this:

from math import sqrt

Here is the complete implementation of the point structure:

from math import sqrt

struct Point :

    var coords: SIMD[DType.float64, 4]

    fn __init__(inout self, x: Float64, y: Float64, z: Float64, t: Float64 = 0.0):

        self.coords = SIMD[DType.float64, 4](x, y, z, t)

    fn x(self) -> Float64:
        return self.coords[0]

    fn y(self) -> Float64:
        return self.coords[1]

    fn z(self) -> Float64:
        return self.coords[2]

    fn t(self) -> Float64:
        return self.coords[3]

    fn distance_to(self, other: Point) -> Float64:

        var delta = self.coords - other.coords
        var sum_of_squares = delta[0]*delta[0] + delta[1]*delta[1] + delta[2]*delta[2];
        return sqrt(sum_of_squares[0])
    
    

And here’s an example of how to use this structure in a main function:

fn main():

    var pt1 = Point(1.0, 2.0, 3.0, 1222.34)
    var pt2 = Point(2.0, 3.0, 3.0, 22.34)
    print(pt1.distance_to(pt2))
      

You can also view this example on GitLab:

Mojo Point Structure Example