use crate::dd::DD;
use crate::dual::{Du, Lane, Ou, Qu};
use crate::dv::DV;
use crate::field::{Fl, MatLane};
use crate::jet::{Jet, OJet, QJet};
use crate::linalg::{DDLu, ODLu, QDLu};
use crate::od::{OdAcc, OD};
use crate::ov::OV;
use crate::qd::{QdAcc, QD};
use crate::qv::QV;
use crate::simd::Simd;
use std::fmt::Debug;
use std::ops::{Add, Mul, Neg, Sub};
pub trait Precision<S: Simd>: Copy + Send + Sync + 'static {
const PARTS: usize;
type Scalar: Copy
+ Send
+ Sync
+ Debug
+ PartialEq
+ Add<Output = Self::Scalar>
+ Sub<Output = Self::Scalar>
+ Mul<Output = Self::Scalar>
+ Neg<Output = Self::Scalar>;
type V: MatLane<S = S>;
type U: MatLane<S = S>;
type J<const N: usize>: MatLane<S = S>;
type Lu: Send + Sync;
const ZERO: Self::Scalar;
const ONE: Self::Scalar;
fn new(p: &[f64]) -> Self::Scalar;
fn from_raw(p: &[f64]) -> Self::Scalar;
fn part(x: Self::Scalar, q: usize) -> f64;
fn from_dd(d: DD) -> Self::Scalar;
fn to_dd(x: Self::Scalar) -> DD;
fn dot(pairs: impl Iterator<Item = (Self::Scalar, Self::Scalar)>) -> Self::Scalar;
fn splat(x: Self::Scalar) -> Self::V;
fn dual(v: Self::V, d: Self::V) -> Self::U;
fn coef<L: Lane<S = S>>(l: L, k: usize) -> Self::V;
fn set_coef<L: Lane<S = S>>(l: &mut L, k: usize, c: Self::V);
fn load_field<const B: usize>(parts: &[&[f64]]) -> Fl<Self::V, B>;
fn field_get<const B: usize>(f: &Fl<Self::V, B>, j: usize) -> Self::Scalar;
fn field_coef<L: Lane<S = S>, const B: usize>(f: &Fl<L, B>, k: usize) -> Fl<Self::V, B>;
fn field_set_coef<L: Lane<S = S>, const B: usize>(f: &mut Fl<L, B>, k: usize, c: &Fl<Self::V, B>);
fn field_dual<const B: usize>(v: &Fl<Self::V, B>, t: &Fl<Self::V, B>) -> Fl<Self::U, B>;
fn lu_factor(n: usize, ld: usize, c: Vec<Vec<f64>>, par: bool) -> Result<Self::Lu, String>;
fn lu_solve(lu: &Self::Lu, b: &mut [Self::Scalar]);
}
#[derive(Clone, Copy, Debug)]
pub struct Double;
#[derive(Clone, Copy, Debug)]
pub struct Quad;
#[derive(Clone, Copy, Debug)]
pub struct Octo;
fn parts_of<const P: usize>(p: &[f64], what: &str) -> [f64; P] {
assert!(p.len() <= P, "{what}: at most {P} parts, not {}", p.len());
std::array::from_fn(|q| if q < p.len() { p[q] } else { 0.0 })
}
impl<S: Simd> Precision<S> for Double {
const PARTS: usize = 2;
type Scalar = DD;
type V = DV<S>;
type U = Du<S>;
type J<const N: usize> = Jet<S, N>;
type Lu = DDLu;
const ZERO: DD = DD::ZERO;
const ONE: DD = DD::ONE;
#[inline]
fn new(p: &[f64]) -> DD {
let x: [f64; 2] = parts_of(p, "Double::new");
DD::new(x[0], x[1])
}
#[inline]
fn from_raw(p: &[f64]) -> DD {
assert_eq!(p.len(), 2, "Double::from_raw: 2 parts");
DD { hi: p[0], lo: p[1] }
}
#[inline]
fn part(x: DD, q: usize) -> f64 {
[x.hi, x.lo][q]
}
#[inline]
fn from_dd(d: DD) -> DD {
d
}
#[inline]
fn to_dd(x: DD) -> DD {
x
}
fn dot(pairs: impl Iterator<Item = (DD, DD)>) -> DD {
let mut acc = DD::ZERO;
for (a, b) in pairs {
acc = acc + a * b;
}
acc
}
#[inline]
fn splat(x: DD) -> DV<S> {
DV::splat(x)
}
#[inline]
fn dual(v: DV<S>, d: DV<S>) -> Du<S> {
Du { v, d }
}
#[inline]
fn coef<L: Lane<S = S>>(l: L, k: usize) -> DV<S> {
l.coef(k)
}
#[inline]
fn set_coef<L: Lane<S = S>>(l: &mut L, k: usize, c: DV<S>) {
l.set_coef(k, c)
}
fn load_field<const B: usize>(parts: &[&[f64]]) -> Fl<DV<S>, B> {
assert_eq!(parts.len(), 2, "Double::load_field: 2 parts");
Fl::load(parts[0], parts[1])
}
#[inline]
fn field_get<const B: usize>(f: &Fl<DV<S>, B>, j: usize) -> DD {
f.get(j)
}
#[inline]
fn field_coef<L: Lane<S = S>, const B: usize>(f: &Fl<L, B>, k: usize) -> Fl<DV<S>, B> {
f.coef(k)
}
#[inline]
fn field_set_coef<L: Lane<S = S>, const B: usize>(f: &mut Fl<L, B>, k: usize, c: &Fl<DV<S>, B>) {
f.set_coef(k, c)
}
#[inline]
fn field_dual<const B: usize>(v: &Fl<DV<S>, B>, t: &Fl<DV<S>, B>) -> Fl<Du<S>, B> {
Fl::<Du<S>, B>::from_vt(v, t)
}
fn lu_factor(n: usize, ld: usize, c: Vec<Vec<f64>>, par: bool) -> Result<DDLu, String> {
let [hi, lo]: [Vec<f64>; 2] = c
.try_into()
.map_err(|c: Vec<Vec<f64>>| format!("Double::lu_factor: 2 parts, not {}", c.len()))?;
DDLu::factor(n, ld, hi, lo, par)
}
#[inline]
fn lu_solve(lu: &DDLu, b: &mut [DD]) {
lu.solve(b)
}
}
impl<S: Simd> Precision<S> for Quad {
const PARTS: usize = 4;
type Scalar = QD;
type V = QV<S>;
type U = Qu<S>;
type J<const N: usize> = QJet<S, N>;
type Lu = QDLu;
const ZERO: QD = QD::ZERO;
const ONE: QD = QD::ONE;
#[inline]
fn new(p: &[f64]) -> QD {
let x: [f64; 4] = parts_of(p, "Quad::new");
QD::new(x[0], x[1], x[2], x[3])
}
#[inline]
fn from_raw(p: &[f64]) -> QD {
assert_eq!(p.len(), 4, "Quad::from_raw: 4 parts");
QD {
x: [p[0], p[1], p[2], p[3]],
}
}
#[inline]
fn part(x: QD, q: usize) -> f64 {
x.x[q]
}
#[inline]
fn from_dd(d: DD) -> QD {
QD::from_dd(d)
}
#[inline]
fn to_dd(x: QD) -> DD {
x.to_dd()
}
fn dot(pairs: impl Iterator<Item = (QD, QD)>) -> QD {
let mut acc = QdAcc::zero();
for (a, b) in pairs {
acc = acc.add_prod(a, b);
}
acc.finish()
}
#[inline]
fn splat(x: QD) -> QV<S> {
QV::splat(x)
}
#[inline]
fn dual(v: QV<S>, d: QV<S>) -> Qu<S> {
Qu { v, d }
}
#[inline]
fn coef<L: Lane<S = S>>(l: L, k: usize) -> QV<S> {
l.coef_qd(k)
}
#[inline]
fn set_coef<L: Lane<S = S>>(l: &mut L, k: usize, c: QV<S>) {
l.set_coef_qd(k, c)
}
fn load_field<const B: usize>(parts: &[&[f64]]) -> Fl<QV<S>, B> {
assert_eq!(parts.len(), 4, "Quad::load_field: 4 parts");
Fl::load_parts(parts[0], parts[1], parts[2], parts[3])
}
#[inline]
fn field_get<const B: usize>(f: &Fl<QV<S>, B>, j: usize) -> QD {
f.get(j)
}
#[inline]
fn field_coef<L: Lane<S = S>, const B: usize>(f: &Fl<L, B>, k: usize) -> Fl<QV<S>, B> {
f.coef_qd(k)
}
#[inline]
fn field_set_coef<L: Lane<S = S>, const B: usize>(f: &mut Fl<L, B>, k: usize, c: &Fl<QV<S>, B>) {
f.set_coef_qd(k, c)
}
#[inline]
fn field_dual<const B: usize>(v: &Fl<QV<S>, B>, t: &Fl<QV<S>, B>) -> Fl<Qu<S>, B> {
Fl::<Qu<S>, B>::from_vt_qd(v, t)
}
fn lu_factor(n: usize, ld: usize, c: Vec<Vec<f64>>, par: bool) -> Result<QDLu, String> {
let c: [Vec<f64>; 4] = c
.try_into()
.map_err(|c: Vec<Vec<f64>>| format!("Quad::lu_factor: 4 parts, not {}", c.len()))?;
QDLu::factor(n, ld, c, par)
}
#[inline]
fn lu_solve(lu: &QDLu, b: &mut [QD]) {
lu.solve(b)
}
}
impl<S: Simd> Precision<S> for Octo {
const PARTS: usize = 8;
type Scalar = OD;
type V = OV<S>;
type U = Ou<S>;
type J<const N: usize> = OJet<S, N>;
type Lu = ODLu;
const ZERO: OD = OD::ZERO;
const ONE: OD = OD::ONE;
#[inline]
fn new(p: &[f64]) -> OD {
OD::new(parts_of(p, "Octo::new"))
}
#[inline]
fn from_raw(p: &[f64]) -> OD {
assert_eq!(p.len(), 8, "Octo::from_raw: 8 parts");
OD {
x: std::array::from_fn(|q| p[q]),
}
}
#[inline]
fn part(x: OD, q: usize) -> f64 {
x.x[q]
}
#[inline]
fn from_dd(d: DD) -> OD {
OD::from_dd(d)
}
#[inline]
fn to_dd(x: OD) -> DD {
x.to_dd()
}
fn dot(pairs: impl Iterator<Item = (OD, OD)>) -> OD {
let mut acc = OdAcc::zero();
for (a, b) in pairs {
acc = acc.add_prod(a, b);
}
acc.finish()
}
#[inline]
fn splat(x: OD) -> OV<S> {
OV::splat(x)
}
#[inline]
fn dual(v: OV<S>, d: OV<S>) -> Ou<S> {
Ou { v, d }
}
#[inline]
fn coef<L: Lane<S = S>>(l: L, k: usize) -> OV<S> {
l.coef_od(k)
}
#[inline]
fn set_coef<L: Lane<S = S>>(l: &mut L, k: usize, c: OV<S>) {
l.set_coef_od(k, c)
}
fn load_field<const B: usize>(parts: &[&[f64]]) -> Fl<OV<S>, B> {
assert_eq!(parts.len(), 8, "Octo::load_field: 8 parts");
Fl::load_parts_od(std::array::from_fn(|q| parts[q]))
}
#[inline]
fn field_get<const B: usize>(f: &Fl<OV<S>, B>, j: usize) -> OD {
f.get(j)
}
#[inline]
fn field_coef<L: Lane<S = S>, const B: usize>(f: &Fl<L, B>, k: usize) -> Fl<OV<S>, B> {
f.coef_od(k)
}
#[inline]
fn field_set_coef<L: Lane<S = S>, const B: usize>(f: &mut Fl<L, B>, k: usize, c: &Fl<OV<S>, B>) {
f.set_coef_od(k, c)
}
#[inline]
fn field_dual<const B: usize>(v: &Fl<OV<S>, B>, t: &Fl<OV<S>, B>) -> Fl<Ou<S>, B> {
Fl::<Ou<S>, B>::from_vt_od(v, t)
}
fn lu_factor(n: usize, ld: usize, c: Vec<Vec<f64>>, par: bool) -> Result<ODLu, String> {
let c: [Vec<f64>; 8] = c
.try_into()
.map_err(|c: Vec<Vec<f64>>| format!("Octo::lu_factor: 8 parts, not {}", c.len()))?;
ODLu::factor(n, ld, c, par)
}
#[inline]
fn lu_solve(lu: &ODLu, b: &mut [OD]) {
lu.solve(b)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::simd::Native;
type N = Native;
fn qd_parts(seed: u64) -> [f64; 4] {
let a = 1.0 + (seed as f64) / 7.0;
let x = QD::from_dd(DD::from_f64(a)) / QD::from_dd(DD::from_f64(3.0 + seed as f64));
x.x
}
fn od_parts(seed: u64) -> [f64; 8] {
let a = 1.0 + (seed as f64) / 7.0;
let x = OD::from_dd(DD::from_f64(a)) / OD::from_dd(DD::from_f64(3.0 + seed as f64));
x.x
}
#[test]
fn precision_scalars_match_their_types() {
for s in 0..20 {
let p = qd_parts(s);
assert_eq!(<Quad as Precision<N>>::new(&p), QD::new(p[0], p[1], p[2], p[3]));
assert_eq!(<Quad as Precision<N>>::new(&p[..2]), QD::new(p[0], p[1], 0.0, 0.0));
assert_eq!(<Quad as Precision<N>>::from_raw(&p).x, p);
let o = od_parts(s);
assert_eq!(<Octo as Precision<N>>::new(&o), OD::new(o));
let mut half = o;
half[4..].fill(0.0);
assert_eq!(<Octo as Precision<N>>::new(&o[..4]), OD::new(half));
assert_eq!(<Octo as Precision<N>>::from_raw(&o).x, o);
for q in 0..8 {
assert_eq!(<Octo as Precision<N>>::part(OD { x: o }, q), o[q]);
}
}
}
#[test]
fn precision_dot_matches_accumulators() {
let xs: Vec<QD> = (0..13).map(|s| QD { x: qd_parts(s) }).collect();
let mut acc = QdAcc::zero();
for w in xs.windows(2) {
acc = acc.add_prod(w[0], w[1]);
}
assert_eq!(
<Quad as Precision<N>>::dot(xs.windows(2).map(|w| (w[0], w[1]))),
acc.finish()
);
let os: Vec<OD> = (0..13).map(|s| OD { x: od_parts(s) }).collect();
let mut acc = OdAcc::zero();
for w in os.windows(2) {
acc = acc.add_prod(w[0], w[1]);
}
assert_eq!(
<Octo as Precision<N>>::dot(os.windows(2).map(|w| (w[0], w[1]))),
acc.finish()
);
}
#[test]
fn precision_fields_match_their_types() {
const B: usize = 2;
let cols: Vec<Vec<f64>> = (0..8)
.map(|q| (0..8 * B).map(|i| od_parts(i as u64)[q]).collect())
.collect();
let refs: Vec<&[f64]> = cols.iter().map(|c| c.as_slice()).collect();
let fq = <Quad as Precision<N>>::load_field::<B>(&refs[..4]);
let gq = Fl::<QV<N>, B>::load_parts(refs[0], refs[1], refs[2], refs[3]);
let fo = <Octo as Precision<N>>::load_field::<B>(&refs);
let go = Fl::<OV<N>, B>::load_parts_od(std::array::from_fn(|q| refs[q]));
let unit_q = Fl::<QV<N>, B>::unit(5);
let unit_o = Fl::<OV<N>, B>::unit(5);
let dq = <Quad as Precision<N>>::field_dual(&fq, &unit_q);
let doo = <Octo as Precision<N>>::field_dual(&fo, &unit_o);
for j in 0..8 * B {
assert_eq!(<Quad as Precision<N>>::field_get(&fq, j), gq.get(j));
assert_eq!(<Octo as Precision<N>>::field_get(&fo, j), go.get(j));
assert_eq!(
<Quad as Precision<N>>::field_coef(&dq, 1).get(j),
Fl::<Qu<N>, B>::from_vt_qd(&gq, &unit_q).coef_qd(1).get(j)
);
assert_eq!(<Octo as Precision<N>>::field_coef(&doo, 0).get(j), go.get(j));
assert_eq!(
<Octo as Precision<N>>::field_coef(&doo, 1).get(j),
if j == 5 { OD::ONE } else { OD::ZERO }
);
}
let mut jq = Fl::<QJet<N, 4>, B>::zero();
<Quad as Precision<N>>::field_set_coef(&mut jq, 2, &gq);
let mut jo = Fl::<OJet<N, 4>, B>::zero();
<Octo as Precision<N>>::field_set_coef(&mut jo, 2, &go);
for j in 0..8 * B {
assert_eq!(<Quad as Precision<N>>::field_coef(&jq, 2).get(j), gq.get(j));
assert_eq!(<Octo as Precision<N>>::field_coef(&jo, 2).get(j), go.get(j));
}
}
#[test]
fn precision_lu_matches_its_type() {
let n = 5;
let ld = 8; let mat = |p: usize, s: u64| -> Vec<Vec<f64>> {
(0..p)
.map(|q| {
(0..ld * n)
.map(|i| {
if q == 0 {
((i * 7 + 3) % 11) as f64 - 5.0 + if i % ld == i / ld { 9.0 } else { 0.0 }
} else {
od_parts(s + i as u64)[q] * 1e-20
}
})
.collect()
})
.collect()
};
let cq = mat(4, 1);
let mut bq: Vec<QD> = (0..n).map(|i| QD { x: qd_parts(i as u64) }).collect();
let mut bq2 = bq.clone();
let lu = <Quad as Precision<N>>::lu_factor(n, ld, cq.clone(), false).unwrap();
<Quad as Precision<N>>::lu_solve(&lu, &mut bq);
QDLu::factor(n, ld, cq.try_into().unwrap(), false)
.unwrap()
.solve(&mut bq2);
assert_eq!(bq, bq2);
let co = mat(8, 2);
let mut bo: Vec<OD> = (0..n).map(|i| OD { x: od_parts(i as u64) }).collect();
let mut bo2 = bo.clone();
let lu = <Octo as Precision<N>>::lu_factor(n, ld, co.clone(), false).unwrap();
<Octo as Precision<N>>::lu_solve(&lu, &mut bo);
ODLu::factor(n, ld, co.try_into().unwrap(), false)
.unwrap()
.solve(&mut bo2);
assert_eq!(bo, bo2);
assert!(<Quad as Precision<N>>::lu_factor(n, ld, mat(8, 1), false).is_err());
}
#[test]
fn precision_double_matches_its_types() {
const B: usize = 2;
type D = Double;
for s in 0..20 {
let q = qd_parts(s);
assert_eq!(<D as Precision<N>>::new(&q[..2]), DD::new(q[0], q[1]));
assert_eq!(<D as Precision<N>>::new(&q[..1]), DD::new(q[0], 0.0));
let r = <D as Precision<N>>::from_raw(&q[..2]);
assert_eq!((r.hi, r.lo), (q[0], q[1]));
assert_eq!(
(<D as Precision<N>>::part(r, 0), <D as Precision<N>>::part(r, 1)),
(q[0], q[1])
);
}
let xs: Vec<DD> = (0..13).map(|s| DD::new(qd_parts(s)[0], qd_parts(s)[1])).collect();
let mut acc = DD::ZERO;
for w in xs.windows(2) {
acc = acc + w[0] * w[1];
}
assert_eq!(<D as Precision<N>>::dot(xs.windows(2).map(|w| (w[0], w[1]))), acc);
let cols: Vec<Vec<f64>> = (0..2)
.map(|q| (0..8 * B).map(|i| qd_parts(i as u64)[q]).collect())
.collect();
let f = <D as Precision<N>>::load_field::<B>(&[&cols[0], &cols[1]]);
let g = Fl::<DV<N>, B>::load(&cols[0], &cols[1]);
let unit = Fl::<DV<N>, B>::unit(3);
let du = <D as Precision<N>>::field_dual(&f, &unit);
let mut jet = Fl::<Jet<N, 4>, B>::zero();
<D as Precision<N>>::field_set_coef(&mut jet, 2, &g);
for j in 0..8 * B {
assert_eq!(<D as Precision<N>>::field_get(&f, j), g.get(j));
assert_eq!(<D as Precision<N>>::field_coef(&du, 0).get(j), g.get(j));
assert_eq!(
<D as Precision<N>>::field_coef(&du, 1).get(j),
if j == 3 { DD::ONE } else { DD::ZERO }
);
assert_eq!(<D as Precision<N>>::field_coef(&jet, 2).get(j), g.get(j));
}
let (n, ld) = (5, 8);
let hi: Vec<f64> = (0..ld * n)
.map(|i| ((i * 7 + 3) % 11) as f64 - 5.0 + if i % ld == i / ld { 9.0 } else { 0.0 })
.collect();
let lo: Vec<f64> = (0..ld * n).map(|i| qd_parts(i as u64)[1] * 1e-3).collect();
let mut b1: Vec<DD> = (0..n)
.map(|i| DD::new(qd_parts(i as u64)[0], qd_parts(i as u64)[1]))
.collect();
let mut b2 = b1.clone();
let lu = <D as Precision<N>>::lu_factor(n, ld, vec![hi.clone(), lo.clone()], false).unwrap();
<D as Precision<N>>::lu_solve(&lu, &mut b1);
DDLu::factor(n, ld, hi, lo, false).unwrap().solve(&mut b2);
assert_eq!(b1, b2);
assert!(<D as Precision<N>>::lu_factor(n, ld, vec![vec![0.0; ld * n]; 4], false).is_err());
}
#[test]
fn precision_mutated_differs() {
let p = qd_parts(3);
let swapped = [p[1], p[0], p[2], p[3]];
assert_ne!(<Quad as Precision<N>>::from_raw(&swapped), QD { x: p });
let o = od_parts(3);
let mut so = o;
so.swap(0, 1);
assert_ne!(<Octo as Precision<N>>::from_raw(&so), OD { x: o });
assert_ne!(
<Double as Precision<N>>::from_raw(&[p[1], p[0]]),
DD { hi: p[0], lo: p[1] }
);
}
}