From eebb027ddcda6717c5c007dcb932ebed59c3f668 Mon Sep 17 00:00:00 2001
From: Javier Cladellas <javiercladellas@gmail.com>
Date: Thu, 23 Jan 2025 17:36:22 +0100
Subject: [PATCH] add heatmap figure #156

---
 .../pages/configurationfiles/plots.adoc       | 31 ++++++++++++
 .../reframe/config/configPlots.py             |  2 +-
 .../benchmarking/report/figureFactory.py      | 47 +++++++++++++++++++
 tests/report/test_figureFactory.py            |  3 +-
 4 files changed, 81 insertions(+), 2 deletions(-)

diff --git a/docs/modules/tutorial/pages/configurationfiles/plots.adoc b/docs/modules/tutorial/pages/configurationfiles/plots.adoc
index 4dab7308..95c079c0 100644
--- a/docs/modules/tutorial/pages/configurationfiles/plots.adoc
+++ b/docs/modules/tutorial/pages/configurationfiles/plots.adoc
@@ -284,6 +284,36 @@ fig = figures[0].createFigure(model.master_df)
 fig.show()
 ----
 
+- `heatmap`
+
+For this case, we will consider the `elements` (N) as `color_axis` and `performance_variable` for secondary axis (slider).
+
+[%dynamic%open%hide_code,python]
+----
+figures = FigureFactory.create(Plot(**{
+    "title": "Absolute performance",
+    "plot_types": [ "heatmap" ],
+    "transformation": "performance",
+    "variables": [ "computation_time","communication_time" ],
+    "names": ["Time"],
+    "color_axis":{
+        "parameter":"elements",
+        "label":"N"
+    },
+    "yaxis":{"label":"Execution time (s)"},
+    "secondary_axis":{
+        "parameter":"performance_variable",
+        "label":"Performance variable"
+    },
+    "xaxis":{
+        "parameter":"nb_tasks.tasks",
+        "label":"Number of tasks"
+    }
+}))
+fig = figures[0].createFigure(model.master_df)
+fig.show()
+----
+
 - `table`
 
 [%dynamic%open%hide_code,python]
@@ -312,6 +342,7 @@ fig = figures[0].createFigure(model.master_df)
 fig.show()
 ----
 
+
 == Aggregations
 
 Depending on the dashboard level that we are located at, it might be necessary to aggregate the data on the master dataframe.
diff --git a/src/feelpp/benchmarking/reframe/config/configPlots.py b/src/feelpp/benchmarking/reframe/config/configPlots.py
index 0ac89960..7f90f64b 100644
--- a/src/feelpp/benchmarking/reframe/config/configPlots.py
+++ b/src/feelpp/benchmarking/reframe/config/configPlots.py
@@ -18,7 +18,7 @@ def checkAgg(cls, v):
 
 class Plot(BaseModel):
     title:str
-    plot_types:List[Literal["scatter","table","stacked_bar","grouped_bar"]]
+    plot_types:List[Literal["scatter","table","stacked_bar","grouped_bar","heatmap"]]
     transformation:Literal["performance","relative_performance","speedup"]
     aggregations:Optional[List[Aggregation]] = None
     variables:Optional[List[str]] = None
diff --git a/src/feelpp/benchmarking/report/figureFactory.py b/src/feelpp/benchmarking/report/figureFactory.py
index 68c1152e..22797da7 100644
--- a/src/feelpp/benchmarking/report/figureFactory.py
+++ b/src/feelpp/benchmarking/report/figureFactory.py
@@ -282,6 +282,51 @@ def createSimpleFigure(self,df):
         """
         return go.Figure(self.createTraces(df))
 
+class HeatmapFigure(Figure):
+    def __init__(self, plot_config, transformation_strategy):
+        super().__init__(plot_config, transformation_strategy)
+
+    def createTraces(self, df):
+        """ Creates the Heatmap traces for a given dataframe. Useful for animation creation.
+        Args:
+            - df (pd.DataFrame): The dataframe containing the figure data.
+        Returns: (go.Trace) The Heatmap traces to display in the scatter figure.
+        """
+        return go.Heatmap(
+                x=df.index.astype(str),
+                y=df.columns.astype(str),
+                z=df.T.values,
+                colorbar = dict(
+                    title = self.config.yaxis.label
+                )
+            )
+
+    def createMultiindexFigure(self, df):
+        """ Creates a plotly figure from a multiIndex dataframe
+        Args:
+            df (pd.DataFrame). The transformed dataframe (must be multiindex)
+        Returns:
+            go.Figure: Heatmap animation where the secondary_ axis corresponds to a specified parameter
+        """
+        return self.createSliderAnimation(df)
+
+    def createSimpleFigure(self, df):
+        """ Creates a plotly figure from a given dataframe
+        Args:
+            df (pd.DataFrame). The transformed dataframe
+        Returns:
+            go.Figure: Heatmap plot
+        """
+        return go.Figure(self.createTraces(df))
+
+    def updateLayout(self, fig):
+        """ Sets the title, yaxis and legend attributes of the layout"""
+        fig.update_layout(
+            title=self.config.title,
+            xaxis=dict(title=self.config.xaxis.label),
+            yaxis=dict(title=self.config.color_axis.label)
+        )
+        return fig
 
 class FigureFactory:
     """ Factory class to dispatch concrete figure elements"""
@@ -305,6 +350,8 @@ def create(plot_config):
                 figures.append(StackedBarFigure(plot_config,strategy))
             elif plot_type == "grouped_bar":
                 figures.append(GroupedBarFigure(plot_config,strategy))
+            elif plot_type == "heatmap":
+                figures.append(HeatmapFigure(plot_config,strategy))
             else:
                 raise NotImplementedError
 
diff --git a/tests/report/test_figureFactory.py b/tests/report/test_figureFactory.py
index e8f93149..87c6cc41 100644
--- a/tests/report/test_figureFactory.py
+++ b/tests/report/test_figureFactory.py
@@ -1,4 +1,4 @@
-from feelpp.benchmarking.report.figureFactory import FigureFactory, ScatterFigure, TableFigure, StackedBarFigure, GroupedBarFigure
+from feelpp.benchmarking.report.figureFactory import FigureFactory, ScatterFigure, TableFigure, StackedBarFigure, GroupedBarFigure, HeatmapFigure
 from test_transformationFactory import PlotConfigMocker
 import pytest
 
@@ -10,6 +10,7 @@
     (["table"],[TableFigure]),
     (["stacked_bar"],[StackedBarFigure]),
     (["grouped_bar"],[GroupedBarFigure]),
+    (["heatmap"],[HeatmapFigure]),
     (["unkown"],[])
 ])
 def test_figureFactory(types,expected_classes):