Extending PyVista

A pyvista.DataSet, such as pyvista.PolyData, can be extended by users. For example, if the user wants to keep track of the location of the maximum point in the (1, 0, 1) direction on the mesh.

There are two methods by which users can handle subclassing. One is directly managing the types objects. This may require checking types during filter operations.

The second is automatic managing of types. Users can control whether user defined classes are nearly always used for particular types of DataSets.

Note

This is for advanced usage only. Automatic managing of types will not work in all situations, in particular when a builtin dataset is directly instantiated. See examples below.

from __future__ import annotations

import numpy as np
import vtk

import pyvista

pyvista.set_plot_theme('document')

A user defined subclass of pyvista.PolyData, FooData is defined. It includes a property to keep track of the point on the mesh that is furthest along in the (1, 0, 1) direction.

class FooData(pyvista.PolyData):  # noqa: D101
    @property
    def max_point(self):
        """Returns index of point that is furthest along (1, 0, 1) direction."""
        return np.argmax(np.dot(self.points, (1.0, 0.0, 1.0)))

Directly Managing Types

Now a foo_sphere object is created of type FooData. The index of the point and location of the point of interest can be obtained directly. The sphere has a radius of 0.5, so the maximum extent in the direction (1, 0, 1) is \(0.5\sqrt{0.5}\approx0.354\)

foo_sphere = FooData(pyvista.Sphere(theta_resolution=100, phi_resolution=100))
print('Original foo sphere:')
print(f'Type: {type(foo_sphere)}')
print(f'Maximum point index: {foo_sphere.max_point}')
print(f'Location of maximum point: {foo_sphere.points[foo_sphere.max_point, :]}')

Original foo sphere:
Type: <class '__main__.FooData'>
Maximum point index: 26
Location of maximum point: [0.35634708 0.         0.35073745]

Using an inplace operation like pyvista.DataSet.rotate_y() does not affect the type of the object.

foo_sphere.rotate_y(90, inplace=True)
print('\nRotated foo sphere:')
print(f'Type: {type(foo_sphere)}')
print(f'Maximum point index: {foo_sphere.max_point}')
print(f'Location of maximum point: {foo_sphere.points[foo_sphere.max_point, :]}')

Rotated foo sphere:
Type: <class '__main__.FooData'>
Maximum point index: 4926
Location of maximum point: [ 3.5073745e-01 -1.1460996e-16  3.5634708e-01]

However, filter operations can return different DataSet types including ones that differ from the original type. In this case, the decimate method returns a pyvista.PolyData object.

print('\nDecimated foo sphere:')
decimated_foo_sphere = foo_sphere.decimate(0.5)
print(f'Type: {type(decimated_foo_sphere)}')

Decimated foo sphere:
Type: <class 'pyvista.core.pointset.PolyData'>

It is now required to explicitly wrap the object into FooData.

decimated_foo_sphere = FooData(foo_sphere.decimate(0.5))
print(f'Type: {type(decimated_foo_sphere)}')
print(f'Maximum point index: {decimated_foo_sphere.max_point}')
print(f'Location of maximum point: {foo_sphere.points[foo_sphere.max_point, :]}')

Type: <class '__main__.FooData'>
Maximum point index: 2481
Location of maximum point: [ 3.5073745e-01 -1.1460996e-16  3.5634708e-01]

Automatically Managing Types

The default pyvista.DataSet type can be set using pyvista._wrappers. In general, it is best to use this method when it is expected to primarily use the user defined class.

In this example, all objects that would have been created as pyvista.PolyData would now be created as a FooData object. Note, that the key is the underlying vtk object.

pyvista._wrappers['vtkPolyData'] = FooData

It is no longer necessary to specifically wrap pyvista.PolyData objects to obtain a FooData object.

foo_sphere = pyvista.Sphere(theta_resolution=100, phi_resolution=100)
print('Original foo sphere:')
print(f'Type: {type(foo_sphere)}')
print(f'Maximum point index: {foo_sphere.max_point}')
print(f'Location of maximum point: {foo_sphere.points[foo_sphere.max_point, :]}')

Original foo sphere:
Type: <class '__main__.FooData'>
Maximum point index: 26
Location of maximum point: [0.35634708 0.         0.35073745]

Using an inplace operation like rotate_y does not affect the type of the object.

foo_sphere.rotate_y(90, inplace=True)
print('\nRotated foo sphere:')
print(f'Type: {type(foo_sphere)}')
print(f'Maximum point index: {foo_sphere.max_point}')
print(f'Location of maximum point: {foo_sphere.points[foo_sphere.max_point, :]}')

Rotated foo sphere:
Type: <class '__main__.FooData'>
Maximum point index: 4926
Location of maximum point: [ 3.5073745e-01 -1.1460996e-16  3.5634708e-01]

Filter operations that return pyvista.PolyData now return FooData

print('\nDecimated foo sphere:')
decimated_foo_sphere = foo_sphere.decimate(0.5)
print(f'Type: {type(decimated_foo_sphere)}')
print(f'Maximum point index: {decimated_foo_sphere.max_point}')
print(f'Location of maximum point: {foo_sphere.points[foo_sphere.max_point, :]}')

Decimated foo sphere:
Type: <class '__main__.FooData'>
Maximum point index: 2481
Location of maximum point: [ 3.5073745e-01 -1.1460996e-16  3.5634708e-01]

Users can still create a native pyvista.PolyData object, but using this method may incur unintended consequences. In this case, it is recommended to use the directly managing types method.

poly_object = pyvista.PolyData(vtk.vtkPolyData())
print(f'Type: {type(poly_object)}')
# catch error
try:
    poly_object.rotate_y(90, inplace=True)
except TypeError:
    print('This operation fails')

Type: <class 'pyvista.core.pointset.PolyData'>
This operation fails

Usage of pyvista._wrappers may require resetting the default value to avoid leaking the setting into cases where it is unused.

pyvista._wrappers['vtkPolyData'] = pyvista.PolyData

For instances where a localized usage is preferred, a tear-down method is recommended. One example is a try...finally block.

try:
    pyvista._wrappers['vtkPolyData'] = FooData
    # some operation that sometimes raises an error
finally:
    pyvista._wrappers['vtkPolyData'] = pyvista.PolyData