data2pkl_MCP.py

# -*- coding:utf-8 -*-
import os
import sys
import pickle
import numpy as np

def pkl_save(pkl_data, save_dir):
    fp = open(save_dir, 'wb')
    pickle.dump(pkl_data, fp)
    fp.close()
    return print('Data saved in file: ' + save_dir)

def pssm_load(pssm_dir, start=2, end=22, tail='.pssm'):
    fp = open(pssm_dir + tail, 'r')
    lines = fp.readlines()
    fp.close()
    pssm = []
    for line in lines:
        temp = line.split()
        if len(temp) == 0:
            continue  
        pssm.append(temp[start:end])
    return np.array(pssm, dtype=np.float32)

def fasta_load(fasta_dir, tail='.fasta'):
    fp = open(fasta_dir + tail, 'r')
    lines = fp.readlines()
    fp.close()
    sequence = ''
    for line in lines[1:]:
        sequence = sequence + line.split()[0]
    return sequence

def files2pkl(id_list, pkl_dir, pkl_name, fasta_dir, pssm_dir):
    pkl = []
    report_i = 0
    for pdbid in id_list:
        report_i = report_i + 1
        sequence = fasta_load(fasta_dir+pdbid)
        pssm_ss3_acc = pssm_load(pssm_dir+pdbid, start=0, end=46, tail='.feat')
        pssm = pssm_ss3_acc[:, 0:20]
        acc = pssm_ss3_acc[:, 40:43]
        ss3 = pssm_ss3_acc[:, 43:46]
        pkl_dict = {b'name': pdbid.encode('utf-8'),
                    b'sequence': sequence.encode('utf-8'),
                    b'PSSM': pssm,
                    b'SS3': ss3,
                    b'ACC': acc}
        pkl.append(pkl_dict)
    if not os.path.exists(pkl_dir):
        os.mkdir(pkl_dir)
    return pkl_save(pkl, pkl_dir + '/' + pkl_name)

temp=sys.argv[1]

id_list = temp.split()
files2pkl(id_list, './example/', sys.argv[1]+'.pkl', './example/', './example/')