iqtree benchmark had missing rule in snakemake

Author: benjamin-lieser

Readme.md

@@ -39,7 +39,7 @@ Just `cargo add phylo_grad` (will work after publishing)
 Create and activate a conda environment:
 
 ```
-conda env create -n phylo_grad -c conda-forge python=3.11 pytorch bioconda::snakemake pytest bioconda::newick_utils
+conda env create -n phylo_grad -c conda-forge python=3.11 
 source activate phylo_grad
 ```
 
@@ -50,17 +50,17 @@ pip install ./phylo_grad_py
 pip install ./phylo_grad_gpu
 ```
 
-Run Tests
-
+Install other dependecies:
 ```
-cd benchmark_test
-pytest test.py
+conda install -c conda-forge pytorch bioconda::snakemake pytest bioconda::newick_utils bioconda::iqtree bioconda::emboss
+cargo install phylotree
 ```
 
-Install phylotree (to generate random trees)
+Run Tests
 
 ```
-cargo install phylotree
+cd benchmark_test
+pytest test.py
 ```
 
 Run Benchmarks

benchmark_test/Snakefile

@@ -50,10 +50,22 @@ rule benchmark_iqtree:
         slurm_extra="--exclusive",
         mem_mb=800000,
         runtime=1000
-    shell: "/usr/bin/time -o {output[1]} -f '%E %P' iqtree3 -s {input.alignment} -nt {threads} --model-joint GTR20+FO --redo -te {input.tree} --show-lh -v -quiet"
+    shell: "/usr/bin/time -o {output[1]} -f '%E %P' iqtree -s {input.alignment} -nt {threads} --model-joint GTR20+FO --redo -te {input.tree} --show-lh -v -quiet"
+
+rule benchmark_phylograd:
+    input: tree = "data/random/tree_{num_leafs}.nwk",
+           alignment = "data/sim/alignment_{num_leafs}_{L}.fasta"
+    output: "data/sim/alignment_{num_leafs}_{L}.fasta.phylograd.npz",
+            "data/sim/alignment_{num_leafs}_{L}.fasta.phylograd.time",
+    threads: 64
+    resources:
+        slurm_extra="--exclusive",
+        mem_mb=800000,
+        runtime=1000
+    shell: "/usr/bin/time -o {output[1]} -f '%E %P' python optimize_global_model.py --fasta_amino {input.alignment} --newick {input.tree} --rust --f64 --output {output[0]}"
 
 rule collect_iqtree_benchmark:
-    input: iqtree = expand("data/sim/alignment_{num_leafs}_{L}.fasta.iqtree.time", num_leafs=[2**n for n in [4,6,8,10,12]], L=[100,500]),
-           phylograd = expand("data/sim/alignment_{num_leafs}_{L}.fasta.phylograd.time", num_leafs=[2**n for n in [4,6,8,10,12]], L=[100,500])
+    input: iqtree = expand("data/sim/alignment_{num_leafs}_{L}.fasta.iqtree.time", num_leafs=[2**n for n in [4,6,8]], L=[100,500]),
+           phylograd = expand("data/sim/alignment_{num_leafs}_{L}.fasta.phylograd.time", num_leafs=[2**n for n in [4,6,8]], L=[100,500])
     output: "data/iqtree_benchmark.txt"
     shell: "tail -n +1 {input.iqtree} {input.phylograd} > {output}"

benchmark_test/input.py

@@ -28,6 +28,14 @@ def read_fasta(fasta_file: str) -> dict:
 
     return seq_dict
 
+def read_fasta_numeric(fasta_file: str) -> dict:
+    alignment = AlignIO.read(fasta_file, "fasta")
+    assert isinstance(alignment, MultipleSeqAlignment)
+
+    seq_dict = {seq.id: np.array([amino_mapping[c] for c in seq.seq], dtype=np.int32) for seq in alignment}
+
+    return seq_dict
+
 def read_newick(newick_file: str) -> dict:
     """
         Reads a newick file and returns a parent_list, branch_lengths and number of leaf nodes.

benchmark_test/optimize_global_model.py

@@ -1,14 +1,9 @@
-"""
-    This script is the benchmark, it supports the phylo_grad gradients and pytorch gradients.
-"""
-
 import sys
 
 import torch
 
 import input
 import phylo_grad
-import felsenstein
 import numpy as np
 import argparse
 
@@ -27,8 +22,6 @@
 backend = parser.add_argument_group('Backend')
 exclusive_group = backend.add_mutually_exclusive_group(required=True)
 exclusive_group.add_argument('--rust', action='store_true')
-exclusive_group.add_argument('--pytorch', action='store_true')
-exclusive_group.add_argument('--pytorch_gpu', action='store_true')
 
 fp_precision = parser.add_argument_group('fp precision')
 exclusive_group = fp_precision.add_mutually_exclusive_group(required=True)
@@ -45,42 +38,27 @@
     np_dtype = np.float32
 
 if args.fasta_amino is not None:
-    tree, L = input.read_newick_fasta(args.newick, args.fasta_amino)
+    alignment = input.read_fasta_numeric(args.fasta_amino)
     # Counts amino acids
     counts = torch.zeros(21, dtype=torch.int64)
-    for _, _, seq in tree:
-        if seq is not None:
-            numeric = [input.amino_mapping[c] for c in seq]
-            for i in numeric:
-                counts[i] += 1
+    for seq in alignment.values():
+        L = len(seq)
+        for i in seq:
+            counts[i] += 1
 
     initial_energies = torch.log(counts[:-1])
-
-    if args.pytorch_gpu:
-        tree = [(par, dist, input.amino_to_embedding(seq).cuda() if seq else None) for par, dist, seq in tree]
-    else:
-        tree = [(par, dist, input.amino_to_embedding(seq)) for par, dist, seq in tree]
-    
 
 
 
 #Init random parameters
 torch.manual_seed(0)
 
-if args.pytorch_gpu:
-    shared = torch.zeros(190, requires_grad=True, dtype=dtype, device="cuda")
-    energies = torch.tensor(initial_energies, requires_grad=True, dtype=dtype, device="cuda")
-else:
-    shared = torch.zeros(190, requires_grad=True, dtype=dtype)
-    energies = torch.tensor(initial_energies, requires_grad=True, dtype=dtype)
+shared = torch.zeros(190, requires_grad=True, dtype=dtype)
+energies = initial_energies.clone().to( dtype=dtype).requires_grad_(True)
 
 if args.rust:
-    leaf_log_p = torch.stack([seq for _,_, seq in tree if seq is not None]).transpose(1,0)
-    tree = np.array([(par, dist) for par, dist, _ in tree], dtype=np_dtype)
-    tree = phylo_grad.FelsensteinTree(tree, leaf_log_p.type(dtype).numpy(), 1e-4)
-    
-else:
-    tree = felsenstein.FelsensteinTree(tree)
+    leaf_log_p = input.read_fasta(args.fasta_amino)
+    tree = phylo_grad.FelsensteinTree.from_newick(args.newick, leaf_log_p, np_dtype, 1e-4, gpu = False)
 
 optimizer = torch.optim.Adam([shared, energies], lr=0.1)
 
@@ -133,7 +111,7 @@ def rate_matrix(shared, energies, L):
 
 S, sqrt_pi = rate_matrix(shared, energies, L)
 
-np.savez(args.output, S=S.detach().cpu().numpy()[0], sqrt_pi=sqrt_pi.detach().cpu().numpy()[0], shared=shared.detach().cpu().numpy(), energies=energies.detach().cpu().numpy())
+np.savez(args.output, S=S.detach().numpy()[0], sqrt_pi=sqrt_pi.detach().numpy()[0], shared=shared.detach().numpy(), energies=energies.detach().numpy())
 
 # Print peak memory usage
 import resource