{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MNIST Addition\n",
    "\n",
    "This notebook shows an implementation of [MNIST Addition](https://arxiv.org/abs/1805.10872). In this task, pairs of MNIST handwritten images and their sums are given, alongwith a domain knowledge base containing information on how to perform addition operations. The task is to recognize the digits of handwritten images and accurately determine their sum.\n",
    "\n",
    "Intuitively, we first use a machine learning model (learning part) to convert the input images to digits (we call them pseudo-labels), and then use the knowledge base (reasoning part) to calculate the sum of these digits. Since we do not have ground-truth of the digits, in Abductive Learning, the reasoning part will leverage domain knowledge and revise the initial digits yielded by the learning part through abductive reasoning. This process enables us to further update the machine learning model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary libraries and modules\n",
    "import os.path as osp\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.optim import RMSprop, lr_scheduler\n",
    "\n",
    "from ablkit.bridge import SimpleBridge\n",
    "from ablkit.data.evaluation import ReasoningMetric, SymbolAccuracy\n",
    "from ablkit.learning import ABLModel, BasicNN\n",
    "from ablkit.reasoning import KBBase, Reasoner\n",
    "from ablkit.utils import ABLLogger, print_log\n",
    "\n",
    "from datasets import get_dataset\n",
    "from models.nn import LeNet5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working with Data\n",
    "\n",
    "First, we get the training and testing datasets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = get_dataset(train=True, get_pseudo_label=True)\n",
    "test_data = get_dataset(train=False, get_pseudo_label=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`train_data` and `test_data` share identical structures: tuples with three components: X (list where each element is a list of two images), gt_pseudo_label (list where each element is a list of two digits, i.e., pseudo-labels) and Y (list where each element is the sum of the two digits). The length and structures of datasets are illustrated as follows.\n",
    "\n",
    "Note: ``gt_pseudo_label`` is only used to evaluate the performance of the learning part but not to train the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Both train_data and test_data consist of 3 components: X, gt_pseudo_label, Y\n",
      "\n",
      "Length of X, gt_pseudo_label, Y in train_data: 30000, 30000, 30000\n",
      "Length of X, gt_pseudo_label, Y in test_data: 5000, 5000, 5000\n",
      "\n",
      "X is a list, with each element being a list of 2 Tensor.\n",
      "gt_pseudo_label is a list, with each element being a list of 2 int.\n",
      "Y is a list, with each element being an int.\n"
     ]
    }
   ],
   "source": [
    "print(f\"Both train_data and test_data consist of 3 components: X, gt_pseudo_label, Y\")\n",
    "print()\n",
    "train_X, train_gt_pseudo_label, train_Y = train_data\n",
    "print(\n",
    "    f\"Length of X, gt_pseudo_label, Y in train_data: \"\n",
    "    + f\"{len(train_X)}, {len(train_gt_pseudo_label)}, {len(train_Y)}\"\n",
    ")\n",
    "test_X, test_gt_pseudo_label, test_Y = test_data\n",
    "print(\n",
    "    f\"Length of X, gt_pseudo_label, Y in test_data: \"\n",
    "    + f\"{len(test_X)}, {len(test_gt_pseudo_label)}, {len(test_Y)}\"\n",
    ")\n",
    "print()\n",
    "\n",
    "X_0, gt_pseudo_label_0, Y_0 = train_X[0], train_gt_pseudo_label[0], train_Y[0]\n",
    "print(\n",
    "    f\"X is a {type(train_X).__name__}, \"\n",
    "    + f\"with each element being a {type(X_0).__name__} \"\n",
    "    + f\"of {len(X_0)} {type(X_0[0]).__name__}.\"\n",
    ")\n",
    "print(\n",
    "    f\"gt_pseudo_label is a {type(train_gt_pseudo_label).__name__}, \"\n",
    "    + f\"with each element being a {type(gt_pseudo_label_0).__name__} \"\n",
    "    + f\"of {len(gt_pseudo_label_0)} {type(gt_pseudo_label_0[0]).__name__}.\"\n",
    ")\n",
    "print(f\"Y is a {type(train_Y).__name__}, \" + f\"with each element being an {type(Y_0).__name__}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ith element of X, gt_pseudo_label, and Y together constitute the ith data example. As an illustration, in the first data example of the training set, we have:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X in the first data example (a list of two images):\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAgMAAAD1CAYAAADNj/Z6AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8pXeV/AAAACXBIWXMAAA9hAAAPYQGoP6dpAAAKHklEQVR4nO3dT4hV9f/H8bk22ZRDVJgbI4MKK1rZoqg2kbTIUHAXNURUK6loUSQkVNiuooJw0cboD9EfIiISCiaoRQNTWhhFiyGyFiVkY0oJ6vlufvzgy8/f+1w798694+vx2L6u957v1+7p2YH5TK9pmmYCAIi1YtQXAACMlhgAgHBiAADCiQEACCcGACCcGACAcGIAAMKJAQAIJwYAINxkvy/s9XrDvA6gD8vxwFD3Dhi9tnuHJwMAEE4MAEA4MQAA4cQAAIQTAwAQTgwAQDgxAADhxAAAhBMDABBODABAODEAAOHEAACEEwMAEE4MAEA4MQAA4cQAAIQTAwAQTgwAQDgxAADhxAAAhBMDABBODABAODEAAOHEAACEEwMAEE4MAEA4MQAA4cQAAIQTAwAQTgwAQDgxAADhxAAAhBMDABBODABAODEAAOHEAACEEwMAEE4MAEC4yVFfAADZzj333HJfvXr1El3J/2/Tpk3lvmbNmnL/66+/yv3dd98t9wMHDpR7V54MAEA4MQAA4cQAAIQTAwAQTgwAQDgxAADhxAAAhOs1TdP09cJeb9jXsuytWrWq3F966aVyv/fee8v9xIkT5f7OO++Ue9vf4dtvv13uJ0+eLPeuZmdny/3w4cND/fzloM+v61hx76DNI488Uu7PP//8El3J6CwsLJT75Zdf3un92+4dngwAQDgxAADhxAAAhBMDABBODABAODEAAOHEAACEc87AAK1YUbfV3NxcuW/YsGGQl/N/tP0djvpn2D///PNy37JlS7kvLi4O8nLG0qj/jv4N944z35o1a8r9999/L/e2M1bWrl1b7rfffnu5D8JHH31U7m3/G/fu3dtpbztnpo1zBgCAkhgAgHBiAADCiQEACCcGACCcGACAcGIAAMKJAQAI59ChJfTkk0+W+44dO4b6+eN+6FCb66+/vtzn5+eX6EpGZ9z/jk7FvePMt2fPnnJvO1DsmWeeGeTlcAoOHQIASmIAAMKJAQAIJwYAIJwYAIBwYgAAwokBAAg3OeoLSPLss8+W++LiYqf3v/nmm8v9kksuKfe2n0P94osvyn3btm3lvnLlynIHxtPMzEy5b9y4sdx/+OGHQV4OQ+DJAACEEwMAEE4MAEA4MQAA4cQAAIQTAwAQTgwAQLhe0+cvSPc7ybnsssvK/Ztvvin36enpcv/222/L/YYbbij3Y8eOlfuZoM+v61hx7xh/bd/N/fv3l/vFF19c7tddd125O4dg+NruHZ4MAEA4MQAA4cQAAIQTAwAQTgwAQDgxAADhxAAAhJsc9QWwfLzwwgvl3vazym22bdtW7gnnCMAo7Ny5s9zXrVtX7o8//ni5O0dg/HkyAADhxAAAhBMDABBODABAODEAAOHEAACEEwMAEM45A/yvDRs2lPutt97a6f0PHDhQ7gsLC53eHzi1+++/v9zvu+++Tu+/a9euTn+e0fNkAADCiQEACCcGACCcGACAcGIAAMKJAQAIJwYAIFyvaZqmrxf2esO+FoZs5cqV5X7o0KFyn5qaKvfjx4+X+0033VTu8/Pz5c7ERJ9f17Hi3jF8GzduLPc9e/aU+1lnndXp848cOVLu3333Xblv37693GdnZ0/7mvhvbfcOTwYAIJwYAIBwYgAAwokBAAgnBgAgnBgAgHBiAADCOWcgyPvvv1/umzdv7vT+b775ZrnPzMx0en+cM5Bq9erV5b6wsFDu5513Xrnv27ev3K+88spyX7Gi/u/K6enpct+/f3+5t51Rcvjw4XLHOQMAQAsxAADhxAAAhBMDABBODABAODEAAOHEAACEmxz1BTA4d911V7lv2bKl3Nt+DvXgwYPl/tBDD5U78O+cOHGi3D/77LNy//jjj8t9165dp3tJ/+WCCy4o971795b7tddeW+4PPPBAuT/33HPlTjtPBgAgnBgAgHBiAADCiQEACCcGACCcGACAcGIAAMI5Z2AZueKKK8r9tddeK/euv1d+586d5X7o0KFO7w/j6I033ij3888/v/U97rzzznI/cuRIubd9tzZv3tx6DcN07NixTnubP//8s9Ofp50nAwAQTgwAQDgxAADhxAAAhBMDABBODABAODEAAOGcMzBGrrnmmnL/+uuvy71pmk6f33aOwCuvvNLp/WE5ajtH4I477mh9j3vuuafcX3755dO6pqU2OVn/q2L79u3lvn79+nL/6aefyv2tt94qd7rzZAAAwokBAAgnBgAgnBgAgHBiAADCiQEACCcGACCccwaW0DnnnFPuTz31VLmfffbZnT7/008/Lfenn3663I8fP97p82E52rdvX7n3c87AY489Vu5tP8f/448/lvuvv/5a7mvXri33tnvTo48+Wu433nhjubfdO5544olyP3r0aLnTnScDABBODABAODEAAOHEAACEEwMAEE4MAEA4MQAA4XpN0zR9vbDXG/a1nPFmZmbKfffu3Z3e/48//ij3Sy+9tNz//vvvTp/P8PX5dR0ry/3e0fa9aTuHYGJiYuLCCy8c0NWc2j///FPuU1NTQ/38gwcPlnvbOQWvvvrqIC+HU2i7d3gyAADhxAAAhBMDABBODABAODEAAOHEAACEEwMAEM45AwO0bt26cv/+++/Lve13ip88ebLct27dWu4ffvhhuTP+nDMwfl588cXW1zz44IPlPu7/Hy0sLJT7LbfcUu4///zzIC+Hf8E5AwBASQwAQDgxAADhxAAAhBMDABBODABAODEAAOEmR30By0nbOQA7duzo9OfbvP766+XuHAFYeg8//HDray666KJyv/vuuwd1Oac0Oztb7u+99165f/DBB+X+yy+/nPY1MV48GQCAcGIAAMKJAQAIJwYAIJwYAIBwYgAAwokBAAgnBgAgXK9pmqavF/Z6w76WsTczM1Puu3fv7vT+v/32W7lfffXV5b64uNjp8xl/fX5dx4p7x8TEqlWryn39+vXlvmnTpnL/5JNPyn1ubq7cl+M/V5yetr9jTwYAIJwYAIBwYgAAwokBAAgnBgAgnBgAgHBiAADCOWfgf1x11VWtr/nqq6/KfWpqqtyPHj1a7rfddlu5f/nll+XOmW85/jz4mX7vgOXAOQMAQEkMAEA4MQAA4cQAAIQTAwAQTgwAQDgxAADhJkd9AeNienq69TVt5wi0mZ+fL3fnCAAwCp4MAEA4MQAA4cQAAIQTAwAQTgwAQDgxAADhxAAAhHPOwADNzc2V+9atW5foSgCgf54MAEA4MQAA4cQAAIQTAwAQTgwAQDgxAADhxAAAhOs1TdP09cJeb9jXArTo8+s6Vtw7YPTa7h2eDABAODEAAOHEAACEEwMAEE4MAEA4MQAA4cQAAIQTAwAQTgwAQDgxAADhxAAAhBMDABBODABAODEAAOHEAACE6zXL8RekAwAD48kAAIQTAwAQTgwAQDgxAADhxAAAhBMDABBODABAODEAAOHEAACE+w8bbeVfVjXu7QAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gt_pseudo_label in the first data example (a list of two ground truth pseudo-labels): [7, 5]\n",
      "Y in the first data example (their sum result): 12\n"
     ]
    }
   ],
   "source": [
    "X_0, gt_pseudo_label_0, Y_0 = train_X[0], train_gt_pseudo_label[0], train_Y[0]\n",
    "print(f\"X in the first data example (a list of two images):\")\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.axis(\"off\")\n",
    "plt.imshow(X_0[0].squeeze(), cmap=\"gray\")\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.axis(\"off\")\n",
    "plt.imshow(X_0[1].squeeze(), cmap=\"gray\")\n",
    "plt.show()\n",
    "print(\n",
    "    f\"gt_pseudo_label in the first data example (a list of two ground truth pseudo-labels): {gt_pseudo_label_0}\"\n",
    ")\n",
    "print(f\"Y in the first data example (their sum result): {Y_0}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the Learning Part"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To build the learning part, we need to first build a machine learning base model. We use a simple [LeNet-5 neural network](https://en.wikipedia.org/wiki/LeNet), and encapsulate it within a `BasicNN` object to create the base model. `BasicNN` is a class that encapsulates a PyTorch model, transforming it into a base model with a sklearn-style interface. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "cls = LeNet5(num_classes=10)\n",
    "loss_fn = nn.CrossEntropyLoss(label_smoothing=0.2)\n",
    "optimizer = RMSprop(cls.parameters(), lr=0.0003, alpha=0.9)\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "scheduler = lr_scheduler.OneCycleLR(optimizer, max_lr=0.0003, pct_start=0.15, total_steps=200)\n",
    "\n",
    "base_model = BasicNN(\n",
    "    cls,\n",
    "    loss_fn,\n",
    "    optimizer,\n",
    "    scheduler=scheduler,\n",
    "    device=device,\n",
    "    batch_size=32,\n",
    "    num_epochs=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`BasicNN` offers methods like `predict` and `predict_prob`, which are used to predict the class index and the probabilities of each class for images. As shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted class index for a batch of 32 instances: ndarray with shape (32,)\n",
      "Predicted class probabilities for a batch of 32 instances: ndarray with shape (32, 10)\n"
     ]
    }
   ],
   "source": [
    "data_instances = [torch.randn(1, 28, 28) for _ in range(32)]\n",
    "pred_idx = base_model.predict(X=data_instances)\n",
    "print(\n",
    "    f\"Predicted class index for a batch of 32 instances: \"\n",
    "    + f\"{type(pred_idx).__name__} with shape {pred_idx.shape}\"\n",
    ")\n",
    "pred_prob = base_model.predict_proba(X=data_instances)\n",
    "print(\n",
    "    f\"Predicted class probabilities for a batch of 32 instances: \"\n",
    "    + f\"{type(pred_prob).__name__} with shape {pred_prob.shape}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, the base model built above deals with instance-level data (i.e., individual images), and can not directly deal with example-level data (i.e., a pair of images). Therefore, we wrap the base model into `ABLModel`, which enables the learning part to train, test, and predict on example-level data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ABLModel(base_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an illustration, consider this example of training on example-level data using the `predict` method in `ABLModel`. In this process, the method accepts data examples as input and outputs the class labels and the probabilities of each class for all instances within these data examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted class labels for the 100 data examples: \n",
      "a list of length 100, and each element is a ndarray of shape (2,).\n",
      "\n",
      "Predicted class probabilities for the 100 data examples: \n",
      "a list of length 100, and each element is a ndarray of shape (2, 10).\n"
     ]
    }
   ],
   "source": [
    "from ablkit.data.structures import ListData\n",
    "\n",
    "# ListData is a data structure provided by ABLkit that can be used to organize data examples\n",
    "data_examples = ListData()\n",
    "# We use the first 100 data examples in the training set as an illustration\n",
    "data_examples.X = train_X[:100]\n",
    "data_examples.gt_pseudo_label = train_gt_pseudo_label[:100]\n",
    "data_examples.Y = train_Y[:100]\n",
    "\n",
    "# Perform prediction on the 100 data examples\n",
    "pred_label, pred_prob = model.predict(data_examples)[\"label\"], model.predict(data_examples)[\"prob\"]\n",
    "print(\n",
    "    f\"Predicted class labels for the 100 data examples: \\n\"\n",
    "    + f\"a list of length {len(pred_label)}, and each element is \"\n",
    "    + f\"a {type(pred_label[0]).__name__} of shape {pred_label[0].shape}.\\n\"\n",
    ")\n",
    "print(\n",
    "    f\"Predicted class probabilities for the 100 data examples: \\n\"\n",
    "    + f\"a list of length {len(pred_prob)}, and each element is \"\n",
    "    + f\"a {type(pred_prob[0]).__name__} of shape {pred_prob[0].shape}.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the Reasoning Part"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the reasoning part, we first build a knowledge base which contains information on how to perform addition operations. We build it by creating a subclass of `KBBase`. In the derived subclass, we initialize the `pseudo_label_list` parameter specifying list of possible pseudo-labels, and override the `logic_forward` function defining how to perform (deductive) reasoning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AddKB(KBBase):\n",
    "    def __init__(self, pseudo_label_list=list(range(10))):\n",
    "        super().__init__(pseudo_label_list)\n",
    "\n",
    "    # Implement the deduction function\n",
    "    def logic_forward(self, nums):\n",
    "        return sum(nums)\n",
    "\n",
    "\n",
    "kb = AddKB()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The knowledge base can perform logical reasoning (both deductive reasoning and abductive reasoning). Below is an example of performing (deductive) reasoning, and users can refer to [Documentation]() for details of abductive reasoning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reasoning result of pseudo-labels [1, 2] is 3.\n"
     ]
    }
   ],
   "source": [
    "pseudo_labels = [1, 2]\n",
    "reasoning_result = kb.logic_forward(pseudo_labels)\n",
    "print(f\"Reasoning result of pseudo-labels {pseudo_labels} is {reasoning_result}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: In addition to building a knowledge base based on `KBBase`, we can also establish a knowledge base with a ground KB using `GroundKB`, or a knowledge base implemented based on Prolog files using `PrologKB`. The corresponding code for these implementations can be found in the `main.py` file. Those interested are encouraged to examine it for further insights."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we create a reasoner by instantiating the class ``Reasoner``. Due to the indeterminism of abductive reasoning, there could be multiple candidates compatible with the knowledge base. When this happens, reasoner can minimize inconsistencies between the knowledge base and pseudo-labels predicted by the learning part, and then return only one candidate that has the highest consistency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "reasoner = Reasoner(kb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: During creating reasoner, the definition of \"consistency\" can be customized within the `dist_func` parameter. In the code above, we employ a consistency measurement based on confidence, which calculates the consistency between the data example and candidates based on the confidence derived from the predicted probability. In `main.py`, we provide options for utilizing other forms of consistency measurement.\n",
    "\n",
    "Note: Also, during the process of inconsistency minimization, one can leverage [ZOOpt library](https://github.com/polixir/ZOOpt) for acceleration. Options for this are also available in `main.py`. Those interested are encouraged to explore these features."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building Evaluation Metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we set up evaluation metrics. These metrics will be used to evaluate the model performance during training and testing. Specifically, we use `SymbolAccuracy` and `ReasoningMetric`, which are used to evaluate the accuracy of the machine learning model’s predictions and the accuracy of the final reasoning results, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "metric_list = [SymbolAccuracy(prefix=\"mnist_add\"), ReasoningMetric(kb=kb, prefix=\"mnist_add\")]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bridging Learning and Reasoning\n",
    "\n",
    "Now, the last step is to bridge the learning and reasoning part. We proceed with this step by creating an instance of `SimpleBridge`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "bridge = SimpleBridge(model, reasoner, metric_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform training and testing by invoking the `train` and `test` methods of `SimpleBridge`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build logger\n",
    "print_log(\"Abductive Learning on the MNIST Addition example.\", logger=\"current\")\n",
    "log_dir = ABLLogger.get_current_instance().log_dir\n",
    "weights_dir = osp.join(log_dir, \"weights\")\n",
    "\n",
    "bridge.train(train_data, loops=2, segment_size=0.01, save_interval=1, save_dir=weights_dir)\n",
    "bridge.test(test_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Performance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We present the results of ABL as follows, which include the reasoning accuracy (the proportion of equations that are correctly summed), and the training time used to achieve this accuracy. These results are compared with the following methods:\n",
    "- [**NeurASP**](https://github.com/azreasoners/NeurASP): An extension of answer set programs by treating the neural network output as the probability distribution over atomic facts;\n",
    "- [**DeepProbLog**](https://github.com/ML-KULeuven/deepproblog): An extension of ProbLog by introducing neural predicates in Probabilistic Logic Programming;\n",
    "- [**DeepStochLog**](https://github.com/ML-KULeuven/deepstochlog): A neural-symbolic framework based on stochastic logic program."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "|  Method      | Accuracy | Time to achieve the Acc. (s) |\n",
    "| :----------: | :------: | :--------------------------: |\n",
    "|  NeurASP     |  0.964   |             354              |\n",
    "| DeepProbLog  |  0.965   |             1965             |\n",
    "| DeepStochLog |  0.975   |             727              |\n",
    "|     ABL      |  0.980   |              42              |"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "abl",
   "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.8.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "9c8d454494e49869a4ee4046edcac9a39ff683f7d38abf0769f648402670238e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
