Sl/numba tauchen

quantecon/markov/approximation.py

@@ -9,8 +9,10 @@
 
 """
 
+from math import erfc, sqrt
+
 import numpy as np
-from scipy.stats import norm
+from numba import njit
 
 
 def tauchen(rho, sigma_u, m=3, n=7):
@@ -44,16 +46,15 @@ def tauchen(rho, sigma_u, m=3, n=7):
         of transitioning from x[i] to x[j]
 
     """
-    F = norm(loc=0, scale=sigma_u).cdf
 
     # standard deviation of y_t
-    std_y = np.sqrt(sigma_u**2 / (1-rho**2))
+    std_y = np.sqrt(sigma_u**2 / (1 - rho**2))
 
     # top of discrete state space
     x_max = m * std_y
 
     # bottom of discrete state space
-    x_min = - x_max
+    x_min = -x_max
 
     # discretized state space
     x = np.linspace(x_min, x_max, n)
@@ -62,11 +63,23 @@ def tauchen(rho, sigma_u, m=3, n=7):
     half_step = 0.5 * step
     P = np.empty((n, n))
 
-    for i in range(n):
-        P[i, 0] = F(x[0]-rho * x[i] + half_step)
-        P[i, n-1] = 1 - F(x[n-1] - rho * x[i] - half_step)
-        for j in range(1, n-1):
-            z = x[j] - rho * x[i]
-            P[i, j] = F(z + half_step) - F(z - half_step)
+    _fill_tauchen(x, P, n, rho, sigma_u, half_step)
 
     return x, P
+
+
+@njit
+def std_norm_cdf(x):
+    return 0.5 * erfc(-x / sqrt(2))
+
+
+@njit
+def _fill_tauchen(x, P, n, rho, sigma, half_step):
+    for i in range(n):
+        P[i, 0] = std_norm_cdf((x[0] - rho * x[i] + half_step) / sigma)
+        P[i, n - 1] = 1 - \
+            std_norm_cdf((x[n - 1] - rho * x[i] - half_step) / sigma)
+        for j in range(1, n - 1):
+            z = x[j] - rho * x[i]
+            P[i, j] = (std_norm_cdf((z + half_step) / sigma) -
+                       std_norm_cdf((z - half_step) / sigma))