{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import exafmm.helmholtz as helmholtz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"exafmm's submodule for Helmholtz kernel\""
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "helmholtz.__doc__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. create sources and targets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "nsrcs = 100000\n",
    "ntrgs = 200000\n",
    "\n",
    "# generate random positions for particles\n",
    "src_coords = np.random.random((nsrcs, 3))\n",
    "trg_coords = np.random.random((nsrcs, 3))\n",
    "\n",
    "# generate random charges for sources\n",
    "src_charges = np.zeros(nsrcs, dtype=np.complex_)\n",
    "src_charges.real = np.random.random(nsrcs)\n",
    "src_charges.imag = np.random.random(nsrcs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a list of source instances\n",
    "sources = helmholtz.init_sources(src_coords, src_charges)\n",
    "\n",
    "# create a list of target instances\n",
    "targets = helmholtz.init_targets(trg_coords)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. create a HelmholtzFmm instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "fmm = helmholtz.HelmholtzFmm(p=10, ncrit=300, wavek=10, filename='helmholtz_example_k10.dat')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. setup fmm\n",
    "\n",
    "Given sources, targets and FMM configurations, `setup()` builds the tree and interaction list and pre-compute invariant matrices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree = helmholtz.setup(sources, targets, fmm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "helmholtz_example_k10.dat  test_file.dat\r\n"
     ]
    }
   ],
   "source": [
    "ls *.dat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. evaluate\n",
    "\n",
    "`evaluate()` triggers the evaluation and returns a $n_{trg} \\times 4$ `numpy.ndarray`.\n",
    "The $i$-th row of `trg_values` starts with the potential value of the $i$-th target, followed by its three gradient values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "trg_values = helmholtz.evaluate(tree, fmm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ -221.98960533 -761.47936505j,  -865.25180217+2481.40881884j,\n",
       "        1290.96521265 -389.92609902j, -6337.67344409+2655.85947907j])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trg_values[6]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. check accuracy (optional)\n",
    "\n",
    "`fmm.verify(tree.leafs)` returns L2-norm the relative error of potential and gradient in a list, compared with the values calculated from direct method. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[3.177234172973424e-08, 5.900307387772e-07]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fmm.verify(tree.leafs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6. update charges of sources and run FMM iteratively"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------- iteration 0 ----------\n",
      "Error:  [3.1462242799274535e-08, 5.831925128982927e-07]\n",
      "---------- iteration 1 ----------\n",
      "Error:  [3.1732536842029743e-08, 5.208091857409024e-07]\n",
      "---------- iteration 2 ----------\n",
      "Error:  [3.180717719777336e-08, 5.515657504897553e-07]\n",
      "---------- iteration 3 ----------\n",
      "Error:  [3.1376568787696064e-08, 5.542072823427909e-07]\n"
     ]
    }
   ],
   "source": [
    "niters = 4\n",
    "\n",
    "for i in range(niters):\n",
    "    print('-'*10 + ' iteration {} '.format(i) + '-'*10)  # print divider between iterations\n",
    "    \n",
    "    src_charges.real = np.random.random(nsrcs)     # generate new random charges\n",
    "    src_charges.imag = np.random.random(nsrcs)          \n",
    "    helmholtz.update_charges(tree, src_charges)      # update charges\n",
    "    helmholtz.clear_values(tree)                     # clear values\n",
    "    trg_values = helmholtz.evaluate(tree, fmm)       # evaluate potentials\n",
    "\n",
    "    print(\"Error: \", fmm.verify(tree.leafs))       # check accuracy"
   ]
  }
 ],
 "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}