Updated PCA visualizer to extend Projection Visualizer (#937)

tests/test_features/test_projection.py

@@ -31,22 +31,28 @@
 
 from sklearn.decomposition import PCA
 from sklearn.pipeline import Pipeline
-from sklearn.preprocessing import  StandardScaler
+from sklearn.preprocessing import StandardScaler
 from sklearn.datasets import make_classification, make_regression
 
 
 ##########################################################################
 ## Fixtures
 ##########################################################################
 
+
 @pytest.fixture(scope="class")
 def discrete(request):
     """
     Creare a random classification fixture.
     """
     X, y = make_classification(
-        n_samples=400, n_features=12, n_informative=10, n_redundant=0,
-        n_classes=5, random_state=2019)
+        n_samples=400,
+        n_features=12,
+        n_informative=10,
+        n_redundant=0,
+        n_classes=5,
+        random_state=2019,
+    )
 
     # Set a class attribute for discrete data
     request.cls.discrete = Dataset(X, y)
@@ -69,26 +75,47 @@ def continuous(request):
 ## MockVisualizer
 ##########################################################################
 
+
 class MockVisualizer(ProjectionVisualizer):
     """
     The MockVisualizer implements the ProjectionVisualizer interface using
     PCA as an internal transformer. This visualizer is used to directly test
     how subclasses interact with the ProjectionVisualizer base class.
     """
 
-    def __init__(self, ax=None, features=None, classes=None, colors=None,
-             colormap=None, target_type="auto", projection=2,
-             alpha=0.75,**kwargs):
-
-        super(MockVisualizer, self).__init__(ax=ax,
-                                             features=features, classes=classes,
-                                             colors=colors, colormap=colormap,
-                                             target_type=target_type,
-                                             projection=projection, alpha=alpha,
-                                             **kwargs)
-
-        self.pca_transformer = Pipeline([("scale", StandardScaler()),
-                                    ("pca", PCA(self.projection, random_state=2019))])
+    def __init__(
+        self,
+        ax=None,
+        features=None,
+        classes=None,
+        colors=None,
+        colormap=None,
+        target_type="auto",
+        projection=2,
+        alpha=0.75,
+        colorbar=True,
+        **kwargs
+    ):
+
+        super(MockVisualizer, self).__init__(
+            ax=ax,
+            features=features,
+            classes=classes,
+            colors=colors,
+            colormap=colormap,
+            target_type=target_type,
+            projection=projection,
+            alpha=alpha,
+            colorbar=colorbar,
+            **kwargs
+        )
+
+        self.pca_transformer = Pipeline(
+            [
+                ("scale", StandardScaler()),
+                ("pca", PCA(self.projection, random_state=2019)),
+            ]
+        )
 
     def fit(self, X, y=None):
         super(MockVisualizer, self).fit(X, y)
@@ -108,6 +135,7 @@ def transform(self, X, y=None):
 ## ProjectionVisualizer Tests
 ##########################################################################
 
+
 @pytest.mark.usefixtures("discrete", "continuous")
 class TestProjectionVisualizer(VisualTestCase):
     """
@@ -144,8 +172,9 @@ def test_continuous_when_target_discrete(self):
         """
         _, ax = plt.subplots()
         X, y = self.discrete
-        visualizer = MockVisualizer(ax=ax, projection="2D",
-                                          target_type="continuous", colormap="cool")
+        visualizer = MockVisualizer(
+            ax=ax, projection="2D", target_type="continuous", colormap="cool"
+        )
         visualizer.fit(X, y)
         visualizer.transform(X, y)
         visualizer.finalize()
@@ -169,7 +198,7 @@ def test_discrete_3d(self):
         X, y = self.discrete
 
         classes = ["a", "b", "c", "d", "e"]
-        colors = ["r", "b", "g", "m","c"]
+        colors = ["r", "b", "g", "m", "c"]
         visualizer = MockVisualizer(projection=3, colors=colors, classes=classes)
         visualizer.fit_transform(X, y)
         npt.assert_array_equal(visualizer.classes_, classes)
@@ -222,10 +251,10 @@ def test_target_not_label_encoded(self):
         """
         X, y = self.discrete
         # Multiply every element by 10 to make non-label encoded
-        y = y*10
+        y = y * 10
         visualizer = MockVisualizer()
         msg = "Target needs to be label encoded."
-        with pytest.raises(YellowbrickValueError, match = msg):
+        with pytest.raises(YellowbrickValueError, match=msg):
             visualizer.fit_transform(X, y)
 
     @pytest.mark.parametrize("dataset", ("discrete", "continuous"))
@@ -238,5 +267,15 @@ def test_y_required_for_discrete_and_continuous(self, dataset):
         visualizer.fit(X, y)
 
         msg = "y is required for {} target".format(dataset)
-        with pytest.raises(YellowbrickValueError, match = msg):
-            visualizer.transform(X)
+        with pytest.raises(YellowbrickValueError, match=msg):
+            visualizer.transform(X)
+
+    def test_colorbar_false(self):
+        """
+        Test that colorbar equals false works correctly 
+        """
+        visualizer = MockVisualizer(colorbar=False, colormap="YlOrRd")
+        visualizer.fit_transform(*self.continuous)
+        visualizer.finalize()
+
+        self.assert_images_similar(visualizer)

yellowbrick/features/__init__.py

@@ -22,7 +22,7 @@
 from .radviz import RadialVisualizer, RadViz, radviz
 from .rankd import Rank1D, rank1d, Rank2D, rank2d
 from .jointplot import JointPlot, JointPlotVisualizer, joint_plot
-from .pca import PCADecomposition, pca_decomposition
+from .pca import PCA, PCADecomposition, pca_decomposition
 from .importances import FeatureImportances, feature_importances
 from .rfecv import RFECV, rfecv
 from .manifold import Manifold, manifold_embedding