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 __init__(self, odfs, vertices, faces, indices, scale, norm, radial_scale, shape, global_cm, colormap, opacity, affine=None, B=None): self.vertices = vertices self.faces = faces self.odfs = odfs self.indices = indices self.B = B self.radial_scale = radial_scale self.colormap = colormap self.grid_shape = shape self.global_cm = global_cm # declare a mask to be instantiated in slice_along_axis self.mask = None # If a B matrix is given, odfs are expected to # be in SH basis coefficients. if self.B is not None: # In that case, we need to save our normalisation and scale # to apply them after conversion from SH to SF. self.norm = norm self.scale = scale else: # If our input is in SF coefficients, we can normalise and # scale it only once, here. if norm: self.odfs /= np.abs(self.odfs).max(axis=-1, keepdims=True) self.odfs *= scale # Compute world coordinates of an affine is supplied self.affine = affine if self.affine is not None: self.w_verts = self.vertices.dot(affine[:3, :3]) self.w_pos = apply_affine(affine, np.asarray(self.indices).T) # Initialize mapper and slice to the # middle of the volume along Z axis self.mapper = PolyDataMapper() self.SetMapper(self.mapper) self.slice_along_axis(self.grid_shape[-1] // 2) self.set_opacity(opacity)
def get_polymapper_from_polydata(polydata): """Get vtkPolyDataMapper from a vtkPolyData. Parameters ---------- polydata : vtkPolyData Returns ------- poly_mapper : vtkPolyDataMapper """ poly_mapper = set_input(PolyDataMapper(), polydata) poly_mapper.ScalarVisibilityOn() poly_mapper.InterpolateScalarsBeforeMappingOn() poly_mapper.Update() poly_mapper.StaticOn() return poly_mapper
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 __init__(self, directions, indices, values=None, affine=None, colors=None, lookup_colormap=None, linewidth=1, symmetric=True): if affine is not None: w_pos = apply_affine(affine, np.asarray(indices).T) valid_dirs = directions[indices] num_dirs = len(np.nonzero(np.abs(valid_dirs).max(axis=-1) > 0)[0]) pnts_per_line = 2 points_array = np.empty((num_dirs * pnts_per_line, 3)) centers_array = np.empty_like(points_array, dtype=int) diffs_array = np.empty_like(points_array) line_count = 0 for idx, center in enumerate(zip(indices[0], indices[1], indices[2])): if affine is None: xyz = np.asarray(center) else: xyz = w_pos[idx, :] valid_peaks = np.nonzero( np.abs(valid_dirs[idx, :, :]).max(axis=-1) > 0.)[0] for direction in valid_peaks: if values is not None: pv = values[center][direction] else: pv = 1. if symmetric: point_i = directions[center][direction] * pv + xyz point_e = -directions[center][direction] * pv + xyz else: point_i = directions[center][direction] * pv + xyz point_e = xyz diff = point_e - point_i points_array[line_count * pnts_per_line, :] = point_e points_array[line_count * pnts_per_line + 1, :] = point_i centers_array[line_count * pnts_per_line, :] = center centers_array[line_count * pnts_per_line + 1, :] = center diffs_array[line_count * pnts_per_line, :] = diff diffs_array[line_count * pnts_per_line + 1, :] = diff line_count += 1 vtk_points = numpy_to_vtk_points(points_array) vtk_cells = _points_to_vtk_cells(points_array) colors_tuple = _peaks_colors_from_points(points_array, colors=colors) vtk_colors, colors_are_scalars, self.__global_opacity = colors_tuple poly_data = PolyData() poly_data.SetPoints(vtk_points) poly_data.SetLines(vtk_cells) poly_data.GetPointData().SetScalars(vtk_colors) self.__mapper = PolyDataMapper() self.__mapper.SetInputData(poly_data) self.__mapper.ScalarVisibilityOn() self.__mapper.SetScalarModeToUsePointFieldData() self.__mapper.SelectColorArray('colors') self.__mapper.Update() self.SetMapper(self.__mapper) attribute_to_actor(self, centers_array, 'center') attribute_to_actor(self, diffs_array, 'diff') vs_dec_code = import_fury_shader('peak_dec.vert') vs_impl_code = import_fury_shader('peak_impl.vert') fs_dec_code = import_fury_shader('peak_dec.frag') fs_impl_code = import_fury_shader('peak_impl.frag') shader_to_actor(self, 'vertex', decl_code=vs_dec_code, impl_code=vs_impl_code) shader_to_actor(self, 'fragment', decl_code=fs_dec_code) shader_to_actor(self, 'fragment', impl_code=fs_impl_code, block='light') # Color scale with a lookup table if colors_are_scalars: if lookup_colormap is None: lookup_colormap = colormap_lookup_table() self.__mapper.SetLookupTable(lookup_colormap) self.__mapper.UseLookupTableScalarRangeOn() self.__mapper.Update() self.__lw = linewidth self.GetProperty().SetLineWidth(self.__lw) if self.__global_opacity >= 0: self.GetProperty().SetOpacity(self.__global_opacity) self.__min_centers = np.min(indices, axis=1) self.__max_centers = np.max(indices, axis=1) self.__is_range = True self.__low_ranges = self.__min_centers self.__high_ranges = self.__max_centers self.__cross_section = self.__high_ranges // 2 self.__mapper.AddObserver(Command.UpdateShaderEvent, self.__display_peaks_vtk_callback)
class PeakActor(Actor): """VTK actor for visualizing slices of ODF field. Parameters ---------- directions : ndarray Peak directions. The shape of the array should be (X, Y, Z, D, 3). indices : tuple Indices given in tuple(x_indices, y_indices, z_indices) format for mapping 2D ODF array to 3D voxel grid. values : ndarray, optional Peak values. The shape of the array should be (X, Y, Z, D). affine : array, optional 4x4 transformation array from native coordinates to world coordinates. colors : None or string ('rgb_standard') or tuple (3D or 4D) or array/ndarray (N, 3 or 4) or array/ndarray (K, 3 or 4) or array/ndarray(N, ) or array/ndarray (K, ) If None a standard orientation colormap is used for every line. If one tuple of color is used. Then all streamlines will have the same color. If an array (N, 3 or 4) is given, where N is equal to the number of points. Then every point is colored with a different RGB(A) color. If an array (K, 3 or 4) is given, where K is equal to the number of lines. Then every line is colored with a different RGB(A) color. If an array (N, ) is given, where N is the number of points then these are considered as the values to be used by the colormap. If an array (K,) is given, where K is the number of lines then these are considered as the values to be used by the colormap. lookup_colormap : vtkLookupTable, optional Add a default lookup table to the colormap. Default is None which calls :func:`fury.actor.colormap_lookup_table`. linewidth : float, optional Line thickness. Default is 1. symmetric: bool, optional If True, peaks are drawn for both peaks_dirs and -peaks_dirs. Else, peaks are only drawn for directions given by peaks_dirs. Default is True. """ def __init__(self, directions, indices, values=None, affine=None, colors=None, lookup_colormap=None, linewidth=1, symmetric=True): if affine is not None: w_pos = apply_affine(affine, np.asarray(indices).T) valid_dirs = directions[indices] num_dirs = len(np.nonzero(np.abs(valid_dirs).max(axis=-1) > 0)[0]) pnts_per_line = 2 points_array = np.empty((num_dirs * pnts_per_line, 3)) centers_array = np.empty_like(points_array, dtype=int) diffs_array = np.empty_like(points_array) line_count = 0 for idx, center in enumerate(zip(indices[0], indices[1], indices[2])): if affine is None: xyz = np.asarray(center) else: xyz = w_pos[idx, :] valid_peaks = np.nonzero( np.abs(valid_dirs[idx, :, :]).max(axis=-1) > 0.)[0] for direction in valid_peaks: if values is not None: pv = values[center][direction] else: pv = 1. if symmetric: point_i = directions[center][direction] * pv + xyz point_e = -directions[center][direction] * pv + xyz else: point_i = directions[center][direction] * pv + xyz point_e = xyz diff = point_e - point_i points_array[line_count * pnts_per_line, :] = point_e points_array[line_count * pnts_per_line + 1, :] = point_i centers_array[line_count * pnts_per_line, :] = center centers_array[line_count * pnts_per_line + 1, :] = center diffs_array[line_count * pnts_per_line, :] = diff diffs_array[line_count * pnts_per_line + 1, :] = diff line_count += 1 vtk_points = numpy_to_vtk_points(points_array) vtk_cells = _points_to_vtk_cells(points_array) colors_tuple = _peaks_colors_from_points(points_array, colors=colors) vtk_colors, colors_are_scalars, self.__global_opacity = colors_tuple poly_data = PolyData() poly_data.SetPoints(vtk_points) poly_data.SetLines(vtk_cells) poly_data.GetPointData().SetScalars(vtk_colors) self.__mapper = PolyDataMapper() self.__mapper.SetInputData(poly_data) self.__mapper.ScalarVisibilityOn() self.__mapper.SetScalarModeToUsePointFieldData() self.__mapper.SelectColorArray('colors') self.__mapper.Update() self.SetMapper(self.__mapper) attribute_to_actor(self, centers_array, 'center') attribute_to_actor(self, diffs_array, 'diff') vs_dec_code = import_fury_shader('peak_dec.vert') vs_impl_code = import_fury_shader('peak_impl.vert') fs_dec_code = import_fury_shader('peak_dec.frag') fs_impl_code = import_fury_shader('peak_impl.frag') shader_to_actor(self, 'vertex', decl_code=vs_dec_code, impl_code=vs_impl_code) shader_to_actor(self, 'fragment', decl_code=fs_dec_code) shader_to_actor(self, 'fragment', impl_code=fs_impl_code, block='light') # Color scale with a lookup table if colors_are_scalars: if lookup_colormap is None: lookup_colormap = colormap_lookup_table() self.__mapper.SetLookupTable(lookup_colormap) self.__mapper.UseLookupTableScalarRangeOn() self.__mapper.Update() self.__lw = linewidth self.GetProperty().SetLineWidth(self.__lw) if self.__global_opacity >= 0: self.GetProperty().SetOpacity(self.__global_opacity) self.__min_centers = np.min(indices, axis=1) self.__max_centers = np.max(indices, axis=1) self.__is_range = True self.__low_ranges = self.__min_centers self.__high_ranges = self.__max_centers self.__cross_section = self.__high_ranges // 2 self.__mapper.AddObserver(Command.UpdateShaderEvent, self.__display_peaks_vtk_callback) @calldata_type(VTK_OBJECT) def __display_peaks_vtk_callback(self, caller, event, calldata=None): if calldata is not None: calldata.SetUniformi('isRange', self.__is_range) calldata.SetUniform3f('highRanges', self.__high_ranges) calldata.SetUniform3f('lowRanges', self.__low_ranges) calldata.SetUniform3f('crossSection', self.__cross_section) def display_cross_section(self, x, y, z): if self.__is_range: self.__is_range = False self.__cross_section = [x, y, z] def display_extent(self, x1, x2, y1, y2, z1, z2): if not self.__is_range: self.__is_range = True self.__low_ranges = [x1, y1, z1] self.__high_ranges = [x2, y2, z2] @property def cross_section(self): return self.__cross_section @property def global_opacity(self): return self.__global_opacity @global_opacity.setter def global_opacity(self, opacity): self.__global_opacity = opacity self.GetProperty().SetOpacity(self.__global_opacity) @property def high_ranges(self): return self.__high_ranges @property def is_range(self): return self.__is_range @property def low_ranges(self): return self.__low_ranges @property def linewidth(self): return self.__lw @linewidth.setter def linewidth(self, linewidth): self.__lw = linewidth self.GetProperty().SetLineWidth(self.__lw) @property def max_centers(self): return self.__max_centers @property def min_centers(self): return self.__min_centers
class OdfSlicerActor(Actor): """ VTK actor for visualizing slices of ODF field. Parameters ---------- odfs : ndarray SF or SH coefficients 2-dimensional array. vertices: ndarray The sphere vertices used for SH to SF projection. faces: ndarray Indices of sphere vertices forming triangles. Should be ordered clockwise (see fury.utils.fix_winding_order). indices: tuple Indices given in tuple(x_indices, y_indices, z_indices) format for mapping 2D ODF array to 3D voxel grid. scale : float Multiplicative factor to apply to ODF amplitudes. norm : bool Normalize SF amplitudes so that the maximum ODF amplitude per voxel along a direction is 1. radial_scale : bool Scale sphere points by ODF values. global_cm : bool If True the colormap will be applied in all ODFs. If False it will be applied individually at each voxel. colormap : None or str The name of the colormap to use. Matplotlib colormaps are supported (e.g., 'inferno'). If None then a RGB colormap is used. opacity : float Takes values from 0 (fully transparent) to 1 (opaque). affine : array optional 4x4 transformation array from native coordinates to world coordinates. B : ndarray (n_coeffs, n_vertices) Optional SH to SF matrix for projecting `odfs` given in SH coefficents on the `sphere`. If None, then the input is assumed to be expressed in SF coefficients. """ def __init__(self, odfs, vertices, faces, indices, scale, norm, radial_scale, shape, global_cm, colormap, opacity, affine=None, B=None): self.vertices = vertices self.faces = faces self.odfs = odfs self.indices = indices self.B = B self.radial_scale = radial_scale self.colormap = colormap self.grid_shape = shape self.global_cm = global_cm # declare a mask to be instantiated in slice_along_axis self.mask = None # If a B matrix is given, odfs are expected to # be in SH basis coefficients. if self.B is not None: # In that case, we need to save our normalisation and scale # to apply them after conversion from SH to SF. self.norm = norm self.scale = scale else: # If our input is in SF coefficients, we can normalise and # scale it only once, here. if norm: self.odfs /= np.abs(self.odfs).max(axis=-1, keepdims=True) self.odfs *= scale # Compute world coordinates of an affine is supplied self.affine = affine if self.affine is not None: self.w_verts = self.vertices.dot(affine[:3, :3]) self.w_pos = apply_affine(affine, np.asarray(self.indices).T) # Initialize mapper and slice to the # middle of the volume along Z axis self.mapper = PolyDataMapper() self.SetMapper(self.mapper) self.slice_along_axis(self.grid_shape[-1] // 2) self.set_opacity(opacity) def set_opacity(self, opacity): """ Set opacity value of ODFs to display. """ self.GetProperty().SetOpacity(opacity) def display_extent(self, x1, x2, y1, y2, z1, z2): """ Set visible volume from x1 (inclusive) to x2 (inclusive), y1 (inclusive) to y2 (inclusive), z1 (inclusive) to z2 (inclusive). """ mask = np.zeros(self.grid_shape, dtype=bool) mask[x1:x2 + 1, y1:y2 + 1, z1:z2 + 1] = True self.mask = mask self._update_mapper() def slice_along_axis(self, slice_index, axis='zaxis'): """ Slice ODF field at given `slice_index` along axis in ['xaxis', 'yaxis', zaxis']. """ if axis == 'xaxis': self.display_extent(slice_index, slice_index, 0, self.grid_shape[1] - 1, 0, self.grid_shape[2] - 1) elif axis == 'yaxis': self.display_extent(0, self.grid_shape[0] - 1, slice_index, slice_index, 0, self.grid_shape[2] - 1) elif axis == 'zaxis': self.display_extent(0, self.grid_shape[0] - 1, 0, self.grid_shape[1] - 1, slice_index, slice_index) else: raise ValueError('Invalid axis name {0}.'.format(axis)) def display(self, x=None, y=None, z=None): """ Display a slice along x, y, or z axis. """ if x is None and y is None and z is None: self.slice_along_axis(self.grid_shape[2] // 2) elif x is not None: self.slice_along_axis(x, 'xaxis') elif y is not None: self.slice_along_axis(y, 'yaxis') elif z is not None: self.slice_along_axis(z, 'zaxis') def update_sphere(self, vertices, faces, B): """ Dynamically change the sphere used for SH to SF projection. """ if self.B is None: raise ValueError('Can\'t update sphere when using ' 'SF coefficients.') self.vertices = vertices if self.affine is not None: self.w_verts = self.vertices.dot(self.affine[:3, :3]) self.faces = faces self.B = B # draw ODFs with new sphere self._update_mapper() def _update_mapper(self): """ Map vtkPolyData to the actor. """ polydata = PolyData() offsets = self._get_odf_offsets(self.mask) if len(offsets) == 0: self.mapper.SetInputData(polydata) return None sph_dirs = self._get_sphere_directions() sf = self._get_sf(self.mask) all_vertices = self._get_all_vertices(offsets, sph_dirs, sf) all_faces = self._get_all_faces(len(offsets), len(sph_dirs)) all_colors = self._generate_color_for_vertices(sf) # TODO: There is a lot of deep copy here. # Optimize (see viz_network.py example). set_polydata_triangles(polydata, all_faces) set_polydata_vertices(polydata, all_vertices) set_polydata_colors(polydata, all_colors) self.mapper.SetInputData(polydata) def _get_odf_offsets(self, mask): """ Get the position of non-zero voxels inside `mask`. """ if self.affine is not None: return self.w_pos[mask[self.indices]] return np.asarray(self.indices).T[mask[self.indices]] def _get_sphere_directions(self): """ Get the sphere directions onto which is projected the signal. """ if self.affine is not None: return self.w_verts return self.vertices def _get_sf(self, mask): """ Get SF coefficients inside `mask`. """ # when odfs are expressed in SH coefficients if self.B is not None: sf = self.odfs[mask[self.indices]].dot(self.B) # normalisation and scaling is done on SF coefficients if self.norm: sf /= np.abs(sf).max(axis=-1, keepdims=True) return sf * self.scale # when odfs are in SF coefficients, the normalisation and scaling # are done during initialisation. We simply return them: return self.odfs[mask[self.indices]] def _get_all_vertices(self, offsets, sph_dirs, sf): """ Get array of all the vertices of the ODFs to display. """ if self.radial_scale: # apply SF amplitudes to all sphere # directions and offset each voxel return np.tile(sph_dirs, (len(offsets), 1)) * sf.reshape(-1, 1) +\ np.repeat(offsets, len(sph_dirs), axis=0) # return scaled spheres offsetted by `offsets` return np.tile(sph_dirs, (len(offsets), 1)) * self.scale +\ np.repeat(offsets, len(sph_dirs), axis=0) def _get_all_faces(self, nb_odfs, nb_dirs): """ Get array of all the faces of the ODFs to display. """ return np.tile(self.faces, (nb_odfs, 1)) +\ np.repeat(np.arange(nb_odfs) * nb_dirs, len(self.faces))\ .reshape(-1, 1) def _generate_color_for_vertices(self, sf): """ Get array of all vertices colors. """ if self.global_cm: if self.colormap is None: raise IOError("if global_cm=True, colormap must be defined.") else: all_colors = create_colormap(sf.ravel(), self.colormap) * 255 elif self.colormap is not None: if isinstance(self.colormap, str): # Map ODFs values [min, max] to [0, 1] for each ODF range_sf = sf.max(axis=-1) - sf.min(axis=-1) rescaled = sf - sf.min(axis=-1, keepdims=True) rescaled[range_sf > 0] /= range_sf[range_sf > 0][..., None] all_colors =\ create_colormap(rescaled.ravel(), self.colormap) * 255 else: all_colors = np.tile( np.array(self.colormap).reshape(1, 3), (sf.shape[0] * sf.shape[1], 1)) else: all_colors = np.tile(np.abs(self.vertices) * 255, (len(sf), 1)) return all_colors.astype(np.uint8)
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