fp/matrix/
subspace.rs

1use std::{io, ops::Deref};
2
3use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
4use itertools::Itertools;
5use serde::{Deserialize, Serialize};
6
7use super::Matrix;
8use crate::{
9    prime::ValidPrime,
10    vector::{FpSlice, FpSliceMut, FpVector},
11};
12
13/// A subspace of a vector space.
14///
15/// In general, a method is defined on the [`Subspace`] if it is a meaningful property of the
16/// subspace itself. Otherwise, users can dereference the subspace to gain read-only access to the
17/// underlying matrix object.
18///
19/// # Fields
20///  * `matrix` - A matrix in reduced row echelon, whose number of columns is the dimension of the
21///    ambient space and each row is a basis vector of the subspace.
22#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
23#[repr(transparent)]
24pub struct Subspace {
25    matrix: Matrix,
26}
27
28// We implement `Deref` to make it easier to access the methods of the underlying matrix. Since we
29// don't implement `DerefMut`, we still ensure that the matrix stays row reduced.
30impl Deref for Subspace {
31    type Target = Matrix;
32
33    fn deref(&self) -> &Self::Target {
34        &self.matrix
35    }
36}
37
38impl Subspace {
39    pub fn new(p: ValidPrime, dim: usize) -> Self {
40        // We add an extra row to the matrix to allow for adding vectors to the subspace. This way,
41        // even if the subspace is already the entire ambient space, we still have the space to add
42        // one more vector, which will then be reduced to zero by the row reduction.
43        let mut matrix = Matrix::new(p, dim + 1, dim);
44        matrix.initialize_pivots();
45        Self::from_matrix(matrix)
46    }
47
48    /// Create a new subspace from a matrix. The matrix does not have to be in row echelon form.
49    pub fn from_matrix(mut matrix: Matrix) -> Self {
50        matrix.row_reduce();
51        Self { matrix }
52    }
53
54    /// Run a closure on the matrix and then ensure it is row-reduced.
55    pub fn update_then_row_reduce<T, F: FnOnce(&mut Matrix) -> T>(&mut self, f: F) -> T {
56        let ret = f(&mut self.matrix);
57        self.matrix.row_reduce();
58        ret
59    }
60
61    pub fn from_bytes(p: ValidPrime, data: &mut impl io::Read) -> io::Result<Self> {
62        let rows = data.read_u64::<LittleEndian>()? as usize;
63        let ambient_dimension = data.read_u64::<LittleEndian>()? as usize;
64
65        let mut matrix = Matrix::from_bytes(p, rows, ambient_dimension, data)?;
66
67        matrix.pivots = Matrix::read_pivot(matrix.columns(), data)?;
68
69        Ok(Self { matrix })
70    }
71
72    pub fn to_bytes(&self, buffer: &mut impl io::Write) -> io::Result<()> {
73        buffer.write_u64::<LittleEndian>(self.matrix.rows() as u64)?;
74        buffer.write_u64::<LittleEndian>(self.ambient_dimension() as u64)?;
75
76        self.matrix.to_bytes(buffer)?;
77        Matrix::write_pivot(self.pivots(), buffer)
78    }
79
80    pub fn entire_space(p: ValidPrime, dim: usize) -> Self {
81        let mut result = Self::new(p, dim);
82        for i in 0..dim {
83            result.matrix.row_mut(i).set_entry(i, 1);
84            result.matrix.pivots_mut()[i] = i as isize;
85        }
86        result
87    }
88
89    /// This adds a vector to the subspace. This function assumes that the last row of the
90    /// matrix is zero, i.e. the dimension of the current subspace is strictly less than the number
91    /// of rows. This can be achieved by setting the number of rows to be the dimension plus one
92    /// when creating the subspace.
93    ///
94    /// # Returns
95    /// The new dimension of the subspace
96    pub fn add_vector(&mut self, row: FpSlice) -> usize {
97        let last_row = self.matrix.rows() - 1;
98        self.matrix.row_mut(last_row).assign(row);
99        self.matrix.row_reduce()
100    }
101
102    /// This adds some rows to the subspace
103    ///
104    /// # Arguments
105    ///  - `rows`: A function that writes the row to be added to the given FpSliceMut. This returns
106    ///    `None` if it runs out of rows, `Some(())` otherwise.
107    pub fn add_vectors(&mut self, mut rows: impl for<'a> FnMut(FpSliceMut<'a>) -> Option<()>) {
108        let num_rows = self.matrix.rows();
109        'outer: loop {
110            let first_row = self.dimension();
111            if first_row == num_rows {
112                return;
113            }
114
115            for i in first_row..num_rows {
116                if rows(self.matrix.row_mut(i)).is_none() {
117                    break 'outer;
118                }
119            }
120            self.matrix.row_reduce();
121        }
122        self.matrix.row_reduce();
123    }
124
125    pub fn add_basis_elements(&mut self, mut rows: impl std::iter::Iterator<Item = usize>) {
126        self.add_vectors(|mut row| {
127            row.set_entry(rows.next()?, 1);
128            Some(())
129        });
130    }
131
132    /// Projects a vector to a complement of the subspace. The complement is the set of vectors
133    /// that have a 0 in every column where there is a pivot in `matrix`
134    pub fn reduce(&self, mut vector: FpSliceMut) {
135        assert_eq!(vector.as_slice().len(), self.ambient_dimension());
136        if self.matrix.rows() == 0 {
137            return;
138        }
139        let p = self.prime();
140        let iter = self
141            .pivots()
142            .iter()
143            .enumerate()
144            .filter(|(_, x)| **x >= 0)
145            .map(|(col, _)| col)
146            .zip(self.iter());
147        for (col, row) in iter {
148            let c = vector.as_slice().entry(col);
149            if c != 0 {
150                vector.add(row, p - c);
151            }
152        }
153    }
154
155    pub fn contains(&self, vector: FpSlice) -> bool {
156        let mut vector: FpVector = vector.to_owned();
157        self.reduce(vector.as_slice_mut());
158        vector.is_zero()
159    }
160
161    pub fn contains_space(&self, other: &Self) -> bool {
162        other.iter().all(|row| self.contains(row))
163    }
164
165    pub fn dimension(&self) -> usize {
166        self.pivots()
167            .iter()
168            .rev()
169            .find(|&&i| i >= 0)
170            .map_or(0, |&i| i as usize + 1)
171    }
172
173    /// Whether the subspace is empty. This assumes the subspace is row reduced.
174    pub fn is_empty(&self) -> bool {
175        self.matrix.rows() == 0 || self.matrix.row(0).is_zero()
176    }
177
178    pub fn ambient_dimension(&self) -> usize {
179        self.matrix.columns()
180    }
181
182    /// Returns a basis of the subspace.
183    pub fn basis(&self) -> impl Iterator<Item = FpSlice<'_>> {
184        self.matrix.iter().take(self.dimension())
185    }
186
187    /// Sets the subspace to be the zero subspace.
188    pub fn set_to_zero(&mut self) {
189        self.matrix.set_to_zero();
190        for x in self.matrix.pivots_mut() {
191            *x = -1;
192        }
193    }
194
195    /// Sets the subspace to be the entire subspace.
196    pub fn set_to_entire(&mut self) {
197        self.matrix.set_to_zero();
198        for i in 0..self.matrix.columns() {
199            self.matrix.row_mut(i).set_entry(i, 1);
200            self.matrix.pivots_mut()[i] = i as isize;
201        }
202    }
203
204    pub fn iter(&self) -> impl Iterator<Item = FpSlice<'_>> {
205        self.matrix.iter().take(self.dimension())
206    }
207
208    /// Iterate over all vectors in the subspace.
209    ///
210    /// # Example
211    /// ```
212    /// # use fp::{prime::ValidPrime, matrix::{Matrix, Subspace}, vector::FpVector};
213    /// let subspace = Subspace::from_matrix(Matrix::from_vec(
214    ///     ValidPrime::new(3),
215    ///     &[vec![1, 0, 0], vec![0, 1, 2]],
216    /// ));
217    /// let mut all_vectors = subspace.iter_all_vectors().map(|v| (&v).into());
218    ///
219    /// assert_eq!(all_vectors.next(), Some(vec![0, 0, 0]));
220    /// assert_eq!(all_vectors.next(), Some(vec![0, 1, 2]));
221    /// assert_eq!(all_vectors.next(), Some(vec![0, 2, 1]));
222    /// assert_eq!(all_vectors.next(), Some(vec![1, 0, 0]));
223    /// assert_eq!(all_vectors.next(), Some(vec![1, 1, 2]));
224    /// assert_eq!(all_vectors.next(), Some(vec![1, 2, 1]));
225    /// assert_eq!(all_vectors.next(), Some(vec![2, 0, 0]));
226    /// assert_eq!(all_vectors.next(), Some(vec![2, 1, 2]));
227    /// assert_eq!(all_vectors.next(), Some(vec![2, 2, 1]));
228    /// assert_eq!(all_vectors.next(), None);
229    /// ```
230    pub fn iter_all_vectors(&self) -> impl Iterator<Item = FpVector> + '_ {
231        crate::prime::iter::combinations(self.prime(), self.dimension()).map(|combination| {
232            let mut vector = FpVector::new(self.prime(), self.ambient_dimension());
233            for (&c, v) in combination.iter().zip(self.matrix.iter()) {
234                vector.as_slice_mut().add(v, c);
235            }
236            vector
237        })
238    }
239
240    pub fn sum(&self, other: &Self) -> Self {
241        assert_eq!(self.prime(), other.prime());
242        assert_eq!(self.ambient_dimension(), other.ambient_dimension());
243
244        let mut sum = self.matrix.clone();
245        for other_row in other.iter() {
246            let mut new_row = sum.add_row();
247            new_row.assign(other_row);
248        }
249
250        let mut ret = Self::from_matrix(sum);
251        ret.matrix.trim(0, self.matrix.columns() + 1, 0);
252        ret
253    }
254}
255
256impl std::fmt::Display for Subspace {
257    /// # Example
258    /// ```
259    /// # use expect_test::expect;
260    /// # use fp::matrix::Subspace;
261    /// # use fp::prime::TWO;
262    /// let subspace = Subspace::entire_space(TWO, 3);
263    ///
264    /// expect![[r#"
265    ///     [1, 0, 0]
266    ///     [0, 1, 0]
267    ///     [0, 0, 1]
268    /// "#]]
269    /// .assert_eq(&format!("{}", subspace));
270    ///
271    /// assert_eq!(
272    ///     "[1, 0, 0], [0, 1, 0], [0, 0, 1]",
273    ///     &format!("{:#}", subspace)
274    /// );
275    /// ```
276    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
277        let dim = self.dimension();
278        let mut rows = self.matrix.iter().take(dim);
279
280        let output = if f.alternate() {
281            rows.join(", ")
282        } else {
283            rows.fold(String::new(), |mut output, row| {
284                use std::fmt::Write;
285                let _ = writeln!(output, "{row}");
286                output
287            })
288        };
289
290        write!(f, "{output}")?;
291        Ok(())
292    }
293}
294
295#[cfg(feature = "proptest")]
296pub mod arbitrary {
297    use proptest::prelude::*;
298
299    use super::*;
300    pub use crate::matrix::matrix_inner::arbitrary::MAX_COLUMNS as MAX_DIM;
301    use crate::matrix::matrix_inner::arbitrary::MatrixArbParams;
302
303    #[derive(Debug, Clone)]
304    pub struct SubspaceArbParams {
305        pub p: Option<ValidPrime>,
306        pub dim: BoxedStrategy<usize>,
307    }
308
309    impl Default for SubspaceArbParams {
310        fn default() -> Self {
311            Self {
312                p: None,
313                dim: (0..=MAX_DIM).boxed(),
314            }
315        }
316    }
317
318    impl Arbitrary for Subspace {
319        type Parameters = SubspaceArbParams;
320        type Strategy = BoxedStrategy<Self>;
321
322        fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
323            let p = match args.p {
324                Some(p) => Just(p).boxed(),
325                None => any::<ValidPrime>().boxed(),
326            };
327
328            (p, args.dim)
329                .prop_flat_map(move |(p, dim)| {
330                    Matrix::arbitrary_rref_with(MatrixArbParams {
331                        p: Some(p),
332                        rows: (0..=dim + 1).boxed(),
333                        columns: Just(dim).boxed(),
334                    })
335                })
336                .prop_map(Self::from_matrix)
337                .boxed()
338        }
339    }
340}