Forecast next week's sales

"""
Forecast weekly sales with a moving average and seasonal coefficients.

Rung N0. Standard library only, and short enough to read in one sitting.

The idea is the oldest one in forecasting, and it still carries most small
businesses: sales are a level that moves slowly, multiplied by a shape that
repeats every year. Estimate the shape over the whole history, divide it out,
average what is left over the last few weeks, then put the shape back.

What this deliberately does not model is a trend. A moving average only ever
looks backwards, so it follows a change of regime instead of anticipating it.
That is the limit of this rung, and the test says so out loud.
"""

from statistics import fmean


def seasonal_coefficients(history: list[float], season_length: int) -> list[float]:
    """
    One multiplier per position in the cycle, averaged over the history.

    A coefficient of 1.2 for week 50 means that week 50 usually sells twenty
    per cent above the year's level. The coefficients are rescaled to average
    one, so putting the shape back neither inflates nor deflates the forecast.
    """
    level = fmean(history)
    if level == 0:
        return [1.0] * season_length
    raw = [fmean(history[phase::season_length]) / level for phase in range(season_length)]
    scale = fmean(raw)
    return [coefficient / scale for coefficient in raw]


def forecast(
    history: list[float], season_length: int = 52, window: int = 4, horizon: int = 1
) -> list[float]:
    """
    Predict the next `horizon` weeks from `history`.

    The history and the forecast share one clock: the week after the end of
    the history is position `len(history)` in the cycle. The caller never has
    to align the series on a January.

    `window` is the only real knob. A short window reacts fast and trusts the
    last few weeks; a long one is steadier and slower to notice a change.
    """
    if len(history) < 2 * season_length:
        raise ValueError("a seasonal coefficient needs two full cycles at the very least")
    coefficients = seasonal_coefficients(history, season_length)
    # Divide the season out, so the average below measures the level alone.
    deseasonalised = [value / coefficients[week % season_length] for week, value in enumerate(history)]
    level = fmean(deseasonalised[-window:])
    start = len(history)
    return [level * coefficients[(start + step) % season_length] for step in range(horizon)]

/**
 * Forecast weekly sales with a moving average and seasonal coefficients.
 *
 * Rung N0. No dependency, and short enough to read in one sitting.
 *
 * The idea is the oldest one in forecasting, and it still carries most small
 * businesses: sales are a level that moves slowly, multiplied by a shape that
 * repeats every year. Estimate the shape over the whole history, divide it
 * out, average what is left over the last few weeks, then put the shape back.
 *
 * What this deliberately does not model is a trend. A moving average only
 * ever looks backwards, so it follows a change of regime instead of
 * anticipating it. That is the limit of this rung, and the test says so out
 * loud.
 */

function mean(values) {
  return values.reduce((total, value) => total + value, 0) / values.length;
}

/**
 * One multiplier per position in the cycle, averaged over the history.
 *
 * A coefficient of 1.2 for week 50 means that week 50 usually sells twenty
 * per cent above the year's level. The coefficients are rescaled to average
 * one, so putting the shape back neither inflates nor deflates the forecast.
 */
export function seasonalCoefficients(history, seasonLength) {
  const level = mean(history);
  if (level === 0) return new Array(seasonLength).fill(1);
  const raw = [];
  for (let phase = 0; phase < seasonLength; phase += 1) {
    const sameWeekEveryYear = history.filter((_, week) => week % seasonLength === phase);
    raw.push(mean(sameWeekEveryYear) / level);
  }
  const scale = mean(raw);
  return raw.map((coefficient) => coefficient / scale);
}

/**
 * Predict the next `horizon` weeks from `history`.
 *
 * The history and the forecast share one clock: the week after the end of the
 * history is position `history.length` in the cycle. The caller never has to
 * align the series on a January.
 *
 * `window` is the only real knob. A short window reacts fast and trusts the
 * last few weeks; a long one is steadier and slower to notice a change.
 */
export function forecast(history, { seasonLength = 52, window = 4, horizon = 1 } = {}) {
  if (history.length < 2 * seasonLength) {
    throw new RangeError('a seasonal coefficient needs two full cycles at the very least');
  }
  const coefficients = seasonalCoefficients(history, seasonLength);
  // Divide the season out, so the average below measures the level alone.
  const deseasonalised = history.map((value, week) => value / coefficients[week % seasonLength]);
  const level = mean(deseasonalised.slice(-window));
  const start = history.length;
  const weeks = [];
  for (let step = 0; step < horizon; step += 1) {
    weeks.push(level * coefficients[(start + step) % seasonLength]);
  }
  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.
Two pairs are enough to draw a Christmas peak and a summer dip; more pairs
start drawing the noise as well.

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 year.
"""

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.
    """
    design = [calendar_features(week, season_length, harmonics) for week in range(len(history))]
    if len(history) < len(design[0]):
        raise ValueError("fewer weeks of history than features to estimate")
    coefficients, *_ = np.linalg.lstsq(np.array(design), np.array(history, dtype=float), 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, twenty lines that no dependency is worth. 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 year.
 */

/** 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 design = history.map((_, week) => calendarFeatures(week, seasonLength, harmonics));
  const width = design[0].length;
  if (history.length < width) throw new RangeError('fewer weeks of history than features to estimate');

  // 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;
}