docs: Website makeover

docs/completed-tasks.md

@@ -0,0 +1 @@
+[Completed tasks](roadmap.md#completed-tasks)

docs/css/mkdocstrings.css

@@ -1,3 +1,15 @@
-.md-tabs__link {
-    font-weight: bold;
+  :root {
+    --md-primary-fg-color:        #c96c08;
+    --md-primary-fg-color--light: #94f2f7;
+    --md-primary-fg-color--dark:  #335365;
+  }
+
+  .md-tabs__item {
+    font-weight: bolder;
+  }
+
+  .grid {
+    font-weight: bolder;
+    font-size: 160%;
+    font-family: Georgia, serif;
   }

docs/examples/sparse_finch.ipynb

@@ -0,0 +1,365 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Finch backend for `sparse`\n",
+    "\n",
+    "<a href=\"https://colab.research.google.com/github/pydata/sparse/blob/main/examples/sparse_finch.ipynb\" target=\"_blank\">\n",
+    "    <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\" />\n",
+    "</a> to download and run."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# pip install 'sparse[finch]==0.16.0a9' scipy\n",
+    "# export SPARSE_BACKEND=Finch\n",
+    "\n",
+    "# let's make sure we're using Finch backend\n",
+    "import os\n",
+    "\n",
+    "os.environ[\"SPARSE_BACKEND\"] = \"Finch\"\n",
+    "CI_MODE = os.getenv(\"CI_MODE\", default=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import importlib\n",
+    "import time\n",
+    "\n",
+    "import sparse\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "import numpy as np\n",
+    "import scipy.sparse as sps\n",
+    "import scipy.sparse.linalg as splin\n",
+    "\n",
+    "assert sparse.BackendType.Finch == sparse.BACKEND"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tns = sparse.asarray(np.zeros((10, 10)))  # offers a no-copy constructor for NumPy as scipy.sparse inputs\n",
+    "\n",
+    "s1 = sparse.random((100, 10), density=0.01)  # creates random COO tensor\n",
+    "s2 = sparse.random((100, 100, 10), density=0.01)\n",
+    "s2 = sparse.asarray(s2, format=\"csf\")  # can be used to rewrite tensor to a new format\n",
+    "\n",
+    "result = sparse.tensordot(s1, s2, axes=([0, 1], [0, 2]))\n",
+    "\n",
+    "total = sparse.sum(result * result)\n",
+    "print(total)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Example: least squares - closed form"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "y = sparse.random((100, 1), density=0.08)\n",
+    "X = sparse.random((100, 5), density=0.08)\n",
+    "X = sparse.asarray(X, format=\"csc\")\n",
+    "X_lazy = sparse.lazy(X)\n",
+    "\n",
+    "X_X = sparse.compute(sparse.permute_dims(X_lazy, (1, 0)) @ X_lazy, verbose=True)\n",
+    "\n",
+    "X_X = sparse.asarray(X_X, format=\"csc\")  # move back from dense to CSC format\n",
+    "\n",
+    "inverted = splin.inv(X_X)  # dispatching to scipy.sparse.sparray\n",
+    "\n",
+    "b_hat = (inverted @ sparse.permute_dims(X, (1, 0))) @ y\n",
+    "\n",
+    "print(b_hat.todense())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Benchmark plots"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ITERS = 3\n",
+    "rng = np.random.default_rng(0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.style.use(\"seaborn-v0_8\")\n",
+    "plt.rcParams[\"figure.dpi\"] = 400\n",
+    "plt.rcParams[\"figure.figsize\"] = [8, 4]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def benchmark(func, info, args) -> float:\n",
+    "    start = time.time()\n",
+    "    for _ in range(ITERS):\n",
+    "        func(*args)\n",
+    "    elapsed = time.time() - start\n",
+    "    return elapsed / ITERS"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(\"MTTKRP Example:\\n\")\n",
+    "\n",
+    "os.environ[sparse._ENV_VAR_NAME] = \"Numba\"\n",
+    "importlib.reload(sparse)\n",
+    "\n",
+    "configs = [\n",
+    "    {\"I_\": 100, \"J_\": 25, \"K_\": 10, \"L_\": 10, \"DENSITY\": 0.001},\n",
+    "    {\"I_\": 100, \"J_\": 25, \"K_\": 100, \"L_\": 10, \"DENSITY\": 0.001},\n",
+    "    {\"I_\": 100, \"J_\": 25, \"K_\": 100, \"L_\": 100, \"DENSITY\": 0.001},\n",
+    "    {\"I_\": 1000, \"J_\": 25, \"K_\": 100, \"L_\": 100, \"DENSITY\": 0.001},\n",
+    "    {\"I_\": 1000, \"J_\": 25, \"K_\": 1000, \"L_\": 100, \"DENSITY\": 0.001},\n",
+    "    {\"I_\": 1000, \"J_\": 25, \"K_\": 1000, \"L_\": 1000, \"DENSITY\": 0.001},\n",
+    "]\n",
+    "nonzeros = [10000, 100_000, 1_000_000, 10_000_000, 100_000_000, 1_000_000_000]\n",
+    "\n",
+    "if CI_MODE:\n",
+    "    configs = configs[:1]\n",
+    "    nonzeros = nonzeros[:1]\n",
+    "\n",
+    "finch_times = []\n",
+    "numba_times = []\n",
+    "\n",
+    "for config in configs:\n",
+    "    B_sps = sparse.random((config[\"I_\"], config[\"K_\"], config[\"L_\"]), density=config[\"DENSITY\"], random_state=rng) * 10\n",
+    "    D_sps = rng.random((config[\"L_\"], config[\"J_\"])) * 10\n",
+    "    C_sps = rng.random((config[\"K_\"], config[\"J_\"])) * 10\n",
+    "\n",
+    "    # ======= Finch =======\n",
+    "    os.environ[sparse._ENV_VAR_NAME] = \"Finch\"\n",
+    "    importlib.reload(sparse)\n",
+    "\n",
+    "    B = sparse.asarray(B_sps.todense(), format=\"csf\")\n",
+    "    D = sparse.asarray(np.array(D_sps, order=\"F\"))\n",
+    "    C = sparse.asarray(np.array(C_sps, order=\"F\"))\n",
+    "\n",
+    "    @sparse.compiled\n",
+    "    def mttkrp_finch(B, D, C):\n",
+    "        return sparse.sum(B[:, :, :, None] * D[None, None, :, :] * C[None, :, None, :], axis=(1, 2))\n",
+    "\n",
+    "    # Compile\n",
+    "    result_finch = mttkrp_finch(B, D, C)\n",
+    "    # Benchmark\n",
+    "    time_finch = benchmark(mttkrp_finch, info=\"Finch\", args=[B, D, C])\n",
+    "\n",
+    "    # ======= Numba =======\n",
+    "    os.environ[sparse._ENV_VAR_NAME] = \"Numba\"\n",
+    "    importlib.reload(sparse)\n",
+    "\n",
+    "    B = sparse.asarray(B_sps, format=\"gcxs\")\n",
+    "    D = D_sps\n",
+    "    C = C_sps\n",
+    "\n",
+    "    def mttkrp_numba(B, D, C):\n",
+    "        return sparse.sum(B[:, :, :, None] * D[None, None, :, :] * C[None, :, None, :], axis=(1, 2))\n",
+    "\n",
+    "    # Compile\n",
+    "    result_numba = mttkrp_numba(B, D, C)\n",
+    "    # Benchmark\n",
+    "    time_numba = benchmark(mttkrp_numba, info=\"Numba\", args=[B, D, C])\n",
+    "\n",
+    "    np.testing.assert_allclose(result_finch.todense(), result_numba.todense())\n",
+    "    finch_times.append(time_finch)\n",
+    "    numba_times.append(time_numba)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(nrows=1, ncols=1)\n",
+    "\n",
+    "ax.plot(nonzeros, finch_times, \"o-\", label=\"Finch\")\n",
+    "ax.plot(nonzeros, numba_times, \"o-\", label=\"Numba\")\n",
+    "ax.grid(True)\n",
+    "ax.set_xlabel(\"no. of elements\")\n",
+    "ax.set_ylabel(\"time (sec)\")\n",
+    "ax.set_title(\"MTTKRP\")\n",
+    "ax.set_xscale(\"log\")\n",
+    "ax.set_yscale(\"log\")\n",
+    "ax.legend(loc=\"best\", numpoints=1)\n",
+    "\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(\"SDDMM Example:\\n\")\n",
+    "\n",
+    "configs = [\n",
+    "    {\"LEN\": 10, \"DENSITY\": 0.1},\n",
+    "    {\"LEN\": 50, \"DENSITY\": 0.05},\n",
+    "    {\"LEN\": 100, \"DENSITY\": 0.01},\n",
+    "    {\"LEN\": 500, \"DENSITY\": 0.005},\n",
+    "    {\"LEN\": 1000, \"DENSITY\": 0.001},\n",
+    "    {\"LEN\": 5000, \"DENSITY\": 0.00005},\n",
+    "    {\"LEN\": 10000, \"DENSITY\": 0.00001},\n",
+    "]\n",
+    "size_n = [10, 50, 100, 500, 1000, 5000, 10000]\n",
+    "\n",
+    "if CI_MODE:\n",
+    "    configs = configs[:1]\n",
+    "    size_n = size_n[:1]\n",
+    "\n",
+    "finch_times = []\n",
+    "numba_times = []\n",
+    "scipy_times = []\n",
+    "\n",
+    "for config in configs:\n",
+    "    LEN = config[\"LEN\"]\n",
+    "    DENSITY = config[\"DENSITY\"]\n",
+    "\n",
+    "    a_sps = rng.random((LEN, LEN)) * 10\n",
+    "    b_sps = rng.random((LEN, LEN)) * 10\n",
+    "    s_sps = sps.random(LEN, LEN, format=\"coo\", density=DENSITY, random_state=rng) * 10\n",
+    "    s_sps.sum_duplicates()\n",
+    "\n",
+    "    # ======= Finch =======\n",
+    "    os.environ[sparse._ENV_VAR_NAME] = \"Finch\"\n",
+    "    importlib.reload(sparse)\n",
+    "\n",
+    "    s = sparse.asarray(s_sps)\n",
+    "    a = sparse.asarray(np.array(a_sps, order=\"F\"))\n",
+    "    b = sparse.asarray(np.array(b_sps, order=\"C\"))\n",
+    "\n",
+    "    @sparse.compiled\n",
+    "    def sddmm_finch(s, a, b):\n",
+    "        return sparse.sum(\n",
+    "            s[:, :, None] * (a[:, None, :] * sparse.permute_dims(b, (1, 0))[None, :, :]),\n",
+    "            axis=-1,\n",
+    "        )\n",
+    "\n",
+    "    # Compile\n",
+    "    result_finch = sddmm_finch(s, a, b)\n",
+    "    # Benchmark\n",
+    "    time_finch = benchmark(sddmm_finch, info=\"Finch\", args=[s, a, b])\n",
+    "\n",
+    "    # ======= Numba =======\n",
+    "    os.environ[sparse._ENV_VAR_NAME] = \"Numba\"\n",
+    "    importlib.reload(sparse)\n",
+    "\n",
+    "    s = sparse.asarray(s_sps)\n",
+    "    a = a_sps\n",
+    "    b = b_sps\n",
+    "\n",
+    "    def sddmm_numba(s, a, b):\n",
+    "        return s * (a @ b)\n",
+    "\n",
+    "    # Compile\n",
+    "    result_numba = sddmm_numba(s, a, b)\n",
+    "    # Benchmark\n",
+    "    time_numba = benchmark(sddmm_numba, info=\"Numba\", args=[s, a, b])\n",
+    "\n",
+    "    # ======= SciPy =======\n",
+    "    def sddmm_scipy(s, a, b):\n",
+    "        return s.multiply(a @ b)\n",
+    "\n",
+    "    s = s_sps.asformat(\"csr\")\n",
+    "    a = a_sps\n",
+    "    b = b_sps\n",
+    "\n",
+    "    result_scipy = sddmm_scipy(s, a, b)\n",
+    "    # Benchmark\n",
+    "    time_scipy = benchmark(sddmm_scipy, info=\"SciPy\", args=[s, a, b])\n",
+    "\n",
+    "    finch_times.append(time_finch)\n",
+    "    numba_times.append(time_numba)\n",
+    "    scipy_times.append(time_scipy)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(nrows=1, ncols=1)\n",
+    "\n",
+    "ax.plot(size_n, finch_times, \"o-\", label=\"Finch\")\n",
+    "ax.plot(size_n, numba_times, \"o-\", label=\"Numba\")\n",
+    "ax.plot(size_n, scipy_times, \"o-\", label=\"SciPy\")\n",
+    "\n",
+    "ax.grid(True)\n",
+    "ax.set_xlabel(\"size N\")\n",
+    "ax.set_ylabel(\"time (sec)\")\n",
+    "ax.set_title(\"SDDMM\")\n",
+    "ax.set_xscale(\"log\")\n",
+    "# ax.set_yscale('log')\n",
+    "ax.legend(loc=\"best\", numpoints=1)\n",
+    "\n",
+    "plt.show()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "sparse-dev",
+   "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": 2
+}