w3_1.py

"""AR(1) Process Simulation """
import numpy as np
import matplotlib.pyplot as plt
from statsmodels.graphics.tsaplots import plot_acf


def gen_ar1(phi, N=1000):
    """Simulate AR(1) process"""

    Z = np.random.normal(0, 1, N)  # Gaussian W.N.

    X = np.empty(N)
    X[:] = np.nan
    X[0] = Z[0]

    for i in range(1, N):
      X[i] = Z[i] + phi * X[i - 1] # AR(1) process
    return X


def plot_ts_and_acf(X, phi):
    """Plot timeseries and ACF"""
    nrows = 2
    fig, axes = plt.subplots(figsize=(8, 4 * nrows), nrows=nrows)

    # Plot Time series
    ax = axes[0]
    ax.plot(X)
    title = 'Time series phi1={:.1f}'.format(phi)
    ax.set(title=title)

    # Plot ACF
    ax = axes[1]
    title = 'ACF phi1={:.1f}'.format(phi)
    fig = plot_acf(X, ax=ax, title=title)

    fig.tight_layout()

    return fig, axes


if __name__ == '__main__':

    np.random.seed(1234)

    phi = 1.0
    X = gen_ar1(phi)
    fig, axes = plot_ts_and_acf(X, phi)
    fig.savefig('ar1_phi{:.1f}.png'.format(phi))

    phi = 0.3
    X = gen_ar1(phi)
    fig, axes = plot_ts_and_acf(X, phi)
    fig.savefig('ar1_phi{:.1f}.png'.format(phi))

    phi = 1.1
    X = gen_ar1(phi)
    fig, axes = plot_ts_and_acf(X, phi)
    ig.savefig('ar1_phi{:.1f}.png'.format(phi))