3_generate_family_fasta_files.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "cc96e5a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import\n",
    "from Bio import SeqIO\n",
    "from Bio.Seq import Seq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "44ffee5a-a3a0-46e4-a543-f0f06701b62c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define list of samples to process\n",
    "samples = [\"04.B1.W14.01\", \"04.M1.W09.02\", \n",
    "           \"05.B1.W14.04\", \"05.M1.W08.03\",\n",
    "           \"27.B1.W13.06\", \"27.M1.W10.07\", \n",
    "           \"30.B1.W11.08\", \"30.M1.W04.09\", \n",
    "           \"38.B1.W10.11\", \"38.M1.W03.10\", \n",
    "           \"39.B1.W11.12\", \"39.M1.W03.13\", \"39.M1.W05.14\", \n",
    "           \"53.B1.W14.17\", \"53.M1.W07.16\", \n",
    "           \"56.B1.W09.22\", \"56.M1.W03.21\", \n",
    "           \"63.B1.W09.29\", \"63.M1.W02.30\", \n",
    "           \"66.B1.W09.25\", \"66.M1.W02.24\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "f0a1b96d-5ef6-418b-b3c8-521bd775f994",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_haplotype_frequency(sample):\n",
    "    # Initialize dictionaries to store unique haplotypes and their frequencies\n",
    "    unique_haplotypes = {}\n",
    "    haplotype_frequencies = {}\n",
    "    \n",
    "    # Flag to identify if we are within the \"Haplotype frequencies\" section in the file\n",
    "    is_in_haplotype_section = False\n",
    "\n",
    "    # Open and read the file\n",
    "    with open(\"data/\" + sample + \"/HaROLD/step_2/\" + sample + \".log\", \"r\") as file:\n",
    "        lines = file.readlines()\n",
    "        for line in lines:\n",
    "            # Strip the line to remove leading and trailing whitespaces\n",
    "            stripped_line = line.strip()\n",
    "            \n",
    "            # If we encounter the \"Haplotype frequencies\" line, set the flag to True\n",
    "            if 'Haplotype frequencies' in stripped_line:\n",
    "                is_in_haplotype_section = True\n",
    "                continue\n",
    "\n",
    "            # If we're inside the \"Haplotype frequencies\" section and we encounter \n",
    "            # a blank line or a line starting with a non-digit, break out of the loop\n",
    "            if is_in_haplotype_section and (stripped_line == \"\" or not stripped_line[0].isdigit()):\n",
    "                break\n",
    "\n",
    "            # If we're inside the \"Haplotype frequencies\" section, extract key-value pairs and update the dictionary\n",
    "            if is_in_haplotype_section:\n",
    "                key, value = stripped_line.split()\n",
    "                # Convert the string value to float after replacing comma with period\n",
    "                haplotype_frequencies[sample + \"_H\" + str(key)] = round((float(value.replace(',', '.')) * 100), 1)\n",
    "\n",
    "    # Return the final dictionary of haplotype frequencies\n",
    "    return haplotype_frequencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "b1502c09-222c-414d-82b1-470eb4976564",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a dictionary to store the haplotype frequencies across all samples\n",
    "haplotype_frequencies = {}\n",
    "\n",
    "# Loop over each sample in the list\n",
    "for sample in samples:\n",
    "    # Call the 'get_frequencies' function for the sample and update the global haplotype_frequencies dictionary with the frequencies for this sample\n",
    "    haplotype_frequencies = {**haplotype_frequencies, **get_haplotype_frequency(sample)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "4c325218",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "04.B1.W14.01_04.M1.W09.02_ successfully generated.\n",
      "05.B1.W14.04_05.M1.W08.03_ successfully generated.\n",
      "27.B1.W13.06_27.M1.W10.07_ successfully generated.\n",
      "30.B1.W11.08_30.M1.W04.09_ successfully generated.\n",
      "38.B1.W10.11_38.M1.W03.10_ successfully generated.\n",
      "39.B1.W11.12_39.M1.W03.13_39.M1.W05.14_ successfully generated.\n",
      "53.B1.W14.17_53.M1.W07.16_ successfully generated.\n",
      "56.B1.W09.22_56.M1.W03.21_ successfully generated.\n",
      "63.B1.W09.29_63.M1.W02.30_ successfully generated.\n",
      "66.B1.W09.25_66.M1.W02.24_ successfully generated.\n"
     ]
    }
   ],
   "source": [
    "# Create an empty set to store the unique two-character prefixes\n",
    "families = set()\n",
    "\n",
    "# Iterate through each sample and add the first two characters to the set\n",
    "for sample in samples:\n",
    "    families.add(sample[:2])\n",
    "\n",
    "# Convert the set to a list and sort it\n",
    "families = sorted(list(families))\n",
    "\n",
    "# Load the sequences from the additional fasta files\n",
    "ad169_data = [(record.id, str(record.seq)) for record in SeqIO.parse(\"data/fasta/reference/HCMV_AD169_UL33.fasta\", \"fasta\")]\n",
    "merlin_data = [(record.id, str(record.seq)) for record in SeqIO.parse(\"data/fasta/reference/HCMV_Merlin_UL33.fasta\", \"fasta\")]\n",
    "toledo_data = [(record.id, str(record.seq)) for record in SeqIO.parse(\"data/fasta/reference/HCMV_Toledo_UL33.fasta\", \"fasta\")]\n",
    "towne_data = [(record.id, str(record.seq)) for record in SeqIO.parse(\"data/fasta/reference/HCMV_Towne_UL33.fasta\", \"fasta\")]\n",
    "\n",
    "# Loop over families\n",
    "for family in families:\n",
    "    # To store all sequences and their IDs for the current family\n",
    "    all_data = []\n",
    "    # To store concatenated sample names for the current family\n",
    "    output_file_name = \"\"  \n",
    "    # Loop over samples\n",
    "    for sample in samples:\n",
    "        # Check if the sample belongs to the current family\n",
    "        if sample.startswith(family):\n",
    "            fasta_path = \"data/\" + sample + \"/HaROLD/step_2/\" + sample + \".fasta\"\n",
    "            # Initialize an empty list to store the data\n",
    "            data = []\n",
    "            # Open and parse the FASTA file\n",
    "            for record in SeqIO.parse(fasta_path, \"fasta\"):\n",
    "                # Store each record's ID and sequence as a tuple in the list\n",
    "                if haplotype_frequencies[record.id] > 1: data.append((record.id, str(record.seq)))\n",
    "            all_data.extend(data)\n",
    "            output_file_name += sample + \"_\"\n",
    "    # Add the sequences and IDs from HCMV_AD169_UL33.fasta and HCMV_Merlin_UL33.fasta\n",
    "    all_data.extend(ad169_data)\n",
    "    all_data.extend(merlin_data)\n",
    "    all_data.extend(toledo_data)\n",
    "    all_data.extend(towne_data)\n",
    "    \n",
    "    # Save sequences to a new fasta file\n",
    "    with open(\"data/fasta/initial/\" + output_file_name[:-1] + \".fasta\", \"w\", newline='\\r\\n') as f_out:\n",
    "        for seq_id, seq in all_data: \n",
    "            f_out.write(f\">{seq_id}\\n{seq}\\n\")\n",
    "    print(output_file_name, \"successfully generated.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "c4cf338a-a3f9-4cb1-b1f3-176acaea4430",
   "metadata": {},
   "outputs": [],
   "source": [
    "def truncate_sequences(sample, motif_start, motif_end):\n",
    "    # Read the FASTA file containing aligned sequences based on the sample nam\n",
    "    records = list(SeqIO.parse(\"data/fasta/initial/\" + sample + \".fasta\", \"fasta\"))\n",
    "    truncated_sequences = []\n",
    "    for record in records:\n",
    "        sequence = str(record.seq)\n",
    "        # Find the start motif in the sequence and truncate the sequence from that motif onwards\n",
    "        motif_start_index = sequence.find(motif_start)\n",
    "        if motif_start_index != -1:\n",
    "            sequence = sequence[motif_start_index:]\n",
    "        # Find the end motif and truncate the sequence to only include up to and including this motif\n",
    "        motif_end_index = sequence.find(motif_end)\n",
    "        if motif_end_index != -1:\n",
    "            sequence = sequence[:motif_end_index+len(motif_end)]\n",
    "        # Update the sequence of the record\n",
    "        record.seq = Seq(sequence) \n",
    "        \n",
    "        # Check the percentage of nucleotides that are missing\n",
    "        missing_percentage = round((sequence.count('-') / len(sequence)) * 100, 2)\n",
    "        \n",
    "        # If missing percentage is 5% or less, append the record to the list\n",
    "        if missing_percentage <= 5: truncated_sequences.append(record)\n",
    "        else: print(\"Removed sequence: \",record.id, missing_percentage, \"%\")\n",
    "    \n",
    "    # Write the truncated sequences back to the FASTA file\n",
    "    SeqIO.write(truncated_sequences, \"data/fasta/truncated/\" + sample + \"_truncated.fasta\", \"fasta\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "a6ba195a-8522-43a7-9c98-00de2d893f26",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Truncation of sample: 04.B1.W14.01_04.M1.W09.02\n",
      "Removed sequence:  04.M1.W09.02_H0 6.25 %\n",
      "Truncation of sample: 05.B1.W14.04_05.M1.W08.03\n",
      "Truncation of sample: 27.B1.W13.06_27.M1.W10.07\n",
      "Removed sequence:  27.M1.W10.07_H4 5.08 %\n",
      "Truncation of sample: 30.B1.W11.08_30.M1.W04.09\n",
      "Truncation of sample: 38.B1.W10.11_38.M1.W03.10\n",
      "Truncation of sample: 39.B1.W11.12_39.M1.W03.13_39.M1.W05.14\n",
      "Removed sequence:  39.M1.W05.14_H3 5.08 %\n",
      "Truncation of sample: 53.B1.W14.17_53.M1.W07.16\n",
      "Truncation of sample: 56.B1.W09.22_56.M1.W03.21\n",
      "Truncation of sample: 63.B1.W09.29_63.M1.W02.30\n",
      "Truncation of sample: 66.B1.W09.25_66.M1.W02.24\n",
      "Removed sequence:  66.M1.W02.24_H1 5.08 %\n"
     ]
    }
   ],
   "source": [
    "# Define list of samples to process\n",
    "samples = [\"04.B1.W14.01_04.M1.W09.02\", \"05.B1.W14.04_05.M1.W08.03\", \n",
    "           \"27.B1.W13.06_27.M1.W10.07\",\"30.B1.W11.08_30.M1.W04.09\",\n",
    "           \"38.B1.W10.11_38.M1.W03.10\", \"39.B1.W11.12_39.M1.W03.13_39.M1.W05.14\",\n",
    "           \"53.B1.W14.17_53.M1.W07.16\", \"56.B1.W09.22_56.M1.W03.21\", \n",
    "           \"63.B1.W09.29_63.M1.W02.30\", \"66.B1.W09.25_66.M1.W02.24\"]\n",
    "\n",
    "# Define the start and end motifs\n",
    "motif_start = \"ATGGACACCATCATCCAC\"\n",
    "motif_end = \"GGGTATGA\"\n",
    "\n",
    "for sample in samples: \n",
    "    # Print the sample in process\n",
    "    print(\"Truncation of sample:\", sample)\n",
    "    # Truncation of sequences\n",
    "    truncate_sequences(sample, motif_start, motif_end)"
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}