{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cookbook"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We start with importing the CrossPy library, as well as NumPy and CuPy for demostration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import crosspy as xp\n",
    "import numpy as np\n",
    "import cupy as cp"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create arrays\n",
    "\n",
    "We create a NumPy array and a CuPy and then merge them into a CrossPy array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1, 2, 3])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_cpu = np.array([1, 2, 3])\n",
    "x_cpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([4, 5])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "x_gpu = cp.array([4, 5])\n",
    "x_gpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<CUDA Device 0>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_gpu.device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1, 2, 3]); array([4, 5])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_cross = xp.array([x_cpu, x_gpu], axis=0)\n",
    "x_cross"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(3,)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_cpu.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2,)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_gpu.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(5,)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_cross.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([4, 5]); array([4, 5])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = xp.array([x_gpu, x_gpu], axis=0)\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4,)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([4, 5]); array([4, 5])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = xp.array([x_gpu, x_gpu], axis=0)\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4,)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# xp.array([x_cpu, x_cpu.T])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Indexing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = x_cross[0] # slicing with a single integer\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = x_cross[[0, 3]] # slicing with a list of integers\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# x = x_cross[0:2] # slicing with Python built-in slices\n",
    "# x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from crosspy import cpu, gpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# a = xp.array(range(6), distribution=[cpu(0), gpu(0), gpu(1)])\n",
    "# a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# a.device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from crosspy import PartitionScheme\n",
    "# partition = PartitionScheme(6, default_device=cpu(0))\n",
    "\n",
    "# partition[(0, 4, 5)] = cpu(0)\n",
    "# partition[1:3] = gpu(0)\n",
    "# partition[3] = gpu(1)\n",
    "\n",
    "# a = xp.array(range(6), distribution=partition)\n",
    "# a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# a.device"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Arithmetics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# a[0] = a[2] + a[4]\n",
    "# a"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interoperability with NumPy/CuPy"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "CrossPy arrays can be assigned value(s) from/to NumPy/CuPy arrays.\n",
    "\n",
    "When assigning values from NumPy/CuPy arrays to CrossPy arrays, there are two\n",
    "possible behaviors. The first one scatters the data from the source array to the\n",
    "underlying devices of the CrossPy array, i.e., the heterogeneity of the CrossPy\n",
    "array is unchanged. The second one overwrites the corresponding part of the\n",
    "target array with both the data and the device of the source array. The built-in\n",
    "assignment operation in Python is not overloadable and we chose to implement it\n",
    "with the scatter behavior.\n",
    "\n",
    ">>> x_cross[0] = np.array([5, 4, 3, 2, 1])\n",
    ">>> x_cross # doctest: +NORMALIZE_WHITESPACE\n",
    "array {((0, 1), (0, 3)): array([[5, 4, 3]]),\n",
    "       ((0, 1), (3, 5)): array([[2, 1]])}\n",
    ">>> x_cross[0] = cp.array([6, 7, 8, 9, 0])\n",
    ">>> x_cross # doctest: +NORMALIZE_WHITESPACE\n",
    "array {((0, 1), (0, 3)): array([[6, 7, 8]]),\n",
    "       ((0, 1), (3, 5)): array([[9, 0]])}\n",
    "\n",
    "For assigning values from CrossPy arrays to NumPy/CuPy arrays, since the target\n",
    "is distinguishable by devices (NumPy arrays are always on CPU while CuPy arrays are\n",
    "on GPU devices), we use ``to`` to convert the array. We simply use negative integers\n",
    "as CPU devices and otherwise GPU devices.\n",
    "\n",
    ">>> y_cpu = x_cross.to(-1)\n",
    ">>> y_cpu\n",
    "array([[6, 7, 8, 9, 0]])\n",
    ">>> type(y_cpu)\n",
    "<class 'numpy.ndarray'>\n",
    "\n",
    ">>> y_gpu0 = x_cross[:1, (0, 2, 4)].to(0)\n",
    ">>> y_gpu0\n",
    "array([[6, 8, 0]])\n",
    ">>> type(y_gpu0)\n",
    "<class 'cupy.core.core.ndarray'>\n",
    ">>> y_gpu0.device\n",
    "<CUDA Device 0>\n",
    "\n",
    ">>> y_gpu1 = x_cross.to(1)\n",
    ">>> y_gpu1\n",
    "array([[6, 7, 8, 9, 0]])\n",
    ">>> type(y_gpu1)\n",
    "<class 'cupy.core.core.ndarray'>\n",
    ">>> y_gpu1.device\n",
    "<CUDA Device 1>\n",
    "\n",
    "With ``to``, we can use NumPy/CuPy computational functions as usual.\n",
    "\n",
    ">>> np.linalg.norm(x_cross.to(-1))\n",
    "15.165750888103101\n",
    ">>> cp.linalg.norm(x_cross.to(0))\n",
    "array(15.16575089)\n",
    "\n",
    ".. note::\n",
    "    It seems tedious to have the ugly tail ``to``. However, third-party APIs always\n",
    "    have fixed signatures and those in CuPy for example is inherently not compatible\n",
    "    with CrossPy objects (same with NumPy). Therefore, an explicit operation is\n",
    "    necessary to satisfy the input requirements of third-party APIs."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cupy",
   "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.9.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
