Add conversion of prisms or a prism layer to PyVista objects

.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

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::

doc/conf.py

@@ -6,6 +6,7 @@
 #
 import datetime
 
+import pyvista
 from sphinx_gallery.sorting import ExampleTitleSortKey
 
 import harmonica
@@ -30,9 +31,9 @@
     "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 @@
     "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 @@
     "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>'

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/>`__

env/requirements-docs.txt

@@ -6,3 +6,5 @@ boule
 pyproj
 matplotlib
 cartopy>=0.18
+pyvista
+vtk>=9

env/requirements-tests.txt

@@ -3,3 +3,5 @@ boule
 pytest
 pytest-cov
 coverage
+pyvista
+vtk>=9

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

examples/visualization/README.txt

@@ -0,0 +1,2 @@
+3D visualizations
+-----------------

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()

harmonica/forward/prism_layer.py

@@ -13,6 +13,7 @@
 import verde as vd
 import xarray as xr
 
+from ..visualization.prism import prism_to_pyvista
 from .prism import prism_gravity
 
 
@@ -446,3 +447,22 @@ def get_prism(self, indices):
         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)

harmonica/tests/test_prism_layer.py

@@ -17,6 +17,11 @@
 
 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())