Add function to convert prism to vedo grids

Add method to the prism layer to convert the full prism layer to a vedo
grid.

Author: santisoler

doc/api/index.rst

@@ -127,6 +127,7 @@ Visualization
    :toctree: generated/
 
     visualization.prism_to_pyvista
+    visualization.prism_to_vedo
 
 Utilities
 ---------

env/requirements-docs.txt

@@ -13,6 +13,7 @@ ensaio>=0.5.*
 netcdf4
 matplotlib
 pyvista
+vedo
 vtk>=9
 pygmt
 gmt

environment.yml

@@ -21,6 +21,7 @@ dependencies:
   - pyvista
   - vtk
   - numba-progress
+  - vedo
   # Testing requirements
   - pytest
   - pytest-cov

harmonica/_forward/prism_layer.py

@@ -13,7 +13,7 @@
 import verde as vd
 import xarray as xr
 
-from ..visualization import prism_to_pyvista
+from ..visualization import prism_to_pyvista, prism_to_vedo
 from .prism_gravity import prism_gravity
 
 
@@ -482,7 +482,7 @@ def get_prism(self, indices):
 
     def to_pyvista(self, drop_null_prisms=True):
         """
-        Return a pyvista UnstructuredGrid to plot the PrismLayer
+        Return a :class:`pyvista.UnstructuredGrid` to plot the prism layer.
 
         Parameters
         ----------
@@ -495,10 +495,58 @@ def to_pyvista(self, drop_null_prisms=True):
 
         Returns
         -------
-        pv_grid : :class:`pyvista.UnstructuredGrid`
+        pyvista.UnstructuredGrid
             :class:`pyvista.UnstructuredGrid` containing each prism of the
             layer as a hexahedron along with their properties.
         """
+        return self._to_vtk_object(backend="pyvista", drop_null_prisms=drop_null_prisms)
+
+    def to_vedo(self, drop_null_prisms=True):
+        """
+        Return a :class:`vedo.UnstructuredGrid` to plot the prism layer.
+
+        Parameters
+        ----------
+        drop_null_prisms : bool (optional)
+            If True, prisms with zero volume or with any :class:`numpy.nan` as
+            their top or bottom boundaries won't be included in the
+            :class:`vedo.UnstructuredGrid`.
+            If False, every prism in the layer will be included.
+            Default True.
+
+        Returns
+        -------
+        vedo.UnstructuredGrid
+            :class:`vedo.UnstructuredGrid` containing each prism of the
+            layer as a hexahedron along with their properties.
+        """
+        return self._to_vtk_object(backend="vedo", drop_null_prisms=drop_null_prisms)
+
+    def _to_vtk_object(self, backend: str, drop_null_prisms=True):
+        """
+        Return a pyvista or vedo UnstructuredGrid to plot the prism layer.
+        Return a pyvista UnstructuredGrid to plot the PrismLayer
+
+        Parameters
+        ----------
+        backend : {"pyvista", "vedo"}
+            Choose which type of UnstructuredGrid to return.
+        drop_null_prisms : bool (optional)
+            If True, prisms with zero volume or with any :class:`numpy.nan` as
+            their top or bottom boundaries won't be included in the
+            ``UnstructuredGrid``.
+            If False, every prism in the layer will be included.
+            Default True.
+
+        Returns
+        -------
+        pyvista.UnstructuredGrid or vedo.UnstructuredGrid
+            ``UnstructuredGrid`` containing each prism of the layer as
+            a hexahedron along with their properties.
+        """
+        if backend not in ("vedo", "pyvista"):
+            msg = f"Invalid backend '{backend}'."
+            raise ValueError(msg)
         prisms = self._to_prisms()
         null_prisms = np.zeros_like(prisms[:, 0], dtype=bool)
         if drop_null_prisms:
@@ -514,7 +562,8 @@ def to_pyvista(self, drop_null_prisms=True):
                 ]
                 for data_var in self._obj.data_vars
             }
-        return prism_to_pyvista(prisms, properties=properties)
+        func = prism_to_pyvista if backend == "pyvista" else prism_to_vedo
+        return func(prisms, properties=properties)
 
 
 def _discard_thin_prisms(

harmonica/visualization/__init__.py

@@ -4,4 +4,4 @@
 #
 # This code is part of the Fatiando a Terra project (https://www.fatiando.org)
 #
-from ._prism import prism_to_pyvista
+from ._prism import prism_to_pyvista, prism_to_vedo

harmonica/visualization/_prism.py

@@ -5,18 +5,108 @@
 # This code is part of the Fatiando a Terra project (https://www.fatiando.org)
 #
 """
-Functions for visualizing prisms through pyvista
+Functions for visualizing prisms in 3D.
 """
 import numpy as np
 
 try:
     import pyvista
 except ImportError:
     pyvista = None
-else:
+
+try:
+    import vedo
+except ImportError:
+    vedo = None
+
+if pyvista is not None or vedo is not None:
     import vtk
 
 
+def prism_to_vedo(prisms, properties=None):
+    """
+    Create a :class:`vedo.UnstructuredGrid` out of prisms.
+
+    Builds a :class:`vedo.UnstructuredGrid` out of a set of prisms that
+    could be used to plot a 3D representation through :mod:`vedo`.
+
+    Parameters
+    ----------
+    prisms : list or 2d-array
+        List or 2d-array with the boundaries of the prisms.
+        Each row contains the boundaries of each prism in the following order:
+        ``west``, ``east``, ``south``, ``north``, ``bottom``, ``top``.
+    properties : dict or None (optional)
+        Dictionary with the physical properties of the prisms.
+        Each key should be a string and its corresponding value a 1D array.
+        If None, no celldata will be added to the
+        :class:`vedo.UnstructuredGrid`.
+        Default to None.
+
+    Returns
+    -------
+    vedo.UnstructuredGrid
+        :class:`vedo.UnstructuredGrid` that represents the prisms with their
+        properties (if any).
+
+    Examples
+    --------
+    Define a set of prisms and their densities:
+
+    >>> prisms = [
+    ...     [0, 4, 0, 5, -10, 0],
+    ...     [0, 4, 7, 9, -12, -3],
+    ...     [6, 9, 2, 6, -7, 3],
+    ... ]
+    >>> densities = [2900, 3000, 2670]
+
+    Generate a :class:`vedo.UnstructuredGrid` out of them:
+
+    >>> import harmonica as hm
+    >>> vedo_grid = hm.visualization.prism_to_vedo(
+    ...     prisms, properties={"density": densities}
+    ... )
+    >>> print(vedo_grid) # doctest: +SKIP
+    vedo.grids.UnstructuredGrid at (0x56070a0af5b0)
+    nr. of verts  : 24
+    nr. of cells  : 3
+    cell types    : ['HEXAHEDRON']
+    size          : average=6.64122, diagonal=19.6723
+    center of mass: (3.83333, 4.83333, -4.83333)
+    bounds        : x=(0, 9.00), y=(0, 9.00), z=(-12.0, 3.00)
+    celldata *    : "density" (int64), ndim=1, range=(2.67e+3, 3.00e+3)
+    """
+    # Check if vedo are installed
+    if vedo is None:
+        raise ImportError(
+            "Missing optional dependency 'vedo' required for building vedo grids."
+        )
+    # Get prisms and number of prisms
+    prisms = np.atleast_2d(prisms)
+    n_prisms = prisms.shape[0]
+    # Get the vertices of the prisms
+    vertices = _prisms_boundaries_to_vertices(prisms)
+    # Generate the cells for the hexahedrons
+    cells = np.arange(n_prisms * 8).reshape([n_prisms, 8])
+    # Define celltypes
+    celltypes = [vtk.VTK_HEXAHEDRON for _ in range(n_prisms)]
+    # Build the UnstructuredGrid
+    grid = vedo.UnstructuredGrid([vertices, cells, celltypes])
+    # Add properties to the grid
+    if properties is not None:
+        for name, prop in properties.items():
+            # Check if the property is given as 1d array
+            property = np.atleast_1d(prop)
+            if property.ndim > 1:
+                raise ValueError(
+                    f"Multidimensional array found in '{name}' property. "
+                    + "Please, pass prism properties as 1d arrays."
+                )
+            # Assign the property to the cell_data
+            grid.celldata[name] = property
+    return grid
+
+
 def prism_to_pyvista(prisms, properties=None):
     """
     Create a :class:`pyvista.UnstructuredGrid` out of prisms

pyproject.toml

@@ -47,8 +47,9 @@ dependencies = [
 "Source Code" = "https://github.com/fatiando/harmonica"
 
 [project.optional-dependencies]
-visualizations = ["pyvista>=0.27", "vtk>=9"]
+pyvista = ["pyvista>=0.27", "vtk>=9"]
 numba_progress = ["numba_progress"]
+vedo = ["vedo"]
 
 [build-system]
 requires = ["setuptools>=61", "wheel", "setuptools_scm[toml]>=8.0.3"]