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):
[docs] """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):
[docs] """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):
[docs] """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):
[docs] """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):
[docs] """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)