Repository: fatiando/harmonica
Pull request #291 (2021-12-01 19:22:15)

Add conversion of prisms or a prism layer to PyVista objects
Add a new `visualization` module that hosts `prism_to_pyvista`: a function to
convert a set of prisms into a `pyvista.UnstructuredGrid`. Include the new
module and this function in the API Reference. Add a new `to_pyvista()` method
to the `PrismLayer` accessor that converts a prism layer into a `pyvista` grid,
making it easier to plot it in 3D. The `UnstructuredGrid` has the information
about each prism as hexahedrons, along with their physical properties as cell
data. Add tests for the new features. Add `pyvista` and `vtk` as optional
dependencies to `environment.yml` and `setup.cfg`. Add a new example for
plotting a `PrismLayer`. Configure Sphinx to show pyvista plots in the gallery
and to use the `pyvista-plot` directive in docstrings.


**Reminders**:

- [ ] Run `make format` and `make check` to make sure the code follows the style guide.
- [ ] Add tests for new features or tests that would have caught the bug that you're fixing.
- [ ] Add new public functions/methods/classes to `doc/api/index.rst` and the base `__init__.py` file for the package.
- [ ] Write detailed docstrings for all functions/classes/methods. It often helps to design better code if you write the docstrings first.
- [ ] If adding new functionality, add an example to the docstring, gallery, and/or tutorials.
- [ ] Add your full name, affiliation, and [ORCID](https://orcid.org) (optional) to the `AUTHORS.md` file (if you haven't already) in case you'd like to be listed as an author on the [Zenodo](https://zenodo.org/communities/fatiando) archive of the next release.


leouieda: @santisoler how about posting this to the PyVista forum to get feedback from the devs? https://github.com/pyvista/pyvista/discussions
santisoler: Sure! I would love to have feedback from them.
santisoler: Thanks a lot @banesullivan for the review! I really appreciate it!

I'm stepping ahead of time, but after this PR is merged, I want to start drafting a similar way to plot tesseroids with pyvista. That might be more challenging since two of their faces are actually spherical. I'll surely ping you when I have some draft :eyes:
leouieda: @santisoler I'm wondering if a new subpackage is really needed here. We're not planning on adding a bunch of visualizations to this package and this is a bit of an exception since we couldn't avoid it. Would it be better to leave that code in the `prism_layer.py` module for now? Or at least this subpackage should be private (`_visualization`) since we don't want people importing it directly.
santisoler: I'm retaking this PR after a **long** time!

> @santisoler I'm wondering if a new subpackage is really needed here. We're not planning on adding a bunch of visualizations to this package and this is a bit of an exception since we couldn't avoid it. Would it be better to leave that code in the prism_layer.py module for now? Or at least this subpackage should be private (_visualization) since we don't want people importing it directly.

The way the `prisms_to_pyvista` function is written, it would allow users to pass any set of prisms and create an UnstructuredGrid that represent them.
To do so, we should allow them to use this function as `hm.prisms_to_pyvista` or (not a fan) `hm.visualization.prisms_to_pyvista`.

Also, as mentioned before, I'm planning to add visualization functions for tesseroids as well (converstion to PyVista objects), so the `visualization` module sounds a good place for them as well.

If we decide to only allow these conversions through the methods of `PrismLayer` (and the future `TesseroidLayer`), so I'm ok with making `_visualization` private. But since it's already defined for the general purpose of any random set of prisms, I think we should keep it public. What do you think?
santisoler: I've made the `visualization` module public after what we decided on the [Fatiando Call of 2022-04-22](https://github.com/fatiando/community/blob/main/development-calls/2022.md#2022-04-22).

Patch:
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 62cda98..dd6e0aa 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -111,7 +111,17 @@ jobs:
         run: python -m pip install --no-deps dist/*.whl
 
       - name: Build the documentation
-        run: make -C doc clean all
+        run: |
+          # Install xvfb and run some commands to allow pyvista to run on
+          # a headless system.
+          sudo apt-get install xvfb
+          export DISPLAY=:99.0
+          export PYVISTA_OFF_SCREEN=true
+          export PYVISTA_USE_IPYVTK=true
+          Xvfb :99 -screen 0 1024x768x24 > /dev/null 2>&1 &
+          sleep 3
+          # Build the docs
+          make -C doc clean all
 
       # Store the docs as a build artifact so we can deploy it later
       - name: Upload HTML documentation as an artifact
diff --git a/doc/api/index.rst b/doc/api/index.rst
index 40365b8..44f5ab9 100644
--- a/doc/api/index.rst
+++ b/doc/api/index.rst
@@ -55,6 +55,14 @@ Input and Output
 
     load_icgem_gdf
 
+Visualization
+-------------
+
+.. autosummary::
+   :toctree: generated/
+
+    visualization.prism_to_pyvista
+
 Synthetic models and surveys
 ----------------------------
 .. autosummary::
diff --git a/doc/conf.py b/doc/conf.py
index a0c05de..47a82c9 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -6,6 +6,7 @@
 #
 import datetime
 
+import pyvista
 from sphinx_gallery.sorting import ExampleTitleSortKey
 
 import harmonica
@@ -30,9 +31,9 @@ extensions = [
     "sphinx.ext.viewcode",
     "sphinx.ext.extlinks",
     "sphinx.ext.intersphinx",
-    "matplotlib.sphinxext.plot_directive",
     "sphinx.ext.napoleon",
     "sphinx_gallery.gen_gallery",
+    "pyvista.ext.plot_directive",
 ]
 
 # Configuration to include links to other project docs when referencing
@@ -49,6 +50,7 @@ intersphinx_mapping = {
     "verde": ("https://www.fatiando.org/verde/latest/", None),
     "boule": ("https://www.fatiando.org/boule/latest/", None),
     "matplotlib": ("https://matplotlib.org/", None),
+    "pyvista": ("https://docs.pyvista.org", None),
 }
 
 # Autosummary pages will be generated by sphinx-autogen instead of sphinx-build
@@ -94,9 +96,22 @@ sphinx_gallery_conf = {
     "doc_module": "harmonica",
     # Insert links to documentation of objects in the examples
     "reference_url": {"harmonica": None},
+    # Add pyvista to the image scrapers
+    "image_scrapers": ("pyvista", "matplotlib"),
 }
 
 
+# Pyvista configurations
+# -----------------------------------------------------------------------------
+# necessary when building the sphinx gallery
+pyvista.BUILDING_GALLERY = True
+pyvista.OFF_SCREEN = True
+
+# Optional - set parameters like theme or window size
+pyvista.set_plot_theme("document")
+pyvista.global_theme.window_size = (1024 * 2, 768 * 2)
+
+
 # HTML output configuration
 # -----------------------------------------------------------------------------
 html_title = f'{project} <span class="project-version">{version}</span>'
diff --git a/doc/install.rst b/doc/install.rst
index 75e0154..503fbec 100644
--- a/doc/install.rst
+++ b/doc/install.rst
@@ -19,6 +19,12 @@ doesn't interfere with any other Python installations in your system.
 Dependencies
 ------------
 
+The required dependencies should be installed automatically when you install
+Harmonica using ``conda`` or ``pip``. Optional dependencies have to be
+installed manually.
+
+Required:
+
 * `numpy <http://www.numpy.org/>`__
 * `pandas <http://pandas.pydata.org/>`__
 * `numba <https://numba.pydata.org/>`__
@@ -28,6 +34,12 @@ Dependencies
 * `pooch <http://www.fatiando.org/pooch/>`__
 * `verde <http://www.fatiando.org/verde/>`__
 
+Optional:
+
+* `pyvista <https://www.pyvista.org/>`__ and
+  `vtk <https://vtk.org/>`__ (>= 9): for 3D visualizations.
+  See :meth:`harmonica.DatasetAccessorPrismLayer.to_pyvista`.
+
 The examples in the :ref:`gallery` also use:
 
 * `boule <http://www.fatiando.org/boule/>`__
diff --git a/env/requirements-docs.txt b/env/requirements-docs.txt
index c291df0..d8b868a 100644
--- a/env/requirements-docs.txt
+++ b/env/requirements-docs.txt
@@ -6,3 +6,5 @@ boule
 pyproj
 matplotlib
 cartopy>=0.18
+pyvista
+vtk>=9
diff --git a/environment.yml b/environment.yml
index 1e6e22f..afbec49 100644
--- a/environment.yml
+++ b/environment.yml
@@ -19,6 +19,9 @@ dependencies:
   - pooch>=0.7.0
   - verde>=1.5.0
   - xarray
+  # Optional requirements
+  - pyvista
+  - vtk>=9
   # Testing requirements
   - pytest
   - pytest-cov
diff --git a/examples/visualization/README.txt b/examples/visualization/README.txt
new file mode 100644
index 0000000..e684657
--- /dev/null
+++ b/examples/visualization/README.txt
@@ -0,0 +1,2 @@
+3D visualizations
+-----------------
diff --git a/examples/visualization/prism_layer_pyvista.py b/examples/visualization/prism_layer_pyvista.py
new file mode 100644
index 0000000..7406095
--- /dev/null
+++ b/examples/visualization/prism_layer_pyvista.py
@@ -0,0 +1,89 @@
+# Copyright (c) 2018 The Harmonica Developers.
+# Distributed under the terms of the BSD 3-Clause License.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# This code is part of the Fatiando a Terra project (https://www.fatiando.org)
+#
+"""
+Plot a prism layer in 3D
+========================
+
+The :func:`harmonica.prism_layer` allows to easily create a layer of prisms
+whose top and bottom boundaries might drape a certain surface, like topography,
+bathymetry or the Moho discontinuity. It returns a :class:`xarray.Dataset`
+object with the horizontal coordinates of the center of each prism, their top
+and bottom boundaries and their physical properties, like their density.
+Through the ``prism_layer`` accessor (see
+:class:`harmonica.DatasetAccessorPrismLayer`) we can call some methods for our
+prism layer. For example, the
+:meth:`harmonica.DatasetAccessorPrismLayer.gravity` method lets us compute the
+gravitational fields of the layer on any set of observation points.
+Another interesting method is the
+:meth:`harmonica.DatasetAccessorPrismLayer.to_pyvista`: it converts the prism
+layer into a :class:`pyvista.UnstructuredGrid` that could be easily plotted
+through ``pyvista``.
+
+In this example we will show how we can build a prism layer out of a topography
+and bathymetry grid for South Africa and how we can visualize the layer as a 3D
+plot using ``pyvista``.
+
+"""
+import pyproj
+import pyvista as pv
+import verde as vd
+
+import harmonica as hm
+
+# Read South Africa topography
+south_africa_topo = hm.datasets.fetch_south_africa_topography()
+
+# Project the grid
+projection = pyproj.Proj(proj="merc", lat_ts=south_africa_topo.latitude.values.mean())
+south_africa_topo = vd.project_grid(south_africa_topo.topography, projection=projection)
+
+# Create a 2d array with the density of the prisms Points above the geoid will
+# have a density of 2670 kg/m^3 Points below the geoid will have a density
+# contrast equal to the difference between the density of the ocean and the
+# density of the upper crust: # 1000 kg/m^3 - 2900 kg/m^3
+density = south_africa_topo.copy()  # copy topography to a new xr.DataArray
+density.values[:] = 2670.0  # replace every value for the density of the topography
+# Change density values of ocean points
+density = density.where(south_africa_topo >= 0, 1000 - 2900)
+
+# Create layer of prisms
+prisms = hm.prism_layer(
+    (south_africa_topo.easting, south_africa_topo.northing),
+    surface=south_africa_topo,
+    reference=0,
+    properties={"density": density},
+)
+
+# Create a pyvista UnstructuredGrid from the prism layer
+pv_grid = prisms.prism_layer.to_pyvista()
+pv_grid
+
+###############################################################################
+
+# Plot with pyvista
+plotter = pv.Plotter(lighting="three_lights", window_size=(1000, 800))
+plotter.add_mesh(pv_grid, scalars="density")
+plotter.set_scale(zscale=75)  # exaggerate the vertical coordinate
+plotter.camera_position = "xz"
+plotter.camera.elevation = 20
+plotter.camera.azimuth = 35
+plotter.camera.zoom(1.2)
+
+# Add a ceiling light
+west, east, south, north = vd.get_region((prisms.easting, prisms.northing))
+easting_center, northing_center = (east + west) / 2, (north + south) / 2
+light = pv.Light(
+    position=(easting_center, northing_center, 10e3),
+    focal_point=(easting_center, northing_center, 0),
+    intensity=0.3,
+    light_type="scene light",  # the light doesn't move with the camera
+    positional=False,  # the light comes from infinity
+)
+plotter.add_light(light)
+
+plotter.show_axes()
+plotter.show()
diff --git a/harmonica/forward/prism_layer.py b/harmonica/forward/prism_layer.py
index 8063f67..a23a21c 100644
--- a/harmonica/forward/prism_layer.py
+++ b/harmonica/forward/prism_layer.py
@@ -13,6 +13,7 @@ import numpy as np
 import verde as vd
 import xarray as xr
 
+from ..visualization.prism import prism_to_pyvista
 from .prism import prism_gravity
 
 
@@ -446,3 +447,22 @@ class DatasetAccessorPrismLayer:
         bottom = self._obj.bottom.values[indices]
         top = self._obj.top.values[indices]
         return west, east, south, north, bottom, top
+
+    def to_pyvista(self):
+        """
+        Return a pyvista UnstructuredGrid to plot the PrismLayer
+
+        Returns
+        -------
+        pv_grid : :class:`pyvista.UnstructuredGrid`
+            :class:`pyvista.UnstructuredGrid` containing each prism of the
+            layer as a hexahedron along with their properties.
+        """
+        prisms = self._to_prisms()
+        properties = None
+        if self._obj.data_vars:
+            properties = {
+                data_var: np.asarray(self._obj[data_var]).ravel()
+                for data_var in self._obj.data_vars
+            }
+        return prism_to_pyvista(prisms, properties=properties)
diff --git a/harmonica/visualization/__init__.py b/harmonica/visualization/__init__.py
new file mode 100644
index 0000000..d08adbd
--- /dev/null
+++ b/harmonica/visualization/__init__.py
@@ -0,0 +1,7 @@
+# Copyright (c) 2018 The Harmonica Developers.
+# Distributed under the terms of the BSD 3-Clause License.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# This code is part of the Fatiando a Terra project (https://www.fatiando.org)
+#
+from .prism import prism_to_pyvista
diff --git a/harmonica/visualization/prism.py b/harmonica/visualization/prism.py
new file mode 100644
index 0000000..1553bdd
--- /dev/null
+++ b/harmonica/visualization/prism.py
@@ -0,0 +1,215 @@
+# Copyright (c) 2018 The Harmonica Developers.
+# Distributed under the terms of the BSD 3-Clause License.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# This code is part of the Fatiando a Terra project (https://www.fatiando.org)
+#
+"""
+Functions for visualizing prisms through pyvista
+"""
+import numpy as np
+
+try:
+    import pyvista
+except ImportError:
+    pyvista = None
+else:
+    import vtk
+
+
+def prism_to_pyvista(prisms, properties=None):
+    """
+    Create a :class:`pyvista.UnstructuredGrid` out of prisms
+
+    Builds a :class:`pyvista.UnstructuredGrid` out of a set of prisms that
+    could be used to plot a 3D representation through :mod:`pyvista`.
+
+    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 property will be added to the
+        :class:`pyvista.UnstructuredGrid`.
+        Default to None.
+
+    Returns
+    -------
+    pv_grid : :class:`pyvista.UnstructuredGrid`
+        :class:`pyvista.UnstructuredGrid` that represents the prisms with their
+        properties (if any).
+
+    Examples
+    --------
+
+    .. pyvista-plot::
+
+       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:`pyvista.UnstructuredGrid` out of them:
+
+       >>> import harmonica as hm
+       >>> pv_grid = hm.visualization.prism_to_pyvista(
+       ...     prisms, properties={"density": densities}
+       ... )
+       >>> pv_grid # doctest: +SKIP
+
+       Plot it using :mod:`pyvista`:
+
+       >>> pv_grid.plot() # doctest: +SKIP
+
+    """
+    # Check if pyvista are installed
+    if pyvista is None:
+        raise ImportError(
+            "Missing optional dependency 'pyvista' required for building pyvista 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])
+    # Build the UnstructuredGrid
+    pv_grid = pyvista.UnstructuredGrid({vtk.VTK_HEXAHEDRON: cells}, vertices)
+    # 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
+            prop = np.atleast_1d(prop)
+            if prop.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
+            pv_grid.cell_data[name] = prop
+    return pv_grid
+
+
+def _prisms_boundaries_to_vertices(prisms):
+    """
+    Converts prisms boundaries to sets of vertices for each prism
+
+    The vertices for each prism will be in the following order
+
+    .. code-block::
+
+                                7-------6
+                               /|      /|
+                              / |     / |
+                             4-------5  |
+        up   northing        |  |    |  |
+        ^   ~                |  3----|--2
+        |  /                 | /     | /
+        | /                  |/      |/
+         ------> easting     0-------1
+
+    So the vertices of a single prism should be like:
+
+    .. code-block:: python
+
+        [w, s, bottom],
+        [e, s, bottom],
+        [e, n, bottom],
+        [w, n, bottom],
+        [w, s, top],
+        [e, s, top],
+        [e, n, top],
+        [w, n, top],
+
+
+    Parameters
+    ----------
+    prisms : 2d-array
+        2d-array with the boundaries of a set of prisms. Each row of the array
+        should contain the boundaries of a single prism in the following
+        order: ``west``, ``east``, ``south``, ``north``, ``bottom``, ``top``.
+
+    Returns
+    -------
+    vertices : 2d-array
+        2D array with the coordinates of the vertices of the prism. Each row of
+        the array corresponds to the coordinate of a single vertex in the
+        following order: ``easting``, ``northing``, ``upward``.
+        The shape of this array is ``(M, 3)``, where ``M`` is the total number
+        of vertices in the whole set of prisms (number of prisms times 8).
+        The order of the vertices is fixed to be compatible with VTK.
+
+    Examples
+    --------
+    >>> _prisms_boundaries_to_vertices(np.array([[-1, 1, -2, 2, -3, 3]]))
+    array([[-1., -2., -3.],
+           [ 1., -2., -3.],
+           [ 1.,  2., -3.],
+           [-1.,  2., -3.],
+           [-1., -2.,  3.],
+           [ 1., -2.,  3.],
+           [ 1.,  2.,  3.],
+           [-1.,  2.,  3.]])
+    >>> _prisms_boundaries_to_vertices(
+    ...     np.array([[-1, 1, -2, 2, -3, 3], [-4, 4, -5, 5, -6, 6]])
+    ... )
+    array([[-1., -2., -3.],
+           [ 1., -2., -3.],
+           [ 1.,  2., -3.],
+           [-1.,  2., -3.],
+           [-1., -2.,  3.],
+           [ 1., -2.,  3.],
+           [ 1.,  2.,  3.],
+           [-1.,  2.,  3.],
+           [-4., -5., -6.],
+           [ 4., -5., -6.],
+           [ 4.,  5., -6.],
+           [-4.,  5., -6.],
+           [-4., -5.,  6.],
+           [ 4., -5.,  6.],
+           [ 4.,  5.,  6.],
+           [-4.,  5.,  6.]])
+    """
+    # Get number of prisms
+    n_prisms = prisms.shape[0]
+
+    # Allocate vertices array
+    vertices = np.empty((n_prisms, 8, 3))
+
+    # Define a dictionary with the indices of the vertices that contain each
+    # boundary of the prism.
+    # For example, the west boundary is present only in the vertices
+    # number 0, 3, 4 and 7.
+    indices = {
+        "west": (0, 3, 4, 7),
+        "east": (1, 2, 5, 6),
+        "south": (0, 1, 4, 5),
+        "north": (2, 3, 6, 7),
+        "bottom": (0, 1, 2, 3),
+        "top": (4, 5, 6, 7),
+    }
+
+    # Assign the values to each vertex
+    for i, boundary in enumerate(indices):
+        # Determine at which component of the vertices should the current
+        # boundary be assigned to.
+        #   The west and east (i = 0 and i = 1) should go to the component 0.
+        #   The south and north (i = 2 and i = 3) should go to the component 1.
+        #   The bottom and top (i = 4 and i = 5) should go to the component 2.
+        component = i // 2
+        # Assign vertices components
+        for vertex in indices[boundary]:
+            vertices[:, vertex, component] = prisms[:, i]
+
+    # Reshape the vertices array so it has (M, 3) shape, where M is the total
+    # number of vertices.
+    return vertices.reshape((n_prisms * 8, 3))
diff --git a/setup.cfg b/setup.cfg
index 8c7b7eb..299de10 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -49,6 +49,11 @@ install_requires =
     xarray
     verde>=1.5.0
 
+[options.extras_require]
+visualizations =
+    pyvista
+    vtk>=9
+
 [options.package_data]
 harmonica.tests = data/*, baseline/*
 harmonica.datasets = registry.txt
@@ -82,6 +87,9 @@ rst-roles =
     mod,
     meth,
     ref,
+# Add some directives used in our docstrings
+rst-directives =
+    pyvista-plot,
 # Ignore "Unknown target name" raised on citations
 extend-ignore = RST306
 


Test patch:
diff --git a/env/requirements-tests.txt b/env/requirements-tests.txt
index a968055..605da36 100644
--- a/env/requirements-tests.txt
+++ b/env/requirements-tests.txt
@@ -3,3 +3,5 @@ boule
 pytest
 pytest-cov
 coverage
+pyvista
+vtk>=9
diff --git a/harmonica/tests/test_prism_layer.py b/harmonica/tests/test_prism_layer.py
index e531d5e..ddcbe57 100644
--- a/harmonica/tests/test_prism_layer.py
+++ b/harmonica/tests/test_prism_layer.py
@@ -17,6 +17,11 @@ import xarray as xr
 
 from .. import prism_gravity, prism_layer
 
+try:
+    import pyvista
+except ImportError:
+    pyvista = None
+
 
 @pytest.fixture(params=("numpy", "xarray"))
 def dummy_layer(request):
@@ -385,3 +390,35 @@ def test_prism_layer_gravity_density_nans(field, dummy_layer, prism_layer_with_h
         result,
         prism_gravity(coordinates, prisms, rho, field=field),
     )
+
+
+@pytest.mark.skipif(pyvista is None, reason="requires pyvista")
+@pytest.mark.parametrize("properties", (False, True))
+def test_to_pyvista(dummy_layer, properties):
+    """
+    Test the conversion of the prism layer to pyvista.UnstructuredGrid
+    """
+    (easting, northing), surface, reference, density = dummy_layer
+    # Build the layer with or without properties
+    if properties:
+        properties = {"density": density}
+    else:
+        properties = None
+    layer = prism_layer((easting, northing), surface, reference, properties=properties)
+    # Convert the layer to pyvista UnstructuredGrid
+    pv_grid = layer.prism_layer.to_pyvista()
+    # Check properties of the pyvista grid
+    assert pv_grid.n_cells == 20
+    assert pv_grid.n_points == 20 * 8
+    # Check coordinates of prisms
+    for i, prism in enumerate(layer.prism_layer._to_prisms()):
+        npt.assert_allclose(prism, pv_grid.cell_bounds(i))
+    # Check properties of the prisms
+    if properties is None:
+        assert pv_grid.n_arrays == 0
+        assert pv_grid.array_names == []
+    else:
+        assert pv_grid.n_arrays == 1
+        assert pv_grid.array_names == ["density"]
+        assert pv_grid.get_array("density").ndim == 1
+        npt.assert_allclose(pv_grid.get_array("density"), layer.density.values.ravel())
diff --git a/harmonica/tests/test_visualizations.py b/harmonica/tests/test_visualizations.py
new file mode 100644
index 0000000..cad20a3
--- /dev/null
+++ b/harmonica/tests/test_visualizations.py
@@ -0,0 +1,128 @@
+# Copyright (c) 2018 The Harmonica Developers.
+# Distributed under the terms of the BSD 3-Clause License.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# This code is part of the Fatiando a Terra project (https://www.fatiando.org)
+#
+"""
+Test functions from the visualization module.
+"""
+from unittest.mock import patch
+
+import numpy as np
+import numpy.testing as npt
+import pytest
+import xarray as xr
+
+from ..visualization.prism import prism_to_pyvista
+
+try:
+    import pyvista
+except ImportError:
+    pyvista = None
+
+
+@pytest.mark.skipif(pyvista is not None, reason="pyvista must be missing")
+def test_prism_to_pyvista_missing_pyvista():
+    """
+    Check error raise after calling prism_to_pyvista when pyvista is missing
+    """
+    prism = [0, 1, 0, 1, 0, 1]
+    with pytest.raises(ImportError) as exception:
+        prism_to_pyvista(prism)
+    assert "'pyvista'" in str(exception.value)
+
+
+@pytest.fixture(name="prisms")
+def fixture_prisms():
+    """
+    Define a set of sample prisms
+    """
+    return [
+        [1e3, 5e3, -2e3, 4e3, -10e3, -7e3],
+        [200, 500, 600, 900, 100, 300],
+        [-2e3, 0, 5e3, 15e3, -4e3, 2e3],
+        [-10e3, -8e3, -10e3, -8e3, -10e3, -8e3],
+    ]
+
+
+@pytest.fixture(name="density", params=("numpy", "xarray"))
+def fixture_density(request):
+    """
+    Return a density array for the sample prisms
+    """
+    density = [2670.0, 2900.0, 3300.0, 3200.0]
+    if request.param == "xarray":
+        density = xr.DataArray(density)
+    return density
+
+
+@pytest.fixture(name="susceptibility", params=("numpy", "xarray"))
+def fixture_susceptibility(request):
+    """
+    Return a susceptibility array for the sample prisms
+    """
+    susceptibility = [2e-4, 2.5e-4, 5e-4, 1e-3]
+    if request.param == "xarray":
+        susceptibility = xr.DataArray(susceptibility)
+    return susceptibility
+
+
+@pytest.mark.skipif(pyvista is None, reason="requires pyvista")
+def test_prism_to_pyvista(prisms, density):
+    """
+    Test output of prism_to_pyvista
+    """
+    pv_grid = prism_to_pyvista(prisms, properties={"density": density})
+    assert pv_grid.n_cells == 4
+    assert pv_grid.n_points == 32
+    # Check coordinates of prisms
+    for i, prism in enumerate(prisms):
+        npt.assert_allclose(prism, pv_grid.cell_bounds(i))
+    # Check properties of the prisms
+    assert pv_grid.n_arrays == 1
+    assert pv_grid.array_names == ["density"]
+    npt.assert_allclose(pv_grid.get_array("density"), density)
+
+
+@pytest.mark.skipif(pyvista is None, reason="requires pyvista")
+@pytest.mark.parametrize("n_props", [0, 1, 2])
+def test_prism_to_pyvista_properties(n_props, prisms, density, susceptibility):
+    """
+    Test prism_to_pyvista for different number of properties
+    """
+    properties = None
+    if n_props == 1:
+        properties = {"density": density}
+    elif n_props == 2:
+        properties = {"density": density, "susceptibility": susceptibility}
+    pv_grid = prism_to_pyvista(prisms, properties=properties)
+    # Check if the grid has the right properties
+    if properties is None:
+        assert pv_grid.n_arrays == 0
+        assert pv_grid.array_names == []
+    else:
+        assert pv_grid.n_arrays == len(properties)
+        assert set(pv_grid.array_names) == set(properties.keys())
+        for prop in properties:
+            npt.assert_allclose(pv_grid.get_array(prop), properties[prop])
+
+
+@pytest.mark.skipif(pyvista is None, reason="requires pyvista")
+def test_prism_to_pyvista_error_2d_property(prisms, density):
+    """
+    Test if prism_to_pyvista raises error on property as 2D array
+    """
+    density_2d = np.array(density).reshape((2, 2))
+    with pytest.raises(ValueError):
+        prism_to_pyvista(prisms, properties={"density": density_2d})
+
+
+@patch("harmonica.visualization.prism.pyvista", None)
+def test_prisms_pyvista_missing_error(prisms, density):
+    """
+    Check if prism_to_pyvista raises error if pyvista is missing
+    """
+    # Check if error is raised
+    with pytest.raises(ImportError):
+        prism_to_pyvista(prisms, properties={"density": density})