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