diff --git a/examples/gaussian_spherical_wave_expansion.py b/examples/gaussian_spherical_wave_expansion.py new file mode 100644 index 0000000..3adcd95 --- /dev/null +++ b/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() diff --git a/examples/plane_wave_spherical_wave_expansion.py b/examples/plane_wave_spherical_wave_expansion.py new file mode 100644 index 0000000..f0c2b66 --- /dev/null +++ b/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() diff --git a/examples/quadratic_potential_spherical_wave_expansion.py b/examples/quadratic_potential_spherical_wave_expansion.py new file mode 100644 index 0000000..da98ad9 --- /dev/null +++ b/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() diff --git a/grid_lib/spherical_coordinates/potentials.py b/grid_lib/spherical_coordinates/potentials.py index 0b84984..a5a9d4e 100644 --- a/grid_lib/spherical_coordinates/potentials.py +++ b/grid_lib/spherical_coordinates/potentials.py @@ -1,7 +1,7 @@ import warnings import numpy as np -from scipy.special import erf +from scipy.special import erf, spherical_in, spherical_jn from .angular_momentum import LM_to_I, number_of_lm_states from .utils import Ylm, cartesian_to_spherical @@ -301,3 +301,266 @@ def clamped_molecular_potential_Poisson( return V_LM.real return V_LM + + +def gaussian_spherical_wave_expansion(r, r0, A, alpha, L_max, M_max=None): + r""" + Compute the spherical-harmonic radial coefficients of a Gaussian + centered at :math:`\mathbf{r}_0`. + + The Gaussian is defined as + + .. math:: + g(\mathbf{r}) = A \, e^{-\alpha |\mathbf{r} - \mathbf{r}_0|^2} + + and is expanded as + + .. math:: + g(\mathbf{r}) = \sum_{L,M} g_{LM}(r) \, Y_{LM}(\Omega), + + where the radial components are given analytically by + + .. math:: + g_{LM}(r) = 4\pi A \, e^{-\alpha(r^2 + R_0^2)} \, + i_L(2\alpha r R_0) \, Y_{LM}^*(\Omega_0), + + with :math:`R_0 = |\mathbf{r}_0|`, :math:`\Omega_0` the solid angle of + :math:`\mathbf{r}_0`, and :math:`i_L` the modified spherical Bessel + function of the first kind. For :math:`R_0 = 0` only the + :math:`(L, M) = (0, 0)` term survives: + :math:`g_{00}(r) = A \sqrt{4\pi} \, e^{-\alpha r^2}`. + + Parameters + ---------- + r : array_like + Radial grid values where the coefficients are evaluated. + r0 : array_like + Cartesian centre of the Gaussian with shape ``(3,)``. + A : float + Amplitude of the Gaussian. + alpha : float + Exponent of the Gaussian (must be positive). + L_max : int + Maximum angular-momentum rank in the expansion. + M_max : int, optional + Maximum magnetic quantum number stored. Defaults to ``L_max``. + + Returns + ------- + np.ndarray + Array of shape ``(n_LM, len(r))`` storing :math:`g_{LM}(r)` in the + repo's ``LM_to_I`` ordering. Returns a real array if the imaginary + part is negligible, otherwise complex. + """ + r = np.asarray(r, dtype=float) + r0 = np.asarray(r0, dtype=float) + + if L_max < 0: + raise ValueError("L_max must be >= 0") + + if M_max is None: + M_max = L_max + if M_max < 0 or M_max > L_max: + raise ValueError("Require 0 <= M_max <= L_max") + + if r.ndim != 1: + raise ValueError("r must be a one-dimensional array") + if r0.shape != (3,): + raise ValueError("r0 must be a Cartesian vector of shape (3,)") + + n_LM = number_of_lm_states(L_max, M_max) + g_LM = np.zeros((n_LM, len(r)), dtype=complex) + + R0 = np.linalg.norm(r0) + + if R0 == 0: + I_00 = LM_to_I(0, 0, L_max, M_max) + g_LM[I_00] = A * np.sqrt(4 * np.pi) * np.exp(-alpha * r**2) + else: + _, theta_0, phi_0 = cartesian_to_spherical(r0) + z = 2 * alpha * r * R0 + envelope = A * 4 * np.pi * np.exp(-alpha * (r**2 + R0**2)) + + for M in range(-M_max, M_max + 1): + for L in range(abs(M), L_max + 1): + I_LM = LM_to_I(L, M, L_max, M_max) + angular = np.conj(Ylm(L, M, theta_0, phi_0)) + radial = spherical_in(L, z) + g_LM[I_LM] = envelope * radial * angular + + if np.allclose(g_LM.imag, 0.0): + return g_LM.real + + return g_LM + + +def plane_wave_spherical_wave_expansion(r, k, L_max, M_max=None, sign=1): + r""" + Compute the spherical-harmonic radial coefficients of a plane wave. + + The plane wave is expanded as + + .. math:: + e^{i\sigma\mathbf{k}\cdot\mathbf{r}} = + \sum_{L,M} f_{LM}(r)\, Y_{LM}(\hat{r}), + + where :math:`\sigma = \pm 1` is controlled by ``sign`` and + + .. math:: + f_{LM}(r) = 4\pi\,(i\sigma)^L\, j_L(kr)\, Y_{LM}^*(\hat{k}), + + with :math:`k = |\mathbf{k}|`, :math:`\hat{k}` the unit vector of + :math:`\mathbf{k}`, and :math:`j_L` the spherical Bessel function of + the first kind. + + Parameters + ---------- + r : array_like + Radial grid values where the coefficients are evaluated. + k : array_like + Cartesian wave vector with shape ``(3,)``. + L_max : int + Maximum angular-momentum rank in the expansion. + M_max : int, optional + Maximum magnetic quantum number stored. Defaults to ``L_max``. + sign : {+1, -1}, optional + Sign :math:`\sigma` in the exponent. Use ``+1`` (default) for + :math:`e^{i\mathbf{k}\cdot\mathbf{r}}` and ``-1`` for + :math:`e^{-i\mathbf{k}\cdot\mathbf{r}}`. + + Returns + ------- + np.ndarray + Array of shape ``(n_LM, len(r))`` storing :math:`f_{LM}(r)` in the + repo's ``LM_to_I`` ordering. Returns a real array if the imaginary + part is negligible, otherwise complex. + """ + r = np.asarray(r, dtype=float) + k = np.asarray(k, dtype=float) + + if sign not in (1, -1): + raise ValueError("sign must be +1 or -1") + if L_max < 0: + raise ValueError("L_max must be >= 0") + + if M_max is None: + M_max = L_max + if M_max < 0 or M_max > L_max: + raise ValueError("Require 0 <= M_max <= L_max") + + if r.ndim != 1: + raise ValueError("r must be a one-dimensional array") + if k.shape != (3,): + raise ValueError("k must be a Cartesian vector of shape (3,)") + + K = np.linalg.norm(k) + + n_LM = number_of_lm_states(L_max, M_max) + f_LM = np.zeros((n_LM, len(r)), dtype=complex) + + if K == 0: + I_00 = LM_to_I(0, 0, L_max, M_max) + f_LM[I_00] = np.sqrt(4 * np.pi) * np.ones(len(r)) + else: + _, theta_k, phi_k = cartesian_to_spherical(k) + kr = K * r + + for M in range(-M_max, M_max + 1): + for L in range(abs(M), L_max + 1): + I_LM = LM_to_I(L, M, L_max, M_max) + angular = np.conj(Ylm(L, M, theta_k, phi_k)) + radial = spherical_jn(L, kr) + f_LM[I_LM] = 4 * np.pi * (1j * sign) ** L * radial * angular + + if np.allclose(f_LM.imag, 0.0): + return f_LM.real + + return f_LM + + +def quadratic_potential_spherical_wave_expansion( + r, r0, omega, L_max, M_max=None +): + r""" + Compute the spherical-harmonic radial coefficients of a quadratic + potential centered at :math:`\mathbf{r}_0`. + + The potential is defined as + + .. math:: + V(\mathbf{r};\mathbf{r}_0) = \tfrac{1}{2}\omega^2 + |\mathbf{r} - \mathbf{r}_0|^2 + + and is expanded as + + .. math:: + V(\mathbf{r};\mathbf{r}_0) = \sum_{L,M} V_{LM}(r)\, Y_{LM}(\hat{r}). + + The expansion is exact and contains only :math:`L = 0` and :math:`L = 1` + terms: + + .. math:: + V_{00}(r) &= \tfrac{1}{2}\omega^2\,(r^2 + R_0^2)\,\sqrt{4\pi}, \\ + V_{1M}(r) &= -\omega^2\, r R_0\,\frac{4\pi}{3}\,Y_{1M}^*(\hat{r}_0), + + where :math:`R_0 = |\mathbf{r}_0|`. For :math:`R_0 = 0` only + :math:`V_{00}` survives. + + Parameters + ---------- + r : array_like + Radial grid values where the coefficients are evaluated. + r0 : array_like + Cartesian centre of the potential with shape ``(3,)``. + omega : float + Frequency (force constant is :math:`\omega^2`). + L_max : int + Maximum angular-momentum rank stored. Must be ``>= 1`` when + :math:`R_0 \neq 0` to capture the :math:`L=1` terms. + M_max : int, optional + Maximum magnetic quantum number stored. Defaults to ``L_max``. + + Returns + ------- + np.ndarray + Array of shape ``(n_LM, len(r))`` storing :math:`V_{LM}(r)` in the + repo's ``LM_to_I`` ordering. Returns a real array if the imaginary + part is negligible, otherwise complex. + """ + r = np.asarray(r, dtype=float) + r0 = np.asarray(r0, dtype=float) + + if L_max < 0: + raise ValueError("L_max must be >= 0") + + if M_max is None: + M_max = L_max + if M_max < 0 or M_max > L_max: + raise ValueError("Require 0 <= M_max <= L_max") + + if r.ndim != 1: + raise ValueError("r must be a one-dimensional array") + if r0.shape != (3,): + raise ValueError("r0 must be a Cartesian vector of shape (3,)") + + n_LM = number_of_lm_states(L_max, M_max) + V_LM = np.zeros((n_LM, len(r)), dtype=complex) + + R0 = np.linalg.norm(r0) + + # L = 0 term: always present + I_00 = LM_to_I(0, 0, L_max, M_max) + V_LM[I_00] = 0.5 * omega**2 * (r**2 + R0**2) * np.sqrt(4 * np.pi) + + # L = 1 terms: only when centre is off-origin and L_max >= 1 + if R0 != 0 and L_max >= 1: + _, theta_0, phi_0 = cartesian_to_spherical(r0) + for M in range(max(-1, -M_max), min(1, M_max) + 1): + I_1M = LM_to_I(1, M, L_max, M_max) + angular = np.conj(Ylm(1, M, theta_0, phi_0)) + V_LM[I_1M] = -omega**2 * r * R0 * (4 * np.pi / 3) * angular + + if np.allclose(V_LM.imag, 0.0): + return V_LM.real + + return V_LM