Print Tree on Performance Data (#118)

* Change functionality to print tree on Thicket.dataframe instead of Thicket.statsframe.dataframe

* Add support for list and single data-types

* Make tree printout more detailed

* Add render_legend code from hatchet to modify the legend and add in Indicies

* Add error msg for specific idx

* Replace none values with node name

* Add unit tests for tree

* Add level to error msg

thicket/external/console.py

@@ -9,6 +9,8 @@
 
 
 class ThicketRenderer(ConsoleRenderer):
+    """Extends the Hatchet ConsoleRenderer to support multi-dimensional Thicket data."""
+
     # pylint: disable=W1401
     def render_preamble(self):
         lines = [
@@ -50,6 +52,7 @@ def render(self, roots, dataframe, **kwargs):
         self.colormap_annotations = kwargs["colormap_annotations"]
         self.min_value = kwargs["min_value"]
         self.max_value = kwargs["max_value"]
+        self.indices = kwargs["indices"]
 
         if self.color:
             self.colors = self.colors_enabled
@@ -140,6 +143,81 @@ def render(self, roots, dataframe, **kwargs):
         else:
             return result.encode("utf-8")
 
+    def render_legend(self):
+        def render_label(index, low, high):
+            metric_range = self.max_metric - self.min_metric
+
+            return (
+                self.colors.colormap[index]
+                + "█ "
+                + self.colors.end
+                + "{:.2f}".format(low * metric_range + self.min_metric)
+                + " - "
+                + "{:.2f}".format(high * metric_range + self.min_metric)
+                + "\n"
+            )
+
+        legend = (
+            "\n"
+            + "\033[4m"
+            + "Legend"
+            + self.colors.end
+            + " (Metric: "
+            + str(self.primary_metric)
+            + " Min: {:.2f}".format(self.min_metric)
+            + " Max: {:.2f}".format(self.max_metric)
+            + " indices: "
+            + str(self.indices)
+            + ")\n"
+        )
+
+        legend += render_label(0, 0.9, 1.0)
+        legend += render_label(1, 0.7, 0.9)
+        legend += render_label(2, 0.5, 0.7)
+        legend += render_label(3, 0.3, 0.5)
+        legend += render_label(4, 0.1, 0.3)
+        legend += render_label(5, 0.0, 0.1)
+
+        legend += "\n" + self._ansi_color_for_name("name") + "name" + self.colors.end
+        legend += " User code    "
+
+        legend += self.colors.left + self.lr_arrows["◀"] + self.colors.end
+        legend += " Only in left graph    "
+        legend += self.colors.right + self.lr_arrows["▶"] + self.colors.end
+        legend += " Only in right graph\n"
+
+        if self.annotation_column is not None:
+            # temporal pattern legend customization
+            if "_pattern" in self.annotation_column:
+                score_ranges = [0.0, 0.2, 0.4, 0.6, 1.0]
+                legend += "\nTemporal Pattern"
+                for k in self.temporal_symbols.keys():
+                    if "none" not in k:
+                        legend += "   " + self.temporal_symbols[k] + " " + k
+                legend += "\nTemporal Score  "
+                if self.colormap_annotations:
+                    legend_color_mapping = sorted(score_ranges)
+                    legend_colormap = ColorMaps().get_colors(
+                        self.colormap_annotations, False
+                    )
+                    for i in range(len(score_ranges) - 1):
+                        color = legend_colormap[
+                            legend_color_mapping.index(score_ranges[i + 1])
+                            % len(legend_colormap)
+                        ]
+                        legend += "{}".format(color)
+                        legend += "   {} - {}".format(
+                            score_ranges[i], score_ranges[i + 1]
+                        )
+                        legend += "{}".format(self.colors_annotations.end)
+                else:  # no color map passed in
+                    for i in range(len(score_ranges) - 1):
+                        legend += "   {} - {}".format(
+                            score_ranges[i], score_ranges[i + 1]
+                        )
+
+        return legend
+
     def render_frame(self, node, dataframe, indent="", child_indent=""):
         node_depth = node._depth
         if node_depth < self.depth:

thicket/tests/conftest.py

@@ -154,6 +154,19 @@ def rajaperf_seq_O3_1M_cali(data_dir, tmpdir):
     return [os.path.join(str(tmpdir), f) for f in cali_files]
 
 
+@pytest.fixture
+def rajaperf_unique_tunings(data_dir, tmpdir):
+    """1 trial of each unique tuning in the dataset"""
+    cali_files = [
+        f"{data_dir}/rajaperf/lassen/clang10.0.1_nvcc10.2.89_1048576/1/Base_CUDA-block_128.cali",
+        f"{data_dir}/rajaperf/lassen/clang10.0.1_nvcc10.2.89_1048576/1/Base_CUDA-block_256.cali",
+        f"{data_dir}/rajaperf/quartz/gcc10.3.1_1048576/O3/1/Base_Seq-default.cali",
+    ]
+    for cf in cali_files:
+        shutil.copy(cf, str(tmpdir))
+    return [os.path.join(str(tmpdir), f) for f in cali_files]
+
+
 @pytest.fixture
 def literal_thickets():
     """Returns a list of Thicket objects created from literals."""

thicket/tests/test_tree.py

@@ -0,0 +1,44 @@
+# Copyright 2022 Lawrence Livermore National Security, LLC and other
+# Thicket Project Developers. See the top-level LICENSE file for details.
+#
+# SPDX-License-Identifier: MIT
+
+import pytest
+
+import thicket as th
+
+
+def test_indices(rajaperf_unique_tunings):
+    tk = th.Thicket.from_caliperreader(rajaperf_unique_tunings)
+
+    # No error
+    tk.tree(metric_column="Avg time/rank", indices=tk.profile[0])
+
+    tk.metadata_column_to_perfdata("variant")
+    tk.metadata_column_to_perfdata("tuning")
+
+    # Error because there are duplicate variants. We need to add the tuning to the index as well.
+    tk.dataframe = (
+        tk.dataframe.reset_index().set_index(["node", "variant"]).sort_index()
+    )
+    with pytest.raises(
+        KeyError,
+        match=r"Either dataframe cannot be represented as a single index or provided slice,*",
+    ):
+        tk.tree(metric_column="Avg time/rank")
+
+    # Add tuning to the index to avoid the error.
+    tk.dataframe = (
+        tk.dataframe.reset_index().set_index(["node", "variant", "tuning"]).sort_index()
+    )
+    # No error
+    tk.tree(metric_column="Avg time/rank")
+
+    # No error
+    tk.tree(metric_column="Avg time/rank", indices=["Base_Seq", "default"])
+
+    with pytest.raises(
+        KeyError,
+        match=r"The indices, \{\'tuning\': \'hi\'\}, do not exist in the index \'self.dataframe.index\'",
+    ):
+        tk.tree(metric_column="Avg time/rank", indices=["Base_Seq", "hi"])

thicket/thicket.py

@@ -586,6 +586,7 @@ def tree(
         render_header=True,
         min_value=None,
         max_value=None,
+        indices=None,
     ):
         """Visualize the Thicket as a tree
 
@@ -606,6 +607,7 @@ def tree(
             render_header (bool, optional): Shows the Preamble. Defaults to True.
             min_value (int, optional): Overwrites the min value for the coloring legend. Defaults to None.
             max_value (int, optional): Overwrites the max value for the coloring legend. Defaults to None.
+            indices(tuple, list, optional): Index/indices to display on the DataFrame. Defaults to None.
 
         Returns:
             (str): String representation of the tree, ready to print
@@ -631,9 +633,54 @@ def tree(
         elif sys.version_info.major == 3:
             unicode = True
 
+        if indices is None:
+            # Create slice out of first values found starting after the first index.
+            indices = self.dataframe.index[0][1:]
+        elif isinstance(indices, tuple):
+            pass
+        elif isinstance(indices, list):  # Convert list to tuple
+            indices = tuple(indices)
+        else:  # Support for non-iterable types (int, str, ...)
+            try:
+                indices = tuple([indices])
+            except TypeError:
+                raise TypeError(
+                    f"Value provided to 'indices' = {indices} is an unsupported type {type(indices)}"
+                )
+        # For tree legend
+        idx_dict = {
+            self.dataframe.index.names[k + 1]: indices[k] for k in range(len(indices))
+        }
+        # Slices the DataFrame to simulate a single-level index
+        try:
+            slice_df = (
+                self.dataframe.loc[(slice(None),) + indices, :]
+                .reset_index()
+                .set_index("node")
+            )
+        except KeyError:
+            missing_indices = {
+                list(idx_dict.keys())[i]: idx
+                for i, idx in enumerate(indices)
+                if all(idx not in df_idx[1:] for df_idx in self.dataframe.index)
+            }
+            raise KeyError(
+                f"The indices, {missing_indices}, do not exist in the index 'self.dataframe.index'"
+            )
+        # Check for compatibility
+        if len(slice_df) != len(self.graph):
+            raise KeyError(
+                f"Either dataframe cannot be represented as a single index or provided slice, '{indices}' results in a multi-index. See self.dataframe.loc[(slice(None),)+{indices},{metric_column}]"
+            )
+
+        # Prep DataFrame by filling None rows in the "name" column with the node's name.
+        slice_df["name"] = [
+            n.frame["name"] for n in slice_df.index.get_level_values("node")
+        ]
+
         return ThicketRenderer(unicode=unicode, color=color).render(
             self.graph.roots,
-            self.statsframe.dataframe,
+            slice_df,
             metric_column=metric_column,
             annotation_column=annotation_column,
             precision=precision,
@@ -650,6 +697,7 @@ def tree(
             render_header=render_header,
             min_value=min_value,
             max_value=max_value,
+            indices=idx_dict,
         )
 
     @staticmethod