{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DALI expressions and arithmetic operators\n",
    "\n",
    "In this example, we will see how to use arithmetic operators in DALI Pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import types\n",
    "import collections\n",
    "import numpy as np\n",
    "from nvidia.dali.pipeline import Pipeline\n",
    "import nvidia.dali.ops as ops            \n",
    "import nvidia.dali.types as types\n",
    "\n",
    "batch_size = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining the iterator\n",
    "\n",
    "We will use custom iterator producing small tensors filled with series of numbers, so we can easily inspect the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ExternalInputIterator(object):\n",
    "    def __init__(self, batch_size, left_type, right_type):\n",
    "        self.batch_size = batch_size\n",
    "        self.left_type = left_type\n",
    "        self.right_type = right_type\n",
    "\n",
    "    def __iter__(self):\n",
    "        self.i = 0\n",
    "        self.n = 128\n",
    "        return self\n",
    "\n",
    "    def __next__(self):\n",
    "        left = []\n",
    "        right = []\n",
    "        for sample in range(self.batch_size):\n",
    "            left.append(np.array([sample + self.i], dtype = self.left_type))\n",
    "            right.append(np.array([self.batch_size + self.i + sample], dtype = self.right_type))\n",
    "        self.i = (self.i + 1) % self.n\n",
    "        return (left, right)\n",
    "    \n",
    "    next = __next__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiating the iterators\n",
    "\n",
    "We create instances of the `ExternalInputIterator` with different type combinations. Type promotions for binary operators are described below. They apply to `+`, `-`, `*` and `//`. The `/` always returns a float32 for integer inputs, and applies the rules below when at least one of the inputs is a floating point number.\n",
    "\n",
    "```\n",
    "  T      op T      = T\n",
    "  floatX op T      = floatX           (where T is not a float)\n",
    "  floatX op floatY = float(max(X, Y))\n",
    "  intX   op intY   = int(max(X, Y))\n",
    "  uintX  op uintY  = uint(max(X, Y))\n",
    "  intX   op uintY  = int2Y            (if X <= Y)\n",
    "  intX   op uintY  = intX             (if X > Y)\n",
    " ```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "iterator_u8_u8 = iter(ExternalInputIterator(batch_size, np.uint8, np.uint8))\n",
    "iterator_u8_i32 = iter(ExternalInputIterator(batch_size, np.uint8, np.int32))\n",
    "iterator_i16_u8 = iter(ExternalInputIterator(batch_size, np.int16, np.uint8))\n",
    "iterator_i32_f32 = iter(ExternalInputIterator(batch_size, np.int32, np.float32))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining the pipeline\n",
    "\n",
    "The next step is to define the Pipeline.\n",
    "\n",
    "We override `Pipeline.iter_setup`, a method called by the pipeline before every `Pipeline.run`, to call the iterator\n",
    "and feed the result to `ExternalSource()` operators, referenced by `self.left` and `self.right`, by using `feed_input`.\n",
    "\n",
    "Note, that we do not need to instantiate any additional operators, we can use regular Python arithmetic expression on the results of other operators in the `define_graph` step.\n",
    "\n",
    "Here we return both of the inputs and the result of `self.right + self.right * self.left`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    " class ExternalSourcePipeline(Pipeline):                   \n",
    "    def __init__(self, iterator, batch_size, num_threads, device_id):\n",
    "        super(ExternalSourcePipeline, self).__init__(batch_size, num_threads, device_id, seed=12)\n",
    "        self.left_source = ops.ExternalSource()\n",
    "        self.right_source = ops.ExternalSource()\n",
    "        self.iterator = iterator\n",
    "\n",
    "    def define_graph(self):                                                                \n",
    "        self.left = self.left_source()\n",
    "        self.right = self.right_source()\n",
    "        return self.left, self.right, self.right + self.right * self.left\n",
    "\n",
    "    def iter_setup(self):\n",
    "        (l, r) = self.iterator.next()\n",
    "        self.feed_input(self.left, l)\n",
    "        self.feed_input(self.right, r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using the pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[1]] + [[1]] * [[0]] = [[1]] of type uint8\n",
      "[[1]] + [[1]] * [[0]] = [[1]] of type int32\n",
      "[[1]] + [[1]] * [[0]] = [[1]] of type int16\n",
      "[[1.]] + [[1.]] * [[0]] = [[1.]] of type float32\n"
     ]
    }
   ],
   "source": [
    "for it in [iterator_u8_u8, iterator_u8_i32, iterator_i16_u8, iterator_i32_f32]:\n",
    "    pipe = ExternalSourcePipeline(it, batch_size=batch_size, num_threads=2, device_id = 0)\n",
    "    pipe.build()                                                        \n",
    "    pipe_out = pipe.run()\n",
    "    l = pipe_out[0].as_array()\n",
    "    r = pipe_out[1].as_array()\n",
    "    out = pipe_out[2].as_array()\n",
    "    print(\"{} + {} * {} = {} of type {}\".format(r, r, l, out, out.dtype))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}