""" 3D Visualization Module for RNA Structure Comparison Uses py3Dmol for interactive molecular visualization """ import numpy as np from pathlib import Path 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, ref_name=None, query_name=None, ref_sequence=None, query_sequence=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) ref_name: Reference structure name (optional, for display) query_name: Query structure name (optional, for display) ref_sequence: Reference sequence (optional, for display) query_sequence: Query sequence (optional, for display) Returns: HTML string containing the py3Dmol visualization """ # Extract simple names if not provided if ref_name is None: ref_name = Path(ref_path).stem if query_name is None: query_name = Path(query_path).stem # 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"""

🧬 Structures

Reference
Query (Aligned)

⚙️ Display Options

Structures
Style
Components
Labels
Background
Annotation Font Size
RMSD: {f"{rmsd:.3f}" if rmsd is not None else "N/A"} Å
Reference: {ref_name}
{f'
Seq: {ref_sequence}
' if ref_sequence else ''}
Query: {query_name}
{f'
Seq: {query_sequence}
' if query_sequence else ''}
""" 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)