Added custom models as submodule, added pipleline for trap species handling (up to alignment). Initial progress on custom model rules

2026-10-03 15:16:36 +01:00
parent 3b93ceaf6e
commit f0cd594148
17 changed files with 270 additions and 21 deletions
@@ -0,0 +1,91 @@
+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())