Generalize calculation of principal components

histolab/mixins.py

@@ -26,19 +26,21 @@ def normalize_columns(arr: np.ndarray) -> np.ndarray:
         return arr / np.linalg.norm(arr, axis=0)
 
     @staticmethod
-    def two_principal_components(arr: np.ndarray) -> np.ndarray:
+    def principal_components(arr: np.ndarray, n_components: int) -> np.ndarray:
         """
-        Return the first two principal components of the covariance matrix of ``arr``.
+        Return the first principal components of the covariance matrix of ``arr``.
 
         Parameters
         ----------
         arr : np.ndarray
             Input array
+        n_components : int
+            Number of principal components to return
 
         Returns
         -------
         np.ndarray
-            Two principal components.
+            ``n_components`` principal components.
         """
         _, V = np.linalg.eigh(np.cov(arr, rowvar=False))
-        return V[:, 1:3]
+        return V[:, -n_components:]

histolab/stain_normalizer.py

@@ -209,7 +209,7 @@ def stain_matrix(
         od_hat = od[~np.any(od < beta, axis=1)]
 
         # Calculate principal components and project input
-        V = self.two_principal_components(od_hat)
+        V = self.principal_components(od_hat, n_components=2)
         proj = np.dot(od_hat, V)
 
         # Angular coordinates with repect to the principle, orthogonal eigenvectors

tests/unit/test_mixins.py

@@ -20,7 +20,7 @@ def it_knows_how_to_normalize_columns(self):
         assert isinstance(matrix_normalized, np.ndarray)
         np.testing.assert_array_almost_equal(matrix_normalized, expected_arr_normalized)
 
-    def it_knows_how_to_calculate_two_principal_components(self, request):
+    def it_knows_how_to_calculate_principal_components(self, request):
         arr = np.array([[1, 2, 3, 4], [5, 6, 7, 8]])
         np_cov_ = function_mock(request, "numpy.cov")
         np_cov_return_value = np.array(
@@ -48,14 +48,14 @@ def it_knows_how_to_calculate_two_principal_components(self, request):
         )
         expected_principal_components = np.array(
             [
-                [-3.43150291e-01, -7.41676339e-18],
-                [-6.35281904e-01, -2.66963658e-16],
-                [4.89216098e-01, -7.07106781e-01],
-                [4.89216098e-01, 7.07106781e-01],
+                [-7.416763e-18, 5.000000e-01],
+                [-2.669637e-16, 5.000000e-01],
+                [-7.071068e-01, 5.000000e-01],
+                [7.071068e-01, 5.000000e-01],
             ]
         )
 
-        two_principal_components = LinalgMixin.two_principal_components(arr)
+        two_principal_components = LinalgMixin.principal_components(arr, n_components=2)
 
         assert isinstance(two_principal_components, np.ndarray)
         np.testing.assert_array_almost_equal(