Merge pull request #94 from OndrejSladky/min-frequency

Minimum frequency filtering

Author: OndrejSladky

README.md

@@ -99,6 +99,7 @@ fmsi index -p ms-opt.msfa                                          # Create a k-
 Examples of computing masked superstrings (`ms` subcommand):
 ```
 kmercamel ms -k 31 -o ms.msfa yourfile[.fa|.fa.gz]         # From a (gziped) fasta file, use "-" for stdin
+kmercamel ms -k 31 -o ms.msfa -z 2 yourfile.fa             # Represent only k-mers appearing at least z=2 times
 kmercamel ms -k 31 -o ms.msfa -u yourfile.fa               # Treat k-mer and its reverse complement as distinct
 kmercamel ms -k 31 -o ms.msfa -M maxonemask.m yourfile.fa  # Also store mask with maximum ones
 kmercamel ms -k 31 -o ms.msfa -a streaming yourfile.fa     # Use streaming instead of global for lower memory footprint (likely worse result)

src/khash_utils.h

@@ -27,12 +27,16 @@
 // Use 128-bit integers for extra large k-mers to allow for larger k.
 KHASH_SET_INIT_INT256(S256)
 KHASH_MAP_INIT_INT256(P256, size_t)
+KHASH_MAP_INIT_INT256(Q256, uint8_t)
 // Use 128-bit integers for large k-mers to allow for larger k.
 KHASH_SET_INIT_INT128(S128)
 KHASH_MAP_INIT_INT128(P128, size_t)
+KHASH_MAP_INIT_INT128(Q128, uint8_t)
 // Use 64-bits integers for small k-mers for faster operations and less memory usage.
 KHASH_SET_INIT_INT64(S64)
 KHASH_MAP_INIT_INT64(P64, size_t)
+KHASH_MAP_INIT_INT64(Q64, uint8_t)
+
 
 #define INIT_KHASH_WRAPPER(type) \
     struct kmer_dict##type##_t { \
@@ -72,6 +76,18 @@ KHASH_MAP_INIT_INT64(P64, size_t)
         inline void kh_resize_map(kh_P##type##_t *map, khint_t size) { \
             kh_resize_P##type(map, size); \
         }                        \
+        inline void kh_destroy_freq_map(kh_Q##type##_t *map) { \
+            kh_destroy_Q##type(map); \
+        }                        \
+        inline kh_Q##type##_t *kh_init_freq_map() { \
+            return kh_init_Q##type(); \
+        }                         \
+        inline khint_t kh_get_from_freq_map(kh_Q##type##_t *map, kmer##type##_t key) { \
+            return kh_get_Q##type(map, key); \
+        }                        \
+        inline khint_t kh_put_to_freq_map(kh_Q##type##_t *map, kmer##type##_t key, int *ret) { \
+            return kh_put_Q##type(map, key, ret); \
+        }                        \
     };
 
 INIT_KHASH_WRAPPER(64)
@@ -123,6 +139,26 @@ std::vector<kmer_t> kMersToVec(kh_S_t *kMers, [[maybe_unused]] kmer_t _) {
     return res;
 }
 
+/// Construct a vector of the k-mer set in an arbitrary order.
+template <typename kmer_t, typename kh_Q_t>
+std::vector<kmer_t> kMersToVecFiltered(kh_Q_t *kMers, [[maybe_unused]] kmer_t _, uint16_t min_frequency) {
+    std::vector<kmer_t> res;
+    for (auto i = kh_begin(kMers); i != kh_end(kMers); ++i) {
+        if (!kh_exist(kMers, i) || ((uint16_t)kh_val(kMers, i)) + 1 < min_frequency) continue;
+        res.push_back(kh_key(kMers, i));
+    }
+    return res;
+}
+
+template <typename kmer_t, typename kh_S_t, typename kh_wrapper_t>
+void kMersFromVec(kh_S_t *kMers, kh_wrapper_t wrapper, std::vector<kmer_t> &kMerVec) {
+    for (kmer_t kMer : kMerVec) {
+        int ret;
+        wrapper.kh_put_to_set(kMers, kMer, &ret);
+    }
+}
+
+
 /// Add an interval with given index to the given k-mer.
 ///
 /// [intervalsForKMer] store the intervals and [intervals] maps the k-mer to the index in [intervalsForKMer].

src/main.cpp

@@ -73,6 +73,9 @@ int usage_subcommand(std::string subcommand) {
 
     if (subcommand == "compute" || subcommand == "maskopt" || subcommand == "lowerbound")
     std::cerr << "  -u       - treat k-mer and its reverse complement as distinct" << std::endl;
+    
+    if (subcommand == "compute")
+    std::cerr << "  -z INT   - minimum frequency to represent a k-mer; default 1" << std::endl;
 
     if (subcommand == "mssep2ms") {
     std::cerr << "  -m FILE  - input file with mask" << std::endl;
@@ -97,7 +100,7 @@ void Version() {
 /// Run KmerCamel with the given parameters.
 template <typename kmer_t, typename kh_wrapper_t>
 int kmercamel(kh_wrapper_t wrapper, kmer_t kmer_type, std::string path, int k, int d_max, std::ostream *of, std::ostream *maskf, bool complements, bool masks,
-                    std::string algorithm, bool lower_bound, bool assume_simplitigs) {
+                    std::string algorithm, bool lower_bound, bool assume_simplitigs, uint16_t min_frequency) {
     if (masks) {
         WriteLog("Started optimization of a masked superstring from '" + path + "'.");
         int ret = Optimize(wrapper, kmer_type, algorithm, path, *of, k, complements);
@@ -112,7 +115,8 @@ int kmercamel(kh_wrapper_t wrapper, kmer_t kmer_type, std::string path, int k, i
     /* Handle streaming algorithm separately. */
     if (algorithm == "streaming") {
         WriteName(path, algorithm, k, false, !complements, *of);
-        Streaming(wrapper, kmer_type, path, *of,  k , complements);
+        if (min_frequency == 1) Streaming(wrapper, kmer_type, path, *of,  k , complements);
+        else StreamingFiltered(wrapper, kmer_type, path, *of,  k , complements, min_frequency);
         WriteLog("Finished masked superstring computation.");
     }
     /* Handle hash table based algorithms separately so that they consume less memory. */
@@ -121,7 +125,11 @@ int kmercamel(kh_wrapper_t wrapper, kmer_t kmer_type, std::string path, int k, i
         auto *kMers = wrapper.kh_init_set();
         size_t kmer_count;
         if (!assume_simplitigs) {
-            ReadKMers(kMers, wrapper, kmer_type, path, k, complements);
+            if (min_frequency == 1) {
+                ReadKMers(kMers, wrapper, kmer_type, path, k, complements);
+            } else {
+                ReadKMersFiltered(kMers, wrapper, kmer_type, path, k, complements, min_frequency);
+            }
 
             if (!kh_size(kMers)) {
                 std::cerr << "Path '" << path << "' contains no k-mers. Make sure that your file is a FASTA or gzipped FASTA." << std::endl;
@@ -198,8 +206,9 @@ int camel_compute(int argc, char **argv) {
     bool d_set = false;
     bool assume_simplitigs = false;
     int opt;
+    uint16_t min_frequency = 1;
     try {
-        while ((opt = getopt(argc, argv, "k:d:a:o:huxM:S"))  != -1) {
+        while ((opt = getopt(argc, argv, "k:d:a:o:huxM:Sz:"))  != -1) {
             switch(opt) {
                 case 'o':
                     output.open(optarg);
@@ -231,6 +240,9 @@ int camel_compute(int argc, char **argv) {
                 case 'h':
                     usage_subcommand(subcommand);
                     return 0;
+                case 'z':
+                    min_frequency = std::stoi(optarg);
+                    break;
                 default:
                     return usage_subcommand(subcommand);
             }
@@ -263,13 +275,19 @@ int camel_compute(int argc, char **argv) {
     } else if (assume_simplitigs && algorithm != "global") {
         std::cerr << "Optimization for the input being simplitigs is possible only with global." << std::endl;
         return usage_subcommand(subcommand);
+    } else if (min_frequency >= 256 || min_frequency < 1) {
+        std::cerr << "Minimum frequency '-z' must be between 1 and 255." << std::endl;
+        return usage_subcommand(subcommand);
+    } else if (min_frequency != 1 && assume_simplitigs) {
+        std::cerr << "Inputting simplitigs is not compatible with frequency filterring." << std::endl;
+        return usage_subcommand(subcommand); 
     }
     if (k < 32) {
-        return kmercamel(kmer_dict64_t(), kmer64_t(0), path, k, d_max, of, maskf, complements, false, algorithm, false, assume_simplitigs);
+        return kmercamel(kmer_dict64_t(), kmer64_t(0), path, k, d_max, of, maskf, complements, false, algorithm, false, assume_simplitigs, min_frequency);
     } else if (k < 64) {
-        return kmercamel(kmer_dict128_t(), kmer128_t(0), path, k, d_max, of, maskf, complements, false, algorithm, false, assume_simplitigs);
+        return kmercamel(kmer_dict128_t(), kmer128_t(0), path, k, d_max, of, maskf, complements, false, algorithm, false, assume_simplitigs, min_frequency);
     } else {
-        return kmercamel(kmer_dict256_t(), kmer256_t(0), path, k, d_max, of, maskf, complements, false, algorithm, false, assume_simplitigs);
+        return kmercamel(kmer_dict256_t(), kmer256_t(0), path, k, d_max, of, maskf, complements, false, algorithm, false, assume_simplitigs, min_frequency);
     }
 }
 
@@ -327,11 +345,11 @@ int camel_optimize(int argc, char **argv) {
         return usage_subcommand(subcommand);
     }
     if (k < 32) {
-        return kmercamel(kmer_dict64_t(), kmer64_t(0), path, k, 0, of, nullptr, complements, true, algorithm, false, false);
+        return kmercamel(kmer_dict64_t(), kmer64_t(0), path, k, 0, of, nullptr, complements, true, algorithm, false, false, 1);
     } else if (k < 64) {
-        return kmercamel(kmer_dict128_t(), kmer128_t(0), path, k, 0, of, nullptr, complements, true, algorithm, false, false);
+        return kmercamel(kmer_dict128_t(), kmer128_t(0), path, k, 0, of, nullptr, complements, true, algorithm, false, false, 1);
     } else {
-        return kmercamel(kmer_dict256_t(), kmer256_t(0), path, k, 0, of, nullptr, complements, true, algorithm, false, false);
+        return kmercamel(kmer_dict256_t(), kmer256_t(0), path, k, 0, of, nullptr, complements, true, algorithm, false, false, 1);
     }
 }
 
@@ -379,11 +397,11 @@ int camel_lowerbound(int argc, char **argv) {
         return usage_subcommand(subcommand);
     }
     if (k < 32) {
-        return kmercamel(kmer_dict64_t(), kmer64_t(0), path, k, 0, of, nullptr, complements, false, "global", true, false);
+        return kmercamel(kmer_dict64_t(), kmer64_t(0), path, k, 0, of, nullptr, complements, false, "global", true, false, 1);
     } else if (k < 64) {
-        return kmercamel(kmer_dict128_t(), kmer128_t(0), path, k, 0, of, nullptr, complements, false, "global", true, false);
+        return kmercamel(kmer_dict128_t(), kmer128_t(0), path, k, 0, of, nullptr, complements, false, "global", true, false, 1);
     } else {
-        return kmercamel(kmer_dict256_t(), kmer256_t(0), path, k, 0, of, nullptr, complements, false, "global", true, false);
+        return kmercamel(kmer_dict256_t(), kmer256_t(0), path, k, 0, of, nullptr, complements, false, "global", true, false, 1);
     }
 }
 

src/parser.h

@@ -48,6 +48,42 @@ void AddKMers(kh_S_t *kMers, kh_wrapper_t wrapper, [[maybe_unused]] kmer_t _, si
     }
 }
 
+/// Fill the k-mer dictionary with k-mers from the given sequence and their frequencies.
+/// If complements is true, always add the canonical k-mers.
+template <typename kmer_t, typename kh_Q_t, typename kh_wrapper_t>
+void AddKMersWithFrequencies(kh_Q_t *kMers, kh_wrapper_t wrapper, [[maybe_unused]] kmer_t _, size_t sequence_length,
+              const char* sequence, int64_t k, bool complements) {
+    int64_t currentLength = 0;
+    kmer_t currentKMer = 0, reverseComplement = 0;
+    kmer_t mask = (((kmer_t) 1) <<  (2 * k) ) - 1;
+    kmer_t shift = 2 * (k - 1);
+    for (size_t i = 0; i < sequence_length; ++i) {
+        auto data = nucleotideToInt[(uint8_t)sequence[i]];
+        if (data >= 4) {
+            // Restart if "N"-like nucleotide.
+            currentKMer = reverseComplement = 0;
+            currentLength = 0;
+            continue;
+        }
+        currentKMer = ((currentKMer << 2) | data) & mask;
+        reverseComplement = (reverseComplement >> 2) | ((kmer_t(3 ^ data)) << shift);
+        if (++currentLength >= k) {
+            // Add the canonical k-mer to the dictionary.
+            kmer_t canonical = ((!complements) || currentKMer < reverseComplement) ? currentKMer : reverseComplement;
+            auto ptr = wrapper.kh_get_from_freq_map(kMers, canonical);
+            bool contained = ptr != kh_end(kMers);
+            if (contained) {
+                uint16_t count = kh_val(kMers, ptr);
+                kh_value(kMers, ptr) = (uint8_t) std::min(255, count + 1);
+            } else {
+                int ret;
+                ptr = wrapper.kh_put_to_freq_map(kMers, canonical, &ret);
+                kh_value(kMers, ptr) = 0;
+            }
+        }
+    }
+}
+
 /// Return a file/stdin for reading.
 gzFile OpenFile(std::string &path) {
     FILE *in_stream;
@@ -81,6 +117,30 @@ void ReadKMers(kh_S_t *kMers, kh_wrapper_t wrapper, kmer_t _, std::string &path,
     gzclose(fp);
 }
 
+/// Load a dictionary of k-mers from a fasta file.
+/// If complements is true, add the canonical k-mers.
+template <typename kmer_t, typename kh_S_t, typename kh_wrapper_t>
+void ReadKMersFiltered(kh_S_t *kMers, kh_wrapper_t wrapper, kmer_t _, std::string &path, int k, bool complements,
+               uint16_t min_frequency) {
+    gzFile fp = OpenFile(path);
+    kseq_t *seq = kseq_init(fp);
+    
+    auto kMersFrequencies = wrapper.kh_init_freq_map();
+
+    while (kseq_read(seq) >= 0) {
+        AddKMersWithFrequencies(kMersFrequencies, wrapper, _, seq->seq.l, seq->seq.s, k, complements);
+    }
+
+
+    auto kMersVec = kMersToVecFiltered(kMersFrequencies, _, min_frequency);
+    wrapper.kh_destroy_freq_map(kMersFrequencies);
+    kMersFromVec(kMers, wrapper, kMersVec);
+
+    kseq_destroy(seq);
+    gzclose(fp);
+}
+
+
 /// Read the masked superstring from the given path and return it wrapped as a kseq_t.
 kseq_t* ReadMaskedSuperstring(std::string &path) {
     gzFile fp = OpenFile(path);

src/streaming.h

@@ -45,6 +45,63 @@ void Streaming(kh_wrapper_t wrapper, kmer_t kmer_type, std::string &path, std::o
         }
     }
 
+    kseq_destroy(seq);
+    gzclose(fp);
+}
+
+template <typename kmer_t, typename kh_wrapper_t>
+void StreamingFiltered(kh_wrapper_t wrapper, kmer_t kmer_type, std::string &path, std::ostream &of, int k, bool complements, uint16_t min_frequency) {
+    gzFile fp = OpenFile(path);
+    kseq_t *seq = kseq_init(fp);
+
+    auto kMers = wrapper.kh_init_freq_map();
+
+    kmer_t mask = (kmer_t(1)) << (2 * k - 1);
+    mask |= mask - 1;
+
+
+    while (kseq_read(seq) >= 0) {
+        kmer_t kMer = 0;
+        size_t firstIndex = k - 1;
+        int64_t lastOne = -k;
+        for (size_t i = 0; i < seq->seq.l + k - 1; ++i) {
+            kmer_t c = 4;
+            if (i < seq->seq.l) c = nucleotideToInt[(uint8_t)seq->seq.s[i]];
+            if (c >= 4) {
+                kMer = 0;
+                firstIndex = i + k;
+            }
+            kMer <<= 2;
+            kMer |= c;
+            kMer &= mask;
+            auto rc = ReverseComplement(kMer, k);
+            auto canonical = complements ? std::min(kMer, rc) : kMer;
+
+            auto ptr = wrapper.kh_get_from_freq_map(kMers, canonical);
+            bool contained = ptr != kh_end(kMers);
+            uint16_t count = 0;
+            if (i >= firstIndex) {
+                if (contained) {
+                    count = kh_val(kMers, ptr);
+                    count++;
+                    kh_value(kMers, ptr) = (uint8_t)std::min((uint16_t)255, count);                
+                } else {
+                    int ret;
+                    ptr = wrapper.kh_put_to_freq_map(kMers, canonical, &ret);
+                    kh_value(kMers, ptr) = 0;
+                }
+            }
+
+            if (count + 1 == min_frequency) {
+                of << Masked(seq->seq.s[i - k + 1], true);
+                lastOne = i;
+            } else if ((int64_t (i)) <= lastOne + k - 1) {
+                of << Masked(seq->seq.s[i - k + 1], false);
+            }
+        }
+    }
+
+    wrapper.kh_destroy_freq_map(kMers);
     kseq_destroy(seq);
     gzclose(fp);
 }

src/version

@@ -1 +1 @@
-v2.1.1
+v2.2.0

tests/parser_unittest.h

@@ -58,6 +58,33 @@ namespace {
 
             ReadKMers(kMers, wrapper, kmer_t (0), path, t.k, t.complements, t.case_sensitive);
 
+            EXPECT_EQ(t.wantResultSize, kh_size(kMers));
+        }
+    }
+    TEST(Parser, ReadKMersFiltered) {
+        std::string path = std::filesystem::current_path();
+        path += "/tests/testdata/test.fa";
+        struct TestCase {
+            int k;
+            bool complements;
+            uint16_t min_frequency;
+            size_t wantResultSize;
+        };
+        std::vector<TestCase> tests = {
+                {3, true, 2, 6},
+                {3, false, 2, 5},
+                {3, true, 3, 2},
+                {4, true, 2, 3},
+                {1, false, 5, 4},
+                {1, false, 6, 2},
+                {1, false, 10, 1},
+        };
+
+        for (auto &t: tests) {
+            auto kMers  = wrapper.kh_init_set();
+
+            ReadKMersFiltered(kMers, wrapper, kmer_t (0), path, t.k, t.complements, t.min_frequency);
+
             EXPECT_EQ(t.wantResultSize, kh_size(kMers));
         }