Return amse vector (#2)

README.md

@@ -5,11 +5,11 @@
 Emscripten port of the native C package [ctsa](https://github.com/rafat/ctsa) for univariate time series analysis and prediction.
 
 ### API
-Interface of `arima` consists of only one function that takes a 1D vector with observations over time, number of time steps to predict and ARIMA parameters p, d, q. Returns a vector with predictions.
+Interface of `arima` consists of only one function that takes a 1D vector with observations over time, number of time steps to predict and ARIMA parameters p, d, q. Returns two vectors - predictions and mean square errors.
 
 ```javascript
 const arima = require('arima')
-const pred = arima(ts, 20, {
+const [pred, errors] = arima(ts, 20, {
   method: 0, // ARIMA method (Default: 0)
   optimizer: 6, // Optimization method (Default: 6)
   p: 1, // Number of Autoregressive coefficients

example.js

@@ -1,13 +1,22 @@
 const arima = require('.')
-const ts = Array(1000).fill(0).map((v, i) => Math.sin(i + 1))
 
-const pred = arima(ts.slice(0, 980), 20, {
+const n = 1000
+const l = 30
+
+const ts = Array(n).fill(0).map((v, i) => Math.sin(i / 3) + Math.random() / 1.5)
+ts[100] = '--'
+ts[500] = undefined
+
+const [pred, errors] = arima(ts.slice(0, n - l), l, {
   method: 0,
   optimizer: 6,
-  p: 2,
+  p: 10,
   q: 0,
   d: 1,
   verbose: false
 })
 
-console.log(pred, ts.slice(980))
+console.log('Ytrue,Ypred,Err')
+pred.forEach((v, i) => {
+  console.log(ts[n - l + i] + ',' + v + ',' + errors[i])
+})

index.js

@@ -8,6 +8,20 @@ function uintify (arr) {
   return new Uint8Array(Float64Array.from(arr).buffer)
 }
 
+function flat (arr) {
+  return [].concat.apply([], arr)
+}
+
+function prepare (arr) {
+  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,
@@ -19,7 +33,7 @@ const defaults = {
 
 module.exports = function predict (input, length, opts) {
   const options = Object.assign({}, defaults, opts)
-  const ts = uintify(input.flat())
+  const ts = uintify(prepare(input))
   const addr = _arima(
     ts,
     options.p,
@@ -31,9 +45,9 @@ module.exports = function predict (input, length, opts) {
     options.optimizer,
     options.verbose
   )
-  const data = []
-  for (let i = 0; i < length; i++) {
-    data.push(m.HEAPF64[addr / Float64Array.BYTES_PER_ELEMENT + i])
+  const res = [[], []]
+  for (let i = 0; i < length * 2; i++) {
+    res[i < length ? 0 : 1].push(m.HEAPF64[addr / Float64Array.BYTES_PER_ELEMENT + i])
   }
-  return data
+  return res
 }

package.json

@@ -1,6 +1,6 @@
 {
   "name": "arima",
-  "version": "0.0.3",
+  "version": "0.1.0",
   "description": "Univariate ARIMA model",
   "main": "index.js",
   "scripts": {

src/api.c

@@ -11,7 +11,7 @@ double* arima (double* ts, int p, int d, int q, int lin, int lout, int method, i
   */
 
   double *phi, *theta;
-  double *xpred, *amse;
+  double *res, *amse;
 
   theta = (double*)malloc(sizeof(double) * q);
   phi = (double*)malloc(sizeof(double) * p);
@@ -22,19 +22,18 @@ double* arima (double* ts, int p, int d, int q, int lin, int lout, int method, i
   arima_setOptMethod(obj, opt);
   arima_exec(obj, ts);
 
-  xpred = (double*)malloc(sizeof(double) * lout);
-  amse = (double*)malloc(sizeof(double) * lout);
+  res = (double*)malloc(sizeof(double) * lout * 2);
+  amse = res + lout;
 
-  arima_predict(obj, ts, lout, xpred, amse);
+  arima_predict(obj, ts, lout, res, amse);
 
   if (verbose) {
     arima_summary(obj);
   }
 
   free(phi);
   free(theta);
-  free(amse);
 
-  return xpred;
+  return res;
 }