Add plot_waterfall helper for dataframe-based waterfall charts

_doc/api/tools/pandas.rst

@@ -2,6 +2,24 @@
 teachpyx.tools.pandas
 =====================
 
+Example
+-------
+
+.. plot::
+
+ import pandas
+ import matplotlib.pyplot as plt
+ from teachpyx.tools.pandas import plot_waterfall
+
+ plt.close("all")
+
+ df = pandas.DataFrame({"name": ["A", "B", "C"], "delta": [10, -3, 5]})
+ ax, _ = plot_waterfall(df, "delta", "name", total_label="TOTAL")
+ ax.set_title("Example waterfall")
+ plt.xticks(rotation=30, ha="right")
+ plt.tight_layout()
+ plt.show()
+
 .. automodule:: teachpyx.tools.pandas
  :members:
  :no-undoc-members:

_doc/examples/ml/plot_template_data.py

@@ -105,9 +105,7 @@ def make_pipeline(table, cible):
  if num_cols:
  transformers.append(("num", StandardScaler(), num_cols))
  if cat_cols:
- transformers.append(
- ("cats", OneHotEncoder(handle_unknown="ignore"), cat_cols)
- )
+ transformers.append(("cats", OneHotEncoder(handle_unknown="ignore"), cat_cols))
 
  model = Pipeline(
  [

_unittests/ut_tools/test_pandas.py

@@ -1,6 +1,8 @@
 import unittest
+from matplotlib.axes import Axes
+import pandas
 from teachpyx.ext_test_case import ExtTestCase
-from teachpyx.tools.pandas import read_csv_cached
+from teachpyx.tools.pandas import plot_waterfall, read_csv_cached
 
 
 class TestPandas(ExtTestCase):
@@ -14,6 +16,39 @@ def test_read_csv_cached(self):
  self.assertEqual(df.shape, df2.shape)
  self.assertEqual(list(df.columns), list(df2.columns))
 
+ def test_plot_waterfall(self):
+ df = pandas.DataFrame(
+ {
+ "name": ["A", "B", "C"],
+ "delta": [10, -3, 5],
+ }
+ )
+ ax, plot_df = plot_waterfall(df, "delta", "name", total_label="TOTAL")
+ self.assertIsInstance(ax, Axes)
+ self.assertEqual(list(plot_df["label"]), ["A", "B", "C", "TOTAL"])
+ self.assertEqual(list(plot_df["start"]), [0.0, 10.0, 7.0, 0.0])
+ self.assertEqual(list(plot_df["end"]), [10.0, 7.0, 12.0, 12.0])
+
+ def test_plot_waterfall_missing_column(self):
+ df = pandas.DataFrame({"name": ["A"], "delta": [1]})
+ with self.assertRaises(ValueError):
+ plot_waterfall(df, "missing", "name")
+
+ def test_plot_waterfall_missing_label_column(self):
+ df = pandas.DataFrame({"name": ["A"], "delta": [1]})
+ with self.assertRaises(ValueError):
+ plot_waterfall(df, "delta", "missing")
+
+ def test_plot_waterfall_bad_colors(self):
+ df = pandas.DataFrame({"name": ["A"], "delta": [1]})
+ with self.assertRaises(ValueError):
+ plot_waterfall(df, "delta", "name", colors=("r",))
+
+ def test_plot_waterfall_not_numeric(self):
+ df = pandas.DataFrame({"name": ["A"], "delta": ["x"]})
+ with self.assertRaises(ValueError):
+ plot_waterfall(df, "delta", "name")
+
 
 if __name__ == "__main__":
  unittest.main()

teachpyx/tools/pandas.py

@@ -2,6 +2,7 @@
 import os
 import re
 from pathlib import Path
+from typing import Optional, Tuple
 from urllib.parse import urlparse, unquote
 import pandas
 
@@ -46,3 +47,81 @@ def read_csv_cached(
  df = pandas.read_csv(filepath_or_buffer, **kwargs)
  df.to_csv(cache_name, index=False)
  return df
+
+
+def plot_waterfall(
+ data: pandas.DataFrame,
+ value_column: str,
+ label_column: Optional[str] = None,
+ total_label: str = "total",
+ ax=None,
+ colors: Tuple[str, str, str] = ("#2ca02c", "#d62728", "#1f77b4"),
+):
+ """
+ Draws a waterfall chart from a dataframe.
+
+ :param data: dataframe containing increments
+ :param value_column: column with increments
+ :param label_column: column with labels, index is used if None
+ :param total_label: label used for the final total
+ :param ax: existing axis or None to create one
+ :param colors: positive, negative, total colors
+ :return: axis, computed dataframe used to draw the chart
+
+ .. versionadded:: 0.6.1
+ """
+ if value_column not in data.columns:
+ raise ValueError(f"Unable to find column {value_column!r} in dataframe.")
+ if label_column is not None and label_column not in data.columns:
+ raise ValueError(f"Unable to find column {label_column!r} in dataframe.")
+ if len(colors) != 3:
+ raise ValueError(f"colors must contain 3 values, not {len(colors)}.")
+
+ values = pandas.to_numeric(data[value_column], errors="raise").astype(float)
+ labels = data[label_column] if label_column is not None else data.index
+ labels = labels.astype(str)
+
+ starts = values.cumsum().shift(1, fill_value=0.0)
+ plot_df = pandas.DataFrame(
+ {
+ "label": labels,
+ "value": values,
+ "start": starts,
+ "end": starts + values,
+ "kind": "variation",
+ }
+ )
+
+ total = float(values.sum())
+ total_row = pandas.DataFrame(
+ {
+ "label": [total_label],
+ "value": [total],
+ "start": [0.0],
+ "end": [total],
+ "kind": ["total"],
+ }
+ )
+ plot_df = pandas.concat([plot_df, total_row], axis=0, ignore_index=True)
+
+ if ax is None:
+ import matplotlib.pyplot as plt
+
+ _, ax = plt.subplots(1, 1)
+
+ bar_colors = [
+ colors[2] if kind == "total" else (colors[0] if value >= 0 else colors[1])
+ for value, kind in zip(plot_df["value"], plot_df["kind"])
+ ]
+ ax.bar(
+ plot_df["label"],
+ plot_df["value"],
+ bottom=plot_df["start"],
+ color=bar_colors,
+ )
+
+ ax.axhline(0, color="black", linewidth=0.8)
+ ax.set_ylabel(value_column)
+ ax.set_xlabel(label_column or "index")
+
+ return ax, plot_df