test_02.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1dd9799f-7171-4c3d-8da5-6c14dd5d9f15",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numba as nb\n",
    "import numpy as np\n",
    "\n",
    "from math import sin, cos"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5320e1b0-e703-4942-8fab-f76d30dfe5b9",
   "metadata": {
    "tags": []
   },
   "source": [
    "# PoC: New JIT decorator with different default behavior\n",
    "\n",
    "The following example introduces a new decorator, `cjit`, which maintains all semantics of `numba.jit` and therefore remains completely compatible with it. The only difference is that it allows to set new default values, differing from those in `numba`, for arguments in one central place."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "dcdb3328-ae97-44f3-83a1-4763438a70ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cjit(*args, nopython = True, inline = 'always', **kwargs):\n",
    "    \n",
    "    if len(args) == 1 and callable(args[0]):\n",
    "        func = args[0]\n",
    "        args = tuple()\n",
    "    else:\n",
    "        func = None\n",
    "\n",
    "    def wrapper(func):\n",
    "        return nb.jit(\n",
    "            *args,\n",
    "            nopython = nopython,\n",
    "            inline = inline,\n",
    "            **kwargs,\n",
    "        )(func)\n",
    "    \n",
    "    if func is not None:\n",
    "        return wrapper(func)\n",
    "        \n",
    "    return wrapper"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc19f202-6f50-47e9-822d-2e5906f7fc0f",
   "metadata": {},
   "source": [
    "To conform that it is working:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5ab8a075-1bca-4a5c-8107-bb4c5276ea63",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(3,\n",
       " CPUDispatcher(<function foo at 0x7ff08873b520>),\n",
       " 3.0,\n",
       " CPUDispatcher(<function bar at 0x7ff08873a7a0>))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@cjit\n",
    "def foo(x, y):\n",
    "    return x - y\n",
    "\n",
    "@cjit('f8(f8,f8)')\n",
    "def bar(x, y):\n",
    "    return x - y\n",
    "\n",
    "foo(7, 4), foo, bar(7, 4), bar"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4965b3b4-a580-4352-9cfc-fe3796c8c250",
   "metadata": {},
   "source": [
    "# Proposed central `poliastro._jit` module\n",
    "\n",
    "Taking the previous idea one step further, one could introduce literally one central place for switching to different compiler backends, i.e. `cpu` (single-thread), `parallel` (cpu multi-thread) and `cuda` across the entire package. All sub-modules within `poliastro` would import those decorators instead of `jit` and friends from `numba` directly.\n",
    "\n",
    "Change the value of `TARGET`, re-start the kernel (!) and re-run the notebook to see effects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1065cfc7-eac3-451e-8a48-73ded5499389",
   "metadata": {},
   "outputs": [],
   "source": [
    "TARGET = 'cpu'  # os.environ.get('POLIASTRO_TARGET', 'cpu')\n",
    "\n",
    "if TARGET not in ('cpu', 'parallel', 'cuda'):\n",
    "    raise ValueError(f'unknown target \"{TARGET:s}\"')\n",
    "if TARGET == 'cuda':\n",
    "    from numba import cuda  # explicit import required and only performed if target is switched to cuda\n",
    "\n",
    "NOPYTHON = True  # only for debugging, True by default\n",
    "\n",
    "def hjit(*args, **kwargs):\n",
    "    \"\"\"\n",
    "    Scalar helper, pre-configured, internal.\n",
    "    Functions decorated by it can only be called directly if TARGET is cpu or parallel.\n",
    "    \"\"\"\n",
    "    \n",
    "    if len(args) == 1 and callable(args[0]):\n",
    "        func = args[0]\n",
    "        args = tuple()\n",
    "    else:\n",
    "        func = None\n",
    "\n",
    "    def wrapper(func):\n",
    "        \n",
    "        cfg = {}\n",
    "        if TARGET in ('cpu', 'parallel'):\n",
    "            cfg.update({'nopython': NOPYTHON, 'inline': 'always'})\n",
    "        if TARGET == 'cuda':\n",
    "            cfg.update({'device': True, 'inline': True})\n",
    "        cfg.update(kwargs)\n",
    "        \n",
    "        wjit = cuda.jit if TARGET == 'cuda' else nb.jit\n",
    "        \n",
    "        return wjit(\n",
    "            *args,\n",
    "            **cfg,\n",
    "        )(func)\n",
    "    \n",
    "    if func is not None:\n",
    "        return wrapper(func)\n",
    "        \n",
    "    return wrapper\n",
    "\n",
    "def vjit(*args, **kwargs):\n",
    "    \"\"\"\n",
    "    Vectorize on array, pre-configured, user-facing.\n",
    "    Functions decorated by it can always be called directly if needed.\n",
    "    \"\"\"\n",
    "    \n",
    "    if len(args) == 1 and callable(args[0]):\n",
    "        func = args[0]\n",
    "        args = tuple()\n",
    "    else:\n",
    "        func = None\n",
    "\n",
    "    def wrapper(func):\n",
    "        \n",
    "        cfg = {'target': TARGET}\n",
    "        if TARGET in ('cpu', 'parallel'):\n",
    "            cfg.update({'nopython': NOPYTHON})\n",
    "        cfg.update(kwargs)\n",
    "        \n",
    "        return nb.vectorize(\n",
    "            *args,\n",
    "            **cfg,\n",
    "        )(func)\n",
    "    \n",
    "    if func is not None:\n",
    "        return wrapper(func)\n",
    "\n",
    "    return wrapper"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7d2b87e-08da-4181-bd58-0c6547c6eedf",
   "metadata": {},
   "source": [
    "# Example usage\n",
    "\n",
    "In the below example, the function `internal_on_scalar` is strictly speaking not required. Its functionallity could be moved into `user_facing_on_array`. It simply serves to show how a helper function would look like and how it would be decorated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7c22219b-9bd2-4905-85cc-b163105d6917",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11., 12.,\n",
       "        13., 14., 15., 16., 17., 18., 19., 20., 21., 22., 23., 24., 25.,\n",
       "        26., 27., 28., 29., 30., 31., 32., 33., 34., 35., 36., 37., 38.,\n",
       "        39., 40., 41., 42., 43., 44., 45., 46., 47., 48., 49., 50., 51.,\n",
       "        52., 53., 54., 55., 56., 57., 58., 59., 60., 61., 62., 63., 64.,\n",
       "        65., 66., 67., 68., 69., 70., 71., 72., 73., 74., 75., 76., 77.,\n",
       "        78., 79., 80., 81., 82., 83., 84., 85., 86., 87., 88., 89., 90.,\n",
       "        91., 92., 93., 94., 95., 96., 97., 98., 99.]),\n",
       " array([ 0.00000000e+00,  0.00000000e+00, -5.40302306e-01,  3.68995121e-01,\n",
       "         1.35898762e+00,  6.02185122e-01,  3.18522937e-01,  3.91074386e-02,\n",
       "        -7.14794816e-01,  2.74563431e-01,  1.18569369e+00,  6.41672582e-01,\n",
       "         6.37246884e-01,  1.00673966e-01, -8.06772816e-01,  1.83834540e-01,\n",
       "         9.43522453e-01,  6.55619136e-01,  9.30782474e-01,  1.79795228e-01,\n",
       "        -8.08909391e-01,  1.04035860e-01,  6.51765120e-01,  6.42913811e-01,\n",
       "         1.17574683e+00,  2.70168469e-01, -7.21034342e-01,  4.15241080e-02,\n",
       "         3.33662917e-01,  6.04568705e-01,  1.35262623e+00,  3.64594611e-01,\n",
       "        -5.50147747e-01,  1.27893409e-03,  1.45556813e-02,  5.43638367e-01,\n",
       "         1.44733057e+00,  4.55551719e-01, -3.09862333e-01, -1.34937542e-02,\n",
       "        -2.80136687e-01,  4.64976474e-01,  1.45231575e+00,  5.35794204e-01,\n",
       "        -1.93190980e-02, -1.61717285e-03, -5.26939162e-01,  3.74849186e-01,\n",
       "         1.36718466e+00,  5.98929994e-01,  2.98337450e-01,  3.59625964e-02,\n",
       "        -7.06191600e-01,  2.80435992e-01,  1.19871878e+00,  6.39929729e-01,\n",
       "         6.17802973e-01,  9.62519706e-02, -8.03614856e-01,  1.89257792e-01,\n",
       "         9.60338015e-01,  6.55527394e-01,  9.13629030e-01,  1.74448333e-01,\n",
       "        -8.11448249e-01,  1.08577789e-01,  6.71031641e-01,  6.44480487e-01,\n",
       "         1.16225029e+00,  2.64322606e-01, -7.29067774e-01,  4.48229076e-02,\n",
       "         3.53845636e-01,  6.07668999e-01,  1.34386172e+00,  3.58715456e-01,\n",
       "        -5.63035813e-01,  3.07182351e-03,  3.40468552e-02,  5.48025311e-01,\n",
       "         1.44399626e+00,  4.50107604e-01, -3.26578378e-01, -1.33495955e-02,\n",
       "        -2.62889713e-01,  4.70300607e-01,  1.45467725e+00,  5.31218803e-01,\n",
       "        -3.85315313e-02, -3.13322854e-03, -5.13310273e-01,  3.80686390e-01,\n",
       "         1.37505385e+00,  5.95587781e-01,  2.78159080e-01,  3.29070944e-02,\n",
       "        -6.97266467e-01,  2.86321279e-01,  1.21146882e+00,  6.38086943e-01]))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@hjit('f8(f8)')\n",
    "def internal_on_scalar(scalar: float) -> float:\n",
    "    res: float = 0.0\n",
    "    for idx in range(round(scalar)):\n",
    "        if idx % 2 == 0:\n",
    "            res += sin(idx)\n",
    "        else:\n",
    "            res -= cos(idx)\n",
    "    return res\n",
    "\n",
    "@vjit('f8(f8)')\n",
    "def user_facing_on_array(d: float) -> float:\n",
    "    return internal_on_scalar(d)\n",
    "\n",
    "data = np.arange(0, 100, 1, dtype = 'f8')\n",
    "result = user_facing_on_array(data)\n",
    "\n",
    "data, result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7d1a19f-4998-4209-9844-f7c2421959eb",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}