| 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)
|
|
|
|
|
| self._add_model(view)
|
|
|
|
|
| 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'.")
|
|
|
|
|
| view.zoomTo(1.5)
|
|
|
|
|
| return view._make_html()
|
|
|
