train.py

#!/usr/bin/env python3
"""Training program: run gradient descent over the dataset and save the thetas.

Usage:
    python3 train.py [--data data.csv] [--learning-rate 0.5] [--iterations 10000]
"""

from __future__ import annotations

import argparse
import math
import sys

from linear_regression import (
    DEFAULT_ITERS,
    DEFAULT_LR,
    THETA_FILE,
    load_dataset,
    metrics,
    save_model,
    train,
    train_fast,
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--data", default="data.csv", help="dataset CSV (km,price)")
    parser.add_argument("--learning-rate", type=float, default=DEFAULT_LR)
    parser.add_argument("--iterations", type=int, default=DEFAULT_ITERS)
    parser.add_argument("--out", default=THETA_FILE, help="where to write thetas")
    parser.add_argument(
        "--fast",
        action="store_true",
        help="O(iterations) moment-based trainer (identical result, for big datasets)",
    )
    args = parser.parse_args(argv)

    if args.learning_rate <= 0 or args.iterations <= 0:
        print("train: learning-rate and iterations must be positive", file=sys.stderr)
        return 1
    try:
        xs, ys = load_dataset(args.data)
    except (OSError, ValueError) as exc:
        print(f"train: cannot read {args.data}: {exc}", file=sys.stderr)
        return 1

    trainer = train_fast if args.fast else train
    model, history = trainer(xs, ys, args.learning_rate, args.iterations)
    # Divergence shows up as a cost that is non-finite or larger than it started;
    # a huge-but-finite theta would otherwise overflow downstream metrics.
    if not history or not math.isfinite(history[-1]) or history[-1] > history[0]:
        print("train: diverged — try a smaller --learning-rate", file=sys.stderr)
        return 1
    save_model(model, args.out)
    mt = metrics(xs, ys, model)
    print(f"Trained on {len(xs)} points in {len(history)} iterations.")
    print(f"  theta0 = {model.theta0:.6f}")
    print(f"  theta1 = {model.theta1:.8f}")
    print(f"  R^2 = {mt['r2']:.4f}   RMSE = {mt['rmse']:.2f}")
    print(f"Saved thetas to {args.out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())