no-python-overhead.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['PATH'] = ':'.join([os.environ['PATH'], '/home/adr/julia-1.4.0/bin'])\n",
    "\n",
    "import numba\n",
    "import julia\n",
    "jl = julia.Julia(compiled_modules=False, runtime='julia')\n",
    "from julia import Main\n",
    "from diffeqpy import de"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(du,u,p,t):\n",
    "    x, y, z = u\n",
    "    sigma, rho, beta = p\n",
    "    du[0] = sigma * (y - x)\n",
    "    du[1] = x * (rho - z) - y\n",
    "    du[2] = x * y - beta * z\n",
    "\n",
    "numba_f = numba.jit(f)\n",
    "u0 = [1.0,0.0,0.0]\n",
    "tspan = (0., 100.)\n",
    "p = [10.0,28.0,2.66]\n",
    "prob = de.ODEProblem(numba_f, u0, tspan, p)\n",
    "sol = de.solve(prob)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "60.2 ms ± 7.53 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit de.solve(prob)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "c_sig = numba.types.void(\n",
    "    numba.types.int64,\n",
    "    numba.types.int64,\n",
    "    numba.types.CPointer(numba.types.float64),\n",
    "    numba.types.CPointer(numba.types.float64),\n",
    "    numba.types.CPointer(numba.types.float64),\n",
    "    numba.types.float64\n",
    ")\n",
    "\n",
    "@numba.cfunc(c_sig)\n",
    "def f(n_states, n_params, du_, u_, p_, t):\n",
    "    u = numba.carray(u_, n_states)\n",
    "    du = numba.carray(du_, n_states)\n",
    "    p = numba.carray(p_, n_params)\n",
    "    x, y, z = u\n",
    "    sigma, rho, beta = p\n",
    "    du[0] = sigma * (y - x)\n",
    "    du[1] = x * (rho - z) - y\n",
    "    du[2] = x * y - beta * z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "Main.rhs_func_p = f.address"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "jul_f = Main.eval(\"\"\"\n",
    "function f(du,u,p,t)\n",
    "  n_states = length(du)\n",
    "  n_params = length(p)\n",
    "  ccall(\n",
    "      convert(Ptr{Nothing}, rhs_func_p),\n",
    "      Nothing,\n",
    "      (Int64, Int64, Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Float64),\n",
    "      n_states,\n",
    "      n_params,\n",
    "      du,\n",
    "      u,\n",
    "      p,\n",
    "      t\n",
    "    )\n",
    "end\n",
    "\"\"\")\n",
    "u0 = [1.0,0.0,0.0]\n",
    "tspan = (0., 100.)\n",
    "p = [10.0,28.0,2.66]\n",
    "prob = de.ODEProblem(jul_f, u0, tspan, p)\n",
    "sol = de.solve(prob)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.64 ms ± 549 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit de.solve(prob)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}