nutils_poly/
lib.rs

1//! Low-level functions for evaluating and manipulating polynomials.
2//!
3//! The polynomials considered in this crate are [power series] in zero or
4//! more variables centered at zero and truncated to order $p$,
5//!
6//! $$ \\sum_{\\substack{k ∈ ℤ^n \\\\ \\sum_i k_i ≤ p}} c_k \\prod_i x_i^{k_i} $$
7//!
8//! where $c$ is a vector of coefficients, $x$ a vector of $n$ variables and
9//! $p$ a nonnegative integer degree.
10//!
11//! This crate requires the coefficients to be stored in a linear array in
12//! reverse [lexicographic order]: the coefficient for powers $j ∈ ℤ^n$ comes
13//! before the coefficient for powers $k ∈ ℤ^n \\setminus \\{j\\}$ iff $j_i >
14//! k_i$, where $i = \\max_l(j_l ≠ k_l)$, the index of the *last* non-matching
15//! power.
16//!
17//! This crate provides functions for [evaluating polynomials][`eval()`],
18//! computing coefficients for the [partial derivativative][`PartialDerivPlan`]
19//! and [products][`MulPlan`] of polynomials.
20//!
21//! # Examples
22//!
23//! The vector of coefficients for the polynomial $f(x, y) = 3 x y + x^2$ is
24//! `[0, 3, 0, 1, 0, 0]`.
25//!
26//! With [`eval()`] we can evaluate this polynomial:
27//!
28//! ```
29//! use nutils_poly;
30//!
31//! let coeffs = [0, 3, 0, 1, 0, 0];
32//! assert_eq!(nutils_poly::eval(&coeffs, &[1, 0], 2), Ok( 1)); // f(1, 0) =  1
33//! assert_eq!(nutils_poly::eval(&coeffs, &[1, 1], 2), Ok( 4)); // f(1, 1) =  4
34//! assert_eq!(nutils_poly::eval(&coeffs, &[2, 3], 2), Ok(22)); // f(2, 3) = 22
35//! ```
36//!
37//! [`PartialDerivPlan::apply()`] computes the coefficients for the partial
38//! derivative of a polynomial to one of the variables. The partial derivative
39//! of $f$ to $x$, the first variable, is $∂_x f(x, y) = 3 y + 2 x$
40//! (coefficients: `[3, 2, 0]`):
41//!
42//! ```
43//! use nutils_poly::PartialDerivPlan;
44//!
45//! let coeffs = [0, 3, 0, 1, 0, 0];
46//! let pd = PartialDerivPlan::new(
47//!     2, // number of variables
48//!     2, // degree
49//!     0, // variable to compute the partial derivative to
50//! ).unwrap();
51//! assert_eq!(Vec::from_iter(pd.apply(coeffs)?), vec![3, 2, 0]);
52//! # Ok::<_, nutils_poly::Error>(())
53//! ```
54//!
55//! # Nutils project
56//!
57//! This crate is part of the [Nutils project].
58//!
59//! [power series]: https://en.wikipedia.org/wiki/Power_series
60//! [lexicographic order]: https://en.wikipedia.org/wiki/Lexicographic_order
61//! [Nutils project]: https://nutils.org
62
63// Enable unstable feature `test` if we're running benchmarks.
64#![cfg_attr(feature = "bench", feature(test))]
65
66use ndarray::Array2;
67use num_traits::{One, Zero};
68use sqnc::traits::*;
69use std::cmp::Ordering;
70use std::fmt;
71use std::iter::{self, FusedIterator};
72use std::marker::PhantomData;
73use std::ops;
74
75pub type Power = u8;
76
77/// Enum used by [`Error::IncorrectNumberOfCoefficients`].
78#[derive(Debug, Clone, Copy, PartialEq, Eq)]
79pub enum ValueMoreOrLess {
80    Value(usize),
81    More,
82    Less,
83}
84
85use ValueMoreOrLess::{Less, More};
86
87impl From<usize> for ValueMoreOrLess {
88    #[inline]
89    fn from(value: usize) -> Self {
90        Self::Value(value)
91    }
92}
93
94impl fmt::Display for ValueMoreOrLess {
95    #[inline]
96    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
97        match self {
98            Self::Value(value) => value.fmt(f),
99            Self::More => write!(f, "more"),
100            Self::Less => write!(f, "less"),
101        }
102    }
103}
104
105/// The error type for fallible operations in this crate.
106#[derive(Debug, Clone, PartialEq, Eq)]
107#[non_exhaustive]
108pub enum Error {
109    /// The length of the coefficients sequence or iterator is incorrect.
110    IncorrectNumberOfCoefficients {
111        expected: usize,
112        got: ValueMoreOrLess,
113        detail: Option<&'static str>,
114    },
115    TooManyVariables,
116}
117
118impl std::error::Error for Error {}
119
120impl fmt::Display for Error {
121    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
122        match self {
123            Self::IncorrectNumberOfCoefficients {
124                expected,
125                got,
126                detail,
127            } => {
128                if let Some(detail) = detail {
129                    write!(f, "{detail}: ")?;
130                }
131                if *expected == 1 {
132                    write!(f, "Expected 1 coefficient but got {got}.")
133                } else {
134                    write!(f, "Expected {expected} coefficients but got {got}.")
135                }
136            }
137            Self::TooManyVariables => {
138                write!(
139                    f,
140                    "The number of variables exceeds the maximum (`usize::MAX`)."
141                )
142            }
143        }
144    }
145}
146
147#[inline]
148fn check_ncoeffs_sqnc<S: Sequence>(
149    coeffs: &S,
150    expected: usize,
151    detail: Option<&'static str>,
152) -> Result<(), Error> {
153    let got = coeffs.len();
154    if got == expected {
155        Ok(())
156    } else {
157        Err(Error::IncorrectNumberOfCoefficients {
158            expected,
159            got: got.into(),
160            detail,
161        })
162    }
163}
164
165// Function `ncoeffs_impl` computes the number of coefficients for a polynomial
166// of certain degree and number of variables using a loop over all variables.
167// If the number of variables is not know at copmile time, this loop cannot be
168// unrolled. Polynomials of interest typically have only a few variables. To
169// improve the perfomance of the typical case, function `ncoeffs` specializes
170// the calls to `ncoeffs_impl` for zero to three variables. The same holds for
171// function `ncoeffs_sum_impl`.
172
173/// Returns the number of coefficients for a polynomial of given degree and number of variables.
174///
175/// # Related
176///
177/// See [`degree()`] for the inverse operation, [`ncoeffs_iter()`] for an
178/// iterator that yields the number of coefficients for degree zero and up and
179/// [`degree_ncoeffs_iter()`] for an iterator that yields pairs of degrees and
180/// number of coefficients.
181#[inline]
182pub const fn ncoeffs(nvars: usize, degree: Power) -> usize {
183    match nvars {
184        0 => ncoeffs_impl(0, degree),
185        1 => ncoeffs_impl(1, degree),
186        2 => ncoeffs_impl(2, degree),
187        3 => ncoeffs_impl(3, degree),
188        _ => ncoeffs_impl(nvars, degree),
189    }
190}
191
192#[inline]
193const fn ncoeffs_impl(nvars: usize, degree: Power) -> usize {
194    let mut n = 1;
195    let mut i = 0;
196    while i < nvars {
197        i += 1;
198        n = n * (degree as usize + i) / i;
199    }
200    n
201}
202
203/// Returns an iterator of the number of coefficients for a polynomial of degree zero and up.
204///
205/// # Related
206///
207/// See [`degree_ncoeffs_iter()`] for an iterator that yields the degree
208/// together with the number of coefficients and [`ncoeffs()`] for a function
209/// that returns the number of coefficients for a single degree.
210#[inline]
211pub fn ncoeffs_iter(nvars: usize) -> impl Iterator<Item = usize> {
212    degree_ncoeffs_iter(nvars).map(|(_, ncoeffs)| ncoeffs)
213}
214
215/// Returns an iterator of degrees and number of coefficients.
216///
217/// # Related
218///
219/// See [`ncoeffs_iter()`] for an iterator that yields the number of
220/// coefficients only.
221#[inline]
222pub fn degree_ncoeffs_iter(nvars: usize) -> impl Iterator<Item = (Power, usize)> {
223    (0u8..).scan(1, move |ncoeffs, degree| {
224        let current = (degree, *ncoeffs);
225        *ncoeffs = *ncoeffs * (degree as usize + 1 + nvars) / (degree as usize + 1);
226        Some(current)
227    })
228}
229
230/// Returns the sum of the number of coefficients up to (excluding) the given degree.
231#[inline]
232const fn ncoeffs_sum(nvars: usize, degree: Power) -> usize {
233    match nvars {
234        0 => ncoeffs_sum_impl(0, degree),
235        1 => ncoeffs_sum_impl(1, degree),
236        2 => ncoeffs_sum_impl(2, degree),
237        3 => ncoeffs_sum_impl(3, degree),
238        _ => ncoeffs_sum_impl(nvars, degree),
239    }
240}
241
242#[inline]
243const fn ncoeffs_sum_impl(nvars: usize, degree: Power) -> usize {
244    let mut n = 1;
245    let mut i = 0;
246    while i <= nvars {
247        n = n * (degree as usize + i) / (i + 1);
248        i += 1;
249    }
250    n
251}
252
253/// Returns the degree of a polynomial given the number of variables and coefficients.
254///
255/// This is the inverse of [`ncoeffs()`]. If there is no degree for which
256/// `ncoeffs(nvars, degree) == ncoeffs`, this function returns `None`. For
257/// example, a polynomial in two variables has one coefficient for degree zero
258/// and three for degree one, but there is nothing in between.
259///
260/// # Related
261///
262/// See [`ncoeffs()`] for the inverse operation and [`degree_ncoeffs_iter()`]
263/// for an iterator that yields pairs of degrees and number of coefficients for
264/// degree zero and up.
265pub fn degree(nvars: usize, ncoeffs: usize) -> Option<Power> {
266    match nvars {
267        0 => (ncoeffs == 1).then_some(0),
268        1 => ncoeffs
269            .checked_sub(1)
270            .and_then(|degree| degree.try_into().ok()),
271        _ => {
272            // There's nothing else we can do than to iterate over all degrees
273            // starting at zero and to stop as soon as we have found or stepped
274            // over `ncoeffs`.
275            for (tdegree, tncoeffs) in degree_ncoeffs_iter(nvars) {
276                match tncoeffs.cmp(&ncoeffs) {
277                    Ordering::Equal => {
278                        return Some(tdegree);
279                    }
280                    Ordering::Greater => {
281                        return None;
282                    }
283                    _ => {}
284                }
285            }
286            unreachable! {}
287        }
288    }
289}
290
291/// Returns an iterator that yields powers (in reverse order) for the given index.
292///
293/// Returns `None` if the index is larger than the number of coefficients for a
294/// polynomial in the given number of variables and of the given degree.
295///
296/// # Related
297///
298/// The inverse operation is [`powers_rev_iter_to_index()`].
299fn index_to_powers_rev_iter(
300    mut index: usize,
301    nvars: usize,
302    mut degree: Power,
303) -> Option<impl Iterator<Item = Power> + ExactSizeIterator> {
304    // For the last variable the sets of possible indices for each power `k`
305    // are adjacent ranges with lengths `n_k`, ordered from power `k = degree`
306    // to zero. The length of the range for power `k` is the number of
307    // coefficients for a polynomial in one variable less and of degree `degree
308    // - k`. We can find the power `k` of the last variable by finding the
309    // range that contains the index.
310    //
311    // Having found the this power, we can subtract power `k` from `degree`
312    // (now treated as remaining degree), subtract the lower bound of the range
313    // from `index` (now treated as remaining index), and repeat the process.
314    //
315    // For example a polynomial in 2 variables and of degree 3 has the
316    // following ranges of indices for the last variable:
317    //
318    //     {0}          for power 3
319    //     {1,2}        for power 2
320    //     {3,4,5}      for power 1
321    //     {6,7,8,9}    for power 0
322    //
323    // Index 4 corresponds to power 1 for the last variable. The remaining
324    // degree is 2 and the remaining index is 1. Repeating the process gives
325    // power 1 for the first variable.
326
327    (index < ncoeffs(nvars, degree)).then(|| {
328        Iterator::rev(0..nvars).map(move |var| {
329            if var == 0 {
330                // The else branch works fine, but since `ncoeffs(0, ?)` is one
331                // regardless `?` we can simplify the loop over the degree to:
332                degree - index as Power
333            } else {
334                // `degree - i` is the tentative power for variable `var` and
335                // `i` is the tentative remaining degree. We loop over power
336                // from high to low, continuously subtract the length of the
337                // range (`n`) from `index`, and stop just before `index`
338                // would become negative.
339                for (i, n) in degree_ncoeffs_iter(var) {
340                    if index < n {
341                        // We have checked that initial `index` does not exceed
342                        // the number of coefficients above, hence `i` is never
343                        // larger than `degree`.
344                        let power = degree - i;
345                        degree = i;
346                        return power;
347                    }
348                    index -= n;
349                }
350                // The `degree_ncoeffs_iter` never ends.
351                unreachable! {}
352            }
353        })
354    })
355}
356
357/// Returns the index given an iterator of powers in reverse order.
358///
359/// Returns `None` if the sum of the powers exceeds the given `degree`. If the
360/// length of `rev_powers` does not match `nvars`, this function returns an
361/// undefined value.
362///
363/// # Related
364///
365/// The inverse operation is [`index_to_powers_rev_iter()`].
366fn powers_rev_iter_to_index<PowersRevIter>(
367    rev_powers: PowersRevIter,
368    nvars: usize,
369    mut degree: Power,
370) -> Option<usize>
371where
372    PowersRevIter: Iterator<Item = Power>,
373{
374    // For the last variable the range of possible indices for power `k` is
375    // `{i_k,...,i_{k+1}-1}` where `i_k` is the partial sum of `{n_.}`,
376    //
377    //     i_k = Σ_{j | 0 ≤ j < k} n_j
378    //
379    // and `n_k` is the number of coefficients for a polynomial in one variable
380    // less and of degree `degree - k`. We subtract `k` from `degree` and
381    // repeat this process for the remaining variables. The index is the sum of
382    // all `i_k`.
383    let mut index = 0;
384    for (var, power) in iter::zip(Iterator::rev(0..nvars), rev_powers) {
385        degree = degree.checked_sub(power)?;
386        index += ncoeffs_sum(var, degree);
387    }
388    Some(index)
389}
390
391/// Index map relating coefficients from one degree to another.
392///
393/// For each coefficient for a polynomial in `nvars` variables and of degree
394/// `from_degree`, this sequence gives the index of the same coefficient (the
395/// same powers) for a polynomial with degree `to_degree`, where `to_degree`
396/// must be larger or equal to `from_degree`.
397///
398/// The indices are strict monotonic increasing.
399///
400/// # Example
401///
402/// The following example shows how to map a vector of coefficients from one
403/// degree to another.
404///
405/// ```
406/// use nutils_poly::{self, MapDegree};
407/// use sqnc::traits::*;
408///
409/// // Coefficients for a polynomial in two variables and of degree one.
410/// let coeffs1 = [4, 5, 6];
411/// // Index map from degree two to degree three for a polynomial in one variable.
412/// let map = MapDegree::new(2, 1, 2).unwrap();
413/// // Apply the map to obtain a coefficients vector for a polynomial of degree two.
414/// let mut coeffs2 = [0; nutils_poly::ncoeffs(2, 2)];
415/// coeffs2.as_mut_sqnc().select(map).unwrap().assign(coeffs1);
416/// assert_eq!(coeffs2, [0, 0, 4, 0, 5, 6]);
417/// // Both polynomials evaluate to the same value.
418/// assert_eq!(
419///     nutils_poly::eval(coeffs1, &[2, 3], 1)?,
420///     nutils_poly::eval(coeffs2, &[2, 3], 2)?,
421/// );
422/// # Ok::<_, nutils_poly::Error>(())
423/// ```
424#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
425pub struct MapDegree {
426    nvars: usize,
427    from_degree: Power,
428    to_degree: Power,
429}
430
431impl MapDegree {
432    /// Returns an index map or `None` if `to_degree` is smaller than `from_degree`.
433    pub fn new(nvars: usize, from_degree: Power, to_degree: Power) -> Option<Self> {
434        (to_degree >= from_degree).then_some(Self {
435            nvars,
436            from_degree,
437            to_degree,
438        })
439    }
440}
441
442impl<'this> SequenceTypes<'this> for MapDegree {
443    type Item = usize;
444    type Iter = sqnc::derive::Iter<'this, Self>;
445}
446
447impl Sequence for MapDegree {
448    #[inline]
449    fn len(&self) -> usize {
450        ncoeffs(self.nvars, self.from_degree)
451    }
452
453    #[inline]
454    fn get(&self, index: usize) -> Option<usize> {
455        let powers = index_to_powers_rev_iter(index, self.nvars, self.from_degree)?;
456        powers_rev_iter_to_index(powers, self.nvars, self.to_degree)
457    }
458
459    #[inline]
460    fn iter(&self) -> sqnc::derive::Iter<'_, Self> {
461        self.into()
462    }
463}
464
465impl IntoIterator for MapDegree {
466    type Item = usize;
467    type IntoIter = sqnc::derive::IntoIter<Self>;
468
469    #[inline]
470    fn into_iter(self) -> Self::IntoIter {
471        self.into()
472    }
473}
474
475// SAFETY: The coefficients of the polynomials are ordered such that the
476// corresponding powers are in reverse lexicographic order (see crate-level
477// documentation). The powers for a polynomial in `n` variables of some degree
478// `q` is a superset of those for a polynomial in `n` variables of a degree `p
479// ≤ q` and `MapDegree` lists the indices of the powers in polynomial `q` that
480// also exist in polynomial `p`. Due to the lexicographic ordering of both
481// polynomials, this sequence is strict monotonic increasing, hence unique.
482// `MapDegree` cannot be instantiated using public methods if `p > q`.
483unsafe impl UniqueSequence for MapDegree {}
484
485/// Operations for evaluating polynomials.
486///
487/// The polynomials considered in this crate are of the form
488///
489/// $$
490///     s_{0,\[\\ \]} = \\left(
491///         \\sum_{k_{n-1}} x_{n-1}^{k_{n-1}} \\left(
492///             \\cdots
493///             \\sum_{k_1} x_1^{k_1} \\left(
494///                 \\sum_{k_0} x_0^{k_0} c_k
495///             \\right)
496///         \\right)
497///     \\right).
498/// $$
499///
500/// For brevity the bounds of the summations are omited. The scopes, delimited
501/// by parentheses, indicate the order in which the polynomial is being
502/// evaluated numerically. The innermost scope eliminates the first variable,
503/// the next scope the second, etc.
504///
505/// Let $s_{n,k} = c_k$ and let $s_{i,k}$ be the value of the $i$-th scope,
506///
507/// $$ s_{i,k} = \\sum_{j=0}^{m_{i,k}} x_i^j s_{i+1,k+\[j\]}, $$
508///
509/// where $m_{i,k}$ is the appropriate bound and $k+\[j\]$ denotes the
510/// concatenation of vector $k ∈ ℤ^i$ with vector $\[j\]$.
511/// The summation is rewritten such that the powers of $x_i$ are eliminated:
512///
513/// $$ s_{i,k} = ((((s_{i+1,k+\[m_{i,k}\]}) \\cdots) x_i + s_{i+1,k+\[2\]}) x_i + s_{i+1,k+\[1\]}) x_i + s_{i,k+\[0\]} $$
514///
515/// or equivalently,
516///
517/// $$ s_{i,k} = a_{i,k,0} $$
518///
519/// with $a_{i,k,j}$ defined recursively as
520///
521/// $$
522///     a_{i,k,j} = \\begin{cases}
523///         s_{i+1,k+\[j\]} & \\text{if } j = m_{i,k} \\\\
524///         a_{i,k,j+1} x_i + s_{i+1,k+\[j\]} & \\text{if } j < m_{i,k} \\\\
525///     \\end{cases}
526/// $$
527///
528/// Method [`EvalOps::eval()`] uses this recursion to evaluate the
529/// polynomial. The four required methods in this trait perform operations on
530/// accumulator $a_{i,k}$ (the aggregation of $a_{i,k,0},a_{i,k,1},\\ldots$).
531pub trait EvalOps<Coeff, Value: ?Sized> {
532    /// The type of the result of the evaluation.
533    type Output;
534
535    /// Returns an accumulator initialized with zero.
536    ///
537    /// This function performs the following operation:
538    ///
539    /// ```no_compile
540    /// acc = 0
541    /// ```
542    fn init_acc_zero() -> Self::Output;
543
544    /// Returns an accumulator initialized with a polynomial coefficient.
545    ///
546    /// This function performs the following operation:
547    ///
548    /// ```no_compile
549    /// acc = coeff
550    /// ```
551    fn init_acc_coeff(coeff: Coeff) -> Self::Output;
552
553    /// Updates an accumulator with a coefficient.
554    ///
555    /// This function performs the following operation:
556    ///
557    /// ```no_compile
558    /// acc = acc * value + coeff
559    /// ```
560    fn update_acc_coeff(acc: &mut Self::Output, coeff: Coeff, value: &Value);
561
562    /// Updates an accumulator with the result of an inner loop.
563    ///
564    /// This function performs the following operation:
565    ///
566    /// ```no_compile
567    /// acc = acc * value + inner
568    /// ```
569    ///
570    /// where `inner` has the same type as the accumulator.
571    fn update_acc_inner(acc: &mut Self::Output, inner: Self::Output, value: &Value);
572
573    /// Evaluate a polynomial for the given values.
574    ///
575    /// # Errors
576    ///
577    /// This function returns an error if the `coeffs` iterator is too short.
578    #[inline]
579    fn eval<Coeffs, Values>(
580        coeffs: &mut Coeffs,
581        values: &Values,
582        degree: Power,
583    ) -> Result<Self::Output, Error>
584    where
585        Coeffs: Iterator<Item = Coeff>,
586        Values: SequenceRef<OwnedItem = Value> + ?Sized,
587    {
588        <EvalImpl<Self, Coeffs, Values>>::eval_nv(coeffs, values, degree, values.len()).ok_or_else(
589            || Error::IncorrectNumberOfCoefficients {
590                expected: ncoeffs(values.len(), degree),
591                got: Less,
592                detail: None,
593            },
594        )
595    }
596}
597
598// While it would be easier to merge the private methods defined below with
599// [`EvalOps`], we don't want the functions to be public. There are
600// [tricks][Private Methods on a Public Trait] to do this, but this struct does
601// the job just fine.
602//
603// [Private Methods on a Public Trait]: https://web.archive.org/web/20220220002300/https://jack.wrenn.fyi/blog/private-trait-methods/
604
605/// Private methods for [`EvalOps`].
606struct EvalImpl<Ops: ?Sized, Coeffs, Values: ?Sized>(
607    PhantomData<Ops>,
608    PhantomData<Coeffs>,
609    PhantomData<Values>,
610);
611
612impl<Ops, Coeffs, Values> EvalImpl<Ops, Coeffs, Values>
613where
614    Ops: EvalOps<Coeffs::Item, Values::OwnedItem> + ?Sized,
615    Coeffs: Iterator,
616    Values: SequenceRef + ?Sized,
617{
618    /// Evaluate a polynomial in zero variables.
619    #[inline]
620    fn eval_0v(coeffs: &mut Coeffs) -> Option<Ops::Output> {
621        coeffs.next().map(Ops::init_acc_coeff)
622    }
623
624    /// Evaluate a polynomial in one variable.
625    #[inline]
626    fn eval_1v(coeffs: &mut Coeffs, values: &Values, degree: Power) -> Option<Ops::Output> {
627        let mut acc = Ops::init_acc_zero();
628        if let Some(value) = values.get(0) {
629            if coeffs
630                .take(degree as usize + 1)
631                .map(|coeff| Ops::update_acc_coeff(&mut acc, coeff, value))
632                .count()
633                != degree as usize + 1
634            {
635                return None;
636            }
637        }
638        Some(acc)
639    }
640
641    /// Evaluate a polynomial in two variables.
642    #[inline]
643    fn eval_2v(coeffs: &mut Coeffs, values: &Values, degree: Power) -> Option<Ops::Output> {
644        let mut acc = Ops::init_acc_zero();
645        if let Some(value) = values.get(1) {
646            for p in 0..=degree {
647                let inner = Self::eval_1v(coeffs, values, p)?;
648                Ops::update_acc_inner(&mut acc, inner, value);
649            }
650        }
651        Some(acc)
652    }
653
654    /// Evaluate a polynomial in three variables.
655    #[inline]
656    fn eval_3v(coeffs: &mut Coeffs, values: &Values, degree: Power) -> Option<Ops::Output> {
657        let mut acc = Ops::init_acc_zero();
658        if let Some(value) = values.get(2) {
659            for p in 0..=degree {
660                let inner = Self::eval_2v(coeffs, values, p)?;
661                Ops::update_acc_inner(&mut acc, inner, value);
662            }
663        }
664        Some(acc)
665    }
666
667    /// Evaluate a polynomial in any number of variables.
668    #[inline]
669    fn eval_nv(
670        coeffs: &mut Coeffs,
671        values: &Values,
672        degree: Power,
673        nvars: usize,
674    ) -> Option<Ops::Output> {
675        match nvars {
676            0 => Self::eval_0v(coeffs),
677            1 => Self::eval_1v(coeffs, values, degree),
678            2 => Self::eval_2v(coeffs, values, degree),
679            3 => Self::eval_3v(coeffs, values, degree),
680            _ => {
681                let mut acc = Ops::init_acc_zero();
682                if let Some(value) = values.get(nvars - 1) {
683                    for p in 0..=degree {
684                        let inner = Self::eval_nv(coeffs, values, p, nvars - 1)?;
685                        Ops::update_acc_inner(&mut acc, inner, value);
686                    }
687                }
688                Some(acc)
689            }
690        }
691    }
692}
693
694/// Implementation for [`EvalOps`] using [`std::ops::AddAssign`] and [`std::ops::MulAssign`].
695///
696/// This implementation is used by [`eval()`].
697enum DefaultEvalOps {}
698
699impl<Coeff, Value> EvalOps<Coeff, Value> for DefaultEvalOps
700where
701    Value: Zero + ops::AddAssign + ops::AddAssign<Coeff> + for<'v> ops::MulAssign<&'v Value>,
702{
703    type Output = Value;
704
705    #[inline]
706    fn init_acc_coeff(coeff: Coeff) -> Value {
707        let mut acc = Value::zero();
708        acc += coeff;
709        acc
710    }
711
712    #[inline]
713    fn init_acc_zero() -> Value {
714        Value::zero()
715    }
716
717    #[inline]
718    fn update_acc_coeff(acc: &mut Value, coeff: Coeff, value: &Value) {
719        *acc *= value;
720        *acc += coeff;
721    }
722
723    #[inline]
724    fn update_acc_inner(acc: &mut Value, inner: Value, value: &Value) {
725        *acc *= value;
726        *acc += inner;
727    }
728}
729
730/// Evaluates a polynomial for the given values.
731///
732/// The number of coefficients of the polynomial is determined by the length of
733/// the `values` (the number of variables) and the `degree`. The coefficients
734/// have to be iterable. Only the expected number of coefficients are taken
735/// from the iterator.
736///
737/// # Errors
738///
739/// If the iterator of coefficients has fewer elements than expected, an error
740/// is returned.
741///
742/// # Examples
743///
744/// The vector of coefficients for the polynomial $f(x) = x^2 - x + 2$ is
745/// `[1, -1, 2]`. Evaluating $f(0)$, $f(1)$ and $f(2)$:
746///
747/// ```
748/// use nutils_poly;
749///
750/// let coeffs = [1, -1, 2];
751/// assert_eq!(nutils_poly::eval(&coeffs, &[0], 2), Ok(2)); // f(0) = 2
752/// assert_eq!(nutils_poly::eval(&coeffs, &[1], 2), Ok(2)); // f(1) = 2
753/// assert_eq!(nutils_poly::eval(&coeffs, &[2], 2), Ok(4)); // f(2) = 4
754/// ```
755///
756/// Let $g(x) = x^2 + x + 1$ be another polynomial. Since [`eval()`] consumes
757/// only the expected number of coefficients, we can chain the coefficients for
758/// $f$ and $g$ in a single iterator and call [`eval()`] twice to obtain the
759/// values for $f$ and $g$:
760///
761/// ```
762/// use nutils_poly;
763///
764/// let mut coeffs = [1, -1, 2, 1, 1, 1].into_iter();
765/// // Evaluate `f`, consumes the first three coefficients.
766/// assert_eq!(nutils_poly::eval(&mut coeffs, &[2], 2), Ok(4)); // f(2) = 4
767/// // Evaluate `g`, consumes the last three coefficients.
768/// assert_eq!(nutils_poly::eval(&mut coeffs, &[2], 2), Ok(7)); // g(2) = 7
769/// ```
770#[inline]
771pub fn eval<Coeff, Value, Coeffs, Values>(
772    coeffs: Coeffs,
773    values: &Values,
774    degree: Power,
775) -> Result<Value, Error>
776where
777    Value: Zero + ops::AddAssign + ops::AddAssign<Coeff> + for<'v> ops::MulAssign<&'v Value>,
778    Coeffs: IntoIterator<Item = Coeff>,
779    Values: SequenceRef<OwnedItem = Value> + ?Sized,
780{
781    DefaultEvalOps::eval(&mut coeffs.into_iter(), values, degree)
782}
783
784/// Interface for computing a [multiple].
785///
786/// [multiple]: https://en.wikipedia.org/wiki/Multiple_(mathematics)
787pub trait Multiple {
788    /// The return type of [`Multiple::multiple()`].
789    type Output;
790
791    /// Returns the product of `self` and `n`.
792    fn multiple(self, n: usize) -> Self::Output;
793}
794
795macro_rules! impl_int_mul {
796    ($T:ty $(,)?) => {
797        impl Multiple for $T {
798            type Output = $T;
799
800            #[inline]
801            fn multiple(self, n: usize) -> Self::Output {
802                self * n as $T
803            }
804        }
805
806        impl Multiple for &$T {
807            type Output = $T;
808
809            #[inline]
810            fn multiple(self, n: usize) -> Self::Output {
811                self * n as $T
812            }
813        }
814    };
815    ($T:ty, $($tail:tt)*) => {
816        impl_int_mul!{$T}
817        impl_int_mul!{$($tail)*}
818    };
819}
820
821impl_int_mul! {u8, u16, u32, u64, u128, usize, i8, i16, i32, i64, i128, isize, f32, f64}
822
823/// Plan for computing coefficients for partial derivatives.
824///
825/// This struct holds a plan for efficiently evaluating the partial derivative
826/// of polynomials.
827///
828/// # Examples
829///
830/// The partial derivative of $f(x, y) = x y - x^2 + 2$ (coefficients: `[0, 1,
831/// 0, -1, 0, 2]`) to the first variable, $x$:
832///
833/// ```
834/// use nutils_poly::PartialDerivPlan;
835/// use sqnc::traits::*;
836///
837/// let f = [0, 1, 0, -1, 0, 2];
838/// let plan = PartialDerivPlan::new(
839///     2, // number of variables
840///     2, // degree
841///     0, // variable to compute the partial derivative to
842/// ).unwrap();
843/// assert!(Iterator::eq(
844///     plan.apply(f)?.iter(),
845///     [1, -2, 0],
846/// ));
847/// # Ok::<_, nutils_poly::Error>(())
848/// ```
849#[derive(Debug, Clone)]
850pub struct PartialDerivPlan {
851    indices_powers: Box<[(usize, usize)]>,
852    ncoeffs_input: usize,
853    degree_output: Power,
854    nvars: usize,
855}
856
857impl PartialDerivPlan {
858    /// Plan the partial derivative of a polynomial.
859    pub fn new(nvars: usize, degree: Power, var: usize) -> Option<Self> {
860        let n = nvars.checked_sub(var + 1)?;
861        let ncoeffs_input = ncoeffs(nvars, degree);
862        let indices_powers = if degree == 0 {
863            [(0, 0)].into()
864        } else {
865            (0..ncoeffs_input)
866                .filter_map(move |index| {
867                    let power = index_to_powers_rev_iter(index, nvars, degree)?.nth(n)?;
868                    (power > 0).then_some((index, power as usize))
869                })
870                .collect()
871        };
872        Some(Self {
873            indices_powers,
874            ncoeffs_input,
875            degree_output: degree.saturating_sub(1),
876            nvars,
877        })
878    }
879
880    /// Returns the partial derivative of a polynomial.
881    ///
882    /// # Errors
883    ///
884    /// If the number of coefficients of the input doesn't match
885    /// [`PartialDerivPlan::ncoeffs_input()`] then this function returns an
886    /// error.
887    #[inline]
888    pub fn apply<Coeffs>(&self, coeffs: Coeffs) -> Result<PartialDeriv<'_, Coeffs>, Error>
889    where
890        Coeffs: Sequence,
891    {
892        check_ncoeffs_sqnc(&coeffs, self.ncoeffs_input, None)?;
893        Ok(PartialDeriv { plan: self, coeffs })
894    }
895
896    /// Returns the number of coefficients for the input polynomial.
897    #[inline]
898    pub fn ncoeffs_input(&self) -> usize {
899        self.ncoeffs_input
900    }
901
902    /// Returns the number of coefficients for the partial derivative.
903    #[inline]
904    pub fn ncoeffs_output(&self) -> usize {
905        self.indices_powers.len()
906    }
907
908    /// Returns the degree of the partial derivative.
909    #[inline]
910    pub fn degree_output(&self) -> Power {
911        self.degree_output
912    }
913
914    /// Returns the number of variables of both the input and the output.
915    #[inline]
916    pub fn nvars(&self) -> usize {
917        self.nvars
918    }
919}
920
921/// The partial derivative of a polynomial.
922///
923/// This struct is created by [`PartialDerivPlan`]. See its documentation for
924/// more information.
925#[derive(Debug, Clone)]
926pub struct PartialDeriv<'plan, Coeffs> {
927    plan: &'plan PartialDerivPlan,
928    coeffs: Coeffs,
929}
930
931impl<'this, 'plan, Coeffs, OCoeff> SequenceTypes<'this> for PartialDeriv<'plan, Coeffs>
932where
933    for<'a> <Coeffs as SequenceTypes<'a>>::Item: Multiple<Output = OCoeff>,
934    Coeffs: Sequence,
935{
936    type Item = OCoeff;
937    type Iter = PartialDerivIter<'plan, sqnc::Wrapper<&'this Coeffs, ((),)>>;
938}
939
940impl<'plan, Coeffs, OCoeff> Sequence for PartialDeriv<'plan, Coeffs>
941where
942    for<'a> <Coeffs as SequenceTypes<'a>>::Item: Multiple<Output = OCoeff>,
943    Coeffs: Sequence,
944{
945    #[inline]
946    fn len(&self) -> usize {
947        self.plan.indices_powers.len()
948    }
949
950    #[inline]
951    fn get(&self, index: usize) -> Option<OCoeff> {
952        let (index, power) = self.plan.indices_powers.get(index)?;
953        Some(self.coeffs.get(*index)?.multiple(*power))
954    }
955
956    #[inline]
957    fn iter(&self) -> PartialDerivIter<'plan, sqnc::Wrapper<&'_ Coeffs, ((),)>> {
958        PartialDerivIter {
959            indices_powers: self.plan.indices_powers.iter(),
960            coeffs: self.coeffs.as_sqnc(),
961        }
962    }
963}
964
965impl<'plan, Coeffs, OCoeff> IntoIterator for PartialDeriv<'plan, Coeffs>
966where
967    Coeffs: Sequence,
968    for<'a> <Coeffs as SequenceTypes<'a>>::Item: Multiple<Output = OCoeff>,
969{
970    type Item = OCoeff;
971    type IntoIter = PartialDerivIter<'plan, Coeffs>;
972
973    #[inline]
974    fn into_iter(self) -> Self::IntoIter {
975        PartialDerivIter {
976            indices_powers: self.plan.indices_powers.iter(),
977            coeffs: self.coeffs,
978        }
979    }
980}
981
982/// An iterator of the coefficients of a partial derivative.
983///
984/// This struct is created by [`PartialDeriv::iter()`].
985#[derive(Debug, Clone)]
986pub struct PartialDerivIter<'plan, Coeffs> {
987    indices_powers: std::slice::Iter<'plan, (usize, usize)>,
988    coeffs: Coeffs,
989}
990
991impl<'plan, Coeffs, OCoeff> Iterator for PartialDerivIter<'plan, Coeffs>
992where
993    Coeffs: Sequence,
994    for<'a> <Coeffs as SequenceTypes<'a>>::Item: Multiple<Output = OCoeff>,
995{
996    type Item = OCoeff;
997
998    #[inline]
999    fn next(&mut self) -> Option<Self::Item> {
1000        let (index, power) = self.indices_powers.next()?;
1001        Some(self.coeffs.get(*index)?.multiple(*power))
1002    }
1003
1004    #[inline]
1005    fn size_hint(&self) -> (usize, Option<usize>) {
1006        self.indices_powers.size_hint()
1007    }
1008}
1009
1010impl<'plan, Coeffs, OCoeff> DoubleEndedIterator for PartialDerivIter<'plan, Coeffs>
1011where
1012    Coeffs: Sequence,
1013    for<'a> <Coeffs as SequenceTypes<'a>>::Item: Multiple<Output = OCoeff>,
1014{
1015    #[inline]
1016    fn next_back(&mut self) -> Option<Self::Item> {
1017        let (index, power) = self.indices_powers.next_back()?;
1018        Some(self.coeffs.get(*index)?.multiple(*power))
1019    }
1020}
1021
1022impl<'plan, Coeffs, OCoeff> ExactSizeIterator for PartialDerivIter<'plan, Coeffs>
1023where
1024    Coeffs: Sequence,
1025    for<'a> <Coeffs as SequenceTypes<'a>>::Item: Multiple<Output = OCoeff>,
1026{
1027}
1028
1029impl<'plan, Coeffs, OCoeff> FusedIterator for PartialDerivIter<'plan, Coeffs>
1030where
1031    Coeffs: Sequence,
1032    for<'a> <Coeffs as SequenceTypes<'a>>::Item: Multiple<Output = OCoeff>,
1033{
1034}
1035
1036/// Existence of a variable in the operands of a product polynomial.
1037///
1038/// The documentation for [`MulPlan::new()`] explains how to use this enum.
1039#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
1040pub enum MulVar {
1041    Left,
1042    Right,
1043    Both,
1044}
1045
1046/// Plan for computing products of polynomials.
1047///
1048/// This struct holds a plan for efficiently evaluating products of polynomials.
1049///
1050/// # Examples
1051///
1052/// Consider the following polynomials: $f(x) = x^2$ (coefficients: `[1, 0,
1053/// 0]`), $g(x) = 2 x + 1$ (coefficients: `[2, 1]`) and $h(x, y) = x y$
1054/// (coefficients: `[0, 1, 0, 0, 0, 0]`).
1055///
1056/// ```
1057/// use nutils_poly::{MulPlan, MulVar};
1058/// use sqnc::traits::*;
1059///
1060/// let f = [1, 0, 0];          // degree: 2, nvars: 1
1061/// let g = [2, 1];             // degree: 1, nvars: 1
1062/// let h = [0, 1, 0, 0, 0, 0]; // degree: 2, nvars: 2
1063/// ```
1064///
1065/// Computing the coefficients for $x ↦ f(x) g(x)$:
1066///
1067/// ```
1068/// # use nutils_poly::{MulPlan, MulVar};
1069/// use sqnc::traits::*;
1070/// #
1071/// # let f = [1, 0, 0];
1072/// # let g = [2, 1];
1073/// # let h = [0, 1, 0, 0, 0, 0];
1074/// #
1075/// let plan = MulPlan::same_vars(
1076///     1, // number of variables
1077///     2, // degree of left operand
1078///     1, // degree of right operand
1079/// );
1080/// let fx_gx = plan.apply(f.as_sqnc(), g.as_sqnc()).unwrap();
1081/// assert!(fx_gx.iter().eq([2, 1, 0, 0]));
1082/// ```
1083///
1084/// Similarly, but with different variables, $(x, y) ↦ f(x) g(y)$:
1085///
1086/// ```
1087/// # use nutils_poly::{MulPlan, MulVar};
1088/// use sqnc::traits::*;
1089/// #
1090/// # let f = [1, 0, 0];
1091/// # let g = [2, 1];
1092/// # let h = [0, 1, 0, 0, 0, 0];
1093/// #
1094/// let plan = MulPlan::different_vars(
1095///     1, // number of variables of the left operand
1096///     1, // number of variables of the right operand
1097///     2, // degree of left operand
1098///     1, // degree of right operand
1099/// )?;
1100/// let fx_gy = plan.apply(f.as_sqnc(), g.as_sqnc()).unwrap();
1101/// assert!(fx_gy.iter().eq([0, 0, 0, 2, 0, 0, 0, 1, 0, 0]));
1102/// # Ok::<_, nutils_poly::Error>(())
1103/// ```
1104///
1105/// Computing the coefficients for $(x, y) ↦ f(y) h(x,y)$:
1106///
1107/// ```
1108/// # use nutils_poly::{MulPlan, MulVar};
1109/// use sqnc::traits::*;
1110/// #
1111/// # let f = [1, 0, 0];
1112/// # let g = [2, 1];
1113/// # let h = [0, 1, 0, 0, 0, 0];
1114/// #
1115/// let plan = MulPlan::new(
1116///     &[
1117///         MulVar::Right, // variable x, exists only in the right operand
1118///         MulVar::Both,  // variable y, exists in both operands
1119///     ].copied(),
1120///     2, // degree of left operand
1121///     2, // degree of right operand
1122/// );
1123/// let fy_hxy = plan.apply(f.as_sqnc(), h.as_sqnc()).unwrap();
1124/// assert!(fy_hxy.iter().eq([0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]));
1125/// ```
1126#[derive(Debug, Clone)]
1127pub struct MulPlan {
1128    count: Box<[usize]>,
1129    offsets: Box<[usize]>,
1130    indices: Box<[[usize; 2]]>,
1131    ncoeffs_left: usize,
1132    ncoeffs_right: usize,
1133    nvars_output: usize,
1134    degree_output: Power,
1135}
1136
1137impl MulPlan {
1138    /// Plan the product of two polynomials.
1139    ///
1140    /// For each output variable, `vars` lists the relation between output and input variables:
1141    ///
1142    /// *   [`MulVar::Left`] if the variable exists only in the left polynomial,
1143    /// *   [`MulVar::Right`] if the variable exists only in the right polynomial or
1144    /// *   [`MulVar::Both`] if the variable exists in both polynomials.
1145    ///
1146    /// It is not possible to reorder the output variables with respect to the input variables.
1147    #[inline]
1148    pub fn new<'vars, Vars>(vars: &'vars Vars, degree_left: Power, degree_right: Power) -> Self
1149    where
1150        Vars: SequenceOwned<OwnedItem = MulVar>,
1151        <Vars as SequenceTypes<'vars>>::Iter: DoubleEndedIterator,
1152    {
1153        Self::for_output_degree(vars, degree_left, degree_right, degree_left + degree_right)
1154    }
1155
1156    /// Plan the product of two polynomials for the given output degree.
1157    ///
1158    /// If the output degree is smaller than the sum of the degrees of the input operands,
1159    /// then the product is truncated.
1160    pub fn for_output_degree<'vars, Vars>(
1161        vars: &'vars Vars,
1162        degree_left: Power,
1163        degree_right: Power,
1164        degree_output: Power,
1165    ) -> Self
1166    where
1167        Vars: SequenceOwned<OwnedItem = MulVar>,
1168        <Vars as SequenceTypes<'vars>>::Iter: DoubleEndedIterator,
1169    {
1170        let nvars_left = vars.iter().filter(|var| *var != MulVar::Right).count();
1171        let nvars_right = vars.iter().filter(|var| *var != MulVar::Left).count();
1172        let nvars_output = vars.len();
1173        let ncoeffs_left = ncoeffs(nvars_left, degree_left);
1174        let ncoeffs_right = ncoeffs(nvars_right, degree_right);
1175        let ncoeffs_out = ncoeffs(nvars_output, degree_output);
1176
1177        // For every combination of left and right indices, we compute the
1178        // corresponding index of the product polynomial (by computing the
1179        // powers left and right, summing, and computing the index from powers)
1180        // and store the index triple (output first) in `indices`.
1181        let mut indices: Vec<[usize; 3]> = Vec::with_capacity(ncoeffs_left * ncoeffs_right);
1182        for ileft in 0..ncoeffs_left {
1183            for iright in 0..ncoeffs_right {
1184                let mut lpowers = index_to_powers_rev_iter(ileft, nvars_left, degree_left).unwrap();
1185                let mut rpowers =
1186                    index_to_powers_rev_iter(iright, nvars_right, degree_right).unwrap();
1187                let opowers = vars.iter().rev().map(|var| match var {
1188                    MulVar::Left => lpowers.next().unwrap(),
1189                    MulVar::Right => rpowers.next().unwrap(),
1190                    MulVar::Both => lpowers.next().unwrap() + rpowers.next().unwrap(),
1191                });
1192                if let Some(iout) = powers_rev_iter_to_index(opowers, vars.len(), degree_output) {
1193                    indices.push([iout, ileft, iright]);
1194                }
1195            }
1196        }
1197
1198        // Given `indices` we can allocate a vector of zeros as output
1199        // coefficients, walk over `indices` and update the output:
1200        //
1201        //     let mut coeffs_out = vec_of_enough_zeros();
1202        //     for (iout, ileft, iright) in indices {
1203        //         coeffs_out[iout] += coeffs_left[ileft] * coeffs_right[iright];
1204        //     }
1205        //
1206        // However, we also want to be able to get a single coefficient or
1207        // iterate over the coefficients efficiently without allocating the
1208        // output vector first. By sorting `indices` on the output index, we
1209        // can evaluate a single output coefficient by iterating over a slice
1210        // of `indices` and summing the products of coefficients:
1211        //
1212        //     coeffs_out[iout] = indices[offsets[iout]..offsets[iout + 1]].map(|(_, ileft, iright) {
1213        //         coeffs_left[*ileft] * coeffs_right[*iright]
1214        //     }).sum()
1215        //
1216        // with vec `offsets` defined below.
1217        //
1218        // We can use unstable sort because the `ileft` and `iright` pairs are
1219        // unique.
1220        indices.sort_unstable();
1221
1222        let mut oiter = indices.iter().map(|[iout, _, _]| *iout).peekable();
1223        let count: Box<[usize]> = Iterator::map(0..ncoeffs_out, |i| {
1224            let mut n = 0;
1225            while oiter.next_if(|j| i == *j).is_some() {
1226                n += 1;
1227            }
1228            n
1229        })
1230        .collect();
1231
1232        let offsets = iter::once(0)
1233            .chain(count.iter().copied())
1234            .scan(0, |acc, n| {
1235                *acc += n;
1236                Some(*acc)
1237            })
1238            .collect();
1239
1240        // Strip the output index from `indices` because we use `offsets` and
1241        // `count` in `Sequence` instead.
1242        let indices = indices
1243            .into_iter()
1244            .map(|[_, ileft, iright]| [ileft, iright])
1245            .collect();
1246
1247        Self {
1248            count,
1249            offsets,
1250            indices,
1251            ncoeffs_left,
1252            ncoeffs_right,
1253            nvars_output,
1254            degree_output,
1255        }
1256    }
1257
1258    /// Plan the product of two polynomials in the same variables.
1259    #[inline]
1260    pub fn same_vars(nvars: usize, degree_left: Power, degree_right: Power) -> Self {
1261        Self::same_vars_for_output_degree(
1262            nvars,
1263            degree_left,
1264            degree_right,
1265            degree_left + degree_right,
1266        )
1267    }
1268
1269    /// Plan the product of two polynomials in the same variables for the given output degree.
1270    ///
1271    /// If the output degree is smaller than the sum of the degrees of the input operands,
1272    /// then the product is truncated.
1273    #[inline]
1274    pub fn same_vars_for_output_degree(
1275        nvars: usize,
1276        degree_left: Power,
1277        degree_right: Power,
1278        degree_output: Power,
1279    ) -> Self {
1280        Self::for_output_degree(
1281            &[MulVar::Both].copied().repeat(nvars),
1282            degree_left,
1283            degree_right,
1284            degree_output,
1285        )
1286    }
1287
1288    /// Plan the product of two polynomials in different variables.
1289    ///
1290    /// The coefficients returned by [`MulPlan::apply()`] are ordered such that
1291    /// the first `nvars_left` variables are the variables of the left operand and
1292    /// the last `nvars_right` are the variables of the right operand.
1293    #[inline]
1294    pub fn different_vars(
1295        nvars_left: usize,
1296        nvars_right: usize,
1297        degree_left: Power,
1298        degree_right: Power,
1299    ) -> Result<Self, Error> {
1300        Self::different_vars_for_output_degree(
1301            nvars_left,
1302            nvars_right,
1303            degree_left,
1304            degree_right,
1305            degree_left + degree_right,
1306        )
1307    }
1308
1309    /// Plan the product of two polynomials in different variables for the given output degree.
1310    ///
1311    /// The coefficients returned by [`MulPlan::apply()`] are ordered such that
1312    /// the first `nvars_left` variables are the variables of the left operand and
1313    /// the last `nvars_right` are the variables of the right operand.
1314    ///
1315    /// If the output degree is smaller than the sum of the degrees of the input operands,
1316    /// then the product is truncated.
1317    #[inline]
1318    pub fn different_vars_for_output_degree(
1319        nvars_left: usize,
1320        nvars_right: usize,
1321        degree_left: Power,
1322        degree_right: Power,
1323        degree_output: Power,
1324    ) -> Result<Self, Error> {
1325        let vars = [MulVar::Left]
1326            .copied()
1327            .repeat(nvars_left)
1328            .concat([MulVar::Right].copied().repeat(nvars_right))
1329            .ok_or(Error::TooManyVariables)?;
1330        Ok(Self::for_output_degree(
1331            &vars,
1332            degree_left,
1333            degree_right,
1334            degree_output,
1335        ))
1336    }
1337
1338    /// Returns the coefficients for the product of two polynomials.
1339    ///
1340    /// # Errors
1341    ///
1342    /// This function returns an error if the number of coefficients of the
1343    /// left or right polynomial doesn't match [`MulPlan::ncoeffs_left()`] or
1344    /// [`MulPlan::ncoeffs_right()`], respectively.
1345    #[inline]
1346    pub fn apply<LCoeffs, RCoeffs>(
1347        &self,
1348        coeffs_left: LCoeffs,
1349        coeffs_right: RCoeffs,
1350    ) -> Result<Mul<'_, LCoeffs, RCoeffs>, Error>
1351    where
1352        LCoeffs: Sequence,
1353        RCoeffs: Sequence,
1354    {
1355        check_ncoeffs_sqnc(&coeffs_left, self.ncoeffs_left(), Some("left polynomial"))?;
1356        check_ncoeffs_sqnc(
1357            &coeffs_right,
1358            self.ncoeffs_right(),
1359            Some("right polynomial"),
1360        )?;
1361        Ok(Mul {
1362            plan: self,
1363            coeffs_left,
1364            coeffs_right,
1365        })
1366    }
1367
1368    /// Returns the number of coefficients for the left operand.
1369    #[inline]
1370    pub fn ncoeffs_left(&self) -> usize {
1371        self.ncoeffs_left
1372    }
1373
1374    /// Returns the number of coefficients for the left operand.
1375    #[inline]
1376    pub fn ncoeffs_right(&self) -> usize {
1377        self.ncoeffs_right
1378    }
1379
1380    /// Returns the number of coefficients for the product.
1381    #[inline]
1382    pub fn ncoeffs_output(&self) -> usize {
1383        self.count.len()
1384    }
1385
1386    /// Returns the number of variables for the product.
1387    #[inline]
1388    pub fn nvars_output(&self) -> usize {
1389        self.nvars_output
1390    }
1391
1392    /// Returns the degree of the product.
1393    #[inline]
1394    pub fn degree_output(&self) -> Power {
1395        self.degree_output
1396    }
1397}
1398
1399/// The product of two polynomials.
1400///
1401/// This struct is created by [`MulPlan`]. See its documentation for more
1402/// information.
1403#[derive(Debug, Clone)]
1404pub struct Mul<'plan, LCoeffs, RCoeffs> {
1405    plan: &'plan MulPlan,
1406    coeffs_left: LCoeffs,
1407    coeffs_right: RCoeffs,
1408}
1409
1410impl<'this, 'plan, LCoeffs, RCoeffs, OCoeff> SequenceTypes<'this> for Mul<'plan, LCoeffs, RCoeffs>
1411where
1412    for<'a> <LCoeffs as SequenceTypes<'a>>::Item:
1413        ops::Mul<<RCoeffs as SequenceTypes<'a>>::Item, Output = OCoeff>,
1414    OCoeff: ops::AddAssign + Zero,
1415    LCoeffs: Sequence,
1416    RCoeffs: Sequence,
1417{
1418    type Item = OCoeff;
1419    type Iter = sqnc::derive::Iter<'this, Self>;
1420}
1421
1422impl<'plan, LCoeffs, RCoeffs, OCoeff> Sequence for Mul<'plan, LCoeffs, RCoeffs>
1423where
1424    for<'a> <LCoeffs as SequenceTypes<'a>>::Item:
1425        ops::Mul<<RCoeffs as SequenceTypes<'a>>::Item, Output = OCoeff>,
1426    OCoeff: ops::AddAssign + Zero,
1427    LCoeffs: Sequence,
1428    RCoeffs: Sequence,
1429{
1430    #[inline]
1431    fn len(&self) -> usize {
1432        self.plan.ncoeffs_output()
1433    }
1434
1435    #[inline]
1436    fn get(&self, index: usize) -> Option<OCoeff> {
1437        let start = *self.plan.offsets.get(index)?;
1438        let stop = *self.plan.offsets.get(index + 1)?;
1439        let n = stop.saturating_sub(start);
1440        let mut value = OCoeff::zero();
1441        for [l, r] in self.plan.indices.iter().skip(start).take(n) {
1442            value += self.coeffs_left.get(*l)? * self.coeffs_right.get(*r)?;
1443        }
1444        Some(value)
1445    }
1446
1447    #[inline]
1448    fn iter(&self) -> sqnc::derive::Iter<'_, Self> {
1449        self.into()
1450    }
1451}
1452
1453impl<'plan, LCoeffs, RCoeffs, OCoeff> IntoIterator for Mul<'plan, LCoeffs, RCoeffs>
1454where
1455    for<'a> <LCoeffs as SequenceTypes<'a>>::Item:
1456        ops::Mul<<RCoeffs as SequenceTypes<'a>>::Item, Output = OCoeff>,
1457    OCoeff: ops::AddAssign + Zero,
1458    LCoeffs: Sequence,
1459    RCoeffs: Sequence,
1460{
1461    type Item = OCoeff;
1462    type IntoIter = sqnc::derive::IntoIter<Self>;
1463
1464    #[inline]
1465    fn into_iter(self) -> Self::IntoIter {
1466        self.into()
1467    }
1468}
1469
1470/// Returns a matrix that flattens coefficients for a composition of polynomials.
1471///
1472/// Let $f$ be a polynomial in $m$ variables of degree $p$ and $g$ a vector of
1473/// $m$ polynomials in $n$ variables of degree $q$. The composition $f ∘ g$ is
1474/// a polynomial in $n$ variables of degree $p q$.
1475///
1476/// This function returns a matrix $M$ that, when multiplied with a vector of
1477/// coefficients belonging to the outer polynomial $p$, produces the
1478/// coefficients for the composition, for a fixed vector of inner polynomials
1479/// $g$.
1480///
1481/// Argument `inner_coeffs` is a flattened sequence of the coefficients for the
1482/// vector of inner polynomials $g$, with `inner_nvars` the number of inner
1483/// variables $n$ and `inner_degree` the degree $q$.
1484///
1485/// # Errors
1486///
1487/// If the length of `inner_coeffs` sequence differs from the expected length,
1488/// this function returns an error.
1489///
1490/// # Examples
1491///
1492/// Let $f(x) = x_1^2 + 2 x_0 x_1$ (coefficients: `[1, 2, 0, 0, 0, 0]`) and
1493/// $g(y) = [y_1 - y_2, 3 y_0 + 2]$ (flattened coefficients: `[-1, 1, 0, 0, 0,
1494/// 0, 3, 2]`]). The composition is
1495///
1496/// $$\\begin{align}
1497///     f(g(y)) &= (3 y_0 + 2)^2 + 2 (y_1 - y_2) (3 y_0 + 2) \\\\
1498///             &= -6 y_2 y_0 - 4 y_2 + 6 y_1 y_0 + 4 y_1 + 9 y_0^2 + 12 y_0 + 4,
1499/// \\end{align}$$
1500///
1501/// (coefficients: `[0, 0, -6, -4, 0, 6, 4, 9, 12, 4]`).
1502///
1503/// ```
1504/// use ndarray::array;
1505/// use nutils_poly;
1506/// use sqnc::traits::*;
1507///
1508/// let f = array![1, 2, 0, 0, 0, 0];
1509/// let g = array![-1, 1, 0, 0, 0, 0, 3, 2];
1510/// let m = nutils_poly::composition_with_inner_matrix(g.iter().copied(), 3, 1, 2, 2)?;
1511/// assert_eq!(m.dot(&f), array![0, 0, -6, -4, 0, 6, 4, 9, 12, 4]);
1512/// # Ok::<_, nutils_poly::Error>(())
1513/// ```
1514pub fn composition_with_inner_matrix<Coeff, InnerCoeffs>(
1515    mut inner_coeffs: InnerCoeffs,
1516    inner_nvars: usize,
1517    inner_degree: Power,
1518    outer_nvars: usize,
1519    outer_degree: Power,
1520) -> Result<Array2<Coeff>, Error>
1521where
1522    Coeff: One + Zero + Clone + ops::AddAssign,
1523    for<'a> &'a Coeff: ops::Mul<Output = Coeff>,
1524    InnerCoeffs: Iterator<Item = Coeff>,
1525{
1526    let outer_ncoeffs = ncoeffs(outer_nvars, outer_degree);
1527
1528    let result_degree = inner_degree * outer_degree;
1529    let result_ncoeffs = ncoeffs(inner_nvars, result_degree);
1530
1531    let mut matrix: Array2<Coeff> = Array2::zeros((result_ncoeffs, outer_ncoeffs));
1532
1533    // Initialize the column with total degree zero.
1534    matrix[(result_ncoeffs - 1, outer_ncoeffs - 1)] = Coeff::one();
1535    if outer_ncoeffs == 1 {
1536        return Ok(matrix);
1537    }
1538
1539    // Initialize the columns with total degree one.
1540    let mapped_indices = MapDegree::new(inner_nvars, inner_degree, result_degree).unwrap();
1541    for ivar in 0..outer_nvars {
1542        let icolumn = outer_ncoeffs - 1 - ncoeffs(ivar, outer_degree);
1543        for irow in mapped_indices.iter() {
1544            matrix[(irow, icolumn)] =
1545                inner_coeffs
1546                    .next()
1547                    .ok_or_else(|| Error::IncorrectNumberOfCoefficients {
1548                        expected: ncoeffs(inner_nvars, inner_degree) * outer_nvars,
1549                        got: Less,
1550                        detail: None,
1551                    })?;
1552        }
1553    }
1554    if inner_coeffs.next().is_some() {
1555        return Err(Error::IncorrectNumberOfCoefficients {
1556            expected: ncoeffs(inner_nvars, inner_degree) * outer_nvars,
1557            got: More,
1558            detail: None,
1559        });
1560    }
1561
1562    let mul = MulPlan::for_output_degree(
1563        &[MulVar::Both].copied().repeat(inner_nvars),
1564        result_degree,
1565        result_degree,
1566        result_degree,
1567    );
1568    let mut rev_powers: Vec<Power> = iter::repeat(0).take(outer_nvars).collect();
1569    for icol in Iterator::rev(0..outer_ncoeffs) {
1570        iter::zip(
1571            rev_powers.iter_mut(),
1572            index_to_powers_rev_iter(icol, outer_nvars, outer_degree).unwrap(),
1573        )
1574        .for_each(|(d, s)| *d = s);
1575        // Skip the columns that correspond to total degree 0 and 1: we have
1576        // populated these above.
1577        if rev_powers.iter().sum::<Power>() <= 1 {
1578            continue;
1579        }
1580        let rev_ivar = rev_powers.iter().position(|i| *i > 0).unwrap();
1581        let icol1 = outer_ncoeffs - 1 - ncoeffs(outer_nvars - rev_ivar - 1, outer_degree);
1582        rev_powers[rev_ivar] -= 1;
1583        let icol2 = powers_rev_iter_to_index(rev_powers.iter().copied(), outer_nvars, outer_degree)
1584            .unwrap();
1585        let (icol1, icol2) = if icol1 <= icol2 {
1586            (icol1, icol2)
1587        } else {
1588            (icol2, icol1)
1589        };
1590
1591        let mut cols = matrix.columns_mut().into_iter();
1592        let cols = cols.by_ref();
1593        let mut col = cols.nth(icol).unwrap();
1594        let col1 = cols.nth(icol1 - icol - 1).unwrap();
1595        let col1 = col1.as_sqnc();
1596        let col2;
1597        let col2 = if icol2 == icol1 {
1598            col1
1599        } else {
1600            col2 = cols.nth(icol2 - icol1 - 1).unwrap();
1601            col2.as_sqnc()
1602        };
1603        for (dst, src) in iter::zip(col.iter_mut(), mul.apply(col1, col2).unwrap().iter()) {
1604            *dst = src;
1605        }
1606    }
1607
1608    Ok(matrix)
1609}
1610
1611#[cfg(test)]
1612mod tests {
1613    use super::{Error, Less, MapDegree, More, MulPlan, MulVar, Multiple, PartialDerivPlan, Power};
1614    use core::iter;
1615    use ndarray::array;
1616    use sqnc::traits::*;
1617
1618    #[test]
1619    fn error() {
1620        assert_eq!(
1621            Error::IncorrectNumberOfCoefficients {
1622                expected: 2,
1623                got: 3.into(),
1624                detail: None,
1625            }
1626            .to_string(),
1627            "Expected 2 coefficients but got 3.",
1628        );
1629        assert_eq!(
1630            Error::IncorrectNumberOfCoefficients {
1631                expected: 1,
1632                got: More,
1633                detail: None,
1634            }
1635            .to_string(),
1636            "Expected 1 coefficient but got more.",
1637        );
1638        assert_eq!(
1639            Error::IncorrectNumberOfCoefficients {
1640                expected: 2,
1641                got: Less,
1642                detail: Some("left polynomial"),
1643            }
1644            .to_string(),
1645            "left polynomial: Expected 2 coefficients but got less.",
1646        );
1647        assert_eq!(
1648            Error::TooManyVariables.to_string(),
1649            "The number of variables exceeds the maximum (`usize::MAX`).",
1650        );
1651    }
1652
1653    #[test]
1654    fn ncoeffs_degree() {
1655        macro_rules! t {
1656            ($nvars:literal, $ncoeffs_array:expr) => {
1657                let mut ncoeffs_sum = 0;
1658                for (degree, ncoeffs) in $ncoeffs_array.iter().copied().enumerate() {
1659                    let degree = degree as Power;
1660                    assert_eq!(
1661                        super::ncoeffs($nvars, degree),
1662                        ncoeffs,
1663                        "ncoeffs({}, {degree}) != {ncoeffs}",
1664                        $nvars
1665                    );
1666                    assert_eq!(
1667                        super::ncoeffs_sum($nvars, degree),
1668                        ncoeffs_sum,
1669                        "ncoeffs_sum({}, {degree}) != {ncoeffs_sum}",
1670                        $nvars
1671                    );
1672                    if $nvars > 0 || degree == 0 {
1673                        assert_eq!(
1674                            super::degree($nvars, ncoeffs),
1675                            Some(degree),
1676                            "degree({}, {ncoeffs}) != Some({degree})",
1677                            $nvars
1678                        );
1679                    }
1680                    ncoeffs_sum += ncoeffs;
1681                }
1682                assert!(
1683                    super::ncoeffs_iter($nvars)
1684                        .take($ncoeffs_array.len())
1685                        .eq($ncoeffs_array),
1686                    "ncoeffs_iter({}).take({}) != {:?}",
1687                    $nvars,
1688                    $ncoeffs_array.len(),
1689                    $ncoeffs_array
1690                );
1691                assert!(
1692                    super::degree_ncoeffs_iter($nvars)
1693                        .take($ncoeffs_array.len())
1694                        .eq(iter::zip(0.., $ncoeffs_array)),
1695                    "degree_ncoeffs_iter({}).take({}) != {:?}",
1696                    $nvars,
1697                    $ncoeffs_array.len(),
1698                    $ncoeffs_array
1699                );
1700            };
1701        }
1702
1703        t! {0, [1, 1, 1, 1, 1]}
1704        t! {1, [1, 2, 3, 4, 5]}
1705        t! {2, [1, 3, 6, 10, 15]}
1706        t! {3, [1, 4, 10, 20, 35]}
1707        t! {4, [1, 5, 15]}
1708
1709        assert_eq!(super::degree(0, 0), None);
1710        assert_eq!(super::degree(0, 2), None);
1711        assert_eq!(super::degree(1, 0), None);
1712        assert_eq!(super::degree(2, 0), None);
1713        assert_eq!(super::degree(2, 2), None);
1714        assert_eq!(super::degree(2, 4), None);
1715        assert_eq!(super::degree(2, 9), None);
1716        assert_eq!(super::degree(2, 11), None);
1717        assert_eq!(super::degree(3, 0), None);
1718        assert_eq!(super::degree(3, 2), None);
1719        assert_eq!(super::degree(3, 3), None);
1720        assert_eq!(super::degree(4, 3), None);
1721    }
1722
1723    #[test]
1724    fn powers_index_maps() {
1725        macro_rules! assert_index_powers {
1726            ($degree:literal, $desired_powers:tt) => {
1727                let nvars = $desired_powers[0].len();
1728                for (desired_index, desired_powers) in $desired_powers.into_iter().enumerate() {
1729                    assert_eq!(
1730                        super::powers_rev_iter_to_index(
1731                            desired_powers.iter().rev().copied(),
1732                            nvars,
1733                            $degree
1734                        ),
1735                        Some(desired_index)
1736                    );
1737                    assert!(
1738                        super::index_to_powers_rev_iter(desired_index, nvars, $degree)
1739                            .unwrap()
1740                            .eq(desired_powers.iter().rev().copied())
1741                    )
1742                }
1743            };
1744        }
1745
1746        assert_index_powers!(2, [[2], [1], [0]]);
1747        assert_index_powers!(2, [[0, 2], [1, 1], [0, 1], [2, 0], [1, 0], [0, 0]]);
1748        assert_index_powers!(
1749            2,
1750            [
1751                [0, 0, 2],
1752                [0, 1, 1],
1753                [1, 0, 1],
1754                [0, 0, 1],
1755                [0, 2, 0],
1756                [1, 1, 0],
1757                [0, 1, 0],
1758                [2, 0, 0],
1759                [1, 0, 0],
1760                [0, 0, 0]
1761            ]
1762        );
1763    }
1764
1765    #[test]
1766    fn map_degree() {
1767        assert_eq!(MapDegree::new(1, 3, 2), None);
1768        assert!(MapDegree::new(0, 0, 0).unwrap().iter().eq([0]));
1769        assert!(MapDegree::new(1, 2, 2).unwrap().iter().eq([0, 1, 2]));
1770        assert!(MapDegree::new(1, 2, 3).unwrap().iter().eq([1, 2, 3]));
1771        assert!(MapDegree::new(1, 2, 4).unwrap().iter().eq([2, 3, 4]));
1772        assert!(MapDegree::new(2, 1, 2).unwrap().iter().eq([2, 4, 5]));
1773
1774        assert!(MapDegree::new(0, 0, 0).unwrap().into_iter().eq([0]));
1775        assert!(MapDegree::new(1, 2, 2).unwrap().into_iter().eq([0, 1, 2]));
1776    }
1777
1778    #[test]
1779    fn eval_0d() {
1780        let values: [usize; 0] = [];
1781        assert_eq!(super::eval([1], &values, 0), Ok(1));
1782        assert!(super::eval(0..0, &values, 0).is_err());
1783    }
1784
1785    #[test]
1786    fn eval_1d() {
1787        let values = [5];
1788        assert_eq!(super::eval([1], &values, 0), Ok(1));
1789        assert_eq!(super::eval([2, 1], &values, 1), Ok(11));
1790        assert_eq!(super::eval([3, 2, 1], &values, 2), Ok(86));
1791        assert!(super::eval(0..1, &values, 1).is_err());
1792    }
1793
1794    #[test]
1795    fn eval_2d() {
1796        let values = [5, 3];
1797        assert_eq!(super::eval([1], &values, 0), Ok(1));
1798        assert_eq!(super::eval([0, 0, 1], &values, 1), Ok(1));
1799        assert_eq!(super::eval([0, 1, 0], &values, 1), Ok(5));
1800        assert_eq!(super::eval([1, 0, 0], &values, 1), Ok(3));
1801        assert_eq!(super::eval([3, 2, 1], &values, 1), Ok(20));
1802        assert_eq!(super::eval(Iterator::rev(1..7), &values, 2), Ok(227));
1803        assert!(super::eval(0..2, &values, 1).is_err());
1804    }
1805
1806    #[test]
1807    fn eval_3d() {
1808        let values = [5, 3, 2];
1809        assert_eq!(super::eval([1], &values, 0), Ok(1));
1810        assert_eq!(super::eval([0, 0, 0, 1], &values, 1), Ok(1));
1811        assert_eq!(super::eval([0, 0, 1, 0], &values, 1), Ok(5));
1812        assert_eq!(super::eval([0, 1, 0, 0], &values, 1), Ok(3));
1813        assert_eq!(super::eval([1, 0, 0, 0], &values, 1), Ok(2));
1814        assert_eq!(super::eval([4, 3, 2, 1], &values, 1), Ok(28));
1815        assert_eq!(super::eval(Iterator::rev(1..11), &values, 2), Ok(415));
1816        assert!(super::eval(0..3, &values, 1).is_err());
1817    }
1818
1819    #[test]
1820    fn eval_4d() {
1821        let values = [5, 3, 2, 1];
1822        assert_eq!(super::eval([1], &values, 0), Ok(1));
1823        assert_eq!(super::eval([0, 0, 0, 0, 1], &values, 1), Ok(1));
1824        assert_eq!(super::eval([0, 0, 0, 1, 0], &values, 1), Ok(5));
1825        assert_eq!(super::eval([0, 0, 1, 0, 0], &values, 1), Ok(3));
1826        assert_eq!(super::eval([0, 1, 0, 0, 0], &values, 1), Ok(2));
1827        assert_eq!(super::eval([1, 0, 0, 0, 0], &values, 1), Ok(1));
1828        assert!(super::eval(0..4, &values, 1).is_err());
1829    }
1830
1831    #[test]
1832    fn multiple() {
1833        assert_eq!(2.multiple(3), 6);
1834        assert_eq!((&2).multiple(3), 6);
1835    }
1836
1837    #[test]
1838    fn partial_deriv_x0_x() {
1839        let plan = PartialDerivPlan::new(1, 0, 0).unwrap();
1840        assert_eq!(plan.ncoeffs_input(), 1);
1841        assert_eq!(plan.ncoeffs_output(), 1);
1842        assert_eq!(plan.degree_output(), 0);
1843        assert_eq!(plan.nvars(), 1);
1844        let pd = plan.apply([1]).unwrap();
1845        assert!(pd.iter().eq([0]));
1846        assert_eq!(pd.iter().size_hint(), (1, Some(1)));
1847        assert_eq!(pd.get(0), Some(0));
1848        assert_eq!(pd.len(), 1);
1849        let iter = pd.into_iter();
1850        assert_eq!(iter.size_hint(), (1, Some(1)));
1851        assert!(iter.eq([0]));
1852    }
1853
1854    #[test]
1855    fn partial_deriv_x1_x() {
1856        let plan = PartialDerivPlan::new(1, 1, 0).unwrap();
1857        assert_eq!(plan.ncoeffs_input(), 2);
1858        assert_eq!(plan.ncoeffs_output(), 1);
1859        assert_eq!(plan.degree_output(), 0);
1860        assert_eq!(plan.nvars(), 1);
1861        let pd = plan.apply([2, 1]).unwrap();
1862        assert!(pd.iter().eq([2]));
1863        assert_eq!(pd.get(1), None);
1864        assert_eq!(pd.len(), 1);
1865    }
1866
1867    #[test]
1868    fn partial_deriv_x3_x() {
1869        let plan = PartialDerivPlan::new(1, 3, 0).unwrap();
1870        assert_eq!(plan.ncoeffs_input(), 4);
1871        assert_eq!(plan.ncoeffs_output(), 3);
1872        assert_eq!(plan.degree_output(), 2);
1873        assert_eq!(plan.nvars(), 1);
1874        let pd = plan.apply([4, 3, 2, 1]).unwrap();
1875        assert!(pd.iter().eq([12, 6, 2]));
1876        assert_eq!(pd.get(2), Some(2));
1877        assert_eq!(pd.len(), 3);
1878    }
1879
1880    #[test]
1881    fn partial_deriv_xy2_x() {
1882        let plan = PartialDerivPlan::new(2, 2, 0).unwrap();
1883        assert_eq!(plan.ncoeffs_input(), 6);
1884        assert_eq!(plan.ncoeffs_output(), 3);
1885        assert_eq!(plan.degree_output(), 1);
1886        assert_eq!(plan.nvars(), 2);
1887        let pd = plan.apply([6, 5, 4, 3, 2, 1]).unwrap();
1888        assert!(pd.iter().eq([5, 6, 2]));
1889        assert_eq!(pd.get(1), Some(6));
1890        assert_eq!(pd.len(), 3);
1891    }
1892
1893    #[test]
1894    fn partial_deriv_xy2_y() {
1895        let plan = PartialDerivPlan::new(2, 2, 1).unwrap();
1896        assert_eq!(plan.ncoeffs_input(), 6);
1897        assert_eq!(plan.ncoeffs_output(), 3);
1898        assert_eq!(plan.degree_output(), 1);
1899        assert_eq!(plan.nvars(), 2);
1900        let pd = plan.apply([6, 5, 4, 3, 2, 1]).unwrap();
1901        assert!(pd.iter().eq([12, 5, 4]));
1902        assert!(pd.iter().rev().eq([4, 5, 12]));
1903        assert_eq!(pd.get(2), Some(4));
1904        assert_eq!(pd.len(), 3);
1905    }
1906
1907    #[test]
1908    fn partial_deriv_out_of_range() {
1909        assert!(PartialDerivPlan::new(1, 2, 2).is_none());
1910    }
1911
1912    #[test]
1913    fn mul_x1_x1() {
1914        let plan = MulPlan::same_vars(1, 1, 1);
1915        assert_eq!(plan.ncoeffs_left(), 2);
1916        assert_eq!(plan.ncoeffs_right(), 2);
1917        assert_eq!(plan.ncoeffs_output(), 3);
1918        assert_eq!(plan.nvars_output(), 1);
1919        assert_eq!(plan.degree_output(), 2);
1920        // (2 x + 1) (4 x + 3) = 8 x^2 + 10 x + 3
1921        let mul = plan.apply([2, 1], [4, 3]).unwrap();
1922        assert!(mul.iter().eq([8, 10, 3]));
1923        assert!(mul.into_iter().eq([8, 10, 3]));
1924    }
1925
1926    #[test]
1927    fn mul_x1_y1() {
1928        let plan = MulPlan::different_vars(1, 1, 1, 1).unwrap();
1929        assert_eq!(plan.ncoeffs_left(), 2);
1930        assert_eq!(plan.ncoeffs_right(), 2);
1931        assert_eq!(plan.ncoeffs_output(), 6);
1932        assert_eq!(plan.nvars_output(), 2);
1933        assert_eq!(plan.degree_output(), 2);
1934        // (2 x + 1) (4 y + 3) = 8 x y + 4 y + 6 x + 3
1935        let mul = plan.apply([2, 1], [4, 3]).unwrap();
1936        assert!(mul.iter().eq([0, 8, 4, 0, 6, 3]));
1937        assert!(mul.into_iter().eq([0, 8, 4, 0, 6, 3]));
1938    }
1939
1940    #[test]
1941    fn mul_xy1_y2() {
1942        let plan = MulPlan::new(&[MulVar::Left, MulVar::Both].copied(), 1, 2);
1943        assert_eq!(plan.ncoeffs_left(), 3);
1944        assert_eq!(plan.ncoeffs_right(), 3);
1945        assert_eq!(plan.ncoeffs_output(), 10);
1946        assert_eq!(plan.nvars_output(), 2);
1947        assert_eq!(plan.degree_output(), 3);
1948        // (3 y + 2 x + 1) (6 y^2 + 5 y + 4)
1949        // = 18 y^3 + 12 x y^2 + 21 y^2 + 10 x y + 17 y + 8 x + 4
1950        let mul = plan.apply([3, 2, 1], [6, 5, 4]).unwrap();
1951        assert!(mul.iter().eq([18, 12, 21, 0, 10, 17, 0, 0, 8, 4]));
1952        assert!(mul.into_iter().eq([18, 12, 21, 0, 10, 17, 0, 0, 8, 4]));
1953    }
1954
1955    #[test]
1956    fn mul_errors() {
1957        let plan = MulPlan::same_vars(1, 1, 1);
1958        assert_eq!(
1959            plan.apply([2, 1, 3], [4, 3]).unwrap_err(),
1960            Error::IncorrectNumberOfCoefficients {
1961                expected: 2,
1962                got: 3.into(),
1963                detail: Some("left polynomial"),
1964            }
1965        );
1966        assert_eq!(
1967            plan.apply([2, 1], [4, 3, 2]).unwrap_err(),
1968            Error::IncorrectNumberOfCoefficients {
1969                expected: 2,
1970                got: 3.into(),
1971                detail: Some("right polynomial"),
1972            }
1973        );
1974    }
1975
1976    #[test]
1977    fn composition() {
1978        assert_eq!(
1979            super::composition_with_inner_matrix(0..0, 0, 0, 0, 0),
1980            Ok(array![[1]])
1981        );
1982
1983        let f = array![1, 2, 0, 0, 0, 0];
1984        let g = array![-1, 1, 0, 0, 0, 0, 3, 2];
1985        let m = super::composition_with_inner_matrix(g.iter().copied(), 3, 1, 2, 2).unwrap();
1986        assert_eq!(m.dot(&f), array![0, 0, -6, -4, 0, 6, 4, 9, 12, 4]);
1987    }
1988
1989    #[test]
1990    fn composition_errors() {
1991        assert_eq!(
1992            super::composition_with_inner_matrix(0..1, 1, 3, 1, 2),
1993            Err(Error::IncorrectNumberOfCoefficients {
1994                expected: 4,
1995                got: Less,
1996                detail: None,
1997            })
1998        );
1999        assert_eq!(
2000            super::composition_with_inner_matrix(0..9, 1, 3, 1, 2),
2001            Err(Error::IncorrectNumberOfCoefficients {
2002                expected: 4,
2003                got: More,
2004                detail: None,
2005            })
2006        );
2007    }
2008}
2009
2010#[cfg(all(feature = "bench", test))]
2011mod benches {
2012    extern crate test;
2013    use self::test::Bencher;
2014    use super::MulPlan;
2015    use sqnc::traits::*;
2016    use std::iter;
2017
2018    macro_rules! mk_bench_eval {
2019        ($name:ident, $nvars:literal, $degree:literal) => {
2020            #[bench]
2021            fn $name(b: &mut Bencher) {
2022                let coeffs =
2023                    Vec::from_iter(Iterator::map(1..=super::ncoeffs($nvars, $degree), |i| {
2024                        i as f64
2025                    }));
2026                let values: Vec<_> = (1..=$nvars).map(|x| x as f64).collect();
2027                b.iter(|| {
2028                    super::eval(
2029                        test::black_box(coeffs.as_slice()),
2030                        test::black_box(&values[..]),
2031                        $degree,
2032                    )
2033                    .unwrap()
2034                })
2035            }
2036        };
2037    }
2038
2039    mk_bench_eval! {eval_1d_degree1, 1, 1}
2040    mk_bench_eval! {eval_1d_degree2, 1, 2}
2041    mk_bench_eval! {eval_1d_degree3, 1, 3}
2042    mk_bench_eval! {eval_1d_degree4, 1, 4}
2043    mk_bench_eval! {eval_2d_degree1, 2, 1}
2044    mk_bench_eval! {eval_2d_degree2, 2, 2}
2045    mk_bench_eval! {eval_2d_degree3, 2, 3}
2046    mk_bench_eval! {eval_2d_degree4, 2, 4}
2047    mk_bench_eval! {eval_3d_degree1, 3, 1}
2048    mk_bench_eval! {eval_3d_degree2, 3, 2}
2049    mk_bench_eval! {eval_3d_degree3, 3, 3}
2050    mk_bench_eval! {eval_3d_degree4, 3, 4}
2051    mk_bench_eval! {eval_4d_degree1, 4, 1}
2052    mk_bench_eval! {eval_4d_degree2, 4, 2}
2053    mk_bench_eval! {eval_4d_degree3, 4, 3}
2054    mk_bench_eval! {eval_4d_degree4, 4, 4}
2055
2056    #[bench]
2057    fn ncoeffs_3d_degree4(b: &mut Bencher) {
2058        b.iter(|| super::ncoeffs(test::black_box(3), test::black_box(4)));
2059    }
2060
2061    #[bench]
2062    fn mul_xyz4_xyz2(b: &mut Bencher) {
2063        let plan = MulPlan::same_vars(3, 4, 2);
2064        let lcoeffs = Vec::from_iter(Iterator::map(0..super::ncoeffs(3, 4), |i| i as f64));
2065        let rcoeffs = Vec::from_iter(Iterator::map(0..super::ncoeffs(3, 2), |i| i as f64));
2066        let mut ocoeffs = Vec::from_iter(iter::repeat(0f64).take(plan.ncoeffs_output()));
2067        b.iter(|| {
2068            ocoeffs
2069                .as_mut_sqnc()
2070                .assign(
2071                    plan.apply(
2072                        test::black_box(lcoeffs.as_sqnc()),
2073                        test::black_box(rcoeffs.as_sqnc()),
2074                    )
2075                    .unwrap(),
2076                )
2077                .unwrap()
2078        })
2079    }
2080
2081    #[bench]
2082    fn mul_x4_yz2(b: &mut Bencher) {
2083        let plan = MulPlan::different_vars(1, 2, 4, 2).unwrap();
2084        let lcoeffs: Vec<f64> = Iterator::map(0..super::ncoeffs(1, 4), |i| i as f64).collect();
2085        let rcoeffs: Vec<f64> = Iterator::map(0..super::ncoeffs(2, 2), |i| i as f64).collect();
2086        let mut ocoeffs: Vec<f64> = iter::repeat(0f64).take(plan.ncoeffs_output()).collect();
2087        b.iter(|| {
2088            ocoeffs
2089                .as_mut_sqnc()
2090                .assign(
2091                    MulPlan::apply(
2092                        test::black_box(&plan),
2093                        test::black_box(lcoeffs.as_sqnc()),
2094                        test::black_box(rcoeffs.as_sqnc()),
2095                    )
2096                    .unwrap(),
2097                )
2098                .unwrap()
2099        })
2100    }
2101}