Add units to ResultsIndex (#330)

* Expose data unit for ResultIndex

* Update test

* Improve coverage

* Improve readability

Author: PProfizi

src/ansys/dpf/post/index.py

@@ -165,9 +165,31 @@ class ResultsIndex(Index):
     def __init__(
         self,
         values: List[str],
+        units: Union[List[Union[str, None]], None] = None,
     ):
         """Initiate this class."""
-        super().__init__(name=ref_labels.results, values=values, scoping=None)
+        # Initialize results labels
+        result_values = values
+        if units is None:
+            # If no units, initialize as list of None
+            units = [None] * len(values)
+        else:
+            # Truncate units if longer than values
+            if len(units) > len(values):
+                units = units[: len(values)]
+            # Update result labels with unit when necessary
+            for i, unit in enumerate(units):
+                result_values[i] += f" ({unit})" if unit is not None else ""
+            # If units was shorter than values, extend with Nones
+            if len(units) < len(values):
+                units.extend([None] * (len(values) - len(units)))
+        self._units = units
+        super().__init__(name=ref_labels.results, values=result_values, scoping=None)
+
+    @property
+    def units(self) -> List[str]:
+        """Return the list of units for this index."""
+        return self._units
 
     def __repr__(self):
         """Representation of the Index."""

src/ansys/dpf/post/simulation.py

@@ -523,14 +523,15 @@ def _create_dataframe(self, fc, location, columns, comp, base_name, disp_wf=None
                 message=f"Returned Dataframe with columns {columns} is empty.",
                 category=UserWarning,
             )
+        unit = fc[0].unit
         comp_index = None
         if comp is not None:
             comp_index = CompIndex(values=comp)
         row_indexes = [MeshIndex(location=location, fc=fc)]
         if comp_index is not None:
             row_indexes.append(comp_index)
         column_indexes = [
-            ResultsIndex(values=[base_name]),
+            ResultsIndex(values=[base_name], units=[unit]),
             SetIndex(values=fc.get_available_ids_for_label("time")),
         ]
         label_indexes = []

tests/test_dataframe.py

@@ -139,9 +139,15 @@ def test_dataframe_animate(transient_rst):
 
 def test_dataframe_repr(df):
     ref = (
-        "DataFrame<index=MultiIndex<[MeshIndex<name=\"node_ids\", dtype=<class 'int'>>, "
-        "Index<name=\"components\", dtype=<class 'str'>>]>, columns=MultiIndex<[ResultIndex<['U']>, "  # noqa: E501
-        "SetIndex<values=[1]>]>>"
+        "DataFrame<"
+        "index=MultiIndex<["
+        "MeshIndex<name=\"node_ids\", dtype=<class 'int'>>, "
+        "Index<name=\"components\", dtype=<class 'str'>>"
+        "]>, "
+        "columns=MultiIndex<["
+        "ResultIndex<['U (m)']>, "
+        "SetIndex<values=[1]>"
+        "]>>"
     )
     assert repr(df) == ref
 
@@ -151,7 +157,7 @@ def test_dataframe_str(transient_rst):
     df = simulation.displacement(all_sets=True)
     print(df)
     ref = """
-                 results           U                                                                     ...
+                 results       U (m)                                                                     ...
                  set_ids           1           2           3           4           5           6         ...
     node_ids  components                                                                                 ...
          525           X  0.0000e+00  4.8506e-05  2.3022e-04  6.5140e-04  1.4812e-03  2.9324e-03         ...
@@ -166,7 +172,7 @@ def test_dataframe_str(transient_rst):
     df2 = df.select(node_ids=525)
     print(df2)
     ref = """
-                 results           U                                                                     ...
+                 results       U (m)                                                                     ...
                  set_ids           1           2           3           4           5           6         ...
     node_ids  components                                                                                 ...
          525           X  0.0000e+00  4.8506e-05  2.3022e-04  6.5140e-04  1.4812e-03  2.9324e-03         ...
@@ -179,7 +185,7 @@ def test_dataframe_str(transient_rst):
     print(df)
     print(df._fc[0].get_entity_data_by_id(391))
     ref = """
-                 results           S
+                 results      S (Pa)
                  set_ids          35
  element_ids  components            
          391      XX (1) -3.2780e+05
@@ -194,7 +200,7 @@ def test_dataframe_str(transient_rst):
     df2 = df.select(components="YY")
     print(df2)
     ref = """
-                 results           S
+                 results      S (Pa)
                  set_ids          35
  element_ids  components            
          391      YY (1)  1.3601e+06

tests/test_index.py

@@ -0,0 +1,27 @@
+from ansys.dpf.post.index import ResultsIndex
+
+
+class TestResultsIndex:
+    def test_results_index(self):
+        result_index = ResultsIndex(values=["U", "V", "A"], units=None)
+        ref = "ResultIndex<['U', 'V', 'A']>"
+        assert repr(result_index) == ref
+        result_index = ResultsIndex(values=["U", "V", "A"], units=["m", "m/s"])
+        ref = "ResultIndex<['U (m)', 'V (m/s)', 'A']>"
+        assert repr(result_index) == ref
+        result_index = ResultsIndex(values=["U", "V", "A"], units=["m", None, "m/s^2"])
+        ref = "ResultIndex<['U (m)', 'V', 'A (m/s^2)']>"
+        assert repr(result_index) == ref
+        result_index = ResultsIndex(values=["U", "V", "A"], units=["m", "m/s", "m/s^2"])
+        ref = "ResultIndex<['U (m)', 'V (m/s)', 'A (m/s^2)']>"
+        assert repr(result_index) == ref
+
+    def test_results_index_units(self):
+        result_index = ResultsIndex(values=["U", "V", "A"], units=None)
+        assert result_index.units == [None, None, None]
+        result_index = ResultsIndex(values=["U", "V", "A"], units=["m", "m/s"])
+        assert result_index.units == ["m", "m/s", None]
+        result_index = ResultsIndex(values=["U", "V", "A"], units=["m", None, "m/s^2"])
+        assert result_index.units == ["m", None, "m/s^2"]
+        result_index = ResultsIndex(values=["U", "V", "A"], units=["m", "m/s", "m/s^2"])
+        assert result_index.units == ["m", "m/s", "m/s^2"]