benchmark.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import statsmodels.formula.api as smf\n",
    "import fastreg as fr\n",
    "from fastreg import I, R, C, C0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Generate Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "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>x1</th>\n",
       "      <th>x2</th>\n",
       "      <th>yhat0</th>\n",
       "      <th>yhat</th>\n",
       "      <th>id1</th>\n",
       "      <th>id2</th>\n",
       "      <th>y0</th>\n",
       "      <th>y</th>\n",
       "      <th>Eb0</th>\n",
       "      <th>Eb</th>\n",
       "      <th>b0</th>\n",
       "      <th>b</th>\n",
       "      <th>Ep0</th>\n",
       "      <th>Ep</th>\n",
       "      <th>p0</th>\n",
       "      <th>p</th>\n",
       "      <th>nb0</th>\n",
       "      <th>nb</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>-1.429225</td>\n",
       "      <td>1.830295</td>\n",
       "      <td>0.769410</td>\n",
       "      <td>2.149410</td>\n",
       "      <td>H</td>\n",
       "      <td>68</td>\n",
       "      <td>1.791671</td>\n",
       "      <td>1.585510</td>\n",
       "      <td>0.683393</td>\n",
       "      <td>0.895614</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>2.158492</td>\n",
       "      <td>8.579792</td>\n",
       "      <td>4</td>\n",
       "      <td>10</td>\n",
       "      <td>1</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0.687153</td>\n",
       "      <td>0.795373</td>\n",
       "      <td>0.783369</td>\n",
       "      <td>1.403369</td>\n",
       "      <td>C</td>\n",
       "      <td>42</td>\n",
       "      <td>0.295552</td>\n",
       "      <td>3.163529</td>\n",
       "      <td>0.686406</td>\n",
       "      <td>0.802718</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>2.188835</td>\n",
       "      <td>4.068886</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0.764353</td>\n",
       "      <td>-1.060225</td>\n",
       "      <td>-0.306829</td>\n",
       "      <td>-0.036829</td>\n",
       "      <td>C</td>\n",
       "      <td>7</td>\n",
       "      <td>0.381637</td>\n",
       "      <td>-0.639362</td>\n",
       "      <td>0.423889</td>\n",
       "      <td>0.490794</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0.735777</td>\n",
       "      <td>0.963841</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>-0.309887</td>\n",
       "      <td>-0.777701</td>\n",
       "      <td>-0.459587</td>\n",
       "      <td>0.800413</td>\n",
       "      <td>D</td>\n",
       "      <td>96</td>\n",
       "      <td>0.143218</td>\n",
       "      <td>1.139854</td>\n",
       "      <td>0.387084</td>\n",
       "      <td>0.690063</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0.631545</td>\n",
       "      <td>2.226461</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>-0.616042</td>\n",
       "      <td>-0.752266</td>\n",
       "      <td>-0.536172</td>\n",
       "      <td>0.623828</td>\n",
       "      <td>H</td>\n",
       "      <td>46</td>\n",
       "      <td>-0.468395</td>\n",
       "      <td>1.462211</td>\n",
       "      <td>0.369078</td>\n",
       "      <td>0.651089</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0.584983</td>\n",
       "      <td>1.866057</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "         x1        x2     yhat0      yhat id1  id2        y0         y  \\\n",
       "0 -1.429225  1.830295  0.769410  2.149410   H   68  1.791671  1.585510   \n",
       "1  0.687153  0.795373  0.783369  1.403369   C   42  0.295552  3.163529   \n",
       "2  0.764353 -1.060225 -0.306829 -0.036829   C    7  0.381637 -0.639362   \n",
       "3 -0.309887 -0.777701 -0.459587  0.800413   D   96  0.143218  1.139854   \n",
       "4 -0.616042 -0.752266 -0.536172  0.623828   H   46 -0.468395  1.462211   \n",
       "\n",
       "        Eb0        Eb  b0  b       Ep0        Ep  p0   p  nb0  nb  \n",
       "0  0.683393  0.895614   1  1  2.158492  8.579792   4  10    1   5  \n",
       "1  0.686406  0.802718   1  1  2.188835  4.068886   1   2    1   0  \n",
       "2  0.423889  0.490794   1  1  0.735777  0.963841   0   1    2   0  \n",
       "3  0.387084  0.690063   0  1  0.631545  2.226461   0   0    1   3  \n",
       "4  0.369078  0.651089   1  1  0.584983  1.866057   0   1    1   7  "
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "models = ['linear', 'poisson', 'logit', 'negbin']\n",
    "data = fr.dataset(N=1_000_000, K1=10, K2=100, models=models, seed=89320432)\n",
    "data_wide = fr.dataset(N=1_000_000, K1=10, K2=10_000, models=models, seed=89320433)\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# for statsmodels runs\n",
    "data1 = data.copy()\n",
    "data1['id2'] = data1['id2'].astype(str)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Normal OLS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 59.5 ms, sys: 128 ms, total: 187 ms\n",
      "Wall time: 24.4 ms\n"
     ]
    },
    {
     "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>y0</th>\n",
       "      <th>coeff</th>\n",
       "      <th>stderr</th>\n",
       "      <th>low95</th>\n",
       "      <th>high95</th>\n",
       "      <th>pvalue</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>I</th>\n",
       "      <td>0.101859</td>\n",
       "      <td>0.001000</td>\n",
       "      <td>0.099900</td>\n",
       "      <td>0.103818</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x1</th>\n",
       "      <td>0.301647</td>\n",
       "      <td>0.001000</td>\n",
       "      <td>0.299688</td>\n",
       "      <td>0.303606</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x2</th>\n",
       "      <td>0.599408</td>\n",
       "      <td>0.000999</td>\n",
       "      <td>0.597450</td>\n",
       "      <td>0.601366</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "y0     coeff    stderr     low95    high95  pvalue\n",
       "I   0.101859  0.001000  0.099900  0.103818     0.0\n",
       "x1  0.301647  0.001000  0.299688  0.303606     0.0\n",
       "x2  0.599408  0.000999  0.597450  0.601366     0.0"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time fr.ols(y=R.y0, x=I+R.x1+R.x2, data=data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 422 ms, sys: 1.39 s, total: 1.81 s\n",
      "Wall time: 176 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Intercept    0.101859\n",
       "x1           0.301647\n",
       "x2           0.599408\n",
       "dtype: float64"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time smf.ols('y0 ~ 1 + x1 + x2', data=data).fit().params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 455 ms, sys: 577 ms, total: 1.03 s\n",
      "Wall time: 417 ms\n"
     ]
    },
    {
     "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>y</th>\n",
       "      <th>coeff</th>\n",
       "      <th>stderr</th>\n",
       "      <th>low95</th>\n",
       "      <th>high95</th>\n",
       "      <th>pvalue</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>I</th>\n",
       "      <td>0.112463</td>\n",
       "      <td>0.010413</td>\n",
       "      <td>0.092054</td>\n",
       "      <td>0.132872</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x1</th>\n",
       "      <td>0.299557</td>\n",
       "      <td>0.001001</td>\n",
       "      <td>0.297595</td>\n",
       "      <td>0.301518</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x2</th>\n",
       "      <td>0.601840</td>\n",
       "      <td>0.001000</td>\n",
       "      <td>0.599880</td>\n",
       "      <td>0.603800</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=B</th>\n",
       "      <td>0.104734</td>\n",
       "      <td>0.004476</td>\n",
       "      <td>0.095962</td>\n",
       "      <td>0.113507</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=C</th>\n",
       "      <td>0.201965</td>\n",
       "      <td>0.004474</td>\n",
       "      <td>0.193197</td>\n",
       "      <td>0.210734</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=95</th>\n",
       "      <td>0.937461</td>\n",
       "      <td>0.014102</td>\n",
       "      <td>0.909821</td>\n",
       "      <td>0.965101</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=96</th>\n",
       "      <td>0.956624</td>\n",
       "      <td>0.014109</td>\n",
       "      <td>0.928970</td>\n",
       "      <td>0.984278</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=97</th>\n",
       "      <td>0.960951</td>\n",
       "      <td>0.014099</td>\n",
       "      <td>0.933318</td>\n",
       "      <td>0.988584</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=98</th>\n",
       "      <td>0.993398</td>\n",
       "      <td>0.014132</td>\n",
       "      <td>0.965699</td>\n",
       "      <td>1.021097</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=99</th>\n",
       "      <td>0.988670</td>\n",
       "      <td>0.014066</td>\n",
       "      <td>0.961101</td>\n",
       "      <td>1.016239</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>111 rows × 5 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "y          coeff    stderr     low95    high95  pvalue\n",
       "I       0.112463  0.010413  0.092054  0.132872     0.0\n",
       "x1      0.299557  0.001001  0.297595  0.301518     0.0\n",
       "x2      0.601840  0.001000  0.599880  0.603800     0.0\n",
       "id1=B   0.104734  0.004476  0.095962  0.113507     0.0\n",
       "id1=C   0.201965  0.004474  0.193197  0.210734     0.0\n",
       "...          ...       ...       ...       ...     ...\n",
       "id2=95  0.937461  0.014102  0.909821  0.965101     0.0\n",
       "id2=96  0.956624  0.014109  0.928970  0.984278     0.0\n",
       "id2=97  0.960951  0.014099  0.933318  0.988584     0.0\n",
       "id2=98  0.993398  0.014132  0.965699  1.021097     0.0\n",
       "id2=99  0.988670  0.014066  0.961101  1.016239     0.0\n",
       "\n",
       "[111 rows x 5 columns]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time fr.ols(y=R.y, x=I+R.x1+R.x2+C.id1+C.id2, data=data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 38s, sys: 14.6 s, total: 1min 53s\n",
      "Wall time: 10.8 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Intercept    0.112463\n",
       "id1[T.B]     0.104734\n",
       "id1[T.C]     0.201965\n",
       "id1[T.D]     0.300212\n",
       "id1[T.E]     0.403094\n",
       "               ...   \n",
       "id2[T.97]    0.960951\n",
       "id2[T.98]    0.993398\n",
       "id2[T.99]    0.988670\n",
       "x1           0.299557\n",
       "x2           0.601840\n",
       "Length: 111, dtype: float64"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time smf.ols('y ~ 1 + x1 + x2 + id1 + id2', data=data1).fit().params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### High Dimensional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4.1 s, sys: 2.47 s, total: 6.57 s\n",
      "Wall time: 4.53 s\n"
     ]
    },
    {
     "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>y</th>\n",
       "      <th>coeff</th>\n",
       "      <th>stderr</th>\n",
       "      <th>low95</th>\n",
       "      <th>high95</th>\n",
       "      <th>pvalue</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>I</th>\n",
       "      <td>0.143877</td>\n",
       "      <td>0.103699</td>\n",
       "      <td>-0.059369</td>\n",
       "      <td>0.347123</td>\n",
       "      <td>1.653041e-01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x1</th>\n",
       "      <td>0.301630</td>\n",
       "      <td>0.001004</td>\n",
       "      <td>0.299662</td>\n",
       "      <td>0.303597</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x2</th>\n",
       "      <td>0.599284</td>\n",
       "      <td>0.001006</td>\n",
       "      <td>0.597313</td>\n",
       "      <td>0.601255</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=B</th>\n",
       "      <td>0.097642</td>\n",
       "      <td>0.004499</td>\n",
       "      <td>0.088824</td>\n",
       "      <td>0.106461</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=C</th>\n",
       "      <td>0.189008</td>\n",
       "      <td>0.004491</td>\n",
       "      <td>0.180207</td>\n",
       "      <td>0.197810</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=9995</th>\n",
       "      <td>1.045182</td>\n",
       "      <td>0.146196</td>\n",
       "      <td>0.758643</td>\n",
       "      <td>1.331722</td>\n",
       "      <td>8.730794e-13</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=9996</th>\n",
       "      <td>1.029893</td>\n",
       "      <td>0.146587</td>\n",
       "      <td>0.742588</td>\n",
       "      <td>1.317198</td>\n",
       "      <td>2.128075e-12</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=9997</th>\n",
       "      <td>0.892911</td>\n",
       "      <td>0.142021</td>\n",
       "      <td>0.614554</td>\n",
       "      <td>1.171267</td>\n",
       "      <td>3.233234e-10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=9998</th>\n",
       "      <td>0.961372</td>\n",
       "      <td>0.144705</td>\n",
       "      <td>0.677755</td>\n",
       "      <td>1.244989</td>\n",
       "      <td>3.059886e-11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=9999</th>\n",
       "      <td>0.956185</td>\n",
       "      <td>0.149092</td>\n",
       "      <td>0.663969</td>\n",
       "      <td>1.248401</td>\n",
       "      <td>1.423330e-10</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>10011 rows × 5 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "y            coeff    stderr     low95    high95        pvalue\n",
       "I         0.143877  0.103699 -0.059369  0.347123  1.653041e-01\n",
       "x1        0.301630  0.001004  0.299662  0.303597  0.000000e+00\n",
       "x2        0.599284  0.001006  0.597313  0.601255  0.000000e+00\n",
       "id1=B     0.097642  0.004499  0.088824  0.106461  0.000000e+00\n",
       "id1=C     0.189008  0.004491  0.180207  0.197810  0.000000e+00\n",
       "...            ...       ...       ...       ...           ...\n",
       "id2=9995  1.045182  0.146196  0.758643  1.331722  8.730794e-13\n",
       "id2=9996  1.029893  0.146587  0.742588  1.317198  2.128075e-12\n",
       "id2=9997  0.892911  0.142021  0.614554  1.171267  3.233234e-10\n",
       "id2=9998  0.961372  0.144705  0.677755  1.244989  3.059886e-11\n",
       "id2=9999  0.956185  0.149092  0.663969  1.248401  1.423330e-10\n",
       "\n",
       "[10011 rows x 5 columns]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time fr.ols(y=R.y, x=I+R.x1+R.x2+C.id1+C.id2, data=data_wide)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 366 ms, sys: 40.7 ms, total: 407 ms\n",
      "Wall time: 358 ms\n"
     ]
    },
    {
     "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>y</th>\n",
       "      <th>coeff</th>\n",
       "      <th>stderr</th>\n",
       "      <th>low95</th>\n",
       "      <th>high95</th>\n",
       "      <th>pvalue</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>I</th>\n",
       "      <td>0.143877</td>\n",
       "      <td>0.103699</td>\n",
       "      <td>-0.059369</td>\n",
       "      <td>0.347123</td>\n",
       "      <td>1.653041e-01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x1</th>\n",
       "      <td>0.301630</td>\n",
       "      <td>0.001004</td>\n",
       "      <td>0.299662</td>\n",
       "      <td>0.303597</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x2</th>\n",
       "      <td>0.599284</td>\n",
       "      <td>0.001006</td>\n",
       "      <td>0.597313</td>\n",
       "      <td>0.601255</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=B</th>\n",
       "      <td>0.097642</td>\n",
       "      <td>0.004499</td>\n",
       "      <td>0.088824</td>\n",
       "      <td>0.106461</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=C</th>\n",
       "      <td>0.189008</td>\n",
       "      <td>0.004491</td>\n",
       "      <td>0.180207</td>\n",
       "      <td>0.197810</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=9995</th>\n",
       "      <td>1.045182</td>\n",
       "      <td>0.146196</td>\n",
       "      <td>0.758643</td>\n",
       "      <td>1.331722</td>\n",
       "      <td>8.730794e-13</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=9996</th>\n",
       "      <td>1.029893</td>\n",
       "      <td>0.146587</td>\n",
       "      <td>0.742588</td>\n",
       "      <td>1.317198</td>\n",
       "      <td>2.128075e-12</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=9997</th>\n",
       "      <td>0.892911</td>\n",
       "      <td>0.142021</td>\n",
       "      <td>0.614554</td>\n",
       "      <td>1.171267</td>\n",
       "      <td>3.233234e-10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=9998</th>\n",
       "      <td>0.961372</td>\n",
       "      <td>0.144705</td>\n",
       "      <td>0.677755</td>\n",
       "      <td>1.244989</td>\n",
       "      <td>3.059886e-11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=9999</th>\n",
       "      <td>0.956185</td>\n",
       "      <td>0.149092</td>\n",
       "      <td>0.663969</td>\n",
       "      <td>1.248401</td>\n",
       "      <td>1.423330e-10</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>10011 rows × 5 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "y            coeff    stderr     low95    high95        pvalue\n",
       "I         0.143877  0.103699 -0.059369  0.347123  1.653041e-01\n",
       "x1        0.301630  0.001004  0.299662  0.303597  0.000000e+00\n",
       "x2        0.599284  0.001006  0.597313  0.601255  0.000000e+00\n",
       "id1=B     0.097642  0.004499  0.088824  0.106461  0.000000e+00\n",
       "id1=C     0.189008  0.004491  0.180207  0.197810  0.000000e+00\n",
       "...            ...       ...       ...       ...           ...\n",
       "id2=9995  1.045182  0.146196  0.758643  1.331722  8.730794e-13\n",
       "id2=9996  1.029893  0.146587  0.742588  1.317198  2.128075e-12\n",
       "id2=9997  0.892911  0.142021  0.614554  1.171267  3.233234e-10\n",
       "id2=9998  0.961372  0.144705  0.677755  1.244989  3.059886e-11\n",
       "id2=9999  0.956185  0.149092  0.663969  1.248401  1.423330e-10\n",
       "\n",
       "[10011 rows x 5 columns]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time fr.ols(y=R.y, x=I+R.x1+R.x2+C.id1, hdfe=C.id2, data=data_wide)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 765 ms, sys: 1.84 s, total: 2.61 s\n",
      "Wall time: 355 ms\n"
     ]
    },
    {
     "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>y</th>\n",
       "      <th>coeff</th>\n",
       "      <th>stderr</th>\n",
       "      <th>low95</th>\n",
       "      <th>high95</th>\n",
       "      <th>pvalue</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>I</th>\n",
       "      <td>0.605524</td>\n",
       "      <td>0.003035</td>\n",
       "      <td>0.599576</td>\n",
       "      <td>0.611471</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x1</th>\n",
       "      <td>0.301630</td>\n",
       "      <td>0.000998</td>\n",
       "      <td>0.299674</td>\n",
       "      <td>0.303585</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x2</th>\n",
       "      <td>0.599284</td>\n",
       "      <td>0.001014</td>\n",
       "      <td>0.597297</td>\n",
       "      <td>0.601270</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=B</th>\n",
       "      <td>0.097642</td>\n",
       "      <td>0.004531</td>\n",
       "      <td>0.088762</td>\n",
       "      <td>0.106523</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=C</th>\n",
       "      <td>0.189008</td>\n",
       "      <td>0.004545</td>\n",
       "      <td>0.180100</td>\n",
       "      <td>0.197917</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=D</th>\n",
       "      <td>0.295160</td>\n",
       "      <td>0.004506</td>\n",
       "      <td>0.286328</td>\n",
       "      <td>0.303991</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=E</th>\n",
       "      <td>0.402802</td>\n",
       "      <td>0.004510</td>\n",
       "      <td>0.393963</td>\n",
       "      <td>0.411641</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=F</th>\n",
       "      <td>0.495409</td>\n",
       "      <td>0.004493</td>\n",
       "      <td>0.486602</td>\n",
       "      <td>0.504215</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=G</th>\n",
       "      <td>0.596910</td>\n",
       "      <td>0.004467</td>\n",
       "      <td>0.588155</td>\n",
       "      <td>0.605665</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=H</th>\n",
       "      <td>0.694326</td>\n",
       "      <td>0.004490</td>\n",
       "      <td>0.685526</td>\n",
       "      <td>0.703125</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=I</th>\n",
       "      <td>0.791431</td>\n",
       "      <td>0.004541</td>\n",
       "      <td>0.782532</td>\n",
       "      <td>0.800331</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=J</th>\n",
       "      <td>0.889866</td>\n",
       "      <td>0.004498</td>\n",
       "      <td>0.881050</td>\n",
       "      <td>0.898682</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "y         coeff    stderr     low95    high95  pvalue\n",
       "I      0.605524  0.003035  0.599576  0.611471     0.0\n",
       "x1     0.301630  0.000998  0.299674  0.303585     0.0\n",
       "x2     0.599284  0.001014  0.597297  0.601270     0.0\n",
       "id1=B  0.097642  0.004531  0.088762  0.106523     0.0\n",
       "id1=C  0.189008  0.004545  0.180100  0.197917     0.0\n",
       "id1=D  0.295160  0.004506  0.286328  0.303991     0.0\n",
       "id1=E  0.402802  0.004510  0.393963  0.411641     0.0\n",
       "id1=F  0.495409  0.004493  0.486602  0.504215     0.0\n",
       "id1=G  0.596910  0.004467  0.588155  0.605665     0.0\n",
       "id1=H  0.694326  0.004490  0.685526  0.703125     0.0\n",
       "id1=I  0.791431  0.004541  0.782532  0.800331     0.0\n",
       "id1=J  0.889866  0.004498  0.881050  0.898682     0.0"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time fr.ols(y=R.y, x=I+R.x1+R.x2+C.id1, absorb=C.id2, data=data_wide)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Poisson"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[  0] ℓ=-0.63803, g=0.02960, Δβ=0.60303, Δℓ=inf, μR=0.33633, μC=nan\n",
      "[ 31] ℓ=-0.62216, g=0.00010, Δβ=0.00000, Δℓ=0.00000, μR=0.33319, μC=nan\n",
      "CPU times: user 5.55 s, sys: 1.17 s, total: 6.72 s\n",
      "Wall time: 4.54 s\n"
     ]
    },
    {
     "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>p0</th>\n",
       "      <th>coeff</th>\n",
       "      <th>stderr</th>\n",
       "      <th>low95</th>\n",
       "      <th>high95</th>\n",
       "      <th>pvalue</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>I</th>\n",
       "      <td>0.097462</td>\n",
       "      <td>0.001023</td>\n",
       "      <td>0.095458</td>\n",
       "      <td>0.099467</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x1</th>\n",
       "      <td>0.301945</td>\n",
       "      <td>0.000849</td>\n",
       "      <td>0.300281</td>\n",
       "      <td>0.303609</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x2</th>\n",
       "      <td>0.600163</td>\n",
       "      <td>0.000848</td>\n",
       "      <td>0.598502</td>\n",
       "      <td>0.601824</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "p0     coeff    stderr     low95    high95  pvalue\n",
       "I   0.097462  0.001023  0.095458  0.099467     0.0\n",
       "x1  0.301945  0.000849  0.300281  0.303609     0.0\n",
       "x2  0.600163  0.000848  0.598502  0.601824     0.0"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time fr.poisson(y=R.p0, x=I+R.x1+R.x2, data=data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[  0] ℓ=2.39265, g=0.08783, Δβ=0.84455, Δℓ=inf, μR=0.41699, μC=0.37467\n",
      "[ 31] ℓ=2.55489, g=0.00012, Δβ=0.00008, Δℓ=0.00000, μR=0.33582, μC=0.49825\n",
      "CPU times: user 10.2 s, sys: 2.48 s, total: 12.7 s\n",
      "Wall time: 6.92 s\n"
     ]
    },
    {
     "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>p</th>\n",
       "      <th>coeff</th>\n",
       "      <th>stderr</th>\n",
       "      <th>low95</th>\n",
       "      <th>high95</th>\n",
       "      <th>pvalue</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>I</th>\n",
       "      <td>0.104049</td>\n",
       "      <td>0.006941</td>\n",
       "      <td>0.090444</td>\n",
       "      <td>0.117653</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x1</th>\n",
       "      <td>0.303011</td>\n",
       "      <td>0.000508</td>\n",
       "      <td>0.302016</td>\n",
       "      <td>0.304007</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x2</th>\n",
       "      <td>0.600389</td>\n",
       "      <td>0.000506</td>\n",
       "      <td>0.599397</td>\n",
       "      <td>0.601381</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=B</th>\n",
       "      <td>0.094683</td>\n",
       "      <td>0.002832</td>\n",
       "      <td>0.089132</td>\n",
       "      <td>0.100234</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=C</th>\n",
       "      <td>0.206832</td>\n",
       "      <td>0.002770</td>\n",
       "      <td>0.201402</td>\n",
       "      <td>0.212262</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=95</th>\n",
       "      <td>0.949977</td>\n",
       "      <td>0.007787</td>\n",
       "      <td>0.934714</td>\n",
       "      <td>0.965239</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=96</th>\n",
       "      <td>0.952913</td>\n",
       "      <td>0.007813</td>\n",
       "      <td>0.937599</td>\n",
       "      <td>0.968227</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=97</th>\n",
       "      <td>0.967282</td>\n",
       "      <td>0.007822</td>\n",
       "      <td>0.951952</td>\n",
       "      <td>0.982612</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=98</th>\n",
       "      <td>0.997022</td>\n",
       "      <td>0.007779</td>\n",
       "      <td>0.981776</td>\n",
       "      <td>1.012268</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=99</th>\n",
       "      <td>0.994602</td>\n",
       "      <td>0.007727</td>\n",
       "      <td>0.979458</td>\n",
       "      <td>1.009746</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>111 rows × 5 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "p          coeff    stderr     low95    high95  pvalue\n",
       "I       0.104049  0.006941  0.090444  0.117653     0.0\n",
       "x1      0.303011  0.000508  0.302016  0.304007     0.0\n",
       "x2      0.600389  0.000506  0.599397  0.601381     0.0\n",
       "id1=B   0.094683  0.002832  0.089132  0.100234     0.0\n",
       "id1=C   0.206832  0.002770  0.201402  0.212262     0.0\n",
       "...          ...       ...       ...       ...     ...\n",
       "id2=95  0.949977  0.007787  0.934714  0.965239     0.0\n",
       "id2=96  0.952913  0.007813  0.937599  0.968227     0.0\n",
       "id2=97  0.967282  0.007822  0.951952  0.982612     0.0\n",
       "id2=98  0.997022  0.007779  0.981776  1.012268     0.0\n",
       "id2=99  0.994602  0.007727  0.979458  1.009746     0.0\n",
       "\n",
       "[111 rows x 5 columns]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time fr.poisson(y=R.p, x=I+R.x1+R.x2+C.id1+C.id2, data=data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Logit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[  0] ℓ=-0.64539, g=0.00459, Δβ=0.59981, Δℓ=inf, μR=0.34471, μC=nan\n",
      "[ 31] ℓ=-0.64337, g=0.00007, Δβ=0.00000, Δℓ=0.00000, μR=0.33704, μC=nan\n",
      "CPU times: user 5.33 s, sys: 1.13 s, total: 6.47 s\n",
      "Wall time: 3.87 s\n"
     ]
    },
    {
     "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>b0</th>\n",
       "      <th>coeff</th>\n",
       "      <th>stderr</th>\n",
       "      <th>low95</th>\n",
       "      <th>high95</th>\n",
       "      <th>pvalue</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>I</th>\n",
       "      <td>0.101587</td>\n",
       "      <td>0.002104</td>\n",
       "      <td>0.097463</td>\n",
       "      <td>0.105711</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x1</th>\n",
       "      <td>0.304348</td>\n",
       "      <td>0.002146</td>\n",
       "      <td>0.300143</td>\n",
       "      <td>0.308554</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x2</th>\n",
       "      <td>0.605197</td>\n",
       "      <td>0.002269</td>\n",
       "      <td>0.600749</td>\n",
       "      <td>0.609644</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "b0     coeff    stderr     low95    high95  pvalue\n",
       "I   0.101587  0.002104  0.097463  0.105711     0.0\n",
       "x1  0.304348  0.002146  0.300143  0.308554     0.0\n",
       "x2  0.605197  0.002269  0.600749  0.609644     0.0"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time fr.logit(y=R.b0, x=I+R.x1+R.x2, data=data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[  0] ℓ=-0.55010, g=0.00373, Δβ=0.91517, Δℓ=inf, μR=0.40452, μC=0.39485\n",
      "[ 40] ℓ=-0.54383, g=0.00010, Δβ=0.00017, Δℓ=0.00000, μR=0.35234, μC=0.47504\n",
      "[ 73] ℓ=-0.54383, g=0.00010, Δβ=0.00010, Δℓ=0.00000, μR=0.35092, μC=0.47720\n",
      "CPU times: user 19.9 s, sys: 3.57 s, total: 23.5 s\n",
      "Wall time: 13.9 s\n"
     ]
    },
    {
     "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>b</th>\n",
       "      <th>coeff</th>\n",
       "      <th>stderr</th>\n",
       "      <th>low95</th>\n",
       "      <th>high95</th>\n",
       "      <th>pvalue</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>I</th>\n",
       "      <td>0.146217</td>\n",
       "      <td>0.022841</td>\n",
       "      <td>0.101449</td>\n",
       "      <td>0.190985</td>\n",
       "      <td>1.539278e-10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x1</th>\n",
       "      <td>0.312353</td>\n",
       "      <td>0.002382</td>\n",
       "      <td>0.307685</td>\n",
       "      <td>0.317021</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x2</th>\n",
       "      <td>0.594196</td>\n",
       "      <td>0.002508</td>\n",
       "      <td>0.589280</td>\n",
       "      <td>0.599111</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=B</th>\n",
       "      <td>0.081497</td>\n",
       "      <td>0.009835</td>\n",
       "      <td>0.062222</td>\n",
       "      <td>0.100773</td>\n",
       "      <td>2.220446e-16</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=C</th>\n",
       "      <td>0.202393</td>\n",
       "      <td>0.009910</td>\n",
       "      <td>0.182971</td>\n",
       "      <td>0.221816</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=95</th>\n",
       "      <td>0.838382</td>\n",
       "      <td>0.033492</td>\n",
       "      <td>0.772740</td>\n",
       "      <td>0.904024</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=96</th>\n",
       "      <td>0.955043</td>\n",
       "      <td>0.034054</td>\n",
       "      <td>0.888298</td>\n",
       "      <td>1.021788</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=97</th>\n",
       "      <td>0.916935</td>\n",
       "      <td>0.033867</td>\n",
       "      <td>0.850556</td>\n",
       "      <td>0.983314</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=98</th>\n",
       "      <td>0.919533</td>\n",
       "      <td>0.034025</td>\n",
       "      <td>0.852845</td>\n",
       "      <td>0.986220</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=99</th>\n",
       "      <td>0.936677</td>\n",
       "      <td>0.033864</td>\n",
       "      <td>0.870305</td>\n",
       "      <td>1.003048</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>111 rows × 5 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "b          coeff    stderr     low95    high95        pvalue\n",
       "I       0.146217  0.022841  0.101449  0.190985  1.539278e-10\n",
       "x1      0.312353  0.002382  0.307685  0.317021  0.000000e+00\n",
       "x2      0.594196  0.002508  0.589280  0.599111  0.000000e+00\n",
       "id1=B   0.081497  0.009835  0.062222  0.100773  2.220446e-16\n",
       "id1=C   0.202393  0.009910  0.182971  0.221816  0.000000e+00\n",
       "...          ...       ...       ...       ...           ...\n",
       "id2=95  0.838382  0.033492  0.772740  0.904024  0.000000e+00\n",
       "id2=96  0.955043  0.034054  0.888298  1.021788  0.000000e+00\n",
       "id2=97  0.916935  0.033867  0.850556  0.983314  0.000000e+00\n",
       "id2=98  0.919533  0.034025  0.852845  0.986220  0.000000e+00\n",
       "id2=99  0.936677  0.033864  0.870305  1.003048  0.000000e+00\n",
       "\n",
       "[111 rows x 5 columns]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time fr.logit(y=R.b, x=I+R.x1+R.x2+C.id1+C.id2, data=data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ultra Wide"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "id1        10\n",
      "id2    100000\n",
      "dtype: int64\n"
     ]
    }
   ],
   "source": [
    "N = 5_000_000\n",
    "df = pd.DataFrame({ \n",
    "    'x1': np.random.rand(N), \n",
    "    'x2': np.random.rand(N), \n",
    "    'id1': np.ceil(10*np.arange(N)/N+1e-7).astype(int),\n",
    "    'id2': np.random.randint(1, 100_001, size=N)\n",
    "})\n",
    "df['y'] = 1 + 2*df['x1'] + 3*df['x2'] + np.log10(df['id1']) + np.log10(df['id2']) + np.random.randn(N)\n",
    "print(df[['id1', 'id2']].nunique())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.02 s, sys: 776 ms, total: 2.8 s\n",
      "Wall time: 2.03 s\n"
     ]
    },
    {
     "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>y</th>\n",
       "      <th>coeff</th>\n",
       "      <th>stderr</th>\n",
       "      <th>low95</th>\n",
       "      <th>high95</th>\n",
       "      <th>pvalue</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>I</th>\n",
       "      <td>0.943007</td>\n",
       "      <td>0.133720</td>\n",
       "      <td>0.680921</td>\n",
       "      <td>1.205093</td>\n",
       "      <td>1.762146e-12</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x1</th>\n",
       "      <td>2.000551</td>\n",
       "      <td>0.001566</td>\n",
       "      <td>1.997482</td>\n",
       "      <td>2.003620</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x2</th>\n",
       "      <td>2.997936</td>\n",
       "      <td>0.001566</td>\n",
       "      <td>2.994868</td>\n",
       "      <td>3.001004</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=2</th>\n",
       "      <td>0.300425</td>\n",
       "      <td>0.002022</td>\n",
       "      <td>0.296463</td>\n",
       "      <td>0.304387</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1=3</th>\n",
       "      <td>0.475794</td>\n",
       "      <td>0.002022</td>\n",
       "      <td>0.471831</td>\n",
       "      <td>0.479756</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=99996</th>\n",
       "      <td>4.902568</td>\n",
       "      <td>0.193671</td>\n",
       "      <td>4.522980</td>\n",
       "      <td>5.282155</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=99997</th>\n",
       "      <td>5.230915</td>\n",
       "      <td>0.194682</td>\n",
       "      <td>4.849346</td>\n",
       "      <td>5.612484</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=99998</th>\n",
       "      <td>5.331394</td>\n",
       "      <td>0.199103</td>\n",
       "      <td>4.941158</td>\n",
       "      <td>5.721629</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=99999</th>\n",
       "      <td>4.652056</td>\n",
       "      <td>0.200314</td>\n",
       "      <td>4.259447</td>\n",
       "      <td>5.044665</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id2=100000</th>\n",
       "      <td>5.234107</td>\n",
       "      <td>0.197937</td>\n",
       "      <td>4.846158</td>\n",
       "      <td>5.622056</td>\n",
       "      <td>0.000000e+00</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>100011 rows × 5 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "y              coeff    stderr     low95    high95        pvalue\n",
       "I           0.943007  0.133720  0.680921  1.205093  1.762146e-12\n",
       "x1          2.000551  0.001566  1.997482  2.003620  0.000000e+00\n",
       "x2          2.997936  0.001566  2.994868  3.001004  0.000000e+00\n",
       "id1=2       0.300425  0.002022  0.296463  0.304387  0.000000e+00\n",
       "id1=3       0.475794  0.002022  0.471831  0.479756  0.000000e+00\n",
       "...              ...       ...       ...       ...           ...\n",
       "id2=99996   4.902568  0.193671  4.522980  5.282155  0.000000e+00\n",
       "id2=99997   5.230915  0.194682  4.849346  5.612484  0.000000e+00\n",
       "id2=99998   5.331394  0.199103  4.941158  5.721629  0.000000e+00\n",
       "id2=99999   4.652056  0.200314  4.259447  5.044665  0.000000e+00\n",
       "id2=100000  5.234107  0.197937  4.846158  5.622056  0.000000e+00\n",
       "\n",
       "[100011 rows x 5 columns]"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time fr.ols(y=R.y, x=I+R.x1+R.x2+C.id1, hdfe=C.id2, data=df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}