{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "859f09ac-1e0f-40d1-87b3-3f3e16c00272",
   "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",
    "from collections import defaultdict\n",
    "import numpy as np\n",
    "from mayavi import mlab\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "\n",
    "from kifmm_py import (\n",
    "    KiFmm,\n",
    "    LaplaceKernel,\n",
    "    SingleNodeTree,\n",
    "    EvalType,\n",
    "    BlasFieldTranslation,\n",
    "    FftFieldTranslation,\n",
    "    read_stl_triangle_mesh_vertices,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bca79974-3942-4ea4-9325-5314dfb064eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "(sources, faces) = read_stl_triangle_mesh_vertices(\"battleship.STL\")\n",
    "x = sources[:, 0]\n",
    "y = sources[:, 1]\n",
    "z = sources[:, 2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a744f8cc-2d34-4d84-9aec-cabea34cfaa2",
   "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=\"points\", figure=fig\n",
    ")\n",
    "mlab.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5ba2c3af-949d-44a5-b798-2e2ec22c2445",
   "metadata": {},
   "outputs": [],
   "source": [
    "dtype = np.float64\n",
    "\n",
    "# # Set FMM Parameters\n",
    "expansion_order = np.array([6])\n",
    "# n_vec = 1\n",
    "n_crit = 15\n",
    "n_sources = sources.shape[0]\n",
    "n_targets = sources.shape[0]\n",
    "\n",
    "sources = sources.flatten().astype(dtype)\n",
    "charges = np.random.rand(n_sources).astype(dtype)\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 = BlasFieldTranslation(\n",
    "    kernel,\n",
    "    svd_threshold=dtype(1e-16),\n",
    "    surface_diff=2,\n",
    "    random=True,\n",
    "    n_components=100,\n",
    "    n_oversamples=10,\n",
    ")\n",
    "\n",
    "tree = SingleNodeTree(sources, sources, charges, n_crit=n_crit, prune_empty=True)\n",
    "\n",
    "fmm = KiFmm(expansion_order, tree, field_translation, timed=True)\n",
    "\n",
    "fmm.evaluate()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3e52e08-a2c0-473a-9c43-8d5a9fb47b12",
   "metadata": {},
   "source": [
    "# Plot Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f615beab-0210-48e1-8170-dfa360d8d5bb",
   "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[0][:, 0]),\n",
    "    representation=\"surface\",\n",
    ")\n",
    "\n",
    "colorbar = mlab.colorbar(solution, title=\"Log-scaled values\", orientation=\"vertical\")\n",
    "colorbar.title_text_property.font_size = 10\n",
    "colorbar.label_text_property.font_size = 8\n",
    "mlab.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88525acc-aae5-4ccf-a069-0229429584ef",
   "metadata": {},
   "source": [
    "# Statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "91893002-cd58-49e6-b421-0c4d3f6d7c1f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "stats = defaultdict(int)\n",
    "\n",
    "for leaf in fmm.leaves_target_tree:\n",
    "    found = fmm.coordinates_target_tree(leaf)\n",
    "    n = found.shape[0]\n",
    "    stats[n] += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6ae7f253-fc1b-4ee2-8d5c-ced227dba336",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 22696 characters]",
      "text/plain": [
       "<Figure size 800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# # Plot the histogram using seaborn\n",
    "sns.set_theme(style=\"whitegrid\")\n",
    "\n",
    "plt.figure(figsize=(8, 6))\n",
    "\n",
    "# Convert the dictionary to a pandas DataFrame\n",
    "data = []\n",
    "for points, occurrences in stats.items():\n",
    "    data.extend([points] * occurrences)\n",
    "\n",
    "df = pd.DataFrame(data, columns=[\"Points\"])\n",
    "\n",
    "hist_plot = sns.histplot(\n",
    "    df, x=\"Points\", bins=len(stats), color=\"black\", log_scale=(True, False)\n",
    ")\n",
    "\n",
    "# Load CMU font\n",
    "import matplotlib.font_manager as fm\n",
    "\n",
    "font_path = \"/Users/sri/Downloads/cmu/cmunrm.ttf\"\n",
    "font_properties = fm.FontProperties(fname=font_path)\n",
    "\n",
    "hist_plot.set_xlabel(\"Points Per Leaf Box\", fontsize=15, fontproperties=font_properties)\n",
    "hist_plot.set_ylabel(\"Count\", fontsize=15, fontproperties=font_properties)\n",
    "\n",
    "hist_plot.set_ybound(lower=None, upper=117)\n",
    "# Set font properties for tick labels\n",
    "for label in hist_plot.get_xticklabels():\n",
    "    label.set_fontproperties(font_properties)\n",
    "for label in hist_plot.get_yticklabels():\n",
    "    label.set_fontproperties(font_properties)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"stats.svg\", format=\"svg\")\n",
    "plt.show()"
   ]
  }
 ],
 "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
}
