to-jit-or-not-to-jit.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import harmonica as hm\n",
    "from harmonica.forward.prism import jit_prism_gravity, kernel_g_z\n",
    "import numpy as np\n",
    "import verde as vd\n",
    "import numba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/leo/miniconda3/envs/default/lib/python3.8/site-packages/ipykernel/ipkernel.py:287: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
      "  and should_run_async(code)\n"
     ]
    }
   ],
   "source": [
    "coords = vd.grid_coordinates((0, 10, 0, 10), shape=(50, 50), extra_coords=100)\n",
    "coords = tuple(c.ravel() for c in coords)\n",
    "prism = np.empty((10000, 6))\n",
    "prism[:, 0] = 0\n",
    "prism[:, 1] = 1\n",
    "prism[:, 2] = 0\n",
    "prism[:, 3] = 1\n",
    "prism[:, 4] = 0\n",
    "prism[:, 5] = 1\n",
    "density = np.ones(prism.shape[0])\n",
    "field = \"g_z\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.jit(nopython=True)\n",
    "def single_prism(coordinates, prisms, density, kernel, out):\n",
    "    # Iterate over computation points and prisms\n",
    "    for l in range(coordinates[0].size):\n",
    "        # Iterate over the prism boundaries to compute the result of the\n",
    "        # integration (see Nagy et al., 2000)\n",
    "        for i in range(2):\n",
    "            for j in range(2):\n",
    "                for k in range(2):\n",
    "                    shift_east = prisms[1 - i]\n",
    "                    shift_north = prisms[3 - j]\n",
    "                    shift_upward = prisms[5 - k]\n",
    "                    # If i, j or k is 1, the shift_* will refer to the\n",
    "                    # lower boundary, meaning the corresponding term should\n",
    "                    # have a minus sign\n",
    "                    out[l] += (\n",
    "                        density\n",
    "                        * (-1) ** (i + j + k)\n",
    "                        * kernel(\n",
    "                            shift_east - coordinates[0][l],\n",
    "                            shift_north - coordinates[1][l],\n",
    "                            shift_upward - coordinates[2][l],\n",
    "                        )\n",
    "                    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "result_single = np.zeros(coords[0].size)\n",
    "for i in range(density.size):\n",
    "    single_prism(coords, prism[i, :], density[i], kernel_g_z, result_single)\n",
    "    \n",
    "result = np.zeros(coords[0].size)\n",
    "jit_prism_gravity(coords, prism, density, kernel_g_z, result)\n",
    "\n",
    "np.testing.assert_allclose(result, result_single)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/leo/miniconda3/envs/default/lib/python3.8/site-packages/ipykernel/ipkernel.py:287: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
      "  and should_run_async(code)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10.2 s ± 370 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "result_single = np.zeros(coords[0].size)\n",
    "for i in range(density.size):\n",
    "    single_prism(coords, prism[i, :], density[i], kernel_g_z, result_single)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10 s ± 274 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "result = np.zeros(coords[0].size)\n",
    "jit_prism_gravity(coords, prism, density, kernel_g_z, result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:default] *",
   "language": "python",
   "name": "conda-env-default-py"
  },
  "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}