1# -*- coding: utf-8 -*-
2"""
3(Internal) AR(1) model for monthly growth rates aggregated to quarterly freq.
4
5Author: Chad Fulton
6License: BSD-3
7"""
8import warnings
9import numpy as np
10
11from statsmodels.tools.tools import Bunch
12from statsmodels.tsa.statespace import mlemodel, initialization
13from statsmodels.tsa.statespace.kalman_smoother import (
14    SMOOTHER_STATE, SMOOTHER_STATE_COV, SMOOTHER_STATE_AUTOCOV)
15from statsmodels.tsa.statespace.tools import (
16    constrain_stationary_univariate, unconstrain_stationary_univariate)
17
18
19class QuarterlyAR1(mlemodel.MLEModel):
20    r"""
21    AR(1) model for monthly growth rates aggregated to quarterly frequency
22
23    Parameters
24    ----------
25    endog : array_like
26        The observed time-series process :math:`y`
27
28    Notes
29    -----
30    This model is internal, used to estimate starting parameters for the
31    DynamicFactorMQ class. The model is:
32
33    .. math::
34
35        y_t & = \begin{bmatrix} 1 & 2 & 3 & 2 & 1 \end{bmatrix} \alpha_t \\
36        \alpha_t & = \begin{bmatrix}
37            \phi & 0 & 0 & 0 & 0 \\
38               1 & 0 & 0 & 0 & 0 \\
39               0 & 1 & 0 & 0 & 0 \\
40               0 & 0 & 1 & 0 & 0 \\
41               0 & 0 & 0 & 1 & 0 \\
42        \end{bmatrix} +
43        \begin{bmatrix} 1 \\ 0 \\ 0 \\ 0 \\ 0 \end{bmatrix} \varepsilon_t
44
45    The two parameters to be estimated are :math:`\phi` and :math:`\sigma^2`.
46
47    It supports fitting via the usual quasi-Newton methods, as well as using
48    the EM algorithm.
49
50    """
51    def __init__(self, endog):
52        super().__init__(endog, k_states=5, k_posdef=1,
53                         initialization='stationary')
54        self['design'] = [1, 2, 3, 2, 1]
55        self['transition', 1:, :-1] = np.eye(4)
56        self['selection', 0, 0] = 1.
57
58    @property
59    def param_names(self):
60        return ['phi', 'sigma2']
61
62    @property
63    def start_params(self):
64        return np.array([0, np.nanvar(self.endog) / 19])
65
66    def fit(self, *args, **kwargs):
67        # Don't show warnings
68        with warnings.catch_warnings():
69            warnings.simplefilter("ignore")
70            out = super().fit(*args, **kwargs)
71        return out
72
73    def fit_em(self, start_params=None, transformed=True, cov_type='none',
74               cov_kwds=None, maxiter=500, tolerance=1e-6,
75               em_initialization=True, mstep_method=None, full_output=True,
76               return_params=False, low_memory=False):
77        if self._has_fixed_params:
78            raise NotImplementedError('Cannot fit using the EM algorithm while'
79                                      ' holding some parameters fixed.')
80        if low_memory:
81            raise ValueError('Cannot fit using the EM algorithm when using'
82                             ' low_memory option.')
83
84        if start_params is None:
85            start_params = self.start_params
86            transformed = True
87        else:
88            start_params = np.array(start_params, ndmin=1)
89
90        if not transformed:
91            start_params = self.transform_params(start_params)
92
93        # Perform expectation-maximization
94        llf = []
95        params = [start_params]
96        init = None
97        i = 0
98        delta = 0
99        while i < maxiter and (i < 2 or (delta > tolerance)):
100            out = self._em_iteration(params[-1], init=init,
101                                     mstep_method=mstep_method)
102            llf.append(out[0].llf_obs.sum())
103            params.append(out[1])
104            if em_initialization:
105                init = initialization.Initialization(
106                    self.k_states, 'known',
107                    constant=out[0].smoothed_state[..., 0],
108                    stationary_cov=out[0].smoothed_state_cov[..., 0])
109            if i > 0:
110                delta = (2 * (llf[-1] - llf[-2]) /
111                         (np.abs(llf[-1]) + np.abs(llf[-2])))
112            i += 1
113
114        # Just return the fitted parameters if requested
115        if return_params:
116            result = params[-1]
117        # Otherwise construct the results class if desired
118        else:
119            if em_initialization:
120                base_init = self.ssm.initialization
121                self.ssm.initialization = init
122            result = self.smooth(params[-1], transformed=True,
123                                 cov_type=cov_type, cov_kwds=cov_kwds)
124            if em_initialization:
125                self.ssm.initialization = base_init
126
127            # Save the output
128            if full_output:
129                em_retvals = Bunch(**{'params': np.array(params),
130                                      'llf': np.array(llf),
131                                      'iter': i})
132                em_settings = Bunch(**{'tolerance': tolerance,
133                                       'maxiter': maxiter})
134            else:
135                em_retvals = None
136                em_settings = None
137
138            result.mle_retvals = em_retvals
139            result.mle_settings = em_settings
140
141        return result
142
143    def _em_iteration(self, params0, init=None, mstep_method=None):
144        # (E)xpectation step
145        res = self._em_expectation_step(params0, init=init)
146
147        # (M)aximization step
148        params1 = self._em_maximization_step(res, params0,
149                                             mstep_method=mstep_method)
150
151        return res, params1
152
153    def _em_expectation_step(self, params0, init=None):
154        # (E)xpectation step
155        self.update(params0)
156        # Re-initialize state, if new initialization is given
157        if init is not None:
158            base_init = self.ssm.initialization
159            self.ssm.initialization = init
160        # Perform smoothing, only saving what is required
161        res = self.ssm.smooth(
162            SMOOTHER_STATE | SMOOTHER_STATE_COV | SMOOTHER_STATE_AUTOCOV,
163            update_filter=False)
164        res.llf_obs = np.array(
165            self.ssm._kalman_filter.loglikelihood, copy=True)
166        # Reset initialization
167        if init is not None:
168            self.ssm.initialization = base_init
169
170        return res
171
172    def _em_maximization_step(self, res, params0, mstep_method=None):
173        a = res.smoothed_state.T[..., None]
174        cov_a = res.smoothed_state_cov.transpose(2, 0, 1)
175        acov_a = res.smoothed_state_autocov.transpose(2, 0, 1)
176
177        # E[a_t a_t'], t = 0, ..., T
178        Eaa = cov_a.copy() + np.matmul(a, a.transpose(0, 2, 1))
179        # E[a_t a_{t-1}'], t = 1, ..., T
180        Eaa1 = acov_a[:-1] + np.matmul(a[1:], a[:-1].transpose(0, 2, 1))
181
182        # Factor VAR and covariance
183        A = Eaa[:-1, :1, :1].sum(axis=0)
184        B = Eaa1[:, :1, :1].sum(axis=0)
185        C = Eaa[1:, :1, :1].sum(axis=0)
186        nobs = Eaa.shape[0] - 1
187
188        f_A = B / A
189        f_Q = (C - f_A @ B.T) / nobs
190        params1 = np.zeros_like(params0)
191        params1[0] = f_A[0, 0]
192        params1[1] = f_Q[0, 0]
193
194        return params1
195
196    def transform_params(self, unconstrained):
197        return np.array([
198            constrain_stationary_univariate(unconstrained[:1]),
199            unconstrained[1]**2])
200
201    def untransform_params(self, constrained):
202        return np.array([
203            unconstrain_stationary_univariate(constrained[:1]),
204            constrained[1]**0.5])
205
206    def update(self, params, **kwargs):
207        super().update(params, **kwargs)
208
209        self['transition', 0, 0] = params[0]
210        self['state_cov', 0, 0] = params[1]
211