{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "speed_on_your_machine-0",
   "metadata": {},
   "source": [
    "# Speed on your machine\n",
    "\n",
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Abel-ai-lab/andrey/blob/main/examples/speed_on_your_machine.ipynb)\n",
    "\n",
    "Run PC with Andrey and causal-learn on the same observations. The size ladder is deliberately\n",
    "small; increase it to explore your machine's limits. Each fit runs sequentially with one numeric\n",
    "thread. Imports, warm-up, graph conversion, and accuracy scoring are outside the timed region.\n",
    "\n",
    "**The fastest causal discovery package.** That claim comes from\n",
    "the [release cards](https://andrey.abel.ai/docs/benchmarks.html), not a promised result\n",
    "on this machine. The live demo labels its server; your hardware and data determine your times.\n",
    "\n",
    "Run in a fresh kernel. Outside Colab, the install cell adds causal-learn with uv.\n",
    "\n",
    "In Colab, these are timings of the Colab runtime. Download the notebook to measure your own CPU.\n",
    ""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fb27b941602401d91542211134fc71a",
   "metadata": {},
   "source": [
    "## Start\n",
    "\n",
    "Choose a CPU runtime and run this cell first. In Colab it installs `andrey-core==0.9.0` from PyPI;\n",
    "local kernels use the installed environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acae54e37e7d407bbb7b55eff062a284",
   "metadata": {
    "tags": [
     "colab-setup"
    ]
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "IN_COLAB = \"google.colab\" in sys.modules\n",
    "\n",
    "if IN_COLAB:\n",
    "    os.environ[\"ANDREY_DEVICE\"] = \"cpu\"\n",
    "    !pip install --quiet \"andrey-core[viz]==0.9.0\" pandas \"causal-learn==0.1.4.8\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "speed_on_your_machine-install",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not IN_COLAB:\n",
    "    !uv pip install --python {sys.executable} causal-learn==0.1.4.8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "speed_on_your_machine-1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "for variable in (\n",
    "    \"OMP_NUM_THREADS\",\n",
    "    \"OPENBLAS_NUM_THREADS\",\n",
    "    \"MKL_NUM_THREADS\",\n",
    "    \"NUMEXPR_NUM_THREADS\",\n",
    "    \"NUMBA_NUM_THREADS\",\n",
    "    \"ANDREY_NUM_WORKERS\",\n",
    "):\n",
    "    os.environ[variable] = \"1\"\n",
    "os.environ[\"ANDREY_BACKEND\"] = \"numpy\"\n",
    "os.environ[\"ANDREY_DEVICE\"] = \"cpu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "speed_on_your_machine-2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import platform\n",
    "from importlib.metadata import version\n",
    "from time import perf_counter\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from causallearn.search.ConstraintBased.PC import pc as causal_learn_pc\n",
    "from threadpoolctl import threadpool_limits\n",
    "\n",
    "import andrey\n",
    "import andrey.metrics as metrics\n",
    "from andrey.data import benchmarks\n",
    "\n",
    "print(\"Execution environment:\", \"Google Colab runtime\" if IN_COLAB else \"Local kernel\")\n",
    "print(f\"Hardware: {platform.machine()}, {os.cpu_count()} logical CPUs; one thread per fit\")\n",
    "print(f\"Andrey {andrey.__version__}; causal-learn {version('causal-learn')}\")\n",
    "\n",
    "\n",
    "def convert_causal_learn(result):\n",
    "    # causal-learn stores -1=tail, 1=arrow at the same matrix positions as Andrey.\n",
    "    native = result.G.graph\n",
    "    marks = np.zeros(native.shape, dtype=np.uint8)\n",
    "    marks[native == -1] = 1\n",
    "    marks[native == 1] = 2\n",
    "    return andrey.GraphStructure.from_numpy(marks, kind=\"cpdag\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "speed_on_your_machine-3",
   "metadata": {},
   "source": [
    "## Warm both packages, then time fits\n",
    "\n",
    "Both use stable PC, Fisher-Z, and `alpha=0.05`. The same three seeds are used at each size.\n",
    "Speed and accuracy are reported separately: every paired ratio, and each package's SHD beside\n",
    "the empty graph's."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "speed_on_your_machine-4",
   "metadata": {},
   "outputs": [],
   "source": [
    "sizes = [10, 20, 40]\n",
    "seeds = [7, 11, 19]\n",
    "rows = []\n",
    "with threadpool_limits(limits=1):\n",
    "    warm = benchmarks.scm(\"linear_gauss_er\", 5).sample(n=100, seed=0).data\n",
    "    andrey.pc(warm)\n",
    "    causal_learn_pc(warm, alpha=0.05, stable=True, show_progress=False)\n",
    "    for d in sizes:\n",
    "        for seed in seeds:\n",
    "            sample = benchmarks.scm(\"linear_gauss_er\", d).sample(\n",
    "                n=10 * d, seed=seed, scale=\"standardize\"\n",
    "            )\n",
    "            empty = andrey.GraphStructure.from_numpy(np.zeros((d, d)), kind=\"cpdag\")\n",
    "            floor = metrics.score(empty, sample.graph)[\"shd\"]\n",
    "            for package in (\"Andrey\", \"causal-learn\"):\n",
    "                start = perf_counter()\n",
    "                if package == \"Andrey\":\n",
    "                    result = andrey.pc(sample.data, alpha=0.05)\n",
    "                else:\n",
    "                    result = causal_learn_pc(\n",
    "                        sample.data, alpha=0.05, stable=True, show_progress=False\n",
    "                    )\n",
    "                seconds = perf_counter() - start\n",
    "                graph = result.structure if package == \"Andrey\" else convert_causal_learn(result)\n",
    "                shd = metrics.score(graph, sample.graph)[\"shd\"]\n",
    "                rows.append(\n",
    "                    dict(d=d, seed=seed, package=package, seconds=seconds, SHD=shd, empty_SHD=floor)\n",
    "                )\n",
    "measurements = pd.DataFrame(rows)\n",
    "measurements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "speed_on_your_machine-5",
   "metadata": {},
   "outputs": [],
   "source": [
    "summary = measurements.groupby([\"d\", \"package\"])[[\"seconds\", \"SHD\", \"empty_SHD\"]].median()\n",
    "display(summary)\n",
    "paired = measurements.pivot(\n",
    "    index=[\"d\", \"seed\"], columns=\"package\", values=[\"seconds\", \"SHD\", \"empty_SHD\"]\n",
    ")\n",
    "ratios = paired[\"seconds\"][\"causal-learn\"] / paired[\"seconds\"][\"Andrey\"]\n",
    "ratios.groupby(level=\"d\").agg([\"median\", \"count\"]).rename(\n",
    "    columns={\"median\": \"Paired speed ratio\", \"count\": \"Paired seeds\"}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "speed_on_your_machine-6",
   "metadata": {},
   "source": [
    "Ratios are calculated per shared dataset and then aggregated. They need not equal ratios of\n",
    "median times. Increase the ladder only when the preceding sizes finish comfortably;\n",
    "causal-learn can take minutes at larger sizes.\n",
    "\n",
    "The headline workload is PC at 400 variables and 4,000 observations. Add `400` to `sizes` only\n",
    "when ready for a longer run. The benchmark page records the campaign builds and hardware; this\n",
    "notebook measures the installed packages with one numeric thread on your machine."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}