{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "62aefbe6-e4da-4d62-88d0-2994c3872298",
   "metadata": {},
   "source": [
    "# Python API for KiFMM-rs\n",
    "\n",
    "We provide a full, object oriented, Python API for our Rust library. Below are a number of examples of how to achieve different functionality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b3701203-9472-4fb0-ba09-4ee2c84c80c1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "********************************************************************************\n",
      "WARNING: Imported VTK version (9.3) does not match the one used\n",
      "         to build the TVTK classes (9.2). This may cause problems.\n",
      "         Please rebuild TVTK.\n",
      "********************************************************************************\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%gui qt\n",
    "import numpy as np\n",
    "from mayavi import mlab\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from kifmm_py import (\n",
    "    KiFmm,\n",
    "    LaplaceKernel,\n",
    "    HelmholtzKernel,\n",
    "    SingleNodeTree,\n",
    "    EvalType,\n",
    "    BlasFieldTranslation,\n",
    "    FftFieldTranslation,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e8c7b6d-5cd7-4ef7-ba90-fb707c6c5577",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Setting up an FMM requires some coordinate data associated with source and target particles, which we expect as a flat Fortran ordered buffer of shape (n_coordinates, 3), where each column corresponds to the component of each axis.\n",
    "\n",
    "Additionally the user needs to choose some basic parameters, principally\n",
    "\n",
    "## Kernel\n",
    "\n",
    "We support both `LaplaceKernel` and `HelmholtzKernel`, which has been is an active area of development for low-frequencies and for now offers no speed or convergence guarantees.\n",
    "\n",
    "## Field Translation\n",
    "\n",
    "Users must specify the acceleration schemes for the multipole to local (M2L) field translation step of the FMM.\n",
    "\n",
    "We currently offer `FftFieldTranslation`, and `BlasFieldTranslation`, where the matrices associated with `BlasFieldTranslation` can optionally be compressed with a randomised SVD, or a deterministic SVD.\n",
    "\n",
    "## Tree\n",
    "\n",
    "Only `SingleNodeTree`s are provided at present, users can optionally discretise with a specified critical value of particles per leaf box `n_crit`, or by explicitly setting the tree's `depth`. In the latter case `expansion_order` must be set for each level of the tree, allowing for the option to set variable expansion order by level. This can be efficacious for the `HelmholtzKernel`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96cdf276-1357-4b73-bd5e-ab575b9e5796",
   "metadata": {},
   "source": [
    "## Basic Example (Laplace FMM + FFT M2L)\n",
    "\n",
    "Here we use a FFT based M2L field translation, explicitly setting a block size for the field translation algorithm. We discretise the octree with an `n_crit` parameter, which defines the maximum number of particles per leaf box. Here we set `EvalType` to `Value`, meaning that we evaluate only potentials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2df68582-77d0-41f0-8e2a-702264d8b216",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'P2M': 0, 'M2M(2)': 0, 'M2M(1)': 0, 'M2L(2)': 2, 'P2P': 0, 'L2P': 0}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.random.seed(0)\n",
    "\n",
    "dim = 3\n",
    "dtype = np.float32\n",
    "\n",
    "# Set FMM Parameters\n",
    "expansion_order = np.array([6], np.uint64)  # Single expansion order as using n_crit\n",
    "n_vec = 1\n",
    "n_crit = 150\n",
    "n_sources = 1000\n",
    "n_targets = 2000\n",
    "prune_empty = True  # Optionally remove empty leaf boxes, their siblings, and ancestors, from the Tree\n",
    "\n",
    "# Setup source/target/charge data in Fortran order\n",
    "sources = np.random.rand(n_sources * dim).astype(dtype)\n",
    "targets = np.random.rand(n_targets * dim).astype(dtype)\n",
    "charges = np.random.rand(n_sources * n_vec).astype(dtype)\n",
    "\n",
    "eval_type = EvalType.Value\n",
    "\n",
    "# EvalType computes either potentials (EvalType.Value) or potentials + derivatives (EvalType.ValueDeriv)\n",
    "kernel = LaplaceKernel(dtype, eval_type)\n",
    "\n",
    "tree = SingleNodeTree(sources, targets, charges, n_crit=n_crit, prune_empty=prune_empty)\n",
    "\n",
    "field_translation = FftFieldTranslation(kernel, block_size=32)\n",
    "\n",
    "# Create FMM runtime object\n",
    "fmm = KiFmm(expansion_order, tree, field_translation, timed=True)\n",
    "\n",
    "# Evaluate potentials\n",
    "fmm.evaluate()\n",
    "\n",
    "# Examine potentials\n",
    "fmm.all_potentials\n",
    "\n",
    "# Examine operator times rounded in milliseconds\n",
    "fmm.operator_times()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "dc3d5d28-5de2-4b29-8784-200f4109d049",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'source_data': 4, 'target_data': 3, 'source_to_target_data': 256}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fmm.metadata_times()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0f2234b8-f3a2-4cf5-b018-b581ee729b73",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'source_domain': 0, 'target_domain': 0, 'source_tree': 0, 'target_tree': 0}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fmm.communication_times()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d3ea7b0-5c68-4539-8260-59d14f87668d",
   "metadata": {},
   "source": [
    "## Compare with Direct Evaluation\n",
    "\n",
    "Note, as a part of the FMM input coordinate data is sorted into a `Morton' defined order, as a part of tree construction. Therefore to compare with direct evaluation, the evaluated potentials must be unpermuted as a post processing step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1da87e6a-8c63-42e7-8ea9-bf60ca593218",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Examine unsorted potentials\n",
    "fmm.all_potentials_u\n",
    "\n",
    "# Direct evaluation\n",
    "direct = np.zeros(n_targets * eval_type.value).astype(dtype)\n",
    "fmm.evaluate_kernel(EvalType.Value, sources, targets, charges, direct)\n",
    "\n",
    "assert np.allclose(fmm.all_potentials_u[0][:, 0], direct)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5117b3cb-58c8-4230-bce1-2ec53c028d6d",
   "metadata": {},
   "source": [
    "#### Clear Charge Data\n",
    "\n",
    "We offer the option to clear associated charge/potential data and re-run the FMM for the same set of points with new charge data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ba46ad80-72ce-4f5c-a3bd-5f93b030f248",
   "metadata": {},
   "outputs": [],
   "source": [
    "charges = np.random.rand(n_sources * n_vec).astype(dtype)\n",
    "fmm.clear(charges)  # Can be re-evaluated\n",
    "\n",
    "fmm.evaluate()\n",
    "\n",
    "# Direct evaluation\n",
    "direct = np.zeros(n_targets * eval_type.value).astype(dtype)\n",
    "fmm.evaluate_kernel(EvalType.Value, sources, targets, charges, direct)\n",
    "\n",
    "assert np.allclose(fmm.all_potentials_u[0][:, 0], direct)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21fdc077-3656-48e6-8465-cc99a06b92a3",
   "metadata": {},
   "source": [
    "## 3D Plot with MayaVi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "df5236c6-ccb7-495b-958e-b5b7569f1372",
   "metadata": {},
   "outputs": [],
   "source": [
    "from kifmm_py import read_stl_triangle_mesh_vertices"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d55042a5-4515-4128-aeef-077552661d87",
   "metadata": {},
   "source": [
    "#### Read example mesh from a file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "89147a31-182c-4bd0-920e-000fb5ff5f9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "(sources, faces) = read_stl_triangle_mesh_vertices(\"example.stl\")\n",
    "x = sources[:, 0]\n",
    "y = sources[:, 1]\n",
    "z = sources[:, 2]\n",
    "sources = sources.flatten()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cf4a1df-0cbc-4a12-9952-559fb0910b1b",
   "metadata": {},
   "source": [
    "#### Plot Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "333eaf61-cb42-45d2-ae37-382cf866e9d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "mlab.view(azimuth=40, elevation=70, distance=\"auto\", focalpoint=\"auto\")\n",
    "fig = mlab.figure(\n",
    "    size=(960, 1080), bgcolor=(1, 1, 1)\n",
    ")  # This sets the window size for rendering\n",
    "plot = mlab.triangular_mesh(\n",
    "    x, y, z, faces, color=(0.5, 0.5, 0.5), representation=\"surface\", figure=fig\n",
    ")\n",
    "mlab.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5567623c-fd44-47ee-a8c0-13139a7f7df1",
   "metadata": {},
   "source": [
    "#### Setup FMM and Evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "363d8d4b-f967-44dd-b93e-c97192295ff4",
   "metadata": {},
   "outputs": [],
   "source": [
    "dim = 3\n",
    "dtype = np.float32\n",
    "\n",
    "# Set FMM Parameters\n",
    "expansion_order = np.array([6])\n",
    "n_vec = 1\n",
    "n_crit = 150\n",
    "n_sources = len(sources) // dim\n",
    "\n",
    "# Set Random charges\n",
    "charges = np.random.rand(n_sources).astype(dtype)\n",
    "\n",
    "kernel = LaplaceKernel(dtype, EvalType.Value)\n",
    "\n",
    "field_translation = FftFieldTranslation(kernel, block_size=32)\n",
    "tree = SingleNodeTree(sources, sources, charges, n_crit=n_crit, prune_empty=True)\n",
    "\n",
    "# Create FMM runtime object\n",
    "fmm = KiFmm(expansion_order, tree, field_translation, timed=True)\n",
    "\n",
    "fmm.evaluate()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5860ef5f-bbe6-4208-b10b-f8cb79b9a77f",
   "metadata": {},
   "source": [
    "#### Plot Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "25d4692c-d7ab-4f03-b2aa-ca92da29b961",
   "metadata": {},
   "outputs": [],
   "source": [
    "mlab.view(azimuth=40, elevation=70, distance=\"auto\", focalpoint=\"auto\")\n",
    "fig = mlab.figure(\n",
    "    size=(960, 1080), bgcolor=(1, 1, 1)\n",
    ")  # This sets the window size for rendering\n",
    "solution = mlab.triangular_mesh(\n",
    "    x,\n",
    "    y,\n",
    "    z,\n",
    "    faces,\n",
    "    scalars=np.log(fmm.all_potentials_u).flatten(),\n",
    "    representation=\"surface\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f440d57b-26be-47aa-9ba0-4230cf8d151f",
   "metadata": {},
   "source": [
    "# More Complex Examples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "823df83b-e1f9-4fb2-96dd-023d74ddff86",
   "metadata": {},
   "source": [
    "### 1) FMM With Multiple Charge Vectors\n",
    "\n",
    "Here we evaluate multiple charge vectors for the Laplace FMM, with the same set of source/target points. In this case `BlasFieldTranslation`s must be used for the M2L field translation. Additionally we evaluate `EvalType.ValueDeriv`, computing both the potentials and potential derivatives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "cecccd0e-f229-4d1a-9849-7c73c725b58f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Multiple charge vectors only supported with BlasFieldTranslation\n"
     ]
    }
   ],
   "source": [
    "np.random.seed(0)\n",
    "\n",
    "dim = 3\n",
    "dtype = np.float32\n",
    "\n",
    "# Set FMM Parameters\n",
    "expansion_order = np.array([6], np.uint64)  # Single expansion order as using n_crit\n",
    "n_vec = 5\n",
    "n_crit = 150\n",
    "n_sources = 1000\n",
    "n_targets = 2000\n",
    "\n",
    "# Setup source/target/charge data in Fortran order\n",
    "sources = np.random.rand(n_sources * dim).astype(dtype)\n",
    "targets = np.random.rand(n_targets * dim).astype(dtype)\n",
    "charges = np.random.rand(n_sources * n_vec).astype(dtype)\n",
    "\n",
    "svd_threshold = 1e-7\n",
    "\n",
    "eval_type = EvalType.ValueDeriv\n",
    "\n",
    "# EvalType computes either potentials (EvalType.Value) or potentials + derivatives (EvalType.ValueDeriv)\n",
    "kernel = LaplaceKernel(dtype, eval_type)\n",
    "\n",
    "tree = SingleNodeTree(sources, targets, charges, n_crit=n_crit, prune_empty=True)\n",
    "\n",
    "try:\n",
    "    # Will raise TypeError\n",
    "    field_translation = FftFieldTranslation(kernel, block_size=32)\n",
    "    # Create FMM runtime object\n",
    "    fmm = KiFmm(expansion_order, tree, field_translation, timed=True)\n",
    "except TypeError as e:\n",
    "    print(e)\n",
    "    field_translation = BlasFieldTranslation(kernel, svd_threshold, random=False)\n",
    "    # Create FMM runtime object\n",
    "    fmm = KiFmm(expansion_order, tree, field_translation, timed=True)\n",
    "\n",
    "\n",
    "# Run FMM\n",
    "fmm.evaluate()\n",
    "\n",
    "# Check wrt to direct evaluation\n",
    "\n",
    "leaf = fmm.leaves_target_tree[0]\n",
    "leaf_potentials = fmm.leaf_potentials(leaf)\n",
    "leaf_targets = fmm.coordinates_target_tree(leaf)\n",
    "n_targets = len(leaf_targets) // 3\n",
    "\n",
    "for i, evaluated in enumerate(leaf_potentials):\n",
    "    direct = np.zeros(n_targets * eval_type.value).astype(dtype)\n",
    "    fmm.evaluate_kernel(\n",
    "        eval_type,\n",
    "        fmm._tree.sources,\n",
    "        leaf_targets,\n",
    "        fmm._tree.charges[i * n_sources : (i + 1) * n_sources],\n",
    "        direct,\n",
    "    )\n",
    "    assert np.allclose(evaluated, direct, 1e-2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bd93a06-5d86-478d-9ac2-4a93b50a559c",
   "metadata": {},
   "source": [
    "### 2) Helmholtz FMM\n",
    "\n",
    "Support for the low frequency Helmholtz FMM is an active area of development, there are currently no (speed) performance guarantees, and the API is offered as is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "78dff441-27d3-4354-8ac9-766e9150f070",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(0)\n",
    "\n",
    "dim = 3\n",
    "dtype = np.float32\n",
    "ctype = np.complex64\n",
    "\n",
    "# Set FMM Parameters\n",
    "expansion_order = np.array([6], np.uint64)  # Single expansion order as using n_crit\n",
    "n_vec = 1\n",
    "n_crit = 150\n",
    "n_sources = 1000\n",
    "n_targets = 2000\n",
    "wavenumber = dtype(10)\n",
    "prune_empty = True\n",
    "\n",
    "# Setup source/target/charge data in Fortran order\n",
    "sources = np.random.rand(n_sources * dim).astype(dtype)\n",
    "targets = np.random.rand(n_targets * dim).astype(dtype)\n",
    "charges = np.random.rand(n_sources * n_vec).astype(ctype)\n",
    "\n",
    "eval_type = EvalType.Value\n",
    "\n",
    "# EvalType computes either potentials (EvalType.Value) or potentials + derivatives (EvalType.ValueDeriv)\n",
    "kernel = HelmholtzKernel(dtype, wavenumber, eval_type)\n",
    "\n",
    "tree = SingleNodeTree(sources, targets, charges, n_crit=n_crit, prune_empty=prune_empty)\n",
    "\n",
    "field_translation = FftFieldTranslation(kernel, block_size=32)\n",
    "\n",
    "# Create FMM runtime object\n",
    "fmm = KiFmm(expansion_order, tree, field_translation, timed=True)\n",
    "\n",
    "# Evaluate FMM\n",
    "fmm.evaluate()\n",
    "\n",
    "# Direct evaluation check\n",
    "direct = np.zeros(n_targets * eval_type.value).astype(ctype)\n",
    "fmm.evaluate_kernel(EvalType.Value, sources, targets, charges, direct)\n",
    "\n",
    "# assert np.allclose(fmm.all_potentials_u, direct)\n",
    "assert np.allclose(np.abs(direct), np.abs(fmm.all_potentials_u[0].T), 1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24a5eb2e-b718-4b14-9a6b-e270666df62d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3c61f7e-8efa-45da-a754-d02ef29341cb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1e1d1c6-8bf0-4452-b62f-0bf6262fe475",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2155f6db-dc42-4004-bcb1-0e23518e3857",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}