Spaces:

HouBioLab
/

MotifAlign

Sleeping

App Files Files Community

jiehou commited on Nov 21, 2025

Commit

8607410

verified ·

1 Parent(s): a83e037

Upload visualization.py

Browse files

Files changed (1) hide show

visualization.py +915 -0

visualization.py ADDED Viewed

	@@ -0,0 +1,915 @@

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