Add support for external k-mer thresholds for the prefilter

martin-steinegger · martin-steinegger · commit fea8d2037312 · 2022-09-19T18:36:09.000+09:00
diff --git a/src/mmseqs.cpp b/src/mmseqs.cpp
@@ -1,5 +1,6 @@
 #include "Command.h"
 #include "DownloadDatabase.h"
+#include "Prefiltering.h"
 
 const char* binary_name = "mmseqs";
 const char* tool_name = "MMseqs2";
@@ -13,5 +14,6 @@ bool hide_base_commands = false;
 void (*validatorUpdate)(void) = 0;
 std::vector<Command> commands = {};
 std::vector<DatabaseDownload> externalDownloads = {};
+std::vector<KmerThreshold> externalThreshold = {};
 
-bool hide_base_downloads = false;
+bool hide_base_downloads = false;
diff --git a/src/prefiltering/Prefiltering.cpp b/src/prefiltering/Prefiltering.cpp
@@ -1008,6 +1008,12 @@ int Prefiltering::getKmerThreshold(const float sensitivity, const bool isProfile
         return kmerScore.sequence();
     }
     float kmerThrBest = FLT_MAX;
+    int paramType = isProfile ? Parameters::DBTYPE_HMM_PROFILE : Parameters::DBTYPE_AMINO_ACIDS;
+    for(size_t i = 0; i < externalThreshold.size(); i++){
+        if(kmerSize == externalThreshold[i].kmerSize && externalThreshold[i].sequenceType == paramType){
+            return static_cast<int>(externalThreshold[i].base - (externalThreshold[i].sensPerStep * sensitivity));
+        }
+    }
     if (isProfile == true) {
         if (hasContextPseudoCnts == true) {
             if (kmerSize == 5) {
@@ -1131,3 +1137,4 @@ std::pair<int, int> Prefiltering::optimizeSplit(size_t totalMemoryInByte, DBRead
 }
 
 
+
diff --git a/src/prefiltering/Prefiltering.h b/src/prefiltering/Prefiltering.h
@@ -15,6 +15,16 @@
 
 class QueryMatcherTaxonomyHook;
 
+struct KmerThreshold{
+    int sequenceType;
+    int kmerSize;
+    float base;
+    float sensPerStep;
+};
+
+extern std::vector<KmerThreshold> externalThreshold;
+
+
 class Prefiltering {
 public:
     Prefiltering(
@@ -130,3 +140,4 @@ class Prefiltering {
 };
 
 #endif
+
