ocas_py/
matrix.rs

1//! Python `Matrix` class — dense matrices over ℤ, ℚ, or GF(p).
2//!
3//! Wraps [`ocas_poly::matrix::Matrix`] with an enum-erasure strategy,
4//! mirroring the approach used by [`crate::polynomial`].
5
6use crate::domain::DomainKind;
7use ocas_domain::{Domain, FiniteField, Integer, IntegerDomain, Rational, RationalDomain};
8use ocas_poly::matrix::{Matrix, MatrixError};
9use pyo3::exceptions::{PyTypeError, PyValueError};
10use pyo3::prelude::*;
11
12/// Type-erased matrix over one of the three supported domains.
13pub(crate) enum MatrixErased {
14    Int(Matrix<IntegerDomain>),
15    Rat(Matrix<RationalDomain>),
16    Fq(Matrix<FiniteField>),
17}
18
19/// A dense matrix.
20///
21/// `rows` is a list of rows, each a list of coefficients. The `domain`
22/// argument selects the coefficient ring exactly as for `Polynomial`.
23///
24/// ```python
25/// from ocas import Matrix
26///
27/// a = Matrix([[1, 2], [3, 4]])
28/// print(a.determinant())   # -2
29/// print((a @ a).rows())    # [[7, 10], [15, 22]]
30/// ```
31#[pyclass(name = "Matrix", skip_from_py_object)]
32pub struct PyMatrix {
33    pub(crate) inner: MatrixErased,
34}
35
36fn map_matrix_err(e: MatrixError) -> PyErr {
37    PyValueError::new_err(e.to_string())
38}
39
40/// Extract a 2-D integer coefficient matrix from a Python list-of-lists.
41fn extract_int_rows(obj: &Bound<'_, PyAny>) -> PyResult<Vec<Vec<Integer>>> {
42    let data: Vec<Vec<i64>> = obj
43        .extract()
44        .map_err(|_| PyTypeError::new_err("integer matrix entries must be ints"))?;
45    Ok(data
46        .into_iter()
47        .map(|r| r.into_iter().map(Integer::from).collect())
48        .collect())
49}
50
51/// Extract a 2-D rational coefficient matrix. Each entry is an int or a
52/// `(num, denom)` tuple; the whole matrix must be uniformly one form.
53fn extract_rat_rows(obj: &Bound<'_, PyAny>) -> PyResult<Vec<Vec<Rational>>> {
54    if let Ok(int_rows) = obj.extract::<Vec<Vec<i64>>>() {
55        Ok(int_rows
56            .into_iter()
57            .map(|r| r.into_iter().map(|n| Rational::new(n, 1)).collect())
58            .collect())
59    } else {
60        let pair_rows: Vec<Vec<(i64, i64)>> = obj.extract().map_err(|_| {
61            PyTypeError::new_err("rational entries must be ints or (num, denom) tuples")
62        })?;
63        pair_rows
64            .into_iter()
65            .map(|r| {
66                r.into_iter()
67                    .map(|(num, den)| {
68                        if den == 0 {
69                            Err(PyValueError::new_err("rational denominator cannot be zero"))
70                        } else {
71                            Ok(Rational::new(num, den))
72                        }
73                    })
74                    .collect()
75            })
76            .collect()
77    }
78}
79
80/// Build a `PyMatrix` from a Python list-of-lists and a domain kind.
81pub(crate) fn build_matrix(rows: &Bound<'_, PyAny>, domain: &DomainKind) -> PyResult<PyMatrix> {
82    let inner = match domain {
83        DomainKind::Integer => {
84            let r = extract_int_rows(rows)?;
85            MatrixErased::Int(Matrix::from_rows(r, IntegerDomain))
86        }
87        DomainKind::Rational => {
88            let r = extract_rat_rows(rows)?;
89            MatrixErased::Rat(Matrix::from_rows(r, RationalDomain))
90        }
91        DomainKind::FiniteField(p) => {
92            let field = FiniteField::new(p.clone());
93            let data: Vec<Vec<i64>> = rows
94                .extract()
95                .map_err(|_| PyTypeError::new_err("finite-field matrix entries must be ints"))?;
96            let rows: Vec<Vec<_>> = data
97                .into_iter()
98                .map(|r| r.into_iter().map(|v| field.element(v)).collect())
99                .collect();
100            MatrixErased::Fq(Matrix::from_rows(rows, field))
101        }
102    };
103    Ok(PyMatrix { inner })
104}
105
106#[pymethods]
107impl PyMatrix {
108    /// Create a matrix from a list of rows.
109    ///
110    /// `domain` selects the coefficient ring (`"integer"` is the default).
111    #[new]
112    #[pyo3(signature = (rows, domain=None))]
113    fn new(rows: &Bound<'_, PyAny>, domain: Option<&Bound<'_, PyAny>>) -> PyResult<Self> {
114        let kind = match domain {
115            Some(d) => DomainKind::from_py(d)?,
116            None => DomainKind::Integer,
117        };
118        build_matrix(rows, &kind)
119    }
120
121    /// Number of rows.
122    #[getter]
123    fn nrows(&self) -> usize {
124        match &self.inner {
125            MatrixErased::Int(m) => m.nrows(),
126            MatrixErased::Rat(m) => m.nrows(),
127            MatrixErased::Fq(m) => m.nrows(),
128        }
129    }
130
131    /// Number of columns.
132    #[getter]
133    fn ncols(&self) -> usize {
134        match &self.inner {
135            MatrixErased::Int(m) => m.ncols(),
136            MatrixErased::Rat(m) => m.ncols(),
137            MatrixErased::Fq(m) => m.ncols(),
138        }
139    }
140
141    /// Return the shape as `(nrows, ncols)`.
142    fn shape(&self) -> (usize, usize) {
143        (self.nrows(), self.ncols())
144    }
145
146    /// Return `self[(i, j)]`.
147    fn __getitem__(&self, idx: (usize, usize)) -> PyResult<String> {
148        let (i, j) = idx;
149        let out = match &self.inner {
150            MatrixErased::Int(m) => {
151                if i >= m.nrows() || j >= m.ncols() {
152                    return Err(PyValueError::new_err("index out of bounds"));
153                }
154                m[(i, j)].to_string()
155            }
156            MatrixErased::Rat(m) => {
157                if i >= m.nrows() || j >= m.ncols() {
158                    return Err(PyValueError::new_err("index out of bounds"));
159                }
160                m[(i, j)].to_string()
161            }
162            MatrixErased::Fq(m) => {
163                if i >= m.nrows() || j >= m.ncols() {
164                    return Err(PyValueError::new_err("index out of bounds"));
165                }
166                m[(i, j)].value().to_string()
167            }
168        };
169        Ok(out)
170    }
171
172    /// Return all rows as a list of lists of decimal strings (rational
173    /// entries are rendered as `n/d`). Wrap each entry in `int(...)` to
174    /// obtain Python integers.
175    fn rows(&self) -> Vec<Vec<String>> {
176        match &self.inner {
177            MatrixErased::Int(m) => (0..m.nrows())
178                .map(|i| (0..m.ncols()).map(|j| m[(i, j)].to_string()).collect())
179                .collect(),
180            MatrixErased::Rat(m) => (0..m.nrows())
181                .map(|i| (0..m.ncols()).map(|j| m[(i, j)].to_string()).collect())
182                .collect(),
183            MatrixErased::Fq(m) => (0..m.nrows())
184                .map(|i| {
185                    (0..m.ncols())
186                        .map(|j| m[(i, j)].value().to_string())
187                        .collect()
188                })
189                .collect(),
190        }
191    }
192
193    /// Return the transpose.
194    fn transpose(&self) -> PyMatrix {
195        match &self.inner {
196            MatrixErased::Int(m) => PyMatrix {
197                inner: MatrixErased::Int(m.transpose()),
198            },
199            MatrixErased::Rat(m) => PyMatrix {
200                inner: MatrixErased::Rat(m.transpose()),
201            },
202            MatrixErased::Fq(m) => PyMatrix {
203                inner: MatrixErased::Fq(m.transpose()),
204            },
205        }
206    }
207
208    /// Return the trace (sum of the diagonal) of a square matrix.
209    fn trace(&self) -> PyResult<String> {
210        match &self.inner {
211            MatrixErased::Int(m) => Ok(m.trace().map_err(map_matrix_err)?.to_string()),
212            MatrixErased::Rat(m) => Ok(m.trace().map_err(map_matrix_err)?.to_string()),
213            MatrixErased::Fq(m) => Ok(m.trace().map_err(map_matrix_err)?.value().to_string()),
214        }
215    }
216
217    /// Return the rank.
218    fn rank(&self) -> usize {
219        match &self.inner {
220            MatrixErased::Int(m) => m.rank(),
221            MatrixErased::Rat(m) => m.rank(),
222            MatrixErased::Fq(m) => m.rank(),
223        }
224    }
225
226    /// Return the determinant of a square matrix.
227    fn determinant(&self) -> PyResult<String> {
228        match &self.inner {
229            MatrixErased::Int(m) => Ok(m.determinant().map_err(map_matrix_err)?.to_string()),
230            MatrixErased::Rat(m) => Ok(m.determinant().map_err(map_matrix_err)?.to_string()),
231            MatrixErased::Fq(m) => Ok(m.determinant().map_err(map_matrix_err)?.value().to_string()),
232        }
233    }
234
235    /// Return the inverse, or raise `ValueError` if singular/non-square.
236    fn inverse(&self) -> PyResult<PyMatrix> {
237        match &self.inner {
238            MatrixErased::Int(m) => Ok(PyMatrix {
239                inner: MatrixErased::Int(m.inverse().map_err(map_matrix_err)?),
240            }),
241            MatrixErased::Rat(m) => Ok(PyMatrix {
242                inner: MatrixErased::Rat(m.inverse().map_err(map_matrix_err)?),
243            }),
244            MatrixErased::Fq(m) => Ok(PyMatrix {
245                inner: MatrixErased::Fq(m.inverse().map_err(map_matrix_err)?),
246            }),
247        }
248    }
249
250    /// Solve `self * x = rhs` for the vector `rhs`.
251    ///
252    /// `rhs` is a list of ints (integer/fq) or `(num, denom)` tuples
253    /// (rational). Returns the solution as a list of decimal strings.
254    fn solve(&self, rhs: &Bound<'_, PyAny>) -> PyResult<Vec<String>> {
255        match &self.inner {
256            MatrixErased::Int(m) => {
257                let b: Vec<Integer> = extract_int_vector(rhs)?;
258                let sol = m.solve(&b).map_err(map_matrix_err)?;
259                Ok(sol.into_iter().map(|c| c.to_string()).collect())
260            }
261            MatrixErased::Rat(m) => {
262                let b: Vec<Rational> = extract_rat_vector(rhs)?;
263                let sol = m.solve(&b).map_err(map_matrix_err)?;
264                Ok(sol.into_iter().map(|c| c.to_string()).collect())
265            }
266            MatrixErased::Fq(m) => {
267                let field = m.domain().clone();
268                let ints: Vec<i64> = rhs
269                    .extract()
270                    .map_err(|_| PyTypeError::new_err("finite-field rhs entries must be ints"))?;
271                let b: Vec<_> = ints.into_iter().map(|v| field.element(v)).collect();
272                let sol = m.solve(&b).map_err(map_matrix_err)?;
273                Ok(sol.into_iter().map(|c| c.value().to_string()).collect())
274            }
275        }
276    }
277
278    /// Matrix product `self @ other`.
279    fn __matmul__(&self, other: &PyMatrix) -> PyResult<PyMatrix> {
280        match (&self.inner, &other.inner) {
281            (MatrixErased::Int(a), MatrixErased::Int(b)) => Ok(PyMatrix {
282                inner: MatrixErased::Int(a.matmul(b).map_err(map_matrix_err)?),
283            }),
284            (MatrixErased::Rat(a), MatrixErased::Rat(b)) => Ok(PyMatrix {
285                inner: MatrixErased::Rat(a.matmul(b).map_err(map_matrix_err)?),
286            }),
287            (MatrixErased::Fq(a), MatrixErased::Fq(b)) => Ok(PyMatrix {
288                inner: MatrixErased::Fq(a.matmul(b).map_err(map_matrix_err)?),
289            }),
290            _ => Err(PyTypeError::new_err(
291                "@ requires both matrices to share the same coefficient domain",
292            )),
293        }
294    }
295
296    /// Element-wise addition `self + other`.
297    fn __add__(&self, other: &PyMatrix) -> PyResult<PyMatrix> {
298        match (&self.inner, &other.inner) {
299            (MatrixErased::Int(a), MatrixErased::Int(b)) => {
300                check_shape(a.nrows(), a.ncols(), b.nrows(), b.ncols())?;
301                let d = *a.domain();
302                let mut data = Vec::with_capacity(a.nrows() * a.ncols());
303                for i in 0..a.nrows() {
304                    for j in 0..a.ncols() {
305                        data.push(d.add(&a[(i, j)], &b[(i, j)]));
306                    }
307                }
308                Ok(PyMatrix {
309                    inner: MatrixErased::Int(Matrix::new(a.nrows(), a.ncols(), data, d)),
310                })
311            }
312            (MatrixErased::Rat(a), MatrixErased::Rat(b)) => {
313                check_shape(a.nrows(), a.ncols(), b.nrows(), b.ncols())?;
314                let d = *a.domain();
315                let mut data = Vec::with_capacity(a.nrows() * a.ncols());
316                for i in 0..a.nrows() {
317                    for j in 0..a.ncols() {
318                        data.push(d.add(&a[(i, j)], &b[(i, j)]));
319                    }
320                }
321                Ok(PyMatrix {
322                    inner: MatrixErased::Rat(Matrix::new(a.nrows(), a.ncols(), data, d)),
323                })
324            }
325            (MatrixErased::Fq(a), MatrixErased::Fq(b)) => {
326                check_shape(a.nrows(), a.ncols(), b.nrows(), b.ncols())?;
327                let d = a.domain().clone();
328                let mut data = Vec::with_capacity(a.nrows() * a.ncols());
329                for i in 0..a.nrows() {
330                    for j in 0..a.ncols() {
331                        data.push(d.add(&a[(i, j)], &b[(i, j)]));
332                    }
333                }
334                Ok(PyMatrix {
335                    inner: MatrixErased::Fq(Matrix::new(a.nrows(), a.ncols(), data, d)),
336                })
337            }
338            _ => Err(PyTypeError::new_err(
339                "+ requires both matrices to share the same coefficient domain",
340            )),
341        }
342    }
343
344    /// Element-wise subtraction `self - other`.
345    fn __sub__(&self, other: &PyMatrix) -> PyResult<PyMatrix> {
346        match (&self.inner, &other.inner) {
347            (MatrixErased::Int(a), MatrixErased::Int(b)) => {
348                check_shape(a.nrows(), a.ncols(), b.nrows(), b.ncols())?;
349                let d = *a.domain();
350                let mut data = Vec::with_capacity(a.nrows() * a.ncols());
351                for i in 0..a.nrows() {
352                    for j in 0..a.ncols() {
353                        data.push(d.sub(&a[(i, j)], &b[(i, j)]));
354                    }
355                }
356                Ok(PyMatrix {
357                    inner: MatrixErased::Int(Matrix::new(a.nrows(), a.ncols(), data, d)),
358                })
359            }
360            (MatrixErased::Rat(a), MatrixErased::Rat(b)) => {
361                check_shape(a.nrows(), a.ncols(), b.nrows(), b.ncols())?;
362                let d = *a.domain();
363                let mut data = Vec::with_capacity(a.nrows() * a.ncols());
364                for i in 0..a.nrows() {
365                    for j in 0..a.ncols() {
366                        data.push(d.sub(&a[(i, j)], &b[(i, j)]));
367                    }
368                }
369                Ok(PyMatrix {
370                    inner: MatrixErased::Rat(Matrix::new(a.nrows(), a.ncols(), data, d)),
371                })
372            }
373            (MatrixErased::Fq(a), MatrixErased::Fq(b)) => {
374                check_shape(a.nrows(), a.ncols(), b.nrows(), b.ncols())?;
375                let d = a.domain().clone();
376                let mut data = Vec::with_capacity(a.nrows() * a.ncols());
377                for i in 0..a.nrows() {
378                    for j in 0..a.ncols() {
379                        data.push(d.sub(&a[(i, j)], &b[(i, j)]));
380                    }
381                }
382                Ok(PyMatrix {
383                    inner: MatrixErased::Fq(Matrix::new(a.nrows(), a.ncols(), data, d)),
384                })
385            }
386            _ => Err(PyTypeError::new_err(
387                "- requires both matrices to share the same coefficient domain",
388            )),
389        }
390    }
391
392    fn __eq__(&self, other: &PyMatrix) -> bool {
393        match (&self.inner, &other.inner) {
394            (MatrixErased::Int(a), MatrixErased::Int(b)) => a == b,
395            (MatrixErased::Rat(a), MatrixErased::Rat(b)) => a == b,
396            (MatrixErased::Fq(a), MatrixErased::Fq(b)) => a == b,
397            _ => false,
398        }
399    }
400
401    fn __repr__(&self) -> String {
402        let dom = match &self.inner {
403            MatrixErased::Int(_) => "integer",
404            MatrixErased::Rat(_) => "rational",
405            MatrixErased::Fq(_) => "finite-field",
406        };
407        format!(
408            "Matrix({}x{}, domain='{}')",
409            self.nrows(),
410            self.ncols(),
411            dom
412        )
413    }
414}
415
416fn check_shape(r1: usize, c1: usize, r2: usize, c2: usize) -> PyResult<()> {
417    if r1 != r2 || c1 != c2 {
418        Err(PyValueError::new_err(format!(
419            "shape mismatch: {r1}x{c1} vs {r2}x{c2}"
420        )))
421    } else {
422        Ok(())
423    }
424}
425
426/// Extract an integer vector from a Python iterable of ints.
427fn extract_int_vector(obj: &Bound<'_, PyAny>) -> PyResult<Vec<Integer>> {
428    let ints: Vec<i64> = obj
429        .extract()
430        .map_err(|_| PyTypeError::new_err("integer rhs entries must be ints"))?;
431    Ok(ints.into_iter().map(Integer::from).collect())
432}
433
434/// Extract a rational vector from a Python iterable. Entries may be ints or
435/// `(num, denom)` tuples, uniformly.
436fn extract_rat_vector(obj: &Bound<'_, PyAny>) -> PyResult<Vec<Rational>> {
437    if let Ok(ints) = obj.extract::<Vec<i64>>() {
438        Ok(ints.into_iter().map(|n| Rational::new(n, 1)).collect())
439    } else {
440        let pairs: Vec<(i64, i64)> = obj.extract().map_err(|_| {
441            PyTypeError::new_err("rational rhs entries must be ints or (num, denom) tuples")
442        })?;
443        pairs
444            .into_iter()
445            .map(|(num, den)| {
446                if den == 0 {
447                    Err(PyValueError::new_err("rational denominator cannot be zero"))
448                } else {
449                    Ok(Rational::new(num, den))
450                }
451            })
452            .collect()
453    }
454}
ocas_py/matrix.rs

ocas_py/
matrix.rs
