Source code for PVGeo.grids.subset
__all__ = [
'ExtractTopography',
]
__displayname__ = 'Subsetting'
import numpy as np
import pyvista as pv
import vtk
from vtk.numpy_interface import dataset_adapter as dsa
from .. import interface
from ..base import FilterBase
# NOTE: internal import - from scipy.spatial import cKDTree
###############################################################################
[docs]
class ExtractTopography(FilterBase):
"""This filter takes two inputs: any mesh dataset and a set of points for
a topography source. This will add a boolean data array to the cell data of
the input grid on whether that cell should be active (under topographic
layer). A user can also choose to directly extract the data rather than
appending a boolean scalar array via the ``remove`` argument.
Args:
op (str, int, or callable): The operation as a string key, int
index, or callable method
tolerance (float): buffer around the topography surface to include as
part of the decision boundary
offset (float): static value to shift the reference topography surface
ivert (bool): optional to invert the extraction.
remove (bool): Optional parameter to apply a thresholding filter and
return a ``vtkUnstructuredGrid`` object with only the extracted
cells. The ``remove`` option is only available in Python
environments (not available in ParaView). The ``remove`` flag must
be set at the time of instantiation of this algorithm.
This does not actually update the algorithm's output data object
but applies a `PyVista` threshold filter to pass a new data object
after calling ``apply``.
Note:
This currently ignores time varying inputs. We can implement time
variance but need to think about how we would like to do that. Should
the topography surface be static and the volumetric data have time
variance?
"""
__displayname__ = 'Extract Topography'
__category__ = 'filter'
def __init__(
self, op='underneath', tolerance=0.001, offset=0.0, invert=False, remove=False
):
FilterBase.__init__(self, nInputPorts=2, inputType='vtkDataSet', nOutputPorts=1)
self._tolerance = tolerance
self._offset = offset
self._invert = invert
self._remove = remove
self._operation = self._underneath
self.set_operation(op)
# CRITICAL for multiple input ports
[docs]
def FillInputPortInformation(self, port, info):
"""This simply makes sure the user selects the correct inputs"""
typ = 'vtkDataSet'
if port == 1:
typ = 'vtkPointSet' # Make sure topography is some sort of point set
info.Set(self.INPUT_REQUIRED_DATA_TYPE(), typ)
return 1
# THIS IS CRUCIAL to preserve data type through filter
[docs]
def RequestDataObject(self, request, inInfo, outInfo):
"""Constructs the output data object based on the input data object"""
self.OutputType = self.GetInputData(inInfo, 0, 0).GetClassName()
self.FillOutputPortInformation(0, outInfo.GetInformationObject(0))
return 1
#### Extraction Methods ####
[docs]
@staticmethod
def _query(topo_points, data_points):
"""Queries the data points for their closest point on the topography
surface"""
try:
# sklearn's KDTree is faster: use it if available
from sklearn.neighbors import KDTree as Tree
except ImportError:
from scipy.spatial import cKDTree as Tree
tree = Tree(topo_points)
i = tree.query(data_points)[1].ravel()
return topo_points[i]
[docs]
@staticmethod
def _underneath(topo_points, data_points, tolerance):
"""Extract cells underneath the topography surface"""
comp = ExtractTopography._query(topo_points, data_points)
return np.array(data_points[:, 2] < (comp[:, 2] - tolerance), dtype=int)
[docs]
@staticmethod
def _intersection(topo_points, data_points, tolerance):
"""Extract cells intersecting the topography surface"""
comp = ExtractTopography._query(topo_points, data_points)
return np.array(np.abs((data_points[:, 2] - comp[:, 2])) < tolerance, dtype=int)
[docs]
@staticmethod
def get_operations():
"""Returns the extraction operation methods as callable objects in a
dictionary
"""
ops = dict(
underneath=ExtractTopography._underneath,
intersection=ExtractTopography._intersection,
)
return ops
[docs]
@staticmethod
def get_operation_names():
"""Gets a list of the extraction operation keys
Return:
list(str): the keys for getting the operations
"""
ops = ExtractTopography.get_operations()
return list(ops.keys())
[docs]
@staticmethod
def get_operation(idx):
"""Gets a extraction operation based on an index in the keys
Return:
callable: the operation method
"""
if isinstance(idx, str):
return ExtractTopography.get_operations()[idx]
n = ExtractTopography.get_operation_names()[idx]
return ExtractTopography.get_operations()[n]
#### Pipeline Methods ####
[docs]
def RequestData(self, request, inInfo, outInfo):
"""Used by pipeline to generate output"""
# Get input/output of Proxy
igrid = self.GetInputData(inInfo, 0, 0) # Port 0: grid
topo = self.GetInputData(inInfo, 1, 0) # Port 1: topography
grid = self.GetOutputData(outInfo, 0)
grid.DeepCopy(igrid)
# Perform task
ncells = igrid.GetNumberOfCells()
active = np.zeros((ncells), dtype=int)
# Now iterate through the cells in the grid and test if they are beneath the topography
wtopo = dsa.WrapDataObject(topo)
topo_points = np.array(wtopo.Points) # mak sure we do not edit the input
# shift the topography points for the tree
topo_points[:, 2] = topo_points[:, 2] + self._offset
filt = vtk.vtkCellCenters()
filt.SetInputDataObject(igrid)
filt.Update()
data_points = dsa.WrapDataObject(filt.GetOutput(0)).Points
active = self._operation(topo_points, data_points, self._tolerance)
if self._invert:
# NOTE: assumes the given operation produces zeros and ones only
active = 1 - active
# Now add cell data to output
active = interface.convert_array(active, name='Extracted')
grid.GetCellData().AddArray(active)
return 1
[docs]
def apply(self, data, points):
"""Run the algorithm on the input data using the topography points"""
self.SetInputDataObject(0, data)
self.SetInputDataObject(1, points)
self.Update()
output = pv.wrap(self.GetOutput())
if self._remove:
# NOTE: Assumes the given operation produces zeros and ones only
# Also, this does not update the algorithm's output.
# This only sends a new thresholded dataset to the user.
return output.threshold(0.5, scalars='Extracted')
return output
#### Setters/Getters ####
[docs]
def set_tolerance(self, tol):
"""Set the tolerance threshold for the querry"""
if self._tolerance != tol:
self._tolerance = tol
self.Modified()
[docs]
def get_tolerance(self):
"""Get the tolerance threshold for the querry"""
return self._tolerance
[docs]
def set_offset(self, offset):
"""Sets how far off (in Z dir) to slice the data"""
if self._offset != offset:
self._offset = offset
self.Modified()
[docs]
def set_invert(self, flag):
"""Sets the boolean flag on whether to invert the extraction."""
if self._invert != flag:
self._invert = flag
self.Modified()
[docs]
def set_operation(self, op):
"""Set the type of extraction to perform.
Args:
op (str, int, or callable): The operation as a string key, int
index, or callable method
Note:
This can accept a callable method to set a custom operation as long
as its signature is ``<callable>(self, topo_points, data_points)`` and it
strictly produces an integer array of zeros and ones.
"""
if isinstance(op, str):
op = ExtractTopography.get_operations()[op]
elif isinstance(op, int):
op = ExtractTopography.get_operation(op)
if self._operation != op:
self._operation = op
self.Modified()