feat: add sankey flow plotting helper and unit test

Agent-Logs-Url: https://github.com/sdpython/teachpyx/sessions/68526f45-813e-4c5a-959d-84afb8a8544a

Co-authored-by: xadupre <22452781+xadupre@users.noreply.github.com>

Author: Copilot

_unittests/ut_faq/test_faq_python.py

@@ -1,6 +1,14 @@
 import unittest
 import datetime
-from teachpyx.faq.faq_python import get_month_name, get_day_name, class_getitem
+import matplotlib
+from teachpyx.faq.faq_python import (
+    class_getitem,
+    get_day_name,
+    get_month_name,
+    graph_sankey,
+)
+
+matplotlib.use("Agg")
 
 
 class TestFaqPython(unittest.TestCase):
@@ -17,6 +25,17 @@ def test_day_name(self):
     def test_class_getitem(self):
         class_getitem()
 
+    def test_graph_sankey(self):
+        ax, diagrams = graph_sankey(
+            [1, -0.25, -0.75],
+            labels=["input", "loss", "output"],
+            orientations=[0, 1, -1],
+            trunklength=2.0,
+            title="flux",
+        )
+        self.assertEqual(ax.get_title(), "flux")
+        self.assertEqual(len(diagrams), 1)
+
 
 if __name__ == "__main__":
     unittest.main()

teachpyx/faq/faq_python.py

@@ -850,3 +850,45 @@ def __init__(self):
 
     a = A[2]()
     assert a.__class__.__name__ == "A2"
+
+
+def graph_sankey(
+    flows, labels=None, orientations=None, ax=None, title=None, **kwargs
+):
+    """
+    Draws a :epkg:`Sankey` graph to represent flows.
+
+    :param flows: list of positive/negative flows
+    :param labels: labels associated to each flow
+    :param orientations: list of orientations (-1, 0, 1)
+    :param ax: axis to use, if None, creates one
+    :param title: graph title
+    :param kwargs: additional parameters forwarded to
+        ``matplotlib.sankey.Sankey.add``
+    :return: axis, sankey diagrams
+    """
+    if len(flows) < 2:
+        raise ValueError("flows must contain at least two values.")
+    if abs(sum(flows)) > 1e-10:
+        raise ValueError("The sum of all flows must be 0.")
+    if labels is None:
+        labels = [None] * len(flows)
+    if orientations is None:
+        orientations = [0] * len(flows)
+    if len(labels) != len(flows):
+        raise ValueError("labels and flows must have the same length.")
+    if len(orientations) != len(flows):
+        raise ValueError("orientations and flows must have the same length.")
+
+    import matplotlib.pyplot as plt
+    from matplotlib.sankey import Sankey
+
+    if ax is None:
+        _, ax = plt.subplots(1, 1)
+
+    sankey = Sankey(ax=ax)
+    sankey.add(flows=flows, labels=labels, orientations=orientations, **kwargs)
+    diagrams = sankey.finish()
+    if title is not None:
+        ax.set_title(title)
+    return ax, diagrams