Move check for matrix NaN values from matrix to tree module

volkamerlab · Jun 29, 2021 · ddbd62a · ddbd62a
1 parent 46d28fd
commit ddbd62a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/kissim/comparison/matrix.py b/kissim/comparison/matrix.py
@@ -100,13 +100,6 @@ def kinase_distance_matrix(
     if fill_diagonal:
         np.fill_diagonal(matrix.values, 0)
 
-    # If matrix contains missing values, respective rows and columns must be dropped
-    column_has_missing_values = matrix.isna().any()
-    column_names_with_missing_values = column_has_missing_values[column_has_missing_values].index
-    matrix = matrix.drop(column_names_with_missing_values, axis=0).drop(
-        column_names_with_missing_values, axis=1
-    )
-
     # If matrix contains number of structure pairs: NaN > 0, cast to int
     if by == "size":
         matrix = matrix.astype("int64")

diff --git a/kissim/comparison/tree.py b/kissim/comparison/tree.py
@@ -122,6 +122,13 @@ def from_distance_matrix(
     # Curate diagonal - set to 0
     np.fill_diagonal(distance_matrix.values, 0)
 
+    # If matrix contains missing values, respective rows and columns must be dropped
+    column_has_missing_values = distance_matrix.isna().any()
+    column_names_with_missing_values = column_has_missing_values[column_has_missing_values].index
+    distance_matrix = distance_matrix.drop(column_names_with_missing_values, axis=0).drop(
+        column_names_with_missing_values, axis=1
+    )
+
     # Hierarchical clustering
     logger.info(
         f"Clustering (method: {clustering_method}) and "