[0.1] - Ray Tracing

Colab: exercises | solutions

Please send any problems / bugs on the #errata channel in the Slack group, and ask any questions on the dedicated channels for this chapter of material.

If you want to change to dark mode, you can do this by clicking the three horizontal lines in the top-right, then navigating to Settings → Theme.

Links to all other chapters: (0) Fundamentals, (1) Transformer Interpretability, (2) RL.

Introduction

Today we'll be practicing batched matrix operations in PyTorch by writing a basic graphics renderer. We'll start with an extremely simplified case and work up to rendering your very own 3D Pikachu!

We'll also be touching on some general topics which will be important going forwards in this course, such as:

• Using GPT systems to assist your learning and coding
• Typechecking, and good coding practices
• Debugging, with VSCode's built-in run & debug features

Content & Learning Objectives

1️⃣ Rays & Segments

This section introduces the key ideas and concepts in today's exercises, and guides you through some basic functions involving creating & using 2D rays.

Learning Objectives

• Learn how to create PyTorch tensors in a variety of ways
• Understand how to parametrize lines and rays in 2D
• Learn about type annotations and linear operations in PyTorch

2️⃣ Batched Operations

In the next section, you'll extend your work from the first section to perform batched operations, i.e. operations over several different dimensions at once.

Learning Objectives

• Learn about some important concepts related to batched operations, e.g. broadcasting and logical reductions
• Understand and use the einops library
• Apply this knowledge to create & work with a batch of rays

3️⃣ Triangles

In the final section we move into the 2D and 3D realms, and build up to rendering a full 3D mesh as a 2D image.

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

4️⃣ Video & Lighting

In the last set of exercises, you can implement some more mathematically complex forms of raytracing to get things like dynamic lighting & videos. There's also some optional material at the very end which covers the pytest library, and has you write your own tests.

Learning Objectives

• Learn how surface normal vectors can be used to compute lighting effects
• Render your figures in video form (as an animated camera pan)
• (optional) Learn about the pytest library, and write your own tests for the functions you wrote

Setup code

import os
import sys
from functools import partial
from pathlib import Path
from typing import Callable

import einops
import plotly.express as px
import plotly.graph_objects as go
import torch as t
from IPython.display import display
from ipywidgets import interact
from jaxtyping import Bool, Float
from torch import Tensor
from tqdm import tqdm

# Make sure exercises are in the path
chapter = "chapter0_fundamentals"
section = "part1_ray_tracing"
root_dir = next(p for p in Path.cwd().parents if (p / chapter).exists())
exercises_dir = root_dir / chapter / "exercises"
section_dir = exercises_dir / section
if str(exercises_dir) not in sys.path:
    sys.path.append(str(exercises_dir))

import part1_ray_tracing.tests as tests
from part1_ray_tracing.utils import (
    render_lines_with_plotly,
    setup_widget_fig_ray,
    setup_widget_fig_triangle,
)
from plotly_utils import imshow

MAIN = __name__ == "__main__"

with open(section_dir / "pikachu.pt", "rb") as f:
    triangles = t.load(f)