Extract and write fasta files using Biopython based on input from another file

Question

I have the below code that takes a sequence file and another file with a list of contigs and extracts the sequences and writes them to a file, specifically based on the file with the contig list. The code works great..

However, I'm having to do SeqIO.parse('SAMPLE.fasta', 'fasta') inside the for-loop every time making it very slow. If I read the file in earlier using a variable, eg. sample_f (see commented out line), it fails to identify the records.

from Bio import SeqIO
import pandas as pd 

genomes_l = pd.read_csv('test_data.tsv', sep='\t', header=None, names=['anonymous_gsa_id', 'genome_id'])
# sample_f = SeqIO.parse('SAMPLE.fasta', 'fasta')

for i, r in genomes_l.iterrows():
    genome_name = r['anonymous_gsa_id']
    genome_ids = r['genome_id'].split(',')
    genome_contigs = [rec for rec in SeqIO.parse('SAMPLE.fasta', 'fasta') if rec.id in genome_ids]
    print(genome_contigs)
    with open(f'out_dir/{genome_name}_contigs.fasta', 'w') as handle:
        SeqIO.write(genome_contigs, handle, 'fasta')

Would appreciate your help in improving this. Thank you!

Update: adding examples as suggested The contigIDs are comma-separated within the second column (example below)

424182.1        H|S1|C933685,H|S1|C449562,H|S1|C172291,H|S1|C1169825
1217675.1       H|S1|C1168525,H|S1|C573086,H|S1|C357867,H|S1|C85072,H|S1|C965427,H|S1|C1724718
585503.1        H|S1|C874141,H|S1|C529585

I have another file called SAMPLE.fasta that contains contigIDs and the respective sequences in the next line for each contigID (example below)

>H|S1|C933685
GAAAGTTCTTGACCTGTGGACAGGCTGTGAATCGGGTTGGACAAGT
>H|S1|C85072
GGAAACGGCTGCTGCCATCCTTGCCCTTCGCCCAAG
>H|S1|C965427
CTCAAGAAATTCGGTATCACCGGTAACTATGAGGCAGTCGAGGTCG
etc...
etc...
etc..

Based on this information, I would like to create a separate file for each genomeID (example(s) below)

Output_file: 424182.1.fasta

>H|S1|C933685
GAAAGTTCTTGACCTGTGGACAGGCTGTGAATCGGGTTGGACAAGT

Output_file: 1217675.1.fasta

>H|S1|C85072
GGAAACGGCTGCTGCCATCCTTGCCCTTCGCCCAAG
>H|S1|C965427
CTCAAGAAATTCGGTATCACCGGTAACTATGAGGCAGTCGAGGTCG

Answer 1

SeqIO.parse streams through records one at a time, but it doesn't store them, which is why it doesn't work when you do it outside the loop. This is handy when you have giant files (saves a lot of memory), and in general it lets you pick the best storage solution for the task at hand (if you need to store things at all).

In this case, you want to look things up by their ids, so a python dictionary is probably the best option.

If you put the following outside the loop:

contig_lookup = {record.id: record for record in SeqIO.parse('SAMPLE.fasta', 'fasta')}

And then change your contig lookup:

    genome_contigs = [contig_lookup[_id] for _id in genome_ids if _id in contig_lookup]

You should get a decent speedup.