-
Notifications
You must be signed in to change notification settings - Fork 6
/
Copy pathutils.py
142 lines (106 loc) · 4.19 KB
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
import os
import random
import torch
import numpy as np
import pandas as pd
from rdkit import Chem
from rdkit.Chem.rdmolops import PatternFingerprint
UID_COL = 'UniprotID'
SEQ_COL = 'sequence'
RXN_COL = 'CANO_RXN_SMILES'
def tranverse_folder(folder):
all_files = []
for root, dirs, files in os.walk(folder):
for file in files:
all_files.append(os.path.join(root, file))
return all_files
def get_rdkit_mol(mol):
if isinstance(mol, str):
mol = Chem.MolFromSmiles(mol)
return mol
def remove_stereo(mol):
mol = get_rdkit_mol(mol)
Chem.RemoveStereochemistry(mol)
return Chem.MolToSmiles(mol)
def cano_smiles(smiles, remove_stereo=False):
mol = Chem.MolFromSmiles(smiles)
if not mol:
return None
if remove_stereo:
Chem.RemoveStereochemistry(mol)
return Chem.MolToSmiles(mol)
def cano_rxn(rxn, exchange_pos=False, remove_stereo=False):
data = rxn.split('>')
reactants = data[0].split('.')
reactants = [cano_smiles(each, remove_stereo) for each in reactants]
products = data[-1].split('.')
products = [cano_smiles(each, remove_stereo) for each in products]
reactants = sorted(reactants)
products = sorted(products)
if exchange_pos:
new_rxn = f"{'.'.join(products)}>>{'.'.join(reactants)}"
else:
new_rxn = f"{'.'.join(reactants)}>>{'.'.join(products)}"
return new_rxn
def load_feature(model_conf):
if model_conf.auto_load_data:
feat_dir = os.path.join(model_conf.data_dir, f'split_{model_conf.data_split_by}', 'npy_feature', f'{model_conf.rxn_feat_type}_{model_conf.enz_feat_type}')
train_feat_path = os.path.join(feat_dir, 'train_feature.npy')
valid_feat_path = os.path.join(feat_dir, 'valid_feature.npy')
test_feat_path = os.path.join(feat_dir, 'test_feature.npy')
train_feat = np.load(train_feat_path)
valid_feat = np.load(valid_feat_path)
test_feat = np.load(test_feat_path)
else:
train_feat = np.load(model_conf.train_feat)
valid_feat = np.load(model_conf.valid_feat)
test_feat = np.load(model_conf.test_feat)
return train_feat, valid_feat, test_feat
def get_data_path(model_conf, data_type='train'):
return os.path.join(model_conf.data_dir, f'split_{model_conf.data_split_by}', 'csv', f"{data_type}.csv")
def load_origin_data(model_conf):
if model_conf.auto_load_data:
train_path = get_data_path(model_conf, 'train')
valid_path = get_data_path(model_conf, 'valid')
test_path = get_data_path(model_conf, 'test')
origin_train_data = pd.read_csv(train_path)
origin_valid_data = pd.read_csv(valid_path)
origin_test_data = pd.read_csv(test_path)
else:
origin_train_data = pd.read_csv(model_conf.train_path)
origin_valid_data = pd.read_csv(model_conf.valid_path)
origin_test_data = pd.read_csv(model_conf.test_path)
return origin_train_data, origin_valid_data, origin_test_data
def seed_everything(seed=42):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.set_num_threads(5)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
def calc_rxn_center_fp(rxn_center):
# rxn_center就是template
prod_c = rxn_center.split('>>')[-1]
reac_c = rxn_center.split('>>')[0]
prod_mol = Chem.MolFromSmarts(prod_c)
reac_mol = Chem.MolFromSmarts(reac_c)
prod_fp = np.array(PatternFingerprint(prod_mol, 512))
reac_fp = np.array(PatternFingerprint(reac_mol, 512))
return reac_fp, prod_fp
def remove_nan(data):
if isinstance(data, set):
data = {each for each in data if not pd.isna(each)}
elif isinstance(data, list):
data = [each for each in data if not pd.isna(each)]
return data
def check_dir(path):
folder = os.path.dirname(path) if '.' in path.split('/')[-1] else path
if not os.path.exists(folder):
os.makedirs(folder)
print(f'Make new directory: {folder}')
def get_substrate_from_rxn(rxn):
rcts = rxn.split('>>')[0].split('.')
rcts = sorted(rcts, key=lambda x: len(x), reverse=True)
return rcts[0]