Source code for tooluniverse.protein_structure_3d_tool
"""
Protein Structure 3D Visualization Tool
=======================================
Tool for visualizing 3D protein structures using py3Dmol.
Supports PDB IDs, PDB file content, and various visualization styles.
"""
import requests
from typing import Any, Dict
from .visualization_tool import VisualizationTool
from .tool_registry import register_tool
[docs]
@register_tool("ProteinStructure3DTool")
class ProteinStructure3DTool(VisualizationTool):
"""Tool for visualizing 3D protein structures using py3Dmol."""
[docs]
def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
"""Generate 3D protein structure visualization."""
try:
import py3Dmol
# Extract parameters
pdb_id = arguments.get("pdb_id")
pdb_content = arguments.get("pdb_content")
style = arguments.get("style", "cartoon")
color_scheme = arguments.get("color_scheme", "spectrum")
width = arguments.get("width", self.default_width)
height = arguments.get("height", self.default_height)
show_sidechains = arguments.get("show_sidechains", False)
show_surface = arguments.get("show_surface", False)
# Get PDB content
if pdb_content:
pdb_data = pdb_content
elif pdb_id:
pdb_data = self._fetch_pdb_content(pdb_id)
else:
return self.create_error_response(
"Either pdb_id or pdb_content must be provided"
)
# Create py3Dmol viewer
viewer = py3Dmol.view(width=width, height=height)
viewer.addModel(pdb_data, "pdb")
# Apply visualization style
if style == "cartoon":
viewer.setStyle({"cartoon": {"color": color_scheme}})
elif style == "stick":
viewer.setStyle({"stick": {"color": color_scheme}})
elif style == "sphere":
viewer.setStyle({"sphere": {"color": color_scheme}})
elif style == "line":
viewer.setStyle({"line": {"color": color_scheme}})
else:
viewer.setStyle({"cartoon": {"color": color_scheme}})
# Add sidechains if requested
if show_sidechains:
resn_list = [
"ALA",
"ARG",
"ASN",
"ASP",
"CYS",
"GLN",
"GLU",
"GLY",
"HIS",
"ILE",
"LEU",
"LYS",
"MET",
"PHE",
"PRO",
"SER",
"THR",
"TRP",
"TYR",
"VAL",
]
viewer.addStyle(
{"stick": {"radius": 0.1}}, {"and": [{"resn": resn_list}]}
)
# Add surface if requested
if show_surface:
viewer.addSurface(py3Dmol.VDW, {"opacity": 0.7, "color": "white"})
# Zoom to fit
viewer.zoomTo()
# Generate HTML with modern UI
viewer_html = viewer._make_html()
# Create control panel
control_panel = self._create_control_panel(style, color_scheme)
# Create toolbar
toolbar = self._create_toolbar()
# Create info cards
info_cards = self._create_protein_info_cards(pdb_id, pdb_data)
# Generate modern HTML
html_content = self.create_py3dmol_html(
viewer_html,
width,
height,
title=f"Protein Structure: {pdb_id or 'Custom PDB'}",
control_panel=control_panel,
toolbar=toolbar,
info_cards=info_cards,
)
# Add JavaScript controls
html_content = html_content.replace(
"</body>", f"{self.add_3d_controls_script()}</body>"
)
# Prepare metadata
metadata = {
"width": width,
"height": height,
"style": style,
"color_scheme": color_scheme,
"pdb_id": pdb_id,
"show_sidechains": show_sidechains,
"show_surface": show_surface,
}
return self.create_visualization_response(
html_content=html_content,
viz_type="protein_structure_3d",
data={
"pdb_content": (
pdb_data[:500] + "..." if len(pdb_data) > 500 else pdb_data
)
},
metadata=metadata,
)
except ImportError:
return self.create_error_response(
"py3Dmol is not installed. Please install it with: "
"pip install py3Dmol",
"MissingDependency",
)
except Exception as e:
return self.create_error_response(
f"Failed to create protein visualization: {str(e)}"
)
def _fetch_pdb_content(self, pdb_id: str) -> str:
"""Fetch PDB content from RCSB PDB database."""
try:
url = f"https://files.rcsb.org/view/{pdb_id.upper()}.pdb"
response = requests.get(url, timeout=30)
response.raise_for_status()
return response.text
except requests.RequestException as e:
raise ValueError(f"Failed to fetch PDB file for {pdb_id}: {str(e)}")
def _create_control_panel(self, current_style: str, current_color: str) -> str:
"""Create floating control panel HTML."""
return f"""
<div class="control-panel">
<div class="control-group">
<label class="control-label">Style</label>
<select class="control-select" id="styleSelect" onchange="changeStyle()">
<option value="cartoon" {'selected' if current_style == 'cartoon' else ''}>Cartoon</option>
<option value="stick" {'selected' if current_style == 'stick' else ''}>Stick</option>
<option value="sphere" {'selected' if current_style == 'sphere' else ''}>Sphere</option>
<option value="line" {'selected' if current_style == 'line' else ''}>Line</option>
<option value="surface" {'selected' if current_style == 'surface' else ''}>Surface</option>
</select>
</div>
<div class="control-group">
<label class="control-label">Color Scheme</label>
<select class="control-select" id="colorSelect" onchange="changeColor()">
<option value="spectrum" {'selected' if current_color == 'spectrum' else ''}>Spectrum</option>
<option value="rainbow" {'selected' if current_color == 'rainbow' else ''}>Rainbow</option>
<option value="ss" {'selected' if current_color == 'ss' else ''}>Secondary Structure</option>
<option value="chain" {'selected' if current_color == 'chain' else ''}>Chain</option>
<option value="elem" {'selected' if current_color == 'elem' else ''}>Element</option>
</select>
</div>
<div class="control-group">
<label class="control-label">Background</label>
<select class="control-select" id="bgSelect" onchange="changeBackground()">
<option value="white" selected>White</option>
<option value="black">Black</option>
<option value="gray">Gray</option>
</select>
</div>
</div>
"""
def _create_toolbar(self) -> str:
"""Create bottom toolbar HTML."""
return """
<div class="toolbar">
<button class="btn" onclick="resetView()">Reset View</button>
<button class="btn btn-secondary" onclick="downloadScreenshot()">Screenshot</button>
<button class="btn btn-outline" onclick="toggleFullscreen()">Fullscreen</button>
</div>
"""
def _create_protein_info_cards(self, pdb_id: str, pdb_data: str) -> str:
"""Create protein information cards."""
# Parse basic PDB info
lines = pdb_data.split("\n")
title = "Unknown Protein"
organism = "Unknown"
resolution = "N/A"
method = "Unknown"
residue_count = 0
atom_count = 0
for line in lines:
if line.startswith("TITLE"):
title = line[10:].strip()
elif line.startswith("SOURCE"):
if "ORGANISM_SCIENTIFIC" in line:
organism = (
line.split("ORGANISM_SCIENTIFIC:")[1].split(";")[0].strip()
)
elif line.startswith("REMARK 2 RESOLUTION"):
resolution = line.split("RESOLUTION.")[1].split("ANGSTROMS")[0].strip()
elif line.startswith("EXPDTA"):
method = line[10:].strip()
elif line.startswith("ATOM"):
atom_count += 1
if atom_count == 1:
residue_count = 1
elif line[22:26].strip() != lines[atom_count - 2][22:26].strip():
residue_count += 1
return f"""
<div class="card">
<h3 class="card-title">
<svg class="card-icon" viewBox="0 0 24 24">
<path d="M12,2A10,10 0 0,0 2,12A10,10 0 0,0 12,22A10,10 0 0,0 22,12A10,10 0 0,0 12,2M12,4A8,8 0 0,1 20,12A8,8 0 0,1 12,20A8,8 0 0,1 4,12A8,8 0 0,1 12,4M12,6A6,6 0 0,0 6,12A6,6 0 0,0 12,18A6,6 0 0,0 18,12A6,6 0 0,0 12,6M12,8A4,4 0 0,1 16,12A4,4 0 0,1 12,16A4,4 0 0,1 8,12A4,4 0 0,1 12,8Z"/>
</svg>
Protein Information
</h3>
<div class="info-grid">
<div class="info-item">
<span class="info-label">PDB ID</span>
<span class="info-value">{pdb_id or 'Custom'}</span>
</div>
<div class="info-item">
<span class="info-label">Title</span>
<span class="info-value">{title[:30]}{'...' if len(title) > 30 else ''}</span>
</div>
<div class="info-item">
<span class="info-label">Organism</span>
<span class="info-value">{organism[:20]}{'...' if len(organism) > 20 else ''}</span>
</div>
<div class="info-item">
<span class="info-label">Method</span>
<span class="info-value">{method}</span>
</div>
<div class="info-item">
<span class="info-label">Resolution</span>
<span class="info-value">{resolution}</span>
</div>
</div>
</div>
<div class="card">
<h3 class="card-title">
<svg class="card-icon" viewBox="0 0 24 24">
<path d="M12,2A10,10 0 0,0 2,12A10,10 0 0,0 12,22A10,10 0 0,0 22,12A10,10 0 0,0 12,2M12,4A8,8 0 0,1 20,12A8,8 0 0,1 12,20A8,8 0 0,1 4,12A8,8 0 0,1 12,4M12,6A6,6 0 0,0 6,12A6,6 0 0,0 12,18A6,6 0 0,0 18,12A6,6 0 0,0 12,6M12,8A4,4 0 0,1 16,12A4,4 0 0,1 12,16A4,4 0 0,1 8,12A4,4 0 0,1 12,8Z"/>
</svg>
Structure Statistics
</h3>
<div class="info-grid">
<div class="info-item">
<span class="info-label">Residues</span>
<span class="info-value">{residue_count}</span>
</div>
<div class="info-item">
<span class="info-label">Atoms</span>
<span class="info-value">{atom_count}</span>
</div>
<div class="info-item">
<span class="info-label">Chains</span>
<span class="info-value">1</span>
</div>
</div>
</div>
"""