Merge pull request #113 from QuantEcon/arma

BUG: Found bug in solutions notebook

quantecon/arma.py

@@ -6,7 +6,7 @@
 
 """
 import numpy as np
-from numpy import conj, pi, real
+from numpy import conj, pi
 import matplotlib.pyplot as plt
 from scipy.signal import dimpulse, freqz, dlsim
 
@@ -27,9 +27,9 @@ class ARMA(object):
 
         X_t = \phi X_{t-1} + \epsilon_t + \theta \epsilon_{t-1}
 
-    where {epsilon_t} is a white noise process with standard deviation
-    sigma.  If phi and theta are arrays or sequences, then the
-    interpretation is the ARMA(p, q) model
+    where :math:`epsilon_t` is a white noise process with standard
+    deviation :math:`sigma`.  If phi and theta are arrays or sequences,
+    then the interpretation is the ARMA(p, q) model
 
     .. math::
 
@@ -73,6 +73,39 @@ def __init__(self, phi, theta=0, sigma=1):
         self.sigma = sigma
         self.set_params()
 
+    def __repr__(self):
+        m = "ARMA(phi=%s, theta=%s, sigma=%s)"
+        return m % (self.phi, self.theta, self.sigma)
+
+    def __str__(self):
+        m = "An ARMA({p}, {q}) process"
+        p = np.asarray(self.phi).size
+        q = np.asarray(self.theta).size
+        return m.format(p=p, q=q)
+
+    # Special latex print method for working in notebook
+    def _repr_latex_(self):
+        m = r"$X_t = "
+        phi = np.atleast_1d(self.phi)
+        theta = np.atleast_1d(self.theta)
+        rhs = ""
+        for (tm, phi_p) in enumerate(phi):
+            # don't include terms if they are equal to zero
+            if abs(phi_p) > 1e-12:
+                rhs += r"%+g X_{t-%i}" % (phi_p, tm+1)
+
+        if rhs[0] == "+":
+            rhs = rhs[1:]  # remove initial `+` if phi_1 was positive
+
+        rhs += r" + \epsilon_t"
+
+        for (tm, th_q) in enumerate(theta):
+            # don't include terms if they are equal to zero
+            if abs(th_q) > 1e-12:
+                rhs += r"%+g \epsilon_{t-%i}" % (th_q, tm+1)
+
+        return m + rhs + "$"
+
     @property
     def phi(self):
         return self._phi
@@ -166,8 +199,8 @@ def spectral_density(self, two_pi=True, res=1200):
         res : scalar or array_like(int), optional(default=1200)
             If res is a scalar then the spectral density is computed at
             `res` frequencies evenly spaced around the unit circle, but
-            if res is an array then the function computes the response at the
-            frequencies given by the array
+            if res is an array then the function computes the response
+            at the frequencies given by the array
 
         Returns
         -------
@@ -185,9 +218,9 @@ def spectral_density(self, two_pi=True, res=1200):
 
     def autocovariance(self, num_autocov=16):
         """
-        Compute the autocovariance function from the ARMA parameters over the
-        integers range(num_autocov) using the spectral density and the inverse
-        Fourier transform.
+        Compute the autocovariance function from the ARMA parameters
+        over the integers range(num_autocov) using the spectral density
+        and the inverse Fourier transform.
 
         Parameters
         ----------