20251125

2025-11-25 00:28:51 +08:00
commit eb3f16c30e
406 changed files with 91653 additions and 0 deletions
@@ -0,0 +1,146 @@
+#! /usr/bin/env python3
+import os
+from Bio import SeqIO
+from Bio.Seq import Seq
+from Bio.SeqRecord import SeqRecord
+from collections import defaultdict
+import argparse
+
+
+def get_sequence_lengths(fasta_files):
+    """
+    get the lengths of sequences in each FASTA file
+    Assumes all sequences in a file have the same length
+    """
+    file_lengths = {}
+    for fasta_file in fasta_files:
+        try:
+            with open(fasta_file, "r") as f:
+                for record in SeqIO.parse(f, "fasta"):
+                    # get length of the first sequence
+                    file_lengths[fasta_file] = len(record.seq)
+                    break
+        except Exception as e:
+            print(f"Error reading file {fasta_file}: {e}")
+            file_lengths[fasta_file] = 0
+
+    return file_lengths
+
+
+def concatenate_fasta_files(fasta_files, output_file):
+    """
+    Concatenate sequences from multiple FASTA files by name, using "-" for missing sequences.
+    """
+    # Get the sequence lengths for each file
+    file_lengths = get_sequence_lengths(fasta_files)
+
+    # Store all sequence names and their corresponding content
+    sequences_dict = defaultdict(dict)
+    all_sequence_names = set()
+
+    # Read sequences from each file
+    for i, fasta_file in enumerate(fasta_files):
+        try:
+            with open(fasta_file, "r") as f:
+                for record in SeqIO.parse(f, "fasta"):
+                    seq_name = record.id
+                    sequences_dict[seq_name][i] = str(record.seq)
+                    all_sequence_names.add(seq_name)
+        except Exception as e:
+            print(f"Error reading file {fasta_file}: {e}")
+
+    # Create concatenated sequences
+    concatenated_sequences = []
+
+    for seq_name in sorted(all_sequence_names):
+        concatenated_seq = []
+
+        for i, fasta_file in enumerate(fasta_files):
+            if i in sequences_dict[seq_name]:
+                # This file has the sequence, add it directly
+                concatenated_seq.append(sequences_dict[seq_name][i])
+            else:
+                # This file is missing the sequence, use "-" to fill the gap
+                gap_length = file_lengths[fasta_file]
+                concatenated_seq.append("-" * gap_length)
+
+        # Concatenate all parts of the sequence
+        full_sequence = "".join(concatenated_seq)
+
+        # Create a new sequence record
+
+        new_record = SeqRecord(
+            Seq(full_sequence),
+            id=seq_name,
+            description=f"concatenated_from_{len(fasta_files)}_files",
+        )
+        concatenated_sequences.append(new_record)
+
+    with open(output_file, "w") as output_handle:
+        SeqIO.write(concatenated_sequences, output_handle, "fasta")
+
+    print(
+        f"Successfully concatenate {len(concatenated_sequences)} sequences to {output_file}"
+    )
+    print(f"Input file count: {len(fasta_files)}")
+
+    # Output statistics
+    for i, fasta_file in enumerate(fasta_files):
+        seq_count = sum(1 for seqs in sequences_dict.values() if i in seqs)
+        print(
+            f"File {i + 1}: {os.path.basename(fasta_file)} - Sequence count: {seq_count}, Sequence length: {file_lengths[fasta_file]}."
+        )
+
+    print(f"Total output sequence length: {len(concatenated_sequences[0].seq)}.")
+
+
+def get_fasta_files_from_directory(directory, extensions):
+    """
+    get all FASTA files from a directory with specified extensions
+    """
+    fasta_files = []
+    for filename in os.listdir(directory):
+        if any(filename.endswith(ext) for ext in extensions):
+            fasta_files.append(os.path.join(directory, filename))
+    return sorted(fasta_files)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Concatenate multiple FASTA files by sequence names."
+    )
+    parser.add_argument("-i", "--input", nargs="+", help="Input FASTA file list")
+    parser.add_argument("-d", "--directory", help="Directory containing FASTA files")
+    parser.add_argument("-o", "--output", required=True, help="Output file")
+    parser.add_argument(
+        "-e",
+        "--extensions",
+        nargs="+",
+        default=[".fasta", ".fa", ".fna"],
+        help="FASTA file extensions to look for in directory",
+    )
+
+    args = parser.parse_args()
+
+    # 获取输入文件
+    if args.directory:
+        fasta_files = get_fasta_files_from_directory(args.directory, args.extensions)
+        if not fasta_files:
+            print(
+                f"Cannot find FASTA files in {args.directory} with extensions {args.extensions}"
+            )
+            return
+    elif args.input:
+        fasta_files = args.input
+    else:
+        print("Please specify input files or directory")
+        return
+
+    print(f"Found {len(fasta_files)} FASTA files:")
+    
+    # Perform concatenation
+    concatenate_fasta_files(fasta_files, args.output)
+
+
+if __name__ == "__main__":
+    main()