How Lenses Form Images

import urllib.request
from pathlib import Path

import torch
import kornia
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon, Rectangle, Circle, FancyArrowPatch
from matplotlib.collections import LineCollection

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

_ = torch.manual_seed(0)

# Locate the section's two photos regardless of the working directory; download them once if absent.
_CHAPTER_DIR = Path("notebooks/CV/mit-foundations/chapter-6-lenses")
BASE_DIR = _CHAPTER_DIR if _CHAPTER_DIR.is_dir() else Path(".")
ASSETS_DIR = BASE_DIR / "assets"
ASSET_URL = ("https://raw.githubusercontent.com/pantelis/eng-ai-agents/main/notebooks/CV/"
             "mit-foundations/chapter-6-lenses/assets")
for _name in ("dog_scene.png", "ruler_single.png"):
    if not (ASSETS_DIR / _name).exists():
        ASSETS_DIR.mkdir(parents=True, exist_ok=True)
        urllib.request.urlretrieve(f"{ASSET_URL}/{_name}", ASSETS_DIR / _name)

# Section 6.2, physics primitives.
# Snell's law and its paraxial small-angle form, the surface-tilt angle, the
# lensmaker's focal length and conjugate image distance, and the four-angle
# bookkeeping that the thin-lens derivation composes from them. (Book the book.)


def snells_law(n1, n2, theta1):
    """Refraction angle θ₂ at a flat interface (n₁ sin θ₁ = n₂ sin θ₂).

    Angles in radians, measured from the surface normal. Returns NaN past
    the critical angle (total internal reflection).
    """
    n1 = torch.as_tensor(n1, dtype=torch.float32)
    n2 = torch.as_tensor(n2, dtype=torch.float32)
    theta1 = torch.as_tensor(theta1, dtype=torch.float32)
    return torch.asin((n1 / n2) * torch.sin(theta1))


def paraxial_snell(n1, n2, theta1):
    """Paraxial Snell's law: n1·θ1 = n2·θ2, the small-angle form of snells_law.

    Differs from snells_law by sin(θ) ≈ θ. Used throughout §6.2's lensmaker
    derivation where the four ray angles all sit in the paraxial regime.
    Angles in radians, measured from the surface normal.
    """
    n1 = torch.as_tensor(n1, dtype=torch.float32)
    n2 = torch.as_tensor(n2, dtype=torch.float32)
    theta1 = torch.as_tensor(theta1, dtype=torch.float32)
    return (n1 / n2) * theta1


def surface_angle(c, R):
    """Tilt angle θ_S of a spherical surface at height c, radius R.

    Paraxial form: θ_S ≈ c/R. (Book §6.2, Figure 6.5.)
    """
    c = torch.as_tensor(c, dtype=torch.float32)
    R = torch.as_tensor(R, dtype=torch.float32)
    return c / R


def focal_length(n, R1, R2):
    """Focal length of a thin lens via the lensmaker's equation.

    1/f = (n-1) * (1/R1 + 1/R2).  (Book §6.2, Eq. derived from Table 6.1.)
    Sign convention: both R1 and R2 are positive magnitudes. A biconvex lens
    has R1 > 0 and R2 > 0. (Book §6.2, note this differs from Hecht 2016,
    which uses signed radii.)
    """
    n = torch.as_tensor(n, dtype=torch.float32)
    R1 = torch.as_tensor(R1, dtype=torch.float32)
    R2 = torch.as_tensor(R2, dtype=torch.float32)
    return 1.0 / ((n - 1.0) * (1.0 / R1 + 1.0 / R2))


def image_distance(a, f):
    """Image distance b from object distance a and focal length f.

    1/a + 1/b = 1/f  =>  b = a*f / (a - f).  (Book §6.2.)
    Returns inf when a == f (object at the focal point → image at infinity).
    """
    a = torch.as_tensor(a, dtype=torch.float32)
    f = torch.as_tensor(f, dtype=torch.float32)
    return a * f / (a - f)


def lensmaker_angle_bookkeeping(n, R1, R2, c, theta1_deg):
    """Compute the four paraxial ray angles and two surface tilts through a thin lens.

    Applies paraxial Snell's law at each surface using the helpers defined
    above. Returns (theta1_axis, theta2_axis, theta4_axis, thetaS1, thetaS2),
    all in radians measured from the optical axis (rays) or from horizontal
    (surface tilts). (Book §6.2, Table 6.1.)
    """
    theta1_axis = torch.deg2rad(theta1_deg)

    # Surface tilts at the ray's hit height (both radii now positive magnitudes
    # per the book's convention).
    thetaS1 = surface_angle(c, R1)
    thetaS2 = surface_angle(c, R2)

    # Refract entering the glass at the front surface.
    theta1_normal = theta1_axis + thetaS1
    theta2_normal = paraxial_snell(1.0, n, theta1_normal)
    theta2_axis = theta2_normal - thetaS1

    # Refract exiting the glass at the back surface.
    theta3_normal = thetaS2 - theta2_axis
    theta4_normal = paraxial_snell(n, 1.0, theta3_normal)
    theta4_axis = thetaS2 - theta4_normal

    return theta1_axis, theta2_axis, theta4_axis, thetaS1, thetaS2


# Sanity checks, one per function.
# Snell's law: 30° in air → ≈19.47° in glass (n=1.5).
print(f"{torch.rad2deg(snells_law(1.0, 1.5, torch.deg2rad(torch.tensor(30.0)))).item():.2f}°")
# Paraxial Snell: 5° in air → ≈3.33° in glass (n=1.5), close to exact 3.33°.
print(f"{torch.rad2deg(paraxial_snell(1.0, 1.5, torch.deg2rad(torch.tensor(5.0)))).item():.2f}°")
# Surface tilt: radius 50 mm, ray hit at height 5 mm → about 0.1 rad ≈ 5.73°.
print(f"{torch.rad2deg(surface_angle(5.0, 50.0)).item():.2f}°")
# Biconvex lens, n=1.5, R1=R2=50 mm → f = 50 mm. Object at 200 mm → b ≈ 67 mm.
f = focal_length(1.5, 50.0, 50.0)
b = image_distance(200.0, f)
print(f"f = {f.item():.2f} mm, b(a=200) = {b.item():.2f} mm")

19.47°
3.33°
5.73°
f = 50.00 mm, b(a=200) = 66.67 mm