nodes/HowToSD-ComfyUI-Data-Analysis
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added SNS plot nodes that take PyTorch tensors.

Browse Source
This commit is contained in:
HowToSD
2025-03-18 17:46:05 -07:00
parent a8aaa21ce0
commit 9f0304b2ee
18 changed files with 709 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+12
View File
@@ -174,12 +174,16 @@ from .modules.pandas_wrapper.pandas_vertical_concat import PandasVerticalConcat
from .modules.pandas_wrapper.pandas_vertical_split import PandasVerticalSplit
from .modules.pandas_wrapper.pandas_xs import PandasXs
from .modules.seaborn_wrapper.sns_bar import SNSBar
from .modules.seaborn_wrapper.sns_bar_pt import SNSBarPt
from .modules.seaborn_wrapper.sns_boxplot import SNSBoxplot
from .modules.seaborn_wrapper.sns_heatmap import SNSHeatmap
from .modules.seaborn_wrapper.sns_histogram import SNSHistogram
from .modules.seaborn_wrapper.sns_histogram_pt import SNSHistogramPt
from .modules.seaborn_wrapper.sns_line import SNSLine
from .modules.seaborn_wrapper.sns_line_pt import SNSLinePt
from .modules.seaborn_wrapper.sns_pairplot import SNSPairplot
from .modules.seaborn_wrapper.sns_scatter import SNSScatter
from .modules.seaborn_wrapper.sns_scatter_pt import SNSScatterPt
T = TypeVar("T")
@@ -363,12 +367,16 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PyStringToInt": PyStringToInt,
"PyStringToList": PyStringToList,
"SNSBar": SNSBar,
"SNSBarPt": SNSBarPt,
"SNSBoxplot": SNSBoxplot,
"SNSHeatmap": SNSHeatmap,
"SNSHistogram": SNSHistogram,
"SNSHistogramPt": SNSHistogramPt,
"SNSLine": SNSLine,
"SNSLinePt": SNSLinePt,
"SNSPairplot": SNSPairplot,
"SNSScatter": SNSScatter,
"SNSScatterPt": SNSScatterPt,
}
@@ -552,12 +560,16 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PyStringToInt": "Py String To Int",
"PyStringToList": "Py String To List",
"SNSBar": "SNS Bar Chart",
"SNSBarPt": "SNS Bar Chart Pt",
"SNSBoxplot": "SNS Boxplot",
"SNSHeatmap": "SNS Heatmap",
"SNSHistogram": "SNS Histogram",
"SNSHistogramPt": "SNS Histogram Pt",
"SNSLine": "SNS Line Plot",
"SNSLinePt": "SNS Line Plot Pt",
"SNSPairplot": "SNS Pairplot",
"SNSScatter": "SNS Scatter Plot",
"SNSScatterPt": "SNS Scatter Plot Pt",
}
docs/reference/node_reference.md
+4
View File
@@ -170,12 +170,16 @@
| [MPL Pie Chart](mpl_pie_chart.md) | Generates a pie chart from a pandas DataFrame. |
| [MPL Scatter Plot](mpl_scatter.md) | Generates a scatter plot from a pandas DataFrame. |
| [SNS Bar Chart](sns_bar.md) | Generates a bar chart from a Pandas DataFrame using Seaborn. |
| [SNS Bar Chart Pt](sns_bar_pt.md) | Generates a bar chart from a PyTorch tensor using Seaborn. |
| [SNS Boxplot](sns_boxplot.md) | Generates a boxplot from a Pandas DataFrame using Seaborn. |
| [SNS Heatmap](sns_heatmap.md) | Generates a heatmap from a Pandas DataFrame using Seaborn. |
| [SNS Histogram](sns_histogram.md) | Generates a histogram from a Pandas DataFrame using Seaborn. |
| [SNS Histogram Chart Pt](sns_histogram_pt.md) | Generates a histogram from a PyTorch tensor using Seaborn. |
| [SNS Line Plot](sns_line.md) | Generates a line plot from a pandas DataFrame using Seaborn. |
| [SNS Line Plot Pt](sns_line_pt.md) | Generates a line plot from a PyTorch tensor using Seaborn. |
| [SNS Pairplot](sns_pairplot.md) | Generates a pairplot from a Pandas DataFrame using Seaborn. |
| [SNS Scatter Plot](sns_scatter.md) | Generates a scatter plot from a pandas DataFrame using Seaborn. |
| [SNS Scatter Plot Pt](sns_scatter_pt.md) | Generates a scatter plot from a PyTorch tensor using Seaborn. |
## Python wrapper
| Node | Description |
| --- | --- |
docs/reference/pandas_feature_split_to_pt.md
+2 -2
View File
@@ -14,8 +14,8 @@ If output_1d_label is set, label will be output as a 1d array instead of a 2d ar
## Output
| Data type |
|---|
| Ndarray |
| Ndarray |
| Tensor |
| Tensor |
<HR>
Category: Transformation
docs/reference/sns_bar_pt.md
+12
View File
@@ -0,0 +1,12 @@
# SNS Bar Chart Pt
Generates a bar chart from a PyTorch tensor using Seaborn.
## Output
| Data type |
|---|
| Image |
<HR>
Category: Plot
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/sns_histogram_pt.md
+12
View File
@@ -0,0 +1,12 @@
# SNS Histogram Chart Pt
Generates a histogram from a PyTorch tensor using Seaborn.
## Output
| Data type |
|---|
| Image |
<HR>
Category: Plot
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/sns_line_pt.md
+16
View File
@@ -0,0 +1,16 @@
# SNS Line Plot Pt
Generates a line plot from a PyTorch tensor using Seaborn.
For plotting multiple lines, enter a comma-separated column labels in y_axis_dims.
If you want to display a legend for each line, enter a comma-separated legend label
in legend label.
## Output
| Data type |
|---|
| Image |
<HR>
Category: Plot
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/sns_scatter_pt.md
+16
View File
@@ -0,0 +1,16 @@
# SNS Scatter Plot Pt
Generates a scatter plot from a PyTorch tensor using Seaborn.
For plotting multiple scatters, enter a comma-separated column labels in y_axis_dims.
If you want to display a legend for each scatter, enter a comma-separated legend label
in legend label.
## Output
| Data type |
|---|
| Image |
<HR>
Category: Plot
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
modules/seaborn_wrapper/common.py
+45
View File
@@ -20,6 +20,29 @@ def common_input_types():
}
}
def common_input_types_pt():
"""
Defines the input types for a plot function that takes PyTorch tensor.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"x_axis_dim": ("INT", {"default": 0, "min":-100, "max":1e6}),
"y_axis_dims": ("STRING", {"default": ""}),
"title": ("STRING", {"default": ""}),
"x_axis_label": ("STRING", {"default": ""}),
"y_axis_label": ("STRING", {"default": ""}),
"legend_label": ("STRING", {"default": ""}),
"x_tick_as_int": ("BOOLEAN", {"default": False}), # To show 2002 instead of 2002.5 on x-axis
"style": ("STRING", {"default": ""}),
"palette": ("STRING", {"default": ""})
}
}
def common_input_types_y_only_bin():
"""
Defines the input types for a plot function for types of charts that do
@@ -40,3 +63,25 @@ def common_input_types_y_only_bin():
"palette": ("STRING", {"default": ""})
}
}
def common_input_types_y_only_bin_pt():
"""
Defines the input types for a plot function for types of charts that do
not require x-series data as input, but takes bins.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"y_axis_dims": ("STRING", {"default": ""}),
"title": ("STRING", {"default": ""}),
"y_axis_label": ("STRING", {"default": ""}),
# "legend_label": ("STRING", {"default": ""}),
"bins": ("INT", {"default": 10, "min": 2, "max": 200}),
"style": ("STRING", {"default": ""}),
"palette": ("STRING", {"default": ""})
}
}
modules/seaborn_wrapper/sns_bar_pt.py
+101
View File
@@ -0,0 +1,101 @@
import os
import sys
from typing import Any, Dict
from matplotlib import pyplot as plt
import seaborn as sns
import torch
import ast
import pandas as pd
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from seaborn_wrapper.common import common_input_types_pt
from seaborn_wrapper.util import plot_post_steps
from seaborn_wrapper.util import set_up_sns_style_palette
from seaborn_wrapper.util import comma_separated_labels_string_to_list
class SNSBarPt:
"""
SNS Bar Chart Pt:
Generates a bar chart from a PyTorch tensor using Seaborn.
category: Plot
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the `bar_plot` function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return common_input_types_pt()
RETURN_TYPES: tuple = ("IMAGE",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self,
tens: torch.Tensor,
x_axis_dim: int,
y_axis_dims: str,
title: str,
x_axis_label: str,
y_axis_label: str,
legend_label: str,
x_tick_as_int: bool,
style: str=None,
palette: str=None) -> tuple:
"""
Generates a bar chart from a pandas DataFrame.
Args:
tens (torch.Tensor): The tensor.
x_axis_dim (int): The index in tensor for the x-axis.
y_axix_dims (str): The indices in tensor for the y-axis.
title (str): The title of the plot.
x_axis_label (str): The label for the x-axis data.
y_axis_label (str): The label for the y-axis data.
legend_label (str): The legend label of the plot.
x_tick_as_int (bool): If set, use integers for x ticks. For example,
if ticks show floating point numbers for years, e.g. 2002.5,
setting this flag will show integer values instead.
style (str): Style of the plot
palette (str): Palette of the plot
Returns:
tuple: A tuple containing the image tensor representation of the plot.
"""
y_axes = ast.literal_eval(y_axis_dims)
if isinstance(y_axes, int):
y_axes = [y_axes]
if tens.dim() != 2:
raise ValueError("Only rank 2 tensors are supported.")
label_list = comma_separated_labels_string_to_list(legend_label)
if len(y_axes) != len(label_list):
if len(label_list) > len(y_axes):
raise ValueError("More labels were specified than y-axis indices.")
else:
label_list += [""] * (len(y_axes) - len(label_list))
set_up_sns_style_palette(style, palette)
fig, ax = plt.subplots()
for i in range(len(y_axes)):
sns.barplot(
x=tens[:, x_axis_dim].detach().cpu().numpy(),
y=tens[:, y_axes[i]].detach().cpu().numpy(),
ax=ax, label=label_list[i])
return plot_post_steps(
fig,
ax,
plt,
title,
x_axis_label,
y_axis_label,
x_tick_as_int)
modules/seaborn_wrapper/sns_histogram.py
+4 -4
View File
@@ -22,7 +22,7 @@ class SNSHistogram:
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the `bar_chart` function.
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
@@ -30,10 +30,10 @@ class SNSHistogram:
return common_input_types_y_only_bin()
RETURN_TYPES: tuple = ("IMAGE",)
FUNCTION: str = "bar_chart"
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def bar_chart(self,
def f(self,
dataframe: pd.DataFrame,
y_column_name: str,
title: str,
@@ -43,7 +43,7 @@ class SNSHistogram:
style: str=None,
palette: str=None) -> tuple:
"""
Generates a bar chart from a pandas DataFrame.
Generates a histogram from a Pandas DataFrame using Seaborn.
Args:
dataframe (DataFrame): The DataFrame.
modules/seaborn_wrapper/sns_histogram_pt.py
+93
View File
@@ -0,0 +1,93 @@
import os
import sys
from typing import Any, Dict
from matplotlib import pyplot as plt
import seaborn as sns
import torch
import ast
import pandas as pd
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from seaborn_wrapper.common import common_input_types_y_only_bin_pt
from seaborn_wrapper.util import plot_post_steps
from seaborn_wrapper.util import set_up_sns_style_palette
class SNSHistogramPt:
"""
SNS Histogram Chart Pt:
Generates a histogram from a PyTorch tensor using Seaborn.
category: Plot
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return common_input_types_y_only_bin_pt()
RETURN_TYPES: tuple = ("IMAGE",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self,
tens: torch.Tensor,
y_axis_dims: str,
title: str,
y_axis_label: str,
# legend_label: str,
bins: int,
style: str=None,
palette: str=None) -> tuple:
"""
Generates a histogram from a PyTorch tensor using Seaborn.
Args:
tens (torch.Tensor): The tensor.
y_axix_dims (str): The indices in tensor for the y-axis.
y_column_name (str): The column name for the y-axis.
title (str): The title of the plot.
y_axis_label (str): The label for the y-axis data.
bins (int): Number of bins.
style (str): Style of the plot
palette (str): Palette of the plot
Returns:
tuple: A tuple containing the image tensor representation of the plot.
"""
y_axes = ast.literal_eval(y_axis_dims)
if isinstance(y_axes, int):
y_axes = [y_axes]
if tens.dim() > 2:
raise ValueError("Only rank 1 or 2 tensors are supported.")
if tens.dim() == 1:
tens2 = torch.unsqueeze(tens, 1)
else:
tens2 = tens
labels = [""] * len(y_axes)
set_up_sns_style_palette(style, palette)
# Create the plot
fig, ax = plt.subplots()
for i in range(len(y_axes)):
sns.histplot(tens2[:, y_axes[i]].detach().cpu().numpy(),
bins=bins, ax=ax, label=labels[i])
break # TODO: Remove this if we decide support multiple plots
return plot_post_steps(
fig,
ax,
plt,
title,
None,
y_axis_label,
None)
modules/seaborn_wrapper/sns_line_pt.py
+105
View File
@@ -0,0 +1,105 @@
import os
import sys
from typing import Any, Dict
from matplotlib import pyplot as plt
import seaborn as sns
import torch
import ast
import pandas as pd
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from seaborn_wrapper.common import common_input_types_pt
from seaborn_wrapper.util import plot_post_steps
from seaborn_wrapper.util import set_up_sns_style_palette
from seaborn_wrapper.util import comma_separated_labels_string_to_list
class SNSLinePt:
"""
SNS Line Plot Pt:
Generates a line plot from a PyTorch tensor using Seaborn.
For plotting multiple lines, enter a comma-separated column labels in y_axis_dims.
If you want to display a legend for each line, enter a comma-separated legend label
in legend label.
category: Plot
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the `line_plot` function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return common_input_types_pt()
RETURN_TYPES: tuple = ("IMAGE",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self,
tens: torch.Tensor,
x_axis_dim: int,
y_axis_dims: str,
title: str,
x_axis_label: str,
y_axis_label: str,
legend_label: str,
x_tick_as_int: bool,
style: str=None,
palette: str=None) -> tuple:
"""
Generates a line plot from a pandas DataFrame.
Args:
tens (torch.Tensor): The tensor.
x_axis_dim (int): The index in tensor for the x-axis.
y_axix_dims (str): The indices in tensor for the y-axis.
title (str): The title of the plot.
x_axis_label (str): The label for the x-axis data.
y_axis_label (str): The label for the y-axis data.
legend_label (str): The legend label of the plot.
x_tick_as_int (bool): If set, use integers for x ticks. For example,
if ticks show floating point numbers for years, e.g. 2002.5,
setting this flag will show integer values instead.
style (str): Style of the plot
palette (str): Palette of the plot
Returns:
tuple: A tuple containing the image tensor representation of the plot.
"""
y_axes = ast.literal_eval(y_axis_dims)
if isinstance(y_axes, int):
y_axes = [y_axes]
if tens.dim() != 2:
raise ValueError("Only rank 2 tensors are supported.")
label_list = comma_separated_labels_string_to_list(legend_label)
if len(y_axes) != len(label_list):
if len(label_list) > len(y_axes):
raise ValueError("More labels were specified than y-axis indices.")
else:
label_list += [""] * (len(y_axes) - len(label_list))
set_up_sns_style_palette(style, palette)
fig, ax = plt.subplots()
for i in range(len(y_axes)):
sns.lineplot(
x=tens[:, x_axis_dim].detach().cpu().numpy(),
y=tens[:, y_axes[i]].detach().cpu().numpy(),
ax=ax, label=label_list[i])
return plot_post_steps(
fig,
ax,
plt,
title,
x_axis_label,
y_axis_label,
x_tick_as_int)
modules/seaborn_wrapper/sns_scatter_pt.py
+105
View File
@@ -0,0 +1,105 @@
import os
import sys
from typing import Any, Dict
from matplotlib import pyplot as plt
import seaborn as sns
import torch
import ast
import pandas as pd
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from seaborn_wrapper.common import common_input_types_pt
from seaborn_wrapper.util import plot_post_steps
from seaborn_wrapper.util import set_up_sns_style_palette
from seaborn_wrapper.util import comma_separated_labels_string_to_list
class SNSScatterPt:
"""
SNS Scatter Plot Pt:
Generates a scatter plot from a PyTorch tensor using Seaborn.
For plotting multiple scatters, enter a comma-separated column labels in y_axis_dims.
If you want to display a legend for each scatter, enter a comma-separated legend label
in legend label.
category: Plot
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the `scatter_plot` function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return common_input_types_pt()
RETURN_TYPES: tuple = ("IMAGE",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self,
tens: torch.Tensor,
x_axis_dim: int,
y_axis_dims: str,
title: str,
x_axis_label: str,
y_axis_label: str,
legend_label: str,
x_tick_as_int: bool,
style: str=None,
palette: str=None) -> tuple:
"""
Generates a scatter plot from a pandas DataFrame.
Args:
tens (torch.Tensor): The tensor.
x_axis_dim (int): The index in tensor for the x-axis.
y_axix_dims (str): The indices in tensor for the y-axis.
title (str): The title of the plot.
x_axis_label (str): The label for the x-axis data.
y_axis_label (str): The label for the y-axis data.
legend_label (str): The legend label of the plot.
x_tick_as_int (bool): If set, use integers for x ticks. For example,
if ticks show floating point numbers for years, e.g. 2002.5,
setting this flag will show integer values instead.
style (str): Style of the plot
palette (str): Palette of the plot
Returns:
tuple: A tuple containing the image tensor representation of the plot.
"""
y_axes = ast.literal_eval(y_axis_dims)
if isinstance(y_axes, int):
y_axes = [y_axes]
if tens.dim() != 2:
raise ValueError("Only rank 2 tensors are supported.")
label_list = comma_separated_labels_string_to_list(legend_label)
if len(y_axes) != len(label_list):
if len(label_list) > len(y_axes):
raise ValueError("More labels were specified than y-axis indices.")
else:
label_list += [""] * (len(y_axes) - len(label_list))
set_up_sns_style_palette(style, palette)
fig, ax = plt.subplots()
for i in range(len(y_axes)):
sns.scatterplot(
x=tens[:, x_axis_dim].detach().cpu().numpy(),
y=tens[:, y_axes[i]].detach().cpu().numpy(),
ax=ax, label=label_list[i])
return plot_post_steps(
fig,
ax,
plt,
title,
x_axis_label,
y_axis_label,
x_tick_as_int)
modules/seaborn_wrapper/util.py
+7 -2
View File
@@ -1,3 +1,5 @@
import os
import sys
import io
from typing import List, Tuple
import numpy as np
@@ -11,9 +13,12 @@ import torch
import pandas as pd
import seaborn as sns
from typing import Optional
from ..pandas_wrapper.utils import column_labels_string_to_list
from ..pandas_wrapper.utils import comma_separated_labels_string_to_list
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from pandas_wrapper.utils import column_labels_string_to_list
from pandas_wrapper.utils import comma_separated_labels_string_to_list
def process_y_columns_and_labels(df:pd.DataFrame, y_column_name: str, legend_label: str) -> Tuple[List[str], List[str]]:
"""
pyproject.toml
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "data-analysis"
description = "Data analysis custom modules for ComfyUI – Integrates Pandas, Seaborn and Matplotlib into ComfyUI."
version = "3.0.5"
version = "3.0.6"
license = {file = "LICENSE"}
[project.urls]
tests/seaborn_wrapper/sns_bar_pt_test.py
+58
View File
@@ -0,0 +1,58 @@
import os
import sys
import unittest
import torch
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from seaborn_wrapper.sns_bar_pt import SNSBarPt
class TestSNSBarPt(unittest.TestCase):
def setUp(self):
"""Set up test data."""
self.node = SNSBarPt()
self.test_cases = [
(
torch.stack(
[
torch.arange(1, 10, dtype=torch.int64),
torch.arange(1, 10, dtype=torch.int64)*2
],
dim=1
),
"1", "y"
),
(
torch.stack(
[
torch.arange(1, 10, dtype=torch.int64),
torch.arange(1, 10, dtype=torch.int64)*2,
torch.arange(1, 10, dtype=torch.int64)*3
],
dim=1
),
"1,2", "profit,sales"
)
]
def test_1(self):
"""
Tests plot.
"""
for tens, y_axis_dims, labels in self.test_cases:
with self.subTest(data=y_axis_dims):
self.node.f(
tens=tens,
x_axis_dim = 0,
y_axis_dims = y_axis_dims,
title="test",
x_axis_label = "x",
y_axis_label = "y",
legend_label=labels,
x_tick_as_int=True)
if __name__ == "__main__":
unittest.main()
tests/seaborn_wrapper/sns_line_pt_test.py
+58
View File
@@ -0,0 +1,58 @@
import os
import sys
import unittest
import torch
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from seaborn_wrapper.sns_line_pt import SNSLinePt
class TestSNSLinePt(unittest.TestCase):
def setUp(self):
"""Set up test data."""
self.node = SNSLinePt()
self.test_cases = [
(
torch.stack(
[
torch.arange(1, 10, dtype=torch.int64),
torch.arange(1, 10, dtype=torch.int64)*2
],
dim=1
),
"1", "y"
),
(
torch.stack(
[
torch.arange(1, 10, dtype=torch.int64),
torch.arange(1, 10, dtype=torch.int64)*2,
torch.arange(1, 10, dtype=torch.int64)*3
],
dim=1
),
"1,2", "profit,sales"
)
]
def test_1(self):
"""
Tests plot.
"""
for tens, y_axis_dims, labels in self.test_cases:
with self.subTest(data=y_axis_dims):
self.node.f(
tens=tens,
x_axis_dim = 0,
y_axis_dims = y_axis_dims,
title="test",
x_axis_label = "x",
y_axis_label = "y",
legend_label=labels,
x_tick_as_int=True)
if __name__ == "__main__":
unittest.main()
tests/seaborn_wrapper/sns_scatter_pt_test.py
+58
View File
@@ -0,0 +1,58 @@
import os
import sys
import unittest
import torch
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from seaborn_wrapper.sns_scatter_pt import SNSScatterPt
class TestSNSScatterPt(unittest.TestCase):
def setUp(self):
"""Set up test data."""
self.node = SNSScatterPt()
self.test_cases = [
(
torch.stack(
[
torch.arange(1, 10, dtype=torch.int64),
torch.arange(1, 10, dtype=torch.int64)*2
],
dim=1
),
"1", "y"
),
(
torch.stack(
[
torch.arange(1, 10, dtype=torch.int64),
torch.arange(1, 10, dtype=torch.int64)*2,
torch.arange(1, 10, dtype=torch.int64)*3
],
dim=1
),
"1,2", "profit,sales"
)
]
def test_1(self):
"""
Tests plot.
"""
for tens, y_axis_dims, labels in self.test_cases:
with self.subTest(data=y_axis_dims):
self.node.f(
tens=tens,
x_axis_dim = 0,
y_axis_dims = y_axis_dims,
title="test",
x_axis_label = "x",
y_axis_label = "y",
legend_label=labels,
x_tick_as_int=True)
if __name__ == "__main__":
unittest.main()