Add method to export ellipsoids as PyVista objects (#614)

Add a `to_pyvista` method to all ellipsoid classes to export them to
PyVista objects that can be either plotted by it, or exported as VTK
files. Add tests for the new method, running it on all ellipsoid types.

Author: santisoler

harmonica/_forward/ellipsoids/ellipsoids.py

@@ -16,6 +16,11 @@
 
 from ..utils import get_rotation_matrix
 
+try:
+    import pyvista
+except ImportError:
+    pyvista = None
+
 
 def create_ellipsoid(
     a: float,
@@ -258,6 +263,37 @@ def _check_positive_semiaxis(self, value: float, semiaxis: str):
             msg = f"Invalid value of '{semiaxis}' equal to '{value}'. It must be positive."
             raise ValueError(msg)
 
+    def to_pyvista(self, **kwargs):
+        """
+        Export ellipsoid to a :class:`pyvista.PolyData` object.
+
+        .. important::
+
+            The :mod:`pyvista` optional dependency must be installed to use this method.
+
+        Parameters
+        ----------
+        kwargs : dict
+            Keyword arguments passed to :func:`pyvista.ParametricEllipsoid`.
+
+        Returns
+        -------
+        ellipsoid : pyvista.PolyData
+            A PyVista's parametric ellipsoid.
+        """
+        if pyvista is None:
+            msg = (
+                "Missing optional dependency 'pyvista' required for "
+                "exporting ellipsoids to PyVista."
+            )
+            raise ImportError(msg)
+        ellipsoid = pyvista.ParametricEllipsoid(
+            xradius=self.a, yradius=self.b, zradius=self.c, **kwargs
+        )
+        ellipsoid.rotate(rotation=self.rotation_matrix, inplace=True)
+        ellipsoid.translate(self.center, inplace=True)
+        return ellipsoid
+
 
 class TriaxialEllipsoid(BaseEllipsoid):
     """

harmonica/tests/ellipsoids/test_ellipsoids.py

@@ -10,6 +10,7 @@
 
 import itertools
 import re
+from unittest.mock import patch
 
 import numpy as np
 import pytest
@@ -22,6 +23,11 @@
     create_ellipsoid,
 )
 
+try:
+    import pyvista
+except ImportError:
+    pyvista = None
+
 
 class TestProlateEllipsoid:
     """Test the ProlateEllipsoid class."""
@@ -544,3 +550,48 @@ def test_physical_properties(self, a, b, c, physical_property):
         center = (1.0, 2.0, 3.0)
         ellipsoid = create_ellipsoid(a, b, c, center=center, **kwargs)
         np.testing.assert_allclose(getattr(ellipsoid, physical_property), value)
+
+
+@pytest.mark.skipif(pyvista is None, reason="requires pyvista")
+class TestToPyvista:
+    """Test exporting ellipsoids to PyVista objects."""
+
+    @pytest.fixture(params=["oblate", "prolate", "triaxial", "sphere"])
+    def ellipsoid(self, request):
+        a, b, c = 3.0, 2.0, 1.0
+        yaw, pitch, roll = 73.0, 14.0, -35.0
+        center = (43.0, -72.0, 105)
+        match request.param:
+            case "oblate":
+                ellipsoid = OblateEllipsoid(b, a, yaw, pitch, center=center)
+            case "prolate":
+                ellipsoid = ProlateEllipsoid(a, b, yaw, pitch, center=center)
+            case "triaxial":
+                ellipsoid = TriaxialEllipsoid(a, b, c, yaw, pitch, roll, center=center)
+            case "sphere":
+                ellipsoid = Sphere(a, center=center)
+            case _:
+                raise ValueError()
+        return ellipsoid
+
+    @patch("harmonica._forward.ellipsoids.ellipsoids.pyvista", None)
+    def test_pyvista_missing_error(self, ellipsoid):
+        """
+        Check if error is raised when pyvista is not installed.
+        """
+        with pytest.raises(ImportError):
+            ellipsoid.to_pyvista()
+
+    def test_pyvista_object(self, ellipsoid):
+        """
+        Check if method works as expected.
+        """
+        ellipsoid_pv = ellipsoid.to_pyvista()
+        assert isinstance(ellipsoid_pv, pyvista.PolyData)
+        # rtol needed since the parametric ellipsoid is not the exact surface.
+        np.testing.assert_allclose(ellipsoid_pv.center, ellipsoid.center, rtol=1e-4)
+        np.testing.assert_allclose(
+            ellipsoid_pv.volume,
+            4 / 3 * np.pi * ellipsoid.a * ellipsoid.b * ellipsoid.c,
+            rtol=1e-3,
+        )