Add processing funcs (#29)

* add some usefule processing tools

* lint: ruff

---------

Co-authored-by: Choi Minkyu (mc76728) <mc76728@austin.utexas.edu@ece-a52150.austin.utexas.edu>

swarm_visualizer/__init__.py

@@ -1,13 +1,25 @@
-from .barplot import plot_grouped_barplot, plot_stacked_barplot, plot_sns_grouped_barplot
+from .barplot import (
+    plot_grouped_barplot,
+    plot_stacked_barplot,
+    plot_sns_grouped_barplot,
+)
 from .boxplot import plot_grouped_boxplot, plot_paired_boxplot
 from .lineplot import plot_basic_lineplot, plot_overlaid_lineplot
 from .violinplot import plot_grouped_violinplot, plot_paired_violinplot
 from .scatterplot import plot_basic_scatterplot, plot_scatter_pdf_plot
 from .gridplot import plot_grid
 
-__all__ = [ "plot_grouped_barplot", "plot_stacked_barplot", "plot_sns_grouped_barplot",
-            "plot_grouped_boxplot", "plot_paired_boxplot",
-            "plot_basic_lineplot", "plot_overlaid_lineplot",
-            "plot_grouped_violinplot", "plot_paired_violinplot",
-            "plot_basic_scatterplot", "plot_scatter_pdf_plot",
-            "plot_grid" ]
+__all__ = [
+    "plot_grouped_barplot",
+    "plot_stacked_barplot",
+    "plot_sns_grouped_barplot",
+    "plot_grouped_boxplot",
+    "plot_paired_boxplot",
+    "plot_basic_lineplot",
+    "plot_overlaid_lineplot",
+    "plot_grouped_violinplot",
+    "plot_paired_violinplot",
+    "plot_basic_scatterplot",
+    "plot_scatter_pdf_plot",
+    "plot_grid",
+]

swarm_visualizer/barplot.py

@@ -2,6 +2,7 @@
 
 from swarm_visualizer.utility import set_axis_infos
 
+
 def plot_grouped_barplot(
     ax,
     df=None,
@@ -11,7 +12,7 @@ def plot_grouped_barplot(
     title_str=None,
     pal=None,
     y_label=None,
-    **kwargs
+    **kwargs,
 ) -> None:
     """Plots a grouped barplot. In this case, there are multiple y-var for each x-var.
 
@@ -30,7 +31,13 @@ def plot_grouped_barplot(
     if pal:
         colors = [pal(i) for i in range(len(x_var))]
         df.plot(
-            kind="bar", stacked=False, ax=ax, x=x_var, y=y_var, colors=colors, **kwargs
+            kind="bar",
+            stacked=False,
+            ax=ax,
+            x=x_var,
+            y=y_var,
+            colors=colors,
+            **kwargs,
         )
 
     ### set y label
@@ -40,6 +47,7 @@ def plot_grouped_barplot(
     # Set axis infos
     set_axis_infos(ax, ylim=ylim, title_str=title_str)
 
+
 def plot_sns_grouped_barplot(
     ax,
     df=None,
@@ -50,7 +58,7 @@ def plot_sns_grouped_barplot(
     title_str=None,
     pal=None,
     y_label=None,
-    **kwargs
+    **kwargs,
 ) -> None:
     """Plots a grouped barplot with sns. hue specifies the group.
 
@@ -81,7 +89,6 @@ def plot_sns_grouped_barplot(
     set_axis_infos(ax, ylim=ylim, title_str=title_str)
 
 
-
 def plot_stacked_barplot(
     ax,
     df=None,
@@ -91,7 +98,7 @@ def plot_stacked_barplot(
     title_str=None,
     pal=None,
     y_label=None,
-    **kwargs
+    **kwargs,
 ) -> None:
     """Plots a grouped barplot.
 
@@ -110,7 +117,13 @@ def plot_stacked_barplot(
     if pal:
         colors = [pal(i) for i in range(len(x_var))]
         df.plot(
-            kind="bar", stacked=True, ax=ax, x=x_var, y=y_var, colors=colors, **kwargs
+            kind="bar",
+            stacked=True,
+            ax=ax,
+            x=x_var,
+            y=y_var,
+            colors=colors,
+            **kwargs,
         )
 
     ### set y label

swarm_visualizer/boxplot.py

@@ -17,7 +17,7 @@ def plot_paired_boxplot(
     order_list=None,
     pal=None,
     hue=None,
-    **kwargs
+    **kwargs,
 ) -> None:
     """Plots a paired boxplot.
 
@@ -37,11 +37,23 @@ def plot_paired_boxplot(
         # Plots a boxplot with order
         if order_list:
             sns.boxplot(
-                x=x_var, y=y_var, data=df, order=order_list, hue=hue, ax=ax, **kwargs
+                x=x_var,
+                y=y_var,
+                data=df,
+                order=order_list,
+                hue=hue,
+                ax=ax,
+                **kwargs,
             )
         else:
             sns.boxplot(
-                x=x_var, y=y_var, data=df, order=order_list, hue=hue, ax=ax, **kwargs
+                x=x_var,
+                y=y_var,
+                data=df,
+                order=order_list,
+                hue=hue,
+                ax=ax,
+                **kwargs,
             )
 
     # Plots a boxplot with palette
@@ -55,7 +67,7 @@ def plot_paired_boxplot(
                 palette=pal,
                 hue=hue,
                 ax=ax,
-                **kwargs
+                **kwargs,
             )
         else:
             sns.boxplot(
@@ -66,7 +78,7 @@ def plot_paired_boxplot(
                 palette=pal,
                 hue=hue,
                 ax=ax,
-                **kwargs
+                **kwargs,
             )
 
     # Set axis infos
@@ -87,7 +99,7 @@ def plot_grouped_boxplot(
     title_str=None,
     order_list=None,
     pal=None,
-    **kwargs
+    **kwargs,
 ) -> None:
     """Plots a grouped boxplot.
 
@@ -103,18 +115,34 @@ def plot_grouped_boxplot(
     """
     if not pal:
         if order_list:
-            sns.boxplot(x=x_var, y=y_var, data=df, order=order_list, ax=ax, **kwargs)
+            sns.boxplot(
+                x=x_var, y=y_var, data=df, order=order_list, ax=ax, **kwargs
+            )
         else:
-            sns.boxplot(x=x_var, y=y_var, data=df, order=order_list, ax=ax, **kwargs)
+            sns.boxplot(
+                x=x_var, y=y_var, data=df, order=order_list, ax=ax, **kwargs
+            )
 
     if pal:
         if order_list:
             sns.boxplot(
-                x=x_var, y=y_var, data=df, order=order_list, palette=pal, ax=ax, **kwargs
+                x=x_var,
+                y=y_var,
+                data=df,
+                order=order_list,
+                palette=pal,
+                ax=ax,
+                **kwargs,
             )
         else:
             sns.boxplot(
-                x=x_var, y=y_var, data=df, order=order_list, palette=pal, ax=ax, **kwargs
+                x=x_var,
+                y=y_var,
+                data=df,
+                order=order_list,
+                palette=pal,
+                ax=ax,
+                **kwargs,
             )
 
     # Set axis infos

swarm_visualizer/gridplot.py

@@ -10,7 +10,7 @@ def plot_grid(
     plot_file: str = None,
     lw: float = 3.0,
     xlabel: str = None,
-    **kwargs
+    **kwargs,
 ) -> None:
     """Plot grid of time series.
 

swarm_visualizer/histogram.py

@@ -5,7 +5,7 @@
 
 
 def plot_pdf(
-     ax, data=None, xlabel: str = None, title_str: str = None, **kwargs
+    ax, data=None, xlabel: str = None, title_str: str = None, **kwargs
 ) -> None:
     """Plot PDF of a data.
 
@@ -30,7 +30,7 @@ def plot_pdf(
         alpha=0.4,
         edgecolor=(1, 1, 1, 0.4),
         ax=ax,
-        **kwargs
+        **kwargs,
     )
 
     # Set axis infos
@@ -39,16 +39,16 @@ def plot_pdf(
 
 def plot_several_pdf(
     ax,
-    data_list: list[np.ndarray]=None,
+    data_list: list[np.ndarray] = None,
     xlabel: str = None,
     title_str: str = None,
     legend=None,
     ylabel: str = None,
     xlim=None,
     kde: bool = False,
-    bins = "auto",
-    binwidth = None,
-    **kwargs
+    bins="auto",
+    binwidth=None,
+    **kwargs,
 ) -> None:
     """Plot PDF of a data list.
 
@@ -75,7 +75,7 @@ def plot_several_pdf(
             ax=ax,
             bins=bins,
             binwidth=binwidth,
-            **kwargs
+            **kwargs,
         )
 
     # Set axis infos

swarm_visualizer/lineplot.py

@@ -13,7 +13,7 @@ def plot_basic_lineplot(
     lw: float = 3.0,
     ylim=None,
     xlabel: str = "x",
-    **kwargs
+    **kwargs,
 ) -> None:
     """Basic lineplot.
 
@@ -30,7 +30,12 @@ def plot_basic_lineplot(
     ax.plot(y, lw=lw)
 
     set_axis_infos(
-        ax, xlabel=xlabel, ylabel=ylabel, ylim=ylim, title_str=title_str, **kwargs
+        ax,
+        xlabel=xlabel,
+        ylabel=ylabel,
+        ylim=ylim,
+        title_str=title_str,
+        **kwargs,
     )
 
 
@@ -46,10 +51,10 @@ def plot_overlaid_lineplot(
     legend_present: bool = True,
     DEFAULT_MARKERSIZE: float = 15,
     delete_yticks: bool = False,
-    **kwargs
+    **kwargs,
 ) -> None:
     """Overlaid line plot.
-    
+
     :param ax: axis to plot on
     :param normalized_dict: dictionary with values to plot
     :param title_str: title of the plot
@@ -112,7 +117,7 @@ def plot_overlaid_lineplot(
                     ms=DEFAULT_MARKERSIZE,
                     color=color,
                     zorder=zorder,
-                    **kwargs
+                    **kwargs,
                 )
             else:
                 ax.plot(
@@ -124,7 +129,7 @@ def plot_overlaid_lineplot(
                     alpha=alpha,
                     color=color,
                     zorder=zorder,
-                    **kwargs
+                    **kwargs,
                 )
         # Plot without x-axis if x is not specified
         else:
@@ -139,7 +144,7 @@ def plot_overlaid_lineplot(
                     ms=DEFAULT_MARKERSIZE,
                     color=color,
                     zorder=zorder,
-                    **kwargs
+                    **kwargs,
                 )
             else:
                 ax.plot(
@@ -150,7 +155,7 @@ def plot_overlaid_lineplot(
                     alpha=alpha,
                     color=color,
                     zorder=zorder,
-                    **kwargs
+                    **kwargs,
                 )
 
         i += 1

swarm_visualizer/processing/__init__.py

@@ -0,0 +1,12 @@
+"""Package containing for processing data."""
+
+from __future__ import annotations
+
+from .grouping import average_by_group, group_values_by_bound
+from .sorting import sort_array_based_on_reference_array
+
+__all__ = [
+    "group_values_by_bound",
+    "sort_array_based_on_reference_array",
+    "average_by_group",
+]