Spherical wave expansions: quadratic, gaussians and plane waves.

examples/gaussian_spherical_wave_expansion.py

@@ -0,0 +1,87 @@
+import numpy as np
+from matplotlib import pyplot as plt
+
+from grid_lib.pseudospectral_grids.gauss_legendre_lobatto import (
+    GaussLegendreLobatto,
+    Linear_map,
+)
+from grid_lib.spherical_coordinates.angular_momentum import LM_to_I
+from grid_lib.spherical_coordinates.potentials import (
+    gaussian_spherical_wave_expansion,
+)
+
+def g00(r, r0, A, alpha):
+    R0 = np.linalg.norm(r0)
+    prefactor = np.sqrt(np.pi) * A / (2 * alpha * r * R0)
+    return prefactor * (
+        np.exp(-alpha * (r - R0) ** 2) - np.exp(-alpha * (r + R0) ** 2)
+    )
+
+
+# Gaussian parameters
+A = 1.0
+alpha = 1.0
+z0 = 1.5  # centre on the z-axis
+
+r0 = np.array([0.0, 0.0, z0])
+
+L_max = 4
+M_max = 0  # z-axis centre: only M = 0 terms are non-zero
+
+# Radial grid
+N = 150
+r_max = 10.0
+gll = GaussLegendreLobatto(N, Linear_map(r_min=0.0, r_max=r_max))
+r = gll.r[1:-1]
+
+# Compute radial components
+g_LM = gaussian_spherical_wave_expansion(r, r0, A, alpha, L_max, M_max)
+
+# --- Superposition of two Gaussians placed symmetrically on the z-axis ---
+r0_plus = np.array([0.0, 0.0, z0])
+r0_minus = np.array([0.0, 0.0, -z0])
+
+g_LM_plus = gaussian_spherical_wave_expansion(r, r0_plus, A, alpha, L_max, M_max)
+g_LM_minus = gaussian_spherical_wave_expansion(r, r0_minus, A, alpha, L_max, M_max)
+g_LM_super = g_LM_plus + g_LM_minus
+
+fig, (ax, ax2) = plt.subplots(1, 2, figsize=(13, 4), sharey=False)
+
+# Top panel: single Gaussian
+for L in range(L_max + 1):
+    I = LM_to_I(L, 0, L_max, M_max)
+    component = g_LM[I]
+    if np.allclose(component, 0.0):
+        continue
+    ax.plot(r, component, label=rf"$g_{{L={L},M=0}}(r)$")
+
+ax.plot(r, g00(r, r0, A, alpha), "k--", label=r"$g_{00}(r)$ (exact)")
+ax.set_xlabel(r"$r$ (a.u.)")
+ax.set_ylabel(r"$g_{LM}(r)$")
+ax.set_title(
+    rf"Single Gaussian: $e^{{-\alpha|\mathbf{{r}}-\mathbf{{r}}_0|^2}}$, "
+    rf"$\mathbf{{r}}_0 = (0,0,{z0})$, $\alpha={alpha}$"
+)
+ax.legend()
+ax.set_xlim(0, r_max)
+
+# Right panel: symmetric superposition
+for L in range(L_max + 1):
+    I = LM_to_I(L, 0, L_max, M_max)
+    component = g_LM_super[I]
+    if np.allclose(component, 0.0):
+        continue
+    ax2.plot(r, component, label=rf"$g_{{L={L},M=0}}(r)$")
+
+ax2.set_xlabel(r"$r$ (a.u.)")
+ax2.set_ylabel(r"$g_{LM}(r)$")
+ax2.set_xlim(0, r_max)
+ax2.set_title(
+    rf"Superposition: $e^{{-\alpha|\mathbf{{r}}-(0,0,z_0)|^2}} + "
+    rf"e^{{-\alpha|\mathbf{{r}}+(0,0,z_0)|^2}}$, "
+    rf"$z_0={z0}$, $\alpha={alpha}$"
+)
+ax2.legend()
+
+plt.tight_layout()
+plt.show()

examples/plane_wave_spherical_wave_expansion.py

@@ -0,0 +1,82 @@
+import numpy as np
+from matplotlib import pyplot as plt
+
+from grid_lib.pseudospectral_grids.gauss_legendre_lobatto import (
+    GaussLegendreLobatto,
+    Linear_map,
+)
+from grid_lib.spherical_coordinates.angular_momentum import LM_to_I
+from grid_lib.spherical_coordinates.potentials import (
+    plane_wave_spherical_wave_expansion,
+)
+
+# Plane wave parameters
+k = np.array([0.0, 0.0, 2.0])  # wave vector along z-axis
+
+L_max = 3
+M_max = 0  # k along z: only M = 0 terms contribute
+
+# Radial grid
+N = 150
+r_max = 10.0
+gll = GaussLegendreLobatto(N, Linear_map(r_min=0.0, r_max=r_max))
+r = gll.r[1:-1]
+
+# Compute radial components for e^{+ik.r} and e^{-ik.r}
+f_LM_plus = plane_wave_spherical_wave_expansion(r, k, L_max, M_max, sign=+1)
+f_LM_minus = plane_wave_spherical_wave_expansion(r, k, L_max, M_max, sign=-1)
+
+K = np.linalg.norm(k)
+
+fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 4), sharey=False)
+
+# Left panel: e^{+ik.r}
+for L in range(L_max + 1):
+    I = LM_to_I(L, 0, L_max, M_max)
+    data = f_LM_plus[I]
+    if np.isrealobj(data):
+        if np.allclose(data, 0.0):
+            continue
+        ax1.plot(r, data, label=rf"$f_{{L={L},M=0}}(r)$")
+    else:
+        label = rf"$f_{{L={L},M=0}}(r)$"
+        if not np.allclose(data.real, 0.0):
+            ax1.plot(r, data.real, label=label + r" $\mathrm{Re}$")
+        if not np.allclose(data.imag, 0.0):
+            ax1.plot(r, data.imag, "--", label=label + r" $\mathrm{Im}$")
+
+ax1.set_xlabel(r"$r$ (a.u.)")
+ax1.set_ylabel(r"$f_{LM}(r)$")
+ax1.set_title(
+    rf"$e^{{i\mathbf{{k}}\cdot\mathbf{{r}}}}$, "
+    rf"$\mathbf{{k}} = (0,0,{K:.1f})$"
+)
+ax1.legend()
+ax1.set_xlim(0, r_max)
+
+# Right panel: e^{-ik.r}
+for L in range(L_max + 1):
+    I = LM_to_I(L, 0, L_max, M_max)
+    data = f_LM_minus[I]
+    if np.isrealobj(data):
+        if np.allclose(data, 0.0):
+            continue
+        ax2.plot(r, data, label=rf"$f_{{L={L},M=0}}(r)$")
+    else:
+        label = rf"$f_{{L={L},M=0}}(r)$"
+        if not np.allclose(data.real, 0.0):
+            ax2.plot(r, data.real, label=label + r" $\mathrm{Re}$")
+        if not np.allclose(data.imag, 0.0):
+            ax2.plot(r, data.imag, "--", label=label + r" $\mathrm{Im}$")
+
+ax2.set_xlabel(r"$r$ (a.u.)")
+ax2.set_ylabel(r"$f_{LM}(r)$")
+ax2.set_title(
+    rf"$e^{{-i\mathbf{{k}}\cdot\mathbf{{r}}}}$, "
+    rf"$\mathbf{{k}} = (0,0,{K:.1f})$"
+)
+ax2.legend()
+ax2.set_xlim(0, r_max)
+
+plt.tight_layout()
+plt.show()

examples/quadratic_potential_spherical_wave_expansion.py

@@ -0,0 +1,81 @@
+import numpy as np
+from matplotlib import pyplot as plt
+
+from grid_lib.pseudospectral_grids.gauss_legendre_lobatto import (
+    GaussLegendreLobatto,
+    Linear_map,
+)
+from grid_lib.spherical_coordinates.angular_momentum import LM_to_I
+from grid_lib.spherical_coordinates.potentials import (
+    quadratic_potential_spherical_wave_expansion,
+)
+
+# Potential parameters
+omega = 1.0
+z0 = 1.5  # centre on the z-axis
+
+r0 = np.array([0.0, 0.0, z0])
+
+L_max = 1  # exact: only L = 0 and L = 1 are non-zero
+M_max = 0  # z-axis centre: only M = 0 terms contribute
+
+# Radial grid
+N = 150
+r_max = 10.0
+gll = GaussLegendreLobatto(N, Linear_map(r_min=0.0, r_max=r_max))
+r = gll.r[1:-1]
+
+# Compute radial components for single potential
+V_LM = quadratic_potential_spherical_wave_expansion(r, r0, omega, L_max, M_max)
+
+# Superposition: V(r; r0_+) + V(r; r0_-)  -- only even-L terms survive
+r0_plus = np.array([0.0, 0.0, z0])
+r0_minus = np.array([0.0, 0.0, -z0])
+
+V_LM_plus = quadratic_potential_spherical_wave_expansion(
+    r, r0_plus, omega, L_max, M_max
+)
+V_LM_minus = quadratic_potential_spherical_wave_expansion(
+    r, r0_minus, omega, L_max, M_max
+)
+V_LM_super = V_LM_plus + V_LM_minus
+
+fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 4))
+
+# Left panel: single quadratic potential
+for L in range(L_max + 1):
+    I = LM_to_I(L, 0, L_max, M_max)
+    component = V_LM[I]
+    if np.allclose(component, 0.0):
+        continue
+    ax1.plot(r, component, label=rf"$V_{{L={L},M=0}}(r)$")
+
+ax1.set_xlabel(r"$r$ (a.u.)")
+ax1.set_ylabel(r"$V_{LM}(r)$")
+ax1.set_title(
+    rf"Single quadratic: $\frac{{1}}{{2}}\omega^2|\mathbf{{r}}-\mathbf{{r}}_0|^2$, "
+    rf"$\mathbf{{r}}_0=(0,0,{z0})$, $\omega={omega}$"
+)
+ax1.legend()
+ax1.set_xlim(0, r_max)
+
+# Right panel: symmetric superposition
+for L in range(L_max + 1):
+    I = LM_to_I(L, 0, L_max, M_max)
+    component = V_LM_super[I]
+    if np.allclose(component, 0.0):
+        continue
+    ax2.plot(r, component, label=rf"$V_{{L={L},M=0}}(r)$")
+
+ax2.set_xlabel(r"$r$ (a.u.)")
+ax2.set_ylabel(r"$V_{LM}(r)$")
+ax2.set_title(
+    rf"Superposition: $\frac{{1}}{{2}}\omega^2(|\mathbf{{r}}-\mathbf{{r}}_+|^2"
+    rf"+|\mathbf{{r}}-\mathbf{{r}}_-|^2)$, "
+    rf"$z_0=\pm{z0}$, $\omega={omega}$"
+)
+ax2.legend()
+ax2.set_xlim(0, r_max)
+
+plt.tight_layout()
+plt.show()