Skip to main content

sirius/solver/poly/
mod.rs

1//! A multivariable [polynomial ring](https://en.wikipedia.org/wiki/Polynomial_ring) with integer coefficients.
2
3use std::cmp::Ordering;
4
5pub mod coef;
6pub mod mono;
7
8mod arithmetic;
9mod display;
10mod macros;
11
12pub use arithmetic::monomial_div;
13pub use macros::poly;
14
15use coef::Coef;
16pub use mono::{Mono, Pow, Var};
17
18/// A multivariable polynomial with integer coefficients in canonical form, $\mathbb{Z}[x_1,x_2 \dots x_i]$.
19///
20/// For example, $4x^2y^5 - 7xz + 3y + 1$.
21///
22/// ```
23/// use sirius::solver::poly::poly;
24///
25/// assert_eq!(poly!(a^2 - b^2).try_divide(&poly!(a - b)), Some(poly!(a + b)));
26/// assert_eq!(poly!(2*x + 5).substitute('x', &poly!(z^2 + 1)), poly!(2*z^2 + 7));
27/// ```
28#[derive(PartialEq, Eq, Hash, Clone)]
29pub struct Poly {
30    terms: Vec<(Coef, Mono)>,
31}
32
33impl Poly {
34    pub fn zero() -> Self {
35        Self { terms: vec![] }
36    }
37
38    pub fn is_zero(&self) -> bool {
39        self.terms.is_empty()
40    }
41
42    pub fn constant<T: Into<Coef>>(c: T) -> Self {
43        let c = c.into();
44        if c == 0.into() {
45            Poly::zero()
46        } else {
47            Poly {
48                terms: vec![(c, Mono::unit())],
49            }
50        }
51    }
52
53    pub fn var<T: Into<Var>>(v: T, p: Pow) -> Self {
54        if p == 0 {
55            Self::constant(1)
56        } else {
57            Self::term(1, Mono::new(vec![(v.into(), p)]))
58        }
59    }
60
61    pub fn term<T: Into<Coef>>(c: T, m: Mono) -> Self {
62        let c = c.into();
63        if c == 0.into() {
64            Poly::zero()
65        } else {
66            Poly {
67                terms: vec![(c, m)],
68            }
69        }
70    }
71
72    pub fn terms(&self) -> Vec<(Coef, Mono)> {
73        self.terms.clone()
74    }
75
76    pub fn from_terms(terms: impl IntoIterator<Item = (Coef, Mono)>) -> Self {
77        let mut raw: Vec<_> = terms.into_iter().collect();
78        raw.sort_unstable_by(|(_, m1), (_, m2)| m2.cmp(m1));
79        let mut terms = Vec::with_capacity(raw.len());
80        for (c, m) in raw {
81            match terms.last_mut() {
82                Some((lc, lm)) if *lm == m => {
83                    *lc = *lc + c;
84                    if *lc == 0.into() {
85                        terms.pop();
86                    }
87                }
88                _ => {
89                    if c != 0.into() {
90                        terms.push((c, m));
91                    }
92                }
93            }
94        }
95        Self { terms }.debug_checked()
96    }
97
98    pub fn total_degree(&self) -> Pow {
99        self.terms
100            .first()
101            .map(|(_, m)| m.total_degree())
102            .unwrap_or(0)
103    }
104
105    pub fn degree_in<T: Into<Var>>(&self, v: T) -> Pow {
106        let v = v.into();
107        let mut deg = 0;
108        for (_, mono) in &self.terms {
109            let mut term_deg = 0;
110            for (var, pow) in mono.exps() {
111                term_deg += pow;
112                match var.cmp(&v) {
113                    Ordering::Greater => continue,
114                    Ordering::Equal => {
115                        deg = deg.max(*pow);
116                        break;
117                    }
118                    Ordering::Less => break,
119                }
120            }
121            if term_deg < deg {
122                break;
123            }
124        }
125
126        deg
127    }
128
129    pub fn leading_term(&self) -> (Coef, Mono) {
130        match self.terms.first() {
131            Some(m) => m.clone(),
132            None => (Coef::new(0), Mono::unit()),
133        }
134    }
135
136    pub fn eval(&self, mut point: impl FnMut(Var) -> i128) -> Coef {
137        let mut acc = Coef::new(0);
138        for (c, m) in &self.terms {
139            let mut t = *c;
140            for &(v, e) in m.exps() {
141                t = t * Coef::new(point(v).checked_pow(e).unwrap());
142            }
143            acc = acc + t;
144        }
145        acc
146    }
147
148    pub fn map_vars(&self, mut f: impl FnMut(Var) -> Self) -> Self {
149        let mut acc = Poly::zero();
150        for (c, m) in &self.terms {
151            let mut t = Poly::constant(*c);
152            for &(v, e) in m.exps() {
153                t = t.mul(&f(v).pow(e));
154            }
155            acc = acc.add(&t);
156        }
157        acc
158    }
159
160    pub fn substitute<T: Into<Var>>(&self, v: T, q: &Self) -> Self {
161        let v = v.into();
162        self.map_vars(|w| if w == v { q.clone() } else { Self::var(w, 1) })
163    }
164
165    pub fn coef_gcd(&self) -> Coef {
166        self.terms
167            .iter()
168            .fold(Coef::from(0), |g, &(c, _)| g.gcd(&c))
169            .abs()
170    }
171
172    pub fn coef_lcm(&self) -> Coef {
173        self.terms
174            .iter()
175            .fold(Coef::from(1), |g, &(c, _)| g.lcm(&c))
176            .abs()
177    }
178
179    pub fn divide_coefs<T: Into<Coef>>(&self, d: T) -> Option<Self> {
180        let d = d.into();
181        if d.is_zero() {
182            return None;
183        }
184        let mut terms = Vec::with_capacity(self.terms.len());
185        for (c, m) in &self.terms {
186            let (quot, rem) = c.divrem(d);
187            if rem.is_zero() {
188                terms.push((quot, m.clone()));
189            } else {
190                return None;
191            }
192        }
193        Some(Poly { terms }.debug_checked())
194    }
195
196    pub fn as_constant(&self) -> Option<Coef> {
197        match self.terms.as_slice() {
198            [] => Some(Coef::new(0)),
199            [(c, m)] if m.is_unit() => Some(*c),
200            _ => None,
201        }
202    }
203
204    pub fn vars(&self) -> Vec<Var> {
205        let mut vs = vec![];
206
207        for term in &self.terms {
208            for (var, _) in term.1.exps() {
209                if !vs.contains(var) {
210                    vs.push(*var);
211                }
212            }
213        }
214
215        vs.sort_unstable();
216        vs
217    }
218
219    pub fn always_nonneg(&self, nonneg: &dyn Fn(Var) -> bool) -> bool {
220        self.terms()
221            .iter()
222            .all(|(c, m)| c.is_positive() && m.always_nonneg(nonneg))
223    }
224
225    pub fn assert_canonical(&self) {
226        for w in self.terms.windows(2) {
227            assert!(
228                w[0].1 > w[1].1,
229                "terms not strictly descending: {:?} then {:?}",
230                w[0],
231                w[1]
232            );
233        }
234        for (c, m) in &self.terms {
235            assert!(*c != 0.into(), "zero coefficient on {m:?}");
236            m.assert_canonical();
237        }
238    }
239
240    fn debug_checked(self) -> Self {
241        #[cfg(debug_assertions)]
242        self.assert_canonical();
243        self
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use super::{Coef, Poly, monomial_div, poly};
250
251    #[test]
252    fn constants() {
253        let x = poly!(3 * x ^ 2 + 5 * z - 2 + 1);
254        let y = poly!(4 * y * z);
255        assert_eq!(
256            x.mul(&y),
257            poly!(12 * x ^ 2 * y * z + 20 * y * z ^ 2 - 4 * y * z)
258        );
259        assert_eq!(poly!(), poly!(x - x));
260    }
261
262    #[test]
263    fn arithmetic_sanity() {
264        let x = poly!(a ^ 2 + 2 * b + c);
265        let y = poly!(2 * a ^ 2 - c ^ 3 + d);
266        assert!(x.add(&y) == poly!(-1 * c ^ 3 + 3 * a ^ 2 + 2 * b + c + d));
267
268        let x = poly!(a ^ 4 - b ^ 4);
269        let y = poly!(a ^ 2 + b ^ 2);
270        assert_eq!(x.try_divide(&y), Some(poly!(a ^ 2 - b ^ 2)));
271
272        let x = poly!(a ^ 2 - 2 * a * b + b ^ 2);
273        let y = poly!(a - b);
274        assert_eq!(x.try_divide(&y), Some(poly!(a - b)));
275
276        let x = poly!(-4 * b);
277        let y = poly!(a);
278        assert_eq!(x.try_divide(&y), None);
279    }
280
281    #[test]
282    fn arithmetic_fuzz() {
283        use rand::prelude::*;
284
285        use super::Poly;
286
287        let mut rng = SmallRng::seed_from_u64(1);
288
289        fn create_random_poly(rng: &mut SmallRng, term_max: i32) -> Poly {
290            let mut p = Poly::zero();
291
292            for _ in 0..rng.gen_range(0..term_max + 1) {
293                let coef = rng.gen_range(-7..7);
294                let xpow = rng.gen_range(0..2);
295                let ypow = rng.gen_range(0..2);
296                let zpow = rng.gen_range(0..4);
297                let coef = Poly::constant(coef);
298                let xpow = Poly::var('x', xpow);
299                let ypow = Poly::var('y', ypow);
300                let zpow = Poly::var('z', zpow);
301
302                p = p.add(&coef.mul(&xpow).mul(&ypow).mul(&zpow));
303            }
304
305            p
306        }
307
308        for _ in 0..10_000 {
309            let dividend = create_random_poly(&mut rng, 10);
310            let n_divs = rng.gen_range(0..4);
311            let mut divisors: Vec<_> = std::iter::repeat_with(|| create_random_poly(&mut rng, 6))
312                .take(n_divs)
313                .collect();
314
315            let (quotients, rem) = dividend.compound_divide(&mut divisors);
316
317            println!("-------------------------");
318            println!("calculated {:?} / {:?}", dividend, divisors);
319            println!("got {:?} rem {:?}", quotients, rem);
320            println!("-------------------------");
321
322            let calculated_dividend = quotients
323                .into_iter()
324                .zip(divisors.clone())
325                .fold(Poly::zero(), |acc, (x, y)| acc.add(&mut x.mul(&y)))
326                .add(&mut rem.clone());
327
328            assert_eq!(calculated_dividend, dividend);
329        }
330    }
331
332    pub fn s_poly(p: &Poly, q: &Poly) -> Poly {
333        let p_lt = Poly {
334            terms: vec![p.leading_term()],
335        };
336        let q_lt = Poly {
337            terms: vec![q.leading_term()],
338        };
339
340        let lcm_ltp_ltq = Poly {
341            terms: vec![(
342                p_lt.terms[0].0 * q_lt.terms[0].0,
343                p_lt.terms[0].1.lcm(&q_lt.terms[0].1),
344            )],
345        };
346
347        let coef_p = lcm_ltp_ltq.try_divide(&p_lt).unwrap();
348        let coef_q = lcm_ltp_ltq.try_divide(&q_lt).unwrap();
349
350        let a = coef_p.mul(&p);
351        let b = coef_q.mul(&q);
352        let res = a.sub(&b);
353
354        if res.is_zero() {
355            res
356        } else {
357            res.divide_coefs(res.coef_gcd()).unwrap()
358        }
359    }
360
361    pub fn groebner_basis(sys: &[Poly]) -> Vec<Poly> {
362        let mut sys = sys.to_vec();
363
364        // buchberger
365
366        let mut combs = {
367            let mut combs = vec![];
368            for i in 0..sys.len() {
369                for j in 0..sys.len() {
370                    if i != j {
371                        combs.push((sys[i].clone(), sys[j].clone()));
372                    }
373                }
374            }
375
376            combs
377        };
378
379        while let Some((a, b)) = combs.pop() {
380            let s = s_poly(&a, &b);
381            let (_, rem) = s.compound_divide(&sys);
382
383            if !rem.is_zero() {
384                for member in &sys {
385                    combs.push((member.clone(), rem.clone()));
386                }
387                sys.push(rem);
388            }
389        }
390
391        // reduce
392
393        let mut keep = vec![];
394
395        for i in 0..sys.len() {
396            let mut divides_any = false;
397
398            for j in 0..sys.len() {
399                if i != j {
400                    let i_lt = (Coef::new(1), sys[i].leading_term().1);
401                    let j_lt = (Coef::new(1), sys[j].leading_term().1);
402                    if let Some((_, m)) = monomial_div(&i_lt, &j_lt) {
403                        if m.is_unit() {
404                            divides_any = i > j;
405                        } else {
406                            divides_any = true;
407                        }
408
409                        if divides_any {
410                            break;
411                        }
412                    }
413                }
414            }
415
416            if !divides_any {
417                keep.push(sys[i].clone());
418            }
419        }
420
421        keep.sort_by(|p, q| p.leading_term().1.cmp(&q.leading_term().1).reverse());
422        println!("sys: {:?}", keep);
423
424        let mut keep2 = vec![];
425
426        for (i, k) in keep.iter().enumerate() {
427            let all_except = keep
428                .iter()
429                .enumerate()
430                .filter_map(|(j, p)| if j != i { Some(p.clone()) } else { None })
431                .collect::<Vec<_>>();
432
433            let all_except_lcm = all_except
434                .iter()
435                .fold(Coef::new(1), |c, p| c.lcm(&p.coef_lcm()));
436
437            let (_, rem) = k.mul_scalar(all_except_lcm).compound_divide(&all_except);
438            if !rem.is_zero() {
439                keep2.push(rem.divide_coefs(rem.coef_gcd()).unwrap());
440            }
441        }
442
443        for p in keep2.iter_mut() {
444            if !p.leading_term().0.is_positive() {
445                *p = p.mul_scalar(-1);
446            }
447        }
448
449        keep2
450    }
451
452    #[test]
453    fn groebner_basis_validation() {
454        let sys = [
455            poly!(x + y ^ 2 + z),
456            poly!(x - y + 3 * z + 5),
457            poly!(x - 2 * y + 3),
458        ];
459
460        let gb = groebner_basis(&sys);
461
462        let correct = [
463            poly!(9 * z ^ 2 + 7 * z - 3),
464            poly!(x + 6 * z + 7),
465            poly!(y + 3 * z + 2),
466        ];
467
468        println!("{:?}", gb);
469        println!("{:?}", correct);
470
471        for (g, c) in gb.iter().zip(correct.iter()) {
472            assert_eq!(g, c, "{:?} = {:?}", g, c);
473        }
474
475        let sys = [
476            poly!(x ^ 2 * y + 1),
477            poly!(2 * x + y * z - 1),
478            poly!(x - y ^ 2 * z ^ 2 + 1),
479        ];
480
481        let gb = groebner_basis(&sys);
482
483        let correct = [poly!(4 * x - 5), poly!(25 * y + 16), poly!(32 * z - 75)];
484
485        println!("{:?}", gb);
486        println!("{:?}", correct);
487
488        for (g, c) in gb.iter().zip(correct.iter()) {
489            assert_eq!(g, c, "{:?} = {:?}", g, c);
490        }
491    }
492}