def test_partial_dot_mat_mat_block(self):
        mat1 = sp.arange(2 * 3 * 5 * 7 * 11)
        mat1.shape = (2, 3, 5, 7, 11)

        mat1 = matrix.make_mat(mat1,
                               axis_names=('time', 'x', 'y', 'ra', 'z'),
                               row_axes=(0, 1, 3),
                               col_axes=(0, 2, 3, 4))

        mat2 = sp.arange(2 * 13 * 5 * 7 * 17)
        mat2.shape = (2, 13, 7, 5, 17)

        mat2 = matrix.make_mat(mat2,
                               axis_names=('time', 'w', 'ra', 'y', 'freq'),
                               row_axes=(0, 1, 2, 3),
                               col_axes=(1, 2, 4))

        tmp_arr = sp.tensordot(mat1, mat2, ((2, ), (3, )))
        right_ans = sp.empty((7, 13, 2, 3, 11, 17))

        for ii in range(2):
            for jj in range(7):
                this_tmp = tmp_arr[ii, :, jj, :, ii, :, jj, :]
                this_tmp = sp.rollaxis(this_tmp, 2, 0)
                right_ans[jj, :, ii, ...] = this_tmp

        result = dot_products.partial_dot(mat1, mat2)
        self.assertEqual(result.axes, ('ra', 'w', 'time', 'x', 'z', 'freq'))
        self.assertEqual(result.rows, (0, 1, 2, 3))
        self.assertEqual(result.cols, (0, 1, 4, 5))
        self.assertTrue(sp.allclose(right_ans, result))
    def test_partial_dot_mat_mat_block(self):
        mat1 = sp.arange(2 * 3 * 5 * 7 * 11)
        mat1.shape = (2, 3, 5, 7, 11)

        mat1 = matrix.make_mat(mat1, axis_names=('time', 'x', 'y', 'ra', 'z'),
                                row_axes=(0, 1, 3), col_axes=(0, 2, 3, 4))

        mat2 = sp.arange(2 * 13 * 5 * 7 * 17)
        mat2.shape = (2, 13, 7, 5, 17)

        mat2 = matrix.make_mat(mat2,
                                axis_names=('time', 'w', 'ra', 'y', 'freq'),
                                row_axes=(0, 1, 2, 3), col_axes=(1, 2, 4))

        tmp_arr = sp.tensordot(mat1, mat2, ((2, ), (3, )))
        right_ans = sp.empty((7, 13, 2, 3, 11, 17))

        for ii in range(2):
            for jj in range(7):
                this_tmp = tmp_arr[ii, :, jj, :, ii, :, jj, :]
                this_tmp = sp.rollaxis(this_tmp, 2, 0)
                right_ans[jj, :, ii, ...] = this_tmp

        result = dot_products.partial_dot(mat1, mat2)
        self.assertEqual(result.axes, ('ra', 'w', 'time', 'x', 'z', 'freq'))
        self.assertEqual(result.rows, (0, 1, 2, 3))
        self.assertEqual(result.cols, (0, 1, 4, 5))
        self.assertTrue(sp.allclose(right_ans, result))
    def test_partial_dot_mat_mat(self):
        mat1 = sp.asarray(self.mat)
        mat1.shape = (4, 3, 2, 5)

        mat1 = matrix.make_mat(mat1,
                               axis_names=('time', 'x', 'y', 'z'),
                               row_axes=(0, ),
                               col_axes=(1, 2, 3))

        mat2 = sp.asarray(self.mat)
        mat2.shape = (4, 2, 3, 5)

        mat2 = matrix.make_mat(mat2,
                               axis_names=('w', 'y', 'x', 'freq'),
                               row_axes=(0, 1, 2),
                               col_axes=(3, ))

        result = dot_products.partial_dot(mat1, mat2)
        self.assertEqual(result.axes, ('time', 'w', 'z', 'freq'))
        self.assertEqual(result.rows, (0, 1))
        self.assertEqual(result.cols, (2, 3))
        self.assertEqual(result.shape, (4, 4, 5, 5))
        right_ans = sp.tensordot(mat1, mat2, ((1, 2), (2, 1)))
        right_ans = sp.swapaxes(right_ans, 1, 2)
        self.assertTrue(sp.allclose(right_ans, result))
    def test_partial_dot_mat_vect(self):
        self.mat.shape = (4, 6, 5)
        self.mat.rows = (0, 1)
        self.mat.cols = (2,)
        self.mat.axes = ('x', 'y', 'freq')
        new_vect = dot_products.partial_dot(self.mat, self.vect)
        self.assertEqual(new_vect.shape, (4, 6, 2, 3))
        self.assertEqual(new_vect.axes, ('x', 'y', 'a', 'b'))

        numerical_result = sp.dot(sp.reshape(self.mat, (4*6, 5)),
                                  sp.reshape(self.vect, (5, 2*3)))

        self.assertTrue(sp.allclose(numerical_result.flatten(),
                                    new_vect.flatten()))
    def test_partial_dot_mat_vect(self):
        self.mat.shape = (4, 6, 5)
        self.mat.rows = (0, 1)
        self.mat.cols = (2, )
        self.mat.axes = ('x', 'y', 'freq')
        new_vect = dot_products.partial_dot(self.mat, self.vect)
        self.assertEqual(new_vect.shape, (4, 6, 2, 3))
        self.assertEqual(new_vect.axes, ('x', 'y', 'a', 'b'))

        numerical_result = sp.dot(sp.reshape(self.mat, (4 * 6, 5)),
                                  sp.reshape(self.vect, (5, 2 * 3)))

        self.assertTrue(
            sp.allclose(numerical_result.flatten(), new_vect.flatten()))
    def test_transpose_partial_dot(self):
        self.mat.shape = (5, 4, 6)
        self.mat.cols = (1, 2)
        self.mat.rows = (0, )
        self.mat.axes = ('freq', 'x', 'y')
        matT = self.mat.mat_transpose()
        new_vect = dot_products.partial_dot(matT, self.vect)
        self.assertEqual(new_vect.shape, (4, 6, 2, 3))
        self.assertEqual(new_vect.axes, ('x', 'y', 'a', 'b'))

        # Reform origional matrix to get same numerical result.
        mat = sp.reshape(self.mat, (5, 4*6))
        mat = sp.rollaxis(mat, 1, 0)
        numerical_result = sp.dot(mat, sp.reshape(self.vect, (5, 2*3)))
        self.assertTrue(sp.allclose(numerical_result.flatten(),
                                    new_vect.flatten()))
    def test_transpose_partial_dot(self):
        self.mat.shape = (5, 4, 6)
        self.mat.cols = (1, 2)
        self.mat.rows = (0, )
        self.mat.axes = ('freq', 'x', 'y')
        matT = self.mat.mat_transpose()
        new_vect = dot_products.partial_dot(matT, self.vect)
        self.assertEqual(new_vect.shape, (4, 6, 2, 3))
        self.assertEqual(new_vect.axes, ('x', 'y', 'a', 'b'))

        # Reform origional matrix to get same numerical result.
        mat = sp.reshape(self.mat, (5, 4 * 6))
        mat = sp.rollaxis(mat, 1, 0)
        numerical_result = sp.dot(mat, sp.reshape(self.vect, (5, 2 * 3)))
        self.assertTrue(
            sp.allclose(numerical_result.flatten(), new_vect.flatten()))
    def test_partial_dot_mat_mat(self):
        mat1 = sp.asarray(self.mat)
        mat1.shape = (4, 3, 2, 5)

        mat1 = matrix.make_mat(mat1, axis_names=('time', 'x', 'y', 'z'),
                                row_axes=(0, ), col_axes=(1, 2, 3))

        mat2 = sp.asarray(self.mat)
        mat2.shape = (4, 2, 3, 5)

        mat2 = matrix.make_mat(mat2, axis_names=('w', 'y', 'x', 'freq'),
                                row_axes=(0, 1, 2), col_axes=(3, ))

        result = dot_products.partial_dot(mat1, mat2)
        self.assertEqual(result.axes, ('time', 'w', 'z', 'freq'))
        self.assertEqual(result.rows, (0, 1))
        self.assertEqual(result.cols, (2, 3))
        self.assertEqual(result.shape, (4, 4, 5, 5))
        right_ans = sp.tensordot(mat1, mat2, ((1, 2), (2, 1)))
        right_ans = sp.swapaxes(right_ans, 1, 2)
        self.assertTrue(sp.allclose(right_ans, result))