diff --git a/post_processing/spectral_utils_tutorial.ipynb b/post_processing/spectral_utils_tutorial.ipynb
new file mode 100644
index 00000000..2fbb94dd
--- /dev/null
+++ b/post_processing/spectral_utils_tutorial.ipynb
@@ -0,0 +1,1206 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Tutorial: spectral_utils.py"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "This notebook demonstrates how to use the methods in the spectral_utils.py utility. There are four classes (SHT, Fourier, Legendre, and Chebyshev) that are used to perform their respective transformations between physical and spectral space, and to perform angular and radial derivatives.\n",
+ "\n",
+ "# Contents\n",
+ "1. [Generate sample data](#sample)\n",
+ "2. [Spherical harmonic transforms](#SHT)\n",
+ " 1. [Convert Shell Spectra to Shell Slices](#case-1)\n",
+ " 2. [Convert Shell Slices to Shell Spectra](#case-2)\n",
+ " 3. [Legendre transform on Shell Spectra](#case-3)\n",
+ " 4. [Fourier transform on Shell Slices](#case-5)\n",
+ " 5. [Fourier transform on Shell Spectra](#case-7)\n",
+ "3. [Angular derivatives](#derivs)\n",
+ " 1. [$\\sin \\theta \\frac{\\partial}{\\partial\\theta}$ in spectral space](#case-9)\n",
+ " 2. [$\\frac{\\partial}{\\partial \\phi}$ in spectral space using the SHT class](#case-10)\n",
+ " 3. [$\\frac{\\partial}{\\partial \\phi}$ in spectral space using the Fourier class](#case-11)\n",
+ " 4. [$\\frac{\\partial}{\\partial \\theta}$ in physical space](#case-12)\n",
+ " 5. [$\\frac{\\partial}{\\partial \\phi}$ in physical space](#case-13)\n",
+ "4. [Radial transforms and derivatives](#radial)\n",
+ " 1. [Chebyshev transform](#cheby)\n",
+ " 2. [Radial derivatives](#ddr)\n",
+ "\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# import the relevant utility\n",
+ "import spectral_utils\n",
+ "\n",
+ "# import Shell Spectra/Slices classes from Rayleigh\n",
+ "from rayleigh_diagnostics import Shell_Spectra, Shell_Slices\n",
+ "\n",
+ "# import other helpful things\n",
+ "import numpy as np\n",
+ "import matplotlib.pyplot as plt\n",
+ "from scipy.special import sph_harm\n",
+ "from scipy.special import chebyt\n",
+ "\n",
+ "# plotting nonsense\n",
+ "from matplotlib import gridspec\n",
+ "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
+ "from matplotlib import cm,colors\n",
+ "plt.rcParams['image.cmap'] = 'seismic'\n",
+ "plt.rcParams['image.origin'] = 'lower'\n",
+ "plt.rcParams['image.interpolation'] = 'none'"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "# 1. Generate Sample Shell Slice and Sample Shell Spectrum\n",
+ "\n",
+ "If fakedata=True, a sample shell slice and shell spectrum are generated using a combination of spherical harmonics. If you would prefer to use your own Rayleigh outputs, set fakedata=False and set the \"sample_slice_name\", \"sample_slice_path, \"sample_spectrum_name\", and \"sample_spectrum_path\" to the relevant values."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "fakedata = True\n",
+ "sample_slice_name = 'sample_slice'\n",
+ "sample_spectrum_name = 'sample_spectrum'\n",
+ "sample_slice_path = ''\n",
+ "sample_spectrum_path = ''\n",
+ "\n",
+ "# generate fake spherical harmonic data \n",
+ "if fakedata:\n",
+ " \n",
+ " #parameters\n",
+ " ntheta = 384\n",
+ " dealias = 1.5\n",
+ " lmax = int(ntheta/dealias - 1)\n",
+ " nphi = 768\n",
+ " \n",
+ " # Create 2D grid of angular variables\n",
+ " theta = np.linspace(0, np.pi, ntheta)\n",
+ " phi = np.linspace(0, 2*np.pi, nphi)\n",
+ " phi, theta = np.meshgrid(theta, phi)\n",
+ "\n",
+ " # generate spherical harmonics with l,m = (2,2), (5, 3), and (8, 5) \n",
+ " Ylm_1 = sph_harm(2, 2, theta, phi)\n",
+ " Ylm_2 = sph_harm(3, 5, theta, phi)\n",
+ " Ylm_3 = sph_harm(5, 8, theta, phi)\n",
+ " sample_slice = Ylm_1 + Ylm_2 + Ylm_3\n",
+ " sample_slice = np.real(np.transpose(sample_slice))\n",
+ " \n",
+ " # take the spherical harmonic transform for future examples\n",
+ " transform_SHT = spectral_utils.SHT(ntheta, spectral=False)\n",
+ " sample_spectrum = np.zeros((lmax+1, lmax+1), dtype=complex)\n",
+ " sample_spectrum = transform_SHT.to_spectral(sample_slice, th_l_axis=0, phi_m_axis=1)\n",
+ " xlim=10\n",
+ " norm=None\n",
+ " \n",
+ "\n",
+ "else:\n",
+ " # read in sample shell slice \n",
+ " sslice = Shell_Slices(sample_slice_name, path=sample_slice_path)\n",
+ " sample_slice = sslice.vals[:,:,-1,sslice.lut[301],0] # radial vorticity spectrum (301) at the bottom of the radiative interior\n",
+ " sample_slice = np.transpose(sample_slice)\n",
+ " \n",
+ " # read in sample shell spectrum (corresponding to above shell slice)\n",
+ " ss = Shell_Spectra(sample_spectrum_name, path=sample_spectrum_path)\n",
+ " lmax = ss.lmax\n",
+ " sample_spectrum = np.zeros((lmax+1, lmax+1), dtype=complex)\n",
+ " sample_spectrum = ss.vals[:,:,-1,ss.lut[301],0] \n",
+ " \n",
+ " # parameters\n",
+ " ntheta = sslice.ntheta\n",
+ " dealias = ntheta/(lmax + 1)\n",
+ " nphi = sslice.nphi\n",
+ " \n",
+ " xlim=50\n",
+ " norm=colors.LogNorm(vmin=1e-30, vmax=1e-15)\n",
+ "\n",
+ "# \n",
+ "# plot the sample slice and sample spectrum\n",
+ "fig = plt.figure(figsize=(14, 8))\n",
+ "gs = gridspec.GridSpec(1, 2, width_ratios=[2, 1]) \n",
+ "\n",
+ "ax0 = plt.subplot(gs[0])\n",
+ "ax0.set_title('Sample Shell Slice')\n",
+ "pos0=ax0.imshow(sample_slice, extent=[0, 360, -90, 90])\n",
+ "ax0.set_xlabel('Longitude')\n",
+ "ax0.set_ylabel('Latitude')\n",
+ "divider = make_axes_locatable(ax0)\n",
+ "cax = divider.append_axes('right', size='2%', pad=0.05)\n",
+ "fig.colorbar(pos0, cax=cax)\n",
+ "\n",
+ "ax1 = plt.subplot(gs[1])\n",
+ "ax1.set_title('Sample Shell Spectra')\n",
+ "power = sample_spectrum.real**2 + sample_spectrum.imag**2\n",
+ "pos1=ax1.imshow(np.transpose(power), norm=norm)\n",
+ "ax1.set_xlabel('l')\n",
+ "ax1.set_ylabel('m')\n",
+ "ax1.set_xlim(0, xlim)\n",
+ "ax1.set_ylim(0, xlim)\n",
+ "divider = make_axes_locatable(ax1)\n",
+ "cax = divider.append_axes('right', size='2%', pad=0.05)\n",
+ "fig.colorbar(pos1, cax=cax)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "# 2. Spherical harmonic transformations"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "The SHT class performs full spherical harmonic transformations with both the Legendre and Fourier transforms, while the Fourier and class performs only the Fourier transform in $m$. Each class features to_physical and to_spectral functions. Below is a table denoting the currently supported transformations and the best class and function to perform each:\n",
+ "\n",
+ "| link | input | output | transformations | class | function | use case\n",
+ "| --- | --- | --- | --- | --- | --- | ---\n",
+ "| [Case 1](#case-1)| $l$, $m$ | $\\theta$, $\\phi$ | iLT + iFFT | SHT | to_physical(data, th_l_axis=0, phi_m_axis=1) | convert Shell Spectra to Shell Slices\n",
+ "| [Case 2](#case-2) | $\\theta$, $\\phi$ | $l$, $m$ | FFT + LT | SHT | to_spectral(data, th_l_axis=0, phi_m_axis=1) | convert Shell Slices to Shell Spectra\n",
+ "| [Case 3](#case-3) | $l$, $m$ | $\\theta$, $m$ | iLT | SHT | _LT_to_physical(data, axis=0, m_axis=1) | perform inverse Legendre transform of Shell Spectra \n",
+ "| [Case 4](#case-3) | $\\theta$, $m$ | $l$, $m$ | LT | SHT | _LT_to_spectral(data, axis=0, m_axis=1) | reverse of Case 3\n",
+ "| [Case 5](#case-5) | $\\theta$, $\\phi$ | $\\theta$, $m$ | FFT | Fourier | to_spectral(data, axis=0, window=None) | perform Fourier transform of Shell Slice\n",
+ "| [Case 6](#case-5) | $\\theta$, $m$ | $\\theta$, $\\phi$ | iFFT | Fourier | to_physical(data, axis=0) | reverse of Case 5\n",
+ "| [Case 7](#case-7) | $l$, $m$ | $l$, $\\phi$ | iFFT | Fourier | to_physical(data, axis=0) | perform inverse Fourier transform of Shell Spectra\n",
+ "| [Case 8](#case-7) | $l$, $\\phi$ | $l$, $m$ | FFT | Fourier | to_spectral(data, axis=0, window=None) | reverse of Case 7\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## A. Case 1: Convert Shell Spectra to Shell Slices, $(l, m) \\rightarrow (\\theta, \\phi)$\n",
+ "### Instantiate class"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# instantiate SHT class\n",
+ "SHT = spectral_utils.SHT(ntheta, spectral=False, dealias=dealias)\n",
+ "\n",
+ "# alternate instantiation\n",
+ "# SHT = spectral_utils.SHT(lmax, spectral=True, dealias=dealias)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Perform transformation"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# convert from spectral space to physical space\n",
+ "sample_spectra_to_slice = SHT.to_physical(sample_spectrum, th_l_axis=0, phi_m_axis=1)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Plot results"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# plot and compare the transformed spectra to the matching slice\n",
+ "fig, (ax0, ax1) = plt.subplots(1, 2, figsize=(16, 8))\n",
+ "\n",
+ "pos0 = ax0.imshow(sample_spectra_to_slice, extent=[0, 360, -90, 90])\n",
+ "ax0.set_title('Converted Data')\n",
+ "ax0.set_xlabel('Longitude')\n",
+ "ax0.set_ylabel('Latitude')\n",
+ "\n",
+ "divider = make_axes_locatable(ax0)\n",
+ "cax = divider.append_axes('right', size='2%', pad=0.05)\n",
+ "fig.colorbar(pos0, cax=cax)\n",
+ "\n",
+ "\n",
+ "ax1.set_title('Original Slice')\n",
+ "pos1 = ax1.imshow(sample_slice.real, extent=[0, 360, -90, 90])\n",
+ "ax1.set_xlabel('Longitude')\n",
+ "ax1.set_ylabel('Latitude')\n",
+ "\n",
+ "divider = make_axes_locatable(ax1)\n",
+ "cax = divider.append_axes('right', size='2%', pad=0.05)\n",
+ "fig.colorbar(pos1, cax=cax)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## B. Case 2: Convert Shell Slices to Shell Spectra, $(\\theta, \\phi) \\rightarrow (l, m)$\n",
+ "### Instantiate class"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# instantiate SHT class\n",
+ "SHT = spectral_utils.SHT(ntheta, spectral=False, dealias=dealias)\n",
+ "\n",
+ "# alternate instantiation\n",
+ "# SHT = spectral_utils.SHT(lmax, spectral=True, dealias=dealias)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Perform transformation"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# this transformation will produce a complex output, so initialize a complex array\n",
+ "sample_slice_to_spectra = np.zeros((lmax+1, lmax+1), dtype=complex)\n",
+ "\n",
+ "# transform from physical space to spectral space\n",
+ "sample_slice_to_spectra = SHT.to_spectral(sample_slice, th_l_axis=0, phi_m_axis=1)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Plot results"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# calculate power\n",
+ "power = sample_slice_to_spectra.real**2 + sample_slice_to_spectra.imag**2\n",
+ "\n",
+ "# plot the spectra to see how they look\n",
+ "fig, (ax0, ax1) = plt.subplots(1, 2, figsize=(13, 5))\n",
+ "\n",
+ "ax0.set_title('Converted Data')\n",
+ "pos0=ax0.imshow(np.transpose(power), norm=norm)\n",
+ "ax0.set_xlabel('l')\n",
+ "ax0.set_ylabel('m')\n",
+ "ax0.set_xlim(0, xlim)\n",
+ "ax0.set_ylim(0, xlim)\n",
+ "divider = make_axes_locatable(ax0)\n",
+ "cax = divider.append_axes('right', size='2%', pad=0.05)\n",
+ "fig.colorbar(pos0, cax=cax)\n",
+ "\n",
+ "power = sample_spectrum.real**2 + sample_spectrum.imag**2\n",
+ "\n",
+ "ax1.set_title('Original Spectrum')\n",
+ "pos1=ax1.imshow(np.transpose(power), norm=norm)\n",
+ "ax1.set_xlabel('l')\n",
+ "ax1.set_ylabel('m')\n",
+ "ax1.set_xlim(0, xlim)\n",
+ "ax1.set_ylim(0, xlim)\n",
+ "divider = make_axes_locatable(ax1)\n",
+ "cax = divider.append_axes('right', size='2%', pad=0.05)\n",
+ "fig.colorbar(pos1, cax=cax)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## C. Cases 3 and 4: Perform Legendre transform on Shell Spectra $(l, m) \\leftrightarrow (\\theta, m)$\n",
+ "### Instantiate class"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# instantiate SHT class\n",
+ "SHT = spectral_utils.SHT(ntheta, spectral=False, dealias=dealias)\n",
+ "\n",
+ "# alternate initialization\n",
+ "# SHT = spectral_utils.SHT(lmax, spectral=True, dealias=dealias)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Perform inverse Legendre transformation $(l, m) \\rightarrow (\\theta, m)$"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "sample_spectra_iLT = np.zeros((ntheta, lmax+1), dtype=complex)\n",
+ "sample_spectra_iLT = SHT._LT_to_physical(sample_spectrum, axis=0, m_axis=1)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Plot results"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "fig, ax = plt.subplots(figsize=(8, 5))\n",
+ "\n",
+ "pos = ax.imshow(sample_spectra_iLT.real,extent=[-0.5, np.shape(sample_spectra_iLT)[1] - 0.5, -90, 90], aspect='auto',\n",
+ "# vmin=-0.5e-9, vmax=0.5e-9\n",
+ " )\n",
+ "ax.set_xlim(-0.5, 10)\n",
+ "ax.set_xlabel('m')\n",
+ "ax.set_ylabel('Latitude')\n",
+ "fig.colorbar(pos, ax=ax)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Perform Legendre transformation $(\\theta, m) \\rightarrow (l, m)$"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "sample_spectra_LT = SHT._LT_to_spectral(sample_spectra_iLT*nphi, axis=0, m_axis=1)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Plot results"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "power = sample_spectra_LT.real**2 + sample_spectra_LT.imag**2\n",
+ "\n",
+ "fig, ax = plt.subplots(figsize=(8, 5))\n",
+ "\n",
+ "pos = ax.imshow(np.transpose(power), norm=norm)\n",
+ "ax.set_xlim(0, xlim)\n",
+ "ax.set_ylim(0, xlim)\n",
+ "ax.set_xlabel('l')\n",
+ "ax.set_ylabel('m')\n",
+ "divider = make_axes_locatable(ax)\n",
+ "cax = divider.append_axes('right', size='2%', pad=0.05)\n",
+ "fig.colorbar(pos, cax=cax)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## D. Cases 5 and 6: Perform Fourier transform on Shell Slice $(\\theta, \\phi) \\leftrightarrow (\\theta, m)$\n",
+ "### Instantiate class"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "fourier = spectral_utils.Fourier(nphi)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Perform Fourier transformation $(\\theta, \\phi) \\rightarrow (\\theta, m)$"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "sample_slice_FT = fourier.to_spectral(sample_slice, axis=1)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Plot results"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "fig, ax = plt.subplots(figsize=(8, 5))\n",
+ "\n",
+ "pos = ax.imshow(sample_slice_FT.real, aspect='auto', extent=[-0.5, np.shape(sample_slice_FT)[0]-0.5, -90, 90],\n",
+ "# vmin=-1e-9, vmax=1e-9\n",
+ " )\n",
+ "ax.set_xlim(-0.5, 10)\n",
+ "ax.set_xlabel('m')\n",
+ "ax.set_ylabel('Latitude')\n",
+ "fig.colorbar(pos, ax=ax)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Perform inverse Fourier transformation $(\\theta, m) \\rightarrow (\\theta, \\phi)$"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "sample_slice_iFFT = fourier.to_physical(sample_slice_FT, axis=1)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Plot results"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# plot the power to see if it makes sense\n",
+ "fig, ax = plt.subplots(figsize=(8, 5))\n",
+ "pos = ax.imshow(sample_slice_iFFT.real, extent=[0, 360, -90, 90])\n",
+ "ax.set_xlabel('Longitude')\n",
+ "ax.set_ylabel('Latitude')\n",
+ "divider = make_axes_locatable(ax)\n",
+ "cax = divider.append_axes('right', size='2%', pad=0.05)\n",
+ "fig.colorbar(pos, cax=cax)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## E. Cases 7 and 8: Perform Fourier transform on Shell Spectra $(l, m) \\leftrightarrow (l, \\phi)$\n",
+ "### Instantiate class\n",
+ "Because we are starting from fully transformed spherical harmonic spectra, the number of angular values is set by lmax*2 + 1, rather than nphi."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "fourier = spectral_utils.Fourier(lmax*2 + 1)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Perform inverse Fourier transformation $(l, m) \\rightarrow (l, \\phi)$"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "sample_spectra_iFFT = fourier.to_physical(sample_spectrum, axis=1)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Plot results"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "fig, ax = plt.subplots(figsize=(8, 5))\n",
+ "\n",
+ "pos = ax.imshow(sample_spectra_iFFT, aspect='auto',extent=[-0.5, 360, -0.5, lmax-0.5],\n",
+ "# vmin=-1e-9, vmax=1e-9\n",
+ " )\n",
+ "ax.set_ylim(0.5, 15)\n",
+ "ax.set_xlabel('Longitude')\n",
+ "ax.set_ylabel('l')\n",
+ "divider = make_axes_locatable(ax)\n",
+ "cax = divider.append_axes('right', size='2%', pad=0.05)\n",
+ "fig.colorbar(pos, cax=cax)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Perform Fourier transformation $(l, \\phi) \\rightarrow (l, m)$"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# reinstantiate class with the correct number of grid points\n",
+ "fourier = spectral_utils.Fourier(lmax*2)\n",
+ "sample_spectra_FT = fourier.to_spectral(sample_spectra_iFFT, axis=1)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Plot results"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "power = sample_spectra_FT.real**2 + sample_spectra_FT.imag**2\n",
+ "\n",
+ "fig, ax = plt.subplots(figsize=(8, 5))\n",
+ "\n",
+ "pos = ax.imshow(np.transpose(power), norm=norm)\n",
+ "ax.set_xlim(0, xlim)\n",
+ "ax.set_ylim(0, xlim)\n",
+ "ax.set_xlabel('l')\n",
+ "ax.set_ylabel('m')\n",
+ "divider = make_axes_locatable(ax)\n",
+ "cax = divider.append_axes('right', size='2%', pad=0.05)\n",
+ "fig.colorbar(pos, cax=cax)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "# 3. Angular derivatives\n",
+ "\n",
+ "The SHT, Fourier, and Legendre classes can perform angular derivatives in both spectral and physical space. A summary of the available functions is provided in the table below:\n",
+ "\n",
+ "| derivative | domain | class | function\n",
+ "| --- | --- | --- | ---\n",
+ "|[$\\sin \\theta \\frac{\\partial}{\\partial \\theta}$](#case-9) | spectral | SHT | sin_d_dtheta(data, l_axis=0, m_axis=1)\n",
+ "|[$\\frac{\\partial}{\\partial \\phi}$](#case-10) | spectral | SHT | d_dphi(data, m_axis=0)\n",
+ "|[$\\frac{\\partial}{\\partial \\phi}$](#case-11) | spectral | Fourier | d_dphi(data, axis=0, physical=False)\n",
+ "|[$\\frac{\\partial}{\\partial \\theta}$](#case-12) | physical | Legendre | d_dtheta(data, axis=0, physical=True)\n",
+ "|[$\\frac{\\partial}{\\partial \\phi}$](#case-13) | physical | Fourier | d_dphi(data, axis=0, physical=True)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## A. $\\sin \\theta \\frac{\\partial}{\\partial\\theta}$ in spectral space\n",
+ "### Instantiate class"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "SHT = spectral_utils.SHT(ntheta, spectral=False, dealias=dealias)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Take derivative"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "l_derivative = SHT.sin_d_dtheta(sample_spectrum, l_axis=0, m_axis=1)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Plot results"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "power = l_derivative.real**2 + l_derivative.imag**2\n",
+ "\n",
+ "fig, ax = plt.subplots(figsize=(8, 5))\n",
+ "\n",
+ "pos=ax.imshow(np.transpose(power), norm=norm)\n",
+ "ax.set_xlim(0, xlim)\n",
+ "ax.set_ylim(0, xlim)\n",
+ "ax.set_xlabel('l')\n",
+ "ax.set_ylabel('m')\n",
+ "fig.colorbar(pos, ax=ax)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## B. $\\frac{\\partial}{\\partial \\phi}$ in spectral space using the SHT class"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "m_derivative = SHT.d_dphi(sample_spectrum, m_axis=1)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "power = m_derivative.real**2 + m_derivative.imag**2\n",
+ "\n",
+ "fig, ax = plt.subplots(figsize=(8, 5))\n",
+ "\n",
+ "pos=ax.imshow(np.transpose(power), norm=norm)\n",
+ "ax.set_xlim(0, xlim)\n",
+ "ax.set_ylim(0, xlim)\n",
+ "ax.set_xlabel('l')\n",
+ "ax.set_ylabel('m')\n",
+ "fig.colorbar(pos, ax=ax)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## C. $\\frac{\\partial}{\\partial \\phi}$ in spectral space using the Fourier class\n",
+ "### Instantiate class"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "Fourier = spectral_utils.Fourier(lmax*2 + 1)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Take derivative"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "m_derivative = Fourier.d_dphi(sample_spectrum, axis=1, physical=False)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Plot results"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "power = m_derivative.real**2 + m_derivative.imag**2\n",
+ "\n",
+ "fig, ax = plt.subplots(figsize=(8, 5))\n",
+ "\n",
+ "pos=ax.imshow(np.transpose(power), norm=norm)\n",
+ "ax.set_xlim(0, xlim)\n",
+ "ax.set_ylim(0, xlim)\n",
+ "ax.set_xlabel('l')\n",
+ "ax.set_ylabel('m')\n",
+ "fig.colorbar(pos, ax=ax)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## D. $\\frac{\\partial}{\\partial \\theta}$ in physical space\n",
+ "### Instantiate class"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "Legendre = spectral_utils.Legendre(ntheta, spectral=False, dealias=dealias)\n",
+ "\n",
+ "# alternate instantiation\n",
+ "# Legendre_class = spectral_utils.Legendre(lmax, spectral=True, dealias=dealias)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Take derivative"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "th_derivative = Legendre.d_dtheta(sample_slice, axis=0, physical=True)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Plot results"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "fig, ax = plt.subplots(figsize=(8, 5))\n",
+ "\n",
+ "pos=ax.imshow(th_derivative, extent=[0, 360, -90, 90])\n",
+ "ax.set_xlabel('Longitude')\n",
+ "ax.set_ylabel('Latitude')\n",
+ "divider = make_axes_locatable(ax)\n",
+ "cax = divider.append_axes('right', size='2%', pad=0.05)\n",
+ "fig.colorbar(pos, cax=cax)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## E. $\\frac{\\partial}{\\partial \\phi}$ in physical space\n",
+ "### Instantiate class"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "Fourier = spectral_utils.Fourier(nphi)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Take derivative"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "ph_derivative = Fourier.d_dphi(sample_slice, axis=1, physical=True)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Plot results"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "fig, ax = plt.subplots(figsize=(8, 5))\n",
+ "\n",
+ "pos=ax.imshow(ph_derivative, extent=[0, 360, -90, 90])\n",
+ "ax.set_xlabel('Longitude')\n",
+ "ax.set_ylabel('Latitude')\n",
+ "divider = make_axes_locatable(ax)\n",
+ "cax = divider.append_axes('right', size='2%', pad=0.05)\n",
+ "fig.colorbar(pos, cax=cax)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "# 4. Radial transformations and derivatives\n",
+ "The Chebyshev class provides methods to transform between physical and spectral space and perform radial derivatives in physical and spectral space. This example generates sample Chebyshev polynomials for simplicity, but this class can be used on any data output on the full radial grid.\n",
+ "\n",
+ "### Instantiate class\n",
+ "This class currently supports three different types of grids: a) single Chebyshev domain, b) uniform set of N Chebyshev domains, and c) N Chebyshev domains with different resolutions. This example uses a single Chebyshev domain with $n_r = 128$. Instantiation examples of other grid types are listed in the comment below. "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# set rmin and rmax for our fake data\n",
+ "rmin=5e10\n",
+ "rmax=6.83177e10\n",
+ "\n",
+ "# \n",
+ "cheb = spectral_utils.Chebyshev(128, rmin=rmin, rmax=rmax)\n",
+ "\n",
+ "# Single Chebyshev domain with 72 grid points using shell depth & aspect ratio:\n",
+ "# >>> cheb = Chebyshev(72, aspect_ratio=0.2, shell_depth=2)\n",
+ "# >>> cheb.rmin, cheb.rmax\n",
+ "# (0.5, 2.5)\n",
+ "\n",
+ "# Same grid as before, but specifying the minimum/maximum radius:\n",
+ "# >>> cheb = Chebyshev(72, rmin=0.5, rmax=2.5)\n",
+ "# >>> cheb.rmin, cheb.rmax\n",
+ "# (0.5, 2.5)\n",
+ "\n",
+ "# Same grid as before, but specifying the boundaries:\n",
+ "# >>> cheb = Chebyshev(72, boundaries=(0.5, 2.5))\n",
+ "# >>> cheb.rmin, cheb.rmax\n",
+ "# (0.5, 2.5)\n",
+ "# >>> cheb.nr_domains\n",
+ "# [72]\n",
+ "\n",
+ "# Three Chebyshev domains, each with 24 grid points, 72 grid points total:\n",
+ "# >>> cheb = Chebyshev(24, n_uniform_domains=3, aspect_ratio=0.2, shell_depth=2)\n",
+ "# >>> cheb.boundaries\n",
+ "# [2.5, 1.83333, 1.16666, 0.5]\n",
+ "# >>> cheb.nr_domains\n",
+ "# [24, 24, 24]\n",
+ "\n",
+ "# Three Chebyshev domains, nonuniform resolutions, 72 grid points total:\n",
+ "# >>> cheb = Chebyshev([16,36,20], boundaries=[0.5,1.0,2.0,2.4])\n",
+ "# >>> cheb.boundaries\n",
+ "# [2.5, 2.0, 1.0, 0.5]\n",
+ "# >>> cheb.nr_domains\n",
+ "# [20, 36, 16]"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Generate sample data"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# generate radial grid\n",
+ "radius = cheb.radius\n",
+ "\n",
+ "# generate colocation points in the [-1, 1] range\n",
+ "x = cheb.x\n",
+ "\n",
+ "# generate data from a handful of Chebyshev polynomials\n",
+ "rdata = chebyt(2)(x) - 4*chebyt(4)(x) + 3*chebyt(10)(x)\n",
+ "\n",
+ "plt.plot(radius, rdata)\n",
+ "plt.xlabel('Radius')\n",
+ "plt.title('Sample radial data')"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## A. Chebyshev transform\n",
+ "### Transform from physical to spectral space"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "to_spectral = cheb.to_spectral(rdata, axis=0)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Plot results"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# plot the decomposition\n",
+ "reconstructed = np.zeros(np.shape(to_spectral))\n",
+ "reconstructed = to_spectral[0]*chebyt(0)(x)/2 # the coefficient for n=0 is 1/N, as opposed to 2/N\n",
+ "\n",
+ "for i in range(1, len(to_spectral)):\n",
+ " reconstructed += to_spectral[i]*chebyt(i)(x)\n",
+ "\n",
+ "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 4.8))\n",
+ "\n",
+ "ax1.set_title('Spectral coefficients')\n",
+ "ax1.plot(to_spectral)\n",
+ "ax1.set_xlim(-0.5, 15)\n",
+ "ax1.set_xlabel('Order')\n",
+ "ax1.set_ylabel('Coefficient')\n",
+ "\n",
+ "ax2.set_title('Comparison')\n",
+ "ax2.plot(radius, rdata, label='Original')\n",
+ "ax2.scatter(radius, reconstructed, label='Reconstructed', color='orange')\n",
+ "ax2.legend()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Transform from spectral to physical space"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "to_physical = cheb.to_physical(to_spectral, axis=0)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "plt.plot(radius, rdata, label='original')\n",
+ "plt.scatter(radius, to_physical, label='transformed', color='orange')\n",
+ "plt.legend()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## B. Radial derivatives\n",
+ "### $\\frac{\\partial}{\\partial r}$ in physical space"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "ddr_physical = cheb.d_dr(rdata, physical=True)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "plt.plot(radius, ddr_physical, label='Chebyshev d_dr')\n",
+ "plt.scatter(radius, np.gradient(rdata[:,0], radius), label='numpy gradient', color='orange')\n",
+ "plt.legend()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### $\\frac{\\partial}{\\partial r}$ in spectral space"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "ddr_spectral = cheb.d_dr(to_spectral, physical=False)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "plt.plot(ddr_spectral)\n",
+ "plt.xlim(0, 15)\n",
+ "plt.xlabel('Order')\n",
+ "plt.ylabel('Coefficient')"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# transform back and plot against the physical derivative to compare\n",
+ "ddr_spectral_to_physical = cheb.to_physical(ddr_spectral, axis=0)\n",
+ "\n",
+ "plt.plot(radius, ddr_physical, label='physical space')\n",
+ "plt.scatter(radius, ddr_spectral_to_physical, color='orange', label='spectral space')\n",
+ "plt.legend()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.8.3"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}