def test_x_operator(n_col): rng = np.random.RandomState(0) X = rng.randn(11, 8) X_m = rng.randn(8) sqrt_sw = rng.randn(len(X)) Y = rng.randn(11, *n_col) A = rng.randn(9, *n_col) operator = _X_operator(sp.csr_matrix(X), X_m, sqrt_sw) reference_operator = np.hstack( [X - sqrt_sw[:, None] * X_m, sqrt_sw[:, None]]) assert_allclose(reference_operator.dot(A), operator.dot(A)) assert_allclose(reference_operator.T.dot(Y), operator.T.dot(Y))