covid-19/src/utils/regression.ts

import { compose } from '@typed/compose';
import groupBy from 'lodash/groupBy';
import flatten from 'lodash/flatten';
import differenceInDays from 'date-fns/differenceInDays';
import addDays from 'date-fns/addDays';
import { Cases, Data, DataPoint, CountryCase, CountryCases, CountryDataPoint } from '../types';

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

// Number of days to extrapolate the regression fit on the graph
const futureDays = 3;

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

type WithNumericDate = Omit<CountryDataPoint, 'regression'>[];

const withNumericDate = (cases: Cases): WithNumericDate =>
  cases.map(({ date, value }) => ({
    date,
    xValue: date.getTime(),
    value,
  }));

type WithCumulative = (Omit<CountryDataPoint, 'regression'> & {
  valueCumulative: number;
})[];

const withCumulative = (cumulative: boolean) => (cases: WithNumericDate): WithCumulative => {
  const valuesWithCumulative = cases.reduce(
    (last: WithCumulative, { date, xValue, value }): WithCumulative => [
      ...last,
      {
        date,
        xValue,
        value,
        valueCumulative: (value || 0) + (last[last.length - 1]?.valueCumulative || 0),
      },
    ],
    [],
  );

  if (cumulative) {
    return valuesWithCumulative.map(({ valueCumulative, value, ...rest }) => ({
      ...rest,
      value: valueCumulative,
      valueCumulative,
    }));
  }

  return valuesWithCumulative;
};

const logArray = (values: number[]): number[] => values.map(value => Math.log(value));

const withExponentialRegression = (cumulative: boolean, regressionBuffer: number) => (
  cases: WithCumulative,
): CountryDataPoint[] => {
  const startIndex = cases.findIndex(
    ({ valueCumulative = 0 }) => valueCumulative >= regressionStart,
  );
  if (startIndex === -1 || startIndex >= cases.length - regressionBuffer) {
    return cases;
  }

  const casesToRegress = cases.slice(startIndex, cases.length - regressionBuffer);

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

  const xSeries = casesToRegress.map(({ date }) => 1 + differenceInDays(new Date(date), minDate));
  const xBar = mean(xSeries);
  const xVariance = mean(xSeries.map(value => value ** 2)) - xBar ** 2;

  const ySeries = logArray(
    cumulative
      ? casesToRegress.map(({ valueCumulative = 0 }) => valueCumulative)
      : casesToRegress.map(({ value = 0 }) => value),
  );
  const yBar = mean(ySeries);
  const covariance =
    xSeries.reduce((last, value, index) => last + (value - xBar) * (ySeries[index] - yBar), 0) /
    xSeries.length;

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

  const regressionAtDate = (date: Date): number =>
    Math.exp(slope * differenceInDays(date, minDate) + intercept);

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

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

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

function processCountryCases(
  cases: Cases,
  cumulative: boolean,
  regressionBuffer: number,
): CountryDataPoint[] {
  if (!cases.length) {
    return [];
  }

  return compose<Cases, WithNumericDate, WithCumulative, CountryDataPoint[]>(
    withExponentialRegression(cumulative, regressionBuffer),
    withCumulative(cumulative),
    withNumericDate,
  )(cases);
}

function combineData(items: Data[]): Data {
  const groups: {
    [xValue: string]: Data;
  } = groupBy(flatten(items), 'xValue');

  return Object.keys(groups)
    .sort((timeA: string, timeB: string) => Number(timeA) - Number(timeB))
    .map(key =>
      groups[key].reduce(
        (last: DataPoint, item: DataPoint): DataPoint => ({
          ...last,
          value: {
            ...last.value,
            ...item.value,
          },
          regression: {
            ...last.regression,
            ...item.regression,
          },
        }),
        {
          date: groups[key][0].date,
          xValue: groups[key][0].xValue,
          value: {},
          regression: {},
        },
      ),
    );
}

function fillData(countryCases: CountryCases): CountryCases {
  const times: number[] = Array.from(
    new Set(
      flatten(
        countryCases.map(({ dataSource: { cases } }) => cases.map(({ date }) => date.getTime())),
      ),
    ),
  ).sort((timeA, timeB) => timeA - timeB);

  return countryCases.map(
    ({ country, dataSource }): CountryCase => ({
      country,
      dataSource: {
        ...dataSource,
        cases: times.reduce((last: Cases, time: number): Cases => {
          const matchingCase = dataSource.cases.find(({ date }) => date.getTime() === time);
          if (matchingCase) {
            return [...last, matchingCase];
          }

          return [
            ...last,
            {
              date: new Date(time),
              value: 0,
            },
          ];
        }, []),
      },
    }),
  );
}

export function processCases(
  countryCases: CountryCases,
  cumulative = true,
  regressionBuffer = 0, // Don't count the last X days into the regression fit
): Data {
  const filledData = fillData(countryCases);

  const data = filledData.map(({ country, dataSource: { cases } }) =>
    processCountryCases(cases, cumulative, regressionBuffer).map(
      ({ date, xValue, value, regression }) => ({
        date,
        xValue,
        value: {
          [country]: value,
        },
        regression: {
          [country]: regression,
        },
      }),
    ),
  );

  return combineData(data);
}