1use std::ops::{Add, AddAssign, Div, Index, IndexMut, Mul, MulAssign, Sub, SubAssign};
2
3use super::utils::*;
4use nalgebra::{DVector, DVectorView, DVectorViewMut, LpNorm};
5
6use crate::{IndexType, NalgebraContext, NalgebraMat, NalgebraScalar, Scalar, Scale, VectorHost};
7
8use super::{DefaultDenseMatrix, Vector, VectorCommon, VectorIndex, VectorView, VectorViewMut};
9
10#[derive(Debug, Clone, PartialEq)]
11pub struct NalgebraIndex {
12 pub(crate) data: DVector<IndexType>,
13 pub(crate) context: NalgebraContext,
14}
15
16#[derive(Debug, Clone, PartialEq)]
17pub struct NalgebraVec<T: NalgebraScalar> {
18 pub(crate) data: DVector<T>,
19 pub(crate) context: NalgebraContext,
20}
21
22#[derive(Debug, Clone, PartialEq)]
23pub struct NalgebraVecRef<'a, T: NalgebraScalar> {
24 pub(crate) data: DVectorView<'a, T>,
25 pub(crate) context: NalgebraContext,
26}
27
28#[derive(Debug, PartialEq)]
29pub struct NalgebraVecMut<'a, T: NalgebraScalar> {
30 pub(crate) data: DVectorViewMut<'a, T>,
31 pub(crate) context: NalgebraContext,
32}
33
34impl<T: NalgebraScalar> From<DVector<T>> for NalgebraVec<T> {
35 fn from(data: DVector<T>) -> Self {
36 Self {
37 data,
38 context: NalgebraContext::default(),
39 }
40 }
41}
42
43impl<T: NalgebraScalar> DefaultDenseMatrix for NalgebraVec<T> {
44 type M = NalgebraMat<T>;
45}
46
47impl_vector_common!(NalgebraVec<T>, NalgebraContext, DVector<T>, NalgebraScalar);
48impl_vector_common_ref!(
49 NalgebraVecRef<'a, T>,
50 NalgebraContext,
51 DVectorView<'a, T>,
52 NalgebraScalar
53);
54impl_vector_common_ref!(
55 NalgebraVecMut<'a, T>,
56 NalgebraContext,
57 DVectorViewMut<'a, T>,
58 NalgebraScalar
59);
60
61macro_rules! impl_mul_scalar {
62 ($lhs:ty, $out:ty, $scalar:ty) => {
63 impl<T: NalgebraScalar> Mul<Scale<T>> for $lhs {
64 type Output = $out;
65 #[inline]
66 fn mul(self, rhs: Scale<T>) -> Self::Output {
67 let scale: $scalar = rhs.value();
68 Self::Output {
69 data: &self.data * scale,
70 context: self.context,
71 }
72 }
73 }
74 };
75}
76
77macro_rules! impl_div_scalar {
78 ($lhs:ty, $out:ty, $scalar:expr) => {
79 impl<'a, T: NalgebraScalar> Div<Scale<T>> for $lhs {
80 type Output = $out;
81 #[inline]
82 fn div(self, rhs: Scale<T>) -> Self::Output {
83 let inv_rhs: T = T::one() / rhs.value();
84 Self::Output {
85 data: self.data * inv_rhs,
86 context: self.context,
87 }
88 }
89 }
90 };
91}
92
93macro_rules! impl_mul_assign_scalar {
94 ($col_type:ty, $scalar:ty) => {
95 impl<'a, T: NalgebraScalar> MulAssign<Scale<T>> for $col_type {
96 #[inline]
97 fn mul_assign(&mut self, rhs: Scale<T>) {
98 let scale = rhs.value();
99 self.data *= scale;
100 }
101 }
102 };
103}
104
105impl_mul_scalar!(NalgebraVec<T>, NalgebraVec<T>, T);
106impl_mul_scalar!(&NalgebraVec<T>, NalgebraVec<T>, T);
107impl_mul_scalar!(NalgebraVecRef<'_, T>, NalgebraVec<T>, T);
108impl_mul_scalar!(NalgebraVecMut<'_, T>, NalgebraVec<T>, T);
109impl_div_scalar!(NalgebraVec<T>, NalgebraVec<T>, T);
110impl_mul_assign_scalar!(NalgebraVecMut<'a, T>, T);
111impl_mul_assign_scalar!(NalgebraVec<T>, T);
112
113impl_sub_assign!(NalgebraVec<T>, NalgebraVec<T>, NalgebraScalar);
114impl_sub_assign!(NalgebraVec<T>, &NalgebraVec<T>, NalgebraScalar);
115impl_sub_assign!(NalgebraVec<T>, NalgebraVecRef<'_, T>, NalgebraScalar);
116impl_sub_assign!(NalgebraVec<T>, &NalgebraVecRef<'_, T>, NalgebraScalar);
117
118impl_sub_assign!(NalgebraVecMut<'_, T>, NalgebraVec<T>, NalgebraScalar);
119impl_sub_assign!(NalgebraVecMut<'_, T>, &NalgebraVec<T>, NalgebraScalar);
120impl_sub_assign!(NalgebraVecMut<'_, T>, NalgebraVecRef<'_, T>, NalgebraScalar);
121impl_sub_assign!(
122 NalgebraVecMut<'_, T>,
123 &NalgebraVecRef<'_, T>,
124 NalgebraScalar
125);
126
127impl_add_assign!(NalgebraVec<T>, NalgebraVec<T>, NalgebraScalar);
128impl_add_assign!(NalgebraVec<T>, &NalgebraVec<T>, NalgebraScalar);
129impl_add_assign!(NalgebraVec<T>, NalgebraVecRef<'_, T>, NalgebraScalar);
130impl_add_assign!(NalgebraVec<T>, &NalgebraVecRef<'_, T>, NalgebraScalar);
131
132impl_add_assign!(NalgebraVecMut<'_, T>, NalgebraVec<T>, NalgebraScalar);
133impl_add_assign!(NalgebraVecMut<'_, T>, &NalgebraVec<T>, NalgebraScalar);
134impl_add_assign!(NalgebraVecMut<'_, T>, NalgebraVecRef<'_, T>, NalgebraScalar);
135impl_add_assign!(
136 NalgebraVecMut<'_, T>,
137 &NalgebraVecRef<'_, T>,
138 NalgebraScalar
139);
140
141impl_sub_both_ref!(
142 &NalgebraVec<T>,
143 &NalgebraVec<T>,
144 NalgebraVec<T>,
145 NalgebraScalar
146);
147impl_sub_rhs!(
148 &NalgebraVec<T>,
149 NalgebraVec<T>,
150 NalgebraVec<T>,
151 NalgebraScalar
152);
153impl_sub_both_ref!(
154 &NalgebraVec<T>,
155 NalgebraVecRef<'_, T>,
156 NalgebraVec<T>,
157 NalgebraScalar
158);
159impl_sub_both_ref!(
160 &NalgebraVec<T>,
161 &NalgebraVecRef<'_, T>,
162 NalgebraVec<T>,
163 NalgebraScalar
164);
165
166impl_sub_lhs!(
167 NalgebraVec<T>,
168 NalgebraVec<T>,
169 NalgebraVec<T>,
170 NalgebraScalar
171);
172impl_sub_lhs!(
173 NalgebraVec<T>,
174 &NalgebraVec<T>,
175 NalgebraVec<T>,
176 NalgebraScalar
177);
178impl_sub_lhs!(
179 NalgebraVec<T>,
180 NalgebraVecRef<'_, T>,
181 NalgebraVec<T>,
182 NalgebraScalar
183);
184impl_sub_lhs!(
185 NalgebraVec<T>,
186 &NalgebraVecRef<'_, T>,
187 NalgebraVec<T>,
188 NalgebraScalar
189);
190
191impl_sub_rhs!(
192 NalgebraVecRef<'_, T>,
193 NalgebraVec<T>,
194 NalgebraVec<T>,
195 NalgebraScalar
196);
197impl_sub_both_ref!(
198 NalgebraVecRef<'_, T>,
199 &NalgebraVec<T>,
200 NalgebraVec<T>,
201 NalgebraScalar
202);
203impl_sub_both_ref!(
204 NalgebraVecRef<'_, T>,
205 NalgebraVecRef<'_, T>,
206 NalgebraVec<T>,
207 NalgebraScalar
208);
209impl_sub_both_ref!(
210 NalgebraVecRef<'_, T>,
211 &NalgebraVecRef<'_, T>,
212 NalgebraVec<T>,
213 NalgebraScalar
214);
215
216impl_add_both_ref!(
217 &NalgebraVec<T>,
218 &NalgebraVec<T>,
219 NalgebraVec<T>,
220 NalgebraScalar
221);
222impl_add_rhs!(
223 &NalgebraVec<T>,
224 NalgebraVec<T>,
225 NalgebraVec<T>,
226 NalgebraScalar
227);
228impl_add_both_ref!(
229 &NalgebraVec<T>,
230 NalgebraVecRef<'_, T>,
231 NalgebraVec<T>,
232 NalgebraScalar
233);
234impl_add_both_ref!(
235 &NalgebraVec<T>,
236 &NalgebraVecRef<'_, T>,
237 NalgebraVec<T>,
238 NalgebraScalar
239);
240
241impl_add_lhs!(
242 NalgebraVec<T>,
243 NalgebraVec<T>,
244 NalgebraVec<T>,
245 NalgebraScalar
246);
247impl_add_lhs!(
248 NalgebraVec<T>,
249 &NalgebraVec<T>,
250 NalgebraVec<T>,
251 NalgebraScalar
252);
253impl_add_lhs!(
254 NalgebraVec<T>,
255 NalgebraVecRef<'_, T>,
256 NalgebraVec<T>,
257 NalgebraScalar
258);
259impl_add_lhs!(
260 NalgebraVec<T>,
261 &NalgebraVecRef<'_, T>,
262 NalgebraVec<T>,
263 NalgebraScalar
264);
265
266impl_add_rhs!(
267 NalgebraVecRef<'_, T>,
268 NalgebraVec<T>,
269 NalgebraVec<T>,
270 NalgebraScalar
271);
272impl_add_both_ref!(
273 NalgebraVecRef<'_, T>,
274 &NalgebraVec<T>,
275 NalgebraVec<T>,
276 NalgebraScalar
277);
278impl_add_both_ref!(
279 NalgebraVecRef<'_, T>,
280 NalgebraVecRef<'_, T>,
281 NalgebraVec<T>,
282 NalgebraScalar
283);
284impl_add_both_ref!(
285 NalgebraVecRef<'_, T>,
286 &NalgebraVecRef<'_, T>,
287 NalgebraVec<T>,
288 NalgebraScalar
289);
290
291impl_index!(NalgebraVec<T>, NalgebraScalar);
292impl_index_mut!(NalgebraVec<T>, NalgebraScalar);
293
294impl_index!(NalgebraVecRef<'_, T>, NalgebraScalar);
295
296impl VectorIndex for NalgebraIndex {
297 type C = NalgebraContext;
298 fn zeros(len: IndexType, ctx: Self::C) -> Self {
299 let data = DVector::from_element(len, 0);
300 Self { data, context: ctx }
301 }
302 fn len(&self) -> crate::IndexType {
303 self.data.len()
304 }
305 fn from_vec(v: Vec<IndexType>, ctx: Self::C) -> Self {
306 let data = DVector::from_vec(v);
307 Self { data, context: ctx }
308 }
309 fn clone_as_vec(&self) -> Vec<IndexType> {
310 self.data.iter().copied().collect()
311 }
312 fn context(&self) -> &Self::C {
313 &self.context
314 }
315}
316
317impl<'a, T: NalgebraScalar> VectorView<'a> for NalgebraVecRef<'a, T> {
318 type Owned = NalgebraVec<T>;
319
320 fn get_index(&self, index: IndexType) -> Self::T {
321 self.data[index]
322 }
323
324 fn into_owned(self) -> Self::Owned {
325 Self::Owned {
326 data: self.data.into_owned(),
327 context: self.context,
328 }
329 }
330 fn squared_norm(&self, y: &Self::Owned, atol: &Self::Owned, rtol: Self::T) -> Self::T {
331 let mut acc = T::zero();
332 if y.len() != self.data.len() || y.len() != atol.len() {
333 panic!("Vector lengths do not match");
334 }
335 for i in 0..self.data.len() {
336 let yi = unsafe { y.data.get_unchecked(i) };
337 let ai = unsafe { atol.data.get_unchecked(i) };
338 let xi = unsafe { self.data.get_unchecked(i) };
339 let term = *xi / (yi.abs() * rtol + *ai);
340 acc += term * term;
341 }
342 acc / Self::T::from_f64(self.data.len() as f64).unwrap()
343 }
344}
345
346impl<'a, T: NalgebraScalar> VectorViewMut<'a> for NalgebraVecMut<'a, T> {
347 type Owned = NalgebraVec<T>;
348 type View = NalgebraVecRef<'a, T>;
349 type Index = NalgebraIndex;
350 fn copy_from(&mut self, other: &Self::Owned) {
351 self.data.copy_from(&other.data);
352 }
353 fn copy_from_view(&mut self, other: &Self::View) {
354 self.data.copy_from(&other.data);
355 }
356 fn set_index(&mut self, index: IndexType, value: Self::T) {
357 self.data[index] = value;
358 }
359 fn axpy(&mut self, alpha: Self::T, x: &Self::Owned, beta: Self::T) {
360 self.data.axpy(alpha, &x.data, beta);
361 }
362}
363
364impl<T: NalgebraScalar> VectorHost for NalgebraVec<T> {
365 fn as_slice(&self) -> &[Self::T] {
366 self.data.as_slice()
367 }
368 fn as_mut_slice(&mut self) -> &mut [Self::T] {
369 self.data.as_mut_slice()
370 }
371}
372
373impl<T: NalgebraScalar> Vector for NalgebraVec<T> {
374 type View<'a> = NalgebraVecRef<'a, T>;
375 type ViewMut<'a> = NalgebraVecMut<'a, T>;
376 type Index = NalgebraIndex;
377 fn len(&self) -> IndexType {
378 self.data.len()
379 }
380 fn inner_mut(&mut self) -> &mut Self::Inner {
381 &mut self.data
382 }
383 fn context(&self) -> &Self::C {
384 &self.context
385 }
386 fn norm(&self, k: i32) -> Self::T {
387 self.data.apply_norm(&LpNorm(k))
388 }
389 fn get_index(&self, index: IndexType) -> Self::T {
390 self.data[index]
391 }
392 fn set_index(&mut self, index: IndexType, value: Self::T) {
393 self.data[index] = value;
394 }
395 fn squared_norm(&self, y: &Self, atol: &Self, rtol: Self::T) -> Self::T {
396 let mut acc = T::zero();
397 if y.len() != self.len() || y.len() != atol.len() {
398 panic!("Vector lengths do not match");
399 }
400 for i in 0..self.len() {
401 let yi = unsafe { y.data.get_unchecked(i) };
402 let ai = unsafe { atol.data.get_unchecked(i) };
403 let xi = unsafe { self.data.get_unchecked(i) };
404 let term = *xi / (yi.abs() * rtol + *ai);
405 acc += term * term;
406 }
407 acc / Self::T::from_f64(self.len() as f64).unwrap()
408 }
409 fn as_view(&self) -> Self::View<'_> {
410 Self::View {
411 data: self.data.as_view(),
412 context: self.context,
413 }
414 }
415 fn as_view_mut(&mut self) -> Self::ViewMut<'_> {
416 Self::ViewMut {
417 data: self.data.as_view_mut(),
418 context: self.context,
419 }
420 }
421 fn get_batch(&self, batch: usize) -> Self::View<'_> {
422 assert!(
423 batch == 0,
424 "NalgebraVec does not support batching (nbatch > 1)."
425 );
426 self.as_view()
427 }
428 fn get_batch_mut(&mut self, batch: usize) -> Self::ViewMut<'_> {
429 assert!(
430 batch == 0,
431 "NalgebraVec does not support batching (nbatch > 1)."
432 );
433 self.as_view_mut()
434 }
435 fn copy_from(&mut self, other: &Self) {
436 self.data.copy_from(&other.data);
437 }
438 fn fill(&mut self, value: Self::T) {
439 self.data.iter_mut().for_each(|x: &mut _| *x = value);
440 }
441 fn copy_from_view(&mut self, other: &Self::View<'_>) {
442 self.data.copy_from(&other.data);
443 }
444 fn from_element(nstates: usize, value: T, ctx: Self::C) -> Self {
445 let data = DVector::from_element(nstates, value);
446 Self { data, context: ctx }
447 }
448 fn from_vec(vec: Vec<T>, ctx: Self::C) -> Self {
449 let data = DVector::from_vec(vec);
450 Self { data, context: ctx }
451 }
452 fn from_slice(slice: &[T], ctx: Self::C) -> Self {
453 let data = DVector::from_column_slice(slice);
454 Self { data, context: ctx }
455 }
456 fn clone_as_vec(&self) -> Vec<Self::T> {
457 self.data.iter().copied().collect()
458 }
459 fn zeros(nstates: usize, ctx: Self::C) -> Self {
460 let data = DVector::zeros(nstates);
461 Self { data, context: ctx }
462 }
463 fn axpy(&mut self, alpha: T, x: &Self, beta: T) {
464 self.data.axpy(alpha, &x.data, beta);
465 }
466 fn axpy_v(&mut self, alpha: Self::T, x: &Self::View<'_>, beta: Self::T) {
467 self.data.axpy(alpha, &x.data, beta);
468 }
469 fn batched_axpy(&mut self, alpha: &[Self::T], x: &Self, beta: Self::T) {
470 assert_eq!(
471 alpha.len(),
472 1,
473 "NalgebraVec does not support batching (nbatch > 1)."
474 );
475 self.axpy(alpha[0], x, beta);
476 }
477 fn component_div_assign(&mut self, other: &Self) {
478 self.data.component_div_assign(&other.data);
479 }
480 fn component_mul_assign(&mut self, other: &Self) {
481 self.data.component_mul_assign(&other.data);
482 }
483
484 fn root_finding(&self, g1: &Self) -> (bool, Self::T, i32) {
485 let mut max_frac = T::zero();
486 let mut max_frac_index = -1;
487 let mut found_root = false;
488 assert_eq!(self.len(), g1.len(), "Vector lengths do not match");
489 for i in 0..self.len() {
490 let g0 = unsafe { *self.data.get_unchecked(i) };
491 let g1 = unsafe { *g1.data.get_unchecked(i) };
492 if g1 == T::zero() {
493 found_root = true;
494 }
495 if g0 * g1 < T::zero() {
496 let frac = (g1 / (g1 - g0)).abs();
497 if frac > max_frac {
498 max_frac = frac;
499 max_frac_index = i as i32;
500 }
501 }
502 }
503 (found_root, max_frac, max_frac_index)
504 }
505
506 fn assign_at_indices(&mut self, indices: &Self::Index, value: Self::T) {
507 for i in indices.data.iter() {
508 self[*i] = value;
509 }
510 }
511
512 fn copy_from_indices(&mut self, other: &Self, indices: &Self::Index) {
513 for i in indices.data.iter() {
514 self[*i] = other[*i];
515 }
516 }
517
518 fn gather(&mut self, other: &Self, indices: &Self::Index) {
519 assert_eq!(self.len(), indices.len(), "Vector lengths do not match");
520 for (s, o) in self.data.iter_mut().zip(indices.data.iter()) {
521 *s = other[*o];
522 }
523 }
524
525 fn scatter(&self, indices: &Self::Index, other: &mut Self) {
526 assert_eq!(self.len(), indices.len(), "Vector lengths do not match");
527 for (s, o) in self.data.iter().zip(indices.data.iter()) {
528 other[*o] = *s;
529 }
530 }
531}
532
533#[cfg(test)]
534mod tests {
535 use super::*;
536
537 #[test]
538 fn test_error_norm() {
539 let v = NalgebraVec::from_vec(vec![1.0, -2.0, 3.0], Default::default());
540 let y = NalgebraVec::from_vec(vec![1.0, 2.0, 3.0], Default::default());
541 let atol = NalgebraVec::from_vec(vec![0.1, 0.2, 0.3], Default::default());
542 let rtol = 0.1;
543 let mut tmp = y.clone() * Scale(rtol);
544 tmp += &atol;
545 let mut r = v.clone();
546 r.component_div_assign(&tmp);
547 let errorn_check = r.data.norm_squared() / 3.0;
548 assert_eq!(v.squared_norm(&y, &atol, rtol), errorn_check);
549 let vview = v.as_view();
550 assert_eq!(
551 VectorView::squared_norm(&vview, &y, &atol, rtol),
552 errorn_check
553 );
554 }
555
556 #[test]
557 fn test_root_finding() {
558 super::super::tests::test_root_finding::<NalgebraVec<f64>>();
559 }
560
561 #[test]
562 fn test_from_slice() {
563 let slice = [1.0, 2.0, 3.0];
564 let v = NalgebraVec::from_slice(&slice, Default::default());
565 assert_eq!(v.clone_as_vec(), slice);
566 }
567
568 #[test]
569 fn test_into() {
570 let vec = DVector::from_vec(vec![1.0, 2.0, 3.0]);
571 let v: NalgebraVec<f64> = vec.into();
572 assert_eq!(v.clone_as_vec(), vec![1.0, 2.0, 3.0]);
573 }
574
575 super::super::generate_vector_tests_nonbatched!(nalgebra, NalgebraVec<f64>);
576}