{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-01-09T00:23:28.917621Z",
     "iopub.status.busy": "2025-01-09T00:23:28.917332Z",
     "iopub.status.idle": "2025-01-09T00:23:29.477701Z",
     "shell.execute_reply": "2025-01-09T00:23:29.477193Z",
     "shell.execute_reply.started": "2025-01-09T00:23:28.917602Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[CpuDevice(id=0)]"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "\n",
    "jax.devices()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-01-09T00:23:30.540594Z",
     "iopub.status.busy": "2025-01-09T00:23:30.540228Z",
     "iopub.status.idle": "2025-01-09T00:23:30.739685Z",
     "shell.execute_reply": "2025-01-09T00:23:30.739213Z",
     "shell.execute_reply.started": "2025-01-09T00:23:30.540574Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{CpuDevice(id=0)}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.ones(10).devices()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-01-09T00:26:29.239253Z",
     "iopub.status.busy": "2025-01-09T00:26:29.238947Z",
     "iopub.status.idle": "2025-01-09T00:26:29.754752Z",
     "shell.execute_reply": "2025-01-09T00:26:29.754158Z",
     "shell.execute_reply.started": "2025-01-09T00:26:29.239235Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "'nvidia-smi' is not recognized as an internal or external command,\n",
      "operable program or batch file.\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-01-09T00:23:44.232335Z",
     "iopub.status.busy": "2025-01-09T00:23:44.231984Z",
     "iopub.status.idle": "2025-01-09T00:23:45.388035Z",
     "shell.execute_reply": "2025-01-09T00:23:45.387587Z",
     "shell.execute_reply.started": "2025-01-09T00:23:44.232312Z"
    },
    "id": "fHzEldNvR2_K"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "            <div id=\"ykWcWl\"></div>\n",
       "            <script type=\"text/javascript\" data-lets-plot-script=\"library\">\n",
       "                if(!window.letsPlotCallQueue) {\n",
       "                    window.letsPlotCallQueue = [];\n",
       "                }; \n",
       "                window.letsPlotCall = function(f) {\n",
       "                    window.letsPlotCallQueue.push(f);\n",
       "                };\n",
       "                (function() {\n",
       "                    var script = document.createElement(\"script\");\n",
       "                    script.type = \"text/javascript\";\n",
       "                    script.src = \"https://cdn.jsdelivr.net/gh/JetBrains/lets-plot@v4.5.2/js-package/distr/lets-plot.min.js\";\n",
       "                    script.onload = function() {\n",
       "                        window.letsPlotCall = function(f) {f();};\n",
       "                        window.letsPlotCallQueue.forEach(function(f) {f();});\n",
       "                        window.letsPlotCallQueue = [];\n",
       "                        \n",
       "                    };\n",
       "                    script.onerror = function(event) {\n",
       "                        window.letsPlotCall = function(f) {};    // noop\n",
       "                        window.letsPlotCallQueue = [];\n",
       "                        var div = document.createElement(\"div\");\n",
       "                        div.style.color = 'darkred';\n",
       "                        div.textContent = 'Error loading Lets-Plot JS';\n",
       "                        document.getElementById(\"ykWcWl\").appendChild(div);\n",
       "                    };\n",
       "                    var e = document.getElementById(\"ykWcWl\");\n",
       "                    e.appendChild(script);\n",
       "                })()\n",
       "            </script>\n",
       "            "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "            <div id=\"koKgk9\"></div>\n",
       "            <script type=\"text/javascript\" data-lets-plot-script=\"library\">\n",
       "                if(!window.letsPlotCallQueue) {\n",
       "                    window.letsPlotCallQueue = [];\n",
       "                }; \n",
       "                window.letsPlotCall = function(f) {\n",
       "                    window.letsPlotCallQueue.push(f);\n",
       "                };\n",
       "                (function() {\n",
       "                    var script = document.createElement(\"script\");\n",
       "                    script.type = \"text/javascript\";\n",
       "                    script.src = \"https://cdn.jsdelivr.net/gh/JetBrains/lets-plot@v4.5.2/js-package/distr/lets-plot.min.js\";\n",
       "                    script.onload = function() {\n",
       "                        window.letsPlotCall = function(f) {f();};\n",
       "                        window.letsPlotCallQueue.forEach(function(f) {f();});\n",
       "                        window.letsPlotCallQueue = [];\n",
       "                        \n",
       "                    };\n",
       "                    script.onerror = function(event) {\n",
       "                        window.letsPlotCall = function(f) {};    // noop\n",
       "                        window.letsPlotCallQueue = [];\n",
       "                        var div = document.createElement(\"div\");\n",
       "                        div.style.color = 'darkred';\n",
       "                        div.textContent = 'Error loading Lets-Plot JS';\n",
       "                        document.getElementById(\"koKgk9\").appendChild(div);\n",
       "                    };\n",
       "                    var e = document.getElementById(\"koKgk9\");\n",
       "                    e.appendChild(script);\n",
       "                })()\n",
       "            </script>\n",
       "            "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import time\n",
    "from itertools import product\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "from scipy.stats import nbinom\n",
    "from tqdm import tqdm\n",
    "\n",
    "import pyfixest as pf\n",
    "from pyfixest.estimation.demean_ import demean\n",
    "from pyfixest.estimation.demean_jax_ import demean_jax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2025-01-09T00:23:46.290548Z",
     "iopub.status.busy": "2025-01-09T00:23:46.289898Z",
     "iopub.status.idle": "2025-01-09T00:23:46.417097Z",
     "shell.execute_reply": "2025-01-09T00:23:46.416504Z",
     "shell.execute_reply.started": "2025-01-09T00:23:46.290525Z"
    },
    "id": "XQjP2889YJxs",
    "outputId": "3e686d7b-0774-4bb5-c1b9-28e5b9f286a9"
   },
   "outputs": [],
   "source": [
    "# %load_ext watermark\n",
    "# %watermark --iversions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "execution": {
     "iopub.execute_input": "2025-01-09T00:23:49.545271Z",
     "iopub.status.busy": "2025-01-09T00:23:49.545016Z",
     "iopub.status.idle": "2025-01-09T00:23:49.552123Z",
     "shell.execute_reply": "2025-01-09T00:23:49.551676Z",
     "shell.execute_reply.started": "2025-01-09T00:23:49.545253Z"
    },
    "id": "bxMmeyCxR3fb"
   },
   "outputs": [],
   "source": [
    "def generate_test_data(size: int, k: int = 2):\n",
    "    \"\"\"\n",
    "    Generate benchmark data for pyfixest on GPU (similar to the R fixest benchmark data).\n",
    "\n",
    "    Args:\n",
    "        size (int): The number of observations in the data frame.\n",
    "        k (int): The number of covariates in the data frame.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "        pd.DataFrame: The generated data frame for the given size.\n",
    "    \"\"\"\n",
    "    # Constants\n",
    "    all_n = [1000 * 10**i for i in range(5)]\n",
    "    a = 1\n",
    "    b = 0.05\n",
    "\n",
    "    n = all_n[size - 1]\n",
    "\n",
    "    dum_all = []\n",
    "    nb_dum = [n // 20, int(np.sqrt(n)), int(n**0.33)]\n",
    "\n",
    "    dum_all = np.zeros((n, 3))\n",
    "    dum_all[:, 0] = np.random.choice(nb_dum[0], n, replace=True)\n",
    "    dum_all[:, 1] = np.random.choice(nb_dum[1], n, replace=True)\n",
    "    dum_all[:, 2] = np.random.choice(nb_dum[2], n, replace=True)\n",
    "    dum_all = dum_all.astype(int)\n",
    "\n",
    "    X1 = np.random.normal(size=n)\n",
    "    X2 = X1**2\n",
    "\n",
    "    mu = a * X1 + b * X2\n",
    "\n",
    "    for m in range(3):\n",
    "        coef_dum = np.random.normal(size=nb_dum[m])\n",
    "        mu += coef_dum[dum_all[:, m]]\n",
    "\n",
    "    mu = np.exp(mu)\n",
    "    y = nbinom.rvs(0.5, 1 - (mu / (mu + 0.5)), size=n)\n",
    "\n",
    "    X_full = np.column_stack((X1, X2))\n",
    "    base = pd.DataFrame(\n",
    "        {\n",
    "            \"y\": y,\n",
    "            \"ln_y\": np.log(y + 1),\n",
    "            \"X1\": X1,\n",
    "            \"X2\": X2,\n",
    "        }\n",
    "    )\n",
    "\n",
    "    if k > 2:\n",
    "        X = np.random.normal(size=(n, k - 2))\n",
    "        X_df = pd.DataFrame(X, columns=[f\"X{i}\" for i in range(3, k + 1, 1)])\n",
    "        base = pd.concat([base, X_df], axis=1)\n",
    "        X_full = np.column_stack((X_full, X))\n",
    "\n",
    "    for m in range(3):\n",
    "        base[f\"dum_{m + 1}\"] = dum_all[:, m]\n",
    "\n",
    "    weights = np.random.uniform(0, 1, n)\n",
    "    return base, y, X_full, dum_all, weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-01-09T00:23:50.285297Z",
     "iopub.status.busy": "2025-01-09T00:23:50.284967Z",
     "iopub.status.idle": "2025-01-09T00:23:50.460957Z",
     "shell.execute_reply": "2025-01-09T00:23:50.460501Z",
     "shell.execute_reply.started": "2025-01-09T00:23:50.285276Z"
    },
    "id": "nzynhbqwR81H"
   },
   "outputs": [],
   "source": [
    "df, Y, X, f, weights = generate_test_data(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-01-09T00:25:02.873750Z",
     "iopub.status.busy": "2025-01-09T00:25:02.873239Z",
     "iopub.status.idle": "2025-01-09T00:25:03.153458Z",
     "shell.execute_reply": "2025-01-09T00:25:03.153005Z",
     "shell.execute_reply.started": "2025-01-09T00:25:02.873732Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "###\n",
      "\n",
      "Estimation:  OLS\n",
      "Dep. var.: ln_y, Fixed effects: dum_1\n",
      "Inference:  CRV1\n",
      "Observations:  1000\n",
      "\n",
      "| Coefficient   |   Estimate |   Std. Error |   t value |   Pr(>|t|) |   2.5% |   97.5% |\n",
      "|:--------------|-----------:|-------------:|----------:|-----------:|-------:|--------:|\n",
      "| X1            |      0.436 |        0.046 |     9.440 |      0.000 |  0.343 |   0.529 |\n",
      "---\n",
      "RMSE: 1.067 R2: 0.242 R2 Within: 0.131 \n"
     ]
    }
   ],
   "source": [
    "m0 = pf.feols(\"ln_y ~ X1 | dum_1\", df, demeaner_backend=\"numba\")\n",
    "m0.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-01-09T00:25:03.330646Z",
     "iopub.status.busy": "2025-01-09T00:25:03.330367Z",
     "iopub.status.idle": "2025-01-09T00:25:03.571916Z",
     "shell.execute_reply": "2025-01-09T00:25:03.571482Z",
     "shell.execute_reply.started": "2025-01-09T00:25:03.330629Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "###\n",
      "\n",
      "Estimation:  OLS\n",
      "Dep. var.: ln_y, Fixed effects: dum_1\n",
      "Inference:  CRV1\n",
      "Observations:  1000\n",
      "\n",
      "| Coefficient   |   Estimate |   Std. Error |   t value |   Pr(>|t|) |   2.5% |   97.5% |\n",
      "|:--------------|-----------:|-------------:|----------:|-----------:|-------:|--------:|\n",
      "| X1            |      0.436 |        0.046 |     9.440 |      0.000 |  0.343 |   0.529 |\n",
      "---\n",
      "RMSE: 1.067 R2: 0.242 R2 Within: 0.131 \n"
     ]
    }
   ],
   "source": [
    "m1 = pf.feols(\"ln_y ~ X1 | dum_1\", df, demeaner_backend=\"jax\")\n",
    "m1.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-01-09T00:24:06.619552Z",
     "iopub.status.busy": "2025-01-09T00:24:06.619273Z",
     "iopub.status.idle": "2025-01-09T00:24:06.626298Z",
     "shell.execute_reply": "2025-01-09T00:24:06.625727Z",
     "shell.execute_reply.started": "2025-01-09T00:24:06.619534Z"
    },
    "id": "29rZkULUR_A0"
   },
   "outputs": [],
   "source": [
    "def run_standard_benchmark(\n",
    "    fixed_effect,\n",
    "    demeaner_backend,\n",
    "    size=1,\n",
    "    k=1,\n",
    "    solver=\"np.linalg.lstsq\",\n",
    "    skip_demean_benchmark=True,\n",
    "):\n",
    "    \"\"\"\n",
    "    Run the fixest standard benchmark fixed effect models. This is the function the benchmarks\n",
    "    will loop over.\n",
    "\n",
    "    Args:\n",
    "        fixed_effect (str): The fixed effect to use. Must be a list of variables as \"dum_1\", \"dum_1+dum_2\", or \"dum_1+dum_2+dum_3\", etc.\n",
    "        demeaner_backend (str): The backend to use for demeaning. Must be \"numba\" or \"jax\".\n",
    "        size (int): The size of the data to generate. Must be between 1 and 5. For 1, N = 1000, for 2, N = 10000, etc.\n",
    "        k_vals (int): The number of covariates to generate.\n",
    "        solver (str): The solver to use for the estimation. Must be \"np.linalg.lstsq\". \"jax\" currently throws an error.\n",
    "        skip_demean_benchmark (bool): Whether to skip the \"pure\" demean benchmark. Default is True. Only the full call\n",
    "            to feols is benchmarked.\n",
    "\n",
    "    \"\"\"\n",
    "    assert fixed_effect in [\"dum_1\", \"dum_1+dum_2\", \"dum_1+dum_2+dum_3\"]\n",
    "\n",
    "    # one fixed effect\n",
    "    res = []\n",
    "\n",
    "    fml_base = \"ln_y ~ X1\"\n",
    "    fml = f\"{fml_base} | {fixed_effect}\"\n",
    "\n",
    "    # warmup\n",
    "    df, y, X, f, weights = generate_test_data(1)\n",
    "    pf.feols(\n",
    "        fml,\n",
    "        data=df,\n",
    "        demeaner_backend=demeaner_backend,\n",
    "        store_data=False,\n",
    "        copy_data=False,\n",
    "        solver=solver,\n",
    "    )\n",
    "\n",
    "    if k > 1:\n",
    "        xfml = \"+\".join([f\"X{i}\" for i in range(2, k + 1, 1)])\n",
    "        fml = f\"{fml_base} + {xfml} | {fixed_effect}\"\n",
    "    else:\n",
    "        fml = f\"{fml_base} + X1 | {fixed_effect}\"\n",
    "\n",
    "    for rep in range(1, 11):\n",
    "        df, Y, X, f, weights = generate_test_data(size=size, k=k)\n",
    "\n",
    "        tic1 = time.time()\n",
    "        pf.feols(\n",
    "            fml,\n",
    "            data=df,\n",
    "            demeaner_backend=demeaner_backend,\n",
    "            store_data=False,\n",
    "            copy_data=False,\n",
    "            solver=solver,\n",
    "        )\n",
    "        tic2 = time.time()\n",
    "\n",
    "        full_feols_timing = tic2 - tic1\n",
    "\n",
    "        demean_timing = np.nan\n",
    "        if not skip_demean_benchmark:\n",
    "            YX = np.column_stack((Y.reshape(-1, 1), X))\n",
    "            tic3 = time.time()\n",
    "            if demeaner_backend == \"jax\":\n",
    "                _, _ = demean_jax(YX, f, weights, tol=1e-10)\n",
    "            else:\n",
    "                _, _ = demean(YX, f, weights, tol=1e-10)\n",
    "            tic4 = time.time()\n",
    "            demean_timing = tic4 - tic3\n",
    "\n",
    "        res.append(\n",
    "            pd.Series(\n",
    "                {\n",
    "                    \"method\": \"feols\",\n",
    "                    \"solver\": solver,\n",
    "                    \"demeaner_backend\": demeaner_backend,\n",
    "                    \"n_obs\": df.shape[0],\n",
    "                    \"k\": k,\n",
    "                    \"G\": len(fixed_effect.split(\"+\")),\n",
    "                    \"rep\": rep,\n",
    "                    \"full_feols_timing\": full_feols_timing,\n",
    "                    \"demean_timing\": demean_timing,\n",
    "                }\n",
    "            )\n",
    "        )\n",
    "\n",
    "    return pd.concat(res, axis=1).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-01-09T00:28:43.818536Z",
     "iopub.status.busy": "2025-01-09T00:28:43.818246Z",
     "iopub.status.idle": "2025-01-09T00:28:51.489202Z",
     "shell.execute_reply": "2025-01-09T00:28:51.488591Z",
     "shell.execute_reply.started": "2025-01-09T00:28:43.818520Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>method</th>\n",
       "      <th>solver</th>\n",
       "      <th>demeaner_backend</th>\n",
       "      <th>n_obs</th>\n",
       "      <th>k</th>\n",
       "      <th>G</th>\n",
       "      <th>rep</th>\n",
       "      <th>full_feols_timing</th>\n",
       "      <th>demean_timing</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>numba</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0.150473</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>numba</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>0.147583</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>numba</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>0.186491</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>numba</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>4</td>\n",
       "      <td>0.190972</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>numba</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>5</td>\n",
       "      <td>0.162773</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>numba</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>6</td>\n",
       "      <td>0.171777</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>numba</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>7</td>\n",
       "      <td>0.166872</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>numba</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>8</td>\n",
       "      <td>0.158694</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>numba</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>9</td>\n",
       "      <td>0.185547</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>numba</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>10</td>\n",
       "      <td>0.158114</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  method           solver demeaner_backend n_obs  k  G rep full_feols_timing  \\\n",
       "0  feols  np.linalg.lstsq            numba  1000  1  1   1          0.150473   \n",
       "1  feols  np.linalg.lstsq            numba  1000  1  1   2          0.147583   \n",
       "2  feols  np.linalg.lstsq            numba  1000  1  1   3          0.186491   \n",
       "3  feols  np.linalg.lstsq            numba  1000  1  1   4          0.190972   \n",
       "4  feols  np.linalg.lstsq            numba  1000  1  1   5          0.162773   \n",
       "5  feols  np.linalg.lstsq            numba  1000  1  1   6          0.171777   \n",
       "6  feols  np.linalg.lstsq            numba  1000  1  1   7          0.166872   \n",
       "7  feols  np.linalg.lstsq            numba  1000  1  1   8          0.158694   \n",
       "8  feols  np.linalg.lstsq            numba  1000  1  1   9          0.185547   \n",
       "9  feols  np.linalg.lstsq            numba  1000  1  1  10          0.158114   \n",
       "\n",
       "  demean_timing  \n",
       "0           NaN  \n",
       "1           NaN  \n",
       "2           NaN  \n",
       "3           NaN  \n",
       "4           NaN  \n",
       "5           NaN  \n",
       "6           NaN  \n",
       "7           NaN  \n",
       "8           NaN  \n",
       "9           NaN  "
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# test run numba\n",
    "run_standard_benchmark(fixed_effect=\"dum_1\", demeaner_backend=\"numba\", size=1, k=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-01-09T00:28:43.818536Z",
     "iopub.status.busy": "2025-01-09T00:28:43.818246Z",
     "iopub.status.idle": "2025-01-09T00:28:51.489202Z",
     "shell.execute_reply": "2025-01-09T00:28:51.488591Z",
     "shell.execute_reply.started": "2025-01-09T00:28:43.818520Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>method</th>\n",
       "      <th>solver</th>\n",
       "      <th>demeaner_backend</th>\n",
       "      <th>n_obs</th>\n",
       "      <th>k</th>\n",
       "      <th>G</th>\n",
       "      <th>rep</th>\n",
       "      <th>full_feols_timing</th>\n",
       "      <th>demean_timing</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>jax</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0.122831</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>jax</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>0.122887</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>jax</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>0.136041</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>jax</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>4</td>\n",
       "      <td>0.139644</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>jax</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>5</td>\n",
       "      <td>0.136235</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>jax</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>6</td>\n",
       "      <td>0.122477</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>jax</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>7</td>\n",
       "      <td>0.123122</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>jax</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>8</td>\n",
       "      <td>0.119589</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>jax</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>9</td>\n",
       "      <td>0.122247</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>feols</td>\n",
       "      <td>np.linalg.lstsq</td>\n",
       "      <td>jax</td>\n",
       "      <td>1000</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>10</td>\n",
       "      <td>0.118353</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  method           solver demeaner_backend n_obs  k  G rep full_feols_timing  \\\n",
       "0  feols  np.linalg.lstsq              jax  1000  1  1   1          0.122831   \n",
       "1  feols  np.linalg.lstsq              jax  1000  1  1   2          0.122887   \n",
       "2  feols  np.linalg.lstsq              jax  1000  1  1   3          0.136041   \n",
       "3  feols  np.linalg.lstsq              jax  1000  1  1   4          0.139644   \n",
       "4  feols  np.linalg.lstsq              jax  1000  1  1   5          0.136235   \n",
       "5  feols  np.linalg.lstsq              jax  1000  1  1   6          0.122477   \n",
       "6  feols  np.linalg.lstsq              jax  1000  1  1   7          0.123122   \n",
       "7  feols  np.linalg.lstsq              jax  1000  1  1   8          0.119589   \n",
       "8  feols  np.linalg.lstsq              jax  1000  1  1   9          0.122247   \n",
       "9  feols  np.linalg.lstsq              jax  1000  1  1  10          0.118353   \n",
       "\n",
       "  demean_timing  \n",
       "0           NaN  \n",
       "1           NaN  \n",
       "2           NaN  \n",
       "3           NaN  \n",
       "4           NaN  \n",
       "5           NaN  \n",
       "6           NaN  \n",
       "7           NaN  \n",
       "8           NaN  \n",
       "9           NaN  "
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# test run jax\n",
    "run_standard_benchmark(fixed_effect=\"dum_1\", demeaner_backend=\"jax\", size=1, k=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-01-09T00:25:39.248695Z",
     "iopub.status.busy": "2025-01-09T00:25:39.248317Z",
     "iopub.status.idle": "2025-01-09T00:25:39.253652Z",
     "shell.execute_reply": "2025-01-09T00:25:39.253000Z",
     "shell.execute_reply.started": "2025-01-09T00:25:39.248671Z"
    },
    "id": "HxsGRMSlR_jI"
   },
   "outputs": [],
   "source": [
    "def run_all_benchmarks(size_list, k_list):\n",
    "    \"\"\"\n",
    "    Run all the benchmarks.\n",
    "\n",
    "    Args:\n",
    "        size_list (list): The list of sizes to run the benchmarks on. 1-> 1000, 2-> 10000, ..., 5-> 10_000_000\n",
    "        k_list (list): The list of k values to run the benchmarks on.\n",
    "    \"\"\"\n",
    "    res = pd.DataFrame()\n",
    "\n",
    "    all_combinations = list(\n",
    "        product(\n",
    "            [\"numba\", \"jax\"],  # demeaner_backend\n",
    "            [\"dum_1\", \"dum_1+dum_2\", \"dum_1+dum_2+dum_3\"],  # fixef\n",
    "            size_list,  # size\n",
    "            k_list,  # k\n",
    "            [\"np.linalg.lstsq\"],  # solver\n",
    "        )\n",
    "    )\n",
    "\n",
    "    with tqdm(total=len(all_combinations), desc=\"Running Benchmarks\") as pbar:\n",
    "        for demeaner_backend, fixef, size, k, solver in all_combinations:\n",
    "            res = pd.concat(\n",
    "                [\n",
    "                    res,\n",
    "                    run_standard_benchmark(\n",
    "                        solver=solver,\n",
    "                        fixed_effect=fixef,\n",
    "                        demeaner_backend=demeaner_backend,\n",
    "                        size=size,\n",
    "                        k=k,\n",
    "                    ),\n",
    "                ],\n",
    "                axis=0,\n",
    "            )\n",
    "            pbar.update(1)  # Update the progress bar after each iteration\n",
    "\n",
    "    return res"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run Benchmarks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "background_save": true,
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2025-01-09T00:25:39.962554Z",
     "iopub.status.busy": "2025-01-09T00:25:39.962039Z",
     "iopub.status.idle": "2025-01-09T00:26:25.319310Z",
     "shell.execute_reply": "2025-01-09T00:26:25.318687Z",
     "shell.execute_reply.started": "2025-01-09T00:25:39.962536Z"
    },
    "id": "gki1mlqvSEIi",
    "outputId": "3cb40095-df81-4e78-99a6-2410da237884"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Running Benchmarks: 100%|██████████| 24/24 [00:49<00:00,  2.06s/it]\n"
     ]
    }
   ],
   "source": [
    "res_all = run_all_benchmarks(\n",
    "    size_list=[1, 2, 3, 4, 5],  # for N = 1000, 10_000, 100_000, 1_000_000, 10_000_000\n",
    "    k_list=[1, 10, 50, 100],  # for k = 1, 10, 50, 100\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "background_save": true,
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2025-01-09T00:26:25.320894Z",
     "iopub.status.busy": "2025-01-09T00:26:25.320596Z",
     "iopub.status.idle": "2025-01-09T00:26:26.391854Z",
     "shell.execute_reply": "2025-01-09T00:26:26.391236Z",
     "shell.execute_reply.started": "2025-01-09T00:26:25.320871Z"
    },
    "id": "7zEIHj5nXXvq",
    "outputId": "238dedce-a9f5-4a1b-b9a2-b84d4a89c091"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>method</th>\n",
       "      <th>demeaner_backend</th>\n",
       "      <th>k</th>\n",
       "      <th>G</th>\n",
       "      <th>n_obs</th>\n",
       "      <th>full_feols_timing</th>\n",
       "      <th>demean_timing</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>feols</td>\n",
       "      <td>jax</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1000</td>\n",
       "      <td>0.127274</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>feols</td>\n",
       "      <td>jax</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>10000</td>\n",
       "      <td>0.142241</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>feols</td>\n",
       "      <td>jax</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>1000</td>\n",
       "      <td>0.134405</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>feols</td>\n",
       "      <td>jax</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>10000</td>\n",
       "      <td>0.168555</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>feols</td>\n",
       "      <td>jax</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>1000</td>\n",
       "      <td>0.139142</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>feols</td>\n",
       "      <td>jax</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>10000</td>\n",
       "      <td>0.189492</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>feols</td>\n",
       "      <td>jax</td>\n",
       "      <td>5</td>\n",
       "      <td>1</td>\n",
       "      <td>1000</td>\n",
       "      <td>0.144933</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>feols</td>\n",
       "      <td>jax</td>\n",
       "      <td>5</td>\n",
       "      <td>1</td>\n",
       "      <td>10000</td>\n",
       "      <td>0.147049</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>feols</td>\n",
       "      <td>jax</td>\n",
       "      <td>5</td>\n",
       "      <td>2</td>\n",
       "      <td>1000</td>\n",
       "      <td>0.147216</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>feols</td>\n",
       "      <td>jax</td>\n",
       "      <td>5</td>\n",
       "      <td>2</td>\n",
       "      <td>10000</td>\n",
       "      <td>0.175394</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>feols</td>\n",
       "      <td>jax</td>\n",
       "      <td>5</td>\n",
       "      <td>3</td>\n",
       "      <td>1000</td>\n",
       "      <td>0.15034</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>feols</td>\n",
       "      <td>jax</td>\n",
       "      <td>5</td>\n",
       "      <td>3</td>\n",
       "      <td>10000</td>\n",
       "      <td>0.194496</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>feols</td>\n",
       "      <td>numba</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1000</td>\n",
       "      <td>0.153849</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>feols</td>\n",
       "      <td>numba</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>10000</td>\n",
       "      <td>0.183697</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>feols</td>\n",
       "      <td>numba</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>1000</td>\n",
       "      <td>0.17201</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>feols</td>\n",
       "      <td>numba</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>10000</td>\n",
       "      <td>0.16131</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>feols</td>\n",
       "      <td>numba</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>1000</td>\n",
       "      <td>0.169463</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>feols</td>\n",
       "      <td>numba</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>10000</td>\n",
       "      <td>0.169789</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>feols</td>\n",
       "      <td>numba</td>\n",
       "      <td>5</td>\n",
       "      <td>1</td>\n",
       "      <td>1000</td>\n",
       "      <td>0.331457</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>feols</td>\n",
       "      <td>numba</td>\n",
       "      <td>5</td>\n",
       "      <td>1</td>\n",
       "      <td>10000</td>\n",
       "      <td>0.172194</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>feols</td>\n",
       "      <td>numba</td>\n",
       "      <td>5</td>\n",
       "      <td>2</td>\n",
       "      <td>1000</td>\n",
       "      <td>0.162562</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>feols</td>\n",
       "      <td>numba</td>\n",
       "      <td>5</td>\n",
       "      <td>2</td>\n",
       "      <td>10000</td>\n",
       "      <td>0.184233</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>feols</td>\n",
       "      <td>numba</td>\n",
       "      <td>5</td>\n",
       "      <td>3</td>\n",
       "      <td>1000</td>\n",
       "      <td>0.174623</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>23</th>\n",
       "      <td>feols</td>\n",
       "      <td>numba</td>\n",
       "      <td>5</td>\n",
       "      <td>3</td>\n",
       "      <td>10000</td>\n",
       "      <td>0.166674</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   method demeaner_backend  k  G  n_obs full_feols_timing demean_timing\n",
       "0   feols              jax  1  1   1000          0.127274           NaN\n",
       "1   feols              jax  1  1  10000          0.142241           NaN\n",
       "2   feols              jax  1  2   1000          0.134405           NaN\n",
       "3   feols              jax  1  2  10000          0.168555           NaN\n",
       "4   feols              jax  1  3   1000          0.139142           NaN\n",
       "5   feols              jax  1  3  10000          0.189492           NaN\n",
       "6   feols              jax  5  1   1000          0.144933           NaN\n",
       "7   feols              jax  5  1  10000          0.147049           NaN\n",
       "8   feols              jax  5  2   1000          0.147216           NaN\n",
       "9   feols              jax  5  2  10000          0.175394           NaN\n",
       "10  feols              jax  5  3   1000           0.15034           NaN\n",
       "11  feols              jax  5  3  10000          0.194496           NaN\n",
       "12  feols            numba  1  1   1000          0.153849           NaN\n",
       "13  feols            numba  1  1  10000          0.183697           NaN\n",
       "14  feols            numba  1  2   1000           0.17201           NaN\n",
       "15  feols            numba  1  2  10000           0.16131           NaN\n",
       "16  feols            numba  1  3   1000          0.169463           NaN\n",
       "17  feols            numba  1  3  10000          0.169789           NaN\n",
       "18  feols            numba  5  1   1000          0.331457           NaN\n",
       "19  feols            numba  5  1  10000          0.172194           NaN\n",
       "20  feols            numba  5  2   1000          0.162562           NaN\n",
       "21  feols            numba  5  2  10000          0.184233           NaN\n",
       "22  feols            numba  5  3   1000          0.174623           NaN\n",
       "23  feols            numba  5  3  10000          0.166674           NaN"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = (\n",
    "    res_all.drop([\"rep\", \"solver\"], axis=1)\n",
    "    .groupby([\"method\", \"demeaner_backend\", \"k\", \"G\", \"n_obs\"])\n",
    "    .mean()\n",
    "    .reset_index()\n",
    ")\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "VCn6O5MMXlBw"
   },
   "source": [
    "## Visualize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:matplotlib.category:Using categorical units to plot a list of strings that are all parsable as floats or dates. If these strings should be plotted as numbers, cast to the appropriate data type before plotting.\n",
      "INFO:matplotlib.category:Using categorical units to plot a list of strings that are all parsable as floats or dates. If these strings should be plotted as numbers, cast to the appropriate data type before plotting.\n",
      "INFO:matplotlib.category:Using categorical units to plot a list of strings that are all parsable as floats or dates. If these strings should be plotted as numbers, cast to the appropriate data type before plotting.\n",
      "INFO:matplotlib.category:Using categorical units to plot a list of strings that are all parsable as floats or dates. If these strings should be plotted as numbers, cast to the appropriate data type before plotting.\n",
      "INFO:matplotlib.category:Using categorical units to plot a list of strings that are all parsable as floats or dates. If these strings should be plotted as numbers, cast to the appropriate data type before plotting.\n",
      "INFO:matplotlib.category:Using categorical units to plot a list of strings that are all parsable as floats or dates. If these strings should be plotted as numbers, cast to the appropriate data type before plotting.\n",
      "INFO:matplotlib.category:Using categorical units to plot a list of strings that are all parsable as floats or dates. If these strings should be plotted as numbers, cast to the appropriate data type before plotting.\n",
      "INFO:matplotlib.category:Using categorical units to plot a list of strings that are all parsable as floats or dates. If these strings should be plotted as numbers, cast to the appropriate data type before plotting.\n",
      "INFO:matplotlib.category:Using categorical units to plot a list of strings that are all parsable as floats or dates. If these strings should be plotted as numbers, cast to the appropriate data type before plotting.\n",
      "INFO:matplotlib.category:Using categorical units to plot a list of strings that are all parsable as floats or dates. If these strings should be plotted as numbers, cast to the appropriate data type before plotting.\n",
      "INFO:matplotlib.category:Using categorical units to plot a list of strings that are all parsable as floats or dates. If these strings should be plotted as numbers, cast to the appropriate data type before plotting.\n",
      "INFO:matplotlib.category:Using categorical units to plot a list of strings that are all parsable as floats or dates. If these strings should be plotted as numbers, cast to the appropriate data type before plotting.\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 58432 characters]",
      "text/plain": [
       "<Figure size 1587.99x800 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "df[\"G\"] = df[\"G\"].map({1: \"n_fixef = 1\", 2: \"n_fixef = 2\", 3: \"n_fixef = 3\"})\n",
    "df[\"n_obs\"] = df[\"n_obs\"].astype(str)\n",
    "\n",
    "# Dynamically determine unique values for order and hue_order\n",
    "n_obs_order = sorted(df[\"n_obs\"].unique(), key=lambda x: int(x))  # Sort as integers\n",
    "demeaner_backend_order = df[\"demeaner_backend\"].unique()\n",
    "\n",
    "custom_palette = sns.color_palette(\"coolwarm\", n_colors=2)\n",
    "\n",
    "# Create the FacetGrid with reordered columns and rows\n",
    "g = sns.FacetGrid(\n",
    "    df,\n",
    "    col=\"G\",  # G (n_fixef) increases left to right\n",
    "    row=\"k\",  # k increases top to bottom\n",
    "    margin_titles=True,\n",
    "    height=4,\n",
    "    aspect=1.2,\n",
    "    col_order=[\"n_fixef = 1\", \"n_fixef = 2\", \"n_fixef = 3\"],  # Ensure correct order\n",
    ")\n",
    "\n",
    "# Plot the bar chart for each facet with the custom palette\n",
    "g.map(\n",
    "    sns.barplot,\n",
    "    \"n_obs\",\n",
    "    \"full_feols_timing\",\n",
    "    \"demeaner_backend\",\n",
    "    order=n_obs_order,  # Dynamic order for n_obs\n",
    "    hue_order=demeaner_backend_order,  # Dynamic hue order for demeaner_backend\n",
    "    errorbar=None,  # Suppress error bars\n",
    "    palette=custom_palette,\n",
    ")\n",
    "\n",
    "# Add legend and adjust layout\n",
    "g.add_legend(title=\"Demeaner Backend\")\n",
    "g.set_axis_labels(\"Number of Observations\", \"Runtime (seconds)\")\n",
    "g.set_titles(row_template=\"k = {row_name}\", col_template=\"{col_name}\")\n",
    "plt.subplots_adjust(top=0.9)\n",
    "g.fig.suptitle(\"Runtime vs Number of Observations by n_fixef and k\")\n",
    "\n",
    "# Show plot\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "jax",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
