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.),
]
with pytest.warns(FutureWarning):
solution, _ = solve_spherical_system(pde_system,
conditions,
0.0,
1.0,
max_epochs=2,
return_best=True)
assert isinstance(solution, SolutionSpherical)
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_dirichlet_bvp_spherical():
r0, r1 = x0, x1
r2 = (r0 + r1) / 2
no_condition = NoCondition()
# B.C. for the interior boundary (r_min)
net_f = FCNN(2, 1)
f = lambda th, ph: no_condition.enforce(net_f, th, ph)
# B.C. for the exterior boundary (r_max)
net_g = FCNN(2, 1)
g = lambda th, ph: no_condition.enforce(net_g, th, ph)
condition = DirichletBVPSpherical(r_0=r0, f=f, r_1=r1, g=g)
net = FCNN(3, 1)
theta = torch.rand(N_SAMPLES, 1) * np.pi
phi = torch.rand(N_SAMPLES, 1) * 2 * np.pi
r = r0 * ones
assert all_close(condition.enforce(net, r, theta, phi),
f(theta, phi)), "inner Dirichlet BC not satisfied"
r = r1 * ones
assert all_close(condition.enforce(net, r, theta, phi),
g(theta, phi)), "inner Dirichlet BC not satisfied"
condition = DirichletBVPSpherical(r_0=r2, f=f)
r = r2 * ones
assert all_close(condition.enforce(net, r, theta, phi),
f(theta, phi)), "single ended BC not satisfied"
def test_dirichlet_bvp_spherical():
# B.C. for the interior boundary (r_min)
interior = nn.Linear(in_features=2, out_features=1, bias=True)
f = lambda theta, phi: interior(torch.cat([theta, phi], dim=1))
# B.C. for the exterior boundary (r_max)
exterior = nn.Linear(in_features=2, out_features=1, bias=True)
g = lambda theta, phi: exterior(torch.cat([theta, phi], dim=1))
bvp = DirichletBVPSpherical(r_0=0., f=f, r_1=1.0, g=g)
net = nn.Linear(in_features=3, out_features=1, bias=True)
theta = torch.rand(10, 1) * np.pi
phi = torch.rand(10, 1) * 2 * np.pi
r = torch.zeros_like(theta)
v0 = f(theta, phi).detach().cpu().numpy()
u0 = bvp.enforce(net, r, theta, phi).detach().cpu().numpy()
assert np.isclose(v0, u0,
atol=1.e-5).all(), f"Unmatched boundary {v0} != {u0}"
r = torch.ones_like(theta)
v1 = g(theta, phi).detach().cpu().numpy()
u1 = bvp.enforce(net, r, theta, phi).detach().cpu().numpy()
assert np.isclose(v1, u1,
atol=1.e-5).all(), f"Unmatched boundary {v1} != {u1}"
bvp_half = DirichletBVPSpherical(r_0=2., f=f)
r = torch.ones_like(theta) * 2.
v2 = f(theta, phi).detach().cpu().numpy()
u2 = bvp_half.enforce(net, r, theta, phi).detach().cpu().numpy()
assert np.isclose(v2, u2,
atol=1.e-5).all(), f"Unmatched boundary {v2} != {u2}"
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_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 test_monitor_spherical():
f = lambda th, ph: 0.
g = lambda th, ph: 0.
conditions = [DirichletBVPSpherical(r_0=0., f=f, r_1=1., g=g)]
nets = [FCNN(3, 1)]
monitor = MonitorSpherical(0.0, 1.0, check_every=1)
loss_history = {
'train': list(np.random.rand(10)),
'valid': list(np.random.rand(10)),
}
analytic_mse_history = {
'train': list(np.random.rand(10)),
'valid': list(np.random.rand(10)),
}
history = {
'train_loss': list(np.random.rand(10)),
'valid_loss': list(np.random.rand(10)),
'train_foo': list(np.random.rand(10)),
'valid_foo': list(np.random.rand(10)),
'train_bar': list(np.random.rand(10)),
'valid_bar': list(np.random.rand(10)),
}
with pytest.warns(FutureWarning):
monitor.check(nets,
conditions,
history=history,
analytic_mse_history=analytic_mse_history)
with pytest.warns(FutureWarning):
monitor.check(nets, conditions, history=loss_history)
with pytest.warns(FutureWarning):
monitor.check(nets,
conditions,
history=loss_history,
analytic_mse_history=analytic_mse_history)
with pytest.raises(ValueError):
monitor.check(nets,
conditions,
history={
'train_foo': [],
'valid_foo': []
})
monitor.check(nets, conditions, history=history)
def test_monitor_spherical():
f = lambda th, ph: 0.
g = lambda th, ph: 0.
conditions = [DirichletBVPSpherical(r_0=0., f=f, r_1=1., g=g)]
nets = [FCNN(3, 1)]
monitor = MonitorSpherical(0.0, 1.0, check_every=1)
loss_history = {
'train': list(np.random.rand(10)),
'valid': list(np.random.rand(10)),
}
analytic_mse_history = {
'train': list(np.random.rand(10)),
'valid': list(np.random.rand(10)),
}
monitor.check(
nets,
conditions,
loss_history=loss_history,
analytic_mse_history=analytic_mse_history,
)
def test_electric_potential_gaussian_charged_density():
# total charge
Q = 1.
# standard deviation of gaussian
sigma = 1.
# medium permittivity
epsilon = 1.
# Coulomb constant
k = 1 / (4 * np.pi * epsilon)
# coefficient of gaussian term
gaussian_coeff = Q / (sigma**3) / np.power(2 * np.pi, 1.5)
# distribution of charge
rho_f = lambda r: gaussian_coeff * torch.exp(-r.pow(2) / (2 * sigma**2))
# analytic solution, refer to https://en.wikipedia.org/wiki/Poisson%27s_equation
analytic_solution = lambda r, th, ph: (k * Q / r) * torch.erf(r / (np.sqrt(
2) * sigma))
# interior and exterior radius
r_0, r_1 = 0.1, 3.
# values at interior and exterior boundary
v_0 = (k * Q / r_0) * erf(r_0 / (np.sqrt(2) * sigma))
v_1 = (k * Q / r_1) * erf(r_1 / (np.sqrt(2) * sigma))
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
# solving the problem using normal network (subject to the influence of polar singularity of laplacian operator)
pde1 = lambda u, r, th, ph: laplacian_spherical(u, r, th, ph) + rho_f(
r) / epsilon
condition1 = DirichletBVPSpherical(r_0, lambda th, ph: v_0, r_1,
lambda th, ph: v_1)
monitor1 = MonitorSpherical(r_0, r_1, check_every=50)
with pytest.warns(FutureWarning):
solution1, metrics_history = solve_spherical(
pde1,
condition1,
r_0,
r_1,
max_epochs=2,
return_best=True,
analytic_solution=analytic_solution,
monitor=monitor1,
)
validate(solution1)
# solving the problem using spherical harmonics (laplcian computation is optimized)
max_degree = 2
harmonics_fn = RealSphericalHarmonics(max_degree=max_degree)
harmonic_laplacian = HarmonicsLaplacian(max_degree=max_degree)
pde2 = lambda R, r, th, ph: harmonic_laplacian(R, r, th, ph) + rho_f(
r) / epsilon
R_0 = torch.tensor([v_0 * 2] + [0 for _ in range((max_degree + 1)**2 - 1)])
R_1 = torch.tensor([v_1 * 2] + [0 for _ in range((max_degree + 1)**2 - 1)])
def analytic_solution2(r, th, ph):
sol = torch.zeros(r.shape[0], (max_degree + 1)**2)
sol[:, 0:1] = 2 * analytic_solution(r, th, ph)
return sol
condition2 = DirichletBVPSphericalBasis(r_0=r_0,
R_0=R_0,
r_1=r_1,
R_1=R_1,
max_degree=max_degree)
monitor2 = MonitorSphericalHarmonics(r_0,
r_1,
check_every=50,
harmonics_fn=harmonics_fn)
net2 = FCNN(n_input_units=1, n_output_units=(max_degree + 1)**2)
with pytest.warns(FutureWarning):
solution2, metrics_history = solve_spherical(
pde2,
condition2,
r_0,
r_1,
net=net2,
max_epochs=2,
return_best=True,
analytic_solution=analytic_solution2,
monitor=monitor2,
harmonics_fn=harmonics_fn,
)
validate(solution2)
