{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8ad9f0d7",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Dr.Jit quickstart\n",
    "\n",
    "## Overview\n",
    "\n",
    "This short tutorial recaps the basic functionalities and routines of the Dr.Jit library. You can also find more information on the [Dr.Jit documentation](https://drjit.readthedocs.io)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca65dd2f",
   "metadata": {},
   "source": [
    "## Similarity with NumPy\n",
    "\n",
    "On the Python side, the Dr.Jit *syntax* is very similar to NumPy. Moreover, as we will see later, both frameworks are interoperable.\n",
    "\n",
    "Let's first import both NumPy and Dr.Jit using the alias `np` and `dr` respectively"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6cd3093c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import drjit as dr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a51da003",
   "metadata": {},
   "source": [
    "Unlike NumPy, Dr.Jit can perform array arithmetic on both CPU and GPU through various template variants which are exposed in top-level packages:\n",
    "\n",
    "| Variant         | Description                                                        |\n",
    "| --------------- | ------------------------------------------------------------------ |\n",
    "| `drjit.scalar`  | Arrays built on top of scalars (float, int, etc.)                  |\n",
    "| `drjit.llvm`    | Arrays built on top of LLVMArray                                   |\n",
    "| `drjit.cuda`    | Arrays built on top of CUDAArray                                   |\n",
    "| `drjit.llvm.ad` | Similar to `drjit.llvm` but with automatic differentiation support |\n",
    "| `drjit.cuda.ad` | Similar to `drjit.cuda` but with automatic differentiation support |\n",
    "\n",
    "These packages all contains various types like: `Bool, Float, Int, UInt, Array2f, Array2i, Matrix2f Matrix3f, ...`\n",
    "\n",
    "Let's create some arrays using the `drjit.llvm` variants and play around with the NumPy interoperability:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a21e1b72",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "c -> (<class 'drjit.llvm.Float'>) = [9.0, 8.0, 7.0, 6.0]\n",
      "d -> (<class 'numpy.ndarray'>) = [9. 8. 7. 6.]\n"
     ]
    }
   ],
   "source": [
    "from drjit.llvm import Float, UInt32\n",
    "\n",
    "# Create some floating-point arrays\n",
    "a = Float([1.0, 2.0, 3.0, 4.0])\n",
    "b = Float([4.0, 3.0, 2.0, 1.0])\n",
    "\n",
    "# Perform simple arithmetic\n",
    "c = a + 2.0 * b\n",
    "\n",
    "print(f'c -> ({type(c)}) = {c}')\n",
    "\n",
    "# Convert to NumPy array\n",
    "d = np.array(c)\n",
    "\n",
    "print(f'd -> ({type(d)}) = {d}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f6d429d",
   "metadata": {},
   "source": [
    "## Array construction routines\n",
    "\n",
    "This section provides an overview of various Dr.Jit routines (and their NumPy correspondence) to construct arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4df910ea",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dr.zeros: [0.0, 0.0, 0.0, 0.0, 0.0]\n",
      "dr.full: [0.10000000149011612, 0.10000000149011612, 0.10000000149011612, 0.10000000149011612, 0.10000000149011612]\n",
      "dr.arange: [0, 1, 2, 3, 4]\n",
      "dr.linespace: [0.0, 0.5, 1.0, 1.5, 2.0]\n"
     ]
    }
   ],
   "source": [
    "# Initialize floating-point array of size 5 with zeros\n",
    "a = dr.zeros(Float, 5) # np.zeros(5)\n",
    "print(f'dr.zeros: {a}')\n",
    "\n",
    "# Initialize floating-point array of size 5 with a constant value\n",
    "a = dr.full(Float, 0.1, 5) # np.ones(5, 0.4)\n",
    "print(f'dr.full: {a}')\n",
    "\n",
    "a = dr.arange(UInt32, 5) # np.arange(5)\n",
    "print(f'dr.arange: {a}')\n",
    "\n",
    "# Return evenly spaced numbers over a specified interval\n",
    "a = dr.linspace(Float, 0.0, 2.0, 5) # np.linspace(0.0, 2.0, 5)\n",
    "print(f'dr.linespace: {a}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c4ccd24",
   "metadata": {},
   "source": [
    "## Masking\n",
    "\n",
    "Writing codes using Dr.Jit often means working with large arrays at once. Therefore it is not possible to use regular `if .. else ..` statements based on concret values, as different elements in the array might branch differently. This is where **masking** comes to the rescue! \n",
    "\n",
    "A mask (or `Bool`) is an array of boolean values that can be used to disable arithmetic operations on part of an array. It is possible to create such masks with any regular boolean arithmetic (e.g. `>, <, >=, <=`). \n",
    "\n",
    "Often time, we combine masks with the `dr.select(mask, a, b)` statement which correspond to the ternary statement `mask ? a : b`. This is similar to the `np.where` function in NumPy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8a676bc8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x -> (<class 'drjit.llvm.Float'>) [0.0, 1.0, 2.0, 3.0, 4.0]\n",
      "m -> (<class 'drjit.llvm.Bool'>) [False, False, False, True, True]\n",
      "y -> (<class 'drjit.llvm.Float'>) [1.0, 1.0, 1.0, 4.0, 4.0]\n"
     ]
    }
   ],
   "source": [
    "x = dr.arange(Float, 5)\n",
    "m = x > 2.0 # True for all values of a that are greater than 2.0\n",
    "y = dr.select(m, 4.0, 1.0) # Set the values greater than 2.0 to 4.0 otherwise to 1.0\n",
    "print(f'x -> ({type(x)}) {x}')\n",
    "print(f'm -> ({type(m)}) {m}')\n",
    "print(f'y -> ({type(y)}) {y}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "833bf3a5",
   "metadata": {},
   "source": [
    "## Basic math arithmetic\n",
    "\n",
    "All common math operators like `+, -, /, *, *=, +=, %, //, ...` are supported with Dr.Jit arrays.\n",
    "\n",
    "Similarly to NumPy, Dr.Jit provides all kinds of math arithmetic that can be performed on the entire array in a single call. Here is a non-exaustive list of those math functions:\n",
    "`abs, minimum, maximum, sqrt, pow, sin, cos, tan, atan2, sincos, sec, cot, asin, acos, atan, exp, exp2, log, log2, sinh, cosh, tanh, asinh, acosh, atanh, ...`\n",
    "\n",
    "Those routines are present in the root drjit package, hence can be used as follow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "43bae3a6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "m: [0.0, 0.4794255495071411, 0.5403022766113281, 0.07073719799518585, -0.41614681482315063]\n"
     ]
    }
   ],
   "source": [
    "s, c = dr.sincos(a)\n",
    "m = dr.minimum(s, c)\n",
    "print(f'm: {m}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "694011b4",
   "metadata": {},
   "source": [
    "## Horizontal operations\n",
    "\n",
    "Dr.Jit also provides operations that require a pass over the entire array and return a single scalar value. Those operations are expensive as they will trigger a syncronization point, hence it is better to avoid them if possible.\n",
    "\n",
    "The following snippet of code explores a few of those:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "93016119",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a: [1.0, 2.0, 3.0, 4.0, 5.0]\n",
      "dr.sum(a): [15.0]\n",
      "dr.prod(a): [120.0]\n",
      "dr.mean(a): [3.0]\n",
      "m: [False, False, True, True, True]\n",
      "dr.all(m): False\n",
      "dr.any(m): True\n",
      "dr.none(m): False\n"
     ]
    }
   ],
   "source": [
    "a = dr.arange(Float, 5) + 1\n",
    "print(f'a: {a}')\n",
    "\n",
    "# Horizontal sum\n",
    "b = dr.sum(a) # np.sum(a)\n",
    "print(f'dr.sum(a): {b}')\n",
    "\n",
    "# Horizontal product\n",
    "b = dr.prod(a) # np.prod(a)\n",
    "print(f'dr.prod(a): {b}')\n",
    "\n",
    "# Mean value over the entire array\n",
    "b = dr.mean(a) # np.mean(a)\n",
    "print(f'dr.mean(a): {b}')\n",
    "\n",
    "m = a > 2\n",
    "print(f'm: {m}')\n",
    "\n",
    "# True if all value of the mask array are True\n",
    "b = dr.all(m) # np.all(m)\n",
    "print(f'dr.all(m): {b}')\n",
    "\n",
    "# True if any value of the mask array are True\n",
    "b = dr.any(m) # np.any(m)\n",
    "print(f'dr.any(m): {b}')\n",
    "\n",
    "# True if no value of the mask array are True\n",
    "b = dr.none(m) # ~np.any(m)\n",
    "print(f'dr.none(m): {b}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bc55c1a",
   "metadata": {},
   "source": [
    "## `gather` and `scatter` routines\n",
    "\n",
    "In programming languages like C++ or Python, it is possible to access the i-th element of an array using the `array[i]` syntax. This can both be used to read or write values in an array. Similarly, Dr.Jit provides such read/write functionalities through the `dr.gather` and `dr.scatter` functions. Those are much more powerful than the regular array accessors as the index `i` can be an array itself! In which case the read operation (e.g. `dr.gather`) would return a array as well, not just a single value.\n",
    "\n",
    "Here is how one should use the `dr.gather` routine to read entries from an Dr.Jit array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7f5a350f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "source: [0.0, 0.25, 0.5, 0.75, 1.0]\n",
      "indices: [1, 2]\n",
      "result: [0.25, 0.5]\n"
     ]
    }
   ],
   "source": [
    "source = dr.linspace(Float, 0, 1, 5)\n",
    "indices = UInt32([1, 2]) # Only read the 2nd and 3rd elements of the source array\n",
    "result = dr.gather(Float, source, indices)\n",
    "print(f'source: {source}')\n",
    "print(f'indices: {indices}')\n",
    "print(f'result: {result}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07f6203b",
   "metadata": {},
   "source": [
    "And here is how one can write entries at specific indices into a Dr.Jit array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "cce7f415",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "indices: [0, 3, 4]\n",
      "source: [1.0, 2.0, 3.0]\n",
      "target: [1.0, 0.0, 0.0, 2.0, 3.0]\n"
     ]
    }
   ],
   "source": [
    "target = dr.zeros(Float, 5)\n",
    "indices = UInt32([0, 3, 4]) # Write to the first and last two elements of the target array\n",
    "source = Float([1.0, 2.0, 3.0])\n",
    "dr.scatter(target, source, indices)\n",
    "print(f'indices: {indices}')\n",
    "print(f'source: {source}')\n",
    "print(f'target: {target}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}