"""
This Source object allows one to create a vtkData object
vtkStructuredGrid, vtkStructuredPoints, vtkRectilinearGrid,
vtkPolyData, vtkUnstructuredGrid or a corresponding reader object and
use MayaVi to visualize it. This makes it easy to create data on the
fly and then use MayaVi for visualization. The class handles
save_visualization requests intelligently by dumping the used data to
a file that can be used subsequently.
This code is distributed under the conditions of the BSD license. See
LICENSE.txt for details.
Copyright (c) 2001-2003, Prabhu Ramachandran.
"""
__author__ = "Prabhu Ramachandran <prabhu_r@users.sf.net>"
__version__ = "$Revision: 1.10 $"
__date__ = "$Date: 2005/08/02 18:30:14 $"
import string
import Tkinter
import Base.Objects, Common
import vtk, string
import vtkPipeline.vtkMethodParser
import vtkPipeline.ConfigVtkObj
from VtkXMLDataReader import get_array_type
debug = Common.debug
def _get_attribute_list(data):
""" Gets scalar, vector and tensor information from the given data
(either cell or point data). """
debug ("In _get_attribute_list ()")
attr = {'scalars':[], 'vectors':[], 'tensors':[]}
if data:
n = data.GetNumberOfArrays()
for i in range(n):
arr = data.GetArray(i)
type = get_array_type(arr)
name = arr.GetName()
if not name:
name = "%s%d"%(type[:-1], len(attr[type]))
arr.SetName(name)
if type:
attr[type].extend([name])
return attr
def get_attribute_lists (obj):
"""Gets the scalar, vector and tensor attributes that are
available in the given VTK data object."""
debug ("In get_attribute_lists ()")
point_attr = _get_attribute_list(obj.GetPointData())
cell_attr = _get_attribute_list(obj.GetCellData())
return point_attr, cell_attr
class VtkData (Base.Objects.DataSource):
""" This Source object allows one to create a vtkData object
vtkStructuredGrid, vtkStructuredPoints, vtkRectilinearGrid,
vtkPolyData, vtkUnstructuredGrid or a corresponding reader object
and use MayaVi to visualize it. This makes it easy to create data
on the fly and then use MayaVi for visualization. The class
handles save_visualization requests intelligently by dumping the
used data to a file that can be used subsequently."""
def __init__ (self, renwin=None):
debug ("In VtkData::__init__ ()")
Base.Objects.DataSource.__init__ (self)
self.renwin = renwin
self.data = None
self.scalar_lst = [] # avaliable scalars
self.vector_lst = []
self.tensor_lst = []
self.point_attr = None
self.cell_attr = None
self.file_name = "No file"
self.scalar_var = Tkinter.StringVar ()
self.vector_var = Tkinter.StringVar ()
self.tensor_var = Tkinter.StringVar ()
self.data_types = {'vtkStructuredGrid': 'STRUCTURED_GRID',
'vtkStructuredPoints': 'STRUCTURED_POINTS',
'vtkImageData': 'IMAGE_DATA',
'vtkRectilinearGrid': 'RECTILINEAR_GRID',
'vtkPolyData': 'POLYDATA',
'vtkUnstructuredGrid': 'UNSTRUCTURED_GRID'}
self.writers = {}
self.readers = {}
for x in self.data_types:
if x == 'vtkImageData':
self.writers[x] = vtk.vtkStructuredPointsWriter
self.readers[x] = vtk.vtkStructuredPointsReader
else:
self.writers[x] = getattr(vtk, x + 'Writer')
self.readers[x] = getattr(vtk, x + 'Reader')
self.current_type = ''
def _get_type (self, obj):
debug ("In VtkData::_get_type ()")
self.grid_type = ''
for i in self.data_types.keys ():
if obj.IsA (i):
self.grid_type = self.data_types[i]
self.current_type = i
break
if not self.grid_type:
raise Base.Objects.ParseException, \
"Unknown data type: Known data types are one of: %s."\
%(self.data_types.keys ())
def initialize (self, obj):
"""Initialize the object given a valid object."""
debug ("In VtkData::initialize ()")
Common.state.busy ()
if obj.IsA ('vtkDataReader'):
self.data = obj.GetOutput ()
else:
self.data = obj
self._get_type (self.data)
self.setup_defaults()
self.update_references ()
Common.state.idle ()
def setup_defaults (self):
debug ("In VtkData::setup_defaults ()")
pa, ca = get_attribute_lists(self.data)
self.point_attr, self.cell_attr = pa, ca
self.scalar_lst = pa['scalars'] + ca['scalars']
self.vector_lst = pa['vectors'] + ca['vectors']
self.tensor_lst = pa['tensors'] + ca['tensors']
out = self.data
pd = out.GetPointData()
cd = out.GetCellData()
if pa['scalars']:
pd.SetActiveScalars(self.scalar_lst[0])
elif ca['scalars']:
cd.SetActiveScalars(self.scalar_lst[0])
if pa['vectors']:
pd.SetActiveVectors(self.vector_lst[0])
elif ca['vectors']:
cd.SetActiveVectors(self.vector_lst[0])
if pa['tensors']:
pd.SetActiveTensors(self.tensor_lst[0])
elif ca['tensors']:
cd.SetActiveTensors(self.tensor_lst[0])
self.setup_names ()
def setup_names (self):
debug ("In VtkData::setup_names ()")
out = self.data
pd = out.GetPointData()
cd = out.GetCellData()
if not self.scalar_data_name:
if pd:
s = pd.GetScalars()
if s:
self.scalar_data_name = s.GetName()
if not self.scalar_data_name:
if cd:
s = cd.GetScalars()
if s:
self.scalar_data_name = s.GetName()
self.scalar_var.set(self.scalar_data_name)
if not self.vector_data_name:
if pd:
v = pd.GetVectors()
if v:
self.vector_data_name = v.GetName()
if not self.vector_data_name:
if cd:
v = cd.GetVectors()
if v:
self.vector_data_name = v.GetName()
self.vector_var.set(self.vector_data_name)
if not self.tensor_data_name:
if pd:
t = pd.GetTensors()
if t:
self.tensor_data_name = t.GetName()
if not self.tensor_data_name:
if cd:
t = cd.GetTensors()
if t:
self.tensor_data_name = t.GetName()
self.tensor_var.set(self.tensor_data_name)
def set_scalar_name (self, scalar):
debug ("In VtkData::set_scalar_name ()")
assert (scalar in self.scalar_lst), \
"Sorry, no scalar called %s available"%scalar
self.scalar_data_name = scalar
out = self.data
pd = out.GetPointData()
cd = out.GetCellData()
if scalar in self.point_attr['scalars']:
pd.SetActiveScalars(scalar)
else:
cd.SetActiveScalars(scalar)
self.Update ()
self.force_update()
self.update_references ()
def set_vector_name (self, vector):
debug ("In VtkData::set_vector_name ()")
assert (vector in self.vector_lst), \
"Sorry, no vector called %s available"%vector
self.vector_data_name = vector
out = self.data
pd = out.GetPointData()
cd = out.GetCellData()
if vector in self.point_attr['vectors']:
pd.SetActiveVectors(vector)
else:
cd.SetActiveVectors(vector)
self.Update ()
self.force_update()
self.update_references ()
def set_tensor_name (self, tensor):
debug ("In VtkData::set_tensor_name ()")
assert (tensor in self.tensor_lst), \
"Sorry, no tensor called %s available"%tensor
self.tensor_data_name = tensor
out = self.data
pd = out.GetPointData()
cd = out.GetCellData()
if tensor in self.point_attr['tensors']:
pd.SetActiveTensors(tensor)
else:
cd.SetActiveTensors(tensor)
self.Update ()
self.force_update()
self.update_references ()
def GetOutput (self):
"""Get the Data reader's output. """
debug ("In VtkData::GetOutput ()")
return self.data
def get_output (self):
debug ("In VtkData::get_output ()")
return self.data
def get_scalar_list (self):
debug ("In VtkData::get_scalar_list ()")
return self.scalar_lst
def get_vector_list (self):
debug ("In VtkData::get_vector_list ()")
return self.vector_lst
def get_tensor_list (self):
debug ("In VtkData::get_tensor_list ()")
return self.tensor_lst
def force_update(self):
"""Used to force the pipeline to update. This is necessary
since the consumers do not update when the output is
modified."""
debug ("In VtkData::force_update ()")
out = self.GetOutput()
for i in range(out.GetNumberOfConsumers()):
out.GetConsumer(i).Modified()
def Update (self):
debug ("In VtkData::Update ()")
self.data.Update ()
out = self.GetOutput()
out.Modified()
out.Update()
def get_render_window (self):
debug ("In VtkData::get_render_window ()")
return self.renwin
def save_config (self, file):
"""Save the configuration to file."""
debug ("In VtkData::save_config ()")
# write the data to a vtk file.
filename = file.name + '.vtk'
writer = self.writers[self.current_type]()
writer.SetFileName (filename)
writer.SetInput (self.data)
writer.Write ()
# now save the necessary data to the mayavi file.
file.write ("%s, %s\n"%(self.current_type, filename))
file.write ("%s, %s, %s\n"%(self.scalar_data_name,
self.vector_data_name,
self.tensor_data_name))
p = vtkPipeline.vtkMethodParser.VtkPickler ()
p.dump (self.data, file)
def load_config (self, file):
"""Load the saved objects configuration from a file."""
debug ("In VtkData::load_config ()")
val = string.replace (file.readline ()[:-1], ' ', '')
val = string.split (val, ',')
typ, filename = val
val = string.replace (file.readline ()[:-1], ' ', '')
val = string.split (val, ',')
s_n, v_n, t_n = val
self.scalar_data_name = s_n
self.vector_data_name, self.tensor_data_name = v_n, t_n
# set type info.
self.current_type = typ
self.grid_type = self.data_types[typ]
# read data from vtk file.
reader = self.readers[self.current_type]()
reader.SetFileName (filename)
reader.Update ()
if self.current_type == 'vtkImageData':
self.data = vtk.vtkImageData()
self.data.DeepCopy(reader.GetOutput())
else:
self.data = reader.GetOutput ()
# now read the config for the object
p = vtkPipeline.vtkMethodParser.VtkPickler ()
p.load (self.data, file)
pa, ca = get_attribute_lists(self.data)
self.point_attr, self.cell_attr = pa, ca
self.scalar_lst = pa['scalars'] + ca['scalars']
self.vector_lst = pa['vectors'] + ca['vectors']
self.tensor_lst = pa['tensors'] + ca['tensors']
if s_n:
self.set_scalar_name(s_n)
if v_n:
self.set_vector_name(v_n)
if t_n:
self.set_tensor_name(t_n)
self.setup_names ()
self.update_references ()
def scalar_gui (self, master):
debug ("In VtkData::scalar_gui ()")
if not self.scalar_lst:
return
frame = Tkinter.Frame (master, relief='ridge', bd=2)
frame.pack (side='top')
Tkinter.Label (frame, text="Select Scalar").grid (row=0,
column=0,
sticky='ew')
rw = 1
for sc in self.scalar_lst:
rb = Tkinter.Radiobutton (frame, text=sc,
variable=self.scalar_var, value=sc,
command=self.set_scalar_gui)
rb.grid (row=rw, column=0, sticky='w')
rw = rw + 1
def vector_gui (self, master):
debug ("In VtkData::vector_gui ()")
if not self.vector_lst:
return
frame = Tkinter.Frame (master, relief='ridge', bd=2)
frame.pack (side='top')
Tkinter.Label (frame, text="Select Vector").grid (row=0,
column=0,
sticky='ew')
rw = 1
for vec in self.vector_lst:
rb = Tkinter.Radiobutton (frame, text=vec,
variable=self.vector_var, value=vec,
command=self.set_vector_gui)
rb.grid (row=rw, column=0, sticky='w')
rw = rw + 1
def tensor_gui (self, master):
debug ("In VtkData::tensor_gui ()")
if not self.tensor_lst:
return
frame = Tkinter.Frame (master, relief='ridge', bd=2)
frame.pack (side='top')
Tkinter.Label (frame, text="Select Tensor").grid (row=0,
column=0,
sticky='ew')
rw = 1
for ten in self.tensor_lst:
rb = Tkinter.Radiobutton (frame, text=ten,
variable=self.tensor_var, value=ten,
command=self.set_tensor_gui)
rb.grid (row=rw, column=0, sticky='w')
rw = rw + 1
def make_custom_gui (self):
debug ("In VtkData::make_custom_gui ()")
self.make_main_gui ()
self.make_close_button ()
def make_main_gui (self):
debug ("In VtkData::make_main_gui ()")
frame = Tkinter.Frame (self.root, relief='ridge', bd=2)
frame.pack (side='top')
lab = Tkinter.Label (frame,
text="VTK Data object: %s"%self.current_type)
lab.pack (side='top', fill='both', expand=1)
self.scalar_gui (frame)
self.vector_gui (frame)
self.tensor_gui (frame)
but1 = Tkinter.Button (frame, text="More Config options",
command=self.config_data, underline=1)
but1.pack (side='top', fill='both', expand=1)
b = Tkinter.Button (frame, text="Update", underline=0,
command=self.update_data)
b.pack (side='top', fill='both', expand=1)
self.root.bind ("<Alt-o>", self.config_data)
self.root.bind ("<Alt-u>", self.update_data)
def set_scalar_gui (self, event=None):
debug ("In VtkData::set_scalar_gui ()")
scalar = self.scalar_var.get()
if scalar == self.scalar_data_name:
return
Common.state.busy ()
self.set_scalar_name (scalar)
self.renwin.Render ()
Common.state.idle ()
def set_vector_gui (self, event=None):
debug ("In VtkData::set_vector_gui ()")
vector = self.vector_var.get()
if vector == self.vector_data_name:
return
Common.state.busy ()
self.set_vector_name (vector)
self.renwin.Render ()
Common.state.idle ()
def set_tensor_gui (self, event=None):
debug ("In VtkData::set_tensor_gui ()")
tensor = self.tensor_var.get()
if tensor == self.tensor_data_name:
return
Common.state.busy ()
self.set_tensor_name (tensor)
self.renwin.Render ()
Common.state.idle ()
def config_data (self, event=None):
debug ("In VtkData::config_data()")
c = vtkPipeline.ConfigVtkObj.ConfigVtkObj (self.renwin)
c.configure (self.root, self.data)
def update_data(self, event=None):
""" This looks at the input data again. Useful when you have changed
the data and want the visualization to take notice. Warning:
this will not work properly if you change the file format or
the dataset type."""
debug ("In VtkData::update_data ()")
Common.state.busy()
pa, ca = get_attribute_lists(self.data)
self.point_attr, self.cell_attr = pa, ca
self.scalar_lst = pa['scalars'] + ca['scalars']
self.vector_lst = pa['vectors'] + ca['vectors']
self.tensor_lst = pa['tensors'] + ca['tensors']
sc_lst, vec_lst, ten_lst = self.scalar_lst, \
self.vector_lst, self.tensor_lst
if (self.scalar_data_name and \
(self.scalar_data_name not in sc_lst)) or \
(self.vector_data_name and \
(self.vector_data_name not in vec_lst)) or \
(self.tensor_data_name and \
(self.tensor_data_name not in ten_lst)):
msg = "Warning: currently used data names:\n"\
"scalar name: %s\nvector name: %s\ntensor name: %s\n"\
"do not exist in file %s \n"\
"Setting to defaults!"%(self.scalar_data_name,
self.vector_data_name,
self.tensor_data_name,
self.file_name)
Common.print_err(msg)
self.setup_defaults()
self.force_update()
self.update_references()
else:
sd, vd, td = self.scalar_data_name, self.vector_data_name, \
self.tensor_data_name,
if sd:
self.set_scalar_name(sd)
if vd:
self.set_vector_name(vd)
if td:
self.set_tensor_name(td)
self.renwin.Render()
# the following changes the gui so that any changes in the
# data are reflected in the gui.
geom = self.root.geometry()
master = self.root.master
self.close_gui ()
self.configure (master)
self.root.geometry(geom[string.find(geom, '+'):])
Common.state.idle ()
|