Chimie Théorique » QMX

import numpy as np

from ase import Atoms

from vtk import vtkOutlineSource, vtkAxesActor, vtkProperty2D, vtkTextProperty

from ase.visualize.vtk.module import vtkModule, vtkPolyDataModule

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

class vtkUnitCellModule(vtkPolyDataModule):

    def __init__(self, atoms):

        assert isinstance(atoms, Atoms)

        self.pbc = atoms.get_pbc()

        cell = atoms.get_cell()

        """

        if not isinstance(cell, np.ndarray):

            cell = np.array(cell)

        if cell.shape == (3,):

            cell = np.diag(cell)

        assert cell.dtype == float and cell.shape == (3, 3)

        """

        self.vtk_outline = vtkOutlineSource()

        if (cell - np.diag(cell.diagonal())).any():

            corners = np.empty((8,3), dtype=float)

            # edges = [map(int,[(i-1)%2==0,i%4>=2,i>=4]) for i in range(8)]

            for c,a in enumerate([(0, 0, 0), (1, 0, 0), (0, 1, 0), (1, 1, 0), \

                                  (0, 0, 1), (1, 0, 1), (0, 1, 1), (1, 1, 1)]):

                corners[c] = np.dot(a, cell)

            self.bbox = np.array(zip(np.min(corners, axis=0), \

                                     np.max(corners, axis=0))).ravel()

            self.vtk_outline.SetCorners(corners.ravel())

            self.vtk_outline.SetBoxTypeToOriented()

        else:

            self.bbox = np.array(zip(np.zeros(3),cell.diagonal())).ravel()

        self.vtk_outline.SetBounds(self.bbox)

        vtkPolyDataModule.__init__(self, self.vtk_outline)

    def get_bounding_box(self):

        return self.bbox

    def get_size(self):

        return self.bbox[1::2]-self.bbox[0::2]

    def get_pbc(self):

        return self.pbc

    def get_characteristic_length(self):

        return np.prod(self.get_size())**(1.0/3.0)

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

class vtkAxesModule(vtkModule):

    def __init__(self, cell):

        assert isinstance(cell, vtkUnitCellModule)

        self.cell = cell

        l0 = self.cell.get_characteristic_length()

        # Create VTK axes actor (not really a VTK actor though)

        self.vtk_ax = vtkAxesActor()

        self.vtk_ax.SetTipTypeToCone()

        self.vtk_ax.SetConeRadius(5e-2*l0)

        self.vtk_ax.SetShaftTypeToCylinder()

        self.vtk_ax.SetCylinderRadius(5e-3*l0)

        # Create VTK two-dimensional property

        p2d = vtkProperty2D()

        p2d.SetDisplayLocationToBackground()

        vtkModule.__init__(self, self.vtk_ax, p2d)

        # Create VTK text property and apply to axes

        vtk_textproperty = vtkTextProperty()

        vtk_textproperty.SetFontSize(14)

        vtk_textproperty.SetBold(True)

        vtk_textproperty.SetItalic(True)

        vtk_textproperty.SetShadow(True)

        vtk_textproperty.SetJustificationToRight()

        vtk_textproperty.SetVerticalJustificationToCentered()

        self.set_text_property(vtk_textproperty)

    def set_actor(self, vtk_act):

        assert isinstance(vtk_act, vtkAxesActor) #fix for non-vtkActor actor

        self.vtk_act = vtk_act

    def set_property(self, vtk_property):

        assert isinstance(vtk_property, vtkProperty2D)

        for vtk_cap in [self.vtk_ax.GetXAxisCaptionActor2D(),

                        self.vtk_ax.GetYAxisCaptionActor2D(),

                        self.vtk_ax.GetZAxisCaptionActor2D()]:

            #vtk_cap.ThreeDimensionalLeaderOn()

            #vtk_cap.LeaderOn()

            vtk_cap.SetProperty(vtk_property)

            vtk_txt = vtk_cap.GetTextActor()

            vtk_txt.SetProperty(vtk_property)

    def set_text_property(self, vtk_textproperty, scaled=False):

        assert isinstance(vtk_textproperty, vtkTextProperty)

        for vtk_cap in [self.vtk_ax.GetXAxisCaptionActor2D(),

                        self.vtk_ax.GetYAxisCaptionActor2D(),

                        self.vtk_ax.GetZAxisCaptionActor2D()]:

            vtk_txt = vtk_cap.GetTextActor()

            vtk_txt.SetScaledText(scaled)

            vtk_txt.SetTextProperty(vtk_textproperty)