使用 Mojo 中的 Mandelbrot 和 Python 绘图

news2024/12/26 23:53:42

Mojo不仅非常适合编写高性能代码，而且还允许我们利用庞大的Python生态系统中的库和工具。通过无缝的Python互操作性，Mojo可以使用Python来做它擅长的事情，特别是gui，而不会牺牲关键代码的性能。让我们使用经典的Mandelbrot集合算法并在Mojo中实现它。

本教程展示了Mojo的两个方面。首先，证明了Mojo可以为非规则应用开发快速程序;它还展示了如何利用Python来可视化结果。

import benchmark
from math import iota
from sys import num_physical_cores
from algorithm import parallelize, vectorize
from complex import ComplexFloat64, ComplexSIMD
from python import Python

alias float_type = DType.float32
alias int_type = DType.int32
alias simd_width = 2 * simdwidthof[float_type]()
alias unit = benchmark.Unit.ms

首先设置一些参数，你可以尝试更改它们以查看不同的结果:

alias width = 960
alias height = 960
alias MAX_ITERS = 200

alias min_x = -2.0
alias max_x = 0.6
alias min_y = -1.5
alias max_y = 1.5

下面我们定义一个简单的矩阵结构体:

@value
struct Matrix[type: DType, rows: Int, cols: Int]:
    var data: DTypePointer[type]

    fn __init__(inout self):
        self.data = DTypePointer[type].alloc(rows * cols)

    fn __getitem__(self, row: Int, col: Int) -> Scalar[type]:
        return self.data.load(row * cols + col)

    fn store[width: Int = 1](self, row: Int, col: Int, val: SIMD[type, width]):
        self.data.store[width=width](row * cols + col, val)

曼德尔布罗特算法的核心是为每个像素计算一个迭代的复函数，直到它“逃离”半径为2的复圆，并计算逃离的迭代次数:
在这里插入图片描述

# Compute the number of steps to escape.
def mandelbrot_kernel(c: ComplexFloat64) -> Int:
    z = c
    for i in range(MAX_ITERS):
        z = z * z + c