3️⃣ Triangles

Learning Objectives

• Understand how to parametrize triangles in 2D and 3D, and solve for their intersection with rays
• Put everything together, to render your mesh as a 2D image

Triangle Coordinates

The area inside a triangle can be defined by three (non-collinear) points $A$, $B$ and $C$, and can be written algebraically as a convex combination of those three points:

Or equivalently:

These $u, v$ are called "barycentric coordinates".

If we remove the bounds on $u$ and $v$, we get an equation for the plane containing the triangle. Play with the widget to understand the behavior of $u, v$.

one_triangle = t.tensor([[0, 0, 0], [4, 0.5, 0], [2, 3, 0]])
A, B, C = one_triangle
x, y, z = one_triangle.T

fig: go.FigureWidget = setup_widget_fig_triangle(x, y, z)
display(fig)


@interact(u=(-0.5, 1.5, 0.01), v=(-0.5, 1.5, 0.01))
def update(u=0.0, v=0.0):
    P = A + u * (B - A) + v * (C - A)
    fig.update_traces({"x": [P[0]], "y": [P[1]]}, 2)

Triangle-Ray Intersection

Given a ray $R(s) = O + sD$ with origin $O$ and direction $D$, our intersection algorithm will consist of three steps:

• Finding the intersection between the line and the plane containing the triangle, by solving the equation $P(u, v) = R(s)$;
• Checking if $u$ and $v$ are within the bounds of the triangle.
  • $0 \leq u,v$ and $u + v \leq 1$.
• Checking if the intersection is in the direction of travel of the ray (you can't see behind you).
  • $s \geq 0$.

Expanding the equation $P(u, v) = R(s)$, we have:

We can therefore find the coordinates s, u, v of the intersection point by solving the linear system above.

Exercise - implement triangle_ray_intersects

Using torch.linalg.solve and torch.stack, implement triangle_ray_intersects(A, B, C, O, D).

A few tips:

• If you have a 0-dimensional tensor with shape () containing a single value, use the item() method to convert it to a plain Python value.
• If you have a tensor of shape tensor.shape = (3, ...), then you can unpack it along the first dimension into three separate tensors the same way that you'd unpack a normal python list: s, u, v = tensor.
  • Note - if the dimension you want to unpack isn't at the start, a nice alternative is s, u, v = tensor.unbind(dim), which does the same thing but along the dimension given by dim rather than the first dimension.
• If your function isn't working, try making a simple ray and triangle with nice round numbers where you can work out manually if it should intersect or not, then debug from there.

Point = Float[Tensor, "points=3"]


def triangle_ray_intersects(A: Point, B: Point, C: Point, O: Point, D: Point) -> bool:
    """
    A: shape (3,), one vertex of the triangle
    B: shape (3,), second vertex of the triangle
    C: shape (3,), third vertex of the triangle
    O: shape (3,), origin point
    D: shape (3,), direction point

    Return True if the ray and the triangle intersect.
    """
    raise NotImplementedError()


tests.test_triangle_ray_intersects(triangle_ray_intersects)

def triangle_ray_intersects(A: Point, B: Point, C: Point, O: Point, D: Point) -> bool:
    """
    A: shape (3,), one vertex of the triangle
    B: shape (3,), second vertex of the triangle
    C: shape (3,), third vertex of the triangle
    O: shape (3,), origin point
    D: shape (3,), direction point

    Return True if the ray and the triangle intersect.
    """
    s, u, v = t.linalg.solve(t.stack([-D, B - A, C - A], dim=-1), O - A)
    return ((s >= 0) & (u >= 0) & (v >= 0) & (u + v <= 1)).item()

Single-Triangle Rendering

Implement raytrace_triangle using only one call to torch.linalg.solve.

Reshape the output and visualize with imshow. It's normal for the edges to look pixelated and jagged - using a small number of pixels is a good way to debug quickly.

If you think it's working, increase the number of pixels and verify that it looks less pixelated at higher resolution.

Views and Copies

It's critical to know when you are making a copy of a Tensor, versus making a view of it that shares the data with the original tensor. It's preferable to use a view whenever possible to avoid copying memory unnecessarily. On the other hand, modifying a view modifies the original tensor which can be unintended and surprising. Consult the documentation if you're unsure if a function returns a view. A short reference of common functions:

• Tensor.expand: always returns a view
• Tensor.view: always returns a view
• Tensor.detach: always returns a view
• Tensor.repeat: always copies
• torch.clone: always copies
• torch.flip: always copies (different than numpy.flip which returns a view)
• torch.tensor: always copies, but PyTorch recommends using .clone().detach() instead.
• torch.Tensor.contiguous: returns self if possible, otherwise a copy
• torch.transpose: returns a view if possible, otherwise (sparse tensor) a copy
• torch.reshape: returns a view if possible, otherwise a copy
• torch.flatten: returns a view if possible, otherwise a copy (different than numpy.flatten which returns a copy)
• einops.repeat: returns a view if possible, otherwise a copy
• einops.rearrange: returns a view if possible, otherwise a copy
• Basic indexing returns a view, while advanced indexing returns a copy.

Storage Objects

Calling storage() on a Tensor returns a Python object wrapping the underlying C++ array. This array is 1D regardless of the dimensionality of the Tensor. This allows you to look inside the Tensor abstraction and see how the actual data is laid out in RAM.

Note that a new Python wrapper object is generated each time you call storage(), and both x.storage() == x.storage() and x.storage() is x.storage() evaluates to False.

If you want to check if two Tensors share an underlying C++ array, you can compare their storage().data_ptr() fields. This can be useful for debugging.

Tensor._base

If x is a view, you can access the original Tensor with x._base. This is an undocumented internal feature that's useful to know. Consider the following code:

x = t.zeros(1024*1024*1024)
y = x[0]
del x

Here, y was created through basic indexing, so y is a view and y._base refers to x. This means del x won't actually deallocate the 4GB of memory, and that memory will remain in use which can be quite surprising. y = x[0].clone() would be an alternative here that does allow reclaiming the memory.

Exercise - implement raytrace_triangle

Below, you should implement raytrace_triangle, a function which checks whether each ray in rays intersects a single given triangle.

def raytrace_triangle(
    rays: Float[Tensor, "nrays rayPoints=2 dims=3"],
    triangle: Float[Tensor, "trianglePoints=3 dims=3"],
) -> Bool[Tensor, " nrays"]:
    """
    For each ray, return True if the triangle intersects that ray.
    """
    raise NotImplementedError()


tests.test_raytrace_triangle(raytrace_triangle)

A = t.tensor([1, 0.0, -0.5])
B = t.tensor([1, -0.5, 0.0])
C = t.tensor([1, 0.5, 0.5])
num_pixels_y = num_pixels_z = 15
y_limit = z_limit = 0.5

# Plot triangle & rays
test_triangle = t.stack([A, B, C], dim=0)
rays2d = make_rays_2d(num_pixels_y, num_pixels_z, y_limit, z_limit)
triangle_lines = t.stack([A, B, C, A, B, C], dim=0).reshape(-1, 2, 3)
render_lines_with_plotly(rays2d, triangle_lines)

# Calculate and display intersections
intersects = raytrace_triangle(rays2d, test_triangle)
img = intersects.reshape(num_pixels_y, num_pixels_z).int()
imshow(img, origin="lower", width=600, title="Triangle (as intersected by rays)")

def raytrace_triangle(
    rays: Float[Tensor, "nrays rayPoints=2 dims=3"],
    triangle: Float[Tensor, "trianglePoints=3 dims=3"],
) -> Bool[Tensor, " nrays"]:
    """
    For each ray, return True if the triangle intersects that ray.
    """
    NR = rays.size(0)

    # Triangle is [[Ax, Ay, Az], [Bx, By, Bz], [Cx, Cy, Cz]]
    A, B, C = einops.repeat(triangle, "pts dims -> pts NR dims", NR=NR)
    assert A.shape == (NR, 3)

    # Each element of `rays` is [[Ox, Oy, Oz], [Dx, Dy, Dz]]
    O, D = rays.unbind(dim=1)
    assert O.shape == (NR, 3)

    # Define matrix on left hand side of equation
    mat: Float[Tensor, "NR 3 3"] = t.stack([-D, B - A, C - A], dim=-1)

    # Get boolean of where matrix is singular, and replace it with the identity in these positions
    # Note - this works because mat[is_singular] has shape (NR_where_singular, 3, 3), so we
    # can broadcast the identity matrix to that shape.
    dets: Float[Tensor, "NR"] = t.linalg.det(mat)
    is_singular = dets.abs() < 1e-6
    mat[is_singular] = t.eye(3)

    # Define vector on the right hand side of equation
    vec = O - A

    # Solve eqns
    sol: Float[Tensor, "NR 3"] = t.linalg.solve(mat, vec)
    s, u, v = sol.unbind(dim=-1)

    # Return boolean of (matrix is nonsingular) && (solution is in correct range implying intersection)
    return (s >= 0) & (u >= 0) & (v >= 0) & (u + v <= 1) & ~is_singular

Mesh Loading

Use the given code to load the triangles for your Pikachu. By convention, files written with torch.save end in the .pt extension, but these are actually just zip files.

triangles = t.load(section_dir / "pikachu.pt", weights_only=True)

Mesh Rendering

For our purposes, a mesh is just a group of triangles, so to render it we'll intersect all rays and all triangles at once. We previously just returned a boolean for whether a given ray intersects the triangle, but now it's possible that more than one triangle intersects a given ray.

For each ray (pixel) we will return a float representing the minimum distance to a triangle if applicable, otherwise the special value float('inf') representing infinity. We won't return which triangle was intersected for now.

Note - by distance to a triangle, we specifically mean distance along the x-dimension, not Euclidean distance.

Exercise - implement raytrace_mesh

Implement raytrace_mesh and as before, reshape and visualize the output. Your Pikachu is centered on (0, 0, 0), so you'll want to slide the ray origin back to at least x=-2 to see it properly.

Reminder - t.linalg.solve (and most batched operations) can accept multiple dimensions as being batch dims. Previously, you've just used NR (the number of rays) as the batch dimension, but you can also use (NR, NT) (the number of rays and triangles) as your batch dimensions, so you can solve for all rays and triangles at once.

def raytrace_mesh(
    rays: Float[Tensor, "nrays rayPoints=2 dims=3"],
    triangles: Float[Tensor, "ntriangles trianglePoints=3 dims=3"],
) -> Float[Tensor, " nrays"]:
    """
    For each ray, return the distance to the closest intersecting triangle, or infinity.
    """
    raise NotImplementedError()


tests.test_raytrace_mesh(raytrace_mesh)

num_pixels_y = 120
num_pixels_z = 120
y_limit = z_limit = 1

rays = make_rays_2d(num_pixels_y, num_pixels_z, y_limit, z_limit)
rays[:, 0] = t.tensor([-2, 0.0, 0.0])
dists = raytrace_mesh(rays, triangles)
intersects = t.isfinite(dists).view(num_pixels_y, num_pixels_z)
dists_square = dists.view(num_pixels_y, num_pixels_z)
img = t.stack([intersects, dists_square], dim=0)

fig = px.imshow(img, facet_col=0, origin="lower", color_continuous_scale="magma", width=1000)
fig.update_layout(coloraxis_showscale=False)
for i, text in enumerate(["Intersects", "Distance"]):
    fig.layout.annotations[i]["text"] = text
fig.show()

def raytrace_mesh(
    rays: Float[Tensor, "nrays rayPoints=2 dims=3"],
    triangles: Float[Tensor, "ntriangles trianglePoints=3 dims=3"],
) -> Float[Tensor, " nrays"]:
    """
    For each ray, return the distance to the closest intersecting triangle, or infinity.
    """
    NR = rays.size(0)
    NT = triangles.size(0)

    # Each triangle is [[Ax, Ay, Az], [Bx, By, Bz], [Cx, Cy, Cz]]
    triangles = einops.repeat(triangles, "NT pts dims -> pts NR NT dims", NR=NR)
    A, B, C = triangles
    assert A.shape == (NR, NT, 3)

    # Each ray is [[Ox, Oy, Oz], [Dx, Dy, Dz]]
    rays = einops.repeat(rays, "NR pts dims -> pts NR NT dims", NT=NT)
    O, D = rays
    assert O.shape == (NR, NT, 3)

    # Define matrix on left hand side of equation
    mat: Float[Tensor, "NR NT 3 3"] = t.stack([-D, B - A, C - A], dim=-1)
    # Get boolean of where matrix is singular, and replace it with the identity in these positions
    dets: Float[Tensor, "NR NT"] = t.linalg.det(mat)
    is_singular = dets.abs() < 1e-6
    mat[is_singular] = t.eye(3)

    # Define vector on the right hand side of equation
    vec: Float[Tensor, "NR NT 3"] = O - A

    # Solve eqns (note, s is the distance along ray)
    sol: Float[Tensor, "NR NT 3"] = t.linalg.solve(mat, vec)
    s, u, v = sol.unbind(-1)

    # Scale s by Dx to get the distance along ray
    s *= D[..., 0]

    # Get boolean of intersects, and use it to set distance = infinity when there is no intersection
    intersects = (s >= 0) & (u >= 0) & (v >= 0) & (u + v <= 1) & ~is_singular
    s[~intersects] = float("inf")  # t.inf

    # Get the minimum distance (over all triangles) for each ray
    return einops.reduce(s, "NR NT -> NR", "min")


tests.test_raytrace_mesh(raytrace_mesh)

num_pixels_y = 120
num_pixels_z = 120
y_limit = z_limit = 1

rays = make_rays_2d(num_pixels_y, num_pixels_z, y_limit, z_limit)
rays[:, 0] = t.tensor([-2, 0.0, 0.0])
dists = raytrace_mesh(rays, triangles)
intersects = t.isfinite(dists).view(num_pixels_y, num_pixels_z)
dists_square = dists.view(num_pixels_y, num_pixels_z)
img = t.stack([intersects, dists_square], dim=0)

fig = px.imshow(img, facet_col=0, origin="lower", color_continuous_scale="magma", width=1000)
fig.update_layout(coloraxis_showscale=False)
for i, text in enumerate(["Intersects", "Distance"]):
    fig.layout.annotations[i]["text"] = text
fig.show()

Congratulations, you've now got to the end of the exercises! Read on for some more bonus material.