Homology_plot.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33c0d3c6-f399-4cbb-bb53-4faa032612d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "import torch\n",
    "import torchmetrics\n",
    "import torch.nn.functional as F\n",
    "import esm\n",
    "from torch.utils.data import Dataset\n",
    "import pytorch_lightning as pl\n",
    "import plotly.express as px # plotting\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from tqdm import tqdm\n",
    "import pandas as pd\n",
    "import umap\n",
    "from typing import List, Tuple\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad8abdfb-c98f-4a0c-a820-dd77764e1827",
   "metadata": {},
   "outputs": [],
   "source": [
    "def random_seq_emb(p: Path, layer: int):\n",
    "    embeddings = []\n",
    "    for i in tqdm(p.glob(\"*.pt\")):\n",
    "        t = torch.load(i)\n",
    "        emb = t[\"mean_representations\"][layer]\n",
    "        embeddings.append(emb)\n",
    "    return embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4a0a19d-b3b5-4efc-820f-e0d88826bd73",
   "metadata": {},
   "outputs": [],
   "source": [
    "p_normal = Path(\"/shared/homology/esm2_t33_650M_UR50D/sequence_normal\")\n",
    "p_shuffled = Path(\"/shared/homology/esm2_t33_650M_UR50D/sequence_shuffled\")\n",
    "p_shuffled_by_every_aminoacids = Path(\"/shared/homology/esm2_t33_650M_UR50D/sequence_shuffled_by_every_aminoacids\")\n",
    "\n",
    "emb_normal = pd.DataFrame(random_seq_emb(p_normal, 33))\n",
    "emb_shuffled = pd.DataFrame(random_seq_emb(p_shuffled, 33))\n",
    "emb_shuffled_by_every_aminoacids = pd.DataFrame(random_seq_emb(p_shuffled_by_every_aminoacids, 33))\n",
    "\n",
    "emb = pd.concat([emb_normal, emb_shuffled, emb_shuffled_by_every_aminoacids])\n",
    "#print(emb_normal.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a17ad554-53a3-4759-8618-9724d22d4cec",
   "metadata": {},
   "outputs": [],
   "source": [
    "embeds = umap.UMAP(n_components=2, n_neighbors=5, min_dist=0.1, metric=\"euclidean\", random_state = 42).fit_transform(emb)\n",
    "embeds = pd.DataFrame(embeds, columns = [\"x\", \"y\"])\n",
    "embeds['source'] = ['normal'] * emb_normal.shape[0] + ['shuffled'] * emb_shuffled.shape[0] +['shuffled_by_every_aminoacid'] * emb_shuffled_by_every_aminoacids.shape[0]\n",
    "#print(embeds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "125a9ad8-22d0-4b39-90f6-e9e48dec29ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "#fig = px.scatter(ds, x=0, y=1, color_continuous_scale=px.colors.sequential.Rainbow)\n",
    "fig = px.scatter(embeds, x='x', y='y', color='source', height = 600, width = 800, opacity=0.8)\n",
    "fig.update_traces(textposition='top center', marker=dict(size=5))\n",
    "fig.write_image(\"normal_vs_shuffled_33.png\", scale = 5)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f746dc2-1461-4ce5-8d7c-848c6baf67d9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}