{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Example 04: Optical Raytracing (Advanced)\n",
    "\n",
    "This example demonstrates advanced coordinate frame usage in optical raytracing. We'll trace parallel rays through lens surfaces and see how frame transformations simplify optical system design.\n",
    "\n",
    "**Prerequisites**: Basic optics (Snell's Law), numerical methods  \n",
    "**Skip ahead**: See \"TL;DR\" section below for the key insight\n",
    "\n",
    "**Key demonstrations**:\n",
    "1. **Frame transformations simplify optics**: Position/rotate lenses by modifying frames - the raytracing algorithm stays unchanged\n",
    "2. **Interchangeable profiles**: Swap optical surfaces (spherical, hyperbolic) without changing the tracing code\n",
    "\n",
    "**Note**: This is a simplified raytracer for demonstration. Production implementations would handle:\n",
    "  - Rays missing the surface (returns error instead of graceful fallback)\n",
    "  - Rays starting inside the surface\n",
    "  - Numerical stability near grazing incidence\n",
    "  - concave surfaces"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "## TL;DR - The Key Insight\n",
    "\n",
    "If you're skipping the technical details, here's what matters:\n",
    "\n",
    "**Scenario**: We want to trace light rays through a lens. The lens can be perfectly aligned or tilted at some angle.\n",
    "\n",
    "**The traditional approach**: Manually rotate all geometric calculations - intersection points, surface normals, ray directions - every time you change the lens orientation. Error-prone and tedious.\n",
    "\n",
    "**The hazy approach**: Put the lens in its own coordinate frame. All geometric calculations happen in the lens's local frame. Want to tilt the lens 5°? Change **one line**:\n",
    "\n",
    "```python\n",
    "# Aligned lens\n",
    "lens_frame = root.make_child(\"lens\").translate(x=-25)\n",
    "\n",
    "# Tilted lens - ONLY the frame definition changes!\n",
    "lens_frame = root.make_child(\"lens\").translate(x=-25).rotate_euler(y=5, degrees=True)\n",
    "```\n",
    "\n",
    "The entire raytracing algorithm (intersection finding, refraction, propagation) stays **completely unchanged**. Scroll down to see this in action with spherical and hyperbolic lenses.\n",
    "\n",
    "---\n",
    "\n",
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "from collections import defaultdict\n",
    "from dataclasses import dataclass, field\n",
    "from typing import Callable\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from scipy.optimize import brentq\n",
    "\n",
    "from hazy import Frame, Point, Vector"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "## Define optical primitives\n",
    "\n",
    "We need classes for rays, ray history tracking, and axisymmetric surfaces (rotationally symmetric around the optical axis)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class RayState:\n",
    "    \"\"\"Snapshot of ray state at a point in space.\n",
    "\n",
    "    Args:\n",
    "        point: Position in global coordinates\n",
    "        direction: Direction vector in global coordinates\n",
    "    \"\"\"\n",
    "\n",
    "    point: Point\n",
    "    direction: Vector\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class RayHistory:\n",
    "    \"\"\"Tracks ray trajectory through optical system.\"\"\"\n",
    "\n",
    "    history: list[RayState] = field(default_factory=list)\n",
    "\n",
    "    def append(self, ray: Ray) -> RayHistory:\n",
    "        \"\"\"Add ray state to history (converts to global coordinates).\n",
    "\n",
    "        Args:\n",
    "            ray: Ray to record\n",
    "\n",
    "        Returns:\n",
    "            Self for method chaining\n",
    "        \"\"\"\n",
    "        self.history.append(\n",
    "            RayState(point=ray.origin.to_global(), direction=ray.direction.to_global())\n",
    "        )\n",
    "        return self\n",
    "\n",
    "    def plot(self, ax, **plot_args):\n",
    "        \"\"\"Plot ray trajectory.\n",
    "\n",
    "        Args:\n",
    "            ax: Matplotlib 3D axis\n",
    "            **plot_args: Additional plotting arguments (color, linewidth, etc.)\n",
    "        \"\"\"\n",
    "        points = np.array([p.point for p in self.history])\n",
    "        ax.plot(*points.T, **plot_args)\n",
    "\n",
    "\n",
    "@dataclass(frozen=True)\n",
    "class Ray:\n",
    "    \"\"\"Immutable ray with origin and direction.\n",
    "\n",
    "    Rays are geometric primitives similar to Points and Vectors.\n",
    "    All operations return new Ray instances.\n",
    "\n",
    "    Args:\n",
    "        origin: Ray starting point\n",
    "        direction: Ray direction (automatically normalized)\n",
    "    \"\"\"\n",
    "\n",
    "    origin: Point\n",
    "    direction: Vector\n",
    "\n",
    "    def __post_init__(self):\n",
    "        object.__setattr__(self, \"direction\", self.direction.copy().normalize())\n",
    "\n",
    "    def to_frame(self, frame: Frame) -> Ray:\n",
    "        \"\"\"Transform ray to different coordinate frame.\n",
    "\n",
    "        Args:\n",
    "            frame: Target coordinate frame\n",
    "\n",
    "        Returns:\n",
    "            New ray in target frame\n",
    "        \"\"\"\n",
    "        return Ray(\n",
    "            origin=self.origin.to_frame(frame), direction=self.direction.to_frame(frame)\n",
    "        )\n",
    "\n",
    "    def to_global(self) -> Ray:\n",
    "        \"\"\"Transform ray to global coordinates.\n",
    "\n",
    "        Returns:\n",
    "            New ray in global frame\n",
    "        \"\"\"\n",
    "        return Ray(origin=self.origin.to_global(), direction=self.direction.to_global())\n",
    "\n",
    "    def propagate(self, t: float) -> Ray:\n",
    "        \"\"\"Propagate ray along direction.\n",
    "\n",
    "        Args:\n",
    "            t: Distance to propagate\n",
    "\n",
    "        Returns:\n",
    "            New ray at propagated position\n",
    "        \"\"\"\n",
    "        return Ray(origin=self.origin + self.direction * t, direction=self.direction)\n",
    "\n",
    "    def refract(self, normal: Vector, n1: float, n2: float) -> Ray | None:\n",
    "        \"\"\"Refract ray at surface using Snell's law (vector form).\n",
    "\n",
    "        Args:\n",
    "            normal: Surface normal (pointing toward incident medium)\n",
    "            n1: Refractive index of incident medium\n",
    "            n2: Refractive index of transmitted medium\n",
    "\n",
    "        Returns:\n",
    "            Refracted ray, or None if total internal reflection\n",
    "\n",
    "        Examples:\n",
    "            >>> ray_refracted = ray.refract(normal, n1=1.0, n2=1.5)\n",
    "        \"\"\"\n",
    "        normal = normal.to_frame(self.direction.frame).copy().normalize()\n",
    "\n",
    "        ratio = n1 / n2\n",
    "        cos_i = -self.direction.dot(normal)\n",
    "\n",
    "        # Flip normal if ray hits from behind\n",
    "        if cos_i < 0:\n",
    "            normal = -normal\n",
    "            cos_i = -cos_i\n",
    "\n",
    "        sin2_t = ratio**2 * (1 - cos_i**2)\n",
    "\n",
    "        # Total internal reflection\n",
    "        if sin2_t > 1:\n",
    "            return None\n",
    "\n",
    "        cos_t = np.sqrt(1 - sin2_t)\n",
    "\n",
    "        # Snell's law (vector form)\n",
    "        d_refracted = ratio * self.direction + (ratio * cos_i - cos_t) * normal\n",
    "\n",
    "        return Ray(\n",
    "            origin=self.origin.copy(), direction=self.origin.frame.vector(d_refracted)\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class AxisymmetricSurface:\n",
    "    \"\"\"Rotationally symmetric surface defined by 1D profile x = f(r).\n",
    "\n",
    "    Coordinate system:\n",
    "    - x: optical axis\n",
    "    - r = sqrt(y^2 + z^2): radial distance from axis\n",
    "\n",
    "    The surface is created by revolving the 2D profile around the x-axis.\n",
    "\n",
    "    Args:\n",
    "        frame: Local coordinate frame of the surface\n",
    "        profile: Function mapping radial distance to axial position: x = f(r)\n",
    "        aperture_radius: Maximum radial extent of the surface\n",
    "        profile_derivative: Optional derivative dx/dr for faster normal computation\n",
    "    \"\"\"\n",
    "\n",
    "    frame: Frame\n",
    "    profile: Callable[[np.ndarray], np.ndarray]\n",
    "    aperture_radius: float\n",
    "    profile_derivative: Callable[[np.ndarray], np.ndarray] | None = None\n",
    "\n",
    "    def evaluate(self, r: np.ndarray) -> np.ndarray:\n",
    "        \"\"\"Evaluate profile at radial positions.\n",
    "\n",
    "        Args:\n",
    "            r: Radial distance(s) from optical axis\n",
    "\n",
    "        Returns:\n",
    "            Axial position(s) x = f(r)\n",
    "        \"\"\"\n",
    "        return self.profile(np.asarray(r))\n",
    "\n",
    "    def intersect(self, ray: Ray) -> tuple[float, Point, Vector]:\n",
    "        \"\"\"Find ray-surface intersection.\n",
    "\n",
    "        Strategy: Transform the ray to the optics local coordinate system (frame).\n",
    "        Reduce the problem to 2d by exploiting rotational symmetry of the optical element in its frame.\n",
    "\n",
    "        Args:\n",
    "            ray: Incident ray\n",
    "\n",
    "        Returns:\n",
    "            Tuple of (propagation distance, intersection point, surface normal)\n",
    "            or None if no intersection exists\n",
    "        \"\"\"\n",
    "        local_ray = ray.to_frame(self.frame)\n",
    "\n",
    "        def objective_function(t: float) -> float:\n",
    "            point = local_ray.origin + t * local_ray.direction\n",
    "            r = np.sqrt(point.y**2 + point.z**2)\n",
    "\n",
    "            if r > self.aperture_radius:\n",
    "                return 1e10 if t > 0 else -1e10\n",
    "\n",
    "            x_surface = self.profile(np.array([r]))[0]\n",
    "\n",
    "            if np.isnan(x_surface):\n",
    "                return 1e10 if t > 0 else -1e10\n",
    "\n",
    "            return point.x - x_surface\n",
    "\n",
    "        # Search for sign change between t=0 and t=100\n",
    "        try:\n",
    "            t_hit, res = brentq(\n",
    "                objective_function, 1e-3, 100.0, xtol=1e-10, full_output=True\n",
    "            )\n",
    "        except ValueError as err:\n",
    "            raise ValueError(\n",
    "                \"No valid intersection between ray and surface found\"\n",
    "            ) from err\n",
    "\n",
    "        propagated_ray = local_ray.propagate(t_hit)\n",
    "        hit_local = propagated_ray.origin.to_frame(self.frame)\n",
    "\n",
    "        # Compute surface normal in local frame\n",
    "        normal_local = self.normal_at(hit_local)\n",
    "\n",
    "        return t_hit, hit_local, normal_local\n",
    "\n",
    "    def normal_at(self, point: Point) -> Vector:\n",
    "        \"\"\"Compute outward surface normal at point.\n",
    "\n",
    "        Args:\n",
    "            point: Surface point (in surface frame)\n",
    "\n",
    "        Returns:\n",
    "            Outward-pointing unit normal vector\n",
    "        \"\"\"\n",
    "        p = point.to_frame(self.frame)\n",
    "        r = np.sqrt(p.y**2 + p.z**2)\n",
    "\n",
    "        # On-axis: normal points along -x\n",
    "        if r < 1e-12:\n",
    "            return self.frame.vector(-1.0, 0.0, 0.0)\n",
    "\n",
    "        # Compute derivative dx/dr\n",
    "        if self.profile_derivative is not None:\n",
    "            dxdr = self.profile_derivative(np.array([r]))[0]\n",
    "        else:\n",
    "            # Numerical derivative\n",
    "            eps = 1e-8\n",
    "            dxdr = (\n",
    "                self.profile(np.array([r + eps]))[0]\n",
    "                - self.profile(np.array([r - eps]))[0]\n",
    "            ) / (2 * eps)\n",
    "\n",
    "        # Normal in cylindrical coordinates: (-1, dx/dr * y/r, dx/dr * z/r)\n",
    "        return self.frame.vector(-1.0, dxdr * p.y / r, dxdr * p.z / r).normalize()\n",
    "\n",
    "    def plot_profile(self, ax=None, n_points: int = 500, **kwargs):\n",
    "        \"\"\"Plot 2D profile in meridional plane.\n",
    "\n",
    "        Args:\n",
    "            ax: Matplotlib axis (creates new if None)\n",
    "            n_points: Number of points to plot\n",
    "            **kwargs: Additional plot arguments\n",
    "\n",
    "        Returns:\n",
    "            Matplotlib axis\n",
    "        \"\"\"\n",
    "        if ax is None:\n",
    "            fig, ax = plt.subplots(figsize=(8, 6))\n",
    "\n",
    "        y = np.linspace(-self.aperture_radius, self.aperture_radius, n_points)\n",
    "        x = self.profile(np.abs(y))\n",
    "\n",
    "        defaults: dict = {\"linewidth\": 2}\n",
    "        defaults.update(kwargs)\n",
    "        ax.plot(x, y, **defaults)\n",
    "\n",
    "        ax.set_xlabel(\"x (optical axis)\")\n",
    "        ax.set_ylabel(\"y (radial)\")\n",
    "        ax.set_aspect(\"equal\")\n",
    "        ax.grid(True, alpha=0.3)\n",
    "\n",
    "        return ax\n",
    "\n",
    "    def plot_surface(\n",
    "        self, ax=None, n_radial: int = 50, n_azimuthal: int = 60, **surface_kwargs\n",
    "    ):\n",
    "        \"\"\"Plot revolved 3D surface.\n",
    "\n",
    "        Args:\n",
    "            ax: Matplotlib 3D axis (creates new if None)\n",
    "            n_radial: Number of radial sample points\n",
    "            n_azimuthal: Number of azimuthal sample points\n",
    "            **surface_kwargs: Additional surface plot arguments\n",
    "\n",
    "        Returns:\n",
    "            Matplotlib 3D axis\n",
    "        \"\"\"\n",
    "        if ax is None:\n",
    "            fig = plt.figure(figsize=(10, 8))\n",
    "            ax = fig.add_subplot(111, projection=\"3d\")\n",
    "\n",
    "        # Generate surface mesh\n",
    "        r = np.linspace(0, self.aperture_radius, n_radial)\n",
    "        phi = np.linspace(0, 2 * np.pi, n_azimuthal)\n",
    "        r_grid, phi_grid = np.meshgrid(r, phi)\n",
    "\n",
    "        # Convert to Cartesian in surface frame\n",
    "        y_grid = r_grid * np.cos(phi_grid)\n",
    "        z_grid = r_grid * np.sin(phi_grid)\n",
    "        x_grid = self.profile(r_grid)\n",
    "\n",
    "        shape = x_grid.shape\n",
    "\n",
    "        # Transform to global coordinates\n",
    "        x, y, z = self.frame.batch_transform_points_global(\n",
    "            np.stack([x_grid.ravel(), y_grid.ravel(), z_grid.ravel()], axis=-1)\n",
    "        ).T\n",
    "\n",
    "        defaults = {\"alpha\": 0.3, \"edgecolor\": \"none\"}\n",
    "        defaults.update(surface_kwargs)\n",
    "        ax.plot_surface(\n",
    "            x.reshape(shape), y.reshape(shape), z.reshape(shape), **defaults\n",
    "        )\n",
    "\n",
    "        ax.set_xlabel(\"x (optical axis)\")\n",
    "        ax.set_ylabel(\"y\")\n",
    "        ax.set_zlabel(\"z\")\n",
    "\n",
    "        return ax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "## Surface profile functions\n",
    "\n",
    "Optical surfaces are defined by 1D profiles `x = f(r)` that are revolved around the x-axis (optical axis)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def spherical_profile(\n",
    "    R: float,\n",
    ") -> tuple[Callable[[np.ndarray], np.ndarray], Callable[[np.ndarray], np.ndarray]]:\n",
    "    \"\"\"Create spherical surface profile and its derivative.\n",
    "\n",
    "    Args:\n",
    "        R: Radius of curvature\n",
    "           R > 0: convex (bulging toward +x)\n",
    "           R < 0: concave (bulging toward -x)\n",
    "\n",
    "    Returns:\n",
    "        Tuple of (profile function, derivative function)\n",
    "\n",
    "    Examples:\n",
    "        >>> profile, deriv = spherical_profile(R=50)\n",
    "        >>> x = profile(np.array([0, 10, 20]))  # x positions at r = 0, 10, 20\n",
    "    \"\"\"\n",
    "\n",
    "    def profile(r: np.ndarray) -> np.ndarray:\n",
    "        r = np.asarray(r)\n",
    "        arg = R**2 - r**2\n",
    "        valid = arg >= 0\n",
    "        x = np.full_like(r, np.nan, dtype=float)\n",
    "        x[valid] = R - np.sign(R) * np.sqrt(arg[valid])\n",
    "        return x\n",
    "\n",
    "    def derivative(r: np.ndarray) -> np.ndarray:\n",
    "        r = np.asarray(r)\n",
    "        arg = R**2 - r**2\n",
    "        valid = arg > 0\n",
    "        dxdr = np.full_like(r, np.nan, dtype=float)\n",
    "        dxdr[valid] = r[valid] / (np.sign(R) * np.sqrt(arg[valid]))\n",
    "        return dxdr\n",
    "\n",
    "    return profile, derivative\n",
    "\n",
    "\n",
    "def hyperbolic_profile(\n",
    "    R: float, conic_constant: float = -2.0\n",
    ") -> tuple[Callable[[np.ndarray], np.ndarray], Callable[[np.ndarray], np.ndarray]]:\n",
    "    \"\"\"Create hyperbolic surface profile and its derivative.\n",
    "\n",
    "    General conic surface: x = (r^2/R) / (1 + sqrt(1 - (1+K)*(r/R)^2))\n",
    "    where K is the conic constant.\n",
    "\n",
    "    For K < -1: hyperbola (stronger curvature at edges than sphere)\n",
    "    For K = -1: parabola\n",
    "    For -1 < K < 0: ellipse\n",
    "    For K = 0: sphere\n",
    "\n",
    "    Args:\n",
    "        R: Radius of curvature at apex (r=0)\n",
    "        conic_constant: K value (default -2.0 for hyperbola)\n",
    "\n",
    "    Returns:\n",
    "        Tuple of (profile function, derivative function)\n",
    "\n",
    "    Examples:\n",
    "        >>> profile, deriv = hyperbolic_profile(R=25, conic_constant=-2.0)\n",
    "        >>> x = profile(np.array([0, 10, 20]))\n",
    "    \"\"\"\n",
    "\n",
    "    def profile(r: np.ndarray) -> np.ndarray:\n",
    "        r = np.asarray(r)\n",
    "        r2 = r**2\n",
    "        arg = 1 - (1 + conic_constant) * r2 / (R**2)\n",
    "\n",
    "        valid = arg > 0\n",
    "        x = np.full_like(r, np.nan, dtype=float)\n",
    "        x[valid] = r2[valid] / (R * (1 + np.sqrt(arg[valid])))\n",
    "        return x\n",
    "\n",
    "    def derivative(r: np.ndarray) -> np.ndarray:\n",
    "        r = np.asarray(r)\n",
    "        r2 = r**2\n",
    "        arg = 1 - (1 + conic_constant) * r2 / (R**2)\n",
    "\n",
    "        valid = arg > 0\n",
    "        dxdr = np.full_like(r, np.nan, dtype=float)\n",
    "\n",
    "        sqrt_arg = np.sqrt(arg[valid])\n",
    "        denominator = R * (1 + sqrt_arg) ** 2\n",
    "        numerator = (\n",
    "            2\n",
    "            * r[valid]\n",
    "            * (1 + sqrt_arg + (1 + conic_constant) * r2[valid] / (2 * R**2 * sqrt_arg))\n",
    "        )\n",
    "\n",
    "        dxdr[valid] = numerator / denominator\n",
    "        return dxdr\n",
    "\n",
    "    return profile, derivative\n",
    "\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(12, 5))\n",
    "fig.suptitle(\"Lens Profiles\", fontsize=14, fontweight=\"bold\")\n",
    "root = Frame.make_root(\"dummy-frame\")\n",
    "\n",
    "axs[0].set_title(\"Spherical\", fontsize=14, fontweight=\"bold\")\n",
    "profile, profile_derivative = spherical_profile(R=50)\n",
    "AxisymmetricSurface(\n",
    "    frame=root,\n",
    "    profile=profile,\n",
    "    profile_derivative=profile_derivative,\n",
    "    aperture_radius=40,\n",
    ").plot_profile(axs[0])\n",
    "\n",
    "axs[1].set_title(\"Hyperbolic\", fontsize=14, fontweight=\"bold\")\n",
    "profile, profile_derivative = hyperbolic_profile(R=50)\n",
    "AxisymmetricSurface(\n",
    "    frame=root,\n",
    "    profile=profile,\n",
    "    profile_derivative=profile_derivative,\n",
    "    aperture_radius=40,\n",
    ").plot_profile(axs[1])\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "## Raytracing function\n",
    "\n",
    "This function traces parallel rays through a surface, showing the 3D trajectory and final footprint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def trace_rays_through_surface(\n",
    "    surface: AxisymmetricSurface,\n",
    "    n_rays: int = 10,\n",
    "    inner_radius: float = 10.0,\n",
    "    outer_radius: float = 20.0,\n",
    "    n1: float = 1.0,\n",
    "    n2: float = 1.5,\n",
    "    final_x: float = 38.0,\n",
    ") -> tuple[dict[int, RayHistory], tuple]:\n",
    "    \"\"\"Trace parallel rays through an optical surface.\n",
    "\n",
    "    Args:\n",
    "        surface: Optical surface to trace through\n",
    "        n_rays: Number of rays to trace\n",
    "        inner_radius: Radial distance of inner ray bundle from optical axis\n",
    "        outer_radius: Radial distance of outer ray bundle from optical axis\n",
    "        n1: Refractive index before surface\n",
    "        n2: Refractive index after surface\n",
    "        final_x: Final x position to propagate rays to\n",
    "\n",
    "    Returns:\n",
    "        Tuple of (ray histories dict, (fig, ax_3d, ax_2d))\n",
    "    \"\"\"\n",
    "    fig = plt.figure(figsize=(14, 6))\n",
    "    ax_3d = fig.add_subplot(121, projection=\"3d\")\n",
    "    ax_2d = fig.add_subplot(122)\n",
    "\n",
    "    # Plot surface\n",
    "    surface.plot_surface(ax_3d)\n",
    "    ax_3d.set_title(\"3D Ray Trace\", fontsize=14, fontweight=\"bold\")\n",
    "    ax_3d.set_xlim(-30, 30)\n",
    "    ax_3d.set_ylim(-30, 30)\n",
    "    ax_3d.set_zlim(-30, 30)\n",
    "\n",
    "    # Generate rays in circular pattern\n",
    "    phis = np.linspace(0, 2 * np.pi, n_rays, endpoint=False)\n",
    "\n",
    "    for radius, cmap in zip([inner_radius, outer_radius], [plt.cm.Reds, plt.cm.Blues]):\n",
    "        rays = defaultdict(RayHistory)\n",
    "        colors = cmap(np.linspace(0.2, 0.8, n_rays))\n",
    "        for i, phi in enumerate(phis):\n",
    "            z = radius * np.cos(phi)\n",
    "            y = radius * np.sin(phi)\n",
    "\n",
    "            # Create initial ray\n",
    "            ray = Ray(surface.frame.root.point(-30, y, z), surface.frame.root.x_axis)\n",
    "            rays[i].append(ray)\n",
    "\n",
    "            # Intersect with surface\n",
    "            t, point, normal = surface.intersect(ray)\n",
    "            ray = ray.propagate(t)\n",
    "            rays[i].append(ray)\n",
    "\n",
    "            # Refract at surface\n",
    "            ray = ray.refract(normal, n1, n2)\n",
    "            rays[i].append(ray)\n",
    "\n",
    "            # Propagate after refraction\n",
    "            final_x_point = (\n",
    "                surface.frame.root.origin + surface.frame.root.x_axis * final_x\n",
    "            )\n",
    "            propagation_distance = surface.frame.root.x_axis.dot(\n",
    "                final_x_point - ray.origin\n",
    "            ) / surface.frame.root.x_axis.dot(ray.direction)\n",
    "            ray = ray.propagate(propagation_distance)\n",
    "\n",
    "            rays[i].append(ray)\n",
    "\n",
    "        for (i, ray_history), color in zip(rays.items(), colors):\n",
    "            ray_history.plot(ax_3d, c=color, linewidth=2, alpha=0.8)\n",
    "\n",
    "            # Plot start and end points in 2D (y-z plane)\n",
    "            ax_2d.scatter(\n",
    "                *ray_history.history[0].point[[1, 2]],\n",
    "                color=color,\n",
    "                s=80,\n",
    "                alpha=0.5,\n",
    "                edgecolors=\"none\",\n",
    "            )\n",
    "            ax_2d.scatter(\n",
    "                *ray_history.history[-1].point[[1, 2]],\n",
    "                color=color,\n",
    "                s=80,\n",
    "                edgecolors=\"black\",\n",
    "                linewidth=1.5,\n",
    "            )\n",
    "\n",
    "    ax_2d.set_aspect(\"equal\")\n",
    "    ax_2d.set_xlabel(\"y\", fontsize=12)\n",
    "    ax_2d.set_ylabel(\"z\", fontsize=12)\n",
    "    ax_2d.set_title(\n",
    "        \"Ray Footprint (start: faded, end: solid)\", fontsize=14, fontweight=\"bold\"\n",
    "    )\n",
    "    ax_2d.grid(True, alpha=0.3)\n",
    "\n",
    "    return rays, (fig, ax_3d, ax_2d)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "## Example: Raytracing a lens\n",
    "\n",
    "This example demonstrates how coordinate frame transformations simplify optical raytracing.\n",
    "\n",
    "We'll trace parallel rays through a spherical lens surface and observe spherical aberration - rays at different radial distances focus at different points.\n",
    "\n",
    "The key insight: by placing the lens in its own coordinate frame, we can **reposition and reorient the entire optical system by changing a single line of code**, while all geometric calculations remain unchanged."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "### Create root frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "root = Frame.make_root(\"root\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "### Aligned spherical lens\n",
    "\n",
    "Start with a perfectly aligned spherical lens. Notice the **spherical aberration**: rays at different radial distances don't focus to the same point - the footprint spreads out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "lens_frame = root.make_child(\"lens\").translate(x=-25)\n",
    "\n",
    "profile, derivative = spherical_profile(R=25)\n",
    "\n",
    "surface = AxisymmetricSurface(\n",
    "    frame=lens_frame,\n",
    "    profile=profile,\n",
    "    profile_derivative=derivative,\n",
    "    aperture_radius=40.0,\n",
    ")\n",
    "\n",
    "rays, (fig, ax_3d, ax_2d) = trace_rays_through_surface(surface)\n",
    "\n",
    "fig.suptitle(\n",
    "    \"Aligned Spherical Lens: Spherical Aberration\", fontsize=14, fontweight=\"bold\"\n",
    ")\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "### Misalign the lens - only one line changes!\n",
    "\n",
    "Now rotate the lens 5 degrees around the y-axis. **Look closely**: Only the `lens_frame` definition changes - the entire raytracing algorithm (intersection, refraction, propagation) stays identical.\n",
    "\n",
    "This asymmetry creates **coma aberration** - the ray footprint becomes comet-shaped. This is a fundamental optical aberration that occurs when lenses are tilted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Misaligned case: lens rotated 5 degrees\n",
    "# ↓↓↓ ONLY THIS LINE CHANGES - everything else is identical! ↓↓↓\n",
    "lens_frame = root.make_child(\"lens\").translate(x=-20).rotate_euler(y=10, degrees=True)\n",
    "\n",
    "profile, derivative = spherical_profile(R=25)\n",
    "\n",
    "surface = AxisymmetricSurface(\n",
    "    frame=lens_frame,\n",
    "    profile=profile,\n",
    "    profile_derivative=derivative,\n",
    "    aperture_radius=40.0,\n",
    ")\n",
    "\n",
    "rays, (fig, ax_3d, ax_2d) = trace_rays_through_surface(surface)\n",
    "\n",
    "fig.suptitle(\"Misaligned Lens: Coma Aberration\", fontsize=14, fontweight=\"bold\")\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "### Swap the surface profile - only the profile changes!\n",
    "\n",
    "Replace the spherical lens with a hyperbolic one. Again, **only the profile function changes** - frame setup and raytracing code stay identical.\n",
    "\n",
    "Hyperbolic surfaces (conic constant K < -1) have stronger edge curvature than spheres, which affects how rays are focused. Compare the footprint pattern to the spherical case above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperbolic lens: different aberration characteristics\n",
    "# ↓↓↓ ONLY THE PROFILE CHANGES - frame, raytracing, everything else identical! ↓↓↓\n",
    "lens_frame = root.make_child(\"hyperbolic_lens\").translate(x=-25)\n",
    "\n",
    "profile, derivative = hyperbolic_profile(R=25, conic_constant=-2)\n",
    "\n",
    "surface = AxisymmetricSurface(\n",
    "    frame=lens_frame,\n",
    "    profile=profile,\n",
    "    profile_derivative=derivative,\n",
    "    aperture_radius=40.0,\n",
    ")\n",
    "\n",
    "rays, (fig, ax_3d, ax_2d) = trace_rays_through_surface(surface)\n",
    "\n",
    "fig.suptitle(\"Hyperbolic Lens: Enhanced Edge Curvature\", fontsize=14, fontweight=\"bold\")\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "## Key Takeaways\n",
    "\n",
    "**Frame transformations for optics**: Optical elements (lenses, mirrors) live in their own frames. Repositioning or rotating them requires changing only the frame definition - all geometric calculations (intersection finding, normal computation, refraction) work in local coordinates and remain unchanged.\n",
    "\n",
    "**Interchangeable profiles**: Each surface profile (`spherical`, `hyperbolic`, etc.) is just a function `x = f(r)`. Swapping profiles means changing the function - the `AxisymmetricSurface` class and raytracing algorithm stay identical. This modularity allows easy exploration of different optical designs.\n",
    "\n",
    "**Immutable rays**: Like `Point` and `Vector`, `Ray` is immutable - all operations (`propagate`, `refract`, `to_frame`) return new instances. This matches hazy's philosophy: geometric primitives are values, not mutable state.\n",
    "\n",
    "**Without hazy**: Manually compose transformation matrices for each lens position, track multiplication order for nested frames, explicitly transform every point and vector, rewrite intersection code when repositioning optics.\n",
    "\n",
    "**With hazy**: Define surfaces in natural local coordinates (`origin=(0,0,0)`, x-axis = optical axis), modify frames to position/orient elements, call `.to_frame()` for automatic transformation through arbitrarily deep hierarchies."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hazy-frames",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}