blog.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# EZKL Workflow\n",
    "\n",
    "To prove the inference of a trained model using EZKL, we need to follow the steps below. As an example to illustrate the process, let's consider that we have just trained a simple perceptron model using PyTorch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install --upgrade --force-reinstall -r requirements.txt\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader\n",
    "import torchvision\n",
    "\n",
    "# MNIST dataset\n",
    "train, test = (torchvision.datasets.MNIST(\n",
    "    './data', \n",
    "    train=is_train,\n",
    "    transform=torchvision.transforms.Compose([\n",
    "        torchvision.transforms.ToTensor(),\n",
    "        torchvision.transforms.Normalize((0.1307,), (0.3081,)),\n",
    "        torchvision.transforms.Lambda(lambda x: x.view(-1))\n",
    "    ]),\n",
    "    download=True\n",
    ") for is_train in [True, False])\n",
    "\n",
    "input_size, output_size = 28 * 28, 10\n",
    "\n",
    "# Define the model\n",
    "perceptron = nn.Sequential(\n",
    "    nn.Linear(input_size, output_size),\n",
    ")\n",
    "\n",
    "# Create a dataset and data loader\n",
    "train_loader, test_loader = (DataLoader(\n",
    "    dataset, \n",
    "    batch_size=32, \n",
    "    shuffle=True\n",
    ") for dataset in [train, test])\n",
    "\n",
    "# Define the loss function and optimizer\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(perceptron.parameters(), lr=0.01)\n",
    "\n",
    "# Train the model\n",
    "perceptron.train()\n",
    "for data, label in train_loader:\n",
    "    output = perceptron(data)\n",
    "    loss = criterion(output, label)\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 85.69%\n"
     ]
    }
   ],
   "source": [
    "# Test the model\n",
    "perceptron.eval()\n",
    "with torch.no_grad():\n",
    "    correct, total = 0, 0\n",
    "    for data, label in test_loader:\n",
    "        output = perceptron(data)\n",
    "        _, predicted = torch.max(output.data, 1)\n",
    "        total += label.size(0)\n",
    "        correct += (predicted == label).sum().item()\n",
    "    print(f'Accuracy: {100 * correct / total:.2f}%')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. **Model Conversion**\n",
    "\n",
    "Convert the trained model to the ONNX format. In this case, PyTorch provides the function torch.onnx.export. Other frameworks also have similar functions or external tools to convert models to ONNX (e.g., TensorFlow's tf2onnx). Nevertheless, Sklearn models are slighly more complicated to convert to suitable ONNX format, so we must first convert the model to a PyTorch using hummingbird.ml and then convert it to ONNX. We won't cover this process in this article but you can find more information in one of EZKL's notebooks.\n",
    "\n",
    "Before converting our model to ONNX format, we need to tell the converter the input shape of the model. This can be done by passing a dummy input tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "============= Diagnostic Run torch.onnx.export version 2.0.1+cu117 =============\n",
      "verbose: False, log level: Level.ERROR\n",
      "======================= 0 NONE 0 NOTE 0 WARNING 0 ERROR ========================\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Choose any valid input tensor (1st input of the test dataset)\n",
    "input_sample = next(iter(test_loader))[0][0].unsqueeze(0)\n",
    "\n",
    "torch.onnx.export(\n",
    "    perceptron,\n",
    "    input_sample,\n",
    "    \"data/models/perceptron.onnx\",\n",
    "    opset_version=10,\n",
    "    export_params=True,                # Store the trained parameter weights inside the model file\n",
    "    do_constant_folding=True,          # Optimize constant values in the model graph\n",
    "    input_names = ['input'],             # Input and output labels to appear in the ONNX graph \n",
    "    output_names = ['output'],\n",
    "    dynamic_axes={\n",
    "        'input' : {0 : 'batch_size'},    # Variable length axes\n",
    "        'output' : {0 : 'batch_size'}\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. **Setup**\n",
    "\n",
    "EZKL has several setup functions in their exposed API, namely gen_settings, calibrate_settings, compile_circuit, get_srs, setup, and gen_witness, we've group them together in this bullet point to describe the high level setup process. The signature of each function should be self-explanatory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using 6 columns for non-linearity table.\n",
      "Using 11 columns for non-linearity table.\n",
      "calibration failed extended k is too large to accommodate the quotient polynomial with logrows 6\n",
      "Using 11 columns for non-linearity table.\n",
      "calibration failed extended k is too large to accommodate the quotient polynomial with logrows 6\n",
      "calibration failed max lookup input (-478801996, 243305849) is too large\n",
      "calibration failed max lookup input (-478877341, 243282263) is too large\n",
      "calibration failed max lookup input (-957655723, 486630371) is too large\n",
      "calibration failed max lookup input (-1915080799, 973077036) is too large\n",
      "\n",
      "\n",
      " <------------- Numerical Fidelity Report (input_scale: 13, param_scale: 13, scale_input_multiplier: 10) ------------->\n",
      "\n",
      "+---------------+--------------+-------------+---------------+----------------+------------------+---------------+---------------+--------------------+--------------------+------------------------+\n",
      "| mean_error    | median_error | max_error   | min_error     | mean_abs_error | median_abs_error | max_abs_error | min_abs_error | mean_squared_error | mean_percent_error | mean_abs_percent_error |\n",
      "+---------------+--------------+-------------+---------------+----------------+------------------+---------------+---------------+--------------------+--------------------+------------------------+\n",
      "| 0.00013546944 | 0.001152277  | 0.001152277 | -0.0013656616 | 0.0006680012   | 0.001152277      | 0.0013656616  | 0.00021839142 | 0.00000056112486   | 0.00000057066757   | 0.00012321016          |\n",
      "+---------------+--------------+-------------+---------------+----------------+------------------+---------------+---------------+--------------------+--------------------+------------------------+\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import ezkl\n",
    "import json\n",
    "\n",
    "def create_file(filename: str) -> str:\n",
    "    # If the file already exists, clear it\n",
    "    if os.path.exists(filename):\n",
    "        os.remove(filename)\n",
    "\n",
    "    open(filename, 'w').close()\n",
    "    return filename\n",
    "\n",
    "# We have to create empty files manually before running the setup\n",
    "INPUT = create_file(\"data/ezkl/input_data.json\")\n",
    "SETTINGS = create_file(\"data/ezkl/settings.json\")\n",
    "CALIBRATION = create_file(\"data/ezkl/calibration.json\")\n",
    "WITNESS = create_file(\"data/ezkl/witness.json\")\n",
    "COMPILED_MODEL = create_file(\"data/ezkl/model.compiled\")\n",
    "VK = create_file(\"data/ezkl/vk.json\")\n",
    "PK = create_file(\"data/ezkl/pk.json\")\n",
    "PROOF = create_file(\"data/ezkl/proof.json\")\n",
    "\n",
    "def setup(onnx_file: str, model: nn.Module, input_sample: torch.Tensor):\n",
    "\n",
    "    # Create empty files for each of the required inputs\n",
    "    for filename in [INPUT, SETTINGS, CALIBRATION, WITNESS, COMPILED_MODEL, VK, PK, PROOF]:\n",
    "        create_file(filename)\n",
    "\n",
    "    # Save the input data to a file in the expected format\n",
    "    input_data = {\n",
    "        'input_shapes': list(input_sample.shape),\n",
    "        'input_data': input_sample.detach().numpy().tolist(),\n",
    "        \"output_data\": model(input_sample).detach().numpy().tolist()\n",
    "    }\n",
    "\n",
    "    json.dump(\n",
    "        input_data,\n",
    "        open(INPUT, 'w')\n",
    "    )\n",
    "\n",
    "    # Run each setup function and verify that it succeeded\n",
    "    assert ezkl.gen_settings(\n",
    "        onnx_file,\n",
    "        SETTINGS\n",
    "    )\n",
    "\n",
    "    calibration_data = {\n",
    "        'input_data': torch.randn(20, input_sample.shape[1]).numpy().tolist()\n",
    "    }\n",
    "\n",
    "    json.dump(\n",
    "        calibration_data,\n",
    "        open(CALIBRATION, 'w')\n",
    "    )\n",
    "\n",
    "    assert ezkl.calibrate_settings(\n",
    "        INPUT,\n",
    "        onnx_file,\n",
    "        SETTINGS, \n",
    "        \"resources\"\n",
    "    )\n",
    "\n",
    "    assert ezkl.compile_circuit(\n",
    "        onnx_file,\n",
    "        COMPILED_MODEL,\n",
    "        SETTINGS\n",
    "    )\n",
    "\n",
    "    assert ezkl.get_srs(\n",
    "        SETTINGS\n",
    "    )\n",
    "\n",
    "    ezkl.gen_witness(\n",
    "        \"data/ezkl/input_data.json\",\n",
    "        COMPILED_MODEL,\n",
    "        WITNESS\n",
    "    )\n",
    "\n",
    "    assert ezkl.setup(\n",
    "        COMPILED_MODEL,\n",
    "        VK,\n",
    "        PK\n",
    "    )\n",
    "\n",
    "setup(\"data/models/perceptron.onnx\", perceptron, input_sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. **Proof Generation**: Generate the proof using the `gen_proof` function. This function takes the arithmetization of the model, the witness, the public key, and the proof file as inputs and writes the proof to the specified file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'instances': [['2ae4a17d93f5e1439170b97948e833285d588181b64550b829a031e1724e6430',\n",
      "                '2108b8e093f5e1439170b97948e833285d588181b64550b829a031e1724e6430',\n",
      "                'ffebd7c293f5e1439170b97948e833285d588181b64550b829a031e1724e6430',\n",
      "                '015c5e1600000000000000000000000000000000000000000000000000000000',\n",
      "                'a50ed01f00000000000000000000000000000000000000000000000000000000',\n",
      "                '0b92b0e693f5e1439170b97948e833285d588181b64550b829a031e1724e6430',\n",
      "                '79034d9193f5e1439170b97948e833285d588181b64550b829a031e1724e6430',\n",
      "                'e9f13a2300000000000000000000000000000000000000000000000000000000',\n",
      "                '44202f0d00000000000000000000000000000000000000000000000000000000',\n",
      "                'd0cfdb3900000000000000000000000000000000000000000000000000000000']],\n",
      " 'proof': '0x2427cae5f23ce51de1000fdbcd3b4c931d04e172947d227ce649344cd21eaabd1399dfd586f1b1f8d6f179b28bc2b6a3eefa3e5d0474d805cd55ed36c706f9ee034824fe1e70cbabbd6ffd829b012571eba08781aab4c50b420669e71a875ff5210c00730ff7876bc5e32d946b62448cdafbb7544aa7fcf9450f1ba66322bff52f96ce53cccc1f8f156858407506111bc524f09127466f5cd6264369dbaf27ed0ecb76474a5706dfcb55f2b2c9ae9f9a9c17cdbf5002df87c6dbd516c5feedd607e7cde06a431dcc8fe52e38f92e3f22f76eedb941d49f40bb961863e93cb94e1051a71850a4be374434e2b50d5f266d4ddbe7e5409d563ad23ca8d89a4371920477bd7c55669f0456fbb531fc5fe1e7bda7d1da049eaee0c234bb91a9923c3e055edde39e35a5e2137041236ed7146944a6a02d126067a96b565908bc94268728b183640ccaec3bde1092bdc188b44352164c6c127c01aadf98cf7bd4cf1ae62ae7db62372c38df7df0f2927bdd0b457545afc29bec37727b604a841de2285a153003eb430841c06e184f9d15277156684d712915e1e2cc5c40b367256e23da1a905bf6fbf204bf51c0d9b2cf26626ce0395d04c31e70d54eed9802afa355fd1039f12fc14b9c163aab33fd9e6fc33295b0bcd8ac9286a107d8b83d70b30dc70992d8810ade0e13454c07b372323dedaaf4f2949b2545b18030267f428202c804252ff1d92ab58cff7ee601fa906c332746096a53fd34b7f93e7304de1d45082b6bf1bf0867b5437a2139c09560f28cb5525881f39a9c0cddf7bdd69e1a091b0d8190957556fd580db71bc095df3a64fded7dae927cefda33cf971dc3deb9e32a66e5f09f027edf3d6aaa17b30b13d6e873933ab9d963736c4f5f56902fb7910e13abec6a145988b357876938b3293ecb6075dbfe2fabdcd38981c0e8aa2e170ee723f2493fabd5f4c43485607d8406ce1ce366b3ffd33670d67a58a23228741c6dc5a59c53570e8272c7b1905df0efacee869225bcc884bc4ad18f8dcb736607803951ffc618540aa46dfd1a69fc656403f70edb770e7bc8cf84138631d44914f24af468a522a5bf62a80098d9b0765a262e9032ee6b2f4c1e7b22638e1b721bfeb50297a669841a85218f77985f1f1d09846170024fa56fcdc68efc3574a325a0733f4cb0a847bfa71a08575c38d3a581c9dc579b722ec9263435a077a5e423f812420e5626c1a74eebaa5ed5816505597c692013aba03b582f0f154526160071a78ebeec57cd9a21d48538062e4eb718d37ed9e6831a2ecf6c42e859bc4d1f4ae069ad173daeb870d8f11e9682d10569009b44ec595b05685258add0495e095717dfe231b9dcd0da37053721d747f097e18387e879860e765e355cef61d00ff953d07c2fe5dda84b62d1023ce35ab60f1e20da4df5661d74b9168ee82c0602c460b890607b5cf2085a0a807d02d03cd6e8c0b4b92ebfe2b33de79c2b361b28983b43db506b081ad0fcb2bfb20080800313083c495a72541357f946993a43296cf8131fd47d4be5ec3ca7c0c655c0545a6abc50545ccbe563da89b265d0e021e6c3e1ad149db17e19c1328ce1e212d3d36ed3dd095e8fea2974e58205e9ea000000000000000000000000000000000000000000000000000000000000000021e6c3e1ad149db17e19c1328ce1e212d3d36ed3dd095e8fea2974e58205e9ea2cb607775c6b8da76a0016d423e44e07379c5a89476c9a11b3c538af6f4eb9202c77d71d200a0925150831e74c1652ee834a797ce20d97e1ae80488c440ae1a70000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000002988c40f5157ccbffe449c6b7e864e0919d967d0b78641d50ee26c1fc37e06a0bb893de2cb55c3f8533ae224d8fbcda01675bcc182533da36845443df7b6bb82202cd2954de6556bba86e000037deaaafb2dfd3a27c8985ebeec991f7db3a6416b96cd48c3425e0512ffe63b88cb8829a527ec10e2d2f42b6cc548d24132eef2aa42315697295217fdc12eaead42a0b180e98480572f5e92efa6d49b404994f0c3297ea91fa489ba99545fa247037d6e35e00f0d8fe2faaab72f6732ccfc55e0a4e8c91bff3ba1f06f261a1c75b9b9d00aaa22f3d926bf3bc32eaff7b6c017806d3ae69c405c3a3f19dbbcc38ec03cdca49d7d2141eb0f9e6e281cd60cc13cf0cc3c4bb30d6c4bd07b15a5d5e202c1b3ba86a65131a0105cb8fa8cf354b6d970f54022a53fe759fcecdd1d255401aeaa052bfdd635bc1055d49b36a865294ac04730f884046c3512329920a70d27669c1b2b8025d490d711fc0b22d5ebc3f2a2f067756614d723a76b89675621c39c2ab070c28c52567a0832502ef335468db247a66779c928be83a643dc763664d9672d123a2830ee6a2c4db52e4119c950509c30f1eceb843dccc320134ed365406915584b17c94c15222f8742abbae7bcf08d197054394b1a18c22831cf8ffdd94836753e34f55f4281ad578dfba837e4f25f25cd71ab8ef066a66e59d22fc3fd4d2c2bb09a32c52dcea21d9fa02e9735a26e52c4665ff7827319f8a95c9a0af65bdb275adcda8b02c220dd5b62e4411c01315f65eb4283de32a6b971e120e9c551c3733300b83e992d09182e152c309cb1f3f6e1fbe4dbe9c7ea3c044175618ca5a69ff33ea64e0746641493ea2b6f20624d51effe55f51431a61987f24cebee8b476e21d110f13f36b7d063f85fbebae2a5470ca79177cde5420a68d0e43b4a54dec98c990a05ccae9068b3f3589019e0a4c1ecc729fd57275b123f60237f6c198c00841eadfe9d2a23948093ec5a7d60c228cb84728f544b126d76b5ffd056df2e3b7e6eed64a5e48d5514848650512',\n",
      " 'transcript_type': 'EVM'}\n"
     ]
    }
   ],
   "source": [
    "import pprint\n",
    "\n",
    "proof = ezkl.prove(\n",
    "    WITNESS,\n",
    "    COMPILED_MODEL,\n",
    "    PK,\n",
    "    PROOF,\n",
    "    \"single\",\n",
    ")\n",
    "\n",
    "pprint.pprint(proof)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. **Proof Verification**: ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert ezkl.verify(\n",
    "    PROOF,\n",
    "    SETTINGS,\n",
    "    VK\n",
    ") == True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Token Trend Forecasting\n",
    "\n",
    "We adapted code from one of [GIZA's examples](https://github.com/gizatechxyz/Giza-Hub/tree/token_trend_action/awesome-giza-actions/trend_token_prediction) the idea is to train multiple models with different accuracies and then compare the costs of proving the inference of each model. We first \n",
    "explain the feature extraction process in detail, which is not explained in the original example."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data\n",
    "\n",
    "We will use the [Giza's dataset hub](https://github.com/gizatechxyz/datasets), which contains a collection of datasets that are relevant for blockchain applications. These datasets are publicly available and can be loaded using the `DatasetsLoader` class from the `giza_datasets` package, like so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset tokens-daily-prices-mcap-volume not found in cache. Downloading from GCS.\n",
      "Dataset read from cache.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "polars.config.Config"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from giza.datasets import DatasetsLoader\n",
    "import polars as pl\n",
    "\n",
    "# Load the desired dataset\n",
    "DatasetsLoader().load(\"tokens-daily-prices-mcap-volume\")\n",
    "\n",
    "# For pretty printing\n",
    "pl.Config.set_tbl_hide_column_data_types(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Token Daily Price Data\n",
    "Contains daily price data (price, market capitalization, and volume) for a set of tokens (e.g., WBTC, WETH, etc.)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First few rows of the dataset:\n",
      "Dataset read from cache.\n",
      "Loading dataset tokens-daily-prices-mcap-volume from cache.\n",
      "shape: (3, 5)\n",
      "┌────────────┬─────────────┬────────────┬──────────────────┬───────┐\n",
      "│ date       ┆ price       ┆ market_cap ┆ volumes_last_24h ┆ token │\n",
      "╞════════════╪═════════════╪════════════╪══════════════════╪═══════╡\n",
      "│ 2019-02-01 ┆ 3438.360403 ┆ 0.0        ┆ 20589.040403     ┆ WBTC  │\n",
      "│ 2019-02-02 ┆ 3472.243307 ┆ 0.0        ┆ 12576.723906     ┆ WBTC  │\n",
      "│ 2019-02-03 ┆ 3461.058341 ┆ 0.0        ┆ 1852.526033      ┆ WBTC  │\n",
      "└────────────┴─────────────┴────────────┴──────────────────┴───────┘\n"
     ]
    }
   ],
   "source": [
    "print(\"First few rows of the dataset:\")\n",
    "print(DatasetsLoader().load('tokens-daily-prices-mcap-volume').head(n = 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Top APY per protocol\n",
    "Contains the top Annual Percentage Yield (APY) for each protocol in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First few rows of the dataset:\n",
      "Dataset top-pools-apy-per-protocol not found in cache. Downloading from GCS.\n",
      "Dataset read from cache.\n",
      "shape: (3, 6)\n",
      "┌────────────┬──────────┬─────┬─────────┬──────────────────┬──────────┐\n",
      "│ date       ┆ tvlUsd   ┆ apy ┆ project ┆ underlying_token ┆ chain    │\n",
      "╞════════════╪══════════╪═════╪═════════╪══════════════════╪══════════╡\n",
      "│ 2022-02-28 ┆ 12808    ┆ 0.0 ┆ aave-v2 ┆ STETH            ┆ Ethereum │\n",
      "│ 2022-03-01 ┆ 46045250 ┆ 0.0 ┆ aave-v2 ┆ STETH            ┆ Ethereum │\n",
      "│ 2022-03-02 ┆ 90080754 ┆ 0.0 ┆ aave-v2 ┆ STETH            ┆ Ethereum │\n",
      "└────────────┴──────────┴─────┴─────────┴──────────────────┴──────────┘\n"
     ]
    }
   ],
   "source": [
    "print(\"First few rows of the dataset:\")\n",
    "print(DatasetsLoader().load('top-pools-apy-per-protocol').head(n = 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### TVL per project tokens\n",
    "Contains the Total Value Locked (TVL) for each project in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First few rows of the dataset:\n",
      "Dataset tvl-per-project-tokens not found in cache. Downloading from GCS.\n",
      "Dataset read from cache.\n",
      "shape: (3, 47)\n",
      "┌───────┬──────┬────────┬──────┬───┬──────┬──────┬────────────┬─────────┐\n",
      "│ 1INCH ┆ AAVE ┆ AAVE.E ┆ AMPL ┆ … ┆ YFI  ┆ ZRX  ┆ date       ┆ project │\n",
      "╞═══════╪══════╪════════╪══════╪═══╪══════╪══════╪════════════╪═════════╡\n",
      "│ null  ┆ null ┆ null   ┆ null ┆ … ┆ null ┆ null ┆ 2020-11-29 ┆ aave-v2 │\n",
      "│ null  ┆ null ┆ null   ┆ null ┆ … ┆ null ┆ null ┆ 2020-11-30 ┆ aave-v2 │\n",
      "│ null  ┆ null ┆ null   ┆ null ┆ … ┆ null ┆ null ┆ 2020-12-01 ┆ aave-v2 │\n",
      "└───────┴──────┴────────┴──────┴───┴──────┴──────┴────────────┴─────────┘\n"
     ]
    }
   ],
   "source": [
    "print(\"First few rows of the dataset:\")\n",
    "print(DatasetsLoader().load('tvl-per-project-tokens').head(n = 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset read from cache.\n",
      "Loading dataset tokens-daily-prices-mcap-volume from cache.\n",
      "First few rows of the dataset:\n",
      "shape: (3, 19)\n",
      "┌────────────┬─────────┬────────────┬──────────────────┬───┬───────────────┬─────┬───────┬──────┐\n",
      "│ date       ┆ price   ┆ market_cap ┆ volumes_last_24h ┆ … ┆ trend_30_days ┆ day ┆ month ┆ year │\n",
      "╞════════════╪═════════╪════════════╪══════════════════╪═══╪═══════════════╪═════╪═══════╪══════╡\n",
      "│ 2018-02-14 ┆ 839.535 ┆ 0.0        ┆ 54776.5          ┆ … ┆ null          ┆ 3   ┆ 2     ┆ 2018 │\n",
      "│ 2018-02-15 ┆ 947.358 ┆ 0.0        ┆ 111096.0         ┆ … ┆ null          ┆ 4   ┆ 2     ┆ 2018 │\n",
      "│ 2018-02-16 ┆ 886.961 ┆ 0.0        ┆ 57731.7          ┆ … ┆ null          ┆ 5   ┆ 2     ┆ 2018 │\n",
      "└────────────┴─────────┴────────────┴──────────────────┴───┴───────────────┴─────┴───────┴──────┘\n"
     ]
    }
   ],
   "source": [
    "import itertools\n",
    "\n",
    "TOKEN = \"WETH\"\n",
    "LAG = 1\n",
    "DAYS = [1, 3, 7, 15, 30]\n",
    "START_DATE = pl.datetime(2022, 6, 1)\n",
    "\n",
    "token_data = DatasetsLoader().load('tokens-daily-prices-mcap-volume')\n",
    "\n",
    "# Filter the dataset to only include the target token\n",
    "# Add a column with the labels for the target token\n",
    "# Add columns with the price difference over the specified DAYS\n",
    "# Expand the date column into day_of_week, month_of_year, and year\n",
    "target_token_price_trend = token_data \\\n",
    "    .filter(pl.col(\"token\") == TOKEN) \\\n",
    "    .with_columns(\n",
    "        ((pl.col(\"price\").shift(-LAG) - pl.col(\"price\")) > 0).cast(pl.Int8).alias(\"label\")\n",
    "    ) \\\n",
    "    .with_columns(\n",
    "        pl.col(\"price\").diff(n = days).alias(f\"price_diff_{days}_days\")\n",
    "        for days in DAYS\n",
    "    ) \\\n",
    "    .with_columns(\n",
    "        (pl.col(\"price\") - pl.col(\"price\").shift(days) > 0).cast(pl.Int8).alias(f\"trend_{days}_days\")\n",
    "        for days in DAYS\n",
    "    ) \\\n",
    "    .with_columns([\n",
    "        pl.col(\"date\").dt.weekday().alias(\"day\"),\n",
    "        pl.col(\"date\").dt.month().alias(\"month\"),\n",
    "        pl.col(\"date\").dt.year().alias(\"year\")\n",
    "    ])\n",
    "\n",
    "print(\"First few rows of the dataset:\")\n",
    "print(target_token_price_trend.head(n = 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset read from cache.\n",
      "Loading dataset tokens-daily-prices-mcap-volume from cache.\n"
     ]
    }
   ],
   "source": [
    "token_data = DatasetsLoader().load('tokens-daily-prices-mcap-volume')\n",
    "correlations = {}\n",
    "\n",
    "# List all tokens in the dataset\n",
    "tokens = token_data.get_column(\"token\").unique().to_list()\n",
    "\n",
    "# Calculate the correlation between the target token and all other tokens\n",
    "for token_1, token_2 in itertools.permutations(tokens, r=2):\n",
    "    \n",
    "    # Filter the dataset and get the price and date columns\n",
    "    token_1_data = token_data.filter(pl.col(\"token\") == token_1) \\\n",
    "        .select([\"date\", \"price\"])\n",
    "    token_2_data = token_data.filter(pl.col(\"token\") == token_2) \\\n",
    "        .select([\"date\", \"price\"])\n",
    "    \n",
    "    # Join the datasets on the date column\n",
    "    joined_data = token_1_data.join(token_2_data, on=\"date\", suffix=\"_compare\")\n",
    "\n",
    "    # Nested dictionary to store the correlation between the two tokens\n",
    "    correlations[token_1] = correlations.get(token_1, {}) \n",
    "    correlations[token_1][token_2] = correlations[token_1] \\\n",
    "        .get(token_2, {\n",
    "            day: joined_data \\\n",
    "                    .with_columns(pl.col(\"price_compare\").shift(day)) \\\n",
    "                    .select(pl.corr(\"price\", \"price_compare\").alias(\"correlation\")) \\\n",
    "                    .get_column(\"correlation\")[0]\n",
    "            for day in DAYS\n",
    "        })\n",
    "\n",
    "\n",
    "#pprint.pprint(correlations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First few rows of the dataset:\n",
      "shape: (3, 69)\n",
      "┌────────────┬─────────┬───────────┬───────────┬───┬───────────┬───────────┬───────────┬───────────┐\n",
      "│ date       ┆ price   ┆ market_ca ┆ volumes_l ┆ … ┆ price_tre ┆ price_tre ┆ price_tre ┆ price_tre │\n",
      "│            ┆         ┆ p         ┆ ast_24h   ┆   ┆ nd_GNO_3  ┆ nd_GNO_7  ┆ nd_GNO_15 ┆ nd_GNO_30 │\n",
      "╞════════════╪═════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪═══════════╡\n",
      "│ 2018-02-14 ┆ 839.535 ┆ 0.0       ┆ 54776.5   ┆ … ┆ 0         ┆ 1         ┆ 0         ┆ 0         │\n",
      "│ 2018-02-15 ┆ 947.358 ┆ 0.0       ┆ 111096.0  ┆ … ┆ 1         ┆ 1         ┆ 0         ┆ 0         │\n",
      "│ 2018-02-16 ┆ 886.961 ┆ 0.0       ┆ 57731.7   ┆ … ┆ 1         ┆ 1         ┆ 0         ┆ 0         │\n",
      "└────────────┴─────────┴───────────┴───────────┴───┴───────────┴───────────┴───────────┴───────────┘\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_24261/2132130301.py:39: DeprecationWarning: The default coalesce behavior of left join will change to `False` in the next breaking release. Pass `coalesce=True` to keep the current behavior and silence this warning.\n",
      "  price_dataset = price_dataset.join(token_prices, on=\"date\", how=\"left\")\n"
     ]
    }
   ],
   "source": [
    "K = 5\n",
    "\n",
    "# Dataframe to store the final results\n",
    "price_dataset = target_token_price_trend\n",
    "\n",
    "# Retrive the relevant data from the nested dictionary\n",
    "target_token_correlations = correlations[TOKEN]\n",
    "\n",
    "# Get the top K correlated tokens for each lag\n",
    "top_k_correlated_tokens_by_lag = {\n",
    "    lag: sorted(target_token_correlations.items(), key=lambda x: x[1][lag], reverse=True)[:K]\n",
    "    for lag in DAYS\n",
    "}\n",
    "\n",
    "top_k_correlated_tokens_15_days = top_k_correlated_tokens_by_lag[15]\n",
    "\n",
    "for token, _ in top_k_correlated_tokens_15_days:\n",
    "\n",
    "    # Column names for the price differences\n",
    "    price_diff_columns = [f\"price_diff_{token}_{days}\" for days in DAYS]\n",
    "    price_trend_columns = [f\"price_trend_{token}_{days}\" for days in DAYS]\n",
    "\n",
    "    # Filter the dataset to only include the correlated token\n",
    "    token_prices = token_data.filter(pl.col(\"token\") == token)\n",
    "\n",
    "    # Add columns with the price differences for each day\n",
    "    token_prices = token_prices \\\n",
    "        .with_columns(\n",
    "            pl.col(\"price\").diff(n = days).alias(tag)\n",
    "            for days, tag in zip(DAYS, price_diff_columns)\n",
    "        ) \\\n",
    "        .with_columns([\n",
    "            (pl.col(\"price\") - pl.col(\"price\").shift(days) > 0).cast(pl.Int8).alias(tag)\n",
    "            for days, tag in zip(DAYS, price_trend_columns)\n",
    "        ]) \\\n",
    "        .select([\"date\"] + price_diff_columns + price_trend_columns)\n",
    "\n",
    "    # Join the dataset with the target token dataset\n",
    "    price_dataset = price_dataset.join(token_prices, on=\"date\", how=\"left\")\n",
    "\n",
    "print(\"First few rows of the dataset:\")\n",
    "print(price_dataset.head(n = 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset read from cache.\n",
      "Loading dataset top-pools-apy-per-protocol from cache.\n",
      "First few rows of the dataset:\n",
      "shape: (3, 91)\n",
      "┌───────────┬───────────┬───────────┬───────────┬───┬───────────┬───────────┬───────────┬──────────┐\n",
      "│ date      ┆ tvlUsd_pr ┆ tvlUsd_pr ┆ tvlUsd_pr ┆ … ┆ apy_proje ┆ apy_proje ┆ apy_proje ┆ apy_proj │\n",
      "│           ┆ oject_aav ┆ oject_aav ┆ oject_aav ┆   ┆ ct_uniswa ┆ ct_uniswa ┆ ct_uniswa ┆ ect_year │\n",
      "│           ┆ e-v2_Ethe ┆ e-v2_Poly ┆ e-v2_Aval ┆   ┆ p-v3_Arbi ┆ p-v3_Arbi ┆ p-v3_Ethe ┆ n-financ │\n",
      "│           ┆ reu…      ┆ gon…      ┆ anc…      ┆   ┆ tru…      ┆ tru…      ┆ reu…      ┆ e_Ethe…  │\n",
      "╞═══════════╪═══════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪══════════╡\n",
      "│ 2022-02-1 ┆ 246215633 ┆ 560180650 ┆ 719972444 ┆ … ┆ null      ┆ null      ┆ null      ┆ null     │\n",
      "│ 1         ┆ 5         ┆           ┆           ┆   ┆           ┆           ┆           ┆          │\n",
      "│ 2022-02-1 ┆ 246420416 ┆ 537846447 ┆ 672831429 ┆ … ┆ null      ┆ null      ┆ null      ┆ null     │\n",
      "│ 2         ┆ 1         ┆           ┆           ┆   ┆           ┆           ┆           ┆          │\n",
      "│ 2022-02-1 ┆ 249160040 ┆ 508550871 ┆ 669386509 ┆ … ┆ null      ┆ null      ┆ null      ┆ null     │\n",
      "│ 3         ┆ 0         ┆           ┆           ┆   ┆           ┆           ┆           ┆          │\n",
      "└───────────┴───────────┴───────────┴───────────┴───┴───────────┴───────────┴───────────┴──────────┘\n",
      "Number of rows in the dataset: 711\n"
     ]
    }
   ],
   "source": [
    "top_apy_per_protocol = DatasetsLoader().load(\"top-pools-apy-per-protocol\")\n",
    "\n",
    "# Filter the dataset to only include protocols with the target token\n",
    "\"\"\" unique_token_projects = top_apy_per_protocol \\\n",
    "    .filter(pl.col(\"underlying_token\").str.contains(TOKEN)) \\\n",
    "    .filter(pl.col(\"date\") > START_DATE) \\\n",
    "    .unique(\"project\") \\\n",
    "    .pivot(index=\"date\", columns=\"project\", values=[\"apy\", \"tvlUsd\"]) \"\"\"\n",
    "\n",
    "apy_df = top_apy_per_protocol \\\n",
    "    .filter(pl.col(\"underlying_token\").str.contains(TOKEN)) \\\n",
    "    .with_columns(\n",
    "            pl.col(\"project\") + \"_\" + pl.col(\"chain\") +  pl.col(\"underlying_token\")\n",
    "    ) \\\n",
    "    .drop([\"underlying_token\", \"chain\"])\n",
    "\n",
    "unique_projects = apy_df \\\n",
    "    .filter(pl.col(\"date\") <= START_DATE) \\\n",
    "    .select(\"project\") \\\n",
    "    .unique()\n",
    "\n",
    "apy_df_token = apy_df.join(\n",
    "    unique_projects, \n",
    "    on=\"project\", \n",
    "    how=\"inner\"\n",
    ")\n",
    "\n",
    "unique_token_projects = apy_df_token.pivot(\n",
    "    index=\"date\",\n",
    "    columns=\"project\",\n",
    "    values=[\"tvlUsd\", \"apy\"]\n",
    ")\n",
    "\n",
    "print(\"First few rows of the dataset:\")\n",
    "print(unique_token_projects.head(n = 3))\n",
    "print(\"Number of rows in the dataset:\", len(unique_token_projects))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset read from cache.\n",
      "Loading dataset tvl-per-project-tokens from cache.\n",
      "First few rows of the dataset:\n",
      "shape: (3, 20)\n",
      "┌────────────┬──────────┬──────────┬───────────┬───┬───────────┬───────────┬───────────┬───────────┐\n",
      "│ date       ┆ aave-v2  ┆ aura     ┆ pancakesw ┆ … ┆ rocket-po ┆ uniswap-v ┆ balancer- ┆ uniswap-v │\n",
      "│            ┆          ┆          ┆ ap-amm    ┆   ┆ ol        ┆ 2         ┆ v2        ┆ 3         │\n",
      "╞════════════╪══════════╪══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪═══════════╡\n",
      "│ 2023-05-14 ┆ 4.1997e8 ┆ 1.6514e8 ┆ null      ┆ … ┆ 1.1390e9  ┆ null      ┆ 2.4431e8  ┆ 7.7659e8  │\n",
      "│ 2022-12-07 ┆ 3.8163e8 ┆ 1.1824e8 ┆ null      ┆ … ┆ 4.1207e8  ┆ null      ┆ 1.9926e8  ┆ 6.0181e8  │\n",
      "│ 2022-07-31 ┆ 9.0687e8 ┆ 6.2815e6 ┆ null      ┆ … ┆ 3.5818e8  ┆ null      ┆ 1.6390e8  ┆ 9.0439e8  │\n",
      "└────────────┴──────────┴──────────┴───────────┴───┴───────────┴───────────┴───────────┴───────────┘\n"
     ]
    }
   ],
   "source": [
    "tvl_df = DatasetsLoader().load(\"tvl-per-project-tokens\") \\\n",
    "    .unique(subset=[\"date\", \"project\"]) \\\n",
    "    .filter(pl.col(\"date\") > START_DATE) \n",
    "\n",
    "tvl_per_projects_token = tvl_df[[TOKEN, \"project\", \"date\"]].pivot(\n",
    "    index=\"date\",\n",
    "    columns=\"project\",\n",
    "    values=TOKEN\n",
    ")\n",
    "\n",
    "print(\"First few rows of the dataset:\")\n",
    "print(tvl_per_projects_token.head(n = 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First few rows of the dataset:\n",
      "shape: (3, 162)\n",
      "┌───────────┬───────────┬───────────┬───────────┬───┬───────────┬───────────┬───────────┬──────────┐\n",
      "│ volumes_l ┆ price_dif ┆ price_dif ┆ price_dif ┆ … ┆ apy_proje ┆ apy_proje ┆ apy_proje ┆ apy_proj │\n",
      "│ ast_24h   ┆ f_1_days  ┆ f_3_days  ┆ f_7_days  ┆   ┆ ct_uniswa ┆ ct_uniswa ┆ ct_uniswa ┆ ect_year │\n",
      "│           ┆           ┆           ┆           ┆   ┆ p-v3_Arbi ┆ p-v3_Arbi ┆ p-v3_Ethe ┆ n-financ │\n",
      "│           ┆           ┆           ┆           ┆   ┆ tru…      ┆ tru…      ┆ reu…      ┆ e_Ethe…  │\n",
      "╞═══════════╪═══════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪══════════╡\n",
      "│ 1.024704  ┆ -2.197109 ┆ 0.11474   ┆ -0.843464 ┆ … ┆ -0.576543 ┆ -0.54936  ┆ -0.044265 ┆ -0.92143 │\n",
      "│           ┆           ┆           ┆           ┆   ┆           ┆           ┆           ┆ 2        │\n",
      "│ 0.498189  ┆ 0.064754  ┆ -1.764524 ┆ 0.120361  ┆ … ┆ -0.678317 ┆ -0.580351 ┆ -0.044265 ┆ -0.94637 │\n",
      "│           ┆           ┆           ┆           ┆   ┆           ┆           ┆           ┆ 2        │\n",
      "│ 0.246527  ┆ -1.042419 ┆ -1.890364 ┆ 0.298075  ┆ … ┆ 0.0       ┆ 0.0       ┆ 0.0       ┆ -0.94850 │\n",
      "│           ┆           ┆           ┆           ┆   ┆           ┆           ┆           ┆ 5        │\n",
      "└───────────┴───────────┴───────────┴───────────┴───┴───────────┴───────────┴───────────┴──────────┘\n"
     ]
    }
   ],
   "source": [
    "# Join the datasets by the date column to create the final dataset\n",
    "final_dataset = price_dataset \\\n",
    "    .join(tvl_per_projects_token, on=\"date\", how=\"inner\") \\\n",
    "    .join(unique_token_projects, on=\"date\", how=\"inner\")\n",
    "\n",
    "# Drop unnecessary columns and rows with irrelevant data\n",
    "# - columns with token, market_cap , date and current price: not relevant for the model\n",
    "# - rows with year < 2022: historical data is not relevant\n",
    "final_dataset = final_dataset \\\n",
    "    .filter(pl.col(\"year\") >= 2022) \\\n",
    "    .drop([\"token\", \"market_cap\", \"date\", \"price\", \"month\"])\n",
    "final_dataset = final_dataset.slice(0, len(final_dataset) - 1)\n",
    "# Drop columns if there are lots of missing values\n",
    "THRESHOLD = 0.2\n",
    "max_nulls = THRESHOLD * final_dataset.shape[0]\n",
    "columns_to_keep = [\n",
    "        col_name for col_name in final_dataset.columns if final_dataset[col_name].null_count() <= max_nulls\n",
    "]\n",
    "final_dataset = final_dataset.select(columns_to_keep    )\n",
    "\n",
    "# Split the dataset into features and labels\n",
    "features = final_dataset.drop(\"label\")\n",
    "labels = final_dataset[\"label\"]\n",
    "\n",
    "# Normalize the training dataset and fill missing values\n",
    "for col in features.columns:\n",
    "    mean_val = features[col].mean()\n",
    "    std_dev = features[col].std() if features[col].std() != 0 else 1\n",
    "    features = features.with_columns(((features[col].fill_null(mean_val) - mean_val) / std_dev).alias(col))\n",
    "\n",
    "print(\"First few rows of the dataset:\")\n",
    "print(features.head(n = 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Training\n",
    "\n",
    "##### Splitting the dataset\n",
    "\n",
    "We need to convert the dataframes to torch tensors and split the dataset into training and testing sets. As usual, we choose a reasonable split ratio (e.g., 80% training and 20% testing) and shuffle the data before splitting it: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert the dataset to a PyTorch tensor\n",
    "features_tensor = torch.tensor(features.to_numpy(), dtype=torch.float32)\n",
    "labels_tensor = torch.tensor(labels.to_numpy(), dtype=torch.float32)\n",
    "\n",
    "# Get a random permutation of the indices\n",
    "indices = torch.randperm(len(features_tensor))\n",
    "train_indices = indices[:int(0.75 * len(features_tensor))]\n",
    "test_indices = indices[int(0.75 * len(features_tensor)):]\n",
    "train_features, train_labels = features_tensor[train_indices], labels_tensor[train_indices]\n",
    "test_features, test_labels = features_tensor[test_indices], labels_tensor[test_indices]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Model Definition\n",
    "\n",
    "For the sake of simplicity, we illustrate the training process using a simple perceptron model no hidden layers and a single output neuron. We use the Sigmoid activation function to output the probability of the token price increasing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = nn.Sequential(\n",
    "    nn.Linear(len(features.columns), 1),\n",
    "    nn.Sigmoid()\n",
    ")\n",
    "\n",
    "criterion = nn.BCELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Training the Model\n",
    "\n",
    "We then train the model on the Binary Cross-Entropy loss function and the Adam optimizer. For the hyperparameters, we use a learning rate of 0.01 and 500 epochs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 66.67%\n"
     ]
    }
   ],
   "source": [
    "def train_and_test_model(model, train_features, train_labels, test_features, test_labels, criterion, optimizer):\n",
    "    model.train()\n",
    "    for _ in range(1000):\n",
    "        optimizer.zero_grad()\n",
    "        output = model(train_features)\n",
    "        loss = criterion(output, train_labels.unsqueeze(1))\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        correct, total = 0, 0\n",
    "        output = model(test_features)\n",
    "        predicted = torch.tensor([1 if x > 0.5 else 0 for x in output])\n",
    "        total += test_labels.size(0)\n",
    "        correct += (predicted == test_labels).sum().item()\n",
    "    return 100 * correct / total\n",
    "acc = train_and_test_model(model, train_features, train_labels, test_features, test_labels, criterion, optimizer)\n",
    "print(f'Accuracy: {acc:.2f}%')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This simple Logistic Regression model obtains an accuracy of 0.70 on the test set. In the original example from GIZA, the authors trained a Multilayer Perceptron model with 2 hidden layers of decreasing input size (64 and 32). Using their feature extraction process, their MLP achieved an accuracy of around 0.65. In the final part of this article, we will show how the accuracy can be further improved by training more complex models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Benchmarking Proof Generation\n",
    "\n",
    "Once the model is trained, we can generate a proof of its inference using EZKL. We first convert the model to the ONNX format and set up the proof generation process as described in the previous sections. To simplify the process, let us write a few helper functions to convert the model to ONNX format and time the proof generation process:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import os\n",
    "from typing import Tuple, Union\n",
    "import time\n",
    "import torch\n",
    "import sklearn as sk\n",
    "\n",
    "def to_onnx(model, input_sample, onnx_file):\n",
    "    # if the file already exists, delete it\n",
    "    if os.path.exists(onnx_file):\n",
    "        os.remove(onnx_file)\n",
    "    torch.onnx.export(\n",
    "        model,\n",
    "        input_sample,\n",
    "        onnx_file,\n",
    "        input_names = ['input'],             # Input and output labels to appear in the ONNX graph \n",
    "        output_names = ['output'],\n",
    "        opset_version=10,\n",
    "        do_constant_folding=True,\n",
    "        export_params=True, \n",
    "        dynamic_axes={\n",
    "            'input' : {0 : 'batch_size'},    # Variable length axes\n",
    "            'output' : {0 : 'batch_size'}\n",
    "        }\n",
    "    )\n",
    "\n",
    "def prove():\n",
    "    _= ezkl.prove(\n",
    "        WITNESS,\n",
    "        COMPILED_MODEL,\n",
    "        PK,\n",
    "        PROOF,\n",
    "        \"single\",\n",
    "    )\n",
    "\n",
    "def verify():\n",
    "    assert ezkl.verify(\n",
    "        PROOF,\n",
    "        SETTINGS,\n",
    "        VK,\n",
    "    ) == True\n",
    "\n",
    "import contextlib\n",
    "\n",
    "def bench_ezkl_single_round(\n",
    "    model: Union[torch.nn.Module, sk.base.BaseEstimator],\n",
    "    sample: torch.Tensor, \n",
    ") -> Tuple[float, float, float, int]:\n",
    "    \n",
    "    setup_time = -time.time()\n",
    "    with contextlib.redirect_stderr(None):\n",
    "        setup(\"data/models/perceptron.onnx\", model, sample)\n",
    "    setup_time += time.time()\n",
    "\n",
    "    logrows = json.load(open(SETTINGS, 'r'))[\"run_args\"][\"logrows\"]\n",
    "    \n",
    "    # Sleep for 1 second to make sure Rust has enough time to write the files\n",
    "    time.sleep(1) \n",
    "\n",
    "    prove_time = -time.time()\n",
    "    prove()\n",
    "    prove_time += time.time()\n",
    "\n",
    "    time.sleep(1) \n",
    "\n",
    "    verify_time = -time.time()\n",
    "    verify()\n",
    "    verify_time += time.time()\n",
    "\n",
    "    return setup_time, prove_time, verify_time, logrows\n",
    "\n",
    "def bench_ezkl(\n",
    "    model: Union[torch.nn.Module, sk.base.BaseEstimator],\n",
    "    test_features: torch.Tensor,\n",
    "    rounds: int = 1,\n",
    ") -> Tuple[float, float, float]:\n",
    "    \n",
    "    # Convert the model to ONNX and calibrate it\n",
    "    to_onnx(model, test_features[0].unsqueeze(0), \"data/models/perceptron.onnx\")  \n",
    "\n",
    "    setup_time, prove_time, verify_time, logrows = [], [], [], []\n",
    "    for _ in range(rounds):\n",
    "        # Reload the module to avoid any caching issues\n",
    "\n",
    "        from importlib import reload\n",
    "        import ezkl\n",
    "        reload(ezkl)\n",
    "\n",
    "        # randomly sample a feature from the test dataset\n",
    "        sample = test_features[torch.randint(0, len(test_features), (1,))]\n",
    "        s, p, v, l = bench_ezkl_single_round(model, sample)\n",
    "        setup_time.append(s), prove_time.append(p), verify_time.append(v), logrows.append(l)\n",
    "\n",
    "    # Calculate the average and standard deviation of the timings\n",
    "    avg_setup, avg_prove, avg_verify, avg_logrows = (\n",
    "        sum(setup_time) / rounds, \n",
    "        sum(prove_time) / rounds, \n",
    "        sum(verify_time) / rounds,\n",
    "        sum(logrows) / rounds\n",
    "    )\n",
    "\n",
    "    std_setup, std_prove, std_verify, std_logrows = (\n",
    "        (sum((s - avg_setup) ** 2 for s in setup_time) / rounds) ** 0.5,\n",
    "        (sum((p - avg_prove) ** 2 for p in prove_time) / rounds) ** 0.5,\n",
    "        (sum((v - avg_verify) ** 2 for v in verify_time) / rounds) ** 0.5,\n",
    "        (sum((l - avg_logrows) ** 2 for l in logrows) / rounds) ** 0.5\n",
    "    )\n",
    "\n",
    "    print(f\"Setup time: {str(avg_setup)[:5]} ± {str(std_setup/math.sqrt(rounds))[:5]} [s]\")\n",
    "    print(f\"Prover time: {str(avg_prove)[:5]} ± {str(std_prove/math.sqrt(rounds))[:5]} [s]\")\n",
    "    print(f\"Verifier time: {str(avg_verify)[:5]} ± {str(std_verify/math.sqrt(rounds))[:5]} [s]\")\n",
    "    print(f\"Logrows: {str(avg_logrows)[:5]} ± {str(std_logrows/math.sqrt(rounds))[:5]}\")\n",
    "    \n",
    "    return setup_time, prove_time, verify_time, logrows"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now time the `setup`, `prove`, and `verify` functions by calling the `bench_ezkl` function, which allows us to obtain average times with error margins by specifying the number of `rounds`. Let's benchmark the proof generation process for the simple perceptron model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setup time: 0.491 ± 0.0 [s]\n",
      "Prover time: 0.487 ± 0.0 [s]\n",
      "Verifier time: 0.015 ± 0.0 [s]\n",
      "Logrows: 12.0 ± 0.0\n"
     ]
    }
   ],
   "source": [
    "_ = bench_ezkl(\n",
    "    model,\n",
    "    test_features,\n",
    "    rounds=1\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Accuracy vs. Proving Costs\n",
    "\n",
    "For the main part of this article, we will compare the accuracy of the model with the cost of proving its inference. On the one hand, we increase the number of hidden layers and neurons of the perceptron model to show how a linear increase yields a linear increase in proof cost but a diminishing return in accuracy. On the other hand, we show how different architectures (e.g., Decision Trees, Random Forests, and SVMs) can obtain similar accuracies with varying proof costs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Increasing Model Complexity\n",
    "\n",
    "Let's start by increasing the complexity of the perceptron model. We evaluate perceptrons with one, two and three hidden layers for which we vary the number of neurons per layer as follows. We define the possible number of neurons per layer to be one of the following: [4, 8, 16, 32, 64, 128]. In addition, for any two consecutive layers, the outermost layer must have striclty less neurons than the inner one. We then train the model for each configuration and obtain the accuracies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "# Function to create a MLP model with the specified information\n",
    "def create_mlp_model(layer_info: List[Tuple[int, int]]) -> nn.Module:\n",
    "    layers = (\n",
    "        nn.Linear(in_size, out_size)\n",
    "        for in_size, out_size in layer_info\n",
    "    )\n",
    "    return nn.Sequential(*layers, nn.Sigmoid())\n",
    "\n",
    "\n",
    "# Function to train and return the model accuracy\n",
    "def train_and_return_model(model: nn.Module) -> Tuple[nn.Module, float]:\n",
    "    # Get a random permutation of the indices\n",
    "    indices = torch.randperm(len(features_tensor))\n",
    "    train_indices = indices[:int(0.75 * len(features_tensor))]\n",
    "    test_indices = indices[int(0.75 * len(features_tensor)):]\n",
    "    train_features, train_labels = features_tensor[train_indices], labels_tensor[train_indices]\n",
    "    test_features, test_labels = features_tensor[test_indices], labels_tensor[test_indices]\n",
    "\n",
    "    criterion = nn.BCELoss()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.0012)\n",
    "    acc = train_and_test_model(model, train_features, train_labels, test_features, test_labels, criterion, optimizer)\n",
    "    return model, acc\n",
    "\n",
    "def get_num_params(name: str) -> int:\n",
    "    numbers = [len(features.columns)] + [\n",
    "        int(s) for s in name.split('_')[3:] if s.isdigit()] + [1]\n",
    "    weights = sum([s1 * s2 for s1, s2 in zip(numbers, numbers[1:])])\n",
    "    biases = sum(numbers[1:])\n",
    "    return weights + biases\n",
    "\n",
    "# Given the number of layers, return all possible decreasing configurations\n",
    "# Where the number of neurons is in [4, 8, 16, 32, 64]\n",
    "def get_all_configurations(n_layers: int, in_features: int = len(features.columns)) -> List[List[int]]:\n",
    "    if n_layers == 0:\n",
    "        return [[(in_features, 1)]]\n",
    "\n",
    "    in_sizes = [128, 64, 32, 16, 8, 4]\n",
    "    combinations = list(itertools.combinations_with_replacement(in_sizes, n_layers))\n",
    "    increasing_combinations = [\n",
    "        c for c in combinations\n",
    "        if all(c[i] > c[i + 1] for i in range(len(c) - 1))\n",
    "    ]\n",
    "\n",
    "    sizes =[ \n",
    "        [(in_features, c[0])] + [\n",
    "            (c1, c2) for c1, c2 in zip(c, c[1:])\n",
    "        ] + [(c[-1], 1)]\n",
    "        for c in increasing_combinations\n",
    "    ]\n",
    "                                \n",
    "    return sizes\n",
    "\n",
    "# Dictionary to store the accuracy of each model\n",
    "acc = {}\n",
    "ROUNDS = 100\n",
    "\n",
    "for _ in range(ROUNDS):\n",
    "    for layers in range(0, 4):\n",
    "        for layer_info in get_all_configurations(layers):\n",
    "            model, accuracy = train_and_return_model(\n",
    "                create_mlp_model(layer_info)\n",
    "            )\n",
    "            tag = f\"MLP_{layers}_layers_{'_'.join(str(x[0]) for x in layer_info[1:])}\"\n",
    "            acc[tag] = acc.get(tag, []) + [accuracy]\n",
    "\n",
    "# Print the average accuracy for each model\n",
    "sorted_acc = sorted(acc.items(), key=lambda x: get_num_params(x[0]))\n",
    "for name, accuracies in sorted_acc:\n",
    "    print(f\"{name}: {str(sum(accuracies) / len(accuracies))[:5]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, the accuracy of the models correlates with the number of neurons per layer. However, at some point, the increase in accuracy becomes marginal or even stagnates. This is due to overfitting, as the model becomes too complex and starts to memorize the training data instead of generalizing well to unseen data. We will now observe how the proof costs increase with the number of neurons per layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bench_configuration(layer_info):\n",
    "    \n",
    "    layers = len(layer_info) - 1\n",
    "    tag = f\"MLP_{layers}_layers_{'_'.join(str(x[0]) for x in layer_info[1:])}\"\n",
    "    print(f\"{tag} ({get_num_params(tag)} params)\\n\")\n",
    "    \n",
    "    bench_ezkl(\n",
    "        train_and_return_model(create_mlp_model(layer_info))[0],\n",
    "        test_features,\n",
    "        rounds=50,\n",
    "    )\n",
    "\n",
    "    print(\"\\n\")\n",
    "\n",
    "for layers in range(0, 4):\n",
    "    for layer_info in get_all_configurations(layers):\n",
    "        # bench_configuration(layer_info)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 183,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.77%\n",
      "Accuracy: 0.72%\n",
      "Accuracy: 0.61%\n",
      "Accuracy: 0.70%\n",
      "Accuracy: 0.60%\n"
     ]
    }
   ],
   "source": [
    "from sklearn.svm import SVC\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from sklearn.linear_model import RidgeClassifier\n",
    "from sklearn.tree import DecisionTreeClassifier\n",
    "\n",
    "def train_sklearn_model(model: sk.base.BaseEstimator) -> float:\n",
    "    # Convert the dataset to a PyTorch tensor\n",
    "    features_tensor = torch.tensor(features.to_numpy(), dtype=torch.float32)\n",
    "    labels_tensor = torch.tensor(labels.to_numpy(), dtype=torch.float32)\n",
    "\n",
    "    # Get a random permutation of the indices\n",
    "    indices = torch.randperm(len(features_tensor))\n",
    "    train_indices = indices[:int(0.75 * len(features_tensor))]\n",
    "    test_indices = indices[int(0.75 * len(features_tensor)):]\n",
    "    train_features, train_labels = features_tensor[train_indices], labels_tensor[train_indices]\n",
    "    test_features, test_labels = features_tensor[test_indices], labels_tensor[test_indices]\n",
    "    model.fit(train_features, train_labels)\n",
    "    return model, model.score(test_features, test_labels)\n",
    "max_iter = 50\n",
    "# Train and test the model\n",
    "for model in [SVC(kernel=\"linear\", tol=0.01), LogisticRegression(tol=0.01), RandomForestClassifier(), RidgeClassifier(), DecisionTreeClassifier()]:\n",
    "    _, acc = train_sklearn_model(model)\n",
    "    print(f'Accuracy: {acc:.2f}%')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "accuracies = open(\"benches/accuracy.txt\", \"r\").read()\n",
    "accuracies = accuracies.split(\"\\n\")[1:-1]\n",
    "accuracies = [line.split(\": \") for line in accuracies]\n",
    "accuracies = {name: float(acc) for name, acc in accuracies}\n",
    "\n",
    "def get_num_params(name: str) -> int:\n",
    "    numbers = [162] + [\n",
    "        int(s) for s in name.split('_')[3:] if s.isdigit()] + [1]\n",
    "    weights = sum([s1 * s2 for s1, s2 in zip(numbers, numbers[1:])])\n",
    "    biases = sum(numbers[1:])\n",
    "    return weights + biases\n",
    "\n",
    "sorted_acc = sorted(accuracies.items(), key=lambda x: get_num_params(x[0]))\n",
    "\n",
    "for name, acc in sorted_acc:\n",
    "    num_layers = int(name.split('_')[1])\n",
    "    color = {0: 'black', 1: 'blue', 2: 'red', 3: 'green'}[num_layers]\n",
    "    plt.scatter(\n",
    "        get_num_params(name),\n",
    "        acc,\n",
    "        color=color,\n",
    "        marker='x',\n",
    "    )\n",
    "\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.grid(True)\n",
    "plt.scatter([], [], color='black', label='No hidden layers', marker='x')\n",
    "plt.scatter([], [], color='blue', label='1 hidden layer', marker='x')\n",
    "plt.scatter([], [], color='red', label='2 hidden layers', marker='x')\n",
    "plt.scatter([], [], color='green', label='3 hidden layers', marker='x')\n",
    "plt.xscale('log')\n",
    "plt.legend()\n",
    "plt.xlabel(\"Number of parameters\")\n",
    "plt.ylabel(\"Accuracy\")\n",
    "\n",
    "\n",
    "for name, acc in sorted_acc:\n",
    "    num_layers = int(name.split('_')[1])\n",
    "    color = {0: 'black', 1: 'blue', 2: 'red', 3: 'green'}[num_layers]\n",
    "    sns.regplot(\n",
    "        x=[get_num_params(name)],\n",
    "        y=[acc],\n",
    "        color=color,\n",
    "        marker='x',\n",
    "    )\n",
    "\n",
    "sns.set_style(\"whitegrid\")\n",
    "sns.despine()\n",
    "\n",
    "# Save without borders\n",
    "plt.savefig(\"plots/accuracy_vs_params.png\", bbox_inches='tight', pad_inches=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "timings = open(\"benches/time.txt\", \"r\").read()\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Extract the average times from the benchmark results\n",
    "lines_with_text = lambda text: [line for line in timings.split(\"\\n\") if text in line]\n",
    "layers = [int(line.split(\"_\")[1]) for line in lines_with_text(\"layers\")]\n",
    "params = [get_num_params(line.split()[0]) for line in lines_with_text(\"params\")]\n",
    "setup_times = [(float(line.split()[2]), float(line.split()[4])) for line in lines_with_text(\"Setup time\")]\n",
    "prove_times = [(float(line.split()[2]), float(line.split()[4])) for line in lines_with_text(\"Prover time\")]\n",
    "verify_times = [(float(line.split()[2]), float(line.split()[4])) for line in lines_with_text(\"Verifier time\")]\n",
    "logrows = [(float(line.split()[1]), float(line.split()[3])) for line in lines_with_text(\"Logrows\")]\n",
    "\n",
    "\n",
    "# Plot the four metrics\n",
    "fig, axs = plt.subplots(2, 2, figsize=(15, 10))\n",
    "\n",
    "# Grid lines\n",
    "for ax in axs.flatten():\n",
    "    ax.grid(True)\n",
    "\n",
    "# Plot the setup times\n",
    "for (i, j), l, times in zip([(0, 0), (0, 1), (1, 0), (1, 1)], [\"Setup\", \"Prover\", \"Verifier\", \"Logrows\"],  [setup_times, prove_times, verify_times, logrows]):\n",
    "    axs[i, j].set_title(l)\n",
    "    axs[i, j].set_ylabel(\"Time $[s]$\")\n",
    "    if l == \"Logrows\":\n",
    "        axs[i, j].set_ylabel(\"Logrows\")\n",
    "\n",
    "    if l == \"Verifier\":\n",
    "        axs[i, j].set_ylabel(\"Time $[ms]$\")\n",
    "    colors = {\n",
    "        0: 'black',\n",
    "        1: 'blue',\n",
    "        2: 'red',\n",
    "        3: 'green',\n",
    "    }\n",
    "\n",
    "    # Log scale for the x-axis\n",
    "    axs[i, j].set_xscale('log')\n",
    "\n",
    "    for la, param, (time, std) in zip(layers, params, times):\n",
    "        axs[i, j].errorbar(\n",
    "            param,\n",
    "            [time * 1000 if l == \"Verifier\" else time] ,\n",
    "            marker='x',\n",
    "            yerr=std,\n",
    "            # For thinner lines\n",
    "            elinewidth=0.3,\n",
    "            capsize=1,\n",
    "            color=colors[la],\n",
    "            label=f\"{la} layers\",\n",
    "        )\n",
    "\n",
    "        # Add x-lims\n",
    "        axs[i, j].set_xlim([min(params) / 1.5, max(params) * 1.5])\n",
    "\n",
    "        # Save without borders\n",
    "plt.savefig(\"plots/times_vs_params.png\", bbox_inches='tight', pad_inches=0)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}