Source code for specderiv

import numpy as np
from numpy.polynomial import Polynomial as poly
from scipy.fft import dct, dst
from collections import deque
from warnings import warn, catch_warnings, simplefilter


[docs] def cheb_deriv(y_n: np.ndarray, t_n: np.ndarray, order: int, axis: int=0, filter: callable=None, dct_type=1, calc_endpoints=True): """Evaluate derivatives with Chebyshev polynomials via discrete cosine and sine transforms. Caveats: - Taking the 1st derivative twice with a discrete method like this is not exactly the same as taking the second derivative. - For derivatives over the 4th, this method presently returns :code:`NaN` at the edges of the domain. Be cautious if passing the result to another function. Args: y_n (np.ndarray): one-or-multi-dimensional array, values of a function, sampled at cosine-spaced points in the dimension of differentiation. t_n (np.ndarray): 1D array, where the function :math:`y` is sampled in the dimension of differentation. If you're using canonical Chebyshev points with the DCT-I, this will be :code:`x_n = np.cos(np.arange(N+1) * np.pi / N)` (:math:`x \\in [1, -1]`). With the DCT-II, :code:`x_n = np.concatenate(([1], np.cos((np.arange(N+1) + 0.5) * np.pi/(N+1)), [-1]))`. In either case, on a domain from :math:`a` to :math:`b`, this is stretched to :code:`t_n = x_n * (b - a)/2 + (b + a)/2`. Note the order is high-to-low in the :math:`x` or :math:`t` domain, but low-to-high in :math:`n`. Also note both endpoints are *inclusive*. order (int): The order of differentiation, also called :math:`\\nu`. Must be :math:`\\geq 1`. axis (int, optional): For multi-dimensional :code:`y_n`, the dimension along which to take the derivative. Defaults to the first dimension (axis=0). filter (callable, optional): A function or :code:`lambda` that takes the 1D array of Chebyshev polynomial numbers, :math:`k = [0, ... N]`, and returns a same-shaped array of weights, which get multiplied in to the initial frequency transform of the data, :math:`Y_k`. Can be helpful when taking derivatives of noisy data. The default is to apply #nofilter. dct_type (int, optional): 1 or 2, whether to use DCT-I or DCT-II. Defaults to DCT-I. calc_endpoints (bool, optional): Whether to calculate the endpoints of the answer, in case they are unnecessary for a particular use case. Defaults to :code:`True`. :code:`False` silences the :code:`NaN` warning for :code:`order` :math:`> 4`. :returns: (*np.ndarray*) -- :code:`dy_n`, shaped like :code:`y_n`, samples of the :math:`\\nu^{th}` derivative of the function """ # We only have to care about the number of points in the dimension we're differentiating N = y_n.shape[axis] - 1 if dct_type == 1 else y_n.shape[axis] - 3 # if type is 1, we count [0, ... N], if type 2, the endpoints are tacked on additionally M = 2*N if dct_type == 1 else 2*(N+1) # Normalization factor is larger for DCT-II based on repeats of endpoints in equivalent FFT x_n = np.cos(np.arange(N+1) * np.pi/N) if dct_type == 1 else np.concatenate(([1], np.cos((np.arange(N+1) + 0.5) * np.pi/(N+1)), [-1])) # canonical sampling if order < 1: raise ValueError("derivative order, nu, should be >= 1") if dct_type not in (1, 2): raise ValueError("DCT type must be 1 or 2") if len(t_n.shape) > 1 or t_n.shape[0] != y_n.shape[axis]: raise ValueError("t_n should be 1D and have the same length as y_n along the axis of differentiation") if not np.all(np.diff(t_n) < 0): raise ValueError("The domain, t_n, should be ordered high-to-low, [b, ... a]. Try sampling with `np.cos(np.arange(N+1) * np.pi / N) * (b - a)/2 + (b + a)/2`") scale = (t_n[0] - t_n[-1])/2; offset = (t_n[0] + t_n[-1])/2 # Trying to be helpful, because sampling is tricky to get right if not np.allclose(t_n, x_n * scale + offset, atol=1e-5): raise ValueError(f"""Your function is not sampled appropriately for the DCT-{'I'*dct_type}. Try sampling with {'`np.cos(np.arange(N+1) * np.pi / N) * (b - a)/2 + (b + a)/2`' if dct_type == 1 else '`np.concatenate(([b], np.cos((np.arange(N+1) + 0.5) * np.pi/(N+1)) * (b - a)/2 + (b + a)/2, [a]))`'}""") first = [slice(None) for dim in y_n.shape]; first[axis] = 0; first = tuple(first) # for accessing different parts of data last = [slice(None) for dim in y_n.shape]; last[axis] = -1; last = tuple(last) middle = [slice(None) for dim in y_n.shape]; middle[axis] = slice(1, -1); middle = tuple(middle) s = [np.newaxis for dim in y_n.shape]; s[axis] = slice(None); s = tuple(s) # for elevating vectors to have same dimension as data Y_k = dct(y_n, 1, axis=axis) if dct_type == 1 else dct(y_n[middle], 2, axis=axis) # Transform to frequency domain using the 1st definition of the discrete cosine transform k = np.arange(N+1) # [0, ... N], Chebyshev basis polynomial (in x)/wavenumber (in theta) iterator if filter: Y_k *= filter(k)[s] y_primes = [] # Store all derivatives in theta up to the nu^th, because we need them all for reconstruction. for mu in range(1, order + 1): Y_mu = (1j * k[s])**mu * Y_k if mu % 2: # odd derivative # In DST-I case Y_mu[k=0 and N] = 0 and so are not needed for the DST, so only pass the [middle] entries # In DST-III case, Y_mu[0 and N+1] = 0. roll() shifts to the left, so Y'_0 is treated like Y'_{N+1}, and we pass in starting at k=1 y_primes.append(dst(1j * Y_mu[middle], 1, axis=axis).real / M if dct_type == 1 # d/dtheta y = the inverse transform of DST-1 = 1/M * DST-1. Extra j for equivalence with IFFT. else dst(1j * np.roll(Y_mu, -1), 3, axis=axis).real / M) # inverse of DST-II is 1/M * DST-III. Im{y_prime} = 0 for real y, so just keep real. else: # even derivative y_primes.append(dct(Y_mu, 1, axis=axis)[middle].real / M if dct_type == 1 # the inverse transform of DCT-1 is 1/M * DCT-1. Slice off ends to get same length as DST-I result. else dct(Y_mu, 3, axis=axis).real / M) # inverse of DCT-II is 1/M * DCT-III. Im{y_prime} = 0 for real y, so just keep real. # Calculate the polynomials in x necessary for transforming back to the Chebyshev domain numers = deque([poly([-1])]) # just -1 to start, at order 1 denom = poly([1, 0, -1]) # 1 - x^2 for nu in range(2, order + 1): # initialization takes care of order 1, so iterate from order 2 q = 0 for mu in range(1, nu): # Terms come from the previous derivative, so there are nu - 1 of them here. p = numers.popleft() # c = nu - mu/2 numers.append(denom * p.deriv() + (nu - mu/2 - 1) * poly([0, 2]) * p - q) q = p numers.append(-q) # Calculate x derivative as a sum of x polynomials * theta-domain derivatives dy_n = np.zeros(y_n.shape) # The middle of dy will get filled with a derivative expression in terms of y_primes denom_x = denom(x_n[1:-1]) # only calculate this once; leave off +/-1, because they need to be treated specially anyway for term,(numer,y_prime) in enumerate(zip(numers, y_primes), 1): # iterating from lower derivatives to higher c = order - term/2 # c starts at nu - 1/2 and then loses 1/2 for each subsequent term dy_n[middle] += (numer(x_n[1:-1])/np.power(denom_x, c))[s] * y_prime # Calculate the endpoints if order <= 4 and calc_endpoints: C, D = {1: [(-1,), 1], 2: [(1, 1), 3], 3: [(-4, -5, -1), 15], 4: [(36, 49, 14, 1), 105]}[order] # Constants from the math. See the notebook in the warning. LH = 0 # L'Hôpital numerator terms for i,C_i in enumerate(C, 1): # i starts at 1 LH += 2 * C_i * (-1)**i * np.power(k, 2*i) if dct_type == 1: LH[-1] -= C_i * (-1)**i * np.power(N, 2*i) # because Nth element is outside the 2\sum in the DCT-I dy_n[first] = np.sum(LH[s] * Y_k, axis=axis)/ (D*M) dy_n[last] = np.sum((LH * np.power(-1, k))[s] * Y_k, axis=axis) / ((-1)**order * D*M) else: # For higher derivatives, leave the endpoints uncalculated, but direct the user to my analysis of this problem. if calc_endpoints: warn("""endpoints set to NaN, only calculated for 4th derivatives and below. For help with higher derivatives, see https://github.com/pavelkomarov/spectral-derivatives/blob/main/notebooks/chebyshev_domain_endpoints.ipynb""") dy_n[first] = dy_n[last] = np.nan # The above is agnostic to where the data came from, pretends it came from the domain [-1, 1], but the data may actually be return dy_n/scale**order # smooshed from some other domain. So scale the derivative by the relative size of the t and x intervals.
[docs] def fourier_deriv(y_n: np.ndarray, t_n: np.ndarray, order: int, axis: int=0, filter: callable=None): """Evaluate derivatives with complex exponentials via FFT. Caveats: - Only for use with periodic functions. - Taking the 1st derivative twice with a discrete method like this is not exactly the same as taking the second derivative. Args: y_n (np.ndarray): one-or-multi-dimensional array, values of a period of a periodic function, sampled at equispaced points in the dimension of differentiation. t_n (np.ndarray): 1D array, where the function :math:`y` is sampled in the dimension of differentiation. If you're using canonical Fourier points, this will be :code:`th_n = np.arange(M) * 2*np.pi / M` (:math:`\\theta \\in [0, 2\\pi)`). If you're sampling on a domain from :math:`a` to :math:`b`, this needs to be :code:`t_n = np.arange(0, M)/M * (b - a) + a`. Note the lower, left bound is *inclusive* and the upper, right bound is *exclusive*. order (int): The order of differentiation, also called :math:`\\nu`. Can be positive (derivative) or negative (antiderivative, raises warning). axis (int, optional): For multi-dimensional :code:`y_n`, the dimension along which to take the derivative. Defaults to the first dimension (axis=0). filter (callable, optional): A function or :code:`lambda` that takes the array of wavenumbers, :math:`k = [0, ... \\frac{M}{2} , -\\frac{M}{2} + 1, ... -1]` for even :math:`M` or :math:`k = [0, ... \\lfloor \\frac{M}{2} \\rfloor, -\\lfloor \\frac{M}{2} \\rfloor, ... -1]` for odd :math:`M`, and returns a same-shaped array of weights, which get multiplied in to the initial frequency transform of the data, :math:`Y_k`. Can be helpful when taking derivatives of noisy data. The default is to apply #nofilter. :returns: (*np.ndarray*) -- :code:`dy_n`, shaped like :code:`y_n`, samples of the :math:`\\nu^{th}` derivative of the function """ #No worrying about conversion back from a variable transformation. No special treatment of domain boundaries. if len(t_n.shape) > 1 or t_n.shape[0] != y_n.shape[axis]: raise ValueError("t_n should be 1D and have the same length as y_n along the axis of differentiation") delta_t = np.diff(t_n) if not (np.all(delta_t > 0) and np.allclose(delta_t, delta_t[0])): raise ValueError("The domain, t_n, needs to be equispaced, ordered low-to-high on [a, ... b). Try sampling with `np.arange(0, M)/M * (b - a) + a`") M = y_n.shape[axis] # if M has an even length, then we make k = [0, 1, ... M/2 - 1, 0 or M/2, -M/2 + 1, ... -1] # if M has odd length, k = [0, 1, ... floor(M/2), -floor(M/2), ... -1] k = np.concatenate((np.arange(M//2 + 1), np.arange(-M//2 + 1, 0))) if M % 2 == 0 and order % 2 == 1: k[M//2] = 0 # odd derivatives get the Nyquist element zeroed out, if there is one s = [np.newaxis for dim in y_n.shape]; s[axis] = slice(None); s = tuple(s) # for elevating vectors to have same dimension as data Y_k = np.fft.fft(y_n, axis=axis) if filter: Y_k *= filter(k)[s] with catch_warnings(): simplefilter("ignore", category=RuntimeWarning); Y_nu = (1j * k[s])**order * Y_k # if nu < 0, we're dividing by 0 if order < 0: Y_nu[np.where(k==0)] = 0; warn("+c information lost in antiderivative") # Get rid of NaNs. Enables taking the antiderivative. dy_n = np.fft.ifft(Y_nu, axis=axis).real if not np.iscomplexobj(y_n) else np.fft.ifft(Y_nu, axis=axis) # The above is agnostic to where the data came from, pretends it came from the domain [0, 2pi), but the data may actually scale = (t_n[M-1] + t_n[1] - 2*t_n[0])/(2*np.pi) # be smooshed from some other domain. So scale the derivative by the return dy_n/scale**order # relative size of the t and theta intervals.