Added Seaborn SNS Heatmap, SNS Pairplot nodes.

modules/seaborn_wrapper/sns_heatmap.py

@@ -0,0 +1,67 @@
+from typing import Any, Dict
+from matplotlib import pyplot as plt
+import pandas as pd
+import seaborn as sns
+from .util import plot_post_steps
+
+
+class SNSHeatmap:
+    """
+    SNS Heatmap:
+    Generates a heatmap from a Pandas DataFrame using Seaborn.
+
+    category: Plot
+    """
+    
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the `bar_chart` function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "dataframe": ("DATAFRAME", {}),
+                "title": ("STRING", {"default": ""}),
+                "palette": ("STRING", {"default": ""})
+            }
+        }
+
+    RETURN_TYPES: tuple = ("IMAGE",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Data Analysis"
+
+    def f(self,
+         dataframe: pd.DataFrame,
+         title: str,
+         palette: str=None) -> tuple:
+        """
+        Generates a heatmap from a Pandas DataFrame using Seaborn.
+
+        Args:
+            dataframe (DataFrame): The DataFrame.
+            title (str): The title of the plot.
+            palette (str): Palette of the plot
+        Returns:
+            tuple: A tuple containing the image tensor representation of the plot.
+        """
+        # set_up_sns_style_palette(style, palette) won't work for heatmap
+
+        # Create the plot
+        fig, ax = plt.subplots()
+        corr = dataframe.corr()
+        if palette:
+            sns.heatmap(corr, annot=True, ax=ax, cmap=sns.color_palette(palette, as_cmap=True))
+        else:
+            sns.heatmap(corr, annot=True, ax=ax)
+
+        return plot_post_steps(
+            fig,
+            ax,
+            plt,
+            title,
+            None,
+            None,
+            None)

modules/seaborn_wrapper/sns_pairplot.py

@@ -0,0 +1,79 @@
+from typing import Any, Dict
+from matplotlib import pyplot as plt
+import pandas as pd
+import seaborn as sns
+from .util import plot_post_steps
+from .util import set_up_sns_style_palette
+
+
+class SNSPairplot:
+    """
+    SNS Pairplot:
+    Generates a pairplot from a Pandas DataFrame using Seaborn.
+
+    Specify a column name for the hue parameter to color the data points based on that column.
+
+    category: Plot
+    """
+    
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the `bar_chart` function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "dataframe": ("DATAFRAME", {}),
+                "hue_column_name": ("STRING", {"default": ""}),
+                "title": ("STRING", {"default": ""}),
+                "style": ("STRING", {"default": ""}),
+                "palette": ("STRING", {"default": ""})
+            }
+        }
+
+    RETURN_TYPES: tuple = ("IMAGE",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Data Analysis"
+
+    def f(self,
+                  dataframe: pd.DataFrame,
+                  hue_column_name: str,
+                  title: str,
+                  style: str=None,
+                  palette: str=None) -> tuple:
+        """
+        Generates a pairplot from a pandas DataFrame.
+
+        Args:
+            dataframe (DataFrame): The DataFrame.
+            hue_column_name (str): The category (variable) column name.
+            title (str): The title of the plot.
+            style (str): Style of the plot
+            palette (str): Palette of the plot
+        Returns:
+            tuple: A tuple containing the image tensor representation of the plot.
+        """
+        df = dataframe
+        set_up_sns_style_palette(style, palette)
+
+        # Create the plot
+        if hue_column_name:
+            pp = sns.pairplot(df, hue=hue_column_name)
+        else:
+            pp = sns.pairplot(df)
+        
+        pp.figure.suptitle(title, y=0.95)
+        pp.figure.tight_layout(rect=[0, 0, 1, 0.95]) 
+        fig = pp.figure
+
+        return plot_post_steps(
+            fig,
+            None,
+            plt,
+            title,
+            None,
+            None,
+            None)

modules/seaborn_wrapper/util.py

@@ -89,10 +89,11 @@ def plot_post_steps(
         ax.xaxis.set_major_locator(mticker.MaxNLocator(integer=True))
 
     # Set title and labels
-    ax.set_title(title)
-    if x_axis_label:
-        ax.set_xlabel(x_axis_label)
-    ax.set_ylabel(y_axis_label)
+    if ax:
+        ax.set_title(title)
+        if x_axis_label:
+            ax.set_xlabel(x_axis_label)
+        ax.set_ylabel(y_axis_label)
 
     # Save the plot to a buffer
     buf = io.BytesIO()