Mojo 中的光线追踪

本教程基于热门教程《C++ 中的可理解光线追踪》。该教程很好地描述了数学解释，因此在 Mojo 中实现基本光线追踪器时，我们只会向您指出相应的部分以供参考。

步骤 1：基本定义

我们首先定义一个Vec3f结构体，用于表示 3D 空间中的矢量以及 RGB 像素。我们将使用SIMD矢量的表示形式来启用矢量化操作。该SIMD类型是固定大小的矢量，其大小必须是 2 的幂。因此，我们将使用 4 的大小，并始终用 0 填充底层存储。

from math import rsqrt


@register_passable("trivial")
struct Vec3f:
    var data: SIMD[DType.float32, 4]

    @always_inline
    fn __init__(inout self, x: Float32, y: Float32, z: Float32):
        self.data = SIMD[DType.float32, 4](x, y, z, 0)

    @always_inline
    fn __init__(inout self, data: SIMD[DType.float32, 4]):
        self.data = data

    @always_inline
    @staticmethod
    fn zero() -> Vec3f:
        return Vec3f(0, 0, 0)

    @always_inline
    fn __sub__(self, other: Vec3f) -> Vec3f:
        return self.data - other.data

    @always_inline
    fn __add__(self, other: Vec3f) -> Vec3f:
        return self.data + other.data

    @always_inline
    fn __matmul__(self, other: Vec3f) -> Float32:
        return (self.data * other.data).reduce_add()

    @always_inline
    fn __mul__(self, k: Float32) -> Vec3f:
        return self.data * k

    @always_inline
    fn __neg__(self) -> Vec3f:
        return self.data * -1.0

    @always_inline
    fn __getitem__(self, idx: Int) -> SIMD[DType.float32, 1]:
        return self.data[idx]

    @always_inline
    fn cross(self, other: Vec3f) -> Vec3f:
        var self_zxy = self.data.shuffle[2, 0, 1, 3]()
        var other_zxy = other.data.shuffle[2, 0, 1, 3]()
        return (self_zxy * other.data - self.data * other_zxy).shuffle[
            2, 0, 1, 3
        ]()

    @always_inline
    fn normalize(self) -> Vec3f:
        return self.data * rsqrt(self @ self)

我们现在定义Image结构，它将存储图像的 RGB 像素。它还包含一个将此 Mojo 结构转换为 NumPy 图像的方法，该方法将用于实现简单的显示机制。我们还将实现一个从磁盘加载 PNG 文件的函数。

首先安装所需的库：

%%python
from importlib.util import find_spec
import shutil
import subprocess

fix = """
-------------------------------------------------------------------------
fix following the steps here:
    https://github.com/modularml/mojo/issues/1085#issuecomment-1771403719
-------------------------------------------------------------------------
"""

def install_if_missing(name: str):
    if find_spec(name):
        return

    print(f"{
     name} not found, installing...")
    try:
        if shutil.which('python3'): python = "python3"
        elif shutil.which('python'): python = "python"
        else: raise ("python not on path" + fix)
        subprocess.check_call([python, "-m", "pip", "install", name])
    except:
        raise ImportError(f"{
     name} not found" + fix)

install_if_missing("numpy")
install_if_missing("matplotlib")

from python import Python
from python import PythonObject
from memory.unsafe_pointer import UnsafePointer

struct Image:
    # reference count used to make the object efficiently copyable
    var rc: UnsafePointer[Int]
    # the two dimensional image is represented as a flat array
    var pixels: UnsafePointer[Vec3f]
    var height: Int
    var width: Int

    fn __init__(inout self, height: Int, width: Int):
        self.height = height
        self.width = width
        self.pixels = UnsafePointer[Vec3f].alloc(self.height * self.width)
        self.rc = UnsafePointer[Int].alloc(1)
        self.rc[] = 1

    fn __copyinit__(inout self, other: Self):
        other._inc_rc()
        self.pixels = other.pixels
        self.rc = other.rc
        self.height = other.height
        self.width = other.width

    fn __del__(owned self):
        self._dec_rc()

    fn _dec_rc(self):
        if self.rc[] > 1:
            self.rc[] -= 1
            return
        self._free()

    fn _inc_rc(self):
        self.rc[] += 1

    fn _free(self):
        self.rc.free()
        self.pixels.free()

    @always_inline
    fn set(self, row: Int, col: Int, value: Vec3f) -> None:
        self.pixels[self