write_WDS.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7480fba9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import sys\n",
    "import webdataset as wds\n",
    "import os\n",
    "import tqdm\n",
    "import warnings\n",
    "import gzip\n",
    "from Bio.PDB.PDBExceptions import PDBConstructionWarning\n",
    "from Bio.PDB import PDBParser\n",
    "from Bio.SeqUtils import seq1\n",
    "warnings.simplefilter('ignore', PDBConstructionWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01addc81",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0b7875e3-63e7-43a1-95d2-39d0c5bed44e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_residue(residue):\n",
    "    atoms = ['N', 'CA', 'C', 'O']\n",
    "    coordinates = []\n",
    "    for r in atoms:\n",
    "        coord = residue.child_dict.get(r, None)\n",
    "        if coord is None:\n",
    "            if r == 'O':\n",
    "                coord = residue.child_dict.get('OXT', None)\n",
    "            if coord is None:\n",
    "                return None, None\n",
    "        coordinates.append(np.array(coord.get_coord()))\n",
    "    return np.stack(coordinates), seq1(residue.resname)\n",
    "\n",
    "def process_chain(chain):\n",
    "    coordinates = []\n",
    "    seq = []\n",
    "    for r in chain:\n",
    "        output, residue_name = process_residue(r)\n",
    "        if output is not None:\n",
    "            coordinates.append(output)\n",
    "            seq.append(residue_name)\n",
    "    if len(coordinates) == 0:\n",
    "        return None\n",
    "    coordinates = np.stack(coordinates)\n",
    "    seq = ''.join(seq)\n",
    "    return coordinates, seq\n",
    "\n",
    "def process_chains(chains, pep=False, prot=False):\n",
    "    if pep or prot:\n",
    "        chain_lens = []\n",
    "        chain_ids = []\n",
    "        for chain in chains:\n",
    "            for i, res in enumerate(chain):\n",
    "                continue\n",
    "            chain_lens.append(i)\n",
    "            chain_ids.append(chain.id)\n",
    "        if chain_lens[0] < chain_lens[1]:\n",
    "            pep_id = chain_ids[0]\n",
    "            prot_id = chain_ids[1]\n",
    "        else:\n",
    "            pep_id = chain_ids[1]\n",
    "            prot_id = chain_ids[0]\n",
    "        if pep and isinstance(pep, str): pep_id == pep\n",
    "        if prot and isinstance(prot, str): prot_id == prot\n",
    "    output = []\n",
    "    chain_ids = []\n",
    "    for chain in chains:\n",
    "        if (pep and chain.id != pep_id) or (prot and chain.id != prot_id):\n",
    "            continue\n",
    "        out = process_chain(chain)\n",
    "        if out is not None:\n",
    "            output.append(out)\n",
    "            chain_ids.append(chain.id)\n",
    "    coords = [u[0] for u in output]\n",
    "    seqs = [u[1] for u in output]\n",
    "    return coords, seqs, chain_ids\n",
    "\n",
    "def process_structure(structure, pep=False, prot=False):\n",
    "    for s in structure: # only one structure\n",
    "        return process_chains(s, pep, prot)\n",
    "    return None\n",
    "\n",
    "# +\n",
    "def process_pdb(parser, pdb_filename):\n",
    "    # print(pdb_filename)\n",
    "    with gzip.open(pdb_filename, \"rt\") as file_handle:\n",
    "        structure = parser.get_structure(\"?\", file_handle)\n",
    "        date = structure.header['deposition_date']\n",
    "        return process_structure(structure), date\n",
    "    \n",
    "def process_pdb_raw(parser, pdb_filename, pep=False, prot=False):\n",
    "    s = parser.get_structure(\"?\", pdb_filename)\n",
    "    return process_structure(s, pep, prot)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f2074d9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_input_ids(index_file):\n",
    "    input_ids = []\n",
    "    with open(os.path.join(index_file), 'r') as f:\n",
    "        for line in f:\n",
    "            input_ids += [line.strip()]\n",
    "    return np.array(input_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dd4a9e07",
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_dataset(dataset, tar_name, use_shards=False, max_shard_count=10000):\n",
    "    if use_shards:\n",
    "        os.makedirs(tar_name, exist_ok=True)\n",
    "        sink = wds.ShardWriter(f'{tar_name}/shard-%06d.tar',maxcount=max_shard_count)\n",
    "    else:\n",
    "        sink = wds.TarWriter(tar_name)\n",
    "    for index, (batch, pdb_id) in enumerate(dataset):\n",
    "        if index%1000==0:\n",
    "            print(f\"{index:6d}\", end=\"\\r\", flush=True, file=sys.stderr)\n",
    "        if len(batch[0]) == 0:\n",
    "            continue\n",
    "        sink.write({\n",
    "            \"__key__\": \"sample%06d\" % index,\n",
    "            \"inp.pyd\": dict(coords=batch[0], seqs=batch[1], chain_ids=batch[2], pdb_id=pdb_id),\n",
    "        })\n",
    "    sink.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "bad9bd68-0419-4e7d-9f7c-9de9d543d6f2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 19.23it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 25890.77it/s]\n",
      "     0\r"
     ]
    }
   ],
   "source": [
    "dir_ = \"pdbs\"\n",
    "parser = PDBParser()\n",
    "root_pdb = dir_\n",
    "outputs = []\n",
    "for i, pdb_file in tqdm.tqdm(enumerate(os.listdir(dir_)), total=len(os.listdir(dir_))):\n",
    "    pdb_file = pdb_file.strip()\n",
    "    pdb_file = os.path.join(dir_, pdb_file)\n",
    "    out = process_pdb_raw(parser, pdb_file)\n",
    "    pdb_id = pdb_file.split('.')[0]\n",
    "    outputs.append((out, pdb_id))\n",
    "    # os.remove(pdb_filename)\n",
    "\n",
    "output_dicts = {'train': []} #{'train': [], 'test': [], 'val': []}\n",
    "for o, pdb_id in tqdm.tqdm(outputs):\n",
    "    if o is None:\n",
    "        continue\n",
    "    output_dicts['train'].append((o, pdb_id))\n",
    "\n",
    "for k, dataset in output_dicts.items():\n",
    "    if k == 'train':\n",
    "        tar_name = \"wds/4cay_new.wds\"\n",
    "        write_dataset(dataset, tar_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1afb5b98",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:rla_new]",
   "language": "python",
   "name": "conda-env-rla_new-py"
  },
  "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.14"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}