""" 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 " __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 ("", self.config_data) self.root.bind ("", 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 ()