revert simplification (thought it worked, didnt)

AUMAG · Mar 22, 2024 · 7f64f9b · 7f64f9b
1 parent 5ce1785
commit 7f64f9b
Showing 1 changed file with 65 additions and 45 deletions.
diff --git a/ausankey/ausankey.py b/ausankey/ausankey.py
@@ -157,11 +157,26 @@ def sankey(
     weight_sum = np.empty(num_side)
     num_uniq = np.empty(num_side)
     col_hgt = np.empty(num_side)
+    node_sizes = {}
     nodes_uniq = {}
     for ii in range(num_side):
         nodes_uniq[ii] = pd.Series(data[2 * ii]).unique()
         num_uniq[ii] = len(nodes_uniq[ii])
 
+    # for ii in range(num_side):
+    #     node_sizes[ii] = {}
+    #     for lbl in nodes_uniq[ii]:
+    #         if ii == 0:
+    #             ind_prev = data[2 * ii + 0] == lbl
+    #             ind_this = data[2 * ii + 0] == lbl
+    #             ind_next = data[2 * ii + 2] == lbl
+    #             weight_cont = data[2 * ii + 1][ind_this & ind_prev & ind_next].sum()
+    #             weight_only = data[2 * ii + 1][ind_this & ~ind_prev & ~ind_next].sum()
+    #             weight_stop = data[2 * ii + 1][ind_this & ind_prev & ~ind_next].sum()
+    #             weight_strt = data[2 * ii + 1][ind_this & ~ind_prev & ind_next].sum()
+    #             node_sizes[ii][lbl] = weight_cont + weight_only + max(weight_stop, weight_strt)
+
+    # print(node_sizes)
     for ii in range(num_side):
         if ii == 0:
             ind_prev = data[2 * ii + 1].notnull()
@@ -513,51 +528,56 @@ def draw_bar(x, dx, y, dy, label):
 
     # Plot flows
     for i,lbl_l in enumerate(bar_lr[0]):
-        lbl_r = bar_lr[1][i]
-
-        lbot = barpos[0][lbl_l]["bot"]
-        rbot = barpos[1][lbl_r]["bot"]
-        lbar = barsize[0][lbl_l][lbl_r]
-        rbar = barsize[1][lbl_l][lbl_r]
-
-        ys_d = create_curve(lbot, rbot)
-        ys_u = create_curve(lbot + lbar, rbot + rbar)
-
-        # Update bottom edges at each label
-        # so next strip starts at the right place
-        barpos[0][lbl_l]["bot"] += lbar
-        barpos[1][lbl_r]["bot"] += rbar
-
-        xx = np.linspace(x_left, x_right, len(ys_d))
-        cc = combine_colours(color_dict[lbl_l], color_dict[lbl_r], len(ys_d))
-
-        for jj in range(len(ys_d) - 1):
-            ax.fill_between(
-                xx[[jj, jj + 1]],
-                ys_d[[jj, jj + 1]],
-                ys_u[[jj, jj + 1]],
-                color=cc[:, jj],
-                alpha=alpha,
-                lw=0,
-                edgecolor="none",
-                snap=True,
-            )
-            ax.plot(
-                xx[[jj, jj + 1]],
-                ys_d[[jj, jj + 1]],
-                color=cc[:, jj],
-                alpha=edge_alpha,
-                lw=edge_lw,
-                snap=True,
-            )
-            ax.plot(
-                xx[[jj, jj + 1]],
-                ys_u[[jj, jj + 1]],
-                color=cc[:, jj],
-                alpha=edge_alpha,
-                lw=edge_lw,
-                snap=True,
-            )
+        for j,lbl_r in enumerate(bar_lr[1]):
+
+            lind = labels_lr[0] == lbl_l
+            rind = labels_lr[1] == lbl_r
+            if not any(lind & rind):
+                continue
+
+            lbot = barpos[0][lbl_l]["bot"]
+            rbot = barpos[1][lbl_r]["bot"]
+            lbar = barsize[0][lbl_l][lbl_r]
+            rbar = barsize[1][lbl_l][lbl_r]
+
+            ys_d = create_curve(lbot, rbot)
+            ys_u = create_curve(lbot + lbar, rbot + rbar)
+
+            # Update bottom edges at each label
+            # so next strip starts at the right place
+            barpos[0][lbl_l]["bot"] += lbar
+            barpos[1][lbl_r]["bot"] += rbar
+
+            xx = np.linspace(x_left, x_right, len(ys_d))
+            cc = combine_colours(color_dict[lbl_l], color_dict[lbl_r], len(ys_d))
+
+            for jj in range(len(ys_d) - 1):
+                ax.fill_between(
+                    xx[[jj, jj + 1]],
+                    ys_d[[jj, jj + 1]],
+                    ys_u[[jj, jj + 1]],
+                    color=cc[:, jj],
+                    alpha=alpha,
+                    lw=0,
+                    edgecolor="none",
+                    snap=True,
+                )
+                ax.plot(
+                    xx[[jj, jj + 1]],
+                    ys_d[[jj, jj + 1]],
+                    color=cc[:, jj],
+                    alpha=edge_alpha,
+                    lw=edge_lw,
+                    snap=True,
+                )
+                ax.plot(
+                    xx[[jj, jj + 1]],
+                    ys_u[[jj, jj + 1]],
+                    color=cc[:, jj],
+                    alpha=edge_alpha,
+                    lw=edge_lw,
+                    snap=True,
+                )
 
 
 ###########################################