{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":87793,"databundleVersionId":12276181,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":11056335,"sourceType":"datasetVersion","datasetId":6888333},{"sourceId":11056379,"sourceType":"datasetVersion","datasetId":6888367},{"sourceId":11990166,"sourceType":"datasetVersion","datasetId":7541592}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport time\nimport torch\nimport random\nimport shutil\nimport numpy as np\nimport pandas as pd\n\nfrom Bio.Seq import Seq\nfrom Bio import pairwise2\n\nfrom tqdm import tqdm\nfrom scipy.spatial import distance_matrix\nfrom scipy.spatial.transform import Rotation as R\n\nimport warnings\nwarnings.filterwarnings('ignore')\n\n\ntest_sequences = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/test_sequences.csv\")\n\nis_submission_mode = len(test_sequences) != 12","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:17:34.060013Z","iopub.execute_input":"2025-05-29T13:17:34.060361Z","iopub.status.idle":"2025-05-29T13:17:39.093698Z","shell.execute_reply.started":"2025-05-29T13:17:34.060333Z","shell.execute_reply":"2025-05-29T13:17:39.091772Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_seqs = pd.read_csv('/kaggle/input/rna-all-data/merged_sequences_final.csv')\ntrain_labels = pd.read_csv('/kaggle/input/rna-all-data/merged_labels_final.csv')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:17:47.002295Z","iopub.execute_input":"2025-05-29T13:17:47.002597Z","iopub.status.idle":"2025-05-29T13:18:00.599173Z","shell.execute_reply.started":"2025-05-29T13:17:47.002572Z","shell.execute_reply":"2025-05-29T13:18:00.597686Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Check for new CIF files that need processing\nimport os\nimport pandas as pd\nfrom pathlib import Path\n\n# Get existing target_ids (without chain suffix)\nexisting_pdb_ids = set()\nfor target_id in train_seqs['target_id']:\n pdb_id = target_id.rsplit('_', 1)[0] # Remove chain suffix\n existing_pdb_ids.add(pdb_id.lower())\n\nprint(f\"Existing PDB IDs in train_seqs: {len(existing_pdb_ids)}\")\n\n# Get all CIF files in directory\ncif_dir = '/kaggle/input/stanford-rna-3d-folding/PDB_RNA'\nall_cif_files = [f for f in os.listdir(cif_dir) if f.endswith('.cif')]\nall_pdb_ids = set(Path(f).stem.lower() for f in all_cif_files)\n\nprint(f\"Total CIF files found: {len(all_cif_files)}\")\n\n# Find new files to process\nnew_pdb_ids = all_pdb_ids - existing_pdb_ids\nnew_cif_files = [f\"{pdb_id}.cif\" for pdb_id in new_pdb_ids]\n\nprint(f\"New files to process: {len(new_cif_files)}\")\nprint(f\"New PDB IDs: {sorted(list(new_pdb_ids))}\")\n\nif new_cif_files:\n print(f\"\\nFirst 10 new CIF files: {new_cif_files[:10]}\")\nelse:\n print(\"No new files to process!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:18:00.600785Z","iopub.execute_input":"2025-05-29T13:18:00.601136Z","iopub.status.idle":"2025-05-29T13:18:00.713040Z","shell.execute_reply.started":"2025-05-29T13:18:00.601109Z","shell.execute_reply":"2025-05-29T13:18:00.711914Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Comprehensive modified nucleotide mapping\nnucleotide_mapping = {\n # Standard nucleotides\n 'A': 'A', 'U': 'U', 'G': 'G', 'C': 'C',\n \n # === ADENOSINE MODIFICATIONS ===\n 'I': 'A', # Inosine (hypoxanthine)\n '1MA': 'A', # 1-methyladenosine\n '2MA': 'A', # 2-methyladenosine\n '6MA': 'A', # N6-methyladenosine (m6A)\n 'M2A': 'A', # N2-methyladenosine\n 'MS2': 'A', # 2-methylthio-N6-isopentenyladenosine\n 'AET': 'A', # 2-aminoethylthio-adenosine\n 'A2L': 'A', # 2'-O-methyladenosine\n 'A44': 'A', # Modified adenosine\n '6OP': 'A', # Modified adenosine\n '8XA': 'A', # Modified adenosine\n 'ZAD': 'A', # Modified adenosine\n \n # === URIDINE MODIFICATIONS ===\n 'PSU': 'U', # Pseudouridine (most common)\n 'H2U': 'U', # Dihydrouridine\n '5MU': 'U', # 5-methyluridine (ribothymidine)\n '4SU': 'U', # 4-thiouridine\n '2MU': 'U', # 2'-O-methyluridine\n 'OMU': 'U', # O-methyluridine\n 'T': 'U', # Thymine (in RNA)\n 'RT': 'U', # Ribothymidine\n 'DHU': 'U', # Dihydrouridine\n 'UMS': 'U', # 5-methoxycarbonylmethyluridine\n 'U2L': 'U', # Modified uridine\n 'U36': 'U', # Modified uridine\n 'Y5P': 'U', # Modified uridine\n 'P5P': 'U', # Modified uridine\n 'UFT': 'U', # Modified uridine\n 'F2T': 'U', # Modified uridine\n '0U': 'U', # Modified uridine\n '8XU': 'U', # Modified uridine\n 'ZBU': 'U', # Modified uridine\n 'ZTH': 'U', # Modified uridine\n 'ZHP': 'U', # Modified uridine\n 'SSU': 'U', # Modified uridine\n \n # === GUANOSINE MODIFICATIONS ===\n 'M2G': 'G', # N2-methylguanosine\n 'M7G': 'G', # 7-methylguanosine (cap structure)\n 'OMG': 'G', # O-methylguanosine\n '1MG': 'G', # 1-methylguanosine\n '2MG': 'G', # 2'-O-methylguanosine\n 'YYG': 'G', # Modified guanosine\n 'QUO': 'G', # Queuosine\n 'G7M': 'G', # 7-methylguanosine\n 'GTP': 'G', # Guanosine triphosphate\n 'GDP': 'G', # Guanosine diphosphate\n 'GMP': 'G', # Guanosine monophosphate\n 'G2L': 'G', # Modified guanosine\n 'G48': 'G', # Modified guanosine\n '6OO': 'G', # Modified guanosine\n '0G': 'G', # Modified guanosine\n '8XG': 'G', # Modified guanosine\n 'ZGU': 'G', # Modified guanosine\n 'LCG': 'G', # Modified guanosine\n \n # === CYTIDINE MODIFICATIONS ===\n '5MC': 'C', # 5-methylcytidine\n 'OMC': 'C', # O-methylcytidine\n '2MC': 'C', # 2'-O-methylcytidine\n 'M5C': 'C', # 5-methylcytidine\n 'CBV': 'C', # Carbovir cytidine\n 'C2L': 'C', # Modified cytidine\n 'C43': 'C', # Modified cytidine\n '6NW': 'C', # Modified cytidine\n '0C': 'C', # Modified cytidine\n '8XC': 'C', # Modified cytidine\n 'ZCY': 'C', # Modified cytidine\n 'ZBC': 'C', # Modified cytidine\n \n # === RARE/SYNTHETIC MODIFICATIONS ===\n 'ADP': 'A', # Adenosine diphosphate\n 'ATP': 'A', # Adenosine triphosphate\n 'AMP': 'A', # Adenosine monophosphate\n 'UDP': 'U', # Uridine diphosphate\n 'UTP': 'U', # Uridine triphosphate\n 'UMP': 'U', # Uridine monophosphate\n 'CDP': 'C', # Cytidine diphosphate\n 'CTP': 'C', # Cytidine triphosphate\n 'CMP': 'C', # Cytidine monophosphate\n \n # === WYOSINE DERIVATIVES ===\n 'YW1': 'G', # Wybutosine\n 'YW2': 'G', # Wybutosine derivative\n 'YW3': 'G', # Wybutosine derivative\n \n # === HYPERMODIFIED BASES ===\n 'Q': 'G', # Queuosine\n 'X': 'G', # Xanthosine\n 'D': 'U', # Dihydrouridine\n 'P': 'U', # Pseudouridine\n \n # === METHYLATION VARIANTS ===\n 'M1A': 'A', # 1-methyladenosine\n 'M1G': 'G', # 1-methylguanosine\n 'M3C': 'C', # 3-methylcytidine\n 'M5U': 'U', # 5-methyluridine\n 'M6A': 'A', # N6-methyladenosine\n \n # === THIO MODIFICATIONS ===\n 'S2C': 'C', # 2-thiocytidine\n 'S2U': 'U', # 2-thiouridine\n 'S4U': 'U', # 4-thiouridine\n \n # === CAP STRUCTURES ===\n '7MG': 'G', # 7-methylguanosine (5' cap)\n 'M7G': 'G', # 7-methylguanosine\n 'G7M': 'G', # 7-methylguanosine\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:18:21.951471Z","iopub.execute_input":"2025-05-29T13:18:21.951806Z","iopub.status.idle":"2025-05-29T13:18:21.961115Z","shell.execute_reply.started":"2025-05-29T13:18:21.951779Z","shell.execute_reply":"2025-05-29T13:18:21.959871Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Complete final code for processing new CIF files with comprehensive support\nfrom Bio.PDB import MMCIFParser\nimport pandas as pd\nfrom pathlib import Path\nimport os\nfrom tqdm import tqdm\n\n# Original function (for standard nucleotides)\ndef extract_rna_data_from_cif(cif_file_path):\n \"\"\"Extract unique RNA sequences and C1' coordinates from a CIF file\"\"\"\n parser = MMCIFParser(QUIET=True)\n \n try:\n structure = parser.get_structure('structure', cif_file_path)\n pdb_id = Path(cif_file_path).stem.upper()\n \n sequences_data = []\n coordinates_data = []\n seen_sequences = set() # Track unique sequences\n \n for model in structure:\n for chain in model:\n chain_id = chain.id\n target_id = f\"{pdb_id}_{chain_id}\"\n \n # Check if chain contains RNA residues\n rna_residues = []\n for residue in chain:\n if residue.get_resname() in ['A', 'U', 'G', 'C']: # RNA nucleotides\n rna_residues.append(residue)\n \n if rna_residues: # Only process if RNA residues found\n # Build sequence\n sequence = ''.join([res.get_resname() for res in rna_residues])\n \n # Only add if sequence is unique\n if sequence not in seen_sequences:\n seen_sequences.add(sequence)\n sequences_data.append({\n 'target_id': target_id,\n 'sequence': sequence\n })\n \n # Extract C1' coordinates for this unique sequence\n for i, residue in enumerate(rna_residues, 1):\n if \"C1'\" in residue:\n atom = residue[\"C1'\"]\n coordinates_data.append({\n 'ID': f\"{target_id}_{i}\",\n 'resname': residue.get_resname(),\n 'resid': i,\n 'x_1': atom.coord[0],\n 'y_1': atom.coord[1], \n 'z_1': atom.coord[2]\n })\n \n return sequences_data, coordinates_data\n \n except Exception as e:\n print(f\"Error processing {cif_file_path}: {e}\")\n return [], []\n\n# Disorder-aware glycosidic carbon detection\ndef get_glycosidic_carbon_disorder_aware(residue):\n \"\"\"\n Find C1' or C1{suffix} atoms, handling DisorderedAtom objects\n \"\"\"\n # Get all available atom names\n available_atoms = [atom.get_name() for atom in residue]\n \n # Look for C1' first (most common)\n if \"C1'\" in available_atoms:\n atom = residue[\"C1'\"]\n # Handle DisorderedAtom by getting the first conformation\n if hasattr(atom, 'selected_child'):\n return atom.selected_child\n return atom\n \n # Look for any C1{suffix} pattern\n c1_variants = [atom_name for atom_name in available_atoms if atom_name.startswith('C1') and len(atom_name) > 2]\n \n if c1_variants:\n # If multiple C1 variants, prefer the shortest one\n best_variant = min(c1_variants, key=len)\n atom = residue[best_variant]\n # Handle DisorderedAtom by getting the first conformation\n if hasattr(atom, 'selected_child'):\n return atom.selected_child\n return atom\n \n return None\n\n# Comprehensive extraction function with disorder handling\ndef extract_rna_data_from_cif_comprehensive_final(cif_file_path):\n \"\"\"Extract RNA with comprehensive modified nucleotide recognition and disorder handling\"\"\"\n parser = MMCIFParser(QUIET=True)\n \n try:\n structure = parser.get_structure('structure', cif_file_path)\n pdb_id = Path(cif_file_path).stem.upper()\n \n sequences_data = []\n coordinates_data = []\n seen_sequences = set()\n \n for model in structure:\n for chain in model:\n chain_id = chain.id\n target_id = f\"{pdb_id}_{chain_id}\"\n \n # Check if chain contains RNA residues (including all modified ones)\n rna_residues = []\n for residue in chain:\n res_name = residue.get_resname()\n if res_name in nucleotide_mapping:\n rna_residues.append(residue)\n \n if rna_residues: # Only process if RNA residues found\n # Build sequence using standard nucleotides\n sequence = ''.join([nucleotide_mapping[res.get_resname()] for res in rna_residues])\n \n # Only add if sequence is unique\n if sequence not in seen_sequences:\n seen_sequences.add(sequence)\n sequences_data.append({\n 'target_id': target_id,\n 'sequence': sequence\n })\n \n # Extract coordinates using disorder-aware detection\n for i, residue in enumerate(rna_residues, 1):\n carbon_atom = get_glycosidic_carbon_disorder_aware(residue)\n \n if carbon_atom is not None: # Use 'is not None' to avoid DisorderedAtom issues\n coordinates_data.append({\n 'ID': f\"{target_id}_{i}\",\n 'resname': nucleotide_mapping[residue.get_resname()], # Use standard name\n 'resid': i,\n 'x_1': carbon_atom.coord[0],\n 'y_1': carbon_atom.coord[1], \n 'z_1': carbon_atom.coord[2]\n })\n \n return sequences_data, coordinates_data\n \n except Exception as e:\n print(f\"Error processing {cif_file_path}: {e}\")\n return [], []\n\n# Smart extraction function\ndef extract_rna_data_smart_final(cif_file_path):\n \"\"\"\n Final smart extraction: standard nucleotides first, then comprehensive with disorder handling\n \"\"\"\n # First try the original function (standard nucleotides only)\n sequences_std, coordinates_std = extract_rna_data_from_cif(cif_file_path)\n \n # If we found RNA data with standard function, use it\n if sequences_std:\n return sequences_std, coordinates_std, \"standard\"\n \n # If no standard RNA found, try the comprehensive function with disorder handling\n sequences_mod, coordinates_mod = extract_rna_data_from_cif_comprehensive_final(cif_file_path)\n \n if sequences_mod:\n return sequences_mod, coordinates_mod, \"modified\"\n else:\n return [], [], \"none\"\n\n# Main processing function\ndef process_new_cif_files_final(train_seqs, train_labels, cif_dir, new_cif_files):\n \"\"\"Final processing function with comprehensive nucleotide support and disorder handling\"\"\"\n \n if not new_cif_files:\n print(\"No new files to process - all CIF files have already been processed!\")\n return train_seqs, train_labels\n \n print(f\"Processing {len(new_cif_files)} new CIF files with final comprehensive extraction...\")\n print(f\"Nucleotide mapping includes {len(nucleotide_mapping)} variants\")\n print(f\"Includes disorder handling for DisorderedAtom objects\")\n \n new_sequences = []\n new_coordinates = []\n processing_stats = {\"standard\": 0, \"modified\": 0, \"none\": 0}\n \n for cif_file in tqdm(new_cif_files):\n cif_path = os.path.join(cif_dir, cif_file)\n sequences, coordinates, extraction_type = extract_rna_data_smart_final(cif_path)\n \n processing_stats[extraction_type] += 1\n \n if sequences: # Only add if we found RNA data\n new_sequences.extend(sequences)\n new_coordinates.extend(coordinates)\n print(f\"✅ {cif_file} ({extraction_type}): {len(sequences)} sequences, {len(coordinates)} coordinates\")\n else:\n print(f\"❌ {cif_file}: No RNA data found\")\n \n print(f\"\\nFINAL PROCESSING SUMMARY:\")\n print(f\"Files with standard nucleotides: {processing_stats['standard']}\")\n print(f\"Files with modified nucleotides: {processing_stats['modified']}\")\n print(f\"Files with no RNA data: {processing_stats['none']}\")\n print(f\"Success rate: {processing_stats['standard'] + processing_stats['modified']} / {len(new_cif_files)} = {((processing_stats['standard'] + processing_stats['modified']) / len(new_cif_files) * 100):.1f}%\")\n \n if new_sequences:\n # Create DataFrames for new data\n new_sequences_df = pd.DataFrame(new_sequences)\n new_coordinates_df = pd.DataFrame(new_coordinates)\n \n print(f\"\\nFINAL NEW DATA SUMMARY:\")\n print(f\"Total new sequences: {len(new_sequences)}\")\n print(f\"Total new coordinates: {len(new_coordinates)}\")\n \n # Add to existing dataframes\n train_seqs_updated = pd.concat([train_seqs, new_sequences_df], ignore_index=True)\n train_labels_updated = pd.concat([train_labels, new_coordinates_df], ignore_index=True)\n \n print(f\"\\nFINAL UPDATED DATAFRAMES:\")\n print(f\"train_seqs: {train_seqs.shape} -> {train_seqs_updated.shape}\")\n print(f\"train_labels: {train_labels.shape} -> {train_labels_updated.shape}\")\n \n \n print(f\"\\nFinal improvement summary:\")\n print(f\"Sequences added: {len(train_seqs_updated) - len(train_seqs)}\")\n print(f\"Coordinates added: {len(train_labels_updated) - len(train_labels)}\")\n \n return train_seqs_updated, train_labels_updated\n \n else:\n print(\"No new RNA sequences found in any of the new files\")\n return train_seqs, train_labels\n\n# EXECUTE THE FINAL COMPREHENSIVE PROCESSING\nprint(\"=\"*90)\nprint(\"FINAL COMPREHENSIVE PROCESSING WITH DISORDER HANDLING AND 93 NUCLEOTIDE VARIANTS\")\nprint(\"=\"*90)\n\n# Process the new files\ntrain_seqs_final, train_labels_final = process_new_cif_files_final(\n train_seqs=train_seqs, \n train_labels=train_labels, \n cif_dir=cif_dir, \n new_cif_files=new_cif_files\n)\n\nprint(\"\\n\" + \"=\"*90)\nprint(\"FINAL COMPREHENSIVE PROCESSING COMPLETE\")\nprint(\"=\"*90)\nprint(f\"Final train_seqs shape: {train_seqs_final.shape}\")\nprint(f\"Final train_labels shape: {train_labels_final.shape}\")\n\n# Show which files were successfully processed\nsuccessful_files = []\nfailed_files = []\n\nfor cif_file in new_cif_files:\n cif_path = os.path.join(cif_dir, cif_file)\n sequences, coordinates, extraction_type = extract_rna_data_smart_final(cif_path)\n if sequences:\n successful_files.append(cif_file)\n else:\n failed_files.append(cif_file)\n\nprint(f\"\\nSuccessfully processed files ({len(successful_files)}):\")\nfor f in successful_files:\n print(f\" ✅ {f}\")\n\nif failed_files:\n print(f\"\\nFiles with no RNA data ({len(failed_files)}):\")\n for f in failed_files:\n print(f\" ❌ {f}\")\n\nprint(f\"\\nFinal statistics:\")\nprint(f\"Original dataset: {len(train_seqs)} sequences, {len(train_labels)} coordinates\")\nprint(f\"Final dataset: {len(train_seqs_final)} sequences, {len(train_labels_final)} coordinates\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:18:23.369577Z","iopub.execute_input":"2025-05-29T13:18:23.370047Z","iopub.status.idle":"2025-05-29T13:19:06.455165Z","shell.execute_reply.started":"2025-05-29T13:18:23.370017Z","shell.execute_reply":"2025-05-29T13:19:06.454135Z"},"scrolled":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Set up directories\npredictions_dir = \"/kaggle/working/predictions\"\nos.makedirs(predictions_dir, exist_ok=True)\nfasta_dir = \"/kaggle/working/fasta_files\"\nos.makedirs(fasta_dir, exist_ok=True)\n\n# Set time limit for DRfold2 (in seconds)\nDRFOLD_TIME_LIMIT = 7 * 60 * 60 # 7 hours\nstart_time_global = time.time()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:19:06.456044Z","iopub.execute_input":"2025-05-29T13:19:06.456271Z","iopub.status.idle":"2025-05-29T13:19:06.461204Z","shell.execute_reply.started":"2025-05-29T13:19:06.456250Z","shell.execute_reply":"2025-05-29T13:19:06.460011Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!cp -r /kaggle/input/drfold2-repo/DRfold2 /kaggle/working/\n%cd DRfold2\n%cd Arena\n!make Arena\n%cd ..","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:19:06.462613Z","iopub.execute_input":"2025-05-29T13:19:06.462902Z","iopub.status.idle":"2025-05-29T13:19:19.583287Z","shell.execute_reply.started":"2025-05-29T13:19:06.462876Z","shell.execute_reply":"2025-05-29T13:19:19.582232Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!cp -r /kaggle/input/drfold2/model_hub /kaggle/working/DRfold2/","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:19:19.584823Z","iopub.execute_input":"2025-05-29T13:19:19.585179Z","iopub.status.idle":"2025-05-29T13:19:33.305176Z","shell.execute_reply.started":"2025-05-29T13:19:19.585145Z","shell.execute_reply":"2025-05-29T13:19:33.304026Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile /kaggle/working/DRfold2/DRfold_infer.py\nimport os,sys\nimport torch\nimport numpy as np\nfrom subprocess import Popen, PIPE, STDOUT\n\n# Get the directory where the script is located\nexp_dir = os.path.dirname(os.path.abspath(__file__))\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n# dlexps = ['cfg_95','cfg_96','cfg_97','cfg_99']\ndlexps = ['cfg_97']\n\nprint(f\"[DRfold2] Starting prediction pipeline on {device} device\")\n\n# Get input FASTA file and output directory from command line arguments\nfastafile = os.path.realpath(sys.argv[1])\noutdir = os.path.realpath(sys.argv[2])\n\nprint(f\"[DRfold2] Input: {fastafile}\")\nprint(f\"[DRfold2] Output: {outdir}\")\n\n# Initialize clustering flag\npclu = False\n\n# If third argument is '1', enable clustering\nif len(sys.argv) == 4 and sys.argv[3] == '1': \n print('[DRfold2] Clustering enabled - will generate multiple models')\n pclu = True\nelse:\n print('[DRfold2] Clustering disabled - will generate single model')\n\n# Create output directory if it doesn't exist\nif not os.path.isdir(outdir):\n os.makedirs(outdir)\n print(f\"[DRfold2] Created output directory: {outdir}\")\n\n# Create subdirectories for different outputs\nret_dir = os.path.join(outdir,'rets_dir') # For return files\nif not os.path.isdir(ret_dir):\n os.makedirs(ret_dir)\n print(f\"[DRfold2] Created returns directory: {ret_dir}\")\n\nfolddir = os.path.join(outdir,'folds') # For folded structures\nif not os.path.isdir(folddir):\n os.makedirs(folddir)\n print(f\"[DRfold2] Created folds directory: {folddir}\")\n\nrefdir = os.path.join(outdir,'relax') # For relaxed structures\nif not os.path.isdir(refdir):\n os.makedirs(refdir)\n print(f\"[DRfold2] Created relaxation directory: {refdir}\")\n\n# Helper function to run commands and capture output\ndef run_cmd(cmd, description):\n print(f\"[DRfold2] {description}\")\n print(f\"[DRfold2] Command: {cmd}\")\n \n # Execute the command and capture output in real-time\n process = Popen(cmd, shell=True, stdout=PIPE, stderr=STDOUT, universal_newlines=True, bufsize=1)\n \n # Print output line by line as it becomes available\n for line in iter(process.stdout.readline, ''):\n line = line.strip()\n if line:\n print(f\"[DRfold2 subprocess] {line}\")\n \n # Get return code\n return_code = process.wait()\n if return_code == 0:\n print(f\"[DRfold2] {description} completed successfully\")\n else:\n print(f\"[DRfold2] {description} failed with return code {return_code}\")\n return return_code\n\n# Create paths for model directories and test scripts\ndlmains = [os.path.join(exp_dir, one_exp, 'test_modeldir.py') for one_exp in dlexps]\ndirs = [os.path.join(exp_dir, 'model_hub', one_exp) for one_exp in dlexps]\n\n# Check if processing has been done before\nif not os.path.isfile(ret_dir + '/done'): \n print(\"[DRfold2] Step 1/4: GENERATING INITIAL PREDICTIONS\")\n print(f\"[DRfold2] No previous predictions found, will generate e2e and geo files\")\n \n # Run each model configuration\n for idx, (dlmain, one_exp, mdir) in enumerate(zip(dlmains, dlexps, dirs)):\n # Construct command to run the model\n cmd = f'python {dlmain} {device} {fastafile} {ret_dir}/{one_exp}_ {mdir}'\n description = f\"Running model {idx+1}/{len(dlexps)}: {one_exp}\"\n run_cmd(cmd, description)\n\n # Mark processing as complete\n wfile = open(ret_dir+'/done','w')\n wfile.write('1')\n wfile.close()\n print(\"[DRfold2] Initial predictions generation completed\")\nelse:\n print(\"[DRfold2] Step 1/4: USING EXISTING PREDICTIONS\")\n print(f\"[DRfold2] Found previous predictions in {ret_dir}, using existing e2e and geo files\")\n\n# Helper function to get model PDB file\ndef get_model_pdb(tdir,opt):\n files = os.listdir(tdir)\n files = [afile for afile in files if afile.startswith(opt)][0]\n return files\n\n# Set up directory paths and configuration files\ncso_dir = folddir # Directory for coarse-grained structures\nclufile = os.path.join(folddir,'clu.txt') # Clustering results file\nconfig_sel = os.path.join(exp_dir,'cfg_for_selection.json') # Selection configuration\nfoldconfig = os.path.join(exp_dir,'cfg_for_folding.json') # Folding configuration\nselpython = os.path.join(exp_dir,'PotentialFold','Selection.py') # Selection script\noptpython = os.path.join(exp_dir,'PotentialFold','Optimization.py') # Optimization script\nclupy = os.path.join(exp_dir,'PotentialFold','Clust.py') # Clustering script\narena = os.path.join(exp_dir,'Arena','Arena') # Arena executable for structure refinement\n\n# Set up initial save prefixes for optimization and selection\noptsaveprefix = os.path.join(cso_dir, f'opt_0')\nsave_prefix = os.path.join(cso_dir, f'sel_0')\n\n# Get all .ret files from the return directory\nrets = os.listdir(ret_dir)\nrets = [afile for afile in rets if afile.endswith('.ret')]\nrets = [os.path.join(ret_dir,aret) for aret in rets ]\nret_str = ' '.join(rets)\n\nprint(\"[DRfold2] Step 2/4: SELECTION PROCESS\")\nprint(f\"[DRfold2] Found {len(rets)} return files for selection\")\nprint(f\"[DRfold2] Using selection config: {config_sel}\")\nprint(f\"[DRfold2] Output prefix: {save_prefix}\")\n\n# Run selection process\ncmd = f'python {selpython} {fastafile} {config_sel} {save_prefix} {ret_str}'\nrun_cmd(cmd, \"Running selection process\")\n\nprint(\"[DRfold2] Step 3/4: OPTIMIZATION PROCESS\")\nprint(f\"[DRfold2] Using fold config: {foldconfig}\")\nprint(f\"[DRfold2] Optimization output prefix: {optsaveprefix}\")\n\n# Run optimization process\ncmd = f'python {optpython} {fastafile} {optsaveprefix} {ret_dir} {save_prefix} {foldconfig}'\nrun_cmd(cmd, \"Running optimization process\")\n\n# Get the coarse-grained PDB and save refined structure\ncgpdb = os.path.join(folddir,get_model_pdb(folddir,'opt_0'))\nsavepdb = os.path.join(refdir,'model_1.pdb')\n\nprint(\"[DRfold2] Step 4/4: STRUCTURE REFINEMENT\")\nprint(f\"[DRfold2] Found optimized structure: {cgpdb}\")\nprint(f\"[DRfold2] Final output will be saved to: {savepdb}\")\n\ncmd = f'{arena} {cgpdb} {savepdb} 7'\nrun_cmd(cmd, \"Running structure refinement\")\n\n# If clustering is enabled (pclu=True)\nif pclu:\n print(\"[DRfold2] ADDITIONAL STEP: CLUSTERING\")\n print(f\"[DRfold2] Running clustering process, output: {clufile}\")\n \n # Run clustering process\n cmd = f'python {clupy} {ret_dir} {clufile}'\n run_cmd(cmd, \"Running clustering\")\n\n # Read clustering results\n lines = open(clufile).readlines()\n lines = [aline.strip() for aline in lines]\n lines = [aline for aline in lines if aline]\n \n cluster_count = len(lines) - 1\n print(f\"[DRfold2] Found {cluster_count} additional clusters to process\")\n\n # Process each cluster\n for i in range(1,len(lines)):\n print(f\"[DRfold2] PROCESSING CLUSTER {i}/{cluster_count}\")\n \n # Get return files for this cluster\n rets = lines[i].split()\n rets = [os.path.join(ret_dir,aret.replace('.pdb','.ret')) for aret in rets ]\n ret_str = ' '.join(rets)\n\n # Set up save prefixes for this cluster\n optsaveprefix = os.path.join(cso_dir,f'opt_{str(i+1)}')\n save_prefix = os.path.join(cso_dir,f'sel_{str(i+1)}')\n \n print(f\"[DRfold2] Cluster {i} Selection Process\")\n print(f\"[DRfold2] Found {len(rets)} return files for selection\")\n print(f\"[DRfold2] Selection output prefix: {save_prefix}\")\n\n # Run selection process for this cluster\n cmd = f'python {selpython} {fastafile} {config_sel} {save_prefix} {ret_str}'\n run_cmd(cmd, f\"Running selection for cluster {i}\")\n \n print(f\"[DRfold2] Cluster {i} Optimization Process\")\n print(f\"[DRfold2] Optimization output prefix: {optsaveprefix}\")\n\n # Run optimization process for this cluster\n cmd = f'python {optpython} {fastafile} {optsaveprefix} {ret_dir} {save_prefix} {foldconfig}'\n run_cmd(cmd, f\"Running optimization for cluster {i}\")\n\n # Get the coarse-grained PDB and save refined structure for this cluster\n cgpdb = os.path.join(folddir,get_model_pdb(folddir,f'opt_{str(i+1)}'))\n savepdb = os.path.join(refdir,f'model_{str(i+1)}.pdb')\n \n print(f\"[DRfold2] Cluster {i} Refinement Process\")\n print(f\"[DRfold2] Found optimized structure: {cgpdb}\")\n print(f\"[DRfold2] Final output will be saved to: {savepdb}\")\n\n cmd = f'{arena} {cgpdb} {savepdb} 7'\n run_cmd(cmd, f\"Running refinement for cluster {i}\")\n\nprint(\"[DRfold2] PREDICTION PIPELINE COMPLETED SUCCESSFULLY\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:19:33.306380Z","iopub.execute_input":"2025-05-29T13:19:33.306765Z","iopub.status.idle":"2025-05-29T13:19:33.315256Z","shell.execute_reply.started":"2025-05-29T13:19:33.306731Z","shell.execute_reply":"2025-05-29T13:19:33.313921Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile /kaggle/working/DRfold2/PotentialFold/operations.py\n\"\"\"\noperations.py: Core Mathematical Operations for RNA Structure Analysis\n\nThis module provides essential mathematical operations for manipulating and analyzing\nRNA 3D structures, organized into four main categories:\n\n1. Basic Vector Operations:\n Functions for selecting coordinates and calculating distances between points,\n which form the foundation for all structural calculations.\n\n2. Angle Calculations:\n Functions for computing bond angles and dihedral (torsion) angles between atoms,\n with differentiable implementations suitable for gradient-based optimization.\n\n3. Rigid Body Transformations:\n Functions for determining optimal rotations and translations between sets of\n coordinates, enabling structure alignment and manipulation.\n\n4. Sequence Utilities:\n Functions for converting RNA sequence data into standard 3D coordinate templates,\n allowing sequence-structure mapping.\n\nThese operations support the core functionality of RNA structure prediction, analysis,\nand optimization throughout the codebase.\n\"\"\"\n\nimport os\nimport torch\nimport torch.nn as nn\nimport numpy as np \nimport math, sys, math\nfrom io import BytesIO\nimport torch.nn.functional as F\nfrom torch.autograd import Function\nfrom torch.nn.parameter import Parameter\nfrom subprocess import Popen, PIPE, STDOUT\n\n# Use consistent epsilon value across all functions\nEPS = 1e-8\n\n\n# === Basic Vector Operations ===\ndef coor_selection(coor,mask):\n #[L,n,3],[L,n],byte\n return torch.masked_select(coor,mask.bool()).view(-1,3)\n\ndef pair_distance(x1, x2, eps=1e-6, p=2):\n # Use torch.cdist for p=2 (Euclidean) which is highly optimized\n if p == 2:\n return torch.cdist(x1, x2, p=2)\n \n # For other p-norms, avoid memory expansion with broadcasting\n x1_ = x1.unsqueeze(1) # [n1, 1, dim]\n x2_ = x2.unsqueeze(0) # [1, n2, dim]\n diff = torch.abs(x1_ - x2_)\n out = torch.pow(diff + eps, p).sum(dim=2)\n return torch.pow(out, 1. / p)\n\n\n# === Angle Calculations ===\ndef angle(p0, p1, p2):\n # [b 3] \n b0 = p0-p1\n b1 = p2-p1\n\n b0 = b0 / (torch.norm(b0, dim =-1, keepdim=True) + EPS)\n b1 = b1 / (torch.norm(b1, dim =-1, keepdim=True) + EPS)\n \n recos = torch.sum(b0*b1, -1)\n recos = torch.clamp(recos, -0.9999, 0.9999)\n return torch.acos(recos)\n\nclass torsion(Function):\n #PyTorch class to calculate differentiable torsion angle\n #https://stackoverflow.com/questions/20305272/dihedral-torsion-angle-from-four-points-in-cartesian-coordinates-in-python\n #https://salilab.org/modeller/manual/node492.html\n @staticmethod\n def forward(ctx, p0, p1, p2, p3):\n # Save input points for backward pass\n ctx.save_for_backward(p0, p1, p2, p3)\n\n # Calculate bond vectors\n b0 = p0 - p1\n b1 = p2 - p1\n b2 = p3 - p2\n\n # Normalize the middle bond vector\n b1_norm = torch.norm(b1, dim=-1, keepdim=True) + 1e-8\n b1_unit = b1 / b1_norm\n\n # Project the other bonds onto the plane perpendicular to middle bond\n v = b0 - torch.sum(b0 * b1_unit, dim=-1, keepdim=True) * b1_unit\n w = b2 - torch.sum(b2 * b1_unit, dim=-1, keepdim=True) * b1_unit\n\n # Calculate torsion using the arctan2 formula (more stable than arccos)\n x = torch.sum(v * w, dim=-1) # cosine component\n y = torch.sum(torch.cross(b1_unit, v, dim=-1) * w, dim=-1) # sine component\n\n return torch.atan2(y, x)\n\n \n @staticmethod\n def backward(ctx, grad_output):\n # Retrieve saved tensors from forward pass\n p0, p1, p2, p3 = ctx.saved_tensors\n\n # Calculate bond vectors\n r01 = p0 - p1\n r12 = p2 - p1\n r23 = p3 - p2\n\n # Calculate bond lengths with numerical stability\n d01 = torch.norm(r01, dim=-1, keepdim=True) + 1e-8\n d12 = torch.norm(r12, dim=-1, keepdim=True) + 1e-8\n d23 = torch.norm(r23, dim=-1, keepdim=True) + 1e-8\n\n # Normalize bond vectors\n e01 = r01 / d01\n e12 = r12 / d12\n e23 = r23 / d23\n\n # Calculate normal vectors to the two planes\n n1 = torch.cross(e01, e12, dim=-1)\n n2 = torch.cross(e12, e23, dim=-1)\n\n # Normalize normal vectors\n n1_norm = torch.norm(n1, dim=-1, keepdim=True) + 1e-8\n n2_norm = torch.norm(n2, dim=-1, keepdim=True) + 1e-8\n n1 = n1 / n1_norm\n n2 = n2 / n2_norm\n\n # Calculate gradients for each atom\n # These are based on the analytical derivatives of dihedral angles\n g0 = torch.cross(e01, n1, dim=-1) / d01\n g1 = -g0 - torch.cross(e12, n1, dim=-1) / d12\n g2 = torch.cross(e12, n2, dim=-1) / d12 - torch.cross(e23, n2, dim=-1) / d23\n g3 = torch.cross(e23, n2, dim=-1) / d23\n\n # Apply chain rule with incoming gradient\n g0 = g0 * grad_output.unsqueeze(-1)\n g1 = g1 * grad_output.unsqueeze(-1)\n g2 = g2 * grad_output.unsqueeze(-1)\n g3 = g3 * grad_output.unsqueeze(-1)\n\n return g0, g1, g2, g3\n\n\ndef dihedral(input1, input2, input3, input4):\n return torsion.apply(input1, input2, input3, input4)\n\n\n\n# === Rigid Body Transformations ===\ndef rigidFrom3Points(x): \n x1, x2, x3 = x[:, 0], x[:, 1], x[:, 2]\n v1 = x3 - x2\n v2 = x1 - x2\n \n # Normalize v1 to get e1\n e1 = F.normalize(v1, p=2, dim=-1)\n \n # Project v2 onto e1 and subtract to get the component orthogonal to e1\n u2 = v2 - e1 * (torch.einsum('bn,bn->b', e1, v2)[:, None])\n \n # Normalize u2 to get e2\n e2 = F.normalize(u2, p=2, dim=-1)\n \n # Cross product to get e3\n e3 = torch.cross(e1, e2, dim=-1)\n \n return torch.stack([e1, e2, e3], dim=1)\n\n\n# return the direction from to_q to from_p\ndef Kabsch_rigid(bases,x1,x2,x3):\n # Early return for empty input\n if x1.shape[0] == 0:\n return torch.empty(0, 3, 3), torch.empty(0, 3)\n \n the_dim=1\n to_q = torch.stack([x1,x2,x3],dim=the_dim)\n biasq=torch.mean(to_q,dim=the_dim,keepdim=True)\n q=to_q-biasq\n m = torch.einsum('bnz,bny->bzy',bases,q)\n u, s, v = torch.svd(m)\n vt = torch.transpose(v, 1, 2)\n det = torch.det(torch.matmul(u, vt))\n det = det.view(-1, 1, 1)\n vt = torch.cat((vt[:, :2, :], vt[:, -1:, :] * det), 1)\n r = torch.matmul(u, vt)\n return r,biasq.squeeze()\n\n\n\n# === Sequence Utilities ===\ndef Get_base(seq,basenpy_standard):\n base_num = basenpy_standard.shape[1]\n basenpy = np.zeros([len(seq),base_num,3])\n seqnpy = np.array(list(seq))\n basenpy[seqnpy=='A']=basenpy_standard[0]\n basenpy[seqnpy=='a']=basenpy_standard[0]\n\n basenpy[seqnpy=='G']=basenpy_standard[1]\n basenpy[seqnpy=='g']=basenpy_standard[1]\n\n basenpy[seqnpy=='C']=basenpy_standard[2]\n basenpy[seqnpy=='c']=basenpy_standard[2]\n\n basenpy[seqnpy=='U']=basenpy_standard[3]\n basenpy[seqnpy=='u']=basenpy_standard[3]\n\n basenpy[seqnpy=='T']=basenpy_standard[3]\n basenpy[seqnpy=='t']=basenpy_standard[3]\n \n return torch.from_numpy(basenpy).double()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:19:33.316283Z","iopub.execute_input":"2025-05-29T13:19:33.316543Z","iopub.status.idle":"2025-05-29T13:19:33.341229Z","shell.execute_reply.started":"2025-05-29T13:19:33.316516Z","shell.execute_reply":"2025-05-29T13:19:33.339855Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile /kaggle/working/DRfold2/PotentialFold/Optimization.py\n#! /nfs/amino-home/liyangum/miniconda3/bin/python\nimport torch\nimport random\nimport numpy as np \nimport os, json, sys\n\nimport Cubic, Potential\nimport operations\nimport a2b, rigid\nimport torch.optim as opt\nfrom scipy.optimize import minimize\nimport pickle\n\ntorch.manual_seed(6)\nnp.random.seed(9)\nrandom.seed(9)\n\n\nScale_factor = 1.0\nUSEGEO = False\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n\ndef readconfig(configfile=''):\n config=[]\n expdir=os.path.dirname(os.path.abspath(__file__))\n if configfile=='':\n configfile=os.path.join(expdir,'lib','ddf.json')\n config=json.load(open(configfile,'r'))\n return config \n\n \nclass Structure:\n def __init__(self,fastafile,geofiles,saveprefix,initial_ret,foldconfig):\n self.config=readconfig(foldconfig)\n self.seqfile=fastafile\n self.init_ret = initial_ret\n self.foldconfig = foldconfig\n self.geofiles = geofiles\n self.rets = [pickle.load(open(refile,'rb')) for refile in geofiles]\n self.txs=[]\n for ret in self.rets:\n self.txs.append( torch.from_numpy(ret['coor']).double().to(device))\n self.handle_geo()\n self.pair = []\n for ret in self.rets:\n self.pair.append(torch.from_numpy(ret['plddt']).double().to(device))\n self.saveprefix=saveprefix\n self.seq=open(fastafile).readlines()[1].strip()\n self.L=len(self.seq)\n basenpy = np.load(os.path.join(os.path.dirname(os.path.abspath(__file__)),'lib','base.npy'))\n self.basex = operations.Get_base(self.seq,basenpy).to(device)\n othernpy = np.load(os.path.join(os.path.dirname(os.path.abspath(__file__)),'lib','other2.npy'))\n self.otherx = operations.Get_base(self.seq,othernpy).to(device)\n sidenpy = np.load(os.path.join(os.path.dirname(os.path.abspath(__file__)),'lib','side.npy'))\n self.sidex = operations.Get_base(self.seq,sidenpy).to(device)\n \n self.init_mask()\n self.init_paras()\n self._init_fape()\n self.tx2ds = [td.to(device) for td in self.tx2ds]\n self.local_weight = torch.ones(self.L,self.L).to(device)\n \n for i in range(self.L):\n for j in range(i+1,min(self.L,i+2)):\n self.local_weight[i,j] = self.local_weight[j,i] = 4\n for j in range(i+2,min(self.L,i+3)):\n self.local_weight[i,j] = self.local_weight[j,i] = 3\n for j in range(i+3,min(self.L,i+4)):\n self.local_weight[i,j] = self.local_weight[j,i] = 2\n\n def _init_fape(self):\n self.tx2ds = []\n for tx in self.txs:\n true_rot,true_trans = operations.Kabsch_rigid(self.basex,tx[:,0],tx[:,1],tx[:,2])\n true_x2 = tx[:,None,:,:] - true_trans[None,:,None,:]\n true_x2 = torch.einsum('ijnd,jde->ijne',true_x2,true_rot.transpose(-1,-2))\n self.tx2ds.append(true_x2)\n \n def handle_geo(self):\n oldkeys=['dist_p','dist_c','dist_n']\n newkeys=['pp','cc','nn']\n self.geos=[]\n for ret in self.rets:\n geo = {}\n for nk,ok in zip(newkeys,oldkeys):\n geo[nk] = torch.from_numpy(ret[ok].astype(np.float64)).to(device) + 0\n self.geos.append(geo)\n\n\n def init_mask(self):\n halfmask=np.zeros([self.L,self.L])\n fullmask=np.zeros([self.L,self.L])\n for i in range(self.L):\n for j in range(i+1,self.L):\n halfmask[i,j]=1\n fullmask[i,j]=1\n fullmask[j,i]=1\n self.halfmask=(torch.DoubleTensor(halfmask) > 0.5).to(device)\n self.fullmask=(torch.DoubleTensor(fullmask) > 0.5).to(device)\n self.clash_mask = torch.zeros([self.L,self.L,22,22], device=device)\n for i in range(self.L):\n for j in range(i+1,self.L):\n self.clash_mask[i,j]=1\n\n for i in range(self.L):\n self.clash_mask[i,i,:6,7:]=1\n\n for i in range(self.L-1):\n self.clash_mask[i,i+1,:,0]=0\n self.clash_mask[i,i+1,0,:]=0\n self.clash_mask[i,i+1,:,5]=0\n self.clash_mask[i,i+1,5,:]=0\n\n self.side_mask = rigid.side_mask(self.seq).to(device)\n self.side_mask = (self.side_mask[:,None,:,None] * self.side_mask[None,:,None,:]).to(device)\n self.clash_mask = ((self.clash_mask > 0.5) * (self.side_mask > 0.5)).to(device)\n\n self.geo_confimask_cc = []\n self.geo_confimask_pp = []\n self.geo_confimask_nn = []\n for geo in self.geos:\n confimask_cc = geo['cc'][:,:,-1] < 0.5\n confimask_pp = geo['pp'][:,:,-1] < 0.5\n confimask_nn = geo['nn'][:,:,-1] < 0.5\n self.geo_confimask_cc.append(confimask_cc)\n self.geo_confimask_pp.append(confimask_pp)\n self.geo_confimask_nn.append(confimask_nn)\n\n\n def init_paras(self):\n self.geo_cc = []\n self.geo_pp = []\n self.geo_nn = []\n self.cs_coefs = {'cc': [], 'pp': [], 'nn': []}\n self.cs_knots = {'cc': [], 'pp': [], 'nn': []}\n for geo in self.geos:\n cc_cs,cc_decs=Cubic.dis_cubic(geo['cc'],2,40,36)\n pp_cs,pp_decs=Cubic.dis_cubic(geo['pp'],2,40,36)\n nn_cs,nn_decs=Cubic.dis_cubic(geo['nn'],2,40,36)\n self.geo_cc.append([cc_cs,cc_decs])\n self.geo_pp.append([pp_cs,pp_decs])\n self.geo_nn.append([nn_cs,nn_decs])\n \n L = self.L\n cc_coefs_np = np.stack([[cc_cs[i,j].c for j in range(L)] for i in range(L)], axis=0)\n cc_knots_np = np.stack([[cc_cs[i,j].x for j in range(L)] for i in range(L)], axis=0)\n self.cs_coefs['cc'].append(torch.from_numpy(cc_coefs_np).to(device))\n self.cs_knots['cc'].append(torch.from_numpy(cc_knots_np).to(device))\n \n pp_coefs_np = np.stack([[pp_cs[i,j].c for j in range(L)] for i in range(L)], axis=0)\n pp_knots_np = np.stack([[pp_cs[i,j].x for j in range(L)] for i in range(L)], axis=0)\n self.cs_coefs['pp'].append(torch.from_numpy(pp_coefs_np).to(device))\n self.cs_knots['pp'].append(torch.from_numpy(pp_knots_np).to(device))\n \n nn_coefs_np = np.stack([[nn_cs[i,j].c for j in range(L)] for i in range(L)], axis=0)\n nn_knots_np = np.stack([[nn_cs[i,j].x for j in range(L)] for i in range(L)], axis=0)\n self.cs_coefs['nn'].append(torch.from_numpy(nn_coefs_np).to(device))\n self.cs_knots['nn'].append(torch.from_numpy(nn_knots_np).to(device))\n\n\n def compute_bb_clash(self,coor,other_coor):\n com_coor = torch.cat([coor,other_coor],dim=1)\n com_dis = (com_coor[:,None,:,None,:] - com_coor[None,:,None,:,:]).norm(dim=-1)\n dynamicmask2_vdw= (com_dis <= 3.15) * (self.clash_mask)\n vdw_dynamic = Potential.LJpotential(com_dis[dynamicmask2_vdw],3.15)\n return vdw_dynamic.sum()*self.config['weight_vdw']\n\n def compute_full_clash(self,coor,other_coor,side_coor):\n com_coor = torch.cat([coor[:,:2],other_coor,side_coor],dim=1)\n com_dis = (com_coor[:,None,:,None,:] - com_coor[None,:,None,:,:]).norm(dim=-1)\n dynamicmask2_vdw= (com_dis <= 2.5) * (self.clash_mask)\n vdw_dynamic = Potential.LJpotential(com_dis[dynamicmask2_vdw],2.5)\n return vdw_dynamic.sum()*self.config['weight_vdw']\n\n\n def _cubic_pair_energy(self, atom_map, geo_cs, geo_confimask, weight_key):\n \"\"\"General cubic-spline energy for CC/PP/NN pairs.\"\"\"\n min_dis, max_dis, bin_num = 2, 40, 36\n dev = atom_map.device\n upper_th = max_dis - ((max_dis - min_dis) / bin_num) * 0.5\n lower_th = 2.5\n total = torch.zeros((), device=dev, dtype=torch.double)\n spline_key = weight_key.split('_')[1] # 'cc', 'pp', or 'nn'\n coeffs_list = self.cs_coefs[spline_key]\n knots_list = self.cs_knots[spline_key]\n for block_idx, mask_block in enumerate(geo_confimask):\n mask = (atom_map <= upper_th) & mask_block & self.fullmask & (atom_map >= lower_th)\n idx = mask.nonzero(as_tuple=True)\n if idx[0].numel() > 1:\n coef = coeffs_list[block_idx][idx]\n knots = knots_list[block_idx][idx]\n part1 = Potential.cubic_distance(atom_map[mask], coef, knots, min_dis, max_dis, bin_num).sum() * self.config[weight_key] * 0.5\n else:\n part1 = torch.zeros((), device=dev)\n part2 = ((atom_map <= lower_th) & mask_block & self.fullmask).sum() * self.config[weight_key]\n total = total + part1 + part2\n return total\n\n def compute_cc_energy(self, coor):\n atom_map = operations.pair_distance(coor[:,1], coor[:,1])\n return self._cubic_pair_energy(atom_map, self.geo_cc, self.geo_confimask_cc, 'weight_cc')\n \n def compute_pp_energy(self, coor):\n atom_map = operations.pair_distance(coor[:,0], coor[:,0])\n return self._cubic_pair_energy(atom_map, self.geo_pp, self.geo_confimask_pp, 'weight_pp')\n \n def compute_nn_energy(self, coor):\n atom_map = operations.pair_distance(coor[:,-1], coor[:,-1])\n return self._cubic_pair_energy(atom_map, self.geo_nn, self.geo_confimask_nn, 'weight_nn')\n\n def compute_pccp_energy(self,coor):\n p_atoms=coor[:,0]\n c_atoms=coor[:,1]\n pccpmap=operations.dihedral( p_atoms[self.pccpi], c_atoms[self.pccpi], c_atoms[self.pccpj] ,p_atoms[self.pccpj] )\n neg_log = Potential.cubic_torsion(pccpmap,self.pccp_coe,self.pccp_x,36)\n return neg_log.sum()*self.config['weight_pccp']\n\n def compute_cnnc_energy(self,coor):\n n_atoms=coor[:,-1]\n c_atoms=coor[:,1]\n pccpmap=operations.dihedral( c_atoms[self.cnnci], n_atoms[self.cnnci], n_atoms[self.cnncj] ,c_atoms[self.cnncj] )\n neg_log = Potential.cubic_torsion(pccpmap,self.cnnc_coe,self.cnnc_x,36)\n return neg_log.sum()*self.config['weight_cnnc']\n\n def compute_pnnp_energy(self,coor):\n n_atoms=coor[:,-1]\n p_atoms=coor[:,0]\n pccpmap=operations.dihedral( p_atoms[self.pnnpi], n_atoms[self.pnnpi], n_atoms[self.pnnpj] ,p_atoms[self.pnnpj] )\n neg_log = Potential.cubic_torsion(pccpmap,self.pnnp_coe,self.pnnp_x,36)\n return neg_log.sum()*self.config['weight_pnnp']\n\n def compute_pcc_energy(self,coor):\n p_atoms=coor[:,1]\n c_atoms=coor[:,2]\n pccmap=operations.angle( p_atoms[self.pcci], c_atoms[self.pcci], c_atoms[self.pccj] )\n neg_log = Potential.cubic_angle(pccmap,self.pcc_coe,self.pcc_x,12)\n return neg_log.sum()*self.config['weight_pcc']\n\n def compute_fape_energy(self,coor,ep=1e-3,epmax=20):\n energy= 0\n for tx in self.tx2ds:\n px_mean = coor[:,[1]]\n p_rot = operations.rigidFrom3Points(coor)\n p_tran = px_mean[:,0]\n pred_x2 = coor[:,None,:,:] - p_tran[None,:,None,:] # Lx Lrot N , 3\n pred_x2 = torch.einsum('ijnd,jde->ijne',pred_x2,p_rot.transpose(-1,-2)) # transpose should be equal to inverse\n errmap=torch.sqrt( ((pred_x2 - tx)**2).sum(dim=-1) + ep )\n energy = energy + torch.sum( torch.clamp(errmap,max=epmax) )\n return energy * self.config['weight_fape']\n\n def compute_bond_energy(self,coor,other_coor):\n # 3.87\n o3 = other_coor[:-1,-2]\n p = coor[1:,0]\n dis = (o3-p).norm(dim=-1)\n energy = ((dis-1.607)**2).sum()\n return energy * self.config['weight_bond']\n\n def tooth_func(self,errmap, ep = 0.05):\n return -1/(errmap/10+ep) + (1/ep)\n\n def reweight_func(self,ww):\n reweighting = torch.pow(ww,self.config['pair_weight_power'])\n reweighting[ww < self.config['pair_weight_min']] = 0\n return reweighting\n\n def compute_fape_energy_fromquat(self,x,coor,ep=1e-6,epmax=100):\n energy= 0\n p_rot,px_mean = a2b.Non2rot(x[:,:9],x.shape[0]),x[:,9:]\n pred_x2 = coor[:,None,:,:] - px_mean[None,:,None,:] # Lx Lrot N , 3\n pred_x2 = torch.einsum('ijnd,jde->ijne',pred_x2,p_rot.transpose(-1,-2)) # transpose should be equal to inverse\n for tx,weightplddt in zip(self.tx2ds,self.pair):\n\n tamplate_dist_map = torch.min( tx.norm(dim=-1), dim=2 )[0]\n errmap=torch.sqrt( ((pred_x2 - tx)**2).sum(dim=-1) + ep ) \n energy = energy + torch.sum( ( (torch.clamp(errmap,max=self.config['FAPE_max'])**self.config['pair_error_power']) * self.reweight_func(weightplddt[...,None]) * self.local_weight[...,None] )[tamplate_dist_map>self.config['pair_rest_min_dist']] )\n\n return energy * self.config['weight_fape']\n\n\n def energy(self,rama):\n coor=a2b.quat2b(self.basex,rama[:,9:])\n other_coor = a2b.quat2b(self.otherx,rama[:,9:])\n side_coor = a2b.quat2b(self.sidex,torch.cat([rama[:,:9],coor[:,-1]],dim=-1))\n \n if self.config['weight_cc']>0:\n E_cc= self.compute_cc_energy(coor) / len(self.rets)\n else:\n E_cc=0\n if self.config['weight_pp']>0:\n E_pp= self.compute_pp_energy(coor) / len(self.rets)\n else:\n E_pp=0\n if self.config['weight_nn']>0:\n E_nn= self.compute_nn_energy(coor) / len(self.rets)\n else:\n E_nn=0\n\n if self.config['weight_pccp']>0:\n E_pccp= self.compute_pccp_energy(coor) / len(self.rets)\n else:\n E_pccp=0\n\n if self.config['weight_cnnc']>0:\n E_cnnc= self.compute_cnnc_energy(coor) / len(self.rets)\n else:\n E_cnnc=0\n\n if self.config['weight_pnnp']>0:\n E_pnnp= self.compute_pnnp_energy(coor) / len(self.rets)\n else:\n E_pnnp=0\n\n if self.config['weight_vdw']>0:\n E_vdw= self.compute_full_clash(coor,other_coor,side_coor)\n else:\n E_vdw=0\n\n if self.config['weight_fape']>0:\n E_fape= self.compute_fape_energy_fromquat(rama[:,9:],coor) / len(self.rets)\n else:\n E_fape=0\n if self.config['weight_bond']>0:\n E_bond= self.compute_bond_energy(coor,other_coor)\n else:\n E_bond=0\n return E_vdw + E_fape + E_bond + E_pp + E_cc + E_nn + E_pccp + E_cnnc + E_pnnp\n\n\n def obj_func_grad_np(self,rama_):\n rama=torch.DoubleTensor(rama_)\n rama.requires_grad=True\n if rama.grad:\n rama.grad.zero_()\n f=self.energy(rama.view(self.L,21))*Scale_factor\n grad_value=autograd.grad(f,rama)[0]\n return grad_value.data.numpy().astype(np.float64)\n \n def obj_func_np(self,rama_):\n rama=torch.DoubleTensor(rama_)\n rama=rama.view(self.L,21)\n with torch.no_grad():\n f=self.energy(rama)*Scale_factor\n return f.item()\n\n\n def foldning(self):\n ilter = self.init_ret\n # 1) get initial quaternions (double precision)\n try:\n init_q = self.init_quat(ilter).double()\n except:\n init_q = self.init_quat_safe(ilter).double()\n\n # 2) move to target device (GPU if available), enable grad\n param = init_q.to(device).clone().detach().requires_grad_(True)\n\n # 3) set up PyTorch LBFGS optimizer over `param`\n optimizer = opt.LBFGS(\n [param],\n max_iter=self.config.get('max_iter', 300),\n tolerance_grad=1e-6,\n tolerance_change=1e-9,\n history_size=10,\n line_search_fn='strong_wolfe'\n )\n\n # 4) define the “closure” that LBFGS will call to reevaluate loss + gradients\n def closure():\n optimizer.zero_grad() # clear old grads\n E = self.energy(param.view(self.L,21)) * Scale_factor # compute ∂E/∂param\n E.backward()\n return E\n\n # 5) run LBFGS until convergence (it calls closure repeatedly)\n optimizer.step(closure)\n\n # 6) write out final PDB\n final_energy = self.energy(param.view(self.L,21)).item()\n self.outpdb(param, self.saveprefix + '.pdb', energystr=str(final_energy))\n\n\n def outpdb(self,rama,savefile,start=0,end=10000,energystr=''):\n # bring baseframes and quaternion data onto CPU to prevent device mismatch\n basex_cpu = self.basex.detach().cpu()\n otherx_cpu = self.otherx.detach().cpu()\n sidex_cpu = self.sidex.detach().cpu()\n shaped_rama = rama.view(self.L,21).detach().cpu()\n # compute backbone and other coords\n coor_np = a2b.quat2b(basex_cpu, shaped_rama[:,9:]).detach().cpu().numpy()\n other_np = a2b.quat2b(otherx_cpu, shaped_rama[:,9:]).detach().cpu().numpy()\n coor = torch.FloatTensor(coor_np)\n # compute side atom coords\n side_coor_NP = a2b.quat2b(sidex_cpu, torch.cat([shaped_rama[:,:9], coor[:,-1]], dim=-1)).detach().cpu().numpy()\n \n Atom_name=[' P ',\" C4'\",' N1 ']\n Other_Atom_name = [\" O5'\",\" C5'\",\" C3'\",\" O3'\",\" C1'\"]\n other_last_name = ['O',\"C\",\"C\",\"O\",\"C\"]\n\n side_atoms= [' N1 ',' C2 ',' O2 ',' N2 ',' N3 ',' N4 ',' C4 ',' O4 ',' C5 ',' C6 ',' O6 ',' N6 ',' N7 ',' N8 ',' N9 ']\n side_last_name = ['N', \"C\", \"O\", \"N\", \"N\", 'N', 'C', 'O', 'C', 'C', 'O', 'N', 'N', 'N','N']\n\n base_dict = rigid.base_table()\n last_name=['P','C','N']\n wstr=[f'REMARK {str(energystr)}']\n templet='%6s%5d %4s %3s %1s%4d %8.3f%8.3f%8.3f%6.2f%6.2f %2s%2s'\n count=1\n for i in range(self.L):\n if self.seq[i] in ['a','g','A','G']:\n Atom_name = [' P ',\" C4'\",' N9 ']\n #atoms = ['P','C4']\n\n elif self.seq[i] in ['c','u','C','U']:\n Atom_name = [' P ',\" C4'\",' N1 ']\n for j in range(coor_np.shape[1]):\n outs=('ATOM ',count,Atom_name[j],self.seq[i],'A',i+1,coor_np[i][j][0],coor_np[i][j][1],coor_np[i][j][2],0,0,last_name[j],'')\n if i>=start-1 and i < end:\n wstr.append(templet % outs)\n count+=1\n\n for j in range(other_np.shape[1]):\n outs=('ATOM ',count,Other_Atom_name[j],self.seq[i],'A',i+1,other_np[i][j][0],other_np[i][j][1],other_np[i][j][2],0,0,other_last_name[j],'')\n if i>=start-1 and i < end:\n wstr.append(templet % outs)\n count+=1\n \n wstr='\\n'.join(wstr)\n wfile=open(savefile,'w')\n wfile.write(wstr)\n wfile.close()\n \n def outpdb_coor(self,coor_np,savefile,start=0,end=1000,energystr=''):\n Atom_name=[' P ',\" C4'\",' N1 ']\n last_name=['P','C','N']\n wstr=[f'REMARK {str(energystr)}']\n templet='%6s%5d %4s %3s %1s%4d %8.3f%8.3f%8.3f%6.2f%6.2f %2s%2s'\n count=1\n for i in range(self.L):\n if self.seq[i] in ['a','g','A','G']:\n Atom_name = [' P ',\" C4'\",' N9 ']\n\n elif self.seq[i] in ['c','u','C','U']:\n Atom_name = [' P ',\" C4'\",' N1 ']\n for j in range(coor_np.shape[1]):\n outs=('ATOM ',count,Atom_name[j],self.seq[i],'A',i+1,coor_np[i][j][0],coor_np[i][j][1],coor_np[i][j][2],0,0,last_name[j],'')\n if i>=start-1 and i < end:\n wstr.append(templet % outs)\n count+=1\n \n wstr='\\n'.join(wstr)\n wfile=open(savefile,'w')\n wfile.write(wstr)\n wfile.close()\n\n\n def init_quat(self,ii):\n x = torch.rand([self.L,21])\n x[:,18:] = self.txs[ii].mean(dim=1)\n init_coor = self.txs[ii]\n biasq = torch.mean(init_coor,dim=1,keepdim=True)\n q = init_coor - biasq\n m = torch.einsum('bnz,bny->bzy',self.basex,q).reshape([self.L,-1])\n x[:,:9] = x[:,9:18] = m\n x.requires_grad_()\n return x\n\n def init_quat_safe(self,ii):\n x = torch.rand([self.L,21])\n x[:,18:] = self.txs[ii].mean(dim=1)\n init_coor = self.txs[ii]\n biasq = torch.mean(init_coor,dim=1,keepdim=True)\n q = init_coor - biasq + torch.rand([self.L,3,3])\n m = (torch.einsum('bnz,bny->bzy',self.basex,q) + torch.eye(3)[None,:,:]).reshape([self.L,-1])\n x[:,:9] = x[:,9:18] = m\n x.requires_grad_()\n return x\n\n\nif __name__ == '__main__': \n\n fastafile=sys.argv[1]\n saveprefix=sys.argv[2]\n retdirs =sys.argv[3]\n ret_score = sys.argv[4]\n foldconfig = sys.argv[5]\n\n savepare = os.path.dirname(saveprefix)\n if not os.path.isdir(savepare):\n os.makedirs(savepare)\n\n num_of_models = readconfig(foldconfig)['num_of_models']\n\n score_dict = readconfig(ret_score)\n sorted_items = sorted(score_dict.items(), key=lambda x: x[1])\n lowest_n_keys = [item[0] for item in sorted_items][:num_of_models]\n bestkey = lowest_n_keys[0] + ''\n print(\"Before sort:\", lowest_n_keys)\n lowest_n_keys.sort()\n print(\"After sort:\", lowest_n_keys)\n bestindex = lowest_n_keys.index(bestkey)\n\n current_ret = bestkey\n retfiles = [os.path.join(retdirs, afile) for afile in lowest_n_keys]\n stru = Structure(fastafile, retfiles, saveprefix + '_from_' + current_ret, bestindex, foldconfig)\n stru.foldning()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:19:33.342519Z","iopub.execute_input":"2025-05-29T13:19:33.342851Z","iopub.status.idle":"2025-05-29T13:19:33.368821Z","shell.execute_reply.started":"2025-05-29T13:19:33.342804Z","shell.execute_reply":"2025-05-29T13:19:33.367268Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile /kaggle/working/DRfold2/PotentialFold/Selection.py\n#! /nfs/amino-home/liyangum/miniconda3/bin/python\nimport numpy\nimport torch\nimport torch.autograd as autograd\nimport numpy as np \n\nimport random\nimport Cubic, Potential\nimport operations\nimport os, json, sys\n\nimport a2b, rigid\nimport torch.optim as opt\nfrom torch.nn.parameter import Parameter\nimport torch.nn as nn\nimport math\nfrom scipy.optimize import fmin_l_bfgs_b,fmin_cg,fmin_bfgs\nfrom scipy.optimize import minimize\nimport lbfgs_rosetta\nimport pickle\nimport shutil\n\ntorch.manual_seed(6)\ntorch.set_num_threads(4)\nnp.random.seed(9)\nrandom.seed(9)\n\nScale_factor = 1.0\nUSEGEO = False\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\ndef readconfig(configfile=''):\n config=[]\n expdir=os.path.dirname(os.path.abspath(__file__))\n if configfile=='':\n configfile=os.path.join(expdir,'lib','ddf.json')\n config=json.load(open(configfile,'r'))\n return config \n\n \nclass Structure:\n def __init__(self, fastafile, geofiles, foldconfig, saveprefix):\n # Load Configuration and Inputs\n self.config = readconfig(foldconfig)\n self.seqfile = fastafile\n self.foldconfig = foldconfig\n self.geofiles = geofiles\n\n # Load Model Results\n self.rets = [pickle.load(open(refile, 'rb')) for refile in geofiles]\n \n # Extract Coordinates\n self.txs = []\n for ret in self.rets:\n self.txs.append(torch.from_numpy(ret['coor']).double().to(device))\n \n # Handle Geometrical Data\n self.handle_geo()\n\n # Extract pLDDT Scores\n self.pair = []\n for ret in self.rets:\n self.pair.append( torch.from_numpy(ret['plddt']).double().to(device))\n \n # Store Output and Sequence Info\n self.saveprefix = saveprefix\n self.seq = open(fastafile).readlines()[1].strip()\n self.L = len(self.seq)\n \n # Load Reference Arrays for Structure Construction\n basenpy = np.load(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'lib', 'base.npy'))\n self.basex = operations.Get_base(self.seq, basenpy).double().to(device)\n \n othernpy = np.load(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'lib', 'other2.npy'))\n self.otherx = operations.Get_base(self.seq, othernpy).double().to(device)\n \n sidenpy = np.load(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'lib', 'side.npy'))\n self.sidex = operations.Get_base(self.seq, sidenpy).double().to(device) \n \n # Initialize Masks, Parameters, and FAPE\n self.init_mask()\n self.init_paras()\n self._init_fape()\n \n\n def _init_fape(self):\n self.tx2ds = []\n for tx in self.txs:\n true_rot, true_trans = operations.Kabsch_rigid(self.basex, tx[:, 0], tx[:, 1], tx[:, 2])\n true_x2 = tx[:, None, :, :] - true_trans[None, :, None, :]\n true_x2 = torch.einsum('ijnd,jde->ijne', true_x2, true_rot.transpose(-1,-2))\n self.tx2ds.append(true_x2)\n \n\n def handle_geo(self):\n oldkeys = ['dist_p', 'dist_c', 'dist_n']\n newkeys = ['pp', 'cc', 'nn']\n self.geos = []\n geo = {'pp':0, 'cc':0, 'nn':0}\n \n for ret in self.rets: \n for nk, ok in zip(newkeys, oldkeys):\n geo[nk] = geo[nk] + (ret[ok].astype(np.float64) /(len(self.rets)))\n self.geos.append(geo)\n\n\n def init_mask(self):\n halfmask=np.zeros([self.L,self.L])\n fullmask=np.zeros([self.L,self.L])\n for i in range(self.L):\n for j in range(i+1,self.L):\n halfmask[i,j]=1\n fullmask[i,j]=1\n fullmask[j,i]=1\n self.halfmask=torch.DoubleTensor(halfmask) > 0.5\n self.fullmask=torch.DoubleTensor(fullmask) > 0.5\n self.clash_mask = torch.zeros([self.L,self.L,22,22])\n for i in range(self.L):\n for j in range(i+1,self.L):\n self.clash_mask[i,j]=1\n\n for i in range(self.L):\n self.clash_mask[i,i,:6,7:]=1\n\n for i in range(self.L-1):\n self.clash_mask[i,i+1,:,0]=0\n self.clash_mask[i,i+1,0,:]=0\n self.clash_mask[i,i+1,:,5]=0\n self.clash_mask[i,i+1,5,:]=0\n\n self.side_mask = rigid.side_mask(self.seq)\n self.side_mask = self.side_mask[:,None,:,None] * self.side_mask[None,:,None,:]\n self.clash_mask = (self.clash_mask > 0.5) * (self.side_mask > 0.5)\n\n self.geo_confimask_cc = []\n self.geo_confimask_pp = []\n self.geo_confimask_nn = []\n for geo in self.geos:\n confimask_cc = torch.DoubleTensor(geo['cc'][:,:,-1]) < 0.5\n confimask_pp = torch.DoubleTensor(geo['pp'][:,:,-1]) < 0.5\n confimask_nn = torch.DoubleTensor(geo['nn'][:,:,-1]) < 0.5\n self.geo_confimask_cc.append(confimask_cc)\n self.geo_confimask_pp.append(confimask_pp)\n self.geo_confimask_nn.append(confimask_nn)\n\n # Move masks and confimasks to the GPU/CPU device\n self.halfmask = self.halfmask.to(device)\n self.fullmask = self.fullmask.to(device)\n self.clash_mask = self.clash_mask.to(device)\n self.side_mask = self.side_mask.to(device)\n # geo_confimasks are lists\n self.geo_confimask_cc = [m.to(device) for m in self.geo_confimask_cc]\n self.geo_confimask_pp = [m.to(device) for m in self.geo_confimask_pp]\n self.geo_confimask_nn = [m.to(device) for m in self.geo_confimask_nn]\n\n\n def init_paras(self):\n self.geo_cc = []\n self.geo_pp = []\n self.geo_nn = []\n self.cs_coefs = {'cc': [], 'pp': [], 'nn': []}\n self.cs_knots = {'cc': [], 'pp': [], 'nn': []}\n for geo in self.geos:\n cc_cs, cc_decs = Cubic.dis_cubic(geo['cc'], 2, 40, 36)\n pp_cs, pp_decs = Cubic.dis_cubic(geo['pp'], 2, 40, 36)\n nn_cs, nn_decs = Cubic.dis_cubic(geo['nn'], 2, 40, 36)\n self.geo_cc.append([cc_cs, cc_decs])\n self.geo_pp.append([pp_cs, pp_decs])\n self.geo_nn.append([nn_cs, nn_decs])\n L = self.L\n cc_coefs_np = np.stack([[cc_cs[i,j].c for j in range(L)] for i in range(L)], axis=0)\n cc_knots_np = np.stack([[cc_cs[i,j].x for j in range(L)] for i in range(L)], axis=0)\n self.cs_coefs['cc'].append(torch.from_numpy(cc_coefs_np).to(device))\n self.cs_knots['cc'].append(torch.from_numpy(cc_knots_np).to(device))\n pp_coefs_np = np.stack([[pp_cs[i,j].c for j in range(L)] for i in range(L)], axis=0)\n pp_knots_np = np.stack([[pp_cs[i,j].x for j in range(L)] for i in range(L)], axis=0)\n self.cs_coefs['pp'].append(torch.from_numpy(pp_coefs_np).to(device))\n self.cs_knots['pp'].append(torch.from_numpy(pp_knots_np).to(device))\n nn_coefs_np = np.stack([[nn_cs[i,j].c for j in range(L)] for i in range(L)], axis=0)\n nn_knots_np = np.stack([[nn_cs[i,j].x for j in range(L)] for i in range(L)], axis=0)\n self.cs_coefs['nn'].append(torch.from_numpy(nn_coefs_np).to(device))\n self.cs_knots['nn'].append(torch.from_numpy(nn_knots_np).to(device))\n \n\n def _cubic_pair_energy(self, atom_map, geo_cs, geo_confimask, weight_key):\n \"\"\"General cubic-spline energy for CC/PP/NN pairs.\"\"\"\n min_dis, max_dis, bin_num = 2, 40, 36\n dev = atom_map.device\n upper_th = max_dis - ((max_dis - min_dis) / bin_num) * 0.5\n lower_th = 2.5\n total = torch.zeros((), device=dev, dtype=torch.double)\n spline_key = weight_key.split('_')[1]\n coeffs_list = self.cs_coefs[spline_key]\n knots_list = self.cs_knots[spline_key]\n for block_idx, mask_block in enumerate(geo_confimask):\n mask = (atom_map <= upper_th) & mask_block & self.fullmask & (atom_map >= lower_th)\n idx = mask.nonzero(as_tuple=True)\n if idx[0].numel() > 1:\n coef = coeffs_list[block_idx][idx]\n knots = knots_list[block_idx][idx]\n part1 = Potential.cubic_distance(atom_map[mask], coef, knots, min_dis, max_dis, bin_num).sum() * self.config[weight_key] * 0.5\n else:\n part1 = torch.zeros((), device=dev, dtype=torch.double)\n part2 = ((atom_map <= lower_th) & mask_block & self.fullmask).sum() * self.config[weight_key]\n total = total + part1 + part2\n return total\n\n # GPU-friendly torsion and angle energy helpers\n def _cubic_torsion_energy(self, atom_map, coef, x_vals, weight_key, num_bin):\n energy = Potential.cubic_torsion(atom_map, coef, x_vals, num_bin)\n return energy.sum() * self.config[weight_key]\n\n def _cubic_angle_energy(self, atom_map, coef, x_vals, weight_key, num_bin):\n energy = Potential.cubic_angle(atom_map, coef, x_vals, num_bin)\n return energy.sum() * self.config[weight_key]\n\n def compute_cc_energy(self, coor):\n atom_map = operations.pair_distance(coor[:,1], coor[:,1])\n return self._cubic_pair_energy(atom_map, self.geo_cc, self.geo_confimask_cc, 'weight_cc')\n\n def compute_pp_energy(self, coor):\n atom_map = operations.pair_distance(coor[:,0], coor[:,0])\n return self._cubic_pair_energy(atom_map, self.geo_pp, self.geo_confimask_pp, 'weight_pp')\n\n def compute_nn_energy(self, coor):\n atom_map = operations.pair_distance(coor[:,-1], coor[:,-1])\n return self._cubic_pair_energy(atom_map, self.geo_nn, self.geo_confimask_nn, 'weight_nn')\n\n def compute_pccp_energy(self, coor):\n # P-C-C-P dihedral energy on GPU\n p = coor[:, 0]\n c = coor[:, 1]\n dia = operations.dihedral(\n p[self.pccpi], c[self.pccpi], c[self.pccpj], p[self.pccpj]\n )\n return self._cubic_torsion_energy(dia, self.pccp_coe, self.pccp_x, 'weight_pccp', 36)\n\n def compute_cnnc_energy(self, coor):\n # C-N-N-C dihedral energy on GPU\n n = coor[:, -1]\n c = coor[:, 1]\n dia = operations.dihedral(\n c[self.cnnci], n[self.cnnci], n[self.cnncj], c[self.cnncj]\n )\n return self._cubic_torsion_energy(dia, self.cnnc_coe, self.cnnc_x, 'weight_cnnc', 36)\n\n def compute_pnnp_energy(self, coor):\n # P-N-N-P dihedral energy on GPU\n n = coor[:, -1]\n p = coor[:, 0]\n dia = operations.dihedral(\n p[self.pnnpi], n[self.pnnpi], n[self.pnnpj], p[self.pnnpj]\n )\n return self._cubic_torsion_energy(dia, self.pnnp_coe, self.pnnp_x, 'weight_pnnp', 36)\n\n def compute_pcc_energy(self, coor):\n # P-C-C angle energy on GPU\n p = coor[:, 1]\n c = coor[:, 2]\n ang = operations.angle(\n p[self.pcci], c[self.pcci], c[self.pccj]\n )\n return self._cubic_angle_energy(ang, self.pcc_coe, self.pcc_x, 'weight_pcc', 12)\n\n def compute_fape_energy(self,coor,ep=1e-3,epmax=20):\n energy= 0\n for tx in self.tx2ds:\n px_mean = coor[:,[1]]\n p_rot = operations.rigidFrom3Points(coor)\n p_tran = px_mean[:,0]\n pred_x2 = coor[:,None,:,:] - p_tran[None,:,None,:] # Lx Lrot N , 3\n pred_x2 = torch.einsum('ijnd,jde->ijne',pred_x2,p_rot.transpose(-1,-2)) # transpose should be equal to inverse\n errmap=torch.sqrt( ((pred_x2 - tx)**2).sum(dim=-1) + ep )\n energy = energy + torch.sum( torch.clamp(errmap,max=epmax) )\n return energy * self.config['weight_fape']\n\n def compute_bond_energy(self,coor,other_coor):\n # 3.87\n o3 = other_coor[:-1,-2]\n p = coor[1:,0]\n dis = (o3-p).norm(dim=-1)\n energy = ((dis-1.607)**2).sum()\n return energy * self.config['weight_bond']\n\n def tooth_func(self,errmap, ep = 0.05):\n return -1/(errmap/10+ep) + (1/ep)\n \n def reweight_func(self,ww):\n reweighting = torch.pow(ww,self.config['pair_weight_power'])\n reweighting[ww < self.config['pair_weight_min']] = 0\n return reweighting\n \n def compute_fape_energy_fromquat(self,x,coor,ep=1e-6,epmax=100):\n energy= 0\n p_rot,px_mean = a2b.Non2rot(x[:,:9],x.shape[0]),x[:,9:]\n pred_x2 = coor[:,None,:,:] - px_mean[None,:,None,:] # Lx Lrot N , 3\n pred_x2 = torch.einsum('ijnd,jde->ijne',pred_x2,p_rot.transpose(-1,-2)) # transpose should be equal to inverse\n\n for tx,weightplddt in zip(self.tx2ds,self.pair):\n tamplate_dist_map = torch.min( tx.norm(dim=-1), dim=2 )[0]\n errmap=torch.sqrt( ((pred_x2 - tx)**2).sum(dim=-1) + ep ) \n energy = energy + torch.sum( ( (torch.clamp(errmap,max=self.config['FAPE_max'])**self.config['pair_error_power']) * self.reweight_func(weightplddt[...,None]) )[tamplate_dist_map>self.config['pair_rest_min_dist']] )\n\n return energy * self.config['weight_fape']\n \n def compute_fape_energy_fromcoor(self,coor,ep=1e-6,epmax=100):\n energy= 0\n \n p_rot,px_mean = operations.Kabsch_rigid(self.basex,coor[:,0],coor[:,1],coor[:,2])\n pred_x2 = coor[:,None,:,:] - px_mean[None,:,None,:] # Lx Lrot N , 3\n pred_x2 = torch.einsum('ijnd,jde->ijne',pred_x2,p_rot.transpose(-1,-2)) # transpose should be equal to inverse\n \n for tx,weightplddt in zip(self.tx2ds,self.pair):\n tamplate_dist_map = torch.min( tx.norm(dim=-1), dim=2 )[0]\n errmap=torch.sqrt( ((pred_x2 - tx)**2).sum(dim=-1) + ep ) \n energy = energy + torch.sum( ( (torch.clamp(errmap,max=self.config['FAPE_max'])**self.config['pair_error_power']) * self.reweight_func(weightplddt[...,None]) )[tamplate_dist_map>self.config['pair_rest_min_dist']] )\n\n return energy * self.config['weight_fape']\n \n \n def energy(self, rama):\n coor = a2b.quat2b(self.basex, rama[:, 9:])\n other_coor = a2b.quat2b(self.otherx, rama[:, 9:])\n side_coor = a2b.quat2b(self.sidex, torch.cat([rama[:, :9], coor[:, -1]], dim=-1))\n\n E_cc = self.compute_cc_energy(coor) / len(self.geofiles) if self.config['weight_cc'] > 0 else 0\n E_pp = self.compute_pp_energy(coor) / len(self.geofiles) if self.config['weight_pp'] > 0 else 0\n E_nn = self.compute_nn_energy(coor) / len(self.geofiles) if self.config['weight_nn'] > 0 else 0\n E_pccp = self.compute_pccp_energy(coor) / len(self.geofiles) if self.config['weight_pccp'] > 0 else 0\n E_cnnc = self.compute_cnnc_energy(coor) / len(self.geofiles) if self.config['weight_cnnc'] > 0 else 0\n E_pnnp = self.compute_pnnp_energy(coor) / len(self.geofiles) if self.config['weight_pnnp'] > 0 else 0\n E_vdw = self.compute_full_clash(coor, other_coor, side_coor) if self.config['weight_vdw'] > 0 else 0\n E_fape = self.compute_fape_energy_fromquat(rama[:, 9:], coor) / len(self.geofiles) if self.config['weight_fape'] > 0 else 0\n E_bond = self.compute_bond_energy(coor, other_coor) if self.config['weight_bond'] > 0 else 0\n\n return E_vdw + E_fape + E_bond + E_pp + E_cc + E_nn + E_pccp + E_cnnc + E_pnnp\n\n\n def energy_from_coor(self, coor):\n E_cc = self.compute_cc_energy(coor) if self.config['weight_cc'] > 0 else 0\n E_pp = self.compute_pp_energy(coor) if self.config['weight_pp'] > 0 else 0\n E_nn = self.compute_nn_energy(coor) if self.config['weight_nn'] > 0 else 0\n E_fape = (self.compute_fape_energy_fromcoor(coor) / len(self.geofiles)) if self.config['weight_fape'] > 0 else 0\n print(E_fape, E_pp, E_cc, E_nn)\n return E_fape + E_pp + E_cc + E_nn \n\n def obj_func_grad_np(self,rama_):\n rama=torch.DoubleTensor(rama_)\n rama.requires_grad=True\n if rama.grad:\n rama.grad.zero_()\n f=self.energy(rama.view(self.L,21))*Scale_factor\n grad_value=autograd.grad(f,rama)[0]\n return grad_value.data.numpy().astype(np.float64)\n \n def obj_func_np(self,rama_):\n rama=torch.DoubleTensor(rama_)\n rama=rama.view(self.L,21)\n with torch.no_grad():\n f = self.energy(rama)*Scale_factor\n return f.item()\n\n def saveconfig(self,dict,confile):\n json_object = json.dumps(dict, indent = 4)\n wfile = open(confile,'w')\n wfile.write(json_object)\n wfile.close()\n \n def scoring(self):\n geoscale = self.config['geo_scale']\n self.config['weight_pp'] = geoscale * self.config['weight_pp']\n self.config['weight_cc'] = geoscale * self.config['weight_cc']\n self.config['weight_nn'] = geoscale * self.config['weight_nn']\n self.config['weight_pccp'] = geoscale * self.config['weight_pccp']\n self.config['weight_cnnc'] = geoscale * self.config['weight_cnnc']\n self.config['weight_pnnp'] = geoscale * self.config['weight_pnnp'] \n \n energy_dict = {}\n saveenergy_dict = {}\n \n with torch.no_grad():\n for retfile, tx in zip(self.geofiles, self.txs):\n one = self.energy_from_coor(tx)\n aaretfile = os.path.basename(retfile) \n energy_dict[aaretfile] = one.item()\n saveenergy_dict[retfile] = one.item()\n self.saveconfig(energy_dict, self.saveprefix)\n\n\n def foldning(self):\n minenergy=1e16\n count=0\n for tx in self.txs:\n count+=1\n \n minirama=None\n\n ilter = self.init_ret\n selected_ret = self.geofiles[ilter]\n try:\n rama=self.init_quat(ilter).data.numpy()\n self.config=readconfig(os.path.join(os.path.dirname(os.path.abspath(__file__)),'lib','vdw.json'))\n rama = fmin_l_bfgs_b(func=self.obj_func_np, x0=rama, fprime=self.obj_func_grad_np,iprint=10,maxfun=100)[0]\n rama = rama.flatten()\n except:\n rama=self.init_quat_safe(ilter).data.numpy()\n self.config=readconfig(os.path.join(os.path.dirname(os.path.abspath(__file__)),'lib','vdw.json'))\n rama = fmin_l_bfgs_b(func=self.obj_func_np, x0=rama, fprime=self.obj_func_grad_np,iprint=10,maxfun=100)[0]\n rama = rama.flatten()\n \n self.config=readconfig(self.foldconfig)\n geoscale = self.config['geo_scale']\n self.config['weight_pp'] =geoscale * self.config['weight_pp']\n self.config['weight_cc'] =geoscale * self.config['weight_cc']\n self.config['weight_nn'] =geoscale * self.config['weight_nn']\n self.config['weight_pccp'] =geoscale * self.config['weight_pccp']\n self.config['weight_cnnc'] =geoscale * self.config['weight_cnnc']\n self.config['weight_pnnp'] =geoscale * self.config['weight_pnnp']\n for i in range(3):\n line_min = lbfgs_rosetta.ArmijoLineMinimization(self.obj_func_np,self.obj_func_grad_np,True,len(rama),120)\n lbfgs_opt = lbfgs_rosetta.lbfgs(self.obj_func_np,self.obj_func_grad_np)\n rama=lbfgs_opt.run(rama,256,lbfgs_rosetta.absolute_converge_test,line_min,8000,self.obj_func_np,self.obj_func_grad_np,1e-9)\n newrama=rama+0.0\n newrama=torch.DoubleTensor(newrama) \n current_energy =self.obj_func_np(rama)\n\n if current_energy < minenergy:\n print(current_energy,minenergy)\n minenergy=current_energy\n self.outpdb(newrama,self.saveprefix+'.pdb',energystr=str(current_energy))\n\n\n def outpdb(self,rama,savefile,start=0,end=10000,energystr=''):\n coor_np=a2b.quat2b(self.basex,rama.view(self.L,21)[:,9:]).data.numpy()\n other_np=a2b.quat2b(self.otherx,rama.view(self.L,21)[:,9:]).data.numpy()\n shaped_rama=rama.view(self.L,21)\n coor = torch.FloatTensor(coor_np)\n side_coor_NP = a2b.quat2b(self.sidex,torch.cat([shaped_rama[:,:9],coor[:,-1]],dim=-1)).data.numpy()\n \n Atom_name=[' P ',\" C4'\",' N1 ']\n Other_Atom_name = [\" O5'\",\" C5'\",\" C3'\",\" O3'\",\" C1'\"]\n other_last_name = ['O',\"C\",\"C\",\"O\",\"C\"]\n\n side_atoms= [' N1 ',' C2 ',' O2 ',' N2 ',' N3 ',' N4 ',' C4 ',' O4 ',' C5 ',' C6 ',' O6 ',' N6 ',' N7 ',' N8 ',' N9 ']\n side_last_name = ['N', \"C\", \"O\", \"N\", \"N\", 'N', 'C', 'O', 'C', 'C', 'O', 'N', 'N', 'N','N']\n\n base_dict = rigid.base_table()\n \n last_name=['P','C','N']\n wstr=[f'REMARK {str(energystr)}']\n templet='%6s%5d %4s %3s %1s%4d %8.3f%8.3f%8.3f%6.2f%6.2f %2s%2s'\n count=1\n for i in range(self.L):\n if self.seq[i] in ['a','g','A','G']:\n Atom_name = [' P ',\" C4'\",' N9 ']\n\n elif self.seq[i] in ['c','u','C','U']:\n Atom_name = [' P ',\" C4'\",' N1 ']\n for j in range(coor_np.shape[1]):\n outs=('ATOM ',count,Atom_name[j],self.seq[i],'A',i+1,coor_np[i][j][0],coor_np[i][j][1],coor_np[i][j][2],0,0,last_name[j],'')\n if i>=start-1 and i < end:\n wstr.append(templet % outs)\n count+=1\n\n for j in range(other_np.shape[1]):\n outs=('ATOM ',count,Other_Atom_name[j],self.seq[i],'A',i+1,other_np[i][j][0],other_np[i][j][1],other_np[i][j][2],0,0,other_last_name[j],'')\n if i>=start-1 and i < end:\n wstr.append(templet % outs)\n count+=1\n \n wstr='\\n'.join(wstr)\n wfile=open(savefile,'w')\n wfile.write(wstr)\n wfile.close()\n \n \n def outpdb_coor(self,coor_np,savefile,start=0,end=1000,energystr=''):\n Atom_name=[' P ',\" C4'\",' N1 ']\n last_name=['P','C','N']\n wstr=[f'REMARK {str(energystr)}']\n templet='%6s%5d %4s %3s %1s%4d %8.3f%8.3f%8.3f%6.2f%6.2f %2s%2s'\n count=1\n for i in range(self.L):\n if self.seq[i] in ['a','g','A','G']:\n Atom_name = [' P ',\" C4'\",' N9 ']\n\n elif self.seq[i] in ['c','u','C','U']:\n Atom_name = [' P ',\" C4'\",' N1 ']\n \n for j in range(coor_np.shape[1]):\n outs=('ATOM ',count,Atom_name[j],self.seq[i],'A',i+1,coor_np[i][j][0],coor_np[i][j][1],coor_np[i][j][2],0,0,last_name[j],'')\n if i>=start-1 and i < end:\n wstr.append(templet % outs)\n count+=1\n \n wstr='\\n'.join(wstr)\n wfile=open(savefile,'w')\n wfile.write(wstr)\n wfile.close()\n\n\nif __name__ == '__main__': \n\n fastafile = sys.argv[1]\n foldconfig = sys.argv[2]\n save_prefix = sys.argv[3]\n retfiles = sys.argv[4:]\n\n save_parent_dir = os.path.dirname(save_prefix)\n if not os.path.isdir(save_parent_dir):\n os.makedirs(save_parent_dir)\n\n retfiles.sort()\n print(retfiles)\n\n stru = Structure(fastafile, retfiles, foldconfig, save_prefix) \n stru.scoring()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:19:33.371214Z","iopub.execute_input":"2025-05-29T13:19:33.371545Z","iopub.status.idle":"2025-05-29T13:19:33.399976Z","shell.execute_reply.started":"2025-05-29T13:19:33.371518Z","shell.execute_reply":"2025-05-29T13:19:33.398672Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %%writefile /kaggle/working/DRfold2/cfg_97/test_modeldir.py\n# import random\n# random.seed(0)\n# import numpy as np\n# np.random.seed(0)\n# import os,sys,re,random\n# from numpy import select\n# import torch\n# torch.manual_seed(0)\n# torch.backends.cudnn.deterministic = True\n# torch.backends.cudnn.benchmark = False\n# expdir=os.path.dirname(os.path.abspath(__file__))\n\n\n# import torch.optim as opt\n# from torch.nn import functional as F\n# import data,util\n# import EvoMSA2XYZ,basic\n# import math\n# import pickle\n# Batch_size=3\n# Num_cycle=3\n# TEST_STEP=1000\n# VISION_STEP=50\n# device = sys.argv[1]\n\n\n# expdir=os.path.dirname(os.path.abspath(__file__))\n# expround=expdir.split('_')[-1]\n# model_path=os.path.join(expdir,'others','models')\n\n# testdir=os.path.join(expdir,'others','preds')\n# basenpy_standard= np.load( os.path.join(os.path.dirname(os.path.abspath(__file__)),'base.npy' ) )\n\n# def data_collect(pdb_seq):\n# aa_type = data.parse_seq(pdb_seq)\n# base = data.Get_base(pdb_seq,basenpy_standard)\n# seq_idx = np.arange(len(pdb_seq)) + 1\n\n# msa=aa_type[None,:]\n# msa=torch.from_numpy(msa).to(device)\n# msa=torch.cat([msa,msa],0)\n# msa=F.one_hot(msa.long(),6).float()\n\n# base_x = torch.from_numpy(base).float().to(device)\n# seq_idx = torch.from_numpy(seq_idx).long().to(device)\n# return msa,base_x,seq_idx\n# predxs,plddts = model.pred(msa,seq_idx,ss,base_x,sample_1['alpha_0'])\n\n\n\n# def classifier(infasta,out_prefix,model_dir):\n# with torch.no_grad():\n# lines = open(infasta).readlines()[1:]\n# seqs = [aline.strip() for aline in lines]\n# seq = ''.join(seqs)\n# msa,base_x,seq_idx = data_collect(seq)\n \n# msa_dim=6+1\n# m_dim,s_dim,z_dim = 64,64,64\n# N_ensemble, N_cycle = 6, 12\n# model=EvoMSA2XYZ.MSA2XYZ(msa_dim-1,msa_dim,N_ensemble,N_cycle,m_dim,s_dim,z_dim)\n# model.to(device)\n# model.eval()\n# models = os.listdir( model_dir )\n# models = [amodel for amodel in models if 'model' in amodel and 'opt' not in amodel]\n\n# models.sort()\n\n# for amodel in models:\n# saved_model = os.path.join(model_dir, amodel)\n# model.load_state_dict(torch.load(saved_model, map_location='cpu'), strict=True)\n# ret = model.pred(msa, seq_idx, None, base_x, np.array(list(seq)))\n\n# util.outpdb(ret['coor'], seq_idx, seq, out_prefix+f'{amodel}.pdb')\n# pickle.dump(ret, open(out_prefix+f'{amodel}.ret', 'wb'))\n\n\n# if __name__ == '__main__':\n# infasta, out_prefix, model_dir = sys.argv[2], sys.argv[3], sys.argv[4]\n# classifier(infasta, out_prefix, model_dir)","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-05-29T13:19:33.401326Z","iopub.execute_input":"2025-05-29T13:19:33.401604Z","iopub.status.idle":"2025-05-29T13:19:33.426593Z","shell.execute_reply.started":"2025-05-29T13:19:33.401573Z","shell.execute_reply":"2025-05-29T13:19:33.425499Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile /kaggle/working/DRfold2/PotentialFold/Cubic.py\nimport numpy as np \nfrom scipy.interpolate import CubicSpline,UnivariateSpline\nimport os\nfrom torch.autograd import Function\nimport torch\nimport math\n\ndef fit_dis_cubic(dis_matrix,min_dis,max_dis,num_bin):\n # convert torch Tensor on GPU to numpy array for SciPy\n if isinstance(dis_matrix, torch.Tensor):\n dis_matrix = dis_matrix.detach().cpu().numpy()\n dis_region=np.zeros(num_bin)\n for i in range(num_bin):\n dis_region[i]=min_dis+(i+0.5)*(max_dis-min_dis)*1.0/num_bin\n L=dis_matrix.shape[0]\n csnp=[]\n decsnp=[]\n for i in range(L):\n css=[]\n decss=[]\n for j in range(L):\n y=-np.log( (dis_matrix[i,j,1:-1]+1e-8) / (dis_matrix[i,j,[-2]]+1e-8) )\n x=dis_region\n x[0]=-0.0001\n y[0]= max(10,y[1]+4)\n cs= CubicSpline(x,y)\n decs=cs.derivative()\n css.append(cs)\n decss.append(decs)\n csnp.append(css)\n decsnp.append(decss)\n return np.array(csnp),np.array(decsnp)\n\ndef dis_cubic(out,min_dis,max_dis,num_bin):\n print('fitting cubic distance')\n cs,decs=fit_dis_cubic(out,min_dis,max_dis,num_bin)\n return cs,decs\n\n\n\ndef cubic_matrix_torsion(dis_matrix,min_dis,max_dis,num_bin):\n dis_region=np.zeros(num_bin)\n bin_size=(max_dis-min_dis)/num_bin\n for i in range(num_bin):\n dis_region[i]=min_dis+(i+0.5)*(max_dis-min_dis)*1.0/num_bin\n L=dis_matrix.shape[0]\n csnp=[]\n decsnp=[]\n for i in range(L):\n css=[]\n decss=[]\n for j in range(L):\n y=-np.log( dis_matrix[i,j,:-1]+1e-8 )\n x=dis_region\n x=np.append(x,x[-1]+bin_size)\n y=np.append(y,y[0])\n cs= CubicSpline(x,y,bc_type='periodic')\n decs=cs.derivative()\n css.append(cs)\n decss.append(decs)\n csnp.append(css)\n decsnp.append(decss)\n return np.array(csnp),np.array(decsnp)\ndef torsion_cubic(out,min_dis,max_dis,num_bin):\n print('fitting cubic')\n cs,decs=cubic_matrix_torsion(out,min_dis,max_dis,num_bin)\n return cs,decs\n\ndef cubic_matrix_angle(dis_matrix,min_dis,max_dis,num_bin): # 0 - np.pi 12\n dis_region=np.zeros(num_bin)\n bin_size=(max_dis-min_dis)/num_bin\n for i in range(num_bin):\n dis_region[i]=min_dis+(i+0.5)*(max_dis-min_dis)*1.0/num_bin\n L=dis_matrix.shape[0]\n csnp=[]\n decsnp=[]\n for i in range(L):\n css=[]\n decss=[]\n for j in range(L):\n y=-np.log( dis_matrix[i,j,:-1]+1e-8 )\n x=dis_region\n\n x=np.concatenate([[x[0]-bin_size*3,x[0]-bin_size*2,x[0]-bin_size], x,[x[-1]+bin_size,x[-1]+bin_size*2,x[-1]+bin_size*3] ])\n y=np.concatenate([ [y[2],y[1],y[0]],y,[y[-1],y[-2],y[-3]] ])\n\n cs= CubicSpline(x,y)\n decs=cs.derivative()\n\n css.append(cs)\n decss.append(decs)\n csnp.append(css)\n decsnp.append(decss)\n\n return np.array(csnp),np.array(decsnp)\ndef angle_cubic(out,min_dis,max_dis,num_bin):\n\n print('fitting angle cubic')\n cs,decs=cubic_matrix_angle(out,min_dis,max_dis,num_bin)\n\n return cs,decs","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:19:33.427470Z","iopub.execute_input":"2025-05-29T13:19:33.427729Z","iopub.status.idle":"2025-05-29T13:19:33.453473Z","shell.execute_reply.started":"2025-05-29T13:19:33.427699Z","shell.execute_reply":"2025-05-29T13:19:33.452083Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Hybrid Method Follows \n---","metadata":{}},{"cell_type":"code","source":"# Define an improved function to run DRfold2 that captures output\ndef predict_rna_structures_drfold2(sequence, target_id):\n \"\"\"Use DRfold2 to predict RNA structures with proper output capture\"\"\"\n import subprocess\n from subprocess import PIPE, STDOUT\n \n # Create FASTA file for this sequence\n fasta_path = os.path.join(fasta_dir, f\"{target_id}.fasta\")\n with open(fasta_path, \"w\") as f:\n f.write(f\">{target_id}\\n{sequence}\\n\")\n \n # Run DRfold2 with proper output capture\n output_dir = os.path.join(predictions_dir, target_id)\n cmd = f\"python /kaggle/working/DRfold2/DRfold_infer.py {fasta_path} {output_dir} 1\"\n \n print(f\"Running command: {cmd}\")\n process = subprocess.Popen(\n cmd, \n shell=True, \n stdout=PIPE, \n stderr=STDOUT,\n universal_newlines=True,\n bufsize=1\n )\n \n # Print output in real-time\n for line in iter(process.stdout.readline, ''):\n line = line.strip()\n if line:\n print(line)\n \n # Get return code and check success\n return_code = process.wait()\n if return_code != 0:\n print(f\"DRfold2 failed with return code {return_code}\")\n return None\n \n # Clean up FASTA file to save space\n os.remove(fasta_path)\n \n # Extract coordinates\n relax_dir = os.path.join(output_dir, \"relax\")\n if not os.path.isdir(relax_dir):\n print(f\"Warning: No relax directory found for {target_id}\")\n relax_dir = output_dir\n \n # Get up to 5 PDB files\n pdb_files = sorted([f for f in os.listdir(relax_dir) if f.endswith(\".pdb\")])[:5]\n \n if not pdb_files:\n print(f\"Warning: No PDB files found for {target_id}\")\n # Return None to indicate failure\n return None\n \n # Parse PDB files to extract C1' coordinates\n predictions = []\n for pdb_file in pdb_files:\n file_path = os.path.join(relax_dir, pdb_file)\n \n # Read PDB file\n coords = []\n with open(file_path, \"r\") as f:\n residue_map = {}\n for line in f:\n if line.startswith(\"ATOM\") and \" C1' \" in line:\n parts = line.split()\n resid = int(parts[5]) # Residue ID as integer\n x, y, z = float(parts[6]), float(parts[7]), float(parts[8])\n residue_map[resid] = (x, y, z)\n \n # Ensure we have coordinates for all residues\n for j in range(1, len(sequence) + 1):\n if j in residue_map:\n coords.append(residue_map[j])\n else:\n # If residue not found, use zeros\n print(f\"Warning: Residue {j} not found in {pdb_file} for {target_id}\")\n coords.append((0.0, 0.0, 0.0))\n \n predictions.append(coords)\n \n # Clean up PDB files to save space\n if is_submission_mode:\n shutil.rmtree(output_dir)\n \n # If we have fewer than 5 predictions, duplicate the last one\n while len(predictions) < 5:\n predictions.append(predictions[-1] if predictions else [(0.0, 0.0, 0.0) for _ in range(len(sequence))])\n \n return predictions[:5] # Return exactly 5 predictions","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:19:33.455062Z","iopub.execute_input":"2025-05-29T13:19:33.455370Z","iopub.status.idle":"2025-05-29T13:19:33.477344Z","shell.execute_reply.started":"2025-05-29T13:19:33.455346Z","shell.execute_reply":"2025-05-29T13:19:33.476246Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Vectorized version of process_labels function\ndef process_labels_vectorized(labels_df):\n # Extract target_id from ID column (remove last part after underscore)\n labels_df = labels_df.copy()\n labels_df['target_id'] = labels_df['ID'].str.rsplit('_', n=1).str[0]\n \n # Sort by target_id and resid for proper ordering\n labels_df = labels_df.sort_values(['target_id', 'resid'])\n \n # Group by target_id and convert coordinates to arrays\n coords_dict = {}\n for target_id, group in labels_df.groupby('target_id'):\n # Extract coordinates as numpy array in one operation\n coords_dict[target_id] = group[['x_1', 'y_1', 'z_1']].values\n \n return coords_dict\n\ndef find_similar_sequences(query_seq, train_seqs_df, train_coords_dict, top_n=5):\n similar_seqs = []\n query_seq_obj = Seq(query_seq)\n\n for _, row in train_seqs_df.iterrows():\n target_id = row['target_id']\n train_seq = row['sequence']\n\n # Skip if coordinates not available\n if target_id not in train_coords_dict:\n continue\n\n # Skip if sequence is too different in length (more than 40% difference)\n if abs(len(train_seq) - len(query_seq)) / max(len(train_seq), len(query_seq)) > 0.4:\n continue\n\n # Perform sequence alignment\n alignments = pairwise2.align.globalms(query_seq_obj, train_seq, 2.9, -1, -10, -0.5, one_alignment_only=True)\n\n if alignments:\n alignment = alignments[0]\n similarity_score = alignment.score / (2 * min(len(query_seq), len(train_seq)))\n similar_seqs.append((target_id, train_seq, similarity_score, train_coords_dict[target_id]))\n\n # Sort by similarity score (higher is better) and return top N\n similar_seqs.sort(key=lambda x: x[2], reverse=True)\n return similar_seqs[:top_n]\n\n\ndef adaptive_rna_constraints(coordinates, sequence, confidence=1.0):\n # Make a copy of coordinates to refine\n refined_coords = coordinates.copy()\n n_residues = len(sequence)\n \n # Calculate constraint strength (inverse of confidence)\n # High confidence templates receive gentler constraints\n constraint_strength = 0.8 * (1.0 - min(confidence, 0.8))\n \n # 1. Sequential distance constraints (consecutive nucleotides)\n # More flexible distance range (statistical distribution from PDB)\n seq_min_dist = 5.5 # Minimum sequential distance\n seq_max_dist = 6.5 # Maximum sequential distance\n \n for i in range(n_residues - 1):\n current_pos = refined_coords[i]\n next_pos = refined_coords[i+1]\n \n # Calculate current distance\n current_dist = np.linalg.norm(next_pos - current_pos)\n \n # Only adjust if significantly outside expected range\n if current_dist < seq_min_dist or current_dist > seq_max_dist:\n # Calculate target distance (midpoint of range)\n target_dist = (seq_min_dist + seq_max_dist) / 2\n \n # Get direction vector\n direction = next_pos - current_pos\n direction = direction / (np.linalg.norm(direction) + 1e-10)\n \n # Apply partial adjustment based on constraint strength\n adjustment = (target_dist - current_dist) * constraint_strength\n \n # Only adjust the next position to preserve the overall fold\n refined_coords[i+1] = current_pos + direction * (current_dist + adjustment)\n \n # 2. Steric clash prevention (more conservative)\n min_allowed_distance = 3.8 # Minimum distance between non-consecutive C1' atoms\n \n # Calculate all pairwise distances\n dist_matrix = distance_matrix(refined_coords, refined_coords)\n \n # Find severe clashes (atoms too close)\n severe_clashes = np.where((dist_matrix < min_allowed_distance) & (dist_matrix > 0))\n \n # Fix severe clashes\n for idx in range(len(severe_clashes[0])):\n i, j = severe_clashes[0][idx], severe_clashes[1][idx]\n \n # Skip consecutive nucleotides and previously processed pairs\n if abs(i - j) <= 1 or i >= j:\n continue\n \n # Get current positions and distance\n pos_i = refined_coords[i]\n pos_j = refined_coords[j]\n current_dist = dist_matrix[i, j]\n \n # Calculate necessary adjustment but scale by constraint strength\n direction = pos_j - pos_i\n direction = direction / (np.linalg.norm(direction) + 1e-10)\n \n # Calculate partial adjustment\n adjustment = (min_allowed_distance - current_dist) * constraint_strength\n \n # Move points apart\n refined_coords[i] = pos_i - direction * (adjustment / 2)\n refined_coords[j] = pos_j + direction * (adjustment / 2)\n \n # 3. Very light base-pair constraining (if confidence is low)\n if constraint_strength > 0.3: # Only apply if template confidence is low\n # Simple Watson-Crick base pairs\n pairs = {'A': 'U', 'U': 'A', 'G': 'C', 'C': 'G'}\n \n # Scan for potential base pairs\n for i in range(n_residues):\n base_i = sequence[i]\n complement = pairs.get(base_i)\n \n if not complement:\n continue\n \n # Look for complementary bases within a reasonable range\n for j in range(i + 3, min(i + 20, n_residues)):\n if sequence[j] == complement:\n # Calculate current distance\n current_dist = np.linalg.norm(refined_coords[i] - refined_coords[j])\n \n # Only consider if distance suggests potential pairing\n if 8.0 < current_dist < 14.0:\n # Target 10.5Å as generic base-pair C1'-C1' distance\n target_dist = 10.5\n \n # Calculate very gentle adjustment (scaled by constraint_strength)\n adjustment = (target_dist - current_dist) * (constraint_strength * 0.3)\n \n # Get direction vector\n direction = refined_coords[j] - refined_coords[i]\n direction = direction / (np.linalg.norm(direction) + 1e-10)\n \n # Apply very gentle adjustment to both positions\n refined_coords[i] = refined_coords[i] - direction * (adjustment / 2)\n refined_coords[j] = refined_coords[j] + direction * (adjustment / 2)\n \n # Only consider one potential pair per base (closest match)\n break\n \n return refined_coords\n\ndef adapt_template_to_query(query_seq, template_seq, template_coords, alignment=None):\n if alignment is None:\n from Bio.Seq import Seq\n from Bio import pairwise2\n \n query_seq_obj = Seq(query_seq)\n template_seq_obj = Seq(template_seq)\n alignments = pairwise2.align.globalms(query_seq_obj, template_seq_obj, 2.9, -1, -10, -0.5, one_alignment_only=True)\n \n if not alignments:\n return generate_improved_rna_structure(query_seq)\n \n alignment = alignments[0]\n \n aligned_query = alignment.seqA\n aligned_template = alignment.seqB\n \n query_coords = np.zeros((len(query_seq), 3))\n query_coords.fill(np.nan)\n \n # Map template coordinates to query\n query_idx = 0\n template_idx = 0\n \n for i in range(len(aligned_query)):\n query_char = aligned_query[i]\n template_char = aligned_template[i]\n \n if query_char != '-' and template_char != '-':\n if template_idx < len(template_coords):\n query_coords[query_idx] = template_coords[template_idx]\n template_idx += 1\n query_idx += 1\n elif query_char != '-' and template_char == '-':\n query_idx += 1\n elif query_char == '-' and template_char != '-':\n template_idx += 1\n \n # IMPROVED GAP FILLING - maintains RNA backbone geometry\n backbone_distance = 5.9 # Typical C1'-C1' distance\n \n # Fill gaps by maintaining realistic backbone connectivity\n for i in range(len(query_coords)):\n if np.isnan(query_coords[i, 0]):\n # Find nearest valid neighbors\n prev_valid = next_valid = None\n \n for j in range(i-1, -1, -1):\n if not np.isnan(query_coords[j, 0]):\n prev_valid = j\n break\n \n for j in range(i+1, len(query_coords)):\n if not np.isnan(query_coords[j, 0]):\n next_valid = j\n break\n \n if prev_valid is not None and next_valid is not None:\n # Interpolate along realistic RNA backbone path\n gap_size = next_valid - prev_valid\n total_distance = np.linalg.norm(query_coords[next_valid] - query_coords[prev_valid])\n expected_distance = gap_size * backbone_distance\n \n # If gap is compressed, extend it realistically\n if total_distance < expected_distance * 0.7:\n direction = query_coords[next_valid] - query_coords[prev_valid]\n direction = direction / (np.linalg.norm(direction) + 1e-10)\n \n # Place intermediate points along extended path\n for k, idx in enumerate(range(prev_valid + 1, next_valid)):\n progress = (k + 1) / gap_size\n base_pos = query_coords[prev_valid] + direction * expected_distance * progress\n \n # Add slight curvature for realism\n perpendicular = np.cross(direction, [0, 0, 1])\n if np.linalg.norm(perpendicular) < 1e-6:\n perpendicular = np.cross(direction, [1, 0, 0])\n perpendicular = perpendicular / (np.linalg.norm(perpendicular) + 1e-10)\n \n curve_amplitude = 2.0 * np.sin(progress * np.pi)\n query_coords[idx] = base_pos + perpendicular * curve_amplitude\n else:\n # Linear interpolation for normal gaps\n for k, idx in enumerate(range(prev_valid + 1, next_valid)):\n weight = (k + 1) / gap_size\n query_coords[idx] = (1 - weight) * query_coords[prev_valid] + weight * query_coords[next_valid]\n \n elif prev_valid is not None:\n # Extend from previous position\n if prev_valid > 0 and not np.isnan(query_coords[prev_valid-1, 0]):\n direction = query_coords[prev_valid] - query_coords[prev_valid-1]\n direction = direction / (np.linalg.norm(direction) + 1e-10)\n else:\n direction = np.array([1.0, 0.0, 0.0])\n \n steps_needed = i - prev_valid\n for step in range(1, steps_needed + 1):\n pos_idx = prev_valid + step\n if pos_idx < len(query_coords):\n query_coords[pos_idx] = query_coords[prev_valid] + direction * backbone_distance * step\n \n elif next_valid is not None:\n # Work backwards from next position\n direction = np.array([-1.0, 0.0, 0.0]) # Default backward direction\n steps_needed = next_valid - i\n for step in range(steps_needed, 0, -1):\n pos_idx = next_valid - step\n if pos_idx >= 0:\n query_coords[pos_idx] = query_coords[next_valid] - direction * backbone_distance * step\n \n # Final cleanup\n query_coords = np.nan_to_num(query_coords)\n return query_coords\n\n\n\n\ndef generate_improved_rna_structure(sequence):\n \"\"\"\n Generate a more realistic RNA structure fallback based on sequence patterns\n and basic RNA structure principles.\n \n Args:\n sequence: RNA sequence string\n \n Returns:\n Array of 3D coordinates\n \"\"\"\n n_residues = len(sequence)\n coordinates = np.zeros((n_residues, 3))\n \n # Analyze sequence to predict structural elements\n # Look for complementary regions that could form base pairs\n potential_stems = identify_potential_stems(sequence)\n \n # Default parameters\n radius_helix = 10.0\n radius_loop = 15.0\n rise_per_residue_helix = 2.5\n rise_per_residue_loop = 1.5\n angle_per_residue_helix = 0.6\n angle_per_residue_loop = 0.3\n \n # Assign structural classifications\n structure_types = assign_structure_types(sequence, potential_stems)\n \n # Generate coordinates based on predicted structure\n current_pos = np.array([0.0, 0.0, 0.0])\n current_direction = np.array([0.0, 0.0, 1.0])\n current_angle = 0.0\n \n for i in range(n_residues):\n if structure_types[i] == 'stem':\n # Part of a helical stem\n current_angle += angle_per_residue_helix\n coordinates[i] = [\n radius_helix * np.cos(current_angle), \n radius_helix * np.sin(current_angle), \n current_pos[2] + rise_per_residue_helix\n ]\n current_pos = coordinates[i]\n elif structure_types[i] == 'loop':\n # Part of a loop\n current_angle += angle_per_residue_loop\n z_shift = rise_per_residue_loop * np.sin(current_angle * 0.5)\n coordinates[i] = [\n radius_loop * np.cos(current_angle), \n radius_loop * np.sin(current_angle), \n current_pos[2] + z_shift\n ]\n current_pos = coordinates[i]\n else:\n # Single-stranded region\n # Add some randomness to make it look more realistic\n jitter = np.random.normal(0, 1, 3) * 2.0\n coordinates[i] = current_pos + jitter\n current_pos = coordinates[i]\n \n return coordinates\n\ndef identify_potential_stems(sequence):\n \"\"\"\n Identify potential stem regions by looking for self-complementary segments.\n \n Args:\n sequence: RNA sequence string\n \n Returns:\n List of tuples (start1, end1, start2, end2) representing potentially paired regions\n \"\"\"\n complementary_bases = {'A': 'U', 'U': 'A', 'G': 'C', 'C': 'G'}\n min_stem_length = 3\n potential_stems = []\n \n # Simple stem identification\n for i in range(len(sequence) - min_stem_length):\n for j in range(i + min_stem_length + 3, len(sequence) - min_stem_length + 1):\n # Check if regions could form a stem\n potential_stem_len = min(min_stem_length, len(sequence) - j)\n is_stem = True\n \n for k in range(potential_stem_len):\n if sequence[i+k] not in complementary_bases or \\\n complementary_bases[sequence[i+k]] != sequence[j+potential_stem_len-k-1]:\n is_stem = False\n break\n \n if is_stem:\n potential_stems.append((i, i+potential_stem_len-1, j, j+potential_stem_len-1))\n \n return potential_stems\n\ndef assign_structure_types(sequence, potential_stems):\n \"\"\"\n Assign each nucleotide to a structural element type.\n \n Args:\n sequence: RNA sequence string\n potential_stems: List of tuples representing stem regions\n \n Returns:\n List of structure types ('stem', 'loop', 'single')\n \"\"\"\n structure_types = ['single'] * len(sequence)\n \n # Mark stem regions\n for stem in potential_stems:\n start1, end1, start2, end2 = stem\n for i in range(end1 - start1 + 1):\n structure_types[start1 + i] = 'stem'\n structure_types[end2 - i] = 'stem'\n \n # Mark loop regions (regions between paired regions)\n for i in range(len(potential_stems) - 1):\n _, end1, start2, _ = potential_stems[i]\n next_start1, _, _, _ = potential_stems[i+1]\n \n if next_start1 > end1 + 1 and start2 > next_start1:\n for j in range(end1 + 1, next_start1):\n structure_types[j] = 'loop'\n \n return structure_types\n\n\n\n# Function to create a more realistic RNA structure when no good templates are found\ndef generate_rna_structure(sequence, seed=None):\n if seed is not None:\n np.random.seed(seed)\n random.seed(seed)\n \n n_residues = len(sequence)\n coordinates = np.zeros((n_residues, 3))\n \n # Initialize the first few residues in a helix\n for i in range(min(3, n_residues)):\n angle = i * 0.6\n coordinates[i] = [10.0 * np.cos(angle), 10.0 * np.sin(angle), i * 2.5]\n \n # Add more complex folding patterns\n current_direction = np.array([0.0, 0.0, 1.0]) # Start moving along z-axis\n \n # Define base-pairing tendencies (G-C and A-U pairs)\n for i in range(3, n_residues):\n # Check for potential base-pairing in the sequence\n has_pair = False\n pair_idx = -1\n \n # Simple detection of complementary bases (G-C, A-U)\n complementary = {'G': 'C', 'C': 'G', 'A': 'U', 'U': 'A'}\n current_base = sequence[i]\n \n # Look for potential base-pairing within a window before the current position\n window_size = min(i, 15) # Look back up to 15 bases\n for j in range(i-window_size, i):\n if j >= 0 and sequence[j] == complementary.get(current_base, 'X'):\n # Found a potential pair\n has_pair = True\n pair_idx = j\n break\n \n if has_pair and i - pair_idx <= 10 and random.random() < 0.7:\n # Try to create a base-pair by positioning this nucleotide near its pair\n pair_pos = coordinates[pair_idx]\n \n # Create a position that's roughly opposite to the pair\n random_offset = np.random.normal(0, 1, 3) * 2.0\n base_pair_distance = 10.0 + random.uniform(-1.0, 1.0)\n \n # Calculate a vector from base-pair toward center of structure\n center = np.mean(coordinates[:i], axis=0)\n direction = center - pair_pos\n direction = direction / (np.linalg.norm(direction) + 1e-10)\n \n # Position new nucleotide in the general direction of the \"center\"\n coordinates[i] = pair_pos + direction * base_pair_distance + random_offset\n \n # Update direction for next nucleotide\n current_direction = np.random.normal(0, 0.3, 3)\n current_direction = current_direction / (np.linalg.norm(current_direction) + 1e-10)\n \n else:\n # No base-pairing detected, continue with the current fold direction\n # Randomly rotate current direction to simulate RNA flexibility\n if random.random() < 0.3:\n # More significant direction change\n angle = random.uniform(0.2, 0.6)\n axis = np.random.normal(0, 1, 3)\n axis = axis / (np.linalg.norm(axis) + 1e-10)\n rotation = R.from_rotvec(angle * axis)\n current_direction = rotation.apply(current_direction)\n else:\n # Small random changes in direction\n current_direction += np.random.normal(0, 0.15, 3)\n current_direction = current_direction / (np.linalg.norm(current_direction) + 1e-10)\n \n # Distance between consecutive nucleotides (3.5-4.5Å is typical)\n step_size = random.uniform(3.5, 4.5)\n \n # Update position\n coordinates[i] = coordinates[i-1] + step_size * current_direction\n \n return coordinates\n\n\ndef predict_rna_structures(sequence, target_id, train_seqs_df, train_coords_dict, n_predictions=5):\n predictions = []\n \n # Find similar sequences in the training data\n similar_seqs = find_similar_sequences(sequence, train_seqs_df, train_coords_dict, top_n=n_predictions)\n \n # If we found any similar sequences, use them as templates\n if similar_seqs:\n for i, (template_id, template_seq, similarity_score, template_coords) in enumerate(similar_seqs):\n # Adapt template coordinates to the query sequence\n adapted_coords = adapt_template_to_query(sequence, template_seq, template_coords)\n \n if adapted_coords is not None:\n # Apply adaptive constraints based on template similarity\n # For high similarity templates, apply very gentle constraints\n refined_coords = adaptive_rna_constraints(adapted_coords, sequence, confidence=similarity_score)\n \n # Add some randomness (less for better templates)\n random_scale = max(0.05, 0.8 - similarity_score) # Reduced randomness\n randomized_coords = refined_coords.copy()\n randomized_coords += np.random.normal(0, random_scale, randomized_coords.shape)\n \n predictions.append(randomized_coords)\n \n if len(predictions) >= n_predictions:\n break\n \n # If we don't have enough predictions from templates, generate de novo structures\n while len(predictions) < n_predictions:\n seed_value = hash(target_id) % 10000 + len(predictions) * 1000\n de_novo_coords = generate_rna_structure(sequence, seed=seed_value)\n \n # Apply stronger constraints to de novo structures (lower confidence)\n refined_de_novo = adaptive_rna_constraints(de_novo_coords, sequence, confidence=0.2)\n \n predictions.append(refined_de_novo)\n \n return predictions[:n_predictions]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:20:13.384343Z","iopub.execute_input":"2025-05-29T13:20:13.384681Z","iopub.status.idle":"2025-05-29T13:20:13.423874Z","shell.execute_reply.started":"2025-05-29T13:20:13.384643Z","shell.execute_reply":"2025-05-29T13:20:13.422421Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Initialize counters and range settings\nif is_submission_mode:\n DRFOLD_START_IDX = 14\n DRFOLD_END_IDX = len(test_sequences) - 1\nelse:\n DRFOLD_START_IDX = 0\n DRFOLD_END_IDX = 0\n\ndrfold_processed = 0\ntemplate_processed = 0\n\ntrain_coords_dict = process_labels_vectorized(train_labels_final)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Sort test sequences by length to process shorter ones with DRfold2\ntest_sequences = test_sequences.sort_values(by=['sequence'], key=lambda x: x.str.len())\n\n# List to store all prediction records\nall_predictions = []\n\n# Set up time tracking\nstart_time = time.time()\ntotal_targets = len(test_sequences)\n\n# For each sequence in the test set\nfor idx, row in test_sequences.iterrows():\n target_id = row['target_id']\n sequence = row['sequence']\n \n # Progress tracking\n elapsed = time.time() - start_time\n targets_processed = idx\n if targets_processed > 0:\n avg_time_per_target = elapsed / targets_processed\n est_time_remaining = avg_time_per_target * (total_targets - targets_processed)\n time_left = DRFOLD_TIME_LIMIT - (time.time() - start_time_global)\n print(f\"Processing target {targets_processed+1}/{total_targets}: {target_id} ({len(sequence)} nt), \"\n f\"elapsed: {elapsed:.1f}s, est. remaining: {est_time_remaining:.1f}s, time left: {time_left:.1f}s\")\n \n # Check if we should use DRfold2 or template-based approach\n use_drfold = (DRFOLD_START_IDX <= idx <= DRFOLD_END_IDX and \n (time.time() - start_time_global) < DRFOLD_TIME_LIMIT)\n \n # Generate 5 different structure predictions\n if use_drfold:\n print(f\"Using DRfold2 for target {target_id} (index {idx})\")\n predictions = predict_rna_structures_drfold2(sequence, target_id)\n \n # If DRfold2 fails, fall back to template approach\n if predictions is None:\n print(f\"DRfold2 failed for {target_id}, falling back to template approach\")\n predictions = predict_rna_structures(sequence, target_id, train_seqs_final, train_coords_dict)\n template_processed += 1\n else:\n drfold_processed += 1\n else:\n if idx > DRFOLD_END_IDX:\n reason = \"index out of DRfold range\"\n elif idx < DRFOLD_START_IDX:\n reason = \"index before DRfold start range\"\n else:\n reason = \"time limit reached\"\n print(f\"Using template approach for {target_id} ({reason})\")\n predictions = predict_rna_structures(sequence, target_id, train_seqs_final, train_coords_dict)\n template_processed += 1\n \n # For each residue in the sequence\n for j in range(len(sequence)):\n pred_row = {\n 'ID': f\"{target_id}_{j+1}\",\n 'resname': sequence[j],\n 'resid': j + 1\n }\n \n # Add coordinates from all 5 predictions\n for i in range(5):\n pred_row[f'x_{i+1}'] = predictions[i][j][0]\n pred_row[f'y_{i+1}'] = predictions[i][j][1]\n pred_row[f'z_{i+1}'] = predictions[i][j][2]\n \n all_predictions.append(pred_row)\n \n # Free up memory\n if torch.cuda.is_available():\n torch.cuda.empty_cache()\n\n# Create DataFrame with predictions\nsubmission_df = pd.DataFrame(all_predictions)\n\n# Ensure the submission file has the correct format\ncolumn_order = ['ID', 'resname', 'resid']\nfor i in range(1, 6):\n for coord in ['x', 'y', 'z']:\n column_order.append(f'{coord}_{i}')\n \nsubmission_df = submission_df[column_order]\n\n# Clean the working directory before saving\nprint(\"Cleaning working directory...\")\nfor item in os.listdir(\"/kaggle/working/\"):\n item_path = os.path.join(\"/kaggle/working/\", item)\n if os.path.isfile(item_path) and item != \"submission.csv\":\n os.remove(item_path)\n elif os.path.isdir(item_path) and item != \"predictions\" and item != \"fasta_files\" and item != \"DRfold2\":\n shutil.rmtree(item_path)\n\n# Save the submission\nsubmission_df.to_csv('/kaggle/working/submission.csv', index=False)\nprint(f\"Submission file saved to /kaggle/working/submission.csv\")\nprint(f\"Generated predictions for {len(test_sequences)} RNA sequences\")\nprint(f\"Used DRfold2 for {drfold_processed} targets and template approach for {template_processed} targets\")\nprint(f\"Total runtime: {time.time() - start_time_global:.1f} seconds\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:20:34.395425Z","iopub.status.idle":"2025-05-29T13:20:34.395725Z","shell.execute_reply":"2025-05-29T13:20:34.395615Z"},"scrolled":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:20:34.396347Z","iopub.status.idle":"2025-05-29T13:20:34.396573Z","shell.execute_reply":"2025-05-29T13:20:34.396486Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# !rm -rf ./*","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-20T04:22:58.521782Z","iopub.execute_input":"2025-04-20T04:22:58.522163Z","iopub.status.idle":"2025-04-20T04:22:58.535928Z","shell.execute_reply.started":"2025-04-20T04:22:58.522129Z","shell.execute_reply":"2025-04-20T04:22:58.535128Z"}},"outputs":[],"execution_count":null}]}