{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b6593c0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "import matplotlib.pyplot as plt \n",
    "import networkx as nx \n",
    "import os \n",
    "\n",
    "from matplotlib.patches import FancyArrowPatch\n",
    "from networkx.drawing.nx_pylab import draw_networkx_nodes, draw_networkx_labels\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import DataLoader \n",
    "from sklearn.covariance import graphical_lasso "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "67e309e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "project_dir = '../../'\n",
    "\n",
    "os.chdir(project_dir)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6542bd0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from data_generation.graph import DirectedGraphGenerator \n",
    "from data_generation.structural_models import SEM\n",
    "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": "code",
   "execution_count": 5,
   "id": "a7b40b1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Experiment parameters - data\n",
    "\n",
    "n_nodes = 5\n",
    "exp_density = 2\n",
    "n_samples = 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bf487a71",
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_gen = DirectedGraphGenerator(\n",
    "    nodes=n_nodes, \n",
    "    expected_density=exp_density, \n",
    "    enforce_dag=False\n",
    ")\n",
    "\n",
    "graph = graph_gen()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1503df06",
   "metadata": {},
   "outputs": [],
   "source": [
    "sem = SEM(\n",
    "    graph=graph,\n",
    "    contractive=True,  \n",
    "    confounders=True, \n",
    "    off_diag_nonzeros=2,\n",
    "    noise_scale=0.5\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90645954",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 42068 characters]",
      "text/plain": [
       "<Figure size 1500x400 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axs = plt.subplots(1, 3) \n",
    "fig.set_figwidth(15)\n",
    "fig.set_figheight(4)\n",
    "\n",
    "pos = nx.circular_layout(graph)\n",
    "\n",
    "draw_networkx_nodes(graph, pos, ax=axs[0], node_size=300, node_color='white', edgecolors='black')\n",
    "draw_networkx_labels(graph, pos, ax=axs[0], font_size=8, font_color=\"black\")\n",
    "draw_curved_edges(graph, pos, axs[0], curve_scale=0.1, node_size=300)\n",
    "\n",
    "_ = axs[1].imshow(sem.confounder_sigma, cmap='bwr', vmin=-1, vmax=1)\n",
    "\n",
    "_ = axs[2].imshow(sem.confounder_inverse_cov, cmap='bwr', vmin=-5, vmax=5)\n",
    "\n",
    "axs[0].set_title(\"Graph\")\n",
    "axs[1].set_title(\"Confounder Sigma\")\n",
    "_ = axs[2].set_title(\"Confounder noise (Inverse Cov)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f2f9c619",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_inter = 5\n",
    "intervention_sets = [[None]] + [[i] for i in range(n_inter)]\n",
    "\n",
    "datasets = list()\n",
    "intervention_torch_datasets = list()\n",
    "for targets in intervention_sets:\n",
    "    data = sem.generateData(n_samples=n_samples, intervention_set=targets)\n",
    "    datasets.append(data)\n",
    "\n",
    "    intervention_torch_datasets.append(\n",
    "        InterventionDataset(data, targets)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7f056b3a",
   "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": 11,
   "id": "b6f418d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "implicit_block = imBlock(\n",
    "    nnet_x=g_x, \n",
    "    nnet_z=g_z,\n",
    "    confounders=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f751b600",
   "metadata": {},
   "outputs": [],
   "source": [
    "im_trainer = DCCDTrainer(\n",
    "    implicit_block, \n",
    "    max_epochs=200, \n",
    "    batch_size=512, \n",
    "    lr=1e-1,\n",
    "    lambda_c=1e-3, \n",
    "    rho=1e-2\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a9f2ddaf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 200/200, Intervention: 6/6, log P(X): 1.58\n"
     ]
    }
   ],
   "source": [
    "im_trainer.train(\n",
    "    intervention_datasets=datasets,\n",
    "    intervention_targets=intervention_sets,\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "8b34db6c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'Estimated')"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 10920 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axs = plt.subplots(1, 2)\n",
    "\n",
    "axs[0].imshow(np.abs(sem.weights) > 0, cmap='bwr', vmin=-1, vmax=1)\n",
    "axs[0].set_title(\"Ground Truth\")\n",
    "\n",
    "axs[1].imshow(g_x.get_w_adj().detach().cpu().numpy(), cmap='bwr', vmin=-1, vmax=1)\n",
    "axs[1].set_title(\"Estimated\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "cc4a7d83",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'Estimated')"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 10872 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axs = plt.subplots(1, 2)\n",
    "\n",
    "axs[0].imshow(sem.confounder_sigma, cmap='bwr', vmin=-1, vmax=1)\n",
    "axs[0].set_title(\"Ground truth\")\n",
    "\n",
    "axs[1].imshow(implicit_block.cov_mat.detach().cpu().numpy(), cmap='bwr', vmin=-1, vmax=1)\n",
    "axs[1].set_title(\"Estimated\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac915a66",
   "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
}
