Forecast next week’s sales

"""
Forecast weekly sales with a linear model on calendar features.

Rung N1. The moving average of N0 has no notion of a trend: it can only
repeat the recent past, which is exactly what it gets wrong when the business
is growing or shrinking. Here the level, the slope and the seasonal shape are
estimated together, by least squares, over the whole history.

The features are calendar arithmetic and nothing else: a constant, the number
of cycles elapsed, and a few sine and cosine pairs whose period is the season.
Two pairs are enough to place a Christmas peak and a summer dip; a peak only
a few weeks wide needs more pairs to reach its full height.

Counting the trend in cycles rather than in weeks is not cosmetic. It keeps
the columns of the design matrix on comparable scales, and it makes the
coefficient readable on its own: it is the growth per cycle, per year with
the default 52-week season. Less than one full cycle of history cannot tell
the trend from the season, so it is refused.

`season_length` need not be a whole number: this is trigonometry, nothing is
indexed by week. A year is 365.25 / 7 = 52.18 weeks, and the ISO calendar has
a fifty-third week every five or six years, so over three years a Christmas
peak drifts by about a week. Pass 365.25 / 7 to hold it in place — which is
what Forecasting: Principles and Practice recommends for weekly data, and
which N0, indexing by `week % season_length`, cannot do.
"""

import math

import numpy as np


def calendar_features(week: int, season_length: int, harmonics: int) -> list[float]:
    """The row of the design matrix for one week. This is the whole model."""
    row = [1.0, week / season_length]
    for k in range(1, harmonics + 1):
        angle = 2 * math.pi * k * week / season_length
        row += [math.sin(angle), math.cos(angle)]
    return row


def fit(history: list[float], season_length: int = 52, harmonics: int = 2) -> dict:
    """
    Least squares over the whole history. Returns a model you can inspect.

    `coefficients[1]` is the growth per cycle, in the unit of the series. That
    number is worth reading before any forecast is: a model that has found a
    trend nobody in the company recognises is a model to distrust.
    """
    if len(history) < 2 + 2 * harmonics:
        raise ValueError("fewer weeks of history than features to estimate")
    if len(history) < season_length:
        raise ValueError("a trend cannot be told apart from the season on less than one full cycle")
    sales = np.array(history, dtype=float)  # a missing week (None) becomes NaN here
    if not np.all(np.isfinite(sales)):
        raise ValueError("every week of history needs a finite number")
    design = [calendar_features(week, season_length, harmonics) for week in range(len(history))]
    coefficients, *_ = np.linalg.lstsq(np.array(design), sales, rcond=None)
    return {
        "coefficients": [float(c) for c in coefficients],
        "season_length": season_length,
        "harmonics": harmonics,
        "start": len(history),
    }


def forecast(model: dict, horizon: int = 1) -> list[float]:
    """Predict the `horizon` weeks that follow the history the model was fitted on."""
    weeks = []
    for step in range(horizon):
        row = calendar_features(model["start"] + step, model["season_length"], model["harmonics"])
        weeks.append(sum(c * x for c, x in zip(model["coefficients"], row)))
    return weeks

/**
 * Forecast weekly sales with a linear model on calendar features.
 *
 * Rung N1. The moving average of N0 has no notion of a trend: it can only
 * repeat the recent past, which is exactly what it gets wrong when the
 * business is growing or shrinking. Here the level, the slope and the
 * seasonal shape are estimated together, by least squares, over the whole
 * history.
 *
 * The features are calendar arithmetic and nothing else: a constant, the
 * number of cycles elapsed, and a few sine and cosine pairs whose period is
 * the season.
 *
 * Least squares is solved through the normal equations and a Gaussian
 * elimination, short enough that no dependency is worth it. Counting the
 * trend in cycles rather than in weeks keeps those equations well behaved,
 * and makes the coefficient readable on its own: it is the growth per cycle,
 * per year with the default 52-week season. Less than one full cycle of
 * history cannot tell the trend from the season, so it is refused.
 *
 * `seasonLength` need not be a whole number: this is trigonometry, nothing is
 * indexed by week. A year is 365.25 / 7 = 52.18 weeks, and the ISO calendar has
 * a fifty-third week every five or six years, so over three years a Christmas
 * peak drifts by about a week. Pass 365.25 / 7 to hold it in place — which is
 * what Forecasting: Principles and Practice recommends for weekly data, and
 * which N0, indexing by `week % seasonLength`, cannot do.
 */

/** The row of the design matrix for one week. This is the whole model. */
export function calendarFeatures(week, seasonLength, harmonics) {
  const row = [1, week / seasonLength];
  for (let k = 1; k <= harmonics; k += 1) {
    const angle = (2 * Math.PI * k * week) / seasonLength;
    row.push(Math.sin(angle), Math.cos(angle));
  }
  return row;
}

/** Solve a symmetric system by Gaussian elimination with partial pivoting. */
function solve(matrix, vector) {
  const size = vector.length;
  const rows = matrix.map((row, i) => [...row, vector[i]]);
  for (let column = 0; column < size; column += 1) {
    let pivot = column;
    for (let row = column + 1; row < size; row += 1) {
      if (Math.abs(rows[row][column]) > Math.abs(rows[pivot][column])) pivot = row;
    }
    [rows[column], rows[pivot]] = [rows[pivot], rows[column]];
    for (let row = 0; row < size; row += 1) {
      if (row === column) continue;
      const factor = rows[row][column] / rows[column][column];
      for (let k = column; k <= size; k += 1) rows[row][k] -= factor * rows[column][k];
    }
  }
  return rows.map((row, i) => row[size] / row[i]);
}

/**
 * Least squares over the whole history. Returns a model you can inspect.
 *
 * `coefficients[1]` is the growth per cycle, in the unit of the series. That
 * number is worth reading before any forecast is: a model that has found a
 * trend nobody in the company recognises is a model to distrust.
 */
export function fit(history, { seasonLength = 52, harmonics = 2 } = {}) {
  const width = 2 + 2 * harmonics;
  if (history.length < width) throw new RangeError('fewer weeks of history than features to estimate');
  if (history.length < seasonLength) throw new RangeError('a trend cannot be told apart from the season on less than one full cycle');
  // A missing week must be refused: null would be multiplied as zero, NaN would spread.
  if (!history.every(Number.isFinite)) throw new TypeError('every week of history needs a finite number');
  const design = history.map((_, week) => calendarFeatures(week, seasonLength, harmonics));

  // Normal equations: (Xᵀ X) b = Xᵀ y.
  const square = Array.from({ length: width }, (_, i) => Array.from({ length: width }, (_, j) => (
    design.reduce((total, row) => total + row[i] * row[j], 0)
  )));
  const target = Array.from({ length: width }, (_, i) => (
    design.reduce((total, row, week) => total + row[i] * history[week], 0)
  ));

  return { coefficients: solve(square, target), seasonLength, harmonics, start: history.length };
}

/** Predict the `horizon` weeks that follow the history the model was fitted on. */
export function forecast(model, horizon = 1) {
  const weeks = [];
  for (let step = 0; step < horizon; step += 1) {
    const row = calendarFeatures(model.start + step, model.seasonLength, model.harmonics);
    weeks.push(row.reduce((total, value, i) => total + value * model.coefficients[i], 0));
  }
  return weeks;
}