def test_spherical_laplacian(): N_FLOAT = np.float64 T_FLOAT = torch.float64 n_samples = 10 r_value = np.random.rand(n_samples, 1).astype(N_FLOAT) theta_value = np.random.rand(n_samples, 1).astype(N_FLOAT) phi_value = np.random.rand(n_samples, 1).astype(N_FLOAT) r_net = FCNN(n_input_units=1, n_output_units=25).to(T_FLOAT) # compute laplacians using spherical harmonics property r1 = torch.tensor(r_value, dtype=T_FLOAT, requires_grad=True) theta1 = torch.tensor(theta_value, dtype=T_FLOAT, requires_grad=True) phi1 = torch.tensor(phi_value, dtype=T_FLOAT, requires_grad=True) R1 = r_net(r1) harmonics_laplacian = HarmonicsLaplacian(max_degree=4) lap1 = harmonics_laplacian(R1, r1, theta1, phi1) # compute laplacians using brute force r2 = torch.tensor(r_value, dtype=T_FLOAT, requires_grad=True) theta2 = torch.tensor(theta_value, dtype=T_FLOAT, requires_grad=True) phi2 = torch.tensor(phi_value, dtype=T_FLOAT, requires_grad=True) R2 = r_net(r2) spherical_fn = RealSphericalHarmonics(max_degree=4) harmonics = spherical_fn(theta2, phi2) u = torch.sum(R2 * harmonics, dim=1, keepdim=True) lap2 = laplacian_spherical(u, r2, theta2, phi2) lap1 = lap1.detach().cpu().numpy() lap2 = lap2.detach().cpu().numpy() assert np.isclose(lap2, lap1).all(), \ f'Laplcians computed using spherical harmonics trick differ from brute force solution, {lap1} != {lap2}'
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)