Merge pull request #29 from banesullivan/master

Add PyVista for 3D plotting in Python

AUTHORS.md

@@ -13,3 +13,4 @@ and was created by following people.
 ## Contributors (in order of contributions)
 
 - Falk Heße, Email: <[email protected]>
+- Bane Sullivan, GitHub: [@banesullivan](https://github.com/banesullivan)

README.md

@@ -104,7 +104,7 @@ srf.plot()
 <img src="https://raw.githubusercontent.com/GeoStat-Framework/GSTools/master/docs/source/pics/gau_field.png" alt="Random field" width="600px"/>
 </p>
 
-A similar example but for a three dimensional field is exported to a [VTK](https://vtk.org/) file, which can be visualized with [ParaView](https://www.paraview.org/).
+A similar example but for a three dimensional field is exported to a [VTK](https://vtk.org/) file, which can be visualized with [ParaView](https://www.paraview.org/) or [PyVista](https://docs.pyvista.org) in Python:
 
 ```python
 from gstools import SRF, Gaussian
@@ -114,7 +114,10 @@ x = y = z = range(100)
 model = Gaussian(dim=3, var=0.6, len_scale=20)
 srf = SRF(model)
 srf((x, y, z), mesh_type='structured')
-srf.vtk_export('3d_field')
+srf.vtk_export('3d_field') # Save to a VTK file for ParaView
+
+mesh = srf.to_pyvista() # Create a PyVista mesh for plotting in Python
+mesh.threshold_percent(0.5).plot()
 ```
 
 <p align="center">
@@ -259,21 +262,25 @@ yielding
 [kraichnan_link]: https://doi.org/10.1063/1.1692799
 
 
-## VTK Export
+## VTK/PyVista Export
 
 After you have created a field, you may want to save it to file, so we provide
-a handy [VTK][vtk_link] export routine:
+a handy [VTK][vtk_link] export routine using the `.vtk_export()` or you could
+create a VTK/PyVista dataset for use in Python with to `.to_pyvista()` method:
 
 ```python
 from gstools import SRF, Gaussian
 x = y = range(100)
 model = Gaussian(dim=2, var=1, len_scale=10)
 srf = SRF(model)
 srf((x, y), mesh_type='structured')
-srf.vtk_export("field")
+srf.vtk_export("field") # Saves to a VTK file
+mesh = srf.to_pyvista() # Create a VTK/PyVista dataset in memory
+mesh.plot()
 ```
 
-Which gives a RectilinearGrid VTK file ``field.vtr``.
+Which gives a RectilinearGrid VTK file ``field.vtr`` or creates a PyVista mesh
+in memory for immediate 3D plotting in Python.
 
 
 ## Requirements:
@@ -285,6 +292,11 @@ Which gives a RectilinearGrid VTK file ``field.vtr``.
 - [pyevtk](https://bitbucket.org/pauloh/pyevtk)
 - [six](https://github.com/benjaminp/six)
 
+### Optional
+
+- [matplotlib](https://matplotlib.org)
+- [pyvista](https://docs.pyvista.org/)
+
 
 ## Contact
 

docs/source/index.rst

@@ -91,7 +91,8 @@ with a :any:`Gaussian` covariance model.
 
 A similar example but for a three dimensional field is exported to a
 `VTK <https://vtk.org/>`__ file, which can be visualized with
-`ParaView <https://www.paraview.org/>`_.
+`ParaView <https://www.paraview.org/>`_ or
+`PyVista <https://docs.pyvista.org>`__ in Python:
 
 .. code-block:: python
 
@@ -102,7 +103,10 @@ A similar example but for a three dimensional field is exported to a
     model = Gaussian(dim=3, var=0.6, len_scale=20)
     srf = SRF(model)
     srf((x, y, z), mesh_type='structured')
-    srf.vtk_export('3d_field')
+    srf.vtk_export('3d_field') # Save to a VTK file for ParaView
+
+    mesh = srf.to_pyvista() # Create a PyVista mesh for plotting in Python
+    mesh.threshold_percent(0.5).plot()
 
 .. image:: https://raw.githubusercontent.com/GeoStat-Framework/GSTools/master/docs/source/pics/3d_gau_field.png
    :width: 400px
@@ -304,6 +308,12 @@ Requirements
 - `pyevtk <https://bitbucket.org/pauloh/pyevtk>`_
 - `six <https://github.com/benjaminp/six>`_
 
+Optional
+~~~~~~~~
+
+- `matplotlib <https://matplotlib.org>`_
+- `pyvista <https://docs.pyvista.org>`_
+
 
 License
 =======

docs/source/tutorial_01_srf.rst

@@ -205,6 +205,13 @@ Using the field from `previous example <Using an Unstructured Grid_>`__, it can
 
     srf.vtk_export("field")
 
+Or it could visualized immediately in Python using `PyVista <https://docs.pyvista.org>`__:
+
+.. code-block:: python
+
+    mesh = srf.to_pyvista("field")
+    mesh.plot()
+
 The script can be found in :download:`gstools/examples/04_export.py<../../examples/04_export.py>` and
 in :download:`gstools/examples/06_unstr_srf_export.py<../../examples/06_unstr_srf_export.py>`
 

examples/06_unstr_srf_export.py

@@ -15,6 +15,8 @@
 
 field = srf((x, y))
 srf.vtk_export("field")
+# Or create a PyVista dataset
+# mesh = srf.to_pyvista()
 
 pt.tricontourf(x, y, field.T)
 pt.axes().set_aspect("equal")

gstools/__init__.py

@@ -81,8 +81,11 @@
 .. currentmodule:: gstools.tools
 
 .. autosummary::
+   to_vtk
    vtk_export
+   to_vtk_structured
    vtk_export_structured
+   to_vtk_unstructured
    vtk_export_unstructured
 
 variogram estimation
@@ -149,7 +152,10 @@
 
 __all__ += [
     "SRF",
+    "to_vtk_structured",
     "vtk_export_structured",
+    "to_vtk_unstructured",
     "vtk_export_unstructured",
+    "to_vtk",
     "vtk_export",
 ]

gstools/field/base.py

@@ -17,7 +17,7 @@
 import numpy as np
 
 from gstools.covmodel.base import CovModel
-from gstools.tools.export import vtk_export as vtk_ex
+from gstools.tools.export import to_vtk, vtk_export
 from gstools.field.tools import _get_select
 
 __all__ = ["Field"]
@@ -147,16 +147,20 @@ def mesh(
                 mesh.point_data[name] = field
         return out
 
-    def vtk_export(
-        self, filename, field_select="field", fieldname="field"
+    def _to_vtk_helper(
+        self, filename=None, field_select="field", fieldname="field"
     ):  # pragma: no cover
-        """Export the stored field to vtk.
+        """Create a VTK/PyVista grid of the stored field or save a VTK dataset
+        to a file.
+
+        This is an internal helper that will handle saving or creating objects
 
         Parameters
         ----------
         filename : :class:`str`
             Filename of the file to be saved, including the path. Note that an
-            ending (.vtr or .vtu) will be added to the name.
+            ending (.vtr or .vtu) will be added to the name. If ``None`` is
+            passed, a PyVista dataset of the appropriate type will be returned.
         field_select : :class:`str`, optional
             Field that should be stored. Can be:
             "field", "raw_field", "krige_field", "err_field" or "krige_var".
@@ -181,7 +185,10 @@ def vtk_export(
                 fields = {}
                 for i in range(self.model.dim):
                     fields[fieldname + suf[i]] = field[i]
-                vtk_ex(filename, self.pos, fields, self.mesh_type)
+                if filename is None:
+                    return to_vtk(self.pos, fields, self.mesh_type)
+                else:
+                    return vtk_export(filename, self.pos, fields, self.mesh_type)
         elif self.value_type == "scalar":
             if hasattr(self, field_select):
                 field = getattr(self, field_select)
@@ -190,18 +197,63 @@ def vtk_export(
             if not (
                 self.pos is None or field is None or self.mesh_type is None
             ):
-                vtk_ex(filename, self.pos, {fieldname: field}, self.mesh_type)
+                if filename is None:
+                    return to_vtk(self.pos, {fieldname: field}, self.mesh_type)
+                else:
+                    return vtk_export(filename, self.pos, {fieldname: field}, self.mesh_type)
             else:
                 print(
-                    "Field.vtk_export: No "
-                    + field_select
-                    + " stored in the class."
+                    "Field.to_vtk: No " + field_select + " stored in the class."
                 )
         else:
             raise ValueError(
                 "Unknown field value type: {}".format(self.value_type)
             )
 
+
+    def to_pyvista(
+        self, field_select="field", fieldname="field"
+    ):  # pragma: no cover
+        """Create a VTK/PyVista grid of the stored field.
+
+        Parameters
+        ----------
+        field_select : :class:`str`, optional
+            Field that should be stored. Can be:
+            "field", "raw_field", "krige_field", "err_field" or "krige_var".
+            Default: "field"
+        fieldname : :class:`str`, optional
+            Name of the field in the VTK file. Default: "field"
+        """
+        grid = self._to_vtk_helper(filename=None, field_select=field_select,
+                                   fieldname=fieldname)
+        return grid
+
+
+    def vtk_export(
+        self, filename, field_select="field", fieldname="field"
+    ):  # pragma: no cover
+        """Export the stored field to vtk.
+
+        Parameters
+        ----------
+        filename : :class:`str`
+            Filename of the file to be saved, including the path. Note that an
+            ending (.vtr or .vtu) will be added to the name.
+        field_select : :class:`str`, optional
+            Field that should be stored. Can be:
+            "field", "raw_field", "krige_field", "err_field" or "krige_var".
+            Default: "field"
+        fieldname : :class:`str`, optional
+            Name of the field in the VTK file. Default: "field"
+        """
+        if not isinstance(filename, str):
+            raise TypeError("Please use a string filename.")
+        return self._to_vtk_helper(filename=filename,
+                                   field_select=field_select,
+                                   fieldname=fieldname)
+
+
     def plot(self, field="field", fig=None, ax=None):  # pragma: no cover
         """
         Plot the spatial random field.

gstools/tools/__init__.py

@@ -8,9 +8,12 @@
 ^^^^^^
 
 .. autosummary::
-   vtk_export
+   to_vtk_structured
    vtk_export_structured
+   to_vtk_unstructured
    vtk_export_unstructured
+   to_vtk
+   vtk_export
 
 Special functions
 ^^^^^^^^^^^^^^^^^