1from numpy.testing import assert_array_almost_equal, assert_array_equal
2from pytest import raises as assert_raises
3
4from numpy import array, transpose, dot, conjugate, zeros_like, empty
5from numpy.random import random
6from scipy.linalg import cholesky, cholesky_banded, cho_solve_banded, \
7     cho_factor, cho_solve
8
9from scipy.linalg._testutils import assert_no_overwrite
10
11
12class TestCholesky:
13
14    def test_simple(self):
15        a = [[8, 2, 3], [2, 9, 3], [3, 3, 6]]
16        c = cholesky(a)
17        assert_array_almost_equal(dot(transpose(c), c), a)
18        c = transpose(c)
19        a = dot(c, transpose(c))
20        assert_array_almost_equal(cholesky(a, lower=1), c)
21
22    def test_check_finite(self):
23        a = [[8, 2, 3], [2, 9, 3], [3, 3, 6]]
24        c = cholesky(a, check_finite=False)
25        assert_array_almost_equal(dot(transpose(c), c), a)
26        c = transpose(c)
27        a = dot(c, transpose(c))
28        assert_array_almost_equal(cholesky(a, lower=1, check_finite=False), c)
29
30    def test_simple_complex(self):
31        m = array([[3+1j, 3+4j, 5], [0, 2+2j, 2+7j], [0, 0, 7+4j]])
32        a = dot(transpose(conjugate(m)), m)
33        c = cholesky(a)
34        a1 = dot(transpose(conjugate(c)), c)
35        assert_array_almost_equal(a, a1)
36        c = transpose(c)
37        a = dot(c, transpose(conjugate(c)))
38        assert_array_almost_equal(cholesky(a, lower=1), c)
39
40    def test_random(self):
41        n = 20
42        for k in range(2):
43            m = random([n, n])
44            for i in range(n):
45                m[i, i] = 20*(.1+m[i, i])
46            a = dot(transpose(m), m)
47            c = cholesky(a)
48            a1 = dot(transpose(c), c)
49            assert_array_almost_equal(a, a1)
50            c = transpose(c)
51            a = dot(c, transpose(c))
52            assert_array_almost_equal(cholesky(a, lower=1), c)
53
54    def test_random_complex(self):
55        n = 20
56        for k in range(2):
57            m = random([n, n])+1j*random([n, n])
58            for i in range(n):
59                m[i, i] = 20*(.1+abs(m[i, i]))
60            a = dot(transpose(conjugate(m)), m)
61            c = cholesky(a)
62            a1 = dot(transpose(conjugate(c)), c)
63            assert_array_almost_equal(a, a1)
64            c = transpose(c)
65            a = dot(c, transpose(conjugate(c)))
66            assert_array_almost_equal(cholesky(a, lower=1), c)
67
68
69class TestCholeskyBanded:
70    """Tests for cholesky_banded() and cho_solve_banded."""
71
72    def test_check_finite(self):
73        # Symmetric positive definite banded matrix `a`
74        a = array([[4.0, 1.0, 0.0, 0.0],
75                   [1.0, 4.0, 0.5, 0.0],
76                   [0.0, 0.5, 4.0, 0.2],
77                   [0.0, 0.0, 0.2, 4.0]])
78        # Banded storage form of `a`.
79        ab = array([[-1.0, 1.0, 0.5, 0.2],
80                    [4.0, 4.0, 4.0, 4.0]])
81        c = cholesky_banded(ab, lower=False, check_finite=False)
82        ufac = zeros_like(a)
83        ufac[list(range(4)), list(range(4))] = c[-1]
84        ufac[(0, 1, 2), (1, 2, 3)] = c[0, 1:]
85        assert_array_almost_equal(a, dot(ufac.T, ufac))
86
87        b = array([0.0, 0.5, 4.2, 4.2])
88        x = cho_solve_banded((c, False), b, check_finite=False)
89        assert_array_almost_equal(x, [0.0, 0.0, 1.0, 1.0])
90
91    def test_upper_real(self):
92        # Symmetric positive definite banded matrix `a`
93        a = array([[4.0, 1.0, 0.0, 0.0],
94                   [1.0, 4.0, 0.5, 0.0],
95                   [0.0, 0.5, 4.0, 0.2],
96                   [0.0, 0.0, 0.2, 4.0]])
97        # Banded storage form of `a`.
98        ab = array([[-1.0, 1.0, 0.5, 0.2],
99                    [4.0, 4.0, 4.0, 4.0]])
100        c = cholesky_banded(ab, lower=False)
101        ufac = zeros_like(a)
102        ufac[list(range(4)), list(range(4))] = c[-1]
103        ufac[(0, 1, 2), (1, 2, 3)] = c[0, 1:]
104        assert_array_almost_equal(a, dot(ufac.T, ufac))
105
106        b = array([0.0, 0.5, 4.2, 4.2])
107        x = cho_solve_banded((c, False), b)
108        assert_array_almost_equal(x, [0.0, 0.0, 1.0, 1.0])
109
110    def test_upper_complex(self):
111        # Hermitian positive definite banded matrix `a`
112        a = array([[4.0, 1.0, 0.0, 0.0],
113                   [1.0, 4.0, 0.5, 0.0],
114                   [0.0, 0.5, 4.0, -0.2j],
115                   [0.0, 0.0, 0.2j, 4.0]])
116        # Banded storage form of `a`.
117        ab = array([[-1.0, 1.0, 0.5, -0.2j],
118                    [4.0, 4.0, 4.0, 4.0]])
119        c = cholesky_banded(ab, lower=False)
120        ufac = zeros_like(a)
121        ufac[list(range(4)), list(range(4))] = c[-1]
122        ufac[(0, 1, 2), (1, 2, 3)] = c[0, 1:]
123        assert_array_almost_equal(a, dot(ufac.conj().T, ufac))
124
125        b = array([0.0, 0.5, 4.0-0.2j, 0.2j + 4.0])
126        x = cho_solve_banded((c, False), b)
127        assert_array_almost_equal(x, [0.0, 0.0, 1.0, 1.0])
128
129    def test_lower_real(self):
130        # Symmetric positive definite banded matrix `a`
131        a = array([[4.0, 1.0, 0.0, 0.0],
132                   [1.0, 4.0, 0.5, 0.0],
133                   [0.0, 0.5, 4.0, 0.2],
134                   [0.0, 0.0, 0.2, 4.0]])
135        # Banded storage form of `a`.
136        ab = array([[4.0, 4.0, 4.0, 4.0],
137                    [1.0, 0.5, 0.2, -1.0]])
138        c = cholesky_banded(ab, lower=True)
139        lfac = zeros_like(a)
140        lfac[list(range(4)), list(range(4))] = c[0]
141        lfac[(1, 2, 3), (0, 1, 2)] = c[1, :3]
142        assert_array_almost_equal(a, dot(lfac, lfac.T))
143
144        b = array([0.0, 0.5, 4.2, 4.2])
145        x = cho_solve_banded((c, True), b)
146        assert_array_almost_equal(x, [0.0, 0.0, 1.0, 1.0])
147
148    def test_lower_complex(self):
149        # Hermitian positive definite banded matrix `a`
150        a = array([[4.0, 1.0, 0.0, 0.0],
151                   [1.0, 4.0, 0.5, 0.0],
152                   [0.0, 0.5, 4.0, -0.2j],
153                   [0.0, 0.0, 0.2j, 4.0]])
154        # Banded storage form of `a`.
155        ab = array([[4.0, 4.0, 4.0, 4.0],
156                    [1.0, 0.5, 0.2j, -1.0]])
157        c = cholesky_banded(ab, lower=True)
158        lfac = zeros_like(a)
159        lfac[list(range(4)), list(range(4))] = c[0]
160        lfac[(1, 2, 3), (0, 1, 2)] = c[1, :3]
161        assert_array_almost_equal(a, dot(lfac, lfac.conj().T))
162
163        b = array([0.0, 0.5j, 3.8j, 3.8])
164        x = cho_solve_banded((c, True), b)
165        assert_array_almost_equal(x, [0.0, 0.0, 1.0j, 1.0])
166
167
168class TestOverwrite:
169    def test_cholesky(self):
170        assert_no_overwrite(cholesky, [(3, 3)])
171
172    def test_cho_factor(self):
173        assert_no_overwrite(cho_factor, [(3, 3)])
174
175    def test_cho_solve(self):
176        x = array([[2, -1, 0], [-1, 2, -1], [0, -1, 2]])
177        xcho = cho_factor(x)
178        assert_no_overwrite(lambda b: cho_solve(xcho, b), [(3,)])
179
180    def test_cholesky_banded(self):
181        assert_no_overwrite(cholesky_banded, [(2, 3)])
182
183    def test_cho_solve_banded(self):
184        x = array([[0, -1, -1], [2, 2, 2]])
185        xcho = cholesky_banded(x)
186        assert_no_overwrite(lambda b: cho_solve_banded((xcho, False), b),
187                            [(3,)])
188
189
190class TestEmptyArray:
191    def test_cho_factor_empty_square(self):
192        a = empty((0, 0))
193        b = array([])
194        c = array([[]])
195        d = []
196        e = [[]]
197
198        x, _ = cho_factor(a)
199        assert_array_equal(x, a)
200
201        for x in ([b, c, d, e]):
202            assert_raises(ValueError, cho_factor, x)
203