Improving Spectral (#53)

* spitballing

* my interval was way wrong

* updated spectral-derivatives such that some of the checks I had in this PR are no longer necessary

* added some road signs

* updated to spectral-derivatives 0.7

* this test code encourages not-thorough testing

* fixed caching test

* added tests for chebyshev case and removed caching check for spectral because evidently unnecessary

* added Chebyshev examples

* removed print statement

* bumped version number

* added axis param to tests so they pass now

* improved docstring to account for the fact filtering in chebyshev basis is a little more subtle and bumped spectral-derivatives version number to get latest features

* added back a newline

Author: pavelkomarov

derivative/dglobal.py

@@ -7,40 +7,50 @@
 from scipy import interpolate, sparse
 from scipy.special import legendre
 from sklearn.linear_model import Lasso
+from specderiv import cheb_deriv, fourier_deriv
 
 
 @register("spectral")
 class Spectral(Derivative):
-    def __init__(self, **kwargs):
+    def __init__(self, order=1, axis=0, basis='fourier', filter=None):
         """
-        Compute the numerical derivative by first computing the FFT. In Fourier space, derivatives are multiplication
-        by i*phase; compute the IFFT after.
-
-        Args:
-            **kwargs: Optional keyword arguments.
-
+        Compute the numerical derivative by spectral methods. In Fourier space, derivatives are multiplication
+        by i*phase; compute the inverse transform after. Use either Fourier modes of Chebyshev polynomials as
+        the basis.
+        
         Keyword Args:
-            filter: Optional. A function that takes in frequencies and outputs weights to scale the coefficient at
-                the input frequency in Fourier space. Input frequencies are the discrete fourier transform sample
-                frequencies associated with the domain variable. Look into python signal processing resources in
-                scipy.signal for common filters.
-
+            order (int): order of the derivative, defaults to 1st order
+            axis (int): the dimension of the data along which to differentiate, defaults to first dimension
+            basis (str): 'fourier' or 'chebyshev', the set of basis functions to use for differentiation
+                Note `basis='fourier'` assumes your function is periodic and sampled over a period of its domain,
+                [a, b), and `basis='chebyshev'` assumes your function is sampled at cosine-spaced points on the
+                domain [a, b].
+            filter: Optional. A function that takes in basis function indices and outputs weights, which scale
+                the corresponding modes in the basis-space interpolation before derivatives are taken, e.g.
+                `lambda k: k<10` will keep only the first ten modes. With the Fourier basis, k corresponds to
+                wavenumbers, so common filters from scipy.signal can be used. In the Chebyshev basis, modes do
+                not directly correspond to frequencies, so high frequency noise can not be separated quite as
+                cleanly, however it still may be helpful to dampen higher modes.
         """
-        # Filter function. Default: Identity filter
-        self.filter = kwargs.get('filter', np.vectorize(lambda f: 1))
-        self._x_hat = None
-        self._freq = None
-
-    def _dglobal(self, t, x):
-        self._x_hat = np.fft.fft(x)
-        self._freq = np.fft.fftfreq(t.size, d=(t[1] - t[0]))
+        self.order = order
+        self.axis = axis
+        if basis not in ['chebyshev', 'fourier']:
+            raise ValueError("Only chebyshev and fourier bases are allowed.")
+        self.basis = basis
+        self.filter = filter
+
+    @_memoize_arrays(1) # the memoization is 1 deep, as in only remembers the most recent args
+    def _global(self, t, x):
+        if self.basis == 'chebyshev':
+            return cheb_deriv(x, t, self.order, self.axis, self.filter)
+        else: # self.basis == 'fourier'
+            return fourier_deriv(x, t, self.order, self.axis, self.filter)
 
     def compute(self, t, x, i):
         return next(self.compute_for(t, x, [i]))
 
     def compute_for(self, t, x, indices):
-        self._dglobal(t, x)
-        res = np.fft.ifft(1j * 2 * np.pi * self._freq * self.filter(self._freq) * self._x_hat).real
+        res = self._global(t, x) # cached
         for i in indices:
             yield res[i]
 
@@ -212,7 +222,6 @@ def __init__(self, alpha=None):
         """
         self.alpha = alpha
 
-
     @_memoize_arrays(1)
     def _global(self, t, z, alpha):
         if alpha is None:

derivative/differentiation.py

@@ -21,7 +21,7 @@ def _gen_method(x, t, kind, axis, **kwargs):
     return methods.get(kind)(**kwargs)
 
 
-def dxdt(x, t, kind=None, axis=1, **kwargs):
+def dxdt(x, t, kind=None, axis=0, **kwargs):
     """
     Compute the derivative of x with respect to t along axis using the numerical derivative specified by "kind". This is
     the functional interface of the Derivative class.
@@ -35,7 +35,7 @@ def dxdt(x, t, kind=None, axis=1, **kwargs):
         x (:obj:`ndarray` of float): Ordered measurement values.
         t (:obj:`ndarray` of float): Ordered measurement times.
         kind (string): Derivative method name (see available kinds).
-        axis ({0,1}): Axis of x along which to differentiate. Default 1.
+        axis ({0,1}): Axis of x along which to differentiate. Default 0.
         **kwargs: Keyword arguments for the derivative method "kind".
 
     Available kinds
@@ -56,7 +56,7 @@ def dxdt(x, t, kind=None, axis=1, **kwargs):
         return method.d(x, t, axis=axis)
 
 
-def smooth_x(x, t, kind=None, axis=1, **kwargs):
+def smooth_x(x, t, kind=None, axis=0, **kwargs):
     """
     Compute the smoothed version of x given t along axis using the numerical
     derivative specified by "kind". This is the functional interface of
@@ -71,7 +71,7 @@ def smooth_x(x, t, kind=None, axis=1, **kwargs):
         x (:obj:`ndarray` of float): Ordered measurement values.
         t (:obj:`ndarray` of float): Ordered measurement times.
         kind (string): Derivative method name (see available kinds).
-        axis ({0,1}): Axis of x along which to differentiate. Default 1.
+        axis ({0,1}): Axis of x along which to differentiate. Default 0.
         **kwargs: Keyword arguments for the derivative method "kind".
 
     Available kinds
@@ -100,7 +100,7 @@ def compute(self, t, x, i):
         """
         Compute the derivative of one-dimensional data x with respect to t at the index i of x, (dx/dt)[i].
 
-        Computation of a derivative should fail explicitely if the implementation is unable to compute a derivative at
+        Computation of a derivative should fail explicitly if the implementation is unable to compute a derivative at
         the desired index. Used for global differentiation methods, for example.
 
         This requires that x and t have equal lengths >= 2, and that the index i is a valid index.
@@ -174,7 +174,7 @@ def compute_x_for(self, t, x, indices):
         for i in indices:
             yield self.compute_x(t, x, i)
 
-    def d(self, X, t, axis=1):
+    def d(self, X, t, axis=0):
         """
         Compute the derivative of measurements X taken at times t.
 
@@ -184,7 +184,7 @@ def d(self, X, t, axis=1):
         Args:
             X  (:obj:`ndarray` of float): Ordered measurements values. Multiple measurements allowed.
             t (:obj:`ndarray` of float): Ordered measurement times.
-            axis ({0,1}). axis of X along which to differentiate. default 1.
+            axis ({0,1}). axis of X along which to differentiate. default 0.
 
         Returns:
             :obj:`ndarray` of float: Returns dX/dt along axis.
@@ -202,7 +202,7 @@ def d(self, X, t, axis=1):
         return _restore_axes(dX, axis, flat)
 
 
-    def x(self, X, t, axis=1):
+    def x(self, X, t, axis=0):
         """
         Compute the smoothed X values from measurements X taken at times t.
 
@@ -212,7 +212,7 @@ def x(self, X, t, axis=1):
         Args:
             X  (:obj:`ndarray` of float): Ordered measurements values. Multiple measurements allowed.
             t (:obj:`ndarray` of float): Ordered measurement times.
-            axis ({0,1}). axis of X along which to smooth. default 1.
+            axis ({0,1}). axis of X along which to smooth. default 0.
 
         Returns:
             :obj:`ndarray` of float: Returns dX/dt along axis.
@@ -228,6 +228,8 @@ def x(self, X, t, axis=1):
 
 
 def _align_axes(X, t, axis) -> tuple[NDArray, tuple[int, ...]]:
+    """Reshapes the data so the derivative always happens along axis 1.
+    """
     X = np.array(X)
     orig_shape = X.shape
     # By convention, differentiate axis 1
@@ -244,6 +246,8 @@ def _align_axes(X, t, axis) -> tuple[NDArray, tuple[int, ...]]:
 
 
 def _restore_axes(dX: NDArray, axis: int, orig_shape: tuple[int, ...]) -> NDArray:
+    """Undo the operation of _align_axes, so data can be returned in its original shape
+    """
     if len(orig_shape) == 1:
         return dX.flatten()
     else:

docs/notebooks/Examples.ipynb

@@ -1,20 +1,5 @@
 {
  "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2020-05-25T22:56:27.818221Z",
-     "start_time": "2020-05-25T22:56:27.413152Z"
-    },
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "%matplotlib inline"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -65,17 +50,16 @@
    "outputs": [],
    "source": [
     "def plot_example(diff_method, t, data_f, res_f, sigmas, y_label=None):\n",
-    "    '''\n",
-    "    Utility function for concise plotting of examples.\n",
-    "    '''\n",
+    "    '''Utility function for concise plotting of examples.'''\n",
     "    fig, axes = plt.subplots(1, len(sigmas), figsize=[len(sigmas)*4, 3])\n",
     "    \n",
     "    # Compute the derivative\n",
-    "    res = diff_method.d(np.vstack([data_f(t, s) for s in sigmas]), t)\n",
+    "    res = diff_method.d(np.vstack([data_f(t, s) for s in sigmas]), t, axis=1)\n",
     "    for i, s in enumerate(sigmas):\n",
-    "        axes[i].plot(t, res[i])\n",
     "        axes[i].plot(t, res_f(t))\n",
-    "        axes[i].set_title(\"Noise: $\\sigma$={}\".format(s))\n",
+    "        axes[i].plot(t, res[i])\n",
+    "        axes[i].set_title(r\"Noise: $\\sigma$={}\".format(s))\n",
+    "        axes[i].set_ylim([-1.25,1.3])\n",
     "    if y_label:\n",
     "        axes[0].set_ylabel(y_label, fontsize=12)"
    ]
@@ -133,7 +117,7 @@
     "\n",
     "fig,ax = plt.subplots(1, figsize=[5,3])\n",
     "kind = FiniteDifference(k=1)\n",
-    "ax.plot(t, kind.d(x,t))"
+    "ax.plot(t, kind.d(x,t));"
    ]
   },
   {
@@ -159,7 +143,7 @@
     "from derivative import dxdt\n",
     "\n",
     "fig,ax = plt.subplots(1, figsize=[5,3])\n",
-    "ax.plot(t, dxdt(x, t, \"finite_difference\", k=1))"
+    "ax.plot(t, dxdt(x, t, \"finite_difference\", k=1));"
    ]
   },
   {
@@ -203,10 +187,10 @@
     "sigmas = [0, 0.01, 0.1]\n",
     "fig, ax = plt.subplots(1, len(sigmas), figsize=[len(sigmas)*4, 3])\n",
     "\n",
-    "t = np.linspace(0, 2*np.pi, 50)\n",
+    "t = np.linspace(0, 2*np.pi, 50, endpoint=False)\n",
     "for axs, s in zip(ax, sigmas): \n",
     "    axs.scatter(t, noisy_sin(t, s))\n",
-    "    axs.set_title(\"Noise: $\\sigma$={}\".format(s))"
+    "    axs.set_title(r\"Noise: $\\sigma$={}\".format(s))"
    ]
   },
   {
@@ -295,7 +279,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Spectral method\n",
+    "### Spectral method - Fourier basis\n",
     "Add your own filter!"
    ]
   },
@@ -312,12 +296,35 @@
    "outputs": [],
    "source": [
     "no_filter =  derivative.Spectral()\n",
-    "yes_filter = derivative.Spectral(filter=np.vectorize(lambda f: 1 if abs(f) < 0.5 else 0))\n",
+    "yes_filter = derivative.Spectral(filter=np.vectorize(lambda k: 1 if abs(k) < 3 else 0))\n",
     "\n",
     "plot_example(no_filter, t, noisy_sin, np.cos, sigmas, 'No filter')\n",
     "plot_example(yes_filter, t, noisy_sin, np.cos, sigmas, 'Low-pass filter')"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Spectral method - Chebyshev basis\n",
+    "\n",
+    "Now let's do with the Chebyshev basis, which requires cosine-spaced points on [a, b] rather than equispaced points on [a, b)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "t_cos = np.cos(np.pi * np.arange(50) / 49) * np.pi + np.pi # choose a = 0, b = 2*pi\n",
+    "no_filter = derivative.Spectral(basis='chebyshev')\n",
+    "yes_filter = derivative.Spectral(basis='chebyshev', filter=np.vectorize(lambda k: 1 if abs(k) < 6 else 0))\n",
+    "\n",
+    "plot_example(no_filter, t_cos, noisy_sin, np.cos, sigmas, 'No filter')\n",
+    "plot_example(yes_filter, t_cos, noisy_sin, np.cos, sigmas, 'Low-pass filter')"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -394,7 +401,7 @@
    "outputs": [],
    "source": [
     "def noisy_abs(t, sigma):\n",
-    "    '''Sine with gaussian noise.'''\n",
+    "    '''Abs with gaussian noise.'''\n",
     "    np.random.seed(17)\n",
     "    return np.abs(t) + np.random.normal(loc=0, scale=sigma, size=x.shape)\n",
     "\n",
@@ -406,7 +413,7 @@
     "t = np.linspace(-1, 1, 50)\n",
     "for axs, s in zip(ax, sigmas): \n",
     "    axs.scatter(t, noisy_abs(t, s))\n",
-    "    axs.set_title(\"Noise: $\\sigma$={}\".format(s))"
+    "    axs.set_title(r\"Noise: $\\sigma$={}\".format(s))"
    ]
   },
   {
@@ -482,7 +489,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Spectral Method"
+    "### Spectral method - Fourier basis"
    ]
   },
   {
@@ -497,12 +504,33 @@
    "outputs": [],
    "source": [
     "no_filter =  derivative.Spectral()\n",
-    "yes_filter = derivative.Spectral(filter=np.vectorize(lambda f: 1 if abs(f) < 1 else 0))\n",
+    "yes_filter = derivative.Spectral(filter=np.vectorize(lambda k: 1 if abs(k) < 6 else 0))\n",
     "\n",
     "plot_example(no_filter, t, noisy_abs, d_abs, sigmas, 'No filter')\n",
     "plot_example(yes_filter, t, noisy_abs, d_abs, sigmas, 'Low-pass filter')"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Spectral method - Chebyshev basis"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "t_cos = np.cos(np.pi * np.arange(50)/49)\n",
+    "no_filter = derivative.Spectral(basis='chebyshev')\n",
+    "yes_filter = derivative.Spectral(basis='chebyshev', filter=np.vectorize(lambda k: 1 if abs(k) < 15 else 0))\n",
+    "\n",
+    "plot_example(no_filter, t_cos, noisy_abs, d_abs, sigmas, 'No filter')\n",
+    "plot_example(yes_filter, t_cos, noisy_abs, d_abs, sigmas, 'Low-pass filter')"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -549,18 +577,11 @@
     "kal =  derivative.Kalman(alpha=1.)\n",
     "plot_example(kal, t, noisy_abs, d_abs, sigmas, 'alpha: 1.')"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -574,7 +595,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.10"
+   "version": "3.13.1"
   },
   "latex_envs": {
    "LaTeX_envs_menu_present": true,