Skip to main content

a/math/matrix/
element.rs

1//! The types a matrix can hold.
2
3use super::blas::{Op, gemm};
4use crate::math::complex::{Complex, Float};
5use crate::math::numbers::Number;
6use core::fmt;
7use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
8
9/// What the arithmetic of a matrix needs of its elements, sealed.
10pub(crate) mod sealed {
11    use super::*;
12
13    /// The ring operations, by reference so that a [`Number`] need not be copied.
14    pub trait Element: Clone + PartialEq + fmt::Debug {
15        /// 0.
16        fn zero() -> Self;
17        /// 1.
18        fn one() -> Self;
19        /// Whether it is 0.
20        fn is_zero(&self) -> bool;
21        /// The complex conjugate; the element itself when it is real.
22        fn conj(&self) -> Self;
23        /// `a + b`.
24        fn add(a: &Self, b: &Self) -> Self;
25        /// `a - b`.
26        fn sub(a: &Self, b: &Self) -> Self;
27        /// `a · b`.
28        fn mul(a: &Self, b: &Self) -> Self;
29        /// `-a`.
30        fn neg(a: &Self) -> Self;
31        /// `*a += b`.
32        fn add_assign(a: &mut Self, b: &Self);
33        /// `*a -= b`.
34        fn sub_assign(a: &mut Self, b: &Self);
35        /// `*a *= b`.
36        fn mul_assign(a: &mut Self, b: &Self);
37        /// `*c += a · b`.
38        fn mul_add(c: &mut Self, a: &Self, b: &Self);
39
40        /// `c = a · b` for column-major `a` of `m × k`, `b` of `k × n` and `c` of `m × n`, which
41        /// holds zeros when it is called.
42        fn product(c: &mut [Self], a: &[Self], b: &[Self], m: usize, k: usize, n: usize) {
43            for j in 0..n {
44                let c = &mut c[j * m..(j + 1) * m];
45                for p in 0..k {
46                    let x = &b[p + j * k];
47                    if x.is_zero() {
48                        continue;
49                    }
50                    for (c, a) in c.iter_mut().zip(&a[p * m..(p + 1) * m]) {
51                        Self::mul_add(c, a, x);
52                    }
53                }
54            }
55        }
56    }
57
58    /// The field operations of floats and complex numbers, by value.
59    pub trait Scalar:
60        Element
61        + Copy
62        + Default
63        + Add<Output = Self>
64        + Sub<Output = Self>
65        + Mul<Output = Self>
66        + Div<Output = Self>
67        + Neg<Output = Self>
68        + AddAssign
69        + SubAssign
70        + MulAssign
71        + DivAssign
72        + Send
73        + Sync
74        + 'static
75    {
76        /// The real type of the absolute values, the public trait's `Real`.
77        type R: Float;
78        /// The complex type of the same precision, which eigenvalues of real matrices need.
79        type C: super::Scalar<Real = Self::R>;
80        /// The real type of the same precision as an element, for the real arithmetic inside
81        /// complex algorithms.
82        type Re: super::Scalar<Real = Self::R>;
83        /// Whether the type is complex.
84        const IS_COMPLEX: bool;
85        /// The real part.
86        fn re(self) -> Self::R;
87        /// The imaginary part.
88        fn im(self) -> Self::R;
89        /// The number with the real part `re`.
90        fn from_real(re: Self::R) -> Self;
91        /// The number `re + im·i`, or `re` when the type is real.
92        fn from_parts(re: Self::R, im: Self::R) -> Self;
93        /// The absolute value.
94        fn abs(self) -> Self::R;
95        /// The square of the absolute value.
96        fn abs2(self) -> Self::R;
97        /// `|re| + |im|`, cheaper than the absolute value and within a factor √2 of it.
98        fn abs1(self) -> Self::R;
99        /// The product by a real number.
100        fn scale(self, r: Self::R) -> Self;
101        /// `self - a·b`, with fused multiply-adds: one for a real number, four for a complex one.
102        fn mul_sub(self, a: Self, b: Self) -> Self;
103        /// The principal square root.
104        fn sqrt(self) -> Self;
105        /// The complex number.
106        fn to_complex(self) -> Complex<Self::R>;
107        /// The element as its complex type.
108        fn to_c(self) -> Self::C;
109        /// The element of the complex type `c`, its real part for a real type.
110        fn from_c(c: Self::C) -> Self;
111        /// Whether both parts are finite.
112        fn is_finite(self) -> bool;
113        /// `c = β·c + α·a·b` by `gemm` on blocks of the real type `R`, which code generic over the
114        /// element reaches no other way.
115        #[allow(clippy::too_many_arguments)]
116        fn gemm_real(
117            m: usize,
118            n: usize,
119            k: usize,
120            alpha: Self::R,
121            a: &[Self::R],
122            lda: usize,
123            b: &[Self::R],
124            ldb: usize,
125            beta: Self::R,
126            c: &mut [Self::R],
127            ldc: usize,
128        );
129        /// [`gemm`](super::super::blas::gemm) on the type.
130        #[allow(clippy::too_many_arguments)]
131        fn gemm(
132            m: usize,
133            n: usize,
134            k: usize,
135            alpha: Self,
136            a: &[Self],
137            lda: usize,
138            opa: Op,
139            b: &[Self],
140            ldb: usize,
141            opb: Op,
142            beta: Self,
143            c: &mut [Self],
144            ldc: usize,
145        );
146    }
147}
148
149/// A type a [`Matrix`](super::Matrix) can hold and compute with exactly: `f32`, `f64`,
150/// [`Complex<f32>`], [`Complex<f64>`] and [`Number`]. The trait is sealed.
151pub trait Element: sealed::Element {}
152
153/// An element whose matrices have the numerical algorithms, the decompositions, the eigenvalues
154/// and the functions of matrices: `f32`, `f64`, [`Complex<f32>`] and [`Complex<f64>`]. The trait
155/// is sealed.
156pub trait Scalar: Element + sealed::Scalar<R = <Self as Scalar>::Real> {
157    /// The float type of the parts and the absolute values, as of norms and singular values:
158    /// `f32` for `f32` and `Complex<f32>`, `f64` for `f64` and `Complex<f64>`.
159    type Real: Float;
160}
161
162impl Element for Number {}
163
164impl sealed::Element for Number {
165    #[inline]
166    fn zero() -> Self {
167        Number::default()
168    }
169    #[inline]
170    fn one() -> Self {
171        Number::from(1u8)
172    }
173    #[inline]
174    fn is_zero(&self) -> bool {
175        Number::is_zero(self)
176    }
177    #[inline]
178    fn conj(&self) -> Self {
179        self.clone()
180    }
181    #[inline]
182    fn add(a: &Self, b: &Self) -> Self {
183        a + b
184    }
185    #[inline]
186    fn sub(a: &Self, b: &Self) -> Self {
187        a - b
188    }
189    #[inline]
190    fn mul(a: &Self, b: &Self) -> Self {
191        a * b
192    }
193    #[inline]
194    fn neg(a: &Self) -> Self {
195        -a
196    }
197    #[inline]
198    fn add_assign(a: &mut Self, b: &Self) {
199        *a += b;
200    }
201    #[inline]
202    fn sub_assign(a: &mut Self, b: &Self) {
203        *a -= b;
204    }
205    #[inline]
206    fn mul_assign(a: &mut Self, b: &Self) {
207        *a *= b;
208    }
209    #[inline]
210    fn mul_add(c: &mut Self, a: &Self, b: &Self) {
211        if !a.is_zero() {
212            *c += a * b;
213        }
214    }
215}
216
217/// [`sealed::Scalar::gemm`] by the function `$f` of [`gemm`](super::gemm).
218macro_rules! gemm {
219    ($f:ident) => {
220        #[inline]
221        #[track_caller]
222        fn gemm(
223            m: usize,
224            n: usize,
225            k: usize,
226            alpha: Self,
227            a: &[Self],
228            lda: usize,
229            opa: Op,
230            b: &[Self],
231            ldb: usize,
232            opb: Op,
233            beta: Self,
234            c: &mut [Self],
235            ldc: usize,
236        ) {
237            super::gemm::$f(m, n, k, alpha, a, lda, opa, b, ldb, opb, beta, c, ldc);
238        }
239    };
240}
241
242/// [`sealed::Element`] by the operators of a `Copy` type.
243macro_rules! copy_element {
244    ($zero:expr, $one:expr, $conj:expr) => {
245        #[inline]
246        fn zero() -> Self {
247            $zero
248        }
249        #[inline]
250        fn one() -> Self {
251            $one
252        }
253        #[inline]
254        fn is_zero(&self) -> bool {
255            *self == $zero
256        }
257        #[inline]
258        fn conj(&self) -> Self {
259            $conj(*self)
260        }
261        #[inline]
262        fn add(a: &Self, b: &Self) -> Self {
263            *a + *b
264        }
265        #[inline]
266        fn sub(a: &Self, b: &Self) -> Self {
267            *a - *b
268        }
269        #[inline]
270        fn mul(a: &Self, b: &Self) -> Self {
271            *a * *b
272        }
273        #[inline]
274        fn neg(a: &Self) -> Self {
275            -*a
276        }
277        #[inline]
278        fn add_assign(a: &mut Self, b: &Self) {
279            *a += *b;
280        }
281        #[inline]
282        fn sub_assign(a: &mut Self, b: &Self) {
283            *a -= *b;
284        }
285        #[inline]
286        fn mul_assign(a: &mut Self, b: &Self) {
287            *a *= *b;
288        }
289        #[inline]
290        fn mul_add(c: &mut Self, a: &Self, b: &Self) {
291            *c += *a * *b;
292        }
293        fn product(c: &mut [Self], a: &[Self], b: &[Self], m: usize, k: usize, n: usize) {
294            gemm(m, n, k, $one, a, m, Op::N, b, k, Op::N, $zero, c, m);
295        }
296    };
297}
298
299macro_rules! real {
300    ($t:ident) => {
301        impl Element for $t {}
302        impl Scalar for $t {
303            type Real = $t;
304        }
305
306        impl sealed::Element for $t {
307            copy_element!(0.0, 1.0, core::convert::identity);
308        }
309
310        impl sealed::Scalar for $t {
311            type R = $t;
312            type C = Complex<$t>;
313            type Re = $t;
314            const IS_COMPLEX: bool = false;
315            #[inline]
316            fn re(self) -> $t {
317                self
318            }
319            #[inline]
320            fn im(self) -> $t {
321                0.0
322            }
323            #[inline]
324            fn from_real(re: $t) -> Self {
325                re
326            }
327            #[inline]
328            fn from_parts(re: $t, _: $t) -> Self {
329                re
330            }
331            #[inline]
332            fn abs(self) -> $t {
333                $t::abs(self)
334            }
335            #[inline]
336            fn abs2(self) -> $t {
337                self * self
338            }
339            #[inline]
340            fn abs1(self) -> $t {
341                $t::abs(self)
342            }
343            #[inline]
344            fn scale(self, r: $t) -> Self {
345                self * r
346            }
347            #[inline(always)]
348            fn mul_sub(self, a: Self, b: Self) -> Self {
349                $t::mul_add(-a, b, self)
350            }
351            #[inline]
352            fn sqrt(self) -> Self {
353                $t::sqrt(self)
354            }
355            #[inline]
356            fn to_complex(self) -> Complex<$t> {
357                Complex::new(self, 0.0)
358            }
359            #[inline]
360            fn to_c(self) -> Complex<$t> {
361                Complex::new(self, 0.0)
362            }
363            #[inline]
364            fn from_c(c: Complex<$t>) -> Self {
365                c.re
366            }
367            #[inline]
368            fn is_finite(self) -> bool {
369                $t::is_finite(self)
370            }
371            #[inline]
372            fn gemm_real(
373                m: usize,
374                n: usize,
375                k: usize,
376                alpha: $t,
377                a: &[$t],
378                lda: usize,
379                b: &[$t],
380                ldb: usize,
381                beta: $t,
382                c: &mut [$t],
383                ldc: usize,
384            ) {
385                gemm(m, n, k, alpha, a, lda, Op::N, b, ldb, Op::N, beta, c, ldc);
386            }
387            gemm!(real);
388        }
389
390        impl Element for Complex<$t> {}
391        impl Scalar for Complex<$t> {
392            type Real = $t;
393        }
394
395        impl sealed::Element for Complex<$t> {
396            copy_element!(Complex::<$t>::ZERO, Complex::<$t>::ONE, Complex::conj);
397        }
398
399        impl sealed::Scalar for Complex<$t> {
400            type R = $t;
401            type C = Self;
402            type Re = $t;
403            const IS_COMPLEX: bool = true;
404            #[inline]
405            fn re(self) -> $t {
406                self.re
407            }
408            #[inline]
409            fn im(self) -> $t {
410                self.im
411            }
412            #[inline]
413            fn from_real(re: $t) -> Self {
414                Complex::new(re, 0.0)
415            }
416            #[inline]
417            fn from_parts(re: $t, im: $t) -> Self {
418                Complex::new(re, im)
419            }
420            #[inline]
421            fn abs(self) -> $t {
422                Complex::abs(self)
423            }
424            #[inline]
425            fn abs2(self) -> $t {
426                self.norm_sqr()
427            }
428            #[inline]
429            fn abs1(self) -> $t {
430                self.re.abs() + self.im.abs()
431            }
432            #[inline]
433            fn scale(self, r: $t) -> Self {
434                self * r
435            }
436            #[inline(always)]
437            fn mul_sub(self, a: Self, b: Self) -> Self {
438                let re = $t::mul_add(a.im, b.im, $t::mul_add(-a.re, b.re, self.re));
439                let im = $t::mul_add(-a.im, b.re, $t::mul_add(-a.re, b.im, self.im));
440                Complex::new(re, im)
441            }
442            #[inline]
443            fn sqrt(self) -> Self {
444                Complex::sqrt(self)
445            }
446            #[inline]
447            fn to_complex(self) -> Complex<$t> {
448                self
449            }
450            #[inline]
451            fn to_c(self) -> Self {
452                self
453            }
454            #[inline]
455            fn from_c(c: Self) -> Self {
456                c
457            }
458            #[inline]
459            fn is_finite(self) -> bool {
460                Complex::is_finite(self)
461            }
462            #[inline]
463            fn gemm_real(
464                m: usize,
465                n: usize,
466                k: usize,
467                alpha: $t,
468                a: &[$t],
469                lda: usize,
470                b: &[$t],
471                ldb: usize,
472                beta: $t,
473                c: &mut [$t],
474                ldc: usize,
475            ) {
476                gemm(m, n, k, alpha, a, lda, Op::N, b, ldb, Op::N, beta, c, ldc);
477            }
478            gemm!(complex);
479        }
480    };
481}
482
483real!(f32);
484real!(f64);