04_ode.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3144bb5e-9afe-47e6-a45e-0ba66de386de",
   "metadata": {},
   "source": [
    "# Failure: not all code is faster\n",
    "\n",
    "Let's look at a problem that is slower on a GPU when you convert to CuPy. This is a simple ODE solver."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dcb100d-b4cf-4135-a8e9-1fc58da6d69e",
   "metadata": {},
   "source": [
    "## Classic NumPy code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d72dca1-3074-416b-a5d4-78892de62ad2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb85fd9e-d62c-4de8-8924-a04052a3402f",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_max = 1  # Size of x max\n",
    "v_0 = 0\n",
    "koverm = 1  # k / m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4ac304d-2faa-464f-8f86-caaa5c8ab670",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(t, y):\n",
    "    \"Y has two elements, x and v\"\n",
    "    return np.array([-koverm * y[1], y[0]])\n",
    "\n",
    "\n",
    "def euler_ivp(f, init_y, t):\n",
    "    steps = len(t)\n",
    "    order = len(init_y)  # Number of equations\n",
    "\n",
    "    y = np.empty((steps, order))\n",
    "    y[0] = init_y  # Note that this sets the elements of the first row\n",
    "\n",
    "    for n in range(steps - 1):\n",
    "        h = t[n + 1] - t[n]\n",
    "\n",
    "        # Compute dydt based on *current* position\n",
    "        dydt = f(t[n], y[n])\n",
    "\n",
    "        # Compute next velocity and position\n",
    "        y[n + 1] = y[n] - dydt * h\n",
    "\n",
    "    return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa07c447-c9b5-4cd1-a93b-46e2ec34b0ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = np.linspace(0, 40, 1000 + 1)\n",
    "y = euler_ivp(f, [x_max, v_0], ts)\n",
    "plt.plot(ts, np.cos(ts))\n",
    "plt.plot(ts, y[:, 0], \"--\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc7e4e91-e9cf-4ee7-b144-9cb71733a456",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "y = euler_ivp(f, [x_max, v_0], ts)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b90716ba-151d-4efa-97b4-17cd400c3650",
   "metadata": {},
   "source": [
    "## CuPy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d776345-b11b-4df4-8b84-9f5419804bcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cupy as cp\n",
    "import cupyx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9018b6c-d927-4def-84dd-ca4eb9724e66",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(t, y):\n",
    "    \"Y has two elements, x and v\"\n",
    "    # xp = cp.get_array_module(t)\n",
    "    return cp.array([-koverm * y[1], y[0]])\n",
    "\n",
    "\n",
    "def euler_ivp(f, init_y, t):\n",
    "    # xp = cp.get_array_module(t)\n",
    "    steps = len(t)\n",
    "    order = len(init_y)  # Number of equations\n",
    "\n",
    "    y = cp.empty((steps, order))\n",
    "    y[0] = init_y  # Note that this sets the elements of the first row\n",
    "\n",
    "    for n in range(steps - 1):\n",
    "        h = t[n + 1] - t[n]\n",
    "\n",
    "        # Compute dydt based on *current* position\n",
    "        dydt = f(t[n], y[n])\n",
    "\n",
    "        # Compute next velocity and position\n",
    "        y[n + 1] = y[n] - dydt * h\n",
    "\n",
    "    return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7eaf16b6-1139-4146-bd2a-c9abb3d0efd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = cp.linspace(0, 40, 1000 + 1)\n",
    "y = euler_ivp(f, cp.array([x_max, v_0]), ts)\n",
    "plt.plot(ts.get(), np.cos(ts).get())\n",
    "plt.plot(ts.get(), y[:, 0].get(), \"--\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0a4742c-1548-4031-a3ab-52094781192a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "y = euler_ivp(f, cp.array([x_max, v_0]), ts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4552b388-2109-4982-9a79-0e55351444e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(t, y):\n",
    "    \"Y has two elements, x and v\"\n",
    "    return -koverm * y[1], y[0]\n",
    "\n",
    "\n",
    "def euler_ivp(f, init_y, t):\n",
    "    steps = len(t)\n",
    "    order = len(init_y)  # Number of equations (2)\n",
    "\n",
    "    y = cp.empty((steps, order))\n",
    "    y[0] = init_y  # Note that this sets the elements of the first row\n",
    "\n",
    "    for n in range(steps - 1):\n",
    "        h = t[n + 1] - t[n]\n",
    "\n",
    "        # Compute dydt based on *current* position\n",
    "        dydt_0, dydt_1 = f(t[n], y[n])\n",
    "\n",
    "        # Compute next velocity and position\n",
    "        y[n + 1, 0] = y[n, 0] - dydt_0 * h\n",
    "        y[n + 1, 1] = y[n, 1] - dydt_1 * h\n",
    "\n",
    "    return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40a391d0-37b8-4cef-8e05-f8a8e92c2e77",
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = cp.linspace(0, 40, 1000 + 1)\n",
    "y = euler_ivp(f, cp.array([x_max, v_0]), ts)\n",
    "plt.plot(ts.get(), np.cos(ts).get())\n",
    "plt.plot(ts.get(), y[:, 0].get(), \"--\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfc5b3c0-d238-4e24-91aa-e264505471ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "y = euler_ivp(f, cp.array([x_max, v_0]), ts)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "PyGPU Course 2019/12 [course/pygpu/default]",
   "language": "python",
   "name": "sys_pygpu201912"
  },
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}