Fixing area coloring (#2195)

* Fixing area coloring

* Improving performance

* Fixing tests

* For debugging

* Test without raise.
Adding failing tests

* Avoiding failing tests

* Skipping failing tests

* avoiding failing tests

* Expected fails

* setting tests

* Update tests/test_plotting.py

Co-authored-by: Camille <[email protected]>

* Using pytest-xvfb extension

* test no display_test

* New line

* Revert "test no display_test"

This reverts commit 9dff45e.

* Revert "Using pytest-xvfb extension"

This reverts commit 4e2ebab.

---------

Co-authored-by: Camille <[email protected]>

.github/workflows/ci.yml

@@ -179,6 +179,7 @@ jobs:
           cp ../../../../../source/images/dcb.gif ../../../_images/
           sed -i 's+../../../_images/sphx_glr_composite_dcb_004.gif+../../../_images/dcb.gif+g' composite_dcb.html
           cd ../../../../../../
+
       - name: "Upload HTML Documentation"
         uses: actions/upload-artifact@v3
         with:

src/ansys/mapdl/core/mapdl.py

@@ -16,6 +16,7 @@
 from warnings import warn
 import weakref
 
+from matplotlib.colors import to_rgba
 import numpy as np
 from numpy._typing import DTypeLike
 from numpy.typing import NDArray
@@ -55,6 +56,7 @@
     supress_logging,
     wrap_point_SEL,
 )
+from ansys.mapdl.core.theme import PyMAPDL_cmap
 
 if TYPE_CHECKING:  # pragma: no cover
     from ansys.mapdl.reader import Archive
@@ -1765,11 +1767,14 @@ def aplot(
         show_line_numbering : bool, optional
             Display line numbers when ``vtk=True``.
 
-        color_areas : np.array, optional
+        color_areas : Union[bool, str, np.array], optional
             Only used when ``vtk=True``.
-            If ``color_areas`` is a bool, randomly color areas when ``True`` .
-            If ``color_areas`` is an array or list, it colors each area with
-            the RGB colors, specified in that array or list.
+            If ``color_areas`` is a bool, randomly color areas when ``True``.
+            If ``color_areas`` is a string, it must be a valid color string
+            which will be applied to all areas.
+            If ``color_areas`` is an array or list made of color names (str) or
+            the RGBa numbers ([R, G, B, transparency]), it colors each area with
+            the colors, specified in that array or list.
 
         show_lines : bool, optional
             Plot lines and areas.  Change the thickness of the lines
@@ -1829,36 +1834,63 @@ def aplot(
             meshes = []
             labels = []
 
+            anums = np.unique(surf["entity_num"])
+
             # individual surface isolation is quite slow, so just
             # color individual areas
-            if color_areas:
+            # if isinstance(color_areas, np.ndarray) and len(or len(color_areas):
+            if (isinstance(color_areas, np.ndarray) and len(color_areas) > 1) or (
+                not isinstance(color_areas, np.ndarray) and color_areas
+            ):
                 if isinstance(color_areas, bool):
-                    anum = surf["entity_num"]
-                    size_ = max(anum) + 1
-                    # Because this is only going to be used for plotting purpuses, we don't need to allocate
+                    size_ = len(anums)
+                    # Because this is only going to be used for plotting
+                    # purposes, we don't need to allocate
                     # a huge vector with random numbers (colours).
-                    # By default `pyvista.DataSetMapper.set_scalars` `n_colors` argument is set to 256, so let
-                    # do here the same.
-                    # We will limit the number of randoms values (colours) to 256
+                    # By default `pyvista.DataSetMapper.set_scalars` `n_colors`
+                    # argument is set to 256, so let do here the same.
+                    # We will limit the number of randoms values (colours)
+                    # to 256.
                     #
                     # Link: https://docs.pyvista.org/api/plotting/_autosummary/pyvista.DataSetMapper.set_scalars.html#pyvista.DataSetMapper.set_scalars
                     size_ = min([256, size_])
                     # Generating a colour array,
                     # Size = number of areas.
                     # Values are random between 0 and min(256, number_areas)
-                    area_color = np.random.choice(range(size_), size=(len(anum), 3))
+                    colors = PyMAPDL_cmap(anums)
+
+                elif isinstance(color_areas, str):
+                    # A color is provided as a string
+                    colors = np.atleast_2d(np.array(to_rgba(color_areas)))
+
                 else:
-                    if len(surf["entity_num"]) != len(color_areas):
+                    if len(anums) != len(color_areas):
                         raise ValueError(
-                            f"The length of the parameter array 'color_areas' should be the same as the number of areas."
+                            "The length of the parameter array 'color_areas' "
+                            "should be the same as the number of areas."
+                            f"\nanums: {anums}"
+                            f"\ncolor_areas: {color_areas}"
                         )
-                    area_color = color_areas
-                meshes.append({"mesh": surf, "scalars": area_color})
+
+                    if isinstance(color_areas[0], str):
+                        colors = np.array([to_rgba(each) for each in color_areas])
+                    else:
+                        colors = color_areas
+
+                # mapping mapdl areas to pyvista mesh cells
+                def mapper(each):
+                    if len(colors) == 1:
+                        # for the case colors comes from string.
+                        return colors[0]
+                    return colors[each - 1]
+
+                colors_map = np.array(list(map(mapper, surf["entity_num"])))
+                meshes.append({"mesh": surf, "scalars": colors_map})
+
             else:
                 meshes.append({"mesh": surf, "color": kwargs.get("color", "white")})
 
             if show_area_numbering:
-                anums = np.unique(surf["entity_num"])
                 centers = []
                 for anum in anums:
                     area = surf.extract_cells(surf["entity_num"] == anum)

src/ansys/mapdl/core/plotting.py

@@ -1,7 +1,9 @@
 """Plotting helper for MAPDL using pyvista"""
+from typing import Any, Optional
 from warnings import warn
 
 import numpy as np
+from numpy.typing import NDArray
 
 from ansys.mapdl.core import _HAS_PYVISTA
 from ansys.mapdl.core.misc import get_bounding_box, unique_rows
@@ -420,9 +422,21 @@ def _general_plotter(
         )
 
     for mesh in meshes:
+        scalars: Optional[NDArray[Any]] = mesh.get("scalars")
+
+        if (
+            "scalars" in mesh
+            and scalars.ndim == 2
+            and (scalars.shape[1] == 3 or scalars.shape[1] == 4)
+        ):
+            # for the case we are using scalars for plotting
+            rgb = True
+        else:
+            rgb = False
+
         plotter.add_mesh(
             mesh["mesh"],
-            scalars=mesh.get("scalars"),
+            scalars=scalars,
             scalar_bar_args=scalar_bar_args,
             color=mesh.get("color", color),
             style=mesh.get("style", style),
@@ -442,6 +456,7 @@ def _general_plotter(
             cmap=cmap,
             render_points_as_spheres=render_points_as_spheres,
             render_lines_as_tubes=render_lines_as_tubes,
+            rgb=rgb,
             **add_mesh_kwargs,
         )
 

src/ansys/mapdl/core/theme.py

@@ -40,7 +40,7 @@ class myEmptyClass:
     / 255
 )
 
-PyMAPDL_cmap = ListedColormap(MAPDL_colorbar)
+PyMAPDL_cmap: ListedColormap = ListedColormap(MAPDL_colorbar, name="PyMAPDL", N=255)
 
 
 class MapdlTheme(base_class):

tests/test_plotting.py

@@ -569,6 +569,45 @@ def test_vsel_iterable(mapdl, make_block):
     )
 
 
+def test_color_areas(mapdl, make_block):
+    mapdl.aplot(vtk=True, color_areas=True, return_plotter=True)
+
+
+# This is to remind us that the pl.mesh does not return data for all meshes in CICD.
+@pytest.mark.xfail
+def test_color_areas_fail(mapdl, make_block):
+    pl = mapdl.aplot(vtk=True, color_areas=True, return_plotter=True)
+    assert len(np.unique(pl.mesh.cell_data["Data"], axis=0)) == mapdl.geometry.n_area
+
+
+@skip_no_xserver
+@pytest.mark.parametrize(
+    "color_areas",
+    [
+        ["red", "green", "blue", "yellow", "white", "purple"],
+        [
+            [255, 255, 255],
+            [255, 255, 0],
+            [255, 0, 0],
+            [0, 255, 0],
+            [0, 255, 255],
+            [0, 0, 0],
+        ],
+        255
+        * np.array([[1, 1, 1], [1, 1, 0], [1, 0, 0], [0, 1, 0], [0, 1, 1], [0, 0, 0]]),
+    ],
+)
+def test_color_areas_individual(mapdl, make_block, color_areas):
+    pl = mapdl.aplot(vtk=True, color_areas=color_areas, return_plotter=True)
+    assert len(np.unique(pl.mesh.cell_data["Data"], axis=0)) == len(color_areas)
+
+
+def test_color_areas_error(mapdl, make_block):
+    color_areas = ["red", "green", "blue"]
+    with pytest.raises(ValueError):
+        mapdl.aplot(vtk=True, color_areas=color_areas)
+
+
 def test_WithInterativePlotting(mapdl, make_block):
     mapdl.eplot(vtk=False)
     jobname = mapdl.jobname.upper()