1"""
2tauchen
3-------
4Discretizes Gaussian linear AR(1) processes via Tauchen's method
5
6"""
7
8from math import erfc, sqrt
9from .core import MarkovChain
10
11import numpy as np
12from numba import njit
13
14
15def rouwenhorst(n, ybar, sigma, rho):
16    r"""
17    Takes as inputs n, p, q, psi. It will then construct a markov chain
18    that estimates an AR(1) process of:
19    :math:`y_t = \bar{y} + \rho y_{t-1} + \varepsilon_t`
20    where :math:`\varepsilon_t` is i.i.d. normal of mean 0, std dev of sigma
21
22    The Rouwenhorst approximation uses the following recursive defintion
23    for approximating a distribution:
24
25    .. math::
26
27        \theta_2 =
28        \begin{bmatrix}
29        p     &  1 - p \\
30        1 - q &  q     \\
31        \end{bmatrix}
32
33    .. math::
34
35        \theta_{n+1} =
36        p
37        \begin{bmatrix}
38        \theta_n & 0   \\
39        0        & 0   \\
40        \end{bmatrix}
41        + (1 - p)
42        \begin{bmatrix}
43        0  & \theta_n  \\
44        0  &  0        \\
45        \end{bmatrix}
46        + q
47        \begin{bmatrix}
48        0        & 0   \\
49        \theta_n & 0   \\
50        \end{bmatrix}
51        + (1 - q)
52        \begin{bmatrix}
53        0  &  0        \\
54        0  & \theta_n  \\
55        \end{bmatrix}
56
57
58    Parameters
59    ----------
60    n : int
61        The number of points to approximate the distribution
62
63    ybar : float
64        The value :math:`\bar{y}` in the process.  Note that the mean of this
65        AR(1) process, :math:`y`, is simply :math:`\bar{y}/(1 - \rho)`
66
67    sigma : float
68        The value of the standard deviation of the :math:`\varepsilon` process
69
70    rho : float
71        By default this will be 0, but if you are approximating an AR(1)
72        process then this is the autocorrelation across periods
73
74    Returns
75    -------
76
77    mc : MarkovChain
78        An instance of the MarkovChain class that stores the transition
79        matrix and state values returned by the discretization method
80
81    """
82
83    # Get the standard deviation of y
84    y_sd = sqrt(sigma**2 / (1 - rho**2))
85
86    # Given the moments of our process we can find the right values
87    # for p, q, psi because there are analytical solutions as shown in
88    # Gianluca Violante's notes on computational methods
89    p = (1 + rho) / 2
90    q = p
91    psi = y_sd * np.sqrt(n - 1)
92
93    # Find the states
94    ubar = psi
95    lbar = -ubar
96
97    bar = np.linspace(lbar, ubar, n)
98
99    def row_build_mat(n, p, q):
100        """
101        This method uses the values of p and q to build the transition
102        matrix for the rouwenhorst method
103
104        """
105
106        if n == 2:
107            theta = np.array([[p, 1 - p], [1 - q, q]])
108
109        elif n > 2:
110            p1 = np.zeros((n, n))
111            p2 = np.zeros((n, n))
112            p3 = np.zeros((n, n))
113            p4 = np.zeros((n, n))
114
115            new_mat = row_build_mat(n - 1, p, q)
116
117            p1[:n - 1, :n - 1] = p * new_mat
118            p2[:n - 1, 1:] = (1 - p) * new_mat
119            p3[1:, :-1] = (1 - q) * new_mat
120            p4[1:, 1:] = q * new_mat
121
122            theta = p1 + p2 + p3 + p4
123            theta[1:n - 1, :] = theta[1:n - 1, :] / 2
124
125        else:
126            raise ValueError("The number of states must be positive " +
127                             "and greater than or equal to 2")
128
129        return theta
130
131    theta = row_build_mat(n, p, q)
132
133    bar += ybar / (1 - rho)
134
135    return MarkovChain(theta, bar)
136
137
138def tauchen(rho, sigma_u, b=0., m=3, n=7):
139    r"""
140    Computes a Markov chain associated with a discretized version of
141    the linear Gaussian AR(1) process
142
143    .. math::
144
145        y_{t+1} = b + \rho y_t + u_{t+1}
146
147    using Tauchen's method. Here :math:`{u_t}` is an i.i.d. Gaussian process
148    with zero mean.
149
150    Parameters
151    ----------
152    b : scalar(float)
153        The constant term of {y_t}
154    rho : scalar(float)
155        The autocorrelation coefficient
156    sigma_u : scalar(float)
157        The standard deviation of the random process
158    m : scalar(int), optional(default=3)
159        The number of standard deviations to approximate out to
160    n : scalar(int), optional(default=7)
161        The number of states to use in the approximation
162
163    Returns
164    -------
165
166    mc : MarkovChain
167        An instance of the MarkovChain class that stores the transition
168        matrix and state values returned by the discretization method
169
170    """
171
172    # standard deviation of demeaned y_t
173    std_y = np.sqrt(sigma_u**2 / (1 - rho**2))
174
175    # top of discrete state space for demeaned y_t
176    x_max = m * std_y
177
178    # bottom of discrete state space for demeaned y_t
179    x_min = -x_max
180
181    # discretized state space for demeaned y_t
182    x = np.linspace(x_min, x_max, n)
183
184    step = (x_max - x_min) / (n - 1)
185    half_step = 0.5 * step
186    P = np.empty((n, n))
187
188    # approximate Markov transition matrix for
189    # demeaned y_t
190    _fill_tauchen(x, P, n, rho, sigma_u, half_step)
191
192    # shifts the state values by the long run mean of y_t
193    mu = b / (1 - rho)
194
195    mc = MarkovChain(P, state_values=x+mu)
196
197    return mc
198
199
200@njit
201def std_norm_cdf(x):
202    return 0.5 * erfc(-x / sqrt(2))
203
204
205@njit
206def _fill_tauchen(x, P, n, rho, sigma, half_step):
207    for i in range(n):
208        P[i, 0] = std_norm_cdf((x[0] - rho * x[i] + half_step) / sigma)
209        P[i, n - 1] = 1 - \
210            std_norm_cdf((x[n - 1] - rho * x[i] - half_step) / sigma)
211        for j in range(1, n - 1):
212            z = x[j] - rho * x[i]
213            P[i, j] = (std_norm_cdf((z + half_step) / sigma) -
214                       std_norm_cdf((z - half_step) / sigma))
215