feos-campd 0.2.0

Computer-aided molecular and process design using the FeOs framework.
Documentation
use std::array;
use std::iter::Sum;
use std::ops::{Add, AddAssign, Mul};

#[derive(Clone, Debug, PartialEq)]
#[allow(clippy::type_complexity)]
pub struct Polynomial2<const M: usize, const N: usize>(
    pub [[(Vec<Vec<isize>>, Vec<isize>, isize); N]; M],
);

#[derive(Clone, Debug, PartialEq)]
pub struct Polynomial<const N: usize>(pub [(Vec<isize>, isize); N]);

#[derive(Clone, Copy, Debug)]
pub enum Variable {
    Const(isize),
    Var(i32),
}

impl<const N: usize> Polynomial<N> {
    pub fn zero(variables: usize) -> Self {
        Self(array::from_fn(|_| (vec![0; variables], 0)))
    }

    pub fn new(variables: usize, p: [Variable; N]) -> Self {
        let mut coefs = Self::zero(variables).0;
        for i in 0..N {
            match p[i] {
                Variable::Const(c) => coefs[i].1 += c,
                Variable::Var(p) => coefs[i].0[p as usize] += 1,
            }
            if i + 1 < N {
                match p[i + 1] {
                    Variable::Const(c) => coefs[i].1 -= c,
                    Variable::Var(p) => coefs[i].0[p as usize] -= 1,
                }
            }
        }
        Self(coefs)
    }

    pub fn upgrade<const X: usize>(self) -> Polynomial<X> {
        let variables = self.0[0].0.len();
        Polynomial(array::from_fn(|i| {
            if i + N >= X {
                self.0[i + N - X].clone()
            } else {
                (vec![0; variables], 0)
            }
        }))
    }
}

impl<const M: usize, const N: usize> Polynomial2<M, N> {
    pub fn zero(variables: usize) -> Self {
        Self(array::from_fn(|_| {
            array::from_fn(|_| (vec![vec![0; variables]; variables], vec![0; variables], 0))
        }))
    }
}

impl<const N: usize> Add for Polynomial<N> {
    type Output = Self;
    fn add(mut self, other: Self) -> Self {
        self += &other;
        self
    }
}

impl<const N: usize> AddAssign<&Polynomial<N>> for Polynomial<N> {
    fn add_assign(&mut self, other: &Self) {
        for i in 0..N {
            for (s, &o) in self.0[i].0.iter_mut().zip(&other.0[i].0) {
                *s += o;
            }
            self.0[i].1 += other.0[i].1;
        }
    }
}

impl<const M: usize, const N: usize> AddAssign for Polynomial2<M, N> {
    fn add_assign(&mut self, other: Self) {
        for i in 0..M {
            for j in 0..N {
                for (s, o) in self.0[i][j].0.iter_mut().zip(&other.0[i][j].0) {
                    for (s, &o) in s.iter_mut().zip(o) {
                        *s += o;
                    }
                }
                for (s, o) in self.0[i][j].1.iter_mut().zip(&other.0[i][j].1) {
                    *s += o;
                }
                self.0[i][j].2 += other.0[i][j].2
            }
        }
    }
}

impl<const N: usize> Sum for Polynomial<N> {
    fn sum<I: Iterator<Item = Polynomial<N>>>(iter: I) -> Self {
        iter.reduce(|a, b| a + b).unwrap()
    }
}

impl<const M: usize, const N: usize> Mul<Polynomial<N>> for Polynomial<M> {
    type Output = Polynomial2<M, N>;
    fn mul(self, other: Polynomial<N>) -> Polynomial2<M, N> {
        Polynomial2(array::from_fn(|i| {
            array::from_fn(|j| {
                (
                    self.0[i]
                        .0
                        .iter()
                        .map(|&a| other.0[j].0.iter().map(|&b| a * b).collect())
                        .collect(),
                    self.0[i]
                        .0
                        .iter()
                        .zip(&other.0[j].0)
                        .map(|(&a, &b)| a * other.0[j].1 + self.0[i].1 * b)
                        .collect(),
                    self.0[i].1 * other.0[j].1,
                )
            })
        }))
    }
}