pfam.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce014d42",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import pandas as pd\n",
    "from collections import defaultdict\n",
    "import os\n",
    "import os.path as osp\n",
    "import numpy as np\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5382067a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -lha /scratch/SCRATCH_NVME/ilya/pretrain_data/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5d8d753",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in os.listdir(\"/scratch/SCRATCH_NVME/ilya/pretrain_data/\"):\n",
    "    if \"pfam\" in i:\n",
    "        uniprot = pd.read_csv(\"/scratch/SCRATCH_NVME/ilya/pretrain_data/uniprot.tsv.gz\", sep='\\t').set_index(\"Entry\")[['Protein families', \"Cross-reference (Pfam)\"]]\n",
    "        prots = pd.read_pickle(\"/scratch/SCRATCH_NVME/ilya/pretrain_data/{}\".format(i))\n",
    "        prots.index = prots.index.to_series().apply(lambda x:x.split(\"-\")[1])\n",
    "        prots = prots.join(uniprot)\n",
    "\n",
    "        prots = prots[prots['Cross-reference (Pfam)'].notna()].reset_index().drop_duplicates(\"index\").set_index(\"index\")\n",
    "        prots.drop(\"Protein families\", axis=1, inplace=True)\n",
    "\n",
    "        prots['fam'] = prots['Cross-reference (Pfam)'].apply(lambda x: x.split(\";\")[0])\n",
    "        prots['nnodes'] = prots.data.apply(lambda x: x['x'].size(0))\n",
    "        prots = prots[prots['nnodes'].between(150, 1000)]\n",
    "\n",
    "        vc = prots['fam'].value_counts()\n",
    "        prots = prots[prots['fam'].isin(vc[vc > 10].index)]\n",
    "\n",
    "        def add_fam(row):\n",
    "            data = row['data']\n",
    "            if \"index_mapping\" in data:\n",
    "                del data['index_mapping']\n",
    "            data['fam'] = row['fam']\n",
    "            return data\n",
    "\n",
    "        prots['data'] = prots.apply(add_fam, axis=1)\n",
    "        prots[['data']].to_pickle(\"/scratch/SCRATCH_NVME/ilya/pretrain_data/{}\".format(i.split(\".\")[0] + \"processed.pkl\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a9d8202",
   "metadata": {},
   "outputs": [],
   "source": [
    "uniprot = pd.read_csv(\"/scratch/SCRATCH_NVME/ilya/pretrain_data/uniprot.tsv.gz\", sep='\\t').set_index(\"Entry\")[['Protein families', \"Cross-reference (Pfam)\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b2f0913",
   "metadata": {},
   "outputs": [],
   "source": [
    "uniprot = uniprot[uniprot['Cross-reference (Pfam)'].notna()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d31f3bb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "counts = defaultdict(set)\n",
    "for prot_id, fams in uniprot['Cross-reference (Pfam)'].items():\n",
    "    for fam in fams.rstrip(\";\").split(\";\"):\n",
    "        counts[fam].add(prot_id)\n",
    "\n",
    "counts = pd.Series({k:len(v) for k,v in counts.items()})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e313fcce",
   "metadata": {},
   "outputs": [],
   "source": [
    "accepted_fam = counts[counts > 100].index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edd12553",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_new_fams(fams):\n",
    "    nfam = []\n",
    "    for fam in fams.rstrip(\";\").split(\";\"):\n",
    "        if fam in accepted_fam:\n",
    "            nfam.append(fam)\n",
    "    if nfam:\n",
    "        return \";\".join(nfam)\n",
    "    return np.nan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84c7e622",
   "metadata": {},
   "outputs": [],
   "source": [
    "prots = pd.read_pickle(\"results/prepare_proteins/protein_data_template_label_none.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f33803d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_fam(data):\n",
    "    data['fam'] = random.choice([\"a\", \"b\", \"c\", \"d\"])\n",
    "    return data\n",
    "\n",
    "prots['data'] = prots['data'].apply(set_fam)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8b706a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# prots.index = prots.index.to_series().apply(lambda x:x.split(\"-\")[1])\n",
    "prots = prots.join(uniprot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d40fbb26",
   "metadata": {},
   "outputs": [],
   "source": [
    "prots = prots[prots['Cross-reference (Pfam)'].notna()].reset_index().drop_duplicates(\"index\").set_index(\"index\")\n",
    "prots.drop(\"Protein families\", axis=1, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9fa89a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "prots[['data']].to_pickle(\"test.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfb9ae1f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "",
   "language": "python",
   "name": ""
  },
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}