{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "24d0ff5f",
   "metadata": {},
   "source": [
    "## Construct Gene Regulatory Network\n",
    "This module utilizes layers/heads self attention of genes to determine the importance of gene relationships to each other. This can be utilized to create a directed graph for genes that \"pay attention to each other\" in pretrained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ec416b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from stFormer.network_dynamics.gene_regulatory_network import GeneRegulatoryNetwork"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f218c55",
   "metadata": {},
   "source": [
    "1. Model & Tokenizer: Load pre-trained BERT with output_attention=True to extract per layer/per-head attention scores\n",
    "2.  pad all sequences out to the maximum length in the filtered set (so they form a rectangular tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba75b601",
   "metadata": {},
   "outputs": [],
   "source": [
    "grn = GeneRegulatoryNetwork(\n",
    "    model_dir = 'output/spot/models/250422_102707_stFormer_L6_E3/final',\n",
    "    dataset_path = 'output/spot/visium_spot.dataset',\n",
    "    model_type = 'Pretrained',\n",
    "    metadata_column = 'subtype',\n",
    "    metadata_value = 'TNBC',\n",
    "    device='cuda',\n",
    "    batch_size=24,\n",
    "    nproc = 12\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8978ff29",
   "metadata": {},
   "source": [
    "3. Compute Attention:\n",
    "    - Stack & average over heads and layers → a (batch, seq_len, seq_len) tensor.\n",
    "    - Sum these per-batch attention matrices into one running attn_sum\n",
    "    - Finalize by dividing attn_sum by total examples → your final attention_matrix (NumPy array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d3e30c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "grn.compute_attention()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f05da24",
   "metadata": {},
   "source": [
    "4. Build Graph:\n",
    "    - Build a “counts” matrix C of shape (T, seq_len), where C[i,p] = how many times token i appears at position \n",
    "    - Compute N = C @ attention_matrix @ C.T (a single high-performance mat-mul) and normalize by the outer product of occurrence counts to get the average attention from token i to token j.\n",
    "    - Edge Selection: cutoff, percentile, top_k\n",
    "    - add a directed edge for every selected token pair with aggregated weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d44ba755",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Can take a long time, computing average i,j attention across all samples for all token pairs\n",
    "grn.build_graph(\n",
    "    percentile = 0.99999, #filter node-edges by attention weight (value above percentile)\n",
    "    min_cooccurrence=100 #filter node-edges by number of samples expressing pair\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c46ddbc4",
   "metadata": {},
   "source": [
    "write to file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1af9dc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "grn.plot_network('output/spot/gene_network.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcc957cc",
   "metadata": {},
   "source": [
    "5. Transcription Factor Analysis: Future Updates"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1f059da",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "stFormer",
   "language": "python",
   "name": "stformer"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}