import numpy as _onp
import casadi as _cas
from aerosandbox.numpy.determine_type import is_casadi_type
from aerosandbox.numpy.arithmetic_dyadic import centered_mod as _centered_mod
from aerosandbox.numpy.array import array, concatenate, reshape
from typing import Tuple
[docs]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
Adds one new argument, `period`, which is the period of the data. If provided, the difference is taken assuming
the data "wraps around" at the period (i.e., modulo the period). For example:
>>> diff([345, 355, 5, 15], period=360)
>>> [10, 10, 10, 10]
"""
if period is not None:
return _centered_mod(
diff(a, n=n, axis=axis),
period
)
if not is_casadi_type(a):
return _onp.diff(a, n=n, axis=axis)
else:
if axis != -1:
raise NotImplementedError("This could be implemented, but haven't had the need yet.")
result = a
for i in range(n):
result = _cas.diff(a)
return result
[docs]def gradient(
f,
*varargs,
axis=None,
edge_order=1,
n=1,
period=None,
):
"""
Return the gradient of an N-dimensional array.
The gradient is computed using second order accurate central differences in the interior points and either first
or second order accurate one-sides (forward or backwards) differences at the boundaries. The returned gradient
hence has the same shape as the input array.
See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.gradient.html
Args:
f: The array-like object to take the gradient of.
*varargs: The spacing between the points of f. If a scalar, the spacing is assumed to be uniform in all
dimensions. If an array, the array must have the same shape as f.
axis: The axis along which the difference is taken. If None, the gradient is taken for all axes.
edge_order: The order of the error at the boundaries. 1 means first order, 2 means second order.
n: This is a new argument (not in NumPy) that specifies the order of the derivative to take. 1 is the first
derivative (default), 2 is the second derivative. Doing `np.gradient(f, n=2)` results in less discretization
error than doing `np.gradient(np.gradient(f))`.
period: The period of the data. If provided, the gradient is taken assuming the data "wraps around" at the period
(i.e., modulo the period). See `aerosandbox.numpy.diff()` for more information.
Returns: The gradient of f.
"""
if (
not is_casadi_type(f)
and all([not is_casadi_type(vararg) for vararg in varargs])
and n == 1
and period is None
):
return _onp.gradient(
f,
*varargs,
axis=axis,
edge_order=edge_order
)
else:
f = array(f)
shape = f.shape
# Handle the varargs argument
if len(varargs) == 0:
varargs = (1.,)
if len(varargs) == 1:
varargs = [varargs[0] for i in range(len(shape))]
if len(varargs) != len(shape):
raise ValueError("You must specify either 0, 1, or N varargs, where N is the number of dimensions of f.")
else:
dxes = []
for i, vararg in enumerate(varargs):
if _onp.prod(array(vararg).shape) == 1: # If it's a scalar, you have dx values
dxes.append(
vararg * _onp.ones(shape[i] - 1)
)
else:
dxes.append(
diff(
vararg,
)
)
# Handle the axis argument, with the edge case the CasADi arrays are always 2D
if axis is None:
if is_casadi_type(f, recursive=False) and shape[1] == 1:
axis = 0
elif len(shape) <= 1:
axis = 0
else:
axis = tuple(_onp.arange(len(shape)))
try:
tuple(axis) # See if axis is iterable
axis_is_iterable = True
except TypeError:
axis_is_iterable = False
if axis_is_iterable:
return [
gradient(
f,
varargs[axis_i],
axis=axis_i,
edge_order=edge_order,
n=n,
period=period,
)
for axis_i in axis
]
else:
# Check validity of axis
if axis < 0:
axis = len(shape) + axis
if is_casadi_type(f, recursive=False):
if axis not in [0, 1]:
raise ValueError("axis must be 0 or 1 for CasADi arrays.")
dx = dxes[axis]
dx_shape = [1] * len(shape)
dx_shape[axis] = shape[axis] - 1
dx = reshape(dx, dx_shape)
def get_slice(slice_obj: slice) -> Tuple[slice]:
slices = [slice(None)] * len(shape)
slices[axis] = slice_obj
return tuple(slices)
hm = dx[get_slice(slice(None, -1))]
hp = dx[get_slice(slice(1, None))]
dfp = (
f[get_slice(slice(2, None))]
- f[get_slice(slice(1, -1))]
)
dfm = (
f[get_slice(slice(1, -1))]
- f[get_slice(slice(None, -2))]
)
if period is not None:
dfp = _centered_mod(dfp, period)
dfm = _centered_mod(dfm, period)
if n == 1:
grad_f = (
hm ** 2 * dfp + hp ** 2 * dfm
) / (
hm * hp * (hm + hp)
)
if edge_order == 1:
# First point
df_f = dfm[get_slice(slice(0, 1))]
h_f = hm[get_slice(slice(0, 1))]
grad_f_first = df_f / h_f
# Last point
df_l = dfp[get_slice(slice(-1, None))]
h_l = hp[get_slice(slice(-1, None))]
grad_f_last = df_l / h_l
elif edge_order == 2:
# First point
dfm_f = dfm[get_slice(slice(0, 1))]
dfp_f = dfp[get_slice(slice(0, 1))]
hm_f = hm[get_slice(slice(0, 1))]
hp_f = hp[get_slice(slice(0, 1))]
grad_f_first = (
2 * dfm_f * hm_f * hp_f + dfm_f * hp_f ** 2 - dfp_f * hm_f ** 2
) / (
hm_f * hp_f * (hm_f + hp_f)
)
# Last point
dfm_l = dfm[get_slice(slice(-1, None))]
dfp_l = dfp[get_slice(slice(-1, None))]
hm_l = hm[get_slice(slice(-1, None))]
hp_l = hp[get_slice(slice(-1, None))]
grad_f_last = (
-dfm_l * hp_l ** 2 + dfp_l * hm_l ** 2 + 2 * dfp_l * hm_l * hp_l
) / (
hm_l * hp_l * (hm_l + hp_l)
)
else:
raise ValueError("Invalid edge_order.")
grad_f = concatenate((
grad_f_first,
grad_f,
grad_f_last
), axis=axis)
return grad_f
elif n == 2:
grad_grad_f = (
2 / (hm + hp) * (
dfp / hp - dfm / hm
)
)
grad_grad_f_first = grad_grad_f[get_slice(slice(0, 1))]
grad_grad_f_last = grad_grad_f[get_slice(slice(-1, None))]
grad_grad_f = concatenate((
grad_grad_f_first,
grad_grad_f,
grad_grad_f_last
), axis=axis)
return grad_grad_f
else:
raise ValueError(
"A second-order reconstructor only supports first derivatives (n=1) and second derivatives (n=2).")
[docs]def trapz(x, modify_endpoints=False): # TODO unify with NumPy trapz, this is different
"""
Computes each piece of the approximate integral of `x` via the trapezoidal method with unit spacing.
Can be viewed as the opposite of diff().
Args:
x: The vector-like object (1D np.ndarray, cas.MX) to be integrated.
Returns: A vector of length N-1 with each piece corresponding to the mean value of the function on the interval
starting at index i.
"""
import warnings
warnings.warn(
"trapz() will eventually be deprecated, since NumPy plans to remove it in the upcoming NumPy 2.0 release (2024). \n"
"For discrete intervals, use asb.numpy.integrate_discrete_intervals(f, method=\"trapz\") instead.",
PendingDeprecationWarning)
integral = (
x[1:] + x[:-1]
) / 2
if modify_endpoints:
integral[0] = integral[0] + x[0] * 0.5
integral[-1] = integral[-1] + x[-1] * 0.5
return integral
if __name__ == '__main__':
import aerosandbox as asb
import aerosandbox.numpy as np
import casadi as cas
# print(diff(cas.DM([355, 5]), period=360))
print(
gradient(
np.linspace(45, 55, 11) % 50,
period=50
)
)
#
# # a = np.linspace(-500, 500, 21) % 360 - 180
# # print(diff(a, period=360))
#
# x = np.cumsum(np.arange(10))
# y = x ** 2
#
# print(gradient(y, x, edge_order=1))
# print(gradient(y, x, edge_order=1))
# print(gradient(y, x, edge_order=1, n=2))
#
# opti = asb.Opti()
# x = opti.variable(init_guess=[355, 5])
# d = diff(x, period=360)
# opti.subject_to([
# # x[0] == 3,
# x[0] > 180,
# x[1] < 180,
# d < 20,
# d > -20
# ])
# opti.maximize(np.sum(np.cosd(x)))
# sol = opti.solve(
# behavior_on_failure="return_last"
# )
# print(sol(x))