moving-average-arima/src/load.ts

ARIMA, SARIMA, SARIMAX and AutoARIMA models for time series analysis and forecasting

function uintify(arr: Array<any>) {
  return new Uint8Array(Float64Array.from(arr).buffer);
}

function flat(arr: Array<any>) {
  return [].concat.apply([], arr);
}

function transpose(arr: Array<any>) {
  return arr[0].map((_: any, i: number) => arr.map((x) => x[i]));
}

function prepare(arr: Array<any>) {
  const farr = flat(arr);
  for (let i = 0; i < farr.length - 2; i++) {
    if (isNaN(farr[i + 1])) {
      farr[i + 1] = farr[i];
    }
  }
  return farr;
}

const defaults = {
  method: 0,
  optimizer: 6,
  s: 0,
  verbose: true,
  transpose: false,
  auto: false,
  approximation: 1,
  search: 1,
};

const params = {
  p: 1,
  d: 0,
  q: 1,
  P: 0,
  D: 0,
  Q: 0,
};

const paramsAuto = {
  p: 5,
  d: 2,
  q: 5,
  P: 2,
  D: 1,
  Q: 2,
};

export default function (m: any) {
  const _fit_sarimax = m.cwrap("fit_sarimax", "number", [
    "array",
    "array",
    "number",
    "number",
    "number",
    "number",
    "number",
    "number",
    "number",
    "number",
    "number",
    "number",
    "number",
    "boolean",
  ]);
  const _predict_sarimax = m.cwrap("predict_sarimax", "number", [
    "number",
    "array",
    "array",
    "array",
    "number",
  ]);
  const _fit_autoarima = m.cwrap("fit_autoarima", "number", [
    "array",
    "array",
    "number",
    "number",
    "number",
    "number",
    "number",
    "number",
    "number",
    "number",
    "number",
    "number",
    "number",
    "number",
    "number",
    "boolean",
  ]);
  const _predict_autoarima = m.cwrap("predict_autoarima", "number", [
    "number",
    "array",
    "array",
    "array",
    "number",
  ]);

  function getResults(addr, l) {
    const res = [[], []];
    for (let i = 0; i < l * 2; i++) {
      res[i < l ? 0 : 1].push(
        m.HEAPF64[addr / Float64Array.BYTES_PER_ELEMENT + i]
      );
    }
    return res;
  }

  function ARIMA() {
    // Preserve the old functional API: ARIMA(ts, len, opts)
    if (!(this instanceof ARIMA)) {
      console.warn(
        "Calling ARIMA as a function will be deprecated in the future"
      );
      return new ARIMA(arguments[2]).train(arguments[0]).predict(arguments[1]);
    }
    // A new, class API has opts as the only argument here: new ARIMA (opts)
    const opts = arguments[0];
    const o = Object.assign(
      {},
      defaults,
      opts.auto ? paramsAuto : params,
      opts
    );
    if (
      Math.min(o.method, o.optimizer, o.p, o.d, o.q, o.P, o.D, o.Q, o.s) < 0
    ) {
      throw new Error("Model parameter can't be negative");
    }
    if (o.P + o.D + o.Q === 0) {
      o.s = 0;
    } else if (o.s === 0) {
      o.P = o.D = o.Q = 0;
    }
    this.options = o;
  }

  ARIMA.prototype.train = function (ts, exog = []) {
    const o = this.options;
    if (o.transpose && Array.isArray(exog[0])) {
      exog = transpose(exog);
    }
    this.ts = uintify(prepare(ts));
    this.exog = uintify(prepare(exog));
    this.lin = ts.length;
    this.nexog =
      exog.length > 0 ? (Array.isArray(exog[0]) ? exog.length : 1) : 0;
    this.model = o.auto
      ? _fit_autoarima(
        this.ts,
        this.exog,
        o.p,
        o.d,
        o.q,
        o.P,
        o.D,
        o.Q,
        o.s,
        this.nexog,
        this.lin,
        o.method,
        o.optimizer,
        o.approximation,
        o.search,
        o.verbose
      )
      : _fit_sarimax(
        this.ts,
        this.exog,
        o.p,
        o.d,
        o.q,
        o.P,
        o.D,
        o.Q,
        o.s,
        this.nexog,
        this.lin,
        o.method,
        o.optimizer,
        o.verbose
      );
    return this;
  };

  ARIMA.prototype.fit = function (...a) {
    return this.train(...a);
  };

  ARIMA.prototype.predict = function (l, exog = []) {
    const o = this.options;
    if (o.transpose && Array.isArray(exog[0])) {
      exog = transpose(exog);
    }
    const addr = o.auto
      ? _predict_autoarima(
        this.model,
        this.ts,
        this.exog, // old
        uintify(prepare(exog)), // new
        l
      )
      : _predict_sarimax(
        this.model,
        this.ts,
        this.exog, // old
        uintify(prepare(exog)), // new
        l
      );
    return getResults(addr, l);
  };

  return ARIMA;
}