Expansion Identities.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook checks the expansion identities given in Section 2 and Appendix A in the paper.\n",
    "\n",
    "Notation:\n",
    "\n",
    "* $t$ is the target, $c$ is the center, $s$ is the source\n",
    "* $p$ is the expansion order\n",
    "* $\\cos \\theta = \\frac{(t - c) \\cdot (s - c)}{(|t - c| |s - c|}$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Utilities\n",
    "\n",
    "Define a utility function that checks for convergence of an identity.\n",
    "\n",
    "The kernel is a function of the form $K(s, t)$.\n",
    "\n",
    "The identity to be checked is a function of the form $Id(s, c, t, p)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import numpy.linalg as la\n",
    "import itertools\n",
    "\n",
    "\n",
    "TESTS = 5\n",
    "DIM = 3\n",
    "ORDERS = np.array([3, 4, 5, 6, 7])\n",
    "\n",
    "\n",
    "def grouper(iterable, n, fillvalue=None):\n",
    "    \"Collect data into fixed-length chunks or blocks\"\n",
    "    # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx\"\n",
    "    args = [iter(iterable)] * n\n",
    "    return itertools.zip_longest(*args, fillvalue=fillvalue)\n",
    "\n",
    "\n",
    "def verify_identity(knl, identity):\n",
    "    np.random.seed(0)\n",
    "    samples = 2 * np.random.random((3*TESTS, DIM)) - 1\n",
    "\n",
    "    for s, c, t in grouper(samples, 3):\n",
    "        if la.norm(t - c) > la.norm(s - c):\n",
    "            s, t = t, s\n",
    "\n",
    "        errors = []\n",
    "\n",
    "        true_val = knl(s, t)\n",
    "\n",
    "        for order in ORDERS:\n",
    "            errors.append(np.max(np.abs(true_val - identity(s, c, t, order))))\n",
    "\n",
    "        conv_factor = np.exp(np.polyfit(1 + ORDERS, np.log(errors), 1)[0])\n",
    "        assert conv_factor < min(1, 1.3 * la.norm(t - c) / la.norm(s - c)), conv_factor\n",
    "\n",
    "    print(\"Verified!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And an auxiliary function to compute Legendre polynomial values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def legvals(x, n):\n",
    "    pjm2 = 1\n",
    "    pjm1 = x\n",
    "\n",
    "    vals = np.zeros(1 + n)\n",
    "    derivs = np.zeros(1 + n)\n",
    "\n",
    "    vals[0] = 1\n",
    "\n",
    "    derjm2 = 0\n",
    "    derjm1 = 1\n",
    "\n",
    "    if n == 0:\n",
    "        return vals, derivs\n",
    "\n",
    "    vals[1] = x\n",
    "    derivs[1] = 1\n",
    "\n",
    "    if n == 1:\n",
    "        return vals, derivs\n",
    "\n",
    "    for j in range(2, n + 1):\n",
    "        pj = ( (2*j-1)*x*pjm1-(j-1)*pjm2 ) / j\n",
    "        vals[j] = pj\n",
    "\n",
    "        derj = (2*j-1)*(pjm1+x*derjm1)-(j-1)*derjm2\n",
    "        derj = derj / j\n",
    "        derivs[j] = derj\n",
    "\n",
    "        derjm2 = derjm1\n",
    "        derjm1 = derj\n",
    "\n",
    "        pjm2 = pjm1\n",
    "        pjm1 = pj\n",
    "\n",
    "    return vals, derivs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Spherical Harmonics\n",
    "The spherical harmonic addition theorem states that\n",
    "$$ P_n(\\cos \\theta) = \\frac{4 \\pi}{2n+1} \\sum_{m=-n}^n Y^m_n(\\theta_x, \\phi_x) \\overline{Y^{m}_n(\\theta_y, \\phi_y)}. $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cart2polar(x):\n",
    "    # Uses NumPy convention\n",
    "    r = la.norm(x)\n",
    "    theta = np.arctan2(x[1], x[0]) \n",
    "    phi = np.arccos(x[2] / r)\n",
    "    return (r, theta, phi)\n",
    "\n",
    "np.random.seed(0)\n",
    "x, y = np.random.random((2, 3))\n",
    "\n",
    "x /= la.norm(x)\n",
    "y /= la.norm(y)\n",
    "\n",
    "cos_theta = np.dot(x, y)\n",
    "_, theta_x, phi_x = cart2polar(x)\n",
    "_, theta_y, phi_y = cart2polar(y)\n",
    "\n",
    "n = 5\n",
    "\n",
    "lvals, _ = legvals(cos_theta, n)\n",
    "\n",
    "from scipy.special import sph_harm\n",
    "\n",
    "# This uses NumPy's spherical harmonics, which are different from those defined in\n",
    "# the paper, but still satisfy the addition theorem.\n",
    "sph_harm_x = sph_harm(np.arange(-n, 1+n), n, theta_x, phi_x)\n",
    "sph_harm_y = sph_harm(np.arange(-n, 1+n), n, theta_y, phi_y)\n",
    "\n",
    "assert np.allclose(np.vdot(sph_harm_x, sph_harm_y), lvals[-1] * (2*n+1) / (4*np.pi))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Laplace Identities\n",
    "\n",
    "The Laplace Green's function $\\mathcal{G}$:\n",
    "\n",
    "$$ \\mathcal{G}(t, s) = (4 \\pi)^{-1} |t - s|^{-1}. $$\n",
    "\n",
    "The target-specific expansion for the Laplace Green's function:\n",
    "\n",
    "$$ \\mathcal{G}(t, s) = (4 \\pi)^{-1} \\sum_{n=0}^\\infty \\frac{|t-c|^n}{|s-c|^{n+1}} P_n(\\cos \\theta). $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ts_l3d(s, c, t, p):\n",
    "    r = la.norm(t - c)\n",
    "    rho = la.norm(s - c)\n",
    "    cos_theta = np.dot(s - c, t - c) / (r * rho)\n",
    "    lvals, _ = legvals(cos_theta, p)\n",
    "\n",
    "    return sum(\n",
    "        (4 * np.pi) ** -1\n",
    "        * r**k * rho**-(k+1)\n",
    "        * lvals[k]\n",
    "        for k in range(0, 1+p))\n",
    "\n",
    "def k_l3d(s, t):\n",
    "    return 1 / (4 * np.pi * la.norm(s - t))\n",
    "\n",
    "verify_identity(k_l3d, ts_l3d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The $t$-gradient of the Laplace Green's function:\n",
    "    \n",
    "$$ \\nabla_t \\mathcal{G}(t, s) = -(4 \\pi)^{-1} \\frac{(t-s)}{|t-s|^{3}}. $$\n",
    "\n",
    "The target-specific expansion for the $t$-gradient:\n",
    "\n",
    "$$ \\nabla_t \\mathcal{G}(t, s) = (4 \\pi)^{-1} \\sum_{n=1}^\\infty\n",
    "\\frac{|t-c|^{n-1}}{|s-c|^{n+1}}\n",
    "\\left(\n",
    "n \\frac{t - c}{|t-c|}\n",
    "P_n(\\cos \\theta)\n",
    "+\n",
    "\\left[\n",
    "\\frac{s-c}{|s-c|}\n",
    "- \\frac{t-c}{|t-c|} \\cos(\\theta)\n",
    "\\right]\n",
    "P_n'(\\cos \\theta) \\right)\n",
    ". $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ts_tgrad_l3d(s, c, t, p):\n",
    "    r = la.norm(t - c)\n",
    "    rho = la.norm(s - c)\n",
    "    cos_theta = np.dot(s - c, t - c) / (r * rho)\n",
    "    lvals, derivs = legvals(cos_theta, p)\n",
    "    \n",
    "    return sum(\n",
    "        (4 * np.pi) ** -1 * r**(k-1) * rho**-(k+1)\n",
    "        * (\n",
    "            k * (t - c) / r * lvals[k]\n",
    "            + ((s - c) / rho - ((t - c) / r) * cos_theta) * derivs[k]\n",
    "        ) for k in range(1, 1+p))\n",
    "\n",
    "\n",
    "def k_tgrad_l3d(s, t):\n",
    "    val = -(t-s) / (4*np.pi * la.norm(t-s)**3)\n",
    "    return val\n",
    "\n",
    "verify_identity(k_tgrad_l3d, ts_tgrad_l3d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The $s$-gradient of the Laplace Green's function:\n",
    "    \n",
    "$$ \\nabla_s \\mathcal{G}(t, s) = (4 \\pi)^{-1} \\frac{(t-s)}{|t-s|^{3}}. $$\n",
    "\n",
    "The target-specific expansion for the $s$-gradient:\n",
    "\n",
    "$$ \\nabla_s \\mathcal{G}(t, s) = (4 \\pi)^{-1}\n",
    "\\sum_{n=0}^\\infty\n",
    "\\frac{|t-c|^{n}}{|s-c|^{n+2}}\n",
    "\\left(\n",
    "-(n+1) \\frac{s - c}{|s - c|}\n",
    "P_n(\\cos \\theta)\n",
    "+\n",
    "\\left[\n",
    "\\frac{t-c}{|t-c|}\n",
    "- \\frac{s-c}{|s-c|} \\cos(\\theta)\n",
    "\\right]\n",
    "P_n'(\\cos \\theta) \\right)\n",
    ". $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ts_sgrad_l3d(s, c, t, p):\n",
    "    r = la.norm(t - c)\n",
    "    rho = la.norm(s - c)\n",
    "    cos_theta = np.dot(s - c, t - c) / (r * rho)\n",
    "    lvals, derivs = legvals(cos_theta, p)\n",
    "    \n",
    "    return sum(\n",
    "        (4 * np.pi) ** -1 * r**k * rho**-(k+2)\n",
    "        * (\n",
    "            -(k + 1) * (s - c) / rho * lvals[k]\n",
    "            +((t - c) / r - ((s - c) / rho) * cos_theta) * derivs[k]\n",
    "        ) for k in range(0, 1 + p))\n",
    "\n",
    "\n",
    "def k_sgrad_l3d(s, t):\n",
    "    val = (t-s) / (4*np.pi * la.norm(t-s)**3)\n",
    "    return val\n",
    "\n",
    "verify_identity(k_sgrad_l3d, ts_sgrad_l3d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " # Helmholtz Case\n",
    " \n",
    " The Green's function for the Helmholtz equation:\n",
    " $$ \\mathcal{G}_k(t, s) = \\frac{e^{ik|t-s|}}{4 \\pi |t - s|}. $$\n",
    " \n",
    " The target-specific expansion for the Green's function:\n",
    " $$ \\mathcal{G}_k(t, s) = (4 \\pi)^{-1} ik \\sum_{n=0}^\\infty (2n+1) j_n(k|t-c|) h_n(k|s-c|) P_n(\\cos \\theta). $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HELMHOLTZ_K = 3\n",
    "\n",
    "from scipy.special import spherical_jn as jn, spherical_yn as yn\n",
    "\n",
    "def hn(n, x, derivative=False):\n",
    "    return jn(n, x, derivative) + 1j * yn(n, x, derivative)\n",
    "\n",
    "def ts_h3d(s, c, t, p):\n",
    "    r = la.norm(t - c)\n",
    "    rho = la.norm(s - c)\n",
    "    cos_theta = np.dot(s - c, t - c) / (r * rho)\n",
    "    lvals, _ = legvals(cos_theta, p)\n",
    "\n",
    "    return sum(\n",
    "        (4 * np.pi) ** -1\n",
    "        * 1j * HELMHOLTZ_K\n",
    "        * (2*k + 1)\n",
    "        * jn(k, HELMHOLTZ_K * r)\n",
    "        * hn(k, HELMHOLTZ_K * rho)\n",
    "        * lvals[k]\n",
    "        for k in range(0, 1+p))\n",
    "\n",
    "def k_h3d(s, t):\n",
    "    x = la.norm(t - s)\n",
    "    return np.exp(1j * HELMHOLTZ_K * x) / (4 * np.pi * x)\n",
    "\n",
    "verify_identity(k_h3d, ts_h3d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The $t$-gradient of the Green's function:\n",
    "$$\\nabla_t \\mathcal{G}_k(t, s) = \n",
    "(4\\pi)^{-1} e^{ik|t-s|} \\left( \\frac{ik(t-s)}{|t-s|^2} - \\frac{(t-s)}{|t-s|^3} \\right).\n",
    "$$\n",
    "\n",
    "The target-specific expansion for the $t$-gradient:\n",
    "$$\\nabla_t \\mathcal{G}_k(t, s) = (4 \\pi)^{-1} ik \\sum_{n=0}^\\infty (2n+1) \\frac{h_n(k|s-c|)}{|t-c|} \\left( k(t-c)j_n'(k|t-c|)P_n(\\cos \\theta) + \\left[ \\frac{s-c}{|s-c|} - \\frac{t-c}{|t-c|} \\cos \\theta \\right] j_n(k|t-c|) P_n'(\\cos \\theta) \\right).\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ts_tgrad_h3d(s, c, t, p):\n",
    "    r = la.norm(t - c)\n",
    "    rho = la.norm(s - c)\n",
    "    cos_theta = np.dot(s - c, t - c) / (r * rho)\n",
    "    lvals, derivs = legvals(cos_theta, p)\n",
    "\n",
    "    return sum(\n",
    "        (4 * np.pi) ** -1\n",
    "        * 1j * HELMHOLTZ_K\n",
    "        * (2*k + 1)\n",
    "        * hn(k, HELMHOLTZ_K * rho) / r\n",
    "        * (\n",
    "            HELMHOLTZ_K * (t - c) * jn(k, HELMHOLTZ_K * r, True) * lvals[k]\n",
    "            + (\n",
    "                ((s - c) / rho - (t-c) * cos_theta / r)\n",
    "                * jn(k, HELMHOLTZ_K * r)) * derivs[k]\n",
    "        ) for k in range(0, 1+p))\n",
    "\n",
    "def k_tgrad_h3d(s, t):\n",
    "    x = la.norm(t - s)\n",
    "    return (\n",
    "        1/(4*np.pi)\n",
    "        * np.exp(1j * HELMHOLTZ_K * x)\n",
    "        * (1j * HELMHOLTZ_K * (t - s) / x**2 - (t - s)/x**3))\n",
    "\n",
    "verify_identity(k_tgrad_h3d, ts_tgrad_h3d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The $s$-gradient of the Green's function:\n",
    "$$\\nabla_s \\mathcal{G}_k(t, s) = \n",
    "(4\\pi)^{-1} e^{ik|t-s|} \\left(\\frac{(t-s)}{|t-s|^3} - \\frac{ik(t-s)}{|t-s|^2}  \\right).\n",
    "$$\n",
    "\n",
    "The target-specific expansion for the $s$-gradient:\n",
    "$$\\nabla_s \\mathcal{G}_k(t, s) = (4 \\pi)^{-1} ik \\sum_{n=0}^\\infty (2n+1) \\frac{j_n(k|t-c|)}{|s-c|} \\left( k(s-c)h_n'(k|s-c|)P_n(\\cos \\theta) + \\left[ \\frac{t-c}{|t-c|} - \\frac{s-c}{|s-c|} \\cos \\theta \\right] h_n(k|s-c|) P_n'(\\cos \\theta) \\right).\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ts_sgrad_h3d(s, c, t, p):\n",
    "    r = la.norm(t - c)\n",
    "    rho = la.norm(s - c)\n",
    "    cos_theta = np.dot(s - c, t - c) / (r * rho)\n",
    "    lvals, derivs = legvals(cos_theta, p)\n",
    "\n",
    "    return sum(\n",
    "        (4 * np.pi) ** -1\n",
    "        * 1j * HELMHOLTZ_K\n",
    "        * (2*k + 1)\n",
    "        * jn(k, HELMHOLTZ_K * r) / rho\n",
    "        * (\n",
    "            HELMHOLTZ_K * (s - c) * hn(k, HELMHOLTZ_K * rho, True) * lvals[k]\n",
    "            + (\n",
    "                ((t - c) / r - (s-c) * cos_theta / rho)\n",
    "                * hn(k, HELMHOLTZ_K * rho)) * derivs[k]\n",
    "        ) for k in range(0, 1+p))\n",
    "\n",
    "def k_sgrad_h3d(s, t):\n",
    "    x = la.norm(t - s)\n",
    "    return (\n",
    "        1/(4*np.pi)\n",
    "        * np.exp(1j * HELMHOLTZ_K * x)\n",
    "        * (((t - s)/x**3) - 1j * HELMHOLTZ_K * (t - s) / x**2))\n",
    "\n",
    "verify_identity(k_sgrad_h3d, ts_sgrad_h3d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}