Time series class

python package used to simulate examples of time series

import sys
import numpy as np
import matplotlib.pyplot as plt

class arma_model:

  def __init__(self,phi,theta,size,sigma=1.,burnin=0):
    self.size = size
    self.theta = theta
    self.phi = phi
    self.sigma = sigma
    self.data = self.generator(phi, theta, sigma, size,burnin)

  def generator(self,phi,theta,sigma,size,burnin=0):
    ''' 
      Generator of ARMA process. 
        *phi    - the parameter of AR(p) where p = len(phi)
        *theta  - the parameter of MA(q) where q = len(theta)
        *burnin - specifies the Burn-in period of the process
        *sigma  - standard deviation of the process
        *n      - number of points in the resulting process
    ''' 
    lag = max(len(phi),len(theta))

    noise = np.random.normal(loc=0,scale=sigma,size=burnin+size)
    arma = []
    for i in range(size+burnin):
      if i<lag:
        arma.append(noise[i])
      else:
        val = 0.
        for j in range(len(phi)):
          val += phi[j]*arma[i-j-1] #AR part
        for j in range(len(theta)):
          val += theta[j]*noise[i-j-1] #MA part

        arma.append(val+noise[i])

    self.data = np.array(arma[burnin:])
    return self.data

  def mean(self):
    if any(self.data):
      return np.mean(self.data)
    else:
      return None

  def std(self):
    if any(self.data):
      return np.std(self.data)
    else:
      return None

  def plot(self,save='False'):
    '''Method plots the current ARMA(p,q) model '''  
    plt.figure(1)
    plt.xlabel('Time units')
    plt.ylabel('Values')
    plt.title('ARMA(%d,%d) with (%s,%s)' % (len(self.phi),len(self.theta),self.phi,self.theta))
    plt.grid(True)
    plt.plot(self.data,'o-b') 
    if save:
      plt.savefig('arma.pdf')
    plt.show()