Expose streamlines (#480)

* Small typing fix

* Draft of ansys.dpf.post.tools.streamlines.plot_streamlines()

* Apply suggestions from code review

Co-authored-by: Rafael Canton <[email protected]>

* Implement remarks on sources definition

* Remove __init__.py from ansys/dpf/post/tools

* Revert "Remove __init__.py from ansys/dpf/post/tools"

This reverts commit 6b8cdab.

* Add a conftest.py for Doctest

* Rename sub-package to "helpers" as tools.py already exists in ansys/dpf/post/

* Remove debug

* Update pre-commit

* Move the select method to a new helpers/selections module, ensuring retro-compatibility while preparing deprecation.

* Handle multi-zone with a merge step, and add a set_id argument to select set in transient dataframe.

---------

Co-authored-by: Rafael Canton <[email protected]>

Author: PProfizi

conftest.py

@@ -0,0 +1,39 @@
+"""This runs at the init of the doctest pytest session."""
+import doctest
+from doctest import DocTestRunner
+from unittest import mock
+
+from ansys.dpf.core.misc import module_exists
+import pytest
+
+from ansys.dpf import core
+
+# enable matplotlib off_screen plotting to avoid test interruption
+
+if module_exists("matplotlib"):
+    import matplotlib as mpl
+
+    mpl.use("Agg")
+
+
+# enable off_screen plotting to avoid test interruption
+core.settings.disable_off_screen_rendering()
+core.settings.bypass_pv_opengl_osmesa_crash()
+
+
+class _DPFDocTestRunner(DocTestRunner):
+    def run(self, test, compileflags=None, out=None, clear_globs=True):
+        try:
+            return DocTestRunner.run(self, test, compileflags, out, clear_globs)
+        except doctest.UnexpectedException as e:
+            feature_str = "Feature not supported. Upgrade the server to"
+            if feature_str in str(e.exc_info):
+                pass
+            else:
+                raise e
+
+
+@pytest.fixture(autouse=True)
+def _doctest_runner_dpf():
+    with mock.patch("doctest.DocTestRunner", _DPFDocTestRunner):
+        yield

src/ansys/dpf/post/examples/__init__.py

@@ -2,7 +2,7 @@
 
 Examples
 --------
-This module module exposes PyDPF-Core functionalities.
+This package exposes PyDPF-Core functionalities.
 
 See `here <https://dpf.docs.pyansys.com/dev/api/ansys.dpf.core.examples.html>`_
 for a description of the PyDPF-Core ``example`` API.

src/ansys/dpf/post/helpers/__init__.py

@@ -0,0 +1,6 @@
+"""Tools package.
+
+Tools
+-----
+This package regroups helpers for different common post-treatment functionalities.
+"""

src/ansys/dpf/post/helpers/selections.py

@@ -0,0 +1,42 @@
+"""Module containing helpers to build selections.
+
+Selections
+----------
+
+"""
+from ansys.dpf.post import selection
+
+
+def select(
+    time_freq_indexes=None,
+    time_freq_sets=None,
+    time_freq_values=None,
+    named_selection_names=None,
+    **kwargs,
+):
+    """Creates a ``Selection`` instance allowing to choose the domain on which to evaluate results.
+
+    The results domain defines the time frequency and the spatial selection.
+
+    Parameters
+    ----------
+    time_freq_indexes:
+        Time/freq indexes to select.
+    time_freq_sets:
+        Time/freq sets to select.
+    time_freq_values:
+        Time/freq values to select.
+    named_selection_names:
+        Time/freq named selection to select.
+
+    """
+    current_selection = selection.Selection()
+    if time_freq_indexes:
+        current_selection.select_time_freq_indexes(time_freq_indexes)
+    if time_freq_sets:
+        current_selection.select_time_freq_sets(time_freq_sets)
+    if time_freq_values:
+        current_selection.select_time_freq_values(time_freq_values)
+    if named_selection_names:
+        current_selection.select_named_selection(named_selection_names)
+    return current_selection

src/ansys/dpf/post/helpers/streamlines.py

@@ -0,0 +1,92 @@
+"""Module containing the helpers for streamlines.
+
+Streamlines
+-----------
+
+"""
+from typing import List, Union
+
+from ansys.dpf.core.helpers import streamlines as core_streamlines
+from ansys.dpf.core.plotter import DpfPlotter
+
+from ansys.dpf import core as dpf
+from ansys.dpf import post
+
+
+def plot_streamlines(
+    dataframe: post.DataFrame,
+    sources: List[dict],
+    set_id: int = 1,
+    streamline_thickness: Union[float, List[float]] = 0.01,
+    plot_mesh: bool = True,
+    mesh_opacity: float = 0.3,
+    plot_contour: bool = True,
+    contour_opacity: float = 0.3,
+    **kwargs,
+):
+    """Plot streamlines based on a vector field DataFrame.
+
+    Parameters
+    ----------
+    dataframe:
+        A `post.DataFrame` object containing a vector field.
+        If present, merges the `DataFrame` across its `zone` label before plotting.
+    sources:
+        A list of dictionaries defining spherical point sources for the streamlines.
+        Expected keywords are "center", "radius", "max_time" and "n_points".
+        Keyword "max_time" is for the maximum integration pseudo-time for the streamline
+        computation algorithm, which defines the final length of the lines.
+        More information is available at :func:`pyvista.DataSetFilters.streamlines`.
+    set_id:
+        ID of the set (time-step) for which to compute streamlines if the `DataFrame` object
+        contains temporal data.
+    streamline_thickness:
+        Thickness of the streamlines plotted. Use a list to specify a value for each source.
+    plot_contour:
+        Whether to plot the field's norm contour along with the streamlines.
+    contour_opacity:
+        Opacity to use for the field contour in case "plot_contour=True".
+    plot_mesh:
+        Whether to plot the mesh along the streamlines in case "plot_contour=False".
+    mesh_opacity:
+        Opacity to use for the mesh in case "plot_contour=False" and "plot_mesh=True".
+    **kwargs:
+
+    """
+    # Select data to work with
+    fc = dataframe._fc
+    if "zone" in dataframe.columns.names:
+        fc = dpf.operators.utility.merge_fields_by_label(
+            fields_container=fc,
+            label="zone",
+        ).outputs.fields_container()
+    if set_id not in dataframe.columns.set_index.values:
+        raise ValueError("The set_id requested is not available in this dataframe.")
+    field = fc.get_field_by_time_id(timeid=set_id)
+    meshed_region = field.meshed_region
+
+    # Initialize the plotter
+    plt = DpfPlotter(**kwargs)
+
+    if plot_contour:
+        plt.add_field(field=field, opacity=contour_opacity)
+    elif plot_mesh:
+        plt.add_mesh(meshed_region=meshed_region, opacity=mesh_opacity)
+    if not isinstance(streamline_thickness, list):
+        streamline_thickness = [streamline_thickness] * len(sources)
+    # Add streamlines for each source
+    for i, source in enumerate(sources):
+        pv_streamline, pv_source = core_streamlines.compute_streamlines(
+            meshed_region=meshed_region,
+            field=field,
+            return_source=True,
+            source_radius=source["radius"],
+            source_center=source["center"],
+            n_points=source["n_points"] if "n_points" in source else 100,
+            max_time=source["max_time"] if "max_time" in source else None,
+        )
+        plt.add_streamlines(
+            pv_streamline, source=pv_source, radius=streamline_thickness[i]
+        )
+
+    plt.show_figure(**kwargs)

src/ansys/dpf/post/index.py

@@ -379,20 +379,28 @@ def names(self):
 
     @property
     def results_index(self) -> Union[ResultsIndex, None]:
-        """Returns the available ResultsIndex is present."""
+        """Returns the available ResultsIndex if present."""
         for index in self._indexes:
             if isinstance(index, ResultsIndex):
                 return index
         return None
 
     @property
     def mesh_index(self) -> Union[MeshIndex, None]:
-        """Returns the available ResultsIndex is present."""
+        """Returns the available ResultsIndex if present."""
         for index in self._indexes:
             if isinstance(index, MeshIndex):
                 return index
         return None
 
+    @property
+    def set_index(self) -> Union[SetIndex, None]:
+        """Returns the available SetIndex if present."""
+        for index in self._indexes:
+            if isinstance(index, SetIndex):
+                return index
+        return None
+
     # @property
     # def label_names(self):
     #     """Returns a list with the name of each label Index."""

src/ansys/dpf/post/simulation.py

@@ -811,7 +811,7 @@ def _build_selection(
         node_ids: Union[List[int], None] = None,
         location: Union[locations, str] = locations.nodal,
         external_layer: bool = False,
-        skin: Union[bool] = False,
+        skin: Union[bool, List[int]] = False,
         expand_cyclic: Union[bool, List[Union[int, List[int]]]] = True,
     ) -> Selection:
         tot = (

src/ansys/dpf/post/tools.py

@@ -1,47 +1,11 @@
-"""This module holds factories to create geometry, selections...
+"""This module holds a deprecated link to the legacy `ansys.dpf.post.tools.select` method.
 
 tools
 =====
 
-These factories help creating Objects used to defined which results are evaluated.
+This module has been replaced by the `helpers` sub-package.
+Please use `ansys.dpf.post.helpers.selections.select` instead.
 
 """
 
-from ansys.dpf.post import selection
-
-# from ansys.dpf.core.geometry_factory import *
-
-
-def select(
-    time_freq_indexes=None,
-    time_freq_sets=None,
-    time_freq_values=None,
-    named_selection_names=None,
-    **kwargs,
-):
-    """Creates a ``Selection`` instance allowing to choose the domain on which to evaluate results.
-
-    The results domain defines the time frequency and the spatial selection.
-
-    Parameters
-    ----------
-    time_freq_indexes:
-        Time/freq indexes to select.
-    time_freq_sets:
-        Time/freq sets to select.
-    time_freq_values:
-        Time/freq values to select.
-    named_selection_names:
-        Time/freq named selection to select.
-
-    """
-    current_selection = selection.Selection()
-    if time_freq_indexes:
-        current_selection.select_time_freq_indexes(time_freq_indexes)
-    if time_freq_sets:
-        current_selection.select_time_freq_sets(time_freq_sets)
-    if time_freq_values:
-        current_selection.select_time_freq_values(time_freq_values)
-    if named_selection_names:
-        current_selection.select_named_selection(named_selection_names)
-    return current_selection
+from ansys.dpf.post.helpers.selections import select  # noqa: F401

tests/test_dataframe.py

@@ -1,5 +1,4 @@
 import ansys.dpf.core as core
-from conftest import SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_7_0
 import numpy as np
 import pytest
 from pytest import fixture
@@ -17,6 +16,7 @@
 from ansys.dpf.post.modal_mechanical_simulation import ModalMechanicalSimulation
 from ansys.dpf.post.static_mechanical_simulation import StaticMechanicalSimulation
 from ansys.dpf.post.transient_mechanical_simulation import TransientMechanicalSimulation
+from conftest import SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_7_0
 
 
 @fixture

tests/test_faces.py

@@ -1,12 +1,12 @@
 from ansys.dpf.core import examples
-from conftest import SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_7_0
 import pytest
 from pytest import fixture
 
 from ansys.dpf import core as dpf
 from ansys.dpf import post
 from ansys.dpf.post.faces import Face, FaceListById, FaceListByIndex
 from ansys.dpf.post.nodes import NodeListByIndex
+from conftest import SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_7_0
 
 
 @pytest.mark.skipif(

tests/test_fluid_simulation.py

@@ -1,13 +1,13 @@
 from ansys.dpf.core import examples
-from conftest import (
-    SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_7_0,
-    SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_7_1,
-)
 import pytest
 from pytest import fixture
 
 from ansys.dpf import core as dpf
 from ansys.dpf import post
+from conftest import (
+    SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_7_0,
+    SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_7_1,
+)
 
 
 @pytest.mark.skipif(

tests/test_load_simulation.py

@@ -1,8 +1,4 @@
 import ansys.dpf.core as dpf
-from conftest import (
-    SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_4_0,
-    SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_7_0,
-)
 import pytest
 
 from ansys.dpf import post
@@ -12,6 +8,10 @@
 from ansys.dpf.post.modal_mechanical_simulation import ModalMechanicalSimulation
 from ansys.dpf.post.static_mechanical_simulation import StaticMechanicalSimulation
 from ansys.dpf.post.transient_mechanical_simulation import TransientMechanicalSimulation
+from conftest import (
+    SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_4_0,
+    SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_7_0,
+)
 
 
 def test_load_simulation_static_mechanical(simple_bar, complex_model):

tests/test_mesh.py

@@ -1,12 +1,12 @@
 import ansys.dpf.core as dpf
-from conftest import SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_7_0
 import numpy as np
 import pytest
 from pytest import fixture
 
 from ansys.dpf.post import FluidSimulation, Mesh, StaticMechanicalSimulation
 from ansys.dpf.post.connectivity import ConnectivityListByIndex
 from ansys.dpf.post.faces import Face
+from conftest import SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_7_0
 
 
 @fixture

tests/test_mesh_info.py

@@ -1,10 +1,10 @@
-from conftest import SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_7_0
 import pytest
 from pytest import fixture
 
 from ansys.dpf import core as dpf
 from ansys.dpf import post
 from ansys.dpf.post import examples
+from conftest import SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_7_0
 
 
 @pytest.mark.skipif(

tests/test_named_selection.py

@@ -1,10 +1,10 @@
-import conftest
 import pytest
 from pytest import fixture
 
 from ansys.dpf import core as dpf
 from ansys.dpf.post import StaticMechanicalSimulation
 from ansys.dpf.post.named_selection import NamedSelection
+import conftest
 
 
 @fixture

tests/test_phase.py

@@ -1,10 +1,10 @@
 from ansys.dpf.core import examples
-from conftest import SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_7_0
 import pytest
 from pytest import fixture
 
 from ansys.dpf import core as dpf
 from ansys.dpf import post
+from conftest import SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_7_0
 
 
 @pytest.mark.skipif(