"""
Multi-Structure 3D Visualization Module
Based on working pairwise visualization code
"""
import numpy as np
from rmsd_utils import (
parse_residue_atoms,
get_backbone_sugar_coords_from_residue,
get_base_coords_from_residue
)
def extract_window_pdb(pdb_path, window_indices):
"""
Extract specific residues from a PDB file based on window indices.
Uses the WORKING approach from the original code.
"""
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:
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:
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:
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)
if not result or len(result) < 50:
return ''.join(lines)
return result
def transform_pdb_string(pdb_string, rotation_matrix, query_com, ref_com):
"""
Apply rotation and translation to align query with reference.
Uses the WORKING transformation from original code.
CRITICAL: This is right multiplication (coord @ R), NOT left multiplication (R @ coord)
Args:
pdb_string: PDB format string
rotation_matrix: 3x3 rotation matrix from RMSD calculation
query_com: Center of mass of query structure (translate FROM)
ref_com: Center of mass of reference structure (translate TO)
Returns:
Transformed PDB string with aligned coordinates
"""
lines = pdb_string.split('\n')
transformed_lines = []
for line in lines:
if len(line) < 54:
transformed_lines.append(line)
continue
record = line[0:6].strip()
if record in ['ATOM', 'HETATM', 'HETAT']:
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 is the WORKING approach from original code
coord = np.array([x, y, z])
centered = coord - query_com # Move query to origin
rotated = np.dot(centered, rotation_matrix) # RIGHT multiplication
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)
def create_pairwise_visualization(ref_path, query_path, ref_window, query_window,
rotation_matrix, ref_com, query_com, rmsd,
ref_name="Reference", query_name="Query"):
"""
Create interactive 3D visualization of two aligned structures (pairwise).
Based on working original visualization with enhanced controls.
Args:
ref_path: Path to reference PDB file
query_path: Path to query PDB file
ref_window: List of residue indices for reference window
query_window: List of residue indices for query window
rotation_matrix: Rotation matrix from RMSD
ref_com: Center of mass of reference
query_com: Center of mass of query
rmsd: RMSD value
ref_name: Name of reference structure
query_name: Name of query structure
Returns:
HTML string for py3Dmol visualization
"""
# Extract windows
ref_pdb = extract_window_pdb(ref_path, ref_window)
query_pdb = extract_window_pdb(query_path, query_window)
# Transform query to align with reference
transformed_query_pdb = transform_pdb_string(
query_pdb,
rotation_matrix,
query_com,
ref_com
)
# Escape backticks
ref_pdb_escaped = ref_pdb.replace('`', '\\`')
transformed_query_pdb_escaped = transformed_query_pdb.replace('`', '\\`')
# Create HTML with enhanced controls
html = f"""
🧬 Structures
{ref_name}
{query_name} (Aligned)
RMSD: {rmsd:.3f} Å
⚙️ Display Options
Structures
Style
Components
Labels
Background
Download
"""
return html
def create_multistructure_visualization(ref_path, ref_window, ref_com, query_data_list, ref_name="Reference"):
"""
Create interactive 3D visualization with multiple structures aligned to reference.
Args:
ref_path: Path to reference PDB file
ref_window: List of residue indices for reference window
ref_com: Center of mass of reference window
query_data_list: List of dicts with keys:
- name: Structure name
- path: Path to PDB file
- window: List of residue indices
- rotation: Rotation matrix from RMSD
- query_com: Center of mass
- rmsd: RMSD value
- sequence: Sequence string
ref_name: Name of reference structure
Returns:
HTML string for py3Dmol visualization
"""
# Extract reference window
ref_pdb = extract_window_pdb(ref_path, ref_window)
# Color palette
colors = [
'#E94B3C', # Red
'#50C878', # Green
'#FF9500', # Orange
'#9B59B6', # Purple
'#00CED1', # Cyan
'#FF1493', # Pink
'#FFD700', # Gold
'#8B4513', # Brown
'#708090' # Slate Gray
]
# Transform all query structures
transformed_queries = []
for idx, query_data in enumerate(query_data_list):
# Extract query window
query_pdb = extract_window_pdb(query_data['path'], query_data['window'])
# Transform query to align with reference
transformed_pdb = transform_pdb_string(
query_pdb,
query_data['rotation'],
query_data['query_com'],
ref_com
)
color = colors[idx % len(colors)]
transformed_queries.append({
'pdb': transformed_pdb,
'name': query_data['name'],
'color': color,
'rmsd': query_data['rmsd'],
'sequence': query_data['sequence']
})
# Build JavaScript for model data
models_js = f"const refPDB = `{ref_pdb}`;\n"
for idx, tq in enumerate(transformed_queries):
# Escape backticks
pdb_escaped = tq['pdb'].replace('`', '\\`')
models_js += f"const queryPDB{idx} = `{pdb_escaped}`;\n"
# Build controls HTML
controls_html = f'\n'
for idx, tq in enumerate(transformed_queries):
controls_html += f'\n'
# Build legend HTML
legend_html = f'
{ref_name}
\n'
for tq in transformed_queries:
legend_html += f'