import inspect import pytest import numpy as np from numpy.testing import assert_allclose from scipy import linalg, sparse real_floating = [np.float32, np.float64] complex_floating = [np.complex64, np.complex128] floating = real_floating + complex_floating def get_random(shape, *, dtype, rng): A = rng.random(shape) if np.issubdtype(dtype, np.complexfloating): A = A + rng.random(shape) * 1j return A.astype(dtype) def get_nearly_hermitian(shape, dtype, atol, rng): # Generate a batch of nearly Hermitian matrices with specified # `shape` and `dtype`. `atol` controls the level of noise in # Hermitian-ness to by generated by `rng`. A = rng.random(shape).astype(dtype) At = np.conj(A.swapaxes(-1, -2)) noise = rng.standard_normal(size=A.shape).astype(dtype) * atol return A + At + noise class TestBatch: # Test batch support for most linalg functions def batch_test(self, fun, arrays, *, core_dim=2, n_out=1, kwargs=None, dtype=None, broadcast=True, check_kwargs=True): # Check that all outputs of batched call `fun(A, **kwargs)` are the same # as if we loop over the separate vectors/matrices in `A`. Also check # that `fun` accepts `A` by position or keyword and that results are # identical. This is important because the name of the array argument # is manually specified to the decorator, and it's easy to mess up. # However, this makes it hard to test positional arguments passed # after the array, so we test that separately for a few functions to # make sure the decorator is working as it should. kwargs = {} if kwargs is None else kwargs parameters = list(inspect.signature(fun).parameters.keys()) arrays = (arrays,) if not isinstance(arrays, tuple) else arrays # Identical results when passing argument by keyword or position res2 = fun(*arrays, **kwargs) if check_kwargs: res1 = fun(**dict(zip(parameters, arrays)), **kwargs) for out1, out2 in zip(res1, res2): # even a single array is iterable... np.testing.assert_equal(out1, out2) # Check results vs looping over res = (res2,) if n_out == 1 else res2 # This is not the general behavior (only batch dimensions get # broadcasted by the decorator) but it's easier for testing. if broadcast: arrays = np.broadcast_arrays(*arrays) batch_shape = arrays[0].shape[:-core_dim] for i in range(batch_shape[0]): for j in range(batch_shape[1]): arrays_ij = (array[i, j] for array in arrays) ref = fun(*arrays_ij, **kwargs) ref = ((np.asarray(ref),) if n_out == 1 else tuple(np.asarray(refk) for refk in ref)) for k in range(n_out): assert_allclose(res[k][i, j], ref[k]) assert np.shape(res[k][i, j]) == ref[k].shape for k in range(len(ref)): out_dtype = ref[k].dtype if dtype is None else dtype assert res[k].dtype == out_dtype return res2 # return original, non-tuplized result @pytest.mark.parametrize('dtype', floating) def test_expm_cond(self, dtype): rng = np.random.default_rng(8342310302941288912051) A = rng.random((5, 3, 4, 4)).astype(dtype) self.batch_test(linalg.expm_cond, A) @pytest.mark.parametrize('dtype', floating) def test_issymmetric(self, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_nearly_hermitian((5, 3, 4, 4), dtype, 3e-4, rng) res = self.batch_test(linalg.issymmetric, A, kwargs=dict(atol=1e-3)) assert not np.all(res) # ensure test is not trivial: not all True or False; assert np.any(res) # also confirms that `atol` is passed to issymmetric @pytest.mark.parametrize('dtype', floating) def test_ishermitian(self, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_nearly_hermitian((5, 3, 4, 4), dtype, 3e-4, rng) res = self.batch_test(linalg.ishermitian, A, kwargs=dict(atol=1e-3)) assert not np.all(res) # ensure test is not trivial: not all True or False; assert np.any(res) # also confirms that `atol` is passed to ishermitian @pytest.mark.parametrize('dtype', floating) def test_diagsvd(self, dtype): rng = np.random.default_rng(8342310302941288912051) A = rng.random((5, 3, 4)).astype(dtype) res1 = self.batch_test(linalg.diagsvd, A, kwargs=dict(M=6, N=4), core_dim=1) # test that `M, N` can be passed by position res2 = linalg.diagsvd(A, 6, 4) np.testing.assert_equal(res1, res2) @pytest.mark.parametrize('fun', [linalg.inv, linalg.sqrtm, linalg.signm, linalg.sinm, linalg.cosm, linalg.tanhm, linalg.sinhm, linalg.coshm, linalg.tanhm, linalg.pinv, linalg.pinvh, linalg.orth]) @pytest.mark.parametrize('dtype', floating) def test_matmat(self, fun, dtype): # matrix in, matrix out rng = np.random.default_rng(8342310302941288912051) A = get_random((5, 3, 4, 4), dtype=dtype, rng=rng) # sqrtm can return complex output for real input resulting in i/o type # mismatch. Nudge the eigenvalues to positive side to avoid this. if fun == linalg.sqrtm: A = A + 3*np.eye(4, dtype=dtype) self.batch_test(fun, A) @pytest.mark.parametrize('dtype', floating) def test_null_space(self, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((5, 3, 4, 6), dtype=dtype, rng=rng) self.batch_test(linalg.null_space, A) @pytest.mark.parametrize('dtype', floating) def test_funm(self, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((2, 4, 3, 3), dtype=dtype, rng=rng) self.batch_test(linalg.funm, A, kwargs=dict(func=np.sin)) @pytest.mark.parametrize('dtype', floating) def test_fractional_matrix_power(self, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((2, 4, 3, 3), dtype=dtype, rng=rng) res1 = self.batch_test(linalg.fractional_matrix_power, A, kwargs={'t':1.5}) # test that `t` can be passed by position res2 = linalg.fractional_matrix_power(A, 1.5) np.testing.assert_equal(res1, res2) @pytest.mark.parametrize('dtype', floating) def test_logm(self, dtype): # One test failed absolute tolerance with default random seed rng = np.random.default_rng(89940026998903887141749720079406074936) A = get_random((5, 3, 4, 4), dtype=dtype, rng=rng) A = A + 3*np.eye(4) # avoid complex output for real input res1 = self.batch_test(linalg.logm, A) # test that `disp` can be passed by position res2 = linalg.logm(A) for res1i, res2i in zip(res1, res2): np.testing.assert_equal(res1i, res2i) @pytest.mark.parametrize('dtype', floating) def test_pinv(self, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((5, 3, 4, 4), dtype=dtype, rng=rng) self.batch_test(linalg.pinv, A, n_out=2, kwargs=dict(return_rank=True)) @pytest.mark.parametrize('dtype', floating) def test_matrix_balance(self, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((5, 3, 4, 4), dtype=dtype, rng=rng) self.batch_test(linalg.matrix_balance, A, n_out=2) self.batch_test(linalg.matrix_balance, A, n_out=2, kwargs={'separate':True}) @pytest.mark.parametrize('dtype', floating) def test_bandwidth(self, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((4, 4), dtype=dtype, rng=rng) A = np.asarray([np.triu(A, k) for k in range(-3, 3)]).reshape((2, 3, 4, 4)) self.batch_test(linalg.bandwidth, A, n_out=2) @pytest.mark.parametrize('fun_n_out', [(linalg.cholesky, 1), (linalg.ldl, 3), (linalg.cho_factor, 2)]) @pytest.mark.parametrize('dtype', floating) def test_ldl_cholesky(self, fun_n_out, dtype): rng = np.random.default_rng(8342310302941288912051) fun, n_out = fun_n_out A = get_nearly_hermitian((5, 3, 4, 4), dtype, 0, rng) # exactly Hermitian A = A + 4*np.eye(4, dtype=dtype) # ensure positive definite for Cholesky self.batch_test(fun, A, n_out=n_out) @pytest.mark.parametrize('compute_uv', [False, True]) @pytest.mark.parametrize('dtype', floating) def test_svd(self, compute_uv, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((5, 3, 2, 4), dtype=dtype, rng=rng) n_out = 3 if compute_uv else 1 self.batch_test(linalg.svd, A, n_out=n_out, kwargs=dict(compute_uv=compute_uv)) @pytest.mark.parametrize('fun', [linalg.polar, linalg.qr, linalg.rq]) @pytest.mark.parametrize('dtype', floating) def test_polar_qr_rq(self, fun, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((5, 3, 2, 4), dtype=dtype, rng=rng) self.batch_test(fun, A, n_out=2) @pytest.mark.parametrize('cdim', [(5,), (5, 4), (2, 3, 5, 4)]) @pytest.mark.parametrize('dtype', floating) def test_qr_multiply(self, cdim, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((2, 3, 5, 5), dtype=dtype, rng=rng) c = get_random(cdim, dtype=dtype, rng=rng) res = linalg.qr_multiply(A, c, mode='left') q, r = linalg.qr(A) ref = q @ c atol = 1e-6 if dtype in {np.float32, np.complex64} else 1e-12 assert_allclose(res[0], ref, atol=atol) assert_allclose(res[1], r, atol=atol) @pytest.mark.parametrize('uvdim', [[(5,), (3,)], [(4, 5, 2), (4, 3, 2)]]) @pytest.mark.parametrize('dtype', floating) def test_qr_update(self, uvdim, dtype): rng = np.random.default_rng(8342310302941288912051) udim, vdim = uvdim A = get_random((4, 5, 3), dtype=dtype, rng=rng) u = get_random(udim, dtype=dtype, rng=rng) v = get_random(vdim, dtype=dtype, rng=rng) q, r = linalg.qr(A) res = linalg.qr_update(q, r, u, v) for i in range(4): qi, ri = q[i], r[i] ui, vi = (u, v) if u.ndim == 1 else (u[i], v[i]) ref_i = linalg.qr_update(qi, ri, ui, vi) assert_allclose(res[0][i], ref_i[0]) assert_allclose(res[1][i], ref_i[1]) @pytest.mark.parametrize('udim', [(5,), (4, 3, 5)]) @pytest.mark.parametrize('kdim', [(), (4,)]) @pytest.mark.parametrize('dtype', floating) def test_qr_insert(self, udim, kdim, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((4, 5, 5), dtype=dtype, rng=rng) u = get_random(udim, dtype=dtype, rng=rng) k = rng.integers(0, 5, size=kdim) q, r = linalg.qr(A) res = linalg.qr_insert(q, r, u, k) for i in range(4): qi, ri = q[i], r[i] ki = k if k.ndim == 0 else k[i] ui = u if u.ndim == 1 else u[i] ref_i = linalg.qr_insert(qi, ri, ui, ki) assert_allclose(res[0][i], ref_i[0]) assert_allclose(res[1][i], ref_i[1]) @pytest.mark.parametrize('kdim', [(), (4,)]) @pytest.mark.parametrize('dtype', floating) def test_qr_delete(self, kdim, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((4, 5, 5), dtype=dtype, rng=rng) k = rng.integers(0, 4, size=kdim) q, r = linalg.qr(A) res = linalg.qr_delete(q, r, k) for i in range(4): qi, ri = q[i], r[i] ki = k if k.ndim == 0 else k[i] ref_i = linalg.qr_delete(qi, ri, ki) assert_allclose(res[0][i], ref_i[0]) assert_allclose(res[1][i], ref_i[1]) @pytest.mark.parametrize('fun', [linalg.schur, linalg.lu_factor]) @pytest.mark.parametrize('dtype', floating) def test_schur_lu(self, fun, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((5, 3, 4, 4), dtype=dtype, rng=rng) self.batch_test(fun, A, n_out=2) @pytest.mark.parametrize('calc_q', [False, True]) @pytest.mark.parametrize('dtype', floating) def test_hessenberg(self, calc_q, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((5, 3, 4, 4), dtype=dtype, rng=rng) n_out = 2 if calc_q else 1 self.batch_test(linalg.hessenberg, A, n_out=n_out, kwargs=dict(calc_q=calc_q)) @pytest.mark.parametrize('eigvals_only', [False, True]) @pytest.mark.parametrize('dtype', floating) def test_eig_banded(self, eigvals_only, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((5, 3, 4, 4), dtype=dtype, rng=rng) n_out = 1 if eigvals_only else 2 self.batch_test(linalg.eig_banded, A, n_out=n_out, kwargs=dict(eigvals_only=eigvals_only)) @pytest.mark.parametrize('dtype', floating) def test_eigvals_banded(self, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((5, 3, 4, 4), dtype=dtype, rng=rng) self.batch_test(linalg.eigvals_banded, A) @pytest.mark.parametrize('two_in', [False, True]) @pytest.mark.parametrize('fun_n_nout', [(linalg.eigh, 1), (linalg.eigh, 2), (linalg.eigvalsh, 1), (linalg.eigvals, 1)]) @pytest.mark.parametrize('dtype', floating) def test_eigh(self, two_in, fun_n_nout, dtype): rng = np.random.default_rng(8342310302941288912051) fun, n_out = fun_n_nout A = get_nearly_hermitian((1, 3, 4, 4), dtype, 0, rng) # exactly Hermitian B = get_nearly_hermitian((2, 1, 4, 4), dtype, 0, rng) # exactly Hermitian B = B + 4*np.eye(4).astype(dtype) # needs to be positive definite args = (A, B) if two_in else (A,) kwargs = dict(eigvals_only=True) if (n_out == 1 and fun==linalg.eigh) else {} self.batch_test(fun, args, n_out=n_out, kwargs=kwargs) @pytest.mark.parametrize('compute_expm', [False, True]) @pytest.mark.parametrize('dtype', floating) def test_expm_frechet(self, compute_expm, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((1, 3, 4, 4), dtype=dtype, rng=rng) E = get_random((2, 1, 4, 4), dtype=dtype, rng=rng) n_out = 2 if compute_expm else 1 self.batch_test(linalg.expm_frechet, (A, E), n_out=n_out, kwargs=dict(compute_expm=compute_expm)) @pytest.mark.parametrize('dtype', floating) def test_subspace_angles(self, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((1, 3, 4, 3), dtype=dtype, rng=rng) B = get_random((2, 1, 4, 3), dtype=dtype, rng=rng) self.batch_test(linalg.subspace_angles, (A, B)) # just to show that A and B don't need to be broadcastable M, N, K = 4, 5, 3 A = get_random((1, 3, M, N), dtype=dtype, rng=rng) B = get_random((2, 1, M, K), dtype=dtype, rng=rng) assert linalg.subspace_angles(A, B).shape == (2, 3, min(N, K)) @pytest.mark.parametrize('fun', [linalg.svdvals]) @pytest.mark.parametrize('dtype', floating) def test_svdvals(self, fun, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((2, 3, 4, 5), dtype=dtype, rng=rng) self.batch_test(fun, A) @pytest.mark.parametrize('fun_n_out', [(linalg.orthogonal_procrustes, 2), (linalg.khatri_rao, 1), (linalg.solve_continuous_lyapunov, 1), (linalg.solve_discrete_lyapunov, 1), (linalg.qz, 4), (linalg.ordqz, 6)]) @pytest.mark.parametrize('dtype', floating) def test_two_generic_matrix_inputs(self, fun_n_out, dtype): rng = np.random.default_rng(8342310302941288912051) fun, n_out = fun_n_out A = get_random((2, 3, 4, 4), dtype=dtype, rng=rng) B = get_random((2, 3, 4, 4), dtype=dtype, rng=rng) self.batch_test(fun, (A, B), n_out=n_out) @pytest.mark.parametrize('dtype', floating) def test_cossin(self, dtype): rng = np.random.default_rng(8342310302941288912051) p, q = 3, 4 X = get_random((2, 3, 10, 10), dtype=dtype, rng=rng) x11, x12, x21, x22 = (X[..., :p, :q], X[..., :p, q:], X[..., p:, :q], X[..., p:, q:]) res = linalg.cossin(X, p, q) ref = linalg.cossin((x11, x12, x21, x22)) for res_i, ref_i in zip(res, ref): np.testing.assert_equal(res_i, ref_i) for j in range(2): for k in range(3): ref_jk = linalg.cossin(X[j, k], p, q) for res_i, ref_ijk in zip(res, ref_jk): np.testing.assert_equal(res_i[j, k], ref_ijk) @pytest.mark.parametrize('dtype', floating) def test_sylvester(self, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((2, 3, 5, 5), dtype=dtype, rng=rng) B = get_random((2, 3, 5, 5), dtype=dtype, rng=rng) C = get_random((2, 3, 5, 5), dtype=dtype, rng=rng) self.batch_test(linalg.solve_sylvester, (A, B, C)) @pytest.mark.parametrize('fun', [linalg.solve_continuous_are, linalg.solve_discrete_are]) @pytest.mark.parametrize('dtype', floating) def test_are(self, fun, dtype): rng = np.random.default_rng(8342310302941288912051) a = get_random((2, 3, 5, 5), dtype=dtype, rng=rng) b = get_random((2, 3, 5, 5), dtype=dtype, rng=rng) q = get_nearly_hermitian((2, 3, 5, 5), dtype=dtype, atol=0, rng=rng) r = get_nearly_hermitian((2, 3, 5, 5), dtype=dtype, atol=0, rng=rng) a = a + 5*np.eye(5) # making these positive definite seems to help b = b + 5*np.eye(5) q = q + 5*np.eye(5) r = r + 5*np.eye(5) # can't easily generate valid random e, s self.batch_test(fun, (a, b, q, r)) @pytest.mark.parametrize('dtype', floating) def test_rsf2cs(self, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((2, 3, 4, 4), dtype=dtype, rng=rng) T, Z = linalg.schur(A) self.batch_test(linalg.rsf2csf, (T, Z), n_out=2) @pytest.mark.parametrize('dtype', floating) def test_cholesky_banded(self, dtype): rng = np.random.default_rng(8342310302941288912051) ab = get_random((5, 4, 3, 6), dtype=dtype, rng=rng) ab[..., -1, :] = 10 # make diagonal dominant self.batch_test(linalg.cholesky_banded, ab) @pytest.mark.parametrize('dtype', floating) def test_block_diag(self, dtype): rng = np.random.default_rng(8342310302941288912051) a = get_random((1, 3, 1, 3), dtype=dtype, rng=rng) b = get_random((2, 1, 3, 6), dtype=dtype, rng=rng) c = get_random((1, 1, 3, 2), dtype=dtype, rng=rng) # batch_test doesn't have the logic to broadcast just the batch shapes, # so do it manually. a2 = np.broadcast_to(a, (2, 3, 1, 3)) b2 = np.broadcast_to(b, (2, 3, 3, 6)) c2 = np.broadcast_to(c, (2, 3, 3, 2)) ref = self.batch_test(linalg.block_diag, (a2, b2, c2), check_kwargs=False, broadcast=False) # Check that `block_diag` broadcasts the batch shapes as expected. res = linalg.block_diag(a, b, c) assert_allclose(res, ref) @pytest.mark.parametrize('fun_n_out', [(linalg.eigh_tridiagonal, 2), (linalg.eigvalsh_tridiagonal, 1)]) @pytest.mark.parametrize('dtype', real_floating) # "Only real arrays currently supported" def test_eigh_tridiagonal(self, fun_n_out, dtype): rng = np.random.default_rng(8342310302941288912051) fun, n_out = fun_n_out d = get_random((3, 4, 5), dtype=dtype, rng=rng) e = get_random((3, 4, 4), dtype=dtype, rng=rng) self.batch_test(fun, (d, e), core_dim=1, n_out=n_out, broadcast=False) @pytest.mark.parametrize('bdim', [(5,), (5, 4), (2, 3, 5, 4)]) @pytest.mark.parametrize('dtype', floating) def test_solve(self, bdim, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((2, 3, 5, 5), dtype=dtype, rng=rng) b = get_random(bdim, dtype=dtype, rng=rng) x = linalg.solve(A, b) if len(bdim) == 1: x = x[..., np.newaxis] b = b[..., np.newaxis] assert_allclose(A @ x - b, 0, atol=1.5e-6) assert_allclose(x, np.linalg.solve(A, b), atol=3e-6) @pytest.mark.parametrize('bdim', [(5,), (5, 4), (2, 3, 5, 4)]) @pytest.mark.parametrize('dtype', floating) def test_lu_solve(self, bdim, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((2, 3, 5, 5), dtype=dtype, rng=rng) b = get_random(bdim, dtype=dtype, rng=rng) lu_and_piv = linalg.lu_factor(A) x = linalg.lu_solve(lu_and_piv, b) if len(bdim) == 1: x = x[..., np.newaxis] b = b[..., np.newaxis] assert_allclose(A @ x - b, 0, atol=1.5e-6) assert_allclose(x, np.linalg.solve(A, b), atol=3e-6) @pytest.mark.parametrize('l_and_u', [(1, 1), ([2, 1, 0], [0, 1 , 2])]) @pytest.mark.parametrize('bdim', [(5,), (5, 4), (2, 3, 5, 4)]) @pytest.mark.parametrize('dtype', floating) def test_solve_banded(self, l_and_u, bdim, dtype): rng = np.random.default_rng(8342310302941288912051) l, u = l_and_u ab = get_random((2, 3, 3, 5), dtype=dtype, rng=rng) b = get_random(bdim, dtype=dtype, rng=rng) x = linalg.solve_banded((l, u), ab, b) for i in range(2): for j in range(3): bij = b if len(bdim) <= 2 else b[i, j] lj = l if np.ndim(l) == 0 else l[j] uj = u if np.ndim(u) == 0 else u[j] xij = linalg.solve_banded((lj, uj), ab[i, j], bij) assert_allclose(x[i, j], xij) # Can uncomment when `solve_toeplitz` deprecation is done (SciPy 1.17) # @pytest.mark.parametrize('separate_r', [False, True]) # @pytest.mark.parametrize('bdim', [(5,), (5, 4), (2, 3, 5, 4)]) # @pytest.mark.parametrize('dtype', floating) # def test_solve_toeplitz(self, separate_r, bdim, dtype): # rng = np.random.default_rng(8342310302941288912051) # c = get_random((2, 3, 5), dtype=dtype, rng=rng) # r = get_random((2, 3, 5), dtype=dtype, rng=rng) # c_or_cr = (c, r) if separate_r else c # b = get_random(bdim, dtype=dtype, rng=rng) # x = linalg.solve_toeplitz(c_or_cr, b) # for i in range(2): # for j in range(3): # bij = b if len(bdim) <= 2 else b[i, j] # c_or_cr_ij = (c[i, j], r[i, j]) if separate_r else c[i, j] # xij = linalg.solve_toeplitz(c_or_cr_ij, bij) # assert_allclose(x[i, j], xij) @pytest.mark.parametrize('bdim', [(5,), (5, 4), (2, 3, 5, 4)]) @pytest.mark.parametrize('dtype', floating) def test_cho_solve(self, bdim, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_nearly_hermitian((2, 3, 5, 5), dtype=dtype, atol=0, rng=rng) A = A + 5*np.eye(5) c_and_lower = linalg.cho_factor(A) b = get_random(bdim, dtype=dtype, rng=rng) x = linalg.cho_solve(c_and_lower, b) if len(bdim) == 1: x = x[..., np.newaxis] b = b[..., np.newaxis] assert_allclose(A @ x - b, 0, atol=1e-6) assert_allclose(x, np.linalg.solve(A, b), atol=2e-6) @pytest.mark.parametrize('lower', [False, True]) @pytest.mark.parametrize('bdim', [(5,), (5, 4), (2, 3, 5, 4)]) @pytest.mark.parametrize('dtype', floating) def test_cho_solve_banded(self, lower, bdim, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((2, 3, 3, 5), dtype=dtype, rng=rng) row_diag = 0 if lower else -1 A[:, :, row_diag] = 10 cb = linalg.cholesky_banded(A, lower=lower) b = get_random(bdim, dtype=dtype, rng=rng) x = linalg.cho_solve_banded((cb, lower), b) for i in range(2): for j in range(3): bij = b if len(bdim) <= 2 else b[i, j] xij = linalg.cho_solve_banded((cb[i, j], lower), bij) assert_allclose(x[i, j], xij) @pytest.mark.parametrize('bdim', [(5,), (5, 4), (2, 3, 5, 4)]) @pytest.mark.parametrize('dtype', floating) def test_solveh_banded(self, bdim, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((2, 3, 3, 5), dtype=dtype, rng=rng) A[:, :, -1] = 10 b = get_random(bdim, dtype=dtype, rng=rng) x = linalg.solveh_banded(A, b) for i in range(2): for j in range(3): bij = b if len(bdim) <= 2 else b[i, j] xij = linalg.solveh_banded(A[i, j], bij) assert_allclose(x[i, j], xij) @pytest.mark.parametrize('bdim', [(5,), (5, 4), (2, 3, 5, 4)]) @pytest.mark.parametrize('dtype', floating) def test_solve_triangular(self, bdim, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((2, 3, 5, 5), dtype=dtype, rng=rng) A = np.tril(A) b = get_random(bdim, dtype=dtype, rng=rng) x = linalg.solve_triangular(A, b, lower=True) if len(bdim) == 1: x = x[..., np.newaxis] b = b[..., np.newaxis] atol = 1e-10 if dtype in (np.complex128, np.float64) else 2e-4 assert_allclose(A @ x - b, 0, atol=atol) assert_allclose(x, np.linalg.solve(A, b), atol=5*atol) @pytest.mark.parametrize('bdim', [(4,), (4, 3), (2, 3, 4, 3)]) @pytest.mark.parametrize('dtype', floating) def test_lstsq(self, bdim, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((2, 3, 4, 5), dtype=dtype, rng=rng) b = get_random(bdim, dtype=dtype, rng=rng) res = linalg.lstsq(A, b) x = res[0] if len(bdim) == 1: x = x[..., np.newaxis] b = b[..., np.newaxis] assert_allclose(A @ x - b, 0, atol=2e-6) assert len(res) == 4 @pytest.mark.parametrize('dtype', floating) def test_clarkson_woodruff_transform(self, dtype): rng = np.random.default_rng(8342310302941288912051) A = get_random((5, 3, 4, 6), dtype=dtype, rng=rng) self.batch_test(linalg.clarkson_woodruff_transform, A, kwargs=dict(sketch_size=3, rng=311224)) def test_clarkson_woodruff_transform_sparse(self): rng = np.random.default_rng(8342310302941288912051) A = get_random((5, 3, 4, 6), dtype=np.float64, rng=rng) A = sparse.coo_array(A) message = "Batch support for sparse arrays is not available." with pytest.raises(NotImplementedError, match=message): linalg.clarkson_woodruff_transform(A, sketch_size=3, rng=rng)