structure_parser.py

# 二级结构解析工具

import numpy as np
import torch


def parse_dot_bracket(structure: str) -> dict:
    n = len(structure)
    paired = [False] * n
    pairing = [-1] * n
    stacking = [0] * n
    
    # 配对括号映射（支持多种括号类型）
    bracket_pairs = {
        '(': ')',
        '[': ']',
        '{': '}',
        '<': '>',
    }
    open_brackets = list(bracket_pairs.keys())
    close_brackets = list(bracket_pairs.values())
    
    # 使用栈来匹配括号
    stack = []
    depth = 0
    
    for i, char in enumerate(structure):
        if char in open_brackets:
            stack.append((i, depth))
            depth += 1
            stacking[i] = depth
        elif char in close_brackets:
            if stack:
                open_idx, open_depth = stack.pop()
                paired[open_idx] = True
                paired[i] = True
                pairing[open_idx] = i
                pairing[i] = open_idx
                stacking[open_idx] = open_depth + 1
                stacking[i] = open_depth + 1
                depth = open_depth
        # '.' 或其他字符保持默认值（未配对）
    
    return {
        'paired': paired,
        'pairing': pairing,
        'stacking': stacking,
    }


def dotbracket_to_pairmap(structure: str, seq_len: int = None) -> np.ndarray:
    if seq_len is None:
        seq_len = len(structure)
    
    # 使用详细的解析函数
    parsed = parse_dot_bracket(structure[:seq_len])
    
    pair_map = np.zeros((seq_len, seq_len), dtype=np.float32)
    
    # 根据配对关系填充矩阵
    for i in range(min(seq_len, len(parsed['pairing']))):
        pair_idx = parsed['pairing'][i]
        if pair_idx >= 0 and pair_idx < seq_len:
            pair_map[i, pair_idx] = 1.0
            pair_map[pair_idx, i] = 1.0
    
    return pair_map


def struct_to_feature_vector(
    structure: str,
    seq_len: int,
    max_depth: int = 10,
) -> torch.Tensor:
    # 检查是否是数值字符串（如 icshape）
    if ',' in structure:
        try:
            values = [float(x.strip()) for x in structure.split(',')]
            if len(values) >= seq_len:
                values = values[:seq_len]
            else:
                values.extend([-1.0] * (seq_len - len(values)))
            # 使用数值作为特征，扩展到10维
            features = []
            for v in values:
                feature_vec = [v / 2.0 + 0.5] + [0.0] * 9  # 归一化到0-1，填充
                features.append(feature_vec)
            return torch.tensor(features, dtype=torch.float32)
        except ValueError:
            pass  # 回退到dot-bracket解析
    
    # 解析结构
    parsed = parse_dot_bracket(structure)
    
    # 长度匹配
    if len(structure) != seq_len:
        if len(structure) > seq_len:
            structure = structure[:seq_len]
            parsed['paired'] = parsed['paired'][:seq_len]
            parsed['pairing'] = parsed['pairing'][:seq_len]
            parsed['stacking'] = parsed['stacking'][:seq_len]
        else:
            # 填充
            padding = seq_len - len(structure)
            structure = structure + '.' * padding
            parsed['paired'].extend([False] * padding)
            parsed['pairing'].extend([-1] * padding)
            parsed['stacking'].extend([0] * padding)
    
    features = []
    for i in range(seq_len):
        char = structure[i] if i < len(structure) else '.'
        paired = parsed['paired'][i]
        pairing_idx = parsed['pairing'][i]
        depth = min(parsed['stacking'][i], max_depth)
        
        # 计算配对距离（如果配对）
        if pairing_idx >= 0:
            pairing_dist = abs(i - pairing_idx) / seq_len if seq_len > 0 else 0.0  # 归一化距离
        else:
            pairing_dist = 0.0
        
        # 字符类型 one-hot
        char_features = [
            1.0 if char == '.' else 0.0,           # dot
            1.0 if char == '(' else 0.0,           # open_paren
            1.0 if char == ')' else 0.0,           # close_paren
            1.0 if char in '[]' else 0.0,          # bracket
            1.0 if char in '{}' else 0.0,          # brace
            1.0 if char in '<>' else 0.0,          # angle
            1.0 if char not in '.()[]{}<>' else 0.0, # other
        ]
        
        # 组合特征
        feature_vec = [
            float(paired),           # 是否配对 (1维)
            pairing_dist,            # 配对距离（归一化）(1维)
            float(depth) / max_depth, # 嵌套深度（归一化）(1维)
        ] + char_features            # 字符类型 (7维)
        
        features.append(feature_vec)
    
    # 转换为 Tensor
    features_tensor = torch.tensor(features, dtype=torch.float32)  # [seq_len, 10]
    
    return features_tensor


def create_pair_representation(pair_map: np.ndarray, add_distance: bool = True) -> np.ndarray:
    seq_len = pair_map.shape[0]
    
    if add_distance:
        # 创建距离矩阵
        indices = np.arange(seq_len)
        i_matrix, j_matrix = np.meshgrid(indices, indices, indexing='ij')
        distance_matrix = np.abs(i_matrix - j_matrix).astype(np.float32)
        
        # 归一化距离矩阵
        if seq_len > 0:
            distance_matrix = distance_matrix / (seq_len - 1) if seq_len > 1 else distance_matrix
        
        # 组合配对信息和距离信息
        representation = np.stack([pair_map, distance_matrix], axis=-1)  # [seq_len, seq_len, 2]
    else:
        representation = np.expand_dims(pair_map, axis=-1)  # [seq_len, seq_len, 1]
    
    return representation


def create_pair_tokens_from_sequence(sequence: str, base_range: int = 1, lamda: float = 0.8) -> np.ndarray:
    # RNA 碱基编码映射
    rna_map = {'A': 5, 'a': 5, 'U': 6, 'u': 6, 'T': 6, 't': 6, 
               'C': 7, 'c': 7, 'G': 4, 'g': 4, 'N': 3, 'n': 3}
    
    # 将序列转换为数字编码
    seq_encoded = [rna_map.get(base, 3) for base in sequence]
    seq_len = len(seq_encoded)
    
    # 配对规则
    def paired(x, y, lamda_val=0.8):
        if x == 5 and y == 6:  # A-U
            return 2
        elif x == 4 and y == 7:  # G-C
            return 3
        elif x == 4 and y == 6:  # G-U
            return lamda_val
        elif x == 6 and y == 5:  # U-A
            return 2
        elif x == 7 and y == 4:  # C-G
            return 3
        elif x == 6 and y == 4:  # U-G
            return lamda_val
        else:
            return 0
    
    # 创建配对映射矩阵
    paird_map = np.array([[paired(i, j, lamda) for i in range(30)] for j in range(30)])
    
    # 生成 2D tokens
    tokens = np.zeros((seq_len, seq_len), dtype=np.float32)
    data_index = np.arange(seq_len)
    
    # 前向配对
    for add in range(base_range):
        data_index_x = data_index - add
        data_index_y = data_index + add
        mask = ((data_index_x >= 0)[:, None] & (data_index_y < seq_len)[None, :])
        
        if not mask.any():
            break
            
        data_index_x_clipped = np.clip(data_index_x, 0, seq_len - 1)
        data_index_y_clipped = np.clip(data_index_y, 0, seq_len - 1)
        
        scores = paird_map[
            np.array(seq_encoded)[data_index_x_clipped][:, None],
            np.array(seq_encoded)[data_index_y_clipped][None, :]
        ]
        
        mask = mask & (scores != 0)
        tokens += scores * mask * np.exp(-0.5 * (add ** 2))
    
    # 后向配对
    for add in range(1, base_range):
        data_index_x = data_index + add
        data_index_y = data_index - add
        mask = ((data_index_x < seq_len)[:, None] & (data_index_y >= 0)[None, :])
        
        if not mask.any():
            break
            
        data_index_x_clipped = np.clip(data_index_x, 0, seq_len - 1)
        data_index_y_clipped = np.clip(data_index_y, 0, seq_len - 1)
        
        scores = paird_map[
            np.array(seq_encoded)[data_index_x_clipped][:, None],
            np.array(seq_encoded)[data_index_y_clipped][None, :]
        ]
        
        mask = mask & (scores != 0)
        tokens += scores * mask * np.exp(-0.5 * (add ** 2))
    
    # 添加最后一个维度
    tokens = np.expand_dims(tokens, axis=-1)  # [seq_len, seq_len, 1]
    
    return tokens