feat: Improve Mohr circle display

geos-pv/src/geos/pv/plugins/post_processing/PVMohrCirclePlot.py

@@ -7,6 +7,7 @@
 from pathlib import Path
 from enum import Enum
 from typing import Any, Union, cast
+from typing_extensions import Self
 
 import numpy as np
 import numpy.typing as npt
@@ -18,9 +19,8 @@
 from paraview.detail.loghandler import VTKHandler  # type: ignore[import-not-found]
 # source: https://github.com/Kitware/ParaView/blob/master/Wrapping/Python/paraview/detail/loghandler.py
 
-from typing_extensions import Self
 from vtkmodules.vtkCommonCore import vtkDataArraySelection as vtkDAS
-from vtkmodules.vtkCommonCore import vtkInformation, vtkInformationVector
+from vtkmodules.vtkCommonCore import vtkInformation, vtkInformationVector, vtkStringArray, vtkIntArray
 from vtkmodules.vtkCommonDataModel import vtkUnstructuredGrid
 
 # Update sys.path to load all GEOS Python Package dependencies
@@ -73,7 +73,7 @@
 
 * Extract a few number of cells with the `ExtractSelection` ParaView Filter, then use the `MergeBlocks` ParaView Filter
 * Select the resulting mesh in the pipeline
-* Select the filter: Filters > { FilterCategory.GENERIC_PROCESSING.value } > Plot Mohr's Circle
+* Select the filter: Filters > { FilterCategory.GEOS_POST_PROCESSING.value } > Plot Mohr's Circle
 * Select the cell Ids and time steps you want
 * (Optional) Set rock cohesion and/or friction angle
 * Apply
@@ -84,12 +84,15 @@
     After a first application, select again cells and time steps to display, then
         * Apply again
         * Click on `Refresh Data` (you may have to click twice to refresh the Python view correctly).
+    To visualize the index of the cell used to calculate Mohr circle, use the ParaView tool 'Find Data':
+        * The attribute 'ActiveCellMask' allows to select only the right cells (equal to 1).
+        * The attribute 'CellId' has to be used for the 'Selection Labels'.
 """
 
 
 @smproxy.filter( name="PVMohrCirclePlot", label="Plot Mohr's Circles" )
 @smhint.xml( f"""
-    <ShowInMenu category="{ FilterCategory.GENERIC_PROCESSING.value }"/>
+    <ShowInMenu category="{ FilterCategory.GEOS_POST_PROCESSING.value }"/>
     <View type="PythonView"/>
     """ )
 @smproperty.input( name="Input", port_index=0 )
@@ -104,7 +107,10 @@ def __init__( self: Self ) -> None:
 
         Mohr's circles are plotted using a Python View.
         """
-        super().__init__( nInputPorts=1, nOutputPorts=1, outputType="vtkDataObject" )
+        super().__init__( nInputPorts=1,
+                          nOutputPorts=1,
+                          inputType="vtkUnstructuredGrid",
+                          outputType="vtkUnstructuredGrid" )
 
         # Create a new PythonView
         self.pythonView: Any = buildNewLayoutWithPythonView()
@@ -719,9 +725,9 @@ def RequestDataObject(
         outData = self.GetOutputData( outInfoVec, 0 )
 
         assert inData is not None
-        if ( outData is None ) or ( not outData.IsA( inData.GetClassName() ) ):
-            outData = inData.NewInstance()
-            outInfoVec.GetInformationObject( 0 ).Set( outData.DATA_OBJECT(), outData )
+        if outData is None or ( not outData.IsA( "vtkUnstructuredGrid" ) ):
+            outData = vtkUnstructuredGrid()
+            outInfoVec.GetInformationObject( 0 ).Set( outData.DATA_OBJECT(), outData )  # type: ignore
         return super().RequestDataObject( request, inInfoVec, outInfoVec )  # type: ignore[no-any-return]
 
     def RequestData(
@@ -774,8 +780,33 @@ def RequestData(
                     self.frictionAngle,
                     self._getUserChoices(),
                 )
+
                 Render()
 
+                # Cell indexes annotation
+                nbCells = inputMesh.GetNumberOfCells()
+                outputMesh: vtkUnstructuredGrid = self.GetOutputData( outInfoVec, 0 )
+                outputMesh.ShallowCopy( inputMesh )
+
+                cellId = vtkStringArray()
+                cellId.SetName( "CellId" )
+                cellId.SetNumberOfValues( nbCells )
+
+                activeCellMask = vtkIntArray()
+                activeCellMask.SetName( "ActiveCellMask" )
+                activeCellMask.SetNumberOfValues( nbCells )
+
+                originalCellIds = inputMesh.GetCellData().GetArray( "vtkOriginalCellIds" )
+                for localCellId in range( nbCells ):
+                    if str( originalCellIds.GetValue( localCellId ) ) in self.requestedCellIds:
+                        cellId.SetValue( localCellId, f"{ originalCellIds.GetValue( localCellId ) }" )
+                        activeCellMask.SetValue( localCellId, 1 )
+                    else:
+                        activeCellMask.SetValue( localCellId, 0 )
+
+                outputMesh.GetCellData().AddArray( cellId )
+                outputMesh.GetCellData().AddArray( activeCellMask )
+                outputMesh.Modified()
         except Exception as e:
             self.logger.error( "Mohr circles cannot be plotted due to:" )
             self.logger.error( e )
@@ -812,12 +843,21 @@ def _createMohrCirclesAtTimeStep(
     def _filterMohrCircles( self: Self ) -> list[ MohrCircle ]:
         """Filter the list of all MohrCircle to get those to plot.
 
+        Mohr circles are sort by cell indexes first then by timesteps.
+
         Returns:
             list[MohrCircle]: list of MohrCircle to plot.
         """
         # Circle ids to plot
         circleIds: list[ str ] = self._getCircleIds()
-        return [ mohrCircle for mohrCircle in self.mohrCircles if mohrCircle.getCircleId() in circleIds ]
+        mohrCircleToPlot: list[ MohrCircle ] = [ MohrCircle( "-1" ) for i in range( len( circleIds ) ) ]
+        for mohrCircle in self.mohrCircles:
+            try:
+                mohrCircleToPlot[ circleIds.index( str( mohrCircle.getCircleId() ) ) ] = mohrCircle
+            except ValueError:
+                continue
+
+        return mohrCircleToPlot
 
     def _updateRequestedTimeSteps( self: Self ) -> None:
         """Update the requestedTimeStepsIndexes attribute from user choice."""
@@ -842,13 +882,15 @@ def _getUserChoices( self: Self ) -> dict[ str, Any ]:
     def _getCircleIds( self: Self ) -> list[ str ]:
         """Get circle ids to plot.
 
+        This list of circle indexes is sort by cell indexes first then by timesteps
+
         Returns:
             list[str]: list of circle ids to plot.
         """
         cellIds: list[ str ] = pvt.getArrayChoices( self.a01GetCellIdsDAS() )
         timeSteps: list[ str ] = pvt.getArrayChoices( self.a02GetTimestepsToPlot() )
 
-        return [ mcf.getMohrCircleId( cellId, timeStep ) for timeStep in timeSteps for cellId in cellIds ]
+        return [ mcf.getMohrCircleId( cellId, timeStep ) for cellId in cellIds for timeStep in timeSteps ]
 
     def _defineCurvesAspect( self: Self ) -> None:
         """Add curve aspect parameters according to user choices."""