Spaces:
Sleeping
Sleeping
File size: 2,582 Bytes
ce7bf5b |
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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
# Copyright Generate Biomedicines, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Utilities for rendering protein structures in Jupyter notebooks.
This provides convenience functions for rendering our common structure
datatypes, such as `mst.System` and `XCS` tensors, with nglview.
"""
import tempfile
import uuid
import nglview as nv
class SystemTrajectory(nv.base_adaptor.Trajectory, nv.base_adaptor.Structure):
"""MST multi-state System object adaptor, by analogy to other NGLView adaptor
classes (e.g., MDTrajTrajectory or PyTrajTrajectory in nglview.adaptor).
Example
-------
>>> import nglview as nv
>>> import chroma.data.protein import Protein
>>> protein_trajectory = Protein("multi-state.cif")
>>> t = SystemTrajectory(protein_trajectory.sys)
>>> w = nv.NGLWidget(t)
>>> w
"""
def __init__(self, protein):
self.protein = protein
self.ext = "pdb"
self.params = {}
self.id = str(uuid.uuid4())
def get_coordinates(self, index):
self.protein.sys.swap_model(index)
X, _, _ = self.protein.sys.to_XCS()
self.protein.sys.swap_model(index)
return X.view(-1, 3).numpy()
@property
def n_frames(self):
return self.protein.sys.num_models()
def get_structure_string(self):
return self.protein.sys.to_PDB_string()
def view_gsystem(system, **kwargs):
"""Return an NGL Viewer Widget for an generate System.
Args:
system (System): Structure to view.
Returns:
view: NGL Viewer widget instance that. In a Jupyter notebook
returning this to the notebook will trigger display of a
widget.
"""
temp = tempfile.NamedTemporaryFile(suffix=".pdb")
filename = temp.name
system.to_PDB(filename)
view = nv.show_file(filename)
view.clear_representations()
view.add_representation("cartoon")
view.add_representation("licorice", selection="(sidechain or .CA) and not hydrogen")
view.add_representation("contact")
view.center()
return view
|