Adds performance annotations using the new Hatchet annotation interface

thicket/ensemble.py

@@ -6,6 +6,7 @@
 from collections import OrderedDict
 
 from hatchet import GraphFrame
+from hatchet.util.perf_measure import annotate
 import pandas as pd
 
 import thicket.helpers as helpers
@@ -18,10 +19,14 @@
 )
 
 
+_ensemble_annotate = annotate(fmt="Ensemble.{}")
+
+
 class Ensemble:
     """Operations pertaining to ensembling."""
 
     @staticmethod
+    @_ensemble_annotate
     def _unify(thickets, inplace=False, disable_tqdm=False):
         """Create union graph from list of thickets and sync their DataFrames.
 
@@ -83,6 +88,7 @@ def _unify(thickets, inplace=False, disable_tqdm=False):
         return union_graph, _thickets
 
     @staticmethod
+    @_ensemble_annotate
     def _columns(
         thickets,
         headers=None,
@@ -186,9 +192,9 @@ def _handle_misc():
             combined_th.profile = [new_mappings[prf] for prf in combined_th.profile]
             profile_mapping_cp = combined_th.profile_mapping.copy()
             for k, v in profile_mapping_cp.items():
-                combined_th.profile_mapping[
-                    new_mappings[k]
-                ] = combined_th.profile_mapping.pop(k)
+                combined_th.profile_mapping[new_mappings[k]] = (
+                    combined_th.profile_mapping.pop(k)
+                )
             combined_th.performance_cols = helpers._get_perf_columns(
                 combined_th.dataframe
             )
@@ -331,6 +337,7 @@ def _handle_statsframe():
         return combined_th
 
     @staticmethod
+    @_ensemble_annotate
     def _index(
         thickets, from_statsframes=False, fill_perfdata=True, disable_tqdm=False
     ):

thicket/external/console.py

@@ -4,14 +4,19 @@
 import pandas as pd
 from hatchet.external.console import ConsoleRenderer
 from hatchet.util.colormaps import ColorMaps
+from hatchet.util.perf_measure import annotate
 
 from ..version import __version__
 
 
+_thicket_console_renderer_annotate = annotate(fmt="ThicketRenderer.{}")
+
+
 class ThicketRenderer(ConsoleRenderer):
     """Extends the Hatchet ConsoleRenderer to support multi-dimensional Thicket data."""
 
     # pylint: disable=W1401
+    @_thicket_console_renderer_annotate
     def render_preamble(self):
         lines = [
             r"  _____ _     _      _        _   ",
@@ -25,6 +30,7 @@ def render_preamble(self):
 
         return "\n".join(lines)
 
+    @_thicket_console_renderer_annotate
     def render(self, roots, dataframe, **kwargs):
         self.render_header = kwargs["render_header"]
 
@@ -143,6 +149,7 @@ def render(self, roots, dataframe, **kwargs):
         else:
             return result.encode("utf-8")
 
+    @_thicket_console_renderer_annotate
     def render_legend(self):
         def render_label(index, low, high):
             metric_range = self.max_metric - self.min_metric
@@ -218,6 +225,7 @@ def render_label(index, low, high):
 
         return legend
 
+    @_thicket_console_renderer_annotate
     def render_frame(self, node, dataframe, indent="", child_indent=""):
         node_depth = node._depth
         if node_depth < self.depth:

thicket/groupby.py

@@ -7,12 +7,19 @@
 
 import pandas as pd
 
+from hatchet.util.perf_measure import annotate
+
+
+_groupby_annotate = annotate(fmt="GroupBy.{}")
+
 
 class GroupBy(dict):
+    @_groupby_annotate
     def __init__(self, by=None, *args, **kwargs):
         super(GroupBy, self).__init__(*args, **kwargs)
         self.by = by
 
+    @_groupby_annotate
     def agg(self, func, disable_tqdm=False):
         """Aggregate the Thickets' PerfData numerical columns in a GroupBy object.
 
@@ -32,6 +39,7 @@ def agg(self, func, disable_tqdm=False):
 
         return agg_tk
 
+    @_groupby_annotate
     def aggregate_thicket(self, tk, func):
         """Aggregate a Thicket's numerical columns given a statistical function.
 

thicket/helpers.py

@@ -5,8 +5,10 @@
 
 from more_itertools import powerset
 import pandas as pd
+from hatchet.util.perf_measure import annotate
 
 
+@annotate()
 def _are_synced(gh, df):
     """Check if node objects are equal in graph and dataframe id(graph_node) ==
     id(df_node).
@@ -22,6 +24,7 @@ def _are_synced(gh, df):
     return True
 
 
+@annotate()
 def _missing_nodes_to_list(a_df, b_df):
     """Get a list of node differences between two dataframes. Mainly used for "tree"
     function.
@@ -65,6 +68,7 @@ def _missing_nodes_to_list(a_df, b_df):
     return missing_nodes
 
 
+@annotate()
 def _new_statsframe_df(df, multiindex=False):
     """Generate new aggregated statistics table from a dataframe. This is most commonly
     needed when changes are made to the performance data table's index.
@@ -92,6 +96,7 @@ def _new_statsframe_df(df, multiindex=False):
     return new_df
 
 
+@annotate()
 def _print_graph(graph):
     """Print the nodes in a hatchet graph"""
     i = 0
@@ -101,6 +106,7 @@ def _print_graph(graph):
     return i
 
 
+@annotate()
 def _resolve_missing_indicies(th_list):
     """Resolve indices if at least 1 profile has an index that another doesn't.
 
@@ -124,6 +130,7 @@ def _resolve_missing_indicies(th_list):
                 th.dataframe.set_index(idx, append=True, inplace=True)
 
 
+@annotate()
 def _set_node_ordering(thickets):
     """Set node ordering for each thicket in a list. All thickets must have node ordering on, otherwise it will be set to False.
 
@@ -139,6 +146,7 @@ def _set_node_ordering(thickets):
             tk.graph.enumerate_traverse()
 
 
+@annotate()
 def _get_perf_columns(df):
     """Get list of performance dataframe columns that are numeric.
 
@@ -165,6 +173,7 @@ def _get_perf_columns(df):
         return [x for x in numeric_columns if "nid" not in x]
 
 
+@annotate()
 def _powerset_from_tuple(tup):
     pset = [y for y in powerset(tup)]
     return {x[0] if len(x) == 1 else x for x in pset}

thicket/model_extrap.py

@@ -16,11 +16,16 @@
 )  # For some reason it errors if "Experiment" is not explicitly imported
 from extrap.fileio import io_helper
 from extrap.modelers.model_generator import ModelGenerator
+from hatchet.util.perf_measure import annotate
 
 
 MODEL_TAG = "_extrap-model"
 
 
+_extrap_model_wrapper_annotate = annotate(fmt="ExtrapModelWrapper.{}")
+_extrap_modelling_annotate = annotate(fmt="ExtrapModelling.{}")
+
+
 class ModelWrapper:
     """Wrapper for an Extra-P model.
 
@@ -29,18 +34,22 @@ class ModelWrapper:
     the model.
     """
 
+    @_extrap_model_wrapper_annotate
     def __init__(self, mdl, param_name):
         self.mdl = mdl
         self.param_name = param_name  # Needed for plotting / displaying the model
 
+    @_extrap_model_wrapper_annotate
     def __str__(self):
         """Display self as a function"""
         return str(self.mdl.hypothesis.function)
 
+    @_extrap_model_wrapper_annotate
     def eval(self, val):
         """Evaluate function (self) at val. f(val) = result"""
         return self.mdl.hypothesis.function.evaluate(val)
 
+    @_extrap_model_wrapper_annotate
     def display(self, RSS):
         """Display function
 
@@ -89,6 +98,7 @@ def display(self, RSS):
 class Modeling:
     """Produce models for all the metrics across the given graphframes."""
 
+    @_extrap_modelling_annotate
     def __init__(self, tht, param_name, params=None, chosen_metrics=None):
         """Create a new model object.
 
@@ -128,15 +138,16 @@ def __init__(self, tht, param_name, params=None, chosen_metrics=None):
                 v: k for k, v in self.tht.profile_mapping.items()
             }
             for file_name, value in params.items():
-                self.tht.metadata.at[
-                    profile_mapping_flipped[file_name], param_name
-                ] = value
+                self.tht.metadata.at[profile_mapping_flipped[file_name], param_name] = (
+                    value
+                )
             self.params = tht.metadata[param_name].tolist()
         if not chosen_metrics:
             self.chosen_metrics = self.tht.exc_metrics + self.tht.inc_metrics
         else:
             self.chosen_metrics = chosen_metrics
 
+    @_extrap_modelling_annotate
     def to_html(self, RSS=False):
         def model_to_img_html(model_obj):
             fig, ax = model_obj.display(RSS)
@@ -155,6 +166,7 @@ def model_to_img_html(model_obj):
             [met + MODEL_TAG for met in self.chosen_metrics]
         ].to_html(escape=False, formatters=frm_dict)
 
+    @_extrap_modelling_annotate
     def _add_extrap_statistics(self, node, metric):
         """Insert the Extra-P hypothesis function statistics into the aggregated
             statistics table. Has to be called after "produce_models".
@@ -167,22 +179,23 @@ def _add_extrap_statistics(self, node, metric):
             node, metric + MODEL_TAG
         ].mdl.hypothesis
 
-        self.tht.statsframe.dataframe.at[
-            node, metric + "_RSS" + MODEL_TAG
-        ] = hypothesis_fn.RSS
-        self.tht.statsframe.dataframe.at[
-            node, metric + "_rRSS" + MODEL_TAG
-        ] = hypothesis_fn.rRSS
-        self.tht.statsframe.dataframe.at[
-            node, metric + "_SMAPE" + MODEL_TAG
-        ] = hypothesis_fn.SMAPE
-        self.tht.statsframe.dataframe.at[
-            node, metric + "_AR2" + MODEL_TAG
-        ] = hypothesis_fn.AR2
-        self.tht.statsframe.dataframe.at[
-            node, metric + "_RE" + MODEL_TAG
-        ] = hypothesis_fn.RE
+        self.tht.statsframe.dataframe.at[node, metric + "_RSS" + MODEL_TAG] = (
+            hypothesis_fn.RSS
+        )
+        self.tht.statsframe.dataframe.at[node, metric + "_rRSS" + MODEL_TAG] = (
+            hypothesis_fn.rRSS
+        )
+        self.tht.statsframe.dataframe.at[node, metric + "_SMAPE" + MODEL_TAG] = (
+            hypothesis_fn.SMAPE
+        )
+        self.tht.statsframe.dataframe.at[node, metric + "_AR2" + MODEL_TAG] = (
+            hypothesis_fn.AR2
+        )
+        self.tht.statsframe.dataframe.at[node, metric + "_RE" + MODEL_TAG] = (
+            hypothesis_fn.RE
+        )
 
+    @_extrap_modelling_annotate
     def produce_models(self, agg_func=mean, add_stats=True):
         """Produces an Extra-P model. Models are generated by calling Extra-P's
             ModelGenerator.
@@ -282,6 +295,7 @@ def produce_models(self, agg_func=mean, add_stats=True):
                 if add_stats:
                     self._add_extrap_statistics(node, met)
 
+    @_extrap_modelling_annotate
     def _componentize_function(model_object):
         """Componentize one Extra-P modeling object into a dictionary of its parts
 
@@ -307,6 +321,7 @@ def _componentize_function(model_object):
 
         return term_dict
 
+    @_extrap_modelling_annotate
     def componentize_statsframe(self, columns=None):
         """Componentize multiple Extra-P modeling objects in the aggregated statistics
         table