diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6583da63..520c9d47 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,6 +19,8 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
- Implemented new diagnostic quantity codes for scalar fields. Added buoyancy forces, buoyancy work, and buoyant KE production associated with the active chi fields. These are documented in the "Under Development" section of the User Guide. Additional placeholders have been included for future codes to be added. \[Cian Wilson; 7-13-2026; [#582](https://github.com/geodynamics/Rayleigh/pull/582)\]
- Add support for building with NVHPC compilers as a preparation for using OpenMP on Nvidia systems. \[Philipp Edelmann; 7-17-2026; [#630](https://github.com/geodynamics/Rayleigh/pull/630)\]
+- Augmented spectral_utils_tutorial.ipynb to demonstrate how to use the new TemporalFFT class in rayleigh_diagnostics.py. \[Brad Hindman; 7-23-2026; [#636](https://github.com/geodynamics/Rayleigh/pull/636)\]
+
### Changed
- Containers for the TACC systems are no longer supported. \[Philipp Edelmann; 7-12-2026; [#608](https://github.com/geodynamics/Rayleigh/pull/608)\]
diff --git a/post_processing/spectral_utils_tutorial.ipynb b/post_processing/spectral_utils_tutorial.ipynb
index 2fbb94dd..07453f1b 100644
--- a/post_processing/spectral_utils_tutorial.ipynb
+++ b/post_processing/spectral_utils_tutorial.ipynb
@@ -11,7 +11,7 @@
"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",
+ "This notebook demonstrates how to use the methods in the spectral_utils.py utility. There are five classes (SHT, Fourier, Legendre, and Chebyshev, TemporalFFT) 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",
@@ -30,6 +30,12 @@
"4. [Radial transforms and derivatives](#radial)\n",
" 1. [Chebyshev transform](#cheby)\n",
" 2. [Radial derivatives](#ddr)\n",
+ "5. [Temporal transforms](#temporal)\n",
+ " 1. [Generate time series data on a regular temporal grid](#timeseries)\n",
+ " 2. [Temporal FFT of Shell Slices](#temporalFFTslices)\n",
+ " 3. [Temporal FFT of Shell Spectra](#temporalFFTspectra)\n",
+ " 4. [Using spectral filters to extract a single mode](#spectral_filtering)\n",
+ " 5. [Temporal FFT on a Nonuniform Grid](#nonuniform)\n",
"\n"
]
},
@@ -50,6 +56,7 @@
"import matplotlib.pyplot as plt\n",
"from scipy.special import sph_harm\n",
"from scipy.special import chebyt\n",
+ "from scipy.signal.windows import gaussian\n",
"\n",
"# plotting nonsense\n",
"from matplotlib import gridspec\n",
@@ -57,7 +64,7 @@
"from matplotlib import cm,colors\n",
"plt.rcParams['image.cmap'] = 'seismic'\n",
"plt.rcParams['image.origin'] = 'lower'\n",
- "plt.rcParams['image.interpolation'] = 'none'"
+ "plt.rcParams['image.interpolation'] = 'none'\n"
]
},
{
@@ -109,7 +116,7 @@
" 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",
"\n",
"else:\n",
" # read in sample shell slice \n",
@@ -1174,6 +1181,580 @@
"plt.legend()"
]
},
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "# 5. Temporal Transforms\n",
+ "\n",
+ "In this section we use the TemporalFFT class to generate temporal spectra for both Shell_Slice and Shell_Spectra data."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## A. Generate timeseries data on a regular temporal grid\n",
+ "\n",
+ "Sets of shell slices are generated by adding up a small set of sectoral equatorial Rossby waves. Each Rossby wave consists of a single spherical harmonic component (with $m = \\ell$) that oscillates at the frequency $\\omega = -\\frac{2m\\Omega}{\\ell(\\ell+1)}$, where $\\Omega$ is the star's uniform rotation rate. The waves propagate in the retrograde direction (westward). Times and frequencies are nondimensionalized in terms of rotations, i.e., $t = \\frac{2\\pi}{\\Omega} \\tilde{t}$ and $\\omega = \\frac{\\Omega}{2\\pi} \\tilde{\\omega}$."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Generate a set of artificial data\n",
+ "# The resulting data consists of a single radial shell slice as a function of time\n",
+ "# data[theta, phi, t]\n",
+ "\n",
+ "# generate fake spherical harmonic data (a set of sectoral equatorially trapped Rossby waves) \n",
+ "\n",
+ "#parameters\n",
+ "maxm = 20 # number of sectoral Rossby waves\n",
+ "dealias = 1.5 \n",
+ "ntheta = 384\n",
+ "lmax = int(ntheta/dealias - 1)\n",
+ "nell = 2*ntheta//3\n",
+ "nphi = 768\n",
+ "nm = nell\n",
+ "nt = 128\n",
+ "duration = 25.0 # Number of rotations\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 the temporal array\n",
+ "times = np.linspace(0, duration, nt)\n",
+ "\n",
+ "# instantiate the spherical harmonic transform\n",
+ "transform_SHT = spectral_utils.SHT(ntheta, spectral=False)\n",
+ "\n",
+ "# Generate the fake data\n",
+ "slices = np.ndarray((ntheta,nphi,nt), dtype=np.float64)\n",
+ "spectra = np.ndarray((nell,nm,nt), dtype=np.complex128)\n",
+ "ic = complex(0.0,1.0)\n",
+ "for m in range(1, maxm+1):\n",
+ " Ylm = sph_harm(m, m, theta, phi); freq = -4.0*np.pi/(m+1) # horizontal eigenfunction and frequency of Rossby mode\n",
+ " for i,t in enumerate(times):\n",
+ " mode = np.real(np.exp(ic*freq*t)*Ylm) # The sign of the argument of exp() is chosen such that retrograde phase\n",
+ " slices[:,:,i] += np.transpose(mode) # speeds correspond to negative frequencies\n",
+ "for i,t in enumerate(times): # Take the spherical harmonic transform of each slice\n",
+ " spectra[:,:,i] = transform_SHT.to_spectral(slices[:,:,i], th_l_axis=0, phi_m_axis=1)\n",
+ "\n",
+ "axis = -1 # time corresponds the final dimension in the array\n",
+ "xlim=30\n",
+ "norm=None\n",
+ "\n",
+ "\n",
+ "# plot a sample slice and sample spectrum\n",
+ "fig = plt.figure(figsize=(14, 8))\n",
+ "gs = gridspec.GridSpec(1, 2, width_ratios=[2, 1])\n",
+ "\n",
+ "sample_slice = slices[:,:,nt//4]\n",
+ "sample_spectrum = spectra[:,:,nt//4]\n",
+ "\n",
+ "ax0 = plt.subplot(gs[0])\n",
+ "ax0.set_title('Initial 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('Initial Spatial Spectra')\n",
+ "power = sample_spectrum.real**2 + sample_spectrum.imag**2\n",
+ "pos1=ax1.imshow(np.transpose(power), norm=norm)\n",
+ "ax1.set_xlabel(r'Harmonic Degree $\\ell$')\n",
+ "ax1.set_ylabel(r'Azimuthal Order $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)\n",
+ "\n",
+ "fig.tight_layout()\n",
+ "\n",
+ "\n",
+ "# plot a subsequent slice and spectrum\n",
+ "fig = plt.figure(figsize=(14, 8))\n",
+ "gs = gridspec.GridSpec(1, 2, width_ratios=[2, 1])\n",
+ "\n",
+ "sample_slice = slices[:,:,nt//4+5]\n",
+ "sample_spectrum = spectra[:,:,nt//4+5]\n",
+ "\n",
+ "ax0 = plt.subplot(gs[0])\n",
+ "ax0.set_title('Later 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('Later Spatial Spectra')\n",
+ "power = sample_spectrum.real**2 + sample_spectrum.imag**2\n",
+ "pos1=ax1.imshow(np.transpose(power), norm=norm)\n",
+ "ax1.set_xlabel(r'Harmonic Degree $\\ell$')\n",
+ "ax1.set_ylabel(r'Azimuthal Order $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)\n",
+ "\n",
+ "fig.tight_layout()\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## B. Temporal FFT of Shell Slices\n",
+ "\n",
+ "The temporal FFT is computed using a variety of window functions to illustrate how important windowing is to isolate peaks in the spectrum and to reduce the noise floor."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Compute the temporal transform using several different window functions\n",
+ "# Illustrate the spectra by plotting \n",
+ "\n",
+ "# instantiate the Temporal transform\n",
+ "TemporalFFT = spectral_utils.TemporalFFT\n",
+ "\n",
+ "# Ensure that the data is stored as an array\n",
+ "slices = np.asarray(slices)\n",
+ "\n",
+ "# Generate the Hann window for plotting purposes\n",
+ "hann_window = np.hanning(nt)\n",
+ "\n",
+ "#Take some transforms with different windows\n",
+ "dt = times[1]-times[0] # Temporal cadence (uniform)\n",
+ "idx_t = -1\n",
+ "# axis = -1 # The dimension along which to transform (-1 is the most efficient and the default)\n",
+ "# freq_shift = True # Place the zero frequency in the middle of the output array (True is the default)\n",
+ "# window = None # No window (Bad Dog! No cookie for you!): The default is the Hann window ('hann')\n",
+ "nowindow, freqs = TemporalFFT.uniform(slices, dt, window=None, axis=-1, freq_shift=True).fft() \n",
+ "nowinpower = np.abs(nowindow)**2 # compute the power spectrum from the complex Fourier spectrum\n",
+ "\n",
+ "# window = 'hann' # The default is the Hann window ('hann'): Hann window is the goto window function\n",
+ "hann, freqs = TemporalFFT.uniform(slices, dt).fft()\n",
+ "hannpower = np.abs(hann)**2\n",
+ "\n",
+ "# window = \"hamming\" # Hamming window reduces the nearest sidelobes\n",
+ "hamming, freqs = TemporalFFT.uniform(slices, dt, window='hamming').fft() # Hamming window (reduces the nearest sidelobes)\n",
+ "hammingpower = np.abs(hamming)**2\n",
+ "\n",
+ "fig, ax=plt.subplots(2,4, figsize=(12,5))\n",
+ "ilats = [ntheta//4, ntheta//2]\n",
+ "ilons = [0, nphi//2]\n",
+ "ilat = ntheta//2; ilon = 0\n",
+ "\n",
+ "for i in range(2):\n",
+ " ax[i,0].set_ylabel(r'Signal')\n",
+ " ax[1,i].set_xlabel(r'Time $(\\Omega/2\\pi) t$')\n",
+ " ax[i,2].set_ylabel(r'Power')\n",
+ " ax[1,i+2].set_xlabel(r'Frequency $2\\pi \\omega/\\Omega$')\n",
+ " for j in range(2):\n",
+ " ax[i,j].plot(times, slices[ilats[i],ilons[j],:], color='dodgerblue')\n",
+ " ax[i,j].plot(times, hann_window, color='orangered')\n",
+ " ax[i,j].plot(times, slices[ilats[i],ilons[j],:]*hann_window, color='green')\n",
+ "\n",
+ " ax[i,j+2].plot(freqs, nowinpower[ilats[i],ilons[j],:], color='black')\n",
+ " ax[i,j+2].plot(freqs, hannpower[ilats[i],ilons[j],:], color='orangered', zorder = 10)\n",
+ " ax[i,j+2].plot(freqs, hammingpower[ilats[i],ilons[j],:], color='blue', zorder = 5)\n",
+ "\n",
+ " ax[i,j+2].set_yscale('log')\n",
+ " \n",
+ "\n",
+ "\n",
+ "fig.tight_layout()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "The timeseries on the left are extracted at four different points in the domain ($\\theta= \\pi/4$ and $\\pi/4$, $\\phi=0$ and $\\pi$). The light blue curves show the raw time series, the orange curve is the Hann window, and the green is the windowed data. The Hann window forces the data to be periodic in time by forcing both the function and its derivative to vanish at the ends. The temporal power spectra on the right correspond to the time series on the left, but with three different window functions. The black curve is for no window function (honestly, never do this). The blue curve is for the Hamming window (look it up), which tries to minimize the height of the first sidelobe in the spectral response function (FFT of the window). Its a good window function if you have very closely spaced peaks in the frequency that you want to keep separated. The orange curves were generated with the Hann window which is a great multipurpose window. Note that the Hann window does an excellent job of keeping the power isolated in the peaks."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## C. Temporal FFT of Shell Spectra\n",
+ "\n",
+ "The temporal spectra is computed on data for which a spherical harmonic transform has already been performed."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Compute the temporal transforms of the spectra to reveal the waves\n",
+ "\n",
+ "# Ensure that the spatial spectra are an array\n",
+ "spectra = np.asarray(spectra)\n",
+ "axis = -1 # The temporal dimension is the last one (most efficient)\n",
+ "\n",
+ "# Compute the temporal spectra of the spherical harmonic transform\n",
+ "wspectra, freqs = TemporalFFT.uniform(spectra, dt).fft() # Use the Hann window\n",
+ "power = np.abs(wspectra)**2 # Compute the power\n",
+ "\n",
+ "#sum power over l\n",
+ "lpower = np.sum(power,axis=0) # Power as a function of azimuthal order m and frequency omega\n",
+ "\n",
+ "# Plot!\n",
+ "fig,ax=plt.subplots(1,2, figsize=(12,4))\n",
+ "aspect = 8*(64/nt)\n",
+ "nyquist = np.pi/dt # Nyquist\n",
+ "twopi = 2.0*np.pi # For converting from 2 pi/Omega to 1/Omega temporal units\n",
+ "maxomega = nyquist/twopi\n",
+ "\n",
+ "pos = ax[0].imshow(np.log(lpower[0:maxm+1,:]).T, cmap='plasma', extent = [-0.5,maxm+0.5,-maxomega,maxomega], aspect=aspect)\n",
+ "ax[0].axhline(0.0, color = 'gray')\n",
+ "\n",
+ "cols = ['blue', 'red', 'green']\n",
+ "for i,m in enumerate([1,5,15]):\n",
+ " ax[1].plot(freqs, lpower[m,:], color=cols[i])\n",
+ " ax[0].axvline(m, color=cols[i], linestyle=':')\n",
+ "ax[0].set_title('Log Power')\n",
+ "ax[0].set_xlabel(r'Azimuthal Order $m$')\n",
+ "ax[0].set_ylabel(r'Frequency $\\omega/\\Omega$')\n",
+ "ax[1].set_xlabel(r'Frequency $\\omega/\\Omega$')\n",
+ "ax[1].set_ylabel(r'Power')\n",
+ "ax[1].set_xlim(-1*maxomega,maxomega)\n",
+ "ax[1].set_ylim(1e-12,1e4)\n",
+ "ax[1].set_yscale('log')\n",
+ "fig.colorbar(pos)\n",
+ "\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "The left picture shows the dispersion relation for a lovely set of Rossby waves. Negative frequencies indicate retrograde modes. Modes with positive frequencies (if there were any) would be prograde. You can clearly see that the power doesn't stay put. Every place that there is a mode, we see a cross. This is because the power gets \"leaked\" across wavenumber and across frequency. The right image shows three line profiles of the power at different azimuthal orders $m$. The spikes in the wings are caused by power leaked from nearby $m$s, which is why the same spike appears in the each curve in the same location. In numerical convection simulations, there is always noise from the convection. Convective power is usually broadly spread in frequency and wavenumber. Hence, one would generally not easily see the crosses and other ``obvious\" symptoms of power leakage. (Its still there, but you can't easily identify where it came from). If you don't window, the leakage is large and broadly spread.\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## D. Using Spectral Filters to Extract a Single Mode\n",
+ "\n",
+ "We start fully in spectral space (the data is a function of $m$, $\\ell$, and $\\omega$). To isolate a single peak (i.e., a single mode, we build a narrow Gaussian filter that is located at the frequency of the $m=1$ Rossby wave. We ensure that the filter is broader than the spectral peak, but narrow enough to remove the leaked signal from other $m$s. By multiplying the spectrum by this filter and then inverting all of the transforms, we will obtain the wavefunction for a single mode in physical space ($\\phi$, $\\theta$, and $t$)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Build a Gaussian Filter for the m = 4 Rossby Wave\n",
+ "\n",
+ "mfilt = 1\n",
+ "\n",
+ "width = 5.0*(2.0*maxomega)/nt # width of the filter\n",
+ "freq1 = -2.0/(mfilt+1)\n",
+ "filter = np.exp(-((freqs - freq1)/width)**2)\n",
+ "print(freqs.shape, lpower.shape)\n",
+ "\n",
+ "fig,ax=plt.subplots(1,2, figsize=(9,2))\n",
+ "for i in range(2):\n",
+ " ax[i].plot(freqs, lpower[mfilt,:]/np.max(lpower[mfilt,:]), label='Wave Power')\n",
+ " ax[i].plot(freqs, np.sqrt(filter), label='Filter')\n",
+ " ax[i].plot(freqs, lpower[mfilt,:]*np.sqrt(filter)/np.max(lpower[mfilt,:]), label='Filtered Wave Power')\n",
+ " ax[i].set_xlabel(r'Frequency $(2\\pi/\\Omega) \\omega)$')\n",
+ " ax[i].set_ylabel(r'Power')\n",
+ "ax[1].set_yscale('log')\n",
+ "ax[1].legend(loc='lower left')\n",
+ "\n",
+ "fig.tight_layout()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Since, the complex spectrum is fully in spectral space, $f(\\ell,m,\\omega)$, we can apply a spatial filter for $m=1$, by simply zeroing out the signal in all channels other than $m=1$. The Gaussian frequency filter can be applied using the inverse transforms from the TemporalFFT class (the filter is like a window function, but for frequency space). In the cell below, we perform such filtering and then plot a few timeframes of the resulting shell slices."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Filter the entire dataset and invert the transforms\n",
+ "\n",
+ "#Remove all m other than m = 5\n",
+ "fspectra = wspectra.copy()\n",
+ "for i in range(nm):\n",
+ " if i != mfilt:\n",
+ " fspectra[:,i,:] = 0.0 # kill all power that is at an m we don't want\n",
+ "\n",
+ "# Frequency spacing in the transforms\n",
+ "dom = freqs[1]-freqs[0]\n",
+ "\n",
+ "# Invert the temporal transforms while applying the spectral filter (fslices is the inverse transform, ts are the resulting temporal grid)\n",
+ "fslices, ts = TemporalFFT.uniform_inverse(fspectra, dom, filter=filter, axis=-1).ifft()\n",
+ "\n",
+ "#Invert the spherical harmonic transform\n",
+ "SHT = spectral_utils.SHT(ntheta, spectral=False, dealias=dealias) # Instantiate the transform (may not be necessary depending on the cells you have already run)\n",
+ "filtered = SHT.to_physical(fslices, th_l_axis=0, phi_m_axis=1)\n",
+ "filtered = filtered.real # The data should be real\n",
+ "\n",
+ "\n",
+ "fig, ax = plt.subplots(3,1, figsize=(5,8))\n",
+ "for i in range(3):\n",
+ " ax[i].imshow(filtered[:,:,i*10], extent=[0, 360, -90, 90])\n",
+ " ax[i].set_xlabel('Longitude')\n",
+ " ax[i].set_ylabel('Latitude')\n",
+ "ax[0].set_title('Data filtered for a single mode $m=${}'.format(mfilt))\n",
+ "fig.tight_layout()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## D. Temporal FFT on a Nonuniform Grid\n",
+ "\n",
+ "The previous examples have worked with data on a uniform temporal grid. The TemporalFFT class can also deal with irregular grids. It does so by interpolating the data onto a regular grid and then transforming using the same FFT algorithm. We will illustrate this feature by either generating an artificial data set on a nonuniform grid or be reading in a set of Shell_Slices (which often have a variable temporal output cadence).\n",
+ "\n",
+ "If fakedata=True, sets of shell slices are generated by adding up a small set of sectoral equatorial Rossby waves. Each Rossby wave consists of a single spherical harmonic component (with $m = \\ell$) that oscillates at the frequency $\\omega = -\\frac{2m\\Omega}{\\ell(\\ell+1)}$. The waves propagate in the retrograde direction (westward). Times and frequencies are nondimensionalized in terms of rotations, i.e., $t = \\frac{2\\pi}{\\Omega} \\tilde{t}$ and $\\omega = \\frac{\\Omega}{2\\pi} \\tilde{\\omega}$. The times at which this wavefield is evaluated are pulled from a random distribution.\n",
+ "\n",
+ "If you would prefer to use your own Rayleigh outputs, set fakedata=False and set the \"slice_path\" to the relevant value to indicate the directory where your data is stored.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Either Generate a set of artificial data or read in Shell_Slice files\n",
+ "# The resulting data consists of a single radial shell slice as a function of time\n",
+ "# data[theta, phi, t]\n",
+ "\n",
+ "fakedata = True\n",
+ "slice_path = 'Your_Directory_Here'\n",
+ "\n",
+ "# generate fake spherical harmonic data (a set of sectoral Rossby waves) \n",
+ "if fakedata:\n",
+ " \n",
+ " #parameters\n",
+ " noise = 0.2 # scale size for the random deviation for every point in the temporal grid\n",
+ " maxm = 20 # number of sectoral Rossby waves\n",
+ " dealias = 1.5 \n",
+ " ntheta = 384\n",
+ " nphi = 2*ntheta\n",
+ " nt = 128\n",
+ " duration = 25.0 # Number of rotations\n",
+ " nell = ntheta * 2//3\n",
+ " nm = nell\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 the temporal array with random points\n",
+ " if (noise == 0.0): \n",
+ " times = duration*np.random.rand(nt) # points drawn from a uniform distribution \n",
+ " times = np.sort(times)\n",
+ " else:\n",
+ " times = np.linspace(0,duration, nt) # Evenly spaced points with small random fluctuations\n",
+ " times[1:nt-1] += noise*(duration/nt)*(np.random.rand(nt-2)-0.5)\n",
+ "\n",
+ " # Generate the fake data\n",
+ " nonuniform_slices = np.ndarray((ntheta,nphi,nt), dtype=np.float64)\n",
+ " ic = complex(0.0,1.0)\n",
+ " for m in range(1, maxm+1):\n",
+ " Ylm = sph_harm(m, m, theta, phi); freq = -4.0*np.pi/(m+1) \n",
+ " for i,t in enumerate(times):\n",
+ " mode = np.real(np.exp(ic*freq*t)*Ylm)\n",
+ " nonuniform_slices[:,:,i] += np.transpose(mode)\n",
+ " axis = -1 # time corresponds the final dimension in the array\n",
+ " xlim=30\n",
+ " norm = None\n",
+ " \n",
+ "else:\n",
+ " # Generate a list of files in the chosen directory\n",
+ " \n",
+ " import os\n",
+ " \n",
+ " files = os.listdir(slice_path)\n",
+ " files = np.sort(files)\n",
+ " nfiles = len(files)\n",
+ " sslice = Shell_Slices(files[0], path=slice_path+'/')\n",
+ " \n",
+ " nt = sslice.niter * nfiles\n",
+ " ntheta = sslice.ntheta\n",
+ " nphi = sslice.nphi \n",
+ " \n",
+ " nonuniform_slices = np.ndarray((ntheta, nphi, nt), dtype=np.float64)\n",
+ " times = np.ndarray(nt, dtype=np.float64)\n",
+ " for it in range(sslice.niter):\n",
+ " single_slice = sslice.vals[:,:,-1,sslice.lut[1],it] # radial velocity spectrum (1) at the inner boundary\n",
+ " nonuniform_slices[:,:,it] = single_slice.T\n",
+ " times[it] = sslice.time[it]\n",
+ " \n",
+ " # read in the shell slices\n",
+ " i = sslice.niter\n",
+ " for file in files[1:]: \n",
+ " sslice = Shell_Slices(file, path=slice_path+'/')\n",
+ " for it in range(sslice.niter): \n",
+ " single_slice = sslice.vals[:,:,-1,sslice.lut[1],it]\n",
+ " nonuniform_slices[:,:,i] = slice.T\n",
+ " times[i] = sslice.time[it]\n",
+ " i += 1\n",
+ " nt = i\n",
+ " nonuniform_slices = nonuniform_slices[:,:,:nt-1]\n",
+ " times = times[:nt-1]\n",
+ "\n",
+ " axis = -1 # time corresponds the final dimension in the array \n",
+ " xlim=50\n",
+ " norm=colors.LogNorm(vmin=1e-30, vmax=1e-15)\n",
+ " duration = times[-1]-times[0]\n",
+ "\n",
+ "\n",
+ "# plot a sample slice and sample spectrum\n",
+ "fig = plt.figure(figsize=(14, 8))\n",
+ "gs = gridspec.GridSpec(1, 2, width_ratios=[2, 1])\n",
+ "\n",
+ "# Take the spatial transform of a sample slice\n",
+ "transform_SHT = spectral_utils.SHT(ntheta, spectral=False) # instantiate the spherical harmonic transform\n",
+ "sample_slice = nonuniform_slices[:,:,nt//4]\n",
+ "sample_spectrum = transform_SHT.to_spectral(sample_slice, th_l_axis=0, phi_m_axis=1)\n",
+ "\n",
+ "\n",
+ "ax0 = plt.subplot(gs[0])\n",
+ "ax0.set_title('Initial 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('Initial Spatial Spectra')\n",
+ "power = sample_spectrum.real**2 + sample_spectrum.imag**2\n",
+ "pos1=ax1.imshow(np.transpose(power), norm=norm)\n",
+ "ax1.set_xlabel(r'Harmonic Degree $\\ell$')\n",
+ "ax1.set_ylabel(r'Azimuthal Order $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)\n",
+ "\n",
+ "fig.tight_layout()\n",
+ "\n",
+ "\n",
+ "# plot a subsequent slice and sample spectrum\n",
+ "fig = plt.figure(figsize=(14, 8))\n",
+ "gs = gridspec.GridSpec(1, 2, width_ratios=[2, 1])\n",
+ "\n",
+ "# Take the spatial transform of a sample slice\n",
+ "sample_slice = nonuniform_slices[:,:,nt//4+5]\n",
+ "sample_spectrum = transform_SHT.to_spectral(sample_slice, th_l_axis=0, phi_m_axis=1)\n",
+ "\n",
+ "ax0 = plt.subplot(gs[0])\n",
+ "ax0.set_title('Later 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",
+ "\n",
+ "ax1 = plt.subplot(gs[1])\n",
+ "ax1.set_title('Later 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(r'Harmonic Degree $\\ell$')\n",
+ "ax1.set_ylabel(r'Azimuthal Order $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)\n",
+ "\n",
+ "fig.tight_layout()\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Transform to Spectral Space\n",
+ "\n",
+ "method = 'cubic' # Interpolation types\n",
+ " # Options: linear, cubic, and nearest\n",
+ " # Default: linear\n",
+ "\n",
+ "# Transform from physical spatial space to spectral spatial space (theta, phi) -> (l, m)\n",
+ "spatial_spectra = transform_SHT.to_spectral(nonuniform_slices, th_l_axis=0, phi_m_axis=1)\n",
+ "\n",
+ "\n",
+ "# Transform from time to temporal frequency space t -> omega\n",
+ "dt = duration / nt # Mean cadence\n",
+ "lm_omega, omegas = TemporalFFT.nonuniform(spatial_spectra, times, dt, axis = -1, method=method).fft()\n",
+ "power = np.abs(lm_omega)**2\n",
+ "\n",
+ "# Sum power over l\n",
+ "lpower = np.sum(power,axis=0) # Power as a function of azimuthal order m and frequency omega\n",
+ "\n",
+ "# Plot!\n",
+ "fig,ax=plt.subplots(1,2, figsize=(12,4))\n",
+ "nyquist = np.pi/dt # Nyquist\n",
+ "twopi = 2.0*np.pi # For converting from 2 pi/Omega to 1/Omega temporal units\n",
+ "maxomega = np.max(omegas)\n",
+ "\n",
+ "\n",
+ "pos = ax[0].imshow(np.log(lpower[0:maxm+1,:]).T, cmap='plasma', extent = [-0.5,maxm+0.5,-maxomega,maxomega], aspect='auto')\n",
+ "ax[0].axhline(0.0, color = 'gray')\n",
+ "\n",
+ "cols = ['blue', 'red', 'green']\n",
+ "for i,ell in enumerate([1,5,15]):\n",
+ " ax[1].plot(omegas, lpower[ell,:], color=cols[i])\n",
+ " ax[0].axvline(ell, color=cols[i], linestyle=':')\n",
+ "ax[0].set_title('Log Power')\n",
+ "ax[0].set_xlabel(r'Azimuthal Order $m$')\n",
+ "ax[0].set_ylabel(r'Frequency $\\omega$')\n",
+ "ax[1].set_xlabel(r'Frequency $\\omega$')\n",
+ "ax[1].set_ylabel(r'Power')\n",
+ "ax[1].set_xlim(-1*maxomega,maxomega)\n",
+ "ax[1].set_yscale('log')\n",
+ "fig.colorbar(pos)"
+ ]
+ },
{
"cell_type": "code",
"execution_count": null,
@@ -1184,7 +1765,7 @@
],
"metadata": {
"kernelspec": {
- "display_name": "Python 3",
+ "display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
@@ -1198,7 +1779,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
- "version": "3.8.3"
+ "version": "3.12.12"
}
},
"nbformat": 4,