diff --git a/docs/geos_mesh_docs/processing.rst b/docs/geos_mesh_docs/processing.rst index 844c071af..003cc973d 100644 --- a/docs/geos_mesh_docs/processing.rst +++ b/docs/geos_mesh_docs/processing.rst @@ -33,6 +33,14 @@ geos.mesh.processing.SplitMesh filter -------------------------------------- .. automodule:: geos.mesh.processing.SplitMesh + :members: + :undoc-members: + :show-inheritance: + +geos.mesh.processing.ClipToMainFrame filter +-------------------------------------------- + +.. automodule:: geos.mesh.processing.ClipToMainFrame :members: :undoc-members: :show-inheritance: \ No newline at end of file diff --git a/geos-mesh/src/geos/mesh/processing/ClipToMainFrame.py b/geos-mesh/src/geos/mesh/processing/ClipToMainFrame.py new file mode 100644 index 000000000..ecc099701 --- /dev/null +++ b/geos-mesh/src/geos/mesh/processing/ClipToMainFrame.py @@ -0,0 +1,310 @@ +# SPDX-License-Identifier: Apache 2.0 +# SPDX-FileCopyrightText: Copyright 2023-2025 TotalEnergies +# SPDX-FileContributor: Jacques Franc + +from vtkmodules.numpy_interface import dataset_adapter as dsa +from vtkmodules.vtkCommonCore import vtkPoints +from vtkmodules.vtkCommonMath import vtkMatrix4x4 +from vtkmodules.vtkFiltersGeneral import vtkOBBTree +from vtkmodules.vtkFiltersGeometry import vtkDataSetSurfaceFilter +from vtkmodules.vtkCommonDataModel import vtkUnstructuredGrid, vtkMultiBlockDataSet, vtkDataObjectTreeIterator, vtkPolyData +from vtkmodules.vtkCommonTransforms import vtkLandmarkTransform +from vtkmodules.vtkFiltersGeneral import vtkTransformFilter + +from geos.utils.Logger import logging, Logger, getLogger + +from geos.mesh.utils.genericHelpers import getMultiBlockBounds + +import numpy as np +import numpy.typing as npt + +from typing import Tuple + +__doc__ = """ +Module to clip a mesh to the main frame using rigid body transformation. + +Methods include: + - ClipToMainFrameElement class to compute the transformation + - ClipToMainFrame class to apply the transformation to a mesh + +To use it: + +.. code-block:: python + + from geos.mesh.processing.ClipToMainFrame import ClipToMainFrame + + # Filter inputs. + multiBlockDataSet: vtkMultiBlockDataSet + # Optional Inputs + speHandler : bool + + # Instantiate the filter. + filter: ClipToMainFrame = ClipToMainFrame() + filter.SetInputData( multiBlockDataSet ) + + # Set the handler of yours (only if speHandler is True). + yourHandler: logging.Handler + filter.setLoggerHandler( yourHandler ) + + # Do calculations. + filter.ComputeTransform() + filter.Update() + output: vtkMultiBlockDataSet = filter.GetOutput() + +""" + + +class ClipToMainFrameElement( vtkLandmarkTransform ): + + sourcePts: vtkPoints + targetPts: vtkPoints + mesh: vtkUnstructuredGrid + + def __init__( self, mesh: vtkUnstructuredGrid ) -> None: + """Clip mesh to main frame. + + Args: + mesh (vtkUnstructuredGrid): Mesh to clip. + """ + super().__init__() + self.mesh = mesh + + def Update( self ) -> None: + """Update the transformation.""" + self.sourcePts, self.targetPts = self.__getFramePoints( self.__getOBBTree( self.mesh ) ) + self.SetSourceLandmarks( self.sourcePts ) + self.SetTargetLandmarks( self.targetPts ) + self.SetModeToRigidBody() + super().Update() + + def __str__( self ) -> str: + """String representation of the transformation.""" + return super().__str__() + f"\nSource points: {self.sourcePts}" \ + + f"\nTarget points: {self.targetPts}" \ + + f"\nAngle-Axis: {self.__getAngleAxis()}" \ + + f"\nTranslation: {self.__getTranslation()}" + + def __getAngleAxis( self ) -> tuple[ float, npt.NDArray[ np.double ] ]: + """Get the angle and axis of the rotation. + + tuple[float, npt.NDArray[np.double]]: Angle in degrees and axis of rotation. + """ + matrix: vtkMatrix4x4 = self.GetMatrix() + angle: np.double = np.arccos( + ( matrix.GetElement( 0, 0 ) + matrix.GetElement( 1, 1 ) + matrix.GetElement( 2, 2 ) - 1 ) / 2 ) + if angle == 0: + return 0.0, np.array( [ 1.0, 0.0, 0.0 ] ) + rx: float = matrix.GetElement( 2, 1 ) - matrix.GetElement( 1, 2 ) + ry: float = matrix.GetElement( 0, 2 ) - matrix.GetElement( 2, 0 ) + rz: float = matrix.GetElement( 1, 0 ) - matrix.GetElement( 0, 1 ) + r = np.array( [ rx, ry, rz ] ) + r /= np.linalg.norm( r ) + return np.degrees( angle ), r + + def __getTranslation( self ) -> npt.NDArray[ np.double ]: + """Get the translation vector. + + Returns: + npt.NDArray[ np.double ]: The translation vector. + """ + matrix: vtkMatrix4x4 = self.GetMatrix() + return np.array( [ matrix.GetElement( 0, 3 ), matrix.GetElement( 1, 3 ), matrix.GetElement( 2, 3 ) ] ) + + def __getOBBTree( self, mesh: vtkUnstructuredGrid ) -> vtkPoints: + """Get the OBB tree of the mesh. + + Args: + mesh (vtkUnstructuredGrid): Mesh to get the OBB tree from. + + Returns: + vtkPoints: Points from the 0-level OBB tree of the mesh. Fallback on Axis Aligned Bounding Box + """ + OBBTree = vtkOBBTree() + surfFilter = vtkDataSetSurfaceFilter() + surfFilter.SetInputData( mesh ) + surfFilter.Update() + OBBTree.SetDataSet( surfFilter.GetOutput() ) + OBBTree.BuildLocator() + pdata = vtkPolyData() + OBBTree.GenerateRepresentation( 0, pdata ) + # at level 0 this should return 8 corners of the bounding box or fallback on AABB + if pdata.GetNumberOfPoints() < 3: + return self.__allpoints( mesh.GetBounds() ) + + return pdata.GetPoints() + + def __allpoints( self, bounds: tuple[ float, float, float, float, float, float ] ) -> vtkPoints: + """Get the 8 corners of the bounding box. + + Args: + bounds (tuple[float, float, float, float, float, float]): Bounding box. + + Returns: + vtkPoints: 8 corners of the bounding box. + """ + pts = vtkPoints() + pts.SetNumberOfPoints( 8 ) + pts.SetPoint( 0, [ bounds[ 0 ], bounds[ 2 ], bounds[ 4 ] ] ) + pts.SetPoint( 1, [ bounds[ 1 ], bounds[ 2 ], bounds[ 4 ] ] ) + pts.SetPoint( 2, [ bounds[ 1 ], bounds[ 3 ], bounds[ 4 ] ] ) + pts.SetPoint( 3, [ bounds[ 0 ], bounds[ 3 ], bounds[ 4 ] ] ) + pts.SetPoint( 4, [ bounds[ 0 ], bounds[ 2 ], bounds[ 5 ] ] ) + pts.SetPoint( 5, [ bounds[ 1 ], bounds[ 2 ], bounds[ 5 ] ] ) + pts.SetPoint( 6, [ bounds[ 1 ], bounds[ 3 ], bounds[ 5 ] ] ) + pts.SetPoint( 7, [ bounds[ 0 ], bounds[ 3 ], bounds[ 5 ] ] ) + return pts + + def __getFramePoints( self, vpts: vtkPoints ) -> tuple[ vtkPoints, vtkPoints ]: + """Get the source and target points for the transformation. + + Args: + vpts (vtkPoints): Points to transform. + + Returns: + tuple[vtkPoints, vtkPoints]: Source and target points for the transformation. + """ + pts: npt.NDArray[ np.double ] = dsa.numpy_support.vtk_to_numpy( vpts.GetData() ) + #translate pts so they always lie on the -z,-y,-x quadrant + off: npt.NDArray[ np.double ] = np.asarray( [ + -2 * np.amax( np.abs( pts[ :, 0 ] ) ), -2 * np.amax( np.abs( pts[ :, 1 ] ) ), + -2 * np.amax( np.abs( pts[ :, 2 ] ) ) + ] ) + pts += off + further_ix: np.int_ = np.argmax( np.linalg.norm( + pts, axis=1 ) ) # by default take the min point furthest from origin + org: npt.NDArray = pts[ further_ix, : ] + + # find 3 orthogonal vectors + # we assume points are on a box + dist_indexes: npt.NDArray[ np.int_ ] = np.argsort( np.linalg.norm( pts - org, axis=1 ) ) + # find u,v,w + v1: npt.NDArray[ np.double ] = pts[ dist_indexes[ 1 ], : ] - org + v2: npt.NDArray[ np.double ] = pts[ dist_indexes[ 2 ], : ] - org + v1 /= np.linalg.norm( v1 ) + v2 /= np.linalg.norm( v2 ) + if np.abs( v1[ 0 ] ) > np.abs( v2[ 0 ] ): + v1, v2 = v2, v1 + + # ensure orthogonality + v2 -= np.dot( v2, v1 ) * v1 + v2 /= np.linalg.norm( v2 ) + v3: npt.NDArray[ np.double ] = np.cross( v1, v2 ) + v3 /= np.linalg.norm( v3 ) + + #reorder axis if v3 points downward + if v3[ 2 ] < 0: + v3 = -v3 + v1, v2 = v2, v1 + + sourcePts = vtkPoints() + sourcePts.SetNumberOfPoints( 4 ) + sourcePts.SetPoint( 0, list( org - off ) ) + sourcePts.SetPoint( 1, list( v1 + org - off ) ) + sourcePts.SetPoint( 2, list( v2 + org - off ) ) + sourcePts.SetPoint( 3, list( v3 + org - off ) ) + + targetPts = vtkPoints() + targetPts.SetNumberOfPoints( 4 ) + targetPts.SetPoint( 0, [ 0., 0., 0. ] ) + targetPts.SetPoint( 1, [ 1., 0., 0. ] ) + targetPts.SetPoint( 2, [ 0., 1., 0. ] ) + targetPts.SetPoint( 3, [ 0., 0., 1. ] ) + + return ( sourcePts, targetPts ) + + +loggerTitle: str = "Clip mesh to main frame." + + +class ClipToMainFrame( vtkTransformFilter ): + """Filter to clip a mesh to the main frame using ClipToMainFrame class.""" + + def __init__( self, speHandler: bool = False, **properties: str ) -> None: + """Initialize the ClipToMainFrame Filter with optional speHandler args and forwarding properties to main class. + + Args: + speHandler (bool, optional): True to use a specific handler, False to use the internal handler. + Defaults to False. + properties (kwargs): kwargs forwarded to vtkTransformFilter. + """ + super().__init__( **properties ) + # Logger. + self.logger: Logger + if not speHandler: + self.logger = getLogger( loggerTitle, True ) + else: + self.logger = logging.getLogger( loggerTitle ) + self.logger.setLevel( logging.INFO ) + + def ComputeTransform( self ) -> None: + """Update the transformation.""" + # dispatch to ClipToMainFrame depending on input type + if isinstance( self.GetInput(), vtkMultiBlockDataSet ): + #locate reference point + try: + idBlock = self.__locate_reference_point( self.GetInput() ) + except IndexError: + self.logger.error( "Reference point is not in the domain" ) + + clip = ClipToMainFrameElement( self.GetInput().GetDataSet( idBlock ) ) + else: + clip = ClipToMainFrameElement( self.GetInput() ) + + clip.Update() + self.SetTransform( clip ) + + def SetLoggerHandler( self, handler: logging.Handler ) -> None: + """Set a specific handler for the filter logger. In this filter 4 log levels are use, .info, .error, .warning and .critical, be sure to have at least the same 4 levels. + + Args: + handler (logging.Handler): The handler to add. + """ + if not self.logger.hasHandlers(): + self.logger.addHandler( handler ) + else: + self.logger.warning( + "The logger already has an handler, to use yours set the argument 'speHandler' to True during the filter initialization." + ) + + def __locate_reference_point( self, multiBlockDataSet: vtkMultiBlockDataSet ) -> int: + """Locate the block to use as reference for the transformation. + + Args: + multiBlockDataSet (vtkMultiBlockDataSet): Input multiblock mesh. + + Returns: + int: Index of the block to use as reference. + """ + + def __inside( pt: npt.NDArray[ np.double ], bounds: tuple[ float, float, float, float, float, float ] ) -> bool: + """Check if a point is inside a box. + + Args: + pt (npt.NDArray[np.double]): Point to check. + bounds (tuple[float, float, float, float, float, float]): Bounding box. + + Returns: + bool: True if the point is inside the bounding box, False otherwise. + """ + self.logger.info( f"Checking if point {pt} is inside bounds {bounds}" ) + return ( pt[ 0 ] >= bounds[ 0 ] and pt[ 0 ] <= bounds[ 1 ] and pt[ 1 ] >= bounds[ 2 ] + and pt[ 1 ] <= bounds[ 3 ] and pt[ 2 ] >= bounds[ 4 ] and pt[ 2 ] <= bounds[ 5 ] ) + + DOIterator: vtkDataObjectTreeIterator = vtkDataObjectTreeIterator() + DOIterator.SetDataSet( multiBlockDataSet ) + DOIterator.VisitOnlyLeavesOn() + DOIterator.GoToFirstItem() + xmin: float + ymin: float + zmin: float + xmin, _, ymin, _, zmin, _ = getMultiBlockBounds( multiBlockDataSet ) + while DOIterator.GetCurrentDataObject() is not None: + dataSet: vtkUnstructuredGrid = vtkUnstructuredGrid.SafeDownCast( DOIterator.GetCurrentDataObject() ) + bounds: Tuple[ float, float, float, float, float, float ] = dataSet.GetBounds() + #use the furthest bounds corner as reference point in the all negs quadrant + if __inside( np.asarray( [ xmin, ymin, zmin ] ), bounds ): + self.logger.info( f"Using block {DOIterator.GetCurrentFlatIndex()} as reference for transformation" ) + return DOIterator.GetCurrentFlatIndex() + DOIterator.GoToNextItem() + + raise IndexError diff --git a/geos-mesh/tests/test_clipToMainFrame.py b/geos-mesh/tests/test_clipToMainFrame.py new file mode 100644 index 000000000..fecba2704 --- /dev/null +++ b/geos-mesh/tests/test_clipToMainFrame.py @@ -0,0 +1,142 @@ +# SPDX-FileContributor: Jacques Franc +# SPEDX-FileCopyrightText: Copyright 2023-2025 TotalEnergies +# SPDX-License-Identifier: Apache 2.0 +# ruff: noqa: E402 # disable Module level import not at top of file +# mypy: disable-error-code="operator" +import pytest +import itertools +from dataclasses import dataclass +from typing import Generator, Tuple +from vtkmodules.vtkCommonCore import vtkIdList, vtkPoints +from vtkmodules.vtkCommonDataModel import vtkUnstructuredGrid, vtkHexahedron, vtkMultiBlockDataSet +from vtkmodules.numpy_interface import dataset_adapter as dsa +import numpy as np +import numpy.typing as npt +from vtkmodules.util.vtkConstants import VTK_HEXAHEDRON + +from typing import List + +from geos.mesh.processing.ClipToMainFrame import ClipToMainFrame + +Lx, Ly, Lz = 5, 2, 8 +nx, ny, nz = 10, 10, 10 + + +@dataclass( frozen=True ) +class Expected: + mesh: vtkUnstructuredGrid + + +def __gen_box( Lx: int, Ly: int, Lz: int, nx: int, ny: int, nz: int, multx: int, multy: int, + multz: int ) -> Tuple[ npt.NDArray[ np.double ], npt.NDArray[ np.double ] ]: + off: npt.NDArray[ np.double ] = np.max( [ Lx, Ly, Lz ] ) * np.asarray( [ [ multx, multy, multz ] ] ) + pts: List[ npt.NDArray[ np.double ] ] = [] + x, y, z = np.meshgrid( np.linspace( 0, Lx, nx ), np.linspace( 0, Ly, ny ), np.linspace( 0, Lz, ny ) ) + for i in range( x.shape[ 0 ] ): + for j in range( x.shape[ 1 ] ): + for k in range( x.shape[ 2 ] ): + pts.append( np.asarray( [ x[ i, j, k ], y[ i, j, k ], z[ i, j, k ] ] ) ) + + return ( np.asarray( pts ), off ) + + +def __rotate_box( angles: npt.NDArray[ np.double ], pts: npt.NDArray[ np.double ] ) -> npt.NDArray[ np.double ]: + a: np.double = angles[ 0 ] + RX: npt.NDArray[ np.double ] = np.asarray( [ [ 1., 0, 0 ], [ 0, np.cos( a ), -np.sin( a ) ], + [ 0, np.sin( a ), np.cos( a ) ] ] ) + a = angles[ 1 ] + RY: npt.NDArray[ np.double ] = np.asarray( [ [ np.cos( a ), 0, np.sin( a ) ], [ 0, 1, 0 ], + [ -np.sin( a ), 0, np.cos( a ) ] ] ) + a = angles[ 2 ] + RZ: npt.NDArray[ np.double ] = np.asarray( [ [ np.cos( a ), -np.sin( a ), 0 ], [ np.sin( a ), + np.cos( a ), 0 ], [ 0, 0, 1 ] ] ) + + return np.asarray( ( RZ @ RY @ RX @ pts.transpose() ).transpose() ) + + +def __build_test_mesh( mxx: Tuple[ int, ...] ) -> Generator[ Expected, None, None ]: + # generate random points in a box Lx, Ly, Lz + # np.random.seed(1) # for reproducibility + np.random.default_rng() + + #test all quadrant + multx: int + multy: int + multz: int + multx, multy, multz = mxx + pts: npt.NDArray[ np.double ] + off: npt.NDArray[ np.double ] + pts, off = __gen_box( Lx, Ly, Lz, nx, ny, nz, multx, multy, multz ) + + angles: npt.NDArray[ np.double ] = -2 * np.pi + np.random.randn( 1, 3 ) * np.pi # random angles in rad + pts = __rotate_box( angles[ 0 ], pts ) + pts[ :, 0 ] += off[ 0 ][ 0 ] + pts[ :, 1 ] += off[ 0 ][ 1 ] + pts[ :, 2 ] += off[ 0 ][ 2 ] + + # Creating multiple meshes, each time with a different angles + mesh = vtkUnstructuredGrid() + ( vtps := vtkPoints() ).SetData( dsa.numpy_support.numpy_to_vtk( pts ) ) + mesh.SetPoints( vtps ) + + ids = vtkIdList() + for i in range( nx - 1 ): + for j in range( ny - 1 ): + for k in range( nz - 1 ): + hex = vtkHexahedron() + ids = hex.GetPointIds() + ids.SetId( 0, i + j * nx + k * nx * ny ) + ids.SetId( 1, i + j * nx + k * nx * ny + 1 ) + ids.SetId( 2, i + j * nx + k * nx * ny + nx * ny + 1 ) + ids.SetId( 3, i + j * nx + k * nx * ny + nx * ny ) + ids.SetId( 4, i + j * nx + k * nx * ny + nx ) + ids.SetId( 5, i + j * nx + k * nx * ny + nx + 1 ) + ids.SetId( 6, i + j * nx + k * nx * ny + nx * ny + nx + 1 ) + ids.SetId( 7, i + j * nx + k * nx * ny + nx * ny + nx ) + mesh.InsertNextCell( VTK_HEXAHEDRON, ids ) + + yield Expected( mesh=mesh ) + + +# arg-type: ignore +@pytest.mark.parametrize( + "expected", [ item for t in list( itertools.product( [ -1, 1 ], repeat=3 ) ) for item in __build_test_mesh( t ) ] ) +def test_clipToMainFrame_polyhedron( expected: Expected ) -> None: + """Test the ClipToMainFrameFilter on a rotated and translated box hexa mesh.""" + ( filter := ClipToMainFrame() ).SetInputData( expected.mesh ) + filter.ComputeTransform() + filter.Update() + output_mesh: vtkUnstructuredGrid = filter.GetOutput() + assert output_mesh.GetNumberOfPoints() == expected.mesh.GetNumberOfPoints() + assert output_mesh.GetNumberOfCells() == expected.mesh.GetNumberOfCells() + + assert output_mesh.GetBounds()[ 0 ] == pytest.approx( + 0., abs=1e-4 ) and output_mesh.GetBounds()[ 2 ] == pytest.approx( 0., abs=1e-4 ) and np.max( + [ np.abs( output_mesh.GetBounds()[ 4 ] ), + np.abs( output_mesh.GetBounds()[ 5 ] ) ] ) == pytest.approx( Lz, abs=1e-4 ) + # test diagonal + assert np.linalg.norm( + np.array( [ + output_mesh.GetBounds()[ 1 ] - output_mesh.GetBounds()[ 0 ], + output_mesh.GetBounds()[ 3 ] - output_mesh.GetBounds()[ 2 ], + output_mesh.GetBounds()[ 5 ] - output_mesh.GetBounds()[ 4 ] + ] ) ) == pytest.approx( np.linalg.norm( np.array( [ Lx, Ly, Lz ] ) ), abs=1e-4 ) + # test aligned with axis + v0: npt.NDArray[ np.double ] = np.array( output_mesh.GetPoint( 1 ) ) - np.array( output_mesh.GetPoint( 0 ) ) + v1: npt.NDArray[ np.double ] = np.array( output_mesh.GetPoint( nx ) ) - np.array( output_mesh.GetPoint( 0 ) ) + v2: npt.NDArray[ np.double ] = np.array( output_mesh.GetPoint( nx * ny ) ) - np.array( output_mesh.GetPoint( 0 ) ) + assert np.abs( np.dot( v0, v1 ) ) < 1e-10 + assert np.abs( np.dot( v0, v2 ) ) < 1e-10 + assert np.abs( np.dot( v1, v2 ) ) < 1e-10 + + +def test_clipToMainFrame_generic( dataSetTest: vtkMultiBlockDataSet ) -> None: + """Test the ClipToMainFrameFilter on a MultiBlockDataSet.""" + multiBlockDataSet: vtkMultiBlockDataSet = dataSetTest( "multiblock" ) + ( filter := ClipToMainFrame() ).SetInputData( multiBlockDataSet ) + filter.ComputeTransform() + filter.Update() + print( filter.GetTransform() ) + output_mesh: vtkMultiBlockDataSet = filter.GetOutputDataObject( 0 ) + assert output_mesh.GetNumberOfPoints() == multiBlockDataSet.GetNumberOfPoints() + assert output_mesh.GetNumberOfCells() == multiBlockDataSet.GetNumberOfCells()