test_dot.py 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131
  1. import numpy as np
  2. import pytest
  3. from pandas import (
  4. DataFrame,
  5. Series,
  6. )
  7. import pandas._testing as tm
  8. class DotSharedTests:
  9. @pytest.fixture
  10. def obj(self):
  11. raise NotImplementedError
  12. @pytest.fixture
  13. def other(self) -> DataFrame:
  14. """
  15. other is a DataFrame that is indexed so that obj.dot(other) is valid
  16. """
  17. raise NotImplementedError
  18. @pytest.fixture
  19. def expected(self, obj, other) -> DataFrame:
  20. """
  21. The expected result of obj.dot(other)
  22. """
  23. raise NotImplementedError
  24. @classmethod
  25. def reduced_dim_assert(cls, result, expected):
  26. """
  27. Assertion about results with 1 fewer dimension that self.obj
  28. """
  29. raise NotImplementedError
  30. def test_dot_equiv_values_dot(self, obj, other, expected):
  31. # `expected` is constructed from obj.values.dot(other.values)
  32. result = obj.dot(other)
  33. tm.assert_equal(result, expected)
  34. def test_dot_2d_ndarray(self, obj, other, expected):
  35. # Check ndarray argument; in this case we get matching values,
  36. # but index/columns may not match
  37. result = obj.dot(other.values)
  38. assert np.all(result == expected.values)
  39. def test_dot_1d_ndarray(self, obj, expected):
  40. # can pass correct-length array
  41. row = obj.iloc[0] if obj.ndim == 2 else obj
  42. result = obj.dot(row.values)
  43. expected = obj.dot(row)
  44. self.reduced_dim_assert(result, expected)
  45. def test_dot_series(self, obj, other, expected):
  46. # Check series argument
  47. result = obj.dot(other["1"])
  48. self.reduced_dim_assert(result, expected["1"])
  49. def test_dot_series_alignment(self, obj, other, expected):
  50. result = obj.dot(other.iloc[::-1]["1"])
  51. self.reduced_dim_assert(result, expected["1"])
  52. def test_dot_aligns(self, obj, other, expected):
  53. # Check index alignment
  54. other2 = other.iloc[::-1]
  55. result = obj.dot(other2)
  56. tm.assert_equal(result, expected)
  57. def test_dot_shape_mismatch(self, obj):
  58. msg = "Dot product shape mismatch"
  59. # exception raised is of type Exception
  60. with pytest.raises(Exception, match=msg):
  61. obj.dot(obj.values[:3])
  62. def test_dot_misaligned(self, obj, other):
  63. msg = "matrices are not aligned"
  64. with pytest.raises(ValueError, match=msg):
  65. obj.dot(other.T)
  66. class TestSeriesDot(DotSharedTests):
  67. @pytest.fixture
  68. def obj(self):
  69. return Series(np.random.randn(4), index=["p", "q", "r", "s"])
  70. @pytest.fixture
  71. def other(self):
  72. return DataFrame(
  73. np.random.randn(3, 4), index=["1", "2", "3"], columns=["p", "q", "r", "s"]
  74. ).T
  75. @pytest.fixture
  76. def expected(self, obj, other):
  77. return Series(np.dot(obj.values, other.values), index=other.columns)
  78. @classmethod
  79. def reduced_dim_assert(cls, result, expected):
  80. """
  81. Assertion about results with 1 fewer dimension that self.obj
  82. """
  83. tm.assert_almost_equal(result, expected)
  84. class TestDataFrameDot(DotSharedTests):
  85. @pytest.fixture
  86. def obj(self):
  87. return DataFrame(
  88. np.random.randn(3, 4), index=["a", "b", "c"], columns=["p", "q", "r", "s"]
  89. )
  90. @pytest.fixture
  91. def other(self):
  92. return DataFrame(
  93. np.random.randn(4, 2), index=["p", "q", "r", "s"], columns=["1", "2"]
  94. )
  95. @pytest.fixture
  96. def expected(self, obj, other):
  97. return DataFrame(
  98. np.dot(obj.values, other.values), index=obj.index, columns=other.columns
  99. )
  100. @classmethod
  101. def reduced_dim_assert(cls, result, expected):
  102. """
  103. Assertion about results with 1 fewer dimension that self.obj
  104. """
  105. tm.assert_series_equal(result, expected, check_names=False)
  106. assert result.name is None