def generate_points():
centers = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
colors = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) * 255
vtk_vertices = Points()
# Create the topology of the point (a vertex)
vtk_faces = CellArray()
# Add points
for i in range(len(centers)):
p = centers[i]
id = vtk_vertices.InsertNextPoint(p)
vtk_faces.InsertNextCell(1)
vtk_faces.InsertCellPoint(id)
# Create a polydata object
polydata = PolyData()
# Set the vertices and faces we created as the geometry and topology of the
# polydata
polydata.SetPoints(vtk_vertices)
polydata.SetVerts(vtk_faces)
set_polydata_colors(polydata, colors)
mapper = PolyDataMapper()
mapper.SetInputData(polydata)
mapper.SetVBOShiftScaleMethod(False)
point_actor = Actor()
point_actor.SetMapper(mapper)
return point_actor
def test_shader_callback():
cone = ConeSource()
coneMapper = PolyDataMapper()
coneMapper.SetInputConnection(cone.GetOutputPort())
actor = Actor()
actor.SetMapper(coneMapper)
test_values = []
def callbackLow(_caller, _event, calldata=None):
program = calldata
if program is not None:
test_values.append(0)
id_observer = fs.add_shader_callback(actor, callbackLow, 0)
with pytest.raises(Exception):
fs.add_shader_callback(actor, callbackLow, priority='str')
mapper = actor.GetMapper()
mapper.RemoveObserver(id_observer)
scene = window.Scene()
scene.add(actor)
arr1 = window.snapshot(scene, size=(200, 200))
assert len(test_values) == 0
test_values = []
def callbackHigh(_caller, _event, calldata=None):
program = calldata
if program is not None:
test_values.append(999)
def callbackMean(_caller, _event, calldata=None):
program = calldata
if program is not None:
test_values.append(500)
fs.add_shader_callback(actor, callbackHigh, 999)
fs.add_shader_callback(actor, callbackLow, 0)
id_mean = fs.add_shader_callback(actor, callbackMean, 500)
# check the priority of each call
arr2 = window.snapshot(scene, size=(200, 200))
assert np.abs(
[test_values[0] - 999, test_values[1] - 500,
test_values[2] - 0]).sum() == 0
# check if the correct observer was removed
mapper.RemoveObserver(id_mean)
test_values = []
arr3 = window.snapshot(scene, size=(200, 200))
assert np.abs([test_values[0] - 999, test_values[1] - 0]).sum() == 0
def stick(molecule, colormode='discrete', bond_thickness=0.1):
"""Create an actor for stick molecular representation.
Parameters
----------
molecule : Molecule
The molecule to be rendered.
colormode : string, optional
Set the colormode for coloring the bonds. Two valid color modes:
* 'discrete': Bonds are colored using CPK coloring convention.
* 'single': All bonds are colored with the same color (dark grey)
RGB tuple used for coloring the bonds when 'single' colormode is
selected: (50, 50, 50)
Default: 'discrete'
bond_thickness : float, optional
Used to manipulate the thickness of bonds (i.e. thickness of tubes
which are used to render bonds).
Default: 0.1 (Optimal range: 0.1 - 0.5).
Returns
-------
molecule_actor : vtkActor
Actor created to render the stick representation of the molecule to be
visualized.
"""
if molecule.total_num_bonds == 0:
raise ValueError('No bonding data available for the molecule! Stick '
'model cannot be made!')
colormode = colormode.lower()
mst_mapper = OpenGLMoleculeMapper()
mst_mapper.SetInputData(molecule)
mst_mapper.SetRenderAtoms(True)
mst_mapper.SetRenderBonds(True)
mst_mapper.SetBondRadius(bond_thickness)
mst_mapper.SetAtomicRadiusTypeToUnitRadius()
mst_mapper.SetAtomicRadiusScaleFactor(bond_thickness)
if colormode == 'discrete':
mst_mapper.SetAtomColorMode(1)
mst_mapper.SetBondColorMode(1)
elif colormode == 'single':
mst_mapper.SetAtomColorMode(0)
mst_mapper.SetBondColorMode(0)
else:
mst_mapper.SetAtomColorMode(1)
warnings.warn("Incorrect colormode specified! Using discrete.")
molecule_actor = Actor()
molecule_actor.SetMapper(mst_mapper)
return molecule_actor
def sphere_cpk(molecule, colormode='discrete'):
"""Create an actor for sphere molecular representation. It's also referred
to as CPK model and space-filling model.
Parameters
----------
molecule : Molecule
The molecule to be rendered.
colormode : string, optional
Set the colormode for coloring the atoms. Two valid color modes:
* 'discrete': Atoms are colored using CPK coloring convention.
* 'single': All atoms are colored with same color (grey). RGB tuple
used for coloring the atoms when 'single' colormode is selected:
(150, 150, 150).
Default: 'discrete'
Returns
-------
molecule_actor : vtkActor
Actor created to render the space filling representation of the
molecule to be visualized.
References
----------
Corey R.B.; Pauling L. Molecular Models of Amino Acids,
Peptides, and Proteins
`Review of Scientific Instruments 1953, 24 (8), 621-627.
<https://doi.org/10.1063/1.1770803>`_
"""
colormode = colormode.lower()
msp_mapper = OpenGLMoleculeMapper()
msp_mapper.SetInputData(molecule)
msp_mapper.SetRenderAtoms(True)
msp_mapper.SetRenderBonds(False)
msp_mapper.SetAtomicRadiusTypeToVDWRadius()
msp_mapper.SetAtomicRadiusScaleFactor(1)
if colormode == 'discrete':
msp_mapper.SetAtomColorMode(1)
elif colormode == 'single':
msp_mapper.SetAtomColorMode(0)
else:
msp_mapper.SetAtomColorMode(1)
warnings.warn("Incorrect colormode specified! Using discrete.")
# To-Do manipulate shading properties to make it look aesthetic
molecule_actor = Actor()
molecule_actor.SetMapper(msp_mapper)
return molecule_actor
def get_actor_from_polymapper(poly_mapper):
"""Get actor from a vtkPolyDataMapper.
Parameters
----------
poly_mapper : vtkPolyDataMapper
Returns
-------
actor : actor
"""
actor = Actor()
actor.SetMapper(poly_mapper)
actor.GetProperty().BackfaceCullingOn()
actor.GetProperty().SetInterpolationToPhong()
return actor
def repeat_sources(centers,
colors,
active_scalars=1.,
directions=None,
source=None,
vertices=None,
faces=None,
orientation=None):
"""Transform a vtksource to glyph."""
if source is None and faces is None:
raise IOError("A source or faces should be defined")
if np.array(colors).ndim == 1:
colors = np.tile(colors, (len(centers), 1))
pts = numpy_to_vtk_points(np.ascontiguousarray(centers))
cols = numpy_to_vtk_colors(255 * np.ascontiguousarray(colors))
cols.SetName('colors')
if isinstance(active_scalars, (float, int)):
active_scalars = np.tile(active_scalars, (len(centers), 1))
if isinstance(active_scalars, np.ndarray):
ascalars = numpy_support.numpy_to_vtk(np.asarray(active_scalars),
deep=True,
array_type=VTK_DOUBLE)
ascalars.SetName('active_scalars')
if directions is not None:
directions_fa = numpy_support.numpy_to_vtk(np.asarray(directions),
deep=True,
array_type=VTK_DOUBLE)
directions_fa.SetName('directions')
polydata_centers = PolyData()
polydata_geom = PolyData()
if faces is not None:
set_polydata_vertices(polydata_geom, vertices)
set_polydata_triangles(polydata_geom, faces)
polydata_centers.SetPoints(pts)
polydata_centers.GetPointData().AddArray(cols)
if directions is not None:
polydata_centers.GetPointData().AddArray(directions_fa)
polydata_centers.GetPointData().SetActiveVectors('directions')
if isinstance(active_scalars, np.ndarray):
polydata_centers.GetPointData().AddArray(ascalars)
polydata_centers.GetPointData().SetActiveScalars('active_scalars')
glyph = Glyph3D()
if faces is None:
if orientation is not None:
transform = Transform()
transform.SetMatrix(numpy_to_vtk_matrix(orientation))
rtrans = TransformPolyDataFilter()
rtrans.SetInputConnection(source.GetOutputPort())
rtrans.SetTransform(transform)
source = rtrans
glyph.SetSourceConnection(source.GetOutputPort())
else:
glyph.SetSourceData(polydata_geom)
glyph.SetInputData(polydata_centers)
glyph.SetOrient(True)
glyph.SetScaleModeToScaleByScalar()
glyph.SetVectorModeToUseVector()
glyph.Update()
mapper = PolyDataMapper()
mapper.SetInputData(glyph.GetOutput())
mapper.SetScalarModeToUsePointFieldData()
mapper.SelectColorArray('colors')
actor = Actor()
actor.SetMapper(mapper)
return actor
def ribbon(molecule):
"""Create an actor for ribbon molecular representation.
Parameters
----------
molecule : Molecule
The molecule to be rendered.
Returns
-------
molecule_actor : vtkActor
Actor created to render the rubbon representation of the molecule to be
visualized.
References
----------
Richardson, J.S. The anatomy and taxonomy of protein structure
`Advances in Protein Chemistry, 1981, 34, 167-339.
<https://doi.org/10.1016/S0065-3233(08)60520-3>`_
"""
coords = get_all_atomic_positions(molecule)
all_atomic_numbers = get_all_atomic_numbers(molecule)
num_total_atoms = molecule.total_num_atoms
secondary_structures = np.ones(num_total_atoms)
for i in range(num_total_atoms):
secondary_structures[i] = ord('c')
resi = molecule.residue_seq[i]
for j, _ in enumerate(molecule.sheet):
sheet = molecule.sheet[j]
if molecule.chain[i] != sheet[0] or resi < sheet[1] or \
resi > sheet[3]:
continue
secondary_structures[i] = ord('s')
for j, _ in enumerate(molecule.helix):
helix = molecule.helix[j]
if molecule.chain[i] != helix[0] or resi < helix[1] or \
resi > helix[3]:
continue
secondary_structures[i] = ord('h')
output = PolyData()
# for atomic numbers
atomic_num_arr = nps.numpy_to_vtk(num_array=all_atomic_numbers,
deep=True,
array_type=VTK_ID_TYPE)
# setting the array name to atom_type as vtkProteinRibbonFilter requires
# the array to be named atom_type
atomic_num_arr.SetName("atom_type")
output.GetPointData().AddArray(atomic_num_arr)
# for atom names
atom_names = StringArray()
# setting the array name to atom_types as vtkProteinRibbonFilter requires
# the array to be named atom_types
atom_names.SetName("atom_types")
atom_names.SetNumberOfTuples(num_total_atoms)
for i in range(num_total_atoms):
atom_names.SetValue(i, molecule.atom_names[i])
output.GetPointData().AddArray(atom_names)
# for residue sequences
residue_seq = nps.numpy_to_vtk(num_array=molecule.residue_seq,
deep=True,
array_type=VTK_ID_TYPE)
residue_seq.SetName("residue")
output.GetPointData().AddArray(residue_seq)
# for chain
chain = nps.numpy_to_vtk(num_array=molecule.chain,
deep=True,
array_type=VTK_UNSIGNED_CHAR)
chain.SetName("chain")
output.GetPointData().AddArray(chain)
# for secondary structures
s_s = nps.numpy_to_vtk(num_array=secondary_structures,
deep=True,
array_type=VTK_UNSIGNED_CHAR)
s_s.SetName("secondary_structures")
output.GetPointData().AddArray(s_s)
# for secondary structures begin
newarr = np.ones(num_total_atoms)
s_sb = nps.numpy_to_vtk(num_array=newarr,
deep=True,
array_type=VTK_UNSIGNED_CHAR)
s_sb.SetName("secondary_structures_begin")
output.GetPointData().AddArray(s_sb)
# for secondary structures end
newarr = np.ones(num_total_atoms)
s_se = nps.numpy_to_vtk(num_array=newarr,
deep=True,
array_type=VTK_UNSIGNED_CHAR)
s_se.SetName("secondary_structures_end")
output.GetPointData().AddArray(s_se)
# for is_hetatm
is_hetatm = nps.numpy_to_vtk(num_array=molecule.is_hetatm,
deep=True,
array_type=VTK_UNSIGNED_CHAR)
is_hetatm.SetName("ishetatm")
output.GetPointData().AddArray(is_hetatm)
# for model
model = nps.numpy_to_vtk(num_array=molecule.model,
deep=True,
array_type=VTK_UNSIGNED_INT)
model.SetName("model")
output.GetPointData().AddArray(model)
table = PTable()
# for colors and radii of hetero-atoms
radii = np.ones((num_total_atoms, 3))
rgb = np.ones((num_total_atoms, 3))
for i in range(num_total_atoms):
radii[i] = np.repeat(table.atomic_radius(all_atomic_numbers[i], 'VDW'),
3)
rgb[i] = table.atom_color(all_atomic_numbers[i])
Rgb = nps.numpy_to_vtk(num_array=rgb,
deep=True,
array_type=VTK_UNSIGNED_CHAR)
Rgb.SetName("rgb_colors")
output.GetPointData().SetScalars(Rgb)
Radii = nps.numpy_to_vtk(num_array=radii, deep=True, array_type=VTK_FLOAT)
Radii.SetName("radius")
output.GetPointData().SetVectors(Radii)
# setting the coordinates
points = numpy_to_vtk_points(coords)
output.SetPoints(points)
ribbonFilter = ProteinRibbonFilter()
ribbonFilter.SetInputData(output)
ribbonFilter.SetCoilWidth(0.2)
ribbonFilter.SetDrawSmallMoleculesAsSpheres(0)
mapper = PolyDataMapper()
mapper.SetInputConnection(ribbonFilter.GetOutputPort())
molecule_actor = Actor()
molecule_actor.SetMapper(mapper)
return molecule_actor
def ball_stick(molecule,
colormode='discrete',
atom_scale_factor=0.3,
bond_thickness=0.1,
multiple_bonds=True):
"""Create an actor for ball and stick molecular representation.
Parameters
----------
molecule : Molecule
The molecule to be rendered.
colormode : string, optional
Set the colormode for coloring the atoms. Two valid color modes:
* 'discrete': Atoms and bonds are colored using CPK coloring
convention.
* 'single': All atoms are colored with same color (grey) and all bonds
are colored with same color (dark grey). RGB tuple used for coloring
the atoms when 'single' colormode is selected: (150, 150, 150). RGB
tuple used for coloring the bonds when 'single' colormode is
selected: (50, 50, 50)
Default: 'discrete'
atom_scale_factor : float, optional
Scaling factor. Default: 0.3
bond_thickness : float, optional
Used to manipulate the thickness of bonds (i.e. thickness of tubes
which are used to render bonds)
Default: 0.1 (Optimal range: 0.1 - 0.5).
multiple_bonds : bool, optional
Set whether multiple tubes will be used to represent multiple
bonds. If True, multiple bonds (double, triple) will be shown by using
multiple tubes. If False, all bonds (single, double, triple) will be
shown as single bonds (i.e. shown using one tube each).
Default is True.
Returns
-------
molecule_actor : vtkActor
Actor created to render the ball and stick representation of the
molecule to be visualized.
References
----------
Turner, M. Ball and stick models for organic chemistry
`J. Chem. Educ. 1971, 48, 6, 407.
<https://doi.org/10.1021/ed048p407>`_
"""
if molecule.total_num_bonds == 0:
raise ValueError('No bonding data available for the molecule! Ball '
'and stick model cannot be made!')
colormode = colormode.lower()
bs_mapper = OpenGLMoleculeMapper()
bs_mapper.SetInputData(molecule)
bs_mapper.SetRenderAtoms(True)
bs_mapper.SetRenderBonds(True)
bs_mapper.SetBondRadius(bond_thickness)
bs_mapper.SetAtomicRadiusTypeToVDWRadius()
bs_mapper.SetAtomicRadiusScaleFactor(atom_scale_factor)
if multiple_bonds:
bs_mapper.SetUseMultiCylindersForBonds(1)
else:
bs_mapper.SetUseMultiCylindersForBonds(0)
if colormode == 'discrete':
bs_mapper.SetAtomColorMode(1)
bs_mapper.SetBondColorMode(1)
elif colormode == 'single':
bs_mapper.SetAtomColorMode(0)
bs_mapper.SetBondColorMode(0)
else:
bs_mapper.SetAtomColorMode(1)
warnings.warn("Incorrect colormode specified! Using discrete.")
molecule_actor = Actor()
molecule_actor.SetMapper(bs_mapper)
return molecule_actor