Chimie Théorique » QMX

import numpy as np

from vtk import vtkPolyDataMapper, vtkPolyDataNormals, \

                vtkLinearSubdivisionFilter, vtkSmoothPolyDataFilter, \

                vtkDepthSortPolyData, vtkAlgorithm, vtkAlgorithmOutput

from ase.visualize.vtk.grid import vtkVolumeGrid

# -------------------------------------------------------------------

class vtkPipeline:

    _branchable = False

    def __init__(self, vtkish_data=None):

        self.vtkish_data = None

        self.closed = False

        self.pending = False

        self.filters = tuple()

        if vtkish_data is not None:

            self.set_data(vtkish_data)

    def hasdata(self):

        return (self.vtkish_data is not None)

    def hasfilters(self):

        return len(self.filters)>0

    def set_data(self, vtkish_data):

        assert vtkish_data is not None

        if self.hasdata():

            raise RuntimeError('Pipeline already has input.')

        elif vtkish_data in self:

            raise ValueError('Pipeline loop detected.')

        if isinstance(vtkish_data, vtkPipeline):

            # The embedded pipeline takes over

            vtkish_data.signal_close()

        self.vtkish_data = vtkish_data

        if not self.hasfilters():

            return

        vtkish_inner = self.filters[0]

        if isinstance(vtkish_inner, vtkPipeline):

            vtkish_inner.set_data(vtkish_data) #TODO does this work?

        else:

            assert isinstance(vtkish_inner, vtkAlgorithm)

            if isinstance(vtkish_data, vtkPipeline):

                # The embedded pipeline takes over

                vtkish_data.connect(vtkish_inner)

            else:

                assert isinstance(vtkish_data, vtkAlgorithm)

                vtkish_inner.SetInputConnection(vtkish_data.GetOutputPort())

    def isempty(self):

        return not self.hasdata() and not self.hasfilters()

    def isclosed(self):

        if self.pending:

            raise RuntimeError('Pipeline output port state is pending.')

        return self.closed

    def signal_close(self):

        if self.closed:

            raise RuntimeError('Pipeline output port is already closed.')

        elif not self._branchable:

            self.pending = True

    def get_output_port(self):

        if self.closed:

            raise RuntimeError('Pipeline output port is closed.')

        elif self.pending:

            self.closed = True

            self.pending = False

        if self.hasfilters():

            vtkish_outer = self.filters[-1]

        elif self.hasdata():

            vtkish_outer = self.vtkish_data

        else:

            raise RuntimeError('Pipeline output port unavailable.')

        if isinstance(vtkish_outer, vtkPipeline):

            return vtkish_outer.get_output_port()

        else:

            assert isinstance(vtkish_outer, vtkAlgorithm)

            return vtkish_outer.GetOutputPort()

    def set_input_connection(self, vtk_port):

        assert isinstance(vtk_port, vtkAlgorithmOutput)

        vtkish_inner = self.filters[0]

        if isinstance(vtkish_inner, vtkPipeline):

            # Connect must be passed down

            vtkish_inner.set_input_connection(vtk_port)

        else:

            vtkish_inner.SetInputConnection(vtk_port)

    def connect(self, vtkish_filter):

        if vtkish_filter in self:

            raise ValueError('Pipeline loop detected.')

        if isinstance(vtkish_filter, vtkPipeline):

            # Connection must be passed down

            if not self.isempty():

                vtkish_filter.set_input_connection(self.get_output_port())

            # The containing pipeline takes over

            vtkish_filter.signal_close()

        elif not self.isempty():

            assert isinstance(vtkish_filter, vtkAlgorithm)

            vtkish_filter.SetInputConnection(self.get_output_port())

    def append(self, vtkish_filter):

        self.connect(vtkish_filter)

        self.filters += (vtkish_filter,)

    def extend(self, vtkish_filters):

        map(self.append, vtkish_filters)

    def __contains__(self, vtkish_candidate):

        if vtkish_candidate == self:

            return True

        if self.hasdata():

            if isinstance(self.vtkish_data, vtkPipeline):

                if vtkish_candidate in self.vtkish_data:

                    return True

            else:

                if vtkish_candidate == self.vtkish_data:

                    return True

        if vtkish_candidate in self.filters:

            return True

        for vtkish_filter in self.filters:

            if isinstance(vtkish_filter, vtkPipeline) \

                and vtkish_candidate in vtkish_filter:

                return True

        return False

    def __getitem__(self, i): #TODO XXX this is a hack

        return self.filters[i]

# -------------------------------------------------------------------

class vtkPolyDataPipeline(vtkPipeline):

    def __init__(self, vtkish_polydata=None):

        vtkPipeline.__init__(self, vtkish_polydata)

    def connect(self, vtkish_filter):

        if isinstance(vtkish_filter, vtkPolyDataMapper):

            self.signal_close()

        vtkPipeline.connect(self, vtkish_filter)

class vtkSurfaceSmootherPipeline(vtkPolyDataPipeline):

    def __init__(self, grid, vtkish_polydata=None, angle=15):

        vtkPolyDataPipeline.__init__(self, vtkish_polydata)

        # Make sure grid argument is correct type

        assert isinstance(grid, vtkVolumeGrid)

        self.grid = grid

        # Split polys with intersection angles greater than angle

        vtk_dnorm = vtkPolyDataNormals()

        vtk_dnorm.SetFeatureAngle(angle)

        vtk_dnorm.SplittingOn()

        vtk_dnorm.ComputeCellNormalsOff()

        vtk_dnorm.ComputePointNormalsOff()

        self.append(vtk_dnorm)

        relax = self.grid.get_relaxation_factor()

        if relax is not None:

            print 'relax=',relax

            #vtk_subdiv = vtkButterflySubdivisionFilter()

            vtk_subdiv = vtkLinearSubdivisionFilter()

            self.append(vtk_subdiv)

            # Smooth out some of the sharp points.

            vtk_smooth = vtkSmoothPolyDataFilter()

            vtk_smooth.SetRelaxationFactor(relax)

            self.append(vtk_smooth)

class vtkDepthSortPipeline(vtkPolyDataPipeline):

    def __init__(self, vtk_renderer, vtkish_polydata=None):

        vtkPolyDataPipeline.__init__(self, vtkish_polydata)

        # The depht sort object is set up to generate scalars representing

        # the sort depth. A camera is assigned for the sorting. The camera

        # defines the sort vector (position and focal point).

        vtk_ds = vtkDepthSortPolyData()

        vtk_ds.SetCamera(vtk_renderer.GetActiveCamera())

        vtk_ds.SetDirectionToBackToFront()

        #vtk_ds.SetVector(1, 1, 1)

        #vtk_ds.SortScalarsOn()

        #vtk_ds.Update()

        self.append(vtk_ds)

        vtk_renderer.ResetCamera()