Improve 3D visualization with consistent color handling and scalar bar display (#7)

# Improved 3D Visualization with Enhanced Color Mapping

This PR refactors the 3D plotting utilities to provide better color mapping and visualization. Key improvements include:

- Refactored scalar bar implementation with proper color mapping and annotations
- Added support for custom colormaps in surface point visualization
- Fixed lithology block visualization by properly handling element colors
- Improved actor priority for scalar bar assignment
- Enhanced scalar bar appearance with better formatting and positioning
- Changed default scalar field colormap from 'viridis' to 'magma'
- Fixed vectorization of IDs for consistent color mapping between surface points and orientations
- Optimized structural frame access with local variable assignment
- Improved scalar bar configuration with cleaner annotations and better visual styling
- Added gempy_engine dependency to requirements.txt

These changes provide more consistent and visually appealing 3D visualizations with proper color mapping between different elements in the model.

Author: Leguark

gempy_viewer/API/_plot_3d_API.py

@@ -114,20 +114,21 @@ def plot_3d(
         **kwargs_plotter
     )
 
+    structural_frame = model.structural_frame
     if data_to_show.show_topography[0] is True and model.grid.topography is not None:
         plot_topography_3d(
             gempy_vista=gempy_vista,
             topography=model.grid.topography,
             solution=solutions_raw_arrays,
             topography_scalar_type=topography_scalar_type,
-            elements_colors=model.structural_frame.elements_colors[::-1],
+            elements_colors=structural_frame.elements_colors[::-1],
             contours=kwargs_plot_topography.get('contours', True),
             **kwargs_plot_topography
         )
 
     if data_to_show.show_boundaries[0] is True:
         # Check elements to plot .vertices are not empty
-        elements_to_plot = model.structural_frame.structural_elements
+        elements_to_plot = structural_frame.structural_elements
         for element in elements_to_plot:
             if element.vertices is None:
                 elements_to_plot.remove(element)
@@ -174,7 +175,7 @@ def plot_3d(
             scalar_data_type=ScalarDataType.LITHOLOGY,
             active_scalar_field="lith",
             solution=solutions_raw_arrays,
-            cmap=get_geo_model_cmap(model.structural_frame.elements_colors_volumes),
+            cmap=get_geo_model_cmap(structural_frame.elements_colors),
             **kwargs_plot_structured_grid
         )
 
@@ -186,15 +187,16 @@ def plot_3d(
             scalar_data_type=ScalarDataType.SCALAR_FIELD,
             active_scalar_field=active_scalar_field,
             solution=solutions_raw_arrays,
-            cmap='viridis',
+            cmap='magma',
             **kwargs_plot_structured_grid
         )
 
-    if True: 
+    if True:
         set_scalar_bar(
             gempy_vista=gempy_vista,
-            elements_names = model.structural_frame.elements_names,
-            surfaces_ids=model.structural_frame.elements_ids - 1
+            elements_names=structural_frame.elements_names,
+            surfaces_ids=structural_frame.elements_ids - 1,
+            custom_colors=structural_frame.elements_colors_volumes
         )
 
     if ve is not None:

gempy_viewer/modules/plot_3d/drawer_input_3d.py

@@ -3,12 +3,14 @@
 from gempy.core.data import GeoModel
 from gempy.core.data.orientations import OrientationsTable
 from gempy.core.data.surface_points import SurfacePointsTable
+from matplotlib.colors import ListedColormap
+
 from gempy_viewer.modules.plot_2d.plot_2d_utils import get_geo_model_cmap
 from gempy_viewer.modules.plot_3d.vista import GemPyToVista
 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,96 +21,76 @@ 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,
-        elements_colors=model.structural_frame.elements_colors_contacts,
-        **kwargs
+        element_colors=model.structural_frame.elements_colors
     )
 
     plot_orientations(
         gempy_vista=gempy_vista,
         orientations=orientations_copy,
-        elements_colors=model.structural_frame.elements_colors_orientations,
+        surface_points=surface_points_copy,
         arrows_factor=arrows_factor
     )
 
 
 def plot_surface_points(
         gempy_vista: GemPyToVista,
         surface_points: SurfacePointsTable,
-        elements_colors: list[str],
         render_points_as_spheres=True,
-        point_size=10, 
-        **kwargs
+        element_colors=None,
+        point_size=10
 ):
-    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
 
-    cmap = get_geo_model_cmap(
-        elements_colors=np.array(elements_colors),
-        reverse=False
-    )
+    ids = surface_points.ids
+    if ids.shape[0] == 0:
+        return
+    vectorize_ids = _vectorize_ids(ids, ids)
+    poly['id'] = vectorize_ids
 
+    custom_cmap = ListedColormap(element_colors)
+    
     gempy_vista.surface_points_mesh = poly
     gempy_vista.surface_points_actor = gempy_vista.p.add_mesh(
         mesh=poly,
-        cmap=cmap,  # TODO: Add colors
         scalars='id',
         render_points_as_spheres=render_points_as_spheres,
         point_size=point_size,
-        show_scalar_bar=False
+        show_scalar_bar=False,
+        cmap=custom_cmap,
+        clim=(0, np.unique(vectorize_ids).shape[0])
     )
 
 
 def plot_orientations(
         gempy_vista: GemPyToVista,
         orientations: OrientationsTable,
-        elements_colors: list[str],
+        surface_points: SurfacePointsTable,
         arrows_factor: float,
 ):
     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['vectors'] = orientations_grads
 
-    # TODO: I am still trying to figure out colors and ids in orientations and surface points
-    cmap = get_geo_model_cmap(
-        elements_colors=np.array(elements_colors),
-        reverse=False
+    poly['id'] = _vectorize_ids(
+        mapping_ids=surface_points.ids,
+        ids_to_map=orientations.ids
     )
+    poly['vectors'] = orientations_grads
 
     arrows = poly.glyph(
         orient='vectors',
@@ -118,7 +100,21 @@ def plot_orientations(
 
     gempy_vista.orientations_actor = gempy_vista.p.add_mesh(
         mesh=arrows,
-        cmap=cmap,
+        scalars='id',
         show_scalar_bar=False
     )
     gempy_vista.orientations_mesh = arrows
+
+
+def _vectorize_ids(mapping_ids, ids_to_map):
+    def _mapping_dict(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
+        return mapping_dict
+
+    mapping_dict = _mapping_dict(mapping_ids)
+    mapped_array = np.vectorize(mapping_dict.get)(ids_to_map)  # 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,15 @@ 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,
         # ? scalars=main_scalar, if we prepare the structured grid do we need this arg?
-        show_scalar_bar=True,
-        scalar_bar_args=gempy_vista.scalar_bar_arguments,
+        show_scalar_bar=False,
         interpolate_before_map=True,
         opacity=opacity,
-        clim=_clim,
         **kwargs
     )
 
@@ -98,27 +94,26 @@ 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)
+            block_ = max_lith - (data.lith_block - 1)
+            structured_grid.cell_data['id'] = block_
         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

@@ -2,6 +2,7 @@
 
 import numpy as np
 import pandas as pd
+from matplotlib.colors import ListedColormap
 
 from gempy_viewer.modules.plot_3d.vista import GemPyToVista
 
@@ -23,25 +24,68 @@ def select_surfaces_data(data_df: pd.DataFrame, surfaces: Union[str, list[str]]
     return geometric_data
 
 
-def set_scalar_bar(gempy_vista: GemPyToVista, elements_names: list[str], surfaces_ids: np.ndarray):
+def set_scalar_bar(gempy_vista: GemPyToVista, elements_names: list[str],
+                   surfaces_ids: np.ndarray, custom_colors: list = None):
+    """
+   
+    LookupTable (0x7f3d1dc62e00)
+      Table Range:                (0.0, 2.0)
+      N Values:                   256
+      Above Range Color:          None
+      Below Range Color:          None
+      NAN Color:                  Color(name='darkgray', hex='#a9a9a9ff', opacity=255)
+      Log Scale:                  False
+      Color Map:                  "viridis"
+    """ 
     import pyvista as pv
-    
+
     # Get mapper actor 
-    if gempy_vista.surface_points_actor is not None:
-        mapper_actor: pv.Actor = gempy_vista.surface_points_actor
-    elif gempy_vista.regular_grid_actor is not None:
+    if gempy_vista.regular_grid_actor is not None:
         mapper_actor = gempy_vista.regular_grid_actor
+    elif gempy_vista.surface_points_actor is not None:
+        mapper_actor: pv.Actor = gempy_vista.surface_points_actor
     else:
         return None  # * Not a good mapper for the scalar bar
-    
+
+    # Get the lookup table from the mapper
+    lut = mapper_actor.mapper.lookup_table
+
+    # Create annotations mapping integers to element names
     annotations = {}
-    for e, name in enumerate(elements_names):
-        annotations[e] = name
+    for e, name in enumerate(elements_names[::-1]):
+        # Convert integer to string for the annotation key
+        annotations[str(e)] = name
+
+    # Apply annotations to the lookup table
+    lut.annotations = annotations
+
+    # Set number of colors to match the number of categories
+    n_colors = len(elements_names)
+    lut.n_values = n_colors - 1
 
-    mapper_actor.mapper.lookup_table.annotations = annotations
-    
+    # Apply custom colors if provided
+    if custom_colors is not None:
+        # Check if we have enough colors
+        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)
+        # Apply the custom colormap to the lookup table
+        lut.apply_cmap(cmap=custom_cmap, n_values=n_colors, flip=False)
+
+    else:
+        # Apply a default colormap if no custom colors are provided
+        lut.apply_cmap(cmap='Set1', n_values=n_colors)
+
+    # Configure scalar bar arguments
     sargs = gempy_vista.scalar_bar_arguments
+    min_id, max_id = surfaces_ids.min(), surfaces_ids.max()
+    mapper_actor.mapper.scalar_range = (min_id - .4, max_id + .5)
+
     sargs["mapper"] = mapper_actor.mapper
-    
+    sargs["n_labels"] = 0
+
+    # Add scalar bar
     gempy_vista.p.add_scalar_bar(**sargs)
-    gempy_vista.p.update_scalar_bar_range((surfaces_ids.min(), surfaces_ids.max())) # * This has to be here to now screw the colors with the volumes
+
+    # Update scalar bar range to match surface IDs range

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