[WIP] Getting the mapping for the scalar bar in pyvista correct

Author: Leguark

gempy_viewer/API/_plot_3d_API.py

@@ -175,7 +175,7 @@ def plot_3d(
             scalar_data_type=ScalarDataType.LITHOLOGY,
             active_scalar_field="lith",
             solution=solutions_raw_arrays,
-            cmap=get_geo_model_cmap(structural_frame.elements_colors_volumes),
+            cmap=get_geo_model_cmap(structural_frame.elements_colors),
             **kwargs_plot_structured_grid
         )
 

gempy_viewer/modules/plot_3d/drawer_input_3d.py

@@ -8,7 +8,7 @@
 from gempy_viewer.optional_dependencies import require_pyvista
 
 
-def plot_data(gempy_vista: GemPyToVista, 
+def plot_data(gempy_vista: GemPyToVista,
               model: GeoModel,
               arrows_factor: float,
               transformed_data: bool = False,
@@ -19,7 +19,7 @@ def plot_data(gempy_vista: GemPyToVista,
     else:
         surface_points_copy = model.surface_points_copy
         orientations_copy = model.orientations_copy
-        
+
     plot_surface_points(
         gempy_vista=gempy_vista,
         surface_points=surface_points_copy,
@@ -40,27 +40,21 @@ def plot_surface_points(
         surface_points: SurfacePointsTable,
         elements_colors: list[str],
         render_points_as_spheres=True,
-        point_size=10, 
+        point_size=10,
         **kwargs
 ):
-    ids = surface_points.ids
-    if ids.shape[0] == 0:
-        return
-    unique_values, first_indices = np.unique(ids, return_index=True)  # Find the unique elements and their first indices
-    unique_values_order = unique_values[np.argsort(first_indices)]  # Sort the unique values by their first appearance in `a`
-
-    mapping_dict = {value: i for i, value in enumerate(unique_values_order)}  # Use a dictionary to map the original numbers to new values
-    mapped_array = np.vectorize(mapping_dict.get)(ids)  # Map the original array to the new values
-
     # Selecting the surfaces to plot
     xyz = surface_points.xyz
     if transfromed_data := False:  # TODO: Expose this to user
         xyz = surface_points.model_transform.apply(xyz)
 
     pv = require_pyvista()
     poly = pv.PolyData(xyz)
-    poly['id'] = mapped_array
 
+    ids = surface_points.ids
+    if ids.shape[0] == 0:
+        return
+    poly['id'] = _vectorize_ids(ids)
 
     gempy_vista.surface_points_mesh = poly
     gempy_vista.surface_points_actor = gempy_vista.p.add_mesh(
@@ -80,23 +74,16 @@ def plot_orientations(
 ):
     orientations_xyz = orientations.xyz
     orientations_grads = orientations.grads
-    
+
     if orientations_xyz.shape[0] == 0:
         return
 
     pv = require_pyvista()
     poly = pv.PolyData(orientations_xyz)
-    
+
     ids = orientations.ids
-    if ids.shape[0] == 0:
-        return
-    unique_values, first_indices = np.unique(ids, return_index=True)  # Find the unique elements and their first indices
-    unique_values_order = unique_values[np.argsort(first_indices)]  # Sort the unique values by their first appearance in `a`
 
-    mapping_dict = {value: i for i, value in enumerate(unique_values_order)}  # Use a dictionary to map the original numbers to new values
-    mapped_array = np.vectorize(mapping_dict.get)(ids)  # Map the original array to the new values
-    
-    poly['id'] = mapped_array
+    poly['id'] = _vectorize_ids(ids)
     poly['vectors'] = orientations_grads
 
     # TODO: I am still trying to figure out colors and ids in orientations and surface points
@@ -117,3 +104,13 @@ def plot_orientations(
         show_scalar_bar=False
     )
     gempy_vista.orientations_mesh = arrows
+
+
+def _vectorize_ids(ids):
+    unique_values, first_indices = np.unique(ids, return_index=True)  # Find the unique elements and their first indices
+    unique_values_order = unique_values[np.argsort(first_indices)]  # Sort the unique values by their first appearance in `a`
+    # Flip order to please pyvista vertical scalarbar
+    unique_values_order = unique_values_order[::-1]
+    mapping_dict = {value: i + 1 for i, value in enumerate(unique_values_order)}  # Use a dictionary to map the original numbers to new values
+    mapped_array = np.vectorize(mapping_dict.get)(ids)  # Map the original array to the new values
+    return mapped_array

gempy_viewer/modules/plot_3d/drawer_structured_grid_3d.py

@@ -5,7 +5,7 @@
 
 from gempy_engine.core.data.raw_arrays_solution import RawArraysSolution
 from gempy_viewer.core.scalar_data_type import ScalarDataType
-from gempy.core.data.grid_modules import RegularGrid, Topography
+from gempy.core.data.grid_modules import Topography
 from gempy_viewer.modules.plot_3d.vista import GemPyToVista
 from gempy_viewer.optional_dependencies import require_pyvista
 
@@ -30,7 +30,6 @@ def plot_structured_grid(
     structured_grid = set_scalar_data(
         structured_grid=structured_grid,
         data=solution,
-        resolution=resolution,
         scalar_data_type=scalar_data_type
     )
 
@@ -64,18 +63,19 @@ def add_regular_grid_mesh(
         **kwargs
 ):
     if isinstance(cmap, mcolors.Colormap):
-        _clim = (0, cmap.N - 1)
+        _clim = (0, cmap.N)
     else:
         _clim = None
     gempy_vista.regular_grid_actor = gempy_vista.p.add_mesh(
         mesh=structured_grid,
-        cmap=cmap,
+        # cmap=cmap,
         # ? scalars=main_scalar, if we prepare the structured grid do we need this arg?
         show_scalar_bar=False,
-        scalar_bar_args=gempy_vista.scalar_bar_arguments,
+        # scalar_bar_args=gempy_vista.scalar_bar_arguments,
         interpolate_before_map=True,
         opacity=opacity,
-        clim=_clim,
+        # flip_scalars=True,
+        # clim=(4,0),
         **kwargs
     )
 
@@ -98,27 +98,25 @@ def _mask_topography(structured_grid: "pv.StructuredGrid", topography: Topograph
 def set_scalar_data(
         data: RawArraysSolution,
         structured_grid: "pv.StructuredGrid",
-        resolution: np.ndarray,
         scalar_data_type: ScalarDataType,
 ) -> "pv.StructuredGrid":
-    def _convert_sol_array_to_fortran_order(array: np.ndarray) -> np.ndarray:
-        # ? (Miguel Jun 24) Is this function deprecated?
-        # return array.reshape(*resolution, order='C').ravel(order='F')
-        return array
-
+    
     # Substitute the madness of the previous if with match
     match scalar_data_type:
         case ScalarDataType.LITHOLOGY | ScalarDataType.ALL:
-            structured_grid.cell_data['id'] = _convert_sol_array_to_fortran_order(data.lith_block - 1)
+            max_lith = data.n_surfaces # (for basement)
+            structured_grid.cell_data['id'] = max_lith - (data.lith_block - 1)
         case ScalarDataType.SCALAR_FIELD | ScalarDataType.ALL:
             scalar_field_ = 'sf_'
             for e in range(data.scalar_field_matrix.shape[0]):
                 # TODO: Ideally we will have the group name instead the enumeration
-                structured_grid[scalar_field_ + str(e)] = _convert_sol_array_to_fortran_order(data.scalar_field_matrix[e])
+                array1 = data.scalar_field_matrix[e]
+                structured_grid[scalar_field_ + str(e)] = array1
         case ScalarDataType.VALUES | ScalarDataType.ALL:
             scalar_field_ = 'values_'
             for e in range(data.values_matrix.shape[0]):
-                structured_grid[scalar_field_ + str(e)] = _convert_sol_array_to_fortran_order(data.values_matrix[e])
+                array2 = data.values_matrix[e]
+                structured_grid[scalar_field_ + str(e)] = array2
         case _:
             raise ValueError(f'Unknown scalar data type: {scalar_data_type}')
 

gempy_viewer/modules/plot_3d/plot_3d_utils.py

@@ -41,7 +41,7 @@ def set_scalar_bar(gempy_vista: GemPyToVista, elements_names: list[str],
 
     # Create annotations mapping integers to element names
     annotations = {}
-    for e, name in enumerate(elements_names):
+    for e, name in enumerate(elements_names[::-1]):
         # Convert integer to string for the annotation key
         annotations[str(e)] = name
 
@@ -58,8 +58,7 @@ def set_scalar_bar(gempy_vista: GemPyToVista, elements_names: list[str],
         if len(custom_colors) < n_colors:
             raise ValueError(f"Not enough custom colors provided. Got {len(custom_colors)}, need {n_colors}")
 
-        custom_cmap = ListedColormap(custom_colors).reversed()
-
+        custom_cmap = ListedColormap(custom_colors)
         # Apply the custom colormap to the lookup table
         lut.apply_cmap(cmap=custom_cmap, n_values=n_colors)
 

gempy_viewer/modules/plot_3d/vista.py

@@ -126,13 +126,14 @@ def scalar_bar_arguments(self):
         sargs = dict(
             title_font_size=20,
             label_font_size=16,
-            shadow=True,
+            shadow=False,
             italic=True,
+            bold=True,
             font_family="arial",
             height=0.25,
             vertical=True,
-            position_x=0.15,
-            title="id",
+            position_x=0.1,
+            title="Elements",
             fmt="%.0f",
         )
         return sargs

tests/test_plotting/test_plot_3d.py

@@ -26,7 +26,7 @@ def test_plot_3d_solutions(self, one_fault_model_topo_solution):
             show_topography=False,
             show_scalar=False,
             show_lith=True,
-            show_data=True,
+            show_data=False,
             show_boundaries=True,
             image=True
         )