[feat] Support post processing predict as learnt and predict on cells for surface variables

doc/source/api_reference/post_processings.rst

@@ -56,6 +56,14 @@ Available postprocessings
     :members:
 
 
+.. autoclass:: PostProcessingOnCells()
+    :members:
+
+
+.. autoclass:: PostProcessingAsLearnt()
+    :members:
+
+
 .. autoclass:: VolumeVTU()
     :members:
 
@@ -69,3 +77,6 @@ Helpers
 
 .. autoclass:: DownloadableResult()
     :members:
+
+.. autoclass:: PPSurfaceLocation()
+    :members:

src/ansys/simai/core/data/post_processings.py

@@ -25,6 +25,7 @@
 import numbers
 import sys
 from abc import ABC, abstractmethod
+from enum import Enum
 from inspect import cleandoc
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
 
@@ -274,6 +275,40 @@ class SurfaceVTP(_PostProcessingVTKExport):
     """
 
 
+class PostProcessingOnCells(SurfaceVTP):
+    """Decodes the server response for the on-cells prediction of the surface in VTP format.
+
+    This class is generated through the :meth:`~PredictionPostProcessings.surface_vtp()` method.
+    """
+
+    @classmethod
+    def _api_name(cls):
+        return "SurfaceVTP"
+
+
+class PostProcessingAsLearnt(SurfaceVTP):
+    """Decodes the server response for the as-learnt prediction of the surface in VTP format.
+
+    This class is generated through the :meth:`~PredictionPostProcessings.surface_vtp()` method.
+    """
+
+    @classmethod
+    def _api_name(cls):
+        return "SurfaceVTPTDLocation"
+
+
+class PPSurfaceLocation(Enum):
+    """Enumerates the post-processing options on the data location.
+
+    Args:
+        ON_CELLS: identifies that the post-processing should run on cell data.
+        AS_LEARNT: identifies that the post-processing should run on the data as they are learnt.
+    """
+
+    AS_LEARNT = PostProcessingAsLearnt
+    ON_CELLS = PostProcessingOnCells
+
+
 class CustomVolumePointCloud(PostProcessing):
     """Provides a representation of a CustomVolumePointCloud post-processing.
 
@@ -425,7 +460,9 @@ def slice(
         plane = convert_axis_and_coordinate_to_plane_eq_coeffs(axis, coordinate)
         return self._get_or_run(Slice, {"plane": plane, "output_format": format}, run)
 
-    def surface_vtp(self, run: bool = True) -> Optional[SurfaceVTP]:
+    def surface_vtp(
+        self, run: bool = True, pp_location: PPSurfaceLocation = PPSurfaceLocation.ON_CELLS
+    ) -> Optional[SurfaceVTP]:
         """Compute or get the result of the prediction's surface in VTP format.
 
         This is a non-blocking method. It returns the ``PostProcessingVTP``
@@ -442,24 +479,42 @@ def surface_vtp(self, run: bool = True) -> Optional[SurfaceVTP]:
             run: Boolean indicating whether to compute or get the postprocessing.
                 The default is ``True``. If ``False``, the postprocessing is not
                 computed, and ``None`` is returned if it does not exist yet.
+            pp_location: the post-processing data location. It can be one of
+                :class:`PPSurfaceLocation`, otherwise an InvalidArgument error
+                is raised. Default is PPSurfaceLocation.ON_CELLS.
 
         Returns:
             :class:`SurfaceVTP` object that allows downloading the binary data.
             Returns ``None`` if ``run=False`` and the postprocessing does not exist.
 
         Examples:
-            Run and download a surface VTP.
+            Run and download a surface VTP for data location on cells (predict on cells).
 
             .. code-block:: python
 
                 import ansys.simai.core
+                from ansys.simai.core.data.post_processings import PPSurfaceLocation
 
                 simai = ansys.simai.core.from_config()
                 prediction = simai.predictions.list()[0]
-                surface_vtp = prediction.post.surface_vtp().data.download("/tmp/simai.vtp")
+                surface_vtp = prediction.post.surface_vtp(
+                    pp_location=PPSurfaceLocation.ON_CELLS
+                ).data.download("/tmp/simai.vtp")
+
+            Run and download a surface VTP with data location as it is learnt (predict as learnt).
 
+            .. code-block:: python
+
+                import ansys.simai.core
+                from ansys.simai.core.data.post_processings import PPSurfaceLocation
+
+                simai = ansys.simai.core.from_config()
+                prediction = simai.predictions.list()[0]
+                surface_vtp = prediction.post.surface_vtp(
+                    pp_location=PPSurfaceLocation.AS_LEARNT
+                ).data.download("/tmp/simai.vtp")
 
-            Run a surface VTP and open a plot using PyVista.
+            Run a surface VTP with data location on cells, and open a plot using PyVista.
 
             .. code-block:: python
 
@@ -477,7 +532,9 @@ def surface_vtp(self, run: bool = True) -> Optional[SurfaceVTP]:
                     surface_vtp = pyvista.read(temp_vtp_file.name)
                     surface_vtp.plot()
         """
-        return self._get_or_run(SurfaceVTP, {}, run)
+        if not isinstance(pp_location, PPSurfaceLocation):
+            raise InvalidArguments(f"pp_location should be one of {PPSurfaceLocation}")
+        return self._get_or_run(pp_location.value, {}, run)
 
     def volume_vtu(self, run: bool = True) -> Optional[VolumeVTU]:
         """Compute or get the result of the prediction's volume in VTU format.

src/ansys/simai/core/data/selection_post_processings.py

@@ -25,6 +25,7 @@
 from ansys.simai.core.data.lists import ExportablePPList, PPList
 from ansys.simai.core.data.post_processings import (
     GlobalCoefficients,
+    PPSurfaceLocation,
     Slice,
     SurfaceEvol,
     SurfaceVTP,
@@ -144,7 +145,9 @@ def volume_vtu(self) -> PPList[VolumeVTU]:
         """
         return PPList(selection=self._selection, post=lambda pred: pred.post.volume_vtu())
 
-    def surface_vtp(self) -> PPList[SurfaceVTP]:
+    def surface_vtp(
+        self, pp_location: PPSurfaceLocation = PPSurfaceLocation.ON_CELLS
+    ) -> PPList[SurfaceVTP]:
         """Compute or get the result of each prediction's surface in the VTP format.
 
         This is a non-blocking method. It returns a
@@ -159,8 +162,16 @@ def surface_vtp(self) -> PPList[SurfaceVTP]:
         The computation is launched only on first call of this method.
         Subsequent calls do not relaunch it.
 
+        Args:
+            pp_location: the post-processing data location. It can be one of
+                :class:`~ansys.simai.core.data.post_processings.PPSurfaceLocation`,
+                otherwise an InvalidArgument error is raised.
+                Default is PPSurfaceLocation.ON_CELLS.
+
         Returns:
             :py:class:`~ansys.simai.core.data.lists.PPList` instance of
             :py:class:`~ansys.simai.core.data.post_processings.SurfaceVTP` objects.
         """
-        return PPList(selection=self._selection, post=lambda pred: pred.post.surface_vtp())
+        return PPList(
+            selection=self._selection, post=lambda pred: pred.post.surface_vtp(True, pp_location)
+        )

tests/test_post_processings_results.py

@@ -23,17 +23,22 @@
 import json
 from io import BytesIO
 
+import pytest
 import responses
 
 from ansys.simai.core.data.post_processings import (
     CustomVolumePointCloud,
     DownloadableResult,
     GlobalCoefficients,
+    PostProcessingAsLearnt,
+    PostProcessingOnCells,
+    PPSurfaceLocation,
     Slice,
     SurfaceEvol,
     SurfaceVTP,
     VolumeVTU,
 )
+from ansys.simai.core.errors import InvalidArguments
 
 
 @responses.activate
@@ -344,6 +349,77 @@ def test_post_processing_result_surface_vtp(simai_client):
     assert responses.assert_call_count("https://test.test/post-processings/01010101654", 2)
 
 
+@responses.activate
+def test_post_processing_result_surface_vtp_predict_as_learnt(simai_client):
+    """WHEN Running a surface_vtp post-processing on a prediction
+    WITH predict-as-learnt option (location is PPSurfaceLocation.AS_LEARNT)
+    AND calling its .data field,
+    THEN a GET request is made on the post-processings/SurfaceVTPTDLocation endpoint
+    AND the returned instance is of type PostProcessingAsLearnt.
+    """
+    # Mock request for PP creation
+    responses.add(
+        responses.POST,
+        "https://test.test/predictions/7546/post-processings/SurfaceVTPTDLocation",
+        json={"id": "01010101654", "state": "successful"},
+        status=200,
+    )
+
+    pred = simai_client._prediction_directory._model_from({"id": "7546", "state": "successful"})
+    surface_vtp = pred.post.surface_vtp(pp_location=PPSurfaceLocation.AS_LEARNT)
+    assert responses.assert_call_count(
+        "https://test.test/predictions/7546/post-processings/SurfaceVTPTDLocation", 1
+    )
+    assert isinstance(surface_vtp, PostProcessingAsLearnt)
+
+
+@responses.activate
+def test_post_processing_result_surface_vtp_predict_on_cells(simai_client):
+    """WHEN Running a surface_vtp post-processing on a prediction
+    WITH predict-on-cells option (location is PPSurfaceLocation.ON_CELLS)
+    THEN a GET request is made on the post-processings/SurfaceVTP endpoint
+    AND the returned instance is of type PostProcessingOnCells.
+    """
+    # Mock request for PP creation
+    responses.add(
+        responses.POST,
+        "https://test.test/predictions/7546/post-processings/SurfaceVTP",
+        json={"id": "01010101654", "state": "successful"},
+        status=200,
+    )
+
+    pred = simai_client._prediction_directory._model_from({"id": "7546", "state": "successful"})
+    surface_vtp = pred.post.surface_vtp(pp_location=PPSurfaceLocation.ON_CELLS)
+    assert responses.assert_call_count(
+        "https://test.test/predictions/7546/post-processings/SurfaceVTP", 1
+    )
+    assert isinstance(surface_vtp, PostProcessingOnCells)
+
+
+def test_error_in_wrong_location_in_post_processing_surface_vtp(simai_client):
+    """WHEN Running a surface_vtp post-processing on a prediction
+    WITH wrong pp_location,
+    THEN an InvalidArguments error is raised.
+    """
+
+    pred = simai_client._prediction_directory._model_from({"id": "7546", "state": "successful"})
+
+    class TestClass(SurfaceVTP):
+        pass
+
+    with pytest.raises(InvalidArguments) as ex:
+        _ = pred.post.surface_vtp(pp_location=TestClass)
+        assert list(PPSurfaceLocation) in ex.value
+
+    with pytest.raises(InvalidArguments) as ex:
+        _ = pred.post.surface_vtp(pp_location="a_string")
+        assert list(PPSurfaceLocation) in ex.value
+
+    with pytest.raises(InvalidArguments) as ex:
+        _ = pred.post.surface_vtp(pp_location=str)
+        assert list(PPSurfaceLocation) in ex.value
+
+
 @responses.activate
 def test_post_processing_result_point_cloud(simai_client):
     responses.add(

tests/test_selection_post_processings.py

@@ -25,6 +25,8 @@
 
 from ansys.simai.core.data.post_processings import (
     GlobalCoefficients,
+    PostProcessingAsLearnt,
+    PPSurfaceLocation,
     Slice,
     SurfaceEvol,
     SurfaceVTP,
@@ -181,6 +183,33 @@ def test_selection_post_processing_surface_vtp(test_selection):
         assert isinstance(pp, SurfaceVTP)
 
 
+@responses.activate
+def test_selection_post_processing_surface_vtp_predict_as_learnt(test_selection):
+    """WHEN I call post.post.surface_vtp() on a selection
+    THEN the /SurfaceVTPTDLocation endpoint is called for each prediction in the selection
+    AND I get a list of PostProcessingVTUExport objects in return
+    """
+    assert len(test_selection.points) == 2
+
+    responses.add(
+        responses.POST,
+        "https://test.test/predictions/pred1/post-processings/SurfaceVTPTDLocation",
+        json={"id": "vtu01", "status": "queued"},
+        status=200,
+    )
+    responses.add(
+        responses.POST,
+        "https://test.test/predictions/pred2/post-processings/SurfaceVTPTDLocation",
+        json={"id": "vtu02", "status": "queued"},
+        status=200,
+    )
+    post_processings = test_selection.post.surface_vtp(pp_location=PPSurfaceLocation.AS_LEARNT)
+    assert isinstance(post_processings, PPList)
+    assert len(post_processings) == 2
+    for pp in post_processings:
+        assert isinstance(pp, PostProcessingAsLearnt)
+
+
 @responses.activate
 def test_selection_post_processing_error(test_selection):
     """WHEN I call post.post.volume_vtu() on a selection