需要在linux系统中,要装了seqkit软件

import argparse
import subprocess
import os

def run_command(command, check=True):
    result = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    if result.returncode != 0 and check:
        print(f"Error executing command: {command}\nError message: {result.stderr.decode()}")
    return result.stdout.decode()

def main(prefix, species, column, orthogroups_file, output_dir, pvalue_threshold, gene_prefix, pep_file, seqkit_path):
    sample_number = column - 2
    sample_full_id = f"{species}<{sample_number}>"
    
    significant_dir = os.path.join(output_dir, "significant")
    all_dir = os.path.join(output_dir, "all")
    
    os.makedirs(significant_dir, exist_ok=True)
    os.makedirs(all_dir, exist_ok=True)
    
    # File names based on species
    expanded_file = os.path.join(all_dir, f"{species}.expanded")
    contracted_file = os.path.join(all_dir, f"{species}.contracted")
    significant_trees_file = os.path.join(significant_dir, f"{species}_significant_trees.tre")
    significant_ogs_file = os.path.join(significant_dir, f"{species}_significant.ogs")
    p0_01_significant_ogs_file = os.path.join(significant_dir, "p0.01_significant.ogs")
    sample_p0_01_significant_ogs_file = os.path.join(significant_dir, f"{species}_p0.01_significant.ogs")
    expanded_significant_file = os.path.join(significant_dir, f"{species}.expanded.significant")
    contracted_significant_file = os.path.join(significant_dir, f"{species}.contracted.significant")
    expanded_significant_genes_file = os.path.join(significant_dir, f"{species}.expanded.significant.genes")
    contracted_significant_genes_file = os.path.join(significant_dir, f"{species}.contracted.significant.genes")
    expanded_significant_pep_file = os.path.join(significant_dir, f"{species}.expanded.significant.pep.fa")
    contracted_significant_pep_file = os.path.join(significant_dir, f"{species}.contracted.significant.pep.fa")
    expanded_genes_file = os.path.join(all_dir, f"{species}.expanded.genes")
    contracted_genes_file = os.path.join(all_dir, f"{species}.contracted.genes")
    expanded_pep_file = os.path.join(all_dir, f"{species}.expanded.pep.fa")
    contracted_pep_file = os.path.join(all_dir, f"{species}.contracted.pep.fa")

    # Commands
    commands = [
        # Extract expanded and contracted orthogroups
        f"cat {prefix}_change.tab | cut -f1,{column} | grep '+[1-9]' > {expanded_file}",
        f"cat {prefix}_change.tab | cut -f1,{column} | grep '-' > {contracted_file}",
        
        # Extract significant trees
        f"grep '{sample_full_id}\\*' {prefix}_asr.tre > {significant_trees_file}",
        f"grep -E -o 'OG[0-9]+' {significant_trees_file} > {significant_ogs_file}",
        f"awk '$2 <{pvalue_threshold} {{print $1}}' {prefix}_family_results.txt > {p0_01_significant_ogs_file}",
        f"grep -f {significant_ogs_file} {p0_01_significant_ogs_file} > {sample_p0_01_significant_ogs_file}",
        f"grep -f {sample_p0_01_significant_ogs_file} {expanded_file} | cut -f1 > {expanded_significant_file}",
        f"grep -f {sample_p0_01_significant_ogs_file} {contracted_file} | cut -f1 > {contracted_significant_file}",
        
        # Extract genes for significant expanded and contracted orthogroups
        f"grep -f {expanded_significant_file} {orthogroups_file} | sed 's/ /\\n/g' | grep '{gene_prefix}' | sort -k 1.3n | uniq > {expanded_significant_genes_file}",
        f"grep -f {contracted_significant_file} {orthogroups_file} | sed 's/ /\\n/g' | grep '{gene_prefix}' | sort -k 1.3n | uniq > {contracted_significant_genes_file}",
        f"{seqkit_path} grep -f {expanded_significant_genes_file} {pep_file} > {expanded_significant_pep_file}",
        f"{seqkit_path} grep -f {contracted_significant_genes_file} {pep_file} > {contracted_significant_pep_file}",

        # Extract genes for all expanded and contracted orthogroups
        f"cut -f1 {expanded_file} | grep -v '^$' > {expanded_file}.orthogroups",
        f"cut -f1 {contracted_file} | grep -v '^$' > {contracted_file}.orthogroups",
        f"grep -f {expanded_file}.orthogroups {orthogroups_file} | sed 's/ /\\n/g' | grep '{gene_prefix}' | sort -k 1.3n | uniq > {expanded_genes_file}",
        f"grep -f {contracted_file}.orthogroups {orthogroups_file} | sed 's/ /\\n/g' | grep '{gene_prefix}' | sort -k 1.3n | uniq > {contracted_genes_file}",
        f"{seqkit_path} grep -f {expanded_genes_file} {pep_file} > {expanded_pep_file}",
        f"{seqkit_path} grep -f {contracted_genes_file} {pep_file} > {contracted_pep_file}"
    ]

    # Execute commands
    for command in commands:
        run_command(command)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Automate gene family expansion and contraction analysis.")
    parser.add_argument("--prefix", required=True, help="Prefix for the files (e.g., Gamma)")
    parser.add_argument("--species", required=True, help="Species abbreviation (e.g., sample)")
    parser.add_argument("--column", type=int, required=True, help="Column number to extract (e.g., 15)")
    parser.add_argument("--orthogroups_file", required=True, help="Path to the Orthogroups.txt file")
    parser.add_argument("--output_dir", default="genefamilyexpand", help="Output directory")
    parser.add_argument("--pvalue_threshold", type=float, default=0.01, help="P-value threshold for significance (e.g., 0.01)")
    parser.add_argument("--gene_prefix", default="bv", help="Prefix for gene IDs (e.g., 'bv')")
    parser.add_argument("--pep_file", required=True, help="Path to the sample.pep.fa file")
    parser.add_argument("--seqkit_path", default="seqkit", help="Path to the seqkit executable")

    args = parser.parse_args()
    main(args.prefix, args.species, args.column, args.orthogroups_file, args.output_dir, args.pvalue_threshold, args.gene_prefix, args.pep_file, args.seqkit_path)

使用例子

python extract_change.py  --prefix Base --species MLA  --column 13 --orthogroups_file Orthogroups.txt --output_dir genefamilyexpand3 --pvalue_threshold 0.01 --gene_prefix ML --pep_file MLA.faa --seqkit_path seqkit

提取指定物种MLA的全部收缩扩张基因和显著收缩扩张基因,分别在两个子文件夹中,--column要输入该物种在Orthogroups.txt列表中的所在列数