def test_concat_generator():
size1, size2 = 10, 20
t_min, t_max = 0.5, 1.5
generator1 = Generator1D(size1, t_min=t_min, t_max=t_max)
generator2 = Generator1D(size2, t_min=t_min, t_max=t_max)
concat_generator = ConcatGenerator(generator1, generator2)
x = concat_generator.get_examples()
assert _check_shape_and_grad(concat_generator, size1 + size2, x)
grid1 = (4, 4, 4)
size1, size2, size3 = grid1[0] * grid1[1] * grid1[2], 100, 200
generator1 = Generator3D(grid=grid1)
generator2 = GeneratorSpherical(size2)
generator3 = GeneratorSpherical(size3)
concat_generator = ConcatGenerator(generator1, generator2, generator3)
r, theta, phi = concat_generator.get_examples()
assert _check_shape_and_grad(concat_generator, size1 + size2 + size3, r,
theta, phi)
added_generator = generator1 + generator2 + generator3
r, theta, phi = added_generator.get_examples()
assert _check_shape_and_grad(added_generator, size1 + size2 + size3, r,
theta, phi)
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)
with pytest.warns(FutureWarning):
solution, loss_history = solve_spherical(pde,
condition,
0.0,
1.0,
max_epochs=2,
return_best=True)
assert isinstance(solution, SolutionSpherical)
def test_static_generator():
size = 100
generator = Generator1D(size)
static_generator = StaticGenerator(generator)
x1 = static_generator.get_examples()
x2 = static_generator.get_examples()
assert _check_shape_and_grad(generator, size)
assert _check_shape_and_grad(static_generator, size, x1, x2)
assert (x1 == x2).all()
size = 100
generator = GeneratorSpherical(size)
static_generator = StaticGenerator(generator)
r1, theta1, phi1 = static_generator.get_examples()
r2, theta2, phi2 = static_generator.get_examples()
assert _check_shape_and_grad(generator, size)
assert _check_shape_and_grad(static_generator, size, r1, theta1, phi1, r2, theta2, phi2)
assert (r1 == r2).all() and (theta1 == theta2).all() and (phi1 == phi2).all()
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_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 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
def __init__(self, degrees, harmonics_fn):
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)