# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright 2023-2024 TotalEnergies.
# SPDX-FileContributor: Martin Lemay, Paloma Martinez
import numpy as np
import numpy.typing as npt
from typing import Iterator, List, Sequence, Any, Union
from vtkmodules.util.numpy_support import numpy_to_vtk
from vtkmodules.vtkCommonCore import vtkIdList, vtkPoints, reference
from vtkmodules.vtkCommonDataModel import vtkUnstructuredGrid, vtkMultiBlockDataSet, vtkPolyData, vtkDataSet, vtkDataObject, vtkPlane, vtkCellTypes, vtkIncrementalOctreePointLocator
from vtkmodules.vtkFiltersCore import vtk3DLinearGridPlaneCutter
from geos.mesh.utils.multiblockHelpers import ( getBlockElementIndexesFlatten, getBlockFromFlatIndex )
__doc__ = """
Generic VTK utilities.
These methods include:
- extraction of a surface from a given elevation
- conversion from a list to vtkIdList
- conversion of vtk container into iterable
"""
[docs]
def to_vtk_id_list( data: List[ int ] ) -> vtkIdList:
"""Utility function transforming a list of ids into a vtkIdList.
Args:
data (list[int]): Id list
Returns:
result (vtkIdList): Vtk Id List corresponding to input data
"""
result = vtkIdList()
result.Allocate( len( data ) )
for d in data:
result.InsertNextId( d )
return result
[docs]
def vtk_iter( vtkContainer: vtkIdList | vtkCellTypes ) -> Iterator[ Any ]:
"""Utility function transforming a vtk "container" into an iterable.
Args:
vtkContainer (vtkIdList | vtkCellTypes): A vtk container
Returns:
Iterator[ Any ]: The iterator
"""
if isinstance( vtkContainer, vtkIdList ):
for i in range( vtkContainer.GetNumberOfIds() ):
yield vtkContainer.GetId( i )
elif isinstance( vtkContainer, vtkCellTypes ):
for i in range( vtkContainer.GetNumberOfTypes() ):
yield vtkContainer.GetCellType( i )
[docs]
def getBounds(
input: Union[ vtkUnstructuredGrid,
vtkMultiBlockDataSet ] ) -> tuple[ float, float, float, float, float, float ]:
"""Get bounds of either single of composite data set.
Args:
input (Union[vtkUnstructuredGrid, vtkMultiBlockDataSet]): input mesh
Returns:
tuple[float, float, float, float, float, float]: tuple containing
bounds (xmin, xmax, ymin, ymax, zmin, zmax)
"""
if isinstance( input, vtkMultiBlockDataSet ):
return getMultiBlockBounds( input )
else:
return getMonoBlockBounds( input )
[docs]
def getMonoBlockBounds( input: vtkUnstructuredGrid, ) -> tuple[ float, float, float, float, float, float ]:
"""Get boundary box extrema coordinates for a vtkUnstructuredGrid.
Args:
input (vtkMultiBlockDataSet): input single block mesh
Returns:
tuple[float, float, float, float, float, float]: tuple containing
bounds (xmin, xmax, ymin, ymax, zmin, zmax)
"""
return input.GetBounds()
[docs]
def getMultiBlockBounds( input: vtkMultiBlockDataSet, ) -> tuple[ float, float, float, float, float, float ]:
"""Get boundary box extrema coordinates for a vtkMultiBlockDataSet.
Args:
input (vtkMultiBlockDataSet): input multiblock mesh
Returns:
tuple[float, float, float, float, float, float]: bounds.
"""
xmin, ymin, zmin = 3 * [ np.inf ]
xmax, ymax, zmax = 3 * [ -1.0 * np.inf ]
blockIndexes: list[ int ] = getBlockElementIndexesFlatten( input )
for blockIndex in blockIndexes:
block0: vtkDataObject = getBlockFromFlatIndex( input, blockIndex )
assert block0 is not None, "Mesh is undefined."
block: vtkDataSet = vtkDataSet.SafeDownCast( block0 )
bounds: tuple[ float, float, float, float, float, float ] = block.GetBounds()
xmin = bounds[ 0 ] if bounds[ 0 ] < xmin else xmin
xmax = bounds[ 1 ] if bounds[ 1 ] > xmax else xmax
ymin = bounds[ 2 ] if bounds[ 2 ] < ymin else ymin
ymax = bounds[ 3 ] if bounds[ 3 ] > ymax else ymax
zmin = bounds[ 4 ] if bounds[ 4 ] < zmin else zmin
zmax = bounds[ 5 ] if bounds[ 5 ] > zmax else zmax
return xmin, xmax, ymin, ymax, zmin, zmax
[docs]
def getBoundsFromPointCoords( cellPtsCoord: list[ npt.NDArray[ np.float64 ] ] ) -> Sequence[ float ]:
"""Compute bounding box coordinates of the list of points.
Args:
cellPtsCoord (list[npt.NDArray[np.float64]]): list of points
Returns:
Sequence[float]: bounding box coordinates (xmin, xmax, ymin, ymax, zmin, zmax)
"""
bounds: list[ float ] = [
np.inf,
-np.inf,
np.inf,
-np.inf,
np.inf,
-np.inf,
]
for ptsCoords in cellPtsCoord:
mins: npt.NDArray[ np.float64 ] = np.min( ptsCoords, axis=0 )
maxs: npt.NDArray[ np.float64 ] = np.max( ptsCoords, axis=0 )
for i in range( 3 ):
bounds[ 2 * i ] = float( min( bounds[ 2 * i ], mins[ i ] ) )
bounds[ 2 * i + 1 ] = float( max( bounds[ 2 * i + 1 ], maxs[ i ] ) )
return bounds
[docs]
def createSingleCellMesh( cellType: int, ptsCoord: npt.NDArray[ np.float64 ] ) -> vtkUnstructuredGrid:
"""Create a mesh that consists of a single cell.
Args:
cellType (int): cell type
ptsCoord (1DArray[np.float64]): cell point coordinates
Returns:
vtkUnstructuredGrid: output mesh
"""
nbPoints: int = ptsCoord.shape[ 0 ]
points: npt.NDArray[ np.float64 ] = np.vstack( ( ptsCoord, ) )
# Convert points to vtkPoints object
vtkpts: vtkPoints = vtkPoints()
vtkpts.SetData( numpy_to_vtk( points ) )
# create cells from point ids
cellsID: vtkIdList = vtkIdList()
for j in range( nbPoints ):
cellsID.InsertNextId( j )
# add cell to mesh
mesh: vtkUnstructuredGrid = vtkUnstructuredGrid()
mesh.SetPoints( vtkpts )
mesh.Allocate( 1 )
mesh.InsertNextCell( cellType, cellsID )
return mesh
[docs]
def createMultiCellMesh( cellTypes: list[ int ],
cellPtsCoord: list[ npt.NDArray[ np.float64 ] ],
sharePoints: bool = True ) -> vtkUnstructuredGrid:
"""Create a mesh that consists of multiple cells.
.. WARNING:: the mesh is not check for conformity.
Args:
cellTypes (list[int]): cell type
cellPtsCoord (list[1DArray[np.float64]]): list of cell point coordinates
sharePoints (bool): if True, cells share points, else a new point is created for each cell vertex
Returns:
vtkUnstructuredGrid: output mesh
"""
assert len( cellPtsCoord ) == len( cellTypes ), "The lists of cell types of point coordinates must be of same size."
nbCells: int = len( cellPtsCoord )
mesh: vtkUnstructuredGrid = vtkUnstructuredGrid()
points: vtkPoints
cellVertexMapAll: list[ tuple[ int, ...] ]
points, cellVertexMapAll = createVertices( cellPtsCoord, sharePoints )
assert len( cellVertexMapAll ) == len(
cellTypes ), "The lists of cell types of cell point ids must be of same size."
mesh.SetPoints( points )
mesh.Allocate( nbCells )
# create mesh cells
for cellType, ptsId in zip( cellTypes, cellVertexMapAll, strict=True ):
# create cells from point ids
cellsID: vtkIdList = vtkIdList()
for ptId in ptsId:
cellsID.InsertNextId( ptId )
mesh.InsertNextCell( cellType, cellsID )
return mesh
[docs]
def createVertices( cellPtsCoord: list[ npt.NDArray[ np.float64 ] ],
shared: bool = True ) -> tuple[ vtkPoints, list[ tuple[ int, ...] ] ]:
"""Create vertices from cell point coordinates list.
Args:
cellPtsCoord (list[npt.NDArray[np.float64]]): list of cell point coordinates
shared (bool, optional): If True, collocated points are merged. Defaults to True.
Returns:
tuple[vtkPoints, list[tuple[int, ...]]]: tuple containing points and the
map of cell point ids
"""
# get point bounds
bounds: Sequence[ float ] = getBoundsFromPointCoords( cellPtsCoord )
points: vtkPoints = vtkPoints()
# use point locator to check for colocated points
pointsLocator = vtkIncrementalOctreePointLocator()
pointsLocator.InitPointInsertion( points, bounds )
cellVertexMapAll: list[ tuple[ int, ...] ] = []
ptId: reference = reference( 0 )
ptsCoords: npt.NDArray[ np.float64 ]
for ptsCoords in cellPtsCoord:
cellVertexMap: list[ int ] = []
pt: npt.NDArray[ np.float64 ] # 1DArray
for pt in ptsCoords:
if shared:
pointsLocator.InsertUniquePoint( pt.tolist(), ptId ) # type: ignore[arg-type]
else:
pointsLocator.InsertPointWithoutChecking( pt.tolist(), ptId, 1 ) # type: ignore[arg-type]
cellVertexMap += [ ptId.get() ] # type: ignore
cellVertexMapAll += [ tuple( cellVertexMap ) ]
return points, cellVertexMapAll