1use super::simd::{
5 add_u32_lanes, add_u64_lanes, div_u32_lanes, div_u64_lanes, fma_u32_lanes, fma_u64_lanes,
6 mul_u32_lanes, mul_u64_lanes, sqrt_u32_lanes, sqrt_u64_lanes, sub_u32_lanes, sub_u64_lanes,
7};
8use super::{Ieee32, Ieee64};
9use crate::defs::RoundingMode;
10use crate::Consts;
11use crate::Error;
12use crate::ExactNum;
13use alloc::vec::Vec;
14
15const SVD_ITER_MAX: u32 = 64;
18const EIGEN_ITER_MAX: u32 = 64;
21const FFT_MAX_POINTS: usize = 4096;
23const SVD_CONV_GUARD_BITS: i32 = 4;
26
27trait LaneBits: Copy {
28 fn add_lanes(a: &[Self], b: &[Self]) -> Vec<Self>;
29 fn mul_lanes(a: &[Self], b: &[Self]) -> Vec<Self>;
30 fn sub_lanes(a: &[Self], b: &[Self]) -> Vec<Self>;
31 fn div_lanes(a: &[Self], b: &[Self]) -> Vec<Self>;
32 fn sqrt_lanes(a: &[Self]) -> Vec<Self>;
33 fn fma_lanes(a: &[Self], b: &[Self], c: &[Self]) -> Vec<Self>;
34}
35
36impl LaneBits for u32 {
37 fn add_lanes(a: &[Self], b: &[Self]) -> Vec<Self> {
38 add_u32_lanes(a, b)
39 }
40 fn mul_lanes(a: &[Self], b: &[Self]) -> Vec<Self> {
41 mul_u32_lanes(a, b)
42 }
43 fn sub_lanes(a: &[Self], b: &[Self]) -> Vec<Self> {
44 sub_u32_lanes(a, b)
45 }
46 fn div_lanes(a: &[Self], b: &[Self]) -> Vec<Self> {
47 div_u32_lanes(a, b)
48 }
49 fn sqrt_lanes(a: &[Self]) -> Vec<Self> {
50 sqrt_u32_lanes(a)
51 }
52 fn fma_lanes(a: &[Self], b: &[Self], c: &[Self]) -> Vec<Self> {
53 fma_u32_lanes(a, b, c)
54 }
55}
56
57impl LaneBits for u64 {
58 fn add_lanes(a: &[Self], b: &[Self]) -> Vec<Self> {
59 add_u64_lanes(a, b)
60 }
61 fn mul_lanes(a: &[Self], b: &[Self]) -> Vec<Self> {
62 mul_u64_lanes(a, b)
63 }
64 fn sub_lanes(a: &[Self], b: &[Self]) -> Vec<Self> {
65 sub_u64_lanes(a, b)
66 }
67 fn div_lanes(a: &[Self], b: &[Self]) -> Vec<Self> {
68 div_u64_lanes(a, b)
69 }
70 fn sqrt_lanes(a: &[Self]) -> Vec<Self> {
71 sqrt_u64_lanes(a)
72 }
73 fn fma_lanes(a: &[Self], b: &[Self], c: &[Self]) -> Vec<Self> {
74 fma_u64_lanes(a, b, c)
75 }
76}
77
78#[derive(Clone, Debug)]
80pub struct Ieee32Array {
81 bits: Vec<u32>,
82 rows: usize,
83 cols: usize,
84}
85
86#[derive(Clone, Debug)]
88pub struct Ieee64Array {
89 bits: Vec<u64>,
90 rows: usize,
91 cols: usize,
92}
93
94#[derive(Clone, Debug)]
96pub struct ExactNumArray {
97 p: usize,
98 vals: Vec<ExactNum>,
99 rows: usize,
100 cols: usize,
101}
102
103macro_rules! impl_ieee_array {
104 ($arr:ident, $scalar:ident, $bits:ty) => {
105 impl $arr {
106 pub fn new() -> Self {
108 Self {
109 bits: Vec::new(),
110 rows: 0,
111 cols: 0,
112 }
113 }
114
115 pub fn filled(n: usize, fill: $scalar) -> Self {
117 Self {
118 bits: alloc::vec![fill.to_bits(); n],
119 rows: 1,
120 cols: n,
121 }
122 }
123
124 pub fn filled_2d(rows: usize, cols: usize, fill: $scalar) -> Option<Self> {
126 let n = rows.checked_mul(cols)?;
127 Some(Self {
128 bits: alloc::vec![fill.to_bits(); n],
129 rows,
130 cols,
131 })
132 }
133
134 pub fn from_bits(bits: &[$bits]) -> Self {
136 Self {
137 bits: bits.to_vec(),
138 rows: 1,
139 cols: bits.len(),
140 }
141 }
142
143 pub fn from_shape(rows: usize, cols: usize, vals: &[$scalar]) -> Option<Self> {
145 let n = rows.checked_mul(cols)?;
146 if n != vals.len() {
147 return None;
148 }
149 Some(Self {
150 bits: vals.iter().map(|v| v.to_bits()).collect(),
151 rows,
152 cols,
153 })
154 }
155
156 pub fn from_values(vals: &[$scalar]) -> Self {
158 Self {
159 bits: vals.iter().map(|v| v.to_bits()).collect(),
160 rows: 1,
161 cols: vals.len(),
162 }
163 }
164
165 pub fn shape(&self) -> (usize, usize) {
167 (self.rows, self.cols)
168 }
169
170 pub fn reshape(&self, rows: usize, cols: usize) -> Option<Self> {
172 let n = rows.checked_mul(cols)?;
173 if n != self.bits.len() {
174 return None;
175 }
176 Some(Self {
177 bits: self.bits.clone(),
178 rows,
179 cols,
180 })
181 }
182
183 pub fn len(&self) -> usize {
185 self.bits.len()
186 }
187
188 pub fn is_empty(&self) -> bool {
190 self.bits.is_empty()
191 }
192
193 pub fn as_bits(&self) -> &[$bits] {
195 &self.bits
196 }
197
198 pub fn get(&self, i: usize) -> Option<$scalar> {
200 self.bits.get(i).copied().map($scalar::from_bits)
201 }
202
203 pub fn get2(&self, i: usize, j: usize) -> Option<$scalar> {
205 if i >= self.rows || j >= self.cols {
206 return None;
207 }
208 self.get(i * self.cols + j)
209 }
210
211 pub fn add(&self, rhs: &Self) -> Option<Self> {
214 if self.rows != rhs.rows || self.cols != rhs.cols {
215 return None;
216 }
217 Some(Self {
218 bits: <$bits>::add_lanes(&self.bits, &rhs.bits),
219 rows: self.rows,
220 cols: self.cols,
221 })
222 }
223
224 pub fn add_scalar(&self, s: $scalar) -> Self {
226 self.map(|x| x.add(s))
227 }
228
229 pub fn sub(&self, rhs: &Self) -> Option<Self> {
232 if self.rows != rhs.rows || self.cols != rhs.cols {
233 return None;
234 }
235 Some(Self {
236 bits: <$bits>::sub_lanes(&self.bits, &rhs.bits),
237 rows: self.rows,
238 cols: self.cols,
239 })
240 }
241
242 pub fn mul(&self, rhs: &Self) -> Option<Self> {
245 if self.rows != rhs.rows || self.cols != rhs.cols {
246 return None;
247 }
248 Some(Self {
249 bits: <$bits>::mul_lanes(&self.bits, &rhs.bits),
250 rows: self.rows,
251 cols: self.cols,
252 })
253 }
254
255 pub fn mul_scalar(&self, s: $scalar) -> Self {
257 self.map(|x| x.mul(s))
258 }
259
260 pub fn div(&self, rhs: &Self) -> Option<Self> {
263 if self.rows != rhs.rows || self.cols != rhs.cols {
264 return None;
265 }
266 Some(Self {
267 bits: <$bits>::div_lanes(&self.bits, &rhs.bits),
268 rows: self.rows,
269 cols: self.cols,
270 })
271 }
272
273 pub fn sqrt(&self) -> Self {
275 Self {
276 bits: <$bits>::sqrt_lanes(&self.bits),
277 rows: self.rows,
278 cols: self.cols,
279 }
280 }
281
282 pub fn fma(&self, b: &Self, c: &Self) -> Option<Self> {
285 if self.rows != b.rows
286 || self.cols != b.cols
287 || self.rows != c.rows
288 || self.cols != c.cols
289 {
290 return None;
291 }
292 Some(Self {
293 bits: <$bits>::fma_lanes(&self.bits, &b.bits, &c.bits),
294 rows: self.rows,
295 cols: self.cols,
296 })
297 }
298
299 pub fn sum(&self) -> $scalar {
301 let mut acc = $scalar::ZERO;
302 for &b in &self.bits {
303 acc = acc.add($scalar::from_bits(b));
304 }
305 acc
306 }
307
308 pub fn dot(&self, rhs: &Self) -> Option<$scalar> {
310 if self.len() != rhs.len() {
311 return None;
312 }
313 let mut acc = $scalar::ZERO;
314 for (a, b) in self.bits.iter().zip(rhs.bits.iter()) {
315 acc = acc.add($scalar::from_bits(*a).mul($scalar::from_bits(*b)));
316 }
317 Some(acc)
318 }
319
320 pub fn matmul(&self, rhs: &Self) -> Option<Self> {
322 if self.cols != rhs.rows {
323 return None;
324 }
325 let m = self.rows;
326 let k = self.cols;
327 let n = rhs.cols;
328 let mut bits = alloc::vec![$scalar::ZERO.to_bits(); m.checked_mul(n)?];
329 for i in 0..m {
330 for j in 0..n {
331 let mut acc = $scalar::ZERO;
332 for t in 0..k {
333 let a = $scalar::from_bits(self.bits[i * k + t]);
334 let b = $scalar::from_bits(rhs.bits[t * n + j]);
335 acc = acc.add(a.mul(b));
336 }
337 bits[i * n + j] = acc.to_bits();
338 }
339 }
340 Some(Self {
341 bits,
342 rows: m,
343 cols: n,
344 })
345 }
346
347 fn map(&self, op: impl Fn($scalar) -> $scalar) -> Self {
348 Self {
349 bits: self
350 .bits
351 .iter()
352 .map(|&b| op($scalar::from_bits(b)).to_bits())
353 .collect(),
354 rows: self.rows,
355 cols: self.cols,
356 }
357 }
358
359 pub fn map_exact<F>(&self, p: usize, mut op: F) -> Self
361 where
362 F: FnMut(&ExactNum) -> ExactNum,
363 {
364 Self {
365 bits: self
366 .bits
367 .iter()
368 .map(|&b| {
369 let x = $scalar::from_bits(b).to_exact(p);
370 $scalar::from_exact(&op(&x)).to_bits()
371 })
372 .collect(),
373 rows: self.rows,
374 cols: self.cols,
375 }
376 }
377 }
378
379 impl Default for $arr {
380 fn default() -> Self {
381 Self::new()
382 }
383 }
384 };
385}
386
387impl_ieee_array!(Ieee32Array, Ieee32, u32);
388impl_ieee_array!(Ieee64Array, Ieee64, u64);
389
390macro_rules! ieee_unary_exact {
391 ($p:expr, $($name:ident),+ $(,)?) => {
392 $(
393 #[doc = concat!("Elementwise `", stringify!($name), "` via `ExactNum`.")]
394 pub fn $name(&self, cc: &mut Consts) -> Self {
395 self.map_exact($p, |x| x.$name($p, RoundingMode::ToEven, cc))
396 }
397 )+
398 };
399}
400
401macro_rules! ieee_array_specials {
402 ($arr:ident, $scalar:ident, $p:expr) => {
403 impl $arr {
404 ieee_unary_exact!(
405 $p, exp, exp2, exp10, expm1, ln, log2, log10, log1p, sin, cos, tan, asin, acos,
406 atan, sinh, cosh, tanh, asinh, acosh, atanh, erf, erfc, gamma, ln_gamma, digamma,
407 ei, si, ci, li, fresnel_s, fresnel_c, ai, bi, elliptic_k, rem_pi,
408 );
409
410 pub fn elliptic_e_complete(&self, cc: &mut Consts) -> Self {
412 self.map_exact($p, |x| x.elliptic_e_complete($p, RoundingMode::ToEven, cc))
413 }
414
415 pub fn cbrt(&self) -> Self {
417 self.map_exact($p, |x| x.cbrt($p, RoundingMode::ToEven))
418 }
419
420 pub fn atan2(&self, x: $scalar, cc: &mut Consts) -> Self {
422 let xe = x.to_exact($p);
423 self.map_exact($p, |y| y.atan2(&xe, $p, RoundingMode::ToEven, cc))
424 }
425
426 pub fn hypot(&self, other: $scalar) -> Self {
428 let oe = other.to_exact($p);
429 self.map_exact($p, |x| x.hypot(&oe, $p, RoundingMode::ToEven))
430 }
431
432 pub fn pow(&self, n: $scalar, cc: &mut Consts) -> Self {
434 let ne = n.to_exact($p);
435 self.map_exact($p, |x| x.pow(&ne, $p, RoundingMode::ToEven, cc))
436 }
437
438 pub fn log(&self, base: $scalar, cc: &mut Consts) -> Self {
440 let be = base.to_exact($p);
441 self.map_exact($p, |x| x.log(&be, $p, RoundingMode::ToEven, cc))
442 }
443
444 pub fn gammainc(&self, x: $scalar, cc: &mut Consts) -> Self {
446 let xe = x.to_exact($p);
447 self.map_exact($p, |s| s.gammainc(&xe, $p, RoundingMode::ToEven, cc))
448 }
449
450 pub fn gammainc_upper(&self, x: $scalar, cc: &mut Consts) -> Self {
452 let xe = x.to_exact($p);
453 self.map_exact($p, |s| s.gammainc_upper(&xe, $p, RoundingMode::ToEven, cc))
454 }
455
456 pub fn bessel_j(&self, n: usize, cc: &mut Consts) -> Self {
458 self.map_exact($p, |x| x.bessel_j(n, $p, RoundingMode::ToEven, cc))
459 }
460
461 pub fn bessel_j_nu(&self, nu: $scalar, cc: &mut Consts) -> Self {
463 let n = nu.to_exact($p);
464 self.map_exact($p, |x| x.bessel_j_nu(&n, $p, RoundingMode::ToEven, cc))
465 }
466
467 pub fn bessel_y(&self, nu: $scalar, cc: &mut Consts) -> Self {
469 let n = nu.to_exact($p);
470 self.map_exact($p, |x| x.bessel_y(&n, $p, RoundingMode::ToEven, cc))
471 }
472
473 pub fn bessel_i(&self, nu: $scalar, cc: &mut Consts) -> Self {
475 let n = nu.to_exact($p);
476 self.map_exact($p, |x| x.bessel_i(&n, $p, RoundingMode::ToEven, cc))
477 }
478
479 pub fn bessel_k(&self, nu: $scalar, cc: &mut Consts) -> Self {
481 let n = nu.to_exact($p);
482 self.map_exact($p, |x| x.bessel_k(&n, $p, RoundingMode::ToEven, cc))
483 }
484
485 pub fn legendre_p(&self, n: u32) -> Self {
487 self.map_exact($p, |x| x.legendre_p(n, $p, RoundingMode::ToEven))
488 }
489
490 pub fn assoc_legendre_p(&self, n: u32, m: i32) -> Self {
492 self.map_exact($p, |x| x.assoc_legendre_p(n, m, $p, RoundingMode::ToEven))
493 }
494
495 pub fn elliptic_f(&self, m: $scalar, cc: &mut Consts) -> Self {
497 let me = m.to_exact($p);
498 self.map_exact($p, |x| x.elliptic_f(&me, $p, RoundingMode::ToEven, cc))
499 }
500
501 pub fn elliptic_e(&self, m: $scalar, cc: &mut Consts) -> Self {
503 let me = m.to_exact($p);
504 self.map_exact($p, |x| x.elliptic_e(&me, $p, RoundingMode::ToEven, cc))
505 }
506
507 pub fn jacobi_sn(&self, m: $scalar, cc: &mut Consts) -> Self {
509 let me = m.to_exact($p);
510 self.map_exact($p, |x| x.jacobi_sn(&me, $p, RoundingMode::ToEven, cc))
511 }
512
513 pub fn jacobi_cn(&self, m: $scalar, cc: &mut Consts) -> Self {
515 let me = m.to_exact($p);
516 self.map_exact($p, |x| x.jacobi_cn(&me, $p, RoundingMode::ToEven, cc))
517 }
518
519 pub fn jacobi_dn(&self, m: $scalar, cc: &mut Consts) -> Self {
521 let me = m.to_exact($p);
522 self.map_exact($p, |x| x.jacobi_dn(&me, $p, RoundingMode::ToEven, cc))
523 }
524
525 pub fn jacobi_am(&self, m: $scalar, cc: &mut Consts) -> Self {
527 let me = m.to_exact($p);
528 self.map_exact($p, |x| x.jacobi_am(&me, $p, RoundingMode::ToEven, cc))
529 }
530
531 pub fn jacobi_cd(&self, m: $scalar, cc: &mut Consts) -> Self {
533 let me = m.to_exact($p);
534 self.map_exact($p, |x| x.jacobi_cd(&me, $p, RoundingMode::ToEven, cc))
535 }
536
537 pub fn jacobi_ns(&self, m: $scalar, cc: &mut Consts) -> Self {
539 let me = m.to_exact($p);
540 self.map_exact($p, |x| x.jacobi_ns(&me, $p, RoundingMode::ToEven, cc))
541 }
542
543 pub fn jacobi_nc(&self, m: $scalar, cc: &mut Consts) -> Self {
545 let me = m.to_exact($p);
546 self.map_exact($p, |x| x.jacobi_nc(&me, $p, RoundingMode::ToEven, cc))
547 }
548
549 pub fn jacobi_nd(&self, m: $scalar, cc: &mut Consts) -> Self {
551 let me = m.to_exact($p);
552 self.map_exact($p, |x| x.jacobi_nd(&me, $p, RoundingMode::ToEven, cc))
553 }
554
555 pub fn jacobi_sc(&self, m: $scalar, cc: &mut Consts) -> Self {
557 let me = m.to_exact($p);
558 self.map_exact($p, |x| x.jacobi_sc(&me, $p, RoundingMode::ToEven, cc))
559 }
560
561 pub fn jacobi_sd(&self, m: $scalar, cc: &mut Consts) -> Self {
563 let me = m.to_exact($p);
564 self.map_exact($p, |x| x.jacobi_sd(&me, $p, RoundingMode::ToEven, cc))
565 }
566
567 pub fn jacobi_cs(&self, m: $scalar, cc: &mut Consts) -> Self {
569 let me = m.to_exact($p);
570 self.map_exact($p, |x| x.jacobi_cs(&me, $p, RoundingMode::ToEven, cc))
571 }
572
573 pub fn jacobi_ds(&self, m: $scalar, cc: &mut Consts) -> Self {
575 let me = m.to_exact($p);
576 self.map_exact($p, |x| x.jacobi_ds(&me, $p, RoundingMode::ToEven, cc))
577 }
578
579 pub fn jacobi_dc(&self, m: $scalar, cc: &mut Consts) -> Self {
581 let me = m.to_exact($p);
582 self.map_exact($p, |x| x.jacobi_dc(&me, $p, RoundingMode::ToEven, cc))
583 }
584
585 pub fn elliptic_pi_complete(&self, m: $scalar, cc: &mut Consts) -> Self {
587 let me = m.to_exact($p);
588 self.map_exact($p, |n| {
589 n.elliptic_pi_complete(&me, $p, RoundingMode::ToEven, cc)
590 })
591 }
592
593 pub fn elliptic_pi(&self, x: $scalar, m: $scalar, cc: &mut Consts) -> Self {
595 let xe = x.to_exact($p);
596 let me = m.to_exact($p);
597 self.map_exact($p, |n| {
598 n.elliptic_pi(&xe, &me, $p, RoundingMode::ToEven, cc)
599 })
600 }
601
602 pub fn hypergeom_2f1(
604 &self,
605 b: $scalar,
606 c: $scalar,
607 z: $scalar,
608 cc: &mut Consts,
609 ) -> Self {
610 let be = b.to_exact($p);
611 let ce = c.to_exact($p);
612 let ze = z.to_exact($p);
613 self.map_exact($p, |a| {
614 a.hypergeom_2f1(&be, &ce, &ze, $p, RoundingMode::ToEven, cc)
615 })
616 }
617
618 pub fn betainc(&self, b: $scalar, x: $scalar, cc: &mut Consts) -> Self {
620 let be = b.to_exact($p);
621 let xe = x.to_exact($p);
622 self.map_exact($p, |a| a.betainc(&be, &xe, $p, RoundingMode::ToEven, cc))
623 }
624 }
625 };
626}
627
628ieee_array_specials!(Ieee32Array, Ieee32, 64);
629ieee_array_specials!(Ieee64Array, Ieee64, 128);
630
631impl Ieee64Array {
632 pub(crate) fn from_parts(rows: usize, cols: usize, bits: Vec<u64>) -> Result<Self, Error> {
633 let n = rows.checked_mul(cols).ok_or(Error::InvalidArgument)?;
634 if n != bits.len() {
635 return Err(Error::InvalidArgument);
636 }
637 Ok(Self { bits, rows, cols })
638 }
639}
640
641macro_rules! exact_arr_p_rm_cc {
642 ($($name:ident),+ $(,)?) => {
643 $(
644 #[doc = concat!("Elementwise [`ExactNum::", stringify!($name), "`]. `cc` is the constants cache, not a global.")]
645 pub fn $name(&self, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
646 self.map_at(p, |x| x.$name(p, rm, cc))
647 }
648 )+
649 };
650}
651
652impl ExactNumArray {
653 pub fn new(p: usize) -> Self {
655 Self {
656 p,
657 vals: Vec::new(),
658 rows: 0,
659 cols: 0,
660 }
661 }
662
663 pub fn filled(p: usize, n: usize, fill: &ExactNum) -> Self {
665 let mut v = fill.clone();
666 let _ = v.set_precision(p, RoundingMode::ToEven);
667 Self {
668 p,
669 vals: alloc::vec![v; n],
670 rows: 1,
671 cols: n,
672 }
673 }
674
675 pub fn filled_2d(p: usize, rows: usize, cols: usize, fill: &ExactNum) -> Option<Self> {
677 let n = rows.checked_mul(cols)?;
678 let mut v = fill.clone();
679 let _ = v.set_precision(p, RoundingMode::ToEven);
680 Some(Self {
681 p,
682 vals: alloc::vec![v; n],
683 rows,
684 cols,
685 })
686 }
687
688 pub fn from_values(p: usize, vals: &[ExactNum]) -> Self {
690 Self {
691 p,
692 vals: vals
693 .iter()
694 .map(|x| {
695 let mut y = x.clone();
696 let _ = y.set_precision(p, RoundingMode::ToEven);
697 y
698 })
699 .collect(),
700 rows: 1,
701 cols: vals.len(),
702 }
703 }
704
705 pub fn from_shape(p: usize, rows: usize, cols: usize, vals: &[ExactNum]) -> Option<Self> {
707 let n = rows.checked_mul(cols)?;
708 if n != vals.len() {
709 return None;
710 }
711 Some(Self {
712 p,
713 vals: vals
714 .iter()
715 .map(|x| {
716 let mut y = x.clone();
717 let _ = y.set_precision(p, RoundingMode::ToEven);
718 y
719 })
720 .collect(),
721 rows,
722 cols,
723 })
724 }
725
726 #[cfg(any(test, feature = "random"))]
728 pub fn random_fill(
729 shape: (usize, usize),
730 dist: &crate::RandomDist,
731 p: usize,
732 rm: RoundingMode,
733 cc: &mut Consts,
734 ) -> Option<Self> {
735 let (rows, cols) = shape;
736 let n = rows.checked_mul(cols)?;
737 let mut vals = Vec::with_capacity(n);
738 for _ in 0..n {
739 let v = match dist {
740 crate::RandomDist::Uniform(a, b) => ExactNum::random_uniform(a, b, p, rm),
741 crate::RandomDist::Normal(mu, sigma) => {
742 ExactNum::random_gaussian(mu, sigma, p, rm, cc)
743 }
744 crate::RandomDist::Exponential(lambda) => {
745 ExactNum::random_exponential(lambda, p, rm, cc)
746 }
747 };
748 vals.push(v);
749 }
750 Some(Self {
751 p,
752 vals,
753 rows,
754 cols,
755 })
756 }
757
758 pub(crate) fn from_parts(
759 p: usize,
760 rows: usize,
761 cols: usize,
762 vals: Vec<ExactNum>,
763 ) -> Result<Self, Error> {
764 let n = rows.checked_mul(cols).ok_or(Error::InvalidArgument)?;
765 if n != vals.len() {
766 return Err(Error::InvalidArgument);
767 }
768 Ok(Self {
769 p,
770 vals,
771 rows,
772 cols,
773 })
774 }
775
776 pub fn precision(&self) -> usize {
778 self.p
779 }
780
781 pub fn shape(&self) -> (usize, usize) {
783 (self.rows, self.cols)
784 }
785
786 pub fn reshape(&self, rows: usize, cols: usize) -> Option<Self> {
788 let n = rows.checked_mul(cols)?;
789 if n != self.vals.len() {
790 return None;
791 }
792 Some(Self {
793 p: self.p,
794 vals: self.vals.clone(),
795 rows,
796 cols,
797 })
798 }
799
800 pub fn len(&self) -> usize {
802 self.vals.len()
803 }
804
805 pub fn is_empty(&self) -> bool {
807 self.vals.is_empty()
808 }
809
810 pub fn get(&self, i: usize) -> Option<&ExactNum> {
812 self.vals.get(i)
813 }
814
815 pub fn get2(&self, i: usize, j: usize) -> Option<&ExactNum> {
817 if i >= self.rows || j >= self.cols {
818 return None;
819 }
820 self.get(i * self.cols + j)
821 }
822
823 pub fn as_slice(&self) -> &[ExactNum] {
825 &self.vals
826 }
827
828 pub fn add(&self, rhs: &Self) -> Option<Self> {
830 self.zip(rhs, |a, b| a.add(b, self.p, RoundingMode::ToEven))
831 }
832
833 pub fn add_scalar(&self, s: &ExactNum) -> Self {
835 Self {
836 p: self.p,
837 vals: self
838 .vals
839 .iter()
840 .map(|x| x.add(s, self.p, RoundingMode::ToEven))
841 .collect(),
842 rows: self.rows,
843 cols: self.cols,
844 }
845 }
846
847 pub fn sub(&self, rhs: &Self) -> Option<Self> {
849 self.zip(rhs, |a, b| a.sub(b, self.p, RoundingMode::ToEven))
850 }
851
852 pub fn mul(&self, rhs: &Self) -> Option<Self> {
854 self.zip(rhs, |a, b| a.mul(b, self.p, RoundingMode::ToEven))
855 }
856
857 pub fn mul_scalar(&self, s: &ExactNum) -> Self {
859 Self {
860 p: self.p,
861 vals: self
862 .vals
863 .iter()
864 .map(|x| x.mul(s, self.p, RoundingMode::ToEven))
865 .collect(),
866 rows: self.rows,
867 cols: self.cols,
868 }
869 }
870
871 pub fn div(&self, rhs: &Self) -> Option<Self> {
873 self.zip(rhs, |a, b| a.div(b, self.p, RoundingMode::ToEven))
874 }
875
876 pub fn sum(&self) -> ExactNum {
878 let mut acc = ExactNum::from_u8(0, self.p);
879 for v in &self.vals {
880 acc = acc.add(v, self.p, RoundingMode::ToEven);
881 }
882 acc
883 }
884
885 pub fn dot(&self, rhs: &Self) -> Option<ExactNum> {
887 if self.len() != rhs.len() {
888 return None;
889 }
890 let mut acc = ExactNum::from_u8(0, self.p);
891 for (a, b) in self.vals.iter().zip(rhs.vals.iter()) {
892 let t = a.mul(b, self.p, RoundingMode::ToEven);
893 acc = acc.add(&t, self.p, RoundingMode::ToEven);
894 }
895 Some(acc)
896 }
897
898 pub fn matmul(&self, rhs: &Self) -> Option<Self> {
900 if self.cols != rhs.rows {
901 return None;
902 }
903 let m = self.rows;
904 let k = self.cols;
905 let n = rhs.cols;
906 let mut vals = Vec::with_capacity(m.checked_mul(n)?);
907 for i in 0..m {
908 for j in 0..n {
909 let mut acc = ExactNum::from_u8(0, self.p);
910 for t in 0..k {
911 let prod = self.vals[i * k + t].mul(
912 &rhs.vals[t * n + j],
913 self.p,
914 RoundingMode::ToEven,
915 );
916 acc = acc.add(&prod, self.p, RoundingMode::ToEven);
917 }
918 vals.push(acc);
919 }
920 }
921 Some(Self {
922 p: self.p,
923 vals,
924 rows: m,
925 cols: n,
926 })
927 }
928
929 pub fn int(&self) -> Self {
931 self.map_at(self.p, |x| x.int())
932 }
933 pub fn fract(&self) -> Self {
935 self.map_at(self.p, |x| x.fract())
936 }
937 pub fn ceil(&self) -> Self {
939 self.map_at(self.p, |x| x.ceil())
940 }
941 pub fn floor(&self) -> Self {
943 self.map_at(self.p, |x| x.floor())
944 }
945 pub fn round(&self, n: usize, rm: RoundingMode) -> Self {
947 self.map_at(self.p, |x| x.round(n, rm))
948 }
949 pub fn abs(&self) -> Self {
951 self.map_at(self.p, |x| x.abs())
952 }
953 pub fn signum(&self) -> Self {
955 self.map_at(self.p, |x| x.signum())
956 }
957 pub fn neg(&self) -> Self {
959 self.map_at(self.p, |x| x.neg())
960 }
961 pub fn reciprocal(&self, p: usize, rm: RoundingMode) -> Self {
963 self.map_at(p, |x| x.reciprocal(p, rm))
964 }
965 pub fn nth_root(&self, n: usize, p: usize, rm: RoundingMode) -> Self {
967 self.map_at(p, |x| x.nth_root(n, p, rm))
968 }
969 pub fn powi(&self, n: usize, p: usize, rm: RoundingMode) -> Self {
971 self.map_at(p, |x| x.powi(n, p, rm))
972 }
973 pub fn powsi(&self, n: isize, p: usize, rm: RoundingMode) -> Self {
975 self.map_at(p, |x| x.powsi(n, p, rm))
976 }
977
978 exact_arr_p_rm_cc!(
979 sin,
980 cos,
981 tan,
982 asin,
983 acos,
984 atan,
985 sinh,
986 cosh,
987 tanh,
988 asinh,
989 acosh,
990 atanh,
991 exp,
992 exp2,
993 exp10,
994 expm1,
995 ln,
996 log2,
997 log10,
998 log1p,
999 erf,
1000 erfc,
1001 gamma,
1002 ln_gamma,
1003 digamma,
1004 ei,
1005 si,
1006 ci,
1007 li,
1008 fresnel_s,
1009 fresnel_c,
1010 ai,
1011 bi,
1012 elliptic_k,
1013 elliptic_e_complete,
1014 rem_pi,
1015 );
1016
1017 pub fn sin_cos(&self, p: usize, rm: RoundingMode, cc: &mut Consts) -> (Self, Self) {
1019 let mut s = Vec::with_capacity(self.vals.len());
1020 let mut c = Vec::with_capacity(self.vals.len());
1021 for x in &self.vals {
1022 let (sv, cv) = x.sin_cos(p, rm, cc);
1023 s.push(sv);
1024 c.push(cv);
1025 }
1026 (
1027 Self {
1028 p,
1029 vals: s,
1030 rows: self.rows,
1031 cols: self.cols,
1032 },
1033 Self {
1034 p,
1035 vals: c,
1036 rows: self.rows,
1037 cols: self.cols,
1038 },
1039 )
1040 }
1041
1042 pub fn sinh_cosh(&self, p: usize, rm: RoundingMode, cc: &mut Consts) -> (Self, Self) {
1044 let mut s = Vec::with_capacity(self.vals.len());
1045 let mut c = Vec::with_capacity(self.vals.len());
1046 for x in &self.vals {
1047 let (sv, cv) = x.sinh_cosh(p, rm, cc);
1048 s.push(sv);
1049 c.push(cv);
1050 }
1051 (
1052 Self {
1053 p,
1054 vals: s,
1055 rows: self.rows,
1056 cols: self.cols,
1057 },
1058 Self {
1059 p,
1060 vals: c,
1061 rows: self.rows,
1062 cols: self.cols,
1063 },
1064 )
1065 }
1066
1067 pub fn sqrt(&self, p: usize, rm: RoundingMode) -> Self {
1069 self.map_at(p, |x| x.sqrt(p, rm))
1070 }
1071
1072 pub fn cbrt(&self, p: usize, rm: RoundingMode) -> Self {
1074 self.map_at(p, |x| x.cbrt(p, rm))
1075 }
1076
1077 pub fn legendre_p(&self, n: u32, p: usize, rm: RoundingMode) -> Self {
1079 self.map_at(p, |x| x.legendre_p(n, p, rm))
1080 }
1081
1082 pub fn assoc_legendre_p(&self, n: u32, m: i32, p: usize, rm: RoundingMode) -> Self {
1084 self.map_at(p, |x| x.assoc_legendre_p(n, m, p, rm))
1085 }
1086
1087 pub fn hypot(&self, other: &ExactNum, p: usize, rm: RoundingMode) -> Self {
1089 self.map_at(p, |x| x.hypot(other, p, rm))
1090 }
1091
1092 pub fn atan2(&self, x: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1094 self.map_at(p, |y| y.atan2(x, p, rm, cc))
1095 }
1096
1097 pub fn pow(&self, n: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1099 self.map_at(p, |x| x.pow(n, p, rm, cc))
1100 }
1101
1102 pub fn log(&self, base: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1104 self.map_at(p, |x| x.log(base, p, rm, cc))
1105 }
1106
1107 pub fn gammainc(&self, x: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1109 self.map_at(p, |s| s.gammainc(x, p, rm, cc))
1110 }
1111
1112 pub fn gammainc_upper(
1114 &self,
1115 x: &ExactNum,
1116 p: usize,
1117 rm: RoundingMode,
1118 cc: &mut Consts,
1119 ) -> Self {
1120 self.map_at(p, |s| s.gammainc_upper(x, p, rm, cc))
1121 }
1122
1123 pub fn bessel_j(&self, n: usize, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1125 self.map_at(p, |x| x.bessel_j(n, p, rm, cc))
1126 }
1127
1128 pub fn bessel_j_nu(&self, nu: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1130 self.map_at(p, |x| x.bessel_j_nu(nu, p, rm, cc))
1131 }
1132
1133 pub fn bessel_y(&self, nu: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1135 self.map_at(p, |x| x.bessel_y(nu, p, rm, cc))
1136 }
1137
1138 pub fn bessel_i(&self, nu: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1140 self.map_at(p, |x| x.bessel_i(nu, p, rm, cc))
1141 }
1142
1143 pub fn bessel_k(&self, nu: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1145 self.map_at(p, |x| x.bessel_k(nu, p, rm, cc))
1146 }
1147
1148 pub fn elliptic_f(&self, m: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1150 self.map_at(p, |x| x.elliptic_f(m, p, rm, cc))
1151 }
1152
1153 pub fn jacobi_sn(&self, m: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1155 self.map_at(p, |x| x.jacobi_sn(m, p, rm, cc))
1156 }
1157
1158 pub fn jacobi_cn(&self, m: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1160 self.map_at(p, |x| x.jacobi_cn(m, p, rm, cc))
1161 }
1162
1163 pub fn jacobi_dn(&self, m: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1165 self.map_at(p, |x| x.jacobi_dn(m, p, rm, cc))
1166 }
1167
1168 pub fn jacobi_am(&self, m: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1170 self.map_at(p, |x| x.jacobi_am(m, p, rm, cc))
1171 }
1172
1173 pub fn jacobi_cd(&self, m: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1175 self.map_at(p, |x| x.jacobi_cd(m, p, rm, cc))
1176 }
1177
1178 pub fn jacobi_ns(&self, m: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1180 self.map_at(p, |x| x.jacobi_ns(m, p, rm, cc))
1181 }
1182
1183 pub fn jacobi_nc(&self, m: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1185 self.map_at(p, |x| x.jacobi_nc(m, p, rm, cc))
1186 }
1187
1188 pub fn jacobi_nd(&self, m: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1190 self.map_at(p, |x| x.jacobi_nd(m, p, rm, cc))
1191 }
1192
1193 pub fn jacobi_sc(&self, m: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1195 self.map_at(p, |x| x.jacobi_sc(m, p, rm, cc))
1196 }
1197
1198 pub fn jacobi_sd(&self, m: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1200 self.map_at(p, |x| x.jacobi_sd(m, p, rm, cc))
1201 }
1202
1203 pub fn jacobi_cs(&self, m: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1205 self.map_at(p, |x| x.jacobi_cs(m, p, rm, cc))
1206 }
1207
1208 pub fn jacobi_ds(&self, m: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1210 self.map_at(p, |x| x.jacobi_ds(m, p, rm, cc))
1211 }
1212
1213 pub fn jacobi_dc(&self, m: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1215 self.map_at(p, |x| x.jacobi_dc(m, p, rm, cc))
1216 }
1217
1218 pub fn elliptic_e(&self, m: &ExactNum, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
1220 self.map_at(p, |x| x.elliptic_e(m, p, rm, cc))
1221 }
1222
1223 pub fn elliptic_pi_complete(
1225 &self,
1226 m: &ExactNum,
1227 p: usize,
1228 rm: RoundingMode,
1229 cc: &mut Consts,
1230 ) -> Self {
1231 self.map_at(p, |n| n.elliptic_pi_complete(m, p, rm, cc))
1232 }
1233
1234 pub fn elliptic_pi(
1236 &self,
1237 x: &ExactNum,
1238 m: &ExactNum,
1239 p: usize,
1240 rm: RoundingMode,
1241 cc: &mut Consts,
1242 ) -> Self {
1243 self.map_at(p, |n| n.elliptic_pi(x, m, p, rm, cc))
1244 }
1245
1246 pub fn hypergeom_2f1(
1248 &self,
1249 b: &ExactNum,
1250 c: &ExactNum,
1251 z: &ExactNum,
1252 p: usize,
1253 rm: RoundingMode,
1254 cc: &mut Consts,
1255 ) -> Self {
1256 self.map_at(p, |a| a.hypergeom_2f1(b, c, z, p, rm, cc))
1257 }
1258
1259 pub fn betainc(
1261 &self,
1262 b: &ExactNum,
1263 x: &ExactNum,
1264 p: usize,
1265 rm: RoundingMode,
1266 cc: &mut Consts,
1267 ) -> Self {
1268 self.map_at(p, |a| a.betainc(b, x, p, rm, cc))
1269 }
1270
1271 pub fn lu_decomp(&self, p: usize, rm: RoundingMode) -> Option<(Self, Self, Vec<usize>)> {
1277 let n = self.rows;
1278 let m = self.cols;
1279 if n == 0 || m == 0 {
1280 return None;
1281 }
1282 let kmax = n.min(m);
1283 let mut a = try_clone_vals(&self.vals)?;
1284 for v in &mut a {
1285 let _ = v.set_precision(p, rm);
1286 }
1287 let mut perm = try_alloc_vec(n, 0usize)?;
1288 for (i, slot) in perm.iter_mut().enumerate() {
1289 *slot = i;
1290 }
1291 let mut lvals = try_alloc_vec(n.checked_mul(n)?, ExactNum::from_u8(0, p))?;
1292 for i in 0..n {
1293 lvals[i * n + i] = ExactNum::from_u8(1, p);
1294 }
1295 let ix = |r: usize, c: usize| r * m + c;
1296 for k in 0..kmax {
1297 let mut piv = k;
1298 let mut best = a[ix(k, k)].abs();
1299 for r in (k + 1)..n {
1300 let t = a[ix(r, k)].abs();
1301 if matches!(t.cmp(&best), Some(c) if c > 0) {
1302 best = t;
1303 piv = r;
1304 }
1305 }
1306 if a[ix(piv, k)].is_zero() {
1307 return None;
1308 }
1309 if piv != k {
1310 for c in 0..m {
1311 a.swap(ix(k, c), ix(piv, c));
1312 }
1313 for c in 0..k {
1314 lvals.swap(k * n + c, piv * n + c);
1315 }
1316 perm.swap(k, piv);
1317 }
1318 let akk = a[ix(k, k)].clone();
1319 for i in (k + 1)..n {
1320 let lik = a[ix(i, k)].div(&akk, p, rm);
1321 lvals[i * n + k] = lik.clone();
1322 a[ix(i, k)] = ExactNum::from_u8(0, p);
1323 for j in (k + 1)..m {
1324 let t = lik.mul(&a[ix(k, j)], p, rm);
1325 a[ix(i, j)] = a[ix(i, j)].sub(&t, p, rm);
1326 }
1327 }
1328 }
1329 Some((
1330 Self {
1331 p,
1332 vals: lvals,
1333 rows: n,
1334 cols: n,
1335 },
1336 Self {
1337 p,
1338 vals: a,
1339 rows: n,
1340 cols: m,
1341 },
1342 perm,
1343 ))
1344 }
1345
1346 pub fn transpose(&self) -> Self {
1348 let mut vals = Vec::with_capacity(self.vals.len());
1349 for j in 0..self.cols {
1350 for i in 0..self.rows {
1351 vals.push(self.vals[i * self.cols + j].clone());
1352 }
1353 }
1354 Self {
1355 p: self.p,
1356 vals,
1357 rows: self.cols,
1358 cols: self.rows,
1359 }
1360 }
1361
1362 pub fn qr_decomp(&self, p: usize, rm: RoundingMode) -> Option<(Self, Self)> {
1368 let m = self.rows;
1369 let n = self.cols;
1370 if m == 0 || n == 0 {
1371 return None;
1372 }
1373 let k = m.min(n);
1374 let mut q = try_alloc_vec(m.checked_mul(k)?, ExactNum::from_u8(0, p))?;
1375 let mut r = try_alloc_vec(k.checked_mul(n)?, ExactNum::from_u8(0, p))?;
1376 let a = |row: usize, col: usize| -> ExactNum {
1377 let mut v = self.vals[row * n + col].clone();
1378 let _ = v.set_precision(p, rm);
1379 v
1380 };
1381 for j in 0..n {
1382 let mut v: Vec<ExactNum> = (0..m).map(|i| a(i, j)).collect();
1383 let jlim = j.min(k);
1384 for i in 0..jlim {
1385 let mut dot = ExactNum::from_u8(0, p);
1386 for t in 0..m {
1387 let qi = q[t * k + i].clone();
1388 dot = dot.add(&qi.mul(&v[t], p, rm), p, rm);
1389 }
1390 r[i * n + j] = dot.clone();
1391 for t in 0..m {
1392 let qi = q[t * k + i].clone();
1393 v[t] = v[t].sub(&dot.mul(&qi, p, rm), p, rm);
1394 }
1395 }
1396 if j < k {
1397 let mut nrm = ExactNum::from_u8(0, p);
1398 for t in 0..m {
1399 nrm = nrm.add(&v[t].mul(&v[t], p, rm), p, rm);
1400 }
1401 nrm = nrm.sqrt(p, rm);
1402 r[j * n + j] = nrm.clone();
1403 if !nrm.is_zero() {
1404 for t in 0..m {
1405 q[t * k + j] = v[t].div(&nrm, p, rm);
1406 }
1407 }
1408 }
1409 }
1410 Some((
1411 Self {
1412 p,
1413 vals: q,
1414 rows: m,
1415 cols: k,
1416 },
1417 Self {
1418 p,
1419 vals: r,
1420 rows: k,
1421 cols: n,
1422 },
1423 ))
1424 }
1425
1426 pub fn svd_decomp(&self, p: usize, rm: RoundingMode) -> Option<(Self, Self, Self)> {
1434 let m = self.rows;
1435 let n = self.cols;
1436 if m == 0 || n == 0 {
1437 return None;
1438 }
1439 for v in &self.vals {
1440 if v.is_nan() || v.is_inf() {
1441 return None;
1442 }
1443 }
1444 if m < n {
1445 let (ut, s, vtt) = self.transpose().svd_decomp(p, rm)?;
1446 return Some((vtt.transpose(), s, ut.transpose()));
1447 }
1448 svd_decomp_tall(self, p, rm)
1449 }
1450
1451 pub fn eigen_decomp(&self, p: usize, rm: RoundingMode) -> Option<(Self, Self)> {
1459 let n = self.rows;
1460 if n == 0 || n != self.cols {
1461 return None;
1462 }
1463 for v in &self.vals {
1464 if v.is_nan() || v.is_inf() {
1465 return None;
1466 }
1467 }
1468 for i in 0..n {
1469 for j in 0..i {
1470 let aij = svd_copy_prec(&self.vals[i * n + j], p, rm);
1471 let aji = svd_copy_prec(&self.vals[j * n + i], p, rm);
1472 if aij.cmp(&aji) != Some(0) {
1473 return None;
1474 }
1475 }
1476 }
1477 eigen_decomp_sym(self, p, rm)
1478 }
1479
1480 pub fn fft(&self, p: usize, rm: RoundingMode, cc: &mut Consts) -> Option<Self> {
1487 fft_dit(self, false, p, rm, cc)
1488 }
1489
1490 pub fn ifft(&self, p: usize, rm: RoundingMode, cc: &mut Consts) -> Option<Self> {
1493 fft_dit(self, true, p, rm, cc)
1494 }
1495
1496 fn zip(&self, rhs: &Self, op: impl Fn(&ExactNum, &ExactNum) -> ExactNum) -> Option<Self> {
1497 if self.rows != rhs.rows || self.cols != rhs.cols {
1498 return None;
1499 }
1500 Some(Self {
1501 p: self.p,
1502 vals: self
1503 .vals
1504 .iter()
1505 .zip(rhs.vals.iter())
1506 .map(|(a, b)| op(a, b))
1507 .collect(),
1508 rows: self.rows,
1509 cols: self.cols,
1510 })
1511 }
1512
1513 fn map_at(&self, p: usize, mut op: impl FnMut(&ExactNum) -> ExactNum) -> Self {
1514 Self {
1515 p,
1516 vals: self.vals.iter().map(|x| op(x)).collect(),
1517 rows: self.rows,
1518 cols: self.cols,
1519 }
1520 }
1521}
1522
1523fn try_alloc_vec<T: Clone>(n: usize, fill: T) -> Option<Vec<T>> {
1524 let mut v = Vec::new();
1525 v.try_reserve_exact(n).ok()?;
1526 v.resize(n, fill);
1527 Some(v)
1528}
1529
1530fn try_clone_vals(src: &[ExactNum]) -> Option<Vec<ExactNum>> {
1531 let mut v = Vec::new();
1532 v.try_reserve_exact(src.len()).ok()?;
1533 v.extend(src.iter().cloned());
1534 Some(v)
1535}
1536
1537fn svd_zero(p: usize) -> ExactNum {
1538 ExactNum::from_u8(0, p)
1539}
1540
1541fn svd_one(p: usize) -> ExactNum {
1542 ExactNum::from_u8(1, p)
1543}
1544
1545fn svd_copy_prec(x: &ExactNum, p: usize, rm: RoundingMode) -> ExactNum {
1546 let mut y = x.clone();
1547 let _ = y.set_precision(p, rm);
1548 y
1549}
1550
1551fn svd_identity(n: usize, p: usize) -> Option<Vec<ExactNum>> {
1552 let mut v = try_alloc_vec(n.checked_mul(n)?, svd_zero(p))?;
1553 for i in 0..n {
1554 v[i * n + i] = svd_one(p);
1555 }
1556 Some(v)
1557}
1558
1559fn svd_norm(xs: &[ExactNum], p: usize, rm: RoundingMode) -> ExactNum {
1560 let mut n = svd_zero(p);
1561 for x in xs {
1562 n = n.hypot(x, p, rm);
1563 }
1564 n
1565}
1566
1567fn svd_householder(
1568 x: &[ExactNum],
1569 p: usize,
1570 rm: RoundingMode,
1571) -> Option<(Vec<ExactNum>, ExactNum)> {
1572 if x.is_empty() {
1573 return None;
1574 }
1575 let norm = svd_norm(x, p, rm);
1576 if norm.is_zero() {
1577 return None;
1578 }
1579 let mut v = x.to_vec();
1580 let signed = if v[0].is_negative() { norm.neg() } else { norm };
1581 v[0] = v[0].add(&signed, p, rm);
1582 let mut vtv = svd_zero(p);
1583 for vi in &v {
1584 vtv = vtv.add(&vi.mul(vi, p, rm), p, rm);
1585 }
1586 if vtv.is_zero() {
1587 return None;
1588 }
1589 let beta = ExactNum::from_u8(2, p).div(&vtv, p, rm);
1590 Some((v, beta))
1591}
1592
1593fn svd_apply_house_left(
1594 a: &mut [ExactNum],
1595 cols: usize,
1596 row0: usize,
1597 col0: usize,
1598 v: &[ExactNum],
1599 beta: &ExactNum,
1600 p: usize,
1601 rm: RoundingMode,
1602) {
1603 let vlen = v.len();
1604 for j in col0..cols {
1605 let mut s = svd_zero(p);
1606 for i in 0..vlen {
1607 s = s.add(&v[i].mul(&a[(row0 + i) * cols + j], p, rm), p, rm);
1608 }
1609 s = s.mul(beta, p, rm);
1610 for i in 0..vlen {
1611 let t = s.mul(&v[i], p, rm);
1612 let idx = (row0 + i) * cols + j;
1613 a[idx] = a[idx].sub(&t, p, rm);
1614 }
1615 }
1616}
1617
1618fn svd_apply_house_right(
1619 a: &mut [ExactNum],
1620 rows: usize,
1621 cols: usize,
1622 row0: usize,
1623 col0: usize,
1624 v: &[ExactNum],
1625 beta: &ExactNum,
1626 p: usize,
1627 rm: RoundingMode,
1628) {
1629 let vlen = v.len();
1630 for i in row0..rows {
1631 let mut s = svd_zero(p);
1632 for t in 0..vlen {
1633 s = s.add(&a[i * cols + col0 + t].mul(&v[t], p, rm), p, rm);
1634 }
1635 s = s.mul(beta, p, rm);
1636 for t in 0..vlen {
1637 let tt = s.mul(&v[t], p, rm);
1638 let idx = i * cols + col0 + t;
1639 a[idx] = a[idx].sub(&tt, p, rm);
1640 }
1641 }
1642}
1643
1644fn svd_rotg(
1645 a: &ExactNum,
1646 b: &ExactNum,
1647 p: usize,
1648 rm: RoundingMode,
1649) -> (ExactNum, ExactNum, ExactNum) {
1650 let r = a.hypot(b, p, rm);
1651 if r.is_zero() {
1652 return (svd_one(p), svd_zero(p), r);
1653 }
1654 (a.div(&r, p, rm), b.div(&r, p, rm), r)
1655}
1656
1657fn svd_apply_givens_cols(
1658 mat: &mut [ExactNum],
1659 rows: usize,
1660 cols: usize,
1661 j0: usize,
1662 j1: usize,
1663 cs: &ExactNum,
1664 sn: &ExactNum,
1665 p: usize,
1666 rm: RoundingMode,
1667) {
1668 for i in 0..rows {
1669 let a = mat[i * cols + j0].clone();
1670 let b = mat[i * cols + j1].clone();
1671 mat[i * cols + j0] = cs.mul(&a, p, rm).add(&sn.mul(&b, p, rm), p, rm);
1672 mat[i * cols + j1] = cs.mul(&b, p, rm).sub(&sn.mul(&a, p, rm), p, rm);
1673 }
1674}
1675
1676fn svd_negligible(e: &ExactNum, d0: &ExactNum, d1: &ExactNum, p: usize, rm: RoundingMode) -> bool {
1677 if e.is_zero() {
1678 return true;
1679 }
1680 if e.is_nan() || e.is_inf() {
1681 return false;
1682 }
1683 let scale = d0.abs().add(&d1.abs(), p, rm);
1684 if scale.is_zero() {
1685 return e.is_zero();
1686 }
1687 match (e.abs().exponent(), scale.exponent()) {
1688 (Some(ee), Some(se)) => ee < se - (p as i32 - SVD_CONV_GUARD_BITS),
1689 _ => false,
1690 }
1691}
1692
1693fn svd_wilkinson_shift(
1694 d_prev: &ExactNum,
1695 d_last: &ExactNum,
1696 e_prev: &ExactNum,
1697 e_last: &ExactNum,
1698 p: usize,
1699 rm: RoundingMode,
1700) -> ExactNum {
1701 let two = ExactNum::from_u8(2, p);
1702 let b = d_prev
1703 .add(d_last, p, rm)
1704 .mul(&d_prev.sub(d_last, p, rm), p, rm)
1705 .add(&e_prev.mul(e_prev, p, rm), p, rm)
1706 .div(&two, p, rm);
1707 let t = d_last.mul(e_last, p, rm);
1708 let c = t.mul(&t, p, rm);
1709 if b.is_zero() && c.is_zero() {
1710 return svd_zero(p);
1711 }
1712 let disc = b.mul(&b, p, rm).add(&c, p, rm).sqrt(p, rm);
1713 let signed = if b.is_negative() { disc.neg() } else { disc };
1714 let denom = b.add(&signed, p, rm);
1715 if denom.is_zero() {
1716 return svd_zero(p);
1717 }
1718 c.div(&denom, p, rm)
1719}
1720
1721fn svd_qr_sweep(
1722 d: &mut [ExactNum],
1723 e: &mut [ExactNum],
1724 u: &mut [ExactNum],
1725 v: &mut [ExactNum],
1726 m: usize,
1727 n: usize,
1728 p_blk: usize,
1729 q_blk: usize,
1730 p: usize,
1731 rm: RoundingMode,
1732) {
1733 let last = q_blk - 1;
1734 let e_prev = if last >= p_blk + 2 { e[last - 2].clone() } else { svd_zero(p) };
1735 let shift = svd_wilkinson_shift(&d[last - 1], &d[last], &e_prev, &e[last - 1], p, rm);
1736 let mut f = d[p_blk]
1737 .add(&d[last], p, rm)
1738 .mul(&d[p_blk].sub(&d[last], p, rm), p, rm)
1739 .add(&shift, p, rm);
1740 let mut g = d[p_blk].mul(&e[p_blk], p, rm);
1741 for j in p_blk..last {
1742 let (cs, sn, r) = svd_rotg(&f, &g, p, rm);
1743 if j > p_blk {
1744 e[j - 1] = r;
1745 }
1746 let dj = d[j].clone();
1747 let ej = e[j].clone();
1748 let dj1 = d[j + 1].clone();
1749 f = cs.mul(&dj, p, rm).add(&sn.mul(&ej, p, rm), p, rm);
1750 e[j] = cs.mul(&ej, p, rm).sub(&sn.mul(&dj, p, rm), p, rm);
1751 g = sn.mul(&dj1, p, rm);
1752 d[j + 1] = cs.mul(&dj1, p, rm);
1753 svd_apply_givens_cols(v, n, n, j, j + 1, &cs, &sn, p, rm);
1754
1755 let (cs, sn, r) = svd_rotg(&f, &g, p, rm);
1756 d[j] = r;
1757 let ej = e[j].clone();
1758 let dj1 = d[j + 1].clone();
1759 f = cs.mul(&ej, p, rm).add(&sn.mul(&dj1, p, rm), p, rm);
1760 d[j + 1] = cs.mul(&dj1, p, rm).sub(&sn.mul(&ej, p, rm), p, rm);
1761 if j + 1 < last {
1762 g = sn.mul(&e[j + 1], p, rm);
1763 e[j + 1] = cs.mul(&e[j + 1], p, rm);
1764 }
1765 svd_apply_givens_cols(u, m, m, j, j + 1, &cs, &sn, p, rm);
1766 }
1767 e[last - 1] = f;
1768}
1769
1770fn svd_zero_last_super(
1771 d: &mut [ExactNum],
1772 e: &mut [ExactNum],
1773 v: &mut [ExactNum],
1774 n: usize,
1775 p_blk: usize,
1776 q_blk: usize,
1777 p: usize,
1778 rm: RoundingMode,
1779) {
1780 let k = q_blk - 1;
1781 let mut f = e[k - 1].clone();
1782 e[k - 1] = svd_zero(p);
1783 for j in (p_blk..k).rev() {
1784 let (cs, sn, t) = svd_rotg(&d[j], &f, p, rm);
1785 d[j] = t;
1786 if j > p_blk {
1787 f = sn.neg().mul(&e[j - 1], p, rm);
1788 e[j - 1] = cs.mul(&e[j - 1], p, rm);
1789 }
1790 svd_apply_givens_cols(v, n, n, j, k, &cs, &sn, p, rm);
1791 }
1792}
1793
1794fn svd_zero_first_super(
1795 d: &mut [ExactNum],
1796 e: &mut [ExactNum],
1797 u: &mut [ExactNum],
1798 m: usize,
1799 p_blk: usize,
1800 q_blk: usize,
1801 p: usize,
1802 rm: RoundingMode,
1803) {
1804 let mut f = e[p_blk].clone();
1805 e[p_blk] = svd_zero(p);
1806 for j in (p_blk + 1)..q_blk {
1807 let (cs, sn, t) = svd_rotg(&d[j], &f, p, rm);
1808 d[j] = t;
1809 if j + 1 < q_blk {
1810 f = sn.neg().mul(&e[j], p, rm);
1811 e[j] = cs.mul(&e[j], p, rm);
1812 }
1813 svd_apply_givens_cols(u, m, m, p_blk, j, &cs, &sn, p, rm);
1814 }
1815}
1816
1817fn svd_take_cols(vals: &[ExactNum], rows: usize, cols: usize, k: usize, p: usize) -> ExactNumArray {
1818 let mut out = Vec::with_capacity(rows * k);
1819 for i in 0..rows {
1820 for j in 0..k {
1821 out.push(vals[i * cols + j].clone());
1822 }
1823 }
1824 ExactNumArray {
1825 p,
1826 vals: out,
1827 rows,
1828 cols: k,
1829 }
1830}
1831
1832fn svd_vt_from_v(v: &[ExactNum], n: usize, k: usize, p: usize) -> ExactNumArray {
1833 let mut out = Vec::with_capacity(k * n);
1834 for j in 0..k {
1835 for i in 0..n {
1836 out.push(v[i * n + j].clone());
1837 }
1838 }
1839 ExactNumArray {
1840 p,
1841 vals: out,
1842 rows: k,
1843 cols: n,
1844 }
1845}
1846
1847fn svd_decomp_tall(
1848 a0: &ExactNumArray,
1849 p: usize,
1850 rm: RoundingMode,
1851) -> Option<(ExactNumArray, ExactNumArray, ExactNumArray)> {
1852 let m = a0.rows;
1853 let n = a0.cols;
1854 let mut a: Vec<ExactNum> = a0.vals.iter().map(|x| svd_copy_prec(x, p, rm)).collect();
1855 let mut u = svd_identity(m, p)?;
1856 let mut v = svd_identity(n, p)?;
1857
1858 for k in 0..n {
1859 let x: Vec<ExactNum> = (k..m).map(|i| a[i * n + k].clone()).collect();
1860 if let Some((hv, beta)) = svd_householder(&x, p, rm) {
1861 svd_apply_house_left(&mut a, n, k, k, &hv, &beta, p, rm);
1862 svd_apply_house_right(&mut u, m, m, 0, k, &hv, &beta, p, rm);
1863 }
1864 if k + 1 < n {
1865 let x: Vec<ExactNum> = ((k + 1)..n).map(|j| a[k * n + j].clone()).collect();
1866 if let Some((hv, beta)) = svd_householder(&x, p, rm) {
1867 svd_apply_house_right(&mut a, m, n, k, k + 1, &hv, &beta, p, rm);
1868 svd_apply_house_right(&mut v, n, n, 0, k + 1, &hv, &beta, p, rm);
1869 }
1870 }
1871 }
1872
1873 let mut d: Vec<ExactNum> = (0..n).map(|i| a[i * n + i].clone()).collect();
1874 let mut e: Vec<ExactNum> = if n >= 2 {
1875 (0..n - 1).map(|i| a[i * n + i + 1].clone()).collect()
1876 } else {
1877 Vec::new()
1878 };
1879
1880 let max_sweeps = SVD_ITER_MAX.saturating_mul(n.max(1) as u32);
1881 let mut sweeps = 0u32;
1882 loop {
1883 if n == 1 {
1884 break;
1885 }
1886 for i in 0..n - 1 {
1887 if svd_negligible(&e[i], &d[i], &d[i + 1], p, rm) {
1888 e[i] = svd_zero(p);
1889 }
1890 }
1891 if e.iter().all(|x| x.is_zero()) {
1892 break;
1893 }
1894 if sweeps >= max_sweeps {
1895 return None;
1896 }
1897 let mut q = n;
1898 while q > 1 && e[q - 2].is_zero() {
1899 q -= 1;
1900 }
1901 let mut p_blk = q - 1;
1902 while p_blk > 0 && !e[p_blk - 1].is_zero() {
1903 p_blk -= 1;
1904 }
1905 if q - p_blk < 2 {
1906 break;
1907 }
1908
1909 let mut did_split = false;
1910 for i in p_blk..q {
1911 let el = if i > p_blk { e[i - 1].clone() } else { svd_zero(p) };
1912 let er = if i + 1 < q { e[i].clone() } else { svd_zero(p) };
1913 if svd_negligible(&d[i], &el, &er, p, rm) {
1914 d[i] = svd_zero(p);
1915 if i == q - 1 && i > p_blk {
1916 svd_zero_last_super(&mut d, &mut e, &mut v, n, p_blk, q, p, rm);
1917 } else if i < q - 1 {
1918 svd_zero_first_super(&mut d, &mut e, &mut u, m, i, q, p, rm);
1919 }
1920 did_split = true;
1921 break;
1922 }
1923 }
1924 if did_split {
1925 sweeps += 1;
1926 continue;
1927 }
1928 for di in d.iter().chain(e.iter()) {
1929 if di.is_nan() || di.is_inf() {
1930 return None;
1931 }
1932 }
1933 svd_qr_sweep(&mut d, &mut e, &mut u, &mut v, m, n, p_blk, q, p, rm);
1934 sweeps += 1;
1935 }
1936
1937 for i in 0..n {
1938 if d[i].is_negative() {
1939 d[i] = d[i].neg();
1940 for r in 0..m {
1941 let idx = r * m + i;
1942 u[idx] = u[idx].neg();
1943 }
1944 }
1945 }
1946 for i in 0..n {
1947 let mut best = i;
1948 for j in (i + 1)..n {
1949 if matches!(d[j].cmp(&d[best]), Some(c) if c > 0) {
1950 best = j;
1951 }
1952 }
1953 if best != i {
1954 d.swap(i, best);
1955 for r in 0..m {
1956 u.swap(r * m + i, r * m + best);
1957 }
1958 for r in 0..n {
1959 v.swap(r * n + i, r * n + best);
1960 }
1961 }
1962 }
1963
1964 let k = n;
1965 let mut sigma = try_alloc_vec(k.checked_mul(k)?, svd_zero(p))?;
1966 for i in 0..k {
1967 sigma[i * k + i] = d[i].clone();
1968 }
1969 Some((
1970 svd_take_cols(&u, m, m, k, p),
1971 ExactNumArray {
1972 p,
1973 vals: sigma,
1974 rows: k,
1975 cols: k,
1976 },
1977 svd_vt_from_v(&v, n, k, p),
1978 ))
1979}
1980
1981fn eigen_wilkinson(
1982 a: &ExactNum,
1983 b: &ExactNum,
1984 c: &ExactNum,
1985 p: usize,
1986 rm: RoundingMode,
1987) -> ExactNum {
1988 let half = svd_one(p).div(&ExactNum::from_u8(2, p), p, rm);
1989 let delta = a.sub(c, p, rm).mul(&half, p, rm);
1990 if delta.is_zero() && b.is_zero() {
1991 return c.clone();
1992 }
1993 let h = delta.hypot(b, p, rm);
1994 let signed = if delta.is_negative() { h.neg() } else { h };
1995 let denom = delta.add(&signed, p, rm);
1996 if denom.is_zero() {
1997 return c.sub(&b.abs(), p, rm);
1998 }
1999 c.sub(&b.mul(b, p, rm).div(&denom, p, rm), p, rm)
2000}
2001
2002fn eigen_qr_sweep(
2003 d: &mut [ExactNum],
2004 e: &mut [ExactNum],
2005 q: &mut [ExactNum],
2006 n: usize,
2007 p_blk: usize,
2008 q_blk: usize,
2009 p: usize,
2010 rm: RoundingMode,
2011) {
2012 let last = q_blk - 1;
2013 let mu = eigen_wilkinson(&d[last - 1], &e[last - 1], &d[last], p, rm);
2014 let mut f = d[p_blk].sub(&mu, p, rm);
2015 let mut g = e[p_blk].clone();
2016 let two = ExactNum::from_u8(2, p);
2017 for k in p_blk..last {
2018 let (cs, sn, r) = svd_rotg(&f, &g, p, rm);
2019 if k > p_blk {
2020 e[k - 1] = r;
2021 }
2022 let d0 = d[k].clone();
2023 let ee = e[k].clone();
2024 let d1 = d[k + 1].clone();
2025 let c2 = cs.mul(&cs, p, rm);
2026 let s2 = sn.mul(&sn, p, rm);
2027 let cs2 = cs.mul(&sn, p, rm);
2028 let two_cse = two.mul(&cs2.mul(&ee, p, rm), p, rm);
2029 d[k] = c2
2030 .mul(&d0, p, rm)
2031 .add(&two_cse, p, rm)
2032 .add(&s2.mul(&d1, p, rm), p, rm);
2033 d[k + 1] = s2
2034 .mul(&d0, p, rm)
2035 .sub(&two_cse, p, rm)
2036 .add(&c2.mul(&d1, p, rm), p, rm);
2037 e[k] = cs2
2038 .mul(&d1.sub(&d0, p, rm), p, rm)
2039 .add(&c2.sub(&s2, p, rm).mul(&ee, p, rm), p, rm);
2040 svd_apply_givens_cols(q, n, n, k, k + 1, &cs, &sn, p, rm);
2041 if k + 1 < last {
2042 let ek1 = e[k + 1].clone();
2043 f = e[k].clone();
2044 g = sn.mul(&ek1, p, rm);
2045 e[k + 1] = cs.mul(&ek1, p, rm);
2046 }
2047 }
2048}
2049
2050fn eigen_decomp_sym(
2051 a0: &ExactNumArray,
2052 p: usize,
2053 rm: RoundingMode,
2054) -> Option<(ExactNumArray, ExactNumArray)> {
2055 let n = a0.rows;
2056 let mut a: Vec<ExactNum> = a0.vals.iter().map(|x| svd_copy_prec(x, p, rm)).collect();
2057 let mut q = svd_identity(n, p)?;
2058 for k in 0..n.saturating_sub(2) {
2059 let x: Vec<ExactNum> = ((k + 1)..n).map(|i| a[i * n + k].clone()).collect();
2060 if let Some((hv, beta)) = svd_householder(&x, p, rm) {
2061 svd_apply_house_left(&mut a, n, k + 1, k, &hv, &beta, p, rm);
2062 svd_apply_house_right(&mut a, n, n, 0, k + 1, &hv, &beta, p, rm);
2063 svd_apply_house_right(&mut q, n, n, 0, k + 1, &hv, &beta, p, rm);
2064 }
2065 }
2066 let mut d: Vec<ExactNum> = (0..n).map(|i| a[i * n + i].clone()).collect();
2067 let mut e: Vec<ExactNum> = if n >= 2 {
2068 (0..n - 1).map(|i| a[i * n + i + 1].clone()).collect()
2069 } else {
2070 Vec::new()
2071 };
2072
2073 let max_sweeps = EIGEN_ITER_MAX.saturating_mul(n.max(1) as u32);
2074 let mut sweeps = 0u32;
2075 loop {
2076 if n == 1 {
2077 break;
2078 }
2079 for i in 0..n - 1 {
2080 if svd_negligible(&e[i], &d[i], &d[i + 1], p, rm) {
2081 e[i] = svd_zero(p);
2082 }
2083 }
2084 if e.iter().all(|x| x.is_zero()) {
2085 break;
2086 }
2087 if sweeps >= max_sweeps {
2088 return None;
2089 }
2090 let mut q_blk = n;
2091 while q_blk > 1 && e[q_blk - 2].is_zero() {
2092 q_blk -= 1;
2093 }
2094 let mut p_blk = q_blk - 1;
2095 while p_blk > 0 && !e[p_blk - 1].is_zero() {
2096 p_blk -= 1;
2097 }
2098 if q_blk - p_blk < 2 {
2099 break;
2100 }
2101 for di in d.iter().chain(e.iter()) {
2102 if di.is_nan() || di.is_inf() {
2103 return None;
2104 }
2105 }
2106 eigen_qr_sweep(&mut d, &mut e, &mut q, n, p_blk, q_blk, p, rm);
2107 sweeps += 1;
2108 }
2109
2110 for i in 0..n {
2111 let mut best = i;
2112 for j in (i + 1)..n {
2113 if matches!(d[j].cmp(&d[best]), Some(c) if c > 0) {
2114 best = j;
2115 }
2116 }
2117 if best != i {
2118 d.swap(i, best);
2119 for r in 0..n {
2120 q.swap(r * n + i, r * n + best);
2121 }
2122 }
2123 }
2124 Some((
2125 ExactNumArray {
2126 p,
2127 vals: d,
2128 rows: 1,
2129 cols: n,
2130 },
2131 ExactNumArray {
2132 p,
2133 vals: q,
2134 rows: n,
2135 cols: n,
2136 },
2137 ))
2138}
2139
2140fn fft_from_len(n: usize, p: usize) -> ExactNum {
2141 ExactNum::from_word(n as crate::defs::Word, p)
2142}
2143
2144fn fft_bitrev(mut i: usize, logn: u32) -> usize {
2145 let mut r = 0usize;
2146 for _ in 0..logn {
2147 r = (r << 1) | (i & 1);
2148 i >>= 1;
2149 }
2150 r
2151}
2152
2153fn fft_split(
2154 a: &ExactNumArray,
2155 p: usize,
2156 rm: RoundingMode,
2157) -> Option<(usize, Vec<ExactNum>, Vec<ExactNum>)> {
2158 let (rows, cols) = a.shape();
2159 let pack = |n: usize,
2160 re: Vec<ExactNum>,
2161 im: Vec<ExactNum>|
2162 -> Option<(usize, Vec<ExactNum>, Vec<ExactNum>)> {
2163 if n == 0 || !n.is_power_of_two() || n > FFT_MAX_POINTS {
2164 return None;
2165 }
2166 Some((n, re, im))
2167 };
2168 if rows == 2 && cols > 0 {
2169 let mut re = Vec::with_capacity(cols);
2170 let mut im = Vec::with_capacity(cols);
2171 for j in 0..cols {
2172 let r = svd_copy_prec(&a.vals[j], p, rm);
2173 let i = svd_copy_prec(&a.vals[cols + j], p, rm);
2174 if r.is_nan() || r.is_inf() || i.is_nan() || i.is_inf() {
2175 return None;
2176 }
2177 re.push(r);
2178 im.push(i);
2179 }
2180 pack(cols, re, im)
2181 } else if (rows == 1 && cols > 0) || (cols == 1 && rows > 0) {
2182 let n = a.vals.len();
2183 let mut re = Vec::with_capacity(n);
2184 for v in &a.vals {
2185 let r = svd_copy_prec(v, p, rm);
2186 if r.is_nan() || r.is_inf() {
2187 return None;
2188 }
2189 re.push(r);
2190 }
2191 pack(n, re, alloc::vec![svd_zero(p); n])
2192 } else {
2193 None
2194 }
2195}
2196
2197fn fft_dit(
2198 a: &ExactNumArray,
2199 inverse: bool,
2200 p: usize,
2201 rm: RoundingMode,
2202 cc: &mut Consts,
2203) -> Option<ExactNumArray> {
2204 let (n, mut re, mut im) = fft_split(a, p, rm)?;
2205 let logn = n.trailing_zeros();
2206 for i in 0..n {
2207 let j = fft_bitrev(i, logn);
2208 if j > i {
2209 re.swap(i, j);
2210 im.swap(i, j);
2211 }
2212 }
2213 let two_pi = ExactNum::from_u8(2, p).mul(&cc.pi(p, rm), p, rm);
2214 let mut m = 2usize;
2215 while m <= n {
2216 let ang = two_pi.div(&fft_from_len(m, p), p, rm);
2217 let (sn, cs) = ang.sin_cos(p, rm, cc);
2218 let wm_re = cs;
2219 let wm_im = if inverse { sn } else { sn.neg() };
2220 let half = m / 2;
2221 let mut k = 0usize;
2222 while k < n {
2223 let mut w_re = svd_one(p);
2224 let mut w_im = svd_zero(p);
2225 for j in 0..half {
2226 let t = k + j + half;
2227 let u = k + j;
2228 let tr = w_re.mul(&re[t], p, rm).sub(&w_im.mul(&im[t], p, rm), p, rm);
2229 let ti = w_re.mul(&im[t], p, rm).add(&w_im.mul(&re[t], p, rm), p, rm);
2230 let ur = re[u].clone();
2231 let ui = im[u].clone();
2232 re[u] = ur.add(&tr, p, rm);
2233 im[u] = ui.add(&ti, p, rm);
2234 re[t] = ur.sub(&tr, p, rm);
2235 im[t] = ui.sub(&ti, p, rm);
2236 let nr = w_re.mul(&wm_re, p, rm).sub(&w_im.mul(&wm_im, p, rm), p, rm);
2237 let ni = w_re.mul(&wm_im, p, rm).add(&w_im.mul(&wm_re, p, rm), p, rm);
2238 w_re = nr;
2239 w_im = ni;
2240 }
2241 k += m;
2242 }
2243 m *= 2;
2244 }
2245 if inverse {
2246 let inv_n = svd_one(p).div(&fft_from_len(n, p), p, rm);
2247 for i in 0..n {
2248 re[i] = re[i].mul(&inv_n, p, rm);
2249 im[i] = im[i].mul(&inv_n, p, rm);
2250 }
2251 }
2252 let mut vals = Vec::with_capacity(n.checked_mul(2)?);
2253 vals.extend(re);
2254 vals.extend(im);
2255 Some(ExactNumArray {
2256 p,
2257 vals,
2258 rows: 2,
2259 cols: n,
2260 })
2261}
2262
2263#[cfg(test)]
2264mod tests {
2265 use super::*;
2266
2267 #[test]
2268 fn bin64_array_add_dot() {
2269 let one = Ieee64::from_i32(1);
2270 let two = Ieee64::from_i32(2);
2271 let a = Ieee64Array::from_values(&[one, two]);
2272 let b = Ieee64Array::from_values(&[two, one]);
2273 let s = a.add(&b).unwrap();
2274 assert_eq!(s.get(0).unwrap().to_bits(), Ieee64::from_i32(3).to_bits());
2275 assert_eq!(s.get(1).unwrap().to_bits(), Ieee64::from_i32(3).to_bits());
2276 let d = a.dot(&b).unwrap();
2277 assert_eq!(d.to_bits(), Ieee64::from_i32(4).to_bits());
2278 assert_eq!(a.sum().to_bits(), Ieee64::from_i32(3).to_bits());
2279 let scaled = a.mul_scalar(two);
2280 assert_eq!(scaled.get(0).unwrap().to_bits(), two.to_bits());
2281 }
2282
2283 #[test]
2284 fn bin32_array_sqrt() {
2285 let four = Ieee32::from_i32(4);
2286 let a = Ieee32Array::filled(3, four);
2287 let r = a.sqrt();
2288 assert_eq!(r.get(0).unwrap().to_bits(), Ieee32::from_i32(2).to_bits());
2289 assert_eq!(r.len(), 3);
2290 }
2291
2292 #[test]
2293 fn bin64_array_exp_sin() {
2294 let mut cc = Consts::new().unwrap();
2295 let z = Ieee64Array::from_values(&[Ieee64::ZERO]);
2296 let e = z.exp(&mut cc);
2297 assert_eq!(e.get(0).unwrap().to_bits(), Ieee64::from_i32(1).to_bits());
2298 let s = z.sin(&mut cc);
2299 assert!(s.get(0).unwrap().is_zero());
2300 }
2301
2302 #[test]
2303 fn exact_array_add() {
2304 let p = 64;
2305 let one = ExactNum::from_u8(1, p);
2306 let two = ExactNum::from_u8(2, p);
2307 let a = ExactNumArray::from_values(p, &[one.clone(), two.clone()]);
2308 let b = ExactNumArray::from_values(p, &[two.clone(), one.clone()]);
2309 let s = a.add(&b).unwrap();
2310 assert_eq!(s.get(0).unwrap().cmp(&ExactNum::from_u8(3, p)), Some(0));
2311 let d = a.dot(&b).unwrap();
2312 assert_eq!(d.cmp(&ExactNum::from_u8(4, p)), Some(0));
2313 }
2314
2315 #[test]
2316 fn bin64_matmul_2x2() {
2317 let a = Ieee64Array::from_shape(
2318 2,
2319 2,
2320 &[
2321 Ieee64::from_i32(1),
2322 Ieee64::from_i32(2),
2323 Ieee64::from_i32(3),
2324 Ieee64::from_i32(4),
2325 ],
2326 )
2327 .unwrap();
2328 let b = Ieee64Array::from_shape(
2329 2,
2330 2,
2331 &[
2332 Ieee64::from_i32(5),
2333 Ieee64::from_i32(6),
2334 Ieee64::from_i32(7),
2335 Ieee64::from_i32(8),
2336 ],
2337 )
2338 .unwrap();
2339 let c = a.matmul(&b).unwrap();
2340 assert_eq!(c.shape(), (2, 2));
2341 assert_eq!(
2342 c.get2(0, 0).unwrap().to_bits(),
2343 Ieee64::from_i32(19).to_bits()
2344 );
2345 assert_eq!(
2346 c.get2(0, 1).unwrap().to_bits(),
2347 Ieee64::from_i32(22).to_bits()
2348 );
2349 assert_eq!(
2350 c.get2(1, 0).unwrap().to_bits(),
2351 Ieee64::from_i32(43).to_bits()
2352 );
2353 assert_eq!(
2354 c.get2(1, 1).unwrap().to_bits(),
2355 Ieee64::from_i32(50).to_bits()
2356 );
2357 }
2358
2359 #[test]
2360 fn bin64_matmul_identity() {
2361 let i2 = Ieee64Array::from_shape(
2362 2,
2363 2,
2364 &[Ieee64::from_i32(1), Ieee64::ZERO, Ieee64::ZERO, Ieee64::from_i32(1)],
2365 )
2366 .unwrap();
2367 let a = Ieee64Array::from_shape(
2368 2,
2369 2,
2370 &[
2371 Ieee64::from_i32(1),
2372 Ieee64::from_i32(2),
2373 Ieee64::from_i32(3),
2374 Ieee64::from_i32(4),
2375 ],
2376 )
2377 .unwrap();
2378 let c = i2.matmul(&a).unwrap();
2379 assert_eq!(
2380 c.get2(0, 0).unwrap().to_bits(),
2381 Ieee64::from_i32(1).to_bits()
2382 );
2383 assert_eq!(
2384 c.get2(0, 1).unwrap().to_bits(),
2385 Ieee64::from_i32(2).to_bits()
2386 );
2387 assert_eq!(
2388 c.get2(1, 0).unwrap().to_bits(),
2389 Ieee64::from_i32(3).to_bits()
2390 );
2391 assert_eq!(
2392 c.get2(1, 1).unwrap().to_bits(),
2393 Ieee64::from_i32(4).to_bits()
2394 );
2395 }
2396
2397 #[test]
2398 fn bin64_matmul_shape_mismatch() {
2399 let a = Ieee64Array::from_shape(2, 2, &[Ieee64::from_i32(1); 4]).unwrap();
2400 let b = Ieee64Array::from_shape(3, 1, &[Ieee64::from_i32(1); 3]).unwrap();
2401 assert!(a.matmul(&b).is_none());
2402 assert!(Ieee64Array::from_shape(2, 2, &[Ieee64::from_i32(1)]).is_none());
2403 let row = Ieee64Array::from_values(&[Ieee64::from_i32(1); 4]);
2404 assert_eq!(row.shape(), (1, 4));
2405 let sq = row.reshape(2, 2).unwrap();
2406 assert_eq!(sq.shape(), (2, 2));
2407 assert!(row.add(&sq).is_none());
2408 }
2409
2410 #[test]
2411 fn exact_matmul_2x2() {
2412 let p = 64;
2413 let n = |k: u8| ExactNum::from_u8(k, p);
2414 let a = ExactNumArray::from_shape(p, 2, 2, &[n(1), n(2), n(3), n(4)]).unwrap();
2415 let b = ExactNumArray::from_shape(p, 2, 2, &[n(5), n(6), n(7), n(8)]).unwrap();
2416 let c = a.matmul(&b).unwrap();
2417 assert_eq!(c.shape(), (2, 2));
2418 assert_eq!(c.get2(0, 0).unwrap().cmp(&n(19)), Some(0));
2419 assert_eq!(c.get2(0, 1).unwrap().cmp(&n(22)), Some(0));
2420 assert_eq!(c.get2(1, 0).unwrap().cmp(&n(43)), Some(0));
2421 assert_eq!(c.get2(1, 1).unwrap().cmp(&n(50)), Some(0));
2422 }
2423
2424 #[test]
2425 fn bin32_simd_add_mul_golds() {
2426 let one = Ieee32::from_i32(1);
2427 let two = Ieee32::from_i32(2);
2428 let four = Ieee32::from_i32(4);
2429 let a = Ieee32Array::from_values(&[one, two, one, two, one]);
2430 let b = Ieee32Array::from_values(&[one, two, two, one, one]);
2431 let s = a.add(&b).unwrap();
2432 assert_eq!(s.get(0).unwrap().to_bits(), two.to_bits());
2433 assert_eq!(s.get(1).unwrap().to_bits(), four.to_bits());
2434 assert_eq!(s.get(4).unwrap().to_bits(), two.to_bits());
2435 let half = Ieee32::from_bits(0x3F00_0000);
2436 let t = Ieee32Array::filled(4, two);
2437 let h = Ieee32Array::filled(4, half);
2438 let p = t.mul(&h).unwrap();
2439 assert_eq!(p.get(0).unwrap().to_bits(), one.to_bits());
2440 assert_eq!(p.get(3).unwrap().to_bits(), one.to_bits());
2441 }
2442
2443 #[test]
2444 fn ieee64_simd_1000_add_mul_div_sqrt() {
2445 const N: usize = 1000;
2446 let p = 128;
2447 let rm = RoundingMode::ToEven;
2448 let vals: Vec<Ieee64> = (1..=N as i32).map(Ieee64::from_i32).collect();
2449 let ones: Vec<Ieee64> = (0..N).map(|_| Ieee64::from_i32(1)).collect();
2450 let twos: Vec<Ieee64> = (0..N).map(|_| Ieee64::from_i32(2)).collect();
2451 let a = Ieee64Array::from_values(&vals);
2452 let one = Ieee64Array::from_values(&ones);
2453 let two = Ieee64Array::from_values(&twos);
2454 let add = a.add(&one).unwrap();
2455 let mul = a.mul(&two).unwrap();
2456 let div = a.div(&a).unwrap();
2457 let squares = a.mul(&a).unwrap();
2458 let sq = squares.sqrt();
2459 let sub = a.sub(&one).unwrap();
2460 let fma = a.fma(&one, &one).unwrap();
2461 assert_eq!(add.len(), N);
2462 for i in 0..N {
2463 let ai = a.get(i).unwrap();
2464 let oi = one.get(i).unwrap();
2465 let ti = two.get(i).unwrap();
2466 assert_eq!(add.get(i).unwrap().to_bits(), ai.add(oi).to_bits());
2467 assert_eq!(mul.get(i).unwrap().to_bits(), ai.mul(ti).to_bits());
2468 assert_eq!(div.get(i).unwrap().to_bits(), ai.div(ai).to_bits());
2469 assert_eq!(
2470 sq.get(i).unwrap().to_bits(),
2471 ai.mul(ai).sqrt().to_bits()
2472 );
2473 assert_eq!(sub.get(i).unwrap().to_bits(), ai.sub(oi).to_bits());
2474 assert_eq!(
2475 fma.get(i).unwrap().to_bits(),
2476 ai.mul_add(oi, oi).to_bits()
2477 );
2478 let xa = ai.to_exact(p);
2479 let x1 = oi.to_exact(p);
2480 let x2 = ti.to_exact(p);
2481 assert_eq!(
2482 add.get(i).unwrap().to_bits(),
2483 Ieee64::from_exact(&xa.add(&x1, p, rm)).to_bits()
2484 );
2485 assert_eq!(
2486 mul.get(i).unwrap().to_bits(),
2487 Ieee64::from_exact(&xa.mul(&x2, p, rm)).to_bits()
2488 );
2489 assert_eq!(
2490 div.get(i).unwrap().to_bits(),
2491 Ieee64::from_exact(&xa.div(&xa, p, rm)).to_bits()
2492 );
2493 let sqe = xa.mul(&xa, p, rm).sqrt(p, rm);
2494 assert_eq!(sq.get(i).unwrap().to_bits(), Ieee64::from_exact(&sqe).to_bits());
2495 }
2496 assert!(a.fma(&one, &Ieee64Array::from_values(&vals[..10])).is_none());
2497 }
2498
2499 #[test]
2500 fn array_ufunc_identities() {
2501 let mut cc = Consts::new().unwrap();
2502 let z32 = Ieee32Array::from_values(&[Ieee32::ZERO]);
2503 assert!(z32.asin(&mut cc).get(0).unwrap().is_zero());
2504 assert!(z32.expm1(&mut cc).get(0).unwrap().is_zero());
2505 assert!(z32.log1p(&mut cc).get(0).unwrap().is_zero());
2506 assert_eq!(
2507 z32.bessel_j(0, &mut cc).get(0).unwrap().to_bits(),
2508 Ieee32::from_i32(1).to_bits()
2509 );
2510 let one64 = Ieee64Array::from_values(&[Ieee64::from_i32(1)]);
2511 assert!(one64.ln_gamma(&mut cc).get(0).unwrap().is_zero());
2512 let p = 64;
2513 let rm = RoundingMode::ToEven;
2514 let z = ExactNumArray::from_values(p, &[ExactNum::from_u8(0, p)]);
2515 assert!(z.sinh(p, rm, &mut cc).get(0).unwrap().is_zero());
2516 let one = ExactNumArray::from_values(p, &[ExactNum::from_u8(1, p)]);
2517 assert_eq!(
2518 one.ln_gamma(p, rm, &mut cc)
2519 .get(0)
2520 .unwrap()
2521 .cmp(&ExactNum::from_u8(0, p)),
2522 Some(0)
2523 );
2524 let x = ExactNumArray::from_values(p, &[ExactNum::from_u8(3, p)]);
2525 assert_eq!(
2526 x.legendre_p(0, p, rm)
2527 .get(0)
2528 .unwrap()
2529 .cmp(&ExactNum::from_u8(1, p)),
2530 Some(0)
2531 );
2532 assert_eq!(
2533 x.floor().get(0).unwrap().cmp(&ExactNum::from_u8(3, p)),
2534 Some(0)
2535 );
2536 }
2537
2538 #[test]
2539 fn exact_array_sin_2x3_matches_scalar() {
2540 let p = 128;
2541 let rm = RoundingMode::ToEven;
2542 let mut cc = Consts::new().unwrap();
2543 let n = |k: u8| ExactNum::from_u8(k, p);
2544 let vals = [n(1), n(2), n(3), n(4), n(5), n(6)];
2545 let a = ExactNumArray::from_shape(p, 2, 3, &vals).unwrap();
2546 let s = a.sin(p, rm, &mut cc);
2547 assert_eq!(s.shape(), (2, 3));
2548 for i in 0..2 {
2549 for j in 0..3 {
2550 let want = vals[i * 3 + j].sin(p, rm, &mut cc);
2551 assert_eq!(s.get2(i, j).unwrap().cmp(&want), Some(0));
2552 }
2553 }
2554 let row = ExactNumArray::from_values(p, &[n(1), n(2)]);
2555 assert!(a.add(&row).is_none());
2556 }
2557
2558 #[test]
2559 fn exact_array_bessel_j_nu_matches_scalar() {
2560 let p = 128;
2561 let rm = RoundingMode::ToEven;
2562 let mut cc = Consts::new().unwrap();
2563 let half = ExactNum::from_u8(1, p).div(&ExactNum::from_u8(2, p), p, rm);
2564 let xs = [ExactNum::from_u8(1, p), ExactNum::from_u8(2, p), ExactNum::from_u8(3, p)];
2565 let a = ExactNumArray::from_values(p, &xs);
2566 let out = a.bessel_j_nu(&half, p, rm, &mut cc);
2567 for (i, x) in xs.iter().enumerate() {
2568 let want = x.bessel_j_nu(&half, p, rm, &mut cc);
2569 assert_eq!(out.get(i).unwrap().cmp(&want), Some(0));
2570 }
2571 let nan_in = ExactNumArray::from_values(p, &[ExactNum::nan(None)]);
2572 assert!(nan_in.sin(p, rm, &mut cc).get(0).unwrap().is_nan());
2573 }
2574
2575 #[test]
2576 fn exact_array_jacobi_sn_matches_scalar() {
2577 let p = 128;
2578 let rm = RoundingMode::ToEven;
2579 let mut cc = Consts::new().unwrap();
2580 let half = ExactNum::from_u8(1, p).div(&ExactNum::from_u8(2, p), p, rm);
2581 let xs = [
2582 ExactNum::from_u8(0, p),
2583 ExactNum::from_u8(1, p),
2584 ExactNum::from_u8(2, p),
2585 ];
2586 let a = ExactNumArray::from_values(p, &xs);
2587 let out = a.jacobi_sn(&half, p, rm, &mut cc);
2588 let out_ns = a.jacobi_ns(&half, p, rm, &mut cc);
2589 for (i, x) in xs.iter().enumerate() {
2590 let want = x.jacobi_sn(&half, p, rm, &mut cc);
2591 assert_eq!(out.get(i).unwrap().cmp(&want), Some(0));
2592 if !x.is_zero() {
2593 let want_ns = x.jacobi_ns(&half, p, rm, &mut cc);
2594 assert_eq!(out_ns.get(i).unwrap().cmp(&want_ns), Some(0));
2595 }
2596 }
2597 }
2598
2599 fn perm_rows(a: &ExactNumArray, perm: &[usize]) -> ExactNumArray {
2600 let (n, m) = a.shape();
2601 let mut vals = Vec::with_capacity(n * m);
2602 for &r in perm {
2603 for c in 0..m {
2604 vals.push(a.get2(r, c).unwrap().clone());
2605 }
2606 }
2607 ExactNumArray::from_shape(a.precision(), n, m, &vals).unwrap()
2608 }
2609
2610 #[test]
2611 fn exact_lu_2x2_and_singular() {
2612 let p = 256;
2613 let rm = RoundingMode::ToEven;
2614 let n = |k: u8| ExactNum::from_u8(k, p);
2615 let a = ExactNumArray::from_shape(p, 2, 2, &[n(2), n(1), n(4), n(3)]).unwrap();
2616 let (l, u, perm) = a.lu_decomp(p, rm).expect("LU");
2617 let pa = perm_rows(&a, &perm);
2618 let lu = l.matmul(&u).expect("L*U");
2619 assert_eq!(lu.shape(), (2, 2));
2620 for i in 0..2 {
2621 for j in 0..2 {
2622 assert_eq!(
2623 lu.get2(i, j).unwrap().cmp(pa.get2(i, j).unwrap()),
2624 Some(0),
2625 "PA=LU at {i},{j}"
2626 );
2627 }
2628 }
2629 let sing = ExactNumArray::from_shape(p, 2, 2, &[n(1), n(2), n(2), n(4)]).unwrap();
2630 assert!(sing.lu_decomp(p, rm).is_none());
2631 }
2632
2633 fn near_num(a: &ExactNum, b: &ExactNum, p: usize) -> bool {
2634 let d = a.sub(b, p, RoundingMode::None).abs();
2635 d.is_zero() || d.exponent().unwrap_or(0) < -((p as i32) - 40)
2636 }
2637
2638 #[test]
2639 fn exact_qr_recon_orthog_rankdef() {
2640 let p = 256;
2641 let rm = RoundingMode::ToEven;
2642 let n = |k: u8| ExactNum::from_u8(k, p);
2643 let a = ExactNumArray::from_shape(p, 2, 2, &[n(2), n(1), n(4), n(3)]).unwrap();
2644 let (q, r) = a.qr_decomp(p, rm).expect("QR");
2645 let qr = q.matmul(&r).expect("Q*R");
2646 for i in 0..2 {
2647 for j in 0..2 {
2648 assert!(
2649 near_num(qr.get2(i, j).unwrap(), a.get2(i, j).unwrap(), p),
2650 "QR=A at {i},{j}"
2651 );
2652 }
2653 }
2654 let qtq = q.transpose().matmul(&q).expect("Q^T Q");
2655 let one = n(1);
2656 let zero = n(0);
2657 assert!(near_num(qtq.get2(0, 0).unwrap(), &one, p));
2658 assert!(near_num(qtq.get2(1, 1).unwrap(), &one, p));
2659 assert!(near_num(qtq.get2(0, 1).unwrap(), &zero, p));
2660 assert!(near_num(qtq.get2(1, 0).unwrap(), &zero, p));
2661
2662 let def = ExactNumArray::from_shape(p, 2, 2, &[n(1), n(2), n(2), n(4)]).unwrap();
2663 let (qd, rd) = def.qr_decomp(p, rm).expect("rank-def QR");
2664 let _ = qd;
2665 assert!(rd.get2(1, 1).unwrap().is_zero() || near_num(rd.get2(1, 1).unwrap(), &zero, p));
2666 let recon = qd.matmul(&rd).expect("Qd Rd");
2667 for i in 0..2 {
2668 for j in 0..2 {
2669 assert!(near_num(
2670 recon.get2(i, j).unwrap(),
2671 def.get2(i, j).unwrap(),
2672 p
2673 ));
2674 }
2675 }
2676 }
2677
2678 #[test]
2679 fn exact_svd_diag_3_2() {
2680 let p = 256;
2681 let rm = RoundingMode::ToEven;
2682 let n = |k: u8| ExactNum::from_u8(k, p);
2683 let a = ExactNumArray::from_shape(p, 2, 2, &[n(3), n(0), n(0), n(2)]).unwrap();
2684 let (u, s, vt) = a.svd_decomp(p, rm).expect("SVD");
2685 assert_eq!(s.get2(0, 0).unwrap().cmp(&n(3)), Some(0));
2686 assert_eq!(s.get2(1, 1).unwrap().cmp(&n(2)), Some(0));
2687 assert!(near_num(s.get2(0, 1).unwrap(), &n(0), p));
2688 assert!(near_num(s.get2(1, 0).unwrap(), &n(0), p));
2689 let us = u.matmul(&s).expect("U Σ");
2690 let recon = us.matmul(&vt).expect("U Σ V^T");
2691 for i in 0..2 {
2692 for j in 0..2 {
2693 assert!(
2694 near_num(recon.get2(i, j).unwrap(), a.get2(i, j).unwrap(), p),
2695 "UΣV^T=A at {i},{j}"
2696 );
2697 }
2698 }
2699 }
2700
2701 #[test]
2702 fn exact_svd_recon_orthog() {
2703 let p = 256;
2704 let rm = RoundingMode::ToEven;
2705 let n = |k: u8| ExactNum::from_u8(k, p);
2706 let a = ExactNumArray::from_shape(p, 2, 2, &[n(2), n(1), n(4), n(3)]).unwrap();
2707 let (u, s, vt) = a.svd_decomp(p, rm).expect("SVD");
2708 let us = u.matmul(&s).expect("U Σ");
2709 let recon = us.matmul(&vt).expect("U Σ V^T");
2710 for i in 0..2 {
2711 for j in 0..2 {
2712 assert!(
2713 near_num(recon.get2(i, j).unwrap(), a.get2(i, j).unwrap(), p),
2714 "UΣV^T=A at {i},{j}"
2715 );
2716 }
2717 }
2718 let utu = u.transpose().matmul(&u).expect("U^T U");
2719 let v = vt.transpose();
2720 let vtv = vt.matmul(&v).expect("V^T V");
2721 let one = n(1);
2722 let zero = n(0);
2723 for (name, g) in [("U^T U", &utu), ("V^T V", &vtv)] {
2724 assert!(near_num(g.get2(0, 0).unwrap(), &one, p), "{name}[0,0]");
2725 assert!(near_num(g.get2(1, 1).unwrap(), &one, p), "{name}[1,1]");
2726 assert!(near_num(g.get2(0, 1).unwrap(), &zero, p), "{name}[0,1]");
2727 assert!(near_num(g.get2(1, 0).unwrap(), &zero, p), "{name}[1,0]");
2728 }
2729 assert!(ExactNumArray::from_shape(p, 0, 0, &[])
2730 .unwrap()
2731 .svd_decomp(p, rm)
2732 .is_none());
2733
2734 let wide =
2735 ExactNumArray::from_shape(p, 2, 3, &[n(1), n(0), n(0), n(0), n(2), n(0)]).unwrap();
2736 let (uw, sw, vtw) = wide.svd_decomp(p, rm).expect("wide SVD");
2737 assert_eq!(sw.get2(0, 0).unwrap().cmp(&n(2)), Some(0));
2738 assert_eq!(sw.get2(1, 1).unwrap().cmp(&n(1)), Some(0));
2739 let usw = uw.matmul(&sw).expect("Uw Σw");
2740 let recw = usw.matmul(&vtw).expect("wide recon");
2741 for i in 0..2 {
2742 for j in 0..3 {
2743 assert!(near_num(
2744 recw.get2(i, j).unwrap(),
2745 wide.get2(i, j).unwrap(),
2746 p
2747 ));
2748 }
2749 }
2750 }
2751
2752 fn eigen_diag(evals: &ExactNumArray, p: usize) -> ExactNumArray {
2753 let n = evals.cols;
2754 let mut vals = Vec::with_capacity(n * n);
2755 let z = ExactNum::from_u8(0, p);
2756 for i in 0..n {
2757 for j in 0..n {
2758 if i == j {
2759 vals.push(evals.get(i).unwrap().clone());
2760 } else {
2761 vals.push(z.clone());
2762 }
2763 }
2764 }
2765 ExactNumArray::from_shape(p, n, n, &vals).unwrap()
2766 }
2767
2768 #[test]
2769 fn exact_eigen_sym_2x2() {
2770 let p = 256;
2771 let rm = RoundingMode::ToEven;
2772 let n = |k: u8| ExactNum::from_u8(k, p);
2773 let a = ExactNumArray::from_shape(p, 2, 2, &[n(2), n(1), n(1), n(2)]).unwrap();
2774 let (evals, v) = a.eigen_decomp(p, rm).expect("eigen");
2775 assert_eq!(evals.shape(), (1, 2));
2776 assert!(near_num(evals.get(0).unwrap(), &n(3), p));
2777 assert!(near_num(evals.get(1).unwrap(), &n(1), p));
2778 let vtv = v.transpose().matmul(&v).expect("V^T V");
2779 let one = n(1);
2780 let zero = n(0);
2781 assert!(near_num(vtv.get2(0, 0).unwrap(), &one, p));
2782 assert!(near_num(vtv.get2(1, 1).unwrap(), &one, p));
2783 assert!(near_num(vtv.get2(0, 1).unwrap(), &zero, p));
2784 assert!(near_num(vtv.get2(1, 0).unwrap(), &zero, p));
2785 let av = a.matmul(&v).expect("A V");
2786 let lam = eigen_diag(&evals, p);
2787 let vl = v.matmul(&lam).expect("V Λ");
2788 for i in 0..2 {
2789 for j in 0..2 {
2790 assert!(
2791 near_num(av.get2(i, j).unwrap(), vl.get2(i, j).unwrap(), p),
2792 "Av=λv at {i},{j}"
2793 );
2794 }
2795 }
2796 let vlvt = vl.matmul(&v.transpose()).expect("V Λ V^T");
2797 for i in 0..2 {
2798 for j in 0..2 {
2799 assert!(near_num(vlvt.get2(i, j).unwrap(), a.get2(i, j).unwrap(), p));
2800 }
2801 }
2802 let nosym = ExactNumArray::from_shape(p, 2, 2, &[n(1), n(2), n(0), n(1)]).unwrap();
2803 assert!(nosym.eigen_decomp(p, rm).is_none());
2804
2805 let a3 = ExactNumArray::from_shape(
2806 p,
2807 3,
2808 3,
2809 &[n(2), n(1), n(0), n(1), n(2), n(1), n(0), n(1), n(2)],
2810 )
2811 .unwrap();
2812 let (w3, v3) = a3.eigen_decomp(p, rm).expect("eigen 3");
2813 let s2 = n(2).sqrt(p, rm);
2814 let want = [n(2).add(&s2, p, rm), n(2), n(2).sub(&s2, p, rm)];
2815 for (i, wi) in want.iter().enumerate() {
2816 assert!(near_num(w3.get(i).unwrap(), wi, p), "λ[{i}]");
2817 }
2818 let av3 = a3.matmul(&v3).expect("A3 V");
2819 let vl3 = v3.matmul(&eigen_diag(&w3, p)).expect("V3 Λ");
2820 for i in 0..3 {
2821 for j in 0..3 {
2822 assert!(near_num(
2823 av3.get2(i, j).unwrap(),
2824 vl3.get2(i, j).unwrap(),
2825 p
2826 ));
2827 }
2828 }
2829 }
2830
2831 #[test]
2832 fn exact_fft_impulse_cosine_parseval() {
2833 let p = 256;
2834 let rm = RoundingMode::ToEven;
2835 let mut cc = Consts::new().unwrap();
2836 let n = |k: u8| ExactNum::from_u8(k, p);
2837 let impulse = ExactNumArray::from_values(p, &[n(1), n(0), n(0), n(0)]);
2838 let spec = impulse.fft(p, rm, &mut cc).expect("FFT impulse");
2839 assert_eq!(spec.shape(), (2, 4));
2840 for j in 0..4 {
2841 assert!(near_num(spec.get2(0, j).unwrap(), &n(1), p), "re[{j}]");
2842 assert!(near_num(spec.get2(1, j).unwrap(), &n(0), p), "im[{j}]");
2843 }
2844
2845 let n8 = ExactNum::from_u8(8, p);
2846 let two_pi = n(2).mul(&cc.pi(p, rm), p, rm);
2847 let mut cos_vals = Vec::with_capacity(8);
2848 for k in 0..8u8 {
2849 let kn = ExactNum::from_u8(k, p);
2850 let ang = two_pi.mul(&kn, p, rm).div(&n8, p, rm);
2851 cos_vals.push(ang.cos(p, rm, &mut cc));
2852 }
2853 let cosine = ExactNumArray::from_values(p, &cos_vals);
2854 let cspec = cosine.fft(p, rm, &mut cc).expect("FFT cos");
2855 let four = n(4);
2856 let zero = n(0);
2857 for j in 0..8 {
2858 let re = cspec.get2(0, j).unwrap();
2859 let im = cspec.get2(1, j).unwrap();
2860 if j == 1 || j == 7 {
2861 assert!(near_num(re, &four, p), "cos bin {j} re");
2862 } else {
2863 assert!(near_num(re, &zero, p), "cos bin {j} re");
2864 }
2865 assert!(near_num(im, &zero, p), "cos bin {j} im");
2866 }
2867
2868 let back = cspec.ifft(p, rm, &mut cc).expect("IFFT");
2869 for j in 0..8 {
2870 assert!(near_num(
2871 back.get2(0, j).unwrap(),
2872 cosine.get(j).unwrap(),
2873 p
2874 ));
2875 assert!(near_num(back.get2(1, j).unwrap(), &zero, p));
2876 }
2877
2878 let mut e_t = ExactNum::from_u8(0, p);
2879 let mut e_f = ExactNum::from_u8(0, p);
2880 for j in 0..8 {
2881 let x = cosine.get(j).unwrap();
2882 e_t = e_t.add(&x.mul(x, p, rm), p, rm);
2883 let xr = cspec.get2(0, j).unwrap();
2884 let xi = cspec.get2(1, j).unwrap();
2885 e_f = e_f
2886 .add(&xr.mul(xr, p, rm), p, rm)
2887 .add(&xi.mul(xi, p, rm), p, rm);
2888 }
2889 let parseval = e_f.div(&n8, p, rm);
2890 assert!(near_num(&parseval, &e_t, p));
2891 assert!(ExactNumArray::from_values(p, &[n(1), n(2), n(3)])
2892 .fft(p, rm, &mut cc)
2893 .is_none());
2894 }
2895}