Python linalg.BlockMatrix类代码示例

    def test_from_entry_expr_options(self):
        def build_mt(a):
            data = [{'v': 0, 's': 0, 'x': a[0]},
                    {'v': 0, 's': 1, 'x': a[1]},
                    {'v': 0, 's': 2, 'x': a[2]}]
            ht = hl.Table.parallelize(data, hl.dtype('struct{v: int32, s: int32, x: float64}'))
            mt = ht.to_matrix_table(['v'], ['s'])
            ids = mt.key_cols_by()['s'].collect()
            return mt.choose_cols([ids.index(0), ids.index(1), ids.index(2)])

        def check(expr, mean_impute, center, normalize, expected):
            actual = np.squeeze(BlockMatrix.from_entry_expr(expr,
                                                            mean_impute=mean_impute,
                                                            center=center,
                                                            normalize=normalize).to_numpy())
            assert np.allclose(actual, expected)

        a = np.array([0.0, 1.0, 2.0])

        mt = build_mt(a)
        check(mt.x, False, False, False, a)
        check(mt.x, False, True, False, a - 1.0)
        check(mt.x, False, False, True, a / np.sqrt(5))
        check(mt.x, False, True, True, (a - 1.0) / np.sqrt(2))
        check(mt.x + 1 - 1, False, False, False, a)

        mt = build_mt([0.0, hl.null('float64'), 2.0])
        check(mt.x, True, False, False, a)
        check(mt.x, True, True, False, a - 1.0)
        check(mt.x, True, False, True, a / np.sqrt(5))
        check(mt.x, True, True, True, (a - 1.0) / np.sqrt(2))
        with self.assertRaises(Exception):
            BlockMatrix.from_entry_expr(mt.x)

    def test_sum(self):
        def sums_agree(bm, nd):
            self.assertAlmostEqual(bm.sum(), np.sum(nd))
            self._assert_close(bm.sum(axis=0), np.sum(nd, axis=0, keepdims=True))
            self._assert_close(bm.sum(axis=1), np.sum(nd, axis=1, keepdims=True))

        nd = np.random.normal(size=(11, 13))
        bm = BlockMatrix.from_numpy(nd, block_size=3)

        nd2 = np.zeros(shape=(5, 7))
        nd2[2, 4] = 1.0
        nd2[2, 5] = 2.0
        nd2[3, 4] = 3.0
        nd2[3, 5] = 4.0
        bm2 = BlockMatrix.from_numpy(nd2, block_size=2).sparsify_rectangles([[2, 4, 4, 6]])

        bm3 = BlockMatrix.from_numpy(nd2, block_size=2).sparsify_rectangles([[2, 4, 4, 6], [0, 5, 0, 1]])

        bm4 = BlockMatrix.from_numpy(nd2, block_size=2).sparsify_rectangles([[2, 4, 4, 6], [0, 1, 0, 7]])

        nd5 = np.zeros(shape=(5, 7))
        bm5 = BlockMatrix.fill(5, 7, value=0.0, block_size=2).sparsify_rectangles([])

        sums_agree(bm, nd)
        sums_agree(bm2, nd2)
        sums_agree(bm3, nd2)
        sums_agree(bm4, nd2)
        sums_agree(bm5, nd5)

    def test_matrix_ops(self):
        nm = np.matrix([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
        m = BlockMatrix.from_numpy(nm, block_size=2)

        nrow = np.matrix([[7.0, 8.0, 9.0]])
        row = BlockMatrix.from_numpy(nrow, block_size=2)

        self._assert_eq(m.T, nm.T)
        self._assert_eq(m.T, nm.T)
        self._assert_eq(row.T, nrow.T)

        self._assert_eq(m @ m.T, nm @ nm.T)
        self._assert_eq(m @ nm.T, nm @ nm.T)
        self._assert_eq(row @ row.T, nrow @ nrow.T)
        self._assert_eq(row @ nrow.T, nrow @ nrow.T)

        self._assert_eq(m.T @ m, nm.T @ nm)
        self._assert_eq(m.T @ nm, nm.T @ nm)
        self._assert_eq(row.T @ row, nrow.T @ nrow)
        self._assert_eq(row.T @ nrow, nrow.T @ nrow)

        self.assertRaises(ValueError, lambda: m @ m)
        self.assertRaises(ValueError, lambda: m @ nm)

        self._assert_eq(m.diagonal(), np.array([1.0, 5.0]))
        self._assert_eq(m.T.diagonal(), np.array([1.0, 5.0]))
        self._assert_eq((m @ m.T).diagonal(), np.array([14.0, 77.0]))

    def test_stage_locally(self):
        nd = np.arange(0, 80, dtype=float).reshape(8, 10)
        bm_uri = new_temp_file()
        BlockMatrix.from_numpy(nd, block_size=3).write(bm_uri, stage_locally=True)

        bm = BlockMatrix.read(bm_uri)
        self._assert_eq(nd, bm)

    def test_sparsify_band(self):
        nd = np.array([[ 1.0,  2.0,  3.0,  4.0],
                       [ 5.0,  6.0,  7.0,  8.0],
                       [ 9.0, 10.0, 11.0, 12.0],
                       [13.0, 14.0, 15.0, 16.0]])
        bm = BlockMatrix.from_numpy(nd, block_size=2)

        self._assert_eq(
            bm.sparsify_band(lower=-1, upper=2),
            np.array([[ 1.,  2.,  3.,  0.],
                      [ 5.,  6.,  7.,  8.],
                      [ 0., 10., 11., 12.],
                      [ 0.,  0., 15., 16.]]))

        self._assert_eq(
            bm.sparsify_band(lower=0, upper=0, blocks_only=True),
            np.array([[ 1.,  2.,  0.,  0.],
                      [ 5.,  6.,  0.,  0.],
                      [ 0.,  0., 11., 12.],
                      [ 0.,  0., 15., 16.]]))

        nd2 = np.arange(0, 80, dtype=float).reshape(8, 10)
        bm2 = BlockMatrix.from_numpy(nd2, block_size=3)

        for bounds in [[0, 0], [1, 1], [2, 2], [-5, 5], [-7, 0], [0, 9], [-100, 100]]:
            lower, upper = bounds
            actual = bm2.sparsify_band(lower, upper, blocks_only=False).to_numpy()
            mask = np.fromfunction(lambda i, j: (lower <= j - i) * (j - i <= upper), (8, 10))
            self._assert_eq(actual, nd2 * mask)

    def test_fill(self):
        nd = np.ones((3, 5))
        bm = BlockMatrix.fill(3, 5, 1.0)
        bm2 = BlockMatrix.fill(3, 5, 1.0, block_size=2)

        self.assertTrue(bm.block_size == BlockMatrix.default_block_size())
        self.assertTrue(bm2.block_size == 2)
        self._assert_eq(bm, nd)
        self._assert_eq(bm2, nd)

    def test_svd(self):
        def assert_same_columns_up_to_sign(a, b):
            for j in range(a.shape[1]):
                assert np.allclose(a[:, j], b[:, j]) or np.allclose(-a[:, j], b[:, j])

        x0 = np.array([[-2.0, 0.0, 3.0],
                       [-1.0, 2.0, 4.0]])
        u0, s0, vt0 = np.linalg.svd(x0, full_matrices=False)

        x = BlockMatrix.from_numpy(x0)

        # _svd
        u, s, vt = x.svd()
        assert_same_columns_up_to_sign(u, u0)
        assert np.allclose(s, s0)
        assert_same_columns_up_to_sign(vt.T, vt0.T)

        s = x.svd(compute_uv=False)
        assert np.allclose(s, s0)

        # left _svd_gramian
        u, s, vt = x.svd(complexity_bound=0)
        assert_same_columns_up_to_sign(u, u0)
        assert np.allclose(s, s0)
        assert_same_columns_up_to_sign(vt.to_numpy().T, vt0.T)

        s = x.svd(compute_uv=False, complexity_bound=0)
        assert np.allclose(s, s0)

        # right _svd_gramian
        x = BlockMatrix.from_numpy(x0.T)
        u, s, vt = x.svd(complexity_bound=0)
        assert_same_columns_up_to_sign(u.to_numpy(), vt0.T)
        assert np.allclose(s, s0)
        assert_same_columns_up_to_sign(vt.T, u0)

        s = x.svd(compute_uv=False, complexity_bound=0)
        assert np.allclose(s, s0)

        # left _svd_gramian when dimensions agree
        x = BlockMatrix.from_numpy(x0[:, :2])
        u, s, vt = x.svd(complexity_bound=0)
        assert isinstance(u, np.ndarray)
        assert isinstance(vt, BlockMatrix)

        # rank-deficient X sets negative eigenvalues to 0.0
        a = np.array([[0.0, 1.0, np.e, np.pi, 10.0, 25.0]])
        x0 = a.T @ a  # rank 1
        e, _ = np.linalg.eigh(x0 @ x0.T)

        x = BlockMatrix.from_numpy(x0)
        _, s, _ = x.svd(complexity_bound=0)
        assert np.all(s >= 0.0)

        s = x.svd(compute_uv=False, complexity_bound=0)
        assert np.all(s >= 0)

    def test_export_blocks(self):
        nd = np.ones(shape=(8, 10))
        bm = BlockMatrix.from_numpy(nd, block_size=20)

        bm_path = new_local_temp_dir()
        bm_uri = local_path_uri(bm_path)
        bm.export_blocks(bm_uri, binary=True)
        actual = BlockMatrix.rectangles_to_numpy(bm_path, binary=True)

        self._assert_eq(nd, actual)

    def test_write_overwrite(self):
        path = new_temp_file()

        bm = BlockMatrix.from_numpy(np.array([[0]]))
        bm.write(path)
        self.assertRaises(FatalError, lambda: bm.write(path))

        bm2 = BlockMatrix.from_numpy(np.array([[1]]))
        bm2.write(path, overwrite=True)
        self._assert_eq(BlockMatrix.read(path), bm2)

    def test_random_uniform(self):
        uniform = BlockMatrix.random(10, 10, gaussian=False)

        nuniform = uniform.to_numpy()
        for row in nuniform:
            for entry in row:
                assert entry > 0

    def test_export_rectangles(self):
        nd = np.arange(0, 80, dtype=float).reshape(8, 10)

        rects1 = [[0, 1, 0, 1], [4, 5, 7, 8]]

        rects2 = [[4, 5, 0, 10], [0, 8, 4, 5]]

        rects3 = [[0, 1, 0, 1], [1, 2, 1, 2], [2, 3, 2, 3],
                  [3, 5, 3, 6], [3, 6, 3, 7], [3, 7, 3, 8],
                  [4, 5, 0, 10], [0, 8, 4, 5], [0, 8, 0, 10]]

        for rects in [rects1, rects2, rects3]:
            for block_size in [3, 4, 10]:
                rect_path = new_local_temp_dir()
                rect_uri = local_path_uri(rect_path)

                bm = BlockMatrix.from_numpy(nd, block_size=block_size)
                bm.export_rectangles(rect_uri, rects)

                self._assert_rectangles_eq(nd, rect_path, rects)

                rect_path_bytes = new_local_temp_dir()
                rect_uri_bytes = local_path_uri(rect_path_bytes)

                bm.export_rectangles(rect_uri_bytes, rects, binary=True)
                self._assert_rectangles_eq(nd, rect_path_bytes, rects, binary=True)

    def test_to_matrix_table(self):
        n_partitions = 2
        rows, cols = 2, 5
        bm = BlockMatrix._create(rows, cols, [float(i) for i in range(10)])
        actual = bm.to_matrix_table_row_major(n_partitions)

        expected = hl.utils.range_matrix_table(rows, cols)
        expected = expected.annotate_entries(element=hl.float64(expected.row_idx * cols + expected.col_idx))
        expected = expected.key_cols_by(col_idx=hl.int64(expected.col_idx))
        expected = expected.key_rows_by(row_idx=hl.int64(expected.row_idx))
        assert expected._same(actual)

        bm = BlockMatrix.random(50, 100, block_size=25, seed=0)
        mt = bm.to_matrix_table_row_major(n_partitions)
        mt_round_trip = BlockMatrix.from_entry_expr(mt.element).to_matrix_table_row_major()
        assert mt._same(mt_round_trip)

    def test_sparsify_triangle(self):
        nd = np.array([[ 1.0,  2.0,  3.0,  4.0],
                       [ 5.0,  6.0,  7.0,  8.0],
                       [ 9.0, 10.0, 11.0, 12.0],
                       [13.0, 14.0, 15.0, 16.0]])
        bm = BlockMatrix.from_numpy(nd, block_size=2)

        self.assertFalse(bm.is_sparse)
        self.assertTrue(bm.sparsify_triangle().is_sparse)

        self._assert_eq(
            bm.sparsify_triangle(),
            np.array([[ 1.,  2.,  3.,  4.],
                      [ 0.,  6.,  7.,  8.],
                      [ 0.,  0., 11., 12.],
                      [ 0.,  0.,  0., 16.]]))

        self._assert_eq(
            bm.sparsify_triangle(lower=True),
            np.array([[ 1.,  0.,  0.,  0.],
                      [ 5.,  6.,  0.,  0.],
                      [ 9., 10., 11.,  0.],
                      [13., 14., 15., 16.]]))

        self._assert_eq(
            bm.sparsify_triangle(blocks_only=True),
            np.array([[ 1.,  2.,  3.,  4.],
                      [ 5.,  6.,  7.,  8.],
                      [ 0.,  0., 11., 12.],
                      [ 0.,  0., 15., 16.]]))

    def test_sparsify_row_intervals(self):
        nd = np.array([[ 1.0,  2.0,  3.0,  4.0],
                       [ 5.0,  6.0,  7.0,  8.0],
                       [ 9.0, 10.0, 11.0, 12.0],
                       [13.0, 14.0, 15.0, 16.0]])
        bm = BlockMatrix.from_numpy(nd, block_size=2)

        self._assert_eq(
            bm.sparsify_row_intervals(
                starts=[1, 0, 2, 2],
                stops= [2, 0, 3, 4]),
            np.array([[ 0.,  2.,  0.,  0.],
                      [ 0.,  0.,  0.,  0.],
                      [ 0.,  0., 11.,  0.],
                      [ 0.,  0., 15., 16.]]))

        self._assert_eq(
            bm.sparsify_row_intervals(
                starts=[1, 0, 2, 2],
                stops= [2, 0, 3, 4],
                blocks_only=True),
            np.array([[ 1.,  2.,  0.,  0.],
                      [ 5.,  6.,  0.,  0.],
                      [ 0.,  0., 11., 12.],
                      [ 0.,  0., 15., 16.]]))

        nd2 = np.random.normal(size=(8, 10))
        bm2 = BlockMatrix.from_numpy(nd2, block_size=3)

        for bounds in [[[0, 1, 2, 3, 4, 5, 6, 7],
                        [1, 2, 3, 4, 5, 6, 7, 8]],
                       [[0, 0, 5, 3, 4, 5, 8, 2],
                        [9, 0, 5, 3, 4, 5, 9, 5]],
                       [[0, 5, 10, 8, 7, 6, 5, 4],
                        [0, 5, 10, 9, 8, 7, 6, 5]]]:
            starts, stops = bounds
            actual = bm2.sparsify_row_intervals(starts, stops, blocks_only=False).to_numpy()
            expected = nd2.copy()
            for i in range(0, 8):
                for j in range(0, starts[i]):
                    expected[i, j] = 0.0
                for j in range(stops[i], 10):
                    expected[i, j] = 0.0
            self._assert_eq(actual, expected)

    def test_sum_with_sparsify(self):
        nd = np.zeros(shape=(5, 7))
        nd[2, 4] = 1.0
        nd[2, 5] = 2.0
        nd[3, 4] = 3.0
        nd[3, 5] = 4.0
        bm = BlockMatrix.from_numpy(nd, block_size=2).sparsify_rectangles([[2, 4, 4, 6]])

        bm2 = BlockMatrix.from_numpy(nd, block_size=2).sparsify_rectangles([[2, 4, 4, 6], [0, 5, 0, 1]])

        bm3 = BlockMatrix.from_numpy(nd, block_size=2).sparsify_rectangles([[2, 4, 4, 6], [0, 1, 0, 7]])

        nd4 = np.zeros(shape=(5, 7))
        bm4 = BlockMatrix.fill(5, 7, value=0.0, block_size=2).sparsify_rectangles([])

        self.assert_sums_agree(bm, nd)
        self.assert_sums_agree(bm2, nd)
        self.assert_sums_agree(bm3, nd)
        self.assert_sums_agree(bm4, nd4)