Adds support for tree annotations with coloring and min, max attribut…

…es (#97)

* Adds support for an annotation column and support for multi-indexed datafranes

* Adds optional support for coloring of custom attributes

* Provide dict for color mapping

* Adds min & max parameters; Adds missing node decorators

* Account for changes in Thicket

* Revert changes that indicate NaN nodes

* Adds brackets to easier distinguish annotations

* Adds documentation to tree() plotting

* Formatting

* Reverts changes required for Thicket support

* Updates test

* Remove unused import

* prints also tuples

* comment

* updates timemory tree test

* remove unused package

---------

Co-authored-by: Stephanie Brink <brink2@llnl.gov>

hatchet/external/console.py

@@ -56,6 +56,7 @@ def render(self, roots, dataframe, **kwargs):
             return result
 
         self.metric_columns = kwargs["metric_column"]
+        self.annotation_column = kwargs["annotation_column"]
         self.precision = kwargs["precision"]
         self.name = kwargs["name_column"]
         self.expand = kwargs["expand_name"]
@@ -66,13 +67,31 @@ def render(self, roots, dataframe, **kwargs):
         self.highlight = kwargs["highlight_name"]
         self.colormap = kwargs["colormap"]
         self.invert_colormap = kwargs["invert_colormap"]
+        self.colormap_annotations = kwargs["colormap_annotations"]
+        self.min_value = kwargs["min_value"]
+        self.max_value = kwargs["max_value"]
 
         if self.color:
             self.colors = self.colors_enabled
             # set the colormap based on user input
             self.colors.colormap = ColorMaps().get_colors(
                 self.colormap, self.invert_colormap
             )
+
+            if self.annotation_column and self.colormap_annotations:
+                self.colors_annotations = self.colors_enabled()
+                if isinstance(self.colormap_annotations, (str, list)):
+                    if isinstance(self.colormap_annotations, str):
+                        self.colors_annotations.colormap = ColorMaps().get_colors(
+                            self.colormap_annotations, False
+                        )
+                    elif isinstance(self.colormap_annotations, list):
+                        self.colors_annotations.colormap = self.colormap_annotations
+                    self.colors_annotations_mapping = sorted(
+                        list(dataframe[self.annotation_column].apply(str).unique())
+                    )
+                elif isinstance(self.colormap_annotations, dict):
+                    self.colors_annotations_mapping = self.colormap_annotations
         else:
             self.colors = self.colors_disabled
 
@@ -122,8 +141,8 @@ def render(self, roots, dataframe, **kwargs):
             metric_series.values[isfinite_mask], metric_series.index[isfinite_mask]
         )
 
-        self.max_metric = filtered_series.max()
-        self.min_metric = filtered_series.min()
+        self.max_metric = self.max_value if self.max_value else filtered_series.max()
+        self.min_metric = self.min_value if self.min_value else filtered_series.min()
 
         if self.unicode:
             self.lr_arrows = {"◀": u"◀ ", "▶": u"▶ "}
@@ -175,7 +194,7 @@ def render_label(index, low, high):
             + "Legend"
             + self.colors.end
             + " (Metric: "
-            + self.primary_metric
+            + str(self.primary_metric)
             + " Min: {:.2f}".format(self.min_metric)
             + " Max: {:.2f}".format(self.max_metric)
             + ")\n"
@@ -228,6 +247,26 @@ def render_frame(self, node, dataframe, indent=u"", child_indent=u""):
                     c=self.colors,
                 )
 
+            if self.annotation_column is not None:
+                annotation_content = str(
+                    dataframe.loc[df_index, self.annotation_column]
+                )
+                if self.colormap_annotations:
+                    if isinstance(self.colormap_annotations, dict):
+                        color_annotation = self.colors_annotations_mapping[
+                            annotation_content
+                        ]
+                    else:
+                        color_annotation = self.colors_annotations.colormap[
+                            self.colors_annotations_mapping.index(annotation_content)
+                            % len(self.colors_annotations.colormap)
+                        ]
+                    metric_str += " [{}".format(color_annotation)
+                    metric_str += "{}".format(annotation_content)
+                    metric_str += "{}]".format(self.colors_annotations.end)
+                else:
+                    metric_str += " [{}]".format(annotation_content)
+
             node_name = dataframe.loc[df_index, self.name]
             if self.expand is False:
                 if len(node_name) > 39:

hatchet/graphframe.py

@@ -929,6 +929,7 @@ def unify(self, other):
     def tree(
         self,
         metric_column=None,
+        annotation_column=None,
         precision=3,
         name_column="name",
         expand_name=False,
@@ -939,9 +940,34 @@ def tree(
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        colormap_annotations=None,
         render_header=True,
+        min_value=None,
+        max_value=None,
     ):
-        """Format this graphframe as a tree and return the resulting string."""
+        """Visualize the Hatchet graphframe as a tree
+
+        Arguments:
+            metric_column (str, list, optional): Columns to use the metrics from. Defaults to None.
+            annotation_column (str, optional): Column to use as an annotation. Defaults to None.
+            precision (int, optional): Precision of shown numbers. Defaults to 3.
+            name_column (str, optional): Column of the node name. Defaults to "name".
+            expand_name (bool, optional): Limits the lenght of the node name. Defaults to False.
+            context_column (str, optional): Shows the file this function was called in (Available with HPCToolkit). Defaults to "file".
+            rank (int, optional): Specifies the rank to take the data from. Defaults to 0.
+            thread (int, optional): Specifies the thread to take the data from. Defaults to 0.
+            depth (int, optional): Sets the maximum depth of the tree. Defaults to 10000.
+            highlight_name (bool, optional): Highlights the names of the nodes. Defaults to False.
+            colormap (str, optional): Specifies a colormap to use. Defaults to "RdYlGn".
+            invert_colormap (bool, optional): Reverts the chosen colormap. Defaults to False.
+            colormap_annotations (str, list, dict, optional): Either provide the name of a colormap, a list of colors to use or a dictionary which maps the used annotations to a color. Defaults to None.
+            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.
+
+        Returns:
+            str: String representation of the tree, ready to print
+        """
         color = sys.stdout.isatty()
         shell = None
         if metric_column is None:
@@ -967,6 +993,7 @@ def tree(
             self.graph.roots,
             self.dataframe,
             metric_column=metric_column,
+            annotation_column=annotation_column,
             precision=precision,
             name_column=name_column,
             expand_name=expand_name,
@@ -977,7 +1004,10 @@ def tree(
             highlight_name=highlight_name,
             colormap=colormap,
             invert_colormap=invert_colormap,
+            colormap_annotations=colormap_annotations,
             render_header=render_header,
+            min_value=min_value,
+            max_value=max_value,
         )
 
     def to_dot(self, metric=None, name="name", rank=0, thread=0, threshold=0.0):

hatchet/tests/caliper.py

@@ -13,7 +13,6 @@
 from hatchet import GraphFrame
 from hatchet.readers.caliper_reader import CaliperReader
 from hatchet.util.executable import which
-from hatchet.external.console import ConsoleRenderer
 
 caliperreader_avail = True
 try:
@@ -159,46 +158,18 @@ def test_filter_squash_unify_caliper_data(lulesh_caliper_json):
     squash_gf2.dataframe.set_index(gf2_index_names, inplace=True)
 
 
-def test_tree(lulesh_caliper_json):
+def test_tree(monkeypatch, lulesh_caliper_json):
     """Sanity test a GraphFrame object with known data."""
+    monkeypatch.setattr("sys.stdout.isatty", (lambda: False))
     gf = GraphFrame.from_caliper(str(lulesh_caliper_json))
+    output = gf.tree(metric_column="time")
 
-    output = ConsoleRenderer(unicode=True, color=False).render(
-        gf.graph.roots,
-        gf.dataframe,
-        metric_column="time",
-        precision=3,
-        name_column="name",
-        expand_name=False,
-        context_column="file",
-        rank=0,
-        thread=0,
-        depth=10000,
-        highlight_name=False,
-        colormap="RdYlGn",
-        invert_colormap=False,
-        render_header=True,
-    )
     assert "121489.000 main" in output
     assert "663.000 LagrangeElements" in output
     assert "21493.000 CalcTimeConstraintsForElems" in output
 
-    output = ConsoleRenderer(unicode=True, color=False).render(
-        gf.graph.roots,
-        gf.dataframe,
-        metric_column="time (inc)",
-        precision=3,
-        name_column="name",
-        expand_name=False,
-        context_column="file",
-        rank=0,
-        thread=0,
-        depth=10000,
-        highlight_name=False,
-        colormap="RdYlGn",
-        invert_colormap=False,
-        render_header=True,
-    )
+    output = gf.tree(metric_column="time (inc)")
+
     assert "662712.000 EvalEOSForElems" in output
     assert "2895319.000 LagrangeNodal" in output
 

hatchet/tests/cprofile.py

@@ -7,7 +7,6 @@
 import re
 
 from hatchet import GraphFrame
-from hatchet.external.console import ConsoleRenderer
 
 
 def test_graphframe(hatchet_cycle_pstats):
@@ -27,43 +26,16 @@ def test_graphframe(hatchet_cycle_pstats):
             assert gf.dataframe[col].dtype == object
 
 
-def test_tree(hatchet_cycle_pstats):
+def test_tree(monkeypatch, hatchet_cycle_pstats):
+    monkeypatch.setattr("sys.stdout.isatty", (lambda: False))
     gf = GraphFrame.from_cprofile(str(hatchet_cycle_pstats))
 
-    output = ConsoleRenderer(unicode=True, color=False).render(
-        gf.graph.roots,
-        gf.dataframe,
-        metric_column="time",
-        precision=3,
-        name_column="name",
-        expand_name=False,
-        context_column="file",
-        rank=0,
-        thread=0,
-        depth=10000,
-        highlight_name=False,
-        colormap="RdYlGn",
-        invert_colormap=False,
-        render_header=True,
-    )
+    output = gf.tree(metric_column="time")
+
     assert "g pstats_reader_test.py" in output
     assert "<method 'disable' ...Profiler' objects> ~" in output
 
-    output = ConsoleRenderer(unicode=True, color=False).render(
-        gf.graph.roots,
-        gf.dataframe,
-        metric_column="time (inc)",
-        precision=3,
-        name_column="name",
-        expand_name=False,
-        context_column="file",
-        rank=0,
-        thread=0,
-        depth=10000,
-        highlight_name=False,
-        colormap="RdYlGn",
-        invert_colormap=False,
-        render_header=True,
-    )
+    output = gf.tree(metric_column="time (inc)")
+
     assert "f pstats_reader_test.py" in output
     assert re.match("(.|\n)*recursive(.|\n)*recursive", output)