{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "23d23ed8-0561-4f77-92cb-49e183bee21a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# start coding here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "106abf26-c714-4afb-a4a3-b2339dc23078",
   "metadata": {
    "editable": true,
    "papermill": {
     "duration": 5.432492,
     "end_time": "2025-07-08T14:12:28.182980",
     "exception": false,
     "start_time": "2025-07-08T14:12:22.750488",
     "status": "completed"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import logomaker\n",
    "import torch\n",
    "import re, os\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import textwrap\n",
    "\n",
    "from Bio import motifs\n",
    "from functools import partial\n",
    "from glob import glob\n",
    "from pyjaspar import jaspardb\n",
    "from sklearn.metrics import precision_recall_fscore_support as score\n",
    "from tqdm import tqdm\n",
    "from matplotlib import gridspec\n",
    "from pybedtools import BedTool\n",
    "from gprofiler import GProfiler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1f795f12-92bb-4511-8d9b-a43de85b2a58",
   "metadata": {
    "editable": true,
    "papermill": {
     "duration": 0.259521,
     "end_time": "2025-07-08T14:12:28.458244",
     "exception": false,
     "start_time": "2025-07-08T14:12:28.198723",
     "status": "completed"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.style as mplstyle\n",
    "mplstyle.use('fast')\n",
    "import matplotlib\n",
    "matplotlib.rc('pdf', fonttype=42)\n",
    "from fastcluster import linkage\n",
    "from scipy.spatial import distance\n",
    "from scipy.spatial.distance import pdist\n",
    "from scipy.cluster import hierarchy\n",
    "from scipy.cluster.hierarchy import fcluster, leaves_list\n",
    "from scipy.stats import zscore\n",
    "from collections import defaultdict\n",
    "from sklearn.cluster import AgglomerativeClustering\n",
    "#from cuml.cluster import AgglomerativeClustering "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cfa2f3ab-87cb-4aa1-a2b4-ef3465a2c6a3",
   "metadata": {
    "papermill": {
     "duration": 0.023997,
     "end_time": "2025-07-08T14:12:28.498762",
     "exception": false,
     "start_time": "2025-07-08T14:12:28.474765",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ecfef380-f9dd-43cb-bb96-0fd260e9d552",
   "metadata": {
    "editable": true,
    "papermill": {
     "duration": 0.025713,
     "end_time": "2025-04-01T16:59:29.537950",
     "exception": false,
     "start_time": "2025-04-01T16:59:29.512237",
     "status": "completed"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def compute_tcav(cav, attr):\n",
    "\n",
    "    tcav = torch.matmul(torch.flatten(attr, start_dim=1), cav.T)\n",
    "    print(cav.shape, attr.shape)\n",
    "    tcav_score = ((tcav>0)*1.).mean(axis=0).cpu().numpy()\n",
    "\n",
    "    return tcav, tcav_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "457d1a36-692b-4ac1-ba2b-3dccbedf9ae3",
   "metadata": {
    "editable": true,
    "papermill": {
     "duration": 0.024606,
     "end_time": "2025-04-01T16:59:29.576178",
     "exception": false,
     "start_time": "2025-04-01T16:59:29.551572",
     "status": "completed"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# parameters\n",
    "cav_root_dir = \"./cavs_test/cavs/\"\n",
    "\n",
    "attrs = \"./cavs_test/test.attributions.pt\"\n",
    "\n",
    "activations_projected = \"./cavs_test/test.activations_projected.pt\"\n",
    "activations_residual = \"./cavs_test/test.activations_residual.pt\"\n",
    "\n",
    "regions = \"./cavs_test/test.regions.bed\"\n",
    "\n",
    "output_dir = \"./cavs_test/tcav_report/\"\n",
    "\n",
    "fscore_threshold:float = 0.95"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "7d246c15-3672-4281-bfa1-64cff98fc5cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "cav_dirs = {i: os.path.join(cav_root_dir, i) for i in os.listdir(cav_root_dir)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "c076f0b6-fcd6-48aa-a4ae-21c9e595acd2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'AC0001:GATA-PROP:GATA_control_random_regions': './cavs_test/cavs/AC0001:GATA-PROP:GATA_control_random_regions',\n",
       " 'AC0002:PROP-ALX:Homeodomain_control_random_regions': './cavs_test/cavs/AC0002:PROP-ALX:Homeodomain_control_random_regions',\n",
       " 'AC0003:HNF1A-HNF1B:Homeodomain_control_random_regions': './cavs_test/cavs/AC0003:HNF1A-HNF1B:Homeodomain_control_random_regions',\n",
       " 'AC0004:ZSCAN:C2H2_ZF_control_random_regions': './cavs_test/cavs/AC0004:ZSCAN:C2H2_ZF_control_random_regions',\n",
       " 'AC0005:POU3F-POU1F:Homeodomain,POU_control_random_regions': './cavs_test/cavs/AC0005:POU3F-POU1F:Homeodomain,POU_control_random_regions'}"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cav_dirs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "da79b725-64bb-445e-8139-5ff0c545d05f",
   "metadata": {
    "papermill": {
     "duration": 0.09174,
     "end_time": "2025-07-08T14:12:28.947321",
     "exception": false,
     "start_time": "2025-07-08T14:12:28.855581",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if regions is not None:\n",
    "    regions = np.loadtxt(regions, dtype=str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "cfdaccc8-3b6f-40a0-ba1f-8fb4f232cbf1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2367435/806406243.py:1: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  avs_projected = torch.load(activations_projected, map_location='cpu')\n",
      "/tmp/ipykernel_2367435/806406243.py:2: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  avs_residual = torch.load(activations_residual, map_location='cpu')\n"
     ]
    }
   ],
   "source": [
    "avs_projected = torch.load(activations_projected, map_location='cpu')\n",
    "avs_residual = torch.load(activations_residual, map_location='cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "87d12372-5d53-4354-b609-e838f12366d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "if avs_projected is not None:\n",
    "    activations = torch.concat([avs_projected, avs_residual], dim=1)\n",
    "else:\n",
    "    activations = avs_residual"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "0676be06-2ffb-4c6d-9ac1-428e8bc81100",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(output_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "9a74b33b-51cf-44d0-aaee-ca9be09fa325",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2367435/777408744.py:1: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  attrs = torch.load(attrs)\n"
     ]
    }
   ],
   "source": [
    "attrs = torch.load(attrs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "b51d5840-321d-42b0-8529-1f4080e79d3d",
   "metadata": {
    "papermill": {
     "duration": 0.027998,
     "end_time": "2025-04-01T16:59:30.388033",
     "exception": false,
     "start_time": "2025-04-01T16:59:30.360035",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "tcav_expts = []\n",
    "for cn, cav_dir in cav_dirs.items():\n",
    "    concept, control = cn.split(\"_control_\")\n",
    "    tcav_expts.append((re.sub(r'control_rep[0-9]_insert_', '', concept), control, cav_dir))\n",
    "tcav_expts = pd.DataFrame(tcav_expts, columns=['concept', 'rep', 'cav_dir'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "fdb86eb8-03af-4ca4-a266-a94e8f7e129a",
   "metadata": {
    "papermill": {
     "duration": 0.04468,
     "end_time": "2025-04-01T16:59:30.446445",
     "exception": false,
     "start_time": "2025-04-01T16:59:30.401765",
     "status": "completed"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>concept</th>\n",
       "      <th>rep</th>\n",
       "      <th>cav_dir</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>AC0001:GATA-PROP:GATA</td>\n",
       "      <td>random_regions</td>\n",
       "      <td>./cavs_test/cavs/AC0001:GATA-PROP:GATA_control...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>AC0002:PROP-ALX:Homeodomain</td>\n",
       "      <td>random_regions</td>\n",
       "      <td>./cavs_test/cavs/AC0002:PROP-ALX:Homeodomain_c...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>AC0003:HNF1A-HNF1B:Homeodomain</td>\n",
       "      <td>random_regions</td>\n",
       "      <td>./cavs_test/cavs/AC0003:HNF1A-HNF1B:Homeodomai...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>AC0004:ZSCAN:C2H2_ZF</td>\n",
       "      <td>random_regions</td>\n",
       "      <td>./cavs_test/cavs/AC0004:ZSCAN:C2H2_ZF_control_...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>AC0005:POU3F-POU1F:Homeodomain,POU</td>\n",
       "      <td>random_regions</td>\n",
       "      <td>./cavs_test/cavs/AC0005:POU3F-POU1F:Homeodomai...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                              concept             rep  \\\n",
       "0               AC0001:GATA-PROP:GATA  random_regions   \n",
       "1         AC0002:PROP-ALX:Homeodomain  random_regions   \n",
       "2      AC0003:HNF1A-HNF1B:Homeodomain  random_regions   \n",
       "3                AC0004:ZSCAN:C2H2_ZF  random_regions   \n",
       "4  AC0005:POU3F-POU1F:Homeodomain,POU  random_regions   \n",
       "\n",
       "                                             cav_dir  \n",
       "0  ./cavs_test/cavs/AC0001:GATA-PROP:GATA_control...  \n",
       "1  ./cavs_test/cavs/AC0002:PROP-ALX:Homeodomain_c...  \n",
       "2  ./cavs_test/cavs/AC0003:HNF1A-HNF1B:Homeodomai...  \n",
       "3  ./cavs_test/cavs/AC0004:ZSCAN:C2H2_ZF_control_...  \n",
       "4  ./cavs_test/cavs/AC0005:POU3F-POU1F:Homeodomai...  "
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tcav_expts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "dc4c32e3-9050-4641-9f5c-439d3fa4f6d1",
   "metadata": {
    "editable": true,
    "papermill": {
     "duration": 1.318247,
     "end_time": "2025-04-01T16:59:31.780317",
     "exception": false,
     "start_time": "2025-04-01T16:59:30.462070",
     "status": "completed"
    },
    "scrolled": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "0it [00:00, ?it/s]/tmp/ipykernel_2367435/1255733344.py:26: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  cav = torch.load(f\"{prefix_dir}/classifier_weights.pt\").detach()\n",
      "/tmp/ipykernel_2367435/1255733344.py:26: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  cav = torch.load(f\"{prefix_dir}/classifier_weights.pt\").detach()\n",
      "/tmp/ipykernel_2367435/1255733344.py:26: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  cav = torch.load(f\"{prefix_dir}/classifier_weights.pt\").detach()\n",
      "/tmp/ipykernel_2367435/1255733344.py:26: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  cav = torch.load(f\"{prefix_dir}/classifier_weights.pt\").detach()\n",
      "/tmp/ipykernel_2367435/1255733344.py:26: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  cav = torch.load(f\"{prefix_dir}/classifier_weights.pt\").detach()\n",
      "5it [00:00, 129.84it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 188]) torch.Size([1000, 188])\n",
      "torch.Size([2, 188]) torch.Size([1000, 188])\n",
      "torch.Size([2, 188]) torch.Size([1000, 188])\n",
      "torch.Size([2, 188]) torch.Size([1000, 188])\n",
      "torch.Size([2, 188]) torch.Size([1000, 188])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "760it [00:02, 260.22it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "787it [00:02, 261.06it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "815it [00:03, 265.68it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "846it [00:03, 275.63it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "860it [00:03, 272.85it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "dfs = []\n",
    "cavs = defaultdict(list)\n",
    "tcavs = defaultdict(list)\n",
    "\n",
    "fscores = []\n",
    "for item in tqdm(tcav_expts.itertuples()):\n",
    "\n",
    "    prefix_dir = item.cav_dir\n",
    "    \n",
    "    perform = pd.read_table(f\"{prefix_dir}/classifier_perform_on_test.txt\", \n",
    "                            comment=\"#\")\n",
    "\n",
    "    precision, recall, fscore, support = score(perform.Truth, perform.Pred)\n",
    "\n",
    "    fscores.append(np.mean(fscore))\n",
    "\n",
    "    # here predictions are model outputs, so it's usually expected that concepts 0 has a higher probability of being bound\n",
    "    #model_predictions_0 = np.loadtxt(f\"{prefix_dir}/predictions_concept_0.txt\")\n",
    "    #model_predictions_1 = np.loadtxt(f\"{prefix_dir}/predictions_concept_1.txt\")\n",
    "    \n",
    "    classifier_scores = pd.DataFrame({'precision': precision,\n",
    "                  'recall': recall,\n",
    "                  'fscore': fscore,\n",
    "                  'support': support})\n",
    "\n",
    "    cav = torch.load(f\"{prefix_dir}/classifier_weights.pt\").detach()\n",
    "    tcav, tcav_score = compute_tcav(cav, attrs)\n",
    "    cavs[item.concept].append(cav)\n",
    "    tcavs[item.concept].append(tcav)\n",
    "\n",
    "    dfs.append((np.mean(fscore), support, np.median(tcav.numpy(), axis=0)[0], tcav_score[0], tcav_score[1], item.concept, item.rep))\n",
    "\n",
    "    # dfs.append((fscore[0], support, tcav_score[0], np.median(tcav.numpy(), axis=0)[0], \n",
    "    #             cav, tcav,\n",
    "    #             #model_predictions_0.mean(), \n",
    "    #             item.concept, item.rep, item.layer, 0))\n",
    "    # dfs.append((fscore[1], support, tcav_score[1], np.median(tcav.numpy(), axis=0)[1],\n",
    "    #             cav, tcav,\n",
    "    #             #model_predictions_1.mean(), \n",
    "    #             item.concept, item.rep, item.layer, 1))\n",
    "    \n",
    "# dfs_merge = pd.DataFrame(dfs, columns=['fscore', 'support', 'tcav_score', 'tcav_median', \n",
    "#                                        \"cav\", \"tcav\",\n",
    "#                                        #'model_preds',\n",
    "#                                        'concept', 'rep', 'layer', 'concept_id'])\n",
    "dfs_merge = pd.DataFrame(dfs, columns=['fscore', 'support', 'tcav_score_median_test', 'tcav_pos_ratio_test', 'tcav_pos_ratio_control', \n",
    "                                       'concept', 'rep'])\n",
    "\n",
    "#dfs_merge['concept_name'] = dfs_merge.apply(lambda x: f'Random genome with\\n{x[\"concept\"]} inserted' if x['concept_id']==0 else 'Random genome', axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "4c46072b-f17c-4186-bb86-ea385aac1dbc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Axes: ylabel='Count'>"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 12812 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sns.histplot(x=fscores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "c965f052-8510-4751-ab16-b296678e3f61",
   "metadata": {},
   "outputs": [],
   "source": [
    "dfs_merge\n",
    "dfs_merge.to_csv(f\"{output_dir}/tcav_all_concepts.tsv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "d34fabef-b521-4a95-8a76-cc3176610089",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>fscore</th>\n",
       "      <th>support</th>\n",
       "      <th>tcav_score_median_test</th>\n",
       "      <th>tcav_pos_ratio_test</th>\n",
       "      <th>tcav_pos_ratio_control</th>\n",
       "      <th>concept</th>\n",
       "      <th>rep</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0.514597</td>\n",
       "      <td>[461, 439]</td>\n",
       "      <td>1.188891</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>AC0003:HNF1A-HNF1B:Homeodomain</td>\n",
       "      <td>random_regions</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0.556389</td>\n",
       "      <td>[461, 434]</td>\n",
       "      <td>0.294493</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>AC0004:ZSCAN:C2H2_ZF</td>\n",
       "      <td>random_regions</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     fscore     support  tcav_score_median_test  tcav_pos_ratio_test  \\\n",
       "2  0.514597  [461, 439]                1.188891                  1.0   \n",
       "3  0.556389  [461, 434]                0.294493                  1.0   \n",
       "\n",
       "   tcav_pos_ratio_control                         concept             rep  \n",
       "2                     0.0  AC0003:HNF1A-HNF1B:Homeodomain  random_regions  \n",
       "3                     0.0            AC0004:ZSCAN:C2H2_ZF  random_regions  "
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dfs_merge = dfs_merge.loc[dfs_merge.fscore>fscore_threshold]\n",
    "dfs_merge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "d8cfc13b-fdcf-4496-881b-c97a7acb45b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "motif_clustering_pwms = motifs.parse(open('data/motif-clustering-v2.1beta_consensus_pwms.meme'), fmt='MINIMAL')\n",
    "motif_clustering_pwms = {m.name.replace('/', '-'): m.pwm for m in motif_clustering_pwms}\n",
    "\n",
    "if len(dfs_merge) > 0:\n",
    "    dfs_merge['motif_cluster'] = dfs_merge.concept.apply(lambda x: x if x.startswith('AC') else None).str.replace(r'_rep[0-9]', '', regex=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bed7a289-3d63-4a8f-a316-89f439012e45",
   "metadata": {
    "papermill": {
     "duration": 0.027845,
     "end_time": "2025-04-01T16:59:47.005057",
     "exception": false,
     "start_time": "2025-04-01T16:59:46.977212",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Log ratio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "e76aa30e-43b9-414f-8146-90923efb6a3f",
   "metadata": {
    "editable": true,
    "papermill": {
     "duration": 0.526734,
     "end_time": "2025-04-01T16:59:47.559457",
     "exception": true,
     "start_time": "2025-04-01T16:59:47.032723",
     "status": "failed"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def log_ratio(x):\n",
    "    # the upper and lower bound should be ~5\n",
    "    return np.log((x['tcav_pos_ratio_test']+0.01)/(x['tcav_pos_ratio_control']+0.01))\n",
    "    \n",
    "if len(dfs_merge) > 0:\n",
    "    dfs_merge['log_ratio'] = dfs_merge.apply(log_ratio, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ae4dd89-6cf5-4ba4-b1c6-dd295446f475",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "source": [
    "### CAV correlation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "b8c502ef-5d98-4935-8ef2-625759841d46",
   "metadata": {
    "papermill": {
     "duration": 0.051424,
     "end_time": "2025-07-08T14:18:29.597858",
     "exception": false,
     "start_time": "2025-07-08T14:18:29.546434",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "motif_clustering_pwms = motifs.parse(open('data/motif-clustering-v2.1beta_consensus_pwms.meme'), fmt='MINIMAL')\n",
    "motif_clustering_pwms = {re.findall(r'AC[0-9]*', m.name)[0]: m.pwm for m in motif_clustering_pwms} # index pwm by AC id\n",
    "def seq_logo(cav_key, ax=None, max_len=20):\n",
    "    ac_id = re.findall('AC[0-9]*', cav_key)\n",
    "    if len(ac_id) == 1:\n",
    "        pwm = motif_clustering_pwms.get(ac_id[0])\n",
    "    \n",
    "        pwm_df = pd.DataFrame(pwm)\n",
    "        motif_len = len(pwm_df)\n",
    "        \n",
    "        # Logomaker expects columns as A,C,G,T, so ensure correct order\n",
    "        pwm_df = pwm_df[['A', 'C', 'G', 'T']]\n",
    "    \n",
    "        # Compute information content at each position (2 - entropy)\n",
    "        def compute_ic(pwm_row):\n",
    "            entropy = -sum([p * np.log2(p) if p > 0 else 0 for p in pwm_row])\n",
    "            return 2 - entropy\n",
    "        \n",
    "        # Compute IC matrix: IC_letter = p * IC_total\n",
    "        ic_df = pwm_df.copy()\n",
    "        for i in range(len(pwm_df)):\n",
    "            ic_total = compute_ic(pwm_df.iloc[i])\n",
    "            ic_df.iloc[i] = pwm_df.iloc[i] * ic_total\n",
    "    \n",
    "        \n",
    "        # Plot with logomaker\n",
    "        x0, y0, width, height = ax.get_position().bounds\n",
    "        ax.set_position([x0, y0, width * min(1., motif_len/max_len), height])\n",
    "        logo = logomaker.Logo(ic_df, color_scheme={'A': 'red', 'C': 'blue',\n",
    "                                                  'G': 'orange', 'T': 'green'},\n",
    "                              ax=ax)\n",
    "        logo.ax.axis('off')\n",
    "    else:\n",
    "        ax.axis('off')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "9b08abf5-4d72-47a5-86b8-28659305d7da",
   "metadata": {
    "editable": true,
    "papermill": {
     "duration": 21.95088,
     "end_time": "2025-07-08T14:19:15.506164",
     "exception": false,
     "start_time": "2025-07-08T14:18:53.555284",
     "status": "completed"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Figure size 1300x800 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 82888 characters]",
      "text/plain": [
       "<Figure size 1000x1000 with 8 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "clustermap_width = 0.5\n",
    "barplot_width = 0.2\n",
    "logo_width = 0.3\n",
    "\n",
    "if len(dfs_merge) > 0:\n",
    "    cavs_stack = {}\n",
    "    #pwms_toplot = []\n",
    "    colors_toplot = []\n",
    "    log_ratios = []\n",
    "    for concept, cav_list in cavs.items():\n",
    "        if concept in dfs_merge.concept.tolist():\n",
    "            motif_cluster = dfs_merge.loc[dfs_merge.concept==concept, 'motif_cluster'].item()\n",
    "\n",
    "            cavs_key = concept\n",
    "            cavs_stack[cavs_key] = torch.stack(cav_list).mean(dim=0)[0]\n",
    "            #pwms_toplot.append(motif_clustering_pwms.get(dfs_merge.loc[dfs_merge.concept==concept, 'motif_cluster'].item()))\n",
    "            colors_toplot.append('red' if dfs_merge.loc[dfs_merge.concept==concept, 'tcav_pos_ratio_test'].item()>0.5 else 'blue')\n",
    "            log_ratios.append(dfs_merge.loc[dfs_merge.concept==concept, 'log_ratio'].item())\n",
    "    cavs_keys = list(cavs_stack.keys())\n",
    "    cavs_stack = torch.stack(list(cavs_stack.values()))\n",
    "    \n",
    "    from sklearn.metrics.pairwise import cosine_similarity\n",
    "    if cavs_stack.shape[0] > 1:\n",
    "\n",
    "        fig = plt.figure(figsize=(13, 8))\n",
    "        sfigs = fig.subfigures(2, 1)\n",
    "        \n",
    "        plt.rc('xtick', labelsize=50/np.log(len(dfs_merge)))\n",
    "        plt.rc('ytick', labelsize=50/np.log(len(dfs_merge)))\n",
    "        \n",
    "        matrix_similarity = cosine_similarity(cavs_stack.detach().numpy())\n",
    "        cm = sns.clustermap(matrix_similarity, xticklabels=False, yticklabels=False, cmap='bwr',\n",
    "                       vmin=-1, vmax=1)\n",
    "        labels = fcluster(cm.dendrogram_row.linkage, 10, criterion='maxclust')\n",
    "        labels_toplot = np.array([labels[i] for i in cm.dendrogram_row.reordered_ind])\n",
    "        \n",
    "        colors_toplot = [colors_toplot[i] for i in cm.dendrogram_col.reordered_ind]\n",
    "        for c, text in zip(colors_toplot, cm.ax_heatmap.get_yticklabels()):\n",
    "            text.set_color(c)\n",
    "\n",
    "        cm.gs.update(left=0.05, right=clustermap_width)\n",
    "        cm.ax_cbar.set_position([.01, 0.9, .05, .05])\n",
    "\n",
    "        cavs_keys_reordered = [cavs_keys[i] for i in cm.dendrogram_col.reordered_ind]\n",
    "\n",
    "        #create new gridspec for the log ratio barplot\n",
    "        gs2 = gridspec.GridSpec(1, 1)\n",
    "        # create axes within this new gridspec\n",
    "        ax2 = cm.fig.add_subplot(gs2[0])\n",
    "\n",
    "        heatmap_bbox = cm.ax_heatmap.get_position()\n",
    "        ax2.set_position([0.5, heatmap_bbox.y0, barplot_width, heatmap_bbox.height])\n",
    "\n",
    "        # plot log ratio\n",
    "        ax2.tick_params(axis='y', which='major', pad=fig.get_size_inches()[0] * logo_width * 72)\n",
    "        log_ratios_reordered = [log_ratios[i] for i in cm.dendrogram_col.reordered_ind]\n",
    "        sns.barplot(y=cavs_keys_reordered, x=log_ratios_reordered, orient=\"y\", ax=ax2)\n",
    "        # set color of bar by value\n",
    "        for idx in range(len((ax2.containers[0]))):\n",
    "            if ax2.containers[0].datavalues[idx] > 0:\n",
    "                ax2.containers[0][idx].set_color('red')\n",
    "            else:\n",
    "                ax2.containers[0][idx].set_color('blue')\n",
    "        \n",
    "        ax2.set_xlim(left=-5, right=5)\n",
    "        ax2.yaxis.tick_right()\n",
    "        ax2.set_title('TCAV log ratio')\n",
    "\n",
    "        gs3 = gridspec.GridSpec(len(cavs_keys_reordered), 1)\n",
    "\n",
    "        logo_height = heatmap_bbox.height/len(cavs_keys_reordered)\n",
    "        for i, (cav_key, g) in enumerate(zip(reversed(cavs_keys_reordered), gs3)):\n",
    "            ax3 = plt.subplot(g)\n",
    "            ax3.set_position([0.71, heatmap_bbox.y0+i*logo_height, logo_width+0.01, logo_height])\n",
    "            if cav_key is not None:\n",
    "                seq_logo(cav_key, ax=ax3)\n",
    "            else:\n",
    "                ax3.axis('off')\n",
    "\n",
    "        gs4 = gridspec.GridSpec(1, 1)\n",
    "        # create axes within this new gridspec\n",
    "        ax4 = cm.fig.add_subplot(gs4[0])\n",
    "        ax4.imshow(labels_toplot.reshape(-1, 1), cmap=\"tab10\", aspect='auto')\n",
    "        ax4.axis('off')\n",
    "\n",
    "        ax4.set_position([-0.05, heatmap_bbox.y0, 0.05, heatmap_bbox.height])\n",
    "        \n",
    "        plt.savefig(f\"{output_dir}/tcav_cav_correlation_full.pdf\", bbox_inches=\"tight\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "34707485-04fd-4602-b7f4-1ca97dacbbe9",
   "metadata": {
    "papermill": {
     "duration": 32.652877,
     "end_time": "2025-07-08T14:19:48.276459",
     "exception": false,
     "start_time": "2025-07-08T14:19:15.623582",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1000,)\n",
      "990 regions are positively associated with cluster 1\n"
     ]
    }
   ],
   "source": [
    "if len(dfs_merge) > 0 and regions is not None:\n",
    "    # reuse the stacked cavs\n",
    "    if cavs_stack.shape[0] > 1:\n",
    "        # compute cosine similarity between the activations and cavs\n",
    "        ## reorder according to previous clustermap dendrogram order\n",
    "        cavs_stack_reordered = torch.stack([cavs_stack[i] for i in cm.dendrogram_col.reordered_ind])\n",
    "        ## compute cosine similarity\n",
    "        similarity_avs_cavs = cosine_similarity(activations, cavs_stack_reordered)\n",
    "\n",
    "        posrs = {}\n",
    "        for l in np.unique(labels_toplot):\n",
    "            # subset concepts in the cluster\n",
    "            posr = []\n",
    "            #print(f\"subsample columns: \", np.where(labels_toplot==l)[0])\n",
    "            ss = similarity_avs_cavs[:, np.where(labels_toplot==l)[0]].mean(axis=1) # take avg of cosine similarity among selected concepts\n",
    "            print(ss.shape)\n",
    "            c = 0\n",
    "            ks = [cavs_keys_reordered[i] for i in np.where(labels_toplot==l)[0]]\n",
    "            for idx in np.argsort(ss):\n",
    "                if ss[idx] > 0: # take regions that are positive\n",
    "                    posr.append(regions[idx])\n",
    "                    c+=1\n",
    "            if len(posr) > 0:\n",
    "                bed_df = []\n",
    "                with open(f\"{output_dir}/cluster_{l}_positive_cavs_x_avs.bed\", 'w') as f:\n",
    "                    f.write(f\"#{ks}\")\n",
    "                    for r in posr:\n",
    "                        chrom, start, end = r\n",
    "                        f.write(f\"{chrom}\\t{start}\\t{end}\\n\")\n",
    "                        bed_df.append((chrom, start, end, l))\n",
    "\n",
    "                posrs[l.item()] = pd.DataFrame(bed_df, columns=['chrom', 'start', 'end', 'label'])\n",
    "                print(f\"{len(bed_df)} regions are positively associated with cluster {l}\")\n",
    "            else:\n",
    "                print(f\"No regions are positively associated with cluster {l}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch_latest",
   "language": "python",
   "name": "torch_latest"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
