Finite element dvr

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()

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)

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

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()

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()