mdarray_linalg/naive/matvec/
context.rs

1use std::ops::{Add, Mul};
2
3use mdarray::{Array, Dim, Layout, Shape, Slice};
4use num_complex::ComplexFloat;
5use num_traits::Zero;
6
7use super::simple::naive_outer;
8use crate::{
9    Naive,
10    matvec::{Argmax, MatVec, MatVecBuilder, Outer, OuterBuilder, VecOps},
11    utils::unravel_index,
12};
13
14struct NaiveMatVecBuilder<'a, T, La, Lx, D0, D1>
15where
16    La: Layout,
17    Lx: Layout,
18    D0: Dim,
19    D1: Dim,
20{
21    alpha: T,
22    a: &'a Slice<T, (D0, D1), La>,
23    x: &'a Slice<T, (D1,), Lx>,
24}
25
26impl<'a, T, La, Lx, D0, D1> MatVecBuilder<'a, T, La, Lx, D0, D1>
27    for NaiveMatVecBuilder<'a, T, La, Lx, D0, D1>
28where
29    La: Layout,
30    Lx: Layout,
31    T: ComplexFloat,
32    D0: Dim,
33    D1: Dim,
34{
35    /// `α := α·α'`
36    fn scale(mut self, alpha: T) -> Self {
37        self.alpha = alpha * self.alpha;
38        self
39    }
40
41    fn eval(self) -> Array<T, (D0,)> {
42        let ash = *self.a.shape();
43        let (m, n) = (ash.dim(0), ash.dim(1));
44        let x_len = self.x.shape().dim(0);
45
46        assert!(n == x_len, "Matrix columns must match vector length");
47
48        let result_shape = <(D0,) as Shape>::from_dims(&[m]);
49        let mut result = Array::<T, (D0,)>::from_elem(result_shape, T::zero());
50
51        for i in 0..m {
52            let mut sum = T::zero();
53            for j in 0..n {
54                sum = sum + self.a[[i, j]] * self.x[[j]];
55            }
56            result[[i]] = self.alpha * sum;
57        }
58        result
59    }
60
61    fn write<Ly: Layout>(self, y: &mut Slice<T, (D0,), Ly>) {
62        let ash = *self.a.shape();
63        let (m, n) = (ash.dim(0), ash.dim(1));
64        let x_len = self.x.shape().dim(0);
65        let y_len = y.shape().dim(0);
66
67        assert!(n == x_len, "Matrix columns must match vector length");
68        assert!(m == y_len, "Matrix rows must match y vector length");
69
70        for i in 0..m {
71            let mut sum = T::zero();
72            for j in 0..n {
73                sum = sum + self.a[[i, j]] * self.x[[j]];
74            }
75            y[[i]] = self.alpha * sum;
76        }
77    }
78
79    fn add_to_vec<Ly: Layout>(self, y: &mut Slice<T, (D0,), Ly>) {
80        let ash = *self.a.shape();
81        let (m, n) = (ash.dim(0), ash.dim(1));
82        let x_len = self.x.shape().dim(0);
83        let y_len = y.shape().dim(0);
84
85        assert!(n == x_len, "Matrix columns must match x vector length");
86        assert!(m == y_len, "Matrix rows must match y vector length");
87
88        for i in 0..m {
89            for j in 0..n {
90                y[[i]] = y[[i]] + self.alpha * self.a[[i, j]] * self.x[[j]];
91            }
92        }
93    }
94
95    fn add_to_scaled_vec<Ly: Layout>(self, y: &mut Slice<T, (D0,), Ly>, beta: T) {
96        let ash = *self.a.shape();
97        let (m, n) = (ash.dim(0), ash.dim(1));
98        let x_len = self.x.shape().dim(0);
99        let y_len = y.shape().dim(0);
100
101        assert!(n == x_len, "Matrix columns must match x vector length");
102        assert!(m == y_len, "Matrix rows must match y vector length");
103
104        for i in 0..m {
105            y[[i]] = beta * y[[i]];
106        }
107
108        for i in 0..m {
109            for j in 0..n {
110                y[[i]] = y[[i]] + self.alpha * self.a[[i, j]] * self.x[[j]];
111            }
112        }
113    }
114}
115
116impl<T, D0: Dim, D1: Dim> MatVec<T, D0, D1> for Naive
117where
118    T: ComplexFloat,
119{
120    fn matvec<'a, La, Lx>(
121        &self,
122        a: &'a Slice<T, (D0, D1), La>,
123        x: &'a Slice<T, (D1,), Lx>,
124    ) -> impl MatVecBuilder<'a, T, La, Lx, D0, D1>
125    where
126        La: Layout,
127        Lx: Layout,
128    {
129        NaiveMatVecBuilder {
130            alpha: T::one(),
131            a,
132            x,
133        }
134    }
135}
136
137impl<T: ComplexFloat + Add<Output = T> + Mul<Output = T> + Zero + Copy, D: Dim> VecOps<T, D>
138    for Naive
139{
140    type Real = T::Real;
141
142    fn add_to_scaled<Lx: Layout, Ly: Layout>(
143        &self,
144        alpha: T,
145        x: &Slice<T, (D,), Lx>,
146        y: &mut Slice<T, (D,), Ly>,
147    ) {
148        for (elem_x, elem_y) in std::iter::zip(x, y) {
149            *elem_y = alpha * (*elem_x) + *elem_y;
150        }
151    }
152
153    fn dot<Lx: Layout, Ly: Layout>(&self, x: &Slice<T, (D,), Lx>, y: &Slice<T, (D,), Ly>) -> T {
154        let mut result = T::zero();
155        for (elem_x, elem_y) in std::iter::zip(x, y) {
156            result = result + *elem_x * (*elem_y);
157        }
158        result
159    }
160
161    fn dotc<Lx: Layout, Ly: Layout>(&self, x: &Slice<T, (D,), Lx>, y: &Slice<T, (D,), Ly>) -> T {
162        let mut result = T::zero();
163        for (elem_x, elem_y) in std::iter::zip(x, y) {
164            result = result + elem_x.conj() * (*elem_y);
165        }
166        result
167    }
168
169    fn norm2<Lx: Layout>(&self, x: &Slice<T, (D,), Lx>) -> Self::Real {
170        let mut sum_sq = T::Real::zero();
171        for elem in x.into_iter() {
172            sum_sq = sum_sq + elem.abs().powi(2);
173        }
174        sum_sq.sqrt()
175    }
176
177    fn norm1<Lx: Layout>(&self, x: &Slice<T, (D,), Lx>) -> Self::Real {
178        let mut sum = T::Real::zero();
179        for elem in x.into_iter() {
180            sum = sum + elem.re().abs() + elem.im().abs();
181        }
182        sum
183    }
184
185    fn rot<Lx: Layout, Ly: Layout>(
186        &self,
187        x: &mut Slice<T, (D,), Lx>,
188        y: &mut Slice<T, (D,), Ly>,
189        c: Self::Real,
190        s: T,
191    ) {
192        // Apply a Givens rotation to vectors x and y in-place:
193        //   x[i] =  c * x[i] + s  * y[i]
194        //   y[i] = -s* * x[i] + c * y[i]
195        // where c is real, s is (possibly complex), and c² + |s|² = 1.
196        //
197        // For real types, s* == s and this reduces to the standard real Givens rotation.
198        // This matches the BLAS `drot`/`zrot` convention.
199        for (elem_x, elem_y) in std::iter::zip(x, y) {
200            // Store original x before overwriting it
201            let old_x = *elem_x;
202            let old_y = *elem_y;
203
204            // c is T::Real, so we need to cast it to T for arithmetic with s
205            let c_as_t = T::from(c).unwrap();
206
207            *elem_x = c_as_t * old_x + s * old_y;
208            *elem_y = c_as_t * old_y - s.conj() * old_x;
209        }
210    }
211}
212
213impl<T: ComplexFloat<Real = T> + PartialOrd + Add<Output = T> + Mul<Output = T> + Zero + Copy>
214    Argmax<T> for Naive
215{
216    fn argmax_write<Lx: Layout, S: Shape>(
217        &self,
218        x: &Slice<T, S, Lx>,
219        output: &mut Vec<usize>,
220    ) -> bool {
221        output.clear();
222
223        if x.is_empty() {
224            return false;
225        }
226
227        if x.rank() == 0 {
228            return true;
229        }
230
231        let mut max_flat_idx = 0;
232        let mut max_val = x.iter().next().unwrap();
233
234        for (flat_idx, val) in x.iter().enumerate().skip(1) {
235            if val > max_val {
236                max_val = val;
237                max_flat_idx = flat_idx;
238            }
239        }
240
241        let indices = unravel_index(x, max_flat_idx);
242        output.extend_from_slice(&indices);
243        true
244    }
245
246    fn argmax<Lx: Layout, S: Shape>(&self, x: &Slice<T, S, Lx>) -> Option<Vec<usize>> {
247        let mut result = Vec::new();
248        if self.argmax_write(x, &mut result) {
249            Some(result)
250        } else {
251            None
252        }
253    }
254
255    fn argmax_abs_write<Lx: Layout, S: Shape>(
256        &self,
257        x: &Slice<T, S, Lx>,
258        output: &mut Vec<usize>,
259    ) -> bool {
260        output.clear();
261
262        if x.is_empty() {
263            return false;
264        }
265
266        if x.rank() == 0 {
267            return true;
268        }
269
270        let mut max_flat_idx = 0;
271        let mut max_val = x.iter().next().unwrap().abs();
272
273        for (flat_idx, val) in x.iter().enumerate().skip(1) {
274            if val.abs() > max_val {
275                max_val = val.abs();
276                max_flat_idx = flat_idx;
277            }
278        }
279
280        let indices = unravel_index(x, max_flat_idx);
281        output.extend_from_slice(&indices);
282        true
283    }
284
285    fn argmax_abs<Lx: Layout, S: Shape>(&self, x: &Slice<T, S, Lx>) -> Option<Vec<usize>> {
286        let mut result = Vec::new();
287        if self.argmax_abs_write(x, &mut result) {
288            Some(result)
289        } else {
290            None
291        }
292    }
293}
294
295impl<T, Dx, Dy> Outer<T, Dx, Dy> for Naive
296where
297    T: ComplexFloat,
298    Dx: Dim,
299    Dy: Dim,
300{
301    fn outer<'a, Lx, Ly>(
302        &self,
303        x: &'a Slice<T, (Dx,), Lx>,
304        y: &'a Slice<T, (Dy,), Ly>,
305    ) -> impl OuterBuilder<'a, T, Lx, Ly, Dx, Dy>
306    where
307        Lx: Layout,
308        Ly: Layout,
309    {
310        NaiveOuterBuilder {
311            alpha: T::one(),
312            x,
313            y,
314        }
315    }
316}
317
318struct NaiveOuterBuilder<'a, T, Lx, Ly, Dx, Dy>
319where
320    Lx: Layout,
321    Ly: Layout,
322    Dx: Dim,
323    Dy: Dim,
324{
325    alpha: T,
326    x: &'a Slice<T, (Dx,), Lx>,
327    y: &'a Slice<T, (Dy,), Ly>,
328}
329
330impl<'a, T, Lx, Ly, Dx, Dy> OuterBuilder<'a, T, Lx, Ly, Dx, Dy>
331    for NaiveOuterBuilder<'a, T, Lx, Ly, Dx, Dy>
332where
333    Lx: Layout,
334    Ly: Layout,
335    Dx: Dim,
336    Dy: Dim,
337    T: ComplexFloat,
338{
339    /// `α := α·α'`
340    fn scale(mut self, alpha: T) -> Self {
341        self.alpha = alpha * self.alpha;
342        self
343    }
344
345    /// Returns `α·xy`
346    fn eval(self) -> Array<T, (Dx, Dy)> {
347        let m = self.x.shape().dim(0);
348        let n = self.y.shape().dim(0);
349
350        let a_shape = <(Dx, Dy) as Shape>::from_dims(&[m, n]);
351        let mut a = Array::<T, (Dx, Dy)>::from_elem(a_shape, T::zero());
352
353        naive_outer(&mut a, self.x, self.y, self.alpha);
354
355        a
356    }
357
358    /// `a := α·xy`
359    fn write<La: Layout>(self, a: &mut Slice<T, (Dx, Dy), La>) {
360        let m = self.x.shape().dim(0);
361        let n = self.y.shape().dim(0);
362
363        let ash = *a.shape();
364        let (ma, na) = (ash.dim(0), ash.dim(1));
365
366        assert!(ma == m, "Output shape must match input vector length");
367        assert!(na == n, "Output shape must match input vector length");
368
369        naive_outer(a, self.x, self.y, self.alpha);
370    }
371
372    /// Rank-1 update: `A := α·x·yᵀ + A`
373    fn add_to<La: Layout>(self, a: &mut Slice<T, (Dx, Dy), La>) {
374        let m = self.x.shape().dim(0);
375        let n = self.y.shape().dim(0);
376
377        let ash = *a.shape();
378        let (ma, na) = (ash.dim(0), ash.dim(1));
379
380        assert!(ma == m, "Output shape must match input vector length");
381        assert!(na == n, "Output shape must match input vector length");
382
383        naive_outer(a, self.x, self.y, self.alpha);
384    }
385
386}
mdarray_linalg/naive/matvec/context.rs

mdarray_linalg/naive/matvec/
context.rs
