freq_graph_gcn.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import networkx as nx\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch_geometric.utils.convert import from_networkx\n",
    "import random\n",
    "import networkx as nx\n",
    "from collections import Counter, defaultdict\n",
    "from torch_geometric.loader import DataLoader\n",
    "from torch.nn import Linear\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import GCNConv\n",
    "from torch_geometric.nn import global_mean_pool\n",
    "import src.utils.utils as ut\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Functions for generating sequences and mutations\n",
    "def generate_random_genome_sequence(length):\n",
    "    return ''.join(random.choice('ATCG') for _ in range(length))\n",
    "\n",
    "def mutate_sequence(sequence, num_mutations):\n",
    "    sequence = list(sequence)\n",
    "    length = len(sequence)\n",
    "    for _ in range(num_mutations):\n",
    "        pos = random.randint(0, length - 1)\n",
    "        original_base = sequence[pos]\n",
    "        new_base = random.choice([b for b in 'ATCG' if b != original_base])\n",
    "        sequence[pos] = new_base\n",
    "    return ''.join(sequence)\n",
    "\n",
    "def generate_genome_sequences_with_mutations(n, l, k, num_mutations):\n",
    "    genome_data = []\n",
    "    for i in range(n):\n",
    "        original_sequence = generate_random_genome_sequence(l)\n",
    "        genome_data.append([original_sequence, i])\n",
    "        \n",
    "        for j in range(k):\n",
    "            mutated_sequence = mutate_sequence(original_sequence, num_mutations)\n",
    "            genome_data.append([mutated_sequence, i])\n",
    "    \n",
    "    return genome_data\n",
    "\n",
    "# Function to generate k-mers from a sequence\n",
    "def generate_kmers(sequence, k):\n",
    "    return [sequence[i:i+k] for i in range(len(sequence) - k + 1)]\n",
    "\n",
    "def kmer_to_index(kmer):\n",
    "    \"\"\"Converts a kmer (string) to an index.\"\"\"\n",
    "    base_to_index = {'A': 0, 'C': 1, 'G': 2, 'T': 3}\n",
    "    index = 0\n",
    "    for char in kmer:\n",
    "        index = 4 * index + base_to_index[char]\n",
    "    return index\n",
    "\n",
    "def subkmer_frequencies_in_kmer(kmer, subkmer_length):\n",
    "    \"\"\"Calculates the frequency of each subkmer in a kmer.\"\"\"\n",
    "    subkmer_counts = Counter(kmer[i:i+subkmer_length] for i in range(len(kmer) - subkmer_length + 1))\n",
    "    frequencies = np.zeros(4**subkmer_length)\n",
    "    for subkmer, count in subkmer_counts.items():\n",
    "        index = kmer_to_index(subkmer)\n",
    "        frequencies[index] = count\n",
    "    return frequencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example usage\n",
    "num_classes = 5  # number of original sequences\n",
    "l = 100  # length of each sequence\n",
    "k = 99\n",
    "num_mutations = 20\n",
    "genome_sequences = generate_genome_sequences_with_mutations(num_classes, l, k, num_mutations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "seqs_n = len(genome_sequences)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "graphs = []\n",
    "kmer_len = 4\n",
    "subkmer_len = 2\n",
    "num_features = 4**subkmer_len\n",
    "\n",
    "for seq in genome_sequences:\n",
    "    G = nx.DiGraph()\n",
    "    kmers = generate_kmers(seq[0], kmer_len)\n",
    "    nodes = []\n",
    "    for kmer in kmers:\n",
    "        nodes.append((kmer, {\"x\": torch.as_tensor(subkmer_frequencies_in_kmer(kmer, subkmer_len)/(kmer_len-1), dtype=torch.float32)}))\n",
    "    G.add_nodes_from(nodes)\n",
    "    # edges = []\n",
    "    transition_counts = defaultdict(int)\n",
    "    for i in range(len(kmers)-1):\n",
    "        transition_counts[(kmers[i], kmers[i+1])] += 1\n",
    "    max_count = max(transition_counts.values())\n",
    "    for key in transition_counts.keys():\n",
    "        G.add_edge(key[0], key[1], weight=transition_counts[key]/max_count)\n",
    "    #     edges.append((kmers[i], kmers[i+1]))\n",
    "    # G.add_edges_from(edges)\n",
    "    torch_graph = from_networkx(G)\n",
    "    torch_graph['y'] = torch.tensor([seq[1]])\n",
    "    graphs.append(torch_graph)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "random.shuffle(graphs)\n",
    "split_value = 0.8\n",
    "train_dataset = graphs[:int(seqs_n*split_value)]\n",
    "val_dataset = graphs[int(seqs_n*split_value):]\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
    "val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GCN(torch.nn.Module):\n",
    "    def __init__(self, num_features, hidden_channels, num_classes):\n",
    "        super(GCN, self).__init__()\n",
    "        torch.manual_seed(12345)\n",
    "        self.conv1 = GCNConv(num_features, hidden_channels)\n",
    "        self.conv2 = GCNConv(hidden_channels, hidden_channels)\n",
    "        self.conv3 = GCNConv(hidden_channels, hidden_channels)\n",
    "        self.lin = Linear(hidden_channels, num_classes)\n",
    "\n",
    "    def forward(self, x, edge_index, edge_weight, batch):\n",
    "        # 1. Obtain node embeddings \n",
    "        x = self.conv1(x, edge_index, edge_weight)\n",
    "        x = x.relu()\n",
    "        x = self.conv2(x, edge_index, edge_weight)\n",
    "        x = x.relu()\n",
    "        x = self.conv3(x, edge_index, edge_weight)\n",
    "\n",
    "        # 2. Readout layer\n",
    "        x = global_mean_pool(x, batch)  # [batch_size, hidden_channels]\n",
    "\n",
    "        # 3. Apply a final classifier\n",
    "        x = F.dropout(x, p=0.5, training=self.training)\n",
    "        x = self.lin(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Train Loss: 1.6111, Train Acc: 0.1900, Val Acc: 0.1300\n",
      "Epoch 2, Train Loss: 1.6000, Train Acc: 0.2125, Val Acc: 0.3600\n",
      "Epoch 3, Train Loss: 1.5811, Train Acc: 0.2425, Val Acc: 0.4000\n",
      "Epoch 4, Train Loss: 1.5305, Train Acc: 0.3350, Val Acc: 0.4000\n",
      "Epoch 5, Train Loss: 1.4432, Train Acc: 0.4075, Val Acc: 0.4800\n",
      "Epoch 6, Train Loss: 1.2616, Train Acc: 0.4750, Val Acc: 0.6000\n",
      "Epoch 7, Train Loss: 1.1134, Train Acc: 0.6275, Val Acc: 0.5700\n",
      "Epoch 8, Train Loss: 0.9425, Train Acc: 0.6450, Val Acc: 0.6500\n",
      "Epoch 9, Train Loss: 0.8764, Train Acc: 0.6150, Val Acc: 0.7000\n",
      "Epoch 10, Train Loss: 0.7994, Train Acc: 0.7100, Val Acc: 0.8100\n",
      "Epoch 11, Train Loss: 0.6349, Train Acc: 0.7350, Val Acc: 0.7300\n",
      "Epoch 12, Train Loss: 0.5668, Train Acc: 0.7300, Val Acc: 0.8100\n",
      "Epoch 13, Train Loss: 0.5513, Train Acc: 0.7550, Val Acc: 0.8000\n",
      "Epoch 14, Train Loss: 0.5280, Train Acc: 0.7225, Val Acc: 0.7700\n",
      "Epoch 15, Train Loss: 0.4632, Train Acc: 0.7600, Val Acc: 0.8600\n",
      "Epoch 16, Train Loss: 0.4363, Train Acc: 0.7900, Val Acc: 0.7800\n",
      "Epoch 17, Train Loss: 0.4616, Train Acc: 0.7675, Val Acc: 0.7700\n",
      "Epoch 18, Train Loss: 0.4748, Train Acc: 0.7750, Val Acc: 0.8300\n",
      "Epoch 19, Train Loss: 0.4077, Train Acc: 0.8025, Val Acc: 0.8300\n",
      "Epoch 20, Train Loss: 0.4784, Train Acc: 0.7625, Val Acc: 0.8200\n",
      "Epoch 21, Train Loss: 0.4277, Train Acc: 0.7875, Val Acc: 0.8300\n",
      "Epoch 22, Train Loss: 0.4291, Train Acc: 0.7825, Val Acc: 0.8400\n",
      "Epoch 23, Train Loss: 0.3751, Train Acc: 0.8025, Val Acc: 0.8200\n",
      "Epoch 24, Train Loss: 0.3569, Train Acc: 0.8250, Val Acc: 0.8600\n",
      "Epoch 25, Train Loss: 0.3760, Train Acc: 0.8400, Val Acc: 0.8400\n",
      "Epoch 26, Train Loss: 0.3887, Train Acc: 0.8125, Val Acc: 0.8500\n",
      "Epoch 27, Train Loss: 0.3644, Train Acc: 0.8100, Val Acc: 0.8000\n",
      "Epoch 28, Train Loss: 0.3578, Train Acc: 0.8350, Val Acc: 0.8700\n",
      "Epoch 29, Train Loss: 0.3279, Train Acc: 0.8550, Val Acc: 0.8800\n",
      "Epoch 30, Train Loss: 0.3289, Train Acc: 0.8525, Val Acc: 0.8800\n",
      "Epoch 31, Train Loss: 0.3648, Train Acc: 0.8300, Val Acc: 0.8700\n",
      "Epoch 32, Train Loss: 0.3208, Train Acc: 0.8550, Val Acc: 0.8800\n",
      "Epoch 33, Train Loss: 0.2754, Train Acc: 0.8725, Val Acc: 0.8900\n",
      "Epoch 34, Train Loss: 0.2963, Train Acc: 0.8750, Val Acc: 0.8800\n",
      "Epoch 35, Train Loss: 0.2926, Train Acc: 0.8750, Val Acc: 0.8400\n",
      "Epoch 36, Train Loss: 0.2976, Train Acc: 0.8825, Val Acc: 0.8400\n",
      "Epoch 37, Train Loss: 0.3087, Train Acc: 0.8775, Val Acc: 0.9000\n",
      "Epoch 38, Train Loss: 0.3013, Train Acc: 0.8675, Val Acc: 0.8900\n",
      "Epoch 39, Train Loss: 0.2705, Train Acc: 0.8775, Val Acc: 0.8400\n",
      "Epoch 40, Train Loss: 0.3017, Train Acc: 0.8750, Val Acc: 0.9100\n",
      "Epoch 41, Train Loss: 0.2241, Train Acc: 0.9225, Val Acc: 0.9200\n",
      "Epoch 42, Train Loss: 0.1969, Train Acc: 0.9300, Val Acc: 0.9400\n",
      "Epoch 43, Train Loss: 0.1629, Train Acc: 0.9400, Val Acc: 0.9500\n",
      "Epoch 44, Train Loss: 0.1677, Train Acc: 0.9350, Val Acc: 0.9300\n",
      "Epoch 45, Train Loss: 0.1647, Train Acc: 0.9525, Val Acc: 0.9700\n",
      "Epoch 46, Train Loss: 0.1784, Train Acc: 0.9250, Val Acc: 0.9800\n",
      "Epoch 47, Train Loss: 0.1396, Train Acc: 0.9525, Val Acc: 0.9700\n",
      "Epoch 48, Train Loss: 0.1276, Train Acc: 0.9550, Val Acc: 0.9500\n",
      "Epoch 49, Train Loss: 0.1276, Train Acc: 0.9625, Val Acc: 1.0000\n",
      "Epoch 50, Train Loss: 0.1350, Train Acc: 0.9525, Val Acc: 0.9500\n",
      "Epoch 51, Train Loss: 0.1324, Train Acc: 0.9550, Val Acc: 0.9800\n",
      "Epoch 52, Train Loss: 0.1251, Train Acc: 0.9500, Val Acc: 0.9600\n",
      "Epoch 53, Train Loss: 0.1482, Train Acc: 0.9425, Val Acc: 1.0000\n",
      "Epoch 54, Train Loss: 0.1246, Train Acc: 0.9500, Val Acc: 0.9500\n",
      "Epoch 55, Train Loss: 0.1439, Train Acc: 0.9500, Val Acc: 0.9800\n",
      "Epoch 56, Train Loss: 0.0954, Train Acc: 0.9750, Val Acc: 0.9800\n",
      "Epoch 57, Train Loss: 0.1191, Train Acc: 0.9600, Val Acc: 0.9900\n",
      "Epoch 58, Train Loss: 0.1112, Train Acc: 0.9700, Val Acc: 0.9800\n",
      "Epoch 59, Train Loss: 0.1256, Train Acc: 0.9475, Val Acc: 0.9900\n",
      "Epoch 60, Train Loss: 0.1064, Train Acc: 0.9575, Val Acc: 0.8400\n",
      "Epoch 61, Train Loss: 0.1769, Train Acc: 0.9400, Val Acc: 0.9800\n",
      "Epoch 62, Train Loss: 0.1023, Train Acc: 0.9625, Val Acc: 0.9800\n",
      "Epoch 63, Train Loss: 0.0791, Train Acc: 0.9700, Val Acc: 0.9900\n",
      "Epoch 64, Train Loss: 0.1268, Train Acc: 0.9600, Val Acc: 0.9800\n",
      "Epoch 65, Train Loss: 0.1277, Train Acc: 0.9450, Val Acc: 0.9200\n",
      "Epoch 66, Train Loss: 0.1314, Train Acc: 0.9525, Val Acc: 1.0000\n",
      "Epoch 67, Train Loss: 0.1055, Train Acc: 0.9575, Val Acc: 0.9700\n",
      "Epoch 68, Train Loss: 0.0994, Train Acc: 0.9625, Val Acc: 0.9600\n",
      "Epoch 69, Train Loss: 0.1530, Train Acc: 0.9400, Val Acc: 0.9800\n",
      "Epoch 70, Train Loss: 0.1509, Train Acc: 0.9425, Val Acc: 0.9200\n",
      "Epoch 71, Train Loss: 0.1925, Train Acc: 0.9450, Val Acc: 0.9800\n",
      "Epoch 72, Train Loss: 0.1390, Train Acc: 0.9275, Val Acc: 0.8900\n",
      "Epoch 73, Train Loss: 0.1578, Train Acc: 0.9375, Val Acc: 0.9400\n",
      "Epoch 74, Train Loss: 0.1162, Train Acc: 0.9600, Val Acc: 0.9400\n",
      "Epoch 75, Train Loss: 0.0764, Train Acc: 0.9750, Val Acc: 0.9900\n",
      "Epoch 76, Train Loss: 0.0739, Train Acc: 0.9775, Val Acc: 0.9400\n",
      "Epoch 77, Train Loss: 0.1217, Train Acc: 0.9500, Val Acc: 0.9800\n",
      "Epoch 78, Train Loss: 0.0789, Train Acc: 0.9750, Val Acc: 0.9800\n",
      "Epoch 79, Train Loss: 0.0733, Train Acc: 0.9725, Val Acc: 0.9900\n",
      "Epoch 80, Train Loss: 0.0690, Train Acc: 0.9775, Val Acc: 0.9800\n",
      "Epoch 81, Train Loss: 0.0901, Train Acc: 0.9750, Val Acc: 0.9900\n",
      "Epoch 82, Train Loss: 0.0619, Train Acc: 0.9725, Val Acc: 1.0000\n",
      "Epoch 83, Train Loss: 0.0590, Train Acc: 0.9750, Val Acc: 1.0000\n",
      "Epoch 84, Train Loss: 0.0411, Train Acc: 0.9925, Val Acc: 1.0000\n",
      "Epoch 85, Train Loss: 0.0523, Train Acc: 0.9825, Val Acc: 0.9900\n",
      "Epoch 86, Train Loss: 0.0471, Train Acc: 0.9825, Val Acc: 0.9900\n",
      "Epoch 87, Train Loss: 0.0413, Train Acc: 0.9875, Val Acc: 0.9800\n",
      "Epoch 88, Train Loss: 0.0309, Train Acc: 0.9950, Val Acc: 1.0000\n",
      "Epoch 89, Train Loss: 0.0288, Train Acc: 0.9975, Val Acc: 0.9800\n",
      "Epoch 90, Train Loss: 0.0384, Train Acc: 0.9900, Val Acc: 1.0000\n",
      "Epoch 91, Train Loss: 0.0437, Train Acc: 0.9875, Val Acc: 0.9700\n",
      "Epoch 92, Train Loss: 0.0588, Train Acc: 0.9725, Val Acc: 0.9900\n",
      "Epoch 93, Train Loss: 0.0294, Train Acc: 0.9925, Val Acc: 0.9800\n",
      "Epoch 94, Train Loss: 0.0623, Train Acc: 0.9800, Val Acc: 0.9800\n",
      "Epoch 95, Train Loss: 0.0406, Train Acc: 0.9875, Val Acc: 0.9800\n",
      "Epoch 96, Train Loss: 0.0326, Train Acc: 0.9875, Val Acc: 1.0000\n",
      "Epoch 97, Train Loss: 0.0325, Train Acc: 0.9900, Val Acc: 0.9800\n",
      "Epoch 98, Train Loss: 0.0344, Train Acc: 0.9900, Val Acc: 1.0000\n",
      "Epoch 99, Train Loss: 0.0292, Train Acc: 0.9925, Val Acc: 1.0000\n",
      "Epoch 100, Train Loss: 0.0304, Train Acc: 0.9975, Val Acc: 1.0000\n"
     ]
    }
   ],
   "source": [
    "def train():\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    correct = 0\n",
    "    for data in train_loader:\n",
    "        optimizer.zero_grad()\n",
    "        out = model(data.x, data.edge_index, data.weight, data.batch)\n",
    "        loss = criterion(out, data.y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss += loss.item()\n",
    "        pred = out.argmax(dim=1)\n",
    "        correct += (pred == data.y).sum().item()\n",
    "    return total_loss / len(train_loader), correct / len(train_dataset)\n",
    "\n",
    "# Validation function\n",
    "def validate(loader):\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    for data in loader:\n",
    "        with torch.no_grad():\n",
    "            out = model(data.x, data.edge_index, data.weight, data.batch)\n",
    "            pred = out.argmax(dim=1)\n",
    "            correct += (pred == data.y).sum().item()\n",
    "    return correct / len(loader.dataset)\n",
    "\n",
    "# Initialize the model, optimizer, and loss function\n",
    "model = GCN(num_features=num_features, hidden_channels=32, num_classes=num_classes)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "# Training loop\n",
    "num_epochs = 100\n",
    "for epoch in range(num_epochs):\n",
    "    train_loss, train_acc = train()\n",
    "    val_acc = validate(val_loader)\n",
    "    print(f'Epoch {epoch+1}, Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}, Val Acc: {val_acc:.4f}')\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}