Polar and azimuthal derivatives

Setup

import torch
from torch_harmonics.sht import RealSHT
from torch_harmonics.sht import InverseRealVectorSHT
import matplotlib.pyplot as plt
from torch_harmonics.plotting import plot_sphere

Let’s define an example function \(f = sin(\theta)cos(\phi)\), with derivatives with respect to \(\theta\) and \(\phi\): \(f^*_{\theta}\) and \(f^*_{\phi}\)

# the function f(theta, phi) is the function we want to differentiate
def f(theta, phi):
    return torch.sin(theta) * torch.cos(phi)

# nft(theta, phi) is the derivative of f with respect to theta -> f*_theta
def nft(theta, phi):
    return torch.cos(theta) * torch.cos(phi)

# nfp(theta, phi) is the derivative of f with respect to phi -> f*_phi
def nfp(theta, phi):
    return -torch.sin(theta) * torch.sin(phi)

Here we define a grid of \(\theta\) and \(\phi\) values and we evaluate the functions for each grid point

# Create a grid of theta and phi values
n_theta, n_phi = 128, 256
theta = torch.linspace(0, torch.pi, n_theta)
phi = torch.linspace(0, 2 * torch.pi, n_phi)
theta_grid, phi_grid = torch.meshgrid(theta, phi, indexing='ij')

# Evaluate the function on the grid for f, nft and nfp 
f_grid = f(theta_grid, phi_grid)
nft_grid = nft(theta_grid, phi_grid) 
nfp_grid = nfp(theta_grid, phi_grid)

Computation of the gradient

A convenient approach to the computation of the derivatives along \(\theta\) and \(\phi\) is to compute the Spherical Harmonic Transform (SHT) of the function we want to differentiate and then apply the Inverse Vector Spherical Harmonic Transform (IVSHT) to get \(\nabla f\). We can then proceed to read \(\partial f / \partial \theta\) and \(\partial f / \partial \phi\) from the first and second component.

lmax = 5
sh_coeffs = RealSHT(n_theta, n_phi, lmax=lmax)(f_grid)

# We need to adapt the output of the SHT to be a vector field
vector_coeffs = torch.zeros(2, sh_coeffs.shape[0], sh_coeffs.shape[1], dtype=sh_coeffs.dtype)
vector_coeffs[0] = sh_coeffs 

nabla_f = InverseRealVectorSHT(n_theta, n_phi, lmax=lmax)(vector_coeffs)

# Since from IVSHT we get 1/sin(theta) * d/dphi, we need to multiply by sin(theta) to get the correct gradient
nabla_f[1] = nabla_f[1] * torch.sin(theta_grid)

print(f"Vector coefficients shape: {vector_coeffs.shape}")  # Should be (2, lmax, mmax)
print(f"Gradient field shape: {nabla_f.shape}")  # Should be (2, n_theta, n_phi)

Vector coefficients shape: torch.Size([2, 5, 129])
Gradient field shape: torch.Size([2, 128, 256])

Plotting

fig = plt.figure(layout='constrained', figsize=(12, 8))
subfigs = fig.subfigures(2, 3)
clat, clon = 0, 56
plot_sphere(f_grid,   fig=subfigs[0, 0], central_latitude=clat, central_longitude=clon, projection="orthographic", cmap="turbo", colorbar=True, title=r"$f$")
plot_sphere(nabla_f[0],   fig=subfigs[0, 1], central_latitude=clat, central_longitude=clon, projection="orthographic", cmap="turbo", colorbar=True, title=r"${\partial f } / {\partial \theta}$")
plot_sphere(nabla_f[1],   fig=subfigs[1, 1], central_latitude=clat, central_longitude=clon, projection="orthographic", cmap="turbo", colorbar=True, title=r"${\partial f } / {\partial \phi}$")
plot_sphere(torch.abs(nabla_f[0] - nft_grid), fig=subfigs[0, 2], central_latitude=clat, central_longitude=clon, projection="orthographic", cmap="turbo", colorbar=True, title=r"$|{\partial f }/{\partial \theta} - f^*_{\theta}|$")
plot_sphere(torch.abs(nabla_f[1] - nfp_grid), fig=subfigs[1, 2], central_latitude=clat, central_longitude=clon, projection="orthographic", cmap="turbo", colorbar=True, title=r"$|{\partial f }/{\partial \phi} - f^*_{\phi}|$")
plt.show()