"""
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)