"""Calculate the n-th discrete difference along the given axis."""
from __future__ import annotations
from typing import Optional
import numpy
import numpoly
from ..baseclass import ndpoly, PolyLike
from ..dispatch import implements
[docs]@implements(numpy.diff)
def diff(
a: PolyLike,
n: int = 1,
axis: int = -1,
prepend: Optional[PolyLike] = None,
append: Optional[PolyLike] = None,
) -> ndpoly:
"""
Calculate the n-th discrete difference along the given axis.
The first difference is given by ``out[i] = a[i+1] - a[i]`` along the given
axis, higher differences are calculated by using `diff` recursively.
Args:
a:
Input array
n:
The number of times values are "differenced". If zero, the input is
returned as-is.
axis:
The axis along which the difference is taken, default is the last
axis.
prepend, append:
Values to prepend or append to `a` along axis prior to
performing the difference. Scalar values are expanded to
arrays with length 1 in the direction of axis and the shape
of the input array in along all other axes. Otherwise the
dimension and shape must match `a` except along axis.
Return:
The n-th differences. The shape of the output is the same as `a` except
along `axis` where the dimension is smaller by `n`. The type of the
output is the same as the type of the difference between any two
elements of `a`. This is the same as the type of `a` in most cases.
Example:
>>> q0, q1 = numpoly.variable(2)
>>> poly = numpoly.polynomial([1, q0, q1, q0**2, q1-1])
>>> numpoly.diff(poly)
polynomial([q0-1, q1-q0, q0**2-q1, -q0**2+q1-1])
>>> numpoly.diff(poly, n=2)
polynomial([q1-2*q0+1, q0**2-2*q1+q0, -2*q0**2+2*q1-1])
>>> poly = numpoly.polynomial([[q0, 1], [2, q1]])
>>> numpoly.diff(poly)
polynomial([[-q0+1],
[q1-2]])
>>> numpoly.diff(poly, prepend=7, append=q1)
polynomial([[q0-7, -q0+1, q1-1],
[-5, q1-2, 0]])
"""
if append is not None:
if prepend is not None:
a, append, prepend = numpoly.align_exponents(a, append, prepend)
else:
a, append = numpoly.align_exponents(a, append)
elif prepend is not None:
a, prepend = numpoly.align_exponents(a, prepend)
else:
a = numpoly.aspolynomial(a)
out = None
for key in a.keys:
kwargs = {}
if append is not None:
kwargs["append"] = append.values[key]
if prepend is not None:
kwargs["prepend"] = prepend.values[key]
tmp = numpy.diff(a.values[key], n=n, axis=axis, **kwargs)
if out is None:
out = numpoly.ndpoly(
exponents=a.exponents,
shape=tmp.shape,
names=a.indeterminants,
dtype=tmp.dtype,
)
out.values[key] = tmp
assert out is not None
return numpoly.clean_attributes(out)