Interpolation

Untitled.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numba\n",
    "from numba import jit, guvectorize, float64, generated_jit\n",
    "\n",
    "@jit(nopython=True)\n",
    "def scalar_interp_1d(u,x,y):\n",
    "    K_0 = x.shape[0]\n",
    "    ii0 = np.searchsorted(x,u)\n",
    "    i0 = np.maximum(0, np.minimum(ii0-1, K_0-2))\n",
    "    x0, x1 = (x[i0], x[i0+1])\n",
    "    y0, y1 = (y[i0], y[i0+1])\n",
    "    l0, = ((u-x0)/(x1-x0), ) # barycentric coordinate\n",
    "    res = y0 + l0*(y1-y0)\n",
    "    return res\n",
    "\n",
    "\n",
    "@jit(nopython=True)\n",
    "def vec_interp_1d(u,x,y):\n",
    "    N = u.shape[0]\n",
    "    out = np.zeros(N)\n",
    "    for n in range(N):\n",
    "        out[n] = scalar_interp_1d(u[n], x, y)\n",
    "    return out\n",
    "\n",
    "\n",
    "@generated_jit\n",
    "def interp_1d(u,x,y):\n",
    "    if isinstance(u,  numba.types.Array) and u.ndim == 1:\n",
    "        return vec_interp_1d\n",
    "    elif isinstance(u, numba.types.Float):\n",
    "        return scalar_interp_1d\n",
    "    else:\n",
    "        raise Exception(f\"Unsupported type for u: {u}.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.8611239133554044"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = np.linspace(-4, 4, 10)\n",
    "y = np.cos(x)\n",
    "interp_1d(0.5, x, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-0.65364362, -0.81748713, -0.98133064, -0.83397145, -0.64771193,\n",
       "       -0.29627306,  0.10235992,  0.44606245,  0.76229975,  0.90284967,\n",
       "        0.90284967,  0.76229975,  0.44606245,  0.10235992, -0.29627306,\n",
       "       -0.64771193, -0.83397145, -0.98133064, -0.81748713, -0.65364362])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "u = np.linspace(-4, 4, 20)\n",
    "interp_1d(u, x, y)"
   ]
  },
  {
   "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}