04_tree_construction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Tree Construction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy\n",
    "from treecode_helper import Point, Particle"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the previous notebook, we discussed how to use multi-level multipole expansion to calculate the potential field at target points, when the sources and targets are well separated. So in this notebook, we will think about how to build a hierarchical tree to store the particles in order to achieve a multi-level multipole expansion.\n",
    "\n",
    "In 1986, Barnes and Hut observed that the particles might work in the same way that humans interact with neighboring individuals, more distant villages and larger countries. They came up with an algorithm that the potential on an individual particle from other particles close by is evaluated by direct summation (particle-particle interaction), whereas the potential due to more distant particles is calculated by multipole expansion (particle-cluster interaction). To apply this algorithm, they proposed a tree structure to group the particles in a smart way to facilitate the tree algorithm. Let's follow their idea to construct a tree of particles!\n",
    "\n",
    "Here we demonstrate the idea with a 2d quadtree where each cell can be subdivided into $4$ child cells. We begin with a **root** cell which is big enough to contain all the particles in space. We follow several rules to build this quadtree:\n",
    "* If there is more than one particles in the cell, we need to split this **parent** cell into 4 **child** cells.\n",
    "* If there is only one particle in the cell, which means this cell cannot be split, then we call this a **leaf** cell.\n",
    "* If there is no particles in a cell, we can ignore the cell (do not store the cell in memory).\n",
    "\n",
    "<img src=\"image/division_tree.png\">\n",
    "\n",
    "The figure above illustrates the recursive procedure to build a quadtree based on these rules. In this example, we use the number \"1\" to determine whether the cell needs to be split. We call this number a critical number $n_{crit}$, which implies the maximum number of particles that a **leaf** cell can contain. If a cell accommodates more than $n_{crit}$ number of particles, it has to be split and is no longer a leaf cell. This parameter tells about the granularity of the problem, which has a tremendous effect on the efficiency of treecode.\n",
    "\n",
    "Now let's start to define the problem and our rules to build the tree. Consider there are $n=100$ particles randomly scatterd in the domain $x$, $y$, $z$ $\\in$ $\\left[ 0, 1 \\right]$, each of them is a source and target. To contain all the particles, we define a **root** cubic cell centered at $(0.5,0.5, 0.5)$ with a side length of $1$. For a cubic cell, the radius is the half of side length, thus the root cell's radius $r$ is $0.5$. Then we choose $n_{crit}=10$ as the threshold to split a cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "n = 100          # number of particles\n",
    "particles = [ Particle(m=1.0/n) for i in range(n) ]\n",
    "\n",
    "n_crit = 10      # max number of particles in a single cell"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Defining the class: `Cell`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since each non-empty cell is an instance which has some same properties (eg. cell's center coordinates, its parent, its children, its multipole, particles inside if it's a leaf), we first need to define a class for a cell, and we call this class `Cell`. Every cell is an instance of class `Cell`, and all the instances are stored in a list called `cells`. Therefore, the root cell is `cells[0]`. For those who have not been exposed to object-oriented programming in python, check [this](http://www.tutorialspoint.com/python/python_classes_objects.htm) out as a quick guide.\n",
    "\n",
    "<img src=\"image/cell_class.png\">\n",
    "\n",
    "The figure above shows the \"content\" of a cell element:\n",
    "* $x_c$, $y_c$, $z_c$, $r_c$: center coordinates and radius give the geometry of the cell.\n",
    "* A leaf of a cell is a particle stored in the cell, each leaf corresponds to a particle index from $0$ to $n-1$, and $n_{leaf}$ is number of leaves in the cell\n",
    "* **parent** is the index (in the cells list) of the corresponding parent. \n",
    "* **child** is the array that contains the indices (in the cells list) of the corresponding children.\n",
    "* **nchild** is an 8-bit binary number, each digit represents one of the eight child octants. $1$ denotes non-empty child cell, and $0$ denotes empty child cell in that octant. For example: nchild=00010100 means the current cell only has non-empty child in the fourth and sixth octant.\n",
    "* **multipole**: array of $10$ multipole terms of the cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class Cell():\n",
    "    \"\"\"The class for a cell.\n",
    "    \n",
    "    Arguments:\n",
    "      n_crit: maximum number of particles in a leaf cell\n",
    "    \n",
    "    Attributes:\n",
    "        nleaf (int): number of particles in the cell.\n",
    "        leaf (array of int): array of leaves indices.\n",
    "        nchild (int): 8-bit binary number, used to keep track of the empty child cells\n",
    "        child (array of int): array of children indices.\n",
    "        parent (int): index of parent cell.\n",
    "        x, y, z (float): coordinates of the cell's center.\n",
    "        r (float): radius of the cell (half of the length for cubic cell).\n",
    "        multipole (array of float): array of multipoles' cell.\n",
    "      \n",
    "    \"\"\"\n",
    "    def __init__(self, n_crit):\n",
    "        self.nleaf = 0        # number of particles\n",
    "        self.leaf = numpy.zeros(n_crit, dtype=numpy.int)     # array of leaf index\n",
    "        self.nchild = 0       # binary counter to keep track of empty cells\n",
    "        self.child = numpy.zeros(8, dtype=numpy.int)         # array of child index\n",
    "        self.parent = 0       # index of parent cell\n",
    "        self.x = self.y = self.z = 0.                    # center of the cell\n",
    "        self.r = 0.           # radius of the cell\n",
    "        self.multipole = numpy.zeros(10, dtype=numpy.float)  # multipole array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "root = Cell(n_crit)\n",
    "root.x, root.y, root.z = 0.5, 0.5, 0.5\n",
    "root.r = 0.5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Adding a child"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assuming that we create the root cell, then we add the particles into this root cell. After we put the 10th particle into the root, $n_{crit} = 10$ for the root. If there is a 11th particle coming, we have to split the root cell. So what is the first step if you want to do a split? The answer is to create a new cell instance which is the child of the root cell in the cells list. Thus `add_child` is a dependency of `split_cell`. Now let's think about how to add a child.\n",
    "\n",
    "To begin with, we need to append a new element to the cells list, so now the last element of the cells list should be the new child. First let's find the geometrical relationship between a parent $p$ and its child $c$:\n",
    "* $r_{child}$ = $\\frac{1}{2}r_{parent}$ for a cubic cell\n",
    "* $x_{c_{c}}$, $y_{c_{c}}$, $z_{c_{c}}$ can be determined by its $octant$ and its parent's coordinates $x_{c_{p}}$, $y_{c_{p}}$, $z_{c_{p}}$, refer to the code below and reflect how it works.\n",
    "\n",
    "Then we need to establish a mutual reference between the parent and the child in the cells list. Consider a new child is created in the parent's $5th$ octant. We assign the new child's index to the parent by `parent.child[4]=index_child`, and assign the parent's index to the new child by \"child.parent=index_parent\". Don't forget the 8-bit binary marker `nchild`. Since the new child is at $5th$ octant, the fifth digit from the right should be changed from \"0\" to \"1\". Recall that we always manipulate the binary number with bit shift."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def add_child(octant, p, cells, n_crit):\n",
    "    \"\"\"Add a cell to the end of cells list as a child of p,\n",
    "       initialize the center and radius of the child cell c, \n",
    "       and establish mutual reference between child c and parent p.\n",
    "    \n",
    "    Arguments:\n",
    "        octant: reference to the corresponding octant in 3D structure.\n",
    "        p: is the index (in the cells list) of the corresponding parent.\n",
    "        cells: the list of cells.\n",
    "        n_crit: maximum number of particles in a leaf cell.\n",
    " \n",
    "    \"\"\"\n",
    "    # create a new cell instance and append it to cells list\n",
    "    cells.append(Cell(n_crit))\n",
    "    # the last element of the cells list is the new child c\n",
    "    c = len(cells) - 1\n",
    "    # geometric relationship between parent and child\n",
    "    cells[c].r = cells[p].r / 2\n",
    "    cells[c].x = cells[p].x + cells[c].r * ((octant & 1) * 2 - 1)\n",
    "    cells[c].y = cells[p].y + cells[c].r * ((octant & 2) - 1    )\n",
    "    cells[c].z = cells[p].z + cells[c].r * ((octant & 4) / 2 - 1)\n",
    "    # establish mutual reference in the cells list\n",
    "    cells[c].parent = p\n",
    "    cells[p].child[octant] = c\n",
    "    cells[p].nchild = (cells[p].nchild | (1 << octant))\n",
    "    print('+++cell {} is created as a child of cell {}'.format(c, p))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Splitting a cell"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now Let's focus on splitting a cell. First we need to realize that the `split_cell` function should be recursive because after splitting a cell there is a probability that all the particles are reallocated to the same child. In this scenario, we have to recursively split the child cell again until all the cells satisfy the rule $n_{crit}=10$.\n",
    "\n",
    "Consider now we put $>$10 particles in the root cell. So it is time to split the root. In addition to create new child cells, we have also to \"settle\" down these particles. We loop over the them, and each particle is located in a certain octant from 0 to 7. If there is not a child cell in that octant, then we create one. If there is a child cell already created, we just put this particle in that child. Finally, we check if the number of particles in that child reaches $n_{crit}$, if yes, we split it recursively. \n",
    "\n",
    "<img src=\"image/split_cell.png\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def split_cell(particles, p, cells, n_crit):\n",
    "    \"\"\"Loop in parent p's leaves and reallocate the particles to subcells. \n",
    "    If a subcell has not been created in that octant, it creates one using\n",
    "    add_child. If the subcell nleaf exceeds n_crit, split the\n",
    "    subcell c recursively.\n",
    "    \n",
    "    Arguments: \n",
    "        particles: the list of particles.\n",
    "        p: is the index (in the cells list) of the corresponding parent.\n",
    "        cells: the list of cells.\n",
    "        n_crit: maximum number of particles in a leaf cell.\n",
    "    \n",
    "    \"\"\"\n",
    "    print('======start the split of cell {}======'.format(p))\n",
    "    # loop over the particles in the parent cell that you want to split\n",
    "    for l in cells[p].leaf:\n",
    "        octant = (particles[l].x > cells[p].x) + ((particles[l].y > cells[p].y) << 1) \\\n",
    "               + ((particles[l].z > cells[p].z) << 2)   # finding the particle's octant\n",
    "        # if there is not a child cell in the particle's octant, then create one\n",
    "        if not cells[p].nchild & (1 << octant):\n",
    "            add_child(octant, p, cells, n_crit)\n",
    "        # reallocate the particle in the child cell\n",
    "        c = cells[p].child[octant]\n",
    "        cells[c].leaf[cells[c].nleaf] = l \n",
    "        cells[c].nleaf += 1\n",
    "        print('>>>particle {} is reallocated in cell {}'.format(l, c))\n",
    "        # check if the child reach n_crit\n",
    "        if cells[c].nleaf >= n_crit:\n",
    "            split_cell(particles, c, cells, n_crit)\n",
    "    print('======end split cell {}======'.format(p))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Constructing the tree"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Based on adding and splitting a cell, we can easily form our code to build the tree. We loop over the particles, and find a cell for each particle from top (root cell) to bottom (leaf cell). If the current cell where the particle lives is not a leaf cell ($n_{leaf}=10$), then we recursively split it traverse down (find the octant and go to its child) until we find or create the leaf cell ($n_{leaf}<10$). Finally we can put every particle in a certain leaf cell where $n_{leaf}<10$. After we allocate every particle, we perform a check to determine whether the cell need to be splitted or not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def build_tree(particles, root, n_crit):\n",
    "    \"\"\"Construct a hierarchical octree to store the particles and\n",
    "       return the tree (list) of cells.\n",
    "    \n",
    "    Arguments:\n",
    "        particles: the list of particles.\n",
    "        root: the root cell.\n",
    "        n_crit: maximum number of leaves in a single cell.\n",
    "    \n",
    "    Returns:\n",
    "        cells: the list of cells\n",
    "    \n",
    "    \"\"\"\n",
    "    # set root cell\n",
    "    cells = [root]       # initialize the cells list\n",
    "\n",
    "    # build tree\n",
    "    n = len(particles)\n",
    "    for i in range(n):\n",
    "        # traverse from the root down to a leaf cell\n",
    "        curr = 0\n",
    "        while cells[curr].nleaf >= n_crit:\n",
    "            cells[curr].nleaf += 1\n",
    "            octant = (particles[i].x > cells[curr].x) + ((particles[i].y > cells[curr].y) << 1) \\\n",
    "                   + ((particles[i].z > cells[curr].z) << 2)\n",
    "            # if there is no child cell in the particles octant, then create one\n",
    "            if not cells[curr].nchild & (1 << octant):\n",
    "                add_child(octant, curr, cells, n_crit)\n",
    "            curr = cells[curr].child[octant]\n",
    "        # allocate the particle in the leaf cell\n",
    "        cells[curr].leaf[cells[curr].nleaf] = i\n",
    "        cells[curr].nleaf += 1\n",
    "        print('particle {} is stored in cell {}'.format(i, curr))\n",
    "        # check whether to split or not\n",
    "        if cells[curr].nleaf >= n_crit:\n",
    "            split_cell(particles, curr, cells, n_crit)\n",
    "    \n",
    "    return cells"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's build our tree!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "cells = build_tree(particles, root, n_crit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Eventually, we build the hierarchical octree based on the random particle distribution and our rule that $n_{crit}=10$. By constructing the tree, we put all the cells in a single list, and for each cell element, we can tell its geometry/location parameter ($x_c$, $y_c$, $z_c$, $r$). In addition, we know every parent-child relationship according to the mutual reference, and we know which particles are located in which cell. After we build the tree, we can find there are  two types of cell:\n",
    "* non-leaf cell: $n_{leaf} >= n_{crit}=10$, they have child cells, and they are parents of leaf cells.\n",
    "* leaf cell: $n_{leaf} < n_{crit}=10$, they don't have child cells, and they are the cells where particles live in. And they are the bottom of each branch.\n",
    "\n",
    "In the next notebook, we will discuss how to make full use of this hierarchical tree structure to read the source information and evaluate the multipoles for different cells."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Reference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. R. Yokota, 12 Steps to a Fast Multipole Method on GPUs, Pan-American Advanced Studies Institute, Valparaiso, Chile, 3-14 January, 2011.\n",
    "2. Raykar, V. C., \"[A short primer on the fast multipole method: FMM tutorial](http://www.umiacs.umd.edu/labs/cvl/pirl/vikas/publications/FMM_tutorial.pdf),\", University of Maryland, College Park, Apr. 8, 2006."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from IPython.core.display import HTML\n",
    "def css_styling():\n",
    "    styles = open('./style/fmmstyle.css', 'r').read()\n",
    "    return HTML(styles)\n",
    "css_styling()"
   ]
  }
 ],
 "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.4.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}