def test_zero_order_spherical_harmonics(shape, max_degree):
# note that in scipy, theta is azimuthal angle (0, 2 pi) while phi is polar angle (0, pi)
thetas1 = np.random.rand(*shape) * np.pi * 2
phis1 = np.random.rand(*shape) * np.pi
# in neurodiffeq, theta and phi should be exchanged
thetas2 = torch.tensor(phis1, requires_grad=True)
phis2 = torch.tensor(thetas1, requires_grad=True)
order = 0
y1 = np.concatenate(
[
sph_harm(order, degree, thetas1, phis1)
for degree in range(max_degree + 1)
],
axis=1,
)
assert (np.imag(y1) == 0).all(), f"y1 has non-zero imaginary part: {y1}"
y1 = np.real(y1)
net = ZonalSphericalHarmonics(max_degree)
y2 = net(thetas2, phis2)
assert y2.requires_grad, f"output seems detached from the graph"
y2 = y2.detach().cpu().numpy()
assert isclose(y2, y1, atol=1e-5, rtol=1e-3).all(), \
f"y1 = {y1}, y2 = {y2}, delta = {y1 - y2}, max_delta = {np.max(abs(y1 - y2))}"
def test_zero_order_spherical_harmonics_laplacian():
# Somehow, if changing default dtype to float32, the test fails by a large margin
N_FLOAT = np.float64
T_FLOAT = torch.float64
THETA_EPS = 0.1
r_values = np.random.rand(*shape).astype(N_FLOAT) + 1.1
theta_values = np.random.uniform(THETA_EPS, np.pi - THETA_EPS,
size=shape).astype(N_FLOAT)
phi_values = np.random.rand(*shape).astype(N_FLOAT) * np.pi * 2
net = nn.Sequential(
nn.Linear(1, 10),
nn.Tanh(),
nn.Linear(10, max_degree + 1),
).to(T_FLOAT)
harmonics = ZonalSphericalHarmonics(max_degree=max_degree)
r1 = torch.tensor(r_values, requires_grad=True)
theta1 = torch.tensor(theta_values, requires_grad=True)
phi1 = torch.tensor(phi_values, requires_grad=True)
coeffs1 = net(r1)
us = torch.sum(coeffs1 * harmonics(theta1, phi1), dim=1, keepdim=True)
def laplacian1(u, r, theta, phi):
r_lap = diff(u * r, r, order=2) / r
theta_lap = diff(diff(u, theta) * torch.sin(theta),
theta) / (r**2) / torch.sin(theta)
phi_lap = diff(u, phi, order=2) / (r**2) / torch.sin(theta)**2
return r_lap + theta_lap + phi_lap
lap1 = laplacian1(us, r1, theta1, phi1)
assert lap1.requires_grad, "lap1 seems detached from graph"
r2 = torch.tensor(r_values, requires_grad=True)
theta2 = torch.tensor(theta_values, requires_grad=True)
phi2 = torch.tensor(phi_values, requires_grad=True)
coeffs2 = net(r2)
laplacian2 = ZonalSphericalHarmonicsLaplacian(max_degree=max_degree)
lap2 = laplacian2(coeffs2, r2, theta2, phi2)
assert lap2.requires_grad, "lap2 seems detached from graph"
lap1 = lap1.detach().cpu().numpy()
lap2 = lap2.detach().cpu().numpy()
assert isclose(lap2, lap1).all(), \
f"lap1 = {lap1}\nlap2 = {lap2}\ndelta = {lap1 - lap2}\nmax_delta = {np.max(abs(lap1 - lap2))}"
def u(x):
return HarmonicsNN(degrees, ZonalSphericalHarmonics(degrees=degrees))(*x)
def U(x):
F = [
HarmonicsNN(degrees, ZonalSphericalHarmonics(degrees=degrees))
for _ in range(3)
]
return list(map(lambda f: f(*x), F))
def forward(self, r, theta, phi):
R = self.net_r(r)
Y = self.harmonics_fn(theta, phi)
return (R * Y).sum(dim=1, keepdim=True)
EPS = 1e-4
n_points = 1024
r_min = 1.0
r_max = 1.0
generator = GeneratorSpherical(n_points, r_min=r_min, r_max=r_max)
r, theta, phi = [t.reshape(-1, 1) for t in generator.get_examples()]
degrees = list(range(10))
harmoincs_fn = ZonalSphericalHarmonics(degrees=degrees)
F_r, F_theta, F_phi = [HarmonicsNN(degrees, harmoincs_fn) for _ in range(3)]
vector_u = (F_r(r, theta, phi), F_theta(r, theta, phi), F_phi(r, theta, phi))
f = HarmonicsNN(degrees, harmoincs_fn)
scalar_u = f(r, theta, phi)
curl = lambda a, b, c: spherical_curl(a, b, c, r, theta, phi)
grad = lambda a: spherical_grad(a, r, theta, phi)
div = lambda a, b, c: spherical_div(a, b, c, r, theta, phi)
lap = lambda a: spherical_laplacian(a, r, theta, phi)
vec_lap = lambda a, b, c: spherical_vector_laplacian(a, b, c, r, theta, phi)
def is_zero(t):
def harmonics_fn(degrees):
return ZonalSphericalHarmonics(degrees=degrees)