notebooks/feature_generation.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Feature generation for PDB file inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import atom3.complex as comp\n",
    "import atom3.conservation as con\n",
    "import atom3.neighbors as nb\n",
    "import atom3.pair as pair\n",
    "import atom3.parse as parse\n",
    "import dill as pickle\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "from project.utils.utils import annotate_idr_residues, impute_missing_feature_values, postprocess_pruned_pair, process_raw_file_into_dgl_graphs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Parse PDB file input to pair-wise features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdb_filename = \"../project/datasets/Input/raw/pdb/2g/12gs.pdb1\"  # note: an input PDB must be uncompressed (e.g., not in `.gz` archive format) using e.g., `gunzip`\n",
    "output_pkl = \"../project/datasets/Input/interim/parsed/2g/12gs.pdb1.pkl\"\n",
    "complexes_dill = \"../project/datasets/Input/interim/complexes/complexes.dill\"\n",
    "pairs_dir = \"../project/datasets/Input/interim/pairs\"\n",
    "pkl_filenames = [output_pkl]\n",
    "source_type = \"rcsb\"  # note: this default value will likely work for common use cases (i.e., those concerning bound-state PDB protein complex structure inputs)\n",
    "neighbor_def = \"non_heavy_res\"\n",
    "cutoff = 6  # note: distance threshold (in Angstrom) for classifying inter-chain interactions can be customized here\n",
    "unbound = False  # note: if `source_type` is set to `rcsb`, this value should likely be `False`\n",
    "\n",
    "for item in [\n",
    "    Path(pdb_filename).parent,\n",
    "    Path(output_pkl).parent,\n",
    "    Path(complexes_dill).parent,\n",
    "    pairs_dir,\n",
    "]:\n",
    "    os.makedirs(item, exist_ok=True)\n",
    "\n",
    "# note: the following replicates the logic within `make_dataset.py` for a single PDB file input\n",
    "parse.parse(\n",
    "    # note: assumes the PDB file input (i.e., `pdb_filename`) is not compressed\n",
    "    pdb_filename=pdb_filename,\n",
    "    output_pkl=output_pkl\n",
    ")\n",
    "complexes = comp.get_complexes(filenames=pkl_filenames, type=source_type)\n",
    "comp.write_complexes(complexes=complexes, output_dill=complexes_dill)\n",
    "get_neighbors = nb.build_get_neighbors(criteria=neighbor_def, cutoff=cutoff)\n",
    "get_pairs = pair.build_get_pairs(\n",
    "    neighbor_def=neighbor_def,\n",
    "    type=source_type,\n",
    "    unbound=unbound,\n",
    "    nb_fn=get_neighbors,\n",
    "    full=False\n",
    ")\n",
    "complexes = comp.read_complexes(input_dill=complexes_dill)\n",
    "pair.complex_to_pairs(\n",
    "    complex=list(complexes['data'].values())[0],\n",
    "    source_type=source_type,\n",
    "    get_pairs=get_pairs,\n",
    "    output_dir=pairs_dir\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Compute sequence-based features using external tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psaia_dir = \"~/Programs/PSAIA-1.0/bin/linux/psa\"  # note: replace this with the path to your local installation of PSAIA\n",
    "psaia_config_file = \"../project/datasets/builder/psaia_config_file_dips.txt\"  # note: choose `psaia_config_file_dips.txt` according to the `source_type` selected above\n",
    "file_list_file = os.path.join(\"../project/datasets/Input/interim/external_feats/\", 'PSAIA', source_type.upper(), 'pdb_list.fls')\n",
    "num_cpus = 8\n",
    "pkl_filename = \"../project/datasets/Input/interim/parsed/2g/12gs.pdb1.pkl\"\n",
    "output_filename = \"../project/datasets/Input/interim/external_feats/parsed/2g/12gs.pdb1.pkl\"\n",
    "hhsuite_db = \"~/Data/Databases/pfamA_35.0/pfam\"  # note: substitute the path to your local HHsuite3 database here\n",
    "num_iter = 2\n",
    "msa_only = False\n",
    "\n",
    "for item in [\n",
    "    Path(file_list_file).parent,\n",
    "    Path(output_filename).parent,\n",
    "]:\n",
    "    os.makedirs(item, exist_ok=True)\n",
    "\n",
    "# note: the following replicates the logic within `generate_psaia_features.py` and `generate_hhsuite_features.py` for a single PDB file input\n",
    "with open(file_list_file, 'w') as file:\n",
    "    file.write(f'{pdb_filename}\\n')  # note: references the `pdb_filename` as defined previously\n",
    "con.gen_protrusion_index(\n",
    "    psaia_dir=psaia_dir,\n",
    "    psaia_config_file=psaia_config_file,\n",
    "    file_list_file=file_list_file,\n",
    ")\n",
    "con.map_profile_hmms(\n",
    "    num_cpus=num_cpus,\n",
    "    pkl_filename=pkl_filename,\n",
    "    output_filename=output_filename,\n",
    "    hhsuite_db=hhsuite_db,\n",
    "    source_type=source_type,\n",
    "    num_iter=num_iter,\n",
    "    msa_only=msa_only,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Compute structure-based features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from project.utils.utils import __should_keep_postprocessed\n",
    "\n",
    "\n",
    "raw_pdb_dir = \"../project/datasets/Input/raw/pdb\"\n",
    "pair_filename = \"../project/datasets/Input/interim/pairs/2g/12gs.pdb1_0.dill\"\n",
    "source_type = \"rcsb\"\n",
    "external_feats_dir = \"../project/datasets/Input/interim/external_feats/parsed\"\n",
    "output_filename = \"../project/datasets/Input/final/raw/2g/12gs.pdb1_0.dill\"\n",
    "\n",
    "unprocessed_pair, raw_pdb_filenames, should_keep = __should_keep_postprocessed(raw_pdb_dir, pair_filename, source_type)\n",
    "if should_keep:\n",
    "    # note: save `postprocessed_pair` to local storage within `project/datasets/Input/final/raw` for future reference as desired\n",
    "    postprocessed_pair = postprocess_pruned_pair(\n",
    "        raw_pdb_filenames=raw_pdb_filenames,\n",
    "        external_feats_dir=external_feats_dir,\n",
    "        original_pair=unprocessed_pair,\n",
    "        source_type=source_type,\n",
    "    )\n",
    "    # write into output_filenames if not exist\n",
    "    os.makedirs(Path(output_filename).parent, exist_ok=True)\n",
    "    with open(output_filename, 'wb') as f:\n",
    "        pickle.dump(postprocessed_pair, f)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Embed deep learning-based IDR features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# note: ensures the Docker image for `flDPnn` is available locally before trying to run inference with the model\n",
    "!docker pull docker.io/sinaghadermarzi/fldpnn\n",
    "\n",
    "input_pair_filename = \"../project/datasets/Input/final/raw/2g/12gs.pdb1_0.dill\"\n",
    "pickle_filepaths = [input_pair_filename]\n",
    "\n",
    "annotate_idr_residues(\n",
    "    pickle_filepaths=pickle_filepaths\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Impute missing feature values (optional)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_pair_filename = \"../project/datasets/Input/final/raw/2g/12gs.pdb1_0.dill\"\n",
    "output_pair_filename = \"../project/datasets/Input/final/raw/2g/12gs.pdb1_0_imputed.dill\"\n",
    "impute_atom_features = False\n",
    "advanced_logging = False\n",
    "\n",
    "impute_missing_feature_values(\n",
    "    input_pair_filename=input_pair_filename,\n",
    "    output_pair_filename=output_pair_filename,\n",
    "    impute_atom_features=impute_atom_features,\n",
    "    advanced_logging=advanced_logging,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6. Convert pair-wise features into graph inputs (optional)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_filepath = \"../project/datasets/Input/final/raw/2g/12gs.pdb1_0_imputed.dill\"\n",
    "new_graph_dir = \"../project/datasets/Input/final/processed/2g\"\n",
    "processed_filepath = \"../project/datasets/Input/final/processed/2g/12gs.pdb1.pt\"\n",
    "edge_dist_cutoff = 15.0\n",
    "edge_limit = 5000\n",
    "self_loops = True\n",
    "\n",
    "os.makedirs(new_graph_dir, exist_ok=True)\n",
    "\n",
    "process_raw_file_into_dgl_graphs(\n",
    "    raw_filepath=raw_filepath,\n",
    "    new_graph_dir=new_graph_dir,\n",
    "    processed_filepath=processed_filepath,\n",
    "    edge_dist_cutoff=edge_dist_cutoff,\n",
    "    edge_limit=edge_limit,\n",
    "    self_loops=self_loops,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "DIPS-Plus",
   "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.8.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}