def numpy_to_vtk_cells(data, is_coords=True):
"""Convert numpy array to a vtk cell array.
Parameters
----------
data : ndarray
points coordinate or connectivity array (e.g triangles).
is_coords : ndarray
Select the type of array. default: True.
Returns
-------
vtk_cell : vtkCellArray
connectivity + offset information
"""
data = np.array(data, dtype=object)
nb_cells = len(data)
# Get lines_array in vtk input format
connectivity = data.flatten() if not is_coords else []
offset = [0, ]
current_position = 0
cell_array = CellArray()
if VTK_9_PLUS:
for i in range(nb_cells):
current_len = len(data[i])
offset.append(offset[-1] + current_len)
if is_coords:
end_position = current_position + current_len
connectivity += list(range(current_position, end_position))
current_position = end_position
connectivity = np.array(connectivity, np.intp)
offset = np.array(offset, dtype=connectivity.dtype)
vtk_array_type = numpy_support.get_vtk_array_type(connectivity.dtype)
cell_array.SetData(
numpy_support.numpy_to_vtk(offset, deep=True,
array_type=vtk_array_type),
numpy_support.numpy_to_vtk(connectivity, deep=True,
array_type=vtk_array_type))
else:
for i in range(nb_cells):
current_len = len(data[i])
end_position = current_position + current_len
connectivity += [current_len]
connectivity += list(range(current_position, end_position))
current_position = end_position
connectivity = np.array(connectivity)
cell_array.GetData().DeepCopy(numpy_support.numpy_to_vtk(connectivity))
cell_array.SetNumberOfCells(nb_cells)
return cell_array
def _points_to_vtk_cells(points, points_per_line=2):
"""
Returns the VTK cell array for the peaks given the set of points
coordinates.
Parameters
----------
points : (N, 3) array or ndarray
points coordinates array.
points_per_line : int (1 or 2), optional
number of points per peak direction.
Returns
-------
cell_array : vtkCellArray
connectivity + offset information.
"""
num_pnts = len(points)
num_cells = num_pnts // points_per_line
cell_array = CellArray()
"""
Connectivity is an array that contains the indices of the points that
need to be connected in the visualization. The indices start from 0.
"""
connectivity = np.asarray(list(range(0, num_pnts)), dtype=int)
"""
Offset is an array that contains the indices of the first point of
each line. The indices start from 0 and given the known geometry of
this actor the creation of this array requires a 2 points padding
between indices.
"""
offset = np.asarray(list(range(0, num_pnts + 1, points_per_line)),
dtype=int)
vtk_array_type = numpy_support.get_vtk_array_type(connectivity.dtype)
cell_array.SetData(
numpy_support.numpy_to_vtk(offset,
deep=True,
array_type=vtk_array_type),
numpy_support.numpy_to_vtk(connectivity,
deep=True,
array_type=vtk_array_type))
cell_array.SetNumberOfCells(num_cells)
return cell_array
def attribute_to_actor(actor, arr, attr_name, deep=True):
"""Link a numpy array with vertex attribute.
Parameters
----------
actor : vtkActor
Rendered Object
arr : ndarray
array to link to vertices
attr_name : str
vertex attribute name. the vtk array will take the same name as the
attribute.
deep : bool, optional
If True a deep copy is applied. Otherwise a shallow copy is applied,
by default True
"""
nb_components = arr.shape[1] if arr.ndim > 1 else arr.ndim
vtk_array = numpy_support.numpy_to_vtk(arr, deep=deep)
vtk_array.SetNumberOfComponents(nb_components)
vtk_array.SetName(attr_name)
actor.GetMapper().GetInput().GetPointData().AddArray(vtk_array)
mapper = actor.GetMapper()
mapper.MapDataArrayToVertexAttribute(attr_name, attr_name,
DataObject.FIELD_ASSOCIATION_POINTS,
-1)
def numpy_to_vtk_colors(colors):
"""Convert Numpy color array to a vtk color array.
Parameters
----------
colors: ndarray
Returns
-------
vtk_colors : vtkDataArray
Notes
-----
If colors are not already in UNSIGNED_CHAR you may need to multiply by 255.
Examples
--------
>>> import numpy as np
>>> from fury.utils import numpy_to_vtk_colors
>>> rgb_array = np.random.rand(100, 3)
>>> vtk_colors = numpy_to_vtk_colors(255 * rgb_array)
"""
vtk_colors = numpy_support.numpy_to_vtk(np.asarray(colors),
deep=True,
array_type=VTK_UNSIGNED_CHAR)
return vtk_colors
def test_tangents_from_actor():
my_actor = actor.square(np.array([[0, 0, 0]]))
tangents = tangents_from_actor(my_actor)
npt.assert_equal(tangents, None)
array = np.array([[0, 0, 0], [1, 1, 1], [2, 2, 2], [3, 3, 3]])
my_actor.GetMapper().GetInput().GetPointData().SetTangents(
numpy_support.numpy_to_vtk(array, deep=True, array_type=VTK_FLOAT))
tangents = tangents_from_actor(my_actor)
npt.assert_array_equal(tangents, array)
def test_normals_from_actor():
my_actor = actor.square(np.array([[0, 0, 0]]))
normals = normals_from_actor(my_actor)
npt.assert_equal(normals, None)
array = np.array([[0, 0, 0], [1, 1, 1], [2, 2, 2], [3, 3, 3]])
my_actor.GetMapper().GetInput().GetPointData().SetNormals(
numpy_support.numpy_to_vtk(array, deep=True))
normals = normals_from_actor(my_actor)
npt.assert_array_equal(normals, array)
def test_get_polydata_tangents():
my_polydata = PolyData()
tangents = get_polydata_tangents(my_polydata)
npt.assert_equal(tangents, None)
array = np.array([[0, 0, 0], [1, 1, 1]])
my_polydata.GetPointData().SetTangents(
numpy_support.numpy_to_vtk(array, deep=True, array_type=VTK_FLOAT))
tangents = get_polydata_tangents(my_polydata)
npt.assert_array_equal(tangents, array)
def test_get_polydata_field():
my_polydata = PolyData()
field_data = get_polydata_field(my_polydata, 'Test')
npt.assert_equal(field_data, None)
data = 1
field_name = 'Test'
vtk_data = numpy_support.numpy_to_vtk(data)
vtk_data.SetName(field_name)
my_polydata.GetFieldData().AddArray(vtk_data)
field_data = get_polydata_field(my_polydata, field_name)
npt.assert_equal(field_data, data)
def test_skybox():
# Test scene created without skybox
scene = window.Scene()
npt.assert_warns(UserWarning, scene.skybox)
# Test removing automatically shown skybox
test_tex = Texture()
test_tex.CubeMapOn()
checker_arr = np.array([[1, 0], [0, 1]], dtype=np.uint8) * 255
for i in range(6):
vtk_img = ImageData()
vtk_img.SetDimensions(2, 2, 1)
img_arr = np.zeros((2, 2, 3), dtype=np.uint8)
img_arr[:, :, i // 2] = checker_arr
vtk_arr = numpy_support.numpy_to_vtk(img_arr.reshape((-1, 3),
order='F'))
vtk_arr.SetName('Image')
img_point_data = vtk_img.GetPointData()
img_point_data.AddArray(vtk_arr)
img_point_data.SetActiveScalars('Image')
test_tex.SetInputDataObject(i, vtk_img)
scene = window.Scene(skybox=test_tex)
report = window.analyze_scene(scene)
npt.assert_equal(report.actors, 1)
ss = window.snapshot(scene)
npt.assert_array_equal(ss[75, 75, :], [0, 0, 255])
npt.assert_array_equal(ss[75, 225, :], [0, 0, 0])
scene.yaw(90)
ss = window.snapshot(scene)
npt.assert_array_equal(ss[75, 75, :], [255, 0, 0])
npt.assert_array_equal(ss[75, 225, :], [0, 0, 0])
scene.pitch(90)
ss = window.snapshot(scene)
npt.assert_array_equal(ss[75, 75, :], [0, 0, 0])
npt.assert_array_equal(ss[75, 225, :], [0, 255, 0])
# Test skybox is not added twice
scene.skybox()
report = window.analyze_scene(scene)
npt.assert_equal(report.actors, 1)
# Test make skybox invisible
scene.skybox(visible=False)
report = window.analyze_scene(scene)
npt.assert_equal(report.actors, 0)
ss = window.snapshot(scene)
npt.assert_array_equal(ss[75, 75, :], [0, 0, 0])
scene.yaw(90)
ss = window.snapshot(scene)
npt.assert_array_equal(ss[75, 75, :], [0, 0, 0])
scene.pitch(90)
ss = window.snapshot(scene)
npt.assert_array_equal(ss[75, 225, :], [0, 0, 0])
# Test make skybox visible again
scene.skybox(visible=True)
report = window.analyze_scene(scene)
npt.assert_equal(report.actors, 1)
def set_polydata_normals(polydata, normals):
"""Set polydata normals with a numpy array (ndarrays Nx3 int).
Parameters
----------
polydata : vtkPolyData
normals : normals, represented as 2D ndarrays (Nx3) (one per vertex)
"""
vtk_normals = numpy_support.numpy_to_vtk(normals, deep=True)
polydata.GetPointData().SetNormals(vtk_normals)
return polydata
def set_polydata_vertices(polydata, vertices):
"""Set polydata vertices with a numpy array (ndarrays Nx3 int).
Parameters
----------
polydata : vtkPolyData
vertices : vertices, represented as 2D ndarrays (Nx3)
"""
vtk_points = Points()
vtk_points.SetData(numpy_support.numpy_to_vtk(vertices, deep=True))
polydata.SetPoints(vtk_points)
return polydata
def set_polydata_normals(polydata, normals):
"""Set polydata normals with a numpy array (ndarrays Nx3 int).
Parameters
----------
polydata : vtkPolyData
normals : normals, represented as 2D ndarrays (Nx3) (one per vertex)
"""
vtk_normals = numpy_support.numpy_to_vtk(normals, deep=True)
# VTK does not require a specific name for the normals array, however, for
# readability purposes, we set it to "Normals"
vtk_normals.SetName('Normals')
polydata.GetPointData().SetNormals(vtk_normals)
return polydata
def numpy_to_vtk_points(points):
"""Convert Numpy points array to a vtk points array.
Parameters
----------
points : ndarray
Returns
-------
vtk_points : vtkPoints()
"""
vtk_points = Points()
vtk_points.SetData(
numpy_support.numpy_to_vtk(np.asarray(points), deep=True))
return vtk_points
def set_polydata_colors(polydata, colors, array_name="colors"):
"""Set polydata colors with a numpy array (ndarrays Nx3 int).
Parameters
----------
polydata : vtkPolyData
colors : colors, represented as 2D ndarrays (Nx3)
colors are uint8 [0,255] RGB for each points
"""
vtk_colors = numpy_support.numpy_to_vtk(colors, deep=True,
array_type=VTK_UNSIGNED_CHAR)
nb_components = colors.shape[1]
vtk_colors.SetNumberOfComponents(nb_components)
vtk_colors.SetName(array_name)
polydata.GetPointData().SetScalars(vtk_colors)
return polydata
def set_polydata_tangents(polydata, tangents):
"""Set polydata tangents with a numpy array (ndarrays Nx3 int).
Parameters
----------
polydata : vtkPolyData
tangents : tangents, represented as 2D ndarrays (Nx3) (one per vertex)
"""
vtk_tangents = numpy_support.numpy_to_vtk(tangents,
deep=True,
array_type=VTK_FLOAT)
# VTK does not require a specific name for the tangents array, however, for
# readability purposes, we set it to "Tangents"
vtk_tangents.SetName('Tangents')
polydata.GetPointData().SetTangents(vtk_tangents)
return polydata
def add_polydata_numeric_field(polydata,
field_name,
field_data,
array_type=VTK_INT):
"""Add a field to a vtk polydata.
Parameters
----------
polydata : vtkPolyData
field_name : str
field_data : bool, int, float, double, numeric array or ndarray
array_type : vtkArrayType
"""
vtk_field_data = numpy_support.numpy_to_vtk(field_data,
deep=True,
array_type=array_type)
vtk_field_data.SetName(field_name)
polydata.GetFieldData().AddArray(vtk_field_data)
return polydata
def numpy_to_vtk_image_data(array,
spacing=(1.0, 1.0, 1.0),
origin=(0.0, 0.0, 0.0),
deep=True):
"""Convert numpy array to a vtk image data.
Parameters
----------
array : ndarray
pixel coordinate and colors values.
spacing : (float, float, float) (optional)
sets the size of voxel (unit of space in each direction x,y,z)
origin : (float, float, float) (optional)
sets the origin at the given point
deep : bool (optional)
decides the type of copy(ie. deep or shallow)
Returns
-------
vtk_image : vtkImageData
"""
if array.ndim not in [2, 3]:
raise IOError("only 2D (L, RGB, RGBA) or 3D image available")
vtk_image = ImageData()
depth = 1 if array.ndim == 2 else array.shape[2]
vtk_image.SetDimensions(array.shape[1], array.shape[0], depth)
vtk_image.SetExtent(0, array.shape[1] - 1, 0, array.shape[0] - 1, 0, 0)
vtk_image.SetSpacing(spacing)
vtk_image.SetOrigin(origin)
temp_arr = np.flipud(array)
temp_arr = temp_arr.reshape(array.shape[1] * array.shape[0], depth)
temp_arr = np.ascontiguousarray(temp_arr, dtype=array.dtype)
vtk_array_type = numpy_support.get_vtk_array_type(array.dtype)
uchar_array = numpy_support.numpy_to_vtk(temp_arr,
deep=deep,
array_type=vtk_array_type)
vtk_image.GetPointData().SetScalars(uchar_array)
return vtk_image
def rgb_to_vtk(data):
"""RGB or RGBA images to VTK arrays.
Parameters
----------
data : ndarray
Shape can be (X, Y, 3) or (X, Y, 4)
Returns
-------
vtkImageData
"""
grid = ImageData()
grid.SetDimensions(data.shape[1], data.shape[0], 1)
nd = data.shape[-1]
vtkarr = numpy_support.numpy_to_vtk(
np.flip(data.swapaxes(0, 1), axis=1).reshape((-1, nd), order='F'))
vtkarr.SetName('Image')
grid.GetPointData().AddArray(vtkarr)
grid.GetPointData().SetActiveScalars('Image')
grid.GetPointData().Update()
return grid
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 save_image(arr,
filename,
compression_quality=75,
compression_type='deflation',
use_pillow=True):
"""Save a 2d or 3d image.
Expect an image with the following shape: (H, W) or (H, W, 1) or
(H, W, 3) or (H, W, 4).
Parameters
----------
arr : ndarray
array to save
filename : string
should be png, bmp, jpeg or jpg files
compression_quality : int, optional
compression_quality for jpeg data.
0 = Low quality, 100 = High quality
compression_type : str, optional
compression type for tiff file
select between: None, lzw, deflation (default)
use_pillow : bool, optional
Use imageio python library to save the files.
"""
if arr.ndim > 3:
raise IOError("Image Dimensions should be <=3")
d_writer = {
".png": PNGWriter,
".bmp": BMPWriter,
".jpeg": JPEGWriter,
".jpg": JPEGWriter,
".tiff": TIFFWriter,
".tif": TIFFWriter,
}
extension = os.path.splitext(os.path.basename(filename).lower())[1]
if extension.lower() not in d_writer.keys():
raise IOError(
"Impossible to save the file {0}: Unknown extension {1}".format(
filename, extension))
if use_pillow:
arr = np.flipud(arr)
im = Image.fromarray(arr)
im.save(filename, quality=compression_quality)
return
if arr.ndim == 2:
arr = arr[..., None]
shape = arr.shape
if extension.lower() in [
'.png',
]:
arr = arr.astype(np.uint8)
arr = arr.reshape((shape[1] * shape[0], shape[2]))
arr = np.ascontiguousarray(arr, dtype=arr.dtype)
vtk_array_type = numpy_support.get_vtk_array_type(arr.dtype)
vtk_array = numpy_support.numpy_to_vtk(num_array=arr,
deep=True,
array_type=vtk_array_type)
# Todo, look the following link for managing png 16bit
# https://stackoverflow.com/questions/15667947/vtkpngwriter-printing-out-black-images
vtk_data = ImageData()
vtk_data.SetDimensions(shape[1], shape[0], shape[2])
vtk_data.SetExtent(0, shape[1] - 1, 0, shape[0] - 1, 0, 0)
vtk_data.SetSpacing(1.0, 1.0, 1.0)
vtk_data.SetOrigin(0.0, 0.0, 0.0)
vtk_data.GetPointData().SetScalars(vtk_array)
writer = d_writer.get(extension)()
writer.SetFileName(filename)
writer.SetInputData(vtk_data)
if extension.lower() in [".jpg", ".jpeg"]:
writer.ProgressiveOn()
writer.SetQuality(compression_quality)
if extension.lower() in [".tif", ".tiff"]:
compression_type = compression_type or 'nocompression'
l_compression = ['nocompression', 'packbits', 'jpeg', 'deflate', 'lzw']
if compression_type.lower() in l_compression:
comp_id = l_compression.index(compression_type.lower())
writer.SetCompression(comp_id)
else:
writer.SetCompressionToDeflate()
writer.Write()
def lines_to_vtk_polydata(lines, colors=None):
"""Create a vtkPolyData with lines and colors.
Parameters
----------
lines : list
list of N curves represented as 2D ndarrays
colors : array (N, 3), list of arrays, tuple (3,), array (K,)
If None or False, a standard orientation colormap is used for every
line.
If one tuple of color is used. Then all streamlines will have the same
colour.
If an array (N, 3) is given, where N is equal to the number of lines.
Then every line is coloured with a different RGB color.
If a list of RGB arrays is given then every point of every line takes
a different color.
If an array (K, 3) is given, where K is the number of points of all
lines then every point is colored with a different RGB color.
If an array (K,) is given, where K is the number of points of all
lines then these are considered as the values to be used by the
colormap.
If an array (L,) is given, where L is the number of streamlines then
these are considered as the values to be used by the colormap per
streamline.
If an array (X, Y, Z) or (X, Y, Z, 3) is given then the values for the
colormap are interpolated automatically using trilinear interpolation.
Returns
-------
poly_data : vtkPolyData
color_is_scalar : bool, true if the color array is a single scalar
Scalar array could be used with a colormap lut
None if no color was used
"""
# Get the 3d points_array
points_array = np.vstack(lines)
# Set Points to vtk array format
vtk_points = numpy_to_vtk_points(points_array)
# Set Lines to vtk array format
vtk_cell_array = numpy_to_vtk_cells(lines)
# Create the poly_data
poly_data = PolyData()
poly_data.SetPoints(vtk_points)
poly_data.SetLines(vtk_cell_array)
# Get colors_array (reformat to have colors for each points)
# - if/else tested and work in normal simple case
nb_points = len(points_array)
nb_lines = len(lines)
lines_range = range(nb_lines)
points_per_line = [len(lines[i]) for i in lines_range]
points_per_line = np.array(points_per_line, np.intp)
color_is_scalar = False
if colors is None or colors is False:
# set automatic rgb colors
cols_arr = line_colors(lines)
colors_mapper = np.repeat(lines_range, points_per_line, axis=0)
vtk_colors = numpy_to_vtk_colors(255 * cols_arr[colors_mapper])
else:
cols_arr = np.asarray(colors)
if cols_arr.dtype == object: # colors is a list of colors
vtk_colors = numpy_to_vtk_colors(255 * np.vstack(colors))
else:
if len(cols_arr) == nb_points:
if cols_arr.ndim == 1: # values for every point
vtk_colors = numpy_support.numpy_to_vtk(cols_arr,
deep=True)
color_is_scalar = True
elif cols_arr.ndim == 2: # map color to each point
vtk_colors = numpy_to_vtk_colors(255 * cols_arr)
elif cols_arr.ndim == 1:
if len(cols_arr) == nb_lines: # values for every streamline
cols_arrx = []
for (i, value) in enumerate(colors):
cols_arrx += lines[i].shape[0] * [value]
cols_arrx = np.array(cols_arrx)
vtk_colors = numpy_support.numpy_to_vtk(cols_arrx,
deep=True)
color_is_scalar = True
else: # the same colors for all points
vtk_colors = numpy_to_vtk_colors(
np.tile(255 * cols_arr, (nb_points, 1)))
elif cols_arr.ndim == 2: # map color to each line
colors_mapper = np.repeat(lines_range, points_per_line, axis=0)
vtk_colors = numpy_to_vtk_colors(255 * cols_arr[colors_mapper])
else: # colormap
# get colors for each vertex
cols_arr = map_coordinates_3d_4d(cols_arr, points_array)
vtk_colors = numpy_support.numpy_to_vtk(cols_arr, deep=True)
color_is_scalar = True
vtk_colors.SetName("colors")
poly_data.GetPointData().SetScalars(vtk_colors)
return poly_data, color_is_scalar
def _peaks_colors_from_points(points, colors=None, points_per_line=2):
"""
Returns a VTK scalar array containing colors information for each one of
the peaks according to the policy defined by the parameter colors.
Parameters
----------
points : (N, 3) array or ndarray
points coordinates array.
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.
points_per_line : int (1 or 2), optional
number of points per peak direction.
Returns
-------
color_array : vtkDataArray
vtk scalar array with name 'colors'.
colors_are_scalars : bool
indicates whether or not the colors are scalars to be interpreted by a
colormap.
global_opacity : float
returns 1 if the colors array doesn't contain opacity otherwise -1.
"""
num_pnts = len(points)
num_lines = num_pnts // points_per_line
colors_are_scalars = False
global_opacity = 1
if colors is None or colors == 'rgb_standard':
# Automatic RGB colors
colors = np.asarray((0, 0, 0))
color_array = numpy_to_vtk_colors(np.tile(255 * colors, (num_pnts, 1)))
elif type(colors) is tuple:
global_opacity = 1 if len(colors) == 3 else -1
colors = np.asarray(colors)
color_array = numpy_to_vtk_colors(np.tile(255 * colors, (num_pnts, 1)))
else:
colors = np.asarray(colors)
if len(colors) == num_lines:
pnts_colors = np.repeat(colors, points_per_line, axis=0)
if colors.ndim == 1: # Scalar per line
color_array = numpy_support.numpy_to_vtk(pnts_colors,
deep=True)
colors_are_scalars = True
elif colors.ndim == 2: # RGB(A) color per line
global_opacity = 1 if colors.shape[1] == 3 else -1
color_array = numpy_to_vtk_colors(255 * pnts_colors)
elif len(colors) == num_pnts:
if colors.ndim == 1: # Scalar per point
color_array = numpy_support.numpy_to_vtk(colors, deep=True)
colors_are_scalars = True
elif colors.ndim == 2: # RGB(A) color per point
global_opacity = 1 if colors.shape[1] == 3 else -1
color_array = numpy_to_vtk_colors(255 * colors)
color_array.SetName('colors')
return color_array, colors_are_scalars, global_opacity
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,
atomic_numbers=None,
coords=None,
atom_names=None,
model=None,
residue_seq=None,
chain=None,
sheet=None,
helix=None,
is_hetatm=None):
"""Send the atomic data to the molecule.
Parameters
----------
atomic_numbers : ndarray of integers, optional
The shape of the array must be (N, ) where N is the total number of
atoms present in the molecule.
Array having atomic number corresponding to each atom of the
molecule.
coords : ndarray of floats, optional
The shape of the array must be (N, 3) where N is the total number
of atoms present in the molecule.
Array having coordinates corresponding to each atom of the
molecule.
atom_names : ndarray of strings, optional
The shape of the array must be (N, ) where N is the total number of
atoms present in the molecule.
Array having the names of atoms.
model : ndarray of integers, optional
The shape of the array must be (N, ) where N is the total number of
atoms present in the molecule.
Array having the model number corresponding to each atom.
residue_seq : ndarray of integers, optional
The shape of the array must be (N, ) where N is the total number of
atoms present in the molecule.
Array having the residue sequence number corresponding to each atom
of the molecule.
chain : ndarray of integers, optional
The shape of the array must be (N, ) where N is the total number of
atoms present in the molecule.
Array having the chain number corresponding to each atom.
sheet : ndarray of integers, optional
The shape of the array must be (S, 4) where S is the total number
of sheets present in the molecule.
Array containing information about sheets present in the molecule.
helix : ndarray of integers, optional
The shape of the array must be (H, 4) where H is the total number
of helices present in the molecule.
Array containing information about helices present in the molecule.
is_hetatm : ndarray of bools, optional
The shape of the array must be (N, ) where N is the total number of
atoms present in the molecule.
Array containing a bool value to indicate if an atom is a
heteroatom.
"""
if atomic_numbers is None and coords is None:
self.Initialize()
elif not isinstance(atomic_numbers, np.ndarray) \
or \
not isinstance(coords, np.ndarray):
raise ValueError('atom_types and coords must be numpy arrays.')
elif len(atomic_numbers) == len(coords):
self.atom_names = atom_names
self.model = model
self.residue_seq = residue_seq
self.chain = chain
self.sheet = sheet
self.helix = helix
self.is_hetatm = is_hetatm
coords = numpy_to_vtk_points(coords)
atom_nums = nps.numpy_to_vtk(atomic_numbers,
array_type=VTK_UNSIGNED_SHORT)
atom_nums.SetName("Atomic Numbers")
fieldData = DataSetAttributes()
fieldData.AddArray(atom_nums)
self.Initialize(coords, fieldData)
else:
n1 = len(coords)
n2 = len(atomic_numbers)
raise ValueError('Mismatch in length of atomic_numbers({0}) and '
'length of atomic_coords({1}).'.format(n1, n2))
def load_image(filename, as_vtktype=False, use_pillow=True):
"""Load an image.
Parameters
----------
filename: str
should be png, bmp, jpeg or jpg files
as_vtktype: bool, optional
if True, return vtk output otherwise an ndarray. Default False.
use_pillow: bool, optional
Use pillow python library to load the files. Default True
Returns
-------
image: ndarray or vtk output
desired image array
"""
is_url = filename.lower().startswith('http://') \
or filename.lower().startswith('https://')
if is_url:
image_name = os.path.basename(filename)
if len(image_name.split('.')) < 2:
raise IOError(f'{filename} is not a valid image URL')
urlretrieve(filename, image_name)
filename = image_name
if use_pillow:
with Image.open(filename) as pil_image:
if pil_image.mode in ['RGBA', 'RGB', 'L']:
image = np.asarray(pil_image)
elif pil_image.mode.startswith('I;16'):
raw = pil_image.tobytes('raw', pil_image.mode)
dtype = '>u2' if pil_image.mode.endswith('B') else '<u2'
image = np.frombuffer(raw, dtype=dtype)
image.reshape(pil_image.size[::-1]).astype('=u2')
else:
try:
image = pil_image.convert('RGBA')
except ValueError:
raise RuntimeError('Unknown image mode {}'.format(
pil_image.mode))
image = np.asarray(pil_image)
image = np.flipud(image)
if as_vtktype:
if image.ndim not in [2, 3]:
raise IOError("only 2D (L, RGB, RGBA) or 3D image available")
vtk_image = ImageData()
depth = 1 if image.ndim == 2 else image.shape[2]
# width, height
vtk_image.SetDimensions(image.shape[1], image.shape[0], depth)
vtk_image.SetExtent(0, image.shape[1] - 1, 0, image.shape[0] - 1,
0, 0)
vtk_image.SetSpacing(1.0, 1.0, 1.0)
vtk_image.SetOrigin(0.0, 0.0, 0.0)
image = image.reshape(image.shape[1] * image.shape[0], depth)
image = np.ascontiguousarray(image, dtype=image.dtype)
vtk_array_type = numpy_support.get_vtk_array_type(image.dtype)
uchar_array = numpy_support.numpy_to_vtk(image,
deep=True,
array_type=vtk_array_type)
vtk_image.GetPointData().SetScalars(uchar_array)
image = vtk_image
if is_url:
os.remove(filename)
return image
d_reader = {
".png": PNGReader,
".bmp": BMPReader,
".jpeg": JPEGReader,
".jpg": JPEGReader,
".tiff": TIFFReader,
".tif": TIFFReader
}
extension = os.path.splitext(os.path.basename(filename).lower())[1]
if extension.lower() not in d_reader.keys():
raise IOError(
"Impossible to read the file {0}: Unknown extension {1}".format(
filename, extension))
reader = d_reader.get(extension)()
reader.SetFileName(filename)
reader.Update()
reader.GetOutput().GetPointData().GetArray(0).SetName("original")
if not as_vtktype:
w, h, _ = reader.GetOutput().GetDimensions()
vtk_array = reader.GetOutput().GetPointData().GetScalars()
components = vtk_array.GetNumberOfComponents()
image = numpy_support.vtk_to_numpy(vtk_array).reshape(h, w, components)
if is_url:
os.remove(filename)
return reader.GetOutput() if as_vtktype else image
def test_scene():
scene = window.Scene()
# Scene size test
npt.assert_equal(scene.size(), (0, 0))
# Color background test
# Background color for scene (1, 0.5, 0). 0.001 added here to remove
# numerical errors when moving from float to int values
bg_float = (1, 0.501, 0)
# That will come in the image in the 0-255 uint scale
bg_color = tuple((np.round(255 * np.array(bg_float))).astype('uint8'))
scene.background(bg_float)
arr = window.snapshot(scene)
report = window.analyze_snapshot(arr, bg_color=bg_color,
colors=[bg_color, (0, 127, 0)])
npt.assert_equal(report.objects, 0)
npt.assert_equal(report.colors_found, [True, False])
# Add actor to scene to test the remove actor function by analyzing a
# snapshot
axes = actor.axes()
scene.add(axes)
arr = window.snapshot(scene)
report = window.analyze_snapshot(arr, bg_color)
npt.assert_equal(report.objects, 1)
# Test remove actor function by analyzing a snapshot
scene.rm(axes)
arr = window.snapshot(scene)
report = window.analyze_snapshot(arr, bg_color)
npt.assert_equal(report.objects, 0)
# Add actor to scene to test the remove all actors function by analyzing a
# snapshot
scene.add(axes)
arr = window.snapshot(scene)
report = window.analyze_snapshot(arr, bg_color)
npt.assert_equal(report.objects, 1)
# Test remove all actors function by analyzing a snapshot
scene.rm_all()
arr = window.snapshot(scene)
report = window.analyze_snapshot(arr, bg_color)
npt.assert_equal(report.objects, 0)
# Test change background color from scene by analyzing the scene
ren2 = window.Scene(bg_float)
ren2.background((0, 0, 0.))
report = window.analyze_scene(ren2)
npt.assert_equal(report.bg_color, (0, 0, 0))
# Add actor to scene to test the remove actor function by analyzing the
# scene
ren2.add(axes)
report = window.analyze_scene(ren2)
npt.assert_equal(report.actors, 1)
# Test remove actor function by analyzing the scene
ren2.rm(axes)
report = window.analyze_scene(ren2)
npt.assert_equal(report.actors, 0)
# Test camera information retrieving
with captured_output() as (out, err):
scene.camera_info()
npt.assert_equal(out.getvalue().strip(),
'# Active Camera\n '
'Position (0.00, 0.00, 1.00)\n '
'Focal Point (0.00, 0.00, 0.00)\n '
'View Up (0.00, 1.00, 0.00)')
npt.assert_equal(err.getvalue().strip(), '')
# Test skybox
scene = window.Scene()
npt.assert_equal(scene.GetUseImageBasedLighting(), False)
npt.assert_equal(scene.GetAutomaticLightCreation(), 1)
npt.assert_equal(scene.GetSphericalHarmonics(), None)
npt.assert_equal(scene.GetEnvironmentTexture(), None)
test_tex = Texture()
scene = window.Scene(skybox=test_tex)
npt.assert_equal(scene.GetUseImageBasedLighting(), True)
npt.assert_equal(scene.GetAutomaticLightCreation(), 0)
npt.assert_equal(scene.GetSphericalHarmonics(), None)
npt.assert_equal(scene.GetEnvironmentTexture(), test_tex)
# Test automatically shown skybox
test_tex = Texture()
test_tex.CubeMapOn()
checker_arr = np.array([[1, 0], [0, 1]], dtype=np.uint8) * 255
for i in range(6):
vtk_img = ImageData()
vtk_img.SetDimensions(2, 2, 1)
img_arr = np.zeros((2, 2, 3), dtype=np.uint8)
img_arr[:, :, i // 2] = checker_arr
vtk_arr = numpy_support.numpy_to_vtk(img_arr.reshape((-1, 3),
order='F'))
vtk_arr.SetName('Image')
img_point_data = vtk_img.GetPointData()
img_point_data.AddArray(vtk_arr)
img_point_data.SetActiveScalars('Image')
test_tex.SetInputDataObject(i, vtk_img)
scene = window.Scene(skybox=test_tex)
report = window.analyze_scene(scene)
npt.assert_equal(report.actors, 1)
ss = window.snapshot(scene)
npt.assert_array_equal(ss[75, 75, :], [0, 0, 255])
npt.assert_array_equal(ss[75, 225, :], [0, 0, 0])
scene.yaw(90)
ss = window.snapshot(scene)
npt.assert_array_equal(ss[75, 75, :], [255, 0, 0])
npt.assert_array_equal(ss[75, 225, :], [0, 0, 0])
scene.pitch(90)
ss = window.snapshot(scene)
npt.assert_array_equal(ss[75, 75, :], [0, 0, 0])
npt.assert_array_equal(ss[75, 225, :], [0, 255, 0])
def _points_to_vtk_cells(points, points_per_line=2):
"""
Returns the VTK cell array for the peaks given the set of points
coordinates.
Parameters
----------
points : (N, 3) array or ndarray
points coordinates array.
points_per_line : int (1 or 2), optional
number of points per peak direction.
Returns
-------
cell_array : vtkCellArray
connectivity + offset information.
"""
num_pnts = len(points)
num_cells = num_pnts // points_per_line
cell_array = CellArray()
if VTK_9_PLUS:
"""
Connectivity is an array that contains the indices of the points that
need to be connected in the visualization. The indices start from 0.
"""
connectivity = np.asarray(list(range(0, num_pnts)), dtype=int)
"""
Offset is an array that contains the indices of the first point of
each line. The indices start from 0 and given the known geometry of
this actor the creation of this array requires a 2 points padding
between indices.
"""
offset = np.asarray(list(range(0, num_pnts + 1, points_per_line)),
dtype=int)
vtk_array_type = numpy_support.get_vtk_array_type(connectivity.dtype)
cell_array.SetData(
numpy_support.numpy_to_vtk(offset,
deep=True,
array_type=vtk_array_type),
numpy_support.numpy_to_vtk(connectivity,
deep=True,
array_type=vtk_array_type))
else:
connectivity = np.array([], dtype=int)
i_pos = 0
while i_pos < num_pnts:
e_pos = i_pos + points_per_line
"""
In old versions of VTK (<9.0) the connectivity array should include
the length of each line and immediately after the indices of the
points in each line.
"""
connectivity = np.append(connectivity,
[points_per_line, i_pos, e_pos - 1])
i_pos = e_pos
cell_array.GetData().DeepCopy(numpy_support.numpy_to_vtk(connectivity))
cell_array.SetNumberOfCells(num_cells)
return cell_array