{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":87793,"databundleVersionId":12024591,"sourceType":"competition"}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Stanford RNA 3D Folding Competition\n# Revised Notebook for RNA 3D Structure Prediction (Improved V3)\n\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nimport os\nfrom scipy.spatial.distance import cdist\nimport warnings\nimport random\nimport torch.nn.functional as F\nwarnings.filterwarnings('ignore')\n\nprint(\"Starting Stanford RNA 3D Folding notebook...\")\n\n# 1. Data Loading and Exploration\nprint(\"Loading datasets...\")\n# 경로를 환경에 맞게 수정해주세요.\n# 예: '/kaggle/input/stanford-rna-3d-folding/' 또는 './stanford-rna-3d-folding/'\n# 이 코드는 Kaggle 환경을 가정합니다. 로컬에서 실행 시 경로를 확인하세요.\ntry:\n    base_path = '/kaggle/input/stanford-rna-3d-folding/'\n    if not os.path.exists(base_path):\n        print(f\"Warning: Kaggle input path '{base_path}' not found. Trying current directory.\")\n        # 로컬 테스트를 위해 현재 디렉토리를 사용하거나, 사용자 정의 경로를 설정할 수 있습니다.\n        # 예를 들어, 데이터셋이 'data' 폴더에 있다면 base_path = './data/' 로 설정\n        base_path = './' # 현재 디렉토리 또는 사용자가 지정한 경로로 수정\n        # 로컬 테스트 시에는 파일들이 실제로 해당 경로에 있는지 확인해야 합니다.\n        # train_sequences.csv, train_labels.csv 등이 base_path 아래에 있어야 합니다.\n        if not os.path.exists(os.path.join(base_path, 'train_sequences.csv')):\n            print(\"Error: train_sequences.csv not found in the specified base_path. Please check your data paths.\")\n            exit()\n\n\n    train_sequences = pd.read_csv(os.path.join(base_path, 'train_sequences.csv'))\n    train_labels = pd.read_csv(os.path.join(base_path, 'train_labels.csv'))\n    validation_sequences = pd.read_csv(os.path.join(base_path, 'validation_sequences.csv'))\n    validation_labels = pd.read_csv(os.path.join(base_path, 'validation_labels.csv'))\n    test_sequences = pd.read_csv(os.path.join(base_path, 'test_sequences.csv'))\n    sample_submission = pd.read_csv(os.path.join(base_path, 'sample_submission.csv'))\nexcept FileNotFoundError as e:\n    print(f\"Error loading data files: {e}\")\n    print(\"Please ensure the data files are in the correct path.\")\n    print(\"If running locally, you might need to download the dataset and adjust the paths.\")\n    exit()\n\n\n# Fill missing coordinate values to avoid NaNs.\ntrain_labels.fillna(0, inplace=True)\nvalidation_labels.fillna(0, inplace=True)\n\nprint(\"\\nBasic dataset information:\")\nprint(f\"Training sequences: {train_sequences.shape}\")\nprint(f\"Training labels: {train_labels.shape}\")\nprint(f\"Validation sequences: {validation_sequences.shape}\")\nprint(f\"Validation labels: {validation_labels.shape}\")\nprint(f\"Test sequences: {test_sequences.shape}\")\nprint(f\"Sample submission: {sample_submission.shape}\")\n\n# 2. Data Analysis (optional visualizations)\nprint(\"\\nAnalyzing RNA sequence lengths...\")\ntrain_sequences['length'] = train_sequences['sequence'].str.len()\nplt.figure(figsize=(12, 6))\nsns.histplot(train_sequences['length'], bins=50)\nplt.title('Distribution of RNA Sequence Lengths')\nplt.xlabel('Sequence Length')\nplt.ylabel('Count')\n# 저장 경로를 워킹 디렉토리로 변경 (Kaggle에서는 /kaggle/working/)\noutput_viz_path = 'sequence_length_distribution.png'\nif 'KAGGLE_WORKING_DIR' in os.environ:\n    output_viz_path = os.path.join(os.environ['KAGGLE_WORKING_DIR'], output_viz_path)\nplt.savefig(output_viz_path)\nplt.close()\nprint(f\"Sequence length distribution plot saved to {output_viz_path}\")\n\n# 3. Data Preprocessing\ndef preprocess_sequence_data(sequences_df, labels_df=None, is_train=True):\n    nucleotide_map = {'A': 0, 'C': 1, 'G': 2, 'U': 3, 'T': 3} # T를 U로 매핑\n    processed_data = []\n\n    for idx, row in sequences_df.iterrows():\n        seq_id = row['target_id']\n        sequence = row['sequence']\n        numerical_seq = [nucleotide_map.get(nuc, 4) for nuc in sequence] # 4 for unknown\n\n        structures = None\n        if labels_df is not None: # is_train 조건 제거, labels_df 유무로 판단\n            # 레이블 파일의 ID는 'target_id_residue_number' 형식이므로, target_id로 시작하는지 확인\n            sequence_labels = labels_df[labels_df['ID'].str.startswith(seq_id + '_')].copy() # SettingWithCopyWarning 방지\n\n            if not sequence_labels.empty:\n                # resname, resid 컬럼 제외하고 x_i, y_i, z_i 컬럼만 선택\n                coord_cols = [col for col in sequence_labels.columns if col.startswith(('x_', 'y_', 'z_'))]\n                num_structures = len(coord_cols) // 3 # 각 구조당 x,y,z 3개 좌표\n\n                structures = []\n                for i in range(1, num_structures + 1):\n                    coords_list = []\n                    # 레이블 파일은 residue 순서대로 정렬되어 있다고 가정\n                    # 'resid' 컬럼이 있다면 이를 기준으로 정렬하는 것이 더 안전\n                    if 'resid' in sequence_labels.columns:\n                        sequence_labels.sort_values('resid', inplace=True)\n\n                    current_structure_coords = sequence_labels[[f'x_{i}', f'y_{i}', f'z_{i}']].values\n                    \n                    # 시퀀스 길이와 좌표 개수가 일치하는지 확인\n                    if len(current_structure_coords) != len(sequence):\n                        # print(f\"Warning: Mismatch in length for {seq_id}, structure {i}. Seq_len: {len(sequence)}, Coords_len: {len(current_structure_coords)}. Skipping this structure.\")\n                        # 패딩 또는 다른 처리 방법 고려 가능\n                        # 간단히 길이만큼만 사용하거나, 0으로 채우거나, 또는 이 구조를 건너뛸 수 있음\n                        # 여기서는 길이만큼만 사용 (만약 좌표가 더 짧다면) 또는 시퀀스 길이로 자름 (좌표가 더 길다면)\n                        min_len = min(len(current_structure_coords), len(sequence))\n                        current_structure_coords = current_structure_coords[:min_len]\n                        if len(current_structure_coords) < len(sequence): # 좌표가 시퀀스보다 짧으면 패딩\n                             padding = np.zeros((len(sequence) - len(current_structure_coords), 3))\n                             current_structure_coords = np.vstack((current_structure_coords, padding))\n\n                    coords = np.array(current_structure_coords)\n\n                    # Normalize coordinates per sequence (center and scale)\n                    mean = np.mean(coords, axis=0)\n                    std = np.std(coords, axis=0) + 1e-8 # 분모 0 방지\n                    coords_norm = (coords - mean) / std\n                    structures.append(coords_norm)\n            # else:\n                # print(f\"No labels found for {seq_id}\")\n        # else:\n            # print(f\"No labels_df provided for {seq_id} (e.g., for test data)\")\n\n        processed_data.append({\n            'id': seq_id,\n            'sequence': numerical_seq,\n            'structures': structures # structures가 None일 수 있음 (테스트 데이터 또는 레이블 없는 학습 데이터)\n        })\n    return processed_data\n\n\nprint(\"Preprocessing training data...\")\ntrain_data = preprocess_sequence_data(train_sequences, train_labels)\nprint(\"Preprocessing validation data...\")\n# validation_labels가 train_labels와 형식이 다를 수 있음에 유의\nvalidation_data = preprocess_sequence_data(validation_sequences, validation_labels)\nprint(\"Preprocessing test data...\")\ntest_data = preprocess_sequence_data(test_sequences, labels_df=None) # 테스트 데이터에는 레이블 없음\n\n# 4. Feature Engineering\ndef extract_sequence_features(sequence_numerical): # 입력은 수치화된 시퀀스\n    one_hot = np.zeros((len(sequence_numerical), 5)) # 0,1,2,3 (ACGU) + 4 (Unknown)\n    for i, nucleotide_code in enumerate(sequence_numerical):\n        one_hot[i, nucleotide_code] = 1\n\n    gc_content_feature = []\n    window_size = 5\n    for i in range(len(sequence_numerical)):\n        start = max(0, i - window_size // 2)\n        end = min(len(sequence_numerical), i + window_size // 2 + 1)\n        window = sequence_numerical[start:end]\n        # G is 2, C is 1\n        gc_count = sum(1 for n_code in window if n_code == 1 or n_code == 2)\n        gc_content_feature.append(gc_count / len(window) if len(window) > 0 else 0)\n\n    positions = np.array([[i / len(sequence_numerical)] for i in range(len(sequence_numerical))])\n    features = np.hstack((one_hot, positions, np.array(gc_content_feature).reshape(-1, 1)))\n    return features\n\nprint(\"Extracting sequence features...\")\nfor i, data_item in enumerate(train_data):\n    train_data[i]['features'] = extract_sequence_features(data_item['sequence'])\nfor i, data_item in enumerate(validation_data):\n    validation_data[i]['features'] = extract_sequence_features(data_item['sequence'])\nfor i, data_item in enumerate(test_data):\n    test_data[i]['features'] = extract_sequence_features(data_item['sequence'])\n\n# 5. RNA Secondary Structure Prediction (simple rule-based)\ndef predict_rna_secondary_structure(sequence_numerical): # 입력은 수치화된 시퀀스\n    nucleotide_map_inv = {0: 'A', 1: 'C', 2: 'G', 3: 'U', 4: 'X'} # X for unknown\n    seq_chars = [nucleotide_map_inv.get(n_code, 'X') for n_code in sequence_numerical]\n    structure = ['.' for _ in range(len(seq_chars))]\n    complementary = {'A': 'U', 'U': 'A', 'G': 'C', 'C': 'G', 'X': None}\n\n    # 더 나은 로직: 스택 기반 또는 동적 프로그래밍의 단순화된 버전\n    # 여기서는 최소 루프 크기 (예: 3) 고려\n    min_loop_length = 3\n    for i in range(len(seq_chars)):\n        if structure[i] != '.':\n            continue\n        for j in range(len(seq_chars) - 1, i + min_loop_length, -1): # j > i + min_loop_length\n            if structure[j] != '.':\n                continue\n            # 슈도넛 방지를 위해 간단히 가장 바깥쪽 쌍부터 찾음\n            # (더 정교한 알고리즘은 슈도넛을 허용하거나 다른 방식으로 처리)\n            can_pair = True\n            # for k in range(i + 1, j): # 현재 쌍 내부에 다른 쌍이 있는지 (간단한 슈도넛 체크)\n            #     if structure[k] != '.': # 이미 페어링된 뉴클레오티드가 있다면\n            #         # 이것만으로는 슈도넛을 완벽히 막을 수 없음. 더 정교한 로직 필요.\n            #         # 여기서는 슈도넛을 고려하지 않는 매우 단순한 페어링\n            #         pass\n\n            if complementary.get(seq_chars[i]) == seq_chars[j] and can_pair:\n                structure[i] = '('\n                structure[j] = ')'\n                # 쌍을 찾으면 내부 루프 중단 (가장 멀리 있는 쌍 우선)\n                break\n    return ''.join(structure)\n\ndef enhance_features_with_ss(data_list): # 변수명 변경\n    for i, item in enumerate(data_list):\n        seq_num = item['sequence']\n        ss = predict_rna_secondary_structure(seq_num)\n        ss_features = np.zeros((len(ss), 3)) # '.', '(', ')'\n        for j, char_ss in enumerate(ss): # 변수명 변경\n            if char_ss == '.':\n                ss_features[j, 0] = 1\n            elif char_ss == '(':\n                ss_features[j, 1] = 1\n            elif char_ss == ')':\n                ss_features[j, 2] = 1\n        data_list[i]['features'] = np.hstack((item['features'], ss_features))\n    return data_list\n\nprint(\"Enhancing features with secondary structure information...\")\ntrain_data = enhance_features_with_ss(train_data)\nvalidation_data = enhance_features_with_ss(validation_data)\ntest_data = enhance_features_with_ss(test_data)\n\n\n# 6. PyTorch Dataset and DataLoader\nclass RNADataset(Dataset):\n    def __init__(self, data, augment=False, num_structures_to_use=1):\n        self.data = [item for item in data if item['structures'] is not None and len(item['structures']) > 0]\n        if not self.data and data: # 원본 데이터는 있었는데 필터링 후 비었다면 경고\n            print(\"Warning: No valid structures found in the provided data for RNADataset. DataLoader might be empty.\")\n        self.augment = augment\n        self.num_structures_to_use = num_structures_to_use # 사용할 구조의 수 (주로 첫 번째)\n\n    def __len__(self):\n        return len(self.data)\n\n    def __getitem__(self, idx):\n        item = self.data[idx]\n        features = item['features'].copy() # 원본 수정을 피하기 위해 복사\n\n        # Data augmentation\n        if self.augment and random.random() < 0.5: # 증강 확률\n            # Mutation\n            if random.random() < 0.2: # 돌연변이 확률\n                if len(features) > 0:\n                    mut_idx = random.randint(0, len(features)-1)\n                    # one-hot 부분 (앞 5개 컬럼)만 변경\n                    new_nucleotide_one_hot = np.eye(5)[random.choice([0,1,2,3])] # ACGU 중 하나로 변경\n                    features[mut_idx, :5] = new_nucleotide_one_hot\n            # Reversing a segment\n            if len(features) > 20 and random.random() < 0.2: # 세그먼트 반전 확률\n                start = random.randint(0, len(features) - 10)\n                segment_len = random.randint(5, 10)\n                end = min(start + segment_len, len(features))\n                features[start:end] = features[start:end][::-1]\n\n        features_tensor = torch.tensor(features, dtype=torch.float32)\n\n        # 타겟 구조 선택 (여러 구조가 있을 경우 첫 번째 또는 랜덤하게 선택 가능)\n        # preprocess_sequence_data 에서 여러 구조를 로드했다면, 그 중 하나를 선택\n        # 여기서는 첫 번째 구조를 사용한다고 가정 (item['structures'][0])\n        if item['structures'] and len(item['structures']) > 0:\n            # 사용할 구조의 인덱스 (예: 항상 첫 번째, 또는 랜덤)\n            # 여기서는 항상 첫 번째 구조를 사용\n            struct_idx_to_use = 0\n            target_coords = item['structures'][struct_idx_to_use]\n            target_tensor = torch.tensor(target_coords, dtype=torch.float32)\n        else:\n            # 테스트 데이터셋의 경우 타겟이 없을 수 있음\n            # 또는 학습 데이터셋인데 특정 아이템에 구조가 없는 경우 (이런 경우는 필터링 되어야 함)\n            target_tensor = torch.empty(0,3, dtype=torch.float32) # 빈 텐서 또는 None\n\n        return {\n            'features': features_tensor,\n            'target': target_tensor,\n            'length': features_tensor.shape[0],\n            'id': item['id']\n        }\n\nclass RNADatasetTest(Dataset): # 테스트 데이터셋을 위한 별도 클래스\n    def __init__(self, data):\n        self.data = data # 테스트 데이터에는 'structures'가 없을 수 있음\n\n    def __len__(self):\n        return len(self.data)\n\n    def __getitem__(self, idx):\n        item = self.data[idx]\n        features = item['features']\n        features_tensor = torch.tensor(features, dtype=torch.float32)\n\n        return {\n            'features': features_tensor,\n            'target': None, # 테스트 시에는 타겟이 없음\n            'length': features_tensor.shape[0],\n            'id': item['id']\n        }\n\n\ndef collate_fn(batch):\n    # 길이 기준으로 내림차순 정렬 (pack_padded_sequence에 필요)\n    # RNADataset에서 이미 length를 int로 반환하므로 .item() 불필요\n    sorted_batch = sorted(batch, key=lambda x: x['length'], reverse=True)\n\n    features_list = [x['features'] for x in sorted_batch]\n    lengths_list = [x['length'] for x in sorted_batch] # 이미 int list\n    ids_list = [x['id'] for x in sorted_batch]\n\n    # features 패딩\n    # features_list[0].shape[1] 은 특성 차원\n    padded_features = nn.utils.rnn.pad_sequence(features_list, batch_first=True, padding_value=0.0)\n\n    targets_exist = all(x['target'] is not None and x['target'].nelement() > 0 for x in sorted_batch)\n\n    if targets_exist:\n        targets_list = [x['target'] for x in sorted_batch]\n        padded_targets = nn.utils.rnn.pad_sequence(targets_list, batch_first=True, padding_value=0.0) # 0으로 패딩\n    else:\n        padded_targets = None # 타겟이 없는 경우 (예: 테스트 데이터)\n\n    return {\n        'features': padded_features,\n        'targets': padded_targets,\n        'lengths': torch.tensor(lengths_list, dtype=torch.long), # LSTM에 사용될 길이 (정수형 텐서)\n        'ids': ids_list\n    }\n\n# 데이터셋 인스턴스 생성 전에 train_data와 validation_data에 유효한 구조가 있는지 확인\n# preprocess_sequence_data에서 structures가 None이거나 비어있을 수 있음.\n# RNADataset은 유효한 structures가 있는 아이템만 사용하도록 수정됨.\ntrain_dataset = RNADataset(train_data, augment=True)\nvalidation_dataset = RNADataset(validation_data) # 검증 시에는 증강 안 함\ntest_dataset = RNADatasetTest(test_data) # 테스트용 데이터셋\n\nprint(f\"Number of training samples: {len(train_dataset)}\")\nprint(f\"Number of validation samples: {len(validation_dataset)}\")\nprint(f\"Number of test samples: {len(test_dataset)}\")\n\n# DataLoader 인스턴스 생성\n# num_workers는 환경에 따라 조절. Kaggle에서는 보통 2 또는 4.\n# persistent_workers는 DataLoader가 워커 프로세스를 유지하도록 하여 에폭 간 오버헤드 줄임\nnum_w = 2 if 'KAGGLE_KERNEL_RUN_TYPE' in os.environ else 0 # 로컬에서는 0이 더 안정적일 수 있음\n\ntrain_loader = DataLoader(\n    train_dataset,\n    batch_size=8, # 배치 크기 조절 가능\n    shuffle=True,\n    collate_fn=collate_fn,\n    pin_memory=True,\n    num_workers=num_w,\n    persistent_workers=True if num_w > 0 else False\n)\n\nvalidation_loader = DataLoader(\n    validation_dataset,\n    batch_size=8, # 검증 시에도 동일한 배치 크기 사용 가능\n    shuffle=False, # 검증 시에는 셔플 불필요\n    collate_fn=collate_fn,\n    pin_memory=True,\n    num_workers=num_w,\n    persistent_workers=True if num_w > 0 else False\n)\n\ntest_loader = DataLoader(\n    test_dataset,\n    batch_size=4, # 테스트 시에는 예측 안정성을 위해 더 작은 배치 사용 가능\n    shuffle=False,\n    collate_fn=collate_fn,\n    pin_memory=True,\n    num_workers=num_w,\n    persistent_workers=True if num_w > 0 else False\n)\n\n\n# 7. Model Architecture\nclass RNAFoldingModel(nn.Module):\n    def __init__(self, input_dim, hidden_dim=256, num_lstm_layers=3, num_transformer_layers=3, nhead_transformer=8): # 파라미터 조정\n        super().__init__()\n        self.hidden_dim = hidden_dim\n        self.num_lstm_layers = num_lstm_layers\n\n        # BiLSTM Encoder\n        self.lstm = nn.LSTM(\n            input_dim, hidden_dim,\n            num_layers=num_lstm_layers,\n            bidirectional=True,\n            batch_first=True,\n            dropout=0.2 if num_lstm_layers > 1 else 0 # 마지막 레이어 제외한 드롭아웃\n        )\n\n        # Transformer Encoder\n        transformer_encoder_layer = nn.TransformerEncoderLayer(\n            d_model=2 * hidden_dim, # BiLSTM 출력 차원\n            nhead=nhead_transformer,\n            dim_feedforward=hidden_dim * 4, # 일반적인 설정 (d_model * 4)\n            dropout=0.2,\n            activation='gelu', # GELU 사용\n            batch_first=True # batch_first=True로 설정\n        )\n        self.transformer_encoder = nn.TransformerEncoder(\n            transformer_encoder_layer,\n            num_layers=num_transformer_layers\n        )\n\n        # Prediction Head (좌표 3개 예측)\n        # Transformer 출력 후 바로 Linear 레이어 사용\n        self.fc_out = nn.Linear(2 * hidden_dim, 3)\n\n\n    def forward(self, x, lengths):\n        # x: (batch, seq_len, features)\n        # lengths: (batch) - 각 시퀀스의 실제 길이\n\n        # BiLSTM 처리\n        # lengths를 CPU로 옮기고 리스트로 변환해야 할 수 있음\n        # collate_fn에서 이미 lengths를 LongTensor로 반환함\n        # pack_padded_sequence는 lengths가 CPU에 있는 것을 기대할 수 있음\n        packed_input = nn.utils.rnn.pack_padded_sequence(\n            x, lengths.cpu(), batch_first=True, enforce_sorted=False # collate_fn에서 정렬하므로 enforce_sorted=True 가능\n        )\n        packed_output, (h_n, c_n) = self.lstm(packed_input)\n        lstm_out, _ = nn.utils.rnn.pad_packed_sequence(packed_output, batch_first=True)\n        # lstm_out: (batch, seq_len, 2 * hidden_dim)\n\n        # Transformer Encoder\n        # Transformer는 패딩 마스크가 필요할 수 있음\n        # src_key_padding_mask: (batch, seq_len)\n        # True인 위치는 무시됨\n        max_len = lstm_out.size(1)\n        src_key_padding_mask = torch.arange(max_len, device=x.device)[None, :] >= lengths[:, None]\n\n        transformer_out = self.transformer_encoder(lstm_out, src_key_padding_mask=src_key_padding_mask)\n        # transformer_out: (batch, seq_len, 2 * hidden_dim)\n\n        # 최종 좌표 예측\n        predictions = self.fc_out(transformer_out)\n        # predictions: (batch, seq_len, 3)\n\n        return predictions\n\n# 8. Loss Function and Evaluation Metric\nclass GeometricLoss(nn.Module):\n    def __init__(self, coord_weight=0.6, dist_weight=0.4, bond_len_weight=0.0): # 가중치 조절 가능\n        super().__init__()\n        self.coord_weight = coord_weight\n        self.dist_weight = dist_weight\n        self.bond_len_weight = bond_len_weight # 새로운 항: 결합 길이 손실\n\n        self.coord_loss_fn = nn.SmoothL1Loss(reduction='none') # 각 요소별 손실 계산\n        self.dist_loss_fn = nn.MSELoss(reduction='none')\n        self.bond_len_loss_fn = nn.MSELoss(reduction='none') # 또는 L1Loss\n\n        self.target_bond_length = 3.8 # 대략적인 C1'-C1' 결합 길이 (Angstrom) - 조정 필요\n\n    def forward(self, pred_coords, target_coords, lengths):\n        batch_size = pred_coords.size(0)\n        total_loss = 0\n        num_valid_samples = 0\n\n        for i in range(batch_size):\n            l = lengths[i].item()\n            if l < 2: continue # 유효한 길이를 가진 샘플만 처리\n\n            pred_i = pred_coords[i, :l]    # (len, 3)\n            target_i = target_coords[i, :l]  # (len, 3)\n\n            # 1. Coordinate Loss (SmoothL1)\n            loss_coord = self.coord_loss_fn(pred_i, target_i).mean() # 평균내서 스칼라로\n\n            # 2. Distance Matrix Loss (MSE)\n            pred_dist_sq = torch.cdist(pred_i.unsqueeze(0), pred_i.unsqueeze(0), p=2).squeeze(0) # (len, len)\n            target_dist_sq = torch.cdist(target_i.unsqueeze(0), target_i.unsqueeze(0), p=2).squeeze(0)\n            loss_dist = self.dist_loss_fn(pred_dist_sq, target_dist_sq).mean() # 평균내서 스칼라로\n\n            # 3. (Optional) Bond Length Regularization Loss\n            loss_bond = 0\n            if self.bond_len_weight > 0 and l > 1:\n                # 예측된 구조에서 인접한 C1' 원자 간의 거리 계산\n                pred_bond_lengths = torch.norm(pred_i[1:] - pred_i[:-1], p=2, dim=1)\n                target_bond_lengths_ideal = torch.full_like(pred_bond_lengths, self.target_bond_length)\n                loss_bond = self.bond_len_loss_fn(pred_bond_lengths, target_bond_lengths_ideal).mean()\n\n            sample_loss = (self.coord_weight * loss_coord +\n                           self.dist_weight * loss_dist +\n                           self.bond_len_weight * loss_bond)\n            total_loss += sample_loss\n            num_valid_samples +=1\n\n        return total_loss / num_valid_samples if num_valid_samples > 0 else torch.tensor(0.0, device=pred_coords.device)\n\n\ndef kabsch_align(P, Q):\n    \"\"\"\n    Aligns two sets of points P and Q using the Kabsch algorithm.\n    P, Q: Nx3 numpy arrays\n    Returns R (rotation matrix), t (translation vector) such that P_aligned = P @ R + t is closest to Q.\n    And returns P_aligned.\n    \"\"\"\n    P = P.copy()\n    Q = Q.copy()\n\n    # Centroid\n    centroid_P = np.mean(P, axis=0)\n    centroid_Q = np.mean(Q, axis=0)\n\n    P -= centroid_P\n    Q -= centroid_Q\n\n    # Covariance matrix\n    H = P.T @ Q\n    U, S, Vt = np.linalg.svd(H)\n    R = Vt.T @ U.T\n\n    # Reflection check\n    if np.linalg.det(R) < 0:\n        Vt[2, :] *= -1\n        R = Vt.T @ U.T\n\n    t = centroid_Q - centroid_P @ R\n    P_aligned = P @ R + centroid_Q # P를 Q에 맞춤. P@R은 Q의 중심에 대해 회전된 P.\n                                   # 여기에 Q의 중심을 더해줌.\n                                   # 또는 (P @ R + centroid_P @ R + t - centroid_P @ R)\n                                   # = (P - centroid_P) @ R + centroid_Q\n    return R, t, P_aligned\n\n\ndef calculate_tm_score_usalign_like(pred_coords_np, true_coords_np):\n    \"\"\"\n    Calculates TM-score after Kabsch alignment.\n    This is a simplified version and might not perfectly match US-align.\n    pred_coords_np, true_coords_np: (L, 3) numpy arrays\n    \"\"\"\n    L_ref = true_coords_np.shape[0]\n    L_pred = pred_coords_np.shape[0]\n\n    if L_ref == 0 or L_pred == 0:\n        return 0.0\n\n    # US-align은 L_pred와 L_ref가 달라도 최적의 부분정렬을 찾지만,\n    # 여기서는 간단히 길이가 같은 부분만 고려하거나, 짧은 쪽에 맞춰 자름.\n    # 실제 대회에서는 US-align이 이를 처리. 여기서는 길이가 같다고 가정하거나, 맞춰줌.\n    common_len = min(L_ref, L_pred)\n    if common_len < 2: # TM-score 계산에 의미 없는 길이\n        return 0.0\n\n    pred_coords_aligned_subset = pred_coords_np[:common_len]\n    true_coords_subset = true_coords_np[:common_len]\n\n    # Kabsch alignment\n    # P를 Q에 맞춤. pred_coords_aligned_subset을 true_coords_subset에 맞춤.\n    _, _, pred_aligned = kabsch_align(pred_coords_aligned_subset, true_coords_subset)\n\n    # d0 calculation based on L_ref (original length of reference)\n    if L_ref >= 30:\n        d0 = 1.24 * (L_ref - 15)**(1/3) - 1.8\n    elif L_ref >= 24: d0 = 0.7\n    elif L_ref >= 20: d0 = 0.6\n    elif L_ref >= 16: d0 = 0.5\n    elif L_ref >= 12: d0 = 0.4\n    else: d0 = 0.3 # For L_ref < 12\n    d0 = max(d0, 0.5) # As per AlphaFold\n\n    # Calculate TM-score sum over the aligned common length\n    sum_tm = 0\n    for i in range(common_len):\n        di_sq = np.sum((pred_aligned[i] - true_coords_subset[i])**2) # squared distance\n        sum_tm += 1 / (1 + (di_sq / (d0**2)))\n\n    # Normalize by L_ref (as per standard TM-score definition)\n    tm_score = sum_tm / L_ref\n    return np.clip(tm_score, 0.0, 1.0)\n\n\n# 9. Training Loop\ndef train_model(model, train_loader, val_loader, epochs=50, lr=1e-4, device='cpu', patience=10):\n    model = model.to(device)\n    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-5) # AdamW 추천\n    criterion = GeometricLoss(coord_weight=0.5, dist_weight=0.5, bond_len_weight=0.1).to(device) # 손실 함수도 디바이스로\n\n    # Learning rate scheduler (옵션)\n    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=patience//2, verbose=True)\n\n    best_val_tm_score = 0.0\n    epochs_no_improve = 0\n    history = {'train_loss': [], 'val_loss': [], 'tm_score': []}\n\n    # 모델 저장 경로\n    model_save_path = 'best_model.pth'\n    if 'KAGGLE_WORKING_DIR' in os.environ:\n        model_save_path = os.path.join(os.environ['KAGGLE_WORKING_DIR'], model_save_path)\n\n\n    for epoch in range(epochs):\n        # --- Training Phase ---\n        model.train()\n        total_train_loss = 0\n        num_train_batches = 0\n        for batch_idx, batch in enumerate(train_loader):\n            features = batch['features'].to(device)\n            targets = batch['targets'].to(device) # collate_fn에서 None일 수 있음 (RNADataset에서 필터링)\n            lengths = batch['lengths'].to(device) # pack_padded_sequence 위해 CPU로 옮길 필요 없음 (함수 내에서 처리)\n            # ids = batch['ids'] # 학습 시에는 보통 사용 안 함\n\n            if targets is None: # RNADataset에서 걸러지지만, 안전장치\n                continue\n\n            optimizer.zero_grad()\n            outputs = model(features, lengths) # lengths는 LongTensor 여야 함\n\n            loss = criterion(outputs, targets, lengths)\n            loss.backward()\n            nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # 그래디언트 클리핑\n            optimizer.step()\n\n            total_train_loss += loss.item()\n            num_train_batches += 1\n\n        avg_train_loss = total_train_loss / num_train_batches if num_train_batches > 0 else 0\n        history['train_loss'].append(avg_train_loss)\n\n        # --- Validation Phase ---\n        model.eval()\n        total_val_loss = 0\n        num_val_batches = 0\n        all_tm_scores = []\n        with torch.no_grad():\n            for batch in val_loader:\n                features = batch['features'].to(device)\n                targets = batch['targets'].to(device)\n                lengths = batch['lengths'] # .to(device) 불필요, 어차피 CPU에서 사용하거나 model forward에서 처리\n                # ids = batch['ids']\n\n                if targets is None:\n                    continue\n\n                outputs = model(features, lengths.to(device)) # 모델 forward에는 device에 있는 lengths 필요\n                loss = criterion(outputs, targets, lengths) # criterion 내에서는 lengths.item() 사용\n                total_val_loss += loss.item()\n                num_val_batches += 1\n\n                # TM-Score 계산\n                outputs_np = outputs.cpu().numpy()\n                targets_np = targets.cpu().numpy()\n                current_lengths_list = lengths.cpu().tolist()\n\n                for i in range(outputs_np.shape[0]):\n                    l = current_lengths_list[i]\n                    if l < 5: continue # 매우 짧은 시퀀스 TM-score 계산에서 제외\n\n                    pred_coords = outputs_np[i, :l, :]\n                    true_coords = targets_np[i, :l, :]\n\n                    # 정규화된 좌표로 TM-score 계산 시 의미가 다를 수 있음.\n                    # 이상적으로는 원래 스케일로 복원 후 계산하거나,\n                    # US-align과 같은 도구를 사용해야 함.\n                    # 여기서는 정규화된 상태로 계산 (상대적 비교용)\n                    # tm = calculate_tm_score(pred_coords, true_coords) # 이전 버전\n                    tm = calculate_tm_score_usalign_like(pred_coords, true_coords) # Kabsch 정렬 포함\n                    all_tm_scores.append(tm)\n\n        avg_val_loss = total_val_loss / num_val_batches if num_val_batches > 0 else 0\n        avg_tm_score = np.mean(all_tm_scores) if all_tm_scores else 0.0\n        history['val_loss'].append(avg_val_loss)\n        history['tm_score'].append(avg_tm_score)\n\n        print(f\"Epoch {epoch+1}/{epochs} => \"\n              f\"Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}, Val TM-Score: {avg_tm_score:.4f}\")\n\n        scheduler.step(avg_tm_score) # TM-score 기준으로 LR 스케줄링\n\n        if avg_tm_score > best_val_tm_score:\n            best_val_tm_score = avg_tm_score\n            torch.save(model.state_dict(), model_save_path)\n            print(f\"   Best model saved with TM-Score: {best_val_tm_score:.4f} at epoch {epoch+1}\")\n            epochs_no_improve = 0\n        else:\n            epochs_no_improve += 1\n            if epochs_no_improve >= patience:\n                print(f\"   Early stopping triggered after {patience} epochs without improvement.\")\n                break\n    \n    print(f\"Training finished. Best validation TM-Score: {best_val_tm_score:.4f}\")\n    # 가장 좋은 모델 로드 (만약 저장되었다면)\n    if os.path.exists(model_save_path):\n        model.load_state_dict(torch.load(model_save_path))\n        print(f\"Loaded best model from {model_save_path}\")\n    return model, history\n\n\n# 10. Model Inference and Multiple Structure Generation\ndef generate_diverse_structures(model, features_tensor, seq_length_int, num_structures=5, noise_scale=0.02, device='cpu'):\n    model.eval() # 평가 모드\n    structures_list = []\n    with torch.no_grad():\n        # features_tensor는 (seq_len, feature_dim) 형태의 단일 시퀀스 특성\n        # seq_length_int는 해당 시퀀스의 실제 길이 (정수)\n\n        for i in range(num_structures):\n            current_features = features_tensor.clone() # 원본 변경 방지\n            if i > 0: # 첫 번째 예측은 노이즈 없이\n                noise = torch.randn_like(current_features) * noise_scale\n                current_features += noise\n\n            # 모델 입력 형태로 변경: (batch_size=1, seq_len, feature_dim)\n            current_features_batch = current_features.unsqueeze(0).to(device)\n            # 길이도 텐서로: (batch_size=1)\n            lengths_batch = torch.tensor([seq_length_int], dtype=torch.long).to(device)\n\n            output_coords = model(current_features_batch, lengths_batch)\n            # output_coords: (1, seq_len, 3)\n\n            # 실제 길이만큼만 잘라내고 CPU로 이동, numpy로 변환\n            coords_np = output_coords[0, :seq_length_int, :].cpu().numpy()\n            structures_list.append(coords_np)\n    return structures_list\n\n\ndef generate_predictions_for_submission(model, dataloader, device, num_predictions_per_seq=5):\n    model.to(device)\n    model.eval()\n    all_predictions_dict = {} # key: seq_id, value: list of 5 predicted coord arrays\n\n    with torch.no_grad():\n        for batch_idx, batch in enumerate(dataloader): # test_loader 사용\n            features_batch = batch['features'].to(device) # (batch, max_len, feat_dim)\n            lengths_list = batch['lengths'].cpu().tolist() # 각 시퀀스의 실제 길이 리스트\n            ids_list = batch['ids']\n\n            for i in range(features_batch.size(0)): # 배치 내 각 시퀀스에 대해\n                seq_id = ids_list[i]\n                seq_len_int = lengths_list[i]\n                # 해당 시퀀스의 특성만 추출 (패딩 제외)\n                # features_batch[i]는 (max_len, feat_dim)\n                # 실제 길이만큼만 사용: features_batch[i, :seq_len_int, :]\n                single_seq_features = features_batch[i, :seq_len_int, :].clone() # (seq_len, feat_dim)\n\n                # 이 단일 시퀀스에 대해 다양한 구조 생성\n                predicted_structures_for_seq = generate_diverse_structures(\n                    model,\n                    single_seq_features, # (seq_len, feat_dim) 텐서 전달\n                    seq_len_int,         # 실제 길이 (int) 전달\n                    num_structures=num_predictions_per_seq,\n                    noise_scale=0.02, # 노이즈 스케일 조절 가능\n                    device=device\n                )\n                all_predictions_dict[seq_id] = predicted_structures_for_seq\n            if (batch_idx + 1) % 10 == 0:\n                 print(f\"  Processed batch {batch_idx+1}/{len(dataloader)} for predictions.\")\n\n    return all_predictions_dict\n\n\n# 11. Submission File Generation\ndef create_submission_file(predictions_dict, test_sequences_df, output_filename='submission.csv'):\n    submission_rows_list = []\n\n    # test_sequences_df를 순회하며 각 ID에 대한 예측 찾기\n    for _, row in test_sequences_df.iterrows():\n        seq_id = row['target_id']\n        sequence_str = row['sequence'] # 문자열 시퀀스\n        seq_len = len(sequence_str)\n\n        if seq_id in predictions_dict:\n            pred_structures_for_id = predictions_dict[seq_id] # list of 5 (L,3) numpy arrays\n            num_pred_structures = len(pred_structures_for_id)\n\n            for residue_idx in range(seq_len): # 시퀀스의 각 잔기에 대해\n                submission_row = {\n                    'ID': f\"{seq_id}_{residue_idx + 1}\", # ID_resid 형식\n                    'resname': sequence_str[residue_idx],\n                    'resid': residue_idx + 1\n                }\n\n                for pred_idx in range(5): # 5개 예측 구조에 대해\n                    if pred_idx < num_pred_structures:\n                        # 해당 예측 구조에서 현재 잔기의 좌표\n                        # pred_structures_for_id[pred_idx]는 (L,3) 배열\n                        # residue_idx가 이 배열의 길이를 넘지 않는지 확인\n                        if residue_idx < pred_structures_for_id[pred_idx].shape[0]:\n                            coords = pred_structures_for_id[pred_idx][residue_idx]\n                            submission_row[f'x_{pred_idx+1}'] = coords[0]\n                            submission_row[f'y_{pred_idx+1}'] = coords[1]\n                            submission_row[f'z_{pred_idx+1}'] = coords[2]\n                        else: # 예측된 구조가 실제 시퀀스보다 짧은 경우 (이론상 발생 안해야 함)\n                            # 이전 구조의 값으로 채우거나 0으로 채움\n                            if pred_idx > 0:\n                                submission_row[f'x_{pred_idx+1}'] = submission_row[f'x_{pred_idx}']\n                                submission_row[f'y_{pred_idx+1}'] = submission_row[f'y_{pred_idx}']\n                                submission_row[f'z_{pred_idx+1}'] = submission_row[f'z_{pred_idx}']\n                            else: # 첫번째 예측인데도 짧으면 0\n                                submission_row[f'x_{pred_idx+1}'] = 0.0\n                                submission_row[f'y_{pred_idx+1}'] = 0.0\n                                submission_row[f'z_{pred_idx+1}'] = 0.0\n                    else: # 예측된 구조가 5개 미만인 경우, 마지막 유효한 예측으로 채움\n                        # (generate_diverse_structures가 항상 5개를 반환하도록 설계되어 이 경우는 적음)\n                        # 가장 마지막으로 성공적으로 가져온 좌표 사용\n                        last_valid_pred_idx = num_pred_structures # 1-based index for column name\n                        submission_row[f'x_{pred_idx+1}'] = submission_row[f'x_{last_valid_pred_idx}']\n                        submission_row[f'y_{pred_idx+1}'] = submission_row[f'y_{last_valid_pred_idx}']\n                        submission_row[f'z_{pred_idx+1}'] = submission_row[f'z_{last_valid_pred_idx}']\n                submission_rows_list.append(submission_row)\n        else:\n            print(f\"Warning: No predictions found for seq_id {seq_id}. Filling with zeros or placeholders.\")\n            # 예측이 없는 경우 (예: 오류 발생) - 샘플 제출 파일 형식에 맞게 placeholder 채우기\n            for residue_idx in range(seq_len):\n                submission_row = {\n                    'ID': f\"{seq_id}_{residue_idx + 1}\",\n                    'resname': sequence_str[residue_idx],\n                    'resid': residue_idx + 1\n                }\n                for pred_idx in range(5):\n                    submission_row[f'x_{pred_idx+1}'] = 0.0\n                    submission_row[f'y_{pred_idx+1}'] = 0.0\n                    submission_row[f'z_{pred_idx+1}'] = 0.0\n                submission_rows_list.append(submission_row)\n\n\n    submission_df = pd.DataFrame(submission_rows_list)\n\n    # 파일 저장 경로 (Kaggle 환경 고려)\n    if 'KAGGLE_WORKING_DIR' in os.environ:\n        output_filepath = os.path.join(os.environ['KAGGLE_WORKING_DIR'], output_filename)\n    else:\n        output_filepath = output_filename\n\n    try:\n        submission_df.to_csv(output_filepath, index=False)\n        print(f\"Submission file successfully saved to: {os.path.abspath(output_filepath)}\")\n    except Exception as e:\n        print(f\"Error saving submission file: {e}\")\n        return None\n    return submission_df\n\n\n# 12. Visualization Functions\ndef visualize_3d_structure(coords_np, title=\"RNA_3D_Structure\"):\n    # matplotlib.pyplot을 함수 내에서 import하여 Kaggle 환경에서 GUI 백엔드 문제 방지\n    import matplotlib\n    matplotlib.use('Agg') # 비 GUI 백엔드 사용\n    import matplotlib.pyplot as plt\n    # 3D 플롯을 위해 Axes3D 임포트 필요\n    from mpl_toolkits.mplot3d import Axes3D\n\n\n    fig = plt.figure(figsize=(8, 6))\n    ax = fig.add_subplot(111, projection='3d')\n\n    # C1' 원자 위치 산점도\n    ax.scatter(coords_np[:, 0], coords_np[:, 1], coords_np[:, 2], c='blue', marker='o', s=50, label=\"C1' atoms\", depthshade=True)\n\n    # 원자 간 연결선 (backbone)\n    for i in range(len(coords_np) - 1):\n        ax.plot([coords_np[i, 0], coords_np[i+1, 0]],\n                [coords_np[i, 1], coords_np[i+1, 1]],\n                [coords_np[i, 2], coords_np[i+1, 2]], 'gray', lw=1.5)\n\n    ax.set_title(title, fontsize=15)\n    ax.set_xlabel('X (Å)', fontsize=12)\n    ax.set_ylabel('Y (Å)', fontsize=12)\n    ax.set_zlabel('Z (Å)', fontsize=12)\n    ax.legend()\n    ax.grid(True)\n\n    # 파일명에 사용할 수 없는 문자 제거\n    safe_title = \"\".join(c if c.isalnum() else \"_\" for c in title)\n    viz_filename = f\"{safe_title}.png\"\n    if 'KAGGLE_WORKING_DIR' in os.environ:\n        viz_filepath = os.path.join(os.environ['KAGGLE_WORKING_DIR'], viz_filename)\n    else:\n        viz_filepath = viz_filename\n\n    plt.savefig(viz_filepath)\n    plt.close(fig) # 메모리 해제\n    print(f\"3D structure visualization saved to {viz_filepath}\")\n\n\n# 13. Main Execution\ndef main():\n    print(\"\\n--- Main execution ---\")\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    print(f\"Using device: {device}\")\n\n    # 입력 차원 결정 (첫 번째 학습 데이터의 특성 개수)\n    if not train_dataset:\n        print(\"Error: Training dataset is empty. Cannot determine input_dim or train the model.\")\n        return\n    \n    # RNADataset.__getitem__은 딕셔너리를 반환하므로, features를 직접 접근\n    # 또는 train_data에서 직접 가져올 수 있음\n    # 안전하게 첫 번째 유효한 train_data 아이템에서 특성 차원 가져오기\n    first_valid_train_item = next((item for item in train_data if 'features' in item and item['features'] is not None), None)\n    if first_valid_train_item is None or first_valid_train_item['features'].shape[1] == 0:\n        print(\"Error: Could not determine input_dim from training data. Features might be missing or empty.\")\n        # 예시로 고정된 값을 사용하거나, 오류 처리\n        # 여기서는 임의의 값으로 설정하고 경고 (실제로는 오류로 중단해야 함)\n        input_dim = 10 # 임시 값, 실제로는 오류 발생해야 함\n        print(f\"Warning: Setting input_dim to a default value of {input_dim}. This is likely an error.\")\n    else:\n        input_dim = first_valid_train_item['features'].shape[1]\n    \n    print(f\"Input feature dimension: {input_dim}\")\n\n\n    # 모델 인스턴스 생성\n    model = RNAFoldingModel(\n        input_dim=input_dim,\n        hidden_dim=384,       # LSTM hidden dim, Transformer d_model의 절반\n        num_lstm_layers=3,\n        num_transformer_layers=4,\n        nhead_transformer=8\n    ).to(device)\n\n    print(f\"\\nModel architecture: {model}\")\n    num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n    print(f\"Number of trainable parameters: {num_params:,}\")\n\n\n    print(\"\\nStarting model training...\")\n    # 에폭 수, 학습률 등은 실험을 통해 조절\n    # 실제 대회에서는 더 많은 에폭과 세심한 하이퍼파라미터 튜닝 필요\n    trained_model, history = train_model(\n        model=model,\n        train_loader=train_loader,\n        val_loader=validation_loader,\n        epochs=30, # 예시 에폭 (실제로는 더 많이)\n        lr=5e-5,   # 학습률 (조정 필요)\n        device=device,\n        patience=7 # Early stopping을 위한 patience\n    )\n\n    # 학습 과정 시각화 (옵션)\n    plt.figure(figsize=(12, 4))\n    plt.subplot(1, 2, 1)\n    plt.plot(history['train_loss'], label='Train Loss')\n    plt.plot(history['val_loss'], label='Val Loss')\n    plt.title('Loss History')\n    plt.xlabel('Epoch')\n    plt.ylabel('Loss')\n    plt.legend()\n\n    plt.subplot(1, 2, 2)\n    plt.plot(history['tm_score'], label='Validation TM-Score')\n    plt.title('Validation TM-Score History')\n    plt.xlabel('Epoch')\n    plt.ylabel('TM-Score')\n    plt.legend()\n    plt.tight_layout()\n    history_plot_path = 'training_history.png'\n    if 'KAGGLE_WORKING_DIR' in os.environ:\n        history_plot_path = os.path.join(os.environ['KAGGLE_WORKING_DIR'], history_plot_path)\n    plt.savefig(history_plot_path)\n    plt.close()\n    print(f\"Training history plot saved to {history_plot_path}\")\n\n\n    print(\"\\nGenerating predictions on test data...\")\n    # test_loader를 사용하여 예측 생성\n    test_predictions_dict = generate_predictions_for_submission(\n        trained_model, # 학습된 모델 사용\n        test_loader,\n        device,\n        num_predictions_per_seq=5\n    )\n    print(f\"\\nPredictions generated for {len(test_predictions_dict)} test sequences.\")\n\n    print(\"\\nCreating submission file...\")\n    submission_df = create_submission_file(\n        test_predictions_dict,\n        test_sequences, # test_sequences.csv에서 읽은 DataFrame\n        output_filename='submission.csv'\n    )\n\n    if submission_df is not None:\n        print(f\"\\nSubmission file preview (first 5 rows):\")\n        print(submission_df.head())\n    else:\n        print(\"Submission file generation failed.\")\n\n\n    print(\"\\nVisualizing a sample prediction (first test sequence if available)...\")\n    if not test_sequences.empty:\n        first_test_seq_id = test_sequences['target_id'].iloc[0]\n        if first_test_seq_id in test_predictions_dict and test_predictions_dict[first_test_seq_id]:\n            # 첫 번째 예측 구조 시각화\n            sample_predicted_coords = test_predictions_dict[first_test_seq_id][0]\n            visualize_3d_structure(sample_predicted_coords, title=f\"Predicted_3D_Structure_{first_test_seq_id}\")\n        else:\n            print(f\"No prediction found or prediction is empty for the first test sequence ({first_test_seq_id}) for visualization.\")\n    else:\n        print(\"Test sequences data is empty, skipping visualization.\")\n\n    print(\"\\n--- Main execution completed ---\")\n\nif __name__ == '__main__':\n    # 실행 시간 측정 (옵션)\n    import time\n    start_time = time.time()\n    main()\n    end_time = time.time()\n    print(f\"\\nTotal execution time: {end_time - start_time:.2f} seconds.\")\n\nprint(\"\\nNotebook execution finished.\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}