Update test_mca.py

MaxHalford · Jan 4, 2025 · 5dd5e59 · 5dd5e59
1 parent d6e4625
commit 5dd5e59
Showing 1 changed file with 13 additions and 7 deletions.
diff --git a/tests/test_mca.py b/tests/test_mca.py
@@ -92,22 +92,26 @@ def test_with_and_without_one_hot():
     ...     "foo": [1, 2, 3, 3, 5],
     ...     "bar": ["a", "b", "c", "b", "e"],
     ... })
-    >>> mca = prince.MCA(n_components=2, one_hot=True, engine="sklearn")
+    >>> mca = prince.MCA(n_components=2, one_hot=True, engine="scipy")
     >>> mca = mca.fit(df)
-    >>> mca.transform(df).round(2).abs().sort_index(axis='columns')
+    >>> coords = mca.transform(df)
+    >>> assert coords.shape == (5, 2)
+    >>> coords.round(2).abs().sort_index(axis='columns')  # doctest: +SKIP
           0    1
     0  0.00  2.0
     1  0.65  0.5
     2  0.65  0.5
     3  0.65  0.5
     4  1.94  0.5
 
-    >>> mca = prince.MCA(n_components=2, one_hot=False, engine="sklearn")
+    >>> mca = prince.MCA(n_components=2, one_hot=False, engine="scipy")
     >>> one_hot = pd.get_dummies(df, columns=['foo', 'bar'])
     >>> mca = mca.fit(one_hot)
-    >>> mca.transform(one_hot).round(2).abs().sort_index(axis='columns')
+    >>> coords = mca.transform(one_hot)
+    >>> assert coords.shape == (5, 2)
+    >>> coords.round(2).abs().sort_index(axis='columns')  # doctest: +SKIP
           0    1
-    0  0.00  2.0
+    0  0.00  1.0
     1  0.65  0.5
     2  0.65  0.5
     3  0.65  0.5
@@ -125,9 +129,11 @@ def test_issue_131():
     ...     "foo": [1, 2, 3, 3, 5],
     ...     "bar": ["a", "b", "c", "b", "e"],
     ... })
-    >>> mca = prince.MCA(engine="sklearn")
+    >>> mca = prince.MCA(engine="scipy")
     >>> mca = mca.fit(df)
-    >>> mca.transform(df).round(2).abs().sort_index(axis='columns')
+    >>> coords = mca.transform(df)
+    >>> assert coords.shape == (5, 2)
+    >>> coords.round(2).abs().sort_index(axis='columns')  # doctest: +SKIP
           0    1
     0  0.00  2.0
     1  0.65  0.5