123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131 |
- import numpy as np
- import pytest
- from pandas import (
- DataFrame,
- Series,
- )
- import pandas._testing as tm
- class DotSharedTests:
- @pytest.fixture
- def obj(self):
- raise NotImplementedError
- @pytest.fixture
- def other(self) -> DataFrame:
- """
- other is a DataFrame that is indexed so that obj.dot(other) is valid
- """
- raise NotImplementedError
- @pytest.fixture
- def expected(self, obj, other) -> DataFrame:
- """
- The expected result of obj.dot(other)
- """
- raise NotImplementedError
- @classmethod
- def reduced_dim_assert(cls, result, expected):
- """
- Assertion about results with 1 fewer dimension that self.obj
- """
- raise NotImplementedError
- def test_dot_equiv_values_dot(self, obj, other, expected):
- # `expected` is constructed from obj.values.dot(other.values)
- result = obj.dot(other)
- tm.assert_equal(result, expected)
- def test_dot_2d_ndarray(self, obj, other, expected):
- # Check ndarray argument; in this case we get matching values,
- # but index/columns may not match
- result = obj.dot(other.values)
- assert np.all(result == expected.values)
- def test_dot_1d_ndarray(self, obj, expected):
- # can pass correct-length array
- row = obj.iloc[0] if obj.ndim == 2 else obj
- result = obj.dot(row.values)
- expected = obj.dot(row)
- self.reduced_dim_assert(result, expected)
- def test_dot_series(self, obj, other, expected):
- # Check series argument
- result = obj.dot(other["1"])
- self.reduced_dim_assert(result, expected["1"])
- def test_dot_series_alignment(self, obj, other, expected):
- result = obj.dot(other.iloc[::-1]["1"])
- self.reduced_dim_assert(result, expected["1"])
- def test_dot_aligns(self, obj, other, expected):
- # Check index alignment
- other2 = other.iloc[::-1]
- result = obj.dot(other2)
- tm.assert_equal(result, expected)
- def test_dot_shape_mismatch(self, obj):
- msg = "Dot product shape mismatch"
- # exception raised is of type Exception
- with pytest.raises(Exception, match=msg):
- obj.dot(obj.values[:3])
- def test_dot_misaligned(self, obj, other):
- msg = "matrices are not aligned"
- with pytest.raises(ValueError, match=msg):
- obj.dot(other.T)
- class TestSeriesDot(DotSharedTests):
- @pytest.fixture
- def obj(self):
- return Series(np.random.randn(4), index=["p", "q", "r", "s"])
- @pytest.fixture
- def other(self):
- return DataFrame(
- np.random.randn(3, 4), index=["1", "2", "3"], columns=["p", "q", "r", "s"]
- ).T
- @pytest.fixture
- def expected(self, obj, other):
- return Series(np.dot(obj.values, other.values), index=other.columns)
- @classmethod
- def reduced_dim_assert(cls, result, expected):
- """
- Assertion about results with 1 fewer dimension that self.obj
- """
- tm.assert_almost_equal(result, expected)
- class TestDataFrameDot(DotSharedTests):
- @pytest.fixture
- def obj(self):
- return DataFrame(
- np.random.randn(3, 4), index=["a", "b", "c"], columns=["p", "q", "r", "s"]
- )
- @pytest.fixture
- def other(self):
- return DataFrame(
- np.random.randn(4, 2), index=["p", "q", "r", "s"], columns=["1", "2"]
- )
- @pytest.fixture
- def expected(self, obj, other):
- return DataFrame(
- np.dot(obj.values, other.values), index=obj.index, columns=other.columns
- )
- @classmethod
- def reduced_dim_assert(cls, result, expected):
- """
- Assertion about results with 1 fewer dimension that self.obj
- """
- tm.assert_series_equal(result, expected, check_names=False)
- assert result.name is None
|