{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "from sumpy.expansion.local import VolumeTaylorLocalExpansion\n",
    "from sumpy.expansion.multipole import (\n",
    "    LaplaceConformingVolumeTaylorMultipoleExpansion,\n",
    "    LinearPDEConformingVolumeTaylorMultipoleExpansion,\n",
    "    VolumeTaylorMultipoleExpansion,\n",
    ")\n",
    "from sumpy.kernel import HelmholtzKernel, LaplaceKernel, YukawaKernel  # noqa: F401\n",
    "from sumpy.symbolic import make_sym_vector\n",
    "\n",
    "\n",
    "order = 2\n",
    "dim = 2\n",
    "\n",
    "if 0:\n",
    "    knl = LaplaceKernel(dim)\n",
    "    extra_kernel_kwargs = {}\n",
    "    mpole_expn_reduced_class = LaplaceConformingVolumeTaylorMultipoleExpansion\n",
    "\n",
    "else:\n",
    "    helm_k = 1.2\n",
    "    knl = HelmholtzKernel(dim)\n",
    "    extra_kernel_kwargs = {\"k\": helm_k}\n",
    "    mpole_expn_reduced_class = LinearPDEConformingVolumeTaylorMultipoleExpansion\n",
    "\n",
    "mpole_expn_reduced = mpole_expn_reduced_class(knl, order)\n",
    "mpole_expn = VolumeTaylorMultipoleExpansion(knl, order)\n",
    "local_expn = VolumeTaylorLocalExpansion(knl, order)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reduced_wrangler = mpole_expn_reduced.expansion_terms_wrangler\n",
    "full_wrangler = mpole_expn.expansion_terms_wrangler\n",
    "\n",
    "reduced_derivatives = list(\n",
    "    make_sym_vector(\"deriv\", len(reduced_wrangler.stored_identifiers))\n",
    ")\n",
    "full_derivatives = reduced_wrangler.get_full_kernel_derivatives_from_stored(\n",
    "    reduced_derivatives, 1\n",
    ")\n",
    "\n",
    "print(reduced_derivatives)\n",
    "print(full_derivatives)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "full_coeffs = list(\n",
    "    make_sym_vector(\"coeff\", len(reduced_wrangler.get_full_coefficient_identifiers()))\n",
    ")\n",
    "\n",
    "reduced_coeffs = reduced_wrangler.get_stored_mpole_coefficients_from_full(\n",
    "    full_mpole_coefficients=full_coeffs, rscale=1\n",
    ")\n",
    "\n",
    "print(full_coeffs)\n",
    "print(reduced_coeffs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dvec = make_sym_vector(\"d\", dim)\n",
    "translated_reduce_coeffs = mpole_expn_reduced.translate_from(\n",
    "    mpole_expn_reduced, reduced_coeffs, 1, dvec, 1\n",
    ")\n",
    "translated_full_coeffs = mpole_expn.translate_from(mpole_expn, full_coeffs, 1, dvec, 1)\n",
    "translated_full_coeffs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_reduced = sum(a * b for a, b in zip(translated_reduce_coeffs, reduced_derivatives,\n",
    "                                         strict=True))\n",
    "eval_full = sum(a * b for a, b in zip(translated_full_coeffs, full_derivatives,\n",
    "                                      strict=True))\n",
    "\n",
    "(eval_full - eval_reduced).simplify()"
   ]
  }
 ],
 "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
}