434 lines
20 KiB
Python
434 lines
20 KiB
Python
import os
|
||
import pickle
|
||
import numpy as np
|
||
import pandas as pd
|
||
from tensorflow.keras.models import load_model
|
||
from .features.sequence import SequenceFeatureExtractor
|
||
from .features.cnn_input import CNNInputProcessor
|
||
from .utils import extract_window_sequences
|
||
import matplotlib.pyplot as plt
|
||
import joblib
|
||
|
||
|
||
class PRFPredictor:
|
||
|
||
def __init__(self, model_dir=None):
|
||
"""
|
||
初始化PRF预测器
|
||
|
||
Args:
|
||
model_dir: 模型目录路径（可选）
|
||
"""
|
||
if model_dir is None:
|
||
from pkg_resources import resource_filename
|
||
model_dir = resource_filename('FScanpy', 'pretrained')
|
||
|
||
try:
|
||
# 加载模型
|
||
self.gb_model = self._load_pickle(os.path.join(model_dir, 'GradientBoosting_all.pkl'))
|
||
|
||
self.cnn_model = self._load_pickle(os.path.join(model_dir, 'BiLSTM-CNN_all.pkl'))
|
||
|
||
self.voting_model = self._load_pickle(os.path.join(model_dir, 'Voting_all.pkl'))
|
||
|
||
# 初始化特征提取器和CNN处理器，使用与训练时相同的序列长度
|
||
self.gb_seq_length = 33 # HistGradientBoosting使用的序列长度
|
||
self.cnn_seq_length = 399 # BiLSTM-CNN使用的序列长度
|
||
|
||
# 初始化特征提取器和CNN输入处理器
|
||
self.feature_extractor = SequenceFeatureExtractor(seq_length=self.gb_seq_length)
|
||
self.cnn_processor = CNNInputProcessor(max_length=self.cnn_seq_length)
|
||
|
||
except FileNotFoundError as e:
|
||
raise FileNotFoundError(f"无法找到模型文件: {str(e)}")
|
||
except Exception as e:
|
||
raise Exception(f"加载模型出错: {str(e)}")
|
||
|
||
def _load_pickle(self, path):
|
||
return joblib.load(path)
|
||
|
||
def predict_single_position(self, fs_period, full_seq, gb_threshold=0.1):
|
||
'''
|
||
预测单个位置的PRF状态
|
||
|
||
Args:
|
||
fs_period: 33bp序列 (将根据gb_seq_length处理)
|
||
full_seq: 完整序列 (将根据cnn_seq_length处理)
|
||
gb_threshold: GB模型的概率阈值 (默认为0.1)
|
||
Returns:
|
||
dict: 包含预测概率的字典
|
||
'''
|
||
try:
|
||
# 处理序列长度
|
||
if len(fs_period) > self.gb_seq_length:
|
||
fs_period = self.feature_extractor.trim_sequence(fs_period, self.gb_seq_length)
|
||
|
||
# GB模型预测 - 确保输入是二维数组
|
||
try:
|
||
gb_features = self.feature_extractor.extract_features(fs_period)
|
||
|
||
# 检查特征结构并确保是一维数组
|
||
if isinstance(gb_features, np.ndarray):
|
||
# 如果是多维数组，进行扁平化处理
|
||
if gb_features.ndim > 1:
|
||
print(f"警告: 特征是{gb_features.ndim}维数组，进行扁平化处理")
|
||
gb_features = gb_features.flatten()
|
||
|
||
# 明确将特征转换为二维数组，正确形状为(1, n_features)
|
||
gb_features_2d = np.array([gb_features])
|
||
|
||
# 再次检查维度
|
||
if gb_features_2d.ndim != 2:
|
||
raise ValueError(f"处理后特征仍为{gb_features_2d.ndim}维，需要二维数组")
|
||
|
||
gb_prob = self.gb_model.predict_proba(gb_features_2d)[0][1]
|
||
except Exception as e:
|
||
print(f"GB模型预测时出错: {str(e)}")
|
||
# 出错时设置概率为0
|
||
gb_prob = 0.0
|
||
|
||
# 如果GB概率低于阈值，则跳过CNN模型
|
||
if gb_prob < gb_threshold:
|
||
return {
|
||
'GB_Probability': gb_prob,
|
||
'CNN_Probability': 0.0,
|
||
'Voting_Probability': 0.0
|
||
}
|
||
|
||
# CNN模型预测
|
||
try:
|
||
# 首先检查CNN模型的类型 - 通过尝试识别模型类型
|
||
is_sklearn_model = False
|
||
|
||
# 检测模型类型的方法
|
||
if hasattr(self.cnn_model, 'predict_proba'):
|
||
# 这可能是一个scikit-learn模型
|
||
is_sklearn_model = True
|
||
|
||
if is_sklearn_model:
|
||
# 如果是sklearn模型 (如HistGradientBoostingClassifier)，使用与GB相同的特征提取
|
||
# 为CNN模型使用相同的特征提取方法，但从399bp序列中提取
|
||
cnn_features = self.feature_extractor.extract_features(full_seq)
|
||
if isinstance(cnn_features, np.ndarray) and cnn_features.ndim > 1:
|
||
cnn_features = cnn_features.flatten()
|
||
# 转为二维数组
|
||
cnn_features_2d = np.array([cnn_features])
|
||
cnn_pred = self.cnn_model.predict_proba(cnn_features_2d)
|
||
cnn_prob = cnn_pred[0][1]
|
||
else:
|
||
# 假设是深度学习模型，需要三维输入
|
||
cnn_input = self.cnn_processor.prepare_sequence(full_seq)
|
||
# 尝试不同的预测方法
|
||
try:
|
||
# 先尝试不带参数
|
||
cnn_pred = self.cnn_model.predict(cnn_input)
|
||
except TypeError:
|
||
try:
|
||
# 再尝试带verbose参数
|
||
cnn_pred = self.cnn_model.predict(cnn_input, verbose=0)
|
||
except Exception:
|
||
# 最后尝试将输入重塑为2D
|
||
reshaped_input = cnn_input.reshape(1, -1)
|
||
cnn_pred = self.cnn_model.predict(reshaped_input)
|
||
|
||
# 处理预测结果
|
||
if isinstance(cnn_pred, list):
|
||
cnn_pred = cnn_pred[0]
|
||
|
||
# 提取概率值
|
||
if hasattr(cnn_pred, 'shape') and len(cnn_pred.shape) > 1 and cnn_pred.shape[1] > 1:
|
||
cnn_prob = cnn_pred[0][1]
|
||
else:
|
||
cnn_prob = cnn_pred[0][0] if hasattr(cnn_pred[0], '__getitem__') else cnn_pred[0]
|
||
except Exception as e:
|
||
print(f"CNN模型预测时出错: {str(e)}")
|
||
# 出错时设置概率为0
|
||
cnn_prob = 0.0
|
||
|
||
# 投票模型预测
|
||
try:
|
||
# 确保投票模型输入是二维数组 (1, n_features)
|
||
voting_input = np.array([[gb_prob, cnn_prob]])
|
||
voting_prob = self.voting_model.predict_proba(voting_input)[0][1]
|
||
except Exception as e:
|
||
print(f"投票模型预测时出错: {str(e)}")
|
||
# 出错时使用两个模型的平均值
|
||
voting_prob = (gb_prob + cnn_prob) / 2
|
||
|
||
return {
|
||
'GB_Probability': gb_prob,
|
||
'CNN_Probability': cnn_prob,
|
||
'Voting_Probability': voting_prob
|
||
}
|
||
|
||
except Exception as e:
|
||
raise Exception(f"预测过程出错: {str(e)}")
|
||
|
||
def predict_full(self, sequence, window_size=3, gb_threshold=0.1, plot=False):
|
||
"""
|
||
预测完整序列中的PRF位点
|
||
|
||
Args:
|
||
sequence: 输入DNA序列
|
||
window_size: 滑动窗口大小 (默认为3)
|
||
gb_threshold: GB模型概率阈值 (默认为0.1)
|
||
plot: 是否绘制预测结果图表 (默认为False)
|
||
|
||
Returns:
|
||
if plot=False:
|
||
pd.DataFrame: 包含预测结果的DataFrame
|
||
if plot=True:
|
||
tuple: (pd.DataFrame, matplotlib.figure.Figure)
|
||
"""
|
||
if window_size < 1:
|
||
raise ValueError("窗口大小必须大于等于1")
|
||
if gb_threshold < 0:
|
||
raise ValueError("GB阈值必须大于等于0")
|
||
|
||
results = []
|
||
|
||
try:
|
||
# 确保序列为字符串并转换为大写
|
||
sequence = str(sequence).upper()
|
||
|
||
# 滑动窗口预测
|
||
for pos in range(0, len(sequence) - 2, window_size):
|
||
# 提取窗口序列 - 使用与训练时相同的窗口大小
|
||
fs_period, full_seq = extract_window_sequences(sequence, pos)
|
||
|
||
if fs_period is None or full_seq is None:
|
||
continue
|
||
|
||
# 预测并记录结果
|
||
pred = self.predict_single_position(fs_period, full_seq, gb_threshold)
|
||
pred.update({
|
||
'Position': pos,
|
||
'Codon': sequence[pos:pos+3],
|
||
'33bp': fs_period,
|
||
'399bp': full_seq
|
||
})
|
||
results.append(pred)
|
||
|
||
# 创建结果DataFrame
|
||
results_df = pd.DataFrame(results)
|
||
|
||
# 如需绘图
|
||
if plot:
|
||
# 创建图形
|
||
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(15, 10), height_ratios=[2, 1])
|
||
|
||
# 绘制折线图
|
||
ax1.plot(results_df['Position'], results_df['GB_Probability'],
|
||
label='GB模型', alpha=0.7, linewidth=1.5)
|
||
ax1.plot(results_df['Position'], results_df['CNN_Probability'],
|
||
label='CNN模型', alpha=0.7, linewidth=1.5)
|
||
ax1.plot(results_df['Position'], results_df['Voting_Probability'],
|
||
label='投票模型', linewidth=2, color='red')
|
||
|
||
ax1.set_xlabel('序列位置')
|
||
ax1.set_ylabel('移码概率')
|
||
ax1.set_title('移码预测概率')
|
||
ax1.legend()
|
||
ax1.grid(True, alpha=0.3)
|
||
|
||
# 准备热图数据
|
||
positions = results_df['Position'].values
|
||
probabilities = results_df['Voting_Probability'].values
|
||
|
||
# 创建热图矩阵
|
||
heatmap_matrix = np.zeros((1, len(positions)))
|
||
heatmap_matrix[0, :] = probabilities
|
||
|
||
# 绘制热图
|
||
im = ax2.imshow(heatmap_matrix, aspect='auto', cmap='YlOrRd',
|
||
extent=[min(positions), max(positions), 0, 1])
|
||
|
||
# 添加颜色条
|
||
cbar = plt.colorbar(im, ax=ax2)
|
||
cbar.set_label('移码概率')
|
||
|
||
# 设置热图轴标签
|
||
ax2.set_xlabel('序列位置')
|
||
ax2.set_title('移码概率热图')
|
||
ax2.set_yticks([])
|
||
|
||
# 调整布局
|
||
plt.tight_layout()
|
||
|
||
return results_df, fig
|
||
|
||
return results_df
|
||
|
||
except Exception as e:
|
||
raise Exception(f"序列预测过程出错: {str(e)}")
|
||
|
||
def predict_region(self, seq, gb_threshold=0.1):
|
||
'''
|
||
预测区域序列
|
||
|
||
Args:
|
||
seq: 399bp序列或包含399bp序列的DataFrame/Series
|
||
gb_threshold: GB模型概率阈值 (默认为0.1)
|
||
|
||
Returns:
|
||
DataFrame: 包含所有序列预测概率的DataFrame
|
||
'''
|
||
try:
|
||
# 如果输入是DataFrame或Series，转换为列表
|
||
if isinstance(seq, (pd.DataFrame, pd.Series)):
|
||
seq = seq.tolist()
|
||
|
||
# 如果输入是单个字符串，转换为列表
|
||
if isinstance(seq, str):
|
||
seq = [seq]
|
||
|
||
results = []
|
||
for i, seq399 in enumerate(seq):
|
||
try:
|
||
# 从399bp序列中截取中心的33bp (GB模型使用)
|
||
seq33 = self._extract_center_sequence(seq399, target_length=self.gb_seq_length)
|
||
|
||
# GB模型预测 - 确保输入是二维数组
|
||
try:
|
||
gb_features = self.feature_extractor.extract_features(seq33)
|
||
|
||
# 检查特征结构并确保是一维数组
|
||
if isinstance(gb_features, np.ndarray):
|
||
# 如果是多维数组，进行扁平化处理
|
||
if gb_features.ndim > 1:
|
||
print(f"警告: 序列 {i+1} 的特征是{gb_features.ndim}维数组，进行扁平化处理")
|
||
gb_features = gb_features.flatten()
|
||
|
||
# 明确将特征转换为二维数组，正确形状为(1, n_features)
|
||
gb_features_2d = np.array([gb_features])
|
||
|
||
# 再次检查维度
|
||
if gb_features_2d.ndim != 2:
|
||
raise ValueError(f"处理后特征仍为{gb_features_2d.ndim}维，需要二维数组")
|
||
|
||
gb_prob = self.gb_model.predict_proba(gb_features_2d)[0][1]
|
||
except Exception as e:
|
||
print(f"GB模型预测序列 {i+1} 时出错: {str(e)}")
|
||
# 出错时设置概率为0
|
||
gb_prob = 0.0
|
||
|
||
# 如果GB概率低于阈值，添加低概率结果
|
||
if gb_prob < gb_threshold:
|
||
results.append({
|
||
'GB_Probability': gb_prob,
|
||
'CNN_Probability': 0.0,
|
||
'Voting_Probability': 0.0,
|
||
'33bp': seq33,
|
||
'399bp': seq399
|
||
})
|
||
continue
|
||
|
||
# CNN模型预测
|
||
try:
|
||
# 首先检查CNN模型的类型 - 通过尝试识别模型类型
|
||
is_sklearn_model = False
|
||
|
||
# 检测模型类型的方法
|
||
if hasattr(self.cnn_model, 'predict_proba'):
|
||
# 这可能是一个scikit-learn模型
|
||
is_sklearn_model = True
|
||
|
||
if is_sklearn_model:
|
||
# 如果是sklearn模型 (如HistGradientBoostingClassifier)，使用与GB相同的特征提取
|
||
# 为CNN模型使用相同的特征提取方法，但从399bp序列中提取
|
||
cnn_features = self.feature_extractor.extract_features(seq399)
|
||
if isinstance(cnn_features, np.ndarray) and cnn_features.ndim > 1:
|
||
cnn_features = cnn_features.flatten()
|
||
# 转为二维数组
|
||
cnn_features_2d = np.array([cnn_features])
|
||
cnn_pred = self.cnn_model.predict_proba(cnn_features_2d)
|
||
cnn_prob = cnn_pred[0][1]
|
||
else:
|
||
# 假设是深度学习模型，需要三维输入
|
||
cnn_input = self.cnn_processor.prepare_sequence(seq399)
|
||
# 尝试不同的预测方法
|
||
try:
|
||
# 先尝试不带参数
|
||
cnn_pred = self.cnn_model.predict(cnn_input)
|
||
except TypeError:
|
||
try:
|
||
# 再尝试带verbose参数
|
||
cnn_pred = self.cnn_model.predict(cnn_input, verbose=0)
|
||
except Exception:
|
||
# 最后尝试将输入重塑为2D
|
||
reshaped_input = cnn_input.reshape(1, -1)
|
||
cnn_pred = self.cnn_model.predict(reshaped_input)
|
||
|
||
# 处理预测结果
|
||
if isinstance(cnn_pred, list):
|
||
cnn_pred = cnn_pred[0]
|
||
|
||
# 提取概率值
|
||
if hasattr(cnn_pred, 'shape') and len(cnn_pred.shape) > 1 and cnn_pred.shape[1] > 1:
|
||
cnn_prob = cnn_pred[0][1]
|
||
else:
|
||
cnn_prob = cnn_pred[0][0] if hasattr(cnn_pred[0], '__getitem__') else cnn_pred[0]
|
||
except Exception as e:
|
||
print(f"CNN模型预测序列 {i+1} 时出错: {str(e)}")
|
||
# 出错时设置概率为0
|
||
cnn_prob = 0.0
|
||
|
||
# 投票模型预测
|
||
try:
|
||
# 确保投票模型输入是二维数组 (1, n_features)
|
||
voting_input = np.array([[gb_prob, cnn_prob]])
|
||
voting_prob = self.voting_model.predict_proba(voting_input)[0][1]
|
||
except Exception as e:
|
||
print(f"投票模型预测序列 {i+1} 时出错: {str(e)}")
|
||
# 出错时使用两个模型的平均值
|
||
voting_prob = (gb_prob + cnn_prob) / 2
|
||
|
||
results.append({
|
||
'GB_Probability': gb_prob,
|
||
'CNN_Probability': cnn_prob,
|
||
'Voting_Probability': voting_prob,
|
||
'33bp': seq33,
|
||
'399bp': seq399
|
||
})
|
||
|
||
except Exception as e:
|
||
print(f"处理第 {i+1} 个序列时出错: {str(e)}")
|
||
results.append({
|
||
'GB_Probability': 0.0,
|
||
'CNN_Probability': 0.0,
|
||
'Voting_Probability': 0.0,
|
||
'33bp': self._extract_center_sequence(seq399, target_length=self.gb_seq_length) if len(seq399) >= self.gb_seq_length else seq399,
|
||
'399bp': seq399
|
||
})
|
||
|
||
return pd.DataFrame(results)
|
||
|
||
except Exception as e:
|
||
raise Exception(f"区域预测过程出错: {str(e)}")
|
||
|
||
def _extract_center_sequence(self, sequence, target_length=33):
|
||
"""从序列中心位置提取指定长度的子序列"""
|
||
# 确保序列为字符串
|
||
sequence = str(sequence).upper()
|
||
|
||
# 如果序列长度小于目标长度，返回原序列
|
||
if len(sequence) <= target_length:
|
||
return sequence
|
||
|
||
# 计算中心位置
|
||
center = len(sequence) // 2
|
||
half_target = target_length // 2
|
||
|
||
# 提取中心序列
|
||
start = center - half_target
|
||
end = start + target_length
|
||
|
||
# 边界检查
|
||
if start < 0:
|
||
start = 0
|
||
end = target_length
|
||
elif end > len(sequence):
|
||
end = len(sequence)
|
||
start = end - target_length
|
||
|
||
return sequence[start:end] |