add examples/JaxBenchmarks.ipynb

examples/JaxBenchmarks.ipynb

@@ -0,0 +1,231 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from jax import devices, random, jit, vmap, numpy as jnp\n",
+ "from vectorbt import _typing as tp\n",
+ "from numba import njit\n",
+ "from jax import random\n",
+ "import numpy as np\n",
+ "import itertools\n",
+ "from functools import partial\n",
+ "import pandas as pd\n",
+ "import random as pyrandom"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "print(devices(backend=\"gpu\"))\n",
+ "gpu = devices(backend=\"gpu\")[0]"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "@njit(cache=True)\n",
+ "def get_return_nb(input_value: float, output_value: float) -> float:\n",
+ " \"\"\"Calculate return from input and output value.\"\"\"\n",
+ " if input_value == 0:\n",
+ " if output_value == 0:\n",
+ " return 0.\n",
+ " return np.inf * np.sign(output_value)\n",
+ " return_value = (output_value - input_value) / input_value\n",
+ " if input_value < 0:\n",
+ " return_value *= -1\n",
+ " return return_value\n",
+ "\n",
+ "\n",
+ "@njit(cache=True)\n",
+ "def returns_1d_nb(value: tp.Array1d, init_value: float) -> tp.Array1d:\n",
+ " \"\"\"Calculate returns from value.\"\"\"\n",
+ " out = np.empty(value.shape, dtype=np.float_)\n",
+ " input_value = init_value\n",
+ " for i in range(out.shape[0]):\n",
+ " output_value = value[i]\n",
+ " out[i] = get_return_nb(input_value, output_value)\n",
+ " input_value = output_value\n",
+ " return out\n",
+ "\n",
+ "\n",
+ "@njit(cache=True)\n",
+ "def returns_nb(value: tp.Array2d, init_value: tp.Array1d) -> tp.Array2d:\n",
+ " \"\"\"2-dim version of `returns_1d_nb`.\"\"\"\n",
+ " out = np.empty(value.shape, dtype=np.float_)\n",
+ " for col in range(out.shape[1]):\n",
+ " out[:, col] = returns_1d_nb(value[:, col], init_value[col])\n",
+ " return out"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def _returns_1d_jax(value: tp.Array1d, init_value: float):\n",
+ "\tY = jnp.concatenate((jnp.array([init_value]), value))\n",
+ "\treturn jnp.divide(\n",
+ "\t\tjnp.diff(Y), \n",
+ "\t\tY[:1]\n",
+ "\t\t) * jnp.sign(Y[:1])\n",
+ "\n",
+ "returns_1d_jax = jit(_returns_1d_jax, device=gpu)\n",
+ "returns_jax = lambda v, i: vmap(partial(returns_1d_jax, init_value=i), in_axes=1, out_axes=1)(v)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# these functions make random test data, i made one for each of numpy and jax for fairness\n",
+ "def random_walk_nb(start, scale, steps, seed, n_portfolio=None):\n",
+ " rng = np.random.default_rng(seed)\n",
+ " shape = (steps - 1, n_portfolio) if n_portfolio else (steps - 1, )\n",
+ " noise = rng.normal(size=shape) * scale\n",
+ " noise = np.insert(noise, 0, start, axis=0)\n",
+ " walk = np.cumsum(noise, axis=0)\n",
+ " return walk\n",
+ "\n",
+ "# nb_data = random_walk_nb(100, 1, 100000, 1)\n",
+ "\n",
+ "def random_walk_jax(start, scale, steps, seed, n_portfolio=None):\n",
+ " key = random.PRNGKey(seed)\n",
+ " shape = (steps - 1, n_portfolio) if n_portfolio else (steps - 1, )\n",
+ " noise = random.normal(key, shape=shape) * scale\n",
+ " noise = jnp.insert(noise, 0, start, axis=0)\n",
+ " walk = jnp.cumsum(noise, axis=0)\n",
+ " return walk\n",
+ "\n",
+ "# jax_data = random_walk_jax(100, 1, 100000, 1)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# make sure the 1d jax thing works right\n",
+ "# arr = jnp.array([100, 110, 100, 120])\n",
+ "# numerator = jnp.diff(arr)\n",
+ "# print(\"differences\", numerator)\n",
+ "# denominator = arr[:-1]\n",
+ "# print(\"denominators\", denominator)\n",
+ "# returns = numerator / denominator\n",
+ "# print(\"returns\", returns)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# %timeit returns_1d_jax(jax_data, 1.0)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# %timeit returns_1d_nb(nb_data, 1.0)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "n_candles = [2048]\n",
+ "n_portfolios = [2 ** i for i in range(2, 12)]\n",
+ "n_loops = 5\n",
+ "\n",
+ "timings = {\"index\": [], \"jax\": [], \"nb\": []}\n",
+ "testing = True\n",
+ "timing = True\n",
+ "for n_candle, n_portfolio in itertools.product(n_candles, n_portfolios):\n",
+ " size = n_candle*n_portfolio\n",
+ " print(f\"\\ncalculate ({n_candle}, {n_portfolio}) returns (size {size})\\n\")\n",
+ " timings[\"index\"].append(n_portfolio)\n",
+ " def jax_fun(walk=None):\n",
+ " if walk is None:\n",
+ " seed = pyrandom.randint(0, 1000)\n",
+ " walk = random_walk_jax(1, 0.01, n_candle, seed, n_portfolio=n_portfolio)\n",
+ " returns = returns_jax(walk, 1.0).block_until_ready()\n",
+ " return returns\n",
+ " one = np.ones((n_portfolio,))\n",
+ " def nb_fun(walk=None):\n",
+ " if walk is None:\n",
+ " seed = pyrandom.randint(0, 1000)\n",
+ " walk = random_walk_nb(1, 0.01, n_candle, seed, n_portfolio=n_portfolio)\n",
+ " returns = returns_nb(walk, one)\n",
+ " return returns\n",
+ " if timing:\n",
+ " print(\"jax:\")\n",
+ " jax_timings = %timeit -o jax_fun()\n",
+ " timings[\"jax\"].append(jax_timings.average)\n",
+ " print(\"nb:\")\n",
+ " nb_timings = %timeit -o nb_fun()\n",
+ " timings[\"nb\"].append(nb_timings.average)\n",
+ " if testing:\n",
+ " print(\"mean absolute error between jax and numba returns:\")\n",
+ " print(\"jnp.abs(nb_test - jax_test).mean() =\")\n",
+ " maes = []\n",
+ " for i in range(n_loops):\n",
+ " seed = pyrandom.randint(0, 1000)\n",
+ " walk = random_walk_jax(1, 0.01, n_candle, seed, n_portfolio=n_portfolio)\n",
+ " np_walk = np.array(walk)\n",
+ " jax_test = jax_fun(walk)\n",
+ " nb_test = nb_fun(np_walk)\n",
+ " mae = jnp.abs(nb_test - jax_test).mean()\n",
+ " maes.append(mae)\n",
+ " print(maes)\n",
+ "if timing:\n",
+ " df = pd.DataFrame.from_dict(timings)\n",
+ " df.index = df[\"index\"]\n",
+ " df = df.drop(columns=\"index\")\n",
+ " df.plot(title=\"vectorbt returns accelerator benchmark\", xlabel=\"n_portfolios\", ylabel=\"time (lower is better)\")"
+ ]
+ }
+ ],
+ "metadata": {
+ "interpreter": {
+ "hash": "15ccec256c034d9984cedb439e8cbdae72d64b38628afc877d5ed250d5fec1dd"
+ },
+ "kernelspec": {
+ "display_name": "Python 3.8.0 64-bit ('py38': conda)",
+ "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.8.0"
+ },
+ "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}