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)