{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import exafmm.modified_helmholtz as mod_helm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"exafmm's submodule for Modified Helmholtz kernel\""
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mod_helm.__doc__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. create sources and targets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "nsrcs = 100000\n",
    "ntrgs = 100000\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.random.random(nsrcs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a list of source instances\n",
    "sources = mod_helm.init_sources(src_coords, src_charges)\n",
    "\n",
    "# create a list of target instances\n",
    "targets = mod_helm.init_targets(trg_coords)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. create a ModifiedHelmholtzFmm instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "fmm = mod_helm.ModifiedHelmholtzFmm(p=10, ncrit=400, wavek=7.5, filename='modhelm_example.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 = mod_helm.setup(sources, targets, fmm)"
   ]
  },
  {
   "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": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "trg_values = mod_helm.evaluate(tree, fmm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([  652.74894239, -1315.64151095,  -749.84947316,   340.94417624])"
      ]
     },
     "execution_count": 8,
     "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. It only works when `skip_P2P` is set to `False`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1.500146160135269e-10, 9.677558408664847e-09]"
      ]
     },
     "execution_count": 9,
     "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": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------- iteration 0 ----------\n",
      "Error:  [1.5254311868722168e-10, 8.746810807253287e-09]\n",
      "---------- iteration 1 ----------\n",
      "Error:  [1.484783794640184e-10, 1.0294833601137622e-08]\n",
      "---------- iteration 2 ----------\n",
      "Error:  [1.586084602137114e-10, 1.0816053009140436e-08]\n",
      "---------- iteration 3 ----------\n",
      "Error:  [1.534906027004696e-10, 9.33679182423933e-09]\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 = np.random.random(nsrcs)         # generate new random charges          \n",
    "    mod_helm.update_charges(tree, src_charges)      # update charges\n",
    "    mod_helm.clear_values(tree)                     # clear values\n",
    "    trg_values = mod_helm.evaluate(tree, fmm)       # evaluate potentials\n",
    "\n",
    "    print(\"Error: \", fmm.verify(tree.leafs))       # check accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}