From 5d79efffa512da0552871482b743b06315ad78e0 Mon Sep 17 00:00:00 2001 From: haakoek Date: Tue, 1 Jul 2025 11:13:49 +0200 Subject: [PATCH 1/3] Add finite element extension for GaussLobatto DVR. --- grid_methods/pseudospectral_grids/femdvr.py | 173 ++++++++++++++++++++ 1 file changed, 173 insertions(+) create mode 100644 grid_methods/pseudospectral_grids/femdvr.py diff --git a/grid_methods/pseudospectral_grids/femdvr.py b/grid_methods/pseudospectral_grids/femdvr.py new file mode 100644 index 0000000..b4b21f6 --- /dev/null +++ b/grid_methods/pseudospectral_grids/femdvr.py @@ -0,0 +1,173 @@ +import numpy as np +from grid_methods.pseudospectral_grids.gauss_legendre_lobatto import ( + GaussLegendreLobatto, + Linear_map, +) +from scipy.sparse import coo_matrix +import matplotlib.pyplot as plt +from typing import Type + + +class FEMDVR: + """A FEMDVR class. + + The interface is very similar to the GaussLobattto classes. + + """ + + def __init__( + self, + nodes, + n_points, + Mapping, + element_class: Type[GaussLegendreLobatto], + symmetrize=False, + ): + """Constructor. + + Args: + nodes (ndarray): Vector of nodes of boundaries of elements. Number of elements is this len(nodes) - 1. nodes[i] and nodes[i+1] defines the boundary of element i. + n_points (ndarray): Number of *total* grid points per element. The polynomial degree over element i is n_points[i] - 1. + element_class: Class for the elements. A subclass of GaussLobatto (which is an abstract base class). + + The constructor creates the following attributes: + self.D (ndarray): Differentiation matrix + self.D2 (ndarray): Second-order differentiation matrix + self.R (ndarray): Matrix that identifies nodes at element interfaces. + self.x (ndarray): Nodes + self.w (ndarray): Weights + self.dvr (list): List of dvr object instances for each element + self.edge_indices (ndarray): Indices of elements of x corresponding to element boundaries + + """ + + self.n_points = n_points + self.nodes = nodes + self.degrees = self.n_points - 1 + self.n_intervals = len(nodes) - 1 + n_intervals = self.n_intervals + + n_constraints = n_intervals - 1 + dim0 = np.sum( + n_points + ) # dimension of function space without boundary matching + D = np.zeros((dim0, dim0)) + R = np.zeros((dim0, dim0 - n_constraints)) + w = np.zeros((dim0,)) + self.r = np.zeros((dim0 - n_constraints,)) + self.r_dot = np.zeros((dim0 - n_constraints,)) + + # Extract intervals and generate individual DVRs. + a = [] + b = [] + dim = [] + dvr = [] + block_start = [0] # indices for start of original basis + block_start2 = [0] # indices for start of interface matched basis + + # Create DVRs over each element + for i in range(n_intervals): + a.append(nodes[i]) + b.append(nodes[i + 1]) + dvr.append( + element_class( + n_points[i] - 1, Mapping(a[i], b[i]), symmetrize=symmetrize + ) + ) + + for i in range(n_intervals): + # update dimension + dim.append(n_points[i]) + # update block start indices + # block_start[i] is the start of the original basis, block_start2[i] + # is the start of the interface matched basis. + block_start.append(block_start[i] + dim[i]) + block_start2.append(block_start2[i] + dim[i] - 1) + + # if this is the last interval, we need to add one more point to the interface matched basis + # to account for the last point of the last element. + if i == n_intervals - 1: + block_start2[-1] += 1 + + # Fill the differentiation matrix D and the weights w + # D is the differentiation matrix in the original basis, w are the weights. + D[ + block_start[i] : block_start[i + 1], + block_start[i] : block_start[i + 1], + ] += dvr[i].D1 + w[block_start[i] : block_start[i + 1]] = dvr[i].weights + + # Fill the interface matching matrix R + # R is the matrix that identifies nodes at element interfaces. + R[ + block_start[i] : block_start[i + 1], + block_start2[i] : block_start2[i + 1] + 1, + ] = np.eye(dim[i]) + + # Update nodes x + # x are the nodes of the interface matched basis. + self.r[block_start2[i] : block_start2[i + 1]] = dvr[i].r[ + : block_start2[i + 1] - block_start2[i] + ] + self.r_dot[block_start2[i] : block_start2[i + 1]] = dvr[i].r_dot[ + : block_start2[i + 1] - block_start2[i] + ] + + # Compute second derivative matrix. + temp = D @ R + D2 = -temp.T @ np.diag(w) @ temp # D2 is the second derivative matrix + self.D1 = R.T @ D @ R + self.S = R.T @ np.diag(w) @ R # S is the mass matrix + self.weights = w @ R + S_lu = np.linalg.cholesky( + self.S + ) # Cholesky factorization of the mass matrix + # self.D2 = np.linalg.inv(self.S) @ D2 + self.D2 = np.linalg.solve(S_lu, np.linalg.solve(S_lu.T, D2)) + self.R = R + # self.x = x + self.dvr = dvr + self.edge_indices = block_start2.copy() + self.edge_indices[-1] -= 1 + + +if __name__ == "__main__": + + a = -10 + b = 10 + n_elem = 3 + points_per_elem = 21 + + nodes = np.linspace(a, b, n_elem + 1) + # nodes = np.array([-10, -2, 2, 10]) + n_points = np.ones((n_elem,), dtype=int) * points_per_elem + + femdvr = FEMDVR(nodes, n_points, Linear_map, GaussLegendreLobatto) + + # Get nodes and weights and differentiation matrix. + r = femdvr.r + r_dot = femdvr.r_dot + w = femdvr.weights + D = femdvr.D1 + D2 = femdvr.D2 + ei = femdvr.edge_indices + + H_full = -0.5 * D2 + np.diag(0.5 * r**2) + H = H_full[1:-1, 1:-1] + + E, U = np.linalg.eig(H) + i = np.argsort(E) + E = E[i] + U = U[:, i] + + print(E[:5]) + + import matplotlib.pyplot as plt + + plt.figure() + for i in range(3): + y = np.zeros((len(r),)) + y[1:-1] = U[:, i] + plt.plot(r, y, color=f"C{i}", marker=".", linestyle="-") + plt.plot(r[ei], y[ei], color=f"C{i}", marker="o", linestyle="none") + plt.show() From c551e88d9247c1296708f2c00c5d0fa5007a4a65 Mon Sep 17 00:00:00 2001 From: haakoek Date: Wed, 2 Jul 2025 16:25:46 +0200 Subject: [PATCH 2/3] Add test for radial_poisson with femdvr. --- .../gauss_legendre_lobatto.py | 5 +- .../spherical_coordinates/radial_poisson.py | 4 +- tests/test_radial_poisson.py | 68 ++++++++++++++++++- 3 files changed, 71 insertions(+), 6 deletions(-) diff --git a/grid_methods/pseudospectral_grids/gauss_legendre_lobatto.py b/grid_methods/pseudospectral_grids/gauss_legendre_lobatto.py index 4db6db5..f4325b9 100644 --- a/grid_methods/pseudospectral_grids/gauss_legendre_lobatto.py +++ b/grid_methods/pseudospectral_grids/gauss_legendre_lobatto.py @@ -49,12 +49,13 @@ def __init__(self, N, Mapping, symmetrize=True): self.PN_x = legendre.legval(self.x[1:-1], c) self.PN_x2 = legendre.legval(self.x, c) - self.weights = 2 / (N * (N + 1)) + self.weights = 2 / (N * (N + 1)) * np.ones(N + 1) if not symmetrize: self.weights /= self.PN_x2**2 self.r = Mapping.r_x(self.x) self.r_dot = Mapping.dr_dx(self.x) + self.weights *= self.r_dot for i in range(N + 1): for j in range(N + 1): @@ -82,4 +83,4 @@ def __init__(self, N, Mapping, symmetrize=True): if not symmetrize: self.D1[i, j] *= self.PN_x2[i] / self.PN_x2[j] - self.D2 = np.dot(self.D1, self.D1) + self.D2 = np.dot(self.D1, self.D1) \ No newline at end of file diff --git a/grid_methods/spherical_coordinates/radial_poisson.py b/grid_methods/spherical_coordinates/radial_poisson.py index f423634..a1dcb6a 100644 --- a/grid_methods/spherical_coordinates/radial_poisson.py +++ b/grid_methods/spherical_coordinates/radial_poisson.py @@ -74,9 +74,9 @@ def solve_radial_Poisson_dvr(GLL, n_L): tilde_vL_hom = np.zeros((n_r, n_r)) for a in range(n_r): tilde_vL_hom[:, a] = ( - r_dot * r[a] ** L * w_r[a] / r_max ** (2 * L + 1) + r[a] ** L * w_r[a] / r_max ** (2 * L + 1) ) * r ** (L + 1) tilde_vL[L] = tilde_vL_inhom + tilde_vL_hom - return tilde_vL + return tilde_vL \ No newline at end of file diff --git a/tests/test_radial_poisson.py b/tests/test_radial_poisson.py index 4ae5e75..bc22b01 100644 --- a/tests/test_radial_poisson.py +++ b/tests/test_radial_poisson.py @@ -6,6 +6,7 @@ from grid_methods.spherical_coordinates.radial_poisson import ( solve_radial_Poisson_dvr, ) +from grid_methods.pseudospectral_grids.femdvr import FEMDVR from matplotlib import pyplot as plt from scipy.special import erf @@ -47,7 +48,7 @@ def test_radial_poisson(): r_dot = GLL.r_dot[1:-1] u_1s_gaussian = r * np.exp(-(r**2)) - norm_psi = np.dot(w_r, r_dot * u_1s_gaussian**2) + norm_psi = np.dot(w_r, u_1s_gaussian**2) u_1s_gaussian /= np.sqrt(norm_psi) tilde_V0 = np.einsum( @@ -65,7 +66,7 @@ def test_radial_poisson(): assert np.max(np.abs(tilde_V0 - tilde_V0_exact_1s_gaussian)) < 1e-10 u_1s = r * np.exp(-r) - norm_1s = np.dot(w_r, r_dot * u_1s**2) + norm_1s = np.dot(w_r, u_1s**2) u_1s /= np.sqrt(norm_1s) tilde_V0_1s = np.einsum( "b, ab, b->a", u_1s, tilde_V0_dvr[0], u_1s, optimize=True @@ -77,3 +78,66 @@ def test_radial_poisson(): assert np.linalg.norm(tilde_V0_1s - tilde_V0_exact_1s) < 1e-10 assert np.max(np.abs(tilde_V0_1s - tilde_V0_exact_1s)) < 1e-10 + +def test_radial_poisson_fem(): + + """ + The test solves the radial Poisson equation using the FEM-DVR method. + The test is similar to the one in `test_radial_poisson`, but uses the + `FEMDVR` class to solve the radial Poisson equation. + + Test for a for uniform distribution of nodes and a non-uniform distribution of nodes. + """ + + a = 0 + b = 40 + n_elem = 4 + points_per_elem = 51 + + nodes_list = [np.linspace(a, b, n_elem + 1), np.array([0, 4, 16, 28, 40])] + n_points_list = [ + np.ones((n_elem,), dtype=int) * points_per_elem, + np.array([31, 21, 11, 11]), + ] + + for nodes, n_points in zip(nodes_list, n_points_list): + femdvr = FEMDVR(nodes, n_points, Linear_map, GaussLegendreLobatto) + + tilde_V0_dvr = solve_radial_Poisson_dvr(femdvr, n_L=1) + + r = femdvr.r[1:-1] + w_r = femdvr.weights[1:-1] + r_dot = femdvr.r_dot[1:-1] + + u_1s_gaussian = r * np.exp(-(r**2)) + norm_psi = np.dot(w_r, u_1s_gaussian**2) + u_1s_gaussian /= np.sqrt(norm_psi) + + tilde_V0 = np.einsum( + "b, ab, b->a", + u_1s_gaussian, + tilde_V0_dvr[0], + u_1s_gaussian, + optimize=True, + ) + tilde_V0_exact_1s_gaussian = erf( + np.sqrt(2) * r + ) # The exact solution of the radial Poisson equation to a 1s Gaussian + + assert np.linalg.norm(tilde_V0 - tilde_V0_exact_1s_gaussian) < 1e-10 + assert np.max(np.abs(tilde_V0 - tilde_V0_exact_1s_gaussian)) < 1e-10 + + u_1s = r * np.exp(-r) + norm_1s = np.dot(w_r, u_1s**2) + u_1s /= np.sqrt(norm_1s) + + tilde_V0_1s = np.einsum( + "b, ab, b->a", u_1s, tilde_V0_dvr[0], u_1s, optimize=True + ) + + tilde_V0_exact_1s = 1 - (r + 1) * np.exp( + -2 * r + ) # The exact solution of the radial Poisson equation to a 1s Gaussian + + assert np.linalg.norm(tilde_V0_1s - tilde_V0_exact_1s) < 1e-10 + assert np.max(np.abs(tilde_V0_1s - tilde_V0_exact_1s)) < 1e-10 \ No newline at end of file From 14dd6d1341f6c61839414261fa94c004cbb5fd18 Mon Sep 17 00:00:00 2001 From: haakoek Date: Wed, 2 Jul 2025 16:30:32 +0200 Subject: [PATCH 3/3] Make sure tests are compatible with main. --- tests/test_eigenstates.py | 49 +++++++++++++++++++++++++++++++++------ tests/test_quadrature.py | 6 ++--- 2 files changed, 45 insertions(+), 10 deletions(-) diff --git a/tests/test_eigenstates.py b/tests/test_eigenstates.py index afb5779..7caa33a 100644 --- a/tests/test_eigenstates.py +++ b/tests/test_eigenstates.py @@ -4,7 +4,7 @@ GaussLegendreLobatto, Linear_map, ) - +from grid_methods.pseudospectral_grids.femdvr import FEMDVR def test_particle_in_box(): print() @@ -38,12 +38,12 @@ def test_particle_in_box(L): for k in range(1, N_max): eps_k_approx = eps[k - 1] psi_k_approx = U[:, k - 1] - norm_psi_k_approx = np.dot(w, x_dot * psi_k_approx**2) + norm_psi_k_approx = np.dot(w, psi_k_approx**2) psi_k_approx /= np.sqrt(norm_psi_k_approx) eps_k_exact = k**2 * np.pi**2 / (2 * L**2) psi_k_exact = np.sqrt(2 / L) * np.sin(k * np.pi * x / L) - norm_psi_k_exact = np.dot(w, x_dot * psi_k_exact**2) + norm_psi_k_exact = np.dot(w, psi_k_exact**2) np.testing.assert_allclose( eps_k_approx, eps_k_exact, rtol=0.0, atol=1e-12 ) @@ -99,7 +99,7 @@ def test_harmonic_oscillator(omega): eps_k_exact = omega * (k + 0.5) psi_k_approx = U[:, k] - norm_psi_k_approx = np.dot(w, x_dot * psi_k_approx**2) + norm_psi_k_approx = np.dot(w, psi_k_approx**2) psi_k_approx /= np.sqrt(norm_psi_k_approx) psi_k_exact = ( @@ -177,15 +177,15 @@ def test_hydrogenic_sphc(Z): u_n3_l = u_nl_exact(r, n3, l, Z) u_n1_l_approx = U[:, 0] - norm_u_n1_l_approx = np.dot(w, r_dot * u_n1_l_approx**2) + norm_u_n1_l_approx = np.dot(w, u_n1_l_approx**2) u_n1_l_approx /= np.sqrt(norm_u_n1_l_approx) u_n2_l_approx = U[:, 1] - norm_u_n2_l_approx = np.dot(w, r_dot * u_n2_l_approx**2) + norm_u_n2_l_approx = np.dot(w, u_n2_l_approx**2) u_n2_l_approx /= np.sqrt(norm_u_n2_l_approx) u_n3_l_approx = U[:, 2] - norm_u_n3_l_approx = np.dot(w, r_dot * u_n3_l_approx**2) + norm_u_n3_l_approx = np.dot(w, u_n3_l_approx**2) u_n3_l_approx /= np.sqrt(norm_u_n3_l_approx) np.testing.assert_allclose( @@ -215,6 +215,41 @@ def test_hydrogenic_sphc(Z): test_hydrogenic_sphc(4) test_hydrogenic_sphc(10) +def test_ho_fem(): + + a = -10 + b = 10 + n_elem = 3 + points_per_elem = 31 + + nodes_list = [np.linspace(a, b, n_elem + 1), np.array([-10, -2, 2, 10])] + n_points_list = [ + np.ones((n_elem,), dtype=int) * points_per_elem, + np.array([21, 21, 21]), + ] + + for nodes, n_points in zip(nodes_list, n_points_list): + + femdvr = FEMDVR(nodes, n_points, Linear_map, GaussLegendreLobatto) + + # Get nodes and weights and differentiation matrix. + r = femdvr.r + r_dot = femdvr.r_dot + w = femdvr.weights + D = femdvr.D1 + D2 = femdvr.D2 + ei = femdvr.edge_indices + + H_full = -0.5 * D2 + np.diag(0.5 * r**2) + H = H_full[1:-1, 1:-1] + + E, U = np.linalg.eig(H) + i = np.argsort(E) + E = E[i] + U = U[:, i] + + assert np.allclose(E[:5], [0.5, 1.5, 2.5, 3.5, 4.5], rtol=1e-12) + # if __name__ == "__main__": # # test_harmonic_oscillator() diff --git a/tests/test_quadrature.py b/tests/test_quadrature.py index 1508908..6a0d9cb 100644 --- a/tests/test_quadrature.py +++ b/tests/test_quadrature.py @@ -44,7 +44,7 @@ def test_integrate_polynomials_mapped_interval(a, b): r_dot = GLL.r_dot a_tol = 1e-12 for n in range(1, 2 * N): - integral_xn = np.dot(w, r_dot * GLL.r**n) + integral_xn = np.dot(w, GLL.r**n) np.testing.assert_allclose( np.array([integral_xn]), np.array([(b ** (n + 1) - a ** (n + 1)) / (n + 1)]), @@ -84,7 +84,7 @@ def test_integrate_xk_exp_minus_a_x2(a): for k in range(6): f = x**k * np.exp(-a * x**2) - integral_quadrature = np.dot(w, x_dot * f) + integral_quadrature = np.dot(w, f) integral_exact = 0.5 * a ** (-(k + 1) / 2) * gamma((k + 1) / 2) np.testing.assert_allclose( integral_quadrature, integral_exact, rtol=0.0, atol=1e-12 @@ -99,4 +99,4 @@ def test_integrate_xk_exp_minus_a_x2(a): # if __name__ == "__main__": # #test_integrate_polynomials_mapped_interval() -# test_integrate_xk_exp_minus_a_x2() +# test_integrate_xk_exp_minus_a_x2() \ No newline at end of file