use std::iter::Sum;
use nalgebra::{Complex, DMatrix, DVector, RealField};
use num_traits::Float;
use thiserror::Error;
use crate::{Continuous, FrequencyResponse, Polynomial, TransferFunction};
#[derive(Error, Debug)]
pub enum VectorFittingError {
#[error("No sampled data provided")]
EmptyData,
#[error("Matrix is singular or numerically unstable (check order and data)")]
SingularMatrix,
#[error("Poles not set")]
PolesNotSet,
#[error("Failed to find zeros")]
ZerosNotFound,
#[error("Failed to iterate (check data and iterations)")]
IterationError,
#[error("Invalid weights: {len} for {samples} samples, all finite and non-negative needed")]
InvalidWeights { len: usize, samples: usize },
}
#[derive(Debug, Clone)]
pub struct VectorFittingOptions<T> {
pub max_iter: usize,
pub tol: T,
pub fit_d: bool,
pub fit_e: bool,
}
impl<T: Float> Default for VectorFittingOptions<T> {
fn default() -> Self {
Self {
max_iter: 10,
tol: T::from(1e-8).unwrap(),
fit_d: false,
fit_e: false,
}
}
}
#[derive(Debug, Clone)]
pub struct VectorFittingResult<T> {
pub poles: Vec<Complex<T>>,
pub residues: Vec<Complex<T>>,
pub d: T,
pub e: T,
pub rms_errors: Vec<T>,
}
#[derive(Clone, Copy, Debug, PartialEq)]
enum Block<T> {
Real(T),
Pair(Complex<T>),
}
impl<T: Float> Block<T> {
fn size(&self) -> usize {
match self {
Block::Real(_) => 1,
Block::Pair(_) => 2,
}
}
fn pole(&self) -> Complex<T> {
match *self {
Block::Real(p) => Complex::new(p, T::zero()),
Block::Pair(p) => p,
}
}
fn basis(&self, s: Complex<T>) -> [Complex<T>; 2] {
match *self {
Block::Real(p) => [Complex::from(T::one()) / (s - p), Complex::from(T::zero())],
Block::Pair(p) => {
let (u, v) = (Complex::from(T::one()) / (s - p), Complex::from(T::one()) / (s - p.conj()));
let j = Complex::new(T::zero(), T::one());
[u + v, j * (u - v)]
}
}
}
}
fn poles_of<T: Float>(blocks: &[Block<T>]) -> Vec<Complex<T>> {
blocks
.iter()
.flat_map(|b| match *b {
Block::Real(p) => vec![Complex::new(p, T::zero())],
Block::Pair(p) => vec![p, p.conj()],
})
.collect()
}
fn residues_of<T: Float>(blocks: &[Block<T>], c: &[T]) -> Vec<Complex<T>> {
let mut residues = Vec::with_capacity(c.len());
let mut i = 0;
for b in blocks {
match b {
Block::Real(_) => residues.push(Complex::new(c[i], T::zero())),
Block::Pair(_) => {
residues.push(Complex::new(c[i], c[i + 1]));
residues.push(Complex::new(c[i], -c[i + 1]));
}
}
i += b.size();
}
residues
}
fn real_model<T: Float>(poles: &[Complex<T>], residues: &[Complex<T>]) -> (Vec<Block<T>>, Vec<T>) {
let two = T::from(2).unwrap();
let mut used = vec![false; poles.len()];
let (mut blocks, mut c) = (Vec::new(), Vec::new());
for i in 0..poles.len() {
if used[i] {
continue;
}
used[i] = true;
let (p, r) = (poles[i], residues[i]);
if p.im == T::zero() {
blocks.push(Block::Real(p.re));
c.push(r.re);
continue;
}
let partner = (0..poles.len())
.filter(|&j| !used[j] && poles[j].im != T::zero())
.min_by(|&a, &b| (poles[a] - p.conj()).norm().partial_cmp(&(poles[b] - p.conj()).norm()).unwrap_or(std::cmp::Ordering::Equal))
.filter(|&j| (poles[j] - p.conj()).norm() <= T::from(1e-6).unwrap() * p.norm());
let (upper, residue) = match partner {
Some(j) => {
used[j] = true;
let mean = (r + residues[j].conj()) / two;
if p.im > T::zero() { (p, mean) } else { (p.conj(), mean.conj()) }
}
None if p.im > T::zero() => (p, r),
None => (p.conj(), r.conj()),
};
blocks.push(Block::Pair(upper));
c.push(residue.re);
c.push(residue.im);
}
(blocks, c)
}
fn multiply<T: Float>(a: &[T], b: &[T]) -> Vec<T> {
let mut out = vec![T::zero(); a.len() + b.len() - 1];
for (i, &x) in a.iter().enumerate() {
for (j, &y) in b.iter().enumerate() {
out[i + j] = out[i + j] + x * y;
}
}
out
}
fn add_into<T: Float>(a: &mut [T], b: &[T]) {
let offset = a.len() - b.len();
for (i, &c) in b.iter().enumerate() {
a[offset + i] = a[offset + i] + c;
}
}
impl<T: Float + RealField> From<VectorFittingResult<T>> for TransferFunction<T, Continuous> {
fn from(val: VectorFittingResult<T>) -> Self {
let n = val.poles.len();
if n == 0 {
return TransferFunction::from_polynomials(
Polynomial(vec![val.e, val.d]),
Polynomial(vec![T::one()]),
);
}
let two = T::from(2).unwrap();
let (blocks, c) = real_model(&val.poles, &val.residues);
let factors: Vec<Vec<T>> = blocks
.iter()
.map(|b| match *b {
Block::Real(p) => vec![T::one(), -p],
Block::Pair(p) => vec![T::one(), -two * p.re, p.norm_sqr()],
})
.collect();
let denom = factors.iter().fold(vec![T::one()], |acc, f| multiply(&acc, f));
let degree = denom.len() - 1;
let mut numer = vec![T::zero(); degree + 2];
let mut i = 0;
for (k, b) in blocks.iter().enumerate() {
let term = match *b {
Block::Real(_) => vec![c[i]],
Block::Pair(p) => vec![two * c[i], -two * (c[i] * p.re + c[i + 1] * p.im)],
};
let others = factors.iter().enumerate().filter(|&(j, _)| j != k).fold(vec![T::one()], |acc, (_, f)| multiply(&acc, f));
add_into(&mut numer, &multiply(&term, &others));
i += b.size();
}
add_into(&mut numer, &denom.iter().map(|&a| a * val.d).collect::<Vec<T>>());
let mut s_denom: Vec<T> = denom.iter().map(|&a| a * val.e).collect();
s_denom.push(T::zero());
add_into(&mut numer, &s_denom);
let omega_ref = val.poles.iter().map(|p| p.norm()).fold(T::zero(), Float::max);
let omega_ref = if omega_ref > T::zero() { omega_ref } else { T::one() };
let top = numer.len() - 1;
let terms: Vec<T> = numer.iter().enumerate().map(|(i, c)| Float::abs(*c) * Float::powi(omega_ref, (top - i) as i32)).collect();
let largest = terms.iter().copied().fold(T::zero(), Float::max);
let first_valid = terms.iter()
.position(|&t| t > T::from(1e-5).unwrap() * largest)
.unwrap_or(numer.len());
numer.drain(..first_valid);
TransferFunction::from_polynomials(Polynomial(numer), Polynomial(denom))
}
}
pub fn identify<T: Float + RealField + Sum>(
samples: &[FrequencyResponse<T>],
n_poles: usize,
opts: &VectorFittingOptions<T>,
) -> Result<VectorFittingResult<T>, VectorFittingError> {
fit(samples, None, n_poles, opts)
}
pub fn identify_weighted<T: Float + RealField + Sum>(
samples: &[FrequencyResponse<T>],
weights: &[T],
n_poles: usize,
opts: &VectorFittingOptions<T>,
) -> Result<VectorFittingResult<T>, VectorFittingError> {
if weights.len() != samples.len() || weights.iter().any(|w| !(w.is_finite() && *w >= T::zero())) {
return Err(VectorFittingError::InvalidWeights { len: weights.len(), samples: samples.len() });
}
fit(samples, Some(weights), n_poles, opts)
}
fn fit<T: Float + RealField + Sum>(
samples: &[FrequencyResponse<T>],
weights: Option<&[T]>,
n_poles: usize,
opts: &VectorFittingOptions<T>,
) -> Result<VectorFittingResult<T>, VectorFittingError> {
if samples.is_empty() {
return Err(VectorFittingError::EmptyData);
}
let weight = |i: usize| weights.map_or(T::one(), |w| w[i]);
let mut blocks = initial_blocks(samples, n_poles);
let mut rms_errors = Vec::new();
for _iter in 0..opts.max_iter {
let (c, d, e, c_tilde) = least_squares(samples, &weight, &blocks, opts, true)?;
rms_errors.push(rms(samples, &weight, &blocks, &c, d, e));
let new_blocks = zeros_of_sigma(&blocks, &c_tilde)?;
let epsilon = T::from(1e-30).unwrap();
let max_change = if new_blocks.len() == blocks.len() && new_blocks.iter().zip(&blocks).all(|(a, b)| a.size() == b.size()) {
blocks
.iter()
.zip(new_blocks.iter())
.map(|(old, new)| (new.pole() - old.pole()).norm() / (old.pole().norm() + epsilon))
.fold(T::zero(), Float::max)
} else {
T::infinity()
};
blocks = new_blocks;
if max_change < opts.tol {
break;
}
}
let (c, d, e, _) = least_squares(samples, &weight, &blocks, opts, false)?;
Ok(VectorFittingResult { poles: poles_of(&blocks), residues: residues_of(&blocks, &c), d, e, rms_errors })
}
fn initial_blocks<T: Float>(samples: &[FrequencyResponse<T>], n_poles: usize) -> Vec<Block<T>> {
let w_min = samples.iter().fold(T::infinity(), |a, b| T::min(a, b.omega));
let w_max = samples.iter().fold(T::neg_infinity(), |a, b| T::max(a, b.omega));
let n_pairs = n_poles / 2;
let mut blocks = Vec::with_capacity(n_pairs + 1);
for i in 0..n_pairs {
let t = if n_pairs > 1 {
T::from(i as f64 / (n_pairs - 1) as f64).unwrap()
} else {
T::from(0.5).unwrap()
};
let w = w_min * (w_max / w_min).powf(t);
blocks.push(Block::Pair(Complex::new(-T::from(0.01).unwrap() * w, w)));
}
if n_poles % 2 == 1 {
blocks.push(Block::Real(-(w_min * w_max).sqrt()));
}
blocks
}
#[allow(clippy::type_complexity)]
fn least_squares<T: Float + RealField>(
samples: &[FrequencyResponse<T>],
weight: &impl Fn(usize) -> T,
blocks: &[Block<T>],
opts: &VectorFittingOptions<T>,
with_sigma: bool,
) -> Result<(Vec<T>, T, T, Vec<T>), VectorFittingError> {
let n: usize = blocks.iter().map(Block::size).sum();
let n_extra = opts.fit_d as usize + opts.fit_e as usize;
let columns = n + n_extra + if with_sigma { n } else { 0 };
let rows = 2 * samples.len();
let mut a = DMatrix::<T>::zeros(rows, columns);
let mut b = DVector::<T>::zeros(rows);
for (i, sample) in samples.iter().enumerate() {
let w = weight(i);
let s = Complex::new(T::zero(), sample.omega);
let g = sample.value;
let mut row = Vec::with_capacity(columns);
for block in blocks {
row.extend_from_slice(&block.basis(s)[..block.size()]);
}
if opts.fit_d {
row.push(Complex::from(T::one()));
}
if opts.fit_e {
row.push(s);
}
if with_sigma {
for block in blocks {
row.extend(block.basis(s)[..block.size()].iter().map(|&phi| -g * phi));
}
}
for (j, v) in row.iter().enumerate() {
a[(2 * i, j)] = w * v.re;
a[(2 * i + 1, j)] = w * v.im;
}
b[2 * i] = w * g.re;
b[2 * i + 1] = w * g.im;
}
let scale: Vec<T> = (0..columns)
.map(|j| {
let norm = a.column(j).norm();
if norm > T::zero() { norm } else { T::one() }
})
.collect();
for (j, &sc) in scale.iter().enumerate() {
a.column_mut(j).unscale_mut(sc);
}
let x = a
.svd(true, true)
.solve(&b, T::from(1e-12).unwrap())
.map_err(|_| VectorFittingError::SingularMatrix)?;
let x: Vec<T> = x.iter().zip(&scale).map(|(&v, &sc)| v / sc).collect();
if x.iter().any(|v| !v.is_finite()) {
return Err(VectorFittingError::SingularMatrix);
}
let c = x[..n].to_vec();
let mut k = n;
let d = if opts.fit_d { k += 1; x[k - 1] } else { T::zero() };
let e = if opts.fit_e { k += 1; x[k - 1] } else { T::zero() };
let c_tilde = x[k..].to_vec();
Ok((c, d, e, c_tilde))
}
fn zeros_of_sigma<T: Float + RealField>(blocks: &[Block<T>], c_tilde: &[T]) -> Result<Vec<Block<T>>, VectorFittingError> {
let n = c_tilde.len();
if n == 0 {
return Err(VectorFittingError::PolesNotSet);
}
let mut h = DMatrix::<T>::zeros(n, n);
let mut b = DVector::<T>::zeros(n);
let mut i = 0;
for block in blocks {
match *block {
Block::Real(p) => {
h[(i, i)] = p;
b[i] = T::one();
}
Block::Pair(p) => {
h[(i, i)] = p.re;
h[(i, i + 1)] = p.im;
h[(i + 1, i)] = -p.im;
h[(i + 1, i + 1)] = p.re;
b[i] = T::from(2).unwrap();
}
}
i += block.size();
}
for r in 0..n {
for col in 0..n {
h[(r, col)] -= b[r] * c_tilde[col];
}
}
let eigenvalues = nalgebra::linalg::Schur::try_new(h, T::default_epsilon(), 0)
.ok_or(VectorFittingError::ZerosNotFound)?
.complex_eigenvalues();
let mut new_blocks = Vec::with_capacity(n);
for z in eigenvalues.iter() {
let z = Complex::new(z.re, z.im);
if !(z.re.is_finite() && z.im.is_finite()) {
return Err(VectorFittingError::ZerosNotFound);
}
let re = -Float::abs(z.re);
if Float::abs(z.im) <= T::from(1e-9).unwrap() * z.norm() {
new_blocks.push(Block::Real(re));
} else if z.im > T::zero() {
new_blocks.push(Block::Pair(Complex::new(re, z.im)));
}
}
if new_blocks.iter().map(Block::size).sum::<usize>() != n {
return Err(VectorFittingError::ZerosNotFound);
}
new_blocks.sort_by(|a, b| a.pole().norm().partial_cmp(&b.pole().norm()).unwrap_or(std::cmp::Ordering::Equal));
Ok(new_blocks)
}
fn rms<T: Float + RealField + Sum>(
samples: &[FrequencyResponse<T>],
weight: &impl Fn(usize) -> T,
blocks: &[Block<T>],
c: &[T],
d: T,
e: T,
) -> T {
let n = T::from(samples.len()).unwrap();
let sum_sq: T = samples
.iter()
.enumerate()
.map(|(i, sample)| {
let s = Complex::new(T::zero(), sample.omega);
let mut fit = Complex::new(d, T::zero()) + s * e;
let mut k = 0;
for block in blocks {
for phi in &block.basis(s)[..block.size()] {
fit += *phi * c[k];
k += 1;
}
}
(fit - sample.value).norm_sqr() * weight(i) * weight(i)
})
.sum();
Float::sqrt(sum_sq / n)
}
#[cfg(test)]
mod tests {
use super::*;
use num_complex::Complex64;
#[test]
fn test_zeros_of_sigma() {
let blocks = vec![Block::Real(-3.0), Block::Pair(Complex64::new(-1.0, 2.0)), Block::Pair(Complex64::new(-0.5, 10.0))];
let c_tilde = vec![0.7, 1.0, -0.4, 2.0, 0.3];
let poles = poles_of(&blocks);
let residues = residues_of(&blocks, &c_tilde);
assert_eq!(residues[1], residues[2].conj());
let zeros = zeros_of_sigma(&blocks, &c_tilde).unwrap();
assert_eq!(zeros.iter().map(Block::size).sum::<usize>(), 5);
for z in poles_of(&zeros) {
let sigma = |z: Complex64| residues.iter().zip(&poles).fold(Complex64::new(1.0, 0.0), |acc, (r, p)| acc + r / (z - p));
let value = sigma(z).norm().min(sigma(Complex64::new(-z.re, z.im)).norm());
assert!(value < 1e-8, "σ({z}) = {value:.2e}");
}
}
#[test]
fn test_pair_basis() {
let p = Complex64::new(-1.0, 3.0);
let s = Complex64::new(0.0, 2.0);
let (c1, c2) = (0.4, -1.5);
let [phi1, phi2] = Block::Pair(p).basis(s);
let r = Complex64::new(c1, c2);
let expected = r / (s - p) + r.conj() / (s - p.conj());
assert!((phi1 * c1 + phi2 * c2 - expected).norm() < 1e-14);
}
}