def test_invalid():
    # Invalid initializations (also tests for some invalid calls to set)
    assert_raises(ValueError, Initialization, 5, '')
    assert_raises(ValueError, Initialization, 5, 'stationary', constant=[1, 2])
    assert_raises(ValueError, Initialization, 5, 'stationary',
                  stationary_cov=[1, 2])
    assert_raises(ValueError, Initialization, 5, 'known')
    assert_raises(ValueError, Initialization, 5, 'known', constant=[1])
    assert_raises(ValueError, Initialization, 5, 'known', stationary_cov=[0])

    # Invalid set() / unset() calls
    init = Initialization(5)
    assert_raises(ValueError, init.set, -1, 'diffuse')
    assert_raises(ValueError, init.unset, -1)
    assert_raises(ValueError, init.set, 5, 'diffuse')
    assert_raises(ValueError, init.set, 'x', 'diffuse')
    assert_raises(ValueError, init.unset, 'x')
    assert_raises(ValueError, init.set, (1, 2, 3), 'diffuse')
    assert_raises(ValueError, init.unset, (1, 2, 3))
    init.set(None, 'diffuse')
    assert_raises(ValueError, init.set, 1, 'diffuse')
    init.clear()
    init.set(1, 'diffuse')
    assert_raises(ValueError, init.set, None, 'stationary')

    init.clear()
    assert_raises(ValueError, init.unset, 1)

    # Invalid __call__
    init = Initialization(2)
    assert_raises(ValueError, init)
    init = Initialization(2, 'stationary')
    assert_raises(ValueError, init)
def test_invalid():
    # Invalid initializations (also tests for some invalid calls to set)
    assert_raises(ValueError, Initialization, 5, '')
    assert_raises(ValueError, Initialization, 5, 'stationary', constant=[1, 2])
    assert_raises(ValueError,
                  Initialization,
                  5,
                  'stationary',
                  stationary_cov=[1, 2])
    assert_raises(ValueError, Initialization, 5, 'known')
    assert_raises(ValueError, Initialization, 5, 'known', constant=[1])
    assert_raises(ValueError, Initialization, 5, 'known', stationary_cov=[0])

    # Invalid set() / unset() calls
    init = Initialization(5)
    assert_raises(ValueError, init.set, -1, 'diffuse')
    assert_raises(ValueError, init.unset, -1)
    assert_raises(ValueError, init.set, 5, 'diffuse')
    assert_raises(ValueError, init.set, 'x', 'diffuse')
    assert_raises(ValueError, init.unset, 'x')
    assert_raises(ValueError, init.set, (1, 2, 3), 'diffuse')
    assert_raises(ValueError, init.unset, (1, 2, 3))
    init.set(None, 'diffuse')
    assert_raises(ValueError, init.set, 1, 'diffuse')
    init.clear()
    init.set(1, 'diffuse')
    assert_raises(ValueError, init.set, None, 'stationary')

    init.clear()
    assert_raises(ValueError, init.unset, 1)

    # Invalid __call__
    init = Initialization(2)
    assert_raises(ValueError, init)
    init = Initialization(2, 'stationary')
    assert_raises(ValueError, init)
def test_mixed_stationary():
    # More specific tests when one or more blocks are initialized as stationary
    endog = np.zeros(10)
    mod = sarimax.SARIMAX(endog, order=(2, 1, 0))
    phi = [0.5, -0.2]
    sigma2 = 2.
    mod.update(np.r_[phi, sigma2])

    init = Initialization(mod.k_states)
    init.set(0, 'diffuse')
    init.set((1, 3), 'stationary')
    desired_cov = np.zeros((3, 3))
    T = np.array([[0.5, 1],
                  [-0.2, 0]])
    Q = np.diag([sigma2, 0])
    desired_cov[1:, 1:] = solve_discrete_lyapunov(T, Q)
    check_initialization(mod, init, [0, 0, 0], np.diag([1, 0, 0]), desired_cov)

    init.clear()
    init.set(0, 'diffuse')
    init.set(1, 'stationary')
    init.set(2, 'approximate_diffuse')
    T = np.array([[0.5]])
    Q = np.diag([sigma2])
    desired_cov = np.diag([0, np.squeeze(solve_discrete_lyapunov(T, Q)), 1e6])
    check_initialization(mod, init, [0, 0, 0], np.diag([1, 0, 0]), desired_cov)

    init.clear()
    init.set(0, 'diffuse')
    init.set(1, 'stationary')
    init.set(2, 'stationary')
    desired_cov[2, 2] = 0
    check_initialization(mod, init, [0, 0, 0], np.diag([1, 0, 0]), desired_cov)

    # Test with a VAR model
    endog = np.zeros((10, 2))
    mod = varmax.VARMAX(endog, order=(1, 0), )
    intercept = [1.5, -0.1]
    transition = np.array([[0.5, -0.2],
                           [0.1, 0.8]])
    cov = np.array([[1.2, -0.4],
                    [-0.4, 0.4]])
    tril = np.tril_indices(2)
    params = np.r_[intercept, transition.ravel(),
                   np.linalg.cholesky(cov)[tril]]
    mod.update(params)

    # > stationary, global
    init = Initialization(mod.k_states, 'stationary')
    desired_intercept = np.linalg.solve(np.eye(2) - transition, intercept)
    desired_cov = solve_discrete_lyapunov(transition, cov)
    check_initialization(mod, init, desired_intercept, np.diag([0, 0]),
                         desired_cov)

    # > diffuse, global
    init.set(None, 'diffuse')
    check_initialization(mod, init, [0, 0], np.eye(2), np.diag([0, 0]))

    # > stationary, individually
    init.unset(None)
    init.set(0, 'stationary')
    init.set(1, 'stationary')
    a, Pinf, Pstar = init(model=mod)
    desired_intercept = [intercept[0] / (1 - transition[0, 0]),
                         intercept[1] / (1 - transition[1, 1])]
    desired_cov = np.diag([cov[0, 0] / (1 - transition[0, 0]**2),
                           cov[1, 1] / (1 - transition[1, 1]**2)])
    check_initialization(mod, init, desired_intercept, np.diag([0, 0]),
                         desired_cov)
def test_mixed_stationary():
    # More specific tests when one or more blocks are initialized as stationary
    endog = np.zeros(10)
    mod = sarimax.SARIMAX(endog, order=(2, 1, 0))
    phi = [0.5, -0.2]
    sigma2 = 2.
    mod.update(np.r_[phi, sigma2])

    init = Initialization(mod.k_states)
    init.set(0, 'diffuse')
    init.set((1, 3), 'stationary')
    desired_cov = np.zeros((3, 3))
    T = np.array([[0.5, 1], [-0.2, 0]])
    Q = np.diag([sigma2, 0])
    desired_cov[1:, 1:] = solve_discrete_lyapunov(T, Q)
    check_initialization(mod, init, [0, 0, 0], np.diag([1, 0, 0]), desired_cov)

    init.clear()
    init.set(0, 'diffuse')
    init.set(1, 'stationary')
    init.set(2, 'approximate_diffuse')
    T = np.array([[0.5]])
    Q = np.diag([sigma2])
    desired_cov = np.diag([0, np.squeeze(solve_discrete_lyapunov(T, Q)), 1e6])
    check_initialization(mod, init, [0, 0, 0], np.diag([1, 0, 0]), desired_cov)

    init.clear()
    init.set(0, 'diffuse')
    init.set(1, 'stationary')
    init.set(2, 'stationary')
    desired_cov[2, 2] = 0
    check_initialization(mod, init, [0, 0, 0], np.diag([1, 0, 0]), desired_cov)

    # Test with a VAR model
    endog = np.zeros((10, 2))
    mod = varmax.VARMAX(
        endog,
        order=(1, 0),
    )
    intercept = [1.5, -0.1]
    transition = np.array([[0.5, -0.2], [0.1, 0.8]])
    cov = np.array([[1.2, -0.4], [-0.4, 0.4]])
    tril = np.tril_indices(2)
    params = np.r_[intercept,
                   transition.ravel(),
                   np.linalg.cholesky(cov)[tril]]
    mod.update(params)

    # > stationary, global
    init = Initialization(mod.k_states, 'stationary')
    desired_intercept = np.linalg.solve(np.eye(2) - transition, intercept)
    desired_cov = solve_discrete_lyapunov(transition, cov)
    check_initialization(mod, init, desired_intercept, np.diag([0, 0]),
                         desired_cov)

    # > diffuse, global
    init.set(None, 'diffuse')
    check_initialization(mod, init, [0, 0], np.eye(2), np.diag([0, 0]))

    # > stationary, individually
    init.unset(None)
    init.set(0, 'stationary')
    init.set(1, 'stationary')
    a, Pinf, Pstar = init(model=mod)
    desired_intercept = [
        intercept[0] / (1 - transition[0, 0]),
        intercept[1] / (1 - transition[1, 1])
    ]
    desired_cov = np.diag([
        cov[0, 0] / (1 - transition[0, 0]**2),
        cov[1, 1] / (1 - transition[1, 1]**2)
    ])
    check_initialization(mod, init, desired_intercept, np.diag([0, 0]),
                         desired_cov)