1use faer::sparse::{SparseColMatRef, SymbolicSparseColMatRef};
4use faer_traits::math_utils::{copy, from_f64, from_real, max, mul, zero};
5use faer_traits::{ComplexField, Index};
6
7use super::{BoundEstimate, Coeffs, Poly, PolyError, PolyKind, PolyParams};
8use crate::util::spd_bounds::{gershgorin_bounds, power_iteration_max};
9
10const LAMBDA_MIN_FLOOR: f64 = 1e-4;
14
15fn check_square<I: Index, T: ComplexField>(
16 a: SparseColMatRef<'_, I, T>,
17) -> Result<(), PolyError> {
18 if a.nrows() != a.ncols() {
19 return Err(PolyError::NonSquareMatrix {
20 nrows: a.nrows(),
21 ncols: a.ncols(),
22 });
23 }
24 Ok(())
25}
26
27fn resolve_kind<T: ComplexField>(kind: &PolyKind) -> Result<Coeffs<T>, PolyError> {
28 match *kind {
29 PolyKind::Neumann { omega } => {
30 if !omega.is_finite() || omega <= 0.0 {
31 return Err(PolyError::InvalidOmega);
32 }
33 Ok(Coeffs::Neumann {
34 omega: from_f64::<T>(omega),
35 })
36 }
37 PolyKind::Chebyshev {
38 lambda_min,
39 lambda_max,
40 } => {
41 if !lambda_min.is_finite()
42 || !lambda_max.is_finite()
43 || lambda_min <= 0.0
44 || lambda_max <= lambda_min
45 {
46 return Err(PolyError::InvalidBounds);
47 }
48 Ok(Coeffs::Chebyshev {
49 lambda_min: from_f64::<T>(lambda_min),
50 lambda_max: from_f64::<T>(lambda_max),
51 })
52 }
53 }
54}
55
56fn resolve_bounds<I: Index, T: ComplexField>(
58 a: SparseColMatRef<'_, I, T>,
59 estimate: &BoundEstimate,
60) -> Result<(T::Real, T::Real), PolyError> {
61 match *estimate {
62 BoundEstimate::Manual {
63 lambda_min,
64 lambda_max,
65 } => {
66 if !lambda_min.is_finite()
67 || !lambda_max.is_finite()
68 || lambda_min <= 0.0
69 || lambda_max <= lambda_min
70 {
71 return Err(PolyError::InvalidBounds);
72 }
73 Ok((from_f64::<T::Real>(lambda_min), from_f64::<T::Real>(lambda_max)))
74 }
75 BoundEstimate::Gershgorin => {
76 let (glo, ghi) = gershgorin_bounds(a);
77 finalize_bounds::<T>(glo, ghi)
78 }
79 BoundEstimate::PowerIteration { iters } => {
80 let hi = power_iteration_max(a, iters.max(1));
81 let (glo, _) = gershgorin_bounds(a);
82 finalize_bounds::<T>(glo, hi)
83 }
84 }
85}
86
87fn finalize_bounds<T: ComplexField>(glo: T::Real, hi: T::Real) -> Result<(T::Real, T::Real), PolyError> {
91 if hi.partial_cmp(&zero::<T::Real>()) != Some(core::cmp::Ordering::Greater) {
92 return Err(PolyError::InvalidBounds);
93 }
94 let floor = mul(&hi, &from_f64::<T::Real>(LAMBDA_MIN_FLOOR));
95 let lo = max(&glo, &floor);
96 Ok((lo, hi))
97}
98
99impl<I: Index, T: ComplexField> Poly<I, T> {
100 pub fn try_new(a: SparseColMatRef<'_, I, T>, params: PolyParams) -> Result<Self, PolyError> {
109 check_square(a)?;
110 if params.degree == 0 {
111 return Err(PolyError::ZeroDegree);
112 }
113 let coeffs = resolve_kind::<T>(¶ms.kind)?;
114 Ok(Self::assemble(a, params.degree, coeffs, None))
115 }
116
117 pub fn try_new_auto(
129 a: SparseColMatRef<'_, I, T>,
130 degree: usize,
131 estimate: BoundEstimate,
132 ) -> Result<Self, PolyError> {
133 check_square(a)?;
134 if degree == 0 {
135 return Err(PolyError::ZeroDegree);
136 }
137 let (lo, hi) = resolve_bounds::<I, T>(a, &estimate)?;
138 let coeffs = Coeffs::Chebyshev {
139 lambda_min: from_real::<T>(&lo),
140 lambda_max: from_real::<T>(&hi),
141 };
142 Ok(Self::assemble(a, degree, coeffs, Some(estimate)))
143 }
144
145 fn assemble(
146 a: SparseColMatRef<'_, I, T>,
147 degree: usize,
148 coeffs: Coeffs<T>,
149 recompute: Option<BoundEstimate>,
150 ) -> Self {
151 let n = a.ncols();
152 let mut a_col_ptr: Vec<I> = Vec::with_capacity(n + 1);
153 a_col_ptr.push(I::truncate(0));
154 let mut a_row_idx: Vec<I> = Vec::new();
155 let mut a_values: Vec<T> = Vec::new();
156 for j in 0..n {
157 let rows = a.symbolic().row_idx_of_col_raw(j);
158 let vals = a.val_of_col(j);
159 a_row_idx.extend_from_slice(rows);
160 for v in vals {
161 a_values.push(copy(v));
162 }
163 a_col_ptr.push(I::truncate(a_row_idx.len()));
164 }
165 Self {
166 dim: n,
167 degree,
168 a_col_ptr,
169 a_row_idx,
170 a_values,
171 coeffs,
172 recompute,
173 }
174 }
175
176 pub fn refactorize(&mut self, a: SparseColMatRef<'_, I, T>) -> Result<(), PolyError> {
187 let n = self.dim;
188 if a.nrows() != n || a.ncols() != n {
189 return Err(PolyError::PatternMismatch);
190 }
191 for j in 0..n {
192 let vals = a.val_of_col(j);
193 let start = self.a_col_ptr[j].zx();
194 let end = self.a_col_ptr[j + 1].zx();
195 if end - start != vals.len() {
196 return Err(PolyError::PatternMismatch);
197 }
198 for (k, v) in vals.iter().enumerate() {
199 self.a_values[start + k] = copy(v);
200 }
201 }
202 if let Some(est) = self.recompute {
203 let (lo, hi) = resolve_bounds::<I, T>(a, &est)?;
204 self.coeffs = Coeffs::Chebyshev {
205 lambda_min: from_real::<T>(&lo),
206 lambda_max: from_real::<T>(&hi),
207 };
208 }
209 Ok(())
210 }
211
212 #[inline]
214 pub(crate) fn a_view(&self) -> SparseColMatRef<'_, I, T> {
215 let sym = unsafe {
216 SymbolicSparseColMatRef::<'_, I>::new_unchecked(
217 self.dim,
218 self.dim,
219 &self.a_col_ptr,
220 None,
221 &self.a_row_idx,
222 )
223 };
224 SparseColMatRef::new(sym, &self.a_values)
225 }
226}