Adds function that moves perf metric to statsframe (LLNL#198)

* Adds function that moves perf metric to statsframe

* Adds test for move_metrics_to_statsframe

* Changes metric_columns argument to move_column_to_statsframe to a list type

thicket/tests/test_thicket.py

@@ -104,6 +104,46 @@ def test_metadata_column_to_perfdata(mpi_scaling_cali):
         assert metric in values
 
 
+def test_perfdata_column_to_statsframe(literal_thickets, mpi_scaling_cali):
+    th_single = literal_thickets[1].deepcopy()
+
+    with pytest.raises(KeyError):
+        th_single.move_metrics_to_statsframe(["dummy"])
+
+    th_single.move_metrics_to_statsframe(["time"])
+    assert all(
+        th_single.dataframe["time"].values
+        == th_single.statsframe.dataframe["time"].values
+    )
+
+    with pytest.raises(KeyError):
+        th_single.move_metrics_to_statsframe(["time"])
+
+    th_single.move_metrics_to_statsframe(["time", "memory"], override=True)
+    assert all(
+        th_single.dataframe["time"].values
+        == th_single.statsframe.dataframe["time"].values
+    )
+    assert all(
+        th_single.dataframe["memory"].values
+        == th_single.statsframe.dataframe["memory"].values
+    )
+
+    th_mpi = Thicket.from_caliperreader(mpi_scaling_cali)
+    metrics = ["Min time/rank", "Max time/rank", "Avg time/rank", "Total time"]
+    idx = pd.IndexSlice
+
+    with pytest.raises(ValueError):
+        th_mpi.move_metrics_to_statsframe(metrics, profile="fake")
+
+    th_mpi.move_metrics_to_statsframe(metrics, profile=th_mpi.profile[0])
+    for met in metrics:
+        assert all(
+            th_mpi.dataframe.loc[idx[:, th_mpi.profile[0]], :][met].values
+            == th_mpi.statsframe.dataframe[met].values
+        )
+
+
 def test_thicketize_graphframe(rajaperf_seq_O3_1M_cali):
     ht1 = ht.GraphFrame.from_caliperreader(rajaperf_seq_O3_1M_cali[-1])
     th1 = Thicket.thicketize_graphframe(ht1, rajaperf_seq_O3_1M_cali[-1])

thicket/thicket.py

@@ -1482,6 +1482,37 @@ def filter_stats(self, filter_function):
 
         return new_thicket
 
+    def move_metrics_to_statsframe(self, metric_columns, profile=None, override=False):
+        if not isinstance(metric_columns, (list, tuple)):
+            raise TypeError("'metric_columns' must be a list or tuple")
+        profile_list = self.dataframe.index.unique(level="profile").tolist()
+        if profile is None and len(profile_list) != 1:
+            raise ValueError(
+                "Cannot move a metric to statsframe when there are multiple profiles. Set the 'profile' argument to the profile you want to move"
+            )
+        if profile is not None and profile not in profile_list:
+            raise ValueError("Invalid profile: {}".format(profile))
+        df_for_profile = None
+        if profile is None:
+            df_for_profile = self.dataframe.reset_index(level="profile", drop=True)
+        else:
+            df_for_profile = self.dataframe.xs(
+                profile, level="profile", drop_level=True
+            )
+        new_statsframe_df = self.statsframe.dataframe.copy(deep=True)
+        for c in metric_columns:
+            if c in new_statsframe_df and not override:
+                raise KeyError(
+                    "Column {} is already in statsframe. To replace the column, run with 'override=True'".format(
+                        c
+                    )
+                )
+            new_statsframe_df[c] = df_for_profile[c]
+        self.statsframe.dataframe = new_statsframe_df
+        verify_thicket_structures(
+            self.statsframe.dataframe, columns=metric_columns, index=["node"]
+        )
+
     def get_unique_metadata(self):
         """Get unique values per column in metadata.