[WIP] Towards fixing the bar

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_volumes),
             **kwargs_plot_structured_grid
         )
 
@@ -190,11 +191,12 @@ def plot_3d(
             **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/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,56 @@ 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):
     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:
         mapper_actor = gempy_vista.regular_grid_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
+
+    # 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}")
 
-    mapper_actor.mapper.lookup_table.annotations = annotations
-    
+        custom_cmap = ListedColormap(custom_colors)
+
+        # Apply the custom colormap to the lookup table
+        lut.apply_cmap(cmap=custom_cmap, n_values=n_colors)
+
+    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
     sargs["mapper"] = mapper_actor.mapper
-    
+
+    # 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
+    min_id, max_id = surfaces_ids.min(), surfaces_ids.max()
+    gempy_vista.p.update_scalar_bar_range((min_id, max_id))

tests/conftest.py

@@ -32,7 +32,7 @@ def one_fault_model_topo_solution() -> GeoModel:
         topography_resolution=np.array([60, 60])
     )
 
-    gp.compute_model(one_fault_model)
+    gp.compute_model(one_fault_model, gp.data.GemPyEngineConfig(backend=gp.data.AvailableBackends.PYTORCH))
     return one_fault_model
 
 
@@ -43,7 +43,7 @@ def _one_fault_model_generator():
         project_name='fault',
         extent=[0, 2000, 0, 2000, 0, 1000],
         resolution=[20, 5, 20],
-        refinement=4,
+        refinement=1,
         importer_helper=ImporterHelper(
             path_to_surface_points=path_to_data + "model5_surface_points.csv",
             path_to_orientations=path_to_data + "model5_orientations.csv",