variably choose which dimensions to plot

Author: moritzschaefer

scikitplot/decomposition.py

@@ -5,8 +5,8 @@
 properties shared by scikit-learn estimators. The specific requirements are
 documented per function.
 """
-from __future__ import absolute_import, division, print_function, \
-    unicode_literals
+from __future__ import (absolute_import, division, print_function,
+                        unicode_literals)
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -95,6 +95,7 @@ def plot_pca_component_variance(clf, title='PCA Component Explained Variances',
 
 
 def plot_pca_2d_projection(clf, X, y, title='PCA 2-D Projection',
+                           dimensions=[0, 1],
                            biplot=False, feature_labels=None,
                            ax=None, figsize=None, cmap='Spectral',
                            title_fontsize="large", text_fontsize="medium",
@@ -172,31 +173,31 @@ def plot_pca_2d_projection(clf, X, y, title='PCA 2-D Projection',
     colors = plt.cm.get_cmap(cmap)(np.linspace(0, 1, len(classes)))
 
     for label, color in zip(classes, colors):
-        ax.scatter(transformed_X[y == label, 0], transformed_X[y == label, 1],
+        ax.scatter(transformed_X[y == label, dimensions[0]], transformed_X[y == label, dimensions[1]],
                    alpha=0.8, lw=2, label=label, color=color)
 
         if label_dots:
-            for dot in transformed_X[y == label, 0:2]:
+            for dot in transformed_X[y == label, dimensions]:
                 ax.text(*dot, label)
 
     if biplot:
-        xs = transformed_X[:, 0]
-        ys = transformed_X[:, 1]
-        vectors = np.transpose(clf.components_[:2, :])
+        xs = transformed_X[:, dimensions[0]]
+        ys = transformed_X[:, dimensions[1]]
+        vectors = np.transpose(clf.components_[dimensions, :])
         vectors_scaled = vectors * [xs.max(), ys.max()]
         for i in range(vectors.shape[0]):
-            ax.annotate("", xy=(vectors_scaled[i, 0], vectors_scaled[i, 1]),
+            ax.annotate("", xy=(vectors_scaled[i, dimensions[0]], vectors_scaled[i, dimensions[1]]),
                         xycoords='data', xytext=(0, 0), textcoords='data',
                         arrowprops={'arrowstyle': '-|>', 'ec': 'r'})
 
-            ax.text(vectors_scaled[i, 0] * 1.05, vectors_scaled[i, 1] * 1.05,
+            ax.text(vectors_scaled[i, dimensions[0]] * 1.05, vectors_scaled[i, dimensions[1]] * 1.05,
                     feature_labels[i] if feature_labels else "Variable" + str(i),
                     color='b', fontsize=text_fontsize)
 
     ax.legend(loc='best', shadow=False, scatterpoints=1,
               fontsize=text_fontsize)
-    ax.set_xlabel('First Principal Component', fontsize=text_fontsize)
-    ax.set_ylabel('Second Principal Component', fontsize=text_fontsize)
+    ax.set_xlabel(f'Principal Component {dimensions[0]+1}', fontsize=text_fontsize)
+    ax.set_ylabel(f'Principal Component {dimension[1]+1}', fontsize=text_fontsize)
     ax.tick_params(labelsize=text_fontsize)
 
     return ax