Chimie Théorique » QMX

import numpy as np

from vtk import vtkPointData, vtkDataArray, vtkUnstructuredGrid, vtkPoints, \

                vtkIdList, vtkStructuredPoints

from ase.visualize.vtk.cell import vtkUnitCellModule

from ase.visualize.vtk.data import vtkDataArrayFromNumPyBuffer, \

                                   vtkDoubleArrayFromNumPyArray, \

                                   vtkDoubleArrayFromNumPyMultiArray

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

class vtkBaseGrid:

    def __init__(self, npoints, cell):

        self.npoints = npoints

        # Make sure cell argument is correct type

        assert isinstance(cell, vtkUnitCellModule)

        self.cell = cell

        self.vtk_pointdata = None

    def set_point_data(self, vtk_pointdata):

        if self.vtk_pointdata is not None:

            raise RuntimeError('VTK point data already present.')

        assert isinstance(vtk_pointdata, vtkPointData)

        self.vtk_pointdata = vtk_pointdata

        #self.vtk_pointdata.SetCopyScalars(False)

        #self.vtk_pointdata.SetCopyVectors(False)

        #self.vtk_pointdata.SetCopyNormals(False)

    def get_point_data(self):

        if self.vtk_pointdata is None:

            raise RuntimeError('VTK point data missing.')

        return self.vtk_pointdata

    def get_number_of_points(self):

        return self.npoints

    def add_scalar_data_array(self, data, name=None, active=True):

        # Are we converting from NumPy buffer to VTK array?

        if isinstance(data, vtkDataArray):

            vtk_sda = data

        elif isinstance(data, vtkDataArrayFromNumPyBuffer):

            vtk_sda = data.get_output()

        else:

            raise ValueError('Data is not a valid scalar data array.')

        del data

        assert vtk_sda.GetNumberOfComponents() == 1

        assert vtk_sda.GetNumberOfTuples() == self.npoints

        if name is not None:

            vtk_sda.SetName(name)

        # Add VTK array to VTK point data

        self.vtk_pointdata.AddArray(vtk_sda)

        if active:

            self.vtk_pointdata.SetActiveScalars(name)

        return vtk_sda

    def add_vector_data_array(self, data, name=None, active=True):

        # Are we converting from NumPy buffer to VTK array?

        if isinstance(data, vtkDataArray):

            vtk_vda = data

        elif isinstance(data, vtkDataArrayFromNumPyBuffer):

            vtk_vda = data.get_output()

        else:

            raise ValueError('Data is not a valid vector data array.')

        del data

        assert vtk_vda.GetNumberOfComponents() == 3

        assert vtk_vda.GetNumberOfTuples() == self.npoints

        if name is not None:

            vtk_vda.SetName(name)

        # Add VTK array to VTK point data

        self.vtk_pointdata.AddArray(vtk_vda)

        if active:

            self.vtk_pointdata.SetActiveVectors(name)

        return vtk_vda

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

class vtkAtomicPositions(vtkBaseGrid):

    """Provides an interface for adding ``Atoms``-centered data to VTK

    modules. Atomic positions, e.g. obtained using atoms.get_positions(),

    constitute an unstructured grid in VTK, to which scalar and vector

    can be added as point data sets.

    Just like ``Atoms``, instances of ``vtkAtomicPositions`` can be divided

    into subsets, which makes it easy to select atoms and add properties.

    Example:

    >>> cell = vtkUnitCellModule(atoms)

    >>> apos = vtkAtomicPositions(atoms.get_positions(), cell)

    >>> apos.add_scalar_property(atoms.get_charges(), 'charges')

    >>> apos.add_vector_property(atoms.get_forces(), 'forces')

    """

    def __init__(self, pos, cell):

        """Construct basic VTK-representation of a set of atomic positions.

        pos: NumPy array of dtype float and shape ``(n,3)``

            Cartesian positions of the atoms.

        cell: Instance of vtkUnitCellModule of subclass thereof

            Holds information equivalent to that of atoms.get_cell().

        """

        # Make sure position argument is a valid array

        if not isinstance(pos, np.ndarray):

            pos = np.array(pos)

        assert pos.dtype == float and pos.shape[1:] == (3,)

        vtkBaseGrid.__init__(self, len(pos), cell)

        # Convert positions to VTK array

        npy2da = vtkDoubleArrayFromNumPyArray(pos)

        vtk_pda = npy2da.get_output()

        del npy2da

        # Transfer atomic positions to VTK points

        self.vtk_pts = vtkPoints()

        self.vtk_pts.SetData(vtk_pda)

        # Create a VTK unstructured grid of these points

        self.vtk_ugd = vtkUnstructuredGrid()

        self.vtk_ugd.SetWholeBoundingBox(self.cell.get_bounding_box())

        self.vtk_ugd.SetPoints(self.vtk_pts)

        # Extract the VTK point data set

        self.set_point_data(self.vtk_ugd.GetPointData())

    def get_points(self, subset=None):

        """Return (subset of) vtkPoints containing atomic positions.

        subset=None: list of int

            A list of indices into the atomic positions; ignored if None.

        """

        if subset is None:

            return self.vtk_pts

        # Create a list of indices from the subset

        vtk_il = vtkIdList()

        for i in subset:

            vtk_il.InsertNextId(i)

        # Allocate VTK points for subset

        vtk_subpts = vtkPoints()

        vtk_subpts.SetDataType(self.vtk_pts.GetDataType())

        vtk_subpts.SetNumberOfPoints(vtk_il.GetNumberOfIds())

        # Transfer subset of VTK points

        self.vtk_pts.GetPoints(vtk_il, vtk_subpts)

        return vtk_subpts

    def get_unstructured_grid(self, subset=None):

        """Return (subset of) an unstructured grid of the atomic positions.

        subset=None: list of int

            A list of indices into the atomic positions; ignored if None.

        """

        if subset is None:

            return self.vtk_ugd

        # Get subset of VTK points

        vtk_subpts = self.get_points(subset)

        # Create a VTK unstructured grid of these points

        vtk_subugd = vtkUnstructuredGrid()

        vtk_subugd.SetWholeBoundingBox(self.cell.get_bounding_box())

        vtk_subugd.SetPoints(vtk_subpts)

        return vtk_subugd

    def add_scalar_property(self, data, name=None, active=True):

        """Add VTK-representation of scalar data at the atomic positions.

        data: NumPy array of dtype float and shape ``(n,)``

            Scalar values corresponding to the atomic positions.

        name=None: str

            Unique identifier for the scalar data.

        active=True: bool

            Flag indicating whether to use as active scalar data.

        """

        # Make sure data argument is a valid array

        if not isinstance(data, np.ndarray):

            data = np.array(data)

        assert data.dtype == float and data.shape == (self.npoints,)

        # Convert scalar properties to VTK array

        npa2da = vtkDoubleArrayFromNumPyArray(data)

        return vtkBaseGrid.add_scalar_data_array(self, npa2da, name, active)

    def add_vector_property(self, data, name=None, active=True):

        """Add VTK-representation of vector data at the atomic positions.

        data: NumPy array of dtype float and shape ``(n,3)``

            Vector components corresponding to the atomic positions.

        name=None: str

            Unique identifier for the vector data.

        active=True: bool

            Flag indicating whether to use as active vector data.

        """

        # Make sure data argument is a valid array

        if not isinstance(data, np.ndarray):

            data = np.array(data)

        assert data.dtype == float and data.shape == (self.npoints,3,)

        # Convert vector properties to VTK array

        npa2da = vtkDoubleArrayFromNumPyArray(data)

        return vtkBaseGrid.add_vector_data_array(self, npa2da, name, active)

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

class vtkVolumeGrid(vtkBaseGrid):

    def __init__(self, elements, cell, origin=None):

        # Make sure element argument is a valid array

        if not isinstance(elements, np.ndarray):

            elements = np.array(elements)

        assert elements.dtype == int and elements.shape == (3,)

        self.elements = elements

        vtkBaseGrid.__init__(self, np.prod(self.elements), cell)

        # Create a VTK grid of structured points

        self.vtk_spts = vtkStructuredPoints()

        self.vtk_spts.SetWholeBoundingBox(self.cell.get_bounding_box())

        self.vtk_spts.SetDimensions(self.elements)

        self.vtk_spts.SetSpacing(self.get_grid_spacing())

        if origin is not None:

            self.vtk_spts.SetOrigin(origin)

        # Extract the VTK point data set

        self.set_point_data(self.vtk_spts.GetPointData())

    def get_grid_spacing(self):

        # Periodic boundary conditions leave out one boundary along an axis

        # Zero/fixed boundary conditions leave out both boundaries of an axis

        return self.cell.get_size()/(self.elements+1.0-self.cell.get_pbc())

    def get_relaxation_factor(self):

        # The relaxation factor is a floating point value between zero and one.

        # It expresses the need for smoothening (relaxation) e.g. of isosurfaces

        # due to coarse grid spacings. Larger grid spacing -> larger relaxation.

        x = self.get_grid_spacing().mean()/self.cell.get_characteristic_length()

        # The relaxation function f(x) satisfies the following requirements

        # f(x) -> 0 for x -> 0+   and   f(x) -> b for x -> inf

        # f'(x) -> a for x -> 0+  and   f'(x) -> 0 for x -> inf

        # Furthermore, it is a rescaling of arctan, hence we know

        # f(x) = 2 b arctan(a pi x / 2 b) / pi

        # Our reference point is x = r for which medium relaxion is needed

        # f(r) = b/2   <=>   r = 2 b / a pi   <=>   a = 2 b / r pi

        r = 0.025 # corresponding to 0.2 Ang grid spacing in 8 Ang cell

        b = 0.5

        f = 2*b*np.arctan(x/r)/np.pi

        if f > 0.1:

           return f.round(1)

        else:

           return None

    def get_structured_points(self):

        return self.vtk_spts

    def add_scalar_field(self, data, name=None, active=True):

        # Make sure data argument is a valid array

        if not isinstance(data, np.ndarray):

            data = np.array(data)

        assert data.dtype == float and data.shape == tuple(self.elements)

        # Convert scalar field to VTK array

        npa2da = vtkDoubleArrayFromNumPyMultiArray(data[...,np.newaxis])

        return vtkBaseGrid.add_scalar_data_array(self, npa2da, name, active)

    def add_vector_field(self, data, name=None, active=True):

        # Make sure data argument is a valid array

        if not isinstance(data, np.ndarray):

            data = np.array(data)

        assert data.dtype == float and data.shape == tuple(self.elements)+(3,)

        # Convert vector field to VTK array

        npa2da = vtkDoubleArrayFromNumPyMultiArray(data)

        return vtkBaseGrid.add_vector_data_array(self, npa2da, name, active)