Merge pull request #3 from nickderobertis/styled-df

Get styled DataFrames

Pipfile

@@ -23,6 +23,7 @@ pypandoc = "*"
 cruft = "*"
 pandas = "*"
 matplotlib = "*"
+pd_utils = "*"
 
 [requires]
 python_version = "3.7"

conf.py

@@ -17,7 +17,7 @@
 ]
 
 # Package version in the format (major, minor, release)
-PACKAGE_VERSION_TUPLE = (0, 1, 0)
+PACKAGE_VERSION_TUPLE = (0, 1, 1)
 
 # Short description of the package
 PACKAGE_SHORT_DESCRIPTION = "Python Sensitivity Analysis - Gradient DataFrames and Hex-Bin Plots"
@@ -61,6 +61,7 @@
     # 'otherpackage>=1,<2'
     'pandas',
     'matplotlib',
+    'pd_utils',
 ]
 
 # Add any third party packages you use in requirements for optional features of your package here

sensitivity/df.py

@@ -1,7 +1,8 @@
-from typing import Dict, Any, Callable
+from typing import Dict, Any, Callable, Sequence, Optional
 import itertools
 from copy import deepcopy
 import pandas as pd
+import pd_utils
 from pandas.io.formats.style import Styler
 
 from sensitivity.colors import _get_color_map
@@ -33,11 +34,35 @@ def sensitivity_df(sensitivity_values: Dict[str, Any], func: Callable,
         base_param_dict.update({result_name: result})
         df = df.append(pd.DataFrame(pd.Series(base_param_dict)).T)
     df.reset_index(drop=True, inplace=True)
+    df = df.convert_dtypes()
 
     return df
 
 
-def _style_sensitivity_df(df: pd.DataFrame, reverse_colors: bool = False) -> Styler:
+def _two_variable_sensitivity_display_df(df: pd.DataFrame, col1: str, col2: str,
+                                         result_col: str = 'Result') -> pd.DataFrame:
+    selected_df = df[[col1, col2, result_col]]
+    wide_df = pd_utils.long_to_wide(
+        selected_df,
+        col1,
+        result_col,
+        colindex=col2,
+        colindex_only=True
+    ).set_index(col1)
+
+    return wide_df
+
+
+def _style_sensitivity_df(df: pd.DataFrame, col1: str, col2: Optional[str] = None, result_col: str = 'Result',
+                          reverse_colors: bool = False,
+                          col_subset: Optional[Sequence[str]] = None) -> Styler:
+    if col2 is not None:
+        caption = f'{result_col} - {col1} vs. {col2}'
+    else:
+        caption = f'{result_col} vs. {col1}'
+
     color_str = _get_color_map(reverse_colors=reverse_colors)
-    return df.style.background_gradient(cmap=color_str)
+    return df.style.background_gradient(
+        cmap=color_str, subset=col_subset, axis=None
+    ).set_caption(caption)
 

sensitivity/main.py

@@ -1,10 +1,12 @@
+import itertools
 from dataclasses import dataclass
-from typing import Dict, Any, Callable, Optional
+from typing import Dict, Any, Callable, Optional, List, Union
 
 import numpy as np
+from pandas.io.formats.style import Styler
 import matplotlib.pyplot as plt
 
-from sensitivity.df import sensitivity_df, _style_sensitivity_df
+from sensitivity.df import sensitivity_df, _style_sensitivity_df, _two_variable_sensitivity_display_df
 from sensitivity.hexbin import sensitivity_hex_plots, _hex_figure_from_sensitivity_df
 
 
@@ -52,7 +54,7 @@ class SensitivityAnalyzer:
         >>> sa.df
         >>>
         >>> # Styled DataFrame
-        >>> sa.styled_df
+        >>> sa.styled_dfs
         >>>
         >>> # Hex-Bin Plot
         >>> sa.plot
@@ -89,8 +91,50 @@ def plot(self) -> plt.Figure:
         )
 
     @property
-    def styled_df(self):
-        return _style_sensitivity_df(
-            self.df,
-            reverse_colors=self.reverse_colors
-        )
+    def styled_dfs(self) -> Union[Styler, List[Styler]]:
+        # Output a single Styler if only one or two variables
+        sensitivity_cols = list(self.sensitivity_values.keys())
+        if len(sensitivity_cols) == 1:
+            return _style_sensitivity_df(
+                self.df,
+                sensitivity_cols[0],
+                reverse_colors=self.reverse_colors,
+                col_subset=[self.result_name],
+                result_col=self.result_name
+            )
+        elif len(sensitivity_cols) == 2:
+            col1 = sensitivity_cols[0]
+            col2 = sensitivity_cols[1]
+            df = _two_variable_sensitivity_display_df(
+                self.df,
+                col1,
+                col2,
+                result_col=self.result_name
+            )
+            return _style_sensitivity_df(
+                df,
+                col1,
+                col2=col2,
+                reverse_colors=self.reverse_colors,
+                result_col=self.result_name,
+            )
+        elif len(sensitivity_cols) == 0:
+            raise ValueError('must pass sensitivity columns')
+
+        # Length must be greater than 2, need to output multiple, one for each pair of variables
+        results = []
+        for col1, col2 in itertools.combinations(sensitivity_cols, 2):
+            df = _two_variable_sensitivity_display_df(
+                self.df,
+                col1,
+                col2,
+                result_col=self.result_name
+            )
+            results.append(_style_sensitivity_df(
+                df,
+                col1,
+                col2=col2,
+                reverse_colors=self.reverse_colors,
+                result_col=self.result_name,
+            ))
+        return results

tests/test_sensitivity_analyzer.py

@@ -23,9 +23,9 @@ def test_create_df(self):
         sa = self.create_sa()
         assert_frame_equal(sa.df, EXPECT_DF, check_dtype=False)
 
-    def test_create_styled_df(self):
+    def test_create_styled_dfs(self):
         sa = self.create_sa()
-        sa.styled_df
+        sa.styled_dfs
         # TODO [#1]: determine how to test pandas Styler object beyond creation without error
 
     def test_create_plot(self):