algebra/
steenrod_evaluator.rs

1use std::collections::BTreeMap;
2
3use anyhow::anyhow;
4use fp::{
5    prime::{Prime, ValidPrime},
6    vector::FpVector,
7};
8
9use crate::{
10    algebra::{AdemAlgebra, Algebra, MilnorAlgebra, adem_algebra::AdemBasisElement},
11    milnor_algebra::{MilnorBasisElement, PPartEntry},
12    steenrod_parser::*,
13};
14
15pub struct SteenrodEvaluator {
16    pub adem: AdemAlgebra,
17    pub milnor: MilnorAlgebra,
18}
19
20impl SteenrodEvaluator {
21    pub fn new(p: ValidPrime) -> Self {
22        Self {
23            adem: AdemAlgebra::new(p, false),
24            milnor: MilnorAlgebra::new(p, false),
25        }
26    }
27
28    pub fn milnor_to_adem(&self, result: &mut FpVector, coeff: u32, degree: i32, input: &FpVector) {
29        let p = self.prime();
30        for (i, v) in input.iter_nonzero() {
31            self.milnor_to_adem_on_basis(result, (coeff * v) % p, degree, i);
32        }
33    }
34
35    pub fn adem_to_milnor(&self, result: &mut FpVector, coeff: u32, degree: i32, input: &FpVector) {
36        let p = self.prime();
37        for (i, v) in input.iter_nonzero() {
38            self.adem_to_milnor_on_basis(result, (coeff * v) % p, degree, i);
39        }
40    }
41
42    pub fn evaluate_algebra_adem(&self, input: &str) -> anyhow::Result<(i32, FpVector)> {
43        self.evaluate_algebra_node(None, parse_algebra(input)?)
44    }
45
46    pub fn evaluate_algebra_milnor(&self, input: &str) -> anyhow::Result<(i32, FpVector)> {
47        let adem_result = self.evaluate_algebra_adem(input);
48        if let Ok((degree, adem_vector)) = adem_result {
49            let mut milnor_vector = FpVector::new(adem_vector.prime(), adem_vector.len());
50            self.adem_to_milnor(&mut milnor_vector, 1, degree, &adem_vector);
51            Ok((degree, milnor_vector))
52        } else {
53            adem_result
54        }
55    }
56
57    /// # Returns
58    /// This returns a [`BTreeMap`] so that we get deterministic outputs for testing purposes
59    pub fn evaluate_module_adem(
60        &self,
61        items: &str,
62    ) -> anyhow::Result<BTreeMap<String, (i32, FpVector)>> {
63        let mut result: BTreeMap<String, (i32, FpVector)> = BTreeMap::new();
64        if items.is_empty() {
65            return Ok(result);
66        }
67        for (op, g) in parse_module(items)? {
68            if let Some((deg, vec)) = result.get_mut(&g) {
69                let (_, adem_v) = self.evaluate_algebra_node(Some(*deg), op)?;
70                vec.add(&adem_v, 1);
71            } else {
72                let (deg, adem_v) = self.evaluate_algebra_node(None, op)?;
73                result.insert(g, (deg, adem_v));
74            }
75        }
76        Ok(result)
77    }
78
79    fn prime(&self) -> ValidPrime {
80        self.adem.prime()
81    }
82
83    fn compute_basis(&self, degree: i32) {
84        self.adem.compute_basis(degree);
85        self.milnor.compute_basis(degree);
86    }
87
88    fn dimension(&self, degree: i32) -> usize {
89        self.adem.dimension(degree)
90    }
91
92    fn evaluate_algebra_node(
93        &self,
94        mut output_degree: Option<i32>,
95        tree: AlgebraNode,
96    ) -> anyhow::Result<(i32, FpVector)> {
97        let p = self.prime();
98        match tree {
99            AlgebraNode::Sum(left, right) => {
100                let (degree, mut output_left) = self.evaluate_algebra_node(output_degree, *left)?;
101                let (_, output_right) = self.evaluate_algebra_node(Some(degree), *right)?;
102                output_left += &output_right;
103                Ok((degree, output_left))
104            }
105            AlgebraNode::Product(left, right) => {
106                let (degree_left, output_left) = self.evaluate_algebra_node(None, *left)?;
107                if let Some(degree) = output_degree {
108                    if degree < degree_left {
109                        return Err(anyhow!("Mismatched degree"));
110                    }
111                    output_degree = Some(degree - degree_left);
112                }
113                let (degree_right, output_right) =
114                    self.evaluate_algebra_node(output_degree, *right)?;
115                let degree = degree_left + degree_right;
116                self.compute_basis(degree);
117                let mut result = FpVector::new(p, self.adem.dimension(degree));
118                self.adem.multiply_element_by_element(
119                    result.as_slice_mut(),
120                    1,
121                    degree_left,
122                    output_left.as_slice(),
123                    degree_right,
124                    output_right.as_slice(),
125                );
126                Ok((degree, result))
127            }
128            AlgebraNode::BasisElt(basis_elt) => {
129                self.evaluate_basis_element(output_degree, basis_elt)
130            }
131            AlgebraNode::Scalar(x) => {
132                if let Some(degree) = output_degree
133                    && degree != 0
134                {
135                    return Err(anyhow!("Mismatched Degree"));
136                }
137                let mut result = FpVector::new(p, 1);
138                result.set_entry(0, x.rem_euclid(p.as_i32()) as u32);
139                Ok((0, result))
140            }
141        }
142    }
143
144    fn evaluate_basis_element(
145        &self,
146        output_degree: Option<i32>,
147        basis_elt: AlgebraBasisElt,
148    ) -> anyhow::Result<(i32, FpVector)> {
149        let p = self.prime();
150        let q = self.adem.q();
151        let (degree, result) = match basis_elt {
152            AlgebraBasisElt::AList(p_or_b_list) => self.evaluate_p_or_b_list(&p_or_b_list),
153            AlgebraBasisElt::PList(p_list) => {
154                let degree = std::iter::zip(crate::algebra::combinatorics::xi_degrees(p), &p_list)
155                    .map(|(&a, &b)| a * b as i32)
156                    .sum::<i32>()
157                    * q;
158                let elt = MilnorBasisElement {
159                    degree,
160                    p_part: p_list,
161                    q_part: 0,
162                };
163
164                self.compute_basis(degree);
165                let mut result = FpVector::new(p, self.dimension(degree));
166                self.milnor_to_adem_on_basis(
167                    &mut result,
168                    1,
169                    degree,
170                    self.milnor.basis_element_to_index(&elt),
171                );
172                (degree, result)
173            }
174            AlgebraBasisElt::P(x) => {
175                self.compute_basis(x as i32 * q);
176                let (degree, idx) = self.adem.beps_pn(0, x);
177                let mut result = FpVector::new(p, self.dimension(degree));
178                result.set_entry(idx, 1);
179                (degree, result)
180            }
181            AlgebraBasisElt::Q(x) => {
182                let tau_degrees = crate::algebra::combinatorics::tau_degrees(p);
183                let degree = tau_degrees[x as usize];
184                self.compute_basis(degree);
185                let mut result = FpVector::new(p, self.dimension(degree));
186                self.adem_q(&mut result, 1, x);
187                (degree, result)
188            }
189        };
190        if let Some(requested_degree) = output_degree
191            && degree != requested_degree
192        {
193            return Err(anyhow!("Mismatched degree"));
194        }
195        Ok((degree, result))
196    }
197
198    /// Translate from the adem basis to the milnor basis, adding `coeff` times the result to `result`.
199    /// This uses the fact that that $P^n = P(n)$ and $Q_1 = \beta$ and multiplies out the admissible
200    /// monomial.
201    fn adem_to_milnor_on_basis(&self, result: &mut FpVector, coeff: u32, degree: i32, idx: usize) {
202        let elt = self.adem.basis_element_from_index(degree, idx);
203        let p = self.prime();
204        let dim = self.dimension(elt.degree);
205        if dim == 1 {
206            result.set_entry(0, coeff);
207            return;
208        }
209        let mut tmp_vector_a = FpVector::new(p, 1);
210        let mut tmp_vector_b = FpVector::new(p, 0);
211
212        tmp_vector_a.set_entry(0, 1);
213
214        let mut bocksteins = elt.bocksteins;
215        let mut total_degree = 0;
216
217        for &sqn in &elt.ps {
218            let (deg, idx) = self.milnor.beps_pn(bocksteins & 1, sqn as PPartEntry);
219            bocksteins >>= 1;
220            self.compute_basis(total_degree + deg);
221
222            tmp_vector_b.set_scratch_vector_size(self.dimension(total_degree + deg));
223            self.milnor.multiply_element_by_basis_element(
224                tmp_vector_b.as_slice_mut(),
225                1,
226                total_degree,
227                tmp_vector_a.as_slice(),
228                deg,
229                idx,
230            );
231            total_degree += deg;
232            std::mem::swap(&mut tmp_vector_a, &mut tmp_vector_b);
233        }
234        if bocksteins & 1 == 0 {
235            result.add(&tmp_vector_a, coeff);
236        } else {
237            self.milnor.multiply_element_by_basis_element(
238                result.as_slice_mut(),
239                coeff,
240                total_degree,
241                tmp_vector_a.as_slice(),
242                1,
243                0,
244            );
245        }
246    }
247
248    // This is currently pretty inefficient... We should memoize results so that we don't repeatedly
249    // recompute the same inverse.
250    fn milnor_to_adem_on_basis(&self, result: &mut FpVector, coeff: u32, degree: i32, idx: usize) {
251        if self.milnor.generic() {
252            self.milnor_to_adem_on_basis_generic(result, coeff, degree, idx);
253        } else {
254            self.milnor_to_adem_on_basis_2(result, coeff, degree, idx);
255        }
256    }
257
258    fn milnor_to_adem_on_basis_2(
259        &self,
260        result: &mut FpVector,
261        coeff: u32,
262        degree: i32,
263        idx: usize,
264    ) {
265        let elt = self.milnor.basis_element_from_index(degree, idx);
266        let p = self.prime();
267        let dim = self.dimension(elt.degree);
268        if dim == 1 {
269            result.set_entry(0, coeff);
270            return;
271        }
272        let mut t: Vec<u32> = vec![0; elt.p_part.len()];
273        t[elt.p_part.len() - 1] = elt.p_part[elt.p_part.len() - 1];
274        for i in (0..elt.p_part.len() - 1).rev() {
275            t[i] = elt.p_part[i] + 2 * t[i + 1];
276        }
277        let t_idx = self.adem.basis_element_to_index(&AdemBasisElement {
278            degree,
279            bocksteins: 0,
280            ps: t,
281            p_or_sq: p != 2,
282        });
283        let mut tmp_vector_a = FpVector::new(p, dim);
284        self.adem_to_milnor_on_basis(&mut tmp_vector_a, 1, degree, t_idx);
285        assert!(tmp_vector_a.entry(idx) == 1);
286        tmp_vector_a.set_entry(idx, 0);
287        self.milnor_to_adem(result, coeff, degree, &tmp_vector_a);
288        result.add_basis_element(t_idx, coeff);
289    }
290
291    fn milnor_to_adem_on_basis_generic(
292        &self,
293        result: &mut FpVector,
294        coeff: u32,
295        degree: i32,
296        idx: usize,
297    ) {
298        let elt = self.milnor.basis_element_from_index(degree, idx);
299        let p = self.prime();
300        let dim = self.dimension(elt.degree);
301        if dim == 1 {
302            result.set_entry(0, coeff);
303            return;
304        }
305        let t_len = std::cmp::max(
306            elt.p_part.len(),
307            (31u32.saturating_sub(elt.q_part.leading_zeros())) as usize,
308        );
309        let mut t = vec![0; t_len];
310        let last_p_part = if t_len <= elt.p_part.len() {
311            elt.p_part[t_len - 1]
312        } else {
313            0
314        };
315        t[t_len - 1] = last_p_part + ((elt.q_part >> (t_len)) & 1);
316        for i in (0..t_len - 1).rev() {
317            let p_part = if i < elt.p_part.len() {
318                elt.p_part[i]
319            } else {
320                0
321            };
322            t[i] = p_part + ((elt.q_part >> (i + 1)) & 1) + p * t[i + 1];
323        }
324        let t_idx = self.adem.basis_element_to_index(&AdemBasisElement {
325            degree,
326            bocksteins: elt.q_part,
327            ps: t,
328            p_or_sq: p != 2,
329        });
330        let mut tmp_vector_a = FpVector::new(p, dim);
331        self.adem_to_milnor_on_basis(&mut tmp_vector_a, 1, degree, t_idx);
332        assert!(tmp_vector_a.entry(idx) == 1);
333        tmp_vector_a.set_entry(idx, 0);
334        tmp_vector_a.scale(p - 1);
335        self.milnor_to_adem(result, coeff, degree, &tmp_vector_a);
336        result.add_basis_element(t_idx, coeff);
337    }
338
339    /// Express $Q_{qi}$ in the adem basis.
340    fn adem_q(&self, result: &mut FpVector, coeff: u32, qi: u32) {
341        let p = self.prime();
342        let degree = crate::algebra::combinatorics::tau_degrees(p)[qi as usize];
343        let mbe = if self.adem.generic() {
344            MilnorBasisElement {
345                degree,
346                q_part: 1 << qi,
347                p_part: vec![],
348            }
349        } else {
350            let mut p_part = vec![0; qi as usize + 1];
351            p_part[qi as usize] = 1;
352            MilnorBasisElement {
353                degree,
354                q_part: 0,
355                p_part,
356            }
357        };
358        let idx = self.milnor.basis_element_to_index(&mbe);
359        self.milnor_to_adem_on_basis(result, coeff, degree, idx);
360    }
361
362    fn evaluate_p_or_b_list(&self, list: &[BocksteinOrSq]) -> (i32, FpVector) {
363        let p = self.prime();
364        let q = self.adem.q();
365
366        let mut total_degree = 0;
367
368        let mut tmp_vector_a = FpVector::new(p, 1);
369        let mut tmp_vector_b = FpVector::new(p, 0);
370
371        tmp_vector_a.set_entry(0, 1);
372
373        for item in list {
374            let cur_elt = item.to_adem_basis_elt(q);
375
376            self.compute_basis(total_degree + cur_elt.degree);
377            tmp_vector_b.set_scratch_vector_size(self.dimension(total_degree + cur_elt.degree));
378            self.adem.multiply_element_by_basis_element(
379                tmp_vector_b.as_slice_mut(),
380                1,
381                total_degree,
382                tmp_vector_a.as_slice(),
383                cur_elt.degree,
384                self.adem.basis_element_to_index(&cur_elt),
385            );
386            total_degree += cur_elt.degree;
387            std::mem::swap(&mut tmp_vector_a, &mut tmp_vector_b);
388        }
389        (total_degree, tmp_vector_a)
390    }
391}
392
393#[cfg(test)]
394mod tests {
395    use expect_test::{Expect, expect};
396    use rstest::rstest;
397
398    use super::*;
399
400    #[test]
401    fn test_evaluate_2() {
402        let ev = SteenrodEvaluator::new(ValidPrime::new(2));
403
404        let check = |input, adem_output: Expect, milnor_output: Expect| {
405            let (degree, result) = ev.evaluate_algebra_adem(input).unwrap();
406            adem_output.assert_eq(&ev.adem.element_to_string(degree, result.as_slice()));
407
408            let (degree, result) = ev.evaluate_algebra_milnor(input).unwrap();
409            milnor_output.assert_eq(&ev.milnor.element_to_string(degree, result.as_slice()));
410        };
411
412        check(
413            "Sq2 * Sq2",
414            expect![[r#"Sq3 Sq1"#]],
415            expect![[r#"P(1, 1)"#]],
416        );
417        check("A(2 2)", expect![[r#"Sq3 Sq1"#]], expect![[r#"P(1, 1)"#]]);
418        check(
419            "Sq2 * Sq2 * Sq2 + Sq3 * Sq3",
420            expect![[r#"0"#]],
421            expect![[r#"0"#]],
422        );
423        check(
424            "Sq2 * (Sq2 * Sq2 + Sq4)",
425            expect![[r#"Sq6"#]],
426            expect![[r#"P(6)"#]],
427        );
428        check(
429            "Sq7 + Q2",
430            expect![[r#"Sq5 Sq2 + Sq6 Sq1 + Sq4 Sq2 Sq1"#]],
431            expect![[r#"P(7) + P(0, 0, 1)"#]],
432        );
433        check(
434            "(Q2 + Sq7) * Q1",
435            expect![[r#"Sq6 Sq3 Sq1"#]],
436            expect![[r#"P(7, 1) + P(3, 0, 1) + P(0, 1, 1)"#]],
437        );
438    }
439
440    #[test]
441    fn test_evaluate_3() {
442        let ev = SteenrodEvaluator::new(ValidPrime::new(3));
443
444        let check = |input, adem_output: Expect, milnor_output: Expect| {
445            let (degree, result) = ev.evaluate_algebra_adem(input).unwrap();
446            adem_output.assert_eq(&ev.adem.element_to_string(degree, result.as_slice()));
447
448            let (degree, result) = ev.evaluate_algebra_milnor(input).unwrap();
449            milnor_output.assert_eq(&ev.milnor.element_to_string(degree, result.as_slice()));
450        };
451
452        check("P1 * P1", expect![[r#"2 * P2"#]], expect![[r#"2 * P(2)"#]]);
453        check("A(1 1)", expect![[r#"2 * P2"#]], expect![[r#"2 * P(2)"#]]);
454        check(
455            "A(1 b 1)",
456            expect![[r#"b P2 + P2 b"#]],
457            expect![[r#"2 * Q_0 P(2) + Q_1 P(1)"#]],
458        );
459        check(
460            "A(4 2)",
461            expect![[r#"2 * P5 P1"#]],
462            expect![[r#"2 * P(2, 1)"#]],
463        );
464        check("A(5 2)", expect![[r#"0"#]], expect![[r#"0"#]]);
465        check(
466            "A(6 2)",
467            expect![[r#"P6 P2"#]],
468            expect![[r#"P(8) + P(4, 1) + P(0, 2)"#]],
469        );
470    }
471
472    #[rstest(p, max_degree, case(2, 32), case(3, 60))]
473    #[trace]
474    fn test_cob_adem_to_milnor(p: u32, max_degree: i32) {
475        let p = ValidPrime::new(p);
476        let ev = SteenrodEvaluator::new(p);
477        ev.compute_basis(max_degree);
478
479        for degree in 0..max_degree {
480            println!("degree : {degree}");
481            let dim = ev.dimension(degree);
482            let mut milnor_result = FpVector::new(p, dim);
483            let mut adem_result = FpVector::new(p, dim);
484            for i in 0..dim {
485                // println!("i : {}", i);
486                ev.milnor_to_adem_on_basis(&mut adem_result, 1, degree, i);
487                ev.adem_to_milnor(&mut milnor_result, 1, degree, &adem_result);
488                assert!(
489                    milnor_result.entry(i) == 1,
490                    "{} ==> {} ==> {}",
491                    ev.milnor.basis_element_to_string(degree, i),
492                    ev.adem.element_to_string(degree, adem_result.as_slice()),
493                    ev.milnor
494                        .element_to_string(degree, milnor_result.as_slice())
495                );
496                milnor_result.set_entry(i, 0);
497                assert!(
498                    milnor_result.is_zero(),
499                    "{} ==> {} ==> {}",
500                    ev.milnor.basis_element_to_string(degree, i),
501                    ev.adem.element_to_string(degree, adem_result.as_slice()),
502                    ev.milnor
503                        .element_to_string(degree, milnor_result.as_slice())
504                );
505                adem_result.set_to_zero();
506                milnor_result.set_to_zero();
507            }
508        }
509    }
510}