feat: add scalar array selection in geos-trame

geos-trame/src/geos/trame/app/components/alertHandler.py

@@ -26,8 +26,8 @@ def __init__( self ) -> None:
 
         self.state.alerts = []
 
-        self.server.controller.on_add_error.add_task( self.add_error )
-        self.server.controller.on_add_warning.add_task( self.add_warning )
+        self.ctrl.on_add_error.add_task( self.add_error )
+        self.ctrl.on_add_warning.add_task( self.add_warning )
 
         self.generate_alert_ui()
 

geos-trame/src/geos/trame/app/io/data_loader.py

@@ -11,7 +11,7 @@
 from geos.trame.app.geosTrameException import GeosTrameException
 from geos.trame.app.ui.viewer.regionViewer import RegionViewer
 from geos.trame.app.ui.viewer.wellViewer import WellViewer
-from geos.trame.app.utils.pv_utils import read_unstructured_grid
+from geos.trame.app.utils.pv_utils import read_unstructured_grid, split_vector_arrays
 from geos.trame.schema_generated.schema_mod import (
     Vtkmesh,
     Vtkwell,
@@ -97,6 +97,7 @@ def _update_vtkmesh( self, mesh: Vtkmesh, show: bool ) -> None:
 
     def _read_mesh( self, mesh: Vtkmesh ) -> None:
         unstructured_grid = read_unstructured_grid( self.source.get_abs_path( mesh.file ) )
+        split_vector_arrays( unstructured_grid )
         self.region_viewer.add_mesh( unstructured_grid )
 
     def _update_vtkwell( self, well: Vtkwell, path: str, show: bool ) -> None:

geos-trame/src/geos/trame/app/ui/viewer/viewer.py

@@ -8,18 +8,14 @@
 from pyvista.trame.ui import plotter_ui
 from trame.widgets import html
 from trame.widgets import vuetify3 as vuetify
+from vtkmodules.vtkCommonDataModel import vtkUnstructuredGrid
 from vtkmodules.vtkRenderingCore import vtkActor
 
 from geos.trame.app.deck.tree import DeckTree
 from geos.trame.app.ui.viewer.perforationViewer import PerforationViewer
 from geos.trame.app.ui.viewer.regionViewer import RegionViewer
 from geos.trame.app.ui.viewer.wellViewer import WellViewer
-from geos.trame.schema_generated.schema_mod import (
-    Vtkmesh,
-    Vtkwell,
-    Perforation,
-    InternalWell,
-)
+from geos.trame.schema_generated.schema_mod import Vtkmesh, Vtkwell, InternalWell, Perforation
 
 pv.OFF_SCREEN = True
 
@@ -49,13 +45,24 @@ def __init__(
         """
         super().__init__( **kwargs )
 
+        self._picked_actor: vtkActor | None = None
+        self._point_data_array_names: list[ str ] = []
+        self._cell_data_array_names: list[ str ] = []
         self._source = source
         self._pl = pv.Plotter()
+        self._pl.enable_mesh_picking( callback=self._picking_callback,
+                                      show_message=False,
+                                      line_width=2,
+                                      use_actor=True )
 
         self.CUT_PLANE = "on_cut_plane_visibility_change"
         self.ZAMPLIFICATION = "_z_amplification"
-        self.server.state[ self.CUT_PLANE ] = True
-        self.server.state[ self.ZAMPLIFICATION ] = 1
+        self.state[ self.CUT_PLANE ] = True
+        self.state[ self.ZAMPLIFICATION ] = 1
+
+        self.DATA_ARRAYS = "viewer_data_arrays_items"
+        self.SELECTED_DATA_ARRAY = "viewer_selected_data_array"
+        self.state.change( self.SELECTED_DATA_ARRAY )( self._update_actor_array )
 
         self.region_engine = region_viewer
         self.well_engine = well_viewer
@@ -68,8 +75,9 @@ def __init__(
             view = plotter_ui(
                 self._pl,
                 add_menu_items=self.rendering_menu_extra_items,
-                style="position: absolute;",
             )
+            view.menu.style += "; height: 50px; min-width: 50px;"
+            view.menu.children[ 0 ].style += "; justify-content: center;"
             self.ctrl.view_update = view.update
 
     @property
@@ -88,21 +96,65 @@ def rendering_menu_extra_items( self ) -> None:
         For now, adding a button to show/hide all widgets.
         """
         self.state.change( self.CUT_PLANE )( self._on_clip_visibility_change )
-        vuetify.VDivider( vertical=True, classes="mr-1" )
-        with vuetify.VTooltip( location="bottom" ):
-            with (
-                    vuetify.Template( v_slot_activator=( "{ props }", ) ),
-                    html.Div( v_bind=( "props", ) ),
-            ):
-                vuetify.VCheckbox(
-                    v_model=( self.CUT_PLANE, True ),
-                    icon=True,
-                    true_icon="mdi-eye",
-                    false_icon="mdi-eye-off",
-                    dense=True,
-                    hide_details=True,
-                )
-            html.Span( "Show/Hide widgets" )
+        with vuetify.VRow(
+                classes='pa-0 ma-0 align-center fill-height',
+                style="flex-wrap: nowrap",
+        ):
+            vuetify.VDivider( vertical=True, classes="mr-1" )
+            with vuetify.VTooltip( location="bottom" ):
+                with (
+                        vuetify.Template( v_slot_activator=( "{ props }", ) ),
+                        html.Div( v_bind=( "props", ) ),
+                ):
+                    vuetify.VCheckbox(
+                        v_model=( self.CUT_PLANE, True ),
+                        icon=True,
+                        true_icon="mdi-eye",
+                        false_icon="mdi-eye-off",
+                        dense=True,
+                        hide_details=True,
+                    )
+                html.Span( "Show/Hide widgets" )
+            vuetify.VDivider( vertical=True, classes="mr-1" )
+            vuetify.VSelect(
+                hide_details=True,
+                label="Data Array",
+                items=( self.DATA_ARRAYS, [] ),
+                v_model=( self.SELECTED_DATA_ARRAY, None ),
+                min_width="150px",
+            )
+
+    def _picking_callback( self, actor: vtkActor | None ) -> None:
+        """Actor picking callback.
+
+        Get the data arrays from its mesh.
+        """
+        if actor is None:
+            self.state[ self.DATA_ARRAYS ] = []
+        else:
+            mesh = actor.mapper.GetInputDataObject( 0, 0 )
+            self._point_data_array_names = list( mesh.point_data.keys() )
+            self._cell_data_array_names = list( mesh.cell_data.keys() )
+            self.state[ self.DATA_ARRAYS ] = self._point_data_array_names + self._cell_data_array_names
+            self._picked_actor = actor
+
+        self.state[ self.SELECTED_DATA_ARRAY ] = None
+
+    def _update_actor_array( self, **_: Any ) -> None:
+        """Update the picked actor scalar array."""
+        array_name = self.state[ self.SELECTED_DATA_ARRAY ]
+        if array_name is None or self._picked_actor is None:
+            return
+        mapper: pv.DataSetMapper = self._picked_actor.mapper
+        mesh: vtkUnstructuredGrid = mapper.GetInputDataObject( 0, 0 )
+
+        data = mesh.GetPointData() if array_name in self._point_data_array_names else mesh.GetCellData()
+        data.SetActiveScalars( array_name )
+        mapper.scalar_range = data.GetArray( array_name ).GetRange()
+        mapper.scalar_map_mode = "point" if array_name in self._point_data_array_names else "cell"
+
+        self.plotter.scalar_bar.title = array_name
+        self.ctrl.view_update()
 
     def update_viewer( self, active_block: BaseModel, path: str, show_obj: bool ) -> None:
         """Add from path the dataset given by the user.
@@ -205,7 +257,7 @@ def _update_internalwell( self, path: str, show: bool ) -> None:
         tube_actor = self.plotter.add_mesh( self.well_engine.get_tube( self.well_engine.get_last_mesh_idx() ) )
         self.well_engine.append_actor( path, tube_actor )
 
-        self.server.controller.view_update()
+        self.ctrl.view_update()
 
     def _update_vtkwell( self, path: str, show: bool ) -> None:
         """Used to control the visibility of the Vtkwell.
@@ -219,7 +271,7 @@ def _update_vtkwell( self, path: str, show: bool ) -> None:
         tube_actor = self.plotter.add_mesh( self.well_engine.get_tube( self.well_engine.get_last_mesh_idx() ) )
         self.well_engine.append_actor( path, tube_actor )
 
-        self.server.controller.view_update()
+        self.ctrl.view_update()
 
     def _update_vtkmesh( self, show: bool ) -> None:
         """Used to control the visibility of the Vtkmesh.
@@ -233,18 +285,14 @@ def _update_vtkmesh( self, show: bool ) -> None:
             self.plotter.remove_actor( self._clip_mesh )  # type: ignore
             return
 
-        active_scalar = self.region_engine.input.active_scalars_name
         self._clip_mesh: vtkActor = self.plotter.add_mesh_clip_plane(
             self.region_engine.input,
             origin=self.region_engine.input.center,
             normal=[ -1, 0, 0 ],
             crinkle=True,
             show_edges=False,
-            cmap="glasbey_bw",
-            scalars=active_scalar,
         )
-
-        self.server.controller.view_update()
+        self._picking_callback( self._clip_mesh )
 
     def _update_perforation( self, perforation: Perforation, show: bool, path: str ) -> None:
         """Generate VTK dataset from a perforation."""

geos-trame/src/geos/trame/app/utils/pv_utils.py

@@ -2,8 +2,27 @@
 # SPDX-FileCopyrightText: Copyright 2023-2024 TotalEnergies.
 # SPDX-FileContributor: Kitware
 import pyvista as pv
+from vtkmodules.util.numpy_support import vtk_to_numpy, numpy_to_vtk
+from vtkmodules.vtkCommonCore import vtkDataArray
 
 
 def read_unstructured_grid( filename: str ) -> pv.UnstructuredGrid:
     """Read an unstructured grid from a .vtu file."""
     return pv.read( filename ).cast_to_unstructured_grid()
+
+
+def split_vector_arrays( ug: pv.UnstructuredGrid ) -> None:
+    """Create N 1-component arrays from each vector array with N components."""
+    for data in [ ug.GetPointData(), ug.GetCellData() ]:
+        for i in range( data.GetNumberOfArrays() ):
+            array: vtkDataArray = data.GetArray( i )
+            if array.GetNumberOfComponents() != 1:
+                np_array = vtk_to_numpy( array )
+                array_name = array.GetName()
+                data.RemoveArray( array_name )
+                for comp in range( array.GetNumberOfComponents() ):
+                    component = np_array[ :, comp ]
+                    new_array_name = f"{array_name}_{comp}"
+                    new_array = numpy_to_vtk( component, deep=True )
+                    new_array.SetName( new_array_name )
+                    data.AddArray( new_array )

geos-trame/tests/data/geosDeck/geosDeck.xml

@@ -67,7 +67,7 @@
         <!-- Reuse singlePhaseFlow for the testing -->
         <VTKMesh
             name="SyntheticMesh"
-            file="../singlePhaseFlow/synthetic.vtu"
+            file="../singlePhaseFlow/synthetic_random_arrays.vtu"
             logLevel="1" >
 
             <!-- Block to comment to deactivate external well definition -->

geos-trame/tests/test_data_loader.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright 2023-2024 TotalEnergies.
+# SPDX-FileContributor: Kitware
+# ruff: noqa
+import pyvista as pv
+from pathlib import Path
+
+from trame_server import Server
+from trame_server.state import State
+from trame_vuetify.ui.vuetify3 import VAppLayout
+
+from geos.trame.app.core import GeosTrame
+from tests.trame_fixtures import trame_server_layout, trame_state
+
+
+def test_data_loader( trame_server_layout: tuple[ Server, VAppLayout ], trame_state: State ) -> None:
+    root_path = Path( __file__ ).parent.absolute().__str__()
+    file_name = root_path + "/data/geosDeck/geosDeck.xml"
+
+    geos_trame = GeosTrame( trame_server_layout[ 0 ], file_name )
+
+    geos_trame.data_loader.load_vtkmesh_from_id( "Problem/Mesh/0/VTKMesh/0" )
+    ug: pv.UnstructuredGrid = geos_trame.data_loader.region_viewer.input
+    assert ug.GetCellData().HasArray( "attribute" )
+    assert ug.GetPointData().HasArray( "RandomPointScalars" )
+    assert not ug.GetPointData().HasArray( "RandomPointVectors" )
+    assert ug.GetPointData().HasArray( "RandomPointVectors_0" )
+    assert ug.GetPointData().HasArray( "RandomPointVectors_1" )
+    assert ug.GetPointData().HasArray( "RandomPointVectors_2" )
+    assert not ug.GetPointData().HasArray( "RandomPointVectors_3" )