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)
