91 lines
2.9 KiB
Python
91 lines
2.9 KiB
Python
import pysam
|
|
from Bio import SeqIO
|
|
import pandas as pd
|
|
from pydantic import BaseModel
|
|
|
|
class ReferenceMetadata(BaseModel):
|
|
ref_id : str
|
|
genus : str
|
|
species : str
|
|
ref_len : int
|
|
|
|
if __name__ == '__main__':
|
|
input_bamfile = snakemake.input[0]
|
|
reference_genomes = snakemake.input[1]
|
|
output_parquet = snakemake.output[0]
|
|
|
|
# Get Reference Sequence ID -> Taxonomy mapping
|
|
ref_tax_data = {}
|
|
for ref in SeqIO.parse(reference_genomes,'fasta'):
|
|
ref_tax_data[ref.id] = {
|
|
'genus' : ref.description.split(' ')[1].lower(),
|
|
'species' : '_'.join(ref.description.split(' ')[1:3]).lower(),
|
|
'ref_length' : len(ref.seq)
|
|
}
|
|
|
|
# Open Bamfile for reading
|
|
bamfile = pysam.AlignmentFile(input_bamfile,'rb')
|
|
|
|
# Transcribe Taxonomy from Sequnce ID dict to bam header position list for fast lookup
|
|
ref_data = []
|
|
for ref_id in bamfile.header.references:
|
|
ref_data.append(
|
|
ReferenceMetadata(
|
|
ref_id = ref_id,
|
|
genus = ref_tax_data[ref_id]['genus'],
|
|
species = ref_tax_data[ref_id]['species'],
|
|
ref_len = ref_tax_data[ref_id]['ref_length']
|
|
)
|
|
)
|
|
|
|
# Parse alignments
|
|
data = []
|
|
n_records = 0
|
|
n_unmapped = 0
|
|
n_secondary = 0
|
|
n_supplementary = 0
|
|
n_primary = 0
|
|
for record in bamfile.fetch():
|
|
n_records += 1
|
|
if record.is_unmapped:
|
|
n_unmapped += 1
|
|
continue
|
|
if record.is_secondary:
|
|
n_secondary += 1
|
|
continue
|
|
if record.is_supplementary:
|
|
n_supplementary += 1
|
|
continue
|
|
n_primary += 1
|
|
nm_count = -1
|
|
for tag in record.get_tags():
|
|
if tag[0] == 'NM':
|
|
nm_count = tag[1]
|
|
assert nm_count != -1
|
|
assert isinstance(nm_count,int)
|
|
rlen = record.infer_read_length()
|
|
assert rlen is not None
|
|
qlen = record.infer_query_length()
|
|
assert qlen is not None
|
|
ref = ref_data[record.reference_id]
|
|
data.append(
|
|
{
|
|
'read_id' : record.query_name,
|
|
'ref_name' : ref.ref_id,
|
|
'genus' : ref.genus,
|
|
'species' : ref.species,
|
|
'identity' : 1 - (nm_count / qlen),
|
|
'aligned_len' : qlen,
|
|
'query_len' : rlen,
|
|
}
|
|
)
|
|
df = pd.DataFrame(data)
|
|
df.to_parquet(output_parquet,index=False)
|
|
print(f'Processed {input_bamfile}:')
|
|
print(f'\t{n_records} alignment records')
|
|
print(f'\t{n_primary} primary alignments')
|
|
print(f'\t{n_unmapped} unmapped reads')
|
|
print(f'\t{n_secondary+n_supplementary} non-primary alignments')
|
|
print(f'\t{df.shape[0]} alignment records written to disk')
|
|
print('Per-Genus counts:')
|
|
print(df.groupby('genus').aggregate(n_reads=('read_id','nunique')).reset_index()) |