{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0f77701d-6ef5-4c69-ab87-c6b44d713383",
   "metadata": {},
   "source": [
    "Usual notebook magic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "18f82709-89f4-4ae3-b098-be34f04280b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pythran"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5be9f530-f60b-4102-8aeb-a4deb521075f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext pythran.magic"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b93ab90-65e5-4b15-9114-d565f5d6a77c",
   "metadata": {},
   "source": [
    "Define the core of a ufunc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8058809d-6d73-4919-84a2-97583a1609a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%pythran\n",
    "#pythran export ufunc incr(float64)\n",
    "def incr(v):\n",
    "    \"increment a value by one\"\n",
    "    return v + 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03901745-1d32-4ceb-aa84-ff0532de43c6",
   "metadata": {},
   "source": [
    "Ufunc work on scalar:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c0fb2efd-2268-4a28-aba7-02820c191fe2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.0\n"
     ]
    }
   ],
   "source": [
    "print(incr(3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d8c4ab50-c09c-46e4-9acb-5dd99d7c3ab5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"increment a value by one\" in incr.__doc__"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6a41d1b-11dd-4f48-9311-93a48cb2fa60",
   "metadata": {},
   "source": [
    "And on numpy arrays:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e79ffaab-94dd-466d-8de4-cd946d3c9172",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2., 2., 2., 2., 2., 2., 2., 2., 2., 2.])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "incr(np.ones(10))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "664119ec-4c9b-4df7-9a3a-cabdc6d7e4ce",
   "metadata": {},
   "source": [
    "Let's write a naive candidate for broadcasting:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ee1212c2-c8a2-43e0-b4f1-6ff6537f8a55",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%pythran\n",
    "#pythran export ufunc avg(int, int)\n",
    "def avg(i, j):\n",
    "    return type(i)((i + j) // 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5a80fe74-2437-48d8-9996-1979c4f29255",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([5, 5, 5, 5])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "avg(np.ones(4, dtype=int), 9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "69becb9c-5056-48c0-a2ba-b51e95d8cc25",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[3, 3, 3, 3],\n",
       "       [3, 3, 3, 3],\n",
       "       [3, 3, 3, 3],\n",
       "       [3, 3, 3, 3]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "avg(np.ones((4, 4), dtype=int), 5 * np.ones(4, dtype=int))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2744cf85-db77-465b-b669-381c450a0fe2",
   "metadata": {},
   "source": [
    "It's possible to support multiple signatures for the same function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3f1a73f3-4aca-42f5-9415-1eeee4381a2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%pythran\n",
    "#pythran export ufunc decr(int32)\n",
    "#pythran export ufunc decr(uint32)\n",
    "def decr(x):\n",
    "    from numpy import int32, uint32\n",
    "    if isinstance(x, int32):\n",
    "        return int32(x - 1)\n",
    "    else:\n",
    "        return uint32(x if x == 0 else x - 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "65557bed-3910-481b-b95b-7edc508f5e16",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-1\n",
      "0\n"
     ]
    }
   ],
   "source": [
    "print(decr(np.int32(0)))\n",
    "print(decr(np.uint32(0)))"
   ]
  }
 ],
 "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
