debugging the bugfix - only marking the imputed loci with f64::INFINI…

…TY and skipping the loci which were not imputed because they were completely missing across pools
jeffersonfparil · Nov 2, 2023 · 326c28a · 326c28a
1 parent dd6a337
commit 326c28a
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 10 deletions.
diff --git a/src/imputation/adaptive_ld_knn_imputation.rs b/src/imputation/adaptive_ld_knn_imputation.rs
@@ -312,11 +312,15 @@ impl GenotypesAndPhenotypes {
             }
         }
         // Set missing coverages to infinity to mark imputed data
-        for i in 0..self.coverages.nrows() {
-            for j in 0..self.coverages.ncols() {
-                if self.coverages[(i, j)].is_nan() {
-                    self.coverages[(i, j)] = f64::INFINITY
-                };
+        for j in 0..self.coverages.ncols() {
+            // Mark only the imputed loci, i.e. loci which were not completely missing across all pools
+            let n_non_missing = self.coverages.select(Axis(1), &[j]).fold(0, |sum, &x| if x.is_nan()==false{sum + 1}else{sum});
+            if n_non_missing > 0 {
+                for i in 0..self.coverages.nrows() {
+                    if self.coverages[(i, j)].is_nan() {
+                        self.coverages[(i, j)] = f64::INFINITY
+                    };
+                }
             }
         }
         Ok(self)

diff --git a/src/imputation/mean_imputation.rs b/src/imputation/mean_imputation.rs
@@ -40,11 +40,15 @@ impl GenotypesAndPhenotypes {
             }
         }
         // Set missing coverages to infinity to mark imputed data
-        for i in 0..self.coverages.nrows() {
-            for j in 0..self.coverages.ncols() {
-                if self.coverages[(i, j)].is_nan() {
-                    self.coverages[(i, j)] = f64::INFINITY
-                };
+        for j in 0..self.coverages.ncols() {
+            // Mark only the imputed loci, i.e. loci which were not completely missing across all pools
+            let n_non_missing = self.coverages.select(Axis(1), &[j]).fold(0, |sum, &x| if x.is_nan()==false{sum + 1}else{sum});
+            if n_non_missing > 0 {
+                for i in 0..self.coverages.nrows() {
+                    if self.coverages[(i, j)].is_nan() {
+                        self.coverages[(i, j)] = f64::INFINITY
+                    };
+                }
             }
         }
         Ok(self)