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,66 @@
+rule prepare_read_labels:
+    input:
+        bamfile="../data/aligned_reads/{model}_to_genome.sorted.bam",
+        reference="../data/reference_genomes/full_reference.fasta"
+    output:
+        '../data/ml_inputs/{model}_label_store.pq'
+    conda:
+        '../envs/bam2parquet.yaml'
+    threads:
+        1
+    script:
+        '../scripts/bam2annotation.py'
+
+rule build_data_store:
+    input:
+        pod5=expand('../data/raw_pod5/{{model}}/PBK98658_853a956f_57f83f46_{batch}.pod5',batch=range(1,config["pod5_dataset_size"]+1,config["pod5_stride"])),
+        labels='../data/ml_inputs/{model}_label_store.pq'
+    output:
+        directory('../data/ml_inputs/{model}_data_store')
+    conda:
+        config["torch_env"]
+    shell:
+        """
+        python -m custom_models store\
+        --out {output}\
+        --pod5 {input.pod5}\
+        --labels {input.labels}\
+        --reads-per-genus {config[reads_per_genus]}
+        """
+
+rule train_model:
+    input:
+        data_store='../data/ml_inputs/{model}_data_store'
+    output:
+        directory('../data/ml_models/run_{arch}_{params}_{stride}_{embed}_{heads}')
+    conda:
+        config["torch_env"]
+    shell:
+        """
+        python -m custom_models train\
+        --arch {wildcards.arch}\
+        --budget {wildcards.params}\
+        --stride {wildcards.stride}\
+        --d-embed {wildcards.embed}\
+        --n-heads {wildcards.heads}\
+        --budget-tol {config[budget_tolerance]}\
+        --store {input.store}\
+        --out-dir {output}\
+        --run-id {wildcards.arch}_{wildcards.params}_s{wildcards.stride}_g{wildcards.classes}\
+        --stage ladder_point\
+        --seed {config[seed]}
+        """
+       
+rule train_control:
+    input:
+        data_store='../data/ml_inputs/{model}_data_store'
+    output:
+        directory('../data/ml_models/control_run__{model}')
+    conda:
+        config["torch_env"]
+    shell:
+        """
+        python -m custom_models train\
+        --stage control
+
+        """