Spaces:
Sleeping
Sleeping
| """ | |
| 3D Visualization Module for RNA Structure Comparison | |
| Uses py3Dmol for interactive molecular visualization | |
| """ | |
| import numpy as np | |
| from rmsd_utils import ( | |
| parse_residue_atoms, | |
| translate_rotate_coords, | |
| calculate_COM, | |
| get_backbone_sugar_and_selectbase_coords_fixed | |
| ) | |
| def create_structure_visualization(ref_path, query_path, ref_window_indices, query_window_indices, | |
| rotation_matrix, ref_com, query_com, rmsd=None): | |
| """ | |
| Create an interactive 3D visualization of aligned structures. | |
| Args: | |
| ref_path: Path to reference motif PDB file | |
| query_path: Path to query motif PDB file | |
| ref_window_indices: List of residue indices for the reference window | |
| query_window_indices: List of residue indices for the query window | |
| rotation_matrix: Rotation matrix from RMSD calculation | |
| ref_com: Center of mass of reference window | |
| query_com: Center of mass of query window | |
| rmsd: RMSD value (optional, for display) | |
| Returns: | |
| HTML string containing the py3Dmol visualization | |
| """ | |
| # Read PDB files | |
| with open(ref_path) as f: | |
| ref_pdb = f.read() | |
| with open(query_path) as f: | |
| query_pdb_full = f.read() | |
| # Extract only the window residues from both structures | |
| ref_residues = parse_residue_atoms(ref_path) | |
| query_residues = parse_residue_atoms(query_path) | |
| ref_window_pdb = extract_window_pdb(ref_path, ref_window_indices) | |
| query_window_pdb = extract_window_pdb(query_path, query_window_indices) | |
| # Parse window coordinates for transformation | |
| from rmsd_utils import get_backbone_sugar_coords_from_residue, get_base_coords_from_residue | |
| ref_window_coords = [] | |
| for idx in ref_window_indices: | |
| if idx < len(ref_residues): | |
| residue = ref_residues[idx] | |
| backbone_coords = get_backbone_sugar_coords_from_residue(residue) | |
| ref_window_coords.extend(backbone_coords) | |
| base_coords = get_base_coords_from_residue(residue) | |
| ref_window_coords.extend(base_coords) | |
| ref_window_coords = np.asarray(ref_window_coords) | |
| query_window_coords = [] | |
| for idx in query_window_indices: | |
| if idx < len(query_residues): | |
| residue = query_residues[idx] | |
| backbone_coords = get_backbone_sugar_coords_from_residue(residue) | |
| query_window_coords.extend(backbone_coords) | |
| base_coords = get_base_coords_from_residue(residue) | |
| query_window_coords.extend(base_coords) | |
| query_window_coords = np.asarray(query_window_coords) | |
| # Transform query window to align with reference window | |
| # Proper alignment: translate to origin, rotate, translate to reference position | |
| # Note: We need both query_com and ref_com for proper alignment | |
| transformed_query_pdb = transform_pdb_string( | |
| query_window_pdb, | |
| rotation_matrix, | |
| query_com, | |
| ref_com # Add reference COM for proper alignment | |
| ) | |
| # Create py3Dmol visualization | |
| html = f""" | |
| <!DOCTYPE html> | |
| <html> | |
| <head> | |
| <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script> | |
| <style> | |
| #container {{ | |
| width: 100%; | |
| height: 700px; | |
| position: relative; | |
| border: 1px solid #ddd; | |
| }} | |
| .control-panel {{ | |
| position: absolute; | |
| top: 10px; | |
| right: 10px; | |
| background: rgba(255, 255, 255, 0.95); | |
| padding: 15px; | |
| border-radius: 8px; | |
| font-family: Arial, sans-serif; | |
| font-size: 13px; | |
| z-index: 1000; | |
| box-shadow: 0 2px 8px rgba(0,0,0,0.1); | |
| max-width: 220px; | |
| }} | |
| .control-panel h4 {{ | |
| margin: 0 0 10px 0; | |
| font-size: 14px; | |
| color: #333; | |
| }} | |
| .control-section {{ | |
| margin-bottom: 12px; | |
| padding-bottom: 12px; | |
| border-bottom: 1px solid #eee; | |
| }} | |
| .control-section:last-child {{ | |
| border-bottom: none; | |
| margin-bottom: 0; | |
| }} | |
| .control-section label {{ | |
| display: block; | |
| margin: 6px 0; | |
| cursor: pointer; | |
| }} | |
| .control-section input[type="checkbox"] {{ | |
| margin-right: 8px; | |
| }} | |
| .control-section select {{ | |
| width: 100%; | |
| padding: 4px; | |
| margin-top: 5px; | |
| border: 1px solid #ccc; | |
| border-radius: 4px; | |
| }} | |
| .legend {{ | |
| position: absolute; | |
| top: 10px; | |
| left: 10px; | |
| background: rgba(255, 255, 255, 0.95); | |
| padding: 15px; | |
| border-radius: 8px; | |
| font-family: Arial, sans-serif; | |
| font-size: 13px; | |
| z-index: 1000; | |
| box-shadow: 0 2px 8px rgba(0,0,0,0.1); | |
| }} | |
| .legend h4 {{ | |
| margin: 0 0 10px 0; | |
| font-size: 14px; | |
| color: #333; | |
| }} | |
| .legend-item {{ | |
| margin: 6px 0; | |
| display: flex; | |
| align-items: center; | |
| }} | |
| .color-box {{ | |
| width: 24px; | |
| height: 16px; | |
| margin-right: 10px; | |
| border: 1px solid #333; | |
| border-radius: 2px; | |
| }} | |
| .rmsd-info {{ | |
| position: absolute; | |
| bottom: 10px; | |
| left: 10px; | |
| background: rgba(255, 255, 255, 0.95); | |
| padding: 10px 15px; | |
| border-radius: 8px; | |
| font-family: Arial, sans-serif; | |
| font-size: 13px; | |
| z-index: 1000; | |
| box-shadow: 0 2px 8px rgba(0,0,0,0.1); | |
| }} | |
| .section-title {{ | |
| font-weight: bold; | |
| color: #555; | |
| margin-bottom: 5px; | |
| font-size: 12px; | |
| text-transform: uppercase; | |
| }} | |
| </style> | |
| </head> | |
| <body> | |
| <div id="container"></div> | |
| <div class="legend"> | |
| <h4>🧬 Structures</h4> | |
| <div class="legend-item"> | |
| <div class="color-box" style="background: #4A90E2;"></div> | |
| <span>Reference</span> | |
| </div> | |
| <div class="legend-item"> | |
| <div class="color-box" style="background: #E94B3C;"></div> | |
| <span>Query (Aligned)</span> | |
| </div> | |
| </div> | |
| <div class="control-panel"> | |
| <h4>⚙️ Display Options</h4> | |
| <div class="control-section"> | |
| <div class="section-title">Structures</div> | |
| <label> | |
| <input type="checkbox" id="showRef" checked onchange="updateDisplay()"> | |
| Reference | |
| </label> | |
| <label> | |
| <input type="checkbox" id="showQuery" checked onchange="updateDisplay()"> | |
| Query | |
| </label> | |
| </div> | |
| <div class="control-section"> | |
| <div class="section-title">Style</div> | |
| <select id="styleMode" onchange="updateDisplay()"> | |
| <option value="sticks">Sticks</option> | |
| <option value="cartoon">Cartoon</option> | |
| <option value="spheres">Spheres</option> | |
| <option value="lines">Lines</option> | |
| <option value="cartoon_sticks">Cartoon + Sticks</option> | |
| </select> | |
| </div> | |
| <div class="control-section"> | |
| <div class="section-title">Components</div> | |
| <label> | |
| <input type="checkbox" id="showBackbone" checked onchange="updateDisplay()"> | |
| Backbone/Sugar | |
| </label> | |
| <label> | |
| <input type="checkbox" id="showBases" checked onchange="updateDisplay()"> | |
| Bases | |
| </label> | |
| </div> | |
| <div class="control-section"> | |
| <div class="section-title">Labels</div> | |
| <label> | |
| <input type="checkbox" id="showLabels" onchange="updateDisplay()"> | |
| Residue Labels | |
| </label> | |
| <label> | |
| <input type="checkbox" id="showNumbers" onchange="updateDisplay()"> | |
| Residue Numbers | |
| </label> | |
| <label> | |
| <input type="checkbox" id="showAtoms" onchange="updateDisplay()"> | |
| Atom Names | |
| </label> | |
| <select id="atomLabelMode" style="margin-top: 5px; font-size: 11px;" onchange="updateDisplay()"> | |
| <option value="all">All Atoms</option> | |
| <option value="backbone">Backbone Only</option> | |
| <option value="sidechain">Bases Only</option> | |
| </select> | |
| </div> | |
| <div class="control-section"> | |
| <div class="section-title">Background</div> | |
| <select id="bgColor" onchange="updateBackground()"> | |
| <option value="white">White</option> | |
| <option value="black">Black</option> | |
| <option value="gray">Gray</option> | |
| </select> | |
| </div> | |
| </div> | |
| <div class="rmsd-info"> | |
| <strong>RMSD:</strong> <span style="color: #E94B3C; font-weight: bold;">{f"{rmsd:.3f}" if rmsd is not None else "N/A"} Å</span> | |
| </div> | |
| <script> | |
| let viewer = null; | |
| let refModel = null; | |
| let queryModel = null; | |
| const refPDB = `{ref_window_pdb}`; | |
| const queryPDB = `{transformed_query_pdb}`; | |
| // RNA backbone atoms | |
| const backboneAtoms = ['P', 'OP1', 'OP2', "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'"]; | |
| function initViewer() {{ | |
| try {{ | |
| viewer = $3Dmol.createViewer("container", {{ | |
| backgroundColor: 'white' | |
| }}); | |
| if (!refPDB || refPDB.length < 10) {{ | |
| throw new Error("Reference PDB data is empty"); | |
| }} | |
| if (!queryPDB || queryPDB.length < 10) {{ | |
| throw new Error("Query PDB data is empty"); | |
| }} | |
| updateDisplay(); | |
| viewer.zoomTo(); | |
| viewer.render(); | |
| }} catch (error) {{ | |
| console.error("Error initializing viewer:", error); | |
| document.getElementById("container").innerHTML = | |
| '<div style="padding: 20px; color: red; text-align: center;">Error loading visualization: ' + error.message + '</div>'; | |
| }} | |
| }} | |
| function updateBackground() {{ | |
| const bgColor = document.getElementById('bgColor').value; | |
| viewer.setBackgroundColor(bgColor); | |
| viewer.render(); | |
| }} | |
| function updateDisplay() {{ | |
| if (!viewer) return; | |
| try {{ | |
| // Clear everything | |
| viewer.removeAllModels(); | |
| viewer.removeAllLabels(); | |
| const showRef = document.getElementById('showRef').checked; | |
| const showQuery = document.getElementById('showQuery').checked; | |
| const showBackbone = document.getElementById('showBackbone').checked; | |
| const showBases = document.getElementById('showBases').checked; | |
| const showLabels = document.getElementById('showLabels').checked; | |
| const showNumbers = document.getElementById('showNumbers').checked; | |
| const showAtoms = document.getElementById('showAtoms').checked; | |
| const styleMode = document.getElementById('styleMode').value; | |
| // Reference structure (blue) | |
| if (showRef) {{ | |
| refModel = viewer.addModel(refPDB, "pdb"); | |
| applyStyle(refModel, '#4A90E2', '#5BA3F5', styleMode, showBackbone, showBases); | |
| if (showLabels || showNumbers) {{ | |
| addResidueLabels(refModel, '#4A90E2', showLabels, showNumbers); | |
| }} | |
| if (showAtoms) {{ | |
| addAtomLabels(refModel, '#4A90E2'); | |
| }} | |
| }} | |
| // Query structure (red) | |
| if (showQuery) {{ | |
| queryModel = viewer.addModel(queryPDB, "pdb"); | |
| applyStyle(queryModel, '#E94B3C', '#FF6B6B', styleMode, showBackbone, showBases); | |
| if (showLabels || showNumbers) {{ | |
| addResidueLabels(queryModel, '#E94B3C', showLabels, showNumbers); | |
| }} | |
| if (showAtoms) {{ | |
| addAtomLabels(queryModel, '#E94B3C'); | |
| }} | |
| }} | |
| viewer.zoomTo(); | |
| viewer.render(); | |
| }} catch (error) {{ | |
| console.error("Error updating display:", error); | |
| }} | |
| }} | |
| function applyStyle(model, backboneColor, baseColor, styleMode, showBackbone, showBases) {{ | |
| // Clear any existing styles | |
| viewer.setStyle({{model: model}}, {{}}); | |
| if (styleMode === 'cartoon') {{ | |
| // Cartoon representation | |
| viewer.setStyle({{model: model}}, {{ | |
| cartoon: {{ | |
| color: backboneColor, | |
| thickness: 0.5, | |
| opacity: 0.8 | |
| }} | |
| }}); | |
| }} else if (styleMode === 'cartoon_sticks') {{ | |
| // Cartoon + sticks for bases | |
| viewer.setStyle({{model: model}}, {{ | |
| cartoon: {{ | |
| color: backboneColor, | |
| thickness: 0.5, | |
| opacity: 0.7 | |
| }} | |
| }}); | |
| if (showBases) {{ | |
| viewer.addStyle({{model: model, not: {{atom: backboneAtoms}}}}, {{ | |
| stick: {{ | |
| color: baseColor, | |
| radius: 0.15 | |
| }} | |
| }}); | |
| }} | |
| }} else if (styleMode === 'spheres') {{ | |
| // Sphere representation | |
| if (showBackbone) {{ | |
| viewer.setStyle({{model: model, atom: backboneAtoms}}, {{ | |
| sphere: {{ | |
| color: backboneColor, | |
| radius: 0.4 | |
| }} | |
| }}); | |
| }} | |
| if (showBases) {{ | |
| viewer.addStyle({{model: model, not: {{atom: backboneAtoms}}}}, {{ | |
| sphere: {{ | |
| color: baseColor, | |
| radius: 0.35 | |
| }} | |
| }}); | |
| }} | |
| }} else if (styleMode === 'lines') {{ | |
| // Line representation | |
| if (showBackbone) {{ | |
| viewer.setStyle({{model: model, atom: backboneAtoms}}, {{ | |
| line: {{ | |
| color: backboneColor, | |
| linewidth: 2 | |
| }} | |
| }}); | |
| }} | |
| if (showBases) {{ | |
| viewer.addStyle({{model: model, not: {{atom: backboneAtoms}}}}, {{ | |
| line: {{ | |
| color: baseColor, | |
| linewidth: 2 | |
| }} | |
| }}); | |
| }} | |
| }} else {{ | |
| // Stick representation (default) | |
| if (showBackbone) {{ | |
| viewer.setStyle({{model: model, atom: backboneAtoms}}, {{ | |
| stick: {{ | |
| color: backboneColor, | |
| radius: 0.2 | |
| }}, | |
| sphere: {{ | |
| color: backboneColor, | |
| radius: 0.3 | |
| }} | |
| }}); | |
| }} | |
| if (showBases) {{ | |
| viewer.addStyle({{model: model, not: {{atom: backboneAtoms}}}}, {{ | |
| stick: {{ | |
| color: baseColor, | |
| radius: 0.15 | |
| }}, | |
| sphere: {{ | |
| color: baseColor, | |
| radius: 0.25 | |
| }} | |
| }}); | |
| }} | |
| }} | |
| }} | |
| function addResidueLabels(model, color, showLabels, showNumbers) {{ | |
| const atoms = viewer.selectedAtoms({{model: model}}); | |
| const residues = {{}}; | |
| // Group atoms by residue | |
| atoms.forEach(atom => {{ | |
| const key = atom.chain + '_' + atom.resi; | |
| if (!residues[key]) {{ | |
| residues[key] = atom; | |
| }} | |
| }}); | |
| // Add labels for each residue | |
| Object.values(residues).forEach(atom => {{ | |
| let labelText = ''; | |
| if (showLabels && showNumbers) {{ | |
| labelText = atom.resn + atom.resi; | |
| }} else if (showLabels) {{ | |
| labelText = atom.resn; | |
| }} else if (showNumbers) {{ | |
| labelText = atom.resi.toString(); | |
| }} | |
| if (labelText) {{ | |
| viewer.addLabel(labelText, {{ | |
| position: atom, | |
| backgroundColor: color, | |
| backgroundOpacity: 0.7, | |
| fontColor: 'white', | |
| fontSize: 11, | |
| fontWeight: 'bold', | |
| showBackground: true, | |
| borderRadius: 3 | |
| }}); | |
| }} | |
| }}); | |
| }} | |
| function addAtomLabels(model, color) {{ | |
| const atomLabelMode = document.getElementById('atomLabelMode').value; | |
| const atoms = viewer.selectedAtoms({{model: model}}); | |
| // Filter atoms based on mode | |
| let filteredAtoms = atoms; | |
| if (atomLabelMode === 'backbone') {{ | |
| // Only backbone atoms | |
| filteredAtoms = atoms.filter(atom => backboneAtoms.includes(atom.atom)); | |
| }} else if (atomLabelMode === 'sidechain') {{ | |
| // Only base/sidechain atoms (not backbone) | |
| filteredAtoms = atoms.filter(atom => !backboneAtoms.includes(atom.atom)); | |
| }} | |
| // 'all' mode uses all atoms (no filtering) | |
| // Add label for each atom | |
| filteredAtoms.forEach(atom => {{ | |
| // Use atom name (e.g., P, C1', N1, O4, etc.) | |
| const atomName = atom.atom; | |
| viewer.addLabel(atomName, {{ | |
| position: atom, | |
| backgroundColor: color, | |
| backgroundOpacity: 0.6, | |
| fontColor: 'white', | |
| fontSize: 9, | |
| fontWeight: 'normal', | |
| showBackground: true, | |
| borderRadius: 2, | |
| borderThickness: 0.5 | |
| }}); | |
| }}); | |
| }} | |
| // Initialize on load | |
| initViewer(); | |
| </script> | |
| </body> | |
| </html> | |
| """ | |
| return html | |
| def extract_window_pdb(pdb_path, window_indices): | |
| """ | |
| Extract specific residues from a PDB file based on window indices. | |
| Args: | |
| pdb_path: Path to PDB file | |
| window_indices: List of residue indices (0-based) | |
| Returns: | |
| String containing PDB data for only the specified residues | |
| """ | |
| with open(pdb_path) as f: | |
| lines = f.readlines() | |
| # Get all residue numbers from the file | |
| residues = parse_residue_atoms(pdb_path) | |
| if not residues: | |
| # If parsing failed, return original file | |
| return ''.join(lines) | |
| residue_numbers = [res['resnum'] for res in residues] | |
| # Map window indices to actual residue numbers | |
| target_resnums = set() | |
| for idx in window_indices: | |
| if idx < len(residue_numbers): | |
| target_resnums.add(residue_numbers[idx]) | |
| if not target_resnums: | |
| # If no valid residues, return original file | |
| return ''.join(lines) | |
| # Extract lines for these residues | |
| window_lines = [] | |
| for line in lines: | |
| if len(line) < 6: | |
| continue | |
| record = line[0:6].strip() | |
| if record in ['ATOM', 'HETATM', 'HETAT']: | |
| try: | |
| # Handle different PDB formats | |
| resnum_str = line[22:26].strip() | |
| if resnum_str: | |
| resnum = int(resnum_str) | |
| if resnum in target_resnums: | |
| window_lines.append(line) | |
| except (ValueError, IndexError): | |
| continue | |
| elif record in ['HEADER', 'TITLE', 'MODEL', 'ENDMDL']: | |
| window_lines.append(line) | |
| # Always add END record | |
| if window_lines and not any('END' in line for line in window_lines): | |
| window_lines.append('END\n') | |
| result = ''.join(window_lines) | |
| # Debug: print info about extraction | |
| if not result or len(result) < 50: | |
| print(f"Warning: Empty or very small PDB extracted from {pdb_path}") | |
| print(f" Window indices: {window_indices}") | |
| print(f" Target residue numbers: {target_resnums}") | |
| print(f" Result length: {len(result)}") | |
| # Return full structure if extraction failed | |
| return ''.join(lines) | |
| return result | |
| def transform_pdb_string(pdb_string, rotation_matrix, query_com, ref_com=None): | |
| """ | |
| Apply rotation and translation to coordinates in a PDB string to align with reference. | |
| The transformation aligns the query structure to the reference structure: | |
| 1. Translate query to origin (subtract query_com) | |
| 2. Apply rotation matrix | |
| 3. Translate to reference position (add ref_com) | |
| Args: | |
| pdb_string: PDB format string | |
| rotation_matrix: 3x3 rotation matrix | |
| query_com: Center of mass of query structure (to translate FROM) | |
| ref_com: Center of mass of reference structure (to translate TO), optional | |
| Returns: | |
| Transformed PDB string with aligned coordinates | |
| """ | |
| lines = pdb_string.split('\n') | |
| transformed_lines = [] | |
| # If ref_com not provided, just center at origin after rotation | |
| if ref_com is None: | |
| ref_com = np.array([0.0, 0.0, 0.0]) | |
| for line in lines: | |
| if len(line) < 54: | |
| transformed_lines.append(line) | |
| continue | |
| record = line[0:6].strip() | |
| if record in ['ATOM', 'HETATM', 'HETAT']: | |
| # Extract coordinates | |
| try: | |
| x = float(line[30:38].strip()) | |
| y = float(line[38:46].strip()) | |
| z = float(line[46:54].strip()) | |
| # Transform: (coord - query_com) @ rotation_matrix + ref_com | |
| # This aligns query to reference coordinate system | |
| coord = np.array([x, y, z]) | |
| centered = coord - query_com # Move query to origin | |
| rotated = np.dot(centered, rotation_matrix) # Rotate | |
| new_coord = rotated + ref_com # Move to reference position | |
| # Write transformed line | |
| new_line = ( | |
| line[:30] + | |
| f"{new_coord[0]:8.3f}" + | |
| f"{new_coord[1]:8.3f}" + | |
| f"{new_coord[2]:8.3f}" + | |
| line[54:] | |
| ) | |
| transformed_lines.append(new_line) | |
| except (ValueError, IndexError): | |
| transformed_lines.append(line) | |
| else: | |
| transformed_lines.append(line) | |
| return '\n'.join(transformed_lines) | |