Nonlinear AutoRegressive eXogenous (NARX) model

In this example, we illustrate how to build a polynomial NARX model for time series prediction.

# Authors: The fastcan developers
# SPDX-License-Identifier: MIT

Prepare data

First, a simulated time series dataset is generated from the following nonlinear system.

\[y(k) = 0.5y(k-1) + 0.3u_0(k)^2 + 2u_0(k-1)u_0(k-3) + 1.5u_0(k-2)u_1(k-3) + 1\]

where \(k\) is the time index, \(u_0\) and \(u_1\) are input signals, and \(y\) is the output signal.

import numpy as np

rng = np.random.default_rng(12345)
n_samples = 1000
max_delay = 3
e = rng.normal(0, 0.1, n_samples)
u0 = rng.uniform(0, 1, n_samples + max_delay)
u1 = rng.normal(0, 0.1, n_samples + max_delay)
y = np.zeros(n_samples + max_delay)
for i in range(max_delay, n_samples + max_delay):
    y[i] = (
        0.5 * y[i - 1]
        + 0.3 * u0[i] ** 2
        + 2 * u0[i - 1] * u0[i - 3]
        + 1.5 * u0[i - 2] * u1[i - 3]
        + 1
    )
y = y[max_delay:] + e
X = np.c_[u0[max_delay:], u1[max_delay:]]

Build term library

To build a reduced polynomial NARX model, it is normally have two steps:

1. Search the structure of the model, i.e., the terms in the model, e.g., \(u_0(k-1)u_0(k-3)\), \(u_0(k-2)u_1(k-3)\), etc.

2. Learn the coefficients of the terms.

To search the structure of the model, the candidate term library should be constructed by the following two steps.

1. Time-shifted variables: the raw input-output data, i.e., \(u_0(k)\), \(u_1(k)\), and \(y(k)\), are converted into \(u_0(k-1)\), \(u_1(k-2)\), etc.

2. Nonlinear terms: the time-shifted variables are converted to nonlinear terms via polynomial basis functions, e.g., \(u_0(k-1)^2\), \(u_0(k-1)u_0(k-3)\), etc.

References

• “Nonlinear system identification: NARMAX methods in the time, frequency, and spatio-temporal domains” Billings, S. A. John Wiley & Sons, (2013).

Make time-shifted variables

from fastcan.narx import make_time_shift_features, make_time_shift_ids

time_shift_ids = make_time_shift_ids(
    n_features=3,  # Number of inputs (2) and output (1) signals
    max_delay=3,  # Maximum time delays
    include_zero_delay=[True, True, False],  # Whether to include zero delay
    # for each signal. The output signal should not have zero delay.
)

time_shift_vars = make_time_shift_features(np.c_[X, y], time_shift_ids)

Make nonlinear terms

from fastcan.narx import make_poly_features, make_poly_ids

poly_ids = make_poly_ids(
    n_features=time_shift_vars.shape[1],  # Number of time-shifted variables
    degree=2,  # Maximum polynomial degree
)

poly_terms = make_poly_features(time_shift_vars, poly_ids)

Term selection

After the term library is constructed, the terms can be selected by FastCan, whose \(X\) is the nonlinear terms and \(y\) is the output signal.

from fastcan import FastCan
from fastcan.utils import mask_missing_values

# Mask out missing values caused by time-shifting
poly_terms_masked, y_masked = mask_missing_values(poly_terms, y)

selector = FastCan(
    n_features_to_select=4,  # 4 terms should be selected
).fit(poly_terms_masked, y_masked)

support = selector.get_support()
selected_poly_ids = poly_ids[support]

Progress: 1/4, SSC: 0.70832
Progress: 2/4, SSC: 0.92451
Progress: 3/4, SSC: 0.94432
Progress: 4/4, SSC: 0.96589

Build NARX model

As the reduced polynomial NARX is a linear function of the nonlinear terms, the coefficient of each term can be easily estimated by ordinary least squares. In the printed NARX model, it is found that FastCan selects the correct terms and the coefficients are close to the true values.

from fastcan.narx import NARX, print_narx, tp2fd

# Convert poly_ids and time_shift_ids to feat_ids and delay_ids
feat_ids, delay_ids = tp2fd(time_shift_ids, selected_poly_ids)

narx_model = NARX(
    feat_ids=feat_ids,
    delay_ids=delay_ids,
)

narx_model.fit(X, y)

print_narx(narx_model)

| yid |        Term        |   Coef   |
|-----|--------------------|----------|
|  0  |     Intercept      |  1.050   |
|  0  |    y_hat[k-1,0]    |  0.484   |
|  0  |   X[k,0]*X[k,0]    |  0.306   |
|  0  | X[k-1,0]*X[k-3,0]  |  2.000   |
|  0  | X[k-2,0]*X[k-3,1]  |  1.528   |

Automated NARX modelling workflow

We provide narx.make_narx() to automate the workflow above.

from fastcan.narx import make_narx

auto_narx_model = make_narx(
    X=X,
    y=y,
    n_terms_to_select=4,
    max_delay=3,
    poly_degree=2,
    verbose=0,
).fit(X, y)

print_narx(auto_narx_model)

| yid |        Term        |   Coef   |
|-----|--------------------|----------|
|  0  |     Intercept      |  1.050   |
|  0  |    y_hat[k-1,0]    |  0.484   |
|  0  |   X[k,0]*X[k,0]    |  0.306   |
|  0  | X[k-1,0]*X[k-3,0]  |  2.000   |
|  0  | X[k-2,0]*X[k-3,1]  |  1.528   |

Plot NARX prediction performance

import matplotlib.pyplot as plt
from sklearn.metrics import r2_score

y_pred = narx_model.predict(
    X[:100],
    y_init=y[: narx_model.max_delay_],  # Set the initial values of the prediction to
    # the true values
)

plt.plot(y[:100], label="True")
plt.plot(y_pred, label="Predicted")
plt.xlabel("Time index k")
plt.legend()
plt.title(f"NARX prediction R-squared: {r2_score(y[:100], y_pred):.5f}")
plt.show()