covid-19/src/utils/regression.ts

import differenceInDays from 'date-fns/differenceInDays';
import addDays from 'date-fns/addDays';
import { Cases, Data, FutureCase } from '../types';

// Start the regression line after this many cases have been recorded in total
const regressionStart = 50;

const futureDays = 10;

const mean = (values: number[]): number =>
  values.reduce((last, value) => last + value, 0) / values.length;

export function getExponentialRegression(cases: Cases): Data {
  if (!cases.length) {
    return [];
  }

  const { index: startIndex } = cases.reduce(
    ({ sum, index }, { value }, nextIndex) => {
      if (sum >= regressionStart) {
        return { sum, index };
      }

      const nextSum = sum + value;
      if (nextSum < regressionStart) {
        return { sum: nextSum, index };
      }

      return { sum: nextSum, index: nextIndex };
    },
    { sum: 0, index: -1 },
  );

  if (startIndex === -1) {
    return cases.map(({ date, value }) => ({
      date: date.getTime(),
      value,
    }));
  }

  const casesToRegress = cases.slice(startIndex);

  // It's assumed that the input here is ordered by date ascending
  const minDate = casesToRegress[0].date;

  const xSeries = casesToRegress.map(({ date }) => 1 + differenceInDays(date, minDate));
  const ySeries = casesToRegress.map(({ value }) => Math.log(value));

  const xBar = mean(xSeries);
  const yBar = mean(ySeries);

  const covariance =
    xSeries.reduce((last, value, index) => last + (value - xBar) * (ySeries[index] - yBar), 0) /
    xSeries.length;

  const xVariance = mean(xSeries.map(value => value ** 2)) - xBar ** 2;

  const slope = covariance / xVariance;
  const intercept = yBar - slope * xBar;

  const regressionAtDate = (date: Date): number => {
    const xValue = differenceInDays(date, minDate);

    return Math.exp(slope * xValue + intercept);
  };

  const lastDate = cases[cases.length - 1].date;

  const future: FutureCase[] = new Array(futureDays)
    .fill(0)
    .map((_, index) => ({ date: addDays(lastDate, index + 1) }));

  return [...cases, ...future].map(({ date, ...rest }) => ({
    date: date.getTime(),
    regression: regressionAtDate(date),
    ...rest,
  }));
}