Add method to change the number of rows in a `DenseMatrix`

src/dense.rs

@@ -3,31 +3,19 @@ use std::ops::IndexMut;
 
 /// An aligned dense matrix of with a constant number of columns.
 #[derive(Debug, Clone)]
-pub struct DenseMatrix<T: Default, const C: usize = 32> {
+pub struct DenseMatrix<T: Default + Copy, const C: usize = 32> {
     data: Vec<T>,
     indices: Vec<usize>,
 }
 
-impl<T: Default, const C: usize> DenseMatrix<T, C> {
-    /// Create a new matrix
+impl<T: Default + Copy, const C: usize> DenseMatrix<T, C> {
+    /// Create a new matrix with the given number of rows.
     pub fn new(rows: usize) -> Self {
-        // alway over-allocate columns to avoid alignment issues.
-        let c = C.next_power_of_two();
-        // allocate data block
-        let mut data = Vec::with_capacity( (rows+1) * c ); // FIXME
-        data.resize_with(data.capacity(), T::default);
-        // compute offset to have proper alignment
-        let mut offset = 0;
-        while data[offset..].as_ptr() as usize % c > 0 {
-            offset += 1
-        }
-        // record row coordinates (only really useful when aligned)
-        let mut indices = Vec::with_capacity(rows);
-        for i in 0..rows {
-            indices.push(offset + i * c)
-        }
-        // finish matrix
-        Self { data, indices }
+        let data = Vec::new();
+        let indices = Vec::new();
+        let mut matrix = Self { data, indices };
+        matrix.resize(rows);
+        matrix
     }
 
     /// The number of columns of the matrix.
@@ -39,9 +27,43 @@ impl<T: Default, const C: usize> DenseMatrix<T, C> {
     pub fn rows(&self) -> usize {
         self.indices.len()
     }
+
+    /// Change the number of rows of the matrix.
+    pub fn resize(&mut self, rows: usize) {
+
+        // alway over-allocate columns to avoid alignment issues.
+        let c = C.next_power_of_two();
+
+        //
+        let previous_rows = self.rows();
+        let previous_offset = if self.rows() == 0 { 0 } else { self.indices[0] };
+
+        // allocate data block
+        self.data.resize_with((rows + 1) * c, T::default);
+
+        // compute offset to aligned memory
+        let mut offset = 0;
+        while self.data[offset..].as_ptr() as usize % c > 0 {
+            offset += 1
+        }
+
+        // copy data in case alignment offset changed
+        if previous_offset != offset {
+            self.data.as_mut_slice().copy_within(
+                previous_offset..previous_offset+(previous_rows*c),
+                offset
+            );
+        }
+
+        // record row coordinates
+        self.indices.resize(rows, 0);
+        for i in 0..rows {
+            self.indices[i] = offset + i * c;
+        }
+    }
 }
 
-impl<T: Default, const C: usize> Index<usize> for DenseMatrix<T, C> {
+impl<T: Default + Copy, const C: usize> Index<usize> for DenseMatrix<T, C> {
     type Output = [T];
     fn index(&self, index: usize) -> &Self::Output {
         let row = self.indices[index];
@@ -49,9 +71,42 @@ impl<T: Default, const C: usize> Index<usize> for DenseMatrix<T, C> {
     }
 }
 
-impl<T: Default, const C: usize> IndexMut<usize> for DenseMatrix<T, C> {
+impl<T: Default + Copy, const C: usize> IndexMut<usize> for DenseMatrix<T, C> {
     fn index_mut(&mut self, index: usize) -> &mut Self::Output {
         let row = self.indices[index];
         &mut self.data[row..row + C]
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_resize() {
+        let mut dense = DenseMatrix::<u64, 32>::new(4);
+
+        for i in 0..4 {
+            dense[i][0] = (i + 1) as u64;
+        }
+        assert_eq!(dense[0][0], 1);
+        assert_eq!(dense[1][0], 2);
+        assert_eq!(dense[2][0], 3);
+        assert_eq!(dense[3][0], 4);
+        assert_eq!(dense[0].as_ptr() as usize % 4, 0);
+
+        dense.resize(256);
+        assert_eq!(dense[0][0], 1);
+        assert_eq!(dense[1][0], 2);
+        assert_eq!(dense[2][0], 3);
+        assert_eq!(dense[3][0], 4);
+        assert_eq!(dense[0].as_ptr() as usize % 4, 0);
+
+        dense.resize(512);
+        assert_eq!(dense[0][0], 1);
+        assert_eq!(dense[1][0], 2);
+        assert_eq!(dense[2][0], 3);
+        assert_eq!(dense[3][0], 4);
+        assert_eq!(dense[0].as_ptr() as usize % 4, 0);
+    }
+}

src/lib.rs

@@ -19,7 +19,7 @@ pub use pli::Pipeline;
 pub use pwm::Background;
 pub use pwm::CountMatrix;
 pub use pwm::ProbabilityMatrix;
-pub use pwm::WeightMatrix;
 pub use pwm::StripedScores;
+pub use pwm::WeightMatrix;
 pub use seq::EncodedSequence;
 pub use seq::StripedSequence;

src/pli.rs

@@ -124,7 +124,7 @@ impl Pipeline<DnaAlphabet, __m256> {
                 let mut s3 = _mm256_setzero_ps();
                 let mut s4 = _mm256_setzero_ps();
                 for j in 0..pwm.data.rows() {
-                    let x = _mm256_load_si256(seq.data[i+j].as_ptr() as *const __m256i);
+                    let x = _mm256_load_si256(seq.data[i + j].as_ptr() as *const __m256i);
                     let row = pwm.data[j].as_ptr();
                     // compute probabilities using an external lookup table
                     let p1 = _mm256_i32gather_ps(row, _mm256_shuffle_epi8(x, m1), S);

src/pwm.rs

@@ -182,4 +182,4 @@ impl<const C: usize> StripedScores<C> {
         }
         vec
     }
-}
\ No newline at end of file
+}

tests/dna.rs

@@ -34,7 +34,9 @@ fn test_score_generic() {
 
     let cm = CountMatrix::<DnaAlphabet, { DnaAlphabet::K }>::from_sequences(
         "MX000001",
-        PATTERNS.iter().map(|x| EncodedSequence::from_text(x).unwrap()),
+        PATTERNS
+            .iter()
+            .map(|x| EncodedSequence::from_text(x).unwrap()),
     )
     .unwrap();
     let pbm = cm.to_probability([0.1, 0.1, 0.1, 0.1, 0.0]);
@@ -47,10 +49,10 @@ fn test_score_generic() {
     assert_eq!(scores.len(), EXPECTED.len());
     for i in 0..scores.len() {
         assert!(
-            (scores[i] - EXPECTED[i]).abs() < 1e-5, 
-            "{} != {} at position {}", 
-            scores[i], 
-            EXPECTED[i], 
+            (scores[i] - EXPECTED[i]).abs() < 1e-5,
+            "{} != {} at position {}",
+            scores[i],
+            EXPECTED[i],
             i
         );
     }
@@ -66,7 +68,9 @@ fn test_score_avx2() {
 
     let cm = CountMatrix::<DnaAlphabet, { DnaAlphabet::K }>::from_sequences(
         "MX000001",
-        PATTERNS.iter().map(|x| EncodedSequence::from_text(x).unwrap()),
+        PATTERNS
+            .iter()
+            .map(|x| EncodedSequence::from_text(x).unwrap()),
     )
     .unwrap();
     let pbm = cm.to_probability([0.1, 0.1, 0.1, 0.1, 0.0]);
@@ -81,10 +85,10 @@ fn test_score_avx2() {
     // assert_eq!(scores[0], -23.07094); // -23.07094
     for i in 0..EXPECTED.len() {
         assert!(
-            (scores[i] - EXPECTED[i]).abs() < 1e-5, 
-            "{} != {} at position {}", 
-            scores[i], 
-            EXPECTED[i], 
+            (scores[i] - EXPECTED[i]).abs() < 1e-5,
+            "{} != {} at position {}",
+            scores[i],
+            EXPECTED[i],
             i
         );
     }