99 lines
3.1 KiB
Python
99 lines
3.1 KiB
Python
|
import operator
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from pandas import (
|
||
|
|
DataFrame,
|
||
|
|
Index,
|
||
|
|
Series,
|
||
|
|
)
|
||
|
|
import pandas._testing as tm
|
||
|
|
|
||
|
|
|
||
|
|
class TestMatMul:
|
||
|
|
def test_matmul(self):
|
||
|
|
# matmul test is for GH#10259
|
||
|
|
a = DataFrame(
|
||
|
|
np.random.default_rng(2).standard_normal((3, 4)),
|
||
|
|
index=["a", "b", "c"],
|
||
|
|
columns=["p", "q", "r", "s"],
|
||
|
|
)
|
||
|
|
b = DataFrame(
|
||
|
|
np.random.default_rng(2).standard_normal((4, 2)),
|
||
|
|
index=["p", "q", "r", "s"],
|
||
|
|
columns=["one", "two"],
|
||
|
|
)
|
||
|
|
|
||
|
|
# DataFrame @ DataFrame
|
||
|
|
result = operator.matmul(a, b)
|
||
|
|
expected = DataFrame(
|
||
|
|
np.dot(a.values, b.values), index=["a", "b", "c"], columns=["one", "two"]
|
||
|
|
)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
# DataFrame @ Series
|
||
|
|
result = operator.matmul(a, b.one)
|
||
|
|
expected = Series(np.dot(a.values, b.one.values), index=["a", "b", "c"])
|
||
|
|
tm.assert_series_equal(result, expected)
|
||
|
|
|
||
|
|
# np.array @ DataFrame
|
||
|
|
result = operator.matmul(a.values, b)
|
||
|
|
assert isinstance(result, DataFrame)
|
||
|
|
assert result.columns.equals(b.columns)
|
||
|
|
assert result.index.equals(Index(range(3)))
|
||
|
|
expected = np.dot(a.values, b.values)
|
||
|
|
tm.assert_almost_equal(result.values, expected)
|
||
|
|
|
||
|
|
# nested list @ DataFrame (__rmatmul__)
|
||
|
|
result = operator.matmul(a.values.tolist(), b)
|
||
|
|
expected = DataFrame(
|
||
|
|
np.dot(a.values, b.values), index=["a", "b", "c"], columns=["one", "two"]
|
||
|
|
)
|
||
|
|
tm.assert_almost_equal(result.values, expected.values)
|
||
|
|
|
||
|
|
# mixed dtype DataFrame @ DataFrame
|
||
|
|
a["q"] = a.q.round().astype(int)
|
||
|
|
result = operator.matmul(a, b)
|
||
|
|
expected = DataFrame(
|
||
|
|
np.dot(a.values, b.values), index=["a", "b", "c"], columns=["one", "two"]
|
||
|
|
)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
# different dtypes DataFrame @ DataFrame
|
||
|
|
a = a.astype(int)
|
||
|
|
result = operator.matmul(a, b)
|
||
|
|
expected = DataFrame(
|
||
|
|
np.dot(a.values, b.values), index=["a", "b", "c"], columns=["one", "two"]
|
||
|
|
)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
# unaligned
|
||
|
|
df = DataFrame(
|
||
|
|
np.random.default_rng(2).standard_normal((3, 4)),
|
||
|
|
index=[1, 2, 3],
|
||
|
|
columns=range(4),
|
||
|
|
)
|
||
|
|
df2 = DataFrame(
|
||
|
|
np.random.default_rng(2).standard_normal((5, 3)),
|
||
|
|
index=range(5),
|
||
|
|
columns=[1, 2, 3],
|
||
|
|
)
|
||
|
|
|
||
|
|
with pytest.raises(ValueError, match="aligned"):
|
||
|
|
operator.matmul(df, df2)
|
||
|
|
|
||
|
|
def test_matmul_message_shapes(self):
|
||
|
|
# GH#21581 exception message should reflect original shapes,
|
||
|
|
# not transposed shapes
|
||
|
|
a = np.random.default_rng(2).random((10, 4))
|
||
|
|
b = np.random.default_rng(2).random((5, 3))
|
||
|
|
|
||
|
|
df = DataFrame(b)
|
||
|
|
|
||
|
|
msg = r"shapes \(10, 4\) and \(5, 3\) not aligned"
|
||
|
|
with pytest.raises(ValueError, match=msg):
|
||
|
|
a @ df
|
||
|
|
with pytest.raises(ValueError, match=msg):
|
||
|
|
a.tolist() @ df
|