Added measurement type to waterfall colorbar (#293)

* Added measurement type to waterfall colorbar * fix failing test --------- Co-authored-by: matt <[email protected]>
DASDAE · Nov 1, 2023 · b5c0bb3 · b5c0bb3
1 parent 42ca6f5
commit b5c0bb3
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/dascore/viz/waterfall.py b/dascore/viz/waterfall.py
@@ -99,8 +99,11 @@ def waterfall(
     # add color bar with title
     if cmap is not None:
         cb = ax.get_figure().colorbar(im)
+        data_type = str(patch.attrs["data_type"])
         data_units = get_quantity_str(patch.attrs.data_units) or ""
-        cb.ax.set_title(data_units)
+        dunits = f" ({data_units})" if (data_type and data_units) else f"{data_units}"
+        label = f"{data_type}{dunits}"
+        cb.set_label(label)
     ax.invert_yaxis()  # invert y axis so origin is at top
     if show:
         plt.show()

diff --git a/tests/test_viz/test_waterfall.py b/tests/test_viz/test_waterfall.py
@@ -25,8 +25,9 @@ def check_label_units(patch, ax):
         assert coord_name in label_text
     # check colorbar labels
     cax = ax.images[-1].colorbar
-    title = cax.ax.title.get_text()
-    assert str(patch.attrs.data_units.units) in title
+    yaxis_label = cax.ax.yaxis.label.get_text()
+    assert str(patch.attrs.data_units.units) in yaxis_label
+    assert str(patch.attrs.data_type) in yaxis_label
 
 
 @pytest.fixture(scope="session")
@@ -129,6 +130,15 @@ def test_time_no_units(self, patch_two_times):
         assert ax.get_xlabel() == dims[1]
         assert ax.get_ylabel() == dims[0]
 
+    def test_patch_with_data_type(self, random_patch):
+        """Ensure a patch with data_type titles the colorbar."""
+        patch = random_patch.update_attrs(
+            data_type="strain rate",
+            data_units="1/s",
+        )
+        ax = patch.viz.waterfall()
+        check_label_units(patch, ax)
+
     def test_show(self, random_patch, monkeypatch):
         """Ensure show path is callable."""
         monkeypatch.setattr(plt, "show", lambda: None)