CAFE输出结果整理2——提取特定Orthogroups里某个物种的基因名

比较基因组分析很多时候都是用OrthoFinder生成的Orthogroups来进行分析，然后得到一系列OG的ID，对于不同的物种，我们还需要在orthogroup里找到对应的基因名，来进行后续的富集分析等。
*推荐使用第二个脚本。第一个脚本保留的原因是如果需要批量对文件操作可以根据这个来仿写。

一、根据物种名所在位置提取single copy gene序列文件里的对应物种位置的基因名，只适用于单拷贝基因序列

#!/usr/bin/python
# -*- coding: utf-8 -*-
# conda activate python3
"""
    作者：徐诗芬
    内容：1.根据OrthogroupID列表提取单拷贝Orthogroup文件；
         2.根据OrthogroupID列表名读取文件里的基因名；
         3.根据物种对应name.txt的相同位置来提取对应的基因名
    日期：2021.12.16
"""
import sys
import os
import re

def usage():
    print('Usage: python gene_find.py [targetOG.txt] [Single_copy_gene] [name.txt] [species]')

def file_extract(path):
    # 提取基因蛋白序列文件，或者提取mafft.fa文件等修改匹配条件即可
    file_list = []
    for root, dirs, files in os.walk(path):
        for f in files:
            match = re.match(r'OG[0-9]*.fa', f)  # 如果需要其他类型的文件可以修改这里的匹配条件
            if f == match.group(0):
                file = os.path.join(root, f)
                file_list.append(file)
    return file_list

def fasta2dict(inf):
    # 按行读取序列
    # 输入fasta文件，返回名称，序列
    dict = {}
    for line in inf:
        line = line.strip()
        if line.startswith('>'):
            line = line.strip()
            name = line[1:]
        else:
            dict[name] = dict.get(name, '') + line
    return dict

def key_list(d):
    # 生成一个键值的list
    result = []
    for k, v in d.items():
        result.append(k)
    return result

def flag(species, name_list):
    # 定位物种所在的行号
    for num, line in enumerate(name_list):
        while species in line:
            return num
def main():

    Or_index = open(sys.argv[1], 'rt')       # targetOG.txt, 目标 Orthorgroup ID
    Or_list = []
    for i in Or_index.readlines():
        i = i.strip()
        Or_list.append(i)
    path = os.readlink(sys.argv[2])         # Single_copy_gene # Orthorgroup文件所在路径的软链接，软链接注意不要加/
    name_index = open(sys.argv[3], 'rt')    # Orthorfinder 运行时的物种名顺序，依据结果里的Log.txt文件
    name_list = name_index.readlines()
    species = sys.argv[4]                   # 目标物种名
    outfile = open('gene.txt', 'w+')

    file_list = file_extract(path)          # 提取路径中特定的文件

    for file in file_list:
        file_name = os.path.basename(file).split('.')[0]          # 提取文件名
        if file_name in Or_list:                                  # 判断文件名是否在OrID.txt的列表里面
            with open(file, 'rt') as f:
                d = fasta2dict(f)
                gene_list = key_list(d)                           # 提取该OR文件的所有基因名
                i = flag(species, name_list)                      # 提取物种对应gene_list的index
                ln = gene_list[i]                                 # 在这个文件所生成gene_list里找到属于该物种的基因名
                #sequence = d[ln]
                outfile.write(ln + '\n')
                #outfile.write(sequence + '\n')                   #这里还可以输出对应的序列

    outfile.close()
    Or_index.close()
    name_index.close()



try:
    main()
except IndexError:
    usage()

二、在多拷贝序列或者Orthogroup里面根据物种名提取对应Ortho_ID所包含该物种的基因名，直接从Orthogroups.tsv文件里提取。适用于提取该物种扩张和收缩的基因家族里的基因名。

import sys
from itertools import islice

def usage():
    print('Usage: python gene_find2.py [targetOG.txt] [species] [Orthogroups.tsv] [output.txt]')

def fasta2dict(inf):
    # 按行读取序列
    # 输入fasta文件，返回名称，序列
    dict = {}
    for line in inf:
        line = line.strip()
        if line.startswith('>'):
            line = line.strip()
            name = line[1:]
        else:
            dict[name] = dict.get(name, '') + line
    return dict

def main():

    Or_index = open(sys.argv[1], 'rt')       # Sl.txt, 目标 Orthorgroup ID
    Or_list = []
    for i in Or_index.readlines():
        i = i.strip()
        Or_list.append(i)
    species = sys.argv[2]                  # 目标物种名
    outfile = open(sys.argv[4], 'w+')            # 输出文件

    OG_dict = {}
    with open(sys.argv[3], 'rt') as f:
        for head in islice(f, 0, 1):           # 读取文件第一行信息
            head_list = head.strip().split('\t')
        for line in islice(f, 1, None):
            line = line.strip().split('\t')
            name = line[0]
            OG_dict[name] = line[1:]

    for Or in Or_list:
        if Or in list(OG_dict.keys()):                      # 横向找OG的ID所在行
            line = OG_dict[Or]                              # Or所对应的值依旧是一个列表
            flag = head_list.index(species)                 # 找该物种的索引位置
            if len(line) < flag:
                continue  # 设置基因家族收缩时该OG没有基因名的情况,就跳过
            else:
                gene_list = line[flag - 1]  # 找到物种对应该Or的基因集
                for i in gene_list.split(', '):
                    outfile.write(i + '\n')  # 逐一输出基因名

    outfile.close()
    Or_index.close()

try:
    main()
except IndexError:
    usage()