model.rs

//! A two-layer fully connected classifier, written out longhand.
//!
//! Shapes, once, so the rest reads easily. A batch of `B` samples is a matrix
//! `X` of shape `(pixels, B)` — one sample per *column*, which is the layout
//! that makes every step below a plain matrix product:
//!
//! ```text
//!   Z1 = W1 X + b1     (hidden, B)     W1: (hidden, pixels)
//!   A1 = phi(Z1)       (hidden, B)
//!   Z2 = W2 A1 + b2    (classes, B)    W2: (classes, hidden)
//!   A2 = softmax(Z2)   (classes, B)    column-wise
//! ```
//!
//! and backwards, with `Y` the one-hot targets and the `1/B` folded into the
//! first term so it propagates for free:
//!
//! ```text
//!   dZ2 = (A2 - Y)/B   dW2 = dZ2 A1^T   db2 = row sums of dZ2
//!   dZ1 = W2^T dZ2 * phi'(A1)
//!                      dW1 = dZ1 X^T    db1 = row sums of dZ1
//! ```
//!
//! `dZ2 = A2 - Y` is the softmax and cross-entropy Jacobians cancelling each
//! other, which is the entire reason those two are always paired.

use crate::dataloader::MnistData;
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use ndarray::{Array2, ArrayView1, Axis};
use rand::seq::SliceRandom;
use rand::{SeedableRng, rngs::StdRng};
use rand_distr::{Distribution, Normal};
use std::fs::File;
use std::io::{BufReader, BufWriter};
use std::path::Path;
use thiserror::Error;
use tracing::{debug, info};

const MODEL_MAGIC: &[u8; 4] = b"DRNN";
const MODEL_VERSION: u32 = 1;

/// Upper bound on a layer width read out of a model file, so a corrupt header
/// cannot be turned into an allocation request.
const MAX_LAYER: usize = 1 << 24;

#[derive(Debug, Error)]
pub enum ModelError {
    #[error("io error on {path}: {source}")]
    Io {
        path: String,
        #[source]
        source: std::io::Error,
    },
    #[error("{0} is not a digit-recognizer model file")]
    BadMagic(String),
    #[error("{path} is version {found}, this build reads version {MODEL_VERSION}")]
    BadVersion { path: String, found: u32 },
    #[error("{path} names an activation this build does not know (tag {tag})")]
    BadActivation { path: String, tag: u8 },
    #[error("model takes {expected}-pixel inputs, the data has {found}")]
    InputMismatch { expected: usize, found: usize },
    #[error("layer sizes must all be non-zero, got {pixels}x{hidden}x{classes}")]
    EmptyLayer {
        pixels: usize,
        hidden: usize,
        classes: usize,
    },
}

/// Nonlinearity for the hidden layer.
#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)]
pub enum Activation {
    /// `max(0, z)`. The default, and what the initialisation is tuned for.
    Relu,
    /// `1 / (1 + e^-z)`. Slower to train, kept because it is the textbook one.
    Sigmoid,
}

impl Activation {
    fn apply(self, z: &Array2<f64>) -> Array2<f64> {
        match self {
            Activation::Relu => z.mapv(|v| v.max(0.0)),
            Activation::Sigmoid => z.mapv(|v| 1.0 / (1.0 + (-v).exp())),
        }
    }

    /// Derivative with respect to the pre-activation, written in terms of the
    /// activation itself so the forward pass never has to keep `Z1` alive.
    fn derivative(self, a: &Array2<f64>) -> Array2<f64> {
        match self {
            Activation::Relu => a.mapv(|v| if v > 0.0 { 1.0 } else { 0.0 }),
            Activation::Sigmoid => a.mapv(|v| v * (1.0 - v)),
        }
    }

    /// Standard deviation for a layer's initial weights, given its fan-in.
    ///
    /// He (`2/n`) for ReLU, Xavier (`1/n`) for sigmoid. These are not
    /// interchangeable: He assumes the nonlinearity throws away half the
    /// variance, which sigmoid does not do, so pairing He with sigmoid — as
    /// this code used to — starts the hidden units too far out on the flat
    /// parts of the curve, where the gradient is nearly zero.
    fn init_std(self, fan_in: usize) -> f64 {
        let numerator = match self {
            Activation::Relu => 2.0,
            Activation::Sigmoid => 1.0,
        };
        (numerator / fan_in as f64).sqrt()
    }

    fn tag(self) -> u8 {
        match self {
            Activation::Relu => 0,
            Activation::Sigmoid => 1,
        }
    }

    fn from_tag(tag: u8) -> Option<Self> {
        match tag {
            0 => Some(Activation::Relu),
            1 => Some(Activation::Sigmoid),
            _ => None,
        }
    }
}

pub struct TrainConfig {
    pub learning_rate: f64,
    pub epochs: usize,
    pub batch_size: usize,
}

/// One epoch's worth of numbers, so the caller can decide what to print.
#[derive(Debug, Clone, Copy)]
pub struct EpochReport {
    pub epoch: usize,
    pub train_loss: f64,
    pub validation_accuracy: Option<f64>,
}

struct Gradients {
    dw1: Array2<f64>,
    db1: Array2<f64>,
    dw2: Array2<f64>,
    db2: Array2<f64>,
}

pub struct NeuralNetwork {
    w1: Array2<f64>,
    b1: Array2<f64>,
    w2: Array2<f64>,
    b2: Array2<f64>,
    activation: Activation,
    rng: StdRng,
}

impl NeuralNetwork {
    pub fn new(
        pixels: usize,
        hidden: usize,
        classes: usize,
        activation: Activation,
        seed: u64,
    ) -> Result<Self, ModelError> {
        if pixels == 0 || hidden == 0 || classes == 0 {
            return Err(ModelError::EmptyLayer {
                pixels,
                hidden,
                classes,
            });
        }

        let mut rng = StdRng::seed_from_u64(seed);
        let w1 = gaussian(&mut rng, (hidden, pixels), activation.init_std(pixels));
        // The output layer feeds a softmax, not the hidden nonlinearity, so it
        // gets Xavier regardless of what the hidden layer uses.
        let w2 = gaussian(&mut rng, (classes, hidden), (1.0 / hidden as f64).sqrt());

        Ok(Self {
            w1,
            b1: Array2::zeros((hidden, 1)),
            w2,
            b2: Array2::zeros((classes, 1)),
            activation,
            rng,
        })
    }

    pub fn pixels(&self) -> usize {
        self.w1.ncols()
    }

    pub fn hidden(&self) -> usize {
        self.w1.nrows()
    }

    pub fn classes(&self) -> usize {
        self.w2.nrows()
    }

    pub fn activation(&self) -> Activation {
        self.activation
    }

    /// Checks that a data set is the shape this network was built for.
    ///
    /// Worth calling before [`Self::evaluate`] on anything that came out of
    /// [`Self::load`]: a model trained on some other input size would
    /// otherwise get as far as assembling a batch and panic on the shape
    /// mismatch, which is a poor way to say "wrong model file".
    pub fn check_input(&self, data: &MnistData) -> Result<(), ModelError> {
        if data.pixels_per_sample() == self.pixels() {
            Ok(())
        } else {
            Err(ModelError::InputMismatch {
                expected: self.pixels(),
                found: data.pixels_per_sample(),
            })
        }
    }

    /// Returns `(A1, A2)`; `Z1` is not needed downstream, see
    /// [`Activation::derivative`].
    fn forward(&self, x: &Array2<f64>) -> (Array2<f64>, Array2<f64>) {
        let a1 = self.activation.apply(&(self.w1.dot(x) + &self.b1));
        let a2 = softmax_columns(self.w2.dot(&a1) + &self.b2);
        (a1, a2)
    }

    fn backward(
        &self,
        x: &Array2<f64>,
        a1: &Array2<f64>,
        a2: &Array2<f64>,
        y: &Array2<f64>,
    ) -> Gradients {
        let batch = x.ncols() as f64;

        let dz2 = (a2 - y) / batch;
        let dw2 = dz2.dot(&a1.t());
        let db2 = dz2.sum_axis(Axis(1)).insert_axis(Axis(1));

        let dz1 = self.w2.t().dot(&dz2) * self.activation.derivative(a1);
        let dw1 = dz1.dot(&x.t());
        let db1 = dz1.sum_axis(Axis(1)).insert_axis(Axis(1));

        Gradients { dw1, db1, dw2, db2 }
    }

    fn descend(&mut self, g: &Gradients, learning_rate: f64) {
        self.w1.scaled_add(-learning_rate, &g.dw1);
        self.b1.scaled_add(-learning_rate, &g.db1);
        self.w2.scaled_add(-learning_rate, &g.dw2);
        self.b2.scaled_add(-learning_rate, &g.db2);
    }

    /// Trains for `cfg.epochs` passes, reshuffling between each.
    ///
    /// Reshuffling matters more than it looks: without it, plain SGD sees the
    /// same sample sequence every epoch and happily learns the order along
    /// with the digits.
    pub fn train(
        &mut self,
        data: &MnistData,
        validation: Option<&MnistData>,
        cfg: &TrainConfig,
    ) -> Result<Vec<EpochReport>, ModelError> {
        self.check_input(data)?;
        if let Some(set) = validation {
            self.check_input(set)?;
        }

        let batch_size = cfg.batch_size.max(1);
        let mut order: Vec<usize> = (0..data.len()).collect();
        let mut reports = Vec::with_capacity(cfg.epochs);

        for epoch in 1..=cfg.epochs {
            order.shuffle(&mut self.rng);

            let mut total_loss = 0.0;
            for (step, batch) in order.chunks(batch_size).enumerate() {
                let (x, y) = self.assemble_batch(data, batch);

                let (a1, a2) = self.forward(&x);
                // Measured before the update, on the same forward pass that
                // produces the gradient. The old code ran a *second* forward
                // pass after descending, which cost 50% more compute and
                // reported a loss the model had already been corrected for.
                total_loss += cross_entropy(&a2, &y) * batch.len() as f64;

                let gradients = self.backward(&x, &a1, &a2, &y);
                self.descend(&gradients, cfg.learning_rate);

                if (step + 1) % 200 == 0 {
                    debug!(
                        "epoch {}: {}/{} samples, running loss {:.4}",
                        epoch,
                        (step + 1) * batch_size,
                        data.len(),
                        total_loss / ((step + 1) * batch_size) as f64
                    );
                }
            }

            let report = EpochReport {
                epoch,
                train_loss: total_loss / data.len() as f64,
                validation_accuracy: validation.map(|set| self.evaluate(set).1),
            };
            match report.validation_accuracy {
                Some(accuracy) => info!(
                    "epoch {}/{}: loss {:.4}, validation accuracy {:.2}%",
                    epoch,
                    cfg.epochs,
                    report.train_loss,
                    accuracy * 100.0
                ),
                None => info!(
                    "epoch {}/{}: loss {:.4}",
                    epoch, cfg.epochs, report.train_loss
                ),
            }
            reports.push(report);
        }

        Ok(reports)
    }

    /// Gathers the given samples into a `(pixels, B)` input matrix and a
    /// `(classes, B)` one-hot target matrix.
    fn assemble_batch(&self, data: &MnistData, indices: &[usize]) -> (Array2<f64>, Array2<f64>) {
        let mut x = Array2::zeros((self.pixels(), indices.len()));
        let mut y = Array2::zeros((self.classes(), indices.len()));
        for (column, &index) in indices.iter().enumerate() {
            x.column_mut(column).assign(&data.sample(index));
            y[[data.label(index) as usize, column]] = 1.0;
        }
        (x, y)
    }

    /// Classifies a single sample.
    ///
    /// The binary goes through [`Self::evaluate`], which batches; this is kept
    /// because it is the shape of the call anyone reading the module reaches
    /// for first, and the tests hold it to the same answers.
    #[cfg_attr(not(test), allow(dead_code))]
    pub fn predict(&self, sample: ArrayView1<'_, f64>) -> u8 {
        let x = sample
            .to_owned()
            .into_shape_with_order((sample.len(), 1))
            .expect("a length-n vector is an n-by-1 matrix");
        let (_, a2) = self.forward(&x);
        argmax(a2.column(0).iter().copied()) as u8
    }

    /// Classifies a whole set, in batches. Returns the predictions and the
    /// fraction that were right.
    pub fn evaluate(&self, data: &MnistData) -> (Vec<u8>, f64) {
        const EVAL_BATCH: usize = 256;

        let mut predictions = Vec::with_capacity(data.len());
        let indices: Vec<usize> = (0..data.len()).collect();
        for batch in indices.chunks(EVAL_BATCH) {
            let (x, _) = self.assemble_batch(data, batch);
            let (_, a2) = self.forward(&x);
            for column in a2.columns() {
                predictions.push(argmax(column.iter().copied()) as u8);
            }
        }

        let correct = predictions
            .iter()
            .zip(data.labels())
            .filter(|(p, t)| p == t)
            .count();
        let accuracy = if data.is_empty() {
            0.0
        } else {
            correct as f64 / data.len() as f64
        };

        (predictions, accuracy)
    }

    /// Writes the weights out, so `--skip-training` has something to load.
    ///
    /// Deliberately a plain little-endian dump rather than a serialisation
    /// crate: the whole file is a header and four matrices, and being able to
    /// describe the format in one comment is worth more here than being able
    /// to evolve it.
    pub fn save(&self, path: &Path) -> Result<(), ModelError> {
        use std::io::Write;

        let file = File::create(path).map_err(|e| io_error(path, e))?;
        let mut out = BufWriter::new(file);

        let mut header = || -> std::io::Result<()> {
            out.write_all(MODEL_MAGIC)?;
            out.write_u32::<LittleEndian>(MODEL_VERSION)?;
            out.write_u8(self.activation.tag())?;
            out.write_u32::<LittleEndian>(self.pixels() as u32)?;
            out.write_u32::<LittleEndian>(self.hidden() as u32)?;
            out.write_u32::<LittleEndian>(self.classes() as u32)?;
            for matrix in [&self.w1, &self.b1, &self.w2, &self.b2] {
                for &value in matrix.iter() {
                    out.write_f64::<LittleEndian>(value)?;
                }
            }
            out.flush()
        };

        header().map_err(|e| io_error(path, e))
    }

    pub fn load(path: &Path) -> Result<Self, ModelError> {
        use std::io::Read;

        let file = File::open(path).map_err(|e| io_error(path, e))?;
        let mut input = BufReader::new(file);

        let mut magic = [0u8; 4];
        input
            .read_exact(&mut magic)
            .map_err(|e| io_error(path, e))?;
        if &magic != MODEL_MAGIC {
            return Err(ModelError::BadMagic(path.display().to_string()));
        }

        let version = input
            .read_u32::<LittleEndian>()
            .map_err(|e| io_error(path, e))?;
        if version != MODEL_VERSION {
            return Err(ModelError::BadVersion {
                path: path.display().to_string(),
                found: version,
            });
        }

        let tag = input.read_u8().map_err(|e| io_error(path, e))?;
        let activation = Activation::from_tag(tag).ok_or_else(|| ModelError::BadActivation {
            path: path.display().to_string(),
            tag,
        })?;

        let mut dimension = || {
            input
                .read_u32::<LittleEndian>()
                .map(|v| v as usize)
                .map_err(|e| io_error(path, e))
        };
        let pixels = dimension()?;
        let hidden = dimension()?;
        let classes = dimension()?;
        // Bound the dimensions before they are used as allocation sizes: a
        // corrupt file claiming 2^32-1 by 2^32-1 would otherwise reach
        // `Vec::with_capacity` and abort on capacity overflow instead of
        // returning an error anyone can act on.
        let plausible = |value: usize| value > 0 && value <= MAX_LAYER;
        if !plausible(pixels) || !plausible(hidden) || !plausible(classes) {
            return Err(ModelError::EmptyLayer {
                pixels,
                hidden,
                classes,
            });
        }

        let mut read_matrix =
            |rows: usize, cols: usize| -> Result<Array2<f64>, ModelError> {
                let mut values = Vec::with_capacity(rows * cols);
                for _ in 0..rows * cols {
                    values.push(
                        input
                            .read_f64::<LittleEndian>()
                            .map_err(|e| io_error(path, e))?,
                    );
                }
                Ok(Array2::from_shape_vec((rows, cols), values)
                    .expect("read exactly rows*cols values"))
            };

        Ok(Self {
            w1: read_matrix(hidden, pixels)?,
            b1: read_matrix(hidden, 1)?,
            w2: read_matrix(classes, hidden)?,
            b2: read_matrix(classes, 1)?,
            activation,
            // A loaded model carries no training history to resume, and the
            // shuffle order is not part of the weights, so any seed will do.
            rng: StdRng::seed_from_u64(0),
        })
    }
}

fn io_error(path: &Path, source: std::io::Error) -> ModelError {
    ModelError::Io {
        path: path.display().to_string(),
        source,
    }
}

/// Softmax, one column at a time.
///
/// The max subtraction is what keeps `exp` from overflowing. Doing it per
/// column also matters once batches are wider than one: a single global max
/// over the whole matrix still normalises each column to sum to one, but it
/// couples samples in a batch to each other's largest logit.
fn softmax_columns(mut z: Array2<f64>) -> Array2<f64> {
    for mut column in z.axis_iter_mut(Axis(1)) {
        let max = column.fold(f64::NEG_INFINITY, |acc, &v| acc.max(v));
        column.mapv_inplace(|v| (v - max).exp());
        let total = column.sum();
        column.mapv_inplace(|v| v / total);
    }
    z
}

/// Mean per-sample cross entropy, `-(1/B) sum_b sum_i y_i log p_i`.
///
/// The old version divided by the number of *classes* instead of the number of
/// samples, which quietly reported every loss ten times smaller than it was.
fn cross_entropy(probabilities: &Array2<f64>, targets: &Array2<f64>) -> f64 {
    const FLOOR: f64 = 1e-12;

    let batch = probabilities.ncols().max(1) as f64;
    let total: f64 = targets
        .iter()
        .zip(probabilities.iter())
        .filter(|(target, _)| **target != 0.0)
        .map(|(&target, &probability)| target * probability.max(FLOOR).ln())
        .sum();

    -total / batch
}

fn gaussian(rng: &mut StdRng, shape: (usize, usize), std: f64) -> Array2<f64> {
    let dist = Normal::new(0.0, std).expect("std is finite and positive");
    Array2::from_shape_fn(shape, |_| dist.sample(rng))
}

fn argmax(values: impl Iterator<Item = f64>) -> usize {
    values
        .enumerate()
        .fold((0usize, f64::NEG_INFINITY), |(best_i, best_v), (i, v)| {
            if v > best_v { (i, v) } else { (best_i, best_v) }
        })
        .0
}

#[cfg(test)]
mod tests {
    use super::*;

    fn toy_batch() -> (Array2<f64>, Array2<f64>) {
        let x = Array2::from_shape_vec(
            (4, 3),
            vec![
                0.9, -0.4, 0.2, //
                0.1, 0.8, -0.7, //
                -0.5, 0.3, 0.6, //
                0.4, 0.25, -0.15,
            ],
        )
        .unwrap();
        let mut y = Array2::zeros((2, 3));
        y[[0, 0]] = 1.0;
        y[[1, 1]] = 1.0;
        y[[0, 2]] = 1.0;
        (x, y)
    }

    /// The test that actually keeps the backprop honest: compare every
    /// analytic gradient against a central difference of the loss.
    fn gradient_check(activation: Activation) {
        const EPS: f64 = 1e-6;
        const TOLERANCE: f64 = 1e-7;

        let mut net = NeuralNetwork::new(4, 5, 2, activation, 7).unwrap();
        let (x, y) = toy_batch();

        let (a1, a2) = net.forward(&x);
        let analytic = net.backward(&x, &a1, &a2, &y);

        // Perturb one weight at a time and watch the loss move.
        let probes: Vec<(&str, Vec<(usize, usize)>)> = vec![
            ("w1", vec![(0, 0), (2, 3), (4, 1)]),
            ("b1", vec![(0, 0), (3, 0)]),
            ("w2", vec![(0, 0), (1, 4)]),
            ("b2", vec![(0, 0), (1, 0)]),
        ];

        for (name, cells) in probes {
            for cell in cells {
                let loss_at = |net: &mut NeuralNetwork, delta: f64| {
                    let slot = match name {
                        "w1" => &mut net.w1,
                        "b1" => &mut net.b1,
                        "w2" => &mut net.w2,
                        _ => &mut net.b2,
                    };
                    slot[cell] += delta;
                    let (_, a2) = net.forward(&x);
                    let loss = cross_entropy(&a2, &y);
                    let slot = match name {
                        "w1" => &mut net.w1,
                        "b1" => &mut net.b1,
                        "w2" => &mut net.w2,
                        _ => &mut net.b2,
                    };
                    slot[cell] -= delta;
                    loss
                };

                let numeric = (loss_at(&mut net, EPS) - loss_at(&mut net, -EPS)) / (2.0 * EPS);
                let exact = match name {
                    "w1" => analytic.dw1[cell],
                    "b1" => analytic.db1[cell],
                    "w2" => analytic.dw2[cell],
                    _ => analytic.db2[cell],
                };

                let scale = numeric.abs().max(exact.abs()).max(1e-8);
                assert!(
                    (numeric - exact).abs() / scale < TOLERANCE,
                    "{activation:?} {name}{cell:?}: analytic {exact:e}, numeric {numeric:e}"
                );
            }
        }
    }

    #[test]
    fn backprop_matches_finite_differences_for_relu() {
        gradient_check(Activation::Relu);
    }

    #[test]
    fn backprop_matches_finite_differences_for_sigmoid() {
        gradient_check(Activation::Sigmoid);
    }

    #[test]
    fn softmax_normalises_each_column_independently() {
        let z = Array2::from_shape_vec((3, 2), vec![1.0, 40.0, 2.0, 41.0, 3.0, 42.0]).unwrap();
        let p = softmax_columns(z);

        for column in p.columns() {
            assert!((column.sum() - 1.0).abs() < 1e-12, "{column:?}");
        }
        // Both columns are the same logits plus a constant, so the softmax
        // must come out identical. A global max subtraction gets this right
        // too, but a global *sum* would not, and it is cheap to pin down.
        for row in 0..3 {
            assert!((p[[row, 0]] - p[[row, 1]]).abs() < 1e-12);
        }
    }

    #[test]
    fn softmax_survives_logits_that_would_overflow_exp() {
        let z = Array2::from_shape_vec((2, 1), vec![1000.0, 999.0]).unwrap();
        let p = softmax_columns(z);
        assert!(p.iter().all(|v| v.is_finite()), "{p:?}");
        assert!((p.sum() - 1.0).abs() < 1e-12);
    }

    #[test]
    fn cross_entropy_is_per_sample_not_per_class() {
        // Two samples, each with all the mass on the right class. Loss is 0.
        let mut p = Array2::from_elem((10, 2), 0.0);
        let mut y = Array2::zeros((10, 2));
        p[[3, 0]] = 1.0;
        y[[3, 0]] = 1.0;
        p[[7, 1]] = 1.0;
        y[[7, 1]] = 1.0;
        assert!(cross_entropy(&p, &y).abs() < 1e-12);

        // A uniform prediction over 10 classes costs ln(10) per sample. The
        // old implementation divided by the class count and reported ln(10)/10.
        let uniform = Array2::from_elem((10, 2), 0.1);
        let loss = cross_entropy(&uniform, &y);
        assert!(
            (loss - 10f64.ln()).abs() < 1e-12,
            "expected ln(10) = {}, got {loss}",
            10f64.ln()
        );
    }

    #[test]
    fn he_and_xavier_are_not_the_same_number() {
        // He is sqrt(2/n), Xavier sqrt(1/n), so ReLU starts a factor of
        // sqrt(2) wider than sigmoid at the same fan-in.
        assert!((Activation::Sigmoid.init_std(100) - 0.1).abs() < 1e-12);
        assert!((Activation::Relu.init_std(100) - (2.0f64 / 100.0).sqrt()).abs() < 1e-12);
        assert!(
            (Activation::Relu.init_std(784) - Activation::Sigmoid.init_std(784) * 2.0f64.sqrt())
                .abs()
                < 1e-12
        );
    }

    #[test]
    fn argmax_takes_the_first_of_a_tie() {
        assert_eq!(argmax([0.1, 0.5, 0.5, 0.2].into_iter()), 1);
        assert_eq!(
            argmax([f64::NEG_INFINITY, f64::NEG_INFINITY].into_iter()),
            0
        );
    }

    #[test]
    fn one_hot_targets_land_on_the_right_row() {
        let net = NeuralNetwork::new(4, 3, 10, Activation::Relu, 1).unwrap();
        let data = crate::dataloader::tests_support::synthetic(&[2, 9, 0], 4);
        let (x, y) = net.assemble_batch(&data, &[0, 1, 2]);

        assert_eq!(x.dim(), (4, 3));
        assert_eq!(y.dim(), (10, 3));
        for (column, &label) in [2usize, 9, 0].iter().enumerate() {
            assert_eq!(y[[label, column]], 1.0);
            assert_eq!(y.column(column).sum(), 1.0);
        }
    }

    #[test]
    fn a_few_steps_of_descent_reduce_the_loss() {
        let mut net = NeuralNetwork::new(4, 8, 2, Activation::Relu, 3).unwrap();
        let (x, y) = toy_batch();

        let (_, a2) = net.forward(&x);
        let before = cross_entropy(&a2, &y);

        for _ in 0..200 {
            let (a1, a2) = net.forward(&x);
            let g = net.backward(&x, &a1, &a2, &y);
            net.descend(&g, 0.5);
        }

        let (_, a2) = net.forward(&x);
        let after = cross_entropy(&a2, &y);
        assert!(after < before * 0.5, "loss went {before} -> {after}");
    }

    #[test]
    fn training_learns_a_separable_toy_problem() {
        // Two pixels, two classes: bright-left is class 0, bright-right is 1.
        let labels: Vec<u8> = (0..64).map(|i| (i % 2) as u8).collect();
        let data = crate::dataloader::tests_support::two_pixel_problem(&labels);

        let mut net = NeuralNetwork::new(2, 8, 2, Activation::Relu, 11).unwrap();
        let reports = net
            .train(
                &data,
                Some(&data),
                &TrainConfig {
                    learning_rate: 0.5,
                    epochs: 30,
                    batch_size: 8,
                },
            )
            .unwrap();

        assert_eq!(reports.len(), 30);
        assert!(
            reports.last().unwrap().train_loss < reports[0].train_loss,
            "{:?}",
            reports
        );
        assert_eq!(net.evaluate(&data).1, 1.0);
    }

    #[test]
    fn predict_agrees_with_the_batched_evaluate() {
        let labels: Vec<u8> = (0..17).map(|i| (i % 2) as u8).collect();
        let data = crate::dataloader::tests_support::two_pixel_problem(&labels);

        let mut net = NeuralNetwork::new(2, 6, 2, Activation::Relu, 2).unwrap();
        net.train(
            &data,
            None,
            &TrainConfig {
                learning_rate: 0.5,
                epochs: 10,
                batch_size: 4,
            },
        )
        .unwrap();

        // 17 samples against a batch of 256 exercises the ragged last chunk.
        let (batched, _) = net.evaluate(&data);
        let one_at_a_time: Vec<u8> = (0..data.len())
            .map(|i| net.predict(data.sample(i)))
            .collect();
        assert_eq!(batched, one_at_a_time);
    }

    #[test]
    fn train_rejects_data_of_the_wrong_width() {
        let mut net = NeuralNetwork::new(784, 4, 10, Activation::Relu, 1).unwrap();
        let data = crate::dataloader::tests_support::synthetic(&[1], 4);
        let err = net
            .train(
                &data,
                None,
                &TrainConfig {
                    learning_rate: 0.1,
                    epochs: 1,
                    batch_size: 1,
                },
            )
            .unwrap_err();
        assert!(matches!(
            err,
            ModelError::InputMismatch {
                expected: 784,
                found: 4
            }
        ));
    }

    #[test]
    fn check_input_catches_a_model_built_for_another_data_set() {
        let net = NeuralNetwork::new(784, 4, 10, Activation::Relu, 1).unwrap();
        let data = crate::dataloader::tests_support::synthetic(&[1, 2], 4);
        assert!(matches!(
            net.check_input(&data),
            Err(ModelError::InputMismatch {
                expected: 784,
                found: 4
            })
        ));

        let matching = crate::dataloader::tests_support::synthetic(&[1], 784);
        assert!(net.check_input(&matching).is_ok());
    }

    #[test]
    fn new_rejects_empty_layers() {
        assert!(matches!(
            NeuralNetwork::new(784, 0, 10, Activation::Relu, 1),
            Err(ModelError::EmptyLayer { .. })
        ));
    }

    #[test]
    fn a_saved_model_predicts_exactly_as_it_did_before() {
        let dir = std::env::temp_dir().join("digit-recognizer-tests-model");
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("net.bin");

        // Train a little first, so the biases are non-zero and the round trip
        // has something other than the initialiser's output to preserve.
        let mut net = NeuralNetwork::new(4, 6, 2, Activation::Sigmoid, 5).unwrap();
        let (x, y) = toy_batch();
        for _ in 0..20 {
            let (a1, a2) = net.forward(&x);
            let g = net.backward(&x, &a1, &a2, &y);
            net.descend(&g, 0.3);
        }

        net.save(&path).unwrap();
        let loaded = NeuralNetwork::load(&path).unwrap();

        assert_eq!(loaded.activation(), Activation::Sigmoid);
        assert_eq!(
            (loaded.pixels(), loaded.hidden(), loaded.classes()),
            (4, 6, 2)
        );
        assert_eq!(loaded.w1, net.w1);
        assert_eq!(loaded.b1, net.b1);
        assert_eq!(loaded.w2, net.w2);
        assert_eq!(loaded.b2, net.b2);

        // Bit-identical predictions, not merely close ones.
        let (_, expected) = net.forward(&x);
        let (_, actual) = loaded.forward(&x);
        assert_eq!(expected, actual);
    }

    #[test]
    fn load_rejects_a_file_that_is_not_a_model() {
        let dir = std::env::temp_dir().join("digit-recognizer-tests-model");
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("garbage.bin");
        std::fs::write(&path, b"not a model at all").unwrap();
        assert!(matches!(
            NeuralNetwork::load(&path),
            Err(ModelError::BadMagic(_))
        ));
    }
}