Support loco for mdb kinship compute

1 files changed, 26 insertions, 1 deletions

diff --git a/src/gemma_io.cpp b/src/gemma_io.cpp
index e183e3f..3e290f4 100644
--- a/src/gemma_io.cpp
+++ b/src/gemma_io.cpp
@@ -1478,8 +1478,27 @@ void ReadFile_eigenD(const string &file_kd, bool &error, gsl_vector *eval) {
 
 
 // Read bimbam mean genotype file and calculate kinship matrix.
-gsl_matrix *mdb_calc_kin(const string file_geno, bool is_loco, const int k_mode) {
+gsl_matrix *mdb_calc_kin(const string file_geno, const int k_mode, const string loco) {
   checkpoint("mdb-kin",file_geno);
+  bool is_loco = !loco.empty();
+
+  // Convert loco string to what we use in the chrpos index
+  uint8_t loco_chr = 0;
+  if (is_loco) {
+    if (loco == "X") {
+      loco_chr = CHR_X;
+    } else if (loco == "Y") {
+      loco_chr = CHR_Y;
+    } else if (loco == "M") {
+      loco_chr = CHR_M;
+    } else {
+      try {
+        loco_chr = static_cast<uint8_t>(stoi(loco));
+      } catch (...) {
+        loco_chr = 0;
+      }
+    }
+  }
 
   /* Create and open the LMDB environment: */
   auto env = lmdb::env::create();
@@ -1552,6 +1571,12 @@ gsl_matrix *mdb_calc_kin(const string file_geno, bool is_loco, const int k_mode)
     if (indicator_snp[t] == 0)
       continue;
     */
+    const uint8_t* data = reinterpret_cast<const uint8_t*>(key.data());
+    auto chr = static_cast<uint8_t>(data[1]);
+
+    if (is_loco && loco_chr == chr)
+      continue;
+
 #if !defined NDEBUG
     size_t num_floats = value.size() / sizeof(float);
     assert(num_floats == ni_total);