Fix abundances

Author: Guilucand

crates/assembler_pipeline/Cargo.toml

@@ -55,6 +55,7 @@ support_kmer_counters = [
     "io/support_kmer_counters",
     "colors/support_kmer_counters",
     "structs/support_kmer_counters",
+    "sequence_output/support_kmer_counters",
 ]
 
 # Features used to filter ggcat capabilites

crates/assembler_pipeline/src/compute_matchtigs.rs

@@ -514,7 +514,7 @@ pub fn compute_matchtigs_thread<CX: ColorsManager, BK: StructuredSequenceBackend
                 next_sequence,
                 &handle.1,
                 bases_offset,
-                first_data.is_forwards(),
+                edge_data.is_forwards(),
                 &storage.extra_buffer,
                 None,
             );

crates/assembler_pipeline/src/eulertigs.rs

@@ -133,10 +133,17 @@ impl CircularUnitig {
     ) {
         let children_count = self.base_children.len();
 
+        #[cfg(feature = "support_kmer_counters")]
+        let mut counted_orig_indexes = std::collections::HashSet::with_capacity(children_count);
+
         #[cfg(feature = "support_kmer_counters")]
         {
             let first_unitig_entry = unitigs.get(&self.base_children[0].orig_index).unwrap();
-            _out_abundance.first = first_unitig_entry.2.first;
+            _out_abundance.first = if self.base_children[0].rc {
+                first_unitig_entry.2.last
+            } else {
+                first_unitig_entry.2.first
+            };
         }
 
         for child in self.base_children.iter().take(children_count - 1) {
@@ -146,7 +153,9 @@ impl CircularUnitig {
 
             #[cfg(feature = "support_kmer_counters")]
             {
-                _out_abundance.sum += _abundance.sum;
+                if counted_orig_indexes.insert(child.orig_index) {
+                    _out_abundance.sum += _abundance.sum;
+                }
             }
 
             let should_rc = child.rc;
@@ -178,7 +187,9 @@ impl CircularUnitig {
         #[cfg(feature = "support_kmer_counters")]
         {
             let abundance = &last_part_entry.2;
-            _out_abundance.sum += abundance.sum;
+            if counted_orig_indexes.insert(last.orig_index) {
+                _out_abundance.sum += abundance.sum;
+            }
             _out_abundance.last = if last.rc {
                 abundance.first
             } else {
@@ -909,6 +920,11 @@ mod tests {
     use hashes::cn_seqhash::u128::CanonicalSeqHashFactory;
     use io::compressed_read::CompressedRead;
 
+    #[cfg(feature = "support_kmer_counters")]
+    fn build_abundance(first: u64, sum: u64, last: u64) -> SequenceAbundanceType {
+        SequenceAbundanceType { first, sum, last }
+    }
+
     #[test]
     fn test_rc_rotation() {
         let k = 31;
@@ -984,4 +1000,60 @@ mod tests {
             }
         }
     }
+
+    #[cfg(feature = "support_kmer_counters")]
+    #[test]
+    fn rotated_circular_unitig_counts_split_abundance_once() {
+        let k = 5;
+        let mut stream = Vec::new();
+        let sequence = b"AACCGGTTAAC";
+        CompressedRead::compress_from_plain(sequence, |b| stream.extend_from_slice(b));
+        let read = CompressedRead::new_from_compressed(&stream, sequence.len());
+
+        let unitigs_hashmap = DashMap::new();
+        unitigs_hashmap.insert(
+            0,
+            (
+                CompressedReadIndipendent::from_read_inplace(&read, &stream),
+                NonColoredManager,
+                build_abundance(3, 21, 7),
+            ),
+        );
+
+        let mut circular_unitig = CircularUnitig::new();
+        circular_unitig.base_children.push(CircularUnitigPart {
+            orig_index: 0,
+            start_pos: 0,
+            length: sequence.len() - (k - 1),
+            rc: false,
+        });
+
+        circular_unitig.rotate_with_rc(0, 3, false);
+        assert_eq!(
+            circular_unitig
+                .base_children
+                .iter()
+                .filter(|part| part.orig_index == 0)
+                .count(),
+            2
+        );
+
+        let mut writer = AlignedDynamicCompressedRead::new();
+        let dummy_buffer = color_types::PartialUnitigsColorStructure::<NonColoredManager>::new_temp_buffer();
+        let mut colors_buffer = color_types::PartialUnitigsColorStructure::<NonColoredManager>::new_temp_buffer();
+        let mut abundance = SequenceAbundanceType::default();
+
+        circular_unitig.write_packed::<CanonicalSeqHashFactory, NonColoredManager>(
+            &stream,
+            &unitigs_hashmap,
+            &mut writer,
+            &dummy_buffer,
+            &mut colors_buffer,
+            k,
+            true,
+            &mut abundance,
+        );
+
+        assert_eq!(abundance.sum, 21);
+    }
 }

crates/sequence_output/Cargo.toml

@@ -31,4 +31,8 @@ serde = "1.0.228"
 
 
 [features]
-support_kmer_counters = []
+support_kmer_counters = [
+	"io/support_kmer_counters",
+	"colors/support_kmer_counters",
+	"structs/support_kmer_counters",
+]

crates/sequence_output/src/indirect_reads_extractor.rs

@@ -232,6 +232,8 @@ mod tests {
         bundles::multifile_building::ColorBundleMultifileBuilding,
         colors_manager::color_types::PartialUnitigsColorStructure,
     };
+    #[cfg(feature = "support_kmer_counters")]
+    use io::partial_unitigs_extra_data::SequenceAbundance;
     use io::{
         compressed_read::CompressedReadIndipendent,
         concurrent::temp_reads::extra_data::{
@@ -288,7 +290,7 @@ mod tests {
                 colors,
                 mode: PartialUnitigMode::Inline,
                 #[cfg(feature = "support_kmer_counters")]
-                counters: sequence_output::structured_sequences::SequenceAbundance {
+                counters: SequenceAbundance {
                     first: 1,
                     sum: runs.iter().map(|(_, count)| *count as u64).sum(),
                     last: 1,

crates/sequence_output/src/sequences_joiner.rs

@@ -286,13 +286,14 @@ impl<CX: ColorsManager> IndirectSequencesJoiner<CX> {
             } else {
                 (extra.counters.first, extra.counters.last)
             };
+            let overlapping_kmers = (overlapping_characters >= self.k) as u64;
 
             self.counters.first = if self.sequence.bases_count == 0 {
                 other_first
             } else {
                 self.counters.first
             };
-            self.counters.sum += extra.counters.sum - other_first;
+            self.counters.sum += extra.counters.sum - (other_first * overlapping_kmers);
             self.counters.last = other_last;
         }
 
@@ -473,21 +474,29 @@ mod tests {
     };
 
     use super::IndirectSequencesJoiner;
+    #[cfg(feature = "support_kmer_counters")]
+    use io::partial_unitigs_extra_data::SequenceAbundance;
     use io::partial_unitigs_extra_data::{PartialUnitigExtraData, PartialUnitigMode};
 
-    fn inline_extra() -> PartialUnitigExtraData<NonColoredManager> {
+    fn inline_extra_with_multiplicity(
+        multiplicity: u64,
+    ) -> PartialUnitigExtraData<NonColoredManager> {
         PartialUnitigExtraData {
             colors: NonColoredManager,
             mode: PartialUnitigMode::Inline,
             #[cfg(feature = "support_kmer_counters")]
-            counters: sequence_output::structured_sequences::SequenceAbundance {
-                first: 1,
-                sum: 1,
-                last: 1,
+            counters: SequenceAbundance {
+                first: multiplicity,
+                sum: multiplicity,
+                last: multiplicity,
             },
         }
     }
 
+    fn inline_extra() -> PartialUnitigExtraData<NonColoredManager> {
+        inline_extra_with_multiplicity(1)
+    }
+
     fn new_temp_path(prefix: &str) -> PathBuf {
         let now = SystemTime::now()
             .duration_since(UNIX_EPOCH)
@@ -516,6 +525,29 @@ mod tests {
         );
     }
 
+    #[cfg(feature = "support_kmer_counters")]
+    fn append_plain_with_multiplicity(
+        joiner: &mut IndirectSequencesJoiner<NonColoredManager>,
+        seq: &str,
+        overlap: usize,
+        is_rc: bool,
+        multiplicity: u64,
+    ) {
+        let mut storage = Vec::new();
+        let read = CompressedReadIndipendent::from_plain(seq.as_bytes(), &mut storage);
+        let mut extra = inline_extra_with_multiplicity(multiplicity);
+        extra.counters.sum = multiplicity * (seq.len() - joiner.k + 1) as u64;
+        let extra_buffer = <PartialUnitigExtraData<NonColoredManager>>::new_temp_buffer();
+        joiner.append_sequence(
+            read.as_reference(&storage),
+            &extra,
+            overlap,
+            is_rc,
+            &extra_buffer,
+            None,
+        );
+    }
+
     #[test]
     fn inline_join_various_lengths_forward() {
         let k = 5;
@@ -564,4 +596,46 @@ mod tests {
 
         fs::remove_file(&temp_path).unwrap();
     }
+
+    #[cfg(feature = "support_kmer_counters")]
+    #[test]
+    fn abundance_keeps_all_kmers_for_k_minus_one_overlap() {
+        let k = 5;
+        let temp_path = new_temp_path("joiner_abundance_k_minus_one_overlap");
+        let writer = ConcurrentFileWriter::create(&temp_path).unwrap();
+        let mut joiner = IndirectSequencesJoiner::<NonColoredManager>::new(k, writer.clone());
+
+        append_plain_with_multiplicity(&mut joiner, "AACCTGGA", 0, false, 2);
+        append_plain_with_multiplicity(&mut joiner, "TGGAACCC", k - 1, false, 5);
+
+        let joined = joiner.get_sequence();
+
+        assert_eq!(joined.sequence.to_string(), "AACCTGGAACCC");
+        assert_eq!(joined.extra.counters.first, 2);
+        assert_eq!(joined.extra.counters.last, 5);
+        assert_eq!(joined.extra.counters.sum, 28);
+
+        fs::remove_file(&temp_path).unwrap();
+    }
+
+    #[cfg(feature = "support_kmer_counters")]
+    #[test]
+    fn abundance_drops_one_kmer_for_exact_k_overlap() {
+        let k = 5;
+        let temp_path = new_temp_path("joiner_abundance_k_overlap");
+        let writer = ConcurrentFileWriter::create(&temp_path).unwrap();
+        let mut joiner = IndirectSequencesJoiner::<NonColoredManager>::new(k, writer.clone());
+
+        append_plain_with_multiplicity(&mut joiner, "AAACCCCTA", 0, false, 2);
+        append_plain_with_multiplicity(&mut joiner, "CCCTAT", k, false, 5);
+
+        let joined = joiner.get_sequence();
+
+        assert_eq!(joined.sequence.to_string(), "AAACCCCTAT");
+        assert_eq!(joined.extra.counters.first, 2);
+        assert_eq!(joined.extra.counters.last, 5);
+        assert_eq!(joined.extra.counters.sum, 15);
+
+        fs::remove_file(&temp_path).unwrap();
+    }
 }