diff --git a/aerosandbox/numpy/calculus.py b/aerosandbox/numpy/calculus.py index 8f8dad5c..09182054 100644 --- a/aerosandbox/numpy/calculus.py +++ b/aerosandbox/numpy/calculus.py @@ -1,14 +1,25 @@ import numpy as _onp import casadi as _cas from aerosandbox.numpy.determine_type import is_casadi_type +from aerosandbox.numpy.conditionals import where as _where +from aerosandbox.numpy.arithmetic_dyadic import mod as _mod -def diff(a, n=1, axis=-1): +def diff(a, n=1, axis=-1, period=None): """ Calculate the n-th discrete difference along the given axis. See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.diff.html """ + if period is not None: + original_result = diff(a, n=n, axis=axis) + remainder = _mod(original_result, period) + return _where( + remainder > period / 2, + remainder - period, + remainder + ) + if not is_casadi_type(a): return _onp.diff(a, n=n, axis=axis) @@ -42,3 +53,10 @@ def trapz(x, modify_endpoints=False): # TODO unify with NumPy trapz, this is di integral[-1] = integral[-1] + x[-1] * 0.5 return integral + + +if __name__ == '__main__': + import numpy as np + + a = np.linspace(-500, 500, 21) % 360 - 180 + print(diff(a, period=360))