MotifAlign / visualization.py
jiehou's picture
Upload visualization.py
8607410 verified
"""
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"""
<!DOCTYPE html>
<html>
<head>
<script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
<script src="https://cdnjs.cloudflare.com/ajax/libs/html2canvas/1.4.1/html2canvas.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: 12px 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: 450px;
}}
.info-row {{
margin: 4px 0;
line-height: 1.4;
}}
.info-label {{
font-weight: bold;
color: #555;
}}
.info-value {{
color: #333;
font-family: 'Courier New', monospace;
}}
.section-title {{
font-weight: bold;
color: #555;
margin-bottom: 5px;
font-size: 12px;
text-transform: uppercase;
}}
.download-section {{
position: absolute;
bottom: 10px;
right: 10px;
background: rgba(255, 255, 255, 0.95);
padding: 10px;
border-radius: 8px;
font-family: Arial, sans-serif;
z-index: 1000;
box-shadow: 0 2px 8px rgba(0,0,0,0.1);
}}
.download-btn {{
background: #4A90E2;
color: white;
border: none;
padding: 8px 16px;
border-radius: 4px;
cursor: pointer;
font-size: 13px;
font-weight: bold;
}}
.download-btn:hover {{
background: #357ABD;
}}
</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 class="control-section">
<div class="section-title">Annotation Font Size</div>
<select id="annotationFontSize">
<option value="small">Small (18pt/16pt/14pt)</option>
<option value="medium">Medium (22pt/18pt/16pt)</option>
<option value="large" selected>Large (28pt/22pt/18pt)</option>
<option value="xlarge">Extra Large (36pt/28pt/22pt)</option>
</select>
</div>
</div>
<div class="rmsd-info">
<div class="info-row">
<span class="info-label">RMSD:</span>
<span style="color: #E94B3C; font-weight: bold; font-size: 14px;">{f"{rmsd:.3f}" if rmsd is not None else "N/A"} Å</span>
</div>
<div style="margin-top: 8px; padding-top: 8px; border-top: 1px solid #ddd;">
<div class="info-row">
<span class="info-label">Reference:</span>
<span class="info-value">{ref_name}</span>
</div>
{f'<div class="info-row" style="margin-left: 15px; font-size: 12px;"><span class="info-label">Seq:</span> <span class="info-value">{ref_sequence}</span></div>' if ref_sequence else ''}
</div>
<div style="margin-top: 6px;">
<div class="info-row">
<span class="info-label">Query:</span>
<span class="info-value">{query_name}</span>
</div>
{f'<div class="info-row" style="margin-left: 15px; font-size: 12px;"><span class="info-label">Seq:</span> <span class="info-value">{query_sequence}</span></div>' if query_sequence else ''}
</div>
</div>
<div class="download-section">
<button class="download-btn" onclick="downloadImage()">📷 Download PNG</button>
</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
}});
}});
}}
function downloadImage() {{
try {{
// Generate filename with metadata
var refName = "{ref_name}".replace('.pdb', '');
var queryName = "{query_name}".replace('.pdb', '');
var rmsdValue = "{f'{rmsd:.3f}' if rmsd is not None else 'NA'}";
var refSeq = "{ref_sequence if ref_sequence else ''}";
var querySeq = "{query_sequence if query_sequence else ''}";
var filenameOriginal = 'alignment_' + refName + '_' + queryName + '_RMSD_' + rmsdValue + '.png';
var filenameAnnotated = 'annotated_' + refName + '_' + queryName + '_RMSD_' + rmsdValue + '.png';
// Get selected font size
const fontSizeSelect = document.getElementById('annotationFontSize');
const fontSizeOption = fontSizeSelect ? fontSizeSelect.value : 'large';
// Define font sizes based on selection (all values are at 2x scale for high resolution)
let fontSizes;
switch(fontSizeOption) {{
case 'small':
fontSizes = {{ rmsd: 36, name: 32, seq: 28 }}; // 18pt/16pt/14pt at 2x
break;
case 'medium':
fontSizes = {{ rmsd: 44, name: 36, seq: 32 }}; // 22pt/18pt/16pt at 2x
break;
case 'large':
fontSizes = {{ rmsd: 56, name: 44, seq: 36 }}; // 28pt/22pt/18pt at 2x
break;
case 'xlarge':
fontSizes = {{ rmsd: 72, name: 56, seq: 44 }}; // 36pt/28pt/22pt at 2x
break;
default:
fontSizes = {{ rmsd: 56, name: 44, seq: 36 }}; // Default to large
}}
// Ensure viewer is rendered
if (viewer) {{
viewer.render();
}}
// Get the container element
const container = document.getElementById('container');
if (!container) {{
alert('Container not ready. Please wait and try again.');
return;
}}
// Use html2canvas to capture the entire container with overlays
html2canvas(container, {{
backgroundColor: '#ffffff',
scale: 2, // Higher resolution
logging: false,
useCORS: true,
allowTaint: true
}}).then(function(canvas) {{
// Create ANNOTATED version
const annotatedCanvas = document.createElement('canvas');
annotatedCanvas.width = canvas.width;
annotatedCanvas.height = canvas.height;
const ctx = annotatedCanvas.getContext('2d');
// Draw the original image onto new canvas
ctx.drawImage(canvas, 0, 0);
// Add annotations
const margin = 30; // Scaled for 2x resolution
const padding = 24;
const lineSpacing = 16;
// Prepare annotation text with selected font sizes
const annotations = [
{{ text: 'RMSD: ' + rmsdValue + ' Å', fontSize: fontSizes.rmsd, fontFamily: 'bold Arial', color: '#E94B3C' }},
{{ text: '', fontSize: 20, fontFamily: 'Arial', color: '#333' }}, // Spacer
{{ text: 'Reference: ' + refName, fontSize: fontSizes.name, fontFamily: 'Arial', color: '#333' }},
{{ text: ' Seq: ' + refSeq, fontSize: fontSizes.seq, fontFamily: 'Courier New, monospace', color: '#666' }},
{{ text: '', fontSize: 20, fontFamily: 'Arial', color: '#333' }}, // Spacer
{{ text: 'Query: ' + queryName, fontSize: fontSizes.name, fontFamily: 'Arial', color: '#333' }},
{{ text: ' Seq: ' + querySeq, fontSize: fontSizes.seq, fontFamily: 'Courier New, monospace', color: '#666' }}
];
// Calculate box dimensions
let maxWidth = 0;
let totalHeight = padding * 2;
const textMetrics = [];
annotations.forEach(ann => {{
if (ann.text) {{
ctx.font = ann.fontSize + 'px ' + ann.fontFamily;
const metrics = ctx.measureText(ann.text);
const height = ann.fontSize * 1.2; // Approximate height
textMetrics.push({{ width: metrics.width, height: height }});
maxWidth = Math.max(maxWidth, metrics.width);
totalHeight += height + lineSpacing;
}} else {{
textMetrics.push({{ width: 0, height: lineSpacing / 2 }});
totalHeight += lineSpacing / 2;
}}
}});
const boxWidth = maxWidth + padding * 2;
const boxHeight = totalHeight;
// Position box in bottom-left
const boxX = margin;
const boxY = annotatedCanvas.height - boxHeight - margin;
// Draw semi-transparent white background with rounded corners
ctx.fillStyle = 'rgba(255, 255, 255, 0.95)';
const radius = 16;
ctx.beginPath();
ctx.moveTo(boxX + radius, boxY);
ctx.lineTo(boxX + boxWidth - radius, boxY);
ctx.quadraticCurveTo(boxX + boxWidth, boxY, boxX + boxWidth, boxY + radius);
ctx.lineTo(boxX + boxWidth, boxY + boxHeight - radius);
ctx.quadraticCurveTo(boxX + boxWidth, boxY + boxHeight, boxX + boxWidth - radius, boxY + boxHeight);
ctx.lineTo(boxX + radius, boxY + boxHeight);
ctx.quadraticCurveTo(boxX, boxY + boxHeight, boxX, boxY + boxHeight - radius);
ctx.lineTo(boxX, boxY + radius);
ctx.quadraticCurveTo(boxX, boxY, boxX + radius, boxY);
ctx.closePath();
ctx.fill();
// Draw border
ctx.strokeStyle = 'rgba(200, 200, 200, 0.95)';
ctx.lineWidth = 2;
ctx.stroke();
// Draw text
let currentY = boxY + padding;
annotations.forEach((ann, idx) => {{
if (ann.text) {{
ctx.font = ann.fontSize + 'px ' + ann.fontFamily;
ctx.fillStyle = ann.color;
ctx.fillText(ann.text, boxX + padding, currentY + textMetrics[idx].height * 0.8);
currentY += textMetrics[idx].height + lineSpacing;
}} else {{
currentY += textMetrics[idx].height;
}}
}});
// Download ONLY the annotated PNG
const annotatedDataURL = annotatedCanvas.toDataURL('image/png');
const linkAnnotated = document.createElement('a');
linkAnnotated.download = filenameAnnotated;
linkAnnotated.href = annotatedDataURL;
document.body.appendChild(linkAnnotated);
linkAnnotated.click();
document.body.removeChild(linkAnnotated);
}}).catch(function(error) {{
console.error('html2canvas error:', error);
alert('Error creating images. Please try again.');
}});
}} catch (error) {{
console.error('PNG download error:', error);
alert('Error downloading PNG: ' + error.message);
}}
}}
// 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)