Add conversion of prisms or a prism layer to PyVista objects

.github/workflows/docs.yml

@@ -122,7 +122,17 @@ jobs:
           HARMONICA_DATA_DIR: ${{ runner.temp }}/cache/harmonica
 
       - 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

.github/workflows/test.yml

@@ -30,7 +30,7 @@ jobs:
   #############################################################################
   # Run tests and upload to codecov
   test:
-    name: ${{ matrix.os }} py${{ matrix.python }}
+    name: ${{ matrix.os }} py${{ matrix.python }} ${{ matrix.dependencies }}
     runs-on: ${{ matrix.os }}-latest
     strategy:
       # Otherwise, the workflow would stop if a single job fails. We want to
@@ -41,9 +41,11 @@ jobs:
         python: ["3.6", "3.9"]
         # If "optional", will install non-required dependencies in the build
         # environment. Otherwise, only required dependencies are installed.
-        dependencies: [""]
+        dependencies: ["", "optional"]
     env:
       REQUIREMENTS: requirements.txt env/requirements-tests.txt
+      REQUIREMENTS_OPTIONAL: requirements-optional.txt
+      DEPENDENCIES: ${{ matrix.dependencies }}
       # Used to tag codecov submissions
       OS: ${{ matrix.os }}
       PYTHON: ${{ matrix.python }}
@@ -102,6 +104,10 @@ jobs:
                 cat $requirement >> $requirements_file
               done
           fi
+          if [ "$DEPENDENCIES" == "optional" ]; then
+              echo "Capturing optional dependencies from $REQUIREMENTS_OPTIONAL"
+              cat $REQUIREMENTS_OPTIONAL >> $requirements_file
+          fi
           if [ -f $requirements_file ]; then
               echo "Collected dependencies:"
               cat $requirements_file
@@ -159,7 +165,7 @@ jobs:
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
           file: ./coverage.xml
-          env_vars: OS,PYTHON
+          env_vars: OS,PYTHON,DEPENDENCIES
           # Don't mark the job as failed if the upload fails for some reason.
           # It does sometimes but shouldn't be the reason for running
           # everything again unless something else is broken.

doc/conf.py

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

environment.yml

@@ -14,6 +14,9 @@ dependencies:
   - pooch>=0.7.0
   - verde>=1.5.0
   - xarray
+  # Optional requirements
+  - pyvista
+  - vtk>=9
   # Development requirements
   - boule
   - pyproj

examples/visualization/README.txt

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

examples/visualization/prism_layer.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 prisms_to_pyvista
 from .prism import prism_gravity
 
 
@@ -446,3 +447,21 @@ 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: self._obj[data_var] for data_var in self._obj.data_vars
+            }
+        return prisms_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())