improve metricset handling of tracesets

Author: telemin

pypop/metrics/hybrid.py

@@ -37,7 +37,8 @@ class MPI_OpenMP_Ineff_Metrics(MetricSet):
         Metric("IPC Scaling", 2, "IPC Scaling"),
     ]
 
-    _default_metric_key = "Hybrid Layout"
+    _default_metric_key = "Number of Processes"
+    _default_group_key = "Threads per Process"
 
     def _calculate_metrics(self, ref_key=None, sort_keys=True):
         if not ref_key:
@@ -186,7 +187,8 @@ class MPI_OpenMP_Metrics(MetricSet):
         Metric("IPC Scaling", 2, "IPC Scaling"),
     ]
 
-    _default_metric_key = "Hybrid Layout"
+    _default_metric_key = "Number of Processes"
+    _default_group_key = "Threads per Process"
 
     def _calculate_metrics(self, ref_key=None, sort_keys=True):
         if not ref_key:
@@ -344,7 +346,8 @@ class MPI_OpenMP_Multiplicative_Metrics(MetricSet):
         Metric("IPC Scaling", 2, "IPC Scaling"),
     ]
 
-    _default_metric_key = "Hybrid Layout"
+    _default_metric_key = "Number of Processes"
+    _default_group_key = "Threads per Process"
 
     def _calculate_metrics(self, ref_key=None, sort_keys=True):
         if not ref_key:

pypop/metrics/judit.py

@@ -32,7 +32,8 @@ class Judit_Hybrid_Metrics(MetricSet):
         Metric("IPC Scaling", 2, "IPC Scaling"),
     ]
 
-    _default_metric_key = "Hybrid Layout"
+    _default_metric_key = "Number of Processes"
+    _default_group_key = "Threads per Process"
 
     def _calculate_metrics(self, ref_key=None, sort_keys=True):
         if not ref_key:

pypop/metrics/metricset.py

@@ -17,6 +17,7 @@
 import matplotlib.ticker as mtick
 
 from ..trace import Trace
+from ..traceset import TraceSet
 from .._plotsettings import pypop_mpl_params, figparams
 
 __all__ = ["Metric", "MetricSet"]
@@ -54,12 +55,20 @@ class MetricSet:
     """
 
     _default_metric_key = None
+    _default_group_key = None
 
-    def __init__(self, stats_dict, ref_key=None, sort_keys=True):
+    _key_descriptions = {
+        "Number of Processes": "",
+        "Threads per Process": "",
+        "Total Threads": "",
+        "Hybrid Layout": "",
+    }
+
+    def __init__(self, stats_data, ref_key=None, sort_keys=True):
         """
         Parameters
         ----------
-        stats_dict: dict or list of `pd.DataFrame`
+        stats_data: TraceSet instance, dict, iterable or instance of Trace
             Statistics as collected with `collect_statistics()`. Dictionary keys will be
             used as the dataframe index. If a list, a dict will be constructed by
             enumeration.
@@ -73,16 +82,7 @@ def __init__(self, stats_dict, ref_key=None, sort_keys=True):
             If true (default), lexically sort the keys in the returned DataFrame.
         """
 
-        if not isinstance(stats_dict, dict):
-            stats_dict = {k: v for k, v in enumerate(stats_dict)}
-
-        for df in stats_dict.values():
-            if not isinstance(df, Trace):
-                raise ValueError(
-                    "stats_dict must be an iterable of pypop.traceset.RunData"
-                )
-
-        self._stats_dict = stats_dict
+        self._stats_dict = MetricSet._dictify_stats(stats_data)
         self._metric_data = None
         self._ref_key = ref_key
         self._sort_keys = sort_keys
@@ -101,6 +101,22 @@ def metric_data(self):
             self._calculate_metrics(ref_key=self._ref_key)
         return self._metric_data
 
+    @staticmethod
+    def _dictify_stats(stats_data):
+        if isinstance(stats_data, TraceSet):
+            return {k: v for k, v in enumerate(stats_data.traces)}
+        else:
+            if isinstance(stats_data, Trace):
+                return {0: stats_data}
+            if not isinstance(stats_data, dict):
+                stats_data = {k: v for k, v in enumerate(stats_data)}
+
+        for df in stats_data.values():
+            if not isinstance(df, Trace):
+                raise ValueError("stats_dict must be an iterable of pypop.trace.Trace")
+
+            return stats_data
+
     @property
     def metrics(self):
         """List of :py:class:`pypop.metrics.Metric`: List of metrics that will be

pypop/traceset.py

@@ -176,7 +176,7 @@ def by_threads_per_process(self):
 
         This is a helper function equivalent to
 
-        by_key(lambda x: x.metadata.rank_threads[0])
+        by_key(lambda x: x.metadata.threads_per_process[0])
 
         Returns
         -------
@@ -185,3 +185,26 @@ def by_threads_per_process(self):
         """
 
         return self.by_key(lambda x: x.metadata.threads_per_process[0])
+
+    def by_hybrid_layout(self):
+        """Return a dictionary of traces keyed by hybrid layout
+
+        This is a helper function equivalent to
+
+        by_key(
+            lambda x: "{}x{}".format(
+                x.metadata.num_processes, x.metadata.threads_per_process[0]
+            )
+        )
+
+        Returns
+        -------
+        traces_by_key: dict
+            A dictionary of traces organised by the requested key
+        """
+
+        return self.by_key(
+            lambda x: "{}x{}".format(
+                x.metadata.num_processes, x.metadata.threads_per_process[0]
+            )
+        )