My experiment on porting Python to Mojo: supercharging a Physics Simulation

Mojo promises Python’s ease of use with C-like performance. But what does it take to port a real Python project? I converted a hard-sphere collision simulation from Python to Mojo to find out.

Here’s a look at the key changes, Mojo’s Python interoperability, and the lessons I learned along the way.

The code can be found here.

On the Hard-Sphere Collision Problem

The hard-sphere collision problem studies when and how two (or more) perfectly rigid, non‑penetrating spheres collide elastically: predict the collision time (solve when center distance equals sum of radii) then update velocities by conserving momentum and kinetic energy along the line of centers. It is a classic model used in kinetic theory, molecular dynamics, and as an algorithms exercise (priority queues / event-driven simulation).

This is often used as a demonstration in algorithm courses (priority queues). My implementation follows the approach described here: https://introcs.cs.princeton.edu/java/assignments/collisions.html

My Python implementation:
https://github.com/kanhua/hard-sphere-collision/blob/main/python/hard_sphere_demo.py

Now I will briefly describe how I rewrite this python code to Mojo:

Key Changes: From class to struct and Traits

The most significant changes involved moving from Python’s dynamic classes to Mojo’s static, performance-oriented structures.

1. Value Semantics with struct

In CPython, names bind to references (pointers) to heap‑allocated PyObject instances, each carrying a header (type pointer, refcount, etc.). A plain Python list of numbers therefore holds pointers, adding indirection and per‑element overhead that can hurt cache locality (though libraries like array, NumPy, or memoryview achieve contiguous unboxed layouts). In Mojo, a struct is a value type: when you store a struct (or a primitive) in a List its bytes (fields inline, with normal padding/alignment) are laid out contiguously, avoiding per-element object headers and improving locality. This shift from ubiquitous boxed objects to explicit value types is a key pillar of Mojo performance.

• Python:

  class Particle:
      # ...

• Mojo:

  struct Particle:
      var px: float
      # ...

2. The Use of Traits

A plain struct is just a data container. To make it useful—for example, to store it in a List or use it in a priority queue—it must conform to certain behaviors. In Mojo, these behaviors are defined by traits.

Traits are contracts that a type can implement. For our simulation, the Event struct needs to be sortable in a priority queue, so it must implement the Comparable trait, which defines comparison methods like __lt__ and __gt__.

@fieldwise_init
struct Particle(Copyable, Movable):
    var px: float
    var py: float
    var vx: float
    var vy: float

In the snippet above, we explicitly list the trait conformances (Copyable, Movable) and use @fieldwise_init to auto-generate an initializer.

3. A New Memory Mindset

For many Python developers this is the steepest part of the learning curve. In Python we freely pass object references—ownership, aliasing, and lifetime are mostly invisible (ref counting + GC). In Mojo (closer to Rust/C++ conceptually) you must consider: Where does data live? When is it copied vs moved? How does layout affect cache behavior?

The solution was to refactor the Event struct to store the integer indices of particles in the main particle list, rather than direct references. This is a common pattern when porting from a garbage-collected language and, while a hurdle, it forces a more memory-aware design that unlocks performance.

Mojo’s Superpower: Seamless Python Interop

You can import and use existing Python libraries directly—no need to reimplement tooling up front. This project uses Python modules for configuration, logging, and visualization (e.g. matplotlib).

Reading Configuration and Saving Data

Instead of writing a new file parser, we just use Python’s tomllib and csv modules.

def run_headless(...):
    # Import Python modules
    csv = Python.import_module("csv")
    builtins = Python.import_module("builtins")

    # Use them as if they were native
    file_handle = builtins.open(logfile, "w")
    writer = csv.writer(file_handle)
    writer.writerow(["time", "kinetic_energy"])
    # ...

Mojo seamlessly calls into the Python runtime, and we can easily convert data between the two worlds.

Speed comparison

Below is the comparison of the execution time of the two implementations.

Generating Plots with Matplotlib

The Mojo version saves each simulation frame as an image using Python’s matplotlib. This is a fantastic demonstration of interoperability.

def animate_simulation(...):
    plt = Python.import_module("matplotlib.pyplot")
    # ...
    # Pass Mojo lists to Python for plotting
    ax.scatter(
        Python.list(pxs),      # pxs is a Mojo List[float]
        Python.list(pys),      # pys is a Mojo List[float]
        s=Python.list(sizes)
    )
    plt.savefig(filename)

This ability to leverage the vast Python ecosystem is Mojo’s killer feature, making it practical for real-world projects today.

Final Thoughts

The Mojo version runs significantly faster, but just as important, Python interoperability provides a practical migration path: move only the hot loops first, keep ecosystem leverage. The main conceptual shift is adopting value-based structs and explicit traits instead of dynamic, reference-based objects.

Next step: experiment with Mojo’s GPU capabilities to push the hard-sphere simulation further.