{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NeuPI: Loading and Evaluating PGMs\n",
    "\n",
    "This notebook demonstrates the foundational capabilities of the `neupi` library: loading Probabilistic Graphical Models (PGMs) and using them to evaluate the log-likelihood of variable assignments.\n",
    "\n",
    "We will cover:\n",
    "1. Loading a **Markov Network (MN)** from a `.uai` file.\n",
    "2. Loading a **Sum-Product Network (SPN)** from a `.json` file.\n",
    "3. Creating random data assignments.\n",
    "4. Evaluating the assignments to compute their log-likelihood using the loaded PGMs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "First, let's import the necessary components from PyTorch and the `neupi` library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cuda\n",
      "Markov Network path: networks/mn/Segmentation_12.uai\n",
      "Sum-Product Network path: networks/spn/nltcs/spn.json\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from pathlib import Path\n",
    "import os\n",
    "\n",
    "# Import the PGM wrappers from neupi\n",
    "from neupi import MarkovNetwork, SumProductNetwork\n",
    "\n",
    "# Define the device for computation\n",
    "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f\"Using device: {DEVICE}\")\n",
    "\n",
    "# --- Path Setup ---\n",
    "# get root path of the project\n",
    "\n",
    "MN_UAI_PATH = Path(\"networks\") / \"mn\" / \"Segmentation_12.uai\"\n",
    "SPN_JSON_PATH = Path(\"networks\") / \"spn\" / \"nltcs\" / \"spn.json\"\n",
    "\n",
    "print(f\"Markov Network path: {MN_UAI_PATH}\")\n",
    "print(f\"Sum-Product Network path: {SPN_JSON_PATH}\")\n",
    "\n",
    "# Check if files exist\n",
    "assert MN_UAI_PATH.exists(), f\"File not found: {MN_UAI_PATH}\"\n",
    "assert SPN_JSON_PATH.exists(), f\"File not found: {SPN_JSON_PATH}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: Working with a Markov Network (MN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using 1d factors: False\n",
      "PGM is pairwise.\n",
      "Successfully loaded Markov Network.\n",
      "Number of variables: 229\n",
      "Number of factors: 851\n"
     ]
    }
   ],
   "source": [
    "# Load the Markov Network from the .uai file\n",
    "# The MarkovNetwork class parses the file and constructs the necessary tensors for evaluation.\n",
    "mn_evaluator = MarkovNetwork(uai_file=str(MN_UAI_PATH), device=DEVICE)\n",
    "\n",
    "print(f\"Successfully loaded Markov Network.\")\n",
    "print(f\"Number of variables: {mn_evaluator.num_variables}\")\n",
    "print(f\"Number of factors: {len(mn_evaluator.pgm.prob_tables_list)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of a random assignment batch: torch.Size([10, 229])\n",
      "First assignment vector:\n",
      " tensor([1, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 1, 1, 1, 0, 1, 1, 0, 1,\n",
      "        0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 0, 0, 0, 1, 1, 1, 1, 0, 1, 1, 0, 0, 0, 1,\n",
      "        0, 1, 1, 0, 0, 0, 0, 1, 1, 1, 0, 1, 1, 1, 1, 0, 1, 0, 0, 1, 1, 0, 0, 1,\n",
      "        0, 0, 0, 0, 1, 0, 0, 1, 0, 1, 1, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 1,\n",
      "        0, 0, 1, 1, 0, 0, 0, 0, 1, 1, 0, 0, 1, 1, 0, 1, 1, 0, 0, 0, 1, 0, 0, 1,\n",
      "        1, 0, 0, 0, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 1, 0, 1,\n",
      "        1, 1, 1, 0, 1, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 1, 1, 0, 1, 0, 0, 1, 0,\n",
      "        0, 1, 0, 0, 0, 0, 1, 0, 1, 1, 1, 0, 0, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1,\n",
      "        0, 0, 0, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 0, 0, 0, 1, 1, 0,\n",
      "        0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "# Create a batch of random assignments for the variables.\n",
    "# Each assignment is a binary vector where the i-th element represents the state of the i-th variable.\n",
    "batch_size = 10\n",
    "num_variables_mn = mn_evaluator.num_variables\n",
    "\n",
    "# Generate a random tensor of shape (batch_size, num_variables) with values 0 or 1.\n",
    "mn_assignments = torch.randint(0, 2, (batch_size, num_variables_mn), device=DEVICE)\n",
    "\n",
    "print(\"Shape of a random assignment batch:\", mn_assignments.shape)\n",
    "print(\"First assignment vector:\\n\", mn_assignments[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of the output tensor: torch.Size([10])\n",
      "Log-likelihoods for each assignment:\n",
      " [-554.85034 -566.1757  -582.0522  -577.035   -584.3352  -562.16864\n",
      " -543.0751  -570.3635  -563.93616 -561.0249 ]\n",
      "\n",
      "Average Log-Likelihood (MN): -566.5016\n"
     ]
    }
   ],
   "source": [
    "# Evaluate the log-likelihood of the assignments.\n",
    "# The evaluator is a callable object. It returns the log-likelihood (in base e) for each assignment in the batch.\n",
    "with torch.no_grad():\n",
    "    log_likelihoods_mn = mn_evaluator(mn_assignments)\n",
    "\n",
    "print(f\"Shape of the output tensor: {log_likelihoods_mn.shape}\")\n",
    "print(f\"Log-likelihoods for each assignment:\\n {log_likelihoods_mn.cpu().numpy()}\")\n",
    "\n",
    "# We can also compute the average log-likelihood across the batch\n",
    "average_ll_mn = log_likelihoods_mn.mean()\n",
    "print(f\"\\nAverage Log-Likelihood (MN): {average_ll_mn.item():.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: Working with a Sum-Product Network (SPN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Successfully loaded Sum-Product Network.\n",
      "Number of variables: 72\n"
     ]
    }
   ],
   "source": [
    "# Load the Sum-Product Network from the .json file\n",
    "spn_evaluator = SumProductNetwork(json_file=str(SPN_JSON_PATH), device=DEVICE)\n",
    "\n",
    "print(f\"Successfully loaded Sum-Product Network.\")\n",
    "print(f\"Number of variables: {spn_evaluator.num_var}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of a random assignment batch: torch.Size([10, 72])\n",
      "First assignment vector:\n",
      " tensor([1, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0, 0, 0, 0, 1, 0, 1, 1, 1,\n",
      "        0, 0, 1, 1, 0, 0, 0, 1, 1, 1, 0, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 1,\n",
      "        0, 1, 1, 1, 1, 0, 1, 0, 0, 1, 1, 1, 0, 1, 1, 1, 0, 0, 0, 0, 1, 1, 0, 0],\n",
      "       device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "# Create a batch of random assignments for the SPN.\n",
    "batch_size = 10\n",
    "num_variables_spn = spn_evaluator.num_var\n",
    "\n",
    "spn_assignments = torch.randint(0, 2, (batch_size, num_variables_spn), device=DEVICE)\n",
    "\n",
    "print(\"Shape of a random assignment batch:\", spn_assignments.shape)\n",
    "print(\"First assignment vector:\\n\", spn_assignments[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of the output tensor: torch.Size([10])\n",
      "Log-likelihoods for each assignment:\n",
      " [-18.725033 -23.253717 -20.403366 -12.997818 -14.301248 -16.177542\n",
      " -14.136018 -15.460924 -25.527037 -19.932896]\n",
      "\n",
      "Average Log-Likelihood (SPN): -18.0916\n"
     ]
    }
   ],
   "source": [
    "# Evaluate the log-likelihood of the assignments using the SPN.\n",
    "with torch.no_grad():\n",
    "    log_likelihoods_spn = spn_evaluator(spn_assignments)\n",
    "\n",
    "print(f\"Shape of the output tensor: {log_likelihoods_spn.shape}\")\n",
    "print(f\"Log-likelihoods for each assignment:\\n {log_likelihoods_spn.cpu().numpy()}\")\n",
    "\n",
    "# Compute the average log-likelihood\n",
    "average_ll_spn = log_likelihoods_spn.mean()\n",
    "print(f\"\\nAverage Log-Likelihood (SPN): {average_ll_spn.item():.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "This notebook demonstrated the basic process of loading different types of PGMs (`MarkovNetwork`, `SumProductNetwork`) and using them to compute the log-likelihood of given variable assignments. This evaluation capability is the cornerstone of the `neupi` library, as it provides the supervision signal for training neural networks."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "neupi",
   "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.9.23"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}