Label ProportionShift top bars with system names (#47)

setup.py

@@ -8,7 +8,7 @@
 
 setuptools.setup(
     name="shifterator",
-    version="0.5.0",
+    version="0.5.1",
     author="Ryan J. Gallagher",
     author_email="gallagher.r@northeastern.edu",
     description="Interpretable data visualizations for understanding how texts differ at the word level",

shifterator/__init__.py

@@ -8,3 +8,4 @@
 )
 
 name = "shifterator"
+__version__ = "0.5.1"

shifterator/plotting.py

@@ -23,6 +23,7 @@ def get_plot_params(plot_params, show_score_diffs, diff):
         "height": 15,
         "invisible_spines": [],
         "label_fontsize": 13,
+        "label_total_with_system_names": False,
         "missing_symbol": "*",
         "pos_cumulative_inset": [0.19, 0.12, 0.175, 0.175],
         "pos_text_size_inset": [0.81, 0.12, 0.08, 0.08],
@@ -74,6 +75,13 @@ def get_plot_params(plot_params, show_score_diffs, diff):
         "total": r"$\Sigma$",
     }
     defaults.update(plot_params)
+    # Label the total contribution bars with the system names. This is used by
+    # shifts whose +/- contributions correspond to the two systems (e.g.
+    # ProportionShift), so the top bars are labeled like the JSD shift graph
+    # without flipping bar directions via `all_pos_contributions`.
+    if defaults["label_total_with_system_names"]:
+        defaults["symbols"]["neg_total"] = defaults["system_names"][0]
+        defaults["symbols"]["pos_total"] = defaults["system_names"][1]
     return defaults
 
 

shifterator/shifts.py

@@ -124,6 +124,12 @@ def get_shift_graph(
     ):
         if title is None:
             title = ""
+        # A type's contribution is positive or negative depending on which
+        # system it is more prevalent in, so label the total contribution bars
+        # with the system names (see issue #47). We do not set
+        # `all_pos_contributions`, which would force every bar to point the same
+        # direction since proportion shift scores already carry their sign.
+        kwargs.setdefault("label_total_with_system_names", True)
         ax = super().get_shift_graph(
             top_n=top_n,
             text_size_inset=text_size_inset,

tests/test_issues.py

@@ -1,11 +1,11 @@
-"""Tests for fixes to GitHub issues #11, #26, #38."""
+"""Tests for fixes to GitHub issues #11, #26, #38, #47."""
 
 import matplotlib
 matplotlib.use("Agg")
 
 import matplotlib.pyplot as plt
 import pytest
-from shifterator import WeightedAvgShift, EntropyShift
+from shifterator import WeightedAvgShift, EntropyShift, ProportionShift
 
 
 @pytest.fixture(autouse=True)
@@ -138,3 +138,52 @@ def test_default_no_font_change(self):
         shift = WeightedAvgShift(freq1, freq2, type2score_1=scores)
         ax = shift.get_shift_graph(top_n=2, show_plot=False)
         assert ax is not None
+
+
+# ---------------------------------------------------------------------------
+# Issue #47: ProportionShift should label the top bars with the system names
+# ---------------------------------------------------------------------------
+
+class TestIssue47ProportionShiftSystemNames:
+    freq1 = {"happy": 20, "sad": 5, "the": 50, "good": 15, "bad": 10}
+    freq2 = {"happy": 10, "sad": 15, "the": 45, "good": 5, "bad": 20, "angry": 8}
+
+    def test_system_names_labeled(self):
+        """The total contribution bars should be labeled with the system names."""
+        shift = ProportionShift(self.freq1, self.freq2)
+        ax = shift.get_shift_graph(
+            top_n=6, show_plot=False, system_names=["Corpus A", "Corpus B"]
+        )
+        labels = [t.get_text() for t in ax.texts]
+        assert "Corpus A" in labels
+        assert "Corpus B" in labels
+
+    def test_default_system_names_labeled(self):
+        """Falls back to the default system names when none are given."""
+        shift = ProportionShift(self.freq1, self.freq2)
+        ax = shift.get_shift_graph(top_n=6, show_plot=False)
+        labels = [t.get_text() for t in ax.texts]
+        assert "Text 1" in labels
+        assert "Text 2" in labels
+
+    def test_bars_remain_two_sided(self):
+        """Labeling must not flip every bar to one side (the broken workaround
+        of passing all_pos_contributions=True)."""
+        shift = ProportionShift(self.freq1, self.freq2)
+        ax = shift.get_shift_graph(top_n=6, show_plot=False)
+        widths = [p.get_width() for p in ax.patches if p.get_width() != 0]
+        assert any(w < 0 for w in widths)
+        assert any(w > 0 for w in widths)
+
+    def test_other_shifts_unaffected(self):
+        """Shifts that don't opt in keep empty total-bar symbols."""
+        freq1 = {"hello": 20, "world": 10}
+        freq2 = {"hello": 15, "world": 15}
+        scores = {"hello": 7.0, "world": 5.0}
+        shift = WeightedAvgShift(freq1, freq2, type2score_1=scores)
+        ax = shift.get_shift_graph(
+            top_n=2, show_plot=False, system_names=["Corpus A", "Corpus B"]
+        )
+        labels = [t.get_text() for t in ax.texts]
+        assert "Corpus A" not in labels
+        assert "Corpus B" not in labels