def test_generator_spherical(): size = 64 r_min, r_max = 0.0, 1.0 generator = GeneratorSpherical(size, r_min=r_min, r_max=r_max, method='equally-spaced-noisy') r, theta, phi = generator.get_examples() assert _check_shape_and_grad(generator, size, r, theta, phi) assert _check_boundary((r, theta, phi), (r_min, 0.0, 0.0), (r_max, np.pi, np.pi * 2)) generator = GeneratorSpherical(size, r_min=r_min, r_max=r_max, method='equally-radius-noisy') r, theta, phi = generator.get_examples() assert _check_shape_and_grad(generator, size, r, theta, phi) assert _check_boundary((r, theta, phi), (r_min, 0.0, 0.0), (r_max, np.pi, np.pi * 2)) with pytest.raises(ValueError): _ = GeneratorSpherical(64, method='bad_generator') with pytest.raises(ValueError): _ = GeneratorSpherical(64, r_min=-1.0) with pytest.raises(ValueError): _ = GeneratorSpherical(64, r_min=1.0, r_max=0.0) str(generator) repr(generator)
def test_generator_spherical(): size = 64 r_min, r_max = 0.0, 1.0 generator = GeneratorSpherical(size, r_min=r_min, r_max=r_max, method='equally-spaced-noisy') r, theta, phi = generator.get_examples() assert _check_shape_and_grad(generator, size, r, theta, phi) assert _check_boundary((r, theta, phi), (r_min, 0.0, 0.0), (r_max, np.pi, np.pi * 2)) generator = GeneratorSpherical(size, r_min=r_min, r_max=r_max, method='equally-radius-noisy') r, theta, phi = generator.get_examples() assert _check_shape_and_grad(generator, size, r, theta, phi) assert _check_boundary((r, theta, phi), (r_min, 0.0, 0.0), (r_max, np.pi, np.pi * 2))
def test_solve_spherical_system(): # a PDE system that can be decoupled into 2 Laplacian equations :math:`\\nabla^2 u = 0` and :math:`\\nabla^2 v = 0` pde_system = lambda u, v, r, theta, phi: [ laplacian_spherical(u, r, theta, phi) + laplacian_spherical( v, r, theta, phi), laplacian_spherical(u, r, theta, phi) - laplacian_spherical( v, r, theta, phi), ] # constant boundary conditions for u and v; solution should be u = 0 identically and v = 1 identically conditions = [ DirichletBVPSpherical(r_0=0., f=lambda phi, theta: 0., r_1=1., g=lambda phi, theta: 0.), DirichletBVPSpherical(r_0=0., f=lambda phi, theta: 1., r_1=1., g=lambda phi, theta: 1.), ] solution, loss_history = solve_spherical_system(pde_system, conditions, 0.0, 1.0, max_epochs=2, return_best=True) generator = GeneratorSpherical(512, r_min=0., r_max=1.) rs, thetas, phis = generator.get_examples() us, vs = solution(rs, thetas, phis, as_type='np')
def test_train_generator_spherical(): pde = laplacian_spherical condition = NoCondition() train_generator = GeneratorSpherical(size=64, r_min=0., r_max=1., method='equally-spaced-noisy') r, th, ph = train_generator.get_examples() assert (0. < r.min()) and (r.max() < 1.) assert (0. <= th.min()) and (th.max() <= np.pi) assert (0. <= ph.min()) and (ph.max() <= 2 * np.pi) valid_generator = GeneratorSpherical(size=64, r_min=1., r_max=1., method='equally-radius-noisy') r, th, ph = valid_generator.get_examples() assert (r == 1).all() assert (0. <= th.min()) and (th.max() <= np.pi) assert (0. <= ph.min()) and (ph.max() <= 2 * np.pi) solve_spherical(pde, condition, 0.0, 1.0, train_generator=train_generator, valid_generator=valid_generator, max_epochs=1) with raises(ValueError): _ = GeneratorSpherical(64, method='bad_generator') with raises(ValueError): _ = GeneratorSpherical(64, r_min=-1.0) with raises(ValueError): _ = GeneratorSpherical(64, r_min=1.0, r_max=0.0)
def test_solve_spherical(): pde = laplacian_spherical generator = GeneratorSpherical(512) # 0-boundary condition; solution should be u(r, theta, phi) = 0 identically f = lambda th, ph: 0. g = lambda th, ph: 0. condition = DirichletBVPSpherical(r_0=0., f=f, r_1=1., g=g) solution, loss_history = solve_spherical(pde, condition, 0.0, 1.0, max_epochs=2, return_best=True) rs, thetas, phis = generator.get_examples() us = solution(rs, thetas, phis, as_type='np')
def x(): n_points, r_min, r_max = 1024, 1.0, 10.0 g = GeneratorSpherical(n_points, r_min=r_min, r_max=r_max) return [t.reshape(-1, 1) for t in g.get_examples()]
def validate(solution): generator = GeneratorSpherical(512, r_min=r_0, r_max=r_1) rs, thetas, phis = generator.get_examples() us = solution(rs, thetas, phis, to_numpy=True) vs = analytic_solution(rs, thetas, phis).detach().cpu().numpy() assert us.shape == vs.shape
def validate(solution, loss_history, analytical_mse): generator = GeneratorSpherical(512, r_min=r_0, r_max=r_1) rs, thetas, phis = generator.get_examples() us = solution(rs, thetas, phis, as_type="np") vs = analytic_solution(rs, thetas, phis).detach().cpu().numpy() assert us.shape == vs.shape
super(HarmonicsNN, self).__init__() self.net_r = FCNN(1, n_output_units=len(degrees)) self.harmonics_fn = harmonics_fn 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)