diffsol_la/vector/
nalgebra_serial.rs

1use std::ops::{Add, AddAssign, Div, Index, IndexMut, Mul, MulAssign, Sub, SubAssign};
2
3use nalgebra::{Const, DMatrix, DVector, Dyn, LpNorm, MatrixView, MatrixViewMut};
4
5use crate::context::broadcast_batch;
6use crate::{Context, IndexType, NalgebraContext, NalgebraMat, NalgebraScalar, Scale};
7
8use super::{DefaultDenseMatrix, Vector, VectorCommon, VectorIndex, VectorView, VectorViewMut};
9
10/// A batched vector, stored as one column per batch.
11///
12/// Invariant: `data.ncols() == context.nbatch()` and `data.nrows() == len()`, so
13/// `data.as_slice()` matches the CUDA layout `[batch0 states..., batch1 states..., ...]`.
14#[derive(Debug, Clone, PartialEq)]
15pub struct NalgebraVec<T: NalgebraScalar> {
16    pub(crate) data: DMatrix<T>,
17    pub(crate) context: NalgebraContext,
18}
19
20#[derive(Debug, Clone, PartialEq)]
21pub struct NalgebraIndex {
22    pub(crate) data: DVector<IndexType>,
23    pub(crate) context: NalgebraContext,
24}
25
26#[derive(Debug, Clone, PartialEq)]
27pub struct NalgebraVecRef<'a, T: NalgebraScalar> {
28    pub(crate) data: MatrixView<'a, T, Dyn, Dyn, Const<1>, Dyn>,
29    pub(crate) context: NalgebraContext,
30}
31
32#[derive(Debug, PartialEq)]
33pub struct NalgebraVecMut<'a, T: NalgebraScalar> {
34    pub(crate) data: MatrixViewMut<'a, T, Dyn, Dyn, Const<1>, Dyn>,
35    pub(crate) context: NalgebraContext,
36}
37
38impl<T: NalgebraScalar> NalgebraVec<T> {
39    #[inline]
40    pub(crate) fn batch(&self, batch: usize, nbatch: usize) -> usize {
41        broadcast_batch(batch, self.data.ncols(), nbatch)
42    }
43}
44impl<T: NalgebraScalar> NalgebraVecRef<'_, T> {
45    #[inline]
46    pub(crate) fn batch(&self, batch: usize, nbatch: usize) -> usize {
47        broadcast_batch(batch, self.data.ncols(), nbatch)
48    }
49}
50
51impl<T: NalgebraScalar> From<DVector<T>> for NalgebraVec<T> {
52    fn from(data: DVector<T>) -> Self {
53        Self {
54            data: DMatrix::from_column_slice(data.len(), 1, data.as_slice()),
55            context: NalgebraContext::default(),
56        }
57    }
58}
59
60impl<T: NalgebraScalar> DefaultDenseMatrix for NalgebraVec<T> {
61    type M = NalgebraMat<T>;
62}
63
64macro_rules! impl_vector_common {
65    ($t:ty, $inner:ty) => {
66        impl<T: NalgebraScalar> VectorCommon for $t {
67            type T = T;
68            type C = NalgebraContext;
69            type Inner = $inner;
70            fn inner(&self) -> &Self::Inner {
71                &self.data
72            }
73        }
74    };
75}
76impl_vector_common!(NalgebraVec<T>, DMatrix<T>);
77macro_rules! impl_vector_common_ref {
78    ($t:ty, $inner:ty) => {
79        impl<'a, T: NalgebraScalar> VectorCommon for $t {
80            type T = T;
81            type C = NalgebraContext;
82            type Inner = $inner;
83            fn inner(&self) -> &Self::Inner {
84                &self.data
85            }
86        }
87    };
88}
89impl_vector_common_ref!(
90    NalgebraVecRef<'a, T>,
91    MatrixView<'a, T, Dyn, Dyn, Const<1>, Dyn>
92);
93impl_vector_common_ref!(
94    NalgebraVecMut<'a, T>,
95    MatrixViewMut<'a, T, Dyn, Dyn, Const<1>, Dyn>
96);
97
98macro_rules! impl_binary_ref_ref {
99    ($trait:ident, $method:ident, $lhs:ty, $rhs:ty, $op:tt, $binary:tt) => {
100        impl<T: NalgebraScalar> $trait<$rhs> for $lhs {
101            type Output = NalgebraVec<T>;
102
103            fn $method(self, rhs: $rhs) -> Self::Output {
104                // neither operand is owned, so the result is allocated at the left-hand
105                // side's batch count and `rhs` broadcasts into it
106                self.context
107                    .assert_broadcastable_into(rhs.context.nbatch(), stringify!($method));
108                if self.data.ncols() == rhs.data.ncols() {
109                    return NalgebraVec {
110                        data: &self.data $binary &rhs.data,
111                        context: self.context,
112                    };
113                }
114                let nb = self.data.ncols();
115                let mut data = DMatrix::zeros(self.data.nrows(), nb);
116                for b in 0..nb {
117                    let mut column = data.column_mut(b);
118                    column.copy_from(&self.data.column(b));
119                    column $op &rhs.data.column(rhs.batch(b, nb));
120                }
121                NalgebraVec {
122                    data,
123                    context: self.context,
124                }
125            }
126        }
127    };
128}
129macro_rules! impl_assign {
130    ($trait:ident, $method:ident, $lhs:ty, $rhs:ty, $op:tt) => {
131        impl<T: NalgebraScalar> $trait<$rhs> for $lhs {
132            fn $method(&mut self, rhs: $rhs) {
133                self.context
134                    .assert_broadcastable_into(rhs.context.nbatch(), stringify!($method));
135                if self.data.ncols() == rhs.data.ncols() {
136                    self.data $op &rhs.data;
137                    return;
138                }
139                let nb = self.data.ncols();
140                for b in 0..nb {
141                    let mut column = self.data.column_mut(b);
142                    column $op &rhs.data.column(rhs.batch(b, nb));
143                }
144            }
145        }
146    };
147}
148
149/// `self` is the owned operand, so it is the destination -- which makes the in-place op the
150/// whole implementation, broadcast check included.
151macro_rules! impl_binary_owned_lhs {
152    ($trait:ident, $method:ident, $assign_trait:ident, $assign:ident, $rhs:ty) => {
153        impl<T: NalgebraScalar> $trait<$rhs> for NalgebraVec<T> {
154            type Output = NalgebraVec<T>;
155
156            fn $method(mut self, rhs: $rhs) -> Self::Output {
157                $assign_trait::$assign(&mut self, rhs);
158                self
159            }
160        }
161    };
162}
163
164/// `rhs` is the owned operand, so it is the destination.
165///
166/// A commutative op is just the in-place op with the operands swapped.  A non-commutative one
167/// cannot be: `rhs -= lhs` computes `rhs - lhs`, so it writes `combine(&mut rhs_i, lhs_i)`
168/// instead, which gets `lhs - rhs` in the same single pass.
169macro_rules! impl_binary_owned_rhs {
170    (commutes, $trait:ident, $method:ident, $assign_trait:ident, $assign:ident, $lhs:ty,
171     $op:tt, $combine:expr) => {
172        impl<T: NalgebraScalar> $trait<NalgebraVec<T>> for $lhs {
173            type Output = NalgebraVec<T>;
174
175            fn $method(self, mut rhs: NalgebraVec<T>) -> Self::Output {
176                $assign_trait::$assign(&mut rhs, self);
177                rhs
178            }
179        }
180    };
181    (noncommutes, $trait:ident, $method:ident, $assign_trait:ident, $assign:ident, $lhs:ty, $op:tt, $combine:expr) => {
182        impl<T: NalgebraScalar> $trait<NalgebraVec<T>> for $lhs {
183            type Output = NalgebraVec<T>;
184
185            fn $method(self, mut rhs: NalgebraVec<T>) -> Self::Output {
186                rhs.context
187                    .assert_broadcastable_into(self.context.nbatch(), stringify!($method));
188                if self.data.ncols() == rhs.data.ncols() {
189                    rhs.data.zip_apply(&self.data, $combine);
190                    return rhs;
191                }
192                let nb = rhs.data.ncols();
193                for b in 0..nb {
194                    let lhs = self.data.column(self.batch(b, nb));
195                    rhs.data.column_mut(b).zip_apply(&lhs, $combine);
196                }
197                rhs
198            }
199        }
200    };
201}
202
203macro_rules! impl_binary_set {
204    ($trait:ident,$method:ident,$assign_trait:ident,$assign:ident,$commutes:ident,
205     $op:tt,$binary:tt,$combine:expr) => {
206        impl_binary_owned_lhs!($trait, $method, $assign_trait, $assign, NalgebraVec<T>);
207        impl_binary_owned_lhs!($trait, $method, $assign_trait, $assign, &NalgebraVec<T>);
208        impl_binary_owned_lhs!(
209            $trait,
210            $method,
211            $assign_trait,
212            $assign,
213            NalgebraVecRef<'_, T>
214        );
215        impl_binary_owned_lhs!(
216            $trait,
217            $method,
218            $assign_trait,
219            $assign,
220            &NalgebraVecRef<'_, T>
221        );
222        impl_binary_owned_rhs!(
223            $commutes,
224            $trait,
225            $method,
226            $assign_trait,
227            $assign,
228            NalgebraVecRef<'_, T>,
229            $op,
230            $combine
231        );
232        impl_binary_owned_rhs!(
233            $commutes,
234            $trait,
235            $method,
236            $assign_trait,
237            $assign,
238            &NalgebraVec<T>,
239            $op,
240            $combine
241        );
242        impl_binary_ref_ref!(
243            $trait,
244            $method,
245            NalgebraVecRef<'_, T>,
246            &NalgebraVec<T>,
247            $op,
248            $binary
249        );
250        impl_binary_ref_ref!(
251            $trait,
252            $method,
253            NalgebraVecRef<'_, T>,
254            NalgebraVecRef<'_, T>,
255            $op,
256            $binary
257        );
258        impl_binary_ref_ref!(
259            $trait,
260            $method,
261            NalgebraVecRef<'_, T>,
262            &NalgebraVecRef<'_, T>,
263            $op,
264            $binary
265        );
266        impl_binary_ref_ref!(
267            $trait,
268            $method,
269            &NalgebraVec<T>,
270            &NalgebraVec<T>,
271            $op,
272            $binary
273        );
274        impl_binary_ref_ref!(
275            $trait,
276            $method,
277            &NalgebraVec<T>,
278            NalgebraVecRef<'_, T>,
279            $op,
280            $binary
281        );
282        impl_binary_ref_ref!(
283            $trait,
284            $method,
285            &NalgebraVec<T>,
286            &NalgebraVecRef<'_, T>,
287            $op,
288            $binary
289        );
290    };
291}
292macro_rules! impl_assign_set {
293    ($trait:ident,$method:ident,$op:tt) => {
294        impl_assign!($trait, $method, NalgebraVec<T>, NalgebraVec<T>, $op);
295        impl_assign!($trait, $method, NalgebraVec<T>, &NalgebraVec<T>, $op);
296        impl_assign!($trait, $method, NalgebraVec<T>, NalgebraVecRef<'_, T>, $op);
297        impl_assign!($trait, $method, NalgebraVec<T>, &NalgebraVecRef<'_, T>, $op);
298        impl_assign!($trait, $method, NalgebraVecMut<'_, T>, NalgebraVec<T>, $op);
299        impl_assign!($trait, $method, NalgebraVecMut<'_, T>, &NalgebraVec<T>, $op);
300        impl_assign!(
301            $trait,
302            $method,
303            NalgebraVecMut<'_, T>,
304            NalgebraVecRef<'_, T>,
305            $op
306        );
307        impl_assign!(
308            $trait,
309            $method,
310            NalgebraVecMut<'_, T>,
311            &NalgebraVecRef<'_, T>,
312            $op
313        );
314    };
315}
316impl_binary_set!(Add, add, AddAssign, add_assign, commutes, +=, +, |rhs, lhs| *rhs += lhs);
317impl_binary_set!(Sub, sub, SubAssign, sub_assign, noncommutes, -=, -, |rhs, lhs| *rhs = lhs - *rhs);
318impl_assign_set!(AddAssign, add_assign, +=);
319impl_assign_set!(SubAssign, sub_assign, -=);
320
321impl<T: NalgebraScalar> Mul<Scale<T>> for NalgebraVec<T> {
322    type Output = Self;
323    fn mul(mut self, rhs: Scale<T>) -> Self {
324        self.data *= rhs.value();
325        self
326    }
327}
328macro_rules! impl_mul_scalar_alloc {
329    ($t:ty) => {
330        impl<T: NalgebraScalar> Mul<Scale<T>> for $t {
331            type Output = NalgebraVec<T>;
332            fn mul(self, rhs: Scale<T>) -> Self::Output {
333                NalgebraVec {
334                    data: self.data.clone_owned() * rhs.value(),
335                    context: self.context,
336                }
337            }
338        }
339    };
340}
341impl_mul_scalar_alloc!(&NalgebraVec<T>);
342impl_mul_scalar_alloc!(NalgebraVecRef<'_, T>);
343impl_mul_scalar_alloc!(NalgebraVecMut<'_, T>);
344impl<T: NalgebraScalar> Div<Scale<T>> for NalgebraVec<T> {
345    type Output = Self;
346    fn div(self, rhs: Scale<T>) -> Self {
347        Self {
348            data: self.data * (T::one() / rhs.value()),
349            context: self.context,
350        }
351    }
352}
353macro_rules! impl_mul_assign_scalar {
354    ($t:ty) => {
355        impl<T: NalgebraScalar> MulAssign<Scale<T>> for $t {
356            fn mul_assign(&mut self, rhs: Scale<T>) {
357                self.data *= rhs.value();
358            }
359        }
360    };
361}
362impl_mul_assign_scalar!(NalgebraVec<T>);
363impl_mul_assign_scalar!(NalgebraVecMut<'_, T>);
364macro_rules! impl_index {
365    ($t:ty) => {
366        impl<T: NalgebraScalar> Index<IndexType> for $t {
367            type Output = T;
368            fn index(&self, i: IndexType) -> &T {
369                assert_eq!(
370                    self.context.nbatch(),
371                    1,
372                    "indexing not supported for batched vectors"
373                );
374                &self.data[(i, 0)]
375            }
376        }
377    };
378}
379impl_index!(NalgebraVec<T>);
380impl_index!(NalgebraVecRef<'_, T>);
381impl<T: NalgebraScalar> IndexMut<IndexType> for NalgebraVec<T> {
382    fn index_mut(&mut self, i: IndexType) -> &mut T {
383        assert_eq!(
384            self.context.nbatch(),
385            1,
386            "indexing not supported for batched vectors"
387        );
388        &mut self.data[(i, 0)]
389    }
390}
391
392impl VectorIndex for NalgebraIndex {
393    type C = NalgebraContext;
394    fn zeros(len: IndexType, ctx: Self::C) -> Self {
395        Self {
396            data: DVector::zeros(len),
397            context: ctx,
398        }
399    }
400    fn len(&self) -> IndexType {
401        self.data.len()
402    }
403    fn from_vec(v: Vec<IndexType>, ctx: Self::C) -> Self {
404        Self {
405            data: DVector::from_vec(v),
406            context: ctx,
407        }
408    }
409    fn clone_as_vec(&self) -> Vec<IndexType> {
410        self.data.iter().copied().collect()
411    }
412    fn context(&self) -> &Self::C {
413        &self.context
414    }
415}
416
417// Shared bodies for the `Vector` methods below, in the owned and view flavours.
418macro_rules! copy_from_body {
419    ($self:ident, $other:ident, $method:literal) => {
420        $self
421            .context
422            .assert_broadcastable_into($other.context.nbatch(), $method);
423        assert_eq!(
424            $self.data.nrows(),
425            $other.data.nrows(),
426            "copy_from row mismatch"
427        );
428        // the unbatched path is column-wise on purpose: nalgebra's whole-matrix copy_from is
429        // ~30% slower than copying the single column (see lin_alg_ops copy_from/nalgebra)
430        if $self.context.nbatch() == 1 && $other.context.nbatch() == 1 {
431            $self.data.column_mut(0).copy_from(&$other.data.column(0));
432            return;
433        }
434        if $self.data.ncols() == $other.data.ncols() {
435            $self.data.copy_from(&$other.data);
436            return;
437        }
438        let nb = $self.data.ncols();
439        for b in 0..nb {
440            $self
441                .data
442                .column_mut(b)
443                .copy_from(&$other.data.column($other.batch(b, nb)));
444        }
445    };
446}
447
448/// `self_b = alpha_b * x_b + beta * self_b` for every batch of `self`, broadcasting a
449/// single-batch `x`.  Shared by the owned vector and its mutable view, with `x` either an
450/// owned vector or a view, and `alpha` either one scalar or one value per batch.
451macro_rules! axpy_body {
452    ($self:ident, $x:ident, $beta:expr, $op:literal, |$batch:ident| $alpha:expr) => {{
453        $self
454            .context
455            .assert_broadcastable_into($x.context.nbatch(), $op);
456        let nb = $self.data.ncols();
457        // the unbatched path keeps constant column indices and no loop, which is worth
458        // ~30% here (see lin_alg_ops axpy/nalgebra/100)
459        if nb == 1 {
460            let $batch = 0;
461            let alpha = $alpha;
462            $self
463                .data
464                .column_mut(0)
465                .axpy(alpha, &$x.data.column(0), $beta);
466            return;
467        }
468        for $batch in 0..nb {
469            let alpha = $alpha;
470            $self
471                .data
472                .column_mut($batch)
473                .axpy(alpha, &$x.data.column($x.batch($batch, nb)), $beta);
474        }
475    }};
476}
477
478/// Weighted error norm of every batch, reduced by taking the maximum.
479macro_rules! squared_norm_body {
480    ($self:ident, $y:ident, $atol:ident, $rtol:ident) => {{
481        let nrows = $self.data.nrows();
482        assert!(
483            nrows == $y.data.nrows() && nrows == $atol.data.nrows(),
484            "squared_norm row mismatch"
485        );
486        let nstates = T::from_f64(nrows as f64).unwrap();
487        // rows of a column are contiguous whatever the batch stride, so each batch is a slice
488        let batch_norm = |xb: usize, yb: usize, atolb: usize| {
489            let x = $self.data.column(xb);
490            let y = $y.data.column(yb);
491            let atol = $atol.data.column(atolb);
492            let norm = x
493                .as_slice()
494                .iter()
495                .zip(y.as_slice().iter().zip(atol.as_slice().iter()))
496                .fold(T::zero(), |norm, (&x, (&y, &atol))| {
497                    let term = x.algebraic_div(y.abs().algebraic_mul($rtol).algebraic_add(atol));
498                    norm.algebraic_add(term.algebraic_mul(term))
499                });
500            norm / nstates
501        };
502        // the reduction runs over `self`'s batches and broadcasts `y` and `atol` over them
503        let nb = $self.data.ncols();
504        $self
505            .context
506            .assert_broadcastable_into($y.context.nbatch(), "squared_norm");
507        $self
508            .context
509            .assert_broadcastable_into($atol.context.nbatch(), "squared_norm");
510        // the unbatched case keeps constant column indices, which codegens better than the
511        // loop variable below (and folds away the broadcast arithmetic)
512        if nb == 1 {
513            return batch_norm(0, 0, 0);
514        }
515        let mut max_norm = T::zero();
516        for b in 0..nb {
517            max_norm = max_norm.max(batch_norm(b, $y.batch(b, nb), $atol.batch(b, nb)));
518        }
519        max_norm
520    }};
521}
522
523impl<T: NalgebraScalar> Vector for NalgebraVec<T> {
524    type View<'a> = NalgebraVecRef<'a, T>;
525    type ViewMut<'a> = NalgebraVecMut<'a, T>;
526    type Index = NalgebraIndex;
527    fn len(&self) -> IndexType {
528        self.data.nrows()
529    }
530    fn inner_mut(&mut self) -> &mut Self::Inner {
531        &mut self.data
532    }
533    fn context(&self) -> &Self::C {
534        &self.context
535    }
536    fn norm(&self, k: i32) -> T {
537        // LpNorm is a generic elementwise loop, `norm` is the optimised L2 path
538        (0..self.data.ncols())
539            .map(|b| {
540                let column = self.data.column(b);
541                if k == 2 {
542                    column.norm()
543                } else {
544                    column.apply_norm(&LpNorm(k))
545                }
546            })
547            .fold(T::zero(), |a, b| a.max(b))
548    }
549    fn get_index(&self, i: IndexType) -> T {
550        assert_eq!(
551            self.context.nbatch(),
552            1,
553            "get_index not supported for batched vectors"
554        );
555        self.data[(i, 0)]
556    }
557    fn set_index(&mut self, i: IndexType, v: T) {
558        assert_eq!(
559            self.context.nbatch(),
560            1,
561            "set_index not supported for batched vectors, use fill_index"
562        );
563        self.data.row_mut(i).fill(v);
564    }
565    fn fill_index(&mut self, i: IndexType, v: T) {
566        self.data.row_mut(i).fill(v);
567    }
568    fn squared_norm(&self, y: &Self, atol: &Self, rtol: T) -> T {
569        squared_norm_body!(self, y, atol, rtol)
570    }
571    fn as_view(&self) -> Self::View<'_> {
572        NalgebraVecRef {
573            data: self.data.as_view(),
574            context: self.context,
575        }
576    }
577    fn as_view_mut(&mut self) -> Self::ViewMut<'_> {
578        NalgebraVecMut {
579            data: self.data.as_view_mut(),
580            context: self.context,
581        }
582    }
583    fn get_batch(&self, b: usize) -> Self::View<'_> {
584        NalgebraVecRef {
585            data: self.data.columns(b, 1),
586            context: NalgebraContext::default(),
587        }
588    }
589    fn get_batch_mut(&mut self, b: usize) -> Self::ViewMut<'_> {
590        NalgebraVecMut {
591            data: self.data.columns_mut(b, 1),
592            context: NalgebraContext::default(),
593        }
594    }
595    fn for_each_batch_mut_host<const M: usize, const N: usize>(
596        mut mut_args: [&mut Self; M],
597        args: [&Self; N],
598        mut f: impl FnMut([&mut [T]; M], [&[T]; N], usize),
599    ) {
600        assert!(M > 0, "for_each_batch needs at least one mutable operand");
601        let nbatch = mut_args[0].context.nbatch();
602        {
603            let ctx = &mut_args[0].context;
604            for arg in mut_args.iter() {
605                ctx.assert_broadcastable_into(arg.context.nbatch(), "for_each_batch");
606            }
607            for arg in args.iter() {
608                ctx.assert_broadcastable_into(arg.context.nbatch(), "for_each_batch");
609            }
610        }
611        // unbatched is the overwhelmingly common case, and every operand is then a single
612        // column: hand out the whole backing slice, with no lane index or subslicing at all
613        if nbatch == 1 {
614            let ins = args.map(|a| a.data.as_slice());
615            let outs = mut_args.map(|v| v.data.as_mut_slice());
616            f(outs, ins, 0);
617            return;
618        }
619        for b in 0..nbatch {
620            let ins = args.map(|a| a.batch_as_slice(a.batch(b, nbatch)));
621            let outs = mut_args.each_mut().map(|v| {
622                let vb = v.batch(b, nbatch);
623                v.batch_as_mut_slice(vb)
624            });
625            f(outs, ins, b);
626        }
627    }
628    fn copy_from(&mut self, o: &Self) {
629        copy_from_body!(self, o, "copy_from");
630    }
631    fn fill(&mut self, v: T) {
632        self.data.fill(v)
633    }
634    fn copy_from_view(&mut self, o: &Self::View<'_>) {
635        copy_from_body!(self, o, "copy_from_view");
636    }
637    fn from_element(n: usize, v: T, ctx: Self::C) -> Self {
638        Self {
639            data: DMatrix::from_element(n, ctx.nbatch(), v),
640            context: ctx,
641        }
642    }
643    fn from_vec(v: Vec<T>, ctx: Self::C) -> Self {
644        assert!(
645            v.len().is_multiple_of(ctx.nbatch()),
646            "vector length must be divisible by nbatch"
647        );
648        Self {
649            data: DMatrix::from_vec(v.len() / ctx.nbatch(), ctx.nbatch(), v),
650            context: ctx,
651        }
652    }
653    fn from_slice(v: &[T], ctx: Self::C) -> Self {
654        Self::from_vec(v.to_vec(), ctx)
655    }
656    fn clone_as_vec(&self) -> Vec<T> {
657        self.data.iter().copied().collect()
658    }
659    fn zeros(n: usize, ctx: Self::C) -> Self {
660        Self {
661            data: DMatrix::zeros(n, ctx.nbatch()),
662            context: ctx,
663        }
664    }
665    fn axpy(&mut self, a: T, x: &Self, beta: T) {
666        axpy_body!(self, x, beta, "axpy", |_batch| a)
667    }
668    fn axpy_v(&mut self, a: T, x: &Self::View<'_>, beta: T) {
669        axpy_body!(self, x, beta, "axpy_v", |_batch| a)
670    }
671    fn batched_axpy(&mut self, a: &Self, x: &Self, beta: T) {
672        assert_eq!(
673            a.len(),
674            1,
675            "alpha must be a batched scalar, with len() == 1"
676        );
677        assert_eq!(
678            a.context.nbatch(),
679            self.context.nbatch(),
680            "alpha nbatch must equal nbatch"
681        );
682        axpy_body!(self, x, beta, "batched_axpy", |batch| a.data[(0, batch)])
683    }
684    fn component_div_assign(&mut self, o: &Self) {
685        self.context
686            .assert_broadcastable_into(o.context.nbatch(), "component_div_assign");
687        assert_eq!(
688            self.data.nrows(),
689            o.data.nrows(),
690            "component_div_assign row mismatch"
691        );
692        if self.data.ncols() == o.data.ncols() {
693            self.data.component_div_assign(&o.data);
694            return;
695        }
696        let nb = self.data.ncols();
697        for c in 0..nb {
698            self.data
699                .column_mut(c)
700                .component_div_assign(&o.data.column(o.batch(c, nb)));
701        }
702    }
703    fn component_mul_assign(&mut self, o: &Self) {
704        self.context
705            .assert_broadcastable_into(o.context.nbatch(), "component_mul_assign");
706        assert_eq!(
707            self.data.nrows(),
708            o.data.nrows(),
709            "component_mul_assign row mismatch"
710        );
711        // the elementwise unbatched path beats nalgebra's component_mul_assign by ~25%
712        // (see lin_alg_ops component_mul_assign/nalgebra)
713        if self.context.nbatch() == 1 {
714            for i in 0..self.data.nrows() {
715                let lhs = unsafe { self.data.get_unchecked_mut((i, 0)) };
716                let rhs = unsafe { *o.data.get_unchecked((i, 0)) };
717                *lhs *= rhs;
718            }
719            return;
720        }
721        if self.data.ncols() == o.data.ncols() {
722            self.data.component_mul_assign(&o.data);
723            return;
724        }
725        let nb = self.data.ncols();
726        for c in 0..nb {
727            self.data
728                .column_mut(c)
729                .component_mul_assign(&o.data.column(o.batch(c, nb)));
730        }
731    }
732    fn root_finding(&self, g1: &Self) -> (bool, T, i32) {
733        self.context
734            .assert_broadcastable_into(g1.context.nbatch(), "root_finding");
735        assert_eq!(self.len(), g1.len(), "Vector lengths do not match");
736        let mut out = None;
737        // the scan runs over `self`'s batches and broadcasts `g1` over them
738        let nb = self.data.ncols();
739        for b in 0..nb {
740            let mut found = false;
741            let mut frac = T::zero();
742            let mut idx = -1;
743            let g0_column = self.data.column(b);
744            let g1_column = g1.data.column(g1.batch(b, nb));
745            for (i, (&g0, &g)) in g0_column
746                .as_slice()
747                .iter()
748                .zip(g1_column.as_slice())
749                .enumerate()
750            {
751                if g == T::zero() {
752                    found = true
753                }
754                if g0 * g < T::zero() {
755                    let q = (g / (g - g0)).abs();
756                    if q > frac {
757                        frac = q;
758                        idx = i as i32
759                    }
760                }
761            }
762            if let Some(x) = out {
763                assert_eq!(
764                    x,
765                    (found, frac, idx),
766                    "root finding results differ across batches"
767                )
768            }
769            out = Some((found, frac, idx));
770        }
771        out.unwrap_or((false, T::zero(), -1))
772    }
773    fn assign_at_indices(&mut self, idx: &Self::Index, v: T) {
774        for b in 0..self.data.ncols() {
775            for i in idx.data.iter() {
776                self.data[(*i, b)] = v
777            }
778        }
779    }
780    fn copy_from_indices(&mut self, o: &Self, idx: &Self::Index) {
781        self.context
782            .assert_broadcastable_into(o.context.nbatch(), "copy_from_indices");
783        let nb = self.data.ncols();
784        for b in 0..nb {
785            for i in idx.data.iter() {
786                self.data[(*i, b)] = o.data[(*i, o.batch(b, nb))]
787            }
788        }
789    }
790    fn gather(&mut self, o: &Self, idx: &Self::Index) {
791        assert_eq!(self.len(), idx.len());
792        self.context
793            .assert_broadcastable_into(o.context.nbatch(), "gather");
794        let nb = self.data.ncols();
795        for b in 0..nb {
796            for (i, j) in idx.data.iter().enumerate() {
797                self.data[(i, b)] = o.data[(*j, o.batch(b, nb))]
798            }
799        }
800    }
801    fn scatter(&self, idx: &Self::Index, o: &mut Self) {
802        assert_eq!(self.len(), idx.len());
803        // `o` is the destination here, so its batch count governs the loop
804        o.context
805            .assert_broadcastable_into(self.data.ncols(), "scatter");
806        let nb = o.data.ncols();
807        for b in 0..nb {
808            let src = broadcast_batch(b, self.data.ncols(), nb);
809            for (i, j) in idx.data.iter().enumerate() {
810                o.data[(*j, b)] = self.data[(i, src)]
811            }
812        }
813    }
814}
815
816impl<'a, T: NalgebraScalar> VectorView<'a> for NalgebraVecRef<'a, T> {
817    type Owned = NalgebraVec<T>;
818    fn get_index(&self, i: IndexType) -> T {
819        assert_eq!(self.context.nbatch(), 1, "get_index requires nbatch == 1");
820        self.data[(i, 0)]
821    }
822    fn into_owned(self) -> Self::Owned {
823        NalgebraVec {
824            data: self.data.into_owned(),
825            context: self.context,
826        }
827    }
828    fn squared_norm(&self, y: &Self::Owned, atol: &Self::Owned, rtol: T) -> T {
829        squared_norm_body!(self, y, atol, rtol)
830    }
831}
832impl<'a, T: NalgebraScalar> VectorViewMut<'a> for NalgebraVecMut<'a, T> {
833    type Owned = NalgebraVec<T>;
834    type View = NalgebraVecRef<'a, T>;
835    type Index = NalgebraIndex;
836    fn copy_from(&mut self, o: &Self::Owned) {
837        copy_from_body!(self, o, "copy_from");
838    }
839    fn copy_from_view(&mut self, o: &Self::View) {
840        copy_from_body!(self, o, "copy_from_view");
841    }
842    fn set_index(&mut self, i: IndexType, v: T) {
843        assert_eq!(
844            self.context.nbatch(),
845            1,
846            "set_index not supported for batched vectors, use fill_index"
847        );
848        self.data.row_mut(i).fill(v);
849    }
850    fn fill_index(&mut self, i: IndexType, v: T) {
851        self.data.row_mut(i).fill(v);
852    }
853    fn axpy(&mut self, a: T, x: &Self::Owned, beta: T) {
854        axpy_body!(self, x, beta, "axpy", |_batch| a)
855    }
856}
857impl<T: NalgebraScalar> NalgebraVec<T> {
858    /// Get one batch of the vector data as an immutable slice.
859    ///
860    /// Batches are the columns of the backing matrix, and the storage is column-major.
861    pub fn batch_as_slice(&self, batch: usize) -> &[T] {
862        let n = self.data.nrows();
863        &self.data.as_slice()[batch * n..(batch + 1) * n]
864    }
865    /// Get one batch of the vector data as a mutable slice.
866    pub fn batch_as_mut_slice(&mut self, batch: usize) -> &mut [T] {
867        let n = self.data.nrows();
868        &mut self.data.as_mut_slice()[batch * n..(batch + 1) * n]
869    }
870}
871#[cfg(test)]
872mod tests {
873    use super::*;
874
875    #[test]
876    fn test_error_norm() {
877        let v = NalgebraVec::from_vec(vec![1.0, -2.0, 3.0], Default::default());
878        let y = NalgebraVec::from_vec(vec![1.0, 2.0, 3.0], Default::default());
879        let atol = NalgebraVec::from_vec(vec![0.1, 0.2, 0.3], Default::default());
880        let rtol = 0.1;
881        let mut tmp = y.clone() * Scale(rtol);
882        tmp += &atol;
883        let mut r = v.clone();
884        r.component_div_assign(&tmp);
885        let errorn_check = r.data.column(0).norm_squared() / 3.0;
886        assert_eq!(v.squared_norm(&y, &atol, rtol), errorn_check);
887        let vview = v.as_view();
888        assert_eq!(
889            VectorView::squared_norm(&vview, &y, &atol, rtol),
890            errorn_check
891        );
892    }
893
894    #[test]
895    fn test_root_finding() {
896        super::super::tests::test_root_finding::<NalgebraVec<f64>>();
897    }
898
899    #[test]
900    fn test_from_slice() {
901        let slice = [1.0, 2.0, 3.0];
902        let v = NalgebraVec::from_slice(&slice, Default::default());
903        assert_eq!(v.clone_as_vec(), slice);
904    }
905
906    #[test]
907    fn test_into() {
908        let vector = DVector::from_vec(vec![1.0, 2.0, 3.0]);
909        let v: NalgebraVec<f64> = vector.into();
910        assert_eq!(v.clone_as_vec(), vec![1.0, 2.0, 3.0]);
911    }
912
913    #[test]
914    fn test_host_only() {
915        super::super::tests::test_host_only::<NalgebraVec<f64>>();
916    }
917
918    super::super::generate_vector_tests_nonbatched!(nalgebra, NalgebraVec<f64>);
919    super::super::generate_vector_tests_nonbatched!(nalgebra_f32, NalgebraVec<f32>);
920    super::super::generate_vector_tests_batched!(
921        nalgebra,
922        NalgebraVec<f64>,
923        NalgebraContext::with_nbatch(2),
924        NalgebraContext::with_nbatch(3)
925    );
926    super::super::generate_vector_tests_batched!(
927        nalgebra_f32,
928        NalgebraVec<f32>,
929        NalgebraContext::with_nbatch(2),
930        NalgebraContext::with_nbatch(3)
931    );
932}
diffsol_la/vector/nalgebra_serial.rs

diffsol_la/vector/
nalgebra_serial.rs
