phykit 2.1.0__py3-none-any.whl

phykit/services/alignment/create_concatenation_matrix.py

@@ -0,0 +1,254 @@
+import sys
+import os
+from textwrap import dedent
+from typing import Dict, List, Tuple
+from Bio import SeqIO
+from Bio.SeqRecord import SeqRecord
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from functools import partial
+import multiprocessing as mp
+from collections import defaultdict
+
+from .base import Alignment
+from ...helpers.files import read_single_column_file_to_list
+
+
+class CreateConcatenationMatrix(Alignment):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self) -> None:
+        self.create_concatenation_matrix(
+            self.alignment_list_path,
+            self.prefix
+        )
+
+    def process_args(self, args) -> Dict[str, str]:
+        return dict(alignment_list_path=args.alignment_list, prefix=args.prefix)
+
+    def read_alignment_paths(self, alignment_list_path: str) -> List[str]:
+        try:
+            return read_single_column_file_to_list(alignment_list_path)
+        except FileNotFoundError:
+            print("Alignment list file (-a) is not found. Please check pathing.")
+            sys.exit(2)
+
+    @staticmethod
+    def _get_taxa_from_alignment(alignment_path: str) -> set:
+        """Extract taxa names from a single alignment file."""
+        return {seq_record.id for seq_record in SeqIO.parse(alignment_path, "fasta")}
+
+    def get_taxa_names(self, alignment_paths: List[str]) -> List[str]:
+        """Get all unique taxa names from alignment files in parallel."""
+        taxa = set()
+
+        # Process files in parallel if there are many
+        if len(alignment_paths) > 10:
+            with ProcessPoolExecutor(max_workers=min(mp.cpu_count(), len(alignment_paths))) as executor:
+                futures = [executor.submit(self._get_taxa_from_alignment, path) for path in alignment_paths]
+                for future in as_completed(futures):
+                    taxa.update(future.result())
+        else:
+            # Process sequentially for small datasets
+            for alignment_path in alignment_paths:
+                taxa.update(self._get_taxa_from_alignment(alignment_path))
+
+        return sorted(taxa)
+
+    def print_start_message(
+        self,
+        taxa: List[str],
+        alignment_paths: List[str],
+        file_partition: str,
+        fasta_output: str,
+        file_occupancy: str,
+    ) -> None:
+        start_message = dedent(f"""
+            --------------------
+            | General features |
+            --------------------
+            Total number of taxa: {len(taxa)}
+            Total number of alignments: {len(alignment_paths)}
+
+
+            ----------------
+            | Output files |
+            ----------------
+            Partition file output: {file_partition}
+            Concatenated fasta output: {fasta_output}
+            Occupancy report: {file_occupancy}
+        """)
+        print(start_message)
+
+    def get_list_of_taxa_and_records(
+        self, alignment_path: str
+    ) -> Tuple[set, List[SeqRecord]]:
+        records = list(SeqIO.parse(alignment_path, "fasta"))
+        og_taxa = {record.id for record in records}
+        return og_taxa, records
+
+    def create_missing_seq_str(self, records: List[SeqRecord]) -> Tuple[str, int]:
+        """Create a placeholder string for sequences with missing taxa."""
+        if not records:
+            print(f"No sequence records found. Exiting...")
+            sys.exit(2)
+
+        og_len = len(records[0].seq)
+        missing_seq = '?' * og_len
+        return missing_seq, og_len
+
+    def process_taxa_sequences(
+        self,
+        records: List[SeqRecord],
+        taxa: List[str],
+        concatenated_seqs: Dict[str, List[str]],
+        missing_seq: str,
+    ) -> None:
+        present_taxa = {record.id for record in records}
+        missing_taxa = set(taxa) - present_taxa
+
+        # Add sequences for present taxa
+        for record in records:
+            concatenated_seqs[record.id].append(str(record.seq))
+
+        # Add missing sequences for missing taxa
+        for taxon in missing_taxa:
+            concatenated_seqs[taxon].append(missing_seq)
+
+    def add_to_partition_info(
+        self,
+        partition_info: List[str],
+        og_len: int,
+        field_one: str,
+        fasta: str,
+        first_len: int,
+        second_len: int,
+    ) -> Tuple[List[str], int, int]:
+        second_len += og_len
+        partition_info.append(f"{field_one}, {fasta}={first_len}-{second_len}\n")
+        return partition_info, second_len + 1, second_len
+
+    def add_to_occupancy_info(
+        self,
+        occupancy_info: List[str],
+        present_taxa: set,
+        taxa: List[str],
+        fasta: str,
+    ) -> List[str]:
+        missing_taxa = sorted(set(taxa) - present_taxa)
+        num_present = len(present_taxa)
+        num_missing = len(missing_taxa)
+        percent_occupancy = num_present / len(taxa)
+        occupancy_info.append(f"{fasta}\t{num_present}\t{num_missing}\t{percent_occupancy:.4f}\t{';'.join(missing_taxa)}\n")
+        return occupancy_info
+
+    def fasta_file_write(self, fasta_output: str, concatenated_seqs: Dict[str, List[str]]) -> None:
+        """Write concatenated sequences to FASTA file with buffered I/O."""
+        # Use larger buffer for better I/O performance
+        with open(fasta_output, "w", buffering=8192) as final_fasta_file:
+            for taxon, sequences in concatenated_seqs.items():
+                # Join sequences once instead of in the write statement
+                concatenated = ''.join(sequences)
+                final_fasta_file.write(f">{taxon}\n{concatenated}\n")
+
+    def write_occupancy_or_partition_file(self, info: List[str], output_file_name: str) -> None:
+        with open(output_file_name, "w") as f:
+            f.writelines(info)
+
+    @staticmethod
+    def _process_alignment_file(alignment_path: str, taxa: List[str]) -> Tuple[str, Dict[str, str], set, int]:
+        """Process a single alignment file and return its data."""
+        records = list(SeqIO.parse(alignment_path, "fasta"))
+        present_taxa = {record.id for record in records}
+
+        if not records:
+            return alignment_path, {}, present_taxa, 0
+
+        og_len = len(records[0].seq)
+        missing_seq = '?' * og_len
+
+        # Create sequence dict for this alignment
+        seq_dict = {}
+        for taxon in taxa:
+            if taxon in present_taxa:
+                # Find the sequence for this taxon
+                for record in records:
+                    if record.id == taxon:
+                        seq_dict[taxon] = str(record.seq)
+                        break
+            else:
+                seq_dict[taxon] = missing_seq
+
+        return alignment_path, seq_dict, present_taxa, og_len
+
+    def create_concatenation_matrix(self, alignment_list_path: str, prefix: str) -> None:
+        alignment_paths = self.read_alignment_paths(alignment_list_path)
+        taxa = self.get_taxa_names(alignment_paths)
+
+        # Create output directory if needed
+        output_dir = os.path.dirname(prefix)
+        if output_dir and not os.path.exists(output_dir):
+            os.makedirs(output_dir, exist_ok=True)
+
+        # Assign output file names
+        file_partition = f"{prefix}.partition"
+        fasta_output = f"{prefix}.fa"
+        file_occupancy = f"{prefix}.occupancy"
+
+        self.print_start_message(taxa, alignment_paths, file_partition, fasta_output, file_occupancy)
+
+        # Initialize placeholders for partition info
+        first_len, second_len = 1, 0
+        partition_info, occupancy_info = [], []
+        concatenated_seqs = defaultdict(list)
+
+        # Process alignment files in parallel if there are many
+        if len(alignment_paths) > 2:
+            with ProcessPoolExecutor(max_workers=min(mp.cpu_count(), 8)) as executor:
+                process_func = partial(self._process_alignment_file, taxa=taxa)
+                # Keep results indexed by path to maintain order
+                futures = {executor.submit(process_func, path): path for path in alignment_paths}
+                results = {}
+
+                for future in as_completed(futures):
+                    path = futures[future]
+                    results[path] = future.result()
+
+                # Process results in original order
+                for alignment_path in alignment_paths:
+                    _, seq_dict, present_taxa, og_len = results[alignment_path]
+
+                    # Add sequences to concatenated dict
+                    for taxon in taxa:
+                        concatenated_seqs[taxon].append(seq_dict[taxon])
+
+                    # Add to partition and occupancy info
+                    partition_info, first_len, second_len = self.add_to_partition_info(
+                        partition_info, og_len, "AUTO", alignment_path, first_len, second_len
+                    )
+                    occupancy_info = self.add_to_occupancy_info(occupancy_info, present_taxa, taxa, alignment_path)
+        else:
+            # Process sequentially for small datasets
+            for alignment_path in alignment_paths:
+                present_taxa, records = self.get_list_of_taxa_and_records(alignment_path)
+                missing_seq, og_len = self.create_missing_seq_str(records)
+
+                # Process taxa sequences and add to the concatenated sequences
+                self.process_taxa_sequences(records, taxa, concatenated_seqs, missing_seq)
+
+                # Add to partition and occupancy info
+                partition_info, first_len, second_len = self.add_to_partition_info(
+                    partition_info, og_len, "AUTO", alignment_path, first_len, second_len
+                )
+                occupancy_info = self.add_to_occupancy_info(occupancy_info, present_taxa, taxa, alignment_path)
+
+        # Convert defaultdict to regular dict for writing
+        if isinstance(concatenated_seqs, defaultdict):
+            concatenated_seqs = dict(concatenated_seqs)
+
+        # Write output files
+        self.fasta_file_write(fasta_output, concatenated_seqs)
+        self.write_occupancy_or_partition_file(occupancy_info, file_occupancy)
+        self.write_occupancy_or_partition_file(partition_info, file_partition)
+
+        print("Complete!\n")

phykit/services/alignment/dna_threader.py

@@ -0,0 +1,145 @@
+import sys
+from typing import Dict, List
+import numpy as np
+from Bio import SeqIO
+from Bio.Seq import Seq
+from .base import Alignment
+
+
+class DNAThreader(Alignment):
+    """
+    Threads DNA on top of protein alignment
+    """
+
+    def __init__(self, args) -> None:
+        self.process_args(args)
+
+    def process_args(self, args):
+        self.remove_stop_codon = args.stop
+        self.protein_file_path = args.protein
+        self.nucleotide_file_path = args.nucleotide
+        self.clipkit_log_file = args.clipkit_log_file
+
+    @property
+    def clipkit_log_data(self) -> List[List[str]]:
+        if self.clipkit_log_file:
+            with open(self.clipkit_log_file) as f:
+                return [line.rstrip("\n").split(" ") for line in f.readlines()]
+        return None
+
+    def run(self) -> None:
+        prot_records = SeqIO.parse(self.protein_file_path, "fasta")
+        pal2nal = self.thread(prot_records)
+
+        for gene_id, sequence in pal2nal.items():
+            print(f">{gene_id}")
+            print(f"{sequence}")
+
+    def create_mask(self, length: int) -> List[bool]:
+        if not self.clipkit_log_data:
+            return [True] * length
+
+        # create a mask that replicates the 'keep' and 'remove' status 3 times for each amino acid
+        return [
+            True if row[1] == "keep" else False
+            for row in self.clipkit_log_data
+            for _ in range(3)
+        ]
+
+    def normalize_p_seq(self, p_seq: Seq) -> str:
+        # triplicate each amino acid
+        return ''.join(c * 3 for c in p_seq)
+
+    def normalize_n_seq(self, n_seq: Seq, p_seq: Seq) -> str:
+        # Pre-split codons for faster access
+        codons = [str(n_seq[i:i+3]) for i in range(0, len(n_seq), 3)]
+        normalized_n_seq = []
+        gap_chars = {'-', '?', '*', 'X', 'x'}
+
+        codon_idx = 0
+        for aa in p_seq:
+            if aa in gap_chars:
+                normalized_n_seq.append("---")
+            else:
+                if codon_idx < len(codons):
+                    normalized_n_seq.append(codons[codon_idx])
+                    codon_idx += 1
+                else:
+                    normalized_n_seq.append("---")  # fallback in case of misalignment
+
+        return ''.join(normalized_n_seq)
+
+    def thread(self, prot_records) -> Dict[str, str]:
+        pal2nal = dict()
+        prot_dict = SeqIO.to_dict(prot_records)
+
+        if not prot_dict:
+            print("Protein file is empty or incorrectly formatted.")
+            sys.exit(2)
+
+        # Pre-load nucleotide sequences only for proteins we have
+        nucl_records = {}
+        for record in SeqIO.parse(self.nucleotide_file_path, "fasta"):
+            if record.id in prot_dict:
+                nucl_records[record.id] = record
+
+        length = len(next(iter(prot_dict.values())).seq)
+        keep_mask = self.create_mask(length * 3)
+
+        # Convert keep_mask to numpy array for faster operations
+        keep_mask_arr = np.array(keep_mask)
+        gap_chars = {'-', '?', '*', 'X', 'x'}
+
+        for gene_id, protein_seq_record in prot_dict.items():
+            try:
+                if gene_id not in nucl_records:
+                    print(f"Nucleotide sequence for {gene_id} not found.")
+                    sys.exit(2)
+
+                p_seq = protein_seq_record.seq
+                n_seq = nucl_records[gene_id].seq
+
+                # Get normalized sequences
+                normalized_p_seq = self.normalize_p_seq(p_seq)
+                normalized_n_seq = self.normalize_n_seq(n_seq, normalized_p_seq)
+
+                # Convert to numpy arrays for faster operations
+                p_arr = np.array(list(normalized_p_seq), dtype='U1')
+                n_arr = np.array(list(normalized_n_seq), dtype='U1')
+
+                # Create mask for non-gap positions in protein
+                non_gap_mask_protein = ~np.isin(p_arr, list(gap_chars))
+
+                # Expand protein mask to nucleotide positions (each AA = 3 nucleotides)
+                non_gap_mask = np.repeat(non_gap_mask_protein, 3)
+
+                # Ensure masks have same shape
+                min_len = min(len(non_gap_mask), len(keep_mask_arr), len(n_arr))
+                non_gap_mask = non_gap_mask[:min_len]
+                keep_mask_arr_trimmed = keep_mask_arr[:min_len]
+                n_arr = n_arr[:min_len]
+
+                # Combine masks
+                final_mask = keep_mask_arr_trimmed & non_gap_mask
+
+                # Apply masks and build result
+                result = np.where(final_mask, n_arr, '-')
+
+                # Handle stop codon if needed
+                if self.remove_stop_codon and p_seq[-1] == "*":
+                    # Find the last 3 positions that were kept
+                    kept_indices = np.where(keep_mask_arr_trimmed)[0]
+                    if len(kept_indices) >= 3:
+                        last_3_indices = kept_indices[-3:]
+                        for idx in last_3_indices:
+                            if idx < len(result) and idx < len(n_arr):
+                                result[idx] = n_arr[idx]
+
+                # Only keep positions marked in keep_mask
+                pal2nal[gene_id] = ''.join(result[keep_mask_arr_trimmed[:len(result)]])
+
+            except KeyError:
+                print(f"Nucleotide sequence for {gene_id} not found.")
+                sys.exit(2)
+
+        return pal2nal

phykit/services/alignment/evolutionary_rate_per_site.py

@@ -0,0 +1,85 @@
+from collections import Counter
+from typing import List, Dict
+import numpy as np
+
+from Bio.Align import MultipleSeqAlignment
+
+from .base import Alignment
+
+
+class EvolutionaryRatePerSite(Alignment):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self):
+        alignment, _, is_protein = self.get_alignment_and_format()
+        pic_values = self.calculate_evolutionary_rate_per_site(alignment)
+
+        for idx, value in enumerate(pic_values):
+            print(f"{idx + 1}\t{round(value, 4)}")
+
+    def process_args(self, args):
+        return dict(alignment_file_path=args.alignment)
+
+    def remove_gap_characters(self, seq: str, gap_chars: List[str]) -> str:
+        return ''.join([char for char in seq if char not in gap_chars]).upper()
+
+    def get_number_of_occurrences_per_character(
+        self,
+        alignment: MultipleSeqAlignment,
+        idx: int,
+        gap_chars: List[str]
+    ) -> Dict[str, int]:
+        seq_at_position = alignment[:, idx]
+        clean_seq = self.remove_gap_characters(seq_at_position, gap_chars)
+
+        return Counter(clean_seq)
+
+    def calculate_pic(
+        self,
+        num_occurrences: Dict[str, int],
+    ) -> float:
+        total_frequencies = sum(num_occurrences.values())
+        sum_of_frequencies = sum(
+            (frequency / total_frequencies) ** 2
+            for frequency in num_occurrences.values()
+        )
+        return 1 - sum_of_frequencies
+
+    def calculate_evolutionary_rate_per_site(
+        self,
+        alignment: MultipleSeqAlignment,
+    ) -> List[float]:
+        aln_len = alignment.get_alignment_length()
+        gap_chars = set(self.get_gap_chars())
+
+        # Convert alignment to numpy array for vectorized operations
+        alignment_array = np.array([
+            [c.upper() for c in str(record.seq)]
+            for record in alignment
+        ], dtype='U1')
+
+        pic_values = []
+
+        # Process each column
+        for col_idx in range(aln_len):
+            column = alignment_array[:, col_idx]
+
+            # Filter out gaps
+            non_gap_mask = ~np.isin(column, list(gap_chars))
+            filtered_column = column[non_gap_mask]
+
+            if len(filtered_column) > 0:
+                # Count occurrences using numpy
+                unique_chars, counts = np.unique(filtered_column, return_counts=True)
+                total_frequencies = len(filtered_column)
+
+                # Calculate PIC (Probability of Identical Characters)
+                sum_of_frequencies = np.sum((counts / total_frequencies) ** 2)
+                pic = 1 - sum_of_frequencies
+            else:
+                pic = 0
+
+            pic_values.append(pic)
+
+        return pic_values

phykit/services/alignment/faidx.py

@@ -0,0 +1,21 @@
+from typing import Dict
+
+from Bio import SeqIO
+
+from .base import Alignment
+
+
+class Faidx(Alignment):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self) -> None:
+        record_dict = SeqIO.index(self.fasta, "fasta")
+
+        # Split entries and iterate
+        for e in map(str.strip, self.entry.split(",")):
+            record = record_dict[e]
+            print(f">{record.name}\n{record.seq}")
+
+    def process_args(self, args) -> Dict[str, str]:
+        return dict(fasta=args.fasta, entry=args.entry)

phykit/services/alignment/gc_content.py

@@ -0,0 +1,94 @@
+from enum import Enum
+import re
+import sys
+from typing import Dict, Tuple
+from collections import Counter
+import numpy as np
+
+from Bio.Align import MultipleSeqAlignment
+
+from .base import Alignment
+from ...helpers.files import get_alignment_and_format
+
+
+class FileFormat(Enum):
+    fasta = "fasta"
+    clustal = "clustal"
+    maf = "maf"
+    mauve = "mauve"
+    phylip = "phylip"
+    phylip_seq = "phylip-sequential"
+    stockholm = "stockholm"
+
+
+class GCContent(Alignment):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self):
+        records, _, is_protein = get_alignment_and_format(self.fasta)
+
+        if is_protein:
+            print("GC content can't be calculated for protein sequences")
+            sys.exit(2)
+
+        if self.verbose:
+            self.calculate_gc_per_sequence(records)
+        else:
+            self.calculate_gc_total(records)
+
+    def process_args(self, args) -> Dict[str, str]:
+        return dict(fasta=args.fasta, verbose=args.verbose)
+
+    def calculate_gc_per_sequence(self, records: MultipleSeqAlignment) -> None:
+        gap_chars = set(self.get_gap_chars())
+
+        for record in records:
+            # Convert to numpy array for faster operations
+            seq_arr = np.array(list(str(record.seq).upper()), dtype='U1')
+
+            # Filter out gaps
+            non_gap_mask = ~np.isin(seq_arr, list(gap_chars))
+            cleaned_seq = seq_arr[non_gap_mask]
+
+            if len(cleaned_seq) > 0:
+                # Count G and C
+                gc_count = np.sum((cleaned_seq == 'G') | (cleaned_seq == 'C'))
+                gc_content = gc_count / len(cleaned_seq)
+            else:
+                gc_content = 0
+
+            try:
+                print(f"{record.id}\t{round(gc_content, 4)}")
+            except BrokenPipeError:
+                pass
+
+    def calculate_gc_total(self, records: MultipleSeqAlignment) -> None:
+        gap_chars = set(self.get_gap_chars())
+
+        # Combine all sequences into one array
+        all_seqs = [list(str(record.seq).upper()) for record in records]
+        combined_arr = np.concatenate([np.array(seq, dtype='U1') for seq in all_seqs])
+
+        # Filter out gaps
+        non_gap_mask = ~np.isin(combined_arr, list(gap_chars))
+        cleaned_seq = combined_arr[non_gap_mask]
+
+        if len(cleaned_seq) > 0:
+            # Count G and C
+            gc_count = np.sum((cleaned_seq == 'G') | (cleaned_seq == 'C'))
+            gc_content = round(gc_count / len(cleaned_seq), 4)
+            print(gc_content)
+        else:
+            print(
+                "Input file has an unacceptable format. Please check input file argument."
+            )
+            sys.exit(2)
+
+    def remove_gaps_and_count_gc(self, seq: str) -> Tuple[str, float]:
+        gap_chars = self.get_gap_chars()
+        pattern = "[" + "".join(re.escape(char) for char in gap_chars) + "]"
+        cleaned_seq = re.sub(pattern, "", seq)
+        gc_count = Counter(cleaned_seq.upper())["G"] + Counter(cleaned_seq.upper())["C"]
+
+        return cleaned_seq, gc_count