lab_feature.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pickle\n",
    "from alphafold.common import residue_constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['aatype', 'between_segment_residues', 'domain_name', 'residue_index', 'seq_length', 'sequence', 'deletion_matrix_int', 'msa', 'num_alignments', 'msa_species_identifiers', 'template_aatype', 'template_all_atom_masks', 'template_all_atom_positions', 'template_domain_names', 'template_sequence', 'template_sum_probs'])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load feature dictionary\n",
    "with open(\"features.pkl\", \"rb\") as f:\n",
    "    feature_dict = pickle.load(f)\n",
    "\n",
    "feature_dict.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16,\n",
       "       17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33,\n",
       "       34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50,\n",
       "       51, 52, 53, 54, 55])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "feature_dict[\"residue_index\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_fasta(fasta_string: str):\n",
    "    \"\"\"Parses FASTA string and returns list of strings with amino-acid sequences.\n",
    "\n",
    "    Arguments:\n",
    "      fasta_string: The string contents of a FASTA file.\n",
    "\n",
    "    Returns:\n",
    "      A tuple of two lists:\n",
    "      * A list of sequences.\n",
    "      * A list of sequence descriptions taken from the comment lines. In the\n",
    "        same order as the sequences.\n",
    "    \"\"\"\n",
    "    sequences = []\n",
    "    descriptions = []\n",
    "    index = -1\n",
    "    for line in fasta_string.splitlines():\n",
    "        line = line.strip()\n",
    "        if line.startswith('>'):\n",
    "            index += 1\n",
    "            descriptions.append(line[1:])  # Remove the '>' at the beginning.\n",
    "            sequences.append('')\n",
    "            continue\n",
    "        elif not line:\n",
    "            continue  # Skip blank lines.\n",
    "        sequences[index] += line\n",
    "\n",
    "    return sequences, descriptions\n",
    "\n",
    "\n",
    "def make_sequence_features(\n",
    "    sequence: str, description: str, num_res: int):\n",
    "    \"\"\"Constructs a feature dict of sequence features.\"\"\"\n",
    "    features = {}\n",
    "    features['aatype'] = residue_constants.sequence_to_onehot(\n",
    "          sequence=sequence,\n",
    "          mapping=residue_constants.restype_order_with_x,\n",
    "          map_unknown_to_x=True)\n",
    "    features['between_segment_residues'] = np.zeros((num_res,), dtype=np.int32)\n",
    "    features['domain_name'] = np.array([description.encode('utf-8')],\n",
    "                                        dtype=np.object_)\n",
    "    features['residue_index'] = np.array(range(num_res), dtype=np.int32)\n",
    "    features['seq_length'] = np.array([num_res] * num_res, dtype=np.int32)\n",
    "    features['sequence'] = np.array([sequence.encode('utf-8')], dtype=np.object_)\n",
    "    return features\n",
    "\n",
    "\n",
    "\n",
    "def make_empty_msa_features(\n",
    "    sequence: str, num_res: int):\n",
    "    \"\"\"Constructs a feature dict of empty MSA features.\"\"\"\n",
    "    int_msa = []\n",
    "    num_alignments = 1\n",
    "    species_ids = b''\n",
    "\n",
    "    # Add the query sequence.\n",
    "    int_msa.append([residue_constants.HHBLITS_AA_TO_ID[res] for res in sequence])\n",
    "\n",
    "\n",
    "    features = {}\n",
    "    features['deletion_matrix_int'] = np.zeros((1,num_res), dtype=np.int32)\n",
    "    features['msa'] = np.array(int_msa, dtype=np.int32)\n",
    "    features['num_alignments'] = np.array(\n",
    "        [num_alignments] * num_res, dtype=np.int32)\n",
    "    features['msa_species_identifiers'] = np.array(species_ids, dtype=np.object_)\n",
    "    return features\n",
    "\n",
    "\n",
    "TEMPLATE_FEATURES = {\n",
    "    'template_aatype': np.float32,\n",
    "    'template_all_atom_masks': np.float32,\n",
    "    'template_all_atom_positions': np.float32,\n",
    "    'template_domain_names': object,\n",
    "    'template_sequence': object,\n",
    "    'template_sum_probs': np.float32,\n",
    "}\n",
    "\n",
    "\n",
    "def make_empty_template_features(\n",
    "    sequence: str, num_res: int):\n",
    "    \"\"\"Constructs a feature dict of empty template features.\"\"\"\n",
    "    \n",
    "    template_features = {}\n",
    "    for template_feature_name in TEMPLATE_FEATURES:\n",
    "        template_features[template_feature_name] = []\n",
    "\n",
    "    for name in template_features:\n",
    "        # Make sure the feature has correct dtype even if empty.\n",
    "        template_features[name] = np.array([], dtype=TEMPLATE_FEATURES[name])\n",
    "\n",
    "    return template_features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "fasta_path = \"input/mono_set1/GA98.fasta\"\n",
    "\n",
    "with open(fasta_path) as f:\n",
    "    input_fasta_str = f.read()\n",
    "    input_seqs, input_descs = parse_fasta(input_fasta_str)\n",
    "    if len(input_seqs) != 1:\n",
    "      raise ValueError(\n",
    "          f'More than one input sequence found in {fasta_path}.')\n",
    "    input_sequence = input_seqs[0]\n",
    "    input_description = input_descs[0].split()[0]\n",
    "    num_res = len(input_sequence)\n",
    "\n",
    "# Generate feature: sequence part\n",
    "feature_test = make_sequence_features(\n",
    "               sequence=input_sequence,\n",
    "               description=input_description,\n",
    "               num_res=len(input_sequence))\n",
    "\n",
    "# Generate feature: MSA part\n",
    "feature_test.update(make_empty_msa_features(\n",
    "                sequence=input_sequence,\n",
    "                num_res=len(input_sequence)))\n",
    "\n",
    "# Generate feature: template part\n",
    "feature_test.update(make_empty_template_features(\n",
    "                sequence=input_sequence,\n",
    "                num_res=len(input_sequence)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array_equiv(feature_dict[\"msa_species_identifiers\"], feature_test[\"msa_species_identifiers\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# features_output_path = os.path.join(output_dir, 'features.pkl')\n",
    "# with open(features_output_path, 'wb') as f:\n",
    "#     pickle.dump(feature_dict, f, protocol=4)"
   ]
  }
 ],
 "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.8.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}