Möller–Trumbore ray-triangle intersection algorithm (for ray tracing) in Python and Numpy, vectorized

moller_trumbore_vectorized.py

from typing import Tuple

import numpy as np

from numba import jit


@jit
def rays_triangles_intersection(
    ray_origin: np.ndarray, ray_directions: np.ndarray, triangles_vertices: np.ndarray
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Möller–Trumbore intersection algorithm for calculating whether the ray intersects the triangle
    and for which t-value. Based on: https://github.com/kliment/Printrun/blob/master/printrun/stltool.py,
    which is based on:
    http://en.wikipedia.org/wiki/M%C3%B6ller%E2%80%93Trumbore_intersection_algorithm

    Parameters
    ----------
    ray_origin : np.ndarray(3)
        Origin coordinate (x, y, z) from which the ray is fired
    ray_directions : np.ndarray(n, 3)
        Directions (dx, dy, dz) in which the rays are going
    triangle_vertices : np.ndarray(m, 3, 3)
        3D vertices of multiple triangles

    Returns
    -------
    tuple[np.ndarray<bool>(n, m), np.ndarray(n, m)]
        The first array indicates whether or not there was an intersection, the second array
        contains the t-values of the intersections
    """

    output_shape = (len(ray_directions), len(triangles_vertices))

    all_rays_t = np.zeros(output_shape)
    all_rays_intersected = np.full(output_shape, True)

    v1 = triangles_vertices[:, 0]
    v2 = triangles_vertices[:, 1]
    v3 = triangles_vertices[:, 2]

    eps = 0.000001

    edge1 = v2 - v1
    edge2 = v3 - v1

    for i, ray in enumerate(ray_directions):
        all_t = np.zeros((len(triangles_vertices)))
        intersected = np.full((len(triangles_vertices)), True)

        pvec = np.cross(ray, edge2)

        det = np.sum(edge1 * pvec, axis=1)

        non_intersecting_original_indices = np.absolute(det) < eps

        all_t[non_intersecting_original_indices] = np.nan
        intersected[non_intersecting_original_indices] = False

        inv_det = 1.0 / det

        tvec = ray_origin - v1

        u = np.sum(tvec * pvec, axis=1) * inv_det

        non_intersecting_original_indices = (u < 0.0) + (u > 1.0)
        all_t[non_intersecting_original_indices] = np.nan
        intersected[non_intersecting_original_indices] = False

        qvec = np.cross(tvec, edge1)

        v = np.sum(ray * qvec, axis=1) * inv_det

        non_intersecting_original_indices = (v < 0.0) + (u + v > 1.0)

        all_t[non_intersecting_original_indices] = np.nan
        intersected[non_intersecting_original_indices] = False

        t = (
            np.sum(
                edge2 * qvec,
                axis=1,
            )
            * inv_det
        )

        non_intersecting_original_indices = t < eps
        all_t[non_intersecting_original_indices] = np.nan
        intersected[non_intersecting_original_indices] = False

        intersecting_original_indices = np.invert(non_intersecting_original_indices)
        all_t[intersecting_original_indices] = t[intersecting_original_indices]

        all_rays_t[i] = all_t
        all_rays_intersected[i] = intersected

    return all_rays_intersected, all_rays_t

Möller-Trumbore algorithm in Python.md

On my 2014 MacBook Pro (2.5 GHz), this vectorized version is about 250 times faster when used for a large numbers of triangles compared to the original algorithm. You can also parallelize this code over multiple threads, but due to the large overhead of creating a thread pool, this is only really useful when you have a very large number of rays.

I also tried to vectorize the rays, instead of using a loop. I did this by using:

np.reshape(np.array(np.meshgrid(np.arange(0, len(ray_direction), 1), np.arange(0, len(triangles_vertices), 1))), (2, resulting_array_length))

to create all combinations of indices of the rays and the triangles. However, using this process, combined with using np.take() to actually create these combinations is about twice as slow as the code below.

In case you do want to run it in parallel, you can do it like this:

from joblib import Parallel, delayed

res = Parallel(n_jobs=-1)(delayed(rays_triangles_intersection)(ray_origin, np.array([q]), triangles_vertices) for q in ray_directions)

On my machine, for a large (20k) number of rays, this gives about a 2x speed up over running it non-parallel.

For a regular (with parallelization) implementation, see the file at the bottom of this page.

script_example_usage.py

intersected, t_values = rays_triangles_intersection(
    np.array([-1, 0, 0.5]),
    np.array([[0.5, 0, 0],
              [0.7, 0.8, 0]]),
    np.array([[[0, -1, 0], [0, 1, 0], [0, 0, 1]],
              [[1, -1, 0], [2, 1, 0], [1, 2, 1]]]),
)