{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import ipyvolume.pylab as p3\n",
    "from numpy.testing import assert_array_equal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from ipyvolume.pylab import make_triangles_lines as new_make_triangles_lines\n",
    "def old_make_triangles_lines(x,y,z,wrapx=False,wrapy=False):\n",
    "    def dim(x):\n",
    "        d = 0\n",
    "        el = x\n",
    "        while True:\n",
    "            try:\n",
    "                el = el[0]\n",
    "                d += 1\n",
    "            except Exception:\n",
    "                break\n",
    "        return d\n",
    "\n",
    "    if dim(x) == 2:\n",
    "        nx, ny = shape = x.shape\n",
    "    else:\n",
    "        nx, ny = shape = x[0].shape\n",
    "\n",
    "    def reshape(ar):\n",
    "        if dim(ar) == 3:\n",
    "            return [k.reshape(-1) for k in ar]\n",
    "        else:\n",
    "            return ar.reshape(-1)\n",
    "\n",
    "    x = reshape(x)\n",
    "    y = reshape(y)\n",
    "    z = reshape(z)\n",
    "\n",
    "    mx = nx if wrapx else nx - 1\n",
    "    my = ny if wrapy else ny - 1\n",
    "    triangles = np.zeros(((mx) * (my) * 2, 3), dtype=np.uint32)\n",
    "    lines = np.zeros(((mx) * (my) * 4, 2), dtype=np.uint32)\n",
    "    \n",
    "    \n",
    "      \n",
    " \n",
    "    def index_from2d(i, j):\n",
    "        xi = (i % nx)\n",
    "        yi = (j % ny)\n",
    "        return ny * xi + yi\n",
    "        \"\"\"\n",
    "        ^ ydir\n",
    "        |\n",
    "        2 3\n",
    "        0 1  ---> x dir\n",
    "        \"\"\"\n",
    " \n",
    "    for i in range(mx):\n",
    "        for j in range(my):\n",
    "            p0 = index_from2d(i, j)\n",
    "            p1 = index_from2d(i + 1, j)\n",
    "            p2 = index_from2d(i, j + 1)\n",
    "            p3 = index_from2d(i + 1, j + 1)\n",
    "            triangle_index = (i * my) + j\n",
    "            triangles[triangle_index * 2 + 0, :] = [p0, p1, p3]\n",
    "            triangles[triangle_index * 2 + 1, :] = [p0, p3, p2]\n",
    "            lines[triangle_index * 4 + 0, :] = [p0, p1]\n",
    "            lines[triangle_index * 4 + 1, :] = [p0, p2]\n",
    "            lines[triangle_index * 4 + 2, :] = [p2, p3]\n",
    "            lines[triangle_index * 4 + 3, :] = [p1, p3]\n",
    "    \n",
    "    return triangles, lines\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "assert_array_equal(tri_new,tri_old)\n",
    "assert_array_equal(line_new,line_old)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "330 µs ± 574 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n",
      "44.9 ms ± 135 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "# Small data set.\n",
    "nx = 100\n",
    "ny = 50\n",
    "X = np.linspace(-5, 5, nx)\n",
    "Y = np.linspace(-5, 5, ny)\n",
    "X, Y = np.meshgrid(X, Y, indexing='xy')\n",
    "R = np.sqrt(X**2 + Y**2)\n",
    "Z = np.sin(R)\n",
    "\n",
    "tri_new, line_new = new_make_triangles_lines(X,Y,Z)\n",
    "tri_old, line_old = old_make_triangles_lines(X,Y,Z)\n",
    "\n",
    "assert_array_equal(tri_new,tri_old)\n",
    "assert_array_equal(line_new,line_old)\n",
    "\n",
    "\n",
    "%timeit new_make_triangles_lines(X,Y,Z)\n",
    "%timeit old_make_triangles_lines(X,Y,Z)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12.4 ms ± 80.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n",
      "1.56 s ± 2.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "# Larger data set.\n",
    "nx = 400\n",
    "ny = 425\n",
    "X = np.linspace(-5, 5, nx)\n",
    "Y = np.linspace(-5, 5, ny)\n",
    "X, Y = np.meshgrid(X, Y, indexing='xy')\n",
    "R = np.sqrt(X**2 + Y**2)\n",
    "Z = np.sin(R)\n",
    "\n",
    "tri_new, line_new = new_make_triangles_lines(X,Y,Z)\n",
    "tri_old, line_old = old_make_triangles_lines(X,Y,Z)\n",
    "\n",
    "assert_array_equal(tri_new,tri_old)\n",
    "assert_array_equal(line_new,line_old)\n",
    "\n",
    "%timeit tri_new, line_new = new_make_triangles_lines(X,Y,Z)\n",
    "%timeit tri_old, line_old = old_make_triangles_lines(X,Y,Z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}