{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9c696e2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "import torch \n",
    "import numpy as np \n",
    "import networkx as nx \n",
    "\n",
    "from matplotlib import pyplot as plt \n",
    "\n",
    "from matplotlib.patches import FancyArrowPatch\n",
    "from networkx.drawing.nx_pylab import draw_networkx_nodes, draw_networkx_labels\n",
    "\n",
    "from dagma import utils \n",
    "from dagma.nonlinear import DagmaMLP, DagmaNonlinear"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6b8679d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "project_dir = \"../../\"\n",
    "\n",
    "os.chdir(project_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "72724da0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from data_generation.simulation_data import InterventionDataset\n",
    "\n",
    "from models.helper.functions import gumbelSoftMLP, gnet_z\n",
    "from models.nodags.resblock import iResBlock\n",
    "from models.dccd.implicitblock import imBlock\n",
    "from models.helper.layers.mlpLipschitz import linearLipschitz\n",
    "from models.dccd.trainer_dccd import DCCDTrainer\n",
    "\n",
    "from utils.plotting import draw_curved_edges\n",
    "from utils.graph_transformations import *\n",
    "from utils.error_metrics import *"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "419d0c5f",
   "metadata": {},
   "source": [
    "## Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "66208a26",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = \"datasets/sachs\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d8cbc58c",
   "metadata": {},
   "outputs": [],
   "source": [
    "intervention_sets = [[None]]\n",
    "\n",
    "datasets = list()\n",
    "for i in range(len(intervention_sets)):\n",
    "    dataset = np.load(os.path.join(data_dir, f\"dataset-{i}.npy\"))\n",
    "    datasets.append(\n",
    "        dataset - dataset.mean(axis=0)\n",
    "    )\n",
    "\n",
    "datasets = [np.vstack(datasets)]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a6ab61ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_nodes = datasets[0].shape[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d9b65f87",
   "metadata": {},
   "outputs": [],
   "source": [
    "adj_mat = np.load(os.path.join(data_dir, \"adj-mat.npy\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "524d58ea",
   "metadata": {},
   "source": [
    "## Training the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "d968379e",
   "metadata": {},
   "outputs": [],
   "source": [
    "g_x = gumbelSoftMLP(\n",
    "    n_nodes = n_nodes,\n",
    "    lip_constant = 0.9,\n",
    "    activation = \"tanh\"\n",
    ")\n",
    "\n",
    "g_z = gnet_z(\n",
    "    n_nodes = n_nodes\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "b35df0ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "implicit_block = imBlock(\n",
    "    nnet_x = g_x, \n",
    "    nnet_z = g_z, \n",
    "    confounders = False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "1ea85483",
   "metadata": {},
   "outputs": [],
   "source": [
    "im_trainer = DCCDTrainer(\n",
    "    implicit_block, \n",
    "    max_epochs=1000, \n",
    "    batch_size=1024, \n",
    "    lr=1e-2,\n",
    "    lambda_c=1e-2, \n",
    "    rho=250.0\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "6ce952e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1000/1000, Intervention: 1/1, log P(X): -211.56\n"
     ]
    }
   ],
   "source": [
    "im_trainer.train(\n",
    "    intervention_datasets=datasets,\n",
    "    intervention_targets=intervention_sets,\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25f0b537",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, \"names.txt\"), \"r\") as file:\n",
    "    names = file.readlines()\n",
    "    \n",
    "names = [name.strip() for name in names]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "4171c457",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 22404 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axs = plt.subplots(1, 2)\n",
    "\n",
    "plt.rcParams[\"font.family\"] = \"Helvetica\"\n",
    "\n",
    "d = len(names)\n",
    "\n",
    "axs[0].imshow(np.abs(adj_mat) > 0, cmap='bwr', vmin=-1, vmax=1)\n",
    "axs[0].set_xticks(range(d))\n",
    "axs[0].set_yticks(range(d))\n",
    "axs[0].set_xticklabels(names, fontsize=7, rotation=90)\n",
    "axs[0].set_yticklabels(names, fontsize=7)\n",
    "axs[0].set_title(\"True Graph (Directed)\", fontsize=10)\n",
    "\n",
    "axs[1].imshow(g_x.get_w_adj().detach().cpu().numpy() > 0.5, cmap='bwr', vmin=-1, vmax=1)\n",
    "axs[1].set_xticks(range(d))\n",
    "axs[1].set_yticks(range(d))\n",
    "axs[1].set_xticklabels(names, fontsize=7, rotation=90)\n",
    "_ = axs[1].set_yticklabels(names, fontsize=7)\n",
    "axs[1].set_title(\"DCCD-CONF: Estimated Graph (Directed)\", fontsize=10)\n",
    "\n",
    "fig.subplots_adjust(wspace=0.5)\n",
    "\n",
    "fig.savefig(\"../results/sachs-adj-mat.pdf\", bbox_inches=\"tight\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "590a191b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SHD = 18\n"
     ]
    }
   ],
   "source": [
    "shd, _ = compute_shd(np.abs(adj_mat) > 0, g_x.get_w_adj().detach().cpu().numpy() > 0.5)\n",
    "\n",
    "print(f\"SHD = {int(shd)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4ab7513",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "causal-discovery",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
