SHTOOLS and WindsPharm are white paper backed way to calculate Power Spectra, Total Kinetic Energy and Spherical Harmonics of geospatial datasets loaded with Xarray. I verified the validity of SHTools against an internal library we use
And it seems for scalar fields it works a treat 
.
However!!! When I try the same approach using a vector field
and so when calculating the total kinetic energy (divergence, vorticity, etc.)
I get a consistent difference which I believe to have narrowed down to explosion of derivatives near the North Pole and some miscalculation of irregularly gridded global models.
I tried calculating them using MetPy’s divergence and vorticity as well
def plot_power_spectra_vector(
fst_base_path: PosixPath,
pres_level: Union[float, str],
validate: bool,
threshold=1e-9,
max_marker_size=12,
min_marker_size=1,
base_aspect_ratio=(2, 1), # Width-to-height ratio
max_height=12 # Max plot height
):
"""
Plot power spectra for multiple pressure levels with logarithmic marker size scaling and external legend.
Parameters:
- fst_base_path (PosixPath): Path to the base FST file.
- pres_level ([Union[float, str]): Pressure level to plot.
- threshold (float): Threshold for additional processing (not used directly here).
- cmap (str): Colormap name for cycling line colors.
- max_marker_size (int): Maximum marker size (for small x).
- min_marker_size (int): Minimum marker size (for large x).
- base_aspect_ratio (tuple): Width-to-height ratio of the plot.
- max_height (int): Maximum height of the plot.
"""
a_r = 6.371e6
# Load dataset and calculate the first spectrum
ds = load_ds_uv(fst_base_path)
if ds.UU.dims != ("lat", "lon"):
ds.UU = ds.UU.transpose("lat", "lon")
ds.VV = ds.VV.transpose("lat", "lon")
div = divergence(ds.UU * mu.units('m/s'), ds.VV * mu.units('m/s'), latitude=ds.lat, longitude=ds.lon)
vort = vorticity(ds.UU * mu.units('m/s'), ds.VV * mu.units('m/s'), latitude=ds.lat, longitude=ds.lon)
grid_div = pysh.SHGrid.from_xarray(div, grid="GLQ")
grid_vort = pysh.SHGrid.from_xarray(vort, grid="GLQ")
coeffs_div = grid_div.expand()
coeffs_vort = grid_vort.expand()
spectrum_div = coeffs_div.spectrum()
spectrum_vort = coeffs_vort.spectrum()
degrees = np.arange(1, len(spectrum_div) + 1)
leading_factor = (a_r**2) / (4 * np.pi * degrees * (degrees + 1))
S_div = leading_factor * spectrum_div
S_rot = leading_factor * spectrum_vort
S_tot = S_div + S_rot
wavelength = a_r / degrees
# Calculate dynamic figure dimensions
base_width = 10
base_height = 6
extra_height = 0.2 * max(0, 25 - 4) # Modest height scaling for extra levels
fig_height = min(base_height + extra_height, max_height)
fig_width = fig_height * base_aspect_ratio[0] / base_aspect_ratio[1]
# Create the main plot
fig, ax1 = plt.subplots(figsize=(fig_width, fig_height))
x_ref = np.array([16, 128]) # Reference x-range
y_ref_start = 10 # Starting y-value for the reference line
y_ref = y_ref_start * (x_ref / x_ref[0]) ** -3 # Calculate y-values for k^-3 slope
x_mid = np.sqrt(x_ref[0] * x_ref[1]) # Geometric midpoint of the x range
y_mid = y_ref_start * (x_mid / x_ref[0]) ** -3 # Calculate y-value at midpoint
ax1.text(
x_mid, y_mid * 2, # Slightly above the midpoint of the line
r"$k^{-3}$",
fontsize=12,
color='black',
ha='center', # Center align horizontally
va='bottom' # Align text slightly above the line
)
# Calculate marker sizes based on log degrees
log_degrees = np.log(degrees)
marker_sizes = np.interp(
log_degrees,
(log_degrees.min(), log_degrees.max()),
(max_marker_size, min_marker_size)
)
# Plot the KE Rot spectrum
ax1.scatter(
degrees,
S_rot,
label=f"KE Rotational",
color='blue',
s=marker_sizes,
edgecolors='face'
)
ax1.plot(
degrees,
S_rot,
color='blue',
linestyle='-'
)
# Plot the KE Rot spectrum
ax1.scatter(
degrees,
S_div,
label=f"KE Divergent",
color='green',
s=marker_sizes,
edgecolors='face'
)
ax1.plot(
degrees,
S_div,
color='green',
linestyle='-'
)
# Plot the KE Total spectrum
ax1.scatter(
degrees,
S_tot,
label=f"KE Total",
color='orange',
s=marker_sizes,
edgecolors='face'
)
ax1.plot(
degrees,
S_tot,
color='orange',
linestyle='-'
)
# Plot the reference line
ax1.plot(
x_ref,
y_ref,
color='black',
linestyle='--',
linewidth=1.5,
# label=r'$k^{-3}$'
)
if validate:
# Plot /fs/site5/eccc/cmd/s/yor000/projects/power_spectra/field_vector_qe_harmon.txt
ax1.scatter(
degrees,
np.loadtxt(txts[3]),
label=f"KE Rotational Reference",
color='blue',
s=marker_sizes,
marker="1"
)
ax1.plot(
degrees,
np.loadtxt(txts[3]),
color='blue',
linestyle='dashdot'
)
# Plot /fs/site5/eccc/cmd/s/yor000/projects/power_spectra/field_vector_de_harmon.txt
ax1.scatter(
degrees,
np.loadtxt(txts[1]),
label=f"KE Divergent Reference",
color='green',
s=marker_sizes,
marker="1"
)
ax1.plot(
degrees,
np.loadtxt(txts[1]),
color='green',
linestyle='dashdot'
)
# Plot # Plot /fs/site5/eccc/cmd/s/yor000/projects/power_spectra/field_vector_ke_harmon.txt
ax1.scatter(
degrees,
np.loadtxt(txts[2]),
label=f"KE Total Reference",
color='orange',
s=marker_sizes,
marker="1"
)
ax1.plot(
degrees,
np.loadtxt(txts[2]),
color='orange',
linestyle='dashdot'
)
# Plot the reference line
ax1.plot(
x_ref,
y_ref,
color='black',
linestyle='--',
linewidth=1.5,
# label=r'$k^{-3}$'
)
# Set log-log scale for the main plot
ax1.set_xscale('log')
# ax1.set_xscale('log')
# ax1.set_xticks([1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024])
# ax1.set_xticklabels([1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024])
# ax1.set_xlim(1, 1024)
ax1.set_yscale('log')
# Add main labels
ax1.set_xlabel('Total wave number', fontsize=12)
ax1.set_ylabel('Power Spectral Density [J/kg]', fontsize=12)
# Add the third axis (above the plot)
ax3 = ax1.twiny()
ax3.set_xscale("log")
ax3.set_xlim(ax1.get_xlim())
# Custom tick positions and labels for the third axis
custom_ticks = [1, 4, 16, 40, 80, 160, 400]
custom_labels = [40000, 10000, 2500, 1000, 500, 250, 100]
ax3.set_xticks(custom_ticks)
ax3.set_xticklabels(custom_labels)
ax3.set_xlabel("Wave Length [km]", fontsize=12)
# Add grid
ax1.grid(True, which='both', linestyle='--', linewidth=0.5)
plt.suptitle(
f"Vector Field Power Spectra for the {fst_base_path.parts[-1]} file",
fontsize=14,
x=0.01, # Align to the left
y=0.98, # Slightly below the top of the figure
ha='left',
fontweight='bold'
)
ax1.legend(fontsize=10)
ax1.grid(True, which='both', linestyle='--', linewidth=0.5)
plt.tight_layout()
plt.show()
But I still get a difference. I would to find out there is an easy way to effectively calculate differential calculus on irregular grids (ideally sub global). Anyone have a refernce I can use?