{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f270e84",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "import torch \n",
    "import numpy as np \n",
    "import random\n",
    "import pandas as pd\n",
    "\n",
    "from matplotlib import pyplot as plt \n",
    "\n",
    "plt.rcParams[\"font.family\"] = \"Helvetica\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3dce22df",
   "metadata": {},
   "outputs": [],
   "source": [
    "project_dir = \"../../\"\n",
    "\n",
    "os.chdir(project_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bf81ebf3",
   "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": "code",
   "execution_count": 4,
   "id": "512fb398",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = \"datasets/perturb-cite-seq\"\n",
    "\n",
    "condition = \"cocult\"\n",
    "\n",
    "condition_dir = os.path.join(data_dir, condition)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0be2b4d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_datasets = [\n",
    "    np.load(os.path.join(condition_dir + \"/training_data\", f\"dataset-{i}.npy\"))\n",
    "    for i in range(55)\n",
    "]\n",
    "\n",
    "validation_datasets = [\n",
    "    np.load(os.path.join(condition_dir + \"/validation_data\", f\"dataset-{i}.npy\")) for i in range(6)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f35f467",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "tot_dataset = training_datasets + validation_datasets\n",
    "\n",
    "# centering the data\n",
    "tot_dataset_centered = list()\n",
    "for dataset in tot_dataset:\n",
    "    dataset_centered = dataset - dataset.mean(axis=0)\n",
    "    tot_dataset_centered.append(dataset_centered)\n",
    "\n",
    "chosen_genes = list(range(61))\n",
    "total_datasets = [tot_dataset_centered[i][:, chosen_genes] for i in chosen_genes]\n",
    "total_interventions = [[i] for i in range(len(total_datasets))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d03681af",
   "metadata": {},
   "outputs": [],
   "source": [
    "combined = list(zip(total_interventions, total_datasets))\n",
    "random.shuffle(combined)\n",
    "\n",
    "shuffled_interventions, shuffled_datasets = zip(*combined)\n",
    "\n",
    "n_valid_data = 6\n",
    "\n",
    "training_intervention_sets = shuffled_interventions[:-n_valid_data]\n",
    "training_dataset_centered = shuffled_datasets[:-n_valid_data]\n",
    "\n",
    "validation_dataset_centered = shuffled_datasets[-n_valid_data:]\n",
    "validation_intervention_sets = shuffled_interventions[-n_valid_data:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4f5087c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 20/20, Intervention: 55/55, log P(X): -79.218\n"
     ]
    }
   ],
   "source": [
    "causal_mech = gumbelSoftMLP(\n",
    "    n_nodes=61,\n",
    "    lip_constant=0.9,\n",
    "    activation=\"relu\"\n",
    ")\n",
    "\n",
    "g_z = gnet_z(\n",
    "    n_nodes=61\n",
    ")\n",
    "\n",
    "implicit_block = imBlock(\n",
    "    nnet_x=causal_mech, \n",
    "    nnet_z=g_z,\n",
    "    confounders=True\n",
    ")\n",
    "\n",
    "trainer = DCCDTrainer(\n",
    "    implicit_block, \n",
    "    max_epochs=100, \n",
    "    batch_size=1024, \n",
    "    lr=1e-2,\n",
    "    lambda_c=1e-3, \n",
    "    rho=1e-1,\n",
    "    glasso_iters=1\n",
    ")\n",
    "\n",
    "trainer.train(\n",
    "    intervention_datasets=training_dataset_centered,\n",
    "    intervention_targets=training_intervention_sets,\n",
    "    verbose=True,\n",
    "    sleep_time=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "331b861c",
   "metadata": {},
   "outputs": [],
   "source": [
    "genes = pd.read_csv(\"datasets/perturb-cite-seq/genes.csv\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "0c2d742d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 37172 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axs = plt.subplots(1)\n",
    "\n",
    "axs.imshow(implicit_block.nnet_x.get_w_adj().detach().cpu().numpy() > 0.5, cmap=\"bwr\", vmin=-1, vmax=1)\n",
    "\n",
    "axs.set_xticks(range(61))\n",
    "_ = axs.set_yticks(range(61))\n",
    "_ = axs.set_xticklabels(genes[\"Genes\"], rotation=90, fontsize=4.5)\n",
    "_ = axs.set_yticklabels(genes[\"Genes\"], fontsize=4.5)\n",
    "\n",
    "fig.savefig(\"../results/pcs-adj-mat.pdf\", bbox_inches=\"tight\")\n",
    "\n"
   ]
  }
 ],
 "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
}
