{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Usual magic stuff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pythran"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext pythran.magic"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a pythran capsule, not callable as a Python function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%pythran\n",
    "#pythran export capsule f(int32, float64*, float64* )\n",
    "def f(n, x, cp):\n",
    "    c = cp[0]\n",
    "    return c + x[0] - x[1] * x[2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prepare to pass this function to scipy's LowLevel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ctypes\n",
    "c = ctypes.c_double(1.0)\n",
    "user_data = ctypes.cast(ctypes.pointer(c), ctypes.c_void_p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1200.0000000000002, 1.3322676295501882e-11)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from scipy import integrate, LowLevelCallable\n",
    "func = LowLevelCallable(f, user_data, signature=\"double (int, double *, void *)\")\n",
    "dat =  [[0, 10], [-10, 0], [-1, 1]]\n",
    "integrate.nquad(func, dat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Showcase Numpy integration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%pythran\n",
    "\n",
    "# using numpy adaptator\n",
    "#pythran export capsule transform(int64*, float64*, int32, int32, float64*)\n",
    "from numpy.ctypeslib import as_array\n",
    "def transform(output_coordinates, input_coordinates, output_rank, input_rank, user_data):\n",
    "    shift = user_data[0]\n",
    "    input_data = as_array(input_coordinates, input_rank)\n",
    "    output_data = as_array(output_coordinates, output_rank)\n",
    "    input_data[:] = output_data - shift\n",
    "    return 1\n",
    "\n",
    "# same with explicit loops\n",
    "#pythran export capsule transform_basic(int64*, float64*, int32, int32, float64*)\n",
    "def transform_basic(output_coordinates, input_coordinates, output_rank, input_rank, user_data):\n",
    "    shift = user_data[0]\n",
    "    for i in range(input_rank):\n",
    "        input_coordinates[i] = output_coordinates[i] - shift;\n",
    "    return 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Validate output using the Numpy API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.    , 0.    , 0.    ],\n",
       "       [0.    , 1.3625, 2.7375],\n",
       "       [0.    , 4.8125, 6.1875],\n",
       "       [0.    , 8.2625, 9.6375]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import ctypes\n",
    "import numpy as np\n",
    "from scipy import ndimage, LowLevelCallable\n",
    "\n",
    "\n",
    "shift = 0.5\n",
    "\n",
    "user_data = ctypes.c_double(shift)\n",
    "ptr = ctypes.cast(ctypes.pointer(user_data), ctypes.c_void_p)\n",
    "callback = LowLevelCallable(transform, ptr, \"int (npy_intp *, double *, int, int, void *)\")\n",
    "im = np.arange(12).reshape(4, 3).astype(np.float64)\n",
    "\n",
    "out0 = ndimage.geometric_transform(im, callback)\n",
    "out0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And using the explicit looping. Hopefully the results are the same!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.    , 0.    , 0.    ],\n",
       "       [0.    , 1.3625, 2.7375],\n",
       "       [0.    , 4.8125, 6.1875],\n",
       "       [0.    , 8.2625, 9.6375]])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "callback = LowLevelCallable(transform_basic, ptr, \"int (npy_intp *, double *, int, int, void *)\")\n",
    "im = np.arange(12).reshape(4, 3).astype(np.float64)\n",
    "\n",
    "out1 = ndimage.geometric_transform(im, callback)\n",
    "out1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert np.all(out0 == out1)"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
