Repository: statsmodels/statsmodels

Issue:
REF: Add Kalman filter initialization as keyword argument to constructor
Currently, the Kalman filter initialization must be performed after construction. Need to add the ability to specify it during construction via keyword argument. Four new keywords (and then class attributes) would be specified:

```
Parameters
----------

initialization : {'approximate_diffuse','stationary','known'}, optional
    Method
initial_variance : float, optional
    If approximate diffuse initialization is used, the initial variance of states.
initial_state : array_like, optional
    If known initialization is used, the mean of the initial state's distribution.
initial_state_cov : array_like, optional
    If known initialization is used, the covariance matrix of the initial state's distribution.
```

Follows up on #2252.



Discussion:

Resolving patch:
diff --git a/statsmodels/tsa/statespace/representation.py b/statsmodels/tsa/statespace/representation.py
--- a/statsmodels/tsa/statespace/representation.py
+++ b/statsmodels/tsa/statespace/representation.py
@@ -105,6 +105,14 @@ class Representation(object):
     initial_variance : float, optional
         Initial variance used when approximate diffuse initialization is
         specified. Default is 1e6.
+    initialization : {'approximate_diffuse','stationary','known'}, optional
+        Initialization method for the initial state.
+    initial_state : array_like, optional
+        If known initialization is used, the mean of the initial state's
+        distribution.
+    initial_state_cov : array_like, optional
+        If known initialization is used, the covariance matrix of the initial
+        state's distribution.
     nobs : integer, optional
         If an endogenous vector is not given (i.e. `k_endog` is an integer),
         the number of observations can optionally be specified. If not
@@ -306,21 +314,38 @@ def __init__(self, k_endog, k_states, k_posdef=None,
             if scope[name] is not None:
                 setattr(self, name, scope[name])
 
+        # Options
+        self.initial_variance = initial_variance
+
         # State-space initialization data
-        self.initialization = None
+        self.initialization = kwargs.get('initialization', None)
         self._initial_state = None
         self._initial_state_cov = None
         self._initial_variance = None
 
+        if self.initialization == 'approximate_diffuse':
+            self.initialize_approximate_diffuse()
+        elif self.initialization == 'stationary':
+            self.initialize_stationary()
+        elif self.initialization == 'known':
+            if not 'initial_state' in kwargs:
+                raise ValueError('Initial state must be provided when "known"'
+                                 ' is the specified initialization method.')
+            if not 'initial_state_cov' in kwargs:
+                raise ValueError('Initial state covariance matrix must be'
+                                 ' provided when "known" is the specified'
+                                 ' initialization method.')
+            self.initialize_known(kwargs['initial_state'],
+                                  kwargs['initial_state_cov'])
+        elif self.initialization is not None:
+            raise ValueError("Invalid state space initialization method.")
+
         # Matrix representations storage
         self._representations = {}
 
         # Setup the underlying statespace object storage
         self._statespaces = {}
 
-        # Options
-        self.initial_variance = initial_variance
-
     def __getitem__(self, key):
         _type = type(key)
         # If only a string is given then we must be getting an entire matrix
diff --git a/statsmodels/tsa/statespace/sarimax.py b/statsmodels/tsa/statespace/sarimax.py
--- a/statsmodels/tsa/statespace/sarimax.py
+++ b/statsmodels/tsa/statespace/sarimax.py
@@ -522,8 +522,8 @@ def __init__(self, endog, exog=None, order=(1, 0, 0),
         self.selection = self.initial_selection
 
         # If we are estimating a simple ARMA model, then we can use a faster
-        # initialization method.
-        if k_diffuse_states == 0:
+        # initialization method (unless initialization was already specified).
+        if k_diffuse_states == 0 and self.initialization is None:
             self.initialize_stationary()
 
         # update _init_keys attached by super
@@ -541,9 +541,6 @@ def _get_init_kwds(self):
         kwds = super(SARIMAX, self)._get_init_kwds()
         kwds['endog'] = self.orig_endog
         kwds['exog'] = self.orig_exog
-        if self.initialization == 'approximate_diffuse':
-            import warnings
-            warnings.warn('not all init keys or initialization is available yet')
         return kwds
 
 


Test patch:
diff --git a/statsmodels/tsa/statespace/tests/test_sarimax.py b/statsmodels/tsa/statespace/tests/test_sarimax.py
--- a/statsmodels/tsa/statespace/tests/test_sarimax.py
+++ b/statsmodels/tsa/statespace/tests/test_sarimax.py
@@ -158,9 +158,13 @@ def test_bse(self):
 
 
 class TestARIMADiffuse(ARIMA):
-    def __init__(self):
-        super(TestARIMADiffuse, self).__init__(results_sarimax.wpi1_diffuse)
-        self.model.initialize_approximate_diffuse(self.true['initial_variance'])
+    def __init__(self, **kwargs):
+        kwargs['initialization'] = 'approximate_diffuse'
+        kwargs['initial_variance'] = (
+            results_sarimax.wpi1_diffuse['initial_variance']
+        )
+        super(TestARIMADiffuse, self).__init__(results_sarimax.wpi1_diffuse,
+                                               **kwargs)
         self.result = self.model.filter()
 
     def test_bse(self):
@@ -633,8 +637,8 @@ def test_init_keys_replicate(self):
         mod1 = self.model
         mod1.update(self.true_params)  # test with side effect ?
 
-        if mod1.initialization == 'approximate_diffuse':
-            raise SkipTest('known failure: init_keys incomplete')
+        # if mod1.initialization == 'approximate_diffuse':
+        #     raise SkipTest('known failure: init_keys incomplete')
         kwargs = self.model._get_init_kwds()
         # TODO: current limitations #2259 comments
         #endog = mod1.endog.squeeze() # test failures, endog may be transformed
@@ -727,8 +731,9 @@ class Test_ar_diffuse(SARIMAXCoverageTest):
     # save_results 7
     def __init__(self, *args, **kwargs):
         kwargs['order'] = (3,0,0)
+        kwargs['initialization'] = 'approximate_diffuse'
+        kwargs['initial_variance'] = 1e9
         super(Test_ar_diffuse, self).__init__(6, *args, **kwargs)
-        self.model.initialize_approximate_diffuse(1e9)
 
 class Test_ar_no_enforce(SARIMAXCoverageTest):
     # // AR: (p,0,0) x (0,0,0,0)
@@ -739,6 +744,7 @@ def __init__(self, *args, **kwargs):
         kwargs['enforce_stationarity'] = False
         kwargs['enforce_invertibility'] = False
         kwargs['initial_variance'] = 1e9
+        # kwargs['loglikelihood_burn'] = 0
         super(Test_ar_no_enforce, self).__init__(6, *args, **kwargs)
         # Reset loglikelihood burn, which gets automatically set to the number
         # of states if enforce_stationarity = False
@@ -870,8 +876,9 @@ class Test_ma_diffuse(SARIMAXCoverageTest):
     # save_results 15
     def __init__(self, *args, **kwargs):
         kwargs['order'] = (0,0,3)
+        kwargs['initialization'] = 'approximate_diffuse'
+        kwargs['initial_variance'] = 1e9
         super(Test_ma_diffuse, self).__init__(14, *args, **kwargs)
-        self.model.initialize_approximate_diffuse(1e9)
 
 class Test_ma_exogenous(SARIMAXCoverageTest):
     # // MAX
@@ -956,10 +963,9 @@ class Test_arma_diffuse(SARIMAXCoverageTest):
     # save_results 24
     def __init__(self, *args, **kwargs):
         kwargs['order'] = (3,0,2)
-        # TODO: diffuse initialization not supported through keywords #2259
+        kwargs['initialization'] = 'approximate_diffuse'
         kwargs['initial_variance'] = 1e9
         super(Test_arma_diffuse, self).__init__(23, *args, **kwargs)
-        self.model.initialize_approximate_diffuse(1e9)
 
 class Test_arma_exogenous(SARIMAXCoverageTest):
     # // ARMAX
@@ -1050,8 +1056,9 @@ class Test_seasonal_ar_diffuse(SARIMAXCoverageTest):
     def __init__(self, *args, **kwargs):
         kwargs['order'] = (0,0,0)
         kwargs['seasonal_order'] = (3,0,0,4)
+        kwargs['initialization'] = 'approximate_diffuse'
+        kwargs['initial_variance'] = 1e9
         super(Test_seasonal_ar_diffuse, self).__init__(31, *args, **kwargs)
-        self.model.initialize_approximate_diffuse(1e9)
 
 class Test_seasonal_ar_exogenous(SARIMAXCoverageTest):
     # // SARX
@@ -1139,8 +1146,9 @@ class Test_seasonal_ma_diffuse(SARIMAXCoverageTest):
     def __init__(self, *args, **kwargs):
         kwargs['order'] = (0,0,0)
         kwargs['seasonal_order'] = (0,0,3,4)
+        kwargs['initialization'] = 'approximate_diffuse'
+        kwargs['initial_variance'] = 1e9
         super(Test_seasonal_ma_diffuse, self).__init__(39, *args, **kwargs)
-        self.model.initialize_approximate_diffuse(1e9)
 
 class Test_seasonal_ma_exogenous(SARIMAXCoverageTest):
     # // SMAX
@@ -1234,8 +1242,9 @@ def __init__(self, *args, **kwargs):
         kwargs['order'] = (0,0,0)
         kwargs['seasonal_order'] = (3,0,2,4)
         kwargs['decimal'] = 3
+        kwargs['initialization'] = 'approximate_diffuse'
+        kwargs['initial_variance'] = 1e9
         super(Test_seasonal_arma_diffuse, self).__init__(48, *args, **kwargs)
-        self.model.initialize_approximate_diffuse(1e9)
 
 class Test_seasonal_arma_exogenous(SARIMAXCoverageTest):
     # // SARMAX
@@ -1282,8 +1291,9 @@ def __init__(self, *args, **kwargs):
         endog = results_sarimax.wpi1_data
         kwargs['exog'] = (endog - np.floor(endog))**2
         kwargs['decimal'] = 2
+        kwargs['initialization'] = 'approximate_diffuse'
+        kwargs['initial_variance'] = 1e9
         super(Test_sarimax_exogenous_diffuse, self).__init__(51, *args, **kwargs)
-        self.model.initialize_approximate_diffuse(1e9)
 
 class Test_arma_exog_trend_polynomial_missing(SARIMAXCoverageTest):
     # // ARMA and exogenous and trend polynomial and missing
@@ -1328,7 +1338,7 @@ def test_simple_time_varying():
 def test_invalid_time_varying():
     assert_raises(ValueError, sarimax.SARIMAX, endog=[1,2,3], mle_regression=True, time_varying_regression=True)
 
-def test_manual_initialization():
+def test_manual_stationary_initialization():
     endog = results_sarimax.wpi1_data
 
     # Create the first model to compare against
@@ -1343,10 +1353,71 @@ def test_manual_initialization():
     mod2.initialize_state()  # a noop in this case (include for coverage)
     res2 = mod2.filter()
 
+    # Create a third model with "known" initialization, but specified in kwargs
+    mod3 = sarimax.SARIMAX(endog, order=(3,0,0),
+                           initialization='known',
+                           initial_state=res1.initial_state,
+                           initial_state_cov=res1.initial_state_cov)
+    mod3.update([0.5,0.2,0.1,1])
+    res3 = mod3.filter()
+
+    # Create the forth model with stationary initialization specified in kwargs
+    mod4 = sarimax.SARIMAX(endog, order=(3,0,0), initialization='stationary')
+    mod4.update([0.5,0.2,0.1,1])
+    res4 = mod4.filter()
+
     # Just test a couple of things to make sure the results are the same
     assert_almost_equal(res1.llf, res2.llf)
     assert_almost_equal(res1.filtered_state, res2.filtered_state)
 
+    assert_almost_equal(res1.llf, res3.llf)
+    assert_almost_equal(res1.filtered_state, res3.filtered_state)
+
+    assert_almost_equal(res1.llf, res4.llf)
+    assert_almost_equal(res1.filtered_state, res4.filtered_state)
+
+def test_manual_approximate_diffuse_initialization():
+    endog = results_sarimax.wpi1_data
+
+    # Create the first model to compare against
+    mod1 = sarimax.SARIMAX(endog, order=(3,0,0))
+    mod1.update([0.5,0.2,0.1,1])
+    mod1.initialize_approximate_diffuse(1e9)
+    res1 = mod1.filter()
+
+    # Create a second model with "known" initialization
+    mod2 = sarimax.SARIMAX(endog, order=(3,0,0))
+    mod2.update([0.5,0.2,0.1,1])
+    mod2.initialize_known(res1.initial_state, res1.initial_state_cov)
+    mod2.initialize_state()  # a noop in this case (include for coverage)
+    res2 = mod2.filter()
+
+    # Create a third model with "known" initialization, but specified in kwargs
+    mod3 = sarimax.SARIMAX(endog, order=(3,0,0),
+                           initialization='known',
+                           initial_state=res1.initial_state,
+                           initial_state_cov=res1.initial_state_cov)
+    mod3.update([0.5,0.2,0.1,1])
+    res3 = mod3.filter()
+
+    # Create the forth model with approximate diffuse initialization specified
+    # in kwargs
+    mod4 = sarimax.SARIMAX(endog, order=(3,0,0),
+                           initialization='approximate_diffuse',
+                           initial_variance=1e9)
+    mod4.update([0.5,0.2,0.1,1])
+    res4 = mod4.filter()
+
+    # Just test a couple of things to make sure the results are the same
+    assert_almost_equal(res1.llf, res2.llf)
+    assert_almost_equal(res1.filtered_state, res2.filtered_state)
+
+    assert_almost_equal(res1.llf, res3.llf)
+    assert_almost_equal(res1.filtered_state, res3.filtered_state)
+
+    assert_almost_equal(res1.llf, res4.llf)
+    assert_almost_equal(res1.filtered_state, res4.filtered_state)
+
 def test_results():
     endog = results_sarimax.wpi1_data