File size: 2,061 Bytes
5c776e5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | import py3Dmol
class ProteinVisualizer:
def __init__(self, file_data: str, file_extension: str):
"""
Initializes the ProteinVisualizer instance.
:param file_data: The file contents (PDB or CIF format) as a string.
:param file_extension: The extension of the file ("pdb" or "cif").
"""
self.file_data = file_data
self.file_extension = file_extension.lower()
def _add_model(self, view: py3Dmol.view) -> None:
"""
Adds the model to the 3Dmol viewer based on the file extension.
:param view: The 3Dmol view object where the model will be added.
"""
if self.file_extension == "pdb":
view.addModel(self.file_data, "pdb")
elif self.file_extension == "cif":
view.addModel(self.file_data, "cif")
else:
raise ValueError("Unsupported file format. Only PDB and CIF are supported.")
def generate_visualization(self, style: str = "cartoon") -> str:
"""
Generates an interactive 3D molecular visualization.
:param style: The visualization style ("cartoon", "stick", "sphere").
:return: The HTML content for embedding in a web application.
"""
view = py3Dmol.view(width=800, height=600)
# Add the model based on the file type
self._add_model(view)
# Apply the style chosen by the user
if style.lower() == "cartoon":
view.setStyle({"cartoon": {"color": "spectrum"}})
elif style.lower() == "stick":
view.setStyle({"stick": {}})
elif style.lower() == "sphere":
view.setStyle({"sphere": {}})
else:
raise ValueError(f"Unsupported style: {style}. Choose 'cartoon', 'stick', or 'sphere'.")
# Zoom in on the structure and render it
view.zoomTo(1.5)
# Return the HTML content for rendering
return view._make_html()
|