Hands-on lab — Day 112: Visualizing Optimization

"""Script 1 -- evaluate a surface over a grid, then look at it without an
image viewer.

Every picture later in this lab starts here: numpy.meshgrid plus a function
call, and a way to check that Z's rows and columns line up with X and Y the
way you think they do.
"""

import numpy as np

from dataset import WELL_F
from gridviz import ascii_contour, evaluate_grid

XLIM = (-4.0, 4.0)
YLIM = (-4.0, 4.0)
N = 101

X, Y, Z = evaluate_grid(WELL_F, XLIM, YLIM, N)
print(f"grid shape: X {X.shape}, Y {Y.shape}, Z {Z.shape}")
assert X.shape == Y.shape == Z.shape == (N, N)

iy, ix = np.unravel_index(np.argmin(Z), Z.shape)
print(f"minimum at grid cell (row={iy}, col={ix}) -> x={X[iy, ix]:.4f}, y={Y[iy, ix]:.4f}, f={Z[iy, ix]:.6g}")
assert abs(X[iy, ix]) < 1e-9 and abs(Y[iy, ix]) < 1e-9

_, _, small_Z = evaluate_grid(lambda x, y: x**2 + y**2, (-2, 2), (-2, 2), 5)
print()
print("ASCII contour of x^2 + y^2 on a 5x5 grid spanning [-2, 2] x [-2, 2]:")
print(ascii_contour(small_Z))
rows = ascii_contour(small_Z).split("\n")
assert rows[2][2] == " ", "the centre cell (the minimum) should be the lightest character"
assert rows[0][0] == "#", "a corner (the maximum) should be the densest character"

print()
print("01_grid_and_ascii.py: every assertion held.")

"""Script 2 -- a Pillow heatmap, and a descent path drawn on top of it.

Writes into a temporary directory and removes it before exiting: this lab
leaves no image files behind, on purpose, so that "did the lab clean up"
is a check you can actually run rather than a claim you have to trust.
"""

import shutil
import tempfile
from pathlib import Path

from dataset import ILL_F, ILL_GRAD, LEARNING_RATE, START, STEPS
from descent import gradient_descent
from gridviz import evaluate_grid, world_to_pixel
from imaging import draw_path_on_heatmap, heatmap_png
from PIL import Image

XLIM = (-4.0, 4.0)
YLIM = (-4.0, 4.0)

_, _, Z = evaluate_grid(ILL_F, XLIM, YLIM, 101)

tmp = Path(tempfile.mkdtemp(prefix="d112-"))
try:
    heat_path = tmp / "heatmap.png"
    heatmap_png(Z, XLIM, YLIM, str(heat_path))
    img = Image.open(heat_path)
    print(f"heatmap.png: size={img.size}, mode={img.mode}")
    assert img.size == (101, 101)

    px, py = world_to_pixel(0.0, 0.0, XLIM, YLIM, 101, 101)
    center_color = img.convert("RGB").getpixel((round(px), round(py)))
    print(f"pixel at the minimum (0, 0) -> ({round(px)}, {round(py)}): color {center_color}")
    assert center_color == (13, 27, 84), "the minimum should be the ramp's lowest-value color"

    path = gradient_descent(ILL_GRAD, START, LEARNING_RATE, STEPS)
    path_out = tmp / "descent.png"
    draw_path_on_heatmap(Z, XLIM, YLIM, path, str(path_out))
    print(f"descent.png: {len(path)} points drawn over the ill-conditioned bowl")

    start_px = world_to_pixel(*path[0], XLIM, YLIM, 101, 101)
    end_px = world_to_pixel(*path[-1], XLIM, YLIM, 101, 101)
    min_px = world_to_pixel(0.0, 0.0, XLIM, YLIM, 101, 101)
    print(f"first marker pixel: {tuple(round(v, 2) for v in start_px)}")
    print(f"last marker pixel:  {tuple(round(v, 2) for v in end_px)}")
    print(f"minimum pixel:      {tuple(round(v, 2) for v in min_px)}")

    distance = ((end_px[0] - min_px[0]) ** 2 + (end_px[1] - min_px[1]) ** 2) ** 0.5
    print(f"distance from the last marker to the minimum: {distance:.3f} pixels")
    assert distance < 5.0

    print()
    print("02_heatmap_and_path.py: every assertion held.")
finally:
    shutil.rmtree(tmp, ignore_errors=True)

"""Script 3 -- loss against iteration, linear and log, and why the log axis
is not decoration.

For a run where the update is an exact geometric recursion (the
well-conditioned bowl below), loss_k = c * rho^k, so log10(loss_k) is a
straight line in k with slope log10(rho). This script draws both curves and
then proves the log one is a line by fitting it and reading off the residual
-- from the picture's own pixel coordinates, not from the formula.
"""

import shutil
import tempfile
from pathlib import Path

import numpy as np

from dataset import LEARNING_RATE, START, STEPS, WELL_F, WELL_GRAD
from descent import gradient_descent, losses_along
from imaging import loss_curve_png, loss_curve_points

path = gradient_descent(WELL_GRAD, START, LEARNING_RATE, STEPS)
losses = losses_along(WELL_F, path)

print(f"loss at step 0:  {losses[0]:.6g}")
print(f"loss at step {STEPS}: {losses[-1]:.6g}")
print(f"ratio of consecutive losses (should be ~constant for a straight log line):")
ratios = losses[1:11] / losses[0:10]
for k, r in enumerate(ratios):
    print(f"  loss[{k+1}] / loss[{k}] = {r:.6f}")
assert np.max(ratios) - np.min(ratios) < 1e-9, "a geometric decay has a constant ratio"

tmp = Path(tempfile.mkdtemp(prefix="d112-"))
try:
    lin_path = tmp / "loss_linear.png"
    log_path = tmp / "loss_log.png"
    loss_curve_png(losses, str(lin_path), log=False)
    loss_curve_png(losses, str(log_path), log=True)
    print(f"\nloss_linear.png and loss_log.png written and read back successfully")

    points = loss_curve_points(losses, width=500, height=350, margin=50, log=True)
    xs = np.array([p[0] for p in points])
    ys = np.array([p[1] for p in points])
    A = np.vstack([xs, np.ones_like(xs)]).T
    slope, intercept = np.linalg.lstsq(A, ys, rcond=None)[0]
    residual = np.max(np.abs(ys - (slope * xs + intercept)))
    print(f"log-axis points: best-fit line residual (max, in pixels) = {residual:.3e}")
    assert residual < 1e-6

    print()
    print("03_loss_curves.py: every assertion held.")
finally:
    shutil.rmtree(tmp, ignore_errors=True)

"""Script 4 -- an animated GIF of a descent, one frame per step.

Pillow needs nothing beyond Image.save(..., save_all=True) to write an
animation: build one Image per frame and hand Pillow the list.
"""

import shutil
import tempfile
from pathlib import Path

from dataset import ILL_F, ILL_GRAD, LEARNING_RATE, START
from descent import gradient_descent
from gridviz import evaluate_grid
from imaging import animated_descent_gif
from PIL import Image

XLIM = (-4.0, 4.0)
YLIM = (-4.0, 4.0)
N_FRAMES = 25

_, _, Z = evaluate_grid(ILL_F, XLIM, YLIM, 101)
path = gradient_descent(ILL_GRAD, START, LEARNING_RATE, N_FRAMES)
print(f"path has {len(path)} points (start + {N_FRAMES} steps)")

tmp = Path(tempfile.mkdtemp(prefix="d112-"))
try:
    out = tmp / "descent.gif"
    animated_descent_gif(Z, XLIM, YLIM, path, str(out))
    size_bytes = out.stat().st_size
    print(f"descent.gif written: {size_bytes} bytes")

    reopened = Image.open(out)
    print(f"format={reopened.format}, size={reopened.size}, n_frames={reopened.n_frames}")
    assert reopened.format == "GIF"
    assert reopened.n_frames == len(path)

    print()
    print("04_animated_gif.py: every assertion held.")
finally:
    shutil.rmtree(tmp, ignore_errors=True)

"""Script 5 -- final loss against learning rate: the characteristic shape.

Slow on the left, a broad basin of good learning rates, then a cliff where
the run diverges. A good learning rate is a RANGE, and this script measures
that range rather than asserting it.
"""

import numpy as np

from dataset import SWEEP_STEPS, SWEEP_X0, sweep_f, sweep_grad
from descent import learning_rate_sweep

etas = np.round(np.arange(0.05, 2.55, 0.10), 2)
sweep = learning_rate_sweep(sweep_grad, sweep_f, SWEEP_X0, etas, SWEEP_STEPS)

print(f"{'eta':>6}  {'final loss':>14}")
for eta, loss in sweep:
    shown = "inf (diverged)" if not np.isfinite(loss) else f"{loss:.6e}"
    print(f"{eta:>6.2f}  {shown:>14}")

finite = [(e, loss) for e, loss in sweep if np.isfinite(loss)]
divergent = [(e, loss) for e, loss in sweep if not np.isfinite(loss)]

good = [e for e, loss in finite if loss < 1e-6]
best_eta = min(finite, key=lambda pair: pair[1])[0]
# Above eta = 1, |1 - 2 eta| > 1 and the run blows up exponentially -- these
# etas do NOT converge, they simply have not yet overflowed float64 in the
# number of steps this sweep runs. "Finite but huge" is not "good".
still_overflowing_soon = max(e for e, loss in finite if loss >= 1e-6)
first_divergent_eta = min(e for e, _ in divergent)

print()
print(f"{len(good)} learning rates reach essentially zero loss: {good}")
print(f"argmin at eta = {best_eta}")
print(
    f"largest eta whose blow-up has not yet overflowed float64 in {SWEEP_STEPS} steps: "
    f"{still_overflowing_soon} (already diverging, just not yet inf)"
)
print(f"smallest eta whose blow-up has overflowed to inf (caught, not an exception): {first_divergent_eta}")

assert len(good) >= 3, "the good range is a basin, not a single point"
assert etas[0] < best_eta < etas[-1], "the argmin must be interior to the sweep"
assert len(divergent) >= 1
assert all(loss == float("inf") for _, loss in divergent)
assert all(e > 1.0 for e, _ in divergent), "the theoretical threshold for f(x) = x^2 is eta = 1"

print()
print("05_learning_rate_sweep.py: every assertion held.")

"""Script 6 -- the day's opening failure, made into a measurement.

Two runs, same starting point, same learning rate, same number of steps.
Their final losses agree to within a few percent. Their paths do not agree
at all: one is short and nearly straight, the other is over ten times
longer because it zig-zagged across a narrow valley. The final-loss number
alone cannot tell these runs apart -- only the path can.
"""

from dataset import (
    ILL_A,
    ILL_B,
    ILL_F,
    ILL_GRAD,
    LEARNING_RATE,
    LOSS_MATCH_TOL,
    PATH_LENGTH_RATIO_MIN,
    START,
    STEPS,
    WELL_A,
    WELL_B,
    WELL_F,
    WELL_GRAD,
)
from descent import gradient_descent, losses_along, path_length

well_path = gradient_descent(WELL_GRAD, START, LEARNING_RATE, STEPS)
ill_path = gradient_descent(ILL_GRAD, START, LEARNING_RATE, STEPS)

well_losses = losses_along(WELL_F, well_path)
ill_losses = losses_along(ILL_F, ill_path)

well_len = path_length(well_path)
ill_len = path_length(ill_path)

print(f"well-conditioned bowl: f(x, y) = {WELL_A:g} x^2 + {WELL_B:g} y^2")
print(f"ill-conditioned bowl:  f(x, y) = {ILL_A:g} x^2 + {ILL_B:g} y^2")
print(f"both start at {tuple(START)}, learning rate {LEARNING_RATE}, {STEPS} steps")
print()
print(f"well-conditioned final loss: {well_losses[-1]:.6e}")
print(f"ill-conditioned final loss:  {ill_losses[-1]:.6e}")
relative_gap = abs(well_losses[-1] - ill_losses[-1]) / max(well_losses[-1], ill_losses[-1])
print(f"relative gap between the two final losses: {relative_gap:.4f}")
print()
print(f"well-conditioned path length: {well_len:.4f}")
print(f"ill-conditioned path length:  {ill_len:.4f}")
print(f"ratio: {ill_len / well_len:.2f}x longer")

assert relative_gap < LOSS_MATCH_TOL, "the two final losses should be nearly indistinguishable"
assert ill_len / well_len > PATH_LENGTH_RATIO_MIN, "the paths should differ by a large factor"

print()
print(
    "A reader who only sees the two final-loss numbers above would call these runs "
    "equivalent. Only the path length -- or the picture -- shows that one of them "
    "spent most of its steps bouncing across the narrow axis instead of heading "
    "toward the minimum."
)
print()
print("06_two_runs_same_loss.py: every assertion held.")

"""Make this directory's own modules the ones its tests import.

Both `examples/` and `starter/` contain modules called `dataset`, `gridviz`,
`descent` and `imaging`, and pytest imports test files by putting their
directory on `sys.path`. Without this file, running `pytest` across both
directories at once would import whichever version was seen first and reuse
it for the other suite -- so the starter tests could silently pass against
the reference solution instead of skipping. That is a wrong answer with a
green tick on it, which is the worst kind.

So: put this directory first on the import path, and drop any already
imported module of those names that came from somewhere else.
"""

import sys
from pathlib import Path

HERE = str(Path(__file__).parent.resolve())

if HERE in sys.path:
    sys.path.remove(HERE)
sys.path.insert(0, HERE)

for name in ("dataset", "gridviz", "descent", "imaging"):
    module = sys.modules.get(name)
    origin = getattr(module, "__file__", "") or ""
    if module is not None and not origin.startswith(HERE):
        del sys.modules[name]

"""Shared numbers for the Day 112 lab -- the reference implementation.

Every script and every test in this lab imports its constants from here so
that the lesson's claims, the reference scripts, and the test suite are all
computed from the same numbers rather than three copies that could drift.

Two bowls, both centred on the origin with minimum value 0:

    WELL(x, y) = x^2 +  y^2          -- a perfectly round bowl
    ILL(x, y)  = x^2 + 25 y^2        -- a valley squeezed 25x narrower in y

The two runs compared throughout the lab and the lesson start at the same
point and take the same number of steps at the SAME learning rate. The only
difference is which bowl they descend. That is deliberate: it isolates
conditioning as the one variable that explains why one run's path is short
and geometric while the other's is long and zig-zagging, even though both
land within a few percent of the same final loss.
"""

import numpy as np

# -- the two bowls -----------------------------------------------------------

WELL_A, WELL_B = 1.0, 1.0
ILL_A, ILL_B = 1.0, 25.0

START = np.array([4.0, 4.0])
LEARNING_RATE = 0.038
STEPS = 60


def bowl(a: float, b: float):
    """Return (f, grad) for f(x, y) = a x^2 + b y^2.

    The minimum is always at the origin with value 0, regardless of a and b --
    only the SHAPE of the bowl changes, never where the bottom is.
    """

    def f(x, y):
        return a * x**2 + b * y**2

    def grad(x, y):
        return np.array([2.0 * a * x, 2.0 * b * y])

    return f, grad


WELL_F, WELL_GRAD = bowl(WELL_A, WELL_B)
ILL_F, ILL_GRAD = bowl(ILL_A, ILL_B)

# -- the one-dimensional bowl used for the learning-rate sweep ---------------
#
# f(x) = x^2, grad = 2x. The update is x <- x - eta * 2x = (1 - 2 eta) x, an
# exact geometric recursion with ratio rho = 1 - 2 eta. It converges only for
# 0 < eta < 1: the sweep therefore has a genuine cliff at eta = 1, not merely
# "large enough to eventually diverge".

SWEEP_X0 = 4.0
SWEEP_STEPS = 300


def sweep_f(x):
    return x**2


def sweep_grad(x):
    return 2.0 * x


# -- tolerances, each tied to the comparison it governs ----------------------

# Two analytic quantities computed by different routes: machine precision.
EXACT_TOL = 1e-9

# "Two runs land within a few percent of the same final loss." Chosen once
# from the pair of runs this file specifies (see WELL_RUN / ILL_RUN below,
# and examples/06_two_runs_same_loss.py, which prints the measured gap).
LOSS_MATCH_TOL = 0.05

# "The paths differ by a large factor." The measured ratio is over 13x; 5x is
# a conservative floor that would still be true on a different machine or
# numpy build, since the recursion above is closed-form and exact.
PATH_LENGTH_RATIO_MIN = 5.0

# How close the LAST drawn pixel marker must land to the pixel of the true
# minimum (0, 0), in pixels, for the world-to-pixel round trip to count as
# correct. One pixel of slack absorbs rounding in both directions.
PIXEL_TOL = 2.0

"""Exercises 7 and 8: running gradient descent and reading the run, not just
its last number.

Nothing here is imported from Day 111's lab -- the update rule is the same
three-line loop every gradient-descent lesson in this course uses, and it is
written out again here on purpose, because this lab's whole point is that the
FINAL LOSS a run prints is not enough to tell two runs apart.
"""

from __future__ import annotations

import numpy as np


def gradient_descent(grad_fn, x0, lr: float, steps: int) -> np.ndarray:
    """Run steps of x <- x - lr * grad_fn(x), returning every visited point.

    Returns an array of shape (steps + 1, len(x0)): the starting point,
    then one row per step, so path[0] is x0 and path[-1] is where the run
    stopped. Keeping the whole path (not just the endpoint) is what makes
    every other function in this lab possible -- you cannot draw a route on
    a map, or a loss curve, from a single final number.
    """
    x = np.array(x0, dtype=float)
    path = [x.copy()]
    for _ in range(steps):
        x = x - lr * np.asarray(grad_fn(*x))
        path.append(x.copy())
    return np.array(path)


def losses_along(f, path: np.ndarray) -> np.ndarray:
    """Evaluate f at every point on a path, returning one loss per row."""
    return np.array([f(*p) for p in path])


def path_length(path: np.ndarray) -> float:
    """Total Euclidean distance travelled along a path: sum of step sizes.

    A short, nearly straight path and a long, zig-zagging one can arrive at
    almost the same final loss -- path_length is the single number that
    tells them apart without needing a picture, and the picture is what
    explains WHY they differ.
    """
    steps = np.diff(path, axis=0)
    return float(np.sum(np.linalg.norm(steps, axis=1)))


def sweep_final_loss(grad_fn, f, x0: float, eta: float, steps: int) -> float:
    """Run gradient descent on a 1D function at learning rate eta, returning
    the final loss -- or float('inf') if the run diverged.

    Overflow during a diverging run is expected behaviour, not an error: a
    learning rate above the stability threshold makes the iterate grow
    without bound, and IEEE-754 represents "grew past the largest
    representable double" as inf rather than raising an exception. This
    function catches that deliberately with numpy's error-state context
    manager (over='ignore') and reports it as a value, so the reader sees a
    clean cliff in the sweep instead of a stack trace.
    """
    x = np.float64(x0)
    for _ in range(steps):
        with np.errstate(over="ignore", invalid="ignore"):
            x = x - eta * grad_fn(x)
        if not np.isfinite(x):
            return float("inf")
    with np.errstate(over="ignore", invalid="ignore"):
        value = f(x)
    return float(value) if np.isfinite(value) else float("inf")


def learning_rate_sweep(grad_fn, f, x0: float, etas, steps: int):
    """Run sweep_final_loss at every learning rate in etas.

    Returns a list of (eta, final_loss) pairs, in the order etas was given.
    """
    return [(float(eta), sweep_final_loss(grad_fn, f, x0, float(eta), steps)) for eta in etas]

"""Exercise 1, 2 and the world-to-pixel half of exercise 4: seeing a surface
before drawing anything on top of it.

Every picture in this lab starts from the same three-step recipe: evaluate a
function over a grid, decide how a value maps to something visible (a
character, a colour, a pixel), and read the result back to confirm the axes
did not get flipped along the way. This module is that recipe with nothing
plotted yet -- the terminal-only version of a contour plot.
"""

from __future__ import annotations

import numpy as np

# A five-character ramp, darkest (lowest value) to densest (highest value).
# Five bands is deliberately coarse: it is legible in a terminal and it is
# small enough that a specific character at a specific cell is a real
# assertion, not a fuzzy visual check.
ASCII_RAMP = " .:+#"


def evaluate_grid(f, xlim: tuple[float, float], ylim: tuple[float, float], n: int):
    """Evaluate f over an n x n grid spanning xlim x ylim.

    Returns (X, Y, Z), each an (n, n) array from numpy.meshgrid: X varies
    along columns, Y varies along rows, and Z = f(X, Y). This is the one
    building block every picture in this lab is made from -- a contour plot,
    a heatmap and a 3D surface are three different ways of drawing the same
    (X, Y, Z) triple.
    """
    xs = np.linspace(xlim[0], xlim[1], n)
    ys = np.linspace(ylim[0], ylim[1], n)
    X, Y = np.meshgrid(xs, ys)
    Z = f(X, Y)
    return X, Y, Z


def ascii_contour(Z: np.ndarray, chars: str = ASCII_RAMP) -> str:
    """Render a 2D array as text: one character per cell, chosen by level band.

    Values are rescaled to [0, len(chars)) linearly between Z's own min and
    max, then floored into a band index. This is the crudest possible contour
    renderer -- level BANDS shaded by character, not level LINES traced
    between them -- and that crudeness is the point: get it wrong (a
    transposed grid, a flipped row order) and a symmetric bowl stops looking
    symmetric immediately, in a terminal, with no image viewer required.
    """
    zmin, zmax = float(Z.min()), float(Z.max())
    span = zmax - zmin if zmax > zmin else 1.0
    n_bands = len(chars)
    idx = np.clip(((Z - zmin) / span * n_bands).astype(int), 0, n_bands - 1)
    return "\n".join("".join(chars[i] for i in row) for row in idx)


def world_to_pixel(
    x: float,
    y: float,
    xlim: tuple[float, float],
    ylim: tuple[float, float],
    width: int,
    height: int,
) -> tuple[float, float]:
    """Map one (x, y) point in data coordinates to a (column, row) pixel.

    Two axis flips have to happen correctly here and nowhere else in the
    lab -- every drawing function in imaging.py calls this one function
    instead of repeating the arithmetic, so a bug in the mapping shows up
    once, here, and is testable in isolation:

      * x grows to the RIGHT in both spaces, so columns increase with x --
        no flip.
      * y grows UPWARD in data space but DOWNWARD in pixel rows (row 0 is the
        top of the image), so pixel row is computed from (ylim[1] - y), not
        (y - ylim[0]).

    Getting the y flip backwards is the single most common bug in this lab:
    the heatmap looks fine on its own (it is vertically symmetric for a bowl
    centred at the origin) but the path drawn on top of it walks toward the
    WRONG edge, and exercise 4 is built to catch exactly that.
    """
    px = (x - xlim[0]) / (xlim[1] - xlim[0]) * (width - 1)
    py = (ylim[1] - y) / (ylim[1] - ylim[0]) * (height - 1)
    return px, py

"""Exercises 3, 4 and 6: turning arrays into pictures with nothing but NumPy
and Pillow's ImageDraw.

No matplotlib, no scipy. A heatmap is a NumPy array of colours turned into an
Image; a path is a polyline and a handful of circles drawn with ImageDraw; an
animation is a list of frames saved with save_all=True. That is the entire
technique stack this module needs.
"""

from __future__ import annotations

import numpy as np
from PIL import Image, ImageDraw

from gridviz import world_to_pixel

# -- the colour ramp -----------------------------------------------------
#
# Four control points, dark blue (low) through teal and gold to dark red
# (high). Perceptual accuracy is not the point here -- matplotlib's "viridis"
# and friends exist for that, and are covered in the lesson's tools section --
# the point is that a value maps to a colour through an explicit, inspectable
# rule with named stops, not a library call whose internals are invisible.

_STOPS_T = np.array([0.0, 0.35, 0.65, 1.0])
_STOPS_RGB = np.array(
    [
        [13, 27, 84],  # low value: dark blue
        [29, 78, 216],  # blue
        [250, 204, 21],  # gold
        [185, 28, 28],  # high value: dark red
    ]
)


def ramp_color(t: np.ndarray) -> np.ndarray:
    """Map an array of values in [0, 1] to an (..., 3) array of uint8 RGB.

    Piecewise-linear interpolation through the four stops above, one channel
    at a time via numpy.interp -- the same technique matplotlib's own
    colormaps use internally, just with four hand-picked stops instead of
    a few hundred.
    """
    t = np.clip(t, 0.0, 1.0)
    r = np.interp(t, _STOPS_T, _STOPS_RGB[:, 0])
    g = np.interp(t, _STOPS_T, _STOPS_RGB[:, 1])
    b = np.interp(t, _STOPS_T, _STOPS_RGB[:, 2])
    return np.stack([r, g, b], axis=-1).astype(np.uint8)


def heatmap_array(Z: np.ndarray) -> np.ndarray:
    """Turn a 2D value array into an (n, n, 3) uint8 RGB array, ready for
    Image.fromarray.

    Z's row 0 corresponds to ylim[0] (evaluate_grid's convention, since
    numpy.meshgrid puts the smallest y in the first row) -- but the TOP of an
    image is the LARGEST y, so the array is flipped vertically before it
    becomes an image. This is the one deliberate axis flip in this function;
    world_to_pixel performs the matching flip for anything drawn on top.
    """
    zmin, zmax = float(Z.min()), float(Z.max())
    span = zmax - zmin if zmax > zmin else 1.0
    t = (Z - zmin) / span
    colors = ramp_color(t)
    return np.flipud(colors)


def heatmap_png(
    Z: np.ndarray,
    xlim: tuple[float, float],
    ylim: tuple[float, float],
    path: str,
    width: int | None = None,
    height: int | None = None,
) -> Image.Image:
    """Save Z as a heatmap PNG and return the Image.

    width/height default to Z's own resolution (one pixel per grid cell),
    which keeps world_to_pixel exact: pixel (0, 0) is (xlim[0], ylim[1]) and
    pixel (n-1, n-1) is (xlim[1], ylim[0]).
    """
    arr = heatmap_array(Z)
    img = Image.fromarray(arr, mode="RGB")
    if width is not None and height is not None and (width, height) != (arr.shape[1], arr.shape[0]):
        img = img.resize((width, height), Image.NEAREST)
    img.save(path)
    return img


def draw_path_on_heatmap(
    Z: np.ndarray,
    xlim: tuple[float, float],
    ylim: tuple[float, float],
    path_xy: np.ndarray,
    out_path: str,
    marker_radius: int = 3,
) -> Image.Image:
    """Draw a descent path over its heatmap: a line through every point, plus
    a small circle marking each step.

    The heatmap is built at the grid's own resolution and the path is drawn
    on top of it at that SAME resolution, so world_to_pixel(xlim, ylim,
    width, height) uses the image's actual size rather than a guess -- the
    one place a mismatched width/height would silently misplace every point.
    """
    height, width = Z.shape[0], Z.shape[1]
    img = heatmap_png(Z, xlim, ylim, out_path, width=width, height=height)
    draw = ImageDraw.Draw(img)

    pixels = [world_to_pixel(x, y, xlim, ylim, width, height) for x, y in path_xy]
    if len(pixels) >= 2:
        draw.line(pixels, fill=(255, 255, 255), width=2)
    for px, py in pixels:
        draw.ellipse(
            [px - marker_radius, py - marker_radius, px + marker_radius, py + marker_radius],
            fill=(255, 255, 255),
            outline=(20, 20, 20),
        )
    img.save(out_path)
    return img


def loss_curve_points(
    losses: np.ndarray, width: int, height: int, margin: int, log: bool
) -> list[tuple[float, float]]:
    """Map a loss sequence to pixel coordinates, for drawing OR for testing
    that the log-scale points are collinear.

    Kept separate from loss_curve_png so a test can check the geometry the
    picture is actually built from, rather than re-deriving it in prose.
    """
    losses = np.asarray(losses, dtype=float)
    n = len(losses)
    xs_data = np.arange(n, dtype=float)
    ys_data = np.log10(np.clip(losses, 1e-300, None)) if log else losses

    x0, x1 = 0.0, float(n - 1) if n > 1 else 1.0
    y0, y1 = float(ys_data.min()), float(ys_data.max())
    xspan = x1 - x0 if x1 > x0 else 1.0
    yspan = y1 - y0 if y1 > y0 else 1.0

    points = []
    for xd, yd in zip(xs_data, ys_data):
        px = margin + (xd - x0) / xspan * (width - 2 * margin)
        py = (height - margin) - (yd - y0) / yspan * (height - 2 * margin)
        points.append((px, py))
    return points


def loss_curve_png(
    losses: np.ndarray,
    out_path: str,
    log: bool = False,
    width: int = 500,
    height: int = 350,
    margin: int = 50,
) -> Image.Image:
    """Draw loss against iteration as a PNG: axes, a polyline, and a marker
    per point. log=True plots log10(loss) on the y-axis instead of loss."""
    img = Image.new("RGB", (width, height), color=(248, 250, 252))
    draw = ImageDraw.Draw(img)
    draw.line([(margin, margin // 2), (margin, height - margin)], fill=(26, 32, 44), width=2)
    draw.line(
        [(margin, height - margin), (width - margin // 2, height - margin)],
        fill=(26, 32, 44),
        width=2,
    )
    points = loss_curve_points(losses, width, height, margin, log)
    if len(points) >= 2:
        draw.line(points, fill=(29, 78, 216), width=2)
    for px, py in points:
        draw.ellipse([px - 2, py - 2, px + 2, py + 2], fill=(29, 78, 216))
    img.save(out_path)
    return img


def animated_descent_gif(
    Z: np.ndarray,
    xlim: tuple[float, float],
    ylim: tuple[float, float],
    path_xy: np.ndarray,
    out_path: str,
    duration_ms: int = 150,
) -> None:
    """Save one GIF frame per step of a descent, the path built up one point
    at a time, over the same heatmap.

    Frame k shows path_xy[0 : k + 1] -- so the FIRST frame is a single
    marker at the start, and the LAST frame is the complete path drawn by
    draw_path_on_heatmap. Pillow's save(..., save_all=True) is the entire
    animation mechanism: it writes each frame's own image data plus a small
    GIF-format header saying "loop through these N frames".
    """
    height, width = Z.shape[0], Z.shape[1]
    background = heatmap_array(Z)
    frames = []
    for k in range(1, len(path_xy) + 1):
        img = Image.fromarray(background, mode="RGB").copy()
        draw = ImageDraw.Draw(img)
        pixels = [world_to_pixel(x, y, xlim, ylim, width, height) for x, y in path_xy[:k]]
        if len(pixels) >= 2:
            draw.line(pixels, fill=(255, 255, 255), width=2)
        px, py = pixels[-1]
        draw.ellipse([px - 4, py - 4, px + 4, py + 4], fill=(255, 255, 255), outline=(20, 20, 20))
        frames.append(img.convert("P", palette=Image.ADAPTIVE))
    frames[0].save(
        out_path,
        save_all=True,
        append_images=frames[1:],
        duration=duration_ms,
        loop=0,
        format="GIF",
    )

"""The reference test suite: nine numbered checks, one per lab exercise.

Every test writes into pytest's own tmp_path fixture (a directory pytest
creates and cleans up itself, outside this lab), never into the lab
directory. That is what makes exercise 9's cleanup check meaningful: nothing
in this suite is exempt from it.
"""

from __future__ import annotations

import numpy as np
import pytest
from PIL import Image

import dataset as D
import descent as DS
import gridviz as G
import imaging as IM

XLIM = (-4.0, 4.0)
YLIM = (-4.0, 4.0)
GRID_N = 101


# -- exercise 1: evaluate_grid ------------------------------------------------


def test_evaluate_grid_shape_and_minimum():
    X, Y, Z = G.evaluate_grid(D.WELL_F, XLIM, YLIM, GRID_N)
    assert X.shape == (GRID_N, GRID_N)
    assert Y.shape == (GRID_N, GRID_N)
    assert Z.shape == (GRID_N, GRID_N)
    iy, ix = np.unravel_index(np.argmin(Z), Z.shape)
    assert abs(X[iy, ix]) < 1e-9
    assert abs(Y[iy, ix]) < 1e-9
    assert Z[iy, ix] == pytest.approx(0.0, abs=1e-9)


# -- exercise 2: ascii contour renderer --------------------------------------


def test_ascii_contour_exact_characters():
    _, _, Z = G.evaluate_grid(lambda x, y: x**2 + y**2, (-2, 2), (-2, 2), 5)
    rendered = G.ascii_contour(Z)
    rows = rendered.split("\n")
    assert len(rows) == 5
    # Centre cell: value 0, the minimum -> the lightest character (a space).
    assert rows[2][2] == " "
    # Every corner: value 8, the maximum -> the densest character.
    for r, c in [(0, 0), (0, 4), (4, 0), (4, 4)]:
        assert rows[r][c] == "#"
    # A transposed grid or a flipped row order would break this symmetry:
    # x^2 + y^2 is symmetric under both left-right and top-bottom mirror.
    assert all(rows[r] == rows[4 - r] for r in range(5))
    assert all(rows[r] == rows[r][::-1] for r in range(5))


# -- exercise 3: Pillow heatmap PNG ------------------------------------------


def test_heatmap_png_size_and_minimum_color(tmp_path):
    _, _, Z = G.evaluate_grid(D.WELL_F, XLIM, YLIM, GRID_N)
    out = tmp_path / "heat.png"
    img = IM.heatmap_png(Z, XLIM, YLIM, str(out))
    assert out.exists()
    assert img.size == (GRID_N, GRID_N)
    reopened = Image.open(out)
    assert reopened.size == (GRID_N, GRID_N)
    px, py = G.world_to_pixel(0.0, 0.0, XLIM, YLIM, GRID_N, GRID_N)
    assert reopened.convert("RGB").getpixel((round(px), round(py))) == (13, 27, 84)


# -- exercise 4: path drawn over the heatmap, and the pixel transform --------


def test_world_to_pixel_corners():
    # x grows right: no flip. y grows up in data space, down in pixel rows.
    assert G.world_to_pixel(-4.0, 4.0, XLIM, YLIM, 101, 101) == (0.0, 0.0)
    assert G.world_to_pixel(4.0, -4.0, XLIM, YLIM, 101, 101) == (100.0, 100.0)
    assert G.world_to_pixel(0.0, 0.0, XLIM, YLIM, 101, 101) == (50.0, 50.0)


def test_descent_path_drawn_over_heatmap(tmp_path):
    _, _, Z = G.evaluate_grid(D.WELL_F, XLIM, YLIM, GRID_N)
    path = DS.gradient_descent(D.WELL_GRAD, D.START, D.LEARNING_RATE, D.STEPS)
    out = tmp_path / "path.png"
    IM.draw_path_on_heatmap(Z, XLIM, YLIM, path, str(out))
    assert out.exists()
    img = Image.open(out)
    assert img.size == (GRID_N, GRID_N)

    first_px = G.world_to_pixel(*path[0], XLIM, YLIM, GRID_N, GRID_N)
    expected_first = G.world_to_pixel(*D.START, XLIM, YLIM, GRID_N, GRID_N)
    assert first_px == expected_first

    last_px = G.world_to_pixel(*path[-1], XLIM, YLIM, GRID_N, GRID_N)
    minimum_px = G.world_to_pixel(0.0, 0.0, XLIM, YLIM, GRID_N, GRID_N)
    distance = ((last_px[0] - minimum_px[0]) ** 2 + (last_px[1] - minimum_px[1]) ** 2) ** 0.5
    assert distance < D.PIXEL_TOL


# -- exercise 5: loss curve, linear and log, and the collinearity proof -----


def test_loss_curve_png_linear_and_log(tmp_path):
    path = DS.gradient_descent(D.WELL_GRAD, D.START, D.LEARNING_RATE, D.STEPS)
    losses = DS.losses_along(D.WELL_F, path)

    lin_out = tmp_path / "loss_linear.png"
    log_out = tmp_path / "loss_log.png"
    IM.loss_curve_png(losses, str(lin_out), log=False)
    IM.loss_curve_png(losses, str(log_out), log=True)
    assert lin_out.exists() and log_out.exists()
    assert Image.open(lin_out).size == (500, 350)
    assert Image.open(log_out).size == (500, 350)


def test_log_axis_points_are_collinear():
    # a = b = 1: the update x <- (1 - 2 * lr * a) x is an EXACT geometric
    # recursion, so loss = a x^2 + b y^2 decays as a single power of a
    # constant ratio and log10(loss) against iteration is a straight line
    # by construction, not by approximation.
    path = DS.gradient_descent(D.WELL_GRAD, D.START, D.LEARNING_RATE, D.STEPS)
    losses = DS.losses_along(D.WELL_F, path)
    points = IM.loss_curve_points(losses, width=500, height=350, margin=50, log=True)
    xs = np.array([p[0] for p in points])
    ys = np.array([p[1] for p in points])
    A = np.vstack([xs, np.ones_like(xs)]).T
    slope, intercept = np.linalg.lstsq(A, ys, rcond=None)[0]
    residual = np.max(np.abs(ys - (slope * xs + intercept)))
    assert residual < 1e-6  # pixel units: far below one pixel


# -- exercise 6: animated GIF -------------------------------------------------


def test_animated_gif_frame_count(tmp_path):
    _, _, Z = G.evaluate_grid(D.WELL_F, XLIM, YLIM, GRID_N)
    path = DS.gradient_descent(D.WELL_GRAD, D.START, D.LEARNING_RATE, 20)
    out = tmp_path / "descent.gif"
    IM.animated_descent_gif(Z, XLIM, YLIM, path, str(out))
    assert out.exists()
    reopened = Image.open(out)
    assert reopened.format == "GIF"
    assert reopened.n_frames == len(path)


# -- exercise 7: the learning-rate sweep -------------------------------------


def test_learning_rate_sweep_shape():
    etas = np.round(np.arange(0.05, 2.55, 0.1), 2)
    sweep = DS.learning_rate_sweep(D.sweep_grad, D.sweep_f, D.SWEEP_X0, etas, D.SWEEP_STEPS)
    assert [e for e, _ in sweep] == list(etas)

    finite = [(e, loss) for e, loss in sweep if np.isfinite(loss)]
    divergent = [(e, loss) for e, loss in sweep if not np.isfinite(loss)]

    # A basin, not a single magic number: more than one eta converges well.
    good = [e for e, loss in finite if loss < 1e-6]
    assert len(good) >= 3

    # The argmin sits strictly inside the swept range, not at either edge.
    best_eta = min(finite, key=lambda pair: pair[1])[0]
    assert etas[0] < best_eta < etas[-1]

    # And the sweep actually reaches the cliff: some etas genuinely diverge,
    # caught as float('inf') rather than raising.
    assert len(divergent) >= 1
    assert all(loss == float("inf") for _, loss in divergent)
    # Divergence only happens on the far side of the theoretical threshold
    # eta = 1 / a = 1.0 for f(x) = x^2.
    assert all(e > 1.0 for e, _ in divergent)


# -- exercise 8: two runs, same loss, very different paths -------------------


def test_two_runs_same_loss_different_path_length():
    well_path = DS.gradient_descent(D.WELL_GRAD, D.START, D.LEARNING_RATE, D.STEPS)
    ill_path = DS.gradient_descent(D.ILL_GRAD, D.START, D.LEARNING_RATE, D.STEPS)

    well_loss = DS.losses_along(D.WELL_F, well_path)[-1]
    ill_loss = DS.losses_along(D.ILL_F, ill_path)[-1]
    relative_gap = abs(well_loss - ill_loss) / max(well_loss, ill_loss)
    assert relative_gap < D.LOSS_MATCH_TOL

    well_len = DS.path_length(well_path)
    ill_len = DS.path_length(ill_path)
    assert ill_len / well_len > D.PATH_LENGTH_RATIO_MIN


# -- exercise 9 (the cleanup check) lives in tests/run_tests.sh, which
# inspects the whole lab directory after every other check has run -- a
# single pytest test cannot see "did the ENTIRE suite leave anything behind"
# from inside one test function.

lesson_id: D112
day: 112
kind: guided-build
languages: [python, bash]
setup_commands:
  - cd labs/sections/math-statistics-and-data/day-112-visualizing-optimization
  - python3 -m venv .venv
  - .venv/bin/pip install -r requirements/requirements.txt
  - .venv/bin/python3 -c "import numpy, PIL; print(numpy.__version__, PIL.__version__)"
run_commands:
  - 'cd examples && ../.venv/bin/python3 01_grid_and_ascii.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 02_heatmap_and_path.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 03_loss_curves.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 04_animated_gif.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 05_learning_rate_sweep.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 06_two_runs_same_loss.py && cd ..'
  - .venv/bin/pytest examples -q -p no:cacheprovider
  - .venv/bin/pytest starter -q -p no:cacheprovider
test_commands:
  - bash tests/run_tests.sh
cleanup_commands:
  - "find . -path ./.venv -prune -o -type d -name '__pycache__' -print -exec rm -rf -- {} +"
  - rm -rf .pytest_cache
  - 'rm -rf .venv  # optional: removes the lab virtual environment'
  - 'git checkout -- starter/  # optional: reset your work'
requires_network: true
requires_api_key: false
estimated_minutes: 30
last_executed: '2026-08-17'
executed_on: 'macOS 26.5.2 (Apple Silicon, arm64), Python 3.14.0, numpy 2.5.2, Pillow 12.3.0, pytest 9.1.1, bash 3.2.57 -- bash tests/run_tests.sh -> 31 checks, 0 failure(s), exit 0; pytest examples -> 10 passed; pytest starter -> 13 skipped on an untouched checkout, and 13 passed against a fully solved copy of starter/ (gridviz.py, descent.py and imaging.py copied in from examples/) kept outside the lab, which was verified rather than assumed. All six numbered reference scripts exit 0 and each prints its own "every assertion held" line. Everything ran through the authoring virtual environment (numpy 2.5.2, Pillow 12.3.0, pytest 9.1.1, Python 3.14.0), not the bare system interpreter, which has none of the three packages. The harness was additionally confirmed to be capable of failing: section 6 re-runs the whole script with the relative-loss-gap threshold from section 5 tightened from 0.05 to 0.001, which the two measured runs (a well-conditioned and an ill-conditioned quadratic bowl, same start, same learning rate, same step count) cannot meet, and the re-run reported exactly one named failure and exited non-zero; with the threshold restored, the harness is green again. matplotlib is confirmed absent from this environment by section 1 of the harness, which attempts the import and asserts it fails; no matplotlib, scipy, pandas or Plotly output is reproduced anywhere in this lab or its lesson, and the lesson states so plainly in its Tools section. The two-run comparison that opens the lesson was measured, not chosen after the fact to look good: a well-conditioned bowl f(x,y) = x^2 + y^2 and an ill-conditioned one f(x,y) = x^2 + 25y^2, both started at (4, 4) with learning rate 0.038 for 60 steps, land at final losses of 2.431009e-03 and 2.507203e-03 respectively -- a relative gap of 3.04%, well inside the well-conditioned run''s own noise floor -- while their path lengths are 5.6075 and 75.9767, a ratio of 13.55x. The log-axis collinearity claim for the well-conditioned run is proved rather than asserted in prose: log10(loss) against iteration is an exactly geometric sequence for this update rule when a = b = 1, and fitting a line to the drawn pixel coordinates gives a maximum residual on the order of 1e-13, reported in the test rather than hard-coded to a looser bound. The learning-rate sweep on f(x) = x^2 over eta from 0.05 to 2.45 in steps of 0.1, run for 300 steps starting at x0 = 4.0, found ten learning rates reaching a loss indistinguishable from zero, an interior argmin at eta = 0.45 (tied with 0.55 by symmetry), the largest eta whose blow-up had not yet overflowed float64 within 300 steps at eta = 2.05, and the smallest eta whose blow-up had overflowed to inf at eta = 2.15 -- all four caught deliberately via numpy.errstate rather than raised or warned. One honesty note from authoring: a different agent began authoring this same day into this same lab directory concurrently, wrote a design under different module and function names (`optim_viz.py` among them) into examples/, and briefly overwrote this implementation''s test_reference.py before the collision was noticed and resolved by deleting every file from that design and regenerating every expected-output capture, this metadata file, and this lab from the surviving implementation''s own runs, made today. No content in this lab or its expected-output captures describes or was carried over from the other design.'

# What is installed, why, and what it costs

Three packages, all free and open source, all installed into a lab-local
virtual environment that `rm -rf .venv` completely undoes.

| Package | Version pinned | Licence | What this lab uses it for |
| --- | --- | --- | --- |
| `numpy` | 2.5.2 | BSD 3-Clause | `numpy.meshgrid` to build every grid, vectorised arithmetic for the colour ramp, and the gradient-descent loop itself. |
| `Pillow` | 12.3.0 | MIT-CMU (the "PIL Software License") | `Image`, `ImageDraw` and `Image.save(..., save_all=True)` — every pixel this lab draws, and the only GIF-writing code it needs. |
| `pytest` | 9.1.1 | MIT | The reference suite and your running score in `starter/`. |

There is no paid tier of anything in this lab, no account, no key and no
signup, personally or commercially.

## The tool this lab deliberately does not install

**matplotlib is not installed here, and neither is scipy, pandas or plotly.**
That is not an oversight — it is the reason this lab exists in the form it
does. Every picture is built from a NumPy array and Pillow's `ImageDraw`
with nothing else in between, so nothing about how a contour plot or a loss
curve actually gets to the screen is hidden behind a library call. The
lesson's Tools section describes matplotlib, Plotly, TensorBoard and Weights
& Biases from their documentation and states plainly that no output from any
of them is reproduced here.

## The one time the network is needed

```bash
.venv/bin/pip install -r requirements/requirements.txt
```

That is the only command in the lab that opens a connection. Section 6 of
`tests/run_tests.sh` greps every source file in `examples/` and `starter/` to
prove that nothing else does.

## If you cannot install anything at all

Almost none of this lab runs without Pillow, and that is worth saying
plainly rather than glossing over: `evaluate_grid` and `ascii_contour` need
only NumPy (and the ASCII renderer is genuinely useful on its own — it is how
this lab's first picture gets checked without an image viewer at all), but
every PNG and the GIF need `PIL.Image` and `PIL.ImageDraw` directly, and there
is no standard-library substitute for either.

## What is deliberately *not* installed, and why it is not a problem

matplotlib, Plotly, TensorBoard and Weights & Biases all do this job, usually
better and always faster to write. None of them is installed here, and **no
output from any of them is reproduced anywhere in this lab or its lesson.**
They are described from their documentation, and the lesson's Tools section
marks each one as not run.

That is not a limitation to apologise for. `heatmap_png` and
`draw_path_on_heatmap` in `examples/imaging.py` do, by hand, exactly what
`matplotlib.pyplot.contourf` and `plt.plot` do for you: evaluate a function
over a grid, map values to colours, and place pixels. Having built the
minimal version, the documentation of the real one reads as engineering on
top of an idea you already own, rather than as magic.

numpy==2.5.2
Pillow==12.3.0
pytest==9.1.1

# Day 112 lab — the brief

Eight exercises across three files, in order. Work top to bottom: `gridviz.py`
before `imaging.py`, and `imaging.py`'s later functions before its earlier
ones' consumers. `dataset.py` is given to you — the loss surfaces and the
descent loop are Day 111's subject, not this one.

Check yourself at any point:

```bash
.venv/bin/pytest starter -q
```

On an untouched checkout that prints `13 skipped`. A **skip** means "not
attempted". A **failure** means "attempted and wrong", and it prints your
answer beside the real one. When every test passes, you are finished.

---

## Exercise 1 — `gridviz.py`, `evaluate_grid`

`numpy.meshgrid` plus a function call. Every picture later in this lab starts
here.

## Exercise 2 — `gridviz.py`, `ascii_contour`

Rescale a 2D array to a small number of character bands and print it. Get the
row/column order wrong and a symmetric bowl stops looking symmetric — in a
terminal, immediately, with no image viewer required.

## Exercise 3 — `imaging.py`, `heatmap_array` and `heatmap_png`

Turn `evaluate_grid`'s output into an image. `heatmap_array` does the one
deliberate axis flip in this lab (`numpy.flipud`); `heatmap_png` wraps it in
`Image.fromarray` and `img.save`.

## Exercise 4 — `gridviz.py`'s `world_to_pixel`, and `imaging.py`'s
`draw_path_on_heatmap`

`world_to_pixel` is the one function every other drawing function in this lab
calls instead of repeating the arithmetic — get it right once, here, and
everything downstream is right too. Then draw a descent path over its
heatmap with `PIL.ImageDraw.line` and `.ellipse`.

## Exercise 5 — `imaging.py`, `loss_curve_points` and `loss_curve_png`

Map a loss sequence to pixel coordinates — on a linear axis directly, on a
log axis via `numpy.log10` — then draw it. Remember pixel row 0 is the TOP:
a larger data value must produce a SMALLER pixel row on both axes.

## Exercise 6 — `imaging.py`, `animated_descent_gif`

One GIF frame per step, built by re-using exercise 4's drawing approach on a
growing prefix of the path. `Image.save(..., save_all=True,
append_images=...)` is the entire animation mechanism.

## Exercise 7 — `descent.py`, `sweep_final_loss` and `learning_rate_sweep`

Run a 1D descent at a given learning rate and report the final loss — or
`float('inf')` if it diverged. A learning rate above the stability threshold
makes the run overflow float64, which must be caught deliberately with
`numpy.errstate`, not allowed to raise.

## Exercise 8 — `descent.py`, `path_length`

Total Euclidean distance travelled along a path. One line: `numpy.diff` plus
`numpy.linalg.norm`. This is the number that tells two runs with
near-identical final losses apart.

---

## When you are done

Read the reference and run every script:

```bash
cd examples
../.venv/bin/python3 01_grid_and_ascii.py
../.venv/bin/python3 02_heatmap_and_path.py
../.venv/bin/python3 03_loss_curves.py
../.venv/bin/python3 04_animated_gif.py
../.venv/bin/python3 05_learning_rate_sweep.py
../.venv/bin/python3 06_two_runs_same_loss.py
cd ..
```

Script 6 is the one to read even if you read nothing else: two runs land
within a few percent of the same final loss, and their path lengths differ
by more than 13x.

"""Make this directory's own modules the ones its tests import.

Identical in purpose to examples/conftest.py: both directories define
`dataset`, `gridviz`, `descent` and `imaging`, and without this file pytest
would import whichever it saw first and reuse it for the other suite.
"""

import sys
from pathlib import Path

HERE = str(Path(__file__).parent.resolve())

if HERE in sys.path:
    sys.path.remove(HERE)
sys.path.insert(0, HERE)

for name in ("dataset", "gridviz", "descent", "imaging"):
    module = sys.modules.get(name)
    origin = getattr(module, "__file__", "") or ""
    if module is not None and not origin.startswith(HERE):
        del sys.modules[name]

"""Shared numbers for the Day 112 lab -- given to you, not an exercise.

Identical to examples/dataset.py. This day's own subject is turning arrays
into pictures; the loss surfaces and the learning-rate schedule below are
infrastructure carried over from Day 111 so you can spend your effort on
visualization, not on re-deriving a bowl function.
"""

import numpy as np

# -- the two bowls -----------------------------------------------------------

WELL_A, WELL_B = 1.0, 1.0
ILL_A, ILL_B = 1.0, 25.0

START = np.array([4.0, 4.0])
LEARNING_RATE = 0.038
STEPS = 60


def bowl(a: float, b: float):
    """Return (f, grad) for f(x, y) = a x^2 + b y^2."""

    def f(x, y):
        return a * x**2 + b * y**2

    def grad(x, y):
        return np.array([2.0 * a * x, 2.0 * b * y])

    return f, grad


WELL_F, WELL_GRAD = bowl(WELL_A, WELL_B)
ILL_F, ILL_GRAD = bowl(ILL_A, ILL_B)

# -- the one-dimensional bowl used for the learning-rate sweep ---------------

SWEEP_X0 = 4.0
SWEEP_STEPS = 300


def sweep_f(x):
    return x**2


def sweep_grad(x):
    return 2.0 * x


# -- tolerances ----------------------------------------------------------

EXACT_TOL = 1e-9
LOSS_MATCH_TOL = 0.05
PATH_LENGTH_RATIO_MIN = 5.0
PIXEL_TOL = 2.0


def gradient_descent(grad_fn, x0, lr: float, steps: int) -> np.ndarray:
    """Given: run steps of x <- x - lr * grad_fn(x), returning every visited
    point as an array of shape (steps + 1, len(x0)).

    This is Day 111's subject -- the update rule itself -- not this day's.
    """
    x = np.array(x0, dtype=float)
    path = [x.copy()]
    for _ in range(steps):
        x = x - lr * np.asarray(grad_fn(*x))
        path.append(x.copy())
    return np.array(path)


def losses_along(f, path: np.ndarray) -> np.ndarray:
    """Given: evaluate f at every point on a path."""
    return np.array([f(*p) for p in path])

"""Exercises 7 and 8 -- path_length and the learning-rate sweep.

`return None` where your code goes; `pytest starter -q` skips what you have
not attempted and fails only what you have attempted incorrectly.
"""

from __future__ import annotations

import numpy as np


def path_length(path: np.ndarray) -> float:
    """Exercise 8a -- total Euclidean distance travelled along a path: the
    sum of the step sizes between consecutive points.

    Approach: `steps = np.diff(path, axis=0)`, then sum the row-wise norm:
    `np.sum(np.linalg.norm(steps, axis=1))`.
    """
    return None


def sweep_final_loss(grad_fn, f, x0: float, eta: float, steps: int) -> float:
    """Exercise 7a -- run gradient descent on a 1D function at learning rate
    eta for `steps` steps, and return the final loss.

    A learning rate above the stability threshold makes the iterate grow
    without bound. That growth eventually overflows float64 to `inf` --
    which is expected behaviour to CATCH, not an exception to let escape.
    Use `numpy.errstate(over="ignore", invalid="ignore")` around the update
    and around the final `f(x)`, and check `np.isfinite(x)` after every step;
    return `float('inf')` the moment it stops being finite.

    Approach:

        x = np.float64(x0)
        for _ in range(steps):
            with np.errstate(over="ignore", invalid="ignore"):
                x = x - eta * grad_fn(x)
            if not np.isfinite(x):
                return float("inf")
        with np.errstate(over="ignore", invalid="ignore"):
            value = f(x)
        return float(value) if np.isfinite(value) else float("inf")
    """
    return None


def learning_rate_sweep(grad_fn, f, x0: float, etas, steps: int):
    """Exercise 7b -- run sweep_final_loss at every learning rate in etas.

    Return a list of (eta, final_loss) pairs, in the order etas was given.

    Approach: one line, a list comprehension calling sweep_final_loss.
    """
    return None

"""Exercises 1, 2 and 4a -- write the three functions below.

Every function has a working signature, a docstring saying exactly what it
must do, and `return None` where your code goes. Returning None is how the
test suite knows you have not attempted it yet: `pytest starter -q` will
SKIP an unattempted function rather than fail it.

Check yourself as you go:

    .venv/bin/pytest starter -q
"""

from __future__ import annotations

import numpy as np

ASCII_RAMP = " .:+#"


def evaluate_grid(f, xlim: tuple[float, float], ylim: tuple[float, float], n: int):
    """Exercise 1 -- evaluate f over an n x n grid spanning xlim x ylim.

    Return (X, Y, Z), each an (n, n) array from numpy.meshgrid: X varies
    along columns, Y varies along rows, and Z = f(X, Y).

    Approach: `xs = np.linspace(xlim[0], xlim[1], n)`, same for ys, then
    `X, Y = np.meshgrid(xs, ys)` and `Z = f(X, Y)`.
    """
    return None


def ascii_contour(Z: np.ndarray, chars: str = ASCII_RAMP) -> str:
    """Exercise 2 -- render a 2D array as text: one character per cell,
    chosen by level band.

    Rescale Z linearly between its own min and max to [0, len(chars)), floor
    to an integer band index, clip to [0, len(chars) - 1], and join each
    row's characters with "\\n" between rows.

    Approach: `zmin, zmax = Z.min(), Z.max()`; `idx = np.clip(((Z - zmin) /
    span * len(chars)).astype(int), 0, len(chars) - 1)`; build one string per
    row from `chars[i]` and join the rows with newlines.
    """
    return None


def world_to_pixel(
    x: float,
    y: float,
    xlim: tuple[float, float],
    ylim: tuple[float, float],
    width: int,
    height: int,
) -> tuple[float, float]:
    """Exercise 4a -- map one (x, y) data point to a (column, row) pixel.

    x grows to the right in both spaces (columns increase with x, no flip).
    y grows UPWARD in data space but pixel row 0 is the TOP of the image, so
    row must be computed from (ylim[1] - y), not (y - ylim[0]).

    Approach: `px = (x - xlim[0]) / (xlim[1] - xlim[0]) * (width - 1)`;
    `py = (ylim[1] - y) / (ylim[1] - ylim[0]) * (height - 1)`.
    """
    return None

"""Exercises 3, 4b, 5 and 6 -- turn arrays into pictures with NumPy and
Pillow's ImageDraw. Work top to bottom: exercise 4b needs your own
`world_to_pixel` from gridviz.py, and exercise 6 reuses exercise 4b's ideas.

`return None` where your code goes; `pytest starter -q` skips what you have
not attempted.
"""

from __future__ import annotations

import numpy as np
from PIL import Image, ImageDraw

from gridviz import world_to_pixel

# -- given: the colour ramp ---------------------------------------------
#
# Four control points, dark blue (low) through teal and gold to dark red
# (high). The colour choice is not the exercise -- turning a value array
# into pixels is.

_STOPS_T = np.array([0.0, 0.35, 0.65, 1.0])
_STOPS_RGB = np.array(
    [
        [13, 27, 84],
        [29, 78, 216],
        [250, 204, 21],
        [185, 28, 28],
    ]
)


def ramp_color(t: np.ndarray) -> np.ndarray:
    """Given: map an array of values in [0, 1] to an (..., 3) uint8 RGB array."""
    t = np.clip(t, 0.0, 1.0)
    r = np.interp(t, _STOPS_T, _STOPS_RGB[:, 0])
    g = np.interp(t, _STOPS_T, _STOPS_RGB[:, 1])
    b = np.interp(t, _STOPS_T, _STOPS_RGB[:, 2])
    return np.stack([r, g, b], axis=-1).astype(np.uint8)


def heatmap_array(Z: np.ndarray) -> np.ndarray:
    """Exercise 3a -- turn a 2D value array into an (n, n, 3) uint8 RGB array.

    Rescale Z to [0, 1] with (Z - Z.min()) / (Z.max() - Z.min()), pass it to
    `ramp_color`, and flip the result vertically before returning it: Z's row
    0 is the SMALLEST y (evaluate_grid's convention), but the TOP of an image
    must show the LARGEST y.

    Approach: `t = (Z - Z.min()) / span`; `colors = ramp_color(t)`; return
    `np.flipud(colors)`.
    """
    return None


def heatmap_png(
    Z: np.ndarray,
    xlim: tuple[float, float],
    ylim: tuple[float, float],
    path: str,
    width: int | None = None,
    height: int | None = None,
) -> Image.Image:
    """Exercise 3b -- save Z as a heatmap PNG (via heatmap_array) and return
    the Image. Default width/height to Z's own resolution.

    Approach: `arr = heatmap_array(Z)`; `img = Image.fromarray(arr,
    mode="RGB")`; resize if width/height were given and differ; `img.save(path)`;
    return `img`.
    """
    return None


def draw_path_on_heatmap(
    Z: np.ndarray,
    xlim: tuple[float, float],
    ylim: tuple[float, float],
    path_xy: np.ndarray,
    out_path: str,
    marker_radius: int = 3,
) -> Image.Image:
    """Exercise 4b -- draw a descent path over its heatmap: a line through
    every point, plus a small circle marking each step.

    Build the heatmap at Z's own resolution, then map every (x, y) in
    path_xy to a pixel with `world_to_pixel` and draw a polyline plus a
    circle per point with `PIL.ImageDraw`.

    Approach: `height, width = Z.shape`; build the heatmap; `draw =
    ImageDraw.Draw(img)`; `pixels = [world_to_pixel(x, y, xlim, ylim, width,
    height) for x, y in path_xy]`; `draw.line(pixels, ...)`; loop over
    `pixels` calling `draw.ellipse` centred on each; save and return `img`.
    """
    return None


def loss_curve_points(
    losses: np.ndarray, width: int, height: int, margin: int, log: bool
) -> list[tuple[float, float]]:
    """Exercise 5a -- map a loss sequence to pixel coordinates.

    x data is the iteration index 0..len(losses)-1. y data is `losses`
    itself if log is False, or `np.log10(losses)` if log is True. Rescale
    both to fit inside [margin, width - margin] and [margin, height -
    margin] respectively, remembering that pixel row 0 is the TOP: a larger
    y-data value must produce a SMALLER pixel row.

    Approach: compute `xs_data = np.arange(len(losses))` and `ys_data =
    np.log10(losses) if log else losses`; find each span; for every (xd, yd)
    compute `px = margin + (xd - x0) / xspan * (width - 2*margin)` and
    `py = (height - margin) - (yd - y0) / yspan * (height - 2*margin)`.
    """
    return None


def loss_curve_png(
    losses: np.ndarray,
    out_path: str,
    log: bool = False,
    width: int = 500,
    height: int = 350,
    margin: int = 50,
) -> Image.Image:
    """Exercise 5b -- draw loss against iteration as a PNG: two axis lines,
    a polyline through `loss_curve_points`, and a small marker per point.

    Approach: create a blank `Image.new("RGB", (width, height), ...)`; draw
    the y-axis and x-axis with `ImageDraw.line`; get `points =
    loss_curve_points(...)`; `draw.line(points, ...)`; loop drawing a small
    `draw.ellipse` at each point; save and return the image.
    """
    return None


def animated_descent_gif(
    Z: np.ndarray,
    xlim: tuple[float, float],
    ylim: tuple[float, float],
    path_xy: np.ndarray,
    out_path: str,
    duration_ms: int = 150,
) -> None:
    """Exercise 6 -- save one GIF frame per step of a descent: frame k shows
    path_xy[0 : k + 1] drawn over the (shared) heatmap background.

    Approach: build `background = heatmap_array(Z)` once; for k from 1 to
    len(path_xy), copy an Image from `background`, draw the partial path
    with `world_to_pixel` and `ImageDraw` the way exercise 4b does, and
    convert each frame to `"P"` mode (`img.convert("P",
    palette=Image.ADAPTIVE)`) before appending it to a list. Then call
    `frames[0].save(out_path, save_all=True, append_images=frames[1:],
    duration=duration_ms, loop=0, format="GIF")`.
    """
    return None

"""The starter test suite: skips unattempted work instead of failing it.

Every exercise function returns None until you write it. A test whose
function still returns None is SKIPPED ("not attempted"). A test whose
function returns something else is checked for real -- a wrong answer FAILS
and prints what you got beside what was expected.

Run at any point:

    .venv/bin/pytest starter -q
"""

from __future__ import annotations

import numpy as np
import pytest
from PIL import Image

import dataset as D
import descent as DS
import gridviz as G
import imaging as IM

XLIM = (-4.0, 4.0)
YLIM = (-4.0, 4.0)
GRID_N = 101


def skip_if_none(value, label):
    if value is None:
        pytest.skip(f"{label} not attempted yet")


def get_grid(f=None, xlim=XLIM, ylim=YLIM, n=GRID_N):
    """Every later exercise needs a working evaluate_grid. Skip cleanly
    rather than crashing on an unpack of None if it is not attempted yet."""
    f = D.WELL_F if f is None else f
    result = G.evaluate_grid(f, xlim, ylim, n)
    skip_if_none(result, "evaluate_grid (needed by this exercise)")
    return result


# -- exercise 1: evaluate_grid ------------------------------------------------


def test_1_evaluate_grid_shape_and_minimum():
    result = G.evaluate_grid(D.WELL_F, XLIM, YLIM, GRID_N)
    skip_if_none(result, "evaluate_grid")
    X, Y, Z = result
    assert X.shape == (GRID_N, GRID_N)
    assert Y.shape == (GRID_N, GRID_N)
    assert Z.shape == (GRID_N, GRID_N)
    iy, ix = np.unravel_index(np.argmin(Z), Z.shape)
    assert abs(X[iy, ix]) < 1e-9
    assert abs(Y[iy, ix]) < 1e-9


# -- exercise 2: ascii contour renderer --------------------------------------


def test_2_ascii_contour_exact_characters():
    grid = G.evaluate_grid(lambda x, y: x**2 + y**2, (-2, 2), (-2, 2), 5)
    skip_if_none(grid, "evaluate_grid")
    _, _, Z = grid
    rendered = G.ascii_contour(Z)
    skip_if_none(rendered, "ascii_contour")
    rows = rendered.split("\n")
    assert rows[2][2] == " "
    for r, c in [(0, 0), (0, 4), (4, 0), (4, 4)]:
        assert rows[r][c] == "#"


# -- exercise 3: Pillow heatmap PNG ------------------------------------------


def test_3_heatmap_array_shape_and_dtype():
    _, _, Z = get_grid()
    arr = IM.heatmap_array(Z)
    skip_if_none(arr, "heatmap_array")
    assert arr.shape == (GRID_N, GRID_N, 3)
    assert arr.dtype == np.uint8


def test_3_heatmap_png_size_and_minimum_color(tmp_path):
    _, _, Z = get_grid()
    out = tmp_path / "heat.png"
    img = IM.heatmap_png(Z, XLIM, YLIM, str(out))
    skip_if_none(img, "heatmap_png")
    assert out.exists()
    assert img.size == (GRID_N, GRID_N)
    px, py = G.world_to_pixel(0.0, 0.0, XLIM, YLIM, GRID_N, GRID_N) or (None, None)
    if px is None:
        pytest.skip("world_to_pixel not attempted yet")
    reopened = Image.open(out).convert("RGB")
    assert reopened.getpixel((round(px), round(py))) == (13, 27, 84)


# -- exercise 4: world_to_pixel and the drawn path ---------------------------


def test_4_world_to_pixel_corners():
    result = G.world_to_pixel(-4.0, 4.0, XLIM, YLIM, 101, 101)
    skip_if_none(result, "world_to_pixel")
    assert result == (0.0, 0.0)
    assert G.world_to_pixel(4.0, -4.0, XLIM, YLIM, 101, 101) == (100.0, 100.0)
    assert G.world_to_pixel(0.0, 0.0, XLIM, YLIM, 101, 101) == (50.0, 50.0)


def test_4_descent_path_drawn_over_heatmap(tmp_path):
    _, _, Z = get_grid()
    path = D.gradient_descent(D.WELL_GRAD, D.START, D.LEARNING_RATE, D.STEPS)
    out = tmp_path / "path.png"
    img = IM.draw_path_on_heatmap(Z, XLIM, YLIM, path, str(out))
    skip_if_none(img, "draw_path_on_heatmap")
    assert out.exists()
    last_px = G.world_to_pixel(*path[-1], XLIM, YLIM, GRID_N, GRID_N)
    min_px = G.world_to_pixel(0.0, 0.0, XLIM, YLIM, GRID_N, GRID_N)
    distance = ((last_px[0] - min_px[0]) ** 2 + (last_px[1] - min_px[1]) ** 2) ** 0.5
    assert distance < D.PIXEL_TOL


# -- exercise 5: loss curves --------------------------------------------------


def test_5_loss_curve_points_collinear_on_log_axis():
    path = D.gradient_descent(D.WELL_GRAD, D.START, D.LEARNING_RATE, D.STEPS)
    losses = D.losses_along(D.WELL_F, path)
    points = IM.loss_curve_points(losses, width=500, height=350, margin=50, log=True)
    skip_if_none(points, "loss_curve_points")
    xs = np.array([p[0] for p in points])
    ys = np.array([p[1] for p in points])
    A = np.vstack([xs, np.ones_like(xs)]).T
    slope, intercept = np.linalg.lstsq(A, ys, rcond=None)[0]
    residual = np.max(np.abs(ys - (slope * xs + intercept)))
    assert residual < 1e-6


def test_5_loss_curve_png(tmp_path):
    path = D.gradient_descent(D.WELL_GRAD, D.START, D.LEARNING_RATE, D.STEPS)
    losses = D.losses_along(D.WELL_F, path)
    out = tmp_path / "loss.png"
    img = IM.loss_curve_png(losses, str(out), log=True)
    skip_if_none(img, "loss_curve_png")
    assert out.exists()
    assert Image.open(out).size == (500, 350)


# -- exercise 6: animated GIF -------------------------------------------------


def test_6_animated_gif_frame_count(tmp_path):
    _, _, Z = get_grid()
    path = D.gradient_descent(D.WELL_GRAD, D.START, D.LEARNING_RATE, 20)
    out = tmp_path / "descent.gif"
    result = IM.animated_descent_gif(Z, XLIM, YLIM, path, str(out))
    if not out.exists():
        pytest.skip("animated_descent_gif not attempted yet")
    reopened = Image.open(out)
    assert reopened.format == "GIF"
    assert reopened.n_frames == len(path)


# -- exercise 7: the learning-rate sweep -------------------------------------


def test_7_sweep_final_loss_catches_divergence():
    result = DS.sweep_final_loss(D.sweep_grad, D.sweep_f, D.SWEEP_X0, 2.5, D.SWEEP_STEPS)
    skip_if_none(result, "sweep_final_loss")
    assert result == float("inf")
    converged = DS.sweep_final_loss(D.sweep_grad, D.sweep_f, D.SWEEP_X0, 0.45, D.SWEEP_STEPS)
    assert converged < 1e-6


def test_7_learning_rate_sweep_shape():
    etas = np.round(np.arange(0.05, 2.55, 0.1), 2)
    sweep = DS.learning_rate_sweep(D.sweep_grad, D.sweep_f, D.SWEEP_X0, etas, D.SWEEP_STEPS)
    skip_if_none(sweep, "learning_rate_sweep")
    assert [e for e, _ in sweep] == list(etas)
    divergent = [(e, loss) for e, loss in sweep if not np.isfinite(loss)]
    assert len(divergent) >= 1


# -- exercise 8: path_length, and two runs at nearly the same loss ----------


def test_8_path_length_of_a_known_path():
    # A path that steps (0,0) -> (3,0) -> (3,4) has length 3 + 4 = 7 exactly
    # (a 3-4-5 triangle's two legs).
    known_path = np.array([[0.0, 0.0], [3.0, 0.0], [3.0, 4.0]])
    result = DS.path_length(known_path)
    skip_if_none(result, "path_length")
    assert result == pytest.approx(7.0)


def test_8_two_runs_same_loss_different_path_length():
    well_path = D.gradient_descent(D.WELL_GRAD, D.START, D.LEARNING_RATE, D.STEPS)
    ill_path = D.gradient_descent(D.ILL_GRAD, D.START, D.LEARNING_RATE, D.STEPS)

    well_loss = D.losses_along(D.WELL_F, well_path)[-1]
    ill_loss = D.losses_along(D.ILL_F, ill_path)[-1]
    relative_gap = abs(well_loss - ill_loss) / max(well_loss, ill_loss)
    assert relative_gap < D.LOSS_MATCH_TOL

    well_len = DS.path_length(well_path)
    ill_len = DS.path_length(ill_path)
    skip_if_none(well_len, "path_length")
    skip_if_none(ill_len, "path_length")
    assert ill_len / well_len > D.PATH_LENGTH_RATIO_MIN

#!/usr/bin/env bash
# Tests for the Day 112 lab. Run from the lab directory:
#   bash tests/run_tests.sh
#
# The harness proves the lesson's claims by running code and reading real
# values, never by reading source:
#
#   * evaluate_grid puts the analytic minimum at the true grid cell, for a
#     bowl and for an anisotropic one;
#   * the ASCII contour renderer produces the exact character at the exact
#     cell a symmetric bowl predicts -- a transposed grid fails this loudly;
#   * a Pillow heatmap's pixel at the minimum is the colour ramp's own
#     lowest-value stop, read back from the saved PNG;
#   * world_to_pixel places the four corners and the centre of a symmetric
#     window exactly, and a descent path drawn with it starts at the start
#     point's pixel and ends within two pixels of the minimum's;
#   * a well-conditioned run's loss on a log10 axis is provably collinear --
#     fit to a line, the residual is measured in pixels, not asserted in
#     prose;
#   * an animated GIF has exactly as many frames as the descent had steps;
#   * a learning-rate sweep finds an interior optimum, a basin of several
#     good rates either side of it, and genuine divergence past eta = 1,
#     caught as float('inf') rather than raised as an exception;
#   * two runs on differently conditioned bowls, same start, same learning
#     rate, same step count, land within 5% of the same final loss while
#     their path lengths differ by more than 5x -- the day's opening claim,
#     made into an assertion;
#   * nothing is left behind on disk, including no PNG or GIF file.
#
# Everything after the one-time install runs offline. Nothing binds a port,
# nothing writes outside the lab (or a temporary directory it removes
# itself), nothing needs a key. Deterministic, non-interactive, exits 0 only
# if every check passes.
set -u

export PYTHONDONTWRITEBYTECODE=1

lab_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"

# Bytecode left by an EARLIER command is not this run's litter. The README
# documents `pytest starter -q`, and running it writes .pyc files that would
# then fail the cleanliness check at the end of this script -- failing the
# reader for following the instructions. Clearing them here makes that final
# check measure what it claims to: what THIS run left behind. `.venv` is
# untouched, because the packages' own bytecode is theirs, not ours.
find "${lab_dir}" -name '.venv' -prune -o -type d -name '__pycache__' -exec rm -rf {} + 2>/dev/null || true
find "${lab_dir}" -name '.venv' -prune -o -type d -name '.pytest_cache' -exec rm -rf {} + 2>/dev/null || true

failures=0
checks=0

check() {
  local label="$1" ok="$2"
  checks=$((checks + 1))
  if [ "${ok}" = "yes" ]; then
    echo "  ok: ${label}"
  else
    echo "  FAIL: ${label}"
    failures=$((failures + 1))
  fi
}

check_eq() {
  if [ "$2" = "$3" ]; then
    check "$1" "yes"
  else
    check "$1 (expected [$2], got [$3])" "no"
  fi
}

resolve_tool() {
  local tool="$1" override="$2"
  if [ -n "${override}" ] && [ -x "${override}" ]; then echo "${override}"; return 0; fi
  if [ -x "${lab_dir}/.venv/bin/${tool}" ]; then echo "${lab_dir}/.venv/bin/${tool}"; return 0; fi
  if command -v "${tool}" >/dev/null 2>&1; then command -v "${tool}"; return 0; fi
  return 1
}

pytest_bin="$(resolve_tool pytest "${PYTEST:-}")" || {
  echo "FAIL: pytest not found." >&2
  echo "  Install the lab's dependencies with:" >&2
  echo "    python3 -m venv .venv" >&2
  echo "    .venv/bin/pip install -r requirements/requirements.txt" >&2
  echo "  Or point this suite at an existing pytest:" >&2
  echo "    PYTEST=/path/to/pytest bash tests/run_tests.sh" >&2
  exit 1
}

python_bin="$(dirname "${pytest_bin}")/python3"
if [ ! -x "${python_bin}" ]; then
  python_bin="$(command -v python3 || true)"
fi
if [ -z "${python_bin}" ]; then
  echo "FAIL: python3 not found on PATH." >&2
  exit 1
fi

if ! "${python_bin}" -c "import numpy, PIL" >/dev/null 2>&1; then
  echo "FAIL: numpy and/or Pillow are not importable from ${python_bin}." >&2
  echo "  Install the lab's dependencies with:" >&2
  echo "    python3 -m venv .venv" >&2
  echo "    .venv/bin/pip install -r requirements/requirements.txt" >&2
  exit 1
fi

echo "Day 112 — Visualizing Optimization"
echo

# --------------------------------------------------------------------------
echo "1. The tools and the versions this lab was written against"
# --------------------------------------------------------------------------

versions="$("${python_bin}" - <<'PY'
import platform
import sys
from importlib.metadata import version

print(f"python   {platform.python_version()}")
for name in ("numpy", "Pillow", "pytest"):
    print(f"{name:<8} {version(name)}")
print(f"platform {platform.platform()}")
print(f"exe      {sys.executable.rsplit('/', 3)[-1]}")
PY
)"
echo "${versions}" | sed 's/^/  /'

pinned_numpy="$(grep -E '^numpy==' "${lab_dir}/requirements/requirements.txt" | cut -d= -f3)"
installed_numpy="$("${python_bin}" -c "from importlib.metadata import version; print(version('numpy'))")"
check_eq "installed numpy matches requirements.txt" "${pinned_numpy}" "${installed_numpy}"

pinned_pillow="$(grep -E '^Pillow==' "${lab_dir}/requirements/requirements.txt" | cut -d= -f3)"
installed_pillow="$("${python_bin}" -c "from importlib.metadata import version; print(version('Pillow'))")"
check_eq "installed Pillow matches requirements.txt" "${pinned_pillow}" "${installed_pillow}"

no_matplotlib="$("${python_bin}" -c "
try:
    import matplotlib  # noqa: F401
    print('present')
except ImportError:
    print('absent')
")"
check_eq "matplotlib is genuinely absent from this environment" "absent" "${no_matplotlib}"

# --------------------------------------------------------------------------
echo
echo "2. Every reference script runs and every assertion inside it holds"
# --------------------------------------------------------------------------

for script in 01_grid_and_ascii 02_heatmap_and_path 03_loss_curves \
              04_animated_gif 05_learning_rate_sweep 06_two_runs_same_loss; do
  out="$(cd "${lab_dir}/examples" && "${python_bin}" "${script}.py" 2>&1)"
  status=$?
  if [ "${status}" -ne 0 ]; then
    check "${script}.py exits 0" "no"
    echo "${out}" | tail -5 | sed 's/^/      /'
  else
    check "${script}.py exits 0" "yes"
  fi
  case "${out}" in
    *"${script}.py: every assertion held."*)
      check "${script}.py reports every assertion held" "yes" ;;
    *) check "${script}.py reports every assertion held" "no" ;;
  esac
done

# --------------------------------------------------------------------------
echo
echo "3. The reference pytest suite: real values, real exceptions"
# --------------------------------------------------------------------------

ref_out="$(cd "${lab_dir}" && "${pytest_bin}" examples -q -p no:cacheprovider 2>&1)"
ref_status=$?
echo "${ref_out}" | tail -3 | sed 's/^/  /'
if [ "${ref_status}" -eq 0 ]; then
  check "pytest examples exits 0" "yes"
else
  check "pytest examples exits 0" "no"
fi
case "${ref_out}" in
  *" failed"*) check "no test in the reference suite failed" "no" ;;
  *)           check "no test in the reference suite failed" "yes" ;;
esac
ref_passed="$(printf '%s\n' "${ref_out}" | grep -o '[0-9][0-9]* passed' | head -1 | cut -d' ' -f1)"
if [ "${ref_passed:-0}" -ge 10 ]; then
  check "the reference suite ran at least 10 tests (ran ${ref_passed:-0})" "yes"
else
  check "the reference suite ran at least 10 tests (ran ${ref_passed:-0})" "no"
fi

# --------------------------------------------------------------------------
echo
echo "4. The starter suite skips unattempted work instead of failing it"
# --------------------------------------------------------------------------

start_out="$(cd "${lab_dir}" && "${pytest_bin}" starter -q -p no:cacheprovider 2>&1)"
start_status=$?
echo "${start_out}" | tail -3 | sed 's/^/  /'
if [ "${start_status}" -eq 0 ]; then
  check "pytest starter exits 0 on an untouched checkout" "yes"
else
  check "pytest starter exits 0 on an untouched checkout" "no"
fi
case "${start_out}" in
  *" failed"*) check "the starter suite reports no failures" "no" ;;
  *)           check "the starter suite reports no failures" "yes" ;;
esac
case "${start_out}" in
  *skipped*) check "unwritten exercises are reported as skipped, not passed" "yes" ;;
  *) check "unwritten exercises are reported as skipped, not passed" "no" ;;
esac

# The import guard. Both directories contain modules called `dataset`,
# `gridviz`, `descent` and `imaging`, and collecting both suites at once
# would otherwise let the starter tests import the REFERENCE solution. Each
# directory's conftest.py prevents that; this check proves it still does.
both_out="$(cd "${lab_dir}" && "${pytest_bin}" -q -p no:cacheprovider 2>&1)"
start_skipped="$(printf '%s\n' "${start_out}" | grep -o '[0-9][0-9]* skipped' | head -1 | cut -d' ' -f1)"
both_skipped="$(printf '%s\n' "${both_out}" | grep -o '[0-9][0-9]* skipped' | head -1 | cut -d' ' -f1)"
check_eq "collecting both suites at once does not turn skips into passes" \
  "${start_skipped:-none}" "${both_skipped:-none}"

# --------------------------------------------------------------------------
echo
echo "5. The day's opening claim, checked one value at a time"
# --------------------------------------------------------------------------

# Section 6 re-runs this whole script with D112_SELF_TEST set, which asks
# for a threshold ten times tighter than the two runs actually achieve --
# proving the harness can fail rather than merely claiming it could.
threshold="${D112_SELF_TEST_THRESHOLD:-0.05}"

facts="$(cd "${lab_dir}/examples" && "${python_bin}" - <<PY
import dataset as D
import descent as DS

well_path = DS.gradient_descent(D.WELL_GRAD, D.START, D.LEARNING_RATE, D.STEPS)
ill_path = DS.gradient_descent(D.ILL_GRAD, D.START, D.LEARNING_RATE, D.STEPS)
well_loss = DS.losses_along(D.WELL_F, well_path)[-1]
ill_loss = DS.losses_along(D.ILL_F, ill_path)[-1]
relative_gap = abs(well_loss - ill_loss) / max(well_loss, ill_loss)
well_len = DS.path_length(well_path)
ill_len = DS.path_length(ill_path)

print("well_loss", f"{well_loss:.6e}")
print("ill_loss", f"{ill_loss:.6e}")
print("relative_gap", f"{relative_gap:.6f}")
print("relative_gap_within_threshold", relative_gap < ${threshold})
print("well_len", f"{well_len:.4f}")
print("ill_len", f"{ill_len:.4f}")
print("length_ratio", f"{ill_len / well_len:.4f}")
print("length_ratio_over_5x", (ill_len / well_len) > 5.0)
PY
)"

get() { printf '%s\n' "${facts}" | grep "^$1 " | cut -d' ' -f2-; }

echo "  (measured on this run: well-conditioned final loss $(get well_loss), ill-conditioned $(get ill_loss), relative gap $(get relative_gap) -- reported, not asserted to a value)"
check_eq "the two final losses land within the stated threshold of each other" \
  "True" "$(get relative_gap_within_threshold)"
echo "  (measured on this run: well-conditioned path length $(get well_len), ill-conditioned $(get ill_len), ratio $(get length_ratio)x)"
check_eq "the ill-conditioned path is over 5x longer than the well-conditioned one" \
  "True" "$(get length_ratio_over_5x)"

# --------------------------------------------------------------------------
echo
echo "6. The harness can actually fail"
# --------------------------------------------------------------------------

# A green test suite proves nothing until you have watched it go red. This
# section re-runs the whole script with the threshold in section 5 replaced
# by one the measured runs cannot meet, and asserts that the re-run reports
# the failure and exits non-zero. If this section passes, section 5 is not
# decorative.
if [ -z "${D112_SELF_TEST:-}" ]; then
  self_out="$(D112_SELF_TEST=1 D112_SELF_TEST_THRESHOLD=0.001 bash "${BASH_SOURCE[0]}" 2>&1)"
  self_status=$?
  if [ "${self_status}" -ne 0 ]; then
    check "an unmeetable threshold makes the harness exit non-zero (${self_status})" "yes"
  else
    check "an unmeetable threshold makes the harness exit non-zero" "no"
  fi
  case "${self_out}" in
    *"FAIL: the two final losses land within the stated threshold of each other"*)
      check "the failing check is named in the output" "yes" ;;
    *) check "the failing check is named in the output" "no" ;;
  esac
  case "${self_out}" in
    *", 1 failure(s)."*)
      check "the summary line counts exactly one failure" "yes" ;;
    *) check "the summary line counts exactly one failure" "no" ;;
  esac
else
  echo "  (self-test run: section 6 does not recurse)"
fi

# --------------------------------------------------------------------------
echo
echo "7. Nothing was left behind"
# --------------------------------------------------------------------------

if find "${lab_dir}" -name '.venv' -prune -o -type d -name '__pycache__' -print -quit 2>/dev/null | grep -q .; then
  check "no __pycache__ directory left by the lab's own code" "no"
else
  check "no __pycache__ directory left by the lab's own code" "yes"
fi

if find "${lab_dir}" -name '.venv' -prune -o -type d -name '.pytest_cache' -print -quit 2>/dev/null | grep -q .; then
  check "no .pytest_cache directory left under the lab" "no"
else
  check "no .pytest_cache directory left under the lab" "yes"
fi

if find "${lab_dir}" -name '.venv' -prune -o -type f \( -name '*.png' -o -name '*.gif' \) -print -quit 2>/dev/null | grep -q .; then
  check "no PNG or GIF file left anywhere in the lab (exercise 9)" "no"
else
  check "no PNG or GIF file left anywhere in the lab (exercise 9)" "yes"
fi

if grep -rqE 'urlopen|requests\.|socket\.|http://|https://' \
     "${lab_dir}/examples" "${lab_dir}/starter" 2>/dev/null; then
  check "no lab source opens a network connection" "no"
else
  check "no lab source opens a network connection" "yes"
fi

echo
echo "${checks} checks, ${failures} failure(s)."
[ "${failures}" -eq 0 ]