PDE-reduction and translations.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import numpy.linalg as la\n",
    "\n",
    "import pyopencl as cl\n",
    "from pytools import add_tuples\n",
    "\n",
    "import sumpy.toys as t\n",
    "from sumpy.expansion.local import VolumeTaylorLocalExpansion\n",
    "from sumpy.expansion.multipole import VolumeTaylorMultipoleExpansion\n",
    "from sumpy.kernel import HelmholtzKernel, LaplaceKernel, YukawaKernel  # noqa: F401\n",
    "\n",
    "\n",
    "rng = np.random.default_rng(seed=42)\n",
    "order = 4\n",
    "\n",
    "if 0:\n",
    "    knl = LaplaceKernel(2)\n",
    "    pde = [(1, (2, 0)), (1, (0, 2))]\n",
    "    extra_kernel_kwargs = {}\n",
    "\n",
    "else:\n",
    "    helm_k = 1.2\n",
    "    knl = HelmholtzKernel(2)\n",
    "    extra_kernel_kwargs = {\"k\": helm_k}\n",
    "\n",
    "    pde = [(1, (2, 0)), (1, (0, 2)), (helm_k**2, (0, 0))]\n",
    "\n",
    "mpole_expn = VolumeTaylorMultipoleExpansion(knl, order)\n",
    "local_expn = VolumeTaylorLocalExpansion(knl, order)\n",
    "\n",
    "cl_ctx = cl.create_some_context(answers=[\"port\"])\n",
    "\n",
    "tctx = t.ToyContext(\n",
    "    cl_ctx,\n",
    "    knl,\n",
    "    mpole_expn_class=type(mpole_expn),\n",
    "    local_expn_class=type(local_expn),\n",
    "    extra_kernel_kwargs=extra_kernel_kwargs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pt_src = t.PointSources(tctx, rng.uniform(-0.5, 0.5, size=(2, 50)), np.ones(50))\n",
    "\n",
    "mexp = t.multipole_expand(pt_src, [0, 0], order)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mexp.coeffs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_pde_mat(expn, pde):\n",
    "    coeff_ids = expn.get_coefficient_identifiers()\n",
    "    id_to_index = expn._storage_loc_dict\n",
    "\n",
    "    # FIXME: specific to scalar PDEs\n",
    "    pde_mat = np.zeros((len(coeff_ids), len(coeff_ids)))\n",
    "\n",
    "    row = 0\n",
    "    for base_coeff_id in coeff_ids:\n",
    "        valid = True\n",
    "\n",
    "        for pde_coeff, coeff_id_offset in pde:\n",
    "            other_coeff = add_tuples(base_coeff_id, coeff_id_offset)\n",
    "            if other_coeff not in id_to_index:\n",
    "                valid = False\n",
    "                break\n",
    "\n",
    "            pde_mat[row, id_to_index[other_coeff]] = pde_coeff\n",
    "\n",
    "        if valid:\n",
    "            row += 1\n",
    "        else:\n",
    "            pde_mat[row] = 0\n",
    "\n",
    "    return pde_mat[:row]\n",
    "\n",
    "\n",
    "pde_mat = build_pde_mat(mpole_expn, pde)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_nullspace(mat, tol=1e-10):\n",
    "    _u, sig, vt = la.svd(pde_mat, full_matrices=True)\n",
    "    zerosig = np.where(np.abs(sig) < tol)[0]\n",
    "    if zerosig.size:\n",
    "        nullsp_start = zerosig[0]\n",
    "        assert np.array_equal(zerosig, np.arange(nullsp_start, pde_mat.shape[1]))\n",
    "    else:\n",
    "        nullsp_start = pde_mat.shape[0]\n",
    "\n",
    "    return vt[nullsp_start:].T\n",
    "\n",
    "\n",
    "nullsp = find_nullspace(pde_mat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "la.norm(pde_mat @ nullsp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_translation_mat(mexp, to_center):\n",
    "    n = len(mexp.coeffs)\n",
    "    result = np.zeros((n, n))\n",
    "\n",
    "    for j in range(n):\n",
    "        unit_coeffs = np.zeros(n)\n",
    "        unit_coeffs[j] = 1\n",
    "        unit_mexp = mexp.with_coeffs(unit_coeffs)\n",
    "\n",
    "        result[:, j] = t.multipole_expand(unit_mexp, to_center).coeffs\n",
    "\n",
    "    return result\n",
    "\n",
    "\n",
    "new_center = np.array([0, 0.5])\n",
    "tmat = build_translation_mat(mexp, new_center)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(tmat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nullsp.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if 0:\n",
    "    reduction_mat = nullsp.T\n",
    "    expansion_mat = nullsp\n",
    "elif 1:\n",
    "    chosen_indices_and_coeff_ids = [\n",
    "        (i, cid)\n",
    "        for i, cid in enumerate(mpole_expn.get_coefficient_identifiers())\n",
    "        if cid[0] < 2\n",
    "    ]\n",
    "    chosen_indices = [idx for idx, _ in chosen_indices_and_coeff_ids]\n",
    "\n",
    "    expansion_mat = np.zeros((\n",
    "        len(mpole_expn.get_coefficient_identifiers()),\n",
    "        len(chosen_indices_and_coeff_ids),\n",
    "    ))\n",
    "    for i, (idx, _) in enumerate(chosen_indices_and_coeff_ids):\n",
    "        expansion_mat[idx, i] = 1\n",
    "\n",
    "    reduction_mat = (nullsp @ la.inv(nullsp[chosen_indices])).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_coeffs(expn, coeffs, **kwargs):\n",
    "    x = [cid[0] for cid in expn.get_coefficient_identifiers()]\n",
    "    y = [cid[1] for cid in expn.get_coefficient_identifiers()]\n",
    "    plt.scatter(x, y, c=coeffs, **kwargs)\n",
    "    plt.colorbar()\n",
    "\n",
    "    for cid, coeff in zip(expn.get_coefficient_identifiers(), coeffs, strict=True):\n",
    "        plt.text(cid[0], cid[1] + 0.2, f\"{coeff:.1f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proj_mexp = mexp.with_coeffs(expansion_mat @ reduction_mat @ mexp.coeffs)\n",
    "\n",
    "proj_resid = proj_mexp.coeffs - mexp.coeffs\n",
    "\n",
    "plot_coeffs(mpole_expn, np.log10(1e-15 + np.abs(proj_resid)), vmin=-15, vmax=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(t.l_inf(proj_mexp - mexp, 1.2, center=[3, 0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trans_unproj = t.multipole_expand(mexp, new_center)\n",
    "trans_proj = t.multipole_expand(proj_mexp, new_center)\n",
    "\n",
    "print(t.l_inf(trans_unproj - trans_proj, 1.2, center=[3, 0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(trans_proj.coeffs - trans_unproj.coeffs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "la.norm(reduction_mat @ (trans_proj.coeffs - trans_unproj.coeffs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t.l_inf(trans_unproj - pt_src, 1.2, center=[3, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t.l_inf(mexp - pt_src, 1.2, center=[3, 0])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}